diff --git a/src/tools/WorkflowTool/launchWorkflow.owner.test.ts b/src/tools/WorkflowTool/launchWorkflow.owner.test.ts index cb2bd27f..3b79ba0a 100644 --- a/src/tools/WorkflowTool/launchWorkflow.owner.test.ts +++ b/src/tools/WorkflowTool/launchWorkflow.owner.test.ts @@ -9,6 +9,7 @@ import { } from '../../bootstrap/state.js' import type { AppState } from '../../state/AppState.js' import { IDLE_SPECULATION_STATE } from '../../state/AppStateStore.js' +import { isTerminalTaskStatus } from '../../Task.js' import type { ToolUseContext } from '../../Tool.js' import { getEmptyToolPermissionContext } from '../../Tool.js' import type { SessionId } from '../../types/ids.js' @@ -18,6 +19,11 @@ import { } from '../../utils/messageQueueManager.js' import { drainSdkEvents } from '../../utils/sdkEventQueue.js' import { resetProjectForTesting } from '../../utils/sessionStorage.js' +import { + _clearOutputsForTest, + _resetTaskOutputDirForTest, + cleanupTaskOutput, +} from '../../utils/task/diskOutput.js' import { prepareWorkflowScript } from '../../utils/workflows/runtime.js' import { killWorkflowTask, @@ -29,10 +35,13 @@ const originalConfigDir = process.env.CLAUDE_CONFIG_DIR describe('launchWorkflow lifecycle ownership', () => { let configDir = '' + let launchedTaskIds = new Set() beforeEach(async () => { configDir = await fs.mkdtemp(path.join(os.tmpdir(), 'workflow-owner-')) + launchedTaskIds = new Set() process.env.CLAUDE_CONFIG_DIR = configDir + _resetTaskOutputDirForTest() resetProjectForTesting() resetStateForTests() resetCommandQueue() @@ -42,6 +51,9 @@ describe('launchWorkflow lifecycle ownership', () => { }) afterEach(async () => { + await _clearOutputsForTest() + await Promise.all([...launchedTaskIds].map(cleanupTaskOutput)) + _resetTaskOutputDirForTest() drainSdkEvents() resetCommandQueue() resetStateForTests() @@ -56,9 +68,10 @@ describe('launchWorkflow lifecycle ownership', () => { scriptBody: string ownerAgentId?: string }): { - getTask: () => LocalWorkflowTaskState | undefined + published: Promise setAppState: (updater: (state: AppState) => AppState) => void } { + launchedTaskIds.add(options.taskId) const script = [ "export const meta = { name: 'owner-test', description: 'Owner test' }", options.scriptBody, @@ -70,8 +83,25 @@ describe('launchWorkflow lifecycle ownership', () => { toolPermissionContext: getEmptyToolPermissionContext(), speculation: IDLE_SPECULATION_STATE, } as unknown as AppState + let resolvePublished: (task: LocalWorkflowTaskState) => void = () => {} + let didPublish = false + const published = new Promise(resolve => { + resolvePublished = resolve + }) const setAppState = (updater: (current: AppState) => AppState): void => { state = updater(state) + const task = state.tasks[options.taskId] as + | LocalWorkflowTaskState + | undefined + if ( + !didPublish && + task && + isTerminalTaskStatus(task.status) && + task.notified + ) { + didPublish = true + resolvePublished(task) + } } const context = { agentId: options.ownerAgentId, @@ -95,29 +125,18 @@ describe('launchWorkflow lifecycle ownership', () => { isResume: false, }) return { - getTask: () => state.tasks[options.taskId] as LocalWorkflowTaskState | undefined, + published, setAppState, } } - async function waitForTerminal( - getTask: () => LocalWorkflowTaskState | undefined, - ): Promise { - for (let attempt = 0; attempt < 100; attempt++) { - const task = getTask() - if (task && task.status !== 'running') return task - await Bun.sleep(2) - } - throw new Error('workflow did not settle') - } - test('keeps a nested workflow owned from start through progress and completion', async () => { const run = launch({ taskId: 'w-owned-complete', ownerAgentId: 'parent-agent', scriptBody: "phase('Scan')\nreturn 'done'", }) - expect((await waitForTerminal(run.getTask)).status).toBe('completed') + expect((await run.published).status).toBe('completed') const lifecycle = drainSdkEvents().filter(event => ( 'task_id' in event && event.task_id === 'w-owned-complete' @@ -139,7 +158,7 @@ describe('launchWorkflow lifecycle ownership', () => { ownerAgentId: 'parent-agent', scriptBody: "throw new Error('workflow exploded')", }) - expect((await waitForTerminal(failed.getTask)).status).toBe('failed') + expect((await failed.published).status).toBe('failed') const killed = launch({ taskId: 'w-owned-killed', @@ -147,6 +166,7 @@ describe('launchWorkflow lifecycle ownership', () => { scriptBody: "return 'too late'", }) expect(killWorkflowTask('w-owned-killed', killed.setAppState)).toBe(true) + expect((await killed.published).status).toBe('killed') const terminals = drainSdkEvents().filter(event => ( event.subtype === 'task_notification' && @@ -176,7 +196,7 @@ describe('launchWorkflow lifecycle ownership', () => { taskId: 'w-root-complete', scriptBody: "phase('Root')\nreturn 'done'", }) - expect((await waitForTerminal(run.getTask)).status).toBe('completed') + expect((await run.published).status).toBe('completed') const lifecycle = drainSdkEvents().filter(event => ( 'task_id' in event && event.task_id === 'w-root-complete'