fix(plan): preserve fork scope and benchmark redaction boundaries

Author: Sg312

apps/sim/app/api/mothership/chats/[chatId]/fork/route.test.ts

@@ -1,4 +1,4 @@
-import { copilotChats, member } from '@sim/db/schema'
+import { copilotChats, member, workspace } from '@sim/db/schema'
 import { createSessionPrincipal } from '@sim/testing/factories/principal.factory'
 import { createRouteContext } from '@sim/testing/helpers/http'
 import { authBanMock } from '@sim/testing/mocks/auth-ban.mock'
@@ -145,17 +145,26 @@ function createRequest(chatId: string, body?: unknown) {
   })
 }
 
+function mockWorkspaceForkRows(
+  chat: Omit<typeof parentRow, 'resources'> & { resources: unknown[] } = parentRow
+) {
+  resetDbChainMock()
+  queueTableRows(copilotChats, [chat])
+  queueTableRows(copilotChats, [chat])
+  queueTableRows(workspace, [{ organizationId: null, archivedAt: null }])
+  queueTableRows(copilotChats, [{ id: chat.id, memorySpaceId: null }])
+  dbChainMockFns.returning.mockResolvedValue([{ id: 'row-id', workspaceId: 'ws-1' }])
+}
+
 describe('POST /api/mothership/chats/[chatId]/fork', () => {
   beforeEach(() => {
-    resetDbChainMock()
+    mockWorkspaceForkRows()
     setEnv({ COPILOT_API_KEY: undefined })
     copilotHttpMockFns.mockAuthenticateCopilotRequestSessionOnly.mockResolvedValue({
       userId: 'user-1',
       isAuthenticated: true,
       principal: createSessionPrincipal(),
     })
-    dbChainMockFns.limit.mockResolvedValue([parentRow])
-    dbChainMockFns.returning.mockResolvedValue([{ id: 'row-id', workspaceId: 'ws-1' }])
     mockListForkableChatFiles.mockResolvedValue([])
     mockLoadCopilotChatMessages.mockResolvedValue(threeMessages)
     mockPlanChatFileCopies.mockReturnValue({
@@ -180,12 +189,13 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => {
   })
 
   it('forks organization history under current membership without inventing workspace files', async () => {
-    dbChainMockFns.limit.mockReset()
+    resetDbChainMock()
     const chat = { ...parentRow, workspaceId: null, organizationId: 'org-1', resources: [] }
     queueTableRows(copilotChats, [chat])
     queueTableRows(copilotChats, [chat])
     queueTableRows(member, [{ role: 'member' }])
     queueTableRows(copilotChats, [{ id: chat.id }])
+    dbChainMockFns.returning.mockResolvedValue([{ id: 'row-id', workspaceId: null }])
     const attachments = [
       {
         id: 'upload-1',
@@ -220,7 +230,7 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => {
     // The fork shares its attachment keys with the source, and cleanup deletes an unreferenced
     // key only after deleting the source row: publishing without the source would leave the
     // fork pointing at bytes that cleanup is about to delete.
-    dbChainMockFns.limit.mockReset()
+    resetDbChainMock()
     const chat = { ...parentRow, workspaceId: null, organizationId: 'org-1', resources: [] }
     queueTableRows(copilotChats, [chat])
     queueTableRows(copilotChats, [chat])
@@ -239,7 +249,7 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => {
   it.each(['membership', 'capability'] as const)(
     'denies organization forks after %s revocation before copying',
     async (revocation) => {
-      dbChainMockFns.limit.mockReset()
+      resetDbChainMock()
       const chat = { ...parentRow, workspaceId: null, organizationId: 'org-1' }
       queueTableRows(copilotChats, [chat])
       queueTableRows(copilotChats, [chat])
@@ -254,7 +264,7 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => {
   )
 
   it('404s when the chat belongs to another user', async () => {
-    dbChainMockFns.limit.mockResolvedValue([{ ...parentRow, userId: 'someone-else' }])
+    mockWorkspaceForkRows({ ...parentRow, userId: 'someone-else' })
     const res = await POST(createRequest('chat-1'), createRouteContext({ chatId: 'chat-1' }))
     expect(res.status).toBe(404)
     expect(dbChainMockFns.transaction).not.toHaveBeenCalled()
@@ -447,19 +457,17 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => {
 
   it('drops a Sources tab whose response is past the cut', async () => {
     const kept = { type: 'sources', id: 'cited-sources', title: 'Sources' }
-    dbChainMockFns.limit.mockResolvedValue([
-      { ...parentRow, resources: [{ ...kept, sources: { messageId: 'msg-3' } }] },
-    ])
+    mockWorkspaceForkRows({
+      ...parentRow,
+      resources: [{ ...kept, sources: { messageId: 'msg-3' } }],
+    })
     await POST(createRequest('chat-1'), createRouteContext({ chatId: 'chat-1' }))
     expect(dbChainMockFns.values).toHaveBeenCalledWith(expect.objectContaining({ resources: [] }))
 
-    dbChainMockFns.values.mockClear()
-    dbChainMockFns.limit.mockResolvedValue([
-      {
-        ...parentRow,
-        resources: [{ ...kept, sources: { messageId: 'live-id', requestId: 'req-2' } }],
-      },
-    ])
+    mockWorkspaceForkRows({
+      ...parentRow,
+      resources: [{ ...kept, sources: { messageId: 'live-id', requestId: 'req-2' } }],
+    })
     mockLoadCopilotChatMessages.mockResolvedValue([
       threeMessages[0],
       { ...threeMessages[1], requestId: 'req-2' },
@@ -479,18 +487,16 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => {
     // copies the kept upload AND the pre-cut apple; only the post-cut banana
     // stays behind, so only its resource is dropped — not left pointing at
     // the source chat.
-    dbChainMockFns.limit.mockResolvedValue([
-      {
-        ...parentRow,
-        resources: [
-          { type: 'file', id: OLD_FILE_ID, title: 'cat.png' },
-          { type: 'file', id: 'wf_apple', title: 'apple.png' },
-          { type: 'file', id: 'wf_banana', title: 'banana.png' },
-          { type: 'file', id: 'wf_shared', title: 'shared.pdf' },
-          { type: 'workflow', id: 'wflow-1', title: 'My flow' },
-        ],
-      },
-    ])
+    mockWorkspaceForkRows({
+      ...parentRow,
+      resources: [
+        { type: 'file', id: OLD_FILE_ID, title: 'cat.png' },
+        { type: 'file', id: 'wf_apple', title: 'apple.png' },
+        { type: 'file', id: 'wf_banana', title: 'banana.png' },
+        { type: 'file', id: 'wf_shared', title: 'shared.pdf' },
+        { type: 'workflow', id: 'wflow-1', title: 'My flow' },
+      ],
+    })
     // Every chat-owned file of the source chat in the single read; messageId
     // drives the in-memory cut.
     mockListForkableChatFiles.mockResolvedValue([

apps/sim/app/api/mothership/chats/route.test.ts

@@ -31,15 +31,16 @@ vi.mock('@/lib/mothership/chat-status', () => mothershipChatStatusMock)
 
 vi.mock('@/lib/posthog/server', () => posthogServerMock)
 
-import { GET } from '@/app/api/mothership/chats/route'
+import { OrchestrationError } from '@/lib/core/orchestration/types'
+import { GET, POST } from '@/app/api/mothership/chats/route'
 
 function createRequest(workspaceId: string) {
   return createMockRequest({
     url: `http://localhost:3000/api/mothership/chats?workspaceId=${workspaceId}`,
   })
 }
 
-describe('GET /api/mothership/chats', () => {
+describe('/api/mothership/chats', () => {
   beforeEach(() => {
     resetDbChainMock()
     workspaceContextMockFns.mockResolveActiveWorkspaceApplicationContext.mockImplementation(
@@ -121,6 +122,18 @@ describe('GET /api/mothership/chats', () => {
     ])
   })
 
+  it.each(['forbidden', 'not_found'] as const)(
+    'identifies the workspace in a POST %s refusal',
+    async (code) => {
+      workspaceContextMockFns.mockResolveActiveWorkspaceApplicationContext.mockRejectedValueOnce(
+        new OrchestrationError(code, 'Workspace unavailable')
+      )
+      const response = await POST(createMockRequest('POST', { workspaceId: 'ws-denied' }))
+      expect(response.status).toBe(403)
+      expect(await response.json()).toEqual({ error: 'Workspace access denied' })
+    }
+  )
+
   afterAll(() => {
     resetDbChainMock()
   })

apps/sim/app/api/mothership/chats/route.ts

@@ -75,6 +75,7 @@ export const GET = withRouteHandler(async (request: NextRequest) => {
  * Creates an empty mothership chat and returns its ID.
  */
 export const POST = withRouteHandler(async (request: NextRequest) => {
+  let isWorkspaceRequest = false
   try {
     const { userId, isAuthenticated, principal } = await authenticateCopilotRequestSessionOnly()
     if (!isAuthenticated || !userId) {
@@ -84,6 +85,7 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
     const validation = await parseRequest(createMothershipChatContract, request, {})
     if (!validation.success) return validation.response
     const { workspaceId, organizationId, mode } = validation.data.body
+    isWorkspaceRequest = Boolean(workspaceId)
 
     if (organizationId) {
       if (!principal) return createUnauthorizedResponse()
@@ -110,11 +112,13 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
     return NextResponse.json({ success: true, id: chat.id })
   } catch (error) {
     const code = asOrchestrationError(error)?.code
+    if (
+      isWorkspaceAccessDeniedError(error) ||
+      (isWorkspaceRequest && (code === 'not_found' || code === 'forbidden'))
+    )
+      return createForbiddenResponse('Workspace access denied')
     if (code === 'not_found' || code === 'forbidden')
       return createForbiddenResponse('Organization access denied')
-    if (isWorkspaceAccessDeniedError(error)) {
-      return createForbiddenResponse('Workspace access denied')
-    }
     logger.error('Error creating mothership chat:', error)
     return createInternalServerErrorResponse('Failed to create chat')
   }

apps/sim/lib/benchmarks/artifacts.test.ts

@@ -50,6 +50,26 @@ describe('benchmark reference integrity', () => {
     expect(() => validateBenchmarkRedaction(artifacts)).not.toThrow()
   })
 
+  it('rejects a repeated gold answer that survives outside the masks', () => {
+    expect(() =>
+      validateBenchmarkRedaction({
+        ...artifacts,
+        redactedSpec:
+          '[[BLANK:owner]] owns follow-up until Engineering accepts. Support monitors it.',
+      })
+    ).toThrow()
+  })
+
+  it('does not mistake a blank identifier for a surviving gold answer', () => {
+    expect(() =>
+      validateBenchmarkRedaction({
+        referenceSpec: 'owner approves the request.',
+        redactedSpec: '[[BLANK:owner]] approves the request.',
+        blanks: [{ id: 'owner', answer: 'owner' }],
+      })
+    ).not.toThrow()
+  })
+
   it('rejects a redaction that quietly changes an unmasked requirement', () => {
     expect(() =>
       validateBenchmarkRedaction({ ...artifacts, redactedSpec: '[[BLANK:owner]] owns everything.' })

apps/sim/lib/benchmarks/artifacts.ts

@@ -59,6 +59,15 @@ export function validateBenchmarkRedaction(
       'Every blank must appear as [[BLANK:id]], and replacing the masks with their answers must restore the reference exactly'
     )
   }
+  const visiblePassages = artifacts.redactedSpec.split(/\[\[BLANK:[A-Za-z0-9_-]+\]\]/)
+  if (
+    [...answers.values()].some((answer) => visiblePassages.some((text) => text.includes(answer)))
+  ) {
+    throw new OrchestrationError(
+      'validation',
+      'Every occurrence of a selected answer must be masked'
+    )
+  }
 }
 
 /** Editing an upstream artifact atomically discards results derived from its previous value. */

apps/sim/lib/billing/core/usage-log.test.ts

@@ -456,6 +456,7 @@ describe('ledger aggregates', () => {
 
   beforeEach(() => {
     installSharedDbMocks()
+    vi.spyOn(Date, 'now').mockReturnValue(billingPeriod.start.getTime())
   })
 
   for (const aggregate of aggregates) {

apps/sim/lib/billing/webhooks/stripe-sync-convergence.integration.ts

@@ -402,7 +402,7 @@ function startParkedTransaction(
     }
     throw new Error('No transaction ever waited on the parked one')
   }
-  return { done, release, untilBlocking }
+  return { ready: holderPid, done, release, untilBlocking }
 }
 
 describe('cancel_at_period_end sync', () => {
@@ -1031,6 +1031,7 @@ describe('Team activation', () => {
         reason: 'admin-cancel-at-period-end',
       })
     })
+    await cancelling.ready
     const activating = testDatabase.transaction((tx) =>
       ensureTeamOrganizationForAcceptance({
         billingOwnerUserId: owner.id,
@@ -1039,8 +1040,11 @@ describe('Team activation', () => {
         workspaceIdsToAttach: [],
       })
     )
-    await cancelling.untilBlocking()
-    cancelling.release()
+    try {
+      await cancelling.untilBlocking()
+    } finally {
+      cancelling.release()
+    }
     await cancelling.done
     await expect(activating).resolves.toMatchObject({ success: true })
     expect((await storedSubscription(subscriptionId)).cancelAtPeriodEnd).toBe(false)
@@ -1082,9 +1086,13 @@ describe('operator retry', () => {
         })
       }
     )
+    await writing.ready
     const requeuing = requeueFromAdminApi(pauseSync)
-    await writing.untilBlocking()
-    writing.release()
+    try {
+      await writing.untilBlocking()
+    } finally {
+      writing.release()
+    }
 
     await expect(Promise.all([writing.done, requeuing])).resolves.toBeDefined()
     await deliverUnrelatedUpdate(pro.stripeSubscriptionId)
@@ -1123,14 +1131,18 @@ describe('operator retry', () => {
         })
       }
     )
+    await writing.ready
     const retrying = requestDashboardSubscriptionCancellation({
       organizationId: org.organizationId,
       operationId,
       timing: 'period_end',
       actor,
     })
-    await writing.untilBlocking()
-    writing.release()
+    try {
+      await writing.untilBlocking()
+    } finally {
+      writing.release()
+    }
 
     await expect(Promise.all([writing.done, retrying])).resolves.toBeDefined()
     expect((await storedSubscription(org.subscriptionId)).cancelAtPeriodEnd).toBe(true)

apps/sim/lib/mothership/chat/application/fork.ts

@@ -157,7 +157,14 @@ export const forkChat = defineAuthorizedChatUseCase({
 
       /** Publish only after both the file bytes and the worker conversation are prepared. */
       await db.transaction(async (tx) => {
-        if (parent.workspaceId) await lockActiveWorkspace(tx, parent.workspaceId)
+        if (parent.workspaceId) {
+          const currentWorkspace = await lockActiveWorkspace(tx, parent.workspaceId)
+          if (currentWorkspace.organizationId !== context.workspaceOrganizationId)
+            throw new OrchestrationError(
+              'conflict',
+              'Workspace organization changed. Refresh and retry.'
+            )
+        }
         /**
          * The fork can share keys with its source (organization attachments, files whose copy
          * failed), and chat cleanup deletes a shared key once no remaining chat references it,