diff --git a/src/main/runtime/agent-session/structured-agent-session-codex-compact-exit.test.ts b/src/main/runtime/agent-session/structured-agent-session-codex-compact-exit.test.ts index f0cdd69dd2d..ae17eafa822 100644 --- a/src/main/runtime/agent-session/structured-agent-session-codex-compact-exit.test.ts +++ b/src/main/runtime/agent-session/structured-agent-session-codex-compact-exit.test.ts @@ -80,14 +80,16 @@ it('delivers what waited behind the command once its child dies mid-command, wit 'agentSession.create', harness.createIntentParams() ) + const compacting = harness.codex.nextRequest('thread/compact/start') const command = await startCompact(created.fence) - await vi.waitFor(() => expect(calls('thread/compact/start')).toHaveLength(1)) + await compacting await send('after the exit') + const restarted = harness.codex.nextRequest('turn/start') const dead = exitChild() - await vi.waitFor(() => expect(calls('turn/start')).toHaveLength(1)) - expect(calls('turn/start')[0]).not.toBe(dead) + expect(await restarted).not.toBe(dead) + expect(calls('turn/start')).toHaveLength(1) const settled = await snapshot() const turn = settled.items.find( (item) => item.itemId === structuredAgentSessionCommandTurn(command).itemId @@ -103,8 +105,9 @@ it('ends a command Codex has not opened a turn for by stopping its child, and th 'agentSession.create', harness.createIntentParams() ) + const compacting = harness.codex.nextRequest('thread/compact/start') const command = await startCompact(created.fence) - await vi.waitFor(() => expect(calls('thread/compact/start')).toHaveLength(1)) + await compacting const first = threadChild() const { itemId, turnId } = structuredAgentSessionCommandTurn(command) @@ -123,7 +126,8 @@ it('ends a command Codex has not opened a turn for by stopping its child, and th const turn = (await snapshot()).items.find((item) => item.itemId === itemId) expect(readAgentJournalTurn(turn?.body)?.state).toBe('interrupted') }) + const restarted = harness.codex.nextRequest('turn/start') await send('after the stop') - await vi.waitFor(() => expect(calls('turn/start')).toHaveLength(1)) - expect(calls('turn/start')[0]).not.toBe(first) + expect(await restarted).not.toBe(first) + expect(calls('turn/start')).toHaveLength(1) }) diff --git a/src/main/runtime/rpc/methods/structured-sessions/structured-codex-session-rpc-test-harness.ts b/src/main/runtime/rpc/methods/structured-sessions/structured-codex-session-rpc-test-harness.ts index c79a9ea0382..a9553ba160a 100644 --- a/src/main/runtime/rpc/methods/structured-sessions/structured-codex-session-rpc-test-harness.ts +++ b/src/main/runtime/rpc/methods/structured-sessions/structured-codex-session-rpc-test-harness.ts @@ -51,12 +51,15 @@ export type FakeCodex = { connections: FakeCodexConnection[] openConnection: typeof openCodexAppServerConnection live: () => FakeCodexConnection + /** Resolves with the child that next receives `method`. */ + nextRequest: (method: string) => Promise notify: (method: string, params: unknown) => void ask: (id: number, method: string, params: unknown) => void } function fakeCodex(): FakeCodex { const connections: FakeCodexConnection[] = [] + const requestWaiters = new Map void)[]>() const openConnection: typeof openCodexAppServerConnection = async (launch, handlers = {}) => { const connection: FakeCodexConnection = { launch, @@ -68,6 +71,10 @@ function fakeCodex(): FakeCodex { closed: false, request: async (method, params) => { connection.calls.push({ method, params }) + requestWaiters + .get(method) + ?.splice(0) + .forEach((resolve) => resolve(connection)) if (method === 'thread/start') { return { thread: { id: THREAD, path: '/rollouts/integration.jsonl' } } } @@ -123,6 +130,10 @@ function fakeCodex(): FakeCodex { connections, openConnection, live, + nextRequest: (method) => + new Promise((resolve) => { + requestWaiters.set(method, [...(requestWaiters.get(method) ?? []), resolve]) + }), notify: (method, params) => live().handlers.onNotification?.(method, params), ask: (id, method, params) => live().handlers.onServerRequest?.({ id, method, params }) }