improvement(queries): fix a second row-cache collision, and make the memoized workspace read actually dedupe

Author: waleedlatif1

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,

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', {

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) {

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,

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')

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)

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<void> {
   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

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,15 +94,17 @@ 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(
     () =>
       Promise.all([
         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]

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'
 

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,