diff --git a/apps/sim/lib/mothership/chat/application/context.ts b/apps/sim/lib/mothership/chat/application/context.ts index c12c348830f..6b327b3debd 100644 --- a/apps/sim/lib/mothership/chat/application/context.ts +++ b/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 diff --git a/apps/sim/lib/mothership/request/application/controls.ts b/apps/sim/lib/mothership/request/application/controls.ts index e094268ae4e..2bed39336db 100644 --- a/apps/sim/lib/mothership/request/application/controls.ts +++ b/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([ diff --git a/apps/sim/lib/mothership/request/application/removal-stops-runs.integration.ts b/apps/sim/lib/mothership/request/application/removal-stops-runs.integration.ts index 0aeac64fa6f..940c46f0190 100644 --- a/apps/sim/lib/mothership/request/application/removal-stops-runs.integration.ts +++ b/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. diff --git a/apps/sim/lib/mothership/request/tools/executor.test.ts b/apps/sim/lib/mothership/request/tools/executor.test.ts index 75f74b63904..5bff6dff4d3 100644 --- a/apps/sim/lib/mothership/request/tools/executor.test.ts +++ b/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', diff --git a/apps/sim/lib/mothership/request/tools/executor.ts b/apps/sim/lib/mothership/request/tools/executor.ts index 7e7d053dede..361d143963e 100644 --- a/apps/sim/lib/mothership/request/tools/executor.ts +++ b/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 { 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 { + 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 } : {}),