fix(mothership): keep a started call's result after Stop, and let a blocked owner stop their run

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

@@ -19,7 +19,19 @@ export async function resolveOwnedWorkspaceChatContext(principal: Principal, cha
   return context
 }
 
-export async function resolveOwnedChatContext(principal: Principal, chatId: string) {
+export interface OwnedChatContextOptions {
+  /**
+   * Resolves the chat for an owner whose account is blocked. Only Stop sets it: halting the
+   * owner's existing work is never harmful, so it stays available after access is withheld.
+   */
+  allowBlockedOwner?: boolean
+}
+
+export async function resolveOwnedChatContext(
+  principal: Principal,
+  chatId: string,
+  options: OwnedChatContextOptions = {}
+) {
   const [chat] = await db
     .select({
       userId: copilotChats.userId,
@@ -40,7 +52,7 @@ export async function resolveOwnedChatContext(principal: Principal, chatId: stri
   ) {
     throw new OrchestrationError('not_found', 'Chat not found')
   }
-  if ((await getActivelyBannedUserIds([chat.userId])).length > 0) {
+  if (!options.allowBlockedOwner && (await getActivelyBannedUserIds([chat.userId])).length > 0) {
     throw new OrchestrationError('forbidden', 'User account is suspended')
   }
   const owner: ChatOwnerContext = chat.organizationId

apps/sim/lib/mothership/request/application/controls.ts

@@ -14,7 +14,10 @@ import {
   requestRunStop,
 } 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 {
+  type OwnedChatContextOptions,
+  resolveOwnedChatContext,
+} from '@/lib/mothership/chat/application/context'
 import { appendCopilotChatMessages } from '@/lib/mothership/chat/messages-store'
 import { deliverRunStop } from '@/lib/mothership/request/application/stop-delivery'
 import { requestStreamSteering } from '@/lib/mothership/request/session/steer'
@@ -61,12 +64,13 @@ async function resolveRunContext({
 async function resolveAdmittedRunContext(
   principal: SessionPrincipal,
   input: RunControlInput,
-  run: { chatId: string; workspaceId: string | null; organizationId?: string | null }
+  run: { chatId: string; workspaceId: string | null; organizationId?: string | null },
+  options: OwnedChatContextOptions = {}
 ) {
   if (input.chatId && input.chatId !== run.chatId) {
     throw new OrchestrationError('forbidden', 'Stream does not belong to this chat')
   }
-  const chat = await resolveOwnedChatContext(principal, run.chatId)
+  const chat = await resolveOwnedChatContext(principal, run.chatId, options)
   if (
     (run.workspaceId ?? null) !== (chat.workspaceId ?? null) ||
     (run.organizationId ?? null) !== (chat.organizationId ?? null)
@@ -90,7 +94,12 @@ async function resolveAbortContext({
   input: RunControlInput
 }) {
   const run = await getLatestRunForStream(input.streamId, principal.userId)
-  if (run?.chatId) return resolveAdmittedRunContext(principal, input, run)
+  /**
+   * An owner may stop a run already admitted for them even after their account is blocked. A
+   * blocked account cannot admit a run, so a Stop recorded ahead of admission keeps the ban check.
+   */
+  if (run?.chatId)
+    return resolveAdmittedRunContext(principal, input, run, { allowBlockedOwner: true })
   if (input.chatId) {
     const chat = await resolveOwnedChatContext(principal, input.chatId)
     if (input.workspaceId && input.workspaceId !== chat.workspaceId) {
@@ -147,7 +156,8 @@ export const abortRun = defineAuthorizedChatUseCase({
     const { chatId } = await resolveAdmittedRunContext(
       principal,
       { ...input, workspaceId, organizationId },
-      run
+      run,
+      { allowBlockedOwner: true }
     )
     /** Read alongside the signals, which need none of it. */
     const work = Promise.all([

apps/sim/lib/mothership/request/application/removal-stops-runs.integration.ts

@@ -88,6 +88,7 @@ import { createRunSegment } from '@/lib/mothership/async-runs/repository'
 import { admitChatTurn } from '@/lib/mothership/chat/application/admit-turn'
 import { deleteCopilotChat, softDeleteMothershipChat } from '@/lib/mothership/chat/deletion'
 import { chatPubSub } from '@/lib/mothership/chat-status'
+import { abortRun, steerRun } from '@/lib/mothership/request/application/controls'
 import { claimRunController } from '@/lib/mothership/request/lifecycle/controller-ownership'
 import {
   acquirePendingChatStream,
@@ -966,6 +967,74 @@ describe.runIf(Boolean(redisUrl))('removing a run’s chat, owner or access stop
     expect(await runControl(run)).toEqual(STOPPED)
   })
 
+  it('lets a blocked owner Stop their run, and neither steer nor send, nor anyone else Stop it', async () => {
+    const ownerId = await seedUser()
+    const memberId = await seedUser()
+    const workspaceId = await seedWorkspace(ownerId, [memberId])
+    const run = await seedRun({ userId: ownerId, workspaceId })
+    await db.update(user).set({ banned: true }).where(eq(user.id, ownerId))
+    const owner = createSessionPrincipal({ userId: ownerId, sessionId: generateId() })
+    const input = { streamId: run.streamId, chatId: run.chatId, workspaceId }
+
+    await expect(
+      steerRun.execute({
+        principal: owner,
+        input: { ...input, steeringId: generateId(), content: 'Keep going' },
+      })
+    ).rejects.toMatchObject({ code: 'forbidden' })
+    const sendStreamId = generateId()
+    expect(await acquirePendingChatStream(run.chatId, sendStreamId, 0)).toBe(true)
+    const lease = getLocalChatStreamLease(run.chatId, sendStreamId)
+    if (!lease) throw new Error('The send needs its chat lease')
+    const sendClaim = { normalizedKey: `run-removal:${generateId()}`, claimToken: generateId() }
+    await db
+      .insert(idempotencyKey)
+      .values({ key: sendClaim.normalizedKey, result: { claimToken: sendClaim.claimToken } })
+    await expect(
+      admitChatTurn
+        .execute({
+          principal: owner,
+          input: {
+            chatId: run.chatId,
+            runId: generateId(),
+            executionId: generateId(),
+            requestId: generateId(),
+            message: { id: sendStreamId, content: 'Start over' },
+            recovery: {
+              kind: 'interactive_stream',
+              goRoute: '/api/mothership',
+              clientToolPickupExpected: false,
+              request: {
+                userId: ownerId,
+                chatId: run.chatId,
+                messageId: sendStreamId,
+                workspaceId,
+                message: 'Start over',
+              },
+            },
+            lease,
+            sendClaim,
+            notifyWorkspaceStatus: false,
+          },
+        })
+        .finally(() => releasePendingChatStream(run.chatId, sendStreamId, lease))
+    ).rejects.toMatchObject({ code: 'forbidden' })
+    await expect(
+      abortRun.execute({
+        principal: createSessionPrincipal({ userId: memberId, sessionId: generateId() }),
+        input,
+      })
+    ).rejects.toMatchObject({ code: 'not_found' })
+    expect(await hasStop(run)).toBe(false)
+    expect(pushedStops()).toEqual([])
+
+    await expect(abortRun.execute({ principal: owner, input })).resolves.toMatchObject({
+      aborted: true,
+    })
+    expect(await hasStop(run)).toBe(true)
+    expect(pushedStops()).toEqual([run.streamId])
+  })
+
   /**
    * Admission and deletion both write the chat row, so they serialize on its lock; holding it
    * queues them in a chosen order, then releases them together.

apps/sim/lib/mothership/request/tools/executor.test.ts

@@ -627,12 +627,44 @@ describe('executeToolAndReport provenance isolation', () => {
 
   it('refuses a late tool after durable Stop closed admission', async () => {
     claimToolExecution.mockResolvedValueOnce({ outcome: 'closed' })
+    completePendingAsyncToolCall.mockImplementationOnce(async (input) => ({ ...input }))
     const tool = buildPendingToolCall()
     const result = await executeToolAndReport(tool.id, buildStreamingContext(tool), {
       userId: 'user-1',
       workflowId: 'workflow-1',
     })
-    expect(result.status).toBe('cancelled')
+    expect(result).toMatchObject({
+      status: 'cancelled',
+      message: 'Run stopped before tool admission',
+      data: { cancelled: true },
+    })
+    expect(completePendingAsyncToolCall).toHaveBeenCalledExactlyOnceWith({
+      toolCallId: tool.id,
+      status: 'cancelled',
+      result: { cancelled: true },
+      error: 'Run stopped before tool admission',
+    })
+    expect(publishToolConfirmation).toHaveBeenCalledWith(
+      expect.objectContaining({ toolCallId: tool.id, status: 'cancelled' })
+    )
+    expect(completeAsyncToolCall).not.toHaveBeenCalled()
+    expect(executeTool).not.toHaveBeenCalled()
+    expect(settleSimToolExecution).not.toHaveBeenCalled()
+  })
+
+  it('keeps the result of a call another controller started before Stop closed admission', async () => {
+    claimToolExecution.mockResolvedValueOnce({ outcome: 'closed' })
+    completePendingAsyncToolCall.mockResolvedValueOnce(null)
+    waitForToolConfirmation.mockResolvedValueOnce({ status: 'success', data: { ran: true } })
+    const tool = buildPendingToolCall()
+    const result = await executeToolAndReport(tool.id, buildStreamingContext(tool), {
+      userId: 'user-1',
+      workflowId: 'workflow-1',
+    })
+    expect(result.status).toBe('success')
+    expect(tool.result?.output).toEqual({ ran: true })
+    expect(completeAsyncToolCall).not.toHaveBeenCalled()
+    expect(publishToolConfirmation).not.toHaveBeenCalled()
     expect(executeTool).not.toHaveBeenCalled()
     expect(settleSimToolExecution).not.toHaveBeenCalled()
   })
@@ -743,7 +775,9 @@ describe('executeToolAndReport provenance isolation', () => {
       resolvedSecretTraceRegistry: new ResolvedSecretTraceRegistry(),
     })
 
-    expect(completePendingAsyncToolCall).not.toHaveBeenCalled()
+    expect(completePendingAsyncToolCall).toHaveBeenCalledExactlyOnceWith(
+      expect.objectContaining({ status: 'cancelled', error: 'Run stopped before tool admission' })
+    )
     expect(completion).toMatchObject({
       status: 'cancelled',
       message: 'Run stopped before tool admission',

apps/sim/lib/mothership/request/tools/executor.ts

@@ -11,6 +11,7 @@ import type {
 } from '@/lib/mothership/async-runs/lifecycle'
 import {
   type CompleteAsyncToolCallInput,
+  completePendingAsyncToolCall,
   getAsyncToolCall,
   markAsyncToolRunning,
   upsertAsyncToolCall,
@@ -709,6 +710,8 @@ export async function executeToolAndReport(
   )
 }
 
+type CancelledSpanDetail = { cancelReason: string; error?: string | undefined }
+
 async function executeToolAndReportInner(
   toolCall: ToolCallState,
   context: StreamingContext,
@@ -756,7 +759,11 @@ async function executeToolAndReportInner(
   if (record?.executionDeadlineAt) toolCall.deadlineAt = record.executionDeadlineAt.getTime()
   if (context.runId) {
     const claim = await lifetime.claim(context.runId, execContext.userId)
-    if (claim.outcome === 'closed') return settleCancelled('Run stopped before tool admission')
+    if (claim.outcome === 'closed') {
+      /** A call another controller already started keeps its own result, observed below. */
+      const stopped = await settleCancelledIfUnstarted('Run stopped before tool admission')
+      if (stopped) return stopped
+    }
     // Settled as never run instead of claimed, so a crash before the handler refuses to start
     // cannot leave it recovered as an unknown outcome. Losing that settlement leaves another writer's.
     if (
@@ -835,21 +842,39 @@ async function executeToolAndReportInner(
   // never evaluated (and the span does not exist yet there anyway).
   async function settleCancelled(
     message: string,
-    span?: { cancelReason: string; error?: string | undefined }
+    span?: CancelledSpanDetail
   ): Promise<AsyncCompletionSignal> {
     markToolCallCancelled(message)
     const cancellationResult = toolCall.result
     markToolResultSeen(context, toolCall.id)
-    await lifetime.complete({
+    await lifetime.complete(cancellationRecord(message))
+    if (toolCall.result !== cancellationResult) {
+      return terminalCompletionFromToolCall(toolCall)
+    }
+    return publishCancellation(message, span)
+  }
+
+  /** Cancels the call only while nothing has started it; otherwise reports nothing settled. */
+  async function settleCancelledIfUnstarted(
+    message: string
+  ): Promise<AsyncCompletionSignal | undefined> {
+    if (!(await completePendingAsyncToolCall(cancellationRecord(message)))) return undefined
+    markToolCallCancelled(message)
+    markToolResultSeen(context, toolCall.id)
+    return publishCancellation(message)
+  }
+
+  function cancellationRecord(message: string): CompleteAsyncToolCallInput {
+    return {
       toolCallId: toolCall.id,
       ...(toolCall.targetWorkspaceId ? { targetWorkspaceId: toolCall.targetWorkspaceId } : {}),
       status: MothershipStreamV1AsyncToolRecordStatus.cancelled,
       result: { cancelled: true },
       error: message,
-    })
-    if (toolCall.result !== cancellationResult) {
-      return terminalCompletionFromToolCall(toolCall)
     }
+  }
+
+  function publishCancellation(message: string, span?: CancelledSpanDetail): AsyncCompletionSignal {
     publishTerminalToolConfirmation({
       toolCallId: toolCall.id,
       ...(toolCall.targetWorkspaceId ? { targetWorkspaceId: toolCall.targetWorkspaceId } : {}),