fix(mothership): preserve staging authorization and resumed billing

Author: Sg312

apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/usage-upgrade-display.test.tsx

@@ -15,7 +15,7 @@ vi.mock('@/lib/core/config/deployment-shape', async (importOriginal) => ({
   useDeploymentShape: () => ({ hosted: true }),
 }))
 vi.mock('@/app/workspace/[workspaceId]/providers/workspace-host-provider', () => ({
-  useWorkspaceHostContext: () => ({
+  useOptionalWorkspaceHostContext: () => ({
     workspace: { id: 'workspace', billedAccountUserId: 'owner' },
     hostOrganizationId: 'organization',
     viewer: { isHostOrganizationAdmin: false },

apps/sim/lib/mcp/workflow-tool-schema.ts

@@ -26,7 +26,9 @@ export function sanitizeToolName(name: string): string {
  * This converts the workflow's input format definition to JSON Schema format
  * that MCP clients can use to understand tool parameters.
  */
-export function generateToolInputSchema(inputFormat: InputFormatField[]): McpToolSchema {
+export function generateToolInputSchema(
+  inputFormat: InputFormatField[]
+): McpToolSchema & { properties: Record<string, McpToolSchemaProperty> } {
   const schema = z.toJSONSchema(z.object(generateWorkflowInputShape(inputFormat)), {
     target: 'draft-07',
     io: 'input',

apps/sim/lib/mothership/agent-cli/file-provenance.test.ts

@@ -41,6 +41,7 @@ vi.mock('@/lib/workspace-files/application/resolve-rendered-workspace-artifact',
   resolveRenderedWorkspaceArtifact: mocks.render,
 }))
 vi.mock('@/lib/execution/remote-sandbox/session-files', () => ({
+  SESSION_SANDBOX_HOME: '/home/user',
   readSessionSandboxFile: vi.fn(),
   writeSessionSandboxFile: vi.fn(),
 }))

apps/sim/lib/mothership/agent-cli/run-cli.test.ts

@@ -6,6 +6,7 @@ const { readFile, writeFile, embedded } = vi.hoisted(() => ({
   embedded: vi.fn(),
 }))
 vi.mock('@/lib/execution/remote-sandbox/session-files', () => ({
+  SESSION_SANDBOX_HOME: '/home/user',
   readSessionSandboxFile: readFile,
   writeSessionSandboxFile: writeFile,
 }))

apps/sim/lib/mothership/application/authorize-chat-callback.ts

@@ -8,6 +8,8 @@ import {
   type BillingAttributionSnapshot,
   checkAttributedBillingBlocks,
 } from '@/lib/billing/core/billing-attribution'
+import { defineAuthorizedWorkspaceUseCase } from '@/lib/core/application/authorized-workspace-use-case'
+import { OrchestrationError } from '@/lib/core/orchestration/types'
 import { chatOperations } from '@/lib/mothership/application/operations'
 import {
   COPILOT_APPLICATION_DELEGATION_TTL_MS,
@@ -19,8 +21,6 @@ import {
   COPILOT_VALIDATION_PURPOSE,
   type CopilotValidationPurpose,
 } from '@/lib/mothership/generated/billing-protocol-v1'
-import { defineAuthorizedWorkspaceUseCase } from '@/lib/core/application/authorized-workspace-use-case'
-import { OrchestrationError } from '@/lib/core/orchestration/types'
 import { resolveActiveWorkspaceApplicationContext } from '@/lib/workspaces/application/workspace-context'
 
 const CALLBACK_AUDIENCE = 'sim:copilot-callback'
@@ -57,9 +57,10 @@ interface CopilotChatCallbackContext {
   chatId?: string
   delegationId: string
   purpose: Exclude<CopilotValidationPurpose, 'new-turn'>
+  mode?: 'assistant' | 'agent'
 }
 
-/** Reauthorizes the original server-owned scope across a Go lifecycle boundary. */
+/** Reauthorizes the original server-owned scope across a model lifecycle boundary. */
 export async function authorizeCopilotChatCallback(context: CopilotChatCallbackContext) {
   if (context.organizationId) {
     if (!context.chatId || context.workspaceId) {
@@ -75,7 +76,7 @@ export async function authorizeCopilotChatCallback(context: CopilotChatCallbackC
         ttlMs: COPILOT_APPLICATION_DELEGATION_TTL_MS,
       }
     )
-    await authorizeOrganizationChatDelegation.execute({ principal })
+    await authorizeOrganizationChatDelegation.execute({ principal, mode: context.mode })
     return
   }
   if (!context.workspaceId) return

apps/sim/lib/mothership/chat/child-failure-transcript.test.ts

@@ -221,7 +221,8 @@ describe('child failure transcript boundaries', () => {
         error: failure,
       })
     )
-    expect(markup).toContain('Inspect report — Failed')
+    expect(markup).toContain('Inspect report')
+    expect(markup).not.toContain('Inspect report — Failed')
     expect(markup).toContain(failure)
     expect(markup).not.toContain('<button')
     const sibling = renderToStaticMarkup(

apps/sim/lib/mothership/chat/organization-chats.test.ts

@@ -4,8 +4,8 @@ import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
 import { OrchestrationError } from '@/lib/core/orchestration/types'
 import { createTrustedOrganizationCopilotPrincipal } from '@/lib/mothership/auth/application-delegation'
 import {
-  authorizeOrganizationChatCancellation,
   authorizeOrganizationChat,
+  authorizeOrganizationChatCancellation,
   authorizeOrganizationChatDelegation,
   authorizeOrganizationChatEvents,
   createOrganizationChat,
@@ -226,6 +226,18 @@ describe('organization Build admission', () => {
     })
     expect(permissionConfig).not.toHaveBeenCalled()
   })
+  it('rechecks Build permission for a delegated continuation while allowing Search', async () => {
+    authorize.mockResolvedValue({ userId: 'member-1', organizationId: 'org-1', role: 'owner' })
+    permissionConfig.mockResolvedValue({ disableWorkspaceCreation: true })
+    dbChainMockFns.limit.mockResolvedValue([{ id: 'private-chat' }])
+    await expect(
+      authorizeOrganizationChatDelegation.execute({ principal: principal(), mode: 'agent' })
+    ).rejects.toThrow('Build requires permission')
+    await expect(
+      authorizeOrganizationChatDelegation.execute({ principal: principal(), mode: 'assistant' })
+    ).resolves.toMatchObject({ userId: 'member-1' })
+  })
+
   it('checks current membership before the Build permission projection', async () => {
     authorize.mockRejectedValueOnce(new Error('Membership revoked'))
     await expect(

apps/sim/lib/mothership/chat/organization-chats.ts

@@ -192,7 +192,13 @@ export const authorizeOrganizationChatCancellation = {
 
 /** A trusted service may act only on the subject's persisted private organization chat. */
 export const authorizeOrganizationChatDelegation = {
-  async execute({ principal }: { principal: OrganizationDelegatedPrincipal }) {
+  async execute({
+    principal,
+    mode,
+  }: {
+    principal: OrganizationDelegatedPrincipal
+    mode?: 'assistant' | 'agent'
+  }) {
     if (principal.serviceId !== 'copilot')
       throw new OrchestrationError('forbidden', 'Invalid conversation delegation')
     const operation = Object.values(organizationChatDelegationOperations).find(
@@ -216,6 +222,7 @@ export const authorizeOrganizationChatDelegation = {
       )
       .limit(1)
     if (!chat) throw new OrchestrationError('not_found', 'Conversation not found')
+    if (mode === 'agent') await requireBuildPermission(context)
     return context
   },
 }

apps/sim/lib/mothership/request/application/controls-boundary.test.ts

@@ -1,32 +1,109 @@
 /** @vitest-environment node */
-import { queueTableRows, resetDbChainMock, schemaMock } from '@sim/testing'
+import { dbChainMockFns, queueTableRows, resetDbChainMock, schemaMock } from '@sim/testing'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
-const mocks = vi.hoisted(() => ({ latest: vi.fn(), stop: vi.fn(), signal: vi.fn() }))
+const mocks = vi.hoisted(() => ({
+  latest: vi.fn(),
+  stop: vi.fn(),
+  signal: vi.fn(),
+  permissions: vi.fn(),
+}))
 vi.mock('@/lib/mothership/async-runs/repository', () => ({
   getLatestRunForStream: mocks.latest,
   requestRunStop: mocks.stop,
-  areStreamToolExecutionsSettled: vi.fn(),
-  getUnsettledClientWorkflowExecutions: vi.fn(),
-  getUnsettledStreamSandboxProcesses: vi.fn(),
+  areStreamToolExecutionsSettled: vi.fn(async () => true),
+  getUnsettledClientWorkflowExecutions: vi.fn(async () => []),
+  getUnsettledStreamSandboxProcesses: vi.fn(async () => []),
 }))
 vi.mock('@/lib/mothership/request/session/explicit-abort', () => ({
   requestExplicitStreamAbort: mocks.signal,
 }))
+vi.mock('@/lib/mothership/request/session', () => ({
+  abortActiveStream: vi.fn(),
+  waitForPendingChatStream: vi.fn(async () => true),
+  releasePendingChatStream: vi.fn(),
+}))
+vi.mock('@/lib/permission-groups/resolve.server', async (importOriginal) => ({
+  ...(await importOriginal<typeof import('@/lib/permission-groups/resolve.server')>()),
+  getUserPermissionConfigForOrganization: mocks.permissions,
+}))
 vi.mock('@/lib/auth/ban', () => ({ getActivelyBannedUserIds: async () => [] }))
 
 import { abortRun } from '@/lib/mothership/request/application/controls'
 
 beforeEach(() => {
   vi.clearAllMocks()
   resetDbChainMock()
+  mocks.permissions
+    .mockReset()
+    .mockResolvedValue({ hideCopilot: true, disableWorkspaceCreation: true })
+  mocks.stop.mockReset().mockResolvedValue(null)
+  mocks.signal.mockReset().mockResolvedValue({ settled: true })
   mocks.latest.mockResolvedValue({
     chatId: 'chat',
     workspaceId: null,
     organizationId: 'organization',
   })
 })
 describe('abort authorization before service signaling', () => {
+  const principal = { kind: 'session', userId: 'actor', sessionId: 'session' } as const
+  const input = { streamId: 'stream', chatId: 'chat' }
+  const ownedChat = {
+    userId: 'actor',
+    workspaceId: null,
+    organizationId: 'organization',
+    type: 'mothership',
+  }
+
+  it('persists and forwards Stop for a current member after chat and Build capabilities are revoked', async () => {
+    queueTableRows(schemaMock.copilotChats, [ownedChat])
+    queueTableRows(schemaMock.copilotChats, [ownedChat])
+    queueTableRows(schemaMock.member, [{ role: 'member' }])
+    mocks.stop.mockResolvedValue({
+      chatId: 'chat',
+      workspaceId: null,
+      organizationId: 'organization',
+    })
+
+    await expect(abortRun.execute({ principal, input })).resolves.toEqual({
+      aborted: true,
+      settled: true,
+    })
+
+    expect(mocks.stop).toHaveBeenCalledWith({
+      streamId: 'stream',
+      chatId: 'chat',
+      userId: 'actor',
+      organizationId: 'organization',
+      workspaceId: undefined,
+    })
+    expect(mocks.signal).toHaveBeenCalledWith({
+      streamId: 'stream',
+      chatId: 'chat',
+      userId: 'actor',
+      timeoutMs: 3000,
+    })
+    expect(mocks.permissions).not.toHaveBeenCalled()
+  })
+
+  it('rejects a removed member before persisting or forwarding Stop', async () => {
+    queueTableRows(schemaMock.copilotChats, [ownedChat])
+    queueTableRows(schemaMock.member, [])
+    await expect(abortRun.execute({ principal, input })).rejects.toMatchObject({
+      code: 'not_found',
+    })
+    expect(mocks.stop).not.toHaveBeenCalled()
+    expect(mocks.signal).not.toHaveBeenCalled()
+  })
+
+  it('propagates a live membership lookup failure before persisting or forwarding Stop', async () => {
+    const error = new Error('membership database unavailable')
+    dbChainMockFns.limit.mockResolvedValueOnce([ownedChat]).mockRejectedValueOnce(error)
+    await expect(abortRun.execute({ principal, input })).rejects.toBe(error)
+    expect(mocks.stop).not.toHaveBeenCalled()
+    expect(mocks.signal).not.toHaveBeenCalled()
+  })
+
   it('rejects another user’s canonical chat before persisting or forwarding Stop', async () => {
     queueTableRows(schemaMock.copilotChats, [
       { userId: 'other', workspaceId: null, organizationId: 'organization', type: 'mothership' },

apps/sim/lib/mothership/request/application/recover-stream.ts

@@ -8,11 +8,13 @@ import {
 } from '@/lib/billing/core/billing-attribution'
 import { defineWorkspaceOperation } from '@/lib/core/application'
 import { defineOrganizationOperation } from '@/lib/core/application/organization-operation'
+import { isHosted } from '@/lib/core/config/env-flags'
 import { OrchestrationError } from '@/lib/core/orchestration/types'
 import { getLatestRunForStream } from '@/lib/mothership/async-runs/repository'
 import { defineAuthorizedChatUseCase } from '@/lib/mothership/chat/application/authorized-chat-use-case'
 import { resolveOwnedChatContext } from '@/lib/mothership/chat/application/context'
 import { buildOnComplete, buildOnError } from '@/lib/mothership/chat/completion'
+import { restoreBillingAdmission } from '@/lib/mothership/request/lifecycle/admission'
 import { claimRunController } from '@/lib/mothership/request/lifecycle/controller-ownership'
 import { StreamRecoveryConfigSchema } from '@/lib/mothership/request/lifecycle/recovery-config'
 import { createSSEStream } from '@/lib/mothership/request/lifecycle/start'
@@ -93,11 +95,25 @@ export const readChatStream = defineAuthorizedChatUseCase({
         await releasePendingChatStream(chatId, run.streamId, lease)
         return (await getLatestRunForStream(run.streamId, userId)) ?? run
       }
+      if (isHosted && !config.data.billingAdmission)
+        throw new OrchestrationError(
+          'forbidden',
+          'Hosted recovery is missing its original billing admission'
+        )
+      const restoredAdmission = config.data.billingAdmission
+        ? restoreBillingAdmission(config.data.billingAdmission, {
+            userId,
+            workspaceId,
+            organizationId,
+          })
+        : undefined
       const [events, billingAttribution, userPermission] = await Promise.all([
         readEvents(run.streamId, '0'),
-        organizationId
-          ? resolveOrganizationBillingAttribution({ actorUserId: userId, organizationId })
-          : resolveBillingAttribution({ actorUserId: userId, workspaceId: workspaceId! }),
+        restoredAdmission
+          ? Promise.resolve(restoredAdmission.attribution)
+          : organizationId
+            ? resolveOrganizationBillingAttribution({ actorUserId: userId, organizationId })
+            : resolveBillingAttribution({ actorUserId: userId, workspaceId: workspaceId! }),
         workspaceId
           ? getUserEntityPermissions(userId, 'workspace', workspaceId)
           : Promise.resolve(undefined),