From 2516b5fcae507eba28243c4511fee0e62a25d0de Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 13 Aug 2026 00:47:31 -0700 Subject: [PATCH 01/14] fix(prefetch): parse the file-folder seed through its contract MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The audit found this key was the workspaceFilesKeys bug waiting to recur. The manager's record type and workspaceFileFolderSchema are two independent declarations that agree today by coincidence; the seed had no parse, so adding a column to one would have silently cached a shape a client fetch strips — and three of its fields are z.coerce.date(), the exact divergence that put ISO strings under the file-list key. Its sibling is immune because listWorkspaceFilesWithShares parses at the data layer. This does the same at the seed, and adds the shape-parity assertion the key never had. Verified falsifiable: removing the parse turns it red. Doing so also exposed the existing folder test as fixture-thin — a folder with only an id, which the contract rightly rejects — so it now uses a real row. Also points the credential block's fetchQuery at the exported staleTime constant instead of restating 60 * 1000; it was a fifth producer on that key free to drift from the four that share it. --- .../workspace/[workspaceId]/files/prefetch.ts | 13 ++++- .../[workspaceId]/lib/prefetch.test.ts | 58 +++++++++++++++++-- apps/sim/blocks/blocks/credential.ts | 7 ++- 3 files changed, 71 insertions(+), 7 deletions(-) diff --git a/apps/sim/app/workspace/[workspaceId]/files/prefetch.ts b/apps/sim/app/workspace/[workspaceId]/files/prefetch.ts index 6ed88de422d..cb83b3b9d32 100644 --- a/apps/sim/app/workspace/[workspaceId]/files/prefetch.ts +++ b/apps/sim/app/workspace/[workspaceId]/files/prefetch.ts @@ -1,4 +1,5 @@ import type { QueryClient } from '@tanstack/react-query' +import { listWorkspaceFileFoldersContract } from '@/lib/api/contracts/workspace-file-folders' import { listWorkspaceFileFolders } from '@/lib/uploads/contexts/workspace/workspace-file-folder-manager' import { getWorkspaceHostContextForViewer } from '@/lib/workspaces/host-context' import { prefetchResourceListChrome } from '@/app/workspace/[workspaceId]/lib/prefetch-resource-list-chrome' @@ -40,7 +41,17 @@ export async function prefetchFilesBrowser( await Promise.all([ queryClient.prefetchQuery({ queryKey: workspaceFileFolderKeys.list(workspaceId, 'active'), - queryFn: () => listWorkspaceFileFolders(workspaceId, { scope: 'active' }), + /** + * Parsed through the route's own response schema rather than seeded raw. The + * manager's record type and `workspaceFileFolderSchema` are two independent + * declarations that happen to agree today; without this parse, adding a column + * to one silently seeds a shape a client fetch would have stripped — the exact + * divergence that put ISO strings under `workspaceFilesKeys.list`. + */ + queryFn: async () => { + const folders = await listWorkspaceFileFolders(workspaceId, { scope: 'active' }) + return listWorkspaceFileFoldersContract.response.schema.shape.folders.parse(folders) + }, staleTime: WORKSPACE_FILE_FOLDERS_STALE_TIME, }), prefetchResourceListChrome(queryClient, workspaceId, 'file', userId), diff --git a/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts b/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts index dee79520936..da334353073 100644 --- a/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts @@ -290,17 +290,67 @@ describe('workspace list prefetches', () => { }) }) describe('prefetchFilesBrowser', () => { + /** + * The sibling `workspaceFilesKeys.list` once held ISO strings from one producer and + * `Date`s from another because a seed skipped the contract parse. This key is fed by + * a manager whose record type and the contract schema are independent declarations, + * so the parse — and this assertion — are what stop that recurring here. + */ + it('seeds the shape a client fetch caches, not the raw manager row', async () => { + mockListWorkspaceFileFolders.mockResolvedValue([ + { + id: 'folder-1', + workspaceId: WORKSPACE_ID, + userId: USER_ID, + name: 'Docs', + parentId: null, + path: '/Docs', + sortOrder: 0, + deletedAt: null, + createdAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-02T00:00:00.000Z', + serverOnlyColumn: 'should-be-stripped', + }, + ]) + const client = makeClient() + + await prefetchFilesBrowser(client, WORKSPACE_ID, USER_ID) + + const [cached] = client.getQueryData( + workspaceFileFolderKeys.list(WORKSPACE_ID, 'active') + ) as Array> + expect(cached.createdAt).toBeInstanceOf(Date) + expect(cached.updatedAt).toBeInstanceOf(Date) + expect(cached).not.toHaveProperty('serverOnlyColumn') + }) + it('primes the folder key the client hook reads', async () => { - const folders = [{ id: 'folder-1' }] + const folders = [ + { + id: 'folder-1', + workspaceId: WORKSPACE_ID, + userId: USER_ID, + name: 'Docs', + parentId: null, + path: '/Docs', + sortOrder: 0, + deletedAt: null, + createdAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-02T00:00:00.000Z', + }, + ] mockListWorkspaceFileFolders.mockResolvedValue(folders) const client = makeClient() await prefetchFilesBrowser(client, WORKSPACE_ID, USER_ID) expect(mockListWorkspaceFileFolders).toHaveBeenCalledWith(WORKSPACE_ID, { scope: 'active' }) - expect(client.getQueryData(workspaceFileFolderKeys.list(WORKSPACE_ID, 'active'))).toEqual( - folders - ) + /** Shape parity is asserted by the sibling test; this one pins the key and the args. */ + expect( + client.getQueryData(workspaceFileFolderKeys.list(WORKSPACE_ID, 'active')) as Array<{ + id: string + }> + ).toHaveLength(folders.length) }) /** diff --git a/apps/sim/blocks/blocks/credential.ts b/apps/sim/blocks/blocks/credential.ts index 2ea6ccaec85..d1e7baab8f4 100644 --- a/apps/sim/blocks/blocks/credential.ts +++ b/apps/sim/blocks/blocks/credential.ts @@ -3,7 +3,10 @@ import { getServiceConfigByProviderId } from '@/lib/oauth/utils' import { getQueryClient } from '@/app/_shell/providers/get-query-client' import type { BlockConfig } from '@/blocks/types' import { workspaceCredentialKeys } from '@/hooks/queries/utils/credential-keys' -import { fetchWorkspaceCredentialList } from '@/hooks/queries/utils/fetch-workspace-credentials' +import { + fetchWorkspaceCredentialList, + WORKSPACE_CREDENTIAL_LIST_STALE_TIME, +} from '@/hooks/queries/utils/fetch-workspace-credentials' import { useWorkflowRegistry } from '@/stores/workflows/registry/store' interface CredentialBlockOutput { @@ -72,7 +75,7 @@ export const CredentialBlock: BlockConfig = { const credentials = await getQueryClient().fetchQuery({ queryKey: workspaceCredentialKeys.list(workspaceId), queryFn: () => fetchWorkspaceCredentialList(workspaceId), - staleTime: 60 * 1000, + staleTime: WORKSPACE_CREDENTIAL_LIST_STALE_TIME, }) const seen = new Set() From 6df827bb1f837c552a792d86cde3618c8459dc12 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 13 Aug 2026 00:49:49 -0700 Subject: [PATCH 02/14] fix(queries): stop a table cell edit throwing, and an upload gate failing shut MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two functional bugs found auditing the query layer. patchCachedRows walked tableKeys.rowsRoot non-exact, but rowsRoot is a prefix: the find (search results) and write (pending writes) subtrees hang off it with non-paged shapes, and the updater's old.pages.map threw on them. It runs inside onMutate, so the whole cell edit rejected before reaching the server — reachable as soon as a find entry exists, i.e. after the user searches the table once. The sibling isDefaultOrderRowsQuery already excluded those subtrees and its docstring claimed they "never match"; that was only true of the sibling. Both now share one isRowListQueryKey helper so they cannot drift apart again. useCloudStorageConfigured combined staleTime: Infinity, retry: false, and the global retryOnMount: false on a workspace-independent key, so one transient failure left it errored for the tab's lifetime with no way back — navigating or switching workspace cannot change the key, and the upload path fails closed, so cloud-backed uploads stayed disabled until a full reload. useVoiceSettings carries the same three options and already escapes this with retryOnMount: true; this one now matches. Note: hooks/queries/workspace-files.test.tsx cannot load in a git worktree (pre-existing postcss resolution failure), so CI is the first place that file runs against this change. --- apps/sim/hooks/queries/tables.ts | 29 ++++++++++++++++++++--- apps/sim/hooks/queries/workspace-files.ts | 10 ++++++++ 2 files changed, 36 insertions(+), 3 deletions(-) diff --git a/apps/sim/hooks/queries/tables.ts b/apps/sim/hooks/queries/tables.ts index 6d0848b4caf..ada5568bbf9 100644 --- a/apps/sim/hooks/queries/tables.ts +++ b/apps/sim/hooks/queries/tables.ts @@ -829,6 +829,18 @@ function withOptimisticAutoFireExec(groups: WorkflowGroup[], row: TableRow): Tab return { ...row, executions: nextExecutions } } +/** + * Whether a cache key under {@link tableKeys.rowsRoot} holds paged row-list data. + * + * `rowsRoot` is a shared prefix: the `find` (search results) and `write` (pending + * writes) subtrees live beneath it with entirely different shapes, so any bulk + * update walking the prefix has to exclude them or it will be handed a value whose + * `pages` it cannot map. + */ +function isRowListQueryKey(queryKey: readonly unknown[]): boolean { + return !queryKey.includes('find') && !queryKey.includes('write') +} + /** * Apply a row-level transformation to all cached infinite row queries for this * table. Used for cell edits where positions don't change. @@ -839,9 +851,20 @@ function patchCachedRows( patchRow: (row: TableRow) => TableRow ) { queryClient.setQueriesData>( - { queryKey: tableKeys.rowsRoot(tableId), exact: false }, + { + queryKey: tableKeys.rowsRoot(tableId), + exact: false, + /** + * `rowsRoot` is a prefix, not a leaf: `rowWrites` and `find` hang off it and + * hold non-paged shapes. Without this they are handed to the updater below, + * whose `old.pages.map` throws — and because this runs inside `onMutate`, the + * whole cell edit rejects before it reaches the server. Only reachable once a + * `find` entry exists, i.e. after the user has searched the table. + */ + predicate: (query) => isRowListQueryKey(query.queryKey), + }, (old) => { - if (!old) return old + if (!old?.pages) return old return { ...old, pages: old.pages.map((page) => ({ ...page, rows: page.rows.map(patchRow) })), @@ -859,7 +882,7 @@ function patchCachedRows( * `find`/`write` subtrees aren't row-list data and never match. */ function isDefaultOrderRowsQuery(queryKey: readonly unknown[]): boolean { - if (queryKey.includes('find') || queryKey.includes('write')) return false + if (!isRowListQueryKey(queryKey)) return false const last = queryKey[queryKey.length - 1] if (typeof last !== 'string') return false try { diff --git a/apps/sim/hooks/queries/workspace-files.ts b/apps/sim/hooks/queries/workspace-files.ts index a236a623f2f..c1b0b01ad49 100644 --- a/apps/sim/hooks/queries/workspace-files.ts +++ b/apps/sim/hooks/queries/workspace-files.ts @@ -521,6 +521,16 @@ export function useCloudStorageConfigured(enabled = true) { enabled, retry: false, staleTime: CLOUD_STORAGE_CONFIGURED_STALE_TIME, + /** + * Escapes the global `retryOnMount: false`, which an infinite `staleTime` and + * `retry: false` would otherwise turn into a permanent failure: one transient + * error leaves this query errored for the tab's lifetime, and consumers fail + * closed — the upload path treats "unknown" as "not configured", so a single + * blip would disable cloud-backed uploads until a full reload. The key is + * global, so navigating or switching workspace cannot recover it either. + * Matches {@link useVoiceSettings}, which carries the same three options. + */ + retryOnMount: true, }) } From 93e040fc4efb2162a3d9b8e3f054ba60253fdb40 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 13 Aug 2026 01:01:22 -0700 Subject: [PATCH 03/14] perf(server): memoize the request-scoped workspace and entitlement reads The workspace row was read ~3x per workspace route and ~5x on settings, and the same Max-tier entitlement was resolved twice on one render. Memoization is deliberately partial. getWorkspaceWithOwner accepts a transaction and forUpdate, and live callers use both, so only the plain no-options read routes through the memo; a row read inside one caller's transaction or under a lock it alone holds can never be served to a later caller. includeArchived is part of the key so the two variants cannot alias. Three substitutions were considered and rejected as behavior changes, not optimizations: hostContext.ownerBilling resolves subscriptions differently from hasWorkspaceTierAccess and exposes no Max tier, so it cannot answer the Inbox/Sandbox gates; isOrganizationOnEnterprisePlan carries self-host short-circuits ownerBilling has no equivalent for; and widening WorkspaceHostContext to carry the full row would push owner and org ids onto the wire for every viewer to save a server-side read, since that type is a response contract rather than an internal struct. --- apps/sim/lib/billing/core/subscription.ts | 30 ++++++-- apps/sim/lib/permissions/super-user.ts | 9 ++- apps/sim/lib/workspaces/permissions/utils.ts | 81 +++++++++++++++----- 3 files changed, 89 insertions(+), 31 deletions(-) diff --git a/apps/sim/lib/billing/core/subscription.ts b/apps/sim/lib/billing/core/subscription.ts index 44345553e4d..962afd1625a 100644 --- a/apps/sim/lib/billing/core/subscription.ts +++ b/apps/sim/lib/billing/core/subscription.ts @@ -1,3 +1,4 @@ +import { cache } from 'react' import { db } from '@sim/db' import { member, organization, subscription, user } from '@sim/db/schema' import { createLogger } from '@sim/logger' @@ -429,11 +430,7 @@ export async function isEnterpriseOrgAdminOrOwner(userId: string): Promise { +async function resolveOrganizationEnterprisePlan(organizationId: string): Promise { try { if (!isBillingEnabled) { return true @@ -456,6 +453,17 @@ export async function isOrganizationOnEnterprisePlan(organizationId: string): Pr } } +/** + * Check if an organization has an enterprise plan + * Used for Access Control (Permission Groups) feature gating + * + * Request-memoized: settings renders gate several sections on the same + * organization's plan, and the plan cannot change mid-render. Outside a Server + * Component render React evaluates the resolver normally, so routes, tools, and + * background work re-read exactly as before. + */ +export const isOrganizationOnEnterprisePlan = cache(resolveOrganizationEnterprisePlan) + /** * Entitlement for a single org-scoped enterprise feature. * @@ -612,10 +620,16 @@ async function hasWorkspaceTierAccess( * Whether the workspace's payer is on a usable Max-or-Enterprise subscription. * Shared by the inbox (Sim Mailer), live sync, and custom sandboxes, which all * sit on the same entitlement tier. + * + * Request-memoized because those features are gated side by side on the same + * settings render, each otherwise repeating the identical workspace and + * subscription reads. The per-feature deployment and env short-circuits live in + * the exported wrappers and still run per call. Outside a Server Component + * render React evaluates this normally. */ -async function hasMaxTierWorkspaceAccess(workspaceId: string): Promise { - return hasWorkspaceTierAccess(workspaceId, isMaxTier) -} +const hasMaxTierWorkspaceAccess = cache( + (workspaceId: string): Promise => hasWorkspaceTierAccess(workspaceId, isMaxTier) +) /** * Check whether a workspace is entitled to the inbox (Sim Mailer) feature. diff --git a/apps/sim/lib/permissions/super-user.ts b/apps/sim/lib/permissions/super-user.ts index 597ca135e4c..72d9cd71e8f 100644 --- a/apps/sim/lib/permissions/super-user.ts +++ b/apps/sim/lib/permissions/super-user.ts @@ -1,3 +1,4 @@ +import { cache } from 'react' import { db, dbReplica } from '@sim/db' import { settings, user } from '@sim/db/schema' import { eq } from 'drizzle-orm' @@ -41,8 +42,12 @@ export async function verifyEffectiveSuperUser(userId: string): Promise<{ * served from the replica: this gates features, not security-critical auth, so it * tolerates the replica's bounded staleness (admin role rarely changes). Falls back * to the primary when no replica is configured. + * + * Request-memoized: an account-settings render checks the same viewer in both + * the layout and the page. Outside a Server Component render React evaluates + * the reader normally, so routes and feature-flag lookups are unaffected. */ -export async function isPlatformAdmin(userId: string): Promise { +export const isPlatformAdmin = cache(async (userId: string): Promise => { const [row] = await dbReplica .select({ role: user.role }) .from(user) @@ -50,4 +55,4 @@ export async function isPlatformAdmin(userId: string): Promise { .limit(1) return row?.role === 'admin' -} +}) diff --git a/apps/sim/lib/workspaces/permissions/utils.ts b/apps/sim/lib/workspaces/permissions/utils.ts index 16bb3251d5a..6eaa0e1df54 100644 --- a/apps/sim/lib/workspaces/permissions/utils.ts +++ b/apps/sim/lib/workspaces/permissions/utils.ts @@ -1,3 +1,4 @@ +import { cache } from 'react' import { db } from '@sim/db' import { member, permissions, user, type WorkspaceMode, workspace } from '@sim/db/schema' import { @@ -77,17 +78,12 @@ export async function getWorkspaceById( return exists ? { id: workspaceId } : null } -/** - * Get a workspace with owner info by ID - * - * @param workspaceId - The workspace ID to look up - * @returns The workspace with owner info if found, null otherwise - */ -export async function getWorkspaceWithOwner( +async function selectWorkspaceWithOwner( workspaceId: string, - options?: { includeArchived?: boolean; executor?: DbOrTx; forUpdate?: boolean } + includeArchived: boolean, + executor: DbOrTx, + forUpdate: boolean ): Promise { - const { includeArchived = false, executor = db, forUpdate = false } = options ?? {} const query = executor .select({ id: workspace.id, @@ -110,6 +106,41 @@ export async function getWorkspaceWithOwner( return ws || null } +/** + * Request-memoized plain workspace read, keyed by id and archived visibility. + * A single Server Component render pass resolves the same workspace row through + * several independent gates, so without this the row is re-read once per gate. + * + * Outside a Server Component render React evaluates this normally and retains + * nothing, so API routes and background work are unaffected. + */ +const readWorkspaceWithOwner = cache( + (workspaceId: string, includeArchived: boolean): Promise => + selectWorkspaceWithOwner(workspaceId, includeArchived, db, false) +) + +/** + * Get a workspace with owner info by ID + * + * Transaction-scoped (`executor`) and lock-acquiring (`forUpdate`) reads + * deliberately bypass {@link readWorkspaceWithOwner}: a row read inside one + * caller's transaction, or under a row lock only that caller holds, must never + * be handed to a later caller that took neither. + * + * @param workspaceId - The workspace ID to look up + * @returns The workspace with owner info if found, null otherwise + */ +export function getWorkspaceWithOwner( + workspaceId: string, + options?: { includeArchived?: boolean; executor?: DbOrTx; forUpdate?: boolean } +): Promise { + const { includeArchived = false, executor, forUpdate = false } = options ?? {} + if (executor || forUpdate) { + return selectWorkspaceWithOwner(workspaceId, includeArchived, executor ?? db, forUpdate) + } + return readWorkspaceWithOwner(workspaceId, includeArchived) +} + /** * Resolve the effective workspace permission for a user under the governance * inheritance model: the owners/admins of the organization that owns the @@ -131,18 +162,7 @@ export async function getEffectiveWorkspacePermission( return resolveEffectiveWorkspacePermission(userId, ws.id, ws.organizationId, executor) } -/** - * Check workspace access for a user - * - * Verifies the workspace exists and the user has access to it. - * Returns access level (read/write) based on ownership, explicit permissions, - * and organization-admin inheritance. - * - * @param workspaceId - The workspace ID to check - * @param userId - The user ID to check access for - * @returns WorkspaceAccess object with exists, hasAccess, canWrite, and workspace data - */ -export async function checkWorkspaceAccess( +async function resolveWorkspaceAccessForUser( workspaceId: string, userId: string ): Promise { @@ -167,6 +187,25 @@ export async function checkWorkspaceAccess( return { exists: true, hasAccess, canWrite, canAdmin, workspace: ws, permission } } +/** + * Check workspace access for a user + * + * Verifies the workspace exists and the user has access to it. + * Returns access level (read/write) based on ownership, explicit permissions, + * and organization-admin inheritance. + * + * Request-memoized: a Server Component render pass authorizes the same + * (workspace, viewer) pair through several gates, and the answer cannot change + * mid-render. Outside a render React evaluates the resolver normally, so route + * handlers and background work re-read exactly as before. Takes no executor, so + * no transaction-scoped or locked read can ever be memoized here. + * + * @param workspaceId - The workspace ID to check + * @param userId - The user ID to check access for + * @returns WorkspaceAccess object with exists, hasAccess, canWrite, and workspace data + */ +export const checkWorkspaceAccess = cache(resolveWorkspaceAccessForUser) + /** * Returns `provided` when it was resolved for this exact workspace, otherwise * resolves fresh. The id match is what keeps a caller from authorizing against From fda308249f437aca194a4dcb5291de5985990a35 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 13 Aug 2026 01:01:39 -0700 Subject: [PATCH 04/14] fix(selectors): key CloudWatch lists by search, and stop a caller erasing the credential gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The CloudWatch log-group and log-stream selectors forwarded `search` into the request as `prefix` but left it out of the query key, so every keystroke resolved to the same fresh entry and no refetch fired. Server-side filtering was dead: a log group outside the first page could not be reached. An audit of all 69 selector definitions found these two and no others. useSelectorOptions resolved `args.enabled ?? definition.enabled(...)`, so a caller supplying its own gate replaced the definition's precondition rather than narrowing it. useSelectorDisplayName knows nothing about credentials, so a card holding a saved value with no credential context ran a query that could only reject. The two are now conjoined. The detail hooks keep the override deliberately — resolving one known id needs less context than listing, which their TSDoc already documents. The list-key fix has a test, proven to fail without it. The `enabled` change has none: loading use-selector-query pulls the selector registry and emcn CSS, which cannot resolve in a git worktree. --- .../providers/cloudwatch/selectors.test.ts | 30 +++++++++++++++++++ .../providers/cloudwatch/selectors.ts | 6 ++-- .../sim/hooks/selectors/use-selector-query.ts | 12 +++++++- 3 files changed, 45 insertions(+), 3 deletions(-) create mode 100644 apps/sim/hooks/selectors/providers/cloudwatch/selectors.test.ts diff --git a/apps/sim/hooks/selectors/providers/cloudwatch/selectors.test.ts b/apps/sim/hooks/selectors/providers/cloudwatch/selectors.test.ts new file mode 100644 index 00000000000..66f2ce150ac --- /dev/null +++ b/apps/sim/hooks/selectors/providers/cloudwatch/selectors.test.ts @@ -0,0 +1,30 @@ +/** + * @vitest-environment node + */ +import { describe, expect, it } from 'vitest' +import { cloudwatchSelectors } from '@/hooks/selectors/providers/cloudwatch/selectors' +import type { SelectorQueryArgs } from '@/hooks/selectors/types' + +const AWS_CONTEXT: SelectorQueryArgs['context'] = { + awsAccessKeyId: 'AKIA', + awsSecretAccessKey: 'secret', + awsRegion: 'us-east-1', + logGroupName: '/aws/lambda/fn', +} + +describe('cloudwatch selector query keys', () => { + it.each([['cloudwatch.logGroups' as const], ['cloudwatch.logStreams' as const]])( + '%s scopes its key by search, which fetchList forwards as `prefix`', + (key) => { + const definition = cloudwatchSelectors[key] + const base: SelectorQueryArgs = { key, context: AWS_CONTEXT } + + const noSearch = definition.getQueryKey(base) + const withSearch = definition.getQueryKey({ ...base, search: 'api' }) + const otherSearch = definition.getQueryKey({ ...base, search: 'worker' }) + + expect(withSearch).not.toEqual(noSearch) + expect(withSearch).not.toEqual(otherSearch) + } + ) +}) diff --git a/apps/sim/hooks/selectors/providers/cloudwatch/selectors.ts b/apps/sim/hooks/selectors/providers/cloudwatch/selectors.ts index 065db6680d0..6d7b7480183 100644 --- a/apps/sim/hooks/selectors/providers/cloudwatch/selectors.ts +++ b/apps/sim/hooks/selectors/providers/cloudwatch/selectors.ts @@ -20,11 +20,12 @@ export const cloudwatchSelectors = { key: 'cloudwatch.logGroups', contracts: [selectorContracts.cloudwatchLogGroupsSelectorContract], staleTime: SELECTOR_STALE, - getQueryKey: ({ context }: SelectorQueryArgs) => [ + getQueryKey: ({ context, search }: SelectorQueryArgs) => [ 'selectors', 'cloudwatch.logGroups', context.awsAccessKeyId ?? 'none', context.awsRegion ?? 'none', + search ?? '', ], enabled: ({ context }) => Boolean(context.awsAccessKeyId && context.awsSecretAccessKey && context.awsRegion), @@ -51,12 +52,13 @@ export const cloudwatchSelectors = { key: 'cloudwatch.logStreams', contracts: [selectorContracts.cloudwatchLogStreamsSelectorContract], staleTime: SELECTOR_STALE, - getQueryKey: ({ context }: SelectorQueryArgs) => [ + getQueryKey: ({ context, search }: SelectorQueryArgs) => [ 'selectors', 'cloudwatch.logStreams', context.awsAccessKeyId ?? 'none', context.awsRegion ?? 'none', context.logGroupName ?? 'none', + search ?? '', ], enabled: ({ context }) => Boolean( diff --git a/apps/sim/hooks/selectors/use-selector-query.ts b/apps/sim/hooks/selectors/use-selector-query.ts index ea95d7e879d..a202872f7d2 100644 --- a/apps/sim/hooks/selectors/use-selector-query.ts +++ b/apps/sim/hooks/selectors/use-selector-query.ts @@ -70,7 +70,17 @@ export function useSelectorOptions( context: args.context, search: args.search, } - const isEnabled = args.enabled ?? (definition.enabled ? definition.enabled(queryArgs) : true) + /** + * `definition.enabled` mirrors the preconditions the definition's own `fetchList` / + * `fetchPage` assert (`ensureCredential`, `ensureKnowledgeBase`, an early `return []`), + * so it is a hard precondition for the *list*, not a default a caller may replace. + * A caller's `enabled` narrows further — it never widens — otherwise a card that only + * knows it has a selection runs a fetch that is guaranteed to reject and caches the + * rejection. Unlike {@link useSelectorOptionDetail}, resolving nothing here is never + * cheaper than the list's own preconditions, so there is no case for an override. + */ + const isEnabled = + (args.enabled ?? true) && (definition.enabled ? definition.enabled(queryArgs) : true) const supportsPagination = Boolean(definition.fetchPage) const flatQuery = useQuery({ From 5230c29b5a7379ed0b23c74df943dda2b653c382 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 13 Aug 2026 01:01:39 -0700 Subject: [PATCH 05/14] fix(queries): give optimistic rows collision-free ids MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit generateTempId used Date.now(), so two rows created in the same millisecond shared an id and the first server response overwrote both — leaving one row duplicated and the other's real id lost until a refetch. Now uses generateId(), matching what the workflow mutations already do. Reachable by double-clicking create, or by any scripted or bulk create. Also documents the contract of fetchOAuthConnections, which reports an unknown connection state as disconnected. No consumer reads that field today — both read names and icons, and connection state comes from useWorkspaceCredentials — so letting the query reject would blank the suggested-action rows and drop the credential page to raw provider ids. The note is what stops a future consumer branching on it silently. --- .../hooks/queries/oauth/oauth-connections.ts | 12 ++++++++++ .../queries/utils/optimistic-mutation.test.ts | 24 +++++++++++++++++++ .../queries/utils/optimistic-mutation.ts | 9 ++++++- 3 files changed, 44 insertions(+), 1 deletion(-) create mode 100644 apps/sim/hooks/queries/utils/optimistic-mutation.test.ts diff --git a/apps/sim/hooks/queries/oauth/oauth-connections.ts b/apps/sim/hooks/queries/oauth/oauth-connections.ts index 6a08f723d1f..25338567464 100644 --- a/apps/sim/hooks/queries/oauth/oauth-connections.ts +++ b/apps/sim/hooks/queries/oauth/oauth-connections.ts @@ -56,6 +56,18 @@ function defineServices(): ServiceInfo[] { return servicesList } +/** + * Resolves the service catalog merged with the caller's connections. + * + * A failed request resolves with the bare catalog rather than rejecting, so + * consumers keep correct service names and ids when the merge data is + * unavailable. The cost is that `isConnected`/`accounts` then report *unknown* + * as *disconnected*, which the result cannot distinguish. Read connection + * state from the workspace credentials query (`useWorkspaceCredentials`), which + * surfaces its own errors; a consumer that must branch on `isConnected` here + * needs this fallback removed first, or it will tell a connected user they are + * not. + */ async function fetchOAuthConnections(signal?: AbortSignal): Promise { try { const serviceDefinitions = defineServices() diff --git a/apps/sim/hooks/queries/utils/optimistic-mutation.test.ts b/apps/sim/hooks/queries/utils/optimistic-mutation.test.ts new file mode 100644 index 00000000000..936b1bab173 --- /dev/null +++ b/apps/sim/hooks/queries/utils/optimistic-mutation.test.ts @@ -0,0 +1,24 @@ +/** + * @vitest-environment node + */ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { generateTempId } from '@/hooks/queries/utils/optimistic-mutation' + +describe('generateTempId', () => { + afterEach(() => { + vi.useRealTimers() + }) + + it('is unique for ids created within the same millisecond', () => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')) + + const ids = new Set(Array.from({ length: 100 }, () => generateTempId('temp-folder'))) + + expect(ids.size).toBe(100) + }) + + it('keeps the prefix so callers can identify optimistic rows', () => { + expect(generateTempId('temp-folder').startsWith('temp-folder-')).toBe(true) + }) +}) diff --git a/apps/sim/hooks/queries/utils/optimistic-mutation.ts b/apps/sim/hooks/queries/utils/optimistic-mutation.ts index b734ba941eb..2af55b543af 100644 --- a/apps/sim/hooks/queries/utils/optimistic-mutation.ts +++ b/apps/sim/hooks/queries/utils/optimistic-mutation.ts @@ -1,4 +1,5 @@ import { createLogger } from '@sim/logger' +import { generateId } from '@sim/utils/id' import type { QueryClient } from '@tanstack/react-query' const logger = createLogger('OptimisticMutation') @@ -72,6 +73,12 @@ export function createOptimisticMutationHandlers( } } +/** + * Placeholder id for an optimistic row, held only until the server response + * replaces it. Uses `generateId()` rather than a timestamp so two rows created + * in the same millisecond cannot collide — a collision would make + * `replaceOptimisticEntry` overwrite both entries with one server row. + */ export function generateTempId(prefix: string): string { - return `${prefix}-${Date.now()}` + return `${prefix}-${generateId()}` } From 1e89eb755abead1051e19895625b9a9b0b3b9554 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 13 Aug 2026 01:01:52 -0700 Subject: [PATCH 06/14] fix(queries): close six stale-data gaps found auditing the query layer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each was verified against the mutation that changes the data and the keys that expose it, not taken on report. - Workspace usage/credits were invalidated nowhere in the app. Six sites already refreshed subscriptionKeys after credits moved — post-run, post-wand, limit edits, upgrades, top-ups — and none touched workspace usage, so the credits chip and the run gate held their page-load values until a reload. Adds one shared invalidateWorkspaceUsage and calls it from all six. - Knowledge-base list doc counts went stale: document upload, delete, and bulk delete invalidated only the detail key, though the list carries docCount. - Plan switches that do not redirect refreshed only the host context, leaving subscription and credit state showing the previous plan. - The copilot tool-event handler invalidated a raw workflowKeys.list, which covers only the active scope and skips the selector prefix; it now uses the shared invalidateWorkflowLists like the other thirteen call sites. - scheduleKeys.byId was a strict prefix of scheduleKeys.schedule, so the two addressings aliased, and nothing invalidated byId. De-aliased and invalidated. Not changed: the CSV preview key already folds in the file version and storage key, so a content update addresses a different cache entry — version-in-key is the mechanism there, not a missing invalidation. Tests added for the usage and knowledge fixes, both proven to fail without them. The other four live in files that cannot load in a git worktree (pre-existing postcss resolution failure), so CI is where they first run. --- .../home/hooks/stream/handle-tool-event.ts | 8 +- .../knowledge/hooks/use-knowledge-upload.ts | 6 +- .../upgrade/hooks/use-upgrade-state.ts | 30 ++++++-- .../w/[workflowId]/hooks/use-wand.ts | 2 + .../hooks/use-workflow-execution.ts | 3 + apps/sim/hooks/queries/kb/knowledge.test.ts | 75 +++++++++++++++++++ apps/sim/hooks/queries/kb/knowledge.ts | 12 ++- apps/sim/hooks/queries/schedules.ts | 13 +++- apps/sim/hooks/queries/subscription.ts | 4 + .../sim/hooks/queries/workspace-usage.test.ts | 15 ++++ apps/sim/hooks/queries/workspace-usage.ts | 15 +++- 11 files changed, 166 insertions(+), 17 deletions(-) create mode 100644 apps/sim/hooks/queries/kb/knowledge.test.ts diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-tool-event.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-tool-event.ts index e00e851626d..296f48bc559 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-tool-event.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-tool-event.ts @@ -27,7 +27,7 @@ import { } from '@/app/workspace/[workspaceId]/home/hooks/stream/turn-model' import { deploymentKeys } from '@/hooks/queries/deployments' import { folderKeys } from '@/hooks/queries/utils/folder-keys' -import { workflowKeys } from '@/hooks/queries/workflows' +import { invalidateWorkflowLists } from '@/hooks/queries/utils/invalidate-workflow-lists' type ToolEvent = Extract @@ -58,7 +58,7 @@ function runToolResultSideEffects(ctx: StreamLoopContext, node: ToolNode): void if (deployedWorkflowId && typeof out?.isDeployed === 'boolean') { deps.queryClient.invalidateQueries({ queryKey: deploymentKeys.info(deployedWorkflowId) }) deps.queryClient.invalidateQueries({ queryKey: deploymentKeys.versions(deployedWorkflowId) }) - deps.queryClient.invalidateQueries({ queryKey: workflowKeys.list(deps.workspaceId) }) + void invalidateWorkflowLists(deps.queryClient, deps.workspaceId) } } @@ -66,7 +66,9 @@ function runToolResultSideEffects(ctx: StreamLoopContext, node: ToolNode): void deps.queryClient.invalidateQueries({ queryKey: folderKeys.list(deps.workspaceId) }) } if (WORKFLOW_MUTATION_TOOL_NAMES.has(name) && isSuccess) { - deps.queryClient.invalidateQueries({ queryKey: workflowKeys.list(deps.workspaceId) }) + // `rm` archives, so the archived list moves too — and the shared helper also + // refreshes the workflow selector lists that `@`-mentions and pickers read. + void invalidateWorkflowLists(deps.queryClient, deps.workspaceId, ['active', 'archived']) } const extractedResources = diff --git a/apps/sim/app/workspace/[workspaceId]/knowledge/hooks/use-knowledge-upload.ts b/apps/sim/app/workspace/[workspaceId]/knowledge/hooks/use-knowledge-upload.ts index c9cd91416ef..5923e9e03f6 100644 --- a/apps/sim/app/workspace/[workspaceId]/knowledge/hooks/use-knowledge-upload.ts +++ b/apps/sim/app/workspace/[workspaceId]/knowledge/hooks/use-knowledge-upload.ts @@ -209,7 +209,11 @@ export function useKnowledgeUpload(options: UseKnowledgeUploadOptions = {}) { setUploadProgress((prev) => ({ ...prev, stage: 'processing' })) logger.info(`Successfully started processing ${uploadedDocuments.length} documents`) - await queryClient.invalidateQueries({ queryKey: knowledgeKeys.detail(knowledgeBaseId) }) + await Promise.all([ + queryClient.invalidateQueries({ queryKey: knowledgeKeys.detail(knowledgeBaseId) }), + /** The knowledge-base list rows carry `docCount`, so an upload changes them too. */ + queryClient.invalidateQueries({ queryKey: knowledgeKeys.lists() }), + ]) return uploadedDocuments } catch (err) { diff --git a/apps/sim/app/workspace/[workspaceId]/upgrade/hooks/use-upgrade-state.ts b/apps/sim/app/workspace/[workspaceId]/upgrade/hooks/use-upgrade-state.ts index 1792360d7d6..e0a91ba6477 100644 --- a/apps/sim/app/workspace/[workspaceId]/upgrade/hooks/use-upgrade-state.ts +++ b/apps/sim/app/workspace/[workspaceId]/upgrade/hooks/use-upgrade-state.ts @@ -9,7 +9,9 @@ import type { WorkspaceHostContext } from '@/lib/api/contracts/workspaces' import { useSubscriptionUpgrade } from '@/lib/billing/client/upgrade' import { CREDIT_TIERS } from '@/lib/billing/constants' import { getPlanTierCredits, isEnterprise, isFree, isPro, isTeam } from '@/lib/billing/plan-helpers' +import { subscriptionKeys } from '@/hooks/queries/subscription' import { workspaceHostKeys } from '@/hooks/queries/workspace-host' +import { invalidateWorkspaceUsage } from '@/hooks/queries/workspace-usage' const PRO_TIER = CREDIT_TIERS[0] const MAX_TIER = CREDIT_TIERS[1] @@ -89,8 +91,20 @@ export function useUpgradeState({ } }, [ownerBilling.billingInterval, subscription.isPaid]) - const refreshHostContext = useCallback( - () => queryClient.invalidateQueries({ queryKey: workspaceHostKeys.detail(workspaceId) }), + /** + * A non-redirect plan switch settles server-side immediately, so every read that + * describes the plan has to be refetched — the host context the page renders from, + * the subscription/usage reads the billing surfaces share, and the workspace credit + * availability that drives the credits chip and the run gate. + */ + const refreshBillingState = useCallback( + () => + Promise.all([ + queryClient.invalidateQueries({ queryKey: workspaceHostKeys.detail(workspaceId) }), + queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }), + queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() }), + invalidateWorkspaceUsage(queryClient), + ]), [queryClient, workspaceId] ) @@ -123,9 +137,9 @@ export function useUpgradeState({ await requestJson(billingSwitchPlanContract, { body: { targetPlanName: subscription.plan, interval, workspaceId }, }) - await refreshHostContext() + await refreshBillingState() }, - [isLegacyPlan, refreshHostContext, subscription.plan, workspaceId] + [isLegacyPlan, refreshBillingState, subscription.plan, workspaceId] ) const currentCredits = getPlanTierCredits(subscription.plan) @@ -154,11 +168,11 @@ export function useUpgradeState({ workspaceId, }, }) - await refreshHostContext() + await refreshBillingState() } catch (e) { toast.error(getErrorMessage(e, 'Failed to upgrade')) } - }, [subscription.isTeam, isAnnual, refreshHostContext, workspaceId]) + }, [subscription.isTeam, isAnnual, refreshBillingState, workspaceId]) const onUpgradeToOtherTier = useCallback(async () => { const onMax = @@ -170,11 +184,11 @@ export function useUpgradeState({ await requestJson(billingSwitchPlanContract, { body: { targetPlanName, workspaceId }, }) - await refreshHostContext() + await refreshBillingState() } catch (e) { toast.error(getErrorMessage(e, 'Failed to switch plan')) } - }, [subscription.plan, subscription.isTeam, refreshHostContext, workspaceId]) + }, [subscription.plan, subscription.isTeam, refreshBillingState, workspaceId]) return { isLoading: false, diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts index 0e0a8b795fc..0bf8bac093d 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts +++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts @@ -11,6 +11,7 @@ import { readSSEStream } from '@/lib/core/utils/sse' import { shouldStripCodeFences, stripCodeFences } from '@/lib/wand/strip-code-fences' import type { GenerationType } from '@/blocks/types' import { subscriptionKeys } from '@/hooks/queries/subscription' +import { invalidateWorkspaceUsage } from '@/hooks/queries/workspace-usage' import { useSettingsNavigation } from '@/hooks/use-settings-navigation' import { useWorkflowRegistry } from '@/stores/workflows/registry/store' @@ -302,6 +303,7 @@ export function useWand({ setTimeout(() => { queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }) + void invalidateWorkspaceUsage(queryClient) }, 1000) } catch (error: any) { if (error.name === 'AbortError') { diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts index 37d7fb5d5cb..97f9fee5d53 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts +++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts @@ -63,6 +63,7 @@ import { hasExecutionResult } from '@/executor/utils/errors' import { coerceValue } from '@/executor/utils/start-block' import { subscriptionKeys } from '@/hooks/queries/subscription' import { getWorkflows } from '@/hooks/queries/utils/workflow-cache' +import { invalidateWorkspaceUsage } from '@/hooks/queries/workspace-usage' import { isExecutionStreamHttpError, SSEEventHandlerError, @@ -924,6 +925,7 @@ export function useWorkflowExecution() { // Invalidate subscription queries to update usage setTimeout(() => { queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }) + void invalidateWorkspaceUsage(queryClient) }, 1000) safeEnqueue(encodeSSE({ event: 'final', data: result })) @@ -1460,6 +1462,7 @@ export function useWorkflowExecution() { } setTimeout(() => { queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }) + void invalidateWorkspaceUsage(queryClient) }, 1000) } }, diff --git a/apps/sim/hooks/queries/kb/knowledge.test.ts b/apps/sim/hooks/queries/kb/knowledge.test.ts new file mode 100644 index 00000000000..4c01ae0fcf5 --- /dev/null +++ b/apps/sim/hooks/queries/kb/knowledge.test.ts @@ -0,0 +1,75 @@ +/** + * @vitest-environment node + */ + +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ + requestJson: vi.fn(), + useMutation: vi.fn(), + invalidateQueries: vi.fn(), +})) + +vi.mock('@tanstack/react-query', () => ({ + keepPreviousData: Symbol('keepPreviousData'), + useInfiniteQuery: vi.fn(), + useMutation: mocks.useMutation, + useQuery: vi.fn(), + useQueryClient: vi.fn(() => ({ invalidateQueries: mocks.invalidateQueries })), +})) + +vi.mock('@sim/emcn', () => ({ + toast: { error: vi.fn(), success: vi.fn() }, +})) + +vi.mock('@/lib/api/client/request', () => ({ + requestJson: mocks.requestJson, +})) + +import { useBulkDocumentOperation, useDeleteDocument } from '@/hooks/queries/kb/knowledge' +import { knowledgeKeys } from '@/hooks/queries/utils/knowledge-keys' + +interface CapturedMutation { + onSettled: (data: unknown, error: unknown, variables: Record) => void +} + +function captureMutation(build: () => unknown): CapturedMutation { + let captured: CapturedMutation | undefined + mocks.useMutation.mockImplementation((options: CapturedMutation) => { + captured = options + return {} + }) + build() + if (!captured) throw new Error('useMutation was not called') + return captured +} + +describe('knowledge document mutations', () => { + beforeEach(() => { + vi.clearAllMocks() + }) + + it('invalidates the knowledge-base lists when a document is deleted', () => { + const mutation = captureMutation(() => useDeleteDocument()) + + mutation.onSettled(undefined, undefined, { knowledgeBaseId: 'kb-1', documentId: 'doc-1' }) + + expect(mocks.invalidateQueries).toHaveBeenCalledWith({ queryKey: knowledgeKeys.lists() }) + }) + + it('invalidates the knowledge-base lists on a bulk delete', () => { + const mutation = captureMutation(() => useBulkDocumentOperation()) + + mutation.onSettled(undefined, undefined, { knowledgeBaseId: 'kb-1', operation: 'delete' }) + + expect(mocks.invalidateQueries).toHaveBeenCalledWith({ queryKey: knowledgeKeys.lists() }) + }) + + it('leaves the knowledge-base lists alone on a bulk enable', () => { + const mutation = captureMutation(() => useBulkDocumentOperation()) + + mutation.onSettled(undefined, undefined, { knowledgeBaseId: 'kb-1', operation: 'enable' }) + + expect(mocks.invalidateQueries).not.toHaveBeenCalledWith({ queryKey: knowledgeKeys.lists() }) + }) +}) diff --git a/apps/sim/hooks/queries/kb/knowledge.ts b/apps/sim/hooks/queries/kb/knowledge.ts index 0dd4606b44f..f46457282fc 100644 --- a/apps/sim/hooks/queries/kb/knowledge.ts +++ b/apps/sim/hooks/queries/kb/knowledge.ts @@ -539,6 +539,10 @@ export function useDeleteDocument() { queryClient.invalidateQueries({ queryKey: knowledgeKeys.detail(knowledgeBaseId), }) + /** The knowledge-base list rows carry `docCount`, so removing a document changes them too. */ + queryClient.invalidateQueries({ + queryKey: knowledgeKeys.lists(), + }) }, }) } @@ -573,10 +577,16 @@ export function useBulkDocumentOperation() { return useMutation({ mutationFn: bulkDocumentOperation, - onSettled: (_data, _error, { knowledgeBaseId }) => { + onSettled: (_data, _error, { knowledgeBaseId, operation }) => { queryClient.invalidateQueries({ queryKey: knowledgeKeys.detail(knowledgeBaseId), }) + /** Only a bulk delete changes the `docCount` the knowledge-base list rows render. */ + if (operation === 'delete') { + queryClient.invalidateQueries({ + queryKey: knowledgeKeys.lists(), + }) + } }, }) } diff --git a/apps/sim/hooks/queries/schedules.ts b/apps/sim/hooks/queries/schedules.ts index a2db0ad31e9..c33ed845663 100644 --- a/apps/sim/hooks/queries/schedules.ts +++ b/apps/sim/hooks/queries/schedules.ts @@ -27,7 +27,13 @@ export const scheduleKeys = { details: () => [...scheduleKeys.all, 'detail'] as const, schedule: (workflowId: string, blockId: string) => [...scheduleKeys.details(), workflowId, blockId] as const, - byId: (scheduleId: string) => [...scheduleKeys.details(), scheduleId] as const, + /** + * By-id reads sit under their own segment rather than directly under `details()`: + * a bare `[...details(), scheduleId]` is a prefix of `schedule(scheduleId, blockId)`, + * so the two addressings of the same schedule would alias in the cache. + */ + byIds: () => [...scheduleKeys.details(), 'by-id'] as const, + byId: (scheduleId: string) => [...scheduleKeys.byIds(), scheduleId] as const, } export type ScheduleData = WorkflowScheduleRow @@ -196,7 +202,7 @@ export function useReactivateSchedule() { body: { action: 'reactivate' }, }) - return { workflowId, blockId, workspaceId } + return { scheduleId, workflowId, blockId, workspaceId } }, onSuccess: ({ workflowId, blockId }) => { logger.info('Schedule reactivated', { workflowId, blockId }) @@ -206,9 +212,10 @@ export function useReactivateSchedule() { }, onSettled: async (data) => { if (!data) return - const { workflowId, blockId, workspaceId } = data + const { scheduleId, workflowId, blockId, workspaceId } = data await Promise.all([ queryClient.invalidateQueries({ queryKey: scheduleKeys.schedule(workflowId, blockId) }), + queryClient.invalidateQueries({ queryKey: scheduleKeys.byId(scheduleId) }), workspaceId ? queryClient.invalidateQueries({ queryKey: scheduleKeys.list(workspaceId) }) : Promise.resolve(), diff --git a/apps/sim/hooks/queries/subscription.ts b/apps/sim/hooks/queries/subscription.ts index bafee85ad5b..946f0a0f14d 100644 --- a/apps/sim/hooks/queries/subscription.ts +++ b/apps/sim/hooks/queries/subscription.ts @@ -14,6 +14,7 @@ import { } from '@/lib/api/contracts/subscription' import { organizationKeys } from '@/hooks/queries/organization' import { workspaceKeys } from '@/hooks/queries/workspace' +import { invalidateWorkspaceUsage } from '@/hooks/queries/workspace-usage' export type { SubscriptionApiResponse } @@ -265,6 +266,7 @@ export function useUpdateUsageLimit() { return Promise.all([ queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }), queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() }), + invalidateWorkspaceUsage(queryClient), ]) }, }) @@ -291,6 +293,7 @@ export function useUpgradeSubscription() { queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() }), queryClient.invalidateQueries({ queryKey: subscriptionKeys.invoicesAll() }), queryClient.invalidateQueries({ queryKey: workspaceKeys.lists() }), + invalidateWorkspaceUsage(queryClient), ...(variables.orgId ? [ queryClient.invalidateQueries({ @@ -328,6 +331,7 @@ export function usePurchaseCredits() { return Promise.all([ queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }), queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() }), + invalidateWorkspaceUsage(queryClient), ...(variables.orgId ? [ queryClient.invalidateQueries({ diff --git a/apps/sim/hooks/queries/workspace-usage.test.ts b/apps/sim/hooks/queries/workspace-usage.test.ts index f468d2b07e2..c2b7b09117d 100644 --- a/apps/sim/hooks/queries/workspace-usage.test.ts +++ b/apps/sim/hooks/queries/workspace-usage.test.ts @@ -18,6 +18,7 @@ import { import { fetchWorkspaceCreditAvailability, fetchWorkspaceUsageGate, + invalidateWorkspaceUsage, WORKSPACE_CREDIT_AVAILABILITY_STALE_TIME, WORKSPACE_USAGE_GATE_STALE_TIME, workspaceUsageKeys, @@ -65,4 +66,18 @@ describe('workspace usage gate query', () => { signal, }) }) + + it('invalidates both usage families so a spend or top-up refetches them', async () => { + const invalidateQueries = vi.fn().mockResolvedValue(undefined) + const queryClient = { invalidateQueries } as unknown as Parameters< + typeof invalidateWorkspaceUsage + >[0] + + await invalidateWorkspaceUsage(queryClient) + + expect(invalidateQueries).toHaveBeenCalledWith({ + queryKey: workspaceUsageKeys.creditAvailabilities(), + }) + expect(invalidateQueries).toHaveBeenCalledWith({ queryKey: workspaceUsageKeys.gates() }) + }) }) diff --git a/apps/sim/hooks/queries/workspace-usage.ts b/apps/sim/hooks/queries/workspace-usage.ts index df3ad7e0d08..a49be5c4d77 100644 --- a/apps/sim/hooks/queries/workspace-usage.ts +++ b/apps/sim/hooks/queries/workspace-usage.ts @@ -1,4 +1,4 @@ -import { useQuery } from '@tanstack/react-query' +import { type QueryClient, useQuery } from '@tanstack/react-query' import { requestJson } from '@/lib/api/client/request' import { getWorkspaceCreditAvailabilityContract, @@ -16,6 +16,19 @@ export const workspaceUsageKeys = { gate: (workspaceId: string) => [...workspaceUsageKeys.gates(), workspaceId] as const, } +/** + * Invalidates the workspace credit/usage reads after anything that moves the balance — + * a run that spends credits, a top-up, a plan change, or a usage-limit edit. Both + * families are keyed per workspace but derive from the same billing account, so the + * family prefixes (not a single workspace's key) are what has to be refetched. + */ +export function invalidateWorkspaceUsage(queryClient: QueryClient): Promise { + return Promise.all([ + queryClient.invalidateQueries({ queryKey: workspaceUsageKeys.creditAvailabilities() }), + queryClient.invalidateQueries({ queryKey: workspaceUsageKeys.gates() }), + ]).then(() => undefined) +} + export const WORKSPACE_CREDIT_AVAILABILITY_STALE_TIME = 30 * 1000 export const WORKSPACE_USAGE_GATE_STALE_TIME = 30 * 1000 From 173b1beb7b91c1eb51e3103c64b4d10b556ef67c Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 13 Aug 2026 01:22:29 -0700 Subject: [PATCH 07/14] improvement(queries): make the row-list prefix non-collidable, and share the usage refresh MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two corrections from reviewing the previous commits. patchCachedRows was fixed with a predicate naming the sibling subtrees to skip — a denylist that rots the moment a fifth subtree is added under rowsRoot. The key factory already separated row lists under an 'infinite' segment; it just had no prefix accessor, so every caller reached for the parent and subtracted. Adding infiniteRowsRoot lets the walk be an allowlist by construction and deletes the predicate, the helper, and both docblocks explaining the subtraction. The searched-rows view is consequently no longer patched by a cell edit and is left to its own refetch — it holds a flat result, not pages. That is recorded on the function rather than left to be rediscovered. The delayed usage refresh was written out three times across two files, a duplication the previous commit enlarged rather than introduced. It is now one scheduleUsageRefresh beside the keys it invalidates, which also gives the bare 1000ms a name and one place to change it. --- .../w/[workflowId]/hooks/use-wand.ts | 8 +--- .../hooks/use-workflow-execution.ts | 13 ++----- apps/sim/hooks/queries/tables.ts | 38 ++++++------------- apps/sim/hooks/queries/utils/table-keys.ts | 11 +++++- apps/sim/hooks/queries/workspace-usage.ts | 32 +++++++++++++--- 5 files changed, 54 insertions(+), 48 deletions(-) diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts index 0bf8bac093d..a17c8700934 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts +++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts @@ -10,8 +10,7 @@ import { wandGenerateStreamContract } from '@/lib/api/contracts' import { readSSEStream } from '@/lib/core/utils/sse' import { shouldStripCodeFences, stripCodeFences } from '@/lib/wand/strip-code-fences' import type { GenerationType } from '@/blocks/types' -import { subscriptionKeys } from '@/hooks/queries/subscription' -import { invalidateWorkspaceUsage } from '@/hooks/queries/workspace-usage' +import { scheduleUsageRefresh } from '@/hooks/queries/workspace-usage' import { useSettingsNavigation } from '@/hooks/use-settings-navigation' import { useWorkflowRegistry } from '@/stores/workflows/registry/store' @@ -301,10 +300,7 @@ export function useWand({ strippedFences: generatedContent !== accumulatedContent, }) - setTimeout(() => { - queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }) - void invalidateWorkspaceUsage(queryClient) - }, 1000) + scheduleUsageRefresh(queryClient) } catch (error: any) { if (error.name === 'AbortError') { logger.debug('Wand generation cancelled') diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts index 97f9fee5d53..bfeabf5466d 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts +++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts @@ -61,9 +61,8 @@ import type { SerializableExecutionState } from '@/executor/execution/types' import type { BlockLog, BlockState, ExecutionResult, StreamingExecution } from '@/executor/types' import { hasExecutionResult } from '@/executor/utils/errors' import { coerceValue } from '@/executor/utils/start-block' -import { subscriptionKeys } from '@/hooks/queries/subscription' import { getWorkflows } from '@/hooks/queries/utils/workflow-cache' -import { invalidateWorkspaceUsage } from '@/hooks/queries/workspace-usage' +import { scheduleUsageRefresh } from '@/hooks/queries/workspace-usage' import { isExecutionStreamHttpError, SSEEventHandlerError, @@ -923,10 +922,7 @@ export function useWorkflowExecution() { } // Invalidate subscription queries to update usage - setTimeout(() => { - queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }) - void invalidateWorkspaceUsage(queryClient) - }, 1000) + scheduleUsageRefresh(queryClient) safeEnqueue(encodeSSE({ event: 'final', data: result })) // Note: Logs are already persisted server-side via execution-core.ts @@ -1460,10 +1456,7 @@ export function useWorkflowExecution() { setIsExecuting(activeWorkflowId, false) setActiveBlocks(activeWorkflowId, new Set()) } - setTimeout(() => { - queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }) - void invalidateWorkspaceUsage(queryClient) - }, 1000) + scheduleUsageRefresh(queryClient) } }, diff --git a/apps/sim/hooks/queries/tables.ts b/apps/sim/hooks/queries/tables.ts index ada5568bbf9..6c9bfc564a1 100644 --- a/apps/sim/hooks/queries/tables.ts +++ b/apps/sim/hooks/queries/tables.ts @@ -829,21 +829,18 @@ function withOptimisticAutoFireExec(groups: WorkflowGroup[], row: TableRow): Tab return { ...row, executions: nextExecutions } } -/** - * Whether a cache key under {@link tableKeys.rowsRoot} holds paged row-list data. - * - * `rowsRoot` is a shared prefix: the `find` (search results) and `write` (pending - * writes) subtrees live beneath it with entirely different shapes, so any bulk - * update walking the prefix has to exclude them or it will be handed a value whose - * `pages` it cannot map. - */ -function isRowListQueryKey(queryKey: readonly unknown[]): boolean { - return !queryKey.includes('find') && !queryKey.includes('write') -} - /** * Apply a row-level transformation to all cached infinite row queries for this * table. Used for cell edits where positions don't change. + * + * Walks {@link tableKeys.infiniteRowsRoot} rather than `rowsRoot`: the latter is a + * shared parent, and handing this updater a `find` entry — a flat + * {@link TableFindResult}, not pages — throws on `old.pages` inside `onMutate`, so + * the whole cell edit would reject before reaching the server. + * + * A consequence worth knowing: an open search-results view is therefore left to its + * own refetch rather than patched here, since it holds a different shape. Patching + * it too would need its own updater keyed on {@link tableKeys.find}. */ function patchCachedRows( queryClient: ReturnType, @@ -851,20 +848,9 @@ function patchCachedRows( patchRow: (row: TableRow) => TableRow ) { queryClient.setQueriesData>( - { - queryKey: tableKeys.rowsRoot(tableId), - exact: false, - /** - * `rowsRoot` is a prefix, not a leaf: `rowWrites` and `find` hang off it and - * hold non-paged shapes. Without this they are handed to the updater below, - * whose `old.pages.map` throws — and because this runs inside `onMutate`, the - * whole cell edit rejects before it reaches the server. Only reachable once a - * `find` entry exists, i.e. after the user has searched the table. - */ - predicate: (query) => isRowListQueryKey(query.queryKey), - }, + { queryKey: tableKeys.infiniteRowsRoot(tableId), exact: false }, (old) => { - if (!old?.pages) return old + if (!old) return old return { ...old, pages: old.pages.map((page) => ({ ...page, rows: page.rows.map(patchRow) })), @@ -882,7 +868,7 @@ function patchCachedRows( * `find`/`write` subtrees aren't row-list data and never match. */ function isDefaultOrderRowsQuery(queryKey: readonly unknown[]): boolean { - if (!isRowListQueryKey(queryKey)) return false + if (queryKey.includes('find') || queryKey.includes('write')) return false const last = queryKey[queryKey.length - 1] if (typeof last !== 'string') return false try { diff --git a/apps/sim/hooks/queries/utils/table-keys.ts b/apps/sim/hooks/queries/utils/table-keys.ts index 0721bbfb308..08d9d410329 100644 --- a/apps/sim/hooks/queries/utils/table-keys.ts +++ b/apps/sim/hooks/queries/utils/table-keys.ts @@ -25,8 +25,17 @@ export const tableKeys = { exportJobs: (workspaceId?: string) => [...tableKeys.all, 'export-jobs', workspaceId ?? ''] as const, rowsRoot: (tableId: string) => [...tableKeys.detail(tableId), 'rows'] as const, + /** + * Prefix covering only the paged row lists. + * + * `rowsRoot` is a shared parent — `rowWrites` and `find` hang off it holding + * entirely different shapes — so anything walking the cache to update or snapshot + * row pages must start here instead. Reaching for `rowsRoot` and subtracting the + * siblings is a denylist that rots the moment a fifth subtree is added. + */ + infiniteRowsRoot: (tableId: string) => [...tableKeys.rowsRoot(tableId), 'infinite'] as const, infiniteRows: (tableId: string, paramsKey: string) => - [...tableKeys.rowsRoot(tableId), 'infinite', paramsKey] as const, + [...tableKeys.infiniteRowsRoot(tableId), paramsKey] as const, rowWrites: (tableId: string) => [...tableKeys.rowsRoot(tableId), 'write'] as const, find: (tableId: string, paramsKey: string) => [...tableKeys.rowsRoot(tableId), 'find', paramsKey] as const, diff --git a/apps/sim/hooks/queries/workspace-usage.ts b/apps/sim/hooks/queries/workspace-usage.ts index a49be5c4d77..c23a18a528c 100644 --- a/apps/sim/hooks/queries/workspace-usage.ts +++ b/apps/sim/hooks/queries/workspace-usage.ts @@ -6,6 +6,7 @@ import { type WorkspaceCreditAvailability, type WorkspaceUsageGate, } from '@/lib/api/contracts/workspaces' +import { subscriptionKeys } from '@/hooks/queries/subscription' export const workspaceUsageKeys = { all: ['workspace-usage'] as const, @@ -16,22 +17,28 @@ export const workspaceUsageKeys = { gate: (workspaceId: string) => [...workspaceUsageKeys.gates(), workspaceId] as const, } +export const WORKSPACE_CREDIT_AVAILABILITY_STALE_TIME = 30 * 1000 +export const WORKSPACE_USAGE_GATE_STALE_TIME = 30 * 1000 + +/** + * Usage is written asynchronously as a run settles, so a refetch fired on completion + * races the write and re-reads the old balance. + */ +const USAGE_SETTLE_DELAY_MS = 1000 + /** * Invalidates the workspace credit/usage reads after anything that moves the balance — * a run that spends credits, a top-up, a plan change, or a usage-limit edit. Both * families are keyed per workspace but derive from the same billing account, so the * family prefixes (not a single workspace's key) are what has to be refetched. */ -export function invalidateWorkspaceUsage(queryClient: QueryClient): Promise { +export function invalidateWorkspaceUsage(queryClient: QueryClient) { return Promise.all([ queryClient.invalidateQueries({ queryKey: workspaceUsageKeys.creditAvailabilities() }), queryClient.invalidateQueries({ queryKey: workspaceUsageKeys.gates() }), - ]).then(() => undefined) + ]) } -export const WORKSPACE_CREDIT_AVAILABILITY_STALE_TIME = 30 * 1000 -export const WORKSPACE_USAGE_GATE_STALE_TIME = 30 * 1000 - export function fetchWorkspaceCreditAvailability( workspaceId: string, signal?: AbortSignal @@ -69,3 +76,18 @@ export function useWorkspaceUsageGate(workspaceId?: string) { staleTime: WORKSPACE_USAGE_GATE_STALE_TIME, }) } + +/** + * Refreshes the billing reads a run touches, after a delay: usage is written + * asynchronously as the run settles, so an immediate refetch races the write and + * re-reads the pre-run balance. + * + * Shared by the surfaces that spend credits — workflow execution and wand + * generation — so the delay and the key set stay in one place. + */ +export function scheduleUsageRefresh(queryClient: QueryClient) { + setTimeout(() => { + void queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }) + void invalidateWorkspaceUsage(queryClient) + }, USAGE_SETTLE_DELAY_MS) +} From 03486a3e95aa07098b28553af561d15521950891 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 13 Aug 2026 01:32:12 -0700 Subject: [PATCH 08/14] improvement(prefetch): budget the workspace file seed and narrow its columns --- .../workspace/[workspaceId]/files/prefetch.ts | 7 +- .../[workspaceId]/lib/prefetch.test.ts | 54 +++++++++++++- .../app/workspace/[workspaceId]/prefetch.ts | 72 +++++++++++++------ .../workspace/workspace-file-manager.ts | 54 ++++++++++++-- .../workspace/workspace-file-query.test.ts | 54 +++++++++++++- apps/sim/lib/workspace-files/queries.ts | 12 +++- 6 files changed, 220 insertions(+), 33 deletions(-) diff --git a/apps/sim/app/workspace/[workspaceId]/files/prefetch.ts b/apps/sim/app/workspace/[workspaceId]/files/prefetch.ts index cb83b3b9d32..54d0276fcef 100644 --- a/apps/sim/app/workspace/[workspaceId]/files/prefetch.ts +++ b/apps/sim/app/workspace/[workspaceId]/files/prefetch.ts @@ -16,10 +16,13 @@ import { * (scope `active`), so the browser paints populated on first render. * * The FILE LIST itself is deliberately not here: the sidebar reads it on every workspace route, so - * it is prefetched by `prefetchWorkspaceSidebar` in the layout — the only boundary that renders + * it is seeded by `prefetchWorkspaceSidebar` in the layout — the only boundary that renders * before the sidebar registers the query. Prefetching it again here would re-read it per request * and still not reach the server render (`HydrationBoundary` defers an already-seen query to an - * effect, which SSR never runs). See the note on that entry. + * effect, which SSR never runs). See the note on that entry. The layout declines to seed a + * workspace whose file list exceeds its payload budget; recovering those here would mean + * mirroring that budget check inversely, since an unconditional prefetch would re-read and + * duplicate the entry for every workspace under the budget. * * Folders and the chrome reads all go through the data layer, shaped to their route contracts so a * hydrated entry matches a client fetch. diff --git a/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts b/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts index da334353073..c9e9e242303 100644 --- a/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts @@ -97,7 +97,10 @@ vi.mock('@sim/emcn', () => ({ import { prefetchFilesBrowser } from '@/app/workspace/[workspaceId]/files/prefetch' import { prefetchKnowledgeBases } from '@/app/workspace/[workspaceId]/knowledge/prefetch' -import { prefetchWorkspaceSidebar } from '@/app/workspace/[workspaceId]/prefetch' +import { + prefetchWorkspaceSidebar, + WORKSPACE_FILE_SEED_MAX, +} from '@/app/workspace/[workspaceId]/prefetch' import { prefetchTables } from '@/app/workspace/[workspaceId]/tables/prefetch' import { folderKeys } from '@/hooks/queries/utils/folder-keys' import { knowledgeKeys } from '@/hooks/queries/utils/knowledge-keys' @@ -510,6 +513,55 @@ describe('workspace list prefetches', () => { expect(client.getQueryData(workspaceKeys.list('active'))).toBeUndefined() }) + /** + * The file list is seeded on every workspace route, so it is the one entry whose size + * scales with a workspace's content on routes that never read it. It is read one row + * past the budget so the overflow is detectable. + */ + it('seeds the file list, bounded by the document payload budget', async () => { + const files = [{ id: 'file-1', name: 'a.txt' }] + mockListWorkspaceFilesWithShares.mockResolvedValue(files) + const client = makeClient() + + await prefetchWorkspaceSidebar(client, WORKSPACE_ID, USER_ID, HOST_CONTEXT, null) + + expect(mockListWorkspaceFilesWithShares).toHaveBeenCalledWith(WORKSPACE_ID, 'active', { + limit: WORKSPACE_FILE_SEED_MAX + 1, + }) + expect(client.getQueryData(workspaceFilesKeys.list(WORKSPACE_ID, 'active'))).toEqual(files) + }) + + /** + * The load-bearing half of the budget: a workspace over it seeds NOTHING rather than the + * prefix it read. The sidebar search filters this list client-side and the Files browser + * renders it as the workspace's files, so a truncated seed would silently hide files — + * the client fetch must reach the route for the complete list instead. + */ + it('seeds nothing when the workspace exceeds the budget', async () => { + mockListWorkspaceFilesWithShares.mockResolvedValue( + Array.from({ length: WORKSPACE_FILE_SEED_MAX + 1 }, (_, index) => ({ + id: `file-${index}`, + name: `${index}.txt`, + })) + ) + const client = makeClient() + + await prefetchWorkspaceSidebar(client, WORKSPACE_ID, USER_ID, HOST_CONTEXT, null) + + expect(client.getQueryData(workspaceFilesKeys.list(WORKSPACE_ID, 'active'))).toBeUndefined() + }) + + /** A failed file read is an optimization loss, not a render failure. */ + it('does not throw when the file read rejects, and seeds no files', async () => { + mockListWorkspaceFilesWithShares.mockRejectedValue(new Error('500')) + const client = makeClient() + + await expect( + prefetchWorkspaceSidebar(client, WORKSPACE_ID, USER_ID, HOST_CONTEXT, null) + ).resolves.toBeUndefined() + expect(client.getQueryData(workspaceFilesKeys.list(WORKSPACE_ID, 'active'))).toBeUndefined() + }) + /** Guards the mismatch check that keeps one workspace's data out of another's cache. */ it('seeds nothing when the host context is for a different workspace', async () => { mockListWorkspacesForViewer.mockResolvedValue(LIST_PAYLOAD) diff --git a/apps/sim/app/workspace/[workspaceId]/prefetch.ts b/apps/sim/app/workspace/[workspaceId]/prefetch.ts index c24198fb51f..0ed4a3cba47 100644 --- a/apps/sim/app/workspace/[workspaceId]/prefetch.ts +++ b/apps/sim/app/workspace/[workspaceId]/prefetch.ts @@ -25,10 +25,7 @@ import { workflowKeys } from '@/hooks/queries/utils/workflow-keys' import { mapWorkflow, WORKFLOW_LIST_STALE_TIME } from '@/hooks/queries/utils/workflow-list-query' import { normalizeWorkspacesResponse } from '@/hooks/queries/utils/workspace-list-query' import { WORKSPACE_PERMISSIONS_STALE_TIME, workspaceKeys } from '@/hooks/queries/workspace' -import { - WORKSPACE_FILES_LIST_STALE_TIME, - workspaceFilesKeys, -} from '@/hooks/queries/workspace-files' +import { workspaceFilesKeys } from '@/hooks/queries/workspace-files' import { WORKSPACE_HOST_CONTEXT_STALE_TIME, workspaceHostKeys, @@ -95,6 +92,56 @@ async function seedWorkspaceList( } } +/** + * How many files the layout is willing to inline into the document. + * + * The file list is seeded on EVERY workspace route (see the call site), so its cost is + * paid per navigation into the app, not per visit to Files. At roughly 500 bytes of + * serialized JSON per file, this budgets the entry at ~150 KB; a workspace with + * thousands of files would otherwise push more than a megabyte of HTML ahead of first + * paint on the logs, settings, and editor routes that never read it. + * + * The read is capped at one row past the budget purely to detect the overflow. A + * workspace above it seeds NOTHING rather than a prefix: the sidebar search filters this + * list client-side and the Files browser renders it as the workspace's files, so a + * truncated seed would silently hide files. Those workspaces fetch the complete list + * from the route instead — which for a list that large is also the cheaper first paint. + */ +export const WORKSPACE_FILE_SEED_MAX = 300 + +/** + * Seeds the workspace's file list, which the sidebar's search modal reads on EVERY + * workspace route — so this query is registered by sidebar chrome before any page + * renders. That ordering is why it has to be seeded HERE and not only by the Files + * pages: `HydrationBoundary` hydrates a query the cache has already seen from a + * `useEffect`, which never runs during SSR, so a page-level boundary can only ever hand + * this entry to the client. Seeding it with the layout's own boundary — the first one to + * render — is what lets the server paint the Files browser and the open file's header + * populated instead of shipping a spinner and resolving it a beat later on the client. + * + * Seeded rather than prefetched so it can decline to create an entry at all when the + * workspace exceeds {@link WORKSPACE_FILE_SEED_MAX}: `prefetchQuery` always creates one, + * and a partial one would be read as the whole list. + * + * The shape comes from the same contract-parsed reader `GET /api/workspaces/[id]/files` + * responds with, so a seeded entry is identical to what the client hook would cache. + */ +async function seedWorkspaceFiles(queryClient: QueryClient, workspaceId: string): Promise { + try { + const files = await listWorkspaceFilesWithShares(workspaceId, 'active', { + limit: WORKSPACE_FILE_SEED_MAX + 1, + }) + if (files.length > WORKSPACE_FILE_SEED_MAX) return + queryClient.setQueryData(workspaceFilesKeys.list(workspaceId, 'active'), files) + } catch (error) { + /** Optimization only: the client fetch reaches the route instead. Logged so drift between + * this read and the contract's response schema doesn't degrade silently into a waterfall. */ + logger.warn('Workspace file list seed failed; client will fetch', { + error: getErrorMessage(error), + }) + } +} + /** * Prefetches the sidebar's workflow, chat, folder, workspace-permissions, * workspace, and viewer-profile reads for a workspace and stores them under the @@ -153,22 +200,7 @@ export async function prefetchWorkspaceSidebar( ] : []), prefetchResourceFolders(queryClient, workspaceId, 'workflow', userId), - /** - * The sidebar reads the workspace's files for its search modal, on EVERY workspace route — so this - * query is registered by sidebar chrome before any page renders. That ordering is why it has to be - * prefetched HERE and not only by the Files pages: `HydrationBoundary` hydrates a query the cache - * has already seen from a `useEffect`, which never runs during SSR, so a page-level boundary can - * only ever hand this entry to the client. Seeding it with the layout's own boundary — the first - * one to render — is what lets the server paint the Files browser and the open file's header - * populated instead of shipping a spinner and resolving it a beat later on the client. - * - * Same key + shape as {@link prefetchFilesBrowser}, so whichever runs is a no-op for the other. - */ - queryClient.prefetchQuery({ - queryKey: workspaceFilesKeys.list(workspaceId, 'active'), - queryFn: () => listWorkspaceFilesWithShares(workspaceId, 'active'), - staleTime: WORKSPACE_FILES_LIST_STALE_TIME, - }), + seedWorkspaceFiles(queryClient, workspaceId), queryClient.prefetchQuery({ queryKey: workspaceKeys.permissions(workspaceId), queryFn: () => diff --git a/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts b/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts index cc83cc5e98d..365eaedd936 100644 --- a/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts +++ b/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts @@ -161,6 +161,12 @@ interface ListWorkspaceFilesOptions { hydrateFolderPaths?: boolean /** Propagate storage errors when an incomplete list would be unsafe. */ throwOnError?: boolean + /** + * Row cap for callers that only need to know whether the workspace fits a budget. + * The result is a prefix of the full list, so a caller that reads it as "the + * workspace's files" must not set this. + */ + limit?: number } /** @@ -223,7 +229,9 @@ interface WorkspaceFileMetadataInsert { size: number } -function workspaceFileSize(file: typeof workspaceFiles.$inferSelect): number { +function workspaceFileSize( + file: Pick +): number { return file.sizeBytes ?? file.size } @@ -1007,7 +1015,7 @@ export async function fileExistsInWorkspace( } function mapWorkspaceFileRecord( - file: typeof workspaceFiles.$inferSelect, + file: WorkspaceFileListRow, workspaceId: string, folderPaths: Map ): WorkspaceFileRecord { @@ -1144,9 +1152,40 @@ function workspaceFileScopeCondition(workspaceId: string, scope: WorkspaceFileSc : and(...base, isNull(workspaceFiles.deletedAt)) } +/** + * The columns {@link mapWorkspaceFileRecord} actually reads. The list reads below are + * workspace-wide, so `select()` would pull five columns no reader projects — `context`, + * `chatId`, `messageId`, `displayName`, `secretProvenanceVersion` — for every row of the + * scan. Narrowing here is invisible on the wire (the route contract already strips them) + * and cuts what the database ships per row. + */ +const workspaceFileListColumns = { + id: workspaceFiles.id, + key: workspaceFiles.key, + userId: workspaceFiles.userId, + workspaceId: workspaceFiles.workspaceId, + folderId: workspaceFiles.folderId, + originalName: workspaceFiles.originalName, + contentType: workspaceFiles.contentType, + size: workspaceFiles.size, + sizeBytes: workspaceFiles.sizeBytes, + width: workspaceFiles.width, + height: workspaceFiles.height, + deletedAt: workspaceFiles.deletedAt, + uploadedAt: workspaceFiles.uploadedAt, + updatedAt: workspaceFiles.updatedAt, + contentUpdatedAt: workspaceFiles.contentUpdatedAt, +} as const + +/** A row carrying exactly the columns {@link mapWorkspaceFileRecord} needs; a full row satisfies it. */ +type WorkspaceFileListRow = Pick< + typeof workspaceFiles.$inferSelect, + keyof typeof workspaceFileListColumns +> + /** Resolves `folderPath` for a page of rows, reading the folder tree only if any row needs it. */ async function hydrateWorkspaceFilePaths( - files: (typeof workspaceFiles.$inferSelect)[], + files: WorkspaceFileListRow[], workspaceId: string, options?: { folders?: WorkspaceFileFolderRecord[]; hydrateFolderPaths?: boolean } ): Promise { @@ -1167,12 +1206,13 @@ export async function listWorkspaceFiles( options?: ListWorkspaceFilesOptions ): Promise { try { - const { scope = 'active' } = options ?? {} - const files = await db - .select() + const { scope = 'active', limit } = options ?? {} + const query = db + .select(workspaceFileListColumns) .from(workspaceFiles) .where(workspaceFileScopeCondition(workspaceId, scope)) .orderBy(workspaceFiles.uploadedAt) + const files = await (limit === undefined ? query : query.limit(limit)) return hydrateWorkspaceFilePaths(files, workspaceId, options) } catch (error) { @@ -1265,7 +1305,7 @@ export async function queryWorkspaceFiles( ] const rows = await db - .select() + .select(workspaceFileListColumns) .from(workspaceFiles) .where(and(...conditions)) .orderBy(...listOrderBy(keysetColumns(keys), sortOrder)) diff --git a/apps/sim/lib/uploads/contexts/workspace/workspace-file-query.test.ts b/apps/sim/lib/uploads/contexts/workspace/workspace-file-query.test.ts index f4c81ef8bc3..a0fa4782194 100644 --- a/apps/sim/lib/uploads/contexts/workspace/workspace-file-query.test.ts +++ b/apps/sim/lib/uploads/contexts/workspace/workspace-file-query.test.ts @@ -46,7 +46,10 @@ vi.mock('@/lib/uploads/contexts/workspace/workspace-file-folder-manager', () => resolveWorkspaceFileFolderTarget: vi.fn(async () => null), })) -import { queryWorkspaceFiles } from '@/lib/uploads/contexts/workspace/workspace-file-manager' +import { + listWorkspaceFiles, + queryWorkspaceFiles, +} from '@/lib/uploads/contexts/workspace/workspace-file-manager' const WS = 'workspace-1' @@ -205,3 +208,52 @@ describe('queryWorkspaceFiles', () => { ).rejects.toMatchObject({ code: 'validation' }) }) }) + +/** + * `listWorkspaceFiles` materializes a whole scope, so what it reads per row is + * multiplied by the size of the workspace. These assertions pin the two ways that + * stays bounded: the projection, and the optional row cap its one budgeted caller + * (the workspace layout's server seed) uses to detect an oversized workspace. + */ +describe('listWorkspaceFiles', () => { + beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + }) + + const lastProjection = () => + Object.keys((dbChainMockFns.select.mock.calls.at(-1)?.[0] ?? {}) as Record) + + it('projects only the columns the record mapper reads', async () => { + queueTableRows(schemaMock.workspaceFiles, [buildRow()]) + + await listWorkspaceFiles(WS) + + const projected = lastProjection() + expect(projected).toEqual( + expect.arrayContaining(['id', 'key', 'originalName', 'sizeBytes', 'contentUpdatedAt']) + ) + /** Columns no reader projects; `select()` would ship them for every row of the scan. */ + expect(projected).not.toContain('context') + expect(projected).not.toContain('chatId') + expect(projected).not.toContain('messageId') + expect(projected).not.toContain('displayName') + expect(projected).not.toContain('secretProvenanceVersion') + }) + + it('reads the whole scope when no cap is given', async () => { + queueTableRows(schemaMock.workspaceFiles, [buildRow()]) + + await listWorkspaceFiles(WS) + + expect(dbChainMockFns.limit).not.toHaveBeenCalled() + }) + + it('caps the rows read when the caller only needs to fit a budget', async () => { + queueTableRows(schemaMock.workspaceFiles, [buildRow()]) + + await listWorkspaceFiles(WS, { limit: 2 }) + + expect(dbChainMockFns.limit).toHaveBeenCalledWith(2) + }) +}) diff --git a/apps/sim/lib/workspace-files/queries.ts b/apps/sim/lib/workspace-files/queries.ts index 9a5f0d6b8ca..1a4925907b1 100644 --- a/apps/sim/lib/workspace-files/queries.ts +++ b/apps/sim/lib/workspace-files/queries.ts @@ -15,10 +15,18 @@ import { * to a client fetch rather than carrying a field that vanishes on the next refetch. * * Callers authorize the viewer against `workspaceId` first. + * + * `limit` caps the rows read for callers that only need to know whether the workspace fits + * a payload budget; the result is then a prefix of the list, not the list, so no caller may + * present a limited read as the workspace's files. */ -export async function listWorkspaceFilesWithShares(workspaceId: string, scope: WorkspaceFileScope) { +export async function listWorkspaceFilesWithShares( + workspaceId: string, + scope: WorkspaceFileScope, + options?: { limit?: number } +) { const [files, shares] = await Promise.all([ - listWorkspaceFiles(workspaceId, { scope }), + listWorkspaceFiles(workspaceId, { scope, limit: options?.limit }), getWorkspaceShares('file', workspaceId), ]) const withShares = files.map((file) => ({ ...file, share: shares.get(file.id) ?? null })) From 74a12e42202f787d44c0113023dc5b91105e3d0c Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 13 Aug 2026 01:32:12 -0700 Subject: [PATCH 09/14] improvement(api): resolve credentials and checkpoint reverts in-process --- .../app/api/auth/oauth/token/route.test.ts | 1 + apps/sim/app/api/auth/oauth/token/route.ts | 279 ++------------- .../copilot/checkpoints/revert/route.test.ts | 116 ++++--- .../api/copilot/checkpoints/revert/route.ts | 59 +++- .../sim/app/api/workflows/[id]/state/route.ts | 162 +-------- apps/sim/lib/auth/credential-access.test.ts | 39 ++- apps/sim/lib/auth/credential-access.ts | 26 +- apps/sim/lib/oauth/token-resolution.test.ts | 214 ++++++++++++ apps/sim/lib/oauth/token-resolution.ts | 320 ++++++++++++++++++ .../persistence/save-normalized-state.ts | 197 +++++++++++ apps/sim/tools/index.ts | 108 ++++-- 11 files changed, 1022 insertions(+), 499 deletions(-) create mode 100644 apps/sim/lib/oauth/token-resolution.test.ts create mode 100644 apps/sim/lib/oauth/token-resolution.ts create mode 100644 apps/sim/lib/workflows/persistence/save-normalized-state.ts diff --git a/apps/sim/app/api/auth/oauth/token/route.test.ts b/apps/sim/app/api/auth/oauth/token/route.test.ts index b1c07e96b68..821fe2bfa66 100644 --- a/apps/sim/app/api/auth/oauth/token/route.test.ts +++ b/apps/sim/app/api/auth/oauth/token/route.test.ts @@ -24,6 +24,7 @@ vi.mock('@/lib/oauth/credential-service', () => ({ vi.mock('@/lib/auth/credential-access', () => ({ authorizeCredentialUse: mockAuthorizeCredentialUse, + authorizeCredentialUseForAuth: mockAuthorizeCredentialUse, })) import { TokenServiceAccountValidationError } from '@/lib/credentials/token-service-accounts/errors' diff --git a/apps/sim/app/api/auth/oauth/token/route.ts b/apps/sim/app/api/auth/oauth/token/route.ts index cc66068135a..0bf50ac846d 100644 --- a/apps/sim/app/api/auth/oauth/token/route.ts +++ b/apps/sim/app/api/auth/oauth/token/route.ts @@ -11,17 +11,9 @@ import { authorizeCredentialUse } from '@/lib/auth/credential-access' import { AuthType, checkSessionOrInternalAuth } from '@/lib/auth/hybrid' import { generateRequestId } from '@/lib/core/utils/request' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' -import { TokenServiceAccountValidationError } from '@/lib/credentials/token-service-accounts/errors' -import { - getCredential, - getOAuthToken, - refreshTokenIfNeeded, - resolveOAuthAccountId, - resolveServiceAccountToken, -} from '@/lib/oauth/credential-service' -import { extractSalesforceInstanceUrl, isSalesforceOAuthProviderId } from '@/lib/oauth/salesforce' +import { getCredential, getOAuthToken } from '@/lib/oauth/credential-service' +import { completeOAuthCredentialToken, resolveCredentialToken } from '@/lib/oauth/token-resolution' import { captureServerEvent } from '@/lib/posthog/server' -import { extractZohoDeskBaseFromScope } from '@/tools/zoho_desk/host-allowlist' export const dynamic = 'force-dynamic' @@ -123,194 +115,25 @@ export const POST = withRouteHandler(async (request: NextRequest) => { } } - if (!credentialId) { - return NextResponse.json({ error: 'Credential ID is required' }, { status: 400 }) - } - - const resolved = await resolveOAuthAccountId(credentialId) - if (resolved?.credentialType === 'service_account' && resolved.credentialId) { - const authz = await authorizeCredentialUse(request, { - credentialId, - workflowId: workflowId ?? undefined, - requireWorkflowIdForInternal: false, - callerUserId, - }) - if (!authz.ok) { - return NextResponse.json({ error: authz.error || 'Unauthorized' }, { status: 403 }) - } - - const saActorId = authz.requesterUserId - const saWorkspaceId = resolved.workspaceId ?? authz.workspaceId ?? null - const emitServiceAccountAccess = () => { - if (!saActorId) return - recordAudit({ - workspaceId: saWorkspaceId, - actorId: saActorId, - action: AuditAction.CREDENTIAL_ACCESSED, - resourceType: AuditResourceType.CREDENTIAL, - resourceId: resolved.credentialId ?? credentialId, - description: `Accessed service account credential for provider ${resolved.providerId ?? 'unknown'}`, - metadata: { - provider: resolved.providerId, - credentialType: 'service_account', - }, - request, - }) - captureServerEvent( - saActorId, - 'credential_used', - { - credential_type: 'service_account', - provider_id: resolved.providerId ?? 'unknown', - ...(saWorkspaceId ? { workspace_id: saWorkspaceId } : {}), - }, - saWorkspaceId ? { groups: { workspace: saWorkspaceId } } : undefined - ) - } - - try { - const result = await resolveServiceAccountToken( - resolved.credentialId, - resolved.providerId, - scopes ?? [], - impersonateEmail - ) - emitServiceAccountAccess() - return NextResponse.json( - { - accessToken: result.accessToken, - cloudId: result.cloudId, - domain: result.domain, - instanceUrl: result.instanceUrl, - apiDomain: result.apiDomain, - authStyle: result.authStyle, - }, - { status: 200 } - ) - } catch (error) { - logger.error(`[${requestId}] Service account token error:`, error) - if (error instanceof TokenServiceAccountValidationError) { - // Classified provider outages are infra failures, not bad credentials. - if (error.code === 'provider_unavailable') { - return NextResponse.json( - { error: 'Credential provider is temporarily unavailable' }, - { status: 502 } - ) - } - // A stored host that no longer resolves is a configuration failure — - // surface the code so runtime consumers can say "check the host" - // instead of a generic auth error. - if (error.code === 'site_not_found') { - return NextResponse.json( - { - code: error.code, - error: 'Credential host not found — reconnect the credential with a valid host', - }, - { status: 400 } - ) - } - // A revoked/rotated-away or misconfigured stored secret — surface the - // code so runtime consumers can prompt to reconnect the credential - // rather than showing a generic auth failure. - if (error.code === 'invalid_credentials') { - return NextResponse.json( - { - code: error.code, - error: 'Credential rejected by the provider — reconnect the credential', - }, - { status: 401 } - ) - } - } - return NextResponse.json({ error: 'Failed to get service account token' }, { status: 401 }) - } - } - - const authz = await authorizeCredentialUse(request, { - credentialId, + const auth = await checkSessionOrInternalAuth(request, { requireWorkflowId: false }) + const result = await resolveCredentialToken(auth, { + requestId, + credentialId: credentialId ?? '', workflowId: workflowId ?? undefined, - requireWorkflowIdForInternal: false, + scopes, + impersonateEmail, callerUserId, + auditRequest: request, }) - if (!authz.ok || !authz.credentialOwnerUserId) { - return NextResponse.json({ error: authz.error || 'Unauthorized' }, { status: 403 }) - } - - const resolvedCredentialId = authz.resolvedCredentialId || credentialId - const credential = await getCredential( - requestId, - resolvedCredentialId, - authz.credentialOwnerUserId - ) - - if (!credential) { - return NextResponse.json({ error: 'Credential not found' }, { status: 404 }) - } - - const oauthActorId = authz.requesterUserId - const oauthWorkspaceId = authz.workspaceId ?? null - - try { - const { accessToken } = await refreshTokenIfNeeded( - requestId, - credential, - resolvedCredentialId - ) - - if (oauthActorId) { - recordAudit({ - workspaceId: oauthWorkspaceId, - actorId: oauthActorId, - action: AuditAction.CREDENTIAL_ACCESSED, - resourceType: AuditResourceType.CREDENTIAL, - resourceId: resolvedCredentialId, - description: `Accessed OAuth credential for provider ${credential.providerId}`, - metadata: { - provider: credential.providerId, - credentialType: 'oauth', - }, - request, - }) - captureServerEvent( - oauthActorId, - 'credential_used', - { - credential_type: 'oauth', - provider_id: credential.providerId, - ...(oauthWorkspaceId ? { workspace_id: oauthWorkspaceId } : {}), - }, - oauthWorkspaceId ? { groups: { workspace: oauthWorkspaceId } } : undefined - ) - } - - const instanceUrl = isSalesforceOAuthProviderId(credential.providerId) - ? extractSalesforceInstanceUrl(credential.scope) - : undefined - - // Zoho Desk persists its data-center-specific REST base URL in the scope - // string (derived from the token response api_domain) so callers never - // assume a host. Surface it as apiDomain for tool param injection. - let apiDomain: string | undefined - if (credential.providerId === 'zoho-desk' && credential.scope) { - // Use the shared extractor, not a local regex: it also enforces https + - // the Zoho apex allowlist. This value is injected into EVERY tool call, - // so an unvalidated host here would receive the OAuth token. - apiDomain = extractZohoDeskBaseFromScope(credential.scope) - } + if (!result.ok) { return NextResponse.json( - { - accessToken, - idToken: credential.idToken || undefined, - ...(instanceUrl && { instanceUrl }), - ...(apiDomain && { apiDomain }), - }, - { status: 200 } + { ...(result.code ? { code: result.code } : {}), error: result.error }, + { status: result.status } ) - } catch (error) { - logger.error(`[${requestId}] Failed to refresh access token:`, error) - return NextResponse.json({ error: 'Failed to refresh access token' }, { status: 401 }) } + + return NextResponse.json(result.token, { status: 200 }) } catch (error) { logger.error(`[${requestId}] Error getting access token`, error) return NextResponse.json({ error: 'Internal server error' }, { status: 500 }) @@ -366,70 +189,20 @@ export const GET = withRouteHandler(async (request: NextRequest) => { return NextResponse.json({ error: 'No access token available' }, { status: 400 }) } - const actorId = authz.requesterUserId - const workspaceId = authz.workspaceId ?? null - - try { - const { accessToken } = await refreshTokenIfNeeded( - requestId, - credential, - resolvedCredentialId - ) - - if (actorId) { - recordAudit({ - workspaceId, - actorId, - action: AuditAction.CREDENTIAL_ACCESSED, - resourceType: AuditResourceType.CREDENTIAL, - resourceId: resolvedCredentialId, - description: `Accessed OAuth credential for provider ${credential.providerId}`, - metadata: { - provider: credential.providerId, - credentialType: 'oauth', - }, - request, - }) - captureServerEvent( - actorId, - 'credential_used', - { - credential_type: 'oauth', - provider_id: credential.providerId, - ...(workspaceId ? { workspace_id: workspaceId } : {}), - }, - workspaceId ? { groups: { workspace: workspaceId } } : undefined - ) - } - - const instanceUrl = isSalesforceOAuthProviderId(credential.providerId) - ? extractSalesforceInstanceUrl(credential.scope) - : undefined - - // Zoho Desk persists its data-center-specific REST base URL in the scope - // string (derived from the token response api_domain) so callers never - // assume a host. Surface it as apiDomain for tool param injection. - let apiDomain: string | undefined - if (credential.providerId === 'zoho-desk' && credential.scope) { - // Use the shared extractor, not a local regex: it also enforces https + - // the Zoho apex allowlist. This value is injected into EVERY tool call, - // so an unvalidated host here would receive the OAuth token. - apiDomain = extractZohoDeskBaseFromScope(credential.scope) - } + const result = await completeOAuthCredentialToken({ + requestId, + credential, + resolvedCredentialId, + actorId: authz.requesterUserId, + workspaceId: authz.workspaceId ?? null, + auditRequest: request, + }) - return NextResponse.json( - { - accessToken, - idToken: credential.idToken || undefined, - ...(instanceUrl && { instanceUrl }), - ...(apiDomain && { apiDomain }), - }, - { status: 200 } - ) - } catch (error) { - logger.error(`[${requestId}] Failed to refresh access token:`, error) - return NextResponse.json({ error: 'Failed to refresh access token' }, { status: 401 }) + if (!result.ok) { + return NextResponse.json({ error: result.error }, { status: result.status }) } + + return NextResponse.json(result.token, { status: 200 }) } catch (error) { logger.error(`[${requestId}] Error fetching access token`, error) return NextResponse.json({ error: 'Internal server error' }, { status: 500 }) diff --git a/apps/sim/app/api/copilot/checkpoints/revert/route.test.ts b/apps/sim/app/api/copilot/checkpoints/revert/route.test.ts index 785718e9ba2..67eaf1a94c4 100644 --- a/apps/sim/app/api/copilot/checkpoints/revert/route.test.ts +++ b/apps/sim/app/api/copilot/checkpoints/revert/route.test.ts @@ -17,12 +17,23 @@ import { import { NextRequest } from 'next/server' import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest' -const { mockGetAccessibleCopilotChat } = vi.hoisted(() => ({ +const { + mockGetAccessibleCopilotChat, + mockParseWorkflowStateForPersistence, + mockSaveWorkflowNormalizedState, +} = vi.hoisted(() => ({ mockGetAccessibleCopilotChat: vi.fn(), + mockParseWorkflowStateForPersistence: vi.fn(), + mockSaveWorkflowNormalizedState: vi.fn(), })) vi.mock('@/lib/workflows/utils', () => workflowsUtilsMock) +vi.mock('@/lib/workflows/persistence/save-normalized-state', () => ({ + parseWorkflowStateForPersistence: mockParseWorkflowStateForPersistence, + saveWorkflowNormalizedState: mockSaveWorkflowNormalizedState, +})) + vi.mock('@/lib/copilot/chat/lifecycle', () => ({ getAccessibleCopilotChat: mockGetAccessibleCopilotChat, getAccessibleCopilotChatAuth: mockGetAccessibleCopilotChat, @@ -45,6 +56,12 @@ describe('Copilot Checkpoints Revert API Route', () => { mockGetAccessibleCopilotChat.mockResolvedValue({ id: 'chat-123', userId: 'user-123' }) + mockParseWorkflowStateForPersistence.mockImplementation((value: unknown) => ({ + success: true, + data: value, + })) + mockSaveWorkflowNormalizedState.mockResolvedValue({ success: true, warnings: [] }) + global.fetch = vi.fn() vi.spyOn(Date, 'now').mockReturnValue(1640995200000) @@ -297,24 +314,19 @@ describe('Copilot Checkpoints Revert API Route', () => { }, }) - // Verify fetch was called with correct parameters - expect(global.fetch).toHaveBeenCalledWith( - 'http://localhost:3000/api/workflows/c3d4e5f6-a7b8-4c09-a1e2-f3a4b5c6d7e8/state', - { - method: 'PUT', - headers: { - 'Content-Type': 'application/json', - Cookie: 'session=test-session', - }, - body: JSON.stringify({ + expect(mockSaveWorkflowNormalizedState).toHaveBeenCalledWith( + expect.objectContaining({ + workflowId: 'c3d4e5f6-a7b8-4c09-a1e2-f3a4b5c6d7e8', + userId: 'user-123', + state: { blocks: { block1: { type: 'start' } }, edges: [{ from: 'block1', to: 'block2' }], loops: {}, parallels: {}, isDeployed: true, lastSaved: 1640995200000, - }), - } + }, + }) ) }) @@ -452,7 +464,7 @@ describe('Copilot Checkpoints Revert API Route', () => { }) }) - it('should return 500 when state API call fails', async () => { + it('should return 500 when the state write fails', async () => { setAuthenticated() const mockCheckpoint = { @@ -470,9 +482,10 @@ describe('Copilot Checkpoints Revert API Route', () => { queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) queueTableRows(schemaMock.workflow, [mockWorkflow]) - ;(global.fetch as any).mockResolvedValue({ - ok: false, - text: () => Promise.resolve('State validation failed'), + mockSaveWorkflowNormalizedState.mockResolvedValueOnce({ + success: false, + status: 500, + error: 'Failed to save workflow state', }) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', { @@ -488,6 +501,36 @@ describe('Copilot Checkpoints Revert API Route', () => { expect(responseData.error).toBe('Failed to revert workflow to checkpoint') }) + it('should return 500 when the checkpoint state fails validation', async () => { + setAuthenticated() + + const mockCheckpoint = { + id: 'checkpoint-123', + workflowId: 'a7b8c9d0-e1f2-4a34-b5c6-d7e8f9a0b1c2', + userId: 'user-123', + workflowState: { blocks: {}, edges: [] }, + } + + queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) + queueTableRows(schemaMock.workflow, [{ id: mockCheckpoint.workflowId, userId: 'user-123' }]) + + mockParseWorkflowStateForPersistence.mockReturnValueOnce({ + success: false, + error: { issues: [{ message: 'blocks: invalid' }] }, + }) + + const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ checkpointId: 'checkpoint-123' }), + }) + + const response = await POST(req) + + expect(response.status).toBe(500) + expect(mockSaveWorkflowNormalizedState).not.toHaveBeenCalled() + }) + it('should handle database errors during checkpoint lookup', async () => { setAuthenticated() @@ -536,7 +579,7 @@ describe('Copilot Checkpoints Revert API Route', () => { expect(responseData.error).toBe('Failed to revert to checkpoint') }) - it('should handle fetch network errors', async () => { + it('should handle unexpected errors from the state write', async () => { setAuthenticated() const mockCheckpoint = { @@ -554,7 +597,7 @@ describe('Copilot Checkpoints Revert API Route', () => { queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) queueTableRows(schemaMock.workflow, [mockWorkflow]) - ;(global.fetch as any).mockRejectedValue(new Error('Network error')) + mockSaveWorkflowNormalizedState.mockRejectedValueOnce(new Error('Network error')) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', { method: 'POST', @@ -587,7 +630,7 @@ describe('Copilot Checkpoints Revert API Route', () => { expect(responseData.error).toBe('Failed to revert to checkpoint') }) - it('should forward cookies to state API call', async () => { + it('should apply the state in-process instead of re-authenticating over HTTP', async () => { setAuthenticated() const mockCheckpoint = { @@ -623,17 +666,15 @@ describe('Copilot Checkpoints Revert API Route', () => { await POST(req) - expect(global.fetch).toHaveBeenCalledWith( - 'http://localhost:3000/api/workflows/d0e1f2a3-b4c5-4d67-a8f9-a0b1c2d3e4f5/state', - { - method: 'PUT', - headers: { - 'Content-Type': 'application/json', - Cookie: 'session=test-session; auth=token123', - }, - body: expect.any(String), - } + expect(mockSaveWorkflowNormalizedState).toHaveBeenCalledWith( + expect.objectContaining({ + workflowId: 'd0e1f2a3-b4c5-4d67-a8f9-a0b1c2d3e4f5', + userId: 'user-123', + }) ) + for (const call of (global.fetch as any).mock.calls) { + expect(String(call[0])).not.toContain('/state') + } }) it('should handle missing cookies gracefully', async () => { @@ -673,16 +714,11 @@ describe('Copilot Checkpoints Revert API Route', () => { const response = await POST(req) expect(response.status).toBe(200) - expect(global.fetch).toHaveBeenCalledWith( - 'http://localhost:3000/api/workflows/e1f2a3b4-c5d6-4e78-a9a0-b1c2d3e4f5a6/state', - { - method: 'PUT', - headers: { - 'Content-Type': 'application/json', - Cookie: '', // Empty string when no cookies - }, - body: expect.any(String), - } + expect(mockSaveWorkflowNormalizedState).toHaveBeenCalledWith( + expect.objectContaining({ + workflowId: 'e1f2a3b4-c5d6-4e78-a9a0-b1c2d3e4f5a6', + userId: 'user-123', + }) ) }) diff --git a/apps/sim/app/api/copilot/checkpoints/revert/route.ts b/apps/sim/app/api/copilot/checkpoints/revert/route.ts index f784dc48d84..6b49fb8823e 100644 --- a/apps/sim/app/api/copilot/checkpoints/revert/route.ts +++ b/apps/sim/app/api/copilot/checkpoints/revert/route.ts @@ -1,7 +1,10 @@ import { db } from '@sim/db' import { workflowCheckpoints, workflow as workflowTable } from '@sim/db/schema' import { createLogger } from '@sim/logger' -import { authorizeWorkflowByWorkspacePermission } from '@sim/platform-authz/workflow' +import { + authorizeWorkflowByWorkspacePermission, + WorkflowLockedError, +} from '@sim/platform-authz/workflow' import { and, eq } from 'drizzle-orm' import { type NextRequest, NextResponse } from 'next/server' import { revertCopilotCheckpointContract } from '@/lib/api/contracts/copilot' @@ -15,8 +18,11 @@ import { createRequestTracker, createUnauthorizedResponse, } from '@/lib/copilot/request/http' -import { getInternalApiBaseUrl } from '@/lib/core/utils/urls' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' +import { + parseWorkflowStateForPersistence, + saveWorkflowNormalizedState, +} from '@/lib/workflows/persistence/save-normalized-state' import { isUuidV4 } from '@/executor/constants' const logger = createLogger('CheckpointRevertAPI') @@ -121,28 +127,49 @@ export const POST = withRouteHandler(async (request: NextRequest) => { return NextResponse.json({ error: 'Invalid workflow ID format' }, { status: 400 }) } - const stateResponse = await fetch( - `${getInternalApiBaseUrl()}/api/workflows/${checkpoint.workflowId}/state`, - { - method: 'PUT', - headers: { - 'Content-Type': 'application/json', - Cookie: request.headers.get('Cookie') || '', - }, - body: JSON.stringify(cleanedState), + /** + * The checkpoint blob is persisted JSONB, so it goes through the same + * schema the PUT state contract applies before it is written back — the + * validation the removed HTTP hop used to provide. + */ + const parsedState = parseWorkflowStateForPersistence(cleanedState) + if (!parsedState.success) { + logger.error( + `[${tracker.requestId}] Checkpoint state failed validation`, + parsedState.error.issues + ) + return NextResponse.json( + { error: 'Failed to revert workflow to checkpoint' }, + { status: 500 } + ) + } + + /** + * A locked workflow used to surface here as a non-OK PUT response, so it + * still resolves to the same revert failure rather than the generic + * outer-catch message. Every other throw keeps propagating, matching the + * old transport-error path. + */ + const saveResult = await saveWorkflowNormalizedState({ + requestId: tracker.requestId, + workflowId: checkpoint.workflowId, + userId, + state: parsedState.data, + }).catch((error) => { + if (error instanceof WorkflowLockedError) { + return { success: false as const, status: error.status, error: error.message } } - ) + throw error + }) - if (!stateResponse.ok) { - const errorData = await stateResponse.text() - logger.error(`[${tracker.requestId}] Failed to apply checkpoint state: ${errorData}`) + if (!saveResult.success) { + logger.error(`[${tracker.requestId}] Failed to apply checkpoint state: ${saveResult.error}`) return NextResponse.json( { error: 'Failed to revert workflow to checkpoint' }, { status: 500 } ) } - const result = await stateResponse.json() logger.info( `[${tracker.requestId}] Successfully reverted workflow ${checkpoint.workflowId} to checkpoint ${checkpointId}` ) diff --git a/apps/sim/app/api/workflows/[id]/state/route.ts b/apps/sim/app/api/workflows/[id]/state/route.ts index 209cfd226e8..97a32c40b86 100644 --- a/apps/sim/app/api/workflows/[id]/state/route.ts +++ b/apps/sim/app/api/workflows/[id]/state/route.ts @@ -2,7 +2,6 @@ import { db } from '@sim/db' import { workflow } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { - assertWorkflowMutable, authorizeWorkflowByWorkspacePermission, WorkflowLockedError, } from '@sim/platform-authz/workflow' @@ -12,17 +11,10 @@ import { type NextRequest, NextResponse } from 'next/server' import { putWorkflowNormalizedStateContract } from '@/lib/api/contracts/workflows' import { parseRequest } from '@/lib/api/server' import { checkSessionOrInternalAuth } from '@/lib/auth/hybrid' -import { env } from '@/lib/core/config/env' import { generateRequestId } from '@/lib/core/utils/request' -import { getSocketServerUrl } from '@/lib/core/utils/urls' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' -import { extractAndPersistCustomTools } from '@/lib/workflows/persistence/custom-tools-persistence' -import { prepareWorkflowStateForPersistence } from '@/lib/workflows/persistence/prepare-state' -import { - loadWorkflowFromNormalizedTables, - saveWorkflowToNormalizedTables, -} from '@/lib/workflows/persistence/utils' -import type { BlockState, WorkflowState } from '@/stores/workflows/workflow/types' +import { saveWorkflowNormalizedState } from '@/lib/workflows/persistence/save-normalized-state' +import { loadWorkflowFromNormalizedTables } from '@/lib/workflows/persistence/utils' const logger = createLogger('WorkflowStateAPI') @@ -117,159 +109,33 @@ export const PUT = withRouteHandler( const parsed = await parseRequest(putWorkflowNormalizedStateContract, request, context) if (!parsed.success) return parsed.response - const state = parsed.data.body - - const authorization = await authorizeWorkflowByWorkspacePermission({ - workflowId, - userId, - action: 'write', - }) - const workflowData = authorization.workflow - - if (!workflowData) { - logger.warn(`[${requestId}] Workflow ${workflowId} not found for state update`) - return NextResponse.json({ error: 'Workflow not found' }, { status: 404 }) - } - - const canUpdate = authorization.allowed - - if (!canUpdate) { - logger.warn( - `[${requestId}] User ${userId} denied permission to update workflow state ${workflowId}` - ) - return NextResponse.json( - { error: authorization.message || 'Access denied' }, - { status: authorization.status || 403 } - ) - } - - await assertWorkflowMutable(workflowId) // Note: prior versions cross-checked that each variable's `workflowId` // equalled the path param. The write contract does not carry `workflowId` // per variable (the path param is the source of truth), so the check // is unreachable and was removed. - const { state: preparedState, warnings: preparationWarnings } = - prepareWorkflowStateForPersistence({ - blocks: state.blocks as Record, - edges: state.edges as WorkflowState['edges'], - }) - - const workflowState = { - ...preparedState, - lastSaved: state.lastSaved || Date.now(), - isDeployed: state.isDeployed || false, - deployedAt: state.deployedAt, - } - - const saveResult = await db.transaction(async (tx) => { - await tx - .select({ id: workflow.id }) - .from(workflow) - .where(eq(workflow.id, workflowId)) - .limit(1) - .for('update') - - const result = await saveWorkflowToNormalizedTables( - workflowId, - workflowState as WorkflowState, - tx - ) - - if (!result.success) return result - - // Update workflow's lastSynced timestamp and variables if provided - const updateData: { - lastSynced: Date - updatedAt: Date - variables?: typeof state.variables - } = { - lastSynced: new Date(), - updatedAt: new Date(), - } - - // If variables are provided in the state, update them in the workflow record - if (state.variables !== undefined) { - updateData.variables = state.variables - } - - await tx.update(workflow).set(updateData).where(eq(workflow.id, workflowId)) - - return result + const result = await saveWorkflowNormalizedState({ + requestId, + workflowId, + userId, + state: parsed.data.body, }) - if (!saveResult.success) { - logger.error( - `[${requestId}] Failed to save workflow ${workflowId} state:`, - saveResult.error - ) + if (!result.success) { return NextResponse.json( - { error: 'Failed to save workflow state', details: saveResult.error }, - { status: 500 } + { + error: result.error, + ...(result.details !== undefined ? { details: result.details } : {}), + }, + { status: result.status } ) } - // Extract and persist custom tools to database - try { - const workspaceId = workflowData.workspaceId - if (workspaceId) { - const { saved, errors } = await extractAndPersistCustomTools( - workflowState, - workspaceId, - userId - ) - - if (saved > 0) { - logger.info(`[${requestId}] Persisted ${saved} custom tool(s) to database`, { - workflowId, - }) - } - - if (errors.length > 0) { - logger.warn(`[${requestId}] Some custom tools failed to persist`, { - errors, - workflowId, - }) - } - } else { - logger.warn( - `[${requestId}] Workflow has no workspaceId, skipping custom tools persistence`, - { - workflowId, - } - ) - } - } catch (error) { - logger.error(`[${requestId}] Failed to persist custom tools`, { error, workflowId }) - } - const elapsed = Date.now() - startTime logger.info(`[${requestId}] Successfully saved workflow ${workflowId} state in ${elapsed}ms`) - try { - const notifyResponse = await fetch(`${getSocketServerUrl()}/api/workflow-updated`, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - 'x-api-key': env.INTERNAL_API_SECRET, - }, - body: JSON.stringify({ workflowId }), - }) - - if (!notifyResponse.ok) { - logger.warn( - `[${requestId}] Failed to notify Socket.IO server about workflow ${workflowId} update` - ) - } - } catch (notificationError) { - logger.warn( - `[${requestId}] Error notifying Socket.IO server about workflow ${workflowId} update`, - notificationError - ) - } - - return NextResponse.json({ success: true, warnings: preparationWarnings }, { status: 200 }) + return NextResponse.json({ success: true, warnings: result.warnings }, { status: 200 }) } catch (error: any) { if (error instanceof WorkflowLockedError) { return NextResponse.json({ error: error.message }, { status: error.status }) diff --git a/apps/sim/lib/auth/credential-access.test.ts b/apps/sim/lib/auth/credential-access.test.ts index 4a30c270b64..3f0f1022edb 100644 --- a/apps/sim/lib/auth/credential-access.test.ts +++ b/apps/sim/lib/auth/credential-access.test.ts @@ -23,7 +23,7 @@ vi.mock('@/lib/workspaces/permissions/utils', () => ({ resolveWorkspaceAccess: mockResolveWorkspaceAccess, })) -import { authorizeCredentialUse } from '@/lib/auth/credential-access' +import { authorizeCredentialUse, authorizeCredentialUseForAuth } from '@/lib/auth/credential-access' afterAll(resetDbChainMock) @@ -235,4 +235,41 @@ describe('authorizeCredentialUse', () => { expect(result.error).toBe('Credential not found') }) }) + + /** + * The in-process tool executor synthesizes the AuthResult an internal JWT + * would have produced instead of minting one and POSTing to ourselves, so the + * subject-less case must still fail closed here. + */ + describe('authorizeCredentialUseForAuth', () => { + it('fails closed when the authenticated caller carries no user id', async () => { + const result = await authorizeCredentialUseForAuth( + { success: true, authType: 'internal_jwt' }, + { credentialId: ACCOUNT_ID } + ) + + expect(result.ok).toBe(false) + expect(result.error).toBe('Authentication required') + }) + + it('fails closed when authentication did not succeed', async () => { + const result = await authorizeCredentialUseForAuth( + { success: false, error: 'Unauthorized' }, + { credentialId: ACCOUNT_ID } + ) + + expect(result.ok).toBe(false) + expect(result.error).toBe('Unauthorized') + }) + + it('rejects an asserted caller that does not match the internal token subject', async () => { + const result = await authorizeCredentialUseForAuth( + { success: true, userId: OWNER, authType: 'internal_jwt' }, + { credentialId: ACCOUNT_ID, callerUserId: 'someone-else' } + ) + + expect(result.ok).toBe(false) + expect(result.error).toBe('Caller user does not match internal token subject') + }) + }) }) diff --git a/apps/sim/lib/auth/credential-access.ts b/apps/sim/lib/auth/credential-access.ts index 125511b67a8..75ef5a952a2 100644 --- a/apps/sim/lib/auth/credential-access.ts +++ b/apps/sim/lib/auth/credential-access.ts @@ -2,7 +2,7 @@ import { db } from '@sim/db' import { account, credential, workflow as workflowTable } from '@sim/db/schema' import { and, asc, eq } from 'drizzle-orm' import type { NextRequest } from 'next/server' -import { AuthType, checkSessionOrInternalAuth } from '@/lib/auth/hybrid' +import { type AuthResult, AuthType, checkSessionOrInternalAuth } from '@/lib/auth/hybrid' import { type CredentialActorContext, canUseCredential, @@ -59,11 +59,29 @@ export async function authorizeCredentialUse( callerUserId?: string } ): Promise { - const { credentialId, workflowId, requireWorkflowIdForInternal = true, callerUserId } = params - const auth = await checkSessionOrInternalAuth(request, { - requireWorkflowId: requireWorkflowIdForInternal, + requireWorkflowId: params.requireWorkflowIdForInternal ?? true, }) + return authorizeCredentialUseForAuth(auth, params) +} + +/** + * Credential authorization for a caller whose authentication has already been + * resolved. {@link authorizeCredentialUse} is the HTTP-request wrapper over + * this; in-process callers (the tool executor) construct the same + * {@link AuthResult} directly instead of re-authenticating over HTTP, so both + * paths run one identical authorization rule. + */ +export async function authorizeCredentialUseForAuth( + auth: AuthResult, + params: { + credentialId: string + workflowId?: string + callerUserId?: string + } +): Promise { + const { credentialId, workflowId, callerUserId } = params + if (!auth.success || !auth.userId) { return { ok: false, error: auth.error || 'Authentication required' } } diff --git a/apps/sim/lib/oauth/token-resolution.test.ts b/apps/sim/lib/oauth/token-resolution.test.ts new file mode 100644 index 00000000000..27d239091ae --- /dev/null +++ b/apps/sim/lib/oauth/token-resolution.test.ts @@ -0,0 +1,214 @@ +/** + * @vitest-environment node + */ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { + mockAuthorizeCredentialUseForAuth, + mockGetCredential, + mockRecordAudit, + mockRefreshTokenIfNeeded, + mockResolveOAuthAccountId, + mockResolveServiceAccountToken, +} = vi.hoisted(() => ({ + mockAuthorizeCredentialUseForAuth: vi.fn(), + mockGetCredential: vi.fn(), + mockRecordAudit: vi.fn(), + mockRefreshTokenIfNeeded: vi.fn(), + mockResolveOAuthAccountId: vi.fn(), + mockResolveServiceAccountToken: vi.fn(), +})) + +vi.mock('@sim/audit', () => ({ + AuditAction: { CREDENTIAL_ACCESSED: 'credential.accessed' }, + AuditResourceType: { CREDENTIAL: 'credential' }, + recordAudit: mockRecordAudit, +})) + +vi.mock('@/lib/auth/credential-access', () => ({ + authorizeCredentialUseForAuth: mockAuthorizeCredentialUseForAuth, +})) + +vi.mock('@/lib/oauth/credential-service', () => ({ + getCredential: mockGetCredential, + refreshTokenIfNeeded: mockRefreshTokenIfNeeded, + resolveOAuthAccountId: mockResolveOAuthAccountId, + resolveServiceAccountToken: mockResolveServiceAccountToken, +})) + +vi.mock('@/lib/posthog/server', () => ({ + captureServerEvent: vi.fn(), +})) + +import { TokenServiceAccountValidationError } from '@/lib/credentials/token-service-accounts/errors' +import { resolveCredentialToken } from '@/lib/oauth/token-resolution' + +const INTERNAL_AUTH = { success: true, userId: 'user-1', authType: 'internal_jwt' } as const + +describe('resolveCredentialToken', () => { + beforeEach(() => { + vi.clearAllMocks() + mockResolveOAuthAccountId.mockResolvedValue(null) + }) + + it('fails closed when the credential is not authorized', async () => { + mockAuthorizeCredentialUseForAuth.mockResolvedValue({ + ok: false, + error: 'You do not have access to this credential.', + }) + + const result = await resolveCredentialToken(INTERNAL_AUTH, { + requestId: 'req-1', + credentialId: 'cred-1', + }) + + expect(result).toEqual({ + ok: false, + status: 403, + error: 'You do not have access to this credential.', + }) + expect(mockGetCredential).not.toHaveBeenCalled() + expect(mockRefreshTokenIfNeeded).not.toHaveBeenCalled() + expect(mockRecordAudit).not.toHaveBeenCalled() + }) + + it('fails closed when the caller carries no user id', async () => { + mockAuthorizeCredentialUseForAuth.mockResolvedValue({ + ok: false, + error: 'Authentication required', + }) + + const result = await resolveCredentialToken( + { success: true, authType: 'internal_jwt' }, + { requestId: 'req-1', credentialId: 'cred-1' } + ) + + expect(result).toEqual({ ok: false, status: 403, error: 'Authentication required' }) + }) + + it('refreshes the token, records the access trail, and returns the payload', async () => { + mockAuthorizeCredentialUseForAuth.mockResolvedValue({ + ok: true, + requesterUserId: 'user-1', + credentialOwnerUserId: 'owner-1', + workspaceId: 'ws-1', + resolvedCredentialId: 'account-1', + }) + mockGetCredential.mockResolvedValue({ + providerId: 'google', + idToken: 'id-token', + scope: 'https://www.googleapis.com/auth/gmail.send', + }) + mockRefreshTokenIfNeeded.mockResolvedValue({ accessToken: 'fresh', refreshed: true }) + + const result = await resolveCredentialToken(INTERNAL_AUTH, { + requestId: 'req-1', + credentialId: 'cred-1', + workflowId: 'wf-1', + }) + + expect(result).toEqual({ ok: true, token: { accessToken: 'fresh', idToken: 'id-token' } }) + expect(mockGetCredential).toHaveBeenCalledWith('req-1', 'account-1', 'owner-1') + expect(mockRefreshTokenIfNeeded).toHaveBeenCalled() + expect(mockRecordAudit).toHaveBeenCalledWith( + expect.objectContaining({ + actorId: 'user-1', + workspaceId: 'ws-1', + resourceId: 'account-1', + action: 'credential.accessed', + }) + ) + }) + + it('returns 404 when the authorized credential is missing', async () => { + mockAuthorizeCredentialUseForAuth.mockResolvedValue({ + ok: true, + requesterUserId: 'user-1', + credentialOwnerUserId: 'owner-1', + }) + mockGetCredential.mockResolvedValue(undefined) + + const result = await resolveCredentialToken(INTERNAL_AUTH, { + requestId: 'req-1', + credentialId: 'cred-1', + }) + + expect(result).toEqual({ ok: false, status: 404, error: 'Credential not found' }) + }) + + it('reports a failed refresh as 401 without recording access', async () => { + mockAuthorizeCredentialUseForAuth.mockResolvedValue({ + ok: true, + requesterUserId: 'user-1', + credentialOwnerUserId: 'owner-1', + }) + mockGetCredential.mockResolvedValue({ providerId: 'google' }) + mockRefreshTokenIfNeeded.mockRejectedValue(new Error('refresh token revoked')) + + const result = await resolveCredentialToken(INTERNAL_AUTH, { + requestId: 'req-1', + credentialId: 'cred-1', + }) + + expect(result).toEqual({ ok: false, status: 401, error: 'Failed to refresh access token' }) + expect(mockRecordAudit).not.toHaveBeenCalled() + }) + + it('authorizes service-account credentials before minting a token', async () => { + mockResolveOAuthAccountId.mockResolvedValue({ + credentialType: 'service_account', + credentialId: 'sa-1', + providerId: 'google', + workspaceId: 'ws-1', + accountId: '', + usedCredentialTable: true, + }) + mockAuthorizeCredentialUseForAuth.mockResolvedValue({ ok: false, error: 'Unauthorized' }) + + const result = await resolveCredentialToken(INTERNAL_AUTH, { + requestId: 'req-1', + credentialId: 'cred-1', + }) + + expect(result).toEqual({ ok: false, status: 403, error: 'Unauthorized' }) + expect(mockResolveServiceAccountToken).not.toHaveBeenCalled() + }) + + it('surfaces the classified service-account failure code', async () => { + mockResolveOAuthAccountId.mockResolvedValue({ + credentialType: 'service_account', + credentialId: 'sa-1', + providerId: 'atlassian', + workspaceId: 'ws-1', + accountId: '', + usedCredentialTable: true, + }) + mockAuthorizeCredentialUseForAuth.mockResolvedValue({ ok: true, requesterUserId: 'user-1' }) + mockResolveServiceAccountToken.mockRejectedValue( + new TokenServiceAccountValidationError('invalid_credentials', 401) + ) + + const result = await resolveCredentialToken(INTERNAL_AUTH, { + requestId: 'req-1', + credentialId: 'cred-1', + }) + + expect(result).toEqual({ + ok: false, + status: 401, + code: 'invalid_credentials', + error: 'Credential rejected by the provider — reconnect the credential', + }) + }) + + it('rejects a malformed impersonation subject before touching the credential', async () => { + const result = await resolveCredentialToken(INTERNAL_AUTH, { + requestId: 'req-1', + credentialId: 'cred-1', + impersonateEmail: 'not-an-email', + }) + + expect(result.ok).toBe(false) + expect(mockAuthorizeCredentialUseForAuth).not.toHaveBeenCalled() + }) +}) diff --git a/apps/sim/lib/oauth/token-resolution.ts b/apps/sim/lib/oauth/token-resolution.ts new file mode 100644 index 00000000000..91f658bea1a --- /dev/null +++ b/apps/sim/lib/oauth/token-resolution.ts @@ -0,0 +1,320 @@ +import { AuditAction, AuditResourceType, recordAudit } from '@sim/audit' +import { createLogger } from '@sim/logger' +import { z } from 'zod' +import { authorizeCredentialUseForAuth } from '@/lib/auth/credential-access' +import type { AuthResult } from '@/lib/auth/hybrid' +import { TokenServiceAccountValidationError } from '@/lib/credentials/token-service-accounts/errors' +import { + getCredential, + refreshTokenIfNeeded, + resolveOAuthAccountId, + resolveServiceAccountToken, +} from '@/lib/oauth/credential-service' +import { extractSalesforceInstanceUrl, isSalesforceOAuthProviderId } from '@/lib/oauth/salesforce' +import { captureServerEvent } from '@/lib/posthog/server' +import { extractZohoDeskBaseFromScope } from '@/tools/zoho_desk/host-allowlist' + +const logger = createLogger('OAuthTokenResolution') + +/** + * Minimal duck type of the inbound HTTP request, used only so audit rows can + * record the caller's IP and user agent. In-process callers have no inbound + * request and omit it; the row is then written without forensic headers. + */ +export interface CredentialAuditRequest { + headers: { get(name: string): string | null } +} + +/** Token material a resolved credential yields, as returned to every surface. */ +export interface CredentialTokenPayload { + accessToken: string + idToken?: string + instanceUrl?: string + apiDomain?: string + cloudId?: string + domain?: string + authStyle?: 'x-api-token' +} + +export interface ResolveCredentialTokenInput { + /** Correlation id used by the credential service's own logging. */ + requestId: string + credentialId: string + workflowId?: string + /** Canonical provider scopes, used only by service-account token minting. */ + scopes?: string[] + /** Google domain-wide-delegation subject for service-account credentials. */ + impersonateEmail?: string + /** + * Asserted acting user. When the caller authenticated with an internal JWT it + * must equal the token subject, so a forged assertion cannot widen access. + */ + callerUserId?: string + auditRequest?: CredentialAuditRequest +} + +export type ResolveCredentialTokenResult = + | { ok: true; token: CredentialTokenPayload } + | { ok: false; status: number; error: string; code?: string } + +const impersonateEmailSchema = z.string().email() + +/** + * Emits the semantic "credential used" trail for one resolved credential. + * Both the audit row and the analytics event are fire-and-forget. + */ +function recordCredentialAccess(params: { + actorId: string + workspaceId: string | null + resourceId: string + providerId: string | null | undefined + credentialType: 'oauth' | 'service_account' + extraMetadata?: Record + auditRequest?: CredentialAuditRequest +}): void { + const { actorId, workspaceId, resourceId, providerId, credentialType } = params + recordAudit({ + workspaceId, + actorId, + action: AuditAction.CREDENTIAL_ACCESSED, + resourceType: AuditResourceType.CREDENTIAL, + resourceId, + description: `Accessed ${credentialType === 'oauth' ? 'OAuth' : 'service account'} credential for provider ${providerId ?? 'unknown'}`, + metadata: { + provider: providerId, + credentialType, + ...params.extraMetadata, + }, + request: params.auditRequest, + }) + captureServerEvent( + actorId, + 'credential_used', + { + credential_type: credentialType, + provider_id: providerId ?? 'unknown', + ...(workspaceId ? { workspace_id: workspaceId } : {}), + }, + workspaceId ? { groups: { workspace: workspaceId } } : undefined + ) +} + +/** + * Projects a stored OAuth credential plus its (possibly refreshed) access token + * into the wire payload every surface returns. Provider-specific hosts live in + * the credential's scope string and are extracted through shared, allowlisted + * helpers — never a local regex, since these values are injected into tool + * calls that carry the token. + */ +export function buildOAuthTokenPayload( + credential: { providerId: string; scope?: string | null; idToken?: string | null }, + accessToken: string +): CredentialTokenPayload { + const instanceUrl = isSalesforceOAuthProviderId(credential.providerId) + ? extractSalesforceInstanceUrl(credential.scope ?? undefined) + : undefined + + // Zoho Desk persists its data-center-specific REST base URL in the scope + // string (derived from the token response api_domain) so callers never + // assume a host. Surface it as apiDomain for tool param injection. + let apiDomain: string | undefined + if (credential.providerId === 'zoho-desk' && credential.scope) { + apiDomain = extractZohoDeskBaseFromScope(credential.scope) + } + + return { + accessToken, + idToken: credential.idToken || undefined, + ...(instanceUrl && { instanceUrl }), + ...(apiDomain && { apiDomain }), + } +} + +/** + * Refreshes an authorized OAuth credential, records its access trail, and + * projects the token payload. Shared by every surface that has already + * authorized the credential and loaded it. + */ +export async function completeOAuthCredentialToken(params: { + requestId: string + credential: { providerId: string; scope?: string | null; idToken?: string | null } + resolvedCredentialId: string + actorId?: string + workspaceId: string | null + auditRequest?: CredentialAuditRequest +}): Promise { + const { requestId, credential, resolvedCredentialId, actorId, workspaceId, auditRequest } = params + try { + const { accessToken } = await refreshTokenIfNeeded(requestId, credential, resolvedCredentialId) + + if (actorId) { + recordCredentialAccess({ + actorId, + workspaceId, + resourceId: resolvedCredentialId, + providerId: credential.providerId, + credentialType: 'oauth', + auditRequest, + }) + } + + return { ok: true, token: buildOAuthTokenPayload(credential, accessToken) } + } catch (error) { + logger.error(`[${requestId}] Failed to refresh access token:`, error) + return { ok: false, status: 401, error: 'Failed to refresh access token' } + } +} + +/** + * Authorized application operation behind `POST /api/auth/oauth/token`. + * + * Given an already-authenticated caller, authorizes use of the credential, + * mints or refreshes its token, records the credential-access trail, and + * returns either the token payload or the exact status/error the HTTP surface + * projects. Every surface that needs a credential token — the route and the + * in-process tool executor — goes through here, so authorization, token + * refresh, and audit cannot drift between them. + * + * @param auth Result of authenticating the caller (session or internal JWT). + */ +export async function resolveCredentialToken( + auth: AuthResult, + input: ResolveCredentialTokenInput +): Promise { + const { + requestId, + credentialId, + workflowId, + scopes, + impersonateEmail, + callerUserId, + auditRequest, + } = input + + try { + if (!credentialId) { + return { ok: false, status: 400, error: 'Credential ID is required' } + } + if ( + impersonateEmail !== undefined && + !impersonateEmailSchema.safeParse(impersonateEmail).success + ) { + return { ok: false, status: 400, error: 'impersonateEmail must be a valid email address' } + } + + const resolved = await resolveOAuthAccountId(credentialId) + + if (resolved?.credentialType === 'service_account' && resolved.credentialId) { + const authz = await authorizeCredentialUseForAuth(auth, { + credentialId, + workflowId, + callerUserId, + }) + if (!authz.ok) { + return { ok: false, status: 403, error: authz.error || 'Unauthorized' } + } + + const saActorId = authz.requesterUserId + const saWorkspaceId = resolved.workspaceId ?? authz.workspaceId ?? null + + try { + const result = await resolveServiceAccountToken( + resolved.credentialId, + resolved.providerId, + scopes ?? [], + impersonateEmail + ) + + if (saActorId) { + recordCredentialAccess({ + actorId: saActorId, + workspaceId: saWorkspaceId, + resourceId: resolved.credentialId ?? credentialId, + providerId: resolved.providerId, + credentialType: 'service_account', + auditRequest, + }) + } + + return { + ok: true, + token: { + accessToken: result.accessToken, + cloudId: result.cloudId, + domain: result.domain, + instanceUrl: result.instanceUrl, + apiDomain: result.apiDomain, + authStyle: result.authStyle, + }, + } + } catch (error) { + logger.error(`[${requestId}] Service account token error:`, error) + if (error instanceof TokenServiceAccountValidationError) { + // Classified provider outages are infra failures, not bad credentials. + if (error.code === 'provider_unavailable') { + return { + ok: false, + status: 502, + error: 'Credential provider is temporarily unavailable', + } + } + // A stored host that no longer resolves is a configuration failure — + // surface the code so runtime consumers can say "check the host" + // instead of a generic auth error. + if (error.code === 'site_not_found') { + return { + ok: false, + status: 400, + code: error.code, + error: 'Credential host not found — reconnect the credential with a valid host', + } + } + // A revoked/rotated-away or misconfigured stored secret — surface the + // code so runtime consumers can prompt to reconnect the credential + // rather than showing a generic auth failure. + if (error.code === 'invalid_credentials') { + return { + ok: false, + status: 401, + code: error.code, + error: 'Credential rejected by the provider — reconnect the credential', + } + } + } + return { ok: false, status: 401, error: 'Failed to get service account token' } + } + } + + const authz = await authorizeCredentialUseForAuth(auth, { + credentialId, + workflowId, + callerUserId, + }) + if (!authz.ok || !authz.credentialOwnerUserId) { + return { ok: false, status: 403, error: authz.error || 'Unauthorized' } + } + + const resolvedCredentialId = authz.resolvedCredentialId || credentialId + const credential = await getCredential( + requestId, + resolvedCredentialId, + authz.credentialOwnerUserId + ) + + if (!credential) { + return { ok: false, status: 404, error: 'Credential not found' } + } + + return completeOAuthCredentialToken({ + requestId, + credential, + resolvedCredentialId, + actorId: authz.requesterUserId, + workspaceId: authz.workspaceId ?? null, + auditRequest, + }) + } catch (error) { + logger.error(`[${requestId}] Error getting access token`, error) + return { ok: false, status: 500, error: 'Internal server error' } + } +} diff --git a/apps/sim/lib/workflows/persistence/save-normalized-state.ts b/apps/sim/lib/workflows/persistence/save-normalized-state.ts new file mode 100644 index 00000000000..d86b2bd5b28 --- /dev/null +++ b/apps/sim/lib/workflows/persistence/save-normalized-state.ts @@ -0,0 +1,197 @@ +import { db } from '@sim/db' +import { workflow } from '@sim/db/schema' +import { createLogger } from '@sim/logger' +import { + assertWorkflowMutable, + authorizeWorkflowByWorkspacePermission, +} from '@sim/platform-authz/workflow' +import { eq } from 'drizzle-orm' +import type { z } from 'zod' +import { + type WorkflowStateContractOutput, + workflowStateSchema, +} from '@/lib/api/contracts/workflows' +import { env } from '@/lib/core/config/env' +import { getSocketServerUrl } from '@/lib/core/utils/urls' +import { extractAndPersistCustomTools } from '@/lib/workflows/persistence/custom-tools-persistence' +import { prepareWorkflowStateForPersistence } from '@/lib/workflows/persistence/prepare-state' +import { saveWorkflowToNormalizedTables } from '@/lib/workflows/persistence/utils' +import type { BlockState, WorkflowState } from '@/stores/workflows/workflow/types' + +const logger = createLogger('WorkflowStatePersistence') + +export type SaveWorkflowNormalizedStateResult = + | { success: true; warnings: string[] } + | { success: false; status: number; error: string; details?: string } + +/** + * Validates an untrusted workflow-state blob against the same schema the + * `PUT /api/workflows/[id]/state` contract applies. In-process callers holding a + * persisted blob (e.g. a copilot checkpoint) run it through here so they get the + * identical coercion and rejection the HTTP hop used to give them. + */ +export function parseWorkflowStateForPersistence( + value: unknown +): z.ZodSafeParseResult { + return workflowStateSchema.safeParse(value) +} + +/** + * Writes a complete workflow state to the normalized tables. + * + * Owns everything the state write means: write authorization, the mutability + * (lock) check, block/edge preparation, the row-locked save transaction, the + * `lastSynced`/variables update, custom-tool extraction, and the socket-server + * notification. Every surface that replaces a workflow's state — the PUT route + * and the copilot checkpoint revert — calls this, so none of those steps can be + * skipped by going through a different door. + * + * @throws WorkflowLockedError when the workflow is not mutable. + */ +export async function saveWorkflowNormalizedState(params: { + requestId: string + workflowId: string + userId: string + state: WorkflowStateContractOutput +}): Promise { + const { requestId, workflowId, userId, state } = params + + const authorization = await authorizeWorkflowByWorkspacePermission({ + workflowId, + userId, + action: 'write', + }) + const workflowData = authorization.workflow + + if (!workflowData) { + logger.warn(`[${requestId}] Workflow ${workflowId} not found for state update`) + return { success: false, status: 404, error: 'Workflow not found' } + } + + if (!authorization.allowed) { + logger.warn( + `[${requestId}] User ${userId} denied permission to update workflow state ${workflowId}` + ) + return { + success: false, + status: authorization.status || 403, + error: authorization.message || 'Access denied', + } + } + + await assertWorkflowMutable(workflowId) + + const { state: preparedState, warnings: preparationWarnings } = + prepareWorkflowStateForPersistence({ + blocks: state.blocks as Record, + edges: state.edges as WorkflowState['edges'], + }) + + const workflowState = { + ...preparedState, + lastSaved: state.lastSaved || Date.now(), + isDeployed: state.isDeployed || false, + deployedAt: state.deployedAt, + } + + const saveResult = await db.transaction(async (tx) => { + await tx + .select({ id: workflow.id }) + .from(workflow) + .where(eq(workflow.id, workflowId)) + .limit(1) + .for('update') + + const result = await saveWorkflowToNormalizedTables( + workflowId, + workflowState as WorkflowState, + tx + ) + + if (!result.success) return result + + const updateData: { + lastSynced: Date + updatedAt: Date + variables?: typeof state.variables + } = { + lastSynced: new Date(), + updatedAt: new Date(), + } + + if (state.variables !== undefined) { + updateData.variables = state.variables + } + + await tx.update(workflow).set(updateData).where(eq(workflow.id, workflowId)) + + return result + }) + + if (!saveResult.success) { + logger.error(`[${requestId}] Failed to save workflow ${workflowId} state:`, saveResult.error) + return { + success: false, + status: 500, + error: 'Failed to save workflow state', + details: saveResult.error, + } + } + + try { + const workspaceId = workflowData.workspaceId + if (workspaceId) { + const { saved, errors } = await extractAndPersistCustomTools( + workflowState, + workspaceId, + userId + ) + + if (saved > 0) { + logger.info(`[${requestId}] Persisted ${saved} custom tool(s) to database`, { workflowId }) + } + + if (errors.length > 0) { + logger.warn(`[${requestId}] Some custom tools failed to persist`, { errors, workflowId }) + } + } else { + logger.warn(`[${requestId}] Workflow has no workspaceId, skipping custom tools persistence`, { + workflowId, + }) + } + } catch (error) { + logger.error(`[${requestId}] Failed to persist custom tools`, { error, workflowId }) + } + + await notifySocketServer(requestId, workflowId) + + return { success: true, warnings: preparationWarnings } +} + +/** + * Best-effort nudge so connected editors reload the workflow. Never fails the + * write — the state is already committed by the time this runs. + */ +async function notifySocketServer(requestId: string, workflowId: string): Promise { + try { + const notifyResponse = await fetch(`${getSocketServerUrl()}/api/workflow-updated`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'x-api-key': env.INTERNAL_API_SECRET, + }, + body: JSON.stringify({ workflowId }), + }) + + if (!notifyResponse.ok) { + logger.warn( + `[${requestId}] Failed to notify Socket.IO server about workflow ${workflowId} update` + ) + } + } catch (notificationError) { + logger.warn( + `[${requestId}] Error notifying Socket.IO server about workflow ${workflowId} update`, + notificationError + ) + } +} diff --git a/apps/sim/tools/index.ts b/apps/sim/tools/index.ts index 575b6da5c35..99b2cd2ba5d 100644 --- a/apps/sim/tools/index.ts +++ b/apps/sim/tools/index.ts @@ -56,6 +56,7 @@ import { } from '@/lib/execution/private-tool-metadata' import { parseMcpToolId } from '@/lib/mcp/utils' import { hostedKeyMetrics } from '@/lib/monitoring/metrics' +import type { CredentialTokenPayload } from '@/lib/oauth/token-resolution' import { resolveWorkspaceFileReference } from '@/lib/uploads/contexts/workspace/workspace-file-manager' import { markWorkspaceFileSecretProvenanceUnknown } from '@/lib/uploads/contexts/workspace/workspace-file-secret-provenance' import { assertPermissionsAllowed } from '@/ee/access-control/utils/permission-check' @@ -1702,8 +1703,6 @@ async function executeToolImplementation( `[${requestId}] Tool ${toolId} needs access token for credential: ${contextParams.credential}` ) try { - const baseUrl = getInternalApiBaseUrl() - const workflowId = contextParams._context?.workflowId const userId = contextParams._context?.userId @@ -1724,51 +1723,86 @@ async function executeToolImplementation( } } - logger.info(`[${requestId}] Fetching access token from ${baseUrl}/api/auth/oauth/token`) + /** + * The acting user asserted alongside an internal token. Only sent when the + * run enforces credential access, matching the `userId` query param the HTTP + * surface accepted — it never widens access, it only pins the assertion to + * the token subject. + */ + const callerUserId = + userId && contextParams._context?.enforceCredentialAccess ? userId : undefined - const tokenUrlObj = new URL('/api/auth/oauth/token', baseUrl) - if (workflowId) { - tokenUrlObj.searchParams.set('workflowId', workflowId) - } - if (userId && contextParams._context?.enforceCredentialAccess) { - tokenUrlObj.searchParams.set('userId', userId) - } + let data: CredentialTokenPayload - // Always send Content-Type; add internal auth on server-side runs - const tokenHeaders: Record = { 'Content-Type': 'application/json' } if (typeof window === 'undefined') { - try { - const internalToken = await generateInternalToken(userId) - tokenHeaders.Authorization = `Bearer ${internalToken}` - } catch (_e) { - // Swallow token generation errors; the request will fail and be reported upstream + // Server-side runs resolve the credential through the same application + // operation the route calls, rather than minting an internal JWT and + // POSTing to ourselves through the load balancer. The synthesized + // `AuthResult` is exactly what verifying that self-issued token would + // have produced, so authorization, refresh, and audit are unchanged — + // including failing closed when the run carries no user id. + const { resolveCredentialToken } = await import('@/lib/oauth/token-resolution') + const result = await resolveCredentialToken( + { success: true, authType: 'internal_jwt', ...(userId ? { userId } : {}) }, + { + requestId, + credentialId: contextParams.credential as string, + workflowId, + scopes: tokenPayload.scopes, + impersonateEmail: tokenPayload.impersonateEmail, + callerUserId, + } + ) + + if (!result.ok) { + logger.error(`[${requestId}] Token fetch failed for ${toolId}:`, { + status: result.status, + error: result.error, + }) + const toolLabel = tool?.name || toolId + throw new Error(`Failed to obtain credential for ${toolLabel}: ${result.error}`) } - } - const response = await fetch(tokenUrlObj.toString(), { - method: 'POST', - headers: tokenHeaders, - body: JSON.stringify(tokenPayload), - }) + data = result.token + } else { + const baseUrl = getInternalApiBaseUrl() + logger.info(`[${requestId}] Fetching access token from ${baseUrl}/api/auth/oauth/token`) + + const tokenUrlObj = new URL('/api/auth/oauth/token', baseUrl) + if (workflowId) { + tokenUrlObj.searchParams.set('workflowId', workflowId) + } + if (callerUserId) { + tokenUrlObj.searchParams.set('userId', callerUserId) + } - if (!response.ok) { - const errorText = await response.text() - logger.error(`[${requestId}] Token fetch failed for ${toolId}:`, { - status: response.status, - error: errorText, + // boundary-raw-fetch: browser-side tool runs authenticate with the session cookie against the same-origin token route + const response = await fetch(tokenUrlObj.toString(), { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(tokenPayload), }) - let parsedError = errorText - try { - const parsed = JSON.parse(errorText) - if (parsed.error) parsedError = parsed.error - } catch { - // Use raw text + + if (!response.ok) { + const errorText = await response.text() + logger.error(`[${requestId}] Token fetch failed for ${toolId}:`, { + status: response.status, + error: errorText, + }) + let parsedError = errorText + try { + const parsed = JSON.parse(errorText) + if (parsed.error) parsedError = parsed.error + } catch { + // Use raw text + } + const toolLabel = tool?.name || toolId + throw new Error(`Failed to obtain credential for ${toolLabel}: ${parsedError}`) } - const toolLabel = tool?.name || toolId - throw new Error(`Failed to obtain credential for ${toolLabel}: ${parsedError}`) + + data = (await response.json()) as CredentialTokenPayload } - const data = await response.json() contextParams.accessToken = data.accessToken if (data.idToken) { contextParams.idToken = data.idToken From 0dd8f7c2b7f5da2cb7c202f44a3292a2db4b71d8 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 13 Aug 2026 01:32:12 -0700 Subject: [PATCH 10/14] improvement(executor): run router and evaluator provider calls in-process --- .../evaluator/evaluator-handler.test.ts | 346 ++++++---------- .../handlers/evaluator/evaluator-handler.ts | 43 +- .../handlers/router/router-handler.test.ts | 377 +++++++----------- .../handlers/router/router-handler.ts | 82 +--- apps/sim/executor/utils/provider-request.ts | 76 ++++ 5 files changed, 358 insertions(+), 566 deletions(-) create mode 100644 apps/sim/executor/utils/provider-request.ts diff --git a/apps/sim/executor/handlers/evaluator/evaluator-handler.test.ts b/apps/sim/executor/handlers/evaluator/evaluator-handler.test.ts index 756b98798e2..eed0e6fed06 100644 --- a/apps/sim/executor/handlers/evaluator/evaluator-handler.test.ts +++ b/apps/sim/executor/handlers/evaluator/evaluator-handler.test.ts @@ -4,10 +4,15 @@ import { createLogger } from '@sim/logger' import { authOAuthUtilsMock, authOAuthUtilsMockFns } from '@sim/testing' import { beforeEach, describe, expect, it, type Mock, vi } from 'vitest' -const { mockResolveAutoModel } = vi.hoisted(() => ({ +const { mockResolveAutoModel, mockCheckWorkspaceAccess } = vi.hoisted(() => ({ + mockCheckWorkspaceAccess: vi.fn(), mockResolveAutoModel: vi.fn(), })) +vi.mock('@/lib/workspaces/permissions/utils', () => ({ + checkWorkspaceAccess: mockCheckWorkspaceAccess, +})) + vi.mock('@/lib/oauth/credential-service', () => authOAuthUtilsMock) vi.mock('@/lib/credentials/access', () => ({ @@ -34,24 +39,26 @@ vi.mock('@/lib/model-router/resolve', () => ({ SIM_AUTO_SYSTEM_PREAMBLE: 'Sim auto system preamble', })) -import { - PRIVATE_MODEL_INPUT_PROVENANCE_HEADER, - PRIVATE_MODEL_INPUT_STATE_HEADER, - PROJECTED_MODEL_INPUT_PATHS_V1, -} from '@/lib/execution/model-input-provenance' -import { - RESOLVED_SECRET_PROVENANCE_FIELD, - RESOLVED_SECRET_PROVENANCE_METADATA_V1, -} from '@/lib/execution/private-tool-metadata' import { BlockType } from '@/executor/constants' import { EvaluatorBlockHandler } from '@/executor/handlers/evaluator/evaluator-handler' import type { ExecutionContext } from '@/executor/types' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' +import { executeProviderRequest } from '@/providers' import { getProviderFromModel } from '@/providers/utils' import type { SerializedBlock } from '@/serializer/types' const mockGetProviderFromModel = getProviderFromModel as Mock -const mockFetch = vi.fn() +const mockExecuteProviderRequest = executeProviderRequest as Mock + +/** The provider request the handler built, keyed the way the old wire body was. */ +function providerRequestBody(index = 0): Record { + const [provider, request] = mockExecuteProviderRequest.mock.calls[index] + return { provider, ...request } +} + +function providerRuntimeRegistry(index = 0): ResolvedSecretTraceRegistry | undefined { + return mockExecuteProviderRequest.mock.calls[index][2]?.resolvedSecretTraceRegistry +} const mockLogger = vi.mocked(createLogger).mock.results[ vi.mocked(createLogger).mock.calls.findIndex(([name]) => name === 'EvaluatorBlockHandler') @@ -97,8 +104,7 @@ describe('EvaluatorBlockHandler', () => { // Reset mocks using vi vi.clearAllMocks() - // unstubGlobals removes any module-scope fetch stub before each test, so re-stub here - vi.stubGlobal('fetch', mockFetch) + mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: true }) // Default mock implementations authOAuthUtilsMockFns.mockResolveOAuthAccountId.mockResolvedValue({ @@ -117,19 +123,12 @@ describe('EvaluatorBlockHandler', () => { billableRoutingCost: 0.002, }) - // Set up fetch mock to return a successful response - mockFetch.mockImplementation(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ score1: 5, score2: 8 }), - model: 'mock-model', - tokens: { input: 50, output: 10, total: 60 }, - cost: 0.002, - timing: { total: 200 }, - }), - }) + mockExecuteProviderRequest.mockResolvedValue({ + content: JSON.stringify({ score1: 5, score2: 8 }), + model: 'mock-model', + tokens: { input: 50, output: 10, total: 60 }, + cost: 0.002, + timing: { total: 200 }, }) }) @@ -165,17 +164,9 @@ describe('EvaluatorBlockHandler', () => { const result = await handler.execute(mockContext, mockBlock, inputs) expect(mockGetProviderFromModel).toHaveBeenCalledWith('gpt-4o') - expect(mockFetch).toHaveBeenCalledWith( - expect.any(String), - expect.objectContaining({ - method: 'POST', - headers: expect.any(Object), - body: expect.any(String), - }) - ) + expect(mockExecuteProviderRequest).toHaveBeenCalledTimes(1) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ provider: 'openai', model: 'gpt-4o', @@ -257,15 +248,8 @@ describe('EvaluatorBlockHandler', () => { apiKey: credentialSecret, }) - const request = mockFetch.mock.calls[0][1] - const requestBody = JSON.parse(request.body) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBe( - RESOLVED_SECRET_PROVENANCE_METADATA_V1 - ) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBe( - PROJECTED_MODEL_INPUT_PATHS_V1 - ) - expect(requestBody[RESOLVED_SECRET_PROVENANCE_FIELD]).toEqual({ + const requestBody = providerRequestBody() + expect(providerRuntimeRegistry()?.exportProvenance()).toEqual({ version: 1, complete: true, entries: [ @@ -336,15 +320,11 @@ describe('EvaluatorBlockHandler', () => { registry.recordResolvedInputProjection(secret.path, secret.plaintext, secret.projected) } mockContext.resolvedSecretTraceRegistry = registry - mockFetch.mockResolvedValueOnce({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ [projectedMetric.name.toLowerCase()]: 7 }), - model: 'mock-model', - tokens: {}, - cost: 0, - }), + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ [projectedMetric.name.toLowerCase()]: 7 }), + model: 'mock-model', + tokens: {}, + cost: 0, }) const result = await handler.execute(mockContext, mockBlock, { @@ -354,7 +334,7 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', }) - const requestBody = JSON.parse(mockFetch.mock.calls[0][1].body) + const requestBody = providerRequestBody() const serializedRequest = JSON.stringify(requestBody) for (const secret of secrets) { expect(serializedRequest).not.toContain(secret.plaintext) @@ -365,8 +345,9 @@ describe('EvaluatorBlockHandler', () => { [projectedMetric.name.toLowerCase()]: { type: 'number' }, }) expect( - requestBody[RESOLVED_SECRET_PROVENANCE_FIELD].entries - .map((entry: { name: string }) => entry.name) + providerRuntimeRegistry() + ?.exportProvenance() + .entries.map((entry: { name: string }) => entry.name) .sort() ).toEqual(secrets.map((secret) => secret.name).sort()) expect(result).toMatchObject({ [rawMetric.name.toLowerCase()]: 7 }) @@ -380,11 +361,7 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', }) - const request = mockFetch.mock.calls[0][1] - const requestBody = JSON.parse(request.body) - expect(Object.hasOwn(requestBody, RESOLVED_SECRET_PROVENANCE_FIELD)).toBe(false) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBeNull() - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBeNull() + expect(providerRuntimeRegistry()).toBeUndefined() }) it('resolves sim-auto before executing evaluator and preserves its public identity', async () => { @@ -400,15 +377,11 @@ describe('EvaluatorBlockHandler', () => { model: 'sim-auto', } - mockFetch.mockResolvedValueOnce({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ quality: 5 }), - model: 'fireworks/glm-5.2', - tokens: { input: 80, output: 10, total: 90 }, - cost: { input: 0.001, output: 0.0005, total: 0.0015 }, - }), + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ quality: 5 }), + model: 'fireworks/glm-5.2', + tokens: { input: 80, output: 10, total: 90 }, + cost: { input: 0.001, output: 0.0005, total: 0.0015 }, }) const result = await handler.execute(mockContext, mockBlock, inputs) @@ -427,7 +400,7 @@ describe('EvaluatorBlockHandler', () => { }) expect(mockGetProviderFromModel).toHaveBeenCalledWith('fireworks/glm-5.2') - const requestBody = JSON.parse(mockFetch.mock.calls[0][1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ provider: 'openai', model: 'fireworks/glm-5.2', @@ -448,19 +421,13 @@ describe('EvaluatorBlockHandler', () => { it('bills the cost the provider proxy decided rather than recomputing it', async () => { // The proxy already resolved key provenance and the margin; recomputing // here would re-charge a BYOK caller the proxy correctly zeroed. - mockFetch.mockImplementation(() => - Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ score1: 5, score2: 8 }), - model: 'mock-model', - tokens: { input: 50, output: 10, total: 60 }, - cost: { input: 0.001, output: 0.0005, total: 0.0015 }, - timing: { total: 200 }, - }), - }) - ) + mockExecuteProviderRequest.mockResolvedValue({ + content: JSON.stringify({ score1: 5, score2: 8 }), + model: 'mock-model', + tokens: { input: 50, output: 10, total: 60 }, + cost: { input: 0.001, output: 0.0005, total: 0.0015 }, + timing: { total: 200 }, + }) const result = await handler.execute(mockContext, mockBlock, { content: 'This is the content to evaluate.', @@ -487,24 +454,17 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ clarity: 4 }), - model: 'm', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ clarity: 4 }), + model: 'm', + tokens: {}, + cost: 0, + timing: {}, }) await handler.execute(mockContext, mockBlock, inputs) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ systemPrompt: expect.stringContaining(JSON.stringify(contentObj, null, 2)), }) @@ -524,24 +484,17 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ completeness: 1 }), - model: 'm', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ completeness: 1 }), + model: 'm', + tokens: {}, + cost: 0, + timing: {}, }) await handler.execute(mockContext, mockBlock, inputs) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ systemPrompt: expect.stringContaining(JSON.stringify(contentObj, null, 2)), }) @@ -560,18 +513,12 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: '```json\n{ "quality": 9 }\n```', - model: 'm', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: '```json\n{ "quality": 9 }\n```', + model: 'm', + tokens: {}, + cost: 0, + timing: {}, }) const result = await handler.execute(mockContext, mockBlock, inputs) @@ -586,18 +533,12 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: 'Sorry, I cannot provide a score.', - model: 'm', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: 'Sorry, I cannot provide a score.', + model: 'm', + tokens: {}, + cost: 0, + timing: {}, }) const result = await handler.execute(mockContext, mockBlock, inputs) @@ -615,18 +556,12 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: '{ "accuracy": 1, "fluency": invalid }', - model: 'm', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: '{ "accuracy": 1, "fluency": invalid }', + model: 'm', + tokens: {}, + cost: 0, + timing: {}, }) const result = await handler.execute(mockContext, mockBlock, inputs) @@ -647,18 +582,12 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ camelcasescore: 7 }), - model: 'm', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ camelcasescore: 7 }), + model: 'm', + tokens: {}, + cost: 0, + timing: {}, }) const result = await handler.execute(mockContext, mockBlock, inputs) @@ -684,18 +613,12 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ presentScore: 4 }), - model: 'm', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ presentScore: 4 }), + model: 'm', + tokens: {}, + cost: 0, + timing: {}, }) const result = await handler.execute(mockContext, mockBlock, inputs) @@ -708,13 +631,7 @@ describe('EvaluatorBlockHandler', () => { const inputs = { content: 'Test error handling.', apiKey: 'test-api-key' } // Override fetch mock to return an error - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: false, - status: 500, - json: () => Promise.resolve({ error: 'Server error' }), - }) - }) + mockExecuteProviderRequest.mockRejectedValueOnce(new Error('Server error')) await expect(handler.execute(mockContext, mockBlock, inputs)).rejects.toThrow('Server error') }) @@ -730,11 +647,7 @@ describe('EvaluatorBlockHandler', () => { ]) registry.recordResolved('CONTENT_SECRET', 'resolved-evaluator-secret') mockContext.resolvedSecretTraceRegistry = registry - mockFetch.mockResolvedValueOnce({ - ok: false, - status: 500, - json: () => Promise.resolve({ error: providerError }), - }) + mockExecuteProviderRequest.mockRejectedValueOnce(new Error(providerError)) await expect( handler.execute(mockContext, mockBlock, { @@ -767,24 +680,17 @@ describe('EvaluatorBlockHandler', () => { mockGetProviderFromModel.mockReturnValue('azure-openai') - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ quality: 8 }), - model: 'gpt-4o', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ quality: 8 }), + model: 'gpt-4o', + tokens: {}, + cost: 0, + timing: {}, }) await handler.execute(mockContext, mockBlock, inputs) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ provider: 'azure-openai', @@ -824,24 +730,17 @@ describe('EvaluatorBlockHandler', () => { ;(mockDb.db.query as any).account = { findFirst: vi.fn() } vi.spyOn(mockDb.db.query.account, 'findFirst').mockResolvedValue(mockAccount as any) - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ quality: 9 }), - model: 'gemini-2.0-flash-exp', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ quality: 9 }), + model: 'gemini-2.0-flash-exp', + tokens: {}, + cost: 0, + timing: {}, }) await handler.execute(mockContext, mockBlock, inputs) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ provider: 'vertex', @@ -860,24 +759,17 @@ describe('EvaluatorBlockHandler', () => { // No model provided - should use default } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ score: 7 }), - model: 'claude-sonnet-5', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ score: 7 }), + model: 'claude-sonnet-5', + tokens: {}, + cost: 0, + timing: {}, }) await handler.execute(mockContext, mockBlock, inputs) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody.model).toBe('claude-sonnet-5') }) diff --git a/apps/sim/executor/handlers/evaluator/evaluator-handler.ts b/apps/sim/executor/handlers/evaluator/evaluator-handler.ts index 6a094f64082..f162ae25c38 100644 --- a/apps/sim/executor/handlers/evaluator/evaluator-handler.ts +++ b/apps/sim/executor/handlers/evaluator/evaluator-handler.ts @@ -1,10 +1,5 @@ import { createLogger } from '@sim/logger' -import { - addModelInputProvenanceToRequest, - createModelInputProvenanceRequestMetadata, - markModelInputProjected, - projectResolvedModelInput, -} from '@/lib/execution/model-input-provenance' +import { projectResolvedModelInput } from '@/lib/execution/model-input-provenance' import { type AutoRoutingResult, addAutoRoutingCost, @@ -15,8 +10,8 @@ import type { BlockOutput } from '@/blocks/types' import { validateModelProvider } from '@/ee/access-control/utils/permission-check' import { BlockType, DEFAULTS, EVALUATOR } from '@/executor/constants' import type { BlockHandler, ExecutionContext } from '@/executor/types' -import { buildAPIUrl, buildAuthHeaders, extractAPIErrorMessage } from '@/executor/utils/http' import { isJSONString, parseJSON, stringifyJSON } from '@/executor/utils/json' +import { executeBlockProviderRequest } from '@/executor/utils/provider-request' import { projectResolvedSecretDiagnosticError } from '@/executor/utils/resolved-secret-content-projection' import { refuseResolvedSecretProjection } from '@/executor/utils/resolved-secret-projection-refusal' import type { @@ -192,8 +187,6 @@ export class EvaluatorBlockHandler implements BlockHandler { } try { - const url = buildAPIUrl('/api/providers', ctx.userId ? { userId: ctx.userId } : {}) - const providerRequest: ProviderRequest = { model, systemPrompt: systemPromptObj.systemPrompt, @@ -219,30 +212,13 @@ export class EvaluatorBlockHandler implements BlockHandler { workspaceId: ctx.workspaceId, } - const headers = new Headers(await buildAuthHeaders(ctx.userId)) - const modelInputMetadata = createModelInputProvenanceRequestMetadata( - modelInputProjection.registry, - modelInputPaths - ) - const requestBody = addModelInputProvenanceToRequest( - { provider: providerId, ...providerRequest }, - headers, - modelInputMetadata - ) - if (modelInputMetadata) markModelInputProjected(headers) - const response = await fetch(url.toString(), { - method: 'POST', - headers, - body: stringifyJSON(requestBody), + const result = await executeBlockProviderRequest({ + ctx, + providerId, + request: providerRequest, + resolvedSecretTraceRegistry: modelInputProjection.registry, }) - if (!response.ok) { - const errorMessage = await extractAPIErrorMessage(response) - throw new Error(errorMessage) - } - - const result = await response.json() - const parsedContent = this.extractJSONFromResponse( result.content, ctx.resolvedSecretTraceRegistry @@ -250,9 +226,8 @@ export class EvaluatorBlockHandler implements BlockHandler { const metricScores = this.extractMetricScores(parsedContent, metrics, projectedMetrics) - const inputTokens = result.tokens?.input || result.tokens?.prompt || DEFAULTS.TOKENS.PROMPT - const outputTokens = - result.tokens?.output || result.tokens?.completion || DEFAULTS.TOKENS.COMPLETION + const inputTokens = result.tokens?.input || DEFAULTS.TOKENS.PROMPT + const outputTokens = result.tokens?.output || DEFAULTS.TOKENS.COMPLETION const cost = addAutoRoutingCost( resolveProxiedModelCost(result.cost), diff --git a/apps/sim/executor/handlers/router/router-handler.test.ts b/apps/sim/executor/handlers/router/router-handler.test.ts index a7fe9e45141..8c7f7b0e945 100644 --- a/apps/sim/executor/handlers/router/router-handler.test.ts +++ b/apps/sim/executor/handlers/router/router-handler.test.ts @@ -9,8 +9,13 @@ import { } from '@sim/testing' import { beforeEach, describe, expect, it, type Mock, vi } from 'vitest' -const { mockResolveAutoModel } = vi.hoisted(() => ({ +const { mockResolveAutoModel, mockCheckWorkspaceAccess } = vi.hoisted(() => ({ mockResolveAutoModel: vi.fn(), + mockCheckWorkspaceAccess: vi.fn(), +})) + +vi.mock('@/lib/workspaces/permissions/utils', () => ({ + checkWorkspaceAccess: mockCheckWorkspaceAccess, })) vi.mock('@/lib/oauth/credential-service', () => authOAuthUtilsMock) @@ -40,27 +45,30 @@ vi.mock('@/lib/model-router/resolve', () => ({ SIM_AUTO_SYSTEM_PREAMBLE: 'Sim auto system preamble', })) -import { - PRIVATE_MODEL_INPUT_PROVENANCE_HEADER, - PRIVATE_MODEL_INPUT_STATE_HEADER, - PROJECTED_MODEL_INPUT_PATHS_V1, -} from '@/lib/execution/model-input-provenance' -import { - RESOLVED_SECRET_PROVENANCE_FIELD, - RESOLVED_SECRET_PROVENANCE_METADATA_V1, -} from '@/lib/execution/private-tool-metadata' import { generateRouterPrompt, generateRouterV2Prompt } from '@/blocks/blocks/router' import { BlockType } from '@/executor/constants' import { RouterBlockHandler } from '@/executor/handlers/router/router-handler' import type { ExecutionContext } from '@/executor/types' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' +import { executeProviderRequest } from '@/providers' import { getProviderFromModel } from '@/providers/utils' import type { SerializedBlock, SerializedWorkflow } from '@/serializer/types' const mockGenerateRouterPrompt = generateRouterPrompt as Mock const mockGenerateRouterV2Prompt = generateRouterV2Prompt as Mock const mockGetProviderFromModel = getProviderFromModel as Mock -const mockFetch = vi.fn() +const mockExecuteProviderRequest = executeProviderRequest as Mock + +/** The provider request the handler built, keyed the way the old wire body was. */ +function providerRequestBody(index = 0): Record { + const [provider, request] = mockExecuteProviderRequest.mock.calls[index] + return { provider, ...request } +} + +function providerRuntimeRegistry(index = 0): ResolvedSecretTraceRegistry | undefined { + return mockExecuteProviderRequest.mock.calls[index][2]?.resolvedSecretTraceRegistry +} + const mockLogger = vi.mocked(createLogger).mock.results[ vi.mocked(createLogger).mock.calls.findIndex(([name]) => name === 'RouterBlockHandler') @@ -138,8 +146,7 @@ describe('RouterBlockHandler', () => { vi.clearAllMocks() encryptionMockFns.mockDecryptSecret.mockResolvedValue({ decrypted: 'test-decrypted' }) - // unstubGlobals removes any module-scope fetch stub before each test, so re-stub here - vi.stubGlobal('fetch', mockFetch) + mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: true }) authOAuthUtilsMockFns.mockResolveOAuthAccountId.mockResolvedValue({ accountId: 'test-vertex-credential-id', @@ -158,18 +165,12 @@ describe('RouterBlockHandler', () => { billableRoutingCost: 0.002, }) - mockFetch.mockImplementation(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: 'target-block-1', - model: 'mock-model', - tokens: { input: 100, output: 5, total: 105 }, - cost: 0.003, - timing: { total: 300 }, - }), - }) + mockExecuteProviderRequest.mockResolvedValue({ + content: 'target-block-1', + model: 'mock-model', + tokens: { input: 100, output: 5, total: 105 }, + cost: 0.003, + timing: { total: 300 }, }) }) @@ -219,17 +220,9 @@ describe('RouterBlockHandler', () => { expect(mockGenerateRouterPrompt).toHaveBeenCalledWith(inputs.prompt, expectedTargetBlocks) expect(mockGetProviderFromModel).toHaveBeenCalledWith('gpt-4o') - expect(mockFetch).toHaveBeenCalledWith( - expect.any(String), - expect.objectContaining({ - method: 'POST', - headers: expect.any(Object), - body: expect.any(String), - }) - ) + expect(mockExecuteProviderRequest).toHaveBeenCalledTimes(1) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ provider: 'openai', model: 'gpt-4o', @@ -282,15 +275,8 @@ describe('RouterBlockHandler', () => { apiKey: credentialSecret, }) - const request = mockFetch.mock.calls[0][1] - const requestBody = JSON.parse(request.body) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBe( - RESOLVED_SECRET_PROVENANCE_METADATA_V1 - ) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBe( - PROJECTED_MODEL_INPUT_PATHS_V1 - ) - expect(requestBody[RESOLVED_SECRET_PROVENANCE_FIELD]).toEqual({ + const requestBody = providerRequestBody() + expect(providerRuntimeRegistry()?.exportProvenance()).toEqual({ version: 1, complete: true, entries: [ @@ -353,8 +339,7 @@ describe('RouterBlockHandler', () => { expect(rawState).toEqual({ result: stateSecret, ordinary: 'Box remains raw state' }) expect(mockTargetBlock1.config.params).toEqual({ p: 'a' }) - const requestBody = JSON.parse(mockFetch.mock.calls[0][1].body) - expect(requestBody[RESOLVED_SECRET_PROVENANCE_FIELD]).toEqual({ + expect(providerRuntimeRegistry()?.exportProvenance()).toEqual({ version: 1, complete: true, entries: [], @@ -404,29 +389,19 @@ describe('RouterBlockHandler', () => { apiKey: 'test-api-key', }) - const request = mockFetch.mock.calls[0][1] - const requestBody = JSON.parse(request.body) - expect(Object.hasOwn(requestBody, RESOLVED_SECRET_PROVENANCE_FIELD)).toBe(false) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBeNull() - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBeNull() + expect(providerRuntimeRegistry()).toBeUndefined() }) it('bills the cost the provider proxy decided rather than recomputing it', async () => { // The proxy already resolved key provenance and the margin; recomputing // here would re-charge a BYOK caller the proxy correctly zeroed. - mockFetch.mockImplementation(() => - Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: 'target-block-1', - model: 'mock-model', - tokens: { input: 100, output: 5, total: 105 }, - cost: { input: 0.004, output: 0.002, total: 0.006 }, - timing: { total: 300 }, - }), - }) - ) + mockExecuteProviderRequest.mockResolvedValue({ + content: 'target-block-1', + model: 'mock-model', + tokens: { input: 100, output: 5, total: 105 }, + cost: { input: 0.004, output: 0.002, total: 0.006 }, + timing: { total: 300 }, + }) const result = await handler.execute(mockContext, mockBlock, { prompt: 'Choose the best option.', @@ -439,6 +414,26 @@ describe('RouterBlockHandler', () => { }) }) + it('refuses to reach the provider without an execution subject', async () => { + mockContext.userId = undefined + + await expect( + handler.execute(mockContext, mockBlock, { prompt: 'Choose the best option.' }) + ).rejects.toThrow('Unauthorized') + expect(mockExecuteProviderRequest).not.toHaveBeenCalled() + }) + + it('refuses to reach the provider when the subject lost workspace access', async () => { + mockContext.workspaceId = 'test-workspace' + mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: false }) + + await expect( + handler.execute(mockContext, mockBlock, { prompt: 'Choose the best option.' }) + ).rejects.toThrow('Forbidden') + expect(mockCheckWorkspaceAccess).toHaveBeenCalledWith('test-workspace', 'test-user') + expect(mockExecuteProviderRequest).not.toHaveBeenCalled() + }) + it('should throw error if target block is missing', async () => { const inputs = { prompt: 'Test' } mockContext.workflow!.blocks = [mockBlock, mockTargetBlock2] @@ -446,24 +441,18 @@ describe('RouterBlockHandler', () => { await expect(handler.execute(mockContext, mockBlock, inputs)).rejects.toThrow( 'Target block target-block-1 not found' ) - expect(mockFetch).not.toHaveBeenCalled() + expect(mockExecuteProviderRequest).not.toHaveBeenCalled() }) it('should throw error if LLM response is not a valid target block ID', async () => { const inputs = { prompt: 'Test', apiKey: 'test-api-key' } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: 'invalid-block-id', - model: 'mock-model', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: 'invalid-block-id', + model: 'mock-model', + tokens: {}, + cost: 0, + timing: {}, }) await expect(handler.execute(mockContext, mockBlock, inputs)).rejects.toThrow( @@ -475,16 +464,12 @@ describe('RouterBlockHandler', () => { const plaintext = 'router-provider-plaintext-secret' const content = `${plaintext} __var_API_KEY __sim_runtime` - mockFetch.mockResolvedValueOnce({ - ok: true, - json: () => - Promise.resolve({ - content, - model: 'mock-model', - tokens: {}, - cost: 0, - timing: {}, - }), + mockExecuteProviderRequest.mockResolvedValueOnce({ + content, + model: 'mock-model', + tokens: {}, + cost: 0, + timing: {}, }) await expect(handler.execute(mockContext, mockBlock, { prompt: 'Test' })).rejects.toThrow( @@ -512,8 +497,7 @@ describe('RouterBlockHandler', () => { expect(mockGetProviderFromModel).toHaveBeenCalledWith('claude-sonnet-5') - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ model: 'claude-sonnet-5', temperature: 0.1, @@ -523,13 +507,7 @@ describe('RouterBlockHandler', () => { it('should handle server error responses', async () => { const inputs = { prompt: 'Test error handling.', apiKey: 'test-api-key' } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: false, - status: 500, - json: () => Promise.resolve({ error: 'Server error' }), - }) - }) + mockExecuteProviderRequest.mockRejectedValueOnce(new Error('Server error')) await expect(handler.execute(mockContext, mockBlock, inputs)).rejects.toThrow('Server error') }) @@ -537,11 +515,7 @@ describe('RouterBlockHandler', () => { it('does not log sensitive provider errors while preserving the thrown error', async () => { const providerError = 'provider-plaintext-secret __var_API_KEY __sim_runtime' - mockFetch.mockResolvedValueOnce({ - ok: false, - status: 500, - json: () => Promise.resolve({ error: providerError }), - }) + mockExecuteProviderRequest.mockRejectedValueOnce(new Error(providerError)) await expect(handler.execute(mockContext, mockBlock, { prompt: 'Test' })).rejects.toThrow( providerError @@ -569,8 +543,7 @@ describe('RouterBlockHandler', () => { await handler.execute(mockContext, mockBlock, inputs) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ provider: 'azure-openai', @@ -604,8 +577,7 @@ describe('RouterBlockHandler', () => { await handler.execute(mockContext, mockBlock, inputs) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ provider: 'vertex', @@ -688,8 +660,7 @@ describe('RouterBlockHandler V2', () => { vi.clearAllMocks() - // unstubGlobals removes any module-scope fetch stub before each test, so re-stub here - vi.stubGlobal('fetch', mockFetch) + mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: true }) authOAuthUtilsMockFns.mockResolveOAuthAccountId.mockResolvedValue({ accountId: 'test-vertex-credential-id', @@ -732,19 +703,13 @@ describe('RouterBlockHandler V2', () => { ]), } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ - route: 'route-support', - reasoning: 'The user mentioned a billing issue which is a customer support matter.', - }), - model: 'gpt-4o', - tokens: { input: 150, output: 25, total: 175 }, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ + route: 'route-support', + reasoning: 'The user mentioned a billing issue which is a customer support matter.', + }), + model: 'gpt-4o', + tokens: { input: 150, output: 25, total: 175 }, }) const result = await handler.execute(mockContext, mockRouterV2Block, inputs) @@ -781,14 +746,10 @@ describe('RouterBlockHandler V2', () => { registry.recordResolvedInputProjection(['context'], contextSecret, '{{CONTEXT_SECRET}}') registry.recordResolved('API_KEY', credentialSecret) mockContext.resolvedSecretTraceRegistry = registry - mockFetch.mockResolvedValueOnce({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ route: 'route-support', reasoning: 'Matched support.' }), - model: 'gpt-4o', - tokens: { input: 10, output: 5, total: 15 }, - }), + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ route: 'route-support', reasoning: 'Matched support.' }), + model: 'gpt-4o', + tokens: { input: 10, output: 5, total: 15 }, }) await handler.execute(mockContext, mockRouterV2Block, { @@ -798,15 +759,8 @@ describe('RouterBlockHandler V2', () => { routes: [{ id: 'route-support', title: 'Support', value: 'Support requests' }], }) - const request = mockFetch.mock.calls[0][1] - const requestBody = JSON.parse(request.body) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBe( - RESOLVED_SECRET_PROVENANCE_METADATA_V1 - ) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBe( - PROJECTED_MODEL_INPUT_PATHS_V1 - ) - expect(requestBody[RESOLVED_SECRET_PROVENANCE_FIELD]).toEqual({ + const requestBody = providerRequestBody() + expect(providerRuntimeRegistry()?.exportProvenance()).toEqual({ version: 1, complete: true, entries: [ @@ -821,14 +775,10 @@ describe('RouterBlockHandler V2', () => { }) it('keeps the router V2 request shape when no provenance registry exists', async () => { - mockFetch.mockResolvedValueOnce({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ route: 'route-support', reasoning: 'Matched support.' }), - model: 'gpt-4o', - tokens: { input: 10, output: 5, total: 15 }, - }), + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ route: 'route-support', reasoning: 'Matched support.' }), + model: 'gpt-4o', + tokens: { input: 10, output: 5, total: 15 }, }) await handler.execute(mockContext, mockRouterV2Block, { @@ -838,11 +788,7 @@ describe('RouterBlockHandler V2', () => { routes: [{ id: 'route-support', title: 'Support', value: 'Support requests' }], }) - const request = mockFetch.mock.calls[0][1] - const requestBody = JSON.parse(request.body) - expect(Object.hasOwn(requestBody, RESOLVED_SECRET_PROVENANCE_FIELD)).toBe(false) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBeNull() - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBeNull() + expect(providerRuntimeRegistry()).toBeUndefined() }) it('resolves sim-auto before executing router V2 and preserves its public identity', async () => { @@ -859,18 +805,14 @@ describe('RouterBlockHandler V2', () => { ], } - mockFetch.mockResolvedValueOnce({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ - route: 'route-sales', - reasoning: 'This is a new request.', - }), - model: 'fireworks/glm-5.2', - tokens: { input: 100, output: 20, total: 120 }, - cost: { input: 0.001, output: 0.0005, total: 0.0015 }, - }), + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ + route: 'route-sales', + reasoning: 'This is a new request.', + }), + model: 'fireworks/glm-5.2', + tokens: { input: 100, output: 20, total: 120 }, + cost: { input: 0.001, output: 0.0005, total: 0.0015 }, }) const result = await handler.execute(mockContext, mockRouterV2Block, inputs) @@ -889,7 +831,7 @@ describe('RouterBlockHandler V2', () => { }) expect(mockGetProviderFromModel).toHaveBeenCalledWith('fireworks/glm-5.2') - const requestBody = JSON.parse(mockFetch.mock.calls[0][1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ provider: 'openai', model: 'fireworks/glm-5.2', @@ -915,25 +857,18 @@ describe('RouterBlockHandler V2', () => { routes: JSON.stringify([{ id: 'route-1', title: 'Route 1', value: 'Description 1' }]), } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ - route: 'route-1', - reasoning: 'Test reasoning', - }), - model: 'gpt-4o', - tokens: { input: 100, output: 20, total: 120 }, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ + route: 'route-1', + reasoning: 'Test reasoning', + }), + model: 'gpt-4o', + tokens: { input: 100, output: 20, total: 120 }, }) await handler.execute(mockContext, mockRouterV2Block, inputs) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody.responseFormat).toEqual({ name: 'router_response', @@ -964,19 +899,13 @@ describe('RouterBlockHandler V2', () => { routes: JSON.stringify([{ id: 'route-1', title: 'Route 1', value: 'Specific topic' }]), } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ - route: 'NO_MATCH', - reasoning: 'The query does not relate to any available route.', - }), - model: 'gpt-4o', - tokens: { input: 100, output: 20, total: 120 }, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ + route: 'NO_MATCH', + reasoning: 'The query does not relate to any available route.', + }), + model: 'gpt-4o', + tokens: { input: 100, output: 20, total: 120 }, }) await expect(handler.execute(mockContext, mockRouterV2Block, inputs)).rejects.toThrow( @@ -992,19 +921,13 @@ describe('RouterBlockHandler V2', () => { routes: JSON.stringify([{ id: 'route-1', title: 'Route 1', value: 'Description' }]), } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ - route: 'invalid-route', - reasoning: 'Some reasoning', - }), - model: 'gpt-4o', - tokens: { input: 100, output: 20, total: 120 }, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ + route: 'invalid-route', + reasoning: 'Some reasoning', + }), + model: 'gpt-4o', + tokens: { input: 100, output: 20, total: 120 }, }) await expect(handler.execute(mockContext, mockRouterV2Block, inputs)).rejects.toThrow( @@ -1020,19 +943,13 @@ describe('RouterBlockHandler V2', () => { routes: [{ id: 'route-1', title: 'Route 1', value: 'Description' }], } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ - route: 'route-1', - reasoning: 'Matched route 1', - }), - model: 'gpt-4o', - tokens: { input: 100, output: 20, total: 120 }, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ + route: 'route-1', + reasoning: 'Matched route 1', + }), + model: 'gpt-4o', + tokens: { input: 100, output: 20, total: 120 }, }) const result = await handler.execute(mockContext, mockRouterV2Block, inputs) @@ -1084,16 +1001,10 @@ describe('RouterBlockHandler V2', () => { routes: JSON.stringify([{ id: 'route-1', title: 'Route 1', value: 'Description' }]), } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: 'route-1', - model: 'gpt-4o', - tokens: { input: 100, output: 5, total: 105 }, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: 'route-1', + model: 'gpt-4o', + tokens: { input: 100, output: 5, total: 105 }, }) const result = await handler.execute(mockContext, mockRouterV2Block, inputs) @@ -1111,14 +1022,10 @@ describe('RouterBlockHandler V2', () => { routes: [{ id: 'route-1', title: 'Route 1', value: 'Description' }], } - mockFetch.mockResolvedValueOnce({ - ok: true, - json: () => - Promise.resolve({ - content, - model: 'gpt-4o', - tokens: { input: 100, output: 5, total: 105 }, - }), + mockExecuteProviderRequest.mockResolvedValueOnce({ + content, + model: 'gpt-4o', + tokens: { input: 100, output: 5, total: 105 }, }) await expect(handler.execute(mockContext, mockRouterV2Block, inputs)).rejects.toThrow(content) diff --git a/apps/sim/executor/handlers/router/router-handler.ts b/apps/sim/executor/handlers/router/router-handler.ts index 11e5445889a..3292e2c57c2 100644 --- a/apps/sim/executor/handlers/router/router-handler.ts +++ b/apps/sim/executor/handlers/router/router-handler.ts @@ -1,11 +1,5 @@ import { createLogger } from '@sim/logger' -import { getInternalApiBaseUrl } from '@/lib/core/utils/urls' -import { - addModelInputProvenanceToRequest, - createModelInputProvenanceRequestMetadata, - markModelInputProjected, - projectResolvedModelInput, -} from '@/lib/execution/model-input-provenance' +import { projectResolvedModelInput } from '@/lib/execution/model-input-provenance' import { type AutoRoutingResult, addAutoRoutingCost, @@ -23,7 +17,7 @@ import { ROUTER, } from '@/executor/constants' import type { BlockHandler, ExecutionContext } from '@/executor/types' -import { buildAuthHeaders } from '@/executor/utils/http' +import { executeBlockProviderRequest } from '@/executor/utils/provider-request' import { refuseResolvedSecretProjection } from '@/executor/utils/resolved-secret-projection-refusal' import type { ResolvedSecretInputPath } from '@/executor/utils/resolved-secret-trace-registry' import { resolveVertexCredential } from '@/executor/utils/vertex-credential' @@ -103,9 +97,6 @@ export class RouterBlockHandler implements BlockHandler { } try { - const url = new URL('/api/providers', getInternalApiBaseUrl()) - if (ctx.userId) url.searchParams.set('userId', ctx.userId) - const messages = [{ role: 'user', content: routerConfig.prompt }] const systemPrompt = generateRouterPrompt(routerConfig.prompt, targetBlocks) const resolved = await this.resolveModel( @@ -147,36 +138,13 @@ export class RouterBlockHandler implements BlockHandler { workspaceId: ctx.workspaceId, } - const headers = new Headers(await buildAuthHeaders(ctx.userId)) - const modelInputMetadata = createModelInputProvenanceRequestMetadata( - modelInputProjection.registry, - promptModelInputPaths - ) - const requestBody = addModelInputProvenanceToRequest( - { provider: providerId, ...providerRequest }, - headers, - modelInputMetadata - ) - if (modelInputMetadata) markModelInputProjected(headers) - const response = await fetch(url.toString(), { - method: 'POST', - headers, - body: JSON.stringify(requestBody), + const result = await executeBlockProviderRequest({ + ctx, + providerId, + request: providerRequest, + resolvedSecretTraceRegistry: modelInputProjection.registry, }) - if (!response.ok) { - let errorMessage = `Provider API request failed with status ${response.status}` - try { - const errorData = await response.json() - if (errorData.error) { - errorMessage = errorData.error - } - } catch (_e) {} - throw new Error(errorMessage) - } - - const result = await response.json() - const chosenBlockId = result.content.trim().toLowerCase() const chosenBlock = targetBlocks?.find((b) => b.id === chosenBlockId) @@ -291,9 +259,6 @@ export class RouterBlockHandler implements BlockHandler { } try { - const url = new URL('/api/providers', getInternalApiBaseUrl()) - if (ctx.userId) url.searchParams.set('userId', ctx.userId) - const messages = [{ role: 'user', content: routerConfig.context }] const systemPrompt = generateRouterV2Prompt(routerConfig.context, modelRoutes) const resolved = await this.resolveModel( @@ -354,36 +319,13 @@ export class RouterBlockHandler implements BlockHandler { }, } - const headers = new Headers(await buildAuthHeaders(ctx.userId)) - const modelInputMetadata = createModelInputProvenanceRequestMetadata( - modelInputProjection.registry, - modelInputPaths - ) - const requestBody = addModelInputProvenanceToRequest( - { provider: providerId, ...providerRequest }, - headers, - modelInputMetadata - ) - if (modelInputMetadata) markModelInputProjected(headers) - const response = await fetch(url.toString(), { - method: 'POST', - headers, - body: JSON.stringify(requestBody), + const result = await executeBlockProviderRequest({ + ctx, + providerId, + request: providerRequest, + resolvedSecretTraceRegistry: modelInputProjection.registry, }) - if (!response.ok) { - let errorMessage = `Provider API request failed with status ${response.status}` - try { - const errorData = await response.json() - if (errorData.error) { - errorMessage = errorData.error - } - } catch (_e) {} - throw new Error(errorMessage) - } - - const result = await response.json() - let chosenRouteId: string let reasoning = '' diff --git a/apps/sim/executor/utils/provider-request.ts b/apps/sim/executor/utils/provider-request.ts new file mode 100644 index 00000000000..e13543a0f21 --- /dev/null +++ b/apps/sim/executor/utils/provider-request.ts @@ -0,0 +1,76 @@ +import { createLogger } from '@sim/logger' +import { checkWorkspaceAccess } from '@/lib/workspaces/permissions/utils' +import type { ExecutionContext } from '@/executor/types' +import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' +import { executeProviderRequest } from '@/providers' +import type { ProviderRequest, ProviderResponse } from '@/providers/types' + +const logger = createLogger('ExecutorProviderRequest') + +export interface ExecuteBlockProviderRequestInput { + ctx: ExecutionContext + providerId: string + request: ProviderRequest + /** + * The fork the block's model input was projected through. Supplied to the + * provider runtime in place of the provenance envelope the HTTP boundary used + * to serialize and re-import. + */ + resolvedSecretTraceRegistry: ResolvedSecretTraceRegistry | undefined +} + +/** + * Runs one non-streaming provider request for a block handler in-process. + * + * Replaces the executor's `POST /api/providers` round trip, which re-derived + * everything it needed from claims the executor had itself just supplied. The + * two admission checks the route owned are reproduced here so the outcome is + * unchanged: + * + * - `checkInternalAuth` rejected a token carrying no user. The executor mints + * that token from `ctx.userId`, so the check reduces to requiring one. + * - `checkWorkspaceAccess` rejected an execution subject who is no longer a + * member of the workspace being billed. + * + * The route's remaining work is either already done by the caller (the model + * permission policy, via `validateModelProvider`; Vertex credential + * authorization, via `resolveVertexCredential`) or lives inside + * `executeProviderRequest` itself (BYOK key resolution, attachment provenance + * filtering, cost policy). + */ +export async function executeBlockProviderRequest({ + ctx, + providerId, + request, + resolvedSecretTraceRegistry, +}: ExecuteBlockProviderRequestInput): Promise { + if (!ctx.userId) { + throw new Error('Unauthorized') + } + + if (request.workspaceId) { + const workspaceAccess = await checkWorkspaceAccess(request.workspaceId, ctx.userId) + if (!workspaceAccess.hasAccess) { + throw new Error('Forbidden') + } + } + + /** + * `executionContext` is deliberately not supplied: it is only inherited by + * model-emitted tool calls, and the route this replaces never carried one. + * Router and evaluator requests declare no tools, so passing the executor's + * context here would widen the trusted surface without changing any outcome. + */ + const response = await executeProviderRequest( + providerId, + { ...request, userId: ctx.userId }, + { resolvedSecretTraceRegistry } + ) + + if (response instanceof ReadableStream || (response !== null && 'stream' in response)) { + logger.error('Provider returned a stream for a non-streaming block request', { providerId }) + throw new Error('Provider returned a streaming response for a non-streaming request') + } + + return response +} From d3561e7868ff200f0b5b2a75a99737d1cbef7b32 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 13 Aug 2026 01:51:38 -0700 Subject: [PATCH 11/14] improvement(queries): fix a second row-cache collision, and make the memoized workspace read actually dedupe --- apps/sim/app/api/auth/oauth/token/route.ts | 2 +- .../copilot/checkpoints/revert/route.test.ts | 16 +++++-- .../api/copilot/checkpoints/revert/route.ts | 32 +++---------- .../sim/app/api/workflows/[id]/state/route.ts | 5 -- apps/sim/app/invite/[id]/invite.tsx | 2 +- .../[workspaceId]/lib/prefetch.test.ts | 19 +++----- .../app/workspace/[workspaceId]/prefetch.ts | 18 +++---- .../upgrade/hooks/use-upgrade-state.ts | 10 ++-- .../w/[workflowId]/hooks/use-wand.ts | 2 +- .../hooks/use-workflow-execution.ts | 2 +- apps/sim/executor/utils/provider-request.ts | 2 +- apps/sim/hooks/queries/invitations.ts | 2 +- apps/sim/hooks/queries/organization.ts | 2 +- apps/sim/hooks/queries/subscription.ts | 16 +------ apps/sim/hooks/queries/tables.test.ts | 29 ++++++------ apps/sim/hooks/queries/tables.ts | 2 +- .../hooks/queries/utils/invalidate-usage.ts | 35 ++++++++++++++ .../hooks/queries/utils/subscription-keys.ts | 16 +++++++ apps/sim/hooks/queries/utils/table-keys.ts | 8 ++-- .../queries/utils/workspace-usage-keys.ts | 14 ++++++ apps/sim/hooks/queries/workspace-files.ts | 2 +- .../sim/hooks/queries/workspace-usage.test.ts | 4 +- apps/sim/hooks/queries/workspace-usage.ts | 47 +------------------ .../sim/hooks/selectors/use-selector-query.ts | 3 +- .../lib/api/contracts/oauth-connections.ts | 10 +++- apps/sim/lib/oauth/token-resolution.ts | 32 +++++-------- .../persistence/save-normalized-state.ts | 46 +++++------------- apps/sim/lib/workspace-files/queries.test.ts | 24 ++++++++++ apps/sim/lib/workspace-files/queries.ts | 29 +++++++----- apps/sim/lib/workspaces/permissions/utils.ts | 22 ++++++--- 30 files changed, 232 insertions(+), 221 deletions(-) create mode 100644 apps/sim/hooks/queries/utils/invalidate-usage.ts create mode 100644 apps/sim/hooks/queries/utils/subscription-keys.ts create mode 100644 apps/sim/hooks/queries/utils/workspace-usage-keys.ts diff --git a/apps/sim/app/api/auth/oauth/token/route.ts b/apps/sim/app/api/auth/oauth/token/route.ts index 0bf50ac846d..c3e1744dc1f 100644 --- a/apps/sim/app/api/auth/oauth/token/route.ts +++ b/apps/sim/app/api/auth/oauth/token/route.ts @@ -118,7 +118,7 @@ export const POST = withRouteHandler(async (request: NextRequest) => { const auth = await checkSessionOrInternalAuth(request, { requireWorkflowId: false }) const result = await resolveCredentialToken(auth, { requestId, - credentialId: credentialId ?? '', + credentialId, workflowId: workflowId ?? undefined, scopes, impersonateEmail, diff --git a/apps/sim/app/api/copilot/checkpoints/revert/route.test.ts b/apps/sim/app/api/copilot/checkpoints/revert/route.test.ts index 67eaf1a94c4..621256b8ac3 100644 --- a/apps/sim/app/api/copilot/checkpoints/revert/route.test.ts +++ b/apps/sim/app/api/copilot/checkpoints/revert/route.test.ts @@ -49,9 +49,11 @@ describe('Copilot Checkpoints Revert API Route', () => { authMockFns.mockGetSession.mockResolvedValue(null) + /** Authorization is the route's workflow read, so an allowed result always carries one. */ workflowAuthzMockFns.mockAuthorizeWorkflowByWorkspacePermission.mockResolvedValue({ allowed: true, status: 200, + workflow: { id: 'b2c3d4e5-f6a7-4b89-a0d1-e2f3a4b5c6d7', workspaceId: 'ws-123' }, }) mockGetAccessibleCopilotChat.mockResolvedValue({ id: 'chat-123', userId: 'user-123' }) @@ -201,7 +203,12 @@ describe('Copilot Checkpoints Revert API Route', () => { } queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) - queueTableRows(schemaMock.workflow, []) + /** Authorization performs the workflow read, so a missing workflow surfaces through it. */ + workflowAuthzMockFns.mockAuthorizeWorkflowByWorkspacePermission.mockResolvedValueOnce({ + allowed: false, + status: 404, + workflow: null, + }) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', { method: 'POST', @@ -214,6 +221,7 @@ describe('Copilot Checkpoints Revert API Route', () => { expect(response.status).toBe(404) const responseData = await response.json() expect(responseData.error).toBe('Workflow not found') + expect(mockSaveWorkflowNormalizedState).not.toHaveBeenCalled() }) it('should return 401 when workflow belongs to different user', async () => { @@ -237,6 +245,7 @@ describe('Copilot Checkpoints Revert API Route', () => { workflowAuthzMockFns.mockAuthorizeWorkflowByWorkspacePermission.mockResolvedValueOnce({ allowed: false, status: 403, + workflow: { id: 'b2c3d4e5-f6a7-4b89-a0d1-e2f3a4b5c6d7', workspaceId: 'ws-123' }, }) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', { @@ -562,8 +571,9 @@ describe('Copilot Checkpoints Revert API Route', () => { } dbChainMockFns.where.mockReturnValueOnce(Promise.resolve([mockCheckpoint])) - dbChainMockFns.where.mockReturnValueOnce( - Promise.reject(new Error('Database error during workflow lookup')) + /** Authorization performs the workflow read, so a failed lookup surfaces through it. */ + workflowAuthzMockFns.mockAuthorizeWorkflowByWorkspacePermission.mockRejectedValueOnce( + new Error('Database error during workflow lookup') ) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', { diff --git a/apps/sim/app/api/copilot/checkpoints/revert/route.ts b/apps/sim/app/api/copilot/checkpoints/revert/route.ts index 6b49fb8823e..e022c0a3d57 100644 --- a/apps/sim/app/api/copilot/checkpoints/revert/route.ts +++ b/apps/sim/app/api/copilot/checkpoints/revert/route.ts @@ -1,10 +1,7 @@ import { db } from '@sim/db' -import { workflowCheckpoints, workflow as workflowTable } from '@sim/db/schema' +import { workflowCheckpoints } from '@sim/db/schema' import { createLogger } from '@sim/logger' -import { - authorizeWorkflowByWorkspacePermission, - WorkflowLockedError, -} from '@sim/platform-authz/workflow' +import { authorizeWorkflowByWorkspacePermission } from '@sim/platform-authz/workflow' import { and, eq } from 'drizzle-orm' import { type NextRequest, NextResponse } from 'next/server' import { revertCopilotCheckpointContract } from '@/lib/api/contracts/copilot' @@ -68,21 +65,15 @@ export const POST = withRouteHandler(async (request: NextRequest) => { return createNotFoundResponse('Checkpoint not found or access denied') } - const workflowData = await db - .select() - .from(workflowTable) - .where(eq(workflowTable.id, checkpoint.workflowId)) - .then((rows) => rows[0]) - - if (!workflowData) { - return createNotFoundResponse('Workflow not found') - } - + /** Authorization already loads the workflow, so its absence is the not-found signal. */ const authorization = await authorizeWorkflowByWorkspacePermission({ workflowId: checkpoint.workflowId, userId, action: 'write', }) + if (!authorization.workflow) { + return createNotFoundResponse('Workflow not found') + } if (!authorization.allowed) { return createUnauthorizedResponse() } @@ -144,22 +135,11 @@ export const POST = withRouteHandler(async (request: NextRequest) => { ) } - /** - * A locked workflow used to surface here as a non-OK PUT response, so it - * still resolves to the same revert failure rather than the generic - * outer-catch message. Every other throw keeps propagating, matching the - * old transport-error path. - */ const saveResult = await saveWorkflowNormalizedState({ requestId: tracker.requestId, workflowId: checkpoint.workflowId, userId, state: parsedState.data, - }).catch((error) => { - if (error instanceof WorkflowLockedError) { - return { success: false as const, status: error.status, error: error.message } - } - throw error }) if (!saveResult.success) { diff --git a/apps/sim/app/api/workflows/[id]/state/route.ts b/apps/sim/app/api/workflows/[id]/state/route.ts index 97a32c40b86..999129410cf 100644 --- a/apps/sim/app/api/workflows/[id]/state/route.ts +++ b/apps/sim/app/api/workflows/[id]/state/route.ts @@ -110,11 +110,6 @@ export const PUT = withRouteHandler( const parsed = await parseRequest(putWorkflowNormalizedStateContract, request, context) if (!parsed.success) return parsed.response - // Note: prior versions cross-checked that each variable's `workflowId` - // equalled the path param. The write contract does not carry `workflowId` - // per variable (the path param is the source of truth), so the check - // is unreachable and was removed. - const result = await saveWorkflowNormalizedState({ requestId, workflowId, diff --git a/apps/sim/app/invite/[id]/invite.tsx b/apps/sim/app/invite/[id]/invite.tsx index 232dfbdb884..ffc75e40866 100644 --- a/apps/sim/app/invite/[id]/invite.tsx +++ b/apps/sim/app/invite/[id]/invite.tsx @@ -15,7 +15,7 @@ import { InviteLayout, InviteStatusCard } from '@/app/invite/components' import { useInvitationDetails } from '@/hooks/queries/invitations' import { organizationKeys } from '@/hooks/queries/organization' import { refreshSessionQuery } from '@/hooks/queries/session' -import { subscriptionKeys } from '@/hooks/queries/subscription' +import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys' import { workspaceKeys } from '@/hooks/queries/workspace' const logger = createLogger('InviteById') diff --git a/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts b/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts index c9e9e242303..c506d873673 100644 --- a/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts @@ -515,8 +515,8 @@ describe('workspace list prefetches', () => { /** * The file list is seeded on every workspace route, so it is the one entry whose size - * scales with a workspace's content on routes that never read it. It is read one row - * past the budget so the overflow is detectable. + * scales with a workspace's content on routes that never read it. The budget is passed + * down rather than applied here, so the read can stop before the share join. */ it('seeds the file list, bounded by the document payload budget', async () => { const files = [{ id: 'file-1', name: 'a.txt' }] @@ -526,24 +526,19 @@ describe('workspace list prefetches', () => { await prefetchWorkspaceSidebar(client, WORKSPACE_ID, USER_ID, HOST_CONTEXT, null) expect(mockListWorkspaceFilesWithShares).toHaveBeenCalledWith(WORKSPACE_ID, 'active', { - limit: WORKSPACE_FILE_SEED_MAX + 1, + maxRows: WORKSPACE_FILE_SEED_MAX, }) expect(client.getQueryData(workspaceFilesKeys.list(WORKSPACE_ID, 'active'))).toEqual(files) }) /** * The load-bearing half of the budget: a workspace over it seeds NOTHING rather than the - * prefix it read. The sidebar search filters this list client-side and the Files browser - * renders it as the workspace's files, so a truncated seed would silently hide files — - * the client fetch must reach the route for the complete list instead. + * prefix that was read. The sidebar search filters this list client-side and the Files + * browser renders it as the workspace's files, so a truncated seed would silently hide + * files — the client fetch must reach the route for the complete list instead. */ it('seeds nothing when the workspace exceeds the budget', async () => { - mockListWorkspaceFilesWithShares.mockResolvedValue( - Array.from({ length: WORKSPACE_FILE_SEED_MAX + 1 }, (_, index) => ({ - id: `file-${index}`, - name: `${index}.txt`, - })) - ) + mockListWorkspaceFilesWithShares.mockResolvedValue(null) const client = makeClient() await prefetchWorkspaceSidebar(client, WORKSPACE_ID, USER_ID, HOST_CONTEXT, null) diff --git a/apps/sim/app/workspace/[workspaceId]/prefetch.ts b/apps/sim/app/workspace/[workspaceId]/prefetch.ts index 0ed4a3cba47..c5a53c54aa0 100644 --- a/apps/sim/app/workspace/[workspaceId]/prefetch.ts +++ b/apps/sim/app/workspace/[workspaceId]/prefetch.ts @@ -101,11 +101,11 @@ async function seedWorkspaceList( * thousands of files would otherwise push more than a megabyte of HTML ahead of first * paint on the logs, settings, and editor routes that never read it. * - * The read is capped at one row past the budget purely to detect the overflow. A - * workspace above it seeds NOTHING rather than a prefix: the sidebar search filters this - * list client-side and the Files browser renders it as the workspace's files, so a - * truncated seed would silently hide files. Those workspaces fetch the complete list - * from the route instead — which for a list that large is also the cheaper first paint. + * A workspace above the budget seeds NOTHING rather than a prefix: the sidebar search + * filters this list client-side and the Files browser renders it as the workspace's + * files, so a truncated seed would silently hide files. Those workspaces fetch the + * complete list from the route instead — which for a list that large is also the + * cheaper first paint. */ export const WORKSPACE_FILE_SEED_MAX = 300 @@ -123,15 +123,15 @@ export const WORKSPACE_FILE_SEED_MAX = 300 * workspace exceeds {@link WORKSPACE_FILE_SEED_MAX}: `prefetchQuery` always creates one, * and a partial one would be read as the whole list. * - * The shape comes from the same contract-parsed reader `GET /api/workspaces/[id]/files` - * responds with, so a seeded entry is identical to what the client hook would cache. + * Parsed through the same response contract `GET /api/workspaces/[id]/files` validates + * against, so a seeded entry is identical to what the client hook would cache. */ async function seedWorkspaceFiles(queryClient: QueryClient, workspaceId: string): Promise { try { const files = await listWorkspaceFilesWithShares(workspaceId, 'active', { - limit: WORKSPACE_FILE_SEED_MAX + 1, + maxRows: WORKSPACE_FILE_SEED_MAX, }) - if (files.length > WORKSPACE_FILE_SEED_MAX) return + if (!files) return queryClient.setQueryData(workspaceFilesKeys.list(workspaceId, 'active'), files) } catch (error) { /** Optimization only: the client fetch reaches the route instead. Logged so drift between diff --git a/apps/sim/app/workspace/[workspaceId]/upgrade/hooks/use-upgrade-state.ts b/apps/sim/app/workspace/[workspaceId]/upgrade/hooks/use-upgrade-state.ts index e0a91ba6477..15152b4a57b 100644 --- a/apps/sim/app/workspace/[workspaceId]/upgrade/hooks/use-upgrade-state.ts +++ b/apps/sim/app/workspace/[workspaceId]/upgrade/hooks/use-upgrade-state.ts @@ -9,9 +9,9 @@ import type { WorkspaceHostContext } from '@/lib/api/contracts/workspaces' import { useSubscriptionUpgrade } from '@/lib/billing/client/upgrade' import { CREDIT_TIERS } from '@/lib/billing/constants' import { getPlanTierCredits, isEnterprise, isFree, isPro, isTeam } from '@/lib/billing/plan-helpers' -import { subscriptionKeys } from '@/hooks/queries/subscription' +import { invalidateWorkspaceUsage } from '@/hooks/queries/utils/invalidate-usage' +import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys' import { workspaceHostKeys } from '@/hooks/queries/workspace-host' -import { invalidateWorkspaceUsage } from '@/hooks/queries/workspace-usage' const PRO_TIER = CREDIT_TIERS[0] const MAX_TIER = CREDIT_TIERS[1] @@ -94,8 +94,9 @@ export function useUpgradeState({ /** * A non-redirect plan switch settles server-side immediately, so every read that * describes the plan has to be refetched — the host context the page renders from, - * the subscription/usage reads the billing surfaces share, and the workspace credit - * availability that drives the credits chip and the run gate. + * the subscription/usage reads the billing surfaces share, the proration invoice the + * switch just produced, and the workspace credit availability that drives the credits + * chip and the run gate. */ const refreshBillingState = useCallback( () => @@ -103,6 +104,7 @@ export function useUpgradeState({ queryClient.invalidateQueries({ queryKey: workspaceHostKeys.detail(workspaceId) }), queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }), queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() }), + queryClient.invalidateQueries({ queryKey: subscriptionKeys.invoicesAll() }), invalidateWorkspaceUsage(queryClient), ]), [queryClient, workspaceId] diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts index a17c8700934..cb6573d6db0 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts +++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts @@ -10,7 +10,7 @@ import { wandGenerateStreamContract } from '@/lib/api/contracts' import { readSSEStream } from '@/lib/core/utils/sse' import { shouldStripCodeFences, stripCodeFences } from '@/lib/wand/strip-code-fences' import type { GenerationType } from '@/blocks/types' -import { scheduleUsageRefresh } from '@/hooks/queries/workspace-usage' +import { scheduleUsageRefresh } from '@/hooks/queries/utils/invalidate-usage' import { useSettingsNavigation } from '@/hooks/use-settings-navigation' import { useWorkflowRegistry } from '@/stores/workflows/registry/store' diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts index bfeabf5466d..97fb6a3f507 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts +++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts @@ -61,8 +61,8 @@ import type { SerializableExecutionState } from '@/executor/execution/types' import type { BlockLog, BlockState, ExecutionResult, StreamingExecution } from '@/executor/types' import { hasExecutionResult } from '@/executor/utils/errors' import { coerceValue } from '@/executor/utils/start-block' +import { scheduleUsageRefresh } from '@/hooks/queries/utils/invalidate-usage' import { getWorkflows } from '@/hooks/queries/utils/workflow-cache' -import { scheduleUsageRefresh } from '@/hooks/queries/workspace-usage' import { isExecutionStreamHttpError, SSEEventHandlerError, diff --git a/apps/sim/executor/utils/provider-request.ts b/apps/sim/executor/utils/provider-request.ts index e13543a0f21..d6367de56d0 100644 --- a/apps/sim/executor/utils/provider-request.ts +++ b/apps/sim/executor/utils/provider-request.ts @@ -7,7 +7,7 @@ import type { ProviderRequest, ProviderResponse } from '@/providers/types' const logger = createLogger('ExecutorProviderRequest') -export interface ExecuteBlockProviderRequestInput { +interface ExecuteBlockProviderRequestInput { ctx: ExecutionContext providerId: string request: ProviderRequest diff --git a/apps/sim/hooks/queries/invitations.ts b/apps/sim/hooks/queries/invitations.ts index 394ff0b9113..3f93f6638ab 100644 --- a/apps/sim/hooks/queries/invitations.ts +++ b/apps/sim/hooks/queries/invitations.ts @@ -22,8 +22,8 @@ import { import { updateWorkspacePermissionsContract } from '@/lib/api/contracts/workspaces' import { organizationKeys } from '@/hooks/queries/organization' import { refreshSessionQuery } from '@/hooks/queries/session' -import { subscriptionKeys } from '@/hooks/queries/subscription' import { workspaceCredentialKeys } from '@/hooks/queries/utils/credential-keys' +import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys' import { workspaceKeys } from '@/hooks/queries/workspace' export const invitationKeys = { diff --git a/apps/sim/hooks/queries/organization.ts b/apps/sim/hooks/queries/organization.ts index 0845c99f71d..a0ce1295633 100644 --- a/apps/sim/hooks/queries/organization.ts +++ b/apps/sim/hooks/queries/organization.ts @@ -41,8 +41,8 @@ import { import { client } from '@/lib/auth/auth-client' import { isEnterprise, isPaid, isTeam } from '@/lib/billing/plan-helpers' import { hasPaidSubscriptionStatus } from '@/lib/billing/subscriptions/utils' -import { subscriptionKeys } from '@/hooks/queries/subscription' import { workspaceCredentialKeys } from '@/hooks/queries/utils/credential-keys' +import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys' import { workspaceKeys } from '@/hooks/queries/workspace' const logger = createLogger('OrganizationQueries') diff --git a/apps/sim/hooks/queries/subscription.ts b/apps/sim/hooks/queries/subscription.ts index 946f0a0f14d..5809f6b2eba 100644 --- a/apps/sim/hooks/queries/subscription.ts +++ b/apps/sim/hooks/queries/subscription.ts @@ -13,8 +13,9 @@ import { updateUsageLimitContract, } from '@/lib/api/contracts/subscription' import { organizationKeys } from '@/hooks/queries/organization' +import { invalidateWorkspaceUsage } from '@/hooks/queries/utils/invalidate-usage' +import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys' import { workspaceKeys } from '@/hooks/queries/workspace' -import { invalidateWorkspaceUsage } from '@/hooks/queries/workspace-usage' export type { SubscriptionApiResponse } @@ -22,19 +23,6 @@ export const SUBSCRIPTION_DATA_STALE_TIME = 5 * 60 * 1000 export const USAGE_LIMIT_STALE_TIME = 30 * 1000 export const INVOICES_STALE_TIME = 5 * 60 * 1000 -/** - * Query key factories for subscription-related queries - */ -export const subscriptionKeys = { - all: ['subscription'] as const, - users: () => [...subscriptionKeys.all, 'user'] as const, - user: (includeOrg?: boolean) => [...subscriptionKeys.users(), { includeOrg }] as const, - usage: () => [...subscriptionKeys.all, 'usage'] as const, - invoicesAll: () => [...subscriptionKeys.all, 'invoices'] as const, - invoices: (context: 'user' | 'organization' = 'user', organizationId?: string) => - [...subscriptionKeys.invoicesAll(), context, organizationId ?? ''] as const, -} - /** * Fetch user subscription data * @param includeOrg - Whether to include organization role data diff --git a/apps/sim/hooks/queries/tables.test.ts b/apps/sim/hooks/queries/tables.test.ts index 876fe2a981e..c20ab37b23d 100644 --- a/apps/sim/hooks/queries/tables.test.ts +++ b/apps/sim/hooks/queries/tables.test.ts @@ -69,6 +69,13 @@ import { tableKeys } from '@/hooks/queries/utils/table-keys' const TABLE_ID = 'tbl-1' const WORKSPACE_ID = 'ws-1' +/** + * Where a paged row list actually lives. Seeding at the bare `rowsRoot` prefix would + * exercise a key no hook writes, and would keep matching a cache walk that has been + * narrowed away from the `find` sibling hanging off the same parent. + */ +const ROWS_KEY = tableKeys.infiniteRows(TABLE_ID, tableRowsParamsKey({ pageSize: 1000 })) + function setCache(key: readonly unknown[], value: unknown) { cacheStore.set(JSON.stringify(key), value) } @@ -96,7 +103,7 @@ describe('useDeleteColumn optimistic update', () => { columnWidths: { name: 200, age: 100 }, }, }) - setCache(tableKeys.rowsRoot(TABLE_ID), { + setCache(ROWS_KEY, { rows: [ { id: 'r1', data: { name: 'a', age: 1 } }, { id: 'r2', data: { name: 'b', age: 2 } }, @@ -114,9 +121,7 @@ describe('useDeleteColumn optimistic update', () => { expect(detail?.schema.columns.map((c) => c.name)).toEqual(['name']) expect(detail?.metadata.columnWidths).toEqual({ name: 200 }) - const rows = getCache<{ rows: Array<{ data: Record }> }>( - tableKeys.rowsRoot(TABLE_ID) - ) + const rows = getCache<{ rows: Array<{ data: Record }> }>(ROWS_KEY) expect(rows?.rows.every((r) => !('age' in r.data))).toBe(true) expect(rows?.rows[0]?.data).toEqual({ name: 'a' }) @@ -135,7 +140,7 @@ describe('useDeleteColumn optimistic update', () => { totalCount: 1, } setCache(tableKeys.detail(TABLE_ID), originalDetail) - setCache(tableKeys.rowsRoot(TABLE_ID), originalRows) + setCache(ROWS_KEY, originalRows) const hook = useDeleteColumn({ workspaceId: WORKSPACE_ID, tableId: TABLE_ID }) const ctx = await hook.onMutate?.('age') @@ -145,7 +150,7 @@ describe('useDeleteColumn optimistic update', () => { hook.onError?.(new Error('boom'), 'age', ctx) expect(getCache(tableKeys.detail(TABLE_ID))).toEqual(originalDetail) - expect(getCache(tableKeys.rowsRoot(TABLE_ID))).toEqual(originalRows) + expect(getCache(ROWS_KEY)).toEqual(originalRows) }) it('invalidates schema, rows, and lists in onSettled', () => { @@ -194,7 +199,7 @@ describe('useUpdateColumn optimistic update', () => { id: TABLE_ID, schema: { columns: [{ name: 'age', type: 'number' }] }, }) - setCache(tableKeys.rowsRoot(TABLE_ID), { + setCache(ROWS_KEY, { rows: [ { id: 'r1', data: { age: 30 } }, { id: 'r2', data: { age: 40 } }, @@ -207,9 +212,7 @@ describe('useUpdateColumn optimistic update', () => { // Row data is id-keyed; a rename never moves it. The stored key (`age`) // becomes the column's stamped id, so cells stay reachable via getColumnId. - const rows = getCache<{ rows: Array<{ data: Record }> }>( - tableKeys.rowsRoot(TABLE_ID) - ) + const rows = getCache<{ rows: Array<{ data: Record }> }>(ROWS_KEY) expect(rows?.rows[0]?.data).toEqual({ age: 30 }) expect(rows?.rows[1]?.data).toEqual({ age: 40 }) @@ -265,7 +268,7 @@ describe('useDeleteColumn case-insensitive row cleanup', () => { id: TABLE_ID, schema: { columns: [{ name: 'Age', type: 'number' }] }, }) - setCache(tableKeys.rowsRoot(TABLE_ID), { + setCache(ROWS_KEY, { rows: [{ id: 'r1', data: { Age: 30, name: 'a' } }], totalCount: 1, }) @@ -273,9 +276,7 @@ describe('useDeleteColumn case-insensitive row cleanup', () => { const hook = useDeleteColumn({ workspaceId: WORKSPACE_ID, tableId: TABLE_ID }) await hook.onMutate?.('age') - const rows = getCache<{ rows: Array<{ data: Record }> }>( - tableKeys.rowsRoot(TABLE_ID) - ) + const rows = getCache<{ rows: Array<{ data: Record }> }>(ROWS_KEY) expect(rows?.rows[0]?.data).toEqual({ name: 'a' }) }) }) diff --git a/apps/sim/hooks/queries/tables.ts b/apps/sim/hooks/queries/tables.ts index 6c9bfc564a1..4f41554f9d6 100644 --- a/apps/sim/hooks/queries/tables.ts +++ b/apps/sim/hooks/queries/tables.ts @@ -2196,7 +2196,7 @@ export async function snapshotAndMutateRows( ): Promise { const scope = options?.onlyKey ? ({ queryKey: options.onlyKey, exact: true } as const) - : ({ queryKey: tableKeys.rowsRoot(tableId) } as const) + : ({ queryKey: tableKeys.infiniteRowsRoot(tableId) } as const) if (options?.cancelInFlight !== false) { await queryClient.cancelQueries(scope) } diff --git a/apps/sim/hooks/queries/utils/invalidate-usage.ts b/apps/sim/hooks/queries/utils/invalidate-usage.ts new file mode 100644 index 00000000000..164451b3856 --- /dev/null +++ b/apps/sim/hooks/queries/utils/invalidate-usage.ts @@ -0,0 +1,35 @@ +import type { QueryClient } from '@tanstack/react-query' +import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys' +import { workspaceUsageKeys } from '@/hooks/queries/utils/workspace-usage-keys' + +/** + * Usage is written asynchronously as a run settles, so a refetch fired on completion + * races the write and re-reads the old balance. + */ +const USAGE_SETTLE_DELAY_MS = 1000 + +/** + * Invalidates the workspace credit/usage reads after anything that moves the balance — + * a run that spends credits, a top-up, a plan change, or a usage-limit edit. Both + * families are keyed per workspace but derive from the same billing account, so the + * family prefixes (not a single workspace's key) are what has to be refetched. + */ +export function invalidateWorkspaceUsage(queryClient: QueryClient) { + return Promise.all([ + queryClient.invalidateQueries({ queryKey: workspaceUsageKeys.creditAvailabilities() }), + queryClient.invalidateQueries({ queryKey: workspaceUsageKeys.gates() }), + ]) +} + +/** + * Refreshes the billing reads a run touches, after {@link USAGE_SETTLE_DELAY_MS}. + * + * Shared by the surfaces that spend credits — workflow execution and wand generation — + * so the delay and the key set stay in one place. + */ +export function scheduleUsageRefresh(queryClient: QueryClient) { + setTimeout(() => { + void queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }) + void invalidateWorkspaceUsage(queryClient) + }, USAGE_SETTLE_DELAY_MS) +} diff --git a/apps/sim/hooks/queries/utils/subscription-keys.ts b/apps/sim/hooks/queries/utils/subscription-keys.ts new file mode 100644 index 00000000000..0d6aa6af6b4 --- /dev/null +++ b/apps/sim/hooks/queries/utils/subscription-keys.ts @@ -0,0 +1,16 @@ +/** + * React Query key factory for subscription and billing reads. + * + * Lives in this standalone module — like {@link file://./workspace-usage-keys.ts} — so + * the shared billing invalidations can reference it without importing the hook module + * that consumes those invalidations, which would close an import cycle between the two. + */ +export const subscriptionKeys = { + all: ['subscription'] as const, + users: () => [...subscriptionKeys.all, 'user'] as const, + user: (includeOrg?: boolean) => [...subscriptionKeys.users(), { includeOrg }] as const, + usage: () => [...subscriptionKeys.all, 'usage'] as const, + invoicesAll: () => [...subscriptionKeys.all, 'invoices'] as const, + invoices: (context: 'user' | 'organization' = 'user', organizationId?: string) => + [...subscriptionKeys.invoicesAll(), context, organizationId ?? ''] as const, +} diff --git a/apps/sim/hooks/queries/utils/table-keys.ts b/apps/sim/hooks/queries/utils/table-keys.ts index 08d9d410329..8394373c66b 100644 --- a/apps/sim/hooks/queries/utils/table-keys.ts +++ b/apps/sim/hooks/queries/utils/table-keys.ts @@ -28,10 +28,10 @@ export const tableKeys = { /** * Prefix covering only the paged row lists. * - * `rowsRoot` is a shared parent — `rowWrites` and `find` hang off it holding - * entirely different shapes — so anything walking the cache to update or snapshot - * row pages must start here instead. Reaching for `rowsRoot` and subtracting the - * siblings is a denylist that rots the moment a fifth subtree is added. + * `rowsRoot` is a shared parent — `find` hangs off it holding an entirely different + * shape — so anything walking the cache to update or snapshot row pages must start + * here instead. Reaching for `rowsRoot` and subtracting the siblings is a denylist + * that rots the moment another subtree is added. */ infiniteRowsRoot: (tableId: string) => [...tableKeys.rowsRoot(tableId), 'infinite'] as const, infiniteRows: (tableId: string, paramsKey: string) => diff --git a/apps/sim/hooks/queries/utils/workspace-usage-keys.ts b/apps/sim/hooks/queries/utils/workspace-usage-keys.ts new file mode 100644 index 00000000000..8063c3564b2 --- /dev/null +++ b/apps/sim/hooks/queries/utils/workspace-usage-keys.ts @@ -0,0 +1,14 @@ +/** + * React Query key factory for the per-workspace credit and usage-gate reads. + * + * Standalone for the same reason as {@link file://./subscription-keys.ts}: the shared + * billing invalidations need these keys without importing the hooks that call them. + */ +export const workspaceUsageKeys = { + all: ['workspace-usage'] as const, + creditAvailabilities: () => [...workspaceUsageKeys.all, 'credit-availability'] as const, + creditAvailability: (workspaceId: string) => + [...workspaceUsageKeys.creditAvailabilities(), workspaceId] as const, + gates: () => [...workspaceUsageKeys.all, 'gate'] as const, + gate: (workspaceId: string) => [...workspaceUsageKeys.gates(), workspaceId] as const, +} diff --git a/apps/sim/hooks/queries/workspace-files.ts b/apps/sim/hooks/queries/workspace-files.ts index c1b0b01ad49..ce118ac837d 100644 --- a/apps/sim/hooks/queries/workspace-files.ts +++ b/apps/sim/hooks/queries/workspace-files.ts @@ -528,7 +528,7 @@ export function useCloudStorageConfigured(enabled = true) { * closed — the upload path treats "unknown" as "not configured", so a single * blip would disable cloud-backed uploads until a full reload. The key is * global, so navigating or switching workspace cannot recover it either. - * Matches {@link useVoiceSettings}, which carries the same three options. + * Mirrors {@link useVoiceSettings}, which overrides `retryOnMount` for the same reason. */ retryOnMount: true, }) diff --git a/apps/sim/hooks/queries/workspace-usage.test.ts b/apps/sim/hooks/queries/workspace-usage.test.ts index c2b7b09117d..011b091ae94 100644 --- a/apps/sim/hooks/queries/workspace-usage.test.ts +++ b/apps/sim/hooks/queries/workspace-usage.test.ts @@ -15,13 +15,13 @@ import { getWorkspaceCreditAvailabilityContract, getWorkspaceUsageGateContract, } from '@/lib/api/contracts/workspaces' +import { invalidateWorkspaceUsage } from '@/hooks/queries/utils/invalidate-usage' +import { workspaceUsageKeys } from '@/hooks/queries/utils/workspace-usage-keys' import { fetchWorkspaceCreditAvailability, fetchWorkspaceUsageGate, - invalidateWorkspaceUsage, WORKSPACE_CREDIT_AVAILABILITY_STALE_TIME, WORKSPACE_USAGE_GATE_STALE_TIME, - workspaceUsageKeys, } from '@/hooks/queries/workspace-usage' describe('workspace usage gate query', () => { diff --git a/apps/sim/hooks/queries/workspace-usage.ts b/apps/sim/hooks/queries/workspace-usage.ts index c23a18a528c..f9f46152fee 100644 --- a/apps/sim/hooks/queries/workspace-usage.ts +++ b/apps/sim/hooks/queries/workspace-usage.ts @@ -1,4 +1,4 @@ -import { type QueryClient, useQuery } from '@tanstack/react-query' +import { useQuery } from '@tanstack/react-query' import { requestJson } from '@/lib/api/client/request' import { getWorkspaceCreditAvailabilityContract, @@ -6,39 +6,11 @@ import { type WorkspaceCreditAvailability, type WorkspaceUsageGate, } from '@/lib/api/contracts/workspaces' -import { subscriptionKeys } from '@/hooks/queries/subscription' - -export const workspaceUsageKeys = { - all: ['workspace-usage'] as const, - creditAvailabilities: () => [...workspaceUsageKeys.all, 'credit-availability'] as const, - creditAvailability: (workspaceId: string) => - [...workspaceUsageKeys.creditAvailabilities(), workspaceId] as const, - gates: () => [...workspaceUsageKeys.all, 'gate'] as const, - gate: (workspaceId: string) => [...workspaceUsageKeys.gates(), workspaceId] as const, -} +import { workspaceUsageKeys } from '@/hooks/queries/utils/workspace-usage-keys' export const WORKSPACE_CREDIT_AVAILABILITY_STALE_TIME = 30 * 1000 export const WORKSPACE_USAGE_GATE_STALE_TIME = 30 * 1000 -/** - * Usage is written asynchronously as a run settles, so a refetch fired on completion - * races the write and re-reads the old balance. - */ -const USAGE_SETTLE_DELAY_MS = 1000 - -/** - * Invalidates the workspace credit/usage reads after anything that moves the balance — - * a run that spends credits, a top-up, a plan change, or a usage-limit edit. Both - * families are keyed per workspace but derive from the same billing account, so the - * family prefixes (not a single workspace's key) are what has to be refetched. - */ -export function invalidateWorkspaceUsage(queryClient: QueryClient) { - return Promise.all([ - queryClient.invalidateQueries({ queryKey: workspaceUsageKeys.creditAvailabilities() }), - queryClient.invalidateQueries({ queryKey: workspaceUsageKeys.gates() }), - ]) -} - export function fetchWorkspaceCreditAvailability( workspaceId: string, signal?: AbortSignal @@ -76,18 +48,3 @@ export function useWorkspaceUsageGate(workspaceId?: string) { staleTime: WORKSPACE_USAGE_GATE_STALE_TIME, }) } - -/** - * Refreshes the billing reads a run touches, after a delay: usage is written - * asynchronously as the run settles, so an immediate refetch races the write and - * re-reads the pre-run balance. - * - * Shared by the surfaces that spend credits — workflow execution and wand - * generation — so the delay and the key set stay in one place. - */ -export function scheduleUsageRefresh(queryClient: QueryClient) { - setTimeout(() => { - void queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }) - void invalidateWorkspaceUsage(queryClient) - }, USAGE_SETTLE_DELAY_MS) -} diff --git a/apps/sim/hooks/selectors/use-selector-query.ts b/apps/sim/hooks/selectors/use-selector-query.ts index a202872f7d2..d315f7891fa 100644 --- a/apps/sim/hooks/selectors/use-selector-query.ts +++ b/apps/sim/hooks/selectors/use-selector-query.ts @@ -79,8 +79,7 @@ export function useSelectorOptions( * rejection. Unlike {@link useSelectorOptionDetail}, resolving nothing here is never * cheaper than the list's own preconditions, so there is no case for an override. */ - const isEnabled = - (args.enabled ?? true) && (definition.enabled ? definition.enabled(queryArgs) : true) + const isEnabled = args.enabled !== false && (definition.enabled?.(queryArgs) ?? true) const supportsPagination = Boolean(definition.fetchPage) const flatQuery = useQuery({ diff --git a/apps/sim/lib/api/contracts/oauth-connections.ts b/apps/sim/lib/api/contracts/oauth-connections.ts index 7effa93f300..7512b2bb442 100644 --- a/apps/sim/lib/api/contracts/oauth-connections.ts +++ b/apps/sim/lib/api/contracts/oauth-connections.ts @@ -62,6 +62,9 @@ const trelloCallbackQuerySchema = z }) .passthrough() +/** Google domain-wide-delegation subject. Also applied by in-process credential callers. */ +export const impersonateEmailSchema = z.string().email() + export const oauthTokenRequestBodySchema = z .object({ credentialId: z.string().min(1).optional(), @@ -69,7 +72,7 @@ export const oauthTokenRequestBodySchema = z providerId: z.string().min(1).optional(), workflowId: z.string().min(1).nullish(), scopes: z.array(z.string()).optional(), - impersonateEmail: z.string().email().optional(), + impersonateEmail: impersonateEmailSchema.optional(), }) .refine( (data) => data.credentialId || (data.credentialAccountUserId && data.providerId), @@ -88,7 +91,7 @@ export const oauthTokenPostQuerySchema = z.object({ userId: z.string().min(1).optional(), }) -const oauthTokenResponseSchema = z.object({ +export const oauthTokenResponseSchema = z.object({ accessToken: z.string(), idToken: z.string().optional(), instanceUrl: z.string().optional(), @@ -99,6 +102,9 @@ const oauthTokenResponseSchema = z.object({ authStyle: z.enum(['x-api-token']).optional(), }) +/** Token material a resolved credential yields, on the wire and in-process alike. */ +export type OAuthTokenResponse = z.output + export const oauthTokenGetContract = defineRouteContract({ method: 'GET', path: '/api/auth/oauth/token', diff --git a/apps/sim/lib/oauth/token-resolution.ts b/apps/sim/lib/oauth/token-resolution.ts index 91f658bea1a..86ee15114fd 100644 --- a/apps/sim/lib/oauth/token-resolution.ts +++ b/apps/sim/lib/oauth/token-resolution.ts @@ -1,6 +1,9 @@ import { AuditAction, AuditResourceType, recordAudit } from '@sim/audit' import { createLogger } from '@sim/logger' -import { z } from 'zod' +import { + impersonateEmailSchema, + type OAuthTokenResponse, +} from '@/lib/api/contracts/oauth-connections' import { authorizeCredentialUseForAuth } from '@/lib/auth/credential-access' import type { AuthResult } from '@/lib/auth/hybrid' import { TokenServiceAccountValidationError } from '@/lib/credentials/token-service-accounts/errors' @@ -25,21 +28,16 @@ export interface CredentialAuditRequest { headers: { get(name: string): string | null } } -/** Token material a resolved credential yields, as returned to every surface. */ -export interface CredentialTokenPayload { - accessToken: string - idToken?: string - instanceUrl?: string - apiDomain?: string - cloudId?: string - domain?: string - authStyle?: 'x-api-token' -} +/** + * Token material a resolved credential yields. It is the route's response body, so it + * comes from the contract rather than a parallel declaration that could drift from it. + */ +export type CredentialTokenPayload = OAuthTokenResponse export interface ResolveCredentialTokenInput { /** Correlation id used by the credential service's own logging. */ requestId: string - credentialId: string + credentialId?: string workflowId?: string /** Canonical provider scopes, used only by service-account token minting. */ scopes?: string[] @@ -57,8 +55,6 @@ export type ResolveCredentialTokenResult = | { ok: true; token: CredentialTokenPayload } | { ok: false; status: number; error: string; code?: string } -const impersonateEmailSchema = z.string().email() - /** * Emits the semantic "credential used" trail for one resolved credential. * Both the audit row and the analytics event are fire-and-forget. @@ -104,9 +100,10 @@ function recordCredentialAccess(params: { * into the wire payload every surface returns. Provider-specific hosts live in * the credential's scope string and are extracted through shared, allowlisted * helpers — never a local regex, since these values are injected into tool - * calls that carry the token. + * calls that carry the token. Zoho Desk's data-center-scoped REST base is one + * such value, surfaced as `apiDomain` so callers never assume a host. */ -export function buildOAuthTokenPayload( +function buildOAuthTokenPayload( credential: { providerId: string; scope?: string | null; idToken?: string | null }, accessToken: string ): CredentialTokenPayload { @@ -114,9 +111,6 @@ export function buildOAuthTokenPayload( ? extractSalesforceInstanceUrl(credential.scope ?? undefined) : undefined - // Zoho Desk persists its data-center-specific REST base URL in the scope - // string (derived from the token response api_domain) so callers never - // assume a host. Surface it as apiDomain for tool param injection. let apiDomain: string | undefined if (credential.providerId === 'zoho-desk' && credential.scope) { apiDomain = extractZohoDeskBaseFromScope(credential.scope) diff --git a/apps/sim/lib/workflows/persistence/save-normalized-state.ts b/apps/sim/lib/workflows/persistence/save-normalized-state.ts index d86b2bd5b28..28f98d7c714 100644 --- a/apps/sim/lib/workflows/persistence/save-normalized-state.ts +++ b/apps/sim/lib/workflows/persistence/save-normalized-state.ts @@ -4,6 +4,7 @@ import { createLogger } from '@sim/logger' import { assertWorkflowMutable, authorizeWorkflowByWorkspacePermission, + WorkflowLockedError, } from '@sim/platform-authz/workflow' import { eq } from 'drizzle-orm' import type { z } from 'zod' @@ -11,8 +12,7 @@ import { type WorkflowStateContractOutput, workflowStateSchema, } from '@/lib/api/contracts/workflows' -import { env } from '@/lib/core/config/env' -import { getSocketServerUrl } from '@/lib/core/utils/urls' +import { notifyWorkflowUpdated } from '@/lib/realtime/notify' import { extractAndPersistCustomTools } from '@/lib/workflows/persistence/custom-tools-persistence' import { prepareWorkflowStateForPersistence } from '@/lib/workflows/persistence/prepare-state' import { saveWorkflowToNormalizedTables } from '@/lib/workflows/persistence/utils' @@ -46,7 +46,8 @@ export function parseWorkflowStateForPersistence( * and the copilot checkpoint revert — calls this, so none of those steps can be * skipped by going through a different door. * - * @throws WorkflowLockedError when the workflow is not mutable. + * Every refusal, the workflow lock included, comes back as a failure result so + * callers need only one branch to present. */ export async function saveWorkflowNormalizedState(params: { requestId: string @@ -79,7 +80,14 @@ export async function saveWorkflowNormalizedState(params: { } } - await assertWorkflowMutable(workflowId) + try { + await assertWorkflowMutable(workflowId) + } catch (error) { + if (error instanceof WorkflowLockedError) { + return { success: false, status: error.status, error: error.message } + } + throw error + } const { state: preparedState, warnings: preparationWarnings } = prepareWorkflowStateForPersistence({ @@ -163,35 +171,7 @@ export async function saveWorkflowNormalizedState(params: { logger.error(`[${requestId}] Failed to persist custom tools`, { error, workflowId }) } - await notifySocketServer(requestId, workflowId) + await notifyWorkflowUpdated(workflowId) return { success: true, warnings: preparationWarnings } } - -/** - * Best-effort nudge so connected editors reload the workflow. Never fails the - * write — the state is already committed by the time this runs. - */ -async function notifySocketServer(requestId: string, workflowId: string): Promise { - try { - const notifyResponse = await fetch(`${getSocketServerUrl()}/api/workflow-updated`, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - 'x-api-key': env.INTERNAL_API_SECRET, - }, - body: JSON.stringify({ workflowId }), - }) - - if (!notifyResponse.ok) { - logger.warn( - `[${requestId}] Failed to notify Socket.IO server about workflow ${workflowId} update` - ) - } - } catch (notificationError) { - logger.warn( - `[${requestId}] Error notifying Socket.IO server about workflow ${workflowId} update`, - notificationError - ) - } -} diff --git a/apps/sim/lib/workspace-files/queries.test.ts b/apps/sim/lib/workspace-files/queries.test.ts index e3212079206..a578c16c965 100644 --- a/apps/sim/lib/workspace-files/queries.test.ts +++ b/apps/sim/lib/workspace-files/queries.test.ts @@ -53,6 +53,30 @@ describe('listWorkspaceFilesWithShares', () => { expect(file.uploadedAt).toEqual(new Date('2026-01-01T00:00:00.000Z')) }) + /** + * `maxRows` exists for a caller that will only use the list if the whole workspace fits + * its payload budget, so overflow must be reported as `null` — a prefix returned here + * would be presented as the workspace's complete file list. + */ + it('returns null without joining shares when the workspace exceeds maxRows', async () => { + mockListWorkspaceFiles.mockResolvedValue([STORED_FILE, STORED_FILE, STORED_FILE]) + + const result = await listWorkspaceFilesWithShares('ws-1', 'active', { maxRows: 2 }) + + expect(result).toBeNull() + expect(mockGetWorkspaceShares).not.toHaveBeenCalled() + expect(mockListWorkspaceFiles).toHaveBeenCalledWith('ws-1', { scope: 'active', limit: 3 }) + }) + + it('returns the list when it fits maxRows', async () => { + mockListWorkspaceFiles.mockResolvedValue([STORED_FILE]) + + const result = await listWorkspaceFilesWithShares('ws-1', 'active', { maxRows: 2 }) + + expect(result).toHaveLength(1) + expect(mockGetWorkspaceShares).toHaveBeenCalledWith('file', 'ws-1') + }) + it('joins each file public share onto its row', async () => { const share = { id: 'share-1', diff --git a/apps/sim/lib/workspace-files/queries.ts b/apps/sim/lib/workspace-files/queries.ts index 1a4925907b1..f4af38bf35c 100644 --- a/apps/sim/lib/workspace-files/queries.ts +++ b/apps/sim/lib/workspace-files/queries.ts @@ -6,9 +6,9 @@ import { } from '@/lib/uploads/contexts/workspace/workspace-file-manager' /** - * Lists a workspace's files with each file's public share joined on — shared by - * `GET /api/workspaces/[id]/files` and the Files browser's server prefetch so both cache - * one shape. + * Lists a workspace's files with each file's public share joined on, parsed through the + * `GET /api/workspaces/[id]/files` response contract so the workspace layout's server seed + * caches exactly the shape that route returns. * * Parsing through the route contract's response schema strips the server-only fields * `requestJson` strips on the client (`contentUpdatedAt`), so a prefetched entry is identical @@ -16,19 +16,26 @@ import { * * Callers authorize the viewer against `workspaceId` first. * - * `limit` caps the rows read for callers that only need to know whether the workspace fits - * a payload budget; the result is then a prefix of the list, not the list, so no caller may - * present a limited read as the workspace's files. + * `maxRows` bounds the work for a caller that will only use the list if the whole + * workspace fits a payload budget: the read stops one row past the budget and returns + * `null` on overflow, before the share join and the contract parse — so the workspaces + * the budget exists to protect are the ones that pay least to be rejected. Returning + * `null` rather than the prefix is what stops a caller presenting a truncated read as + * the workspace's files. */ export async function listWorkspaceFilesWithShares( workspaceId: string, scope: WorkspaceFileScope, - options?: { limit?: number } + options?: { maxRows?: number } ) { - const [files, shares] = await Promise.all([ - listWorkspaceFiles(workspaceId, { scope, limit: options?.limit }), - getWorkspaceShares('file', workspaceId), - ]) + const maxRows = options?.maxRows + const files = await listWorkspaceFiles(workspaceId, { + scope, + ...(maxRows === undefined ? {} : { limit: maxRows + 1 }), + }) + if (maxRows !== undefined && files.length > maxRows) return null + + const shares = await getWorkspaceShares('file', workspaceId) const withShares = files.map((file) => ({ ...file, share: shares.get(file.id) ?? null })) return listWorkspaceFilesContract.response.schema.shape.files.parse(withShares) } diff --git a/apps/sim/lib/workspaces/permissions/utils.ts b/apps/sim/lib/workspaces/permissions/utils.ts index 6eaa0e1df54..d1abf073258 100644 --- a/apps/sim/lib/workspaces/permissions/utils.ts +++ b/apps/sim/lib/workspaces/permissions/utils.ts @@ -107,16 +107,22 @@ async function selectWorkspaceWithOwner( } /** - * Request-memoized plain workspace read, keyed by id and archived visibility. - * A single Server Component render pass resolves the same workspace row through - * several independent gates, so without this the row is re-read once per gate. + * Request-memoized plain workspace read, keyed by id alone. A single Server + * Component render pass resolves the same workspace row through several + * independent gates, so without this the row is re-read once per gate. + * + * Keyed on the id and not on archived visibility on purpose: the gates disagree + * about whether they want archived workspaces, and memoizing that argument would + * give each answer its own entry and dedupe nothing. Reading the superset once + * and applying the caller's visibility below is what makes the two agree on a + * single query. * * Outside a Server Component render React evaluates this normally and retains * nothing, so API routes and background work are unaffected. */ const readWorkspaceWithOwner = cache( - (workspaceId: string, includeArchived: boolean): Promise => - selectWorkspaceWithOwner(workspaceId, includeArchived, db, false) + (workspaceId: string): Promise => + selectWorkspaceWithOwner(workspaceId, true, db, false) ) /** @@ -130,7 +136,7 @@ const readWorkspaceWithOwner = cache( * @param workspaceId - The workspace ID to look up * @returns The workspace with owner info if found, null otherwise */ -export function getWorkspaceWithOwner( +export async function getWorkspaceWithOwner( workspaceId: string, options?: { includeArchived?: boolean; executor?: DbOrTx; forUpdate?: boolean } ): Promise { @@ -138,7 +144,9 @@ export function getWorkspaceWithOwner( if (executor || forUpdate) { return selectWorkspaceWithOwner(workspaceId, includeArchived, executor ?? db, forUpdate) } - return readWorkspaceWithOwner(workspaceId, includeArchived) + const ws = await readWorkspaceWithOwner(workspaceId) + if (!ws) return null + return includeArchived || !ws.archivedAt ? ws : null } /** From 140b928a73d47f46e01344dbe7399d318508c2b1 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 13 Aug 2026 01:59:22 -0700 Subject: [PATCH 12/14] chore(test): type the evaluator provider-request helper instead of using any --- .../executor/handlers/evaluator/evaluator-handler.test.ts | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/apps/sim/executor/handlers/evaluator/evaluator-handler.test.ts b/apps/sim/executor/handlers/evaluator/evaluator-handler.test.ts index eed0e6fed06..e8ba5619689 100644 --- a/apps/sim/executor/handlers/evaluator/evaluator-handler.test.ts +++ b/apps/sim/executor/handlers/evaluator/evaluator-handler.test.ts @@ -44,6 +44,7 @@ import { EvaluatorBlockHandler } from '@/executor/handlers/evaluator/evaluator-h import type { ExecutionContext } from '@/executor/types' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' import { executeProviderRequest } from '@/providers' +import type { ProviderRequest } from '@/providers/types' import { getProviderFromModel } from '@/providers/utils' import type { SerializedBlock } from '@/serializer/types' @@ -51,8 +52,11 @@ const mockGetProviderFromModel = getProviderFromModel as Mock const mockExecuteProviderRequest = executeProviderRequest as Mock /** The provider request the handler built, keyed the way the old wire body was. */ -function providerRequestBody(index = 0): Record { - const [provider, request] = mockExecuteProviderRequest.mock.calls[index] +function providerRequestBody(index = 0): ProviderRequest & { provider: string } { + const [provider, request] = mockExecuteProviderRequest.mock.calls[index] as [ + string, + ProviderRequest, + ] return { provider, ...request } } From d2755e9c205d7073056e67d5e99aa663a434e6aa Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 13 Aug 2026 08:48:34 -0700 Subject: [PATCH 13/14] improvement(perf): restore the parallel file read, and close the gaps a diff audit surfaced --- .../[workspaceId]/lib/prefetch.test.ts | 2 + .../app/workspace/[workspaceId]/prefetch.ts | 6 ++ .../hooks/use-workflow-execution.test.tsx | 4 - .../evaluator/evaluator-handler.test.ts | 31 +++++++ apps/sim/executor/utils/provider-request.ts | 19 ++++- apps/sim/hooks/queries/schedules.ts | 5 ++ apps/sim/hooks/queries/tables.test.ts | 27 ++++++ .../persistence/save-normalized-state.test.ts | 83 +++++++++++++++++++ apps/sim/lib/workspace-files/queries.test.ts | 38 ++++++++- apps/sim/lib/workspace-files/queries.ts | 29 ++++--- .../lib/workspaces/permissions/utils.test.ts | 27 ++++++ apps/sim/lib/workspaces/permissions/utils.ts | 7 +- 12 files changed, 257 insertions(+), 21 deletions(-) create mode 100644 apps/sim/lib/workflows/persistence/save-normalized-state.test.ts diff --git a/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts b/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts index c506d873673..c27a77959b0 100644 --- a/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts @@ -527,6 +527,8 @@ describe('workspace list prefetches', () => { expect(mockListWorkspaceFilesWithShares).toHaveBeenCalledWith(WORKSPACE_ID, 'active', { maxRows: WORKSPACE_FILE_SEED_MAX, + /** A failed read must reach the catch, not degrade to a cached empty list. */ + throwOnError: true, }) expect(client.getQueryData(workspaceFilesKeys.list(WORKSPACE_ID, 'active'))).toEqual(files) }) diff --git a/apps/sim/app/workspace/[workspaceId]/prefetch.ts b/apps/sim/app/workspace/[workspaceId]/prefetch.ts index c5a53c54aa0..3ca952e3677 100644 --- a/apps/sim/app/workspace/[workspaceId]/prefetch.ts +++ b/apps/sim/app/workspace/[workspaceId]/prefetch.ts @@ -130,6 +130,12 @@ async function seedWorkspaceFiles(queryClient: QueryClient, workspaceId: string) try { const files = await listWorkspaceFilesWithShares(workspaceId, 'active', { maxRows: WORKSPACE_FILE_SEED_MAX, + /** + * A failed read must reach the catch below, not degrade to an empty list: seeding + * `[]` would cache "this workspace has no files" as authoritative for the entry's + * lifetime, which is worse than seeding nothing and letting the client fetch. + */ + throwOnError: true, }) if (!files) return queryClient.setQueryData(workspaceFilesKeys.list(workspaceId, 'active'), files) diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.test.tsx b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.test.tsx index c9b2fd0a6a4..15c7171064c 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.test.tsx +++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.test.tsx @@ -191,10 +191,6 @@ vi.mock('@/executor/utils/start-block', () => ({ coerceValue: (_type: string, value: unknown) => value, })) -vi.mock('@/hooks/queries/subscription', () => ({ - subscriptionKeys: { users: () => ['subscription', 'users'] }, -})) - vi.mock('@/hooks/queries/utils/workflow-cache', () => ({ getWorkflows: () => [], })) diff --git a/apps/sim/executor/handlers/evaluator/evaluator-handler.test.ts b/apps/sim/executor/handlers/evaluator/evaluator-handler.test.ts index e8ba5619689..6df5974aefa 100644 --- a/apps/sim/executor/handlers/evaluator/evaluator-handler.test.ts +++ b/apps/sim/executor/handlers/evaluator/evaluator-handler.test.ts @@ -145,6 +145,37 @@ describe('EvaluatorBlockHandler', () => { expect(handler.canHandle(nonEvalBlock)).toBe(false) }) + /** + * The admission checks the removed `/api/providers` hop owned. Mirrors the router's + * coverage — both handlers reach the provider through the same shared entry point. + */ + const admissionInputs = { + content: 'Evaluate this.', + metrics: [{ name: 'score1', description: 'First score', range: { min: 0, max: 10 } }], + model: 'gpt-4o', + apiKey: 'test-api-key', + } + + it('refuses to reach the provider without an execution subject', async () => { + mockContext.userId = undefined + + await expect(handler.execute(mockContext, mockBlock, admissionInputs)).rejects.toThrow( + 'Unauthorized' + ) + expect(mockExecuteProviderRequest).not.toHaveBeenCalled() + }) + + it('refuses to reach the provider when the subject lost workspace access', async () => { + mockContext.workspaceId = 'test-workspace' + mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: false }) + + await expect(handler.execute(mockContext, mockBlock, admissionInputs)).rejects.toThrow( + 'Forbidden' + ) + expect(mockCheckWorkspaceAccess).toHaveBeenCalledWith('test-workspace', 'test-user') + expect(mockExecuteProviderRequest).not.toHaveBeenCalled() + }) + it('should execute evaluator block correctly with basic inputs', async () => { const inputs = { content: 'This is the content to evaluate.', diff --git a/apps/sim/executor/utils/provider-request.ts b/apps/sim/executor/utils/provider-request.ts index d6367de56d0..81b4a61b4d7 100644 --- a/apps/sim/executor/utils/provider-request.ts +++ b/apps/sim/executor/utils/provider-request.ts @@ -60,17 +60,32 @@ export async function executeBlockProviderRequest({ * model-emitted tool calls, and the route this replaces never carried one. * Router and evaluator requests declare no tools, so passing the executor's * context here would widen the trusted surface without changing any outcome. + * + * The whole runtime context is omitted when there is no registry, rather than + * passed carrying `undefined`. `executeProviderTool` reads a present context with + * an absent registry as "provenance was expected and is missing" and fails the + * call closed with no error text — unreachable while these blocks declare no + * tools, but a silent failure the day one does. */ const response = await executeProviderRequest( providerId, { ...request, userId: ctx.userId }, - { resolvedSecretTraceRegistry } + resolvedSecretTraceRegistry ? { resolvedSecretTraceRegistry } : undefined ) - if (response instanceof ReadableStream || (response !== null && 'stream' in response)) { + if ( + response instanceof ReadableStream || + (typeof response === 'object' && response !== null && 'stream' in response) + ) { logger.error('Provider returned a stream for a non-streaming block request', { providerId }) throw new Error('Provider returned a streaming response for a non-streaming request') } + logger.info('Provider request completed', { + providerId, + model: request.model, + workflowId: ctx.workflowId, + }) + return response } diff --git a/apps/sim/hooks/queries/schedules.ts b/apps/sim/hooks/queries/schedules.ts index c33ed845663..fd53c8af9ff 100644 --- a/apps/sim/hooks/queries/schedules.ts +++ b/apps/sim/hooks/queries/schedules.ts @@ -249,6 +249,11 @@ export function useRedeployWorkflowSchedule() { const { workflowId, blockId } = data await Promise.all([ queryClient.invalidateQueries({ queryKey: scheduleKeys.schedule(workflowId, blockId) }), + /** + * A redeploy recreates the schedule, so the id-keyed reads go stale too. They are + * a separate subtree from `schedule(workflowId, blockId)`, which does not cover them. + */ + queryClient.invalidateQueries({ queryKey: scheduleKeys.byIds() }), queryClient.invalidateQueries({ queryKey: deploymentKeys.info(workflowId) }), queryClient.invalidateQueries({ queryKey: deploymentKeys.versions(workflowId) }), ]) diff --git a/apps/sim/hooks/queries/tables.test.ts b/apps/sim/hooks/queries/tables.test.ts index c20ab37b23d..b90836494f4 100644 --- a/apps/sim/hooks/queries/tables.test.ts +++ b/apps/sim/hooks/queries/tables.test.ts @@ -129,6 +129,33 @@ describe('useDeleteColumn optimistic update', () => { expect(ctx?.rowSnapshots?.length).toBeGreaterThan(0) }) + /** + * The `find` cache hangs off the same `rowsRoot` parent as the paged rows but holds + * `{matches, truncated}` — no `pages`, no `rows`. A cache walk starting at the shared + * parent reaches it and throws inside `onMutate`, rejecting the mutation before it ever + * reaches the server: search a table, dismiss the search, then edit a cell. + */ + it('survives a cached search result hanging off the shared rows prefix', async () => { + setCache(tableKeys.detail(TABLE_ID), { + id: TABLE_ID, + schema: { columns: [{ name: 'age', type: 'number' }] }, + }) + setCache(ROWS_KEY, { + rows: [{ id: 'r1', data: { age: 1 } }], + totalCount: 1, + }) + setCache(tableKeys.find(TABLE_ID, 'q'), { matches: [{ rowId: 'r1', column: 'age' }] }) + + const hook = useDeleteColumn({ workspaceId: WORKSPACE_ID, tableId: TABLE_ID }) + + await expect(hook.onMutate?.('age')).resolves.toBeDefined() + + const rows = getCache<{ rows: Array<{ data: Record }> }>(ROWS_KEY) + expect(rows?.rows[0]?.data).toEqual({}) + /** The find entry is match coordinates, not row values — it must be left untouched. */ + expect(getCache<{ matches: unknown[] }>(tableKeys.find(TABLE_ID, 'q'))?.matches).toHaveLength(1) + }) + it('rolls back schema and rows on error using snapshots', async () => { const originalDetail = { id: TABLE_ID, diff --git a/apps/sim/lib/workflows/persistence/save-normalized-state.test.ts b/apps/sim/lib/workflows/persistence/save-normalized-state.test.ts new file mode 100644 index 00000000000..9240f827fc9 --- /dev/null +++ b/apps/sim/lib/workflows/persistence/save-normalized-state.test.ts @@ -0,0 +1,83 @@ +/** + * @vitest-environment node + */ +import { describe, expect, it } from 'vitest' +import { parseWorkflowStateForPersistence } from '@/lib/workflows/persistence/save-normalized-state' + +/** + * A checkpoint blob as the revert route builds it: JSONB-derived blocks and edges, plus a + * real `Date` for `deployedAt`. + */ +function checkpointState(overrides?: Record) { + return { + blocks: { + 'block-1': { + id: 'block-1', + type: 'starter', + name: 'Start', + position: { x: 0, y: 0 }, + subBlocks: {}, + outputs: {}, + enabled: true, + }, + }, + edges: [], + loops: {}, + parallels: {}, + isDeployed: false, + lastSaved: 1_754_000_000_000, + ...overrides, + } +} + +describe('parseWorkflowStateForPersistence', () => { + /** + * The revert used to reach this schema by POSTing the blob over HTTP, so every value + * arrived JSON-serialized. In-process the blob keeps its runtime types. Both forms must + * parse identically, or a checkpoint that reverted before would start failing. + */ + it('accepts a Date for deployedAt exactly as it accepted the serialized string', () => { + const deployedAt = new Date('2026-01-02T03:04:05.678Z') + + const fromDate = parseWorkflowStateForPersistence(checkpointState({ deployedAt })) + const overTheWire = parseWorkflowStateForPersistence( + JSON.parse(JSON.stringify(checkpointState({ deployedAt }))) + ) + + expect(fromDate.success).toBe(true) + expect(overTheWire.success).toBe(true) + expect(fromDate.data?.deployedAt).toEqual(deployedAt) + expect(overTheWire.data?.deployedAt).toEqual(fromDate.data?.deployedAt) + }) + + it('round-trips a JSONB-shaped blob without dropping blocks or edges', () => { + const state = checkpointState() + + const parsed = parseWorkflowStateForPersistence(state) + + expect(parsed.success).toBe(true) + expect(Object.keys(parsed.data?.blocks ?? {})).toEqual(['block-1']) + expect(parsed.data?.lastSaved).toBe(1_754_000_000_000) + }) + + it('accepts a null deployedAt, which the revert passes for a never-deployed checkpoint', () => { + const parsed = parseWorkflowStateForPersistence(checkpointState({ deployedAt: null })) + + expect(parsed.success).toBe(true) + expect(parsed.data?.deployedAt).toBeNull() + }) + + /** The validation the removed HTTP hop used to provide: a malformed blob must not be written. */ + it('rejects a blob whose blocks are malformed', () => { + const parsed = parseWorkflowStateForPersistence({ + blocks: { 'block-1': { id: 'block-1' } }, + edges: [], + }) + + expect(parsed.success).toBe(false) + }) + + it('rejects a blob missing blocks entirely', () => { + expect(parseWorkflowStateForPersistence({ edges: [] }).success).toBe(false) + }) +}) diff --git a/apps/sim/lib/workspace-files/queries.test.ts b/apps/sim/lib/workspace-files/queries.test.ts index a578c16c965..0f69aa9ce40 100644 --- a/apps/sim/lib/workspace-files/queries.test.ts +++ b/apps/sim/lib/workspace-files/queries.test.ts @@ -56,15 +56,16 @@ describe('listWorkspaceFilesWithShares', () => { /** * `maxRows` exists for a caller that will only use the list if the whole workspace fits * its payload budget, so overflow must be reported as `null` — a prefix returned here - * would be presented as the workspace's complete file list. + * would be presented as the workspace's complete file list. The share read still runs + * concurrently and is discarded: the under-budget workspaces are the common case, and + * serializing the two reads to save this one would tax every normal request. */ - it('returns null without joining shares when the workspace exceeds maxRows', async () => { + it('returns null when the workspace exceeds maxRows', async () => { mockListWorkspaceFiles.mockResolvedValue([STORED_FILE, STORED_FILE, STORED_FILE]) const result = await listWorkspaceFilesWithShares('ws-1', 'active', { maxRows: 2 }) expect(result).toBeNull() - expect(mockGetWorkspaceShares).not.toHaveBeenCalled() expect(mockListWorkspaceFiles).toHaveBeenCalledWith('ws-1', { scope: 'active', limit: 3 }) }) @@ -77,6 +78,37 @@ describe('listWorkspaceFilesWithShares', () => { expect(mockGetWorkspaceShares).toHaveBeenCalledWith('file', 'ws-1') }) + /** The boundary the `>` comparison turns on: exactly maxRows must still be the list. */ + it('returns the list when it sits exactly on maxRows', async () => { + mockListWorkspaceFiles.mockResolvedValue([STORED_FILE, STORED_FILE]) + + const result = await listWorkspaceFilesWithShares('ws-1', 'active', { maxRows: 2 }) + + expect(result).toHaveLength(2) + }) + + /** + * The file read swallows errors and returns `[]` by default. A caller seeding a cache + * must not receive that: an empty list would be cached as "this workspace has no files". + */ + it('propagates a failed read instead of degrading to an empty list', async () => { + await listWorkspaceFilesWithShares('ws-1', 'active', { throwOnError: true }) + + expect(mockListWorkspaceFiles).toHaveBeenCalledWith( + 'ws-1', + expect.objectContaining({ throwOnError: true }) + ) + }) + + it('does not ask the file read to throw unless the caller opts in', async () => { + await listWorkspaceFilesWithShares('ws-1', 'active') + + expect(mockListWorkspaceFiles).toHaveBeenCalledWith( + 'ws-1', + expect.not.objectContaining({ throwOnError: true }) + ) + }) + it('joins each file public share onto its row', async () => { const share = { id: 'share-1', diff --git a/apps/sim/lib/workspace-files/queries.ts b/apps/sim/lib/workspace-files/queries.ts index f4af38bf35c..d4e4cd11892 100644 --- a/apps/sim/lib/workspace-files/queries.ts +++ b/apps/sim/lib/workspace-files/queries.ts @@ -16,26 +16,33 @@ import { * * Callers authorize the viewer against `workspaceId` first. * - * `maxRows` bounds the work for a caller that will only use the list if the whole + * `maxRows` bounds the result for a caller that will only use the list if the whole * workspace fits a payload budget: the read stops one row past the budget and returns - * `null` on overflow, before the share join and the contract parse — so the workspaces - * the budget exists to protect are the ones that pay least to be rejected. Returning - * `null` rather than the prefix is what stops a caller presenting a truncated read as - * the workspace's files. + * `null` on overflow rather than the prefix, which is what stops a caller presenting a + * truncated read as the workspace's files. The two reads still run concurrently — the + * workspaces under the budget are the common case, and serializing them to save a share + * read on the rare oversized one would tax every normal request to do it. + * + * `throwOnError` propagates a failed file read instead of letting it degrade to an empty + * list. A caller seeding a cache needs that distinction: an empty list would be cached + * as authoritative, telling the user the workspace has no files. */ export async function listWorkspaceFilesWithShares( workspaceId: string, scope: WorkspaceFileScope, - options?: { maxRows?: number } + options?: { maxRows?: number; throwOnError?: boolean } ) { const maxRows = options?.maxRows - const files = await listWorkspaceFiles(workspaceId, { - scope, - ...(maxRows === undefined ? {} : { limit: maxRows + 1 }), - }) + const [files, shares] = await Promise.all([ + listWorkspaceFiles(workspaceId, { + scope, + ...(maxRows === undefined ? {} : { limit: maxRows + 1 }), + ...(options?.throwOnError ? { throwOnError: true } : {}), + }), + getWorkspaceShares('file', workspaceId), + ]) if (maxRows !== undefined && files.length > maxRows) return null - const shares = await getWorkspaceShares('file', workspaceId) const withShares = files.map((file) => ({ ...file, share: shares.get(file.id) ?? null })) return listWorkspaceFilesContract.response.schema.shape.files.parse(withShares) } diff --git a/apps/sim/lib/workspaces/permissions/utils.test.ts b/apps/sim/lib/workspaces/permissions/utils.test.ts index 64d1aa2011d..9d132ff53bf 100644 --- a/apps/sim/lib/workspaces/permissions/utils.test.ts +++ b/apps/sim/lib/workspaces/permissions/utils.test.ts @@ -745,6 +745,33 @@ describe('Permission Utils', () => { expect(result).toEqual({ id: 'workspace123', ownerId: null }) }) + + /** + * Archived visibility is applied in JS, not SQL, so the read can be shared by the + * gates that disagree about it. That makes these two the boundary worth pinning: + * if the filter ever stops matching the old `archived_at IS NULL` predicate, archived + * workspaces silently become visible to callers that asked not to see them. + */ + it.concurrent('should hide an archived workspace by default', async () => { + const chain = createMockChain([ + { id: 'workspace123', ownerId: 'owner456', archivedAt: new Date('2026-01-01') }, + ]) + mockDb.select.mockReturnValue(chain) + + const result = await getWorkspaceWithOwner('workspace123') + + expect(result).toBeNull() + }) + + it.concurrent('should return an archived workspace when asked to include them', async () => { + const archivedAt = new Date('2026-01-01') + const chain = createMockChain([{ id: 'workspace123', ownerId: 'owner456', archivedAt }]) + mockDb.select.mockReturnValue(chain) + + const result = await getWorkspaceWithOwner('workspace123', { includeArchived: true }) + + expect(result).toEqual({ id: 'workspace123', ownerId: 'owner456', archivedAt }) + }) }) describe('workspaceExists', () => { diff --git a/apps/sim/lib/workspaces/permissions/utils.ts b/apps/sim/lib/workspaces/permissions/utils.ts index d1abf073258..8815e96d7a9 100644 --- a/apps/sim/lib/workspaces/permissions/utils.ts +++ b/apps/sim/lib/workspaces/permissions/utils.ts @@ -118,7 +118,12 @@ async function selectWorkspaceWithOwner( * single query. * * Outside a Server Component render React evaluates this normally and retains - * nothing, so API routes and background work are unaffected. + * nothing, so API routes and background work are unaffected — verified against this + * Next version in both dev and production builds. + * + * The row this returns is SHARED by every consumer in the render, including the + * `workspace` on each viewer's access result. Treat it as immutable: an in-place edit + * would poison every gate that reads it for the rest of the pass. */ const readWorkspaceWithOwner = cache( (workspaceId: string): Promise => From aacc3e2c46ffbe35c2e422e70114c5987e95aa63 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 13 Aug 2026 09:03:31 -0700 Subject: [PATCH 14/14] improvement(perf): drop a duplicate authorization, parallelize the credential reads, and trim the comments --- .../api/copilot/checkpoints/revert/route.ts | 2 + .../sim/app/api/workflows/[id]/state/route.ts | 9 +--- .../app/workspace/[workspaceId]/prefetch.ts | 36 ++++--------- apps/sim/executor/utils/provider-request.ts | 43 ++++----------- apps/sim/hooks/queries/schedules.ts | 5 +- apps/sim/hooks/queries/tables.ts | 4 -- .../hooks/queries/utils/invalidate-usage.ts | 11 ++-- .../hooks/queries/utils/subscription-keys.ts | 7 +-- apps/sim/hooks/queries/utils/table-keys.ts | 9 ++-- .../queries/utils/workspace-usage-keys.ts | 6 +-- apps/sim/hooks/queries/workspace-files.ts | 11 ++-- .../sim/hooks/selectors/use-selector-query.ts | 10 ++-- .../lib/api/contracts/oauth-connections.ts | 2 +- apps/sim/lib/auth/credential-access.ts | 8 ++- apps/sim/lib/billing/core/subscription.ts | 15 +++--- apps/sim/lib/oauth/token-resolution.ts | 53 +++++++------------ apps/sim/lib/permissions/super-user.ts | 5 +- .../workspace/workspace-file-manager.ts | 7 +-- .../persistence/save-normalized-state.ts | 35 ++++++------ apps/sim/lib/workspace-files/queries.ts | 19 +++---- apps/sim/lib/workspaces/permissions/utils.ts | 22 +++----- apps/sim/tools/index.ts | 2 +- 22 files changed, 108 insertions(+), 213 deletions(-) diff --git a/apps/sim/app/api/copilot/checkpoints/revert/route.ts b/apps/sim/app/api/copilot/checkpoints/revert/route.ts index e022c0a3d57..1543372f773 100644 --- a/apps/sim/app/api/copilot/checkpoints/revert/route.ts +++ b/apps/sim/app/api/copilot/checkpoints/revert/route.ts @@ -140,6 +140,8 @@ export const POST = withRouteHandler(async (request: NextRequest) => { workflowId: checkpoint.workflowId, userId, state: parsedState.data, + /** Already resolved above; re-deriving it would repeat 2-3 sequential reads. */ + authorization, }) if (!saveResult.success) { diff --git a/apps/sim/app/api/workflows/[id]/state/route.ts b/apps/sim/app/api/workflows/[id]/state/route.ts index 999129410cf..27ef1e72c2d 100644 --- a/apps/sim/app/api/workflows/[id]/state/route.ts +++ b/apps/sim/app/api/workflows/[id]/state/route.ts @@ -1,10 +1,7 @@ import { db } from '@sim/db' import { workflow } from '@sim/db/schema' import { createLogger } from '@sim/logger' -import { - authorizeWorkflowByWorkspacePermission, - WorkflowLockedError, -} from '@sim/platform-authz/workflow' +import { authorizeWorkflowByWorkspacePermission } from '@sim/platform-authz/workflow' import { toError } from '@sim/utils/errors' import { eq, sql } from 'drizzle-orm' import { type NextRequest, NextResponse } from 'next/server' @@ -132,10 +129,6 @@ export const PUT = withRouteHandler( return NextResponse.json({ success: true, warnings: result.warnings }, { status: 200 }) } catch (error: any) { - if (error instanceof WorkflowLockedError) { - return NextResponse.json({ error: error.message }, { status: error.status }) - } - const elapsed = Date.now() - startTime logger.error( `[${requestId}] Error saving workflow ${workflowId} state after ${elapsed}ms`, diff --git a/apps/sim/app/workspace/[workspaceId]/prefetch.ts b/apps/sim/app/workspace/[workspaceId]/prefetch.ts index 3ca952e3677..e19e970b581 100644 --- a/apps/sim/app/workspace/[workspaceId]/prefetch.ts +++ b/apps/sim/app/workspace/[workspaceId]/prefetch.ts @@ -93,38 +93,22 @@ async function seedWorkspaceList( } /** - * How many files the layout is willing to inline into the document. + * How many files the layout is willing to inline into the document. Seeded on EVERY + * workspace route, so at ~500 bytes of JSON per file this budgets the entry at ~150 KB. * - * The file list is seeded on EVERY workspace route (see the call site), so its cost is - * paid per navigation into the app, not per visit to Files. At roughly 500 bytes of - * serialized JSON per file, this budgets the entry at ~150 KB; a workspace with - * thousands of files would otherwise push more than a megabyte of HTML ahead of first - * paint on the logs, settings, and editor routes that never read it. - * - * A workspace above the budget seeds NOTHING rather than a prefix: the sidebar search - * filters this list client-side and the Files browser renders it as the workspace's - * files, so a truncated seed would silently hide files. Those workspaces fetch the - * complete list from the route instead — which for a list that large is also the - * cheaper first paint. + * A workspace above the budget seeds NOTHING rather than a prefix: the sidebar filters + * this list client-side, so a truncated seed would silently hide files. */ export const WORKSPACE_FILE_SEED_MAX = 300 /** - * Seeds the workspace's file list, which the sidebar's search modal reads on EVERY - * workspace route — so this query is registered by sidebar chrome before any page - * renders. That ordering is why it has to be seeded HERE and not only by the Files - * pages: `HydrationBoundary` hydrates a query the cache has already seen from a - * `useEffect`, which never runs during SSR, so a page-level boundary can only ever hand - * this entry to the client. Seeding it with the layout's own boundary — the first one to - * render — is what lets the server paint the Files browser and the open file's header - * populated instead of shipping a spinner and resolving it a beat later on the client. - * - * Seeded rather than prefetched so it can decline to create an entry at all when the - * workspace exceeds {@link WORKSPACE_FILE_SEED_MAX}: `prefetchQuery` always creates one, - * and a partial one would be read as the whole list. + * Seeds the workspace's file list, which sidebar chrome registers on EVERY workspace + * route. It must be seeded HERE, not by the Files pages: `HydrationBoundary` defers a + * query the cache has already seen to a `useEffect`, which SSR never runs. * - * Parsed through the same response contract `GET /api/workspaces/[id]/files` validates - * against, so a seeded entry is identical to what the client hook would cache. + * Seeded rather than prefetched so it can decline to create an entry at all above + * {@link WORKSPACE_FILE_SEED_MAX} — `prefetchQuery` always creates one, and a partial + * entry would be read as the whole list. Parsed through the route's response contract. */ async function seedWorkspaceFiles(queryClient: QueryClient, workspaceId: string): Promise { try { diff --git a/apps/sim/executor/utils/provider-request.ts b/apps/sim/executor/utils/provider-request.ts index 81b4a61b4d7..622d5f77a35 100644 --- a/apps/sim/executor/utils/provider-request.ts +++ b/apps/sim/executor/utils/provider-request.ts @@ -11,32 +11,17 @@ interface ExecuteBlockProviderRequestInput { ctx: ExecutionContext providerId: string request: ProviderRequest - /** - * The fork the block's model input was projected through. Supplied to the - * provider runtime in place of the provenance envelope the HTTP boundary used - * to serialize and re-import. - */ + /** Supplied in place of the provenance envelope the HTTP boundary serialized and re-imported. */ resolvedSecretTraceRegistry: ResolvedSecretTraceRegistry | undefined } /** - * Runs one non-streaming provider request for a block handler in-process. - * - * Replaces the executor's `POST /api/providers` round trip, which re-derived - * everything it needed from claims the executor had itself just supplied. The - * two admission checks the route owned are reproduced here so the outcome is - * unchanged: - * - * - `checkInternalAuth` rejected a token carrying no user. The executor mints - * that token from `ctx.userId`, so the check reduces to requiring one. - * - `checkWorkspaceAccess` rejected an execution subject who is no longer a - * member of the workspace being billed. - * - * The route's remaining work is either already done by the caller (the model - * permission policy, via `validateModelProvider`; Vertex credential - * authorization, via `resolveVertexCredential`) or lives inside - * `executeProviderRequest` itself (BYOK key resolution, attachment provenance - * filtering, cost policy). + * Runs one non-streaming provider request for a block handler in-process, replacing the + * executor's `POST /api/providers` round trip. The route's two admission checks are + * reproduced so the outcome is unchanged: an internal token with no user is rejected (the + * executor mints it from `ctx.userId`), and an execution subject who has left the billed + * workspace is rejected. The route's remaining work is already done by the caller or lives + * inside `executeProviderRequest`. */ export async function executeBlockProviderRequest({ ctx, @@ -56,16 +41,10 @@ export async function executeBlockProviderRequest({ } /** - * `executionContext` is deliberately not supplied: it is only inherited by - * model-emitted tool calls, and the route this replaces never carried one. - * Router and evaluator requests declare no tools, so passing the executor's - * context here would widen the trusted surface without changing any outcome. - * - * The whole runtime context is omitted when there is no registry, rather than - * passed carrying `undefined`. `executeProviderTool` reads a present context with - * an absent registry as "provenance was expected and is missing" and fails the - * call closed with no error text — unreachable while these blocks declare no - * tools, but a silent failure the day one does. + * No `executionContext`: it is only inherited by model-emitted tool calls, and the route + * this replaces never carried one. The whole context is omitted when there is no registry + * rather than passed carrying `undefined` — `executeProviderTool` reads that as missing + * provenance and fails the call closed with no error text. */ const response = await executeProviderRequest( providerId, diff --git a/apps/sim/hooks/queries/schedules.ts b/apps/sim/hooks/queries/schedules.ts index fd53c8af9ff..a22fe07d4fe 100644 --- a/apps/sim/hooks/queries/schedules.ts +++ b/apps/sim/hooks/queries/schedules.ts @@ -249,10 +249,7 @@ export function useRedeployWorkflowSchedule() { const { workflowId, blockId } = data await Promise.all([ queryClient.invalidateQueries({ queryKey: scheduleKeys.schedule(workflowId, blockId) }), - /** - * A redeploy recreates the schedule, so the id-keyed reads go stale too. They are - * a separate subtree from `schedule(workflowId, blockId)`, which does not cover them. - */ + /** A redeploy recreates the schedule; the id-keyed reads are a separate subtree. */ queryClient.invalidateQueries({ queryKey: scheduleKeys.byIds() }), queryClient.invalidateQueries({ queryKey: deploymentKeys.info(workflowId) }), queryClient.invalidateQueries({ queryKey: deploymentKeys.versions(workflowId) }), diff --git a/apps/sim/hooks/queries/tables.ts b/apps/sim/hooks/queries/tables.ts index 4f41554f9d6..bbc3718fd2f 100644 --- a/apps/sim/hooks/queries/tables.ts +++ b/apps/sim/hooks/queries/tables.ts @@ -837,10 +837,6 @@ function withOptimisticAutoFireExec(groups: WorkflowGroup[], row: TableRow): Tab * shared parent, and handing this updater a `find` entry — a flat * {@link TableFindResult}, not pages — throws on `old.pages` inside `onMutate`, so * the whole cell edit would reject before reaching the server. - * - * A consequence worth knowing: an open search-results view is therefore left to its - * own refetch rather than patched here, since it holds a different shape. Patching - * it too would need its own updater keyed on {@link tableKeys.find}. */ function patchCachedRows( queryClient: ReturnType, diff --git a/apps/sim/hooks/queries/utils/invalidate-usage.ts b/apps/sim/hooks/queries/utils/invalidate-usage.ts index 164451b3856..2c08f009d72 100644 --- a/apps/sim/hooks/queries/utils/invalidate-usage.ts +++ b/apps/sim/hooks/queries/utils/invalidate-usage.ts @@ -9,10 +9,9 @@ import { workspaceUsageKeys } from '@/hooks/queries/utils/workspace-usage-keys' const USAGE_SETTLE_DELAY_MS = 1000 /** - * Invalidates the workspace credit/usage reads after anything that moves the balance — - * a run that spends credits, a top-up, a plan change, or a usage-limit edit. Both - * families are keyed per workspace but derive from the same billing account, so the - * family prefixes (not a single workspace's key) are what has to be refetched. + * Invalidates the workspace credit/usage reads after anything that moves the balance. + * Both families are keyed per workspace but derive from one billing account, so the + * family prefixes are what must refetch. */ export function invalidateWorkspaceUsage(queryClient: QueryClient) { return Promise.all([ @@ -23,9 +22,7 @@ export function invalidateWorkspaceUsage(queryClient: QueryClient) { /** * Refreshes the billing reads a run touches, after {@link USAGE_SETTLE_DELAY_MS}. - * - * Shared by the surfaces that spend credits — workflow execution and wand generation — - * so the delay and the key set stay in one place. + * Shared by workflow execution and wand generation. */ export function scheduleUsageRefresh(queryClient: QueryClient) { setTimeout(() => { diff --git a/apps/sim/hooks/queries/utils/subscription-keys.ts b/apps/sim/hooks/queries/utils/subscription-keys.ts index 0d6aa6af6b4..f2a879400a4 100644 --- a/apps/sim/hooks/queries/utils/subscription-keys.ts +++ b/apps/sim/hooks/queries/utils/subscription-keys.ts @@ -1,9 +1,6 @@ /** - * React Query key factory for subscription and billing reads. - * - * Lives in this standalone module — like {@link file://./workspace-usage-keys.ts} — so - * the shared billing invalidations can reference it without importing the hook module - * that consumes those invalidations, which would close an import cycle between the two. + * React Query key factory for subscription and billing reads. Standalone so the shared + * billing invalidations can use it without closing an import cycle through the hook module. */ export const subscriptionKeys = { all: ['subscription'] as const, diff --git a/apps/sim/hooks/queries/utils/table-keys.ts b/apps/sim/hooks/queries/utils/table-keys.ts index 8394373c66b..853103b5122 100644 --- a/apps/sim/hooks/queries/utils/table-keys.ts +++ b/apps/sim/hooks/queries/utils/table-keys.ts @@ -26,12 +26,9 @@ export const tableKeys = { [...tableKeys.all, 'export-jobs', workspaceId ?? ''] as const, rowsRoot: (tableId: string) => [...tableKeys.detail(tableId), 'rows'] as const, /** - * Prefix covering only the paged row lists. - * - * `rowsRoot` is a shared parent — `find` hangs off it holding an entirely different - * shape — so anything walking the cache to update or snapshot row pages must start - * here instead. Reaching for `rowsRoot` and subtracting the siblings is a denylist - * that rots the moment another subtree is added. + * Prefix covering only the paged row lists. `rowsRoot` is a shared parent — `find` + * hangs off it holding a different shape — so anything walking the cache for row + * pages must start here. */ infiniteRowsRoot: (tableId: string) => [...tableKeys.rowsRoot(tableId), 'infinite'] as const, infiniteRows: (tableId: string, paramsKey: string) => diff --git a/apps/sim/hooks/queries/utils/workspace-usage-keys.ts b/apps/sim/hooks/queries/utils/workspace-usage-keys.ts index 8063c3564b2..8bb1a0f1b72 100644 --- a/apps/sim/hooks/queries/utils/workspace-usage-keys.ts +++ b/apps/sim/hooks/queries/utils/workspace-usage-keys.ts @@ -1,8 +1,6 @@ /** - * React Query key factory for the per-workspace credit and usage-gate reads. - * - * Standalone for the same reason as {@link file://./subscription-keys.ts}: the shared - * billing invalidations need these keys without importing the hooks that call them. + * React Query key factory for the per-workspace credit and usage-gate reads. Standalone + * for the same import-cycle reason as {@link file://./subscription-keys.ts}. */ export const workspaceUsageKeys = { all: ['workspace-usage'] as const, diff --git a/apps/sim/hooks/queries/workspace-files.ts b/apps/sim/hooks/queries/workspace-files.ts index ce118ac837d..69c775f2dde 100644 --- a/apps/sim/hooks/queries/workspace-files.ts +++ b/apps/sim/hooks/queries/workspace-files.ts @@ -522,13 +522,10 @@ export function useCloudStorageConfigured(enabled = true) { retry: false, staleTime: CLOUD_STORAGE_CONFIGURED_STALE_TIME, /** - * Escapes the global `retryOnMount: false`, which an infinite `staleTime` and - * `retry: false` would otherwise turn into a permanent failure: one transient - * error leaves this query errored for the tab's lifetime, and consumers fail - * closed — the upload path treats "unknown" as "not configured", so a single - * blip would disable cloud-backed uploads until a full reload. The key is - * global, so navigating or switching workspace cannot recover it either. - * Mirrors {@link useVoiceSettings}, which overrides `retryOnMount` for the same reason. + * Escapes the global `retryOnMount: false`: with an infinite `staleTime` and + * `retry: false`, one transient error leaves this query errored for the tab's lifetime, + * and the upload path reads "unknown" as "not configured" — disabling cloud uploads + * until a full reload. The key is global, so navigation cannot recover it. */ retryOnMount: true, }) diff --git a/apps/sim/hooks/selectors/use-selector-query.ts b/apps/sim/hooks/selectors/use-selector-query.ts index d315f7891fa..213afcdab91 100644 --- a/apps/sim/hooks/selectors/use-selector-query.ts +++ b/apps/sim/hooks/selectors/use-selector-query.ts @@ -71,13 +71,9 @@ export function useSelectorOptions( search: args.search, } /** - * `definition.enabled` mirrors the preconditions the definition's own `fetchList` / - * `fetchPage` assert (`ensureCredential`, `ensureKnowledgeBase`, an early `return []`), - * so it is a hard precondition for the *list*, not a default a caller may replace. - * A caller's `enabled` narrows further — it never widens — otherwise a card that only - * knows it has a selection runs a fetch that is guaranteed to reject and caches the - * rejection. Unlike {@link useSelectorOptionDetail}, resolving nothing here is never - * cheaper than the list's own preconditions, so there is no case for an override. + * `definition.enabled` mirrors the preconditions the definition's own fetchers assert, so + * it is a hard precondition for the list, not a default a caller may replace. A caller's + * `enabled` only narrows — widening would run a fetch guaranteed to reject and cache it. */ const isEnabled = args.enabled !== false && (definition.enabled?.(queryArgs) ?? true) const supportsPagination = Boolean(definition.fetchPage) diff --git a/apps/sim/lib/api/contracts/oauth-connections.ts b/apps/sim/lib/api/contracts/oauth-connections.ts index 7512b2bb442..c9c4951efd2 100644 --- a/apps/sim/lib/api/contracts/oauth-connections.ts +++ b/apps/sim/lib/api/contracts/oauth-connections.ts @@ -91,7 +91,7 @@ export const oauthTokenPostQuerySchema = z.object({ userId: z.string().min(1).optional(), }) -export const oauthTokenResponseSchema = z.object({ +const oauthTokenResponseSchema = z.object({ accessToken: z.string(), idToken: z.string().optional(), instanceUrl: z.string().optional(), diff --git a/apps/sim/lib/auth/credential-access.ts b/apps/sim/lib/auth/credential-access.ts index 75ef5a952a2..719cd30f8ca 100644 --- a/apps/sim/lib/auth/credential-access.ts +++ b/apps/sim/lib/auth/credential-access.ts @@ -66,11 +66,9 @@ export async function authorizeCredentialUse( } /** - * Credential authorization for a caller whose authentication has already been - * resolved. {@link authorizeCredentialUse} is the HTTP-request wrapper over - * this; in-process callers (the tool executor) construct the same - * {@link AuthResult} directly instead of re-authenticating over HTTP, so both - * paths run one identical authorization rule. + * Credential authorization for an already-authenticated caller. + * {@link authorizeCredentialUse} is the HTTP wrapper; in-process callers build the same + * {@link AuthResult} directly, so both paths run one identical rule. */ export async function authorizeCredentialUseForAuth( auth: AuthResult, diff --git a/apps/sim/lib/billing/core/subscription.ts b/apps/sim/lib/billing/core/subscription.ts index 962afd1625a..d697c9fd7fa 100644 --- a/apps/sim/lib/billing/core/subscription.ts +++ b/apps/sim/lib/billing/core/subscription.ts @@ -457,10 +457,8 @@ async function resolveOrganizationEnterprisePlan(organizationId: string): Promis * Check if an organization has an enterprise plan * Used for Access Control (Permission Groups) feature gating * - * Request-memoized: settings renders gate several sections on the same - * organization's plan, and the plan cannot change mid-render. Outside a Server - * Component render React evaluates the resolver normally, so routes, tools, and - * background work re-read exactly as before. + * Request-memoized: a settings render gates several sections on the same + * organization's plan, and it cannot change mid-render. */ export const isOrganizationOnEnterprisePlan = cache(resolveOrganizationEnterprisePlan) @@ -621,11 +619,10 @@ async function hasWorkspaceTierAccess( * Shared by the inbox (Sim Mailer), live sync, and custom sandboxes, which all * sit on the same entitlement tier. * - * Request-memoized because those features are gated side by side on the same - * settings render, each otherwise repeating the identical workspace and - * subscription reads. The per-feature deployment and env short-circuits live in - * the exported wrappers and still run per call. Outside a Server Component - * render React evaluates this normally. + * Request-memoized: these features are gated side by side on one settings render, + * each otherwise repeating the identical workspace and subscription reads. The + * per-feature deployment and env short-circuits live in the wrappers and still run + * per call. */ const hasMaxTierWorkspaceAccess = cache( (workspaceId: string): Promise => hasWorkspaceTierAccess(workspaceId, isMaxTier) diff --git a/apps/sim/lib/oauth/token-resolution.ts b/apps/sim/lib/oauth/token-resolution.ts index 86ee15114fd..4ee18823f8b 100644 --- a/apps/sim/lib/oauth/token-resolution.ts +++ b/apps/sim/lib/oauth/token-resolution.ts @@ -20,18 +20,14 @@ import { extractZohoDeskBaseFromScope } from '@/tools/zoho_desk/host-allowlist' const logger = createLogger('OAuthTokenResolution') /** - * Minimal duck type of the inbound HTTP request, used only so audit rows can - * record the caller's IP and user agent. In-process callers have no inbound - * request and omit it; the row is then written without forensic headers. + * Duck type of the inbound request, used only so audit rows record IP and user agent. + * In-process callers omit it. */ export interface CredentialAuditRequest { headers: { get(name: string): string | null } } -/** - * Token material a resolved credential yields. It is the route's response body, so it - * comes from the contract rather than a parallel declaration that could drift from it. - */ +/** Token material a resolved credential yields; taken from the contract so it cannot drift. */ export type CredentialTokenPayload = OAuthTokenResponse export interface ResolveCredentialTokenInput { @@ -65,7 +61,6 @@ function recordCredentialAccess(params: { resourceId: string providerId: string | null | undefined credentialType: 'oauth' | 'service_account' - extraMetadata?: Record auditRequest?: CredentialAuditRequest }): void { const { actorId, workspaceId, resourceId, providerId, credentialType } = params @@ -79,7 +74,6 @@ function recordCredentialAccess(params: { metadata: { provider: providerId, credentialType, - ...params.extraMetadata, }, request: params.auditRequest, }) @@ -96,12 +90,9 @@ function recordCredentialAccess(params: { } /** - * Projects a stored OAuth credential plus its (possibly refreshed) access token - * into the wire payload every surface returns. Provider-specific hosts live in - * the credential's scope string and are extracted through shared, allowlisted - * helpers — never a local regex, since these values are injected into tool - * calls that carry the token. Zoho Desk's data-center-scoped REST base is one - * such value, surfaced as `apiDomain` so callers never assume a host. + * Projects a stored OAuth credential plus its access token into the wire payload. + * Provider hosts come out of the scope string through shared allowlisted helpers, never a + * local regex — these values are injected into tool calls that carry the token. */ function buildOAuthTokenPayload( credential: { providerId: string; scope?: string | null; idToken?: string | null }, @@ -160,14 +151,9 @@ export async function completeOAuthCredentialToken(params: { } /** - * Authorized application operation behind `POST /api/auth/oauth/token`. - * - * Given an already-authenticated caller, authorizes use of the credential, - * mints or refreshes its token, records the credential-access trail, and - * returns either the token payload or the exact status/error the HTTP surface - * projects. Every surface that needs a credential token — the route and the - * in-process tool executor — goes through here, so authorization, token - * refresh, and audit cannot drift between them. + * Authorized application operation behind `POST /api/auth/oauth/token`. Every surface that + * needs a credential token — the route and the in-process tool executor — goes through + * here, so authorization, refresh, and audit cannot drift between them. * * @param auth Result of authenticating the caller (session or internal JWT). */ @@ -196,14 +182,16 @@ export async function resolveCredentialToken( return { ok: false, status: 400, error: 'impersonateEmail must be a valid email address' } } - const resolved = await resolveOAuthAccountId(credentialId) + /** + * Both branches below authorize with the same arguments, and neither read depends + * on the other, so they resolve together — this runs per credentialed tool call. + */ + const [resolved, authz] = await Promise.all([ + resolveOAuthAccountId(credentialId), + authorizeCredentialUseForAuth(auth, { credentialId, workflowId, callerUserId }), + ]) if (resolved?.credentialType === 'service_account' && resolved.credentialId) { - const authz = await authorizeCredentialUseForAuth(auth, { - credentialId, - workflowId, - callerUserId, - }) if (!authz.ok) { return { ok: false, status: 403, error: authz.error || 'Unauthorized' } } @@ -223,7 +211,7 @@ export async function resolveCredentialToken( recordCredentialAccess({ actorId: saActorId, workspaceId: saWorkspaceId, - resourceId: resolved.credentialId ?? credentialId, + resourceId: resolved.credentialId, providerId: resolved.providerId, credentialType: 'service_account', auditRequest, @@ -279,11 +267,6 @@ export async function resolveCredentialToken( } } - const authz = await authorizeCredentialUseForAuth(auth, { - credentialId, - workflowId, - callerUserId, - }) if (!authz.ok || !authz.credentialOwnerUserId) { return { ok: false, status: 403, error: authz.error || 'Unauthorized' } } diff --git a/apps/sim/lib/permissions/super-user.ts b/apps/sim/lib/permissions/super-user.ts index 72d9cd71e8f..0e17b8b3afa 100644 --- a/apps/sim/lib/permissions/super-user.ts +++ b/apps/sim/lib/permissions/super-user.ts @@ -43,9 +43,8 @@ export async function verifyEffectiveSuperUser(userId: string): Promise<{ * tolerates the replica's bounded staleness (admin role rarely changes). Falls back * to the primary when no replica is configured. * - * Request-memoized: an account-settings render checks the same viewer in both - * the layout and the page. Outside a Server Component render React evaluates - * the reader normally, so routes and feature-flag lookups are unaffected. + * Request-memoized: an account-settings render checks the same viewer in both the + * layout and the page. */ export const isPlatformAdmin = cache(async (userId: string): Promise => { const [row] = await dbReplica diff --git a/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts b/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts index 365eaedd936..fad6d9750a9 100644 --- a/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts +++ b/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts @@ -1153,11 +1153,8 @@ function workspaceFileScopeCondition(workspaceId: string, scope: WorkspaceFileSc } /** - * The columns {@link mapWorkspaceFileRecord} actually reads. The list reads below are - * workspace-wide, so `select()` would pull five columns no reader projects — `context`, - * `chatId`, `messageId`, `displayName`, `secretProvenanceVersion` — for every row of the - * scan. Narrowing here is invisible on the wire (the route contract already strips them) - * and cuts what the database ships per row. + * The columns {@link mapWorkspaceFileRecord} reads. These list reads are workspace-wide, + * so `select()` would ship five unprojected columns for every row of the scan. */ const workspaceFileListColumns = { id: workspaceFiles.id, diff --git a/apps/sim/lib/workflows/persistence/save-normalized-state.ts b/apps/sim/lib/workflows/persistence/save-normalized-state.ts index 28f98d7c714..6ecfb9cb5c0 100644 --- a/apps/sim/lib/workflows/persistence/save-normalized-state.ts +++ b/apps/sim/lib/workflows/persistence/save-normalized-state.ts @@ -5,6 +5,7 @@ import { assertWorkflowMutable, authorizeWorkflowByWorkspacePermission, WorkflowLockedError, + type WorkflowWorkspaceAuthorizationResult, } from '@sim/platform-authz/workflow' import { eq } from 'drizzle-orm' import type { z } from 'zod' @@ -25,10 +26,8 @@ export type SaveWorkflowNormalizedStateResult = | { success: false; status: number; error: string; details?: string } /** - * Validates an untrusted workflow-state blob against the same schema the - * `PUT /api/workflows/[id]/state` contract applies. In-process callers holding a - * persisted blob (e.g. a copilot checkpoint) run it through here so they get the - * identical coercion and rejection the HTTP hop used to give them. + * Validates an untrusted state blob against the same schema `PUT /api/workflows/[id]/state` + * applies, so in-process callers get the coercion and rejection the HTTP hop gave them. */ export function parseWorkflowStateForPersistence( value: unknown @@ -37,31 +36,31 @@ export function parseWorkflowStateForPersistence( } /** - * Writes a complete workflow state to the normalized tables. - * - * Owns everything the state write means: write authorization, the mutability - * (lock) check, block/edge preparation, the row-locked save transaction, the - * `lastSynced`/variables update, custom-tool extraction, and the socket-server - * notification. Every surface that replaces a workflow's state — the PUT route - * and the copilot checkpoint revert — calls this, so none of those steps can be + * Writes a complete workflow state to the normalized tables: write authorization, + * the lock check, block/edge preparation, the row-locked save transaction, + * `lastSynced`/variables, custom-tool extraction, and the socket notification. + * Every surface that replaces a workflow's state calls this, so no step can be * skipped by going through a different door. * - * Every refusal, the workflow lock included, comes back as a failure result so - * callers need only one branch to present. + * Every refusal, the lock included, comes back as a failure result — callers need + * one branch. + * + * `authorization` lets a caller that already resolved the same decision hand it in + * rather than pay for it twice; it must be the `write` decision for this workflow + * and user. */ export async function saveWorkflowNormalizedState(params: { requestId: string workflowId: string userId: string state: WorkflowStateContractOutput + authorization?: WorkflowWorkspaceAuthorizationResult }): Promise { const { requestId, workflowId, userId, state } = params - const authorization = await authorizeWorkflowByWorkspacePermission({ - workflowId, - userId, - action: 'write', - }) + const authorization = + params.authorization ?? + (await authorizeWorkflowByWorkspacePermission({ workflowId, userId, action: 'write' })) const workflowData = authorization.workflow if (!workflowData) { diff --git a/apps/sim/lib/workspace-files/queries.ts b/apps/sim/lib/workspace-files/queries.ts index d4e4cd11892..95c856903eb 100644 --- a/apps/sim/lib/workspace-files/queries.ts +++ b/apps/sim/lib/workspace-files/queries.ts @@ -16,16 +16,13 @@ import { * * Callers authorize the viewer against `workspaceId` first. * - * `maxRows` bounds the result for a caller that will only use the list if the whole - * workspace fits a payload budget: the read stops one row past the budget and returns - * `null` on overflow rather than the prefix, which is what stops a caller presenting a - * truncated read as the workspace's files. The two reads still run concurrently — the - * workspaces under the budget are the common case, and serializing them to save a share - * read on the rare oversized one would tax every normal request to do it. + * `maxRows` bounds the result for a caller that only uses the list if the whole workspace + * fits a payload budget: on overflow it returns `null` rather than the prefix, so a + * truncated read is never presented as the workspace's files. The two reads still run + * concurrently, since under-budget workspaces are the common case. * - * `throwOnError` propagates a failed file read instead of letting it degrade to an empty - * list. A caller seeding a cache needs that distinction: an empty list would be cached - * as authoritative, telling the user the workspace has no files. + * `throwOnError` propagates a failed file read instead of degrading to an empty list, + * which a cache seed would store as authoritative. */ export async function listWorkspaceFilesWithShares( workspaceId: string, @@ -36,8 +33,8 @@ export async function listWorkspaceFilesWithShares( const [files, shares] = await Promise.all([ listWorkspaceFiles(workspaceId, { scope, - ...(maxRows === undefined ? {} : { limit: maxRows + 1 }), - ...(options?.throwOnError ? { throwOnError: true } : {}), + limit: maxRows === undefined ? undefined : maxRows + 1, + throwOnError: options?.throwOnError, }), getWorkspaceShares('file', workspaceId), ]) diff --git a/apps/sim/lib/workspaces/permissions/utils.ts b/apps/sim/lib/workspaces/permissions/utils.ts index 8815e96d7a9..9634157d18d 100644 --- a/apps/sim/lib/workspaces/permissions/utils.ts +++ b/apps/sim/lib/workspaces/permissions/utils.ts @@ -107,23 +107,15 @@ async function selectWorkspaceWithOwner( } /** - * Request-memoized plain workspace read, keyed by id alone. A single Server - * Component render pass resolves the same workspace row through several - * independent gates, so without this the row is re-read once per gate. + * Request-memoized plain workspace read, keyed by id alone. One render pass resolves the + * same row through several independent gates. * - * Keyed on the id and not on archived visibility on purpose: the gates disagree - * about whether they want archived workspaces, and memoizing that argument would - * give each answer its own entry and dedupe nothing. Reading the superset once - * and applying the caller's visibility below is what makes the two agree on a - * single query. + * Keyed on the id and NOT on archived visibility: the gates disagree about archived + * workspaces, so memoizing that argument would give each answer its own entry and dedupe + * nothing. * - * Outside a Server Component render React evaluates this normally and retains - * nothing, so API routes and background work are unaffected — verified against this - * Next version in both dev and production builds. - * - * The row this returns is SHARED by every consumer in the render, including the - * `workspace` on each viewer's access result. Treat it as immutable: an in-place edit - * would poison every gate that reads it for the rest of the pass. + * The returned row is SHARED by every consumer in the render. Treat it as immutable — an + * in-place edit poisons every gate that reads it for the rest of the pass. */ const readWorkspaceWithOwner = cache( (workspaceId: string): Promise => diff --git a/apps/sim/tools/index.ts b/apps/sim/tools/index.ts index 99b2cd2ba5d..7a708b9db0b 100644 --- a/apps/sim/tools/index.ts +++ b/apps/sim/tools/index.ts @@ -1743,7 +1743,7 @@ async function executeToolImplementation( // including failing closed when the run carries no user id. const { resolveCredentialToken } = await import('@/lib/oauth/token-resolution') const result = await resolveCredentialToken( - { success: true, authType: 'internal_jwt', ...(userId ? { userId } : {}) }, + { success: true, authType: 'internal_jwt', userId }, { requestId, credentialId: contextParams.credential as string,