Skip to content

Commit 2be49a8

Browse files
committed
test(mothership): drive a workflow tool past its wait and a stop mid-slice through the real controller
Two live cases beside the sliced draft run (opt-in, SIM_HELPERS_SMOKE and MSHIP_WORKER_ROOT): - tool: run_workflow on a Sim declaring workflow-tool-wait-timeout answers its 500 s run still running after 240 s, the worker waits on it as run-wait slices and reads it back, and the model sees the settled run's output. - declared-stop: the Stop button's request while the worker waits on a detached draft run ends the turn and cancels the run. The harness now marks the turn's tool execution and billing as the chat route does, registers and wires the turn's controller as lifecycle/start.ts does, gives the tool case a real AES key to seal its provenance, and keeps logged errors' messages in the evidence.
1 parent 2306df4 commit 2be49a8

1 file changed

Lines changed: 204 additions & 13 deletions

File tree

‎apps/sim/lib/mothership/agent-cli/saved-run-read.live.test.ts‎

Lines changed: 204 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
* Isolated columns plus required defaults/indexes do not prove migrations or all constraints.
1818
*/
1919
import { execFile, spawn } from 'node:child_process'
20+
import { randomBytes } from 'node:crypto'
2021
import { once } from 'node:events'
2122
import { createReadStream } from 'node:fs'
2223
import { mkdir, mkdtemp, readFile, rm, writeFile } from 'node:fs/promises'
@@ -62,7 +63,7 @@ import {
6263
workspaceFiles,
6364
} from '@sim/db/schema'
6465
import type { PermissionType } from '@sim/platform-authz/workspace'
65-
import { authMockFns, envFlagsMock, setEnvFlags } from '@sim/testing'
66+
import { authMockFns, envFlagsMock, setEnv, setEnvFlags } from '@sim/testing'
6667
import { sleep } from '@sim/utils/helpers'
6768
import { generateId } from '@sim/utils/id'
6869
import { eq, is, SQL, sql } from 'drizzle-orm'
@@ -133,11 +134,18 @@ import {
133134
} from '@/lib/mothership/chat/process-contents'
134135
import { finalizeAssistantTurn } from '@/lib/mothership/chat/terminal-state'
135136
import { createRegistryContentGate } from '@/lib/mothership/content-delivery/registry-gate'
136-
import { WORKFLOW_RUN_WAIT_TIMEOUT_CAPABILITY } from '@/lib/mothership/generated/protocol'
137+
import {
138+
WORKFLOW_RUN_WAIT_TIMEOUT_CAPABILITY,
139+
WORKFLOW_TOOL_WAIT_TIMEOUT_CAPABILITY,
140+
} from '@/lib/mothership/generated/protocol'
137141
import { prepareInboxAttachments } from '@/lib/mothership/inbox/attachments'
138142
import { claimRunController } from '@/lib/mothership/request/lifecycle/controller-ownership'
139143
import { runCopilotLifecycle } from '@/lib/mothership/request/lifecycle/run'
140-
import { isToolCallStreamEvent } from '@/lib/mothership/request/session'
144+
import {
145+
isToolCallStreamEvent,
146+
registerActiveStream,
147+
unregisterActiveStream,
148+
} from '@/lib/mothership/request/session'
141149
import { maybeWriteOutputToFile } from '@/lib/mothership/request/tools/files'
142150
import { maybeWriteOutputToTable } from '@/lib/mothership/request/tools/tables'
143151
import type { ExecutionContext } from '@/lib/mothership/request/types'
@@ -154,6 +162,7 @@ import { calculateNextRunTime, getScheduleTimeValues } from '@/lib/workflows/sch
154162
import { fileOperations } from '@/lib/workspace-files/application/operations'
155163
import { readWorkspaceFileArtifact } from '@/lib/workspace-files/application/read-workspace-file-artifact'
156164
import { readWorkspaceFileText } from '@/lib/workspace-files/application/read-workspace-file-text'
165+
import { POST as abortChatRoute } from '@/app/api/copilot/chat/abort/route'
157166
import {
158167
POST as addChatResourceRoute,
159168
DELETE as removeChatResourceRoute,
@@ -230,7 +239,18 @@ vi.mock('@/lib/mothership/inbox/agentmail-client', () => ({
230239
vi.mock('@sim/logger', async () => {
231240
const { loggerMock, createMockLogger } = await import('@sim/testing')
232241
const logger = createMockLogger()
233-
logger.error.mockImplementation((...args: unknown[]) => fixture.errors.push(args))
242+
// Errors keep their message and top frames, which JSON evidence would otherwise drop as `{}`.
243+
logger.error.mockImplementation((...args: unknown[]) =>
244+
fixture.errors.push(
245+
JSON.parse(
246+
JSON.stringify(args, (_key, value) =>
247+
value instanceof Error
248+
? { message: value.message, stack: value.stack?.split('\n').slice(0, 6) }
249+
: value
250+
)
251+
)
252+
)
253+
)
234254
logger.withMetadata.mockReturnValue(logger)
235255
logger.child.mockReturnValue(logger)
236256
return { ...loggerMock, createLogger: () => logger }
@@ -3601,6 +3621,10 @@ describe.skipIf(!process.env.MSHIP_TEST_DATABASE_URL)(
36013621
['undeclared', 260, 0],
36023622
// The model's own --wait-timeout 5 lapses: an exit 4 the worker hands the model.
36033623
['own-budget', 30, 0],
3624+
// run_workflow on a Sim that answers a tool's run still running after the frame's wait.
3625+
['tool', 300, 200],
3626+
// The user stops the chat while the worker waits on a detached run: the run stops too.
3627+
['declared-stop', 300, 300],
36043628
] as const)(
36053629
'draft test run past a wait slice reaches the model settled (capability %s)',
36063630
async (capability, waitA, waitB) => {
@@ -3670,6 +3694,9 @@ describe.skipIf(!process.env.MSHIP_TEST_DATABASE_URL)(
36703694
fixture.saved.set(draftId, { blocks, edges, loops: {}, parallels: {} })
36713695
// A paid plan's synchronous limit, so the run outlasts several slices.
36723696
fixture.executionTimeoutMs = 50 * 60_000
3697+
// A workflow tool call seals its inputs' provenance, which takes a real AES key.
3698+
const encryptionKey = env.ENCRYPTION_KEY
3699+
setEnv({ ENCRYPTION_KEY: randomBytes(32).toString('hex') })
36733700

36743701
const child = spawn('bun', ['tools/probes/src/draft-run-slices.ts'], {
36753702
cwd: workerRoot,
@@ -3686,6 +3713,7 @@ describe.skipIf(!process.env.MSHIP_TEST_DATABASE_URL)(
36863713
MSHIP_PROBE_RUN_FLAGS: JSON.stringify(
36873714
capability === 'own-budget' ? ['--wait-timeout', '5'] : []
36883715
),
3716+
MSHIP_PROBE_TOOL: capability === 'tool' ? 'run_workflow' : '',
36893717
},
36903718
stdio: ['pipe', 'pipe', 'pipe'],
36913719
})
@@ -3734,8 +3762,51 @@ describe.skipIf(!process.env.MSHIP_TEST_DATABASE_URL)(
37343762
workspaceId,
37353763
type: 'mothership',
37363764
title: `Long draft ${capability}`,
3765+
// The Stop route stops the chat's current stream.
3766+
...(capability === 'declared-stop' ? { conversationId: messageId } : {}),
37373767
})
37383768
const startedAt = Date.now()
3769+
/** When the declared-stop case stopped the chat, once the worker waited on a detached run. */
3770+
let stoppedAt: number | undefined
3771+
// Registered as the chat route registers its turn (lifecycle/start.ts), so a Stop
3772+
// aborts it; the route also aborts it on an observed abort marker.
3773+
const turn = new AbortController()
3774+
registerActiveStream(messageId, turn)
3775+
const stopper =
3776+
capability === 'declared-stop'
3777+
? (async () => {
3778+
// Bounded: a run that never reaches a run wait fails the case below.
3779+
for (let read = 0; read < 300; read++) {
3780+
await sleep(2_000)
3781+
const [current] = await db
3782+
.select({ id: copilotRuns.id })
3783+
.from(copilotRuns)
3784+
.where(eq(copilotRuns.streamId, messageId))
3785+
if (!current) continue
3786+
const waiting = await db
3787+
.select({ args: copilotAsyncToolCalls.args })
3788+
.from(copilotAsyncToolCalls)
3789+
.where(eq(copilotAsyncToolCalls.runId, current.id))
3790+
if (!JSON.stringify(waiting).includes('"runs","wait"')) continue
3791+
await sleep(5_000)
3792+
authMockFns.mockGetSession.mockResolvedValueOnce({
3793+
user: { id: 'run-reader' },
3794+
session: { id: 'stop-session' },
3795+
})
3796+
stoppedAt = Date.now()
3797+
// The Stop button's request: it stops the run, not just the stream.
3798+
const response = await abortChatRoute(
3799+
new NextRequest('https://sim.test/api/copilot/chat/abort', {
3800+
method: 'POST',
3801+
headers: { 'content-type': 'application/json' },
3802+
body: JSON.stringify({ chatId, streamId: messageId, workspaceId }),
3803+
})
3804+
)
3805+
return response.status
3806+
}
3807+
return undefined
3808+
})()
3809+
: undefined
37393810
const result = await runCopilotLifecycle(
37403811
{
37413812
message: 'Test-run the draft and tell me its output.',
@@ -3745,9 +3816,16 @@ describe.skipIf(!process.env.MSHIP_TEST_DATABASE_URL)(
37453816
chatId,
37463817
messageId,
37473818
// What lib/mothership/chat/payload.ts declares on every chat turn.
3748-
...(capability !== 'undeclared'
3749-
? { clientCapabilities: [WORKFLOW_RUN_WAIT_TIMEOUT_CAPABILITY] }
3750-
: {}),
3819+
...(capability === 'tool'
3820+
? {
3821+
clientCapabilities: [
3822+
WORKFLOW_RUN_WAIT_TIMEOUT_CAPABILITY,
3823+
WORKFLOW_TOOL_WAIT_TIMEOUT_CAPABILITY,
3824+
],
3825+
}
3826+
: capability !== 'undeclared'
3827+
? { clientCapabilities: [WORKFLOW_RUN_WAIT_TIMEOUT_CAPABILITY] }
3828+
: {}),
37513829
},
37523830
{
37533831
userId: 'run-reader',
@@ -3758,12 +3836,32 @@ describe.skipIf(!process.env.MSHIP_TEST_DATABASE_URL)(
37583836
interactive: false,
37593837
clientToolPickupExpected: false,
37603838
flushAfterEvent: false,
3761-
abortSignal: AbortSignal.timeout(20 * 60_000),
3839+
// Wired as the chat route wires its controller (lifecycle/start.ts): an observed
3840+
// Stop aborts the turn.
3841+
abortSignal: AbortSignal.any([turn.signal, AbortSignal.timeout(20 * 60_000)]),
3842+
onAbortObserved: (reason) => {
3843+
if (!turn.signal.aborted) turn.abort(reason)
3844+
},
37623845
executionContext: {
37633846
userId: 'run-reader',
37643847
workspaceId,
37653848
workflowId: draftId,
37663849
chatId,
3850+
// As the chat route marks every turn's tool execution and bills it
3851+
// (lib/mothership/chat/post.ts); a workflow tool runs in-process under both.
3852+
copilotToolExecution: true,
3853+
billingAttribution: {
3854+
actorUserId: 'run-reader',
3855+
workspaceId,
3856+
organizationId: null,
3857+
billedAccountUserId: 'run-reader',
3858+
billingEntity: { type: 'user', id: 'run-reader' },
3859+
billingPeriod: {
3860+
start: '2026-09-01T00:00:00.000Z',
3861+
end: '2026-10-01T00:00:00.000Z',
3862+
},
3863+
payerSubscription: null,
3864+
},
37673865
userPermission: 'write',
37683866
resolvedSecretTraceRegistry: new ResolvedSecretTraceRegistry([], {
37693867
userId: 'run-reader',
@@ -3773,6 +3871,8 @@ describe.skipIf(!process.env.MSHIP_TEST_DATABASE_URL)(
37733871
}
37743872
)
37753873
const endedAt = Date.now()
3874+
unregisterActiveStream(messageId, turn)
3875+
const stopStatus = await stopper
37763876
const [run] = await db
37773877
.select()
37783878
.from(copilotRuns)
@@ -3783,8 +3883,38 @@ describe.skipIf(!process.env.MSHIP_TEST_DATABASE_URL)(
37833883
.from(copilotAsyncToolCalls)
37843884
.where(eq(copilotAsyncToolCalls.runId, run.id))
37853885
: []
3886+
/** The detached run's own status after a Stop, read until it leaves `running`. */
3887+
let stoppedRun: { executionId?: string; status?: string; endedAt?: string } = {}
3888+
if (capability === 'declared-stop') {
3889+
const receipt = calls.find(
3890+
(call) =>
3891+
(call.args as { request?: { invocation?: { argv?: string[] } } }).request
3892+
?.invocation?.argv?.[1] === 'run'
3893+
)?.result as { stdout?: string } | undefined
3894+
const executionId = String(JSON.parse(receipt?.stdout ?? '{}').runId ?? '')
3895+
for (let read = 0; read < 30; read++) {
3896+
const [log] = await db
3897+
.select({
3898+
status: workflowExecutionLogs.status,
3899+
endedAt: workflowExecutionLogs.endedAt,
3900+
})
3901+
.from(workflowExecutionLogs)
3902+
.where(eq(workflowExecutionLogs.executionId, executionId))
3903+
stoppedRun = {
3904+
executionId,
3905+
status: log?.status,
3906+
endedAt: log?.endedAt?.toISOString(),
3907+
}
3908+
if (log && log.status !== 'running') break
3909+
await sleep(2_000)
3910+
}
3911+
}
37863912
const evidence = {
37873913
capability,
3914+
stopStatus,
3915+
stoppedAt: stoppedAt ? new Date(stoppedAt).toISOString() : undefined,
3916+
stoppedRun,
3917+
errors: fixture.errors,
37883918
waits: [waitA, waitB],
37893919
startedAt: new Date(startedAt).toISOString(),
37903920
totalSeconds: (endedAt - startedAt) / 1000,
@@ -3798,6 +3928,9 @@ describe.skipIf(!process.env.MSHIP_TEST_DATABASE_URL)(
37983928
toolName: call.toolName,
37993929
args: call.args,
38003930
status: call.status,
3931+
createdAt: call.createdAt?.toISOString(),
3932+
result: call.toolName === 'sim_cli' ? undefined : call.result,
3933+
error: call.error,
38013934
exitCode: (call.result as { exitCode?: unknown } | null)?.exitCode,
38023935
stdout: (call.result as { stdout?: unknown } | null)?.stdout,
38033936
startedAt: call.executionStartedAt?.toISOString(),
@@ -3818,13 +3951,24 @@ describe.skipIf(!process.env.MSHIP_TEST_DATABASE_URL)(
38183951
`${process.env.MSHIP_SLICES_REPORT_PATH}.${capability}.json`,
38193952
JSON.stringify(evidence, null, 2)
38203953
)
3954+
if (capability === 'declared-stop') {
3955+
expect(stopStatus).toBe(200)
3956+
// The turn ended at the Stop, not when the run would have settled.
3957+
expect(endedAt - (stoppedAt ?? endedAt)).toBeLessThan(60_000)
3958+
// The run the worker was waiting on stopped with the chat, as a held run does.
3959+
expect(stoppedRun.status, JSON.stringify(evidence)).toBe('cancelled')
3960+
return
3961+
}
38213962
expect(
38223963
result.success,
38233964
JSON.stringify({ evidence, errors: fixture.errors, diagnostics })
38243965
).toBe(true)
38253966
expect(run?.status).toBe('complete')
38263967
const report = JSON.parse(result.content)
3827-
const seen = String(report.output.value)
3968+
const seen =
3969+
typeof report.output.value === 'string'
3970+
? report.output.value
3971+
: JSON.stringify(report.output.value)
38283972
if (capability === 'own-budget') {
38293973
expect(calls).toHaveLength(1)
38303974
expect(
@@ -3840,17 +3984,63 @@ describe.skipIf(!process.env.MSHIP_TEST_DATABASE_URL)(
38403984
}
38413985
// Sim #9469: a still-running (exit 4) slice settles as done, not failed.
38423986
for (const call of calls) expect(call.status, call.toolName).not.toBe('failed')
3843-
// The model saw the settled run, not a detached receipt.
3844-
expect(seen).toContain('"exitCode":0')
3845-
expect(seen).toContain('"status":"completed"')
3846-
expect(seen).toContain('after-waits')
38473987
const waitSeconds = waitA + waitB
38483988
const argvOf = (call: (typeof calls)[number]) =>
38493989
(call.args as { request: { invocation: { argv: string[] } } }).request.invocation.argv
38503990
const outcome = (call: (typeof calls)[number]) =>
38513991
call.result as { exitCode: number; stdout: string }
38523992
const heldSeconds = (call: (typeof calls)[number]) =>
38533993
(Number(call.executionSettledAt) - Number(call.executionStartedAt)) / 1000
3994+
if (capability === 'tool') {
3995+
const [toolCall, ...slices] = calls
3996+
expect(toolCall.toolName).toBe('run_workflow')
3997+
// The tool answered its run still running by the run's id, after one wait.
3998+
const running = toolCall.result as { status?: string; executionId?: string }
3999+
expect(running).toMatchObject({ status: 'running', workflowId: draftId })
4000+
const executionId = String(running.executionId)
4001+
const answeredAfter = (Number(slices[0].createdAt) - Number(toolCall.createdAt)) / 1000
4002+
expect(answeredAfter).toBeGreaterThan(239)
4003+
expect(answeredAfter).toBeLessThan(250)
4004+
const toolWaits = slices.filter(
4005+
(call) => argvOf(call).slice(0, 3).join(' ') === 'workflows runs wait'
4006+
)
4007+
expect(toolWaits.length).toBeGreaterThan(0)
4008+
for (const [index, call] of toolWaits.entries()) {
4009+
expect(argvOf(call)).toEqual([
4010+
'workflows',
4011+
'runs',
4012+
'wait',
4013+
executionId,
4014+
'--workflow',
4015+
draftId,
4016+
'--wait-timeout',
4017+
'240',
4018+
])
4019+
expect(outcome(call).exitCode).toBe(index === toolWaits.length - 1 ? 0 : 4)
4020+
expect(heldSeconds(call)).toBeLessThan(250)
4021+
}
4022+
const readBack = slices.at(-1)
4023+
if (!readBack) throw new Error('The run was never read back')
4024+
expect(argvOf(readBack)).toEqual([
4025+
'workflows',
4026+
'runs',
4027+
'get',
4028+
executionId,
4029+
'--workflow',
4030+
draftId,
4031+
'--include-output',
4032+
])
4033+
expect(slices).toHaveLength(toolWaits.length + 1)
4034+
// The model saw the settled run with its output, not the still-running answer.
4035+
expect(seen).toContain('"status":"completed"')
4036+
expect(seen).toContain('after-waits')
4037+
expect(endedAt - startedAt).toBeGreaterThan(waitSeconds * 1000)
4038+
return
4039+
}
4040+
// The model saw the settled run, not a detached receipt.
4041+
expect(seen).toContain('"exitCode":0')
4042+
expect(seen).toContain('"status":"completed"')
4043+
expect(seen).toContain('after-waits')
38544044
const runs = calls.filter((call) => argvOf(call)[1] === 'run')
38554045
const waits = calls.filter(
38564046
(call) => argvOf(call).slice(0, 3).join(' ') === 'workflows runs wait'
@@ -3919,6 +4109,7 @@ describe.skipIf(!process.env.MSHIP_TEST_DATABASE_URL)(
39194109
expect(endedAt - startedAt).toBeGreaterThan(waitSeconds * 1000)
39204110
} finally {
39214111
clearTimeout(startupTimeout)
4112+
setEnv({ ENCRYPTION_KEY: encryptionKey })
39224113
fixture.executionTimeoutMs = 20_000
39234114
child.kill()
39244115
await exited

0 commit comments

Comments
 (0)