Skip to content

Commit 30fc8e3

Browse files
authored
fix(mothership): keep a started call's result after Stop, and let a blocked owner stop their run (#9114)
* fix(mothership): settle a stopped run's call as never run only while nothing started it * fix(mothership): let a blocked owner stop their admitted run
1 parent 0e0de65 commit 30fc8e3

5 files changed

Lines changed: 165 additions & 15 deletions

File tree

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

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,19 @@ export async function resolveOwnedWorkspaceChatContext(principal: Principal, cha
1919
return context
2020
}
2121

22-
export async function resolveOwnedChatContext(principal: Principal, chatId: string) {
22+
export interface OwnedChatContextOptions {
23+
/**
24+
* Resolves the chat for an owner whose account is blocked. Only Stop sets it: halting the
25+
* owner's existing work is never harmful, so it stays available after access is withheld.
26+
*/
27+
allowBlockedOwner?: boolean
28+
}
29+
30+
export async function resolveOwnedChatContext(
31+
principal: Principal,
32+
chatId: string,
33+
options: OwnedChatContextOptions = {}
34+
) {
2335
const [chat] = await db
2436
.select({
2537
userId: copilotChats.userId,
@@ -40,7 +52,7 @@ export async function resolveOwnedChatContext(principal: Principal, chatId: stri
4052
) {
4153
throw new OrchestrationError('not_found', 'Chat not found')
4254
}
43-
if ((await getActivelyBannedUserIds([chat.userId])).length > 0) {
55+
if (!options.allowBlockedOwner && (await getActivelyBannedUserIds([chat.userId])).length > 0) {
4456
throw new OrchestrationError('forbidden', 'User account is suspended')
4557
}
4658
const owner: ChatOwnerContext = chat.organizationId

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

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,10 @@ import {
1414
requestRunStop,
1515
} from '@/lib/mothership/async-runs/repository'
1616
import { defineAuthorizedChatUseCase } from '@/lib/mothership/chat/application/authorized-chat-use-case'
17-
import { resolveOwnedChatContext } from '@/lib/mothership/chat/application/context'
17+
import {
18+
type OwnedChatContextOptions,
19+
resolveOwnedChatContext,
20+
} from '@/lib/mothership/chat/application/context'
1821
import { appendCopilotChatMessages } from '@/lib/mothership/chat/messages-store'
1922
import { deliverRunStop } from '@/lib/mothership/request/application/stop-delivery'
2023
import { requestStreamSteering } from '@/lib/mothership/request/session/steer'
@@ -61,12 +64,13 @@ async function resolveRunContext({
6164
async function resolveAdmittedRunContext(
6265
principal: SessionPrincipal,
6366
input: RunControlInput,
64-
run: { chatId: string; workspaceId: string | null; organizationId?: string | null }
67+
run: { chatId: string; workspaceId: string | null; organizationId?: string | null },
68+
options: OwnedChatContextOptions = {}
6569
) {
6670
if (input.chatId && input.chatId !== run.chatId) {
6771
throw new OrchestrationError('forbidden', 'Stream does not belong to this chat')
6872
}
69-
const chat = await resolveOwnedChatContext(principal, run.chatId)
73+
const chat = await resolveOwnedChatContext(principal, run.chatId, options)
7074
if (
7175
(run.workspaceId ?? null) !== (chat.workspaceId ?? null) ||
7276
(run.organizationId ?? null) !== (chat.organizationId ?? null)
@@ -90,7 +94,12 @@ async function resolveAbortContext({
9094
input: RunControlInput
9195
}) {
9296
const run = await getLatestRunForStream(input.streamId, principal.userId)
93-
if (run?.chatId) return resolveAdmittedRunContext(principal, input, run)
97+
/**
98+
* An owner may stop a run already admitted for them even after their account is blocked. A
99+
* blocked account cannot admit a run, so a Stop recorded ahead of admission keeps the ban check.
100+
*/
101+
if (run?.chatId)
102+
return resolveAdmittedRunContext(principal, input, run, { allowBlockedOwner: true })
94103
if (input.chatId) {
95104
const chat = await resolveOwnedChatContext(principal, input.chatId)
96105
if (input.workspaceId && input.workspaceId !== chat.workspaceId) {
@@ -147,7 +156,8 @@ export const abortRun = defineAuthorizedChatUseCase({
147156
const { chatId } = await resolveAdmittedRunContext(
148157
principal,
149158
{ ...input, workspaceId, organizationId },
150-
run
159+
run,
160+
{ allowBlockedOwner: true }
151161
)
152162
/** Read alongside the signals, which need none of it. */
153163
const work = Promise.all([

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

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ import { createRunSegment } from '@/lib/mothership/async-runs/repository'
8888
import { admitChatTurn } from '@/lib/mothership/chat/application/admit-turn'
8989
import { deleteCopilotChat, softDeleteMothershipChat } from '@/lib/mothership/chat/deletion'
9090
import { chatPubSub } from '@/lib/mothership/chat-status'
91+
import { abortRun, steerRun } from '@/lib/mothership/request/application/controls'
9192
import { claimRunController } from '@/lib/mothership/request/lifecycle/controller-ownership'
9293
import {
9394
acquirePendingChatStream,
@@ -966,6 +967,74 @@ describe.runIf(Boolean(redisUrl))('removing a run’s chat, owner or access stop
966967
expect(await runControl(run)).toEqual(STOPPED)
967968
})
968969

970+
it('lets a blocked owner Stop their run, and neither steer nor send, nor anyone else Stop it', async () => {
971+
const ownerId = await seedUser()
972+
const memberId = await seedUser()
973+
const workspaceId = await seedWorkspace(ownerId, [memberId])
974+
const run = await seedRun({ userId: ownerId, workspaceId })
975+
await db.update(user).set({ banned: true }).where(eq(user.id, ownerId))
976+
const owner = createSessionPrincipal({ userId: ownerId, sessionId: generateId() })
977+
const input = { streamId: run.streamId, chatId: run.chatId, workspaceId }
978+
979+
await expect(
980+
steerRun.execute({
981+
principal: owner,
982+
input: { ...input, steeringId: generateId(), content: 'Keep going' },
983+
})
984+
).rejects.toMatchObject({ code: 'forbidden' })
985+
const sendStreamId = generateId()
986+
expect(await acquirePendingChatStream(run.chatId, sendStreamId, 0)).toBe(true)
987+
const lease = getLocalChatStreamLease(run.chatId, sendStreamId)
988+
if (!lease) throw new Error('The send needs its chat lease')
989+
const sendClaim = { normalizedKey: `run-removal:${generateId()}`, claimToken: generateId() }
990+
await db
991+
.insert(idempotencyKey)
992+
.values({ key: sendClaim.normalizedKey, result: { claimToken: sendClaim.claimToken } })
993+
await expect(
994+
admitChatTurn
995+
.execute({
996+
principal: owner,
997+
input: {
998+
chatId: run.chatId,
999+
runId: generateId(),
1000+
executionId: generateId(),
1001+
requestId: generateId(),
1002+
message: { id: sendStreamId, content: 'Start over' },
1003+
recovery: {
1004+
kind: 'interactive_stream',
1005+
goRoute: '/api/mothership',
1006+
clientToolPickupExpected: false,
1007+
request: {
1008+
userId: ownerId,
1009+
chatId: run.chatId,
1010+
messageId: sendStreamId,
1011+
workspaceId,
1012+
message: 'Start over',
1013+
},
1014+
},
1015+
lease,
1016+
sendClaim,
1017+
notifyWorkspaceStatus: false,
1018+
},
1019+
})
1020+
.finally(() => releasePendingChatStream(run.chatId, sendStreamId, lease))
1021+
).rejects.toMatchObject({ code: 'forbidden' })
1022+
await expect(
1023+
abortRun.execute({
1024+
principal: createSessionPrincipal({ userId: memberId, sessionId: generateId() }),
1025+
input,
1026+
})
1027+
).rejects.toMatchObject({ code: 'not_found' })
1028+
expect(await hasStop(run)).toBe(false)
1029+
expect(pushedStops()).toEqual([])
1030+
1031+
await expect(abortRun.execute({ principal: owner, input })).resolves.toMatchObject({
1032+
aborted: true,
1033+
})
1034+
expect(await hasStop(run)).toBe(true)
1035+
expect(pushedStops()).toEqual([run.streamId])
1036+
})
1037+
9691038
/**
9701039
* Admission and deletion both write the chat row, so they serialize on its lock; holding it
9711040
* queues them in a chosen order, then releases them together.

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

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -627,12 +627,44 @@ describe('executeToolAndReport provenance isolation', () => {
627627

628628
it('refuses a late tool after durable Stop closed admission', async () => {
629629
claimToolExecution.mockResolvedValueOnce({ outcome: 'closed' })
630+
completePendingAsyncToolCall.mockImplementationOnce(async (input) => ({ ...input }))
630631
const tool = buildPendingToolCall()
631632
const result = await executeToolAndReport(tool.id, buildStreamingContext(tool), {
632633
userId: 'user-1',
633634
workflowId: 'workflow-1',
634635
})
635-
expect(result.status).toBe('cancelled')
636+
expect(result).toMatchObject({
637+
status: 'cancelled',
638+
message: 'Run stopped before tool admission',
639+
data: { cancelled: true },
640+
})
641+
expect(completePendingAsyncToolCall).toHaveBeenCalledExactlyOnceWith({
642+
toolCallId: tool.id,
643+
status: 'cancelled',
644+
result: { cancelled: true },
645+
error: 'Run stopped before tool admission',
646+
})
647+
expect(publishToolConfirmation).toHaveBeenCalledWith(
648+
expect.objectContaining({ toolCallId: tool.id, status: 'cancelled' })
649+
)
650+
expect(completeAsyncToolCall).not.toHaveBeenCalled()
651+
expect(executeTool).not.toHaveBeenCalled()
652+
expect(settleSimToolExecution).not.toHaveBeenCalled()
653+
})
654+
655+
it('keeps the result of a call another controller started before Stop closed admission', async () => {
656+
claimToolExecution.mockResolvedValueOnce({ outcome: 'closed' })
657+
completePendingAsyncToolCall.mockResolvedValueOnce(null)
658+
waitForToolConfirmation.mockResolvedValueOnce({ status: 'success', data: { ran: true } })
659+
const tool = buildPendingToolCall()
660+
const result = await executeToolAndReport(tool.id, buildStreamingContext(tool), {
661+
userId: 'user-1',
662+
workflowId: 'workflow-1',
663+
})
664+
expect(result.status).toBe('success')
665+
expect(tool.result?.output).toEqual({ ran: true })
666+
expect(completeAsyncToolCall).not.toHaveBeenCalled()
667+
expect(publishToolConfirmation).not.toHaveBeenCalled()
636668
expect(executeTool).not.toHaveBeenCalled()
637669
expect(settleSimToolExecution).not.toHaveBeenCalled()
638670
})
@@ -743,7 +775,9 @@ describe('executeToolAndReport provenance isolation', () => {
743775
resolvedSecretTraceRegistry: new ResolvedSecretTraceRegistry(),
744776
})
745777

746-
expect(completePendingAsyncToolCall).not.toHaveBeenCalled()
778+
expect(completePendingAsyncToolCall).toHaveBeenCalledExactlyOnceWith(
779+
expect.objectContaining({ status: 'cancelled', error: 'Run stopped before tool admission' })
780+
)
747781
expect(completion).toMatchObject({
748782
status: 'cancelled',
749783
message: 'Run stopped before tool admission',

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

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import type {
1111
} from '@/lib/mothership/async-runs/lifecycle'
1212
import {
1313
type CompleteAsyncToolCallInput,
14+
completePendingAsyncToolCall,
1415
getAsyncToolCall,
1516
markAsyncToolRunning,
1617
upsertAsyncToolCall,
@@ -709,6 +710,8 @@ export async function executeToolAndReport(
709710
)
710711
}
711712

713+
type CancelledSpanDetail = { cancelReason: string; error?: string | undefined }
714+
712715
async function executeToolAndReportInner(
713716
toolCall: ToolCallState,
714717
context: StreamingContext,
@@ -756,7 +759,11 @@ async function executeToolAndReportInner(
756759
if (record?.executionDeadlineAt) toolCall.deadlineAt = record.executionDeadlineAt.getTime()
757760
if (context.runId) {
758761
const claim = await lifetime.claim(context.runId, execContext.userId)
759-
if (claim.outcome === 'closed') return settleCancelled('Run stopped before tool admission')
762+
if (claim.outcome === 'closed') {
763+
/** A call another controller already started keeps its own result, observed below. */
764+
const stopped = await settleCancelledIfUnstarted('Run stopped before tool admission')
765+
if (stopped) return stopped
766+
}
760767
// Settled as never run instead of claimed, so a crash before the handler refuses to start
761768
// cannot leave it recovered as an unknown outcome. Losing that settlement leaves another writer's.
762769
if (
@@ -835,21 +842,39 @@ async function executeToolAndReportInner(
835842
// never evaluated (and the span does not exist yet there anyway).
836843
async function settleCancelled(
837844
message: string,
838-
span?: { cancelReason: string; error?: string | undefined }
845+
span?: CancelledSpanDetail
839846
): Promise<AsyncCompletionSignal> {
840847
markToolCallCancelled(message)
841848
const cancellationResult = toolCall.result
842849
markToolResultSeen(context, toolCall.id)
843-
await lifetime.complete({
850+
await lifetime.complete(cancellationRecord(message))
851+
if (toolCall.result !== cancellationResult) {
852+
return terminalCompletionFromToolCall(toolCall)
853+
}
854+
return publishCancellation(message, span)
855+
}
856+
857+
/** Cancels the call only while nothing has started it; otherwise reports nothing settled. */
858+
async function settleCancelledIfUnstarted(
859+
message: string
860+
): Promise<AsyncCompletionSignal | undefined> {
861+
if (!(await completePendingAsyncToolCall(cancellationRecord(message)))) return undefined
862+
markToolCallCancelled(message)
863+
markToolResultSeen(context, toolCall.id)
864+
return publishCancellation(message)
865+
}
866+
867+
function cancellationRecord(message: string): CompleteAsyncToolCallInput {
868+
return {
844869
toolCallId: toolCall.id,
845870
...(toolCall.targetWorkspaceId ? { targetWorkspaceId: toolCall.targetWorkspaceId } : {}),
846871
status: MothershipStreamV1AsyncToolRecordStatus.cancelled,
847872
result: { cancelled: true },
848873
error: message,
849-
})
850-
if (toolCall.result !== cancellationResult) {
851-
return terminalCompletionFromToolCall(toolCall)
852874
}
875+
}
876+
877+
function publishCancellation(message: string, span?: CancelledSpanDetail): AsyncCompletionSignal {
853878
publishTerminalToolConfirmation({
854879
toolCallId: toolCall.id,
855880
...(toolCall.targetWorkspaceId ? { targetWorkspaceId: toolCall.targetWorkspaceId } : {}),

0 commit comments

Comments
 (0)