diff --git a/README.md b/README.md index 092f61ba..88a2ad51 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ Use [OpenAI Codex](https://github.com/openai/codex) from [Agent Client Protocol] - Text prompts, embedded context, images, resource links, and additional workspace directories. - Shell command, file change, [permission request](docs/permission-extension.md), MCP tool call, terminal output, reasoning, plan, web search, image generation, image view, token usage, and review events. - [Native ACP subagent sessions](docs/subagent-sessions.md) (after capability negotiation) with separate child histories and root-routed permissions; a legacy tool-call fallback otherwise. +- [Background terminal tasks](docs/async-tasks.md) in AIR, with task status and targeted stop support after capability negotiation. - Session-scoped long-running goals through the provider-neutral [goal extension](docs/goal-extension.md). - A per-turn [agent file-change report](docs/agent-file-change-report.md) after capability negotiation. - Client-provided MCP servers over command-based stdio config and HTTP transport. @@ -82,6 +83,12 @@ Subagent sessions follow the draft [ACP subagent RFD](https://github.com/agentcl See [docs/subagent-sessions.md](docs/subagent-sessions.md) for the negotiation, lifecycle events, `session/load` reconstruction, and legacy fallback details. +### Background terminal tasks + +Codex can keep a shell command running after a turn continues. AIR clients can show this work in the Async Tasks panel and stop one command. + +See [docs/async-tasks.md](docs/async-tasks.md) for the capability, lifecycle events, and stop request. + ## License By contributing, you agree that your contributions will be licensed under the Apache 2.0 License. diff --git a/docs/async-tasks.md b/docs/async-tasks.md new file mode 100644 index 00000000..88f0e98a --- /dev/null +++ b/docs/async-tasks.md @@ -0,0 +1,42 @@ +# Background terminal tasks + +Codex app-server owns shell commands that continue after their initial tool call. The adapter exposes these commands through the AIR async task extension. + +## Negotiation + +The client adds `asyncTasks` to `_meta.jetbrains.air.capabilities`. The adapter advertises the same capability in its `initialize` response. + +The adapter emits no async task updates when the client does not advertise this capability. + +## Lifecycle + +The adapter uses `thread/backgroundTerminals/list` as the source of active processes. It maps each active process to `async_task_spawned`. + +Before the spawn update, the adapter marks the command with `_meta.jetbrains.air.asyncTasks.backgrounded`. AIR can then keep the command card active without duplicating its output. + +For a root command, the command item ID is both the async task ID and the related tool call ID. A child command prefixes its task ID with the child thread ID. The prefix keeps task IDs distinct across native subagent sessions. The related tool call ID remains the command item ID. + +The adapter publishes a child command on its native subagent session. The app-server process ID remains an internal control handle. + +The existing command card owns the command output. Therefore, a background terminal task sets `showInTranscript` to `false`. + +When the command ends, the adapter emits `async_task_state_update` with `completed` or `failed`. + +The active-terminal list repairs a lost completion event. The adapter reports `stopped` when an announced terminal disappears from that list. + +Session loading restores root and child tasks after it replays their command history. A provider restart stops old tasks and moves task control to the new app-server client. + +If the app-server exits, the adapter reports each unfinished task as `failed`. + +## Stop request + +The client sends `_session/async_task/stop` with the ACP session ID and async task ID: + +```json +{ + "sessionId": "thread-id", + "asyncTaskId": "command-item-id" +} +``` + +The adapter resolves the app-server process ID and calls `thread/backgroundTerminals/terminate`. It returns `{ "stopped": true }` after app-server accepts the termination. diff --git a/src/ACPSessionConnection.ts b/src/ACPSessionConnection.ts index e29a6514..c04570ca 100644 --- a/src/ACPSessionConnection.ts +++ b/src/ACPSessionConnection.ts @@ -2,7 +2,7 @@ import * as acp from "@agentclientprotocol/sdk"; import { type AcpSessionUpdate, asSdkSessionNotification, -} from "./subagents/AcpSubagents"; +} from "./AcpSessionExtensions"; export type AcpClientConnection = Pick; diff --git a/src/AcpExtensions.ts b/src/AcpExtensions.ts index 98fd6170..b450c8bd 100644 --- a/src/AcpExtensions.ts +++ b/src/AcpExtensions.ts @@ -11,6 +11,10 @@ import { LEGACY_GOAL_CONTROL_METHOD, type GoalControlRequest, } from "./GoalExtension"; +import { + ASYNC_TASK_STOP_METHOD, + type AsyncTaskStopExtRequest, +} from "./async-tasks/AsyncTaskExtension"; export { AUTH_STATUS_META_KEY, @@ -74,6 +78,7 @@ export type ExtMethodRequest = | LegacySetSessionModelExtRequest | SessionSteeringExtRequest | GoalControlExtRequest + | AsyncTaskStopExtRequest export function isExtMethodRequest(request: { method: string, params: Record }): request is ExtMethodRequest { return request.method === "authentication/status" @@ -81,7 +86,8 @@ export function isExtMethodRequest(request: { method: string, params: Record & { + update: AcpSessionUpdate; +}; + +/** The only cast needed until the ACP SDK publishes the extension updates. */ +export function asSdkSessionNotification( + notification: AcpSessionNotification, +): SessionNotification { + return notification as SessionNotification; +} diff --git a/src/AirExtension.ts b/src/AirExtension.ts index 5fcf61de..03b12749 100644 --- a/src/AirExtension.ts +++ b/src/AirExtension.ts @@ -15,6 +15,8 @@ export const AIR_EXTENSION_CAPABILITIES_KEY = "capabilities"; export const AIR_SESSION_FAILURE_KEY = "sessionFailure"; export const AIR_AGENT_FILE_CHANGE_REPORT_KEY = "agentFileChangeReport"; export const AIR_NATIVE_SUBAGENT_SESSIONS_KEY = "nativeSubagentSessions"; +export const AIR_ASYNC_TASKS_KEY = "asyncTasks"; +export const AIR_ASYNC_TASKS_BACKGROUNDED_KEY = "backgrounded"; export const AIR_AGENT_FILE_CHANGE_REPORT_REQUEST_KEY = "agentFileChangeReportRequest"; export const AIR_EXTENSION_VERSION = 1; diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index 80638733..4dc15e01 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -130,6 +130,7 @@ import {TitleGenerator} from "./TitleGenerator"; import {once} from "node:events"; import { AIR_AGENT_FILE_CHANGE_REPORT_KEY, + AIR_ASYNC_TASKS_KEY, AIR_NATIVE_SUBAGENT_SESSIONS_KEY, AIR_EXTENSION_CAPABILITIES_KEY, AIR_EXTENSION_VERSION, @@ -139,6 +140,8 @@ import { clientSupportsAirCapability, JETBRAINS_META_KEY, } from "./AirExtension"; +import {ASYNC_TASK_STOP_METHOD} from "./async-tasks/AsyncTaskExtension"; +import {CodexBackgroundTerminalTasks} from "./async-tasks/CodexBackgroundTerminalTasks"; import { type AgentFileChangeReport, type AgentFileChangeReportRequest, @@ -178,6 +181,7 @@ export interface SessionState { sessionFailure?: SessionFailure; titleGen?: TitleGenerator; subagents: CodexSubagentEventRouter; + asyncTasks: CodexBackgroundTerminalTasks; } export type SessionFailureCategory = @@ -278,6 +282,7 @@ export class CodexAcpServer { private readonly goalControlGenerations: Map; private readonly permissionLifecycleContexts: WeakMap; private readonly codexProcessState: CodexProcessState | null; + private codexProcessGeneration = 0; private initializeRequest: acp.InitializeRequest | null = null; private providerUpdate: Promise | null = null; @@ -313,6 +318,7 @@ export class CodexAcpServer { this.booleanConfigOptionsSupported = false; this.currentAuthStatus = null; this.availableCommands = this.createAvailableCommands(codexAcpClient); + this.observeCodexProcess(); } private createAvailableCommands(client: CodexAcpClient): CodexCommands { @@ -390,6 +396,7 @@ export class CodexAcpServer { AIR_SESSION_FAILURE_KEY, AIR_AGENT_FILE_CHANGE_REPORT_KEY, AIR_NATIVE_SUBAGENT_SESSIONS_KEY, + AIR_ASYNC_TASKS_KEY, ], }, }, @@ -413,6 +420,18 @@ export class CodexAcpServer { return await this.unstable_setSessionModel(this.parseLegacySetSessionModelParams(methodRequest.params)); case SESSION_STEERING_METHOD: return await this.executeOrQueueSteeringRequest(this.parseSessionSteerParams(methodRequest.params)); + case ASYNC_TASK_STOP_METHOD: { + if (this.providerUpdate !== null) { + await this.providerUpdate; + } + const sessionState = this.sessions.get(methodRequest.params.sessionId); + if (!sessionState) return {stopped: false}; + return { + stopped: await this.runWithProcessCheck( + () => sessionState.asyncTasks.stop(methodRequest.params.asyncTaskId), + ), + }; + } case GOAL_CONTROL_METHOD: case LEGACY_GOAL_CONTROL_METHOD: { const sessionState = this.sessions.get(methodRequest.params.sessionId); @@ -674,6 +693,7 @@ export class CodexAcpServer { clientSupportsSubagents(this.clientCapabilities), new ACPSessionConnection(this.connection, sessionId), ), + asyncTasks: this.createAsyncTasks(sessionId), }; sessionState.titleGen = new TitleGenerator( this.codexAcpClient.appServerClient, @@ -681,7 +701,7 @@ export class CodexAcpServer { sessionState.cwd, () => sessionState.sessionTitleSource, ); - this.sessions.set(sessionId, sessionState); + this.installSessionState(sessionState); resumeSubscribed = false; const canPublishSessionUpdates = operation !== "fork"; @@ -698,6 +718,7 @@ export class CodexAcpServer { } if (operation === "resume") { this.publishCurrentGoalAsync(sessionState, sessionGeneration); + this.publishAsyncTasksAsync(sessionState, sessionGeneration); } const sessionModelState: LegacySessionModelState = this.createModelState(models, currentModelId); const sessionModeState: SessionModeState = sessionState.agentMode.toSessionModeState(); @@ -732,6 +753,20 @@ export class CodexAcpServer { return a === b; } + private createAsyncTasks(sessionId: string): CodexBackgroundTerminalTasks { + return new CodexBackgroundTerminalTasks( + clientSupportsAirCapability(this.clientCapabilities, AIR_ASYNC_TASKS_KEY), + sessionId, + this.codexAcpClient.appServerClient, + new ACPSessionConnection(this.connection, sessionId), + ); + } + + private installSessionState(sessionState: SessionState): void { + this.sessions.get(sessionState.sessionId)?.asyncTasks.clear(); + this.sessions.set(sessionState.sessionId, sessionState); + } + private getAuthProviderForAuthenticateRequest(request: acp.AuthenticateRequest): string | null { if (isCodexAuthRequest(request) && request.methodId === "gateway") { return "custom-gateway"; @@ -752,6 +787,7 @@ export class CodexAcpServer { } = await this.getOrCreateSessionWithHistory(params); await this.streamThreadHistory(sessionId, thread); + await this.getSessionState(sessionId).asyncTasks.reconcile(); logger.log("Session loaded", { sessionId: sessionId, @@ -832,6 +868,7 @@ export class CodexAcpServer { try { if (sessionState) { await this.interruptSessionTurn(sessionState, "Close", true); + sessionState.asyncTasks.clear(); } else { logger.log("Close request received for unknown local session", {sessionId: params.sessionId}); } @@ -998,6 +1035,10 @@ export class CodexAcpServer { } logger.log("Restarting Codex app-server for provider update", {sessionCount: this.sessions.size}); + for (const session of this.sessions.values()) { + session.asyncTasks.prepareForAppServerReplacement(); + } + await this.finishAllAsyncTasks("stopped", "before the provider restart"); const replacement = await this.restartCodexClient(); apply(replacement); if (this.initializeRequest === null) { @@ -1009,6 +1050,7 @@ export class CodexAcpServer { const resumeErrors: unknown[] = []; for (const session of this.sessions.values()) { + session.asyncTasks.setAppServer(replacement.appServerClient); try { await replacement.resumeSession({ sessionId: session.sessionId, @@ -1017,6 +1059,7 @@ export class CodexAcpServer { mcpServers: session.mcpServers ?? [], }); session.authProvider = replacement.getModelProvider(); + session.asyncTasks.refresh(); logger.log("Resumed session after provider restart", {sessionId: session.sessionId}); } catch (error) { resumeErrors.push(error); @@ -1048,6 +1091,16 @@ export class CodexAcpServer { }); } + private observeCodexProcess(): void { + const process = this.codexProcessState?.connection.process; + if (!process) return; + const generation = ++this.codexProcessGeneration; + process.once("exit", () => { + if (generation !== this.codexProcessGeneration) return; + void this.finishAllAsyncTasks("failed", "after the Codex process exited"); + }); + } + private async restartCodexClient(): Promise { const state = this.codexProcessState; if (state === null) { @@ -1055,6 +1108,7 @@ export class CodexAcpServer { } const previous = state.connection; + this.codexProcessGeneration += 1; const exited = previous.process.exitCode === null ? once(previous.process, "exit") : Promise.resolve(); @@ -1071,6 +1125,7 @@ export class CodexAcpServer { state.stderr = ""; state.connection = startCodexConnection(state.codexPath); this.captureStderr(); + this.observeCodexProcess(); return new CodexAcpClient( new CodexAppServerClient(state.connection.connection), state.config, @@ -1723,6 +1778,11 @@ export class CodexAcpServer { void this.publishCurrentGoalBestEffort(sessionState, sessionGeneration, true); } + private publishAsyncTasksAsync(sessionState: SessionState, sessionGeneration: number): void { + if (!this.sessionPublishIsCurrent(sessionState, sessionGeneration)) return; + sessionState.asyncTasks.refresh(); + } + private async publishCurrentGoalBestEffort( sessionState: SessionState, sessionGeneration: number, @@ -1885,6 +1945,7 @@ export class CodexAcpServer { clientSupportsSubagents(this.clientCapabilities), new ACPSessionConnection(this.connection, sessionId), ), + asyncTasks: this.createAsyncTasks(sessionId), }; sessionState.titleGen = new TitleGenerator( this.codexAcpClient.appServerClient, @@ -1892,7 +1953,7 @@ export class CodexAcpServer { sessionState.cwd, () => sessionState.sessionTitleSource, ); - this.sessions.set(sessionId, sessionState); + this.installSessionState(sessionState); subscribed = false; if (requestedMcpServers.length > 0 && mcpServerStartupVersion !== null) { @@ -2002,6 +2063,15 @@ export class CodexAcpServer { new Set([...ancestry, item.agentThreadId]), threadCache, ); + try { + await sessionState.asyncTasks.recover( + item.agentThreadId, + childSessionId, + commandItemIds(childTurn.items), + ); + } catch (error) { + logger.error(`Failed to restore background terminals for ${item.agentThreadId}`, error); + } } } } @@ -3237,6 +3307,7 @@ export class CodexAcpServer { throw new RequestError(requestErrorCode, `VC++ redistributable should be installed`); } if (exitCode !== null) { + await this.finishAllAsyncTasks("failed", "after the Codex process exited"); const stderr = this.getRecentStderr().trim(); const detail = stderr ? `:\n${stderr}` : ""; throw new RequestError(requestErrorCode, `Codex process has exited with code ${exitCode}${detail}`); @@ -3245,6 +3316,16 @@ export class CodexAcpServer { } } + private async finishAllAsyncTasks(state: "failed" | "stopped", reason: string): Promise { + for (const session of this.sessions.values()) { + try { + await session.asyncTasks.finishAll(state); + } catch (error) { + logger.error(`Failed to finish background terminal tasks ${reason}`, error); + } + } + } + async cancel(params: acp.CancelNotification): Promise { const sessionState = this.sessions.get(params.sessionId); if (!sessionState) { @@ -3314,6 +3395,12 @@ function mergeHistoryUpdates( return merged; } +function commandItemIds(items: ThreadItem[]): Set { + return new Set(items + .filter((item): item is Extract => item.type === "commandExecution") + .map(item => item.id)); +} + function historyUpdateKey(update: UpdateSessionEvent): string | null { switch (update.sessionUpdate) { case "user_message_chunk": diff --git a/src/CodexAppServerClient.ts b/src/CodexAppServerClient.ts index 338a97d8..250c76ef 100644 --- a/src/CodexAppServerClient.ts +++ b/src/CodexAppServerClient.ts @@ -82,6 +82,13 @@ import type { PermissionsRequestApprovalResponse, ItemCompletedNotification, } from "./app-server/v2"; +import type { + ThreadBackgroundTerminalsRequest, + ThreadBackgroundTerminalsTerminateParams, + ThreadBackgroundTerminalsTerminateResponse, + ThreadBackgroundTerminalsListParams, + ThreadBackgroundTerminalsListResponse, +} from "./async-tasks/BackgroundTerminalApi"; export interface ApprovalHandler { handleCommandExecution(params: CommandExecutionRequestApprovalParams): Promise; @@ -581,6 +588,14 @@ export class CodexAppServerClient { return await this.sendRequest({ method: "thread/compact/start", params: params }); } + async threadBackgroundTerminalsList(params: ThreadBackgroundTerminalsListParams): Promise { + return await this.sendRequest({method: "thread/backgroundTerminals/list", params}); + } + + async threadBackgroundTerminalsTerminate(params: ThreadBackgroundTerminalsTerminateParams): Promise { + return await this.sendRequest({method: "thread/backgroundTerminals/terminate", params}); + } + async threadGoalSet(params: ThreadGoalSetParams): Promise { return await this.sendRequest({ method: "thread/goal/set", params: params }); } @@ -1021,7 +1036,7 @@ export type CompactionCompletedNotification = | { method: "thread/compacted", params: Extract["params"] } | { method: "item/completed", params: ItemCompletedNotification & { item: Extract } }; -type CodexRequest = DistributiveOmit +type CodexRequest = DistributiveOmit | ThreadBackgroundTerminalsRequest type DistributiveOmit = T extends any ? Omit diff --git a/src/CodexEventHandler.ts b/src/CodexEventHandler.ts index 64f5f60e..7b567541 100644 --- a/src/CodexEventHandler.ts +++ b/src/CodexEventHandler.ts @@ -374,17 +374,32 @@ export class CodexEventHandler { async handleNotification(notification: ServerNotification) { await this.flushPendingErrors(); + const closingChildren = this.subagents.closingChildSessions(notification); + for (const child of closingChildren) { + await this.sessionState.asyncTasks.reconcile(child.threadId, child.sessionId); + } const handledBySubagents = await this.subagents.handle(notification); for (const buffered of this.subagents.takeBufferedNotifications()) { await this.handleNotification(buffered); } - if (handledBySubagents) { - return; + const ignoredBySubagents = !handledBySubagents && this.subagents.shouldIgnore(notification); + let updateEvent: UpdateSessionEvent | null | undefined; + if (!handledBySubagents + && !ignoredBySubagents + && notification.method === "item/started" + && notification.params.item.type === "commandExecution") { + updateEvent = await this.createUpdateEvent(notification); } - if (this.subagents.shouldIgnore(notification)) { - return; + if (!handledBySubagents) { + await this.sessionState.asyncTasks.handleNotification( + notification, + this.subagents.notificationSessionId(notification), + toolCallTitle(updateEvent), + ); } - const updateEvent = await this.createUpdateEvent(notification); + if (handledBySubagents) return; + if (ignoredBySubagents) return; + if (updateEvent === undefined) updateEvent = await this.createUpdateEvent(notification); if (updateEvent) { await this.session.update(updateEvent, this.subagents.notificationSessionId(notification)); } @@ -1338,3 +1353,8 @@ export class CodexEventHandler { return createGuardianApprovalReviewToolCall(params); } } + +function toolCallTitle(update: UpdateSessionEvent | null | undefined): string | undefined { + if (update?.sessionUpdate !== "tool_call") return undefined; + return update.title; +} diff --git a/src/__tests__/CodexACPAgent/async-tasks.test.ts b/src/__tests__/CodexACPAgent/async-tasks.test.ts new file mode 100644 index 00000000..e24c4076 --- /dev/null +++ b/src/__tests__/CodexACPAgent/async-tasks.test.ts @@ -0,0 +1,698 @@ +import {describe, expect, it, vi} from "vitest"; +import {EventEmitter} from "node:events"; +import type {ThreadItem} from "../../app-server/v2"; +import {ACPSessionConnection, type UpdateSessionEvent} from "../../ACPSessionConnection"; +import type {CodexAppServerClient} from "../../CodexAppServerClient"; +import type {CodexConnection} from "../../CodexJsonRpcConnection"; +import {CodexBackgroundTerminalTasks} from "../../async-tasks/CodexBackgroundTerminalTasks"; +import {ASYNC_TASK_STOP_METHOD} from "../../async-tasks/AsyncTaskExtension"; +import {CodexSubagentEventRouter} from "../../subagents/CodexSubagentEventRouter"; +import type { + ThreadBackgroundTerminal, + ThreadBackgroundTerminalsListResponse, +} from "../../async-tasks/BackgroundTerminalApi"; +import { + createCodexMockTestFixture, + createTestSessionState, + setupPromptAndSendNotifications, +} from "../acp-test-utils"; + +type CommandExecutionItem = Extract; + +describe("Codex background terminal tasks", () => { + it("discovers background work from the root session event stream", async () => { + const fixture = createCodexMockTestFixture(); + await fixture.getCodexAcpAgent().initialize({ + protocolVersion: 1, + clientCapabilities: { + _meta: {jetbrains: {air: {version: 1, capabilities: ["asyncTasks"]}}}, + }, + }); + const sessionState = createTestSessionState({sessionId: "thread-1"}); + sessionState.asyncTasks = new CodexBackgroundTerminalTasks( + true, + sessionState.sessionId, + fixture.getCodexAppServerClient(), + new ACPSessionConnection(fixture.getAcpConnection(), sessionState.sessionId), + ); + // @ts-expect-error - register the local session for session-generation checks + fixture.getCodexAcpAgent().sessions.set(sessionState.sessionId, sessionState); + const rawCommand = "/bin/zsh -lc 'npm run build'"; + vi.spyOn(fixture.getCodexAppServerClient(), "threadBackgroundTerminalsList") + .mockResolvedValue(page([terminal({command: rawCommand})])); + + await setupPromptAndSendNotifications(fixture, sessionState.sessionId, sessionState, [ + started(command({command: rawCommand})), + started({type: "reasoning", id: "reasoning-1", summary: [], content: []}), + ]); + + await vi.waitFor(() => { + const taskUpdates = fixture.getAcpConnectionEvents([]) + .filter(event => event.method === "sessionUpdate") + .map(event => event.args[0].update) + .filter(update => update.sessionUpdate === "async_task_spawned"); + expect(taskUpdates).toEqual([expect.objectContaining({ + asyncTaskId: "command-1", + name: "npm run build", + toolCallId: "command-1", + })]); + }); + }); + + it("publishes a background terminal as a task linked to its command", async () => { + const fixture = createFixture(); + fixture.list.mockResolvedValue(page([terminal()])); + + await fixture.tasks.handleNotification(started(command()), "thread-1", "python -m http.server"); + await fixture.tasks.sync(); + + expect(fixture.updates).toEqual([ + { + sessionUpdate: "tool_call_update", + toolCallId: "command-1", + _meta: {jetbrains: {air: {asyncTasks: {backgrounded: true}}}}, + }, + { + sessionUpdate: "async_task_spawned", + asyncTaskId: "command-1", + name: "python -m http.server", + taskType: "shell", + showInTranscript: false, + canStop: true, + toolCallId: "command-1", + }, + ]); + }); + + it("does not publish a command that completes before it becomes background work", async () => { + const fixture = createFixture(); + fixture.list.mockResolvedValue(page([])); + const item = command(); + + await fixture.tasks.handleNotification(started(item), "thread-1"); + await fixture.tasks.handleNotification(completed({...item, status: "completed", exitCode: 0}), "thread-1"); + await fixture.tasks.sync(); + + expect(fixture.updates).toEqual([]); + }); + + it("publishes the terminal state after a background command exits", async () => { + const fixture = createFixture(); + fixture.list.mockResolvedValue(page([terminal()])); + const item = command(); + + await fixture.tasks.handleNotification(started(item), "thread-1"); + await fixture.tasks.sync(); + await fixture.tasks.handleNotification(completed({...item, status: "failed", exitCode: 1}), "thread-1"); + + expect(fixture.updates.at(-1)).toEqual({ + sessionUpdate: "async_task_state_update", + asyncTaskId: "command-1", + state: "failed", + toolCallId: "command-1", + }); + }); + + it("publishes a spawn before a completion that races with discovery", async () => { + const markerStarted = deferred(); + const releaseMarker = deferred(); + const fixture = createFixture(true, async update => { + if (update.sessionUpdate !== "tool_call_update") return; + markerStarted.resolve(); + await releaseMarker.promise; + }); + fixture.list.mockResolvedValue(page([terminal()])); + + const sync = fixture.tasks.sync(); + await markerStarted.promise; + const completion = fixture.tasks.handleNotification( + completed({...command(), status: "completed", exitCode: 0}), + "thread-1", + ); + releaseMarker.resolve(); + await Promise.all([sync, completion]); + + const spawnIndex = fixture.updates.findIndex(update => update.sessionUpdate === "async_task_spawned"); + const terminalIndex = fixture.updates.findIndex(update => update.sessionUpdate === "async_task_state_update"); + expect(spawnIndex).toBeGreaterThanOrEqual(0); + expect(terminalIndex).toBeGreaterThan(spawnIndex); + }); + + it("retries a terminal update that the client rejected", async () => { + let rejectTerminalUpdate = true; + const fixture = createFixture(true, async update => { + if (update.sessionUpdate === "async_task_state_update" && rejectTerminalUpdate) { + rejectTerminalUpdate = false; + throw new Error("client disconnected"); + } + }); + fixture.list.mockResolvedValue(page([terminal()])); + const completedItem = {...command(), status: "completed" as const, exitCode: 0}; + await fixture.tasks.sync(); + + await expect(fixture.tasks.handleNotification(completed(completedItem), "thread-1")) + .rejects.toThrow("client disconnected"); + await expect(fixture.tasks.handleNotification(completed(completedItem), "thread-1")) + .resolves.toBeUndefined(); + + expect(fixture.updates.filter(update => update.sessionUpdate === "async_task_state_update")) + .toEqual([expect.objectContaining({state: "completed"})]); + }); + + it("finishes an announced task that disappears from the live terminal list", async () => { + const fixture = createFixture(); + fixture.list + .mockResolvedValueOnce(page([terminal()])) + .mockResolvedValueOnce(page([])); + + await fixture.tasks.sync(); + await fixture.tasks.sync(); + + expect(fixture.updates.at(-1)).toEqual({ + sessionUpdate: "async_task_state_update", + asyncTaskId: "command-1", + state: "stopped", + toolCallId: "command-1", + }); + }); + + it("stops one task through the app-server process id", async () => { + const fixture = createFixture(); + fixture.list.mockResolvedValue(page([terminal()])); + fixture.terminate.mockResolvedValue({terminated: true}); + await fixture.tasks.sync(); + + await expect(fixture.tasks.stop("command-1")).resolves.toBe(true); + + expect(fixture.terminate).toHaveBeenCalledWith({ + threadId: "thread-1", + processId: "42", + }); + expect(fixture.updates.at(-1)).toEqual({ + sessionUpdate: "async_task_state_update", + asyncTaskId: "command-1", + state: "stopped", + toolCallId: "command-1", + }); + await expect(fixture.tasks.stop("command-1")).resolves.toBe(false); + }); + + it("uses the replacement app-server client for stop requests", async () => { + const fixture = createFixture(); + fixture.list.mockResolvedValue(page([terminal()])); + await fixture.tasks.sync(); + const replacementTerminate = vi.fn().mockResolvedValue({terminated: true}); + fixture.tasks.setAppServer({ + threadBackgroundTerminalsTerminate: replacementTerminate, + } as unknown as CodexAppServerClient); + + await expect(fixture.tasks.stop("command-1")).resolves.toBe(true); + + expect(fixture.terminate).not.toHaveBeenCalled(); + expect(replacementTerminate).toHaveBeenCalledWith({threadId: "thread-1", processId: "42"}); + }); + + it("fails every announced task after the app-server exits", async () => { + const fixture = createFixture(); + fixture.list.mockResolvedValue(page([ + terminal(), + terminal({itemId: "command-2", processId: "84"}), + ])); + await fixture.tasks.sync(); + + await fixture.tasks.finishAll("failed"); + + expect(fixture.updates.filter(update => update.sessionUpdate === "async_task_state_update")) + .toEqual([ + expect.objectContaining({asyncTaskId: "command-1", state: "failed"}), + expect.objectContaining({asyncTaskId: "command-2", state: "failed"}), + ]); + }); + + it("finishes tasks while app-server queries are suspended for replacement", async () => { + const fixture = createFixture(); + fixture.list.mockResolvedValue(page([terminal()])); + await fixture.tasks.sync(); + + fixture.tasks.prepareForAppServerReplacement(); + await fixture.tasks.finishAll("stopped"); + + expect(fixture.updates.at(-1)).toMatchObject({ + sessionUpdate: "async_task_state_update", + asyncTaskId: "command-1", + state: "stopped", + }); + }); + + it("fails announced tasks as soon as the app-server process exits", async () => { + const process = Object.assign(new EventEmitter(), { + stderr: new EventEmitter(), + stdin: {end: vi.fn()}, + exitCode: null, + }) as unknown as CodexConnection["process"]; + const fixture = createCodexMockTestFixture(undefined, process); + const sessionState = createTestSessionState({sessionId: "thread-1"}); + sessionState.asyncTasks = new CodexBackgroundTerminalTasks( + true, + sessionState.sessionId, + fixture.getCodexAppServerClient(), + new ACPSessionConnection(fixture.getAcpConnection(), sessionState.sessionId), + ); + vi.spyOn(fixture.getCodexAppServerClient(), "threadBackgroundTerminalsList") + .mockResolvedValue(page([terminal()])); + await sessionState.asyncTasks.sync(); + // @ts-expect-error - register the local session for process lifecycle checks + fixture.getCodexAcpAgent().sessions.set(sessionState.sessionId, sessionState); + + process.emit("exit", 1); + + await vi.waitFor(() => { + const terminalUpdate = fixture.getAcpConnectionEvents([]) + .filter(event => event.method === "sessionUpdate") + .map(event => event.args[0].update) + .find(update => update.sessionUpdate === "async_task_state_update"); + expect(terminalUpdate).toMatchObject({asyncTaskId: "command-1", state: "failed"}); + }); + }); + + it("keeps a task stoppable when termination is rejected or fails", async () => { + const fixture = createFixture(); + fixture.list.mockResolvedValue(page([terminal()])); + fixture.terminate + .mockResolvedValueOnce({terminated: false}) + .mockRejectedValueOnce(new Error("terminate failed")) + .mockResolvedValueOnce({terminated: true}); + await fixture.tasks.sync(); + + await expect(fixture.tasks.stop("command-1")).resolves.toBe(false); + await expect(fixture.tasks.stop("command-1")).rejects.toThrow("terminate failed"); + await expect(fixture.tasks.stop("command-1")).resolves.toBe(true); + + expect(fixture.terminate).toHaveBeenCalledTimes(3); + }); + + it("keeps a task stoppable when its stopped update fails", async () => { + let rejectStoppedUpdate = true; + const fixture = createFixture(true, async update => { + if (update.sessionUpdate === "async_task_state_update" && rejectStoppedUpdate) { + rejectStoppedUpdate = false; + throw new Error("client disconnected"); + } + }); + fixture.list.mockResolvedValue(page([terminal()])); + fixture.terminate.mockResolvedValueOnce({terminated: true}); + await fixture.tasks.sync(); + + await expect(fixture.tasks.stop("command-1")).rejects.toThrow("client disconnected"); + await expect(fixture.tasks.stop("command-1")).resolves.toBe(true); + + expect(fixture.terminate).toHaveBeenCalledOnce(); + expect(fixture.updates.filter(update => update.sessionUpdate === "async_task_state_update")) + .toEqual([expect.objectContaining({state: "stopped"})]); + }); + + it("does not overwrite completion that races with a stop request", async () => { + const fixture = createFixture(); + const termination = deferred<{terminated: boolean}>(); + fixture.list.mockResolvedValue(page([terminal()])); + fixture.terminate.mockReturnValue(termination.promise); + await fixture.tasks.sync(); + + const stop = fixture.tasks.stop("command-1"); + await fixture.tasks.handleNotification( + completed({...command(), status: "completed", exitCode: 0}), + "thread-1", + ); + termination.resolve({terminated: true}); + + await expect(stop).resolves.toBe(true); + expect(fixture.updates.filter(update => update.sessionUpdate === "async_task_state_update")) + .toEqual([expect.objectContaining({state: "completed"})]); + }); + + it("routes the AIR stop request to the session task runtime", async () => { + const fixture = createCodexMockTestFixture(); + const sessionState = createTestSessionState({sessionId: "thread-1"}); + const stop = vi.spyOn(sessionState.asyncTasks, "stop").mockResolvedValue(true); + // @ts-expect-error - register the local session for the extension request path + fixture.getCodexAcpAgent().sessions.set(sessionState.sessionId, sessionState); + + await expect(fixture.getCodexAcpAgent().extMethod(ASYNC_TASK_STOP_METHOD, { + sessionId: sessionState.sessionId, + asyncTaskId: "command-1", + })).resolves.toEqual({stopped: true}); + expect(stop).toHaveBeenCalledWith("command-1"); + }); + + it("reads every background terminal page", async () => { + const fixture = createFixture(); + fixture.list + .mockResolvedValueOnce(page([terminal()], "42")) + .mockResolvedValueOnce(page([terminal({itemId: "command-2", processId: "84"})])); + + await fixture.tasks.sync(); + + expect(fixture.list).toHaveBeenNthCalledWith(1, { + threadId: "thread-1", + cursor: null, + }); + expect(fixture.list).toHaveBeenNthCalledWith(2, { + threadId: "thread-1", + cursor: "42", + }); + expect(fixture.updates.filter(update => update.sessionUpdate === "async_task_spawned")).toHaveLength(2); + }); + + it("rejects a repeated background terminal cursor", async () => { + const fixture = createFixture(); + fixture.list + .mockResolvedValueOnce(page([], "repeated")) + .mockResolvedValueOnce(page([], "repeated")); + + await expect(fixture.tasks.sync()).rejects.toThrow("repeated background terminal cursor"); + + expect(fixture.list).toHaveBeenCalledTimes(2); + }); + + it("runs a trailing refresh when a list request is already in flight", async () => { + const fixture = createFixture(); + const listing = deferred(); + fixture.list + .mockReturnValueOnce(listing.promise) + .mockResolvedValueOnce(page([terminal()])); + + const first = fixture.tasks.sync(); + const second = fixture.tasks.sync(); + listing.resolve(page([terminal()])); + await Promise.all([first, second]); + + expect(fixture.list).toHaveBeenCalledTimes(2); + expect(fixture.updates.filter(update => update.sessionUpdate === "async_task_spawned")).toHaveLength(1); + }); + + it("runs a trailing refresh after the in-flight request fails", async () => { + const fixture = createFixture(); + const listing = deferred(); + fixture.list + .mockReturnValueOnce(listing.promise) + .mockResolvedValueOnce(page([terminal()])); + + const first = fixture.tasks.sync(); + const second = fixture.tasks.sync(); + listing.reject(new Error("temporary list failure")); + await Promise.all([first, second]); + + expect(fixture.list).toHaveBeenCalledTimes(2); + expect(fixture.updates.filter(update => update.sessionUpdate === "async_task_spawned")).toHaveLength(1); + }); + + it("discards an in-flight list result after app-server replacement", async () => { + const fixture = createFixture(); + const oldListing = deferred(); + fixture.list.mockReturnValue(oldListing.promise); + const replacementList = vi.fn().mockResolvedValue(page([ + terminal({itemId: "command-2", processId: "84"}), + ])); + const replacement = { + threadBackgroundTerminalsList: replacementList, + threadBackgroundTerminalsTerminate: vi.fn(), + } as unknown as CodexAppServerClient; + + const oldSync = fixture.tasks.sync(); + fixture.tasks.prepareForAppServerReplacement(); + fixture.tasks.setAppServer(replacement); + const replacementSync = fixture.tasks.sync(); + oldListing.resolve(page([terminal()])); + await Promise.all([oldSync, replacementSync]); + + expect(replacementList).toHaveBeenCalledOnce(); + expect(fixture.updates.filter(update => update.sessionUpdate === "async_task_spawned")) + .toEqual([expect.objectContaining({asyncTaskId: "command-2"})]); + }); + + it("does not publish an in-flight list result after clear", async () => { + const fixture = createFixture(); + const listing = deferred(); + fixture.list.mockReturnValue(listing.promise); + + const sync = fixture.tasks.sync(); + fixture.tasks.clear(); + listing.resolve(page([terminal()])); + await sync; + + expect(fixture.updates).toEqual([]); + }); + + it.each([ + ["turn completion", () => turnCompleted("child-1")], + ["activity interruption", childInterrupted], + ["collaboration completion", childCompletedByCollaboration], + ])("publishes a child terminal before %s closes its session", async (_name, terminalNotification) => { + const fixture = createCodexMockTestFixture(); + await fixture.getCodexAcpAgent().initialize({ + protocolVersion: 1, + clientCapabilities: { + _meta: {jetbrains: {air: {version: 1, capabilities: ["asyncTasks", "nativeSubagentSessions"]}}}, + }, + }); + const sessionState = createTestSessionState({sessionId: "thread-1"}); + const session = new ACPSessionConnection(fixture.getAcpConnection(), sessionState.sessionId); + sessionState.subagents = new CodexSubagentEventRouter(sessionState.sessionId, true, session); + sessionState.asyncTasks = new CodexBackgroundTerminalTasks( + true, + sessionState.sessionId, + fixture.getCodexAppServerClient(), + session, + ); + // @ts-expect-error - register the local session for session-generation checks + fixture.getCodexAcpAgent().sessions.set(sessionState.sessionId, sessionState); + let childListCount = 0; + vi.spyOn(fixture.getCodexAppServerClient(), "threadBackgroundTerminalsList") + .mockImplementation(async params => page( + params.threadId === "child-1" && childListCount++ === 0 ? [terminal()] : [], + )); + + await setupPromptAndSendNotifications(fixture, sessionState.sessionId, sessionState, [ + childSpawned(), + childMaterialized(), + started(command(), "child-1"), + terminalNotification(), + ]); + + await vi.waitFor(() => { + const spawned = fixture.getAcpConnectionEvents([]) + .filter(event => event.method === "sessionUpdate") + .map(event => event.args[0]) + .find(event => event.update.sessionUpdate === "async_task_spawned"); + expect(spawned).toMatchObject({ + sessionId: "child-1", + update: {asyncTaskId: "child-1:command-1", toolCallId: "command-1"}, + }); + const childTerminalIndex = fixture.getAcpConnectionEvents([]) + .filter(event => event.method === "sessionUpdate") + .map(event => event.args[0]) + .findIndex(event => event.update.sessionUpdate === "subagent_state_update" + && event.update.subagentSessionId === "child-1"); + const taskSpawnIndex = fixture.getAcpConnectionEvents([]) + .filter(event => event.method === "sessionUpdate") + .map(event => event.args[0]) + .findIndex(event => event.update.sessionUpdate === "async_task_spawned"); + expect(taskSpawnIndex).toBeGreaterThanOrEqual(0); + expect(childTerminalIndex).toBeGreaterThan(taskSpawnIndex); + }); + + const terminate = vi.spyOn(fixture.getCodexAppServerClient(), "threadBackgroundTerminalsTerminate") + .mockResolvedValue({terminated: true}); + await expect(fixture.getCodexAcpAgent().extMethod(ASYNC_TASK_STOP_METHOD, { + sessionId: sessionState.sessionId, + asyncTaskId: "child-1:command-1", + })).resolves.toEqual({stopped: true}); + expect(terminate).toHaveBeenCalledWith({threadId: "child-1", processId: "42"}); + }); + + it("routes a child task stop to its owning Codex thread", async () => { + const fixture = createFixture(); + fixture.list.mockResolvedValue(page([terminal()])); + fixture.terminate.mockResolvedValue({terminated: true}); + await fixture.tasks.handleNotification(started(command(), "child-1"), "child-1"); + await fixture.tasks.sync("child-1", "child-1"); + + await expect(fixture.tasks.stop("child-1:command-1")).resolves.toBe(true); + + expect(fixture.terminate).toHaveBeenCalledWith({threadId: "child-1", processId: "42"}); + }); + + it("does nothing when the client did not negotiate async tasks", async () => { + const fixture = createFixture(false); + + await fixture.tasks.handleNotification(started(command()), "thread-1"); + await fixture.tasks.sync(); + await expect(fixture.tasks.stop("command-1")).resolves.toBe(false); + + expect(fixture.list).not.toHaveBeenCalled(); + expect(fixture.terminate).not.toHaveBeenCalled(); + expect(fixture.updates).toEqual([]); + }); +}); + +function createFixture( + enabled = true, + beforeUpdate?: (update: UpdateSessionEvent) => void | Promise, +) { + const updates: UpdateSessionEvent[] = []; + const list = vi.fn<() => Promise>(); + const terminate = vi.fn(); + const appServer = { + threadBackgroundTerminalsList: list, + threadBackgroundTerminalsTerminate: terminate, + } as unknown as CodexAppServerClient; + const session = new ACPSessionConnection({ + notify: vi.fn(async (_method, params) => { + const update = (params as {update: UpdateSessionEvent}).update; + await beforeUpdate?.(update); + updates.push(update); + }), + request: vi.fn(), + }, "thread-1"); + return { + updates, + list, + terminate, + tasks: new CodexBackgroundTerminalTasks(enabled, "thread-1", appServer, session), + }; +} + +function terminal(overrides: Partial = {}): ThreadBackgroundTerminal { + return { + itemId: "command-1", + processId: "42", + command: "python -m http.server", + ...overrides, + }; +} + +function page(data: ThreadBackgroundTerminal[], nextCursor: string | null = null): ThreadBackgroundTerminalsListResponse { + return {data, nextCursor}; +} + +function command(overrides: Partial = {}): CommandExecutionItem { + return { + type: "commandExecution", + id: "command-1", + pluginId: null, + scriptPath: null, + command: "python -m http.server", + cwd: "/workspace", + processId: "42", + source: "unifiedExecStartup", + status: "inProgress", + commandActions: [], + aggregatedOutput: null, + exitCode: null, + durationMs: null, + ...overrides, + }; +} + +function started(item: ThreadItem, threadId = "thread-1") { + return { + method: "item/started" as const, + params: { + threadId, + turnId: "turn-id", + startedAtMs: 0, + item, + }, + }; +} + +function completed(item: ThreadItem, threadId = "thread-1") { + return { + method: "item/completed" as const, + params: { + threadId, + turnId: "turn-id", + completedAtMs: 0, + item, + }, + }; +} + +function turnCompleted(threadId: string) { + return { + method: "turn/completed" as const, + params: { + threadId, + turn: { + id: "turn-id", + items: [], + itemsView: "notLoaded" as const, + status: "completed" as const, + error: null, + startedAt: null, + completedAt: null, + durationMs: null, + }, + }, + }; +} + +function childSpawned() { + return started({ + type: "collabAgentToolCall", + id: "spawn-child", + tool: "spawnAgent", + status: "inProgress", + senderThreadId: "thread-1", + receiverThreadIds: ["child-1"], + prompt: "Run a server", + model: null, + reasoningEffort: null, + agentsStates: {"child-1": {status: "running", message: null}}, + }); +} + +function childMaterialized() { + return started({ + type: "subAgentActivity", + id: "child-activity", + kind: "started", + agentThreadId: "child-1", + agentPath: "/root/worker", + }); +} + +function childInterrupted() { + return started({ + type: "subAgentActivity", + id: "child-activity-terminal", + kind: "interrupted", + agentThreadId: "child-1", + agentPath: "/root/worker", + }); +} + +function childCompletedByCollaboration() { + return started({ + type: "collabAgentToolCall", + id: "wait-child", + tool: "wait", + status: "completed", + senderThreadId: "thread-1", + receiverThreadIds: ["child-1"], + prompt: null, + model: null, + reasoningEffort: null, + agentsStates: {"child-1": {status: "completed", message: null}}, + }); +} + +function deferred() { + let resolve!: (value?: T) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((innerResolve, innerReject) => { + resolve = innerResolve as (value?: T) => void; + reject = innerReject; + }); + return {promise, resolve, reject}; +} diff --git a/src/__tests__/CodexACPAgent/initialize.test.ts b/src/__tests__/CodexACPAgent/initialize.test.ts index 30f5e94d..7fff7721 100644 --- a/src/__tests__/CodexACPAgent/initialize.test.ts +++ b/src/__tests__/CodexACPAgent/initialize.test.ts @@ -78,7 +78,7 @@ describe('CodexACPAgent - initialize', () => { jetbrains: { air: { version: 1, - capabilities: ["sessionFailure", "agentFileChangeReport", "nativeSubagentSessions"], + capabilities: ["sessionFailure", "agentFileChangeReport", "nativeSubagentSessions", "asyncTasks"], }, }, }, diff --git a/src/__tests__/CodexACPAgent/load-session.test.ts b/src/__tests__/CodexACPAgent/load-session.test.ts index f8460836..ae8ce09a 100644 --- a/src/__tests__/CodexACPAgent/load-session.test.ts +++ b/src/__tests__/CodexACPAgent/load-session.test.ts @@ -92,15 +92,32 @@ describe("CodexACPAgent - loadSession", () => { agentPath: "/root/orphan_child", }, ]); - const child = makeThread("child-history", [{ - type: "agentMessage", - id: "child-history-message-1", - text: "Persisted first-generation output", - phase: null, - memoryCitation: null, - delivery: null, - questions: null, - }]); + const child = makeThread("child-history", [ + { + type: "commandExecution", + id: "child-command-1", + pluginId: null, + scriptPath: null, + command: "python -m http.server", + cwd: "/workspace", + processId: "42", + source: "unifiedExecStartup", + status: "inProgress", + commandActions: [], + aggregatedOutput: null, + exitCode: null, + durationMs: null, + }, + { + type: "agentMessage", + id: "child-history-message-1", + text: "Persisted first-generation output", + phase: null, + memoryCitation: null, + delivery: null, + questions: null, + }, + ]); const firstChildTurn = child.turns[0]!; child.turns.push({ id: "turn-child-history-2", @@ -133,11 +150,17 @@ describe("CodexACPAgent - loadSession", () => { if (threadId === "orphan-history") return Promise.reject(new Error("missing child history")); return Promise.resolve({thread: threadId === root.id ? root : child}); }); + appServer.threadBackgroundTerminalsList = vi.fn().mockImplementation(({threadId}) => Promise.resolve({ + data: threadId === "child-history" + ? [{itemId: "child-command-1", processId: "42", command: "python -m http.server"}] + : [], + nextCursor: null, + })); await agent.initialize({ protocolVersion: 1, clientCapabilities: { - _meta: {jetbrains: {air: {version: 1, capabilities: ["nativeSubagentSessions"]}}}, + _meta: {jetbrains: {air: {version: 1, capabilities: ["nativeSubagentSessions", "asyncTasks"]}}}, }, }); await agent.loadSession({sessionId: root.id, cwd: "/workspace", mcpServers: []}); @@ -148,12 +171,18 @@ describe("CodexACPAgent - loadSession", () => { const firstSpawnIndex = updates.findIndex(({update}) => update.subagentSessionId === "child-history" && update.sessionUpdate === "subagent_spawned"); const firstOutputIndex = updates.findIndex(({update}) => update.messageId === "child-history-message-1"); + const childTaskIndex = updates.findIndex(({update}) => update.sessionUpdate === "async_task_spawned" + && update.asyncTaskId === "child-history:child-command-1"); + const firstTerminalIndex = updates.findIndex(({update}) => update.sessionUpdate === "subagent_state_update" + && update.subagentSessionId === "child-history"); const secondSpawnIndex = updates.findIndex(({update}) => update.subagentSessionId === "child-history:generation:2" && update.sessionUpdate === "subagent_spawned"); const secondOutputIndex = updates.findIndex(({update}) => update.messageId === "child-history-message-2"); const orphanTerminalIndex = updates.findIndex(({update}) => update.subagentSessionId === "orphan-history" && update.state === "disconnected"); expect(firstOutputIndex).toBeGreaterThan(firstSpawnIndex); + expect(childTaskIndex).toBeGreaterThan(firstSpawnIndex); + expect(firstTerminalIndex).toBeGreaterThan(childTaskIndex); expect(secondSpawnIndex).toBeGreaterThan(firstOutputIndex); expect(secondOutputIndex).toBeGreaterThan(secondSpawnIndex); expect(orphanTerminalIndex).toBeGreaterThan(secondOutputIndex); diff --git a/src/__tests__/CodexACPAgent/providers.test.ts b/src/__tests__/CodexACPAgent/providers.test.ts index 45bc744c..5a724039 100644 --- a/src/__tests__/CodexACPAgent/providers.test.ts +++ b/src/__tests__/CodexACPAgent/providers.test.ts @@ -211,8 +211,14 @@ describe("Configurable LLM providers (providers/*)", () => { const agent = fixture.getCodexAcpAgent(); await agent.initialize({protocolVersion: acp.PROTOCOL_VERSION}); const sessions = (agent as unknown as {sessions: Map>}).sessions; - sessions.set("thread-1", createTestSessionState({sessionId: "thread-1", cwd: "/one"})); - sessions.set("thread-2", createTestSessionState({sessionId: "thread-2", cwd: "/two"})); + const firstSession = createTestSessionState({sessionId: "thread-1", cwd: "/one"}); + const secondSession = createTestSessionState({sessionId: "thread-2", cwd: "/two"}); + const firstSessionPrepare = vi.spyOn(firstSession.asyncTasks, "prepareForAppServerReplacement"); + const secondSessionPrepare = vi.spyOn(secondSession.asyncTasks, "prepareForAppServerReplacement"); + const firstSessionSetAppServer = vi.spyOn(firstSession.asyncTasks, "setAppServer"); + const secondSessionSetAppServer = vi.spyOn(secondSession.asyncTasks, "setAppServer"); + sessions.set("thread-1", firstSession); + sessions.set("thread-2", secondSession); await agent.setProvider({ providerId: OPENAI_PROVIDER_ID, @@ -224,6 +230,10 @@ describe("Configurable LLM providers (providers/*)", () => { expect(restart).toHaveBeenCalledTimes(1); expect(firstGatewayReplacement.getModelProvider()).toBe(CUSTOM_GATEWAY_PROVIDER_ID); expect(firstGatewayResume).toHaveBeenCalledTimes(2); + expect(firstSessionPrepare).toHaveBeenCalledOnce(); + expect(secondSessionPrepare).toHaveBeenCalledOnce(); + expect(firstSessionSetAppServer).toHaveBeenCalledWith(firstGatewayReplacement.appServerClient); + expect(secondSessionSetAppServer).toHaveBeenCalledWith(firstGatewayReplacement.appServerClient); expect(firstGatewayResume).toHaveBeenCalledWith(expect.objectContaining({sessionId: "thread-1", cwd: "/one"})); expect(firstGatewayResume).toHaveBeenCalledWith(expect.objectContaining({sessionId: "thread-2", cwd: "/two"})); diff --git a/src/__tests__/acp-test-utils.ts b/src/__tests__/acp-test-utils.ts index 599db8ae..f358664c 100644 --- a/src/__tests__/acp-test-utils.ts +++ b/src/__tests__/acp-test-utils.ts @@ -2,8 +2,8 @@ import * as acp from "@agentclientprotocol/sdk"; import type {CreateElicitationResponse, McpServerStdio, RequestPermissionResponse} from "@agentclientprotocol/sdk"; import {CodexAcpClient} from '../CodexAcpClient'; import {CodexAppServerClient, type CodexConnectionEvent} from '../CodexAppServerClient'; -import {startCodexConnection} from "../CodexJsonRpcConnection"; -import {CodexAcpServer, type SessionState} from "../CodexAcpServer"; +import {type CodexConnection, startCodexConnection} from "../CodexJsonRpcConnection"; +import {CodexAcpServer, type CodexProcessState, type SessionState} from "../CodexAcpServer"; import {ACPSessionConnection, type AcpClientConnection} from "../ACPSessionConnection"; import type {ServerNotification} from "../app-server"; import type {MessageConnection} from "vscode-jsonrpc/node"; @@ -15,6 +15,7 @@ import {DEFAULT_COLLABORATION_MODE} from "../CollaborationModeConfig"; import {expect, vi} from "vitest"; import type {Model, ReasoningEffortOption} from "../app-server/v2"; import {CodexSubagentEventRouter} from "../subagents/CodexSubagentEventRouter"; +import {CodexBackgroundTerminalTasks} from "../async-tasks/CodexBackgroundTerminalTasks"; import {AUTH_STATUS_UPDATE_METHOD} from "../AuthStatusMeta"; export type MethodCallEvent = { method: string; args: any[] }; @@ -88,6 +89,7 @@ export interface ConnectionConfig { connection: MessageConnection; getExitCode: () => number | null; acpConnection?: AcpConnectionConfig; + codexProcessState?: CodexProcessState; } export function createBaseTestFixture(config: ConnectionConfig): TestFixture { @@ -107,6 +109,7 @@ export function createBaseTestFixture(config: ConnectionConfig): TestFixture { undefined, config.getExitCode, undefined, + config.codexProcessState, ); const transportEvents: CodexConnectionEvent[] = []; @@ -269,6 +272,7 @@ export interface CodexMockTestFixture extends TestFixture { */ export function createCodexMockTestFixture( restartCodexClient?: () => Promise, + process?: CodexConnection["process"], ): CodexMockTestFixture { let unhandledNotificationHandler: ((notification: any) => void) | null = null; const requestHandlers = new Map Promise>(); @@ -317,6 +321,13 @@ export function createCodexMockTestFixture( const baseFixture = createBaseTestFixture({ connection: mockCodexConnection, getExitCode: () => null, + ...(process ? {codexProcessState: { + connection: {connection: mockCodexConnection, process}, + codexPath: undefined, + config: undefined, + modelProvider: undefined, + stderr: "", + }} : {}), acpConnection: { connection: acpConnection, events: acpConnectionEvents, @@ -415,6 +426,12 @@ export function createTestSessionState(overrides?: Partial): Sessi false, new ACPSessionConnection({notify: vi.fn(), request: vi.fn()} as AcpClientConnection, sessionId), ), + asyncTasks: new CodexBackgroundTerminalTasks( + false, + sessionId, + {} as CodexAppServerClient, + new ACPSessionConnection({notify: vi.fn(), request: vi.fn()} as AcpClientConnection, sessionId), + ), ...overrides, }; } diff --git a/src/async-tasks/AcpAsyncTasks.ts b/src/async-tasks/AcpAsyncTasks.ts new file mode 100644 index 00000000..e658912a --- /dev/null +++ b/src/async-tasks/AcpAsyncTasks.ts @@ -0,0 +1,41 @@ +export type AsyncTaskState = "running" | "paused" | "completed" | "failed" | "stopped"; + +export type AsyncTaskSpawnedUpdate = { + sessionUpdate: "async_task_spawned"; + asyncTaskId: string; + name: string; + taskType: string; + description?: string; + showInTranscript: boolean; + canStop: boolean; + outputFilePath?: string; + toolCallId?: string; + _meta?: Record | null; +}; + +export type AsyncTaskProgressUpdate = { + sessionUpdate: "async_task_progress"; + asyncTaskId: string; + description?: string; + summary?: string; + lastToolName?: string; + usage?: { totalTokens: number; toolUses: number; durationMs: number }; + outputFilePath?: string; + toolCallId?: string; + _meta?: Record | null; +}; + +export type AsyncTaskStateUpdate = { + sessionUpdate: "async_task_state_update"; + asyncTaskId: string; + state: AsyncTaskState; + summary?: string; + outputFilePath?: string; + toolCallId?: string; + _meta?: Record | null; +}; + +export type AsyncTaskUpdate = + | AsyncTaskSpawnedUpdate + | AsyncTaskProgressUpdate + | AsyncTaskStateUpdate; diff --git a/src/async-tasks/AsyncTaskExtension.ts b/src/async-tasks/AsyncTaskExtension.ts new file mode 100644 index 00000000..4ed8d7ab --- /dev/null +++ b/src/async-tasks/AsyncTaskExtension.ts @@ -0,0 +1,17 @@ +import type {SessionId} from "@agentclientprotocol/sdk"; + +export const ASYNC_TASK_STOP_METHOD = "_session/async_task/stop"; + +export type AsyncTaskStopRequest = { + sessionId: SessionId; + asyncTaskId: string; +}; + +export type AsyncTaskStopResponse = { + stopped: boolean; +}; + +export type AsyncTaskStopExtRequest = { + method: typeof ASYNC_TASK_STOP_METHOD; + params: AsyncTaskStopRequest; +}; diff --git a/src/async-tasks/BackgroundTerminalApi.ts b/src/async-tasks/BackgroundTerminalApi.ts new file mode 100644 index 00000000..d3238c2f --- /dev/null +++ b/src/async-tasks/BackgroundTerminalApi.ts @@ -0,0 +1,30 @@ +/** The fields used from an API that stable `generate-ts` output omits. */ +export type ThreadBackgroundTerminal = { + itemId: string; + processId: string; + command: string; +}; + +export type ThreadBackgroundTerminalsListParams = { + threadId: string; + cursor?: string | null; + limit?: number | null; +}; + +export type ThreadBackgroundTerminalsListResponse = { + data: ThreadBackgroundTerminal[]; + nextCursor: string | null; +}; + +export type ThreadBackgroundTerminalsTerminateParams = { + threadId: string; + processId: string; +}; + +export type ThreadBackgroundTerminalsTerminateResponse = { + terminated: boolean; +}; + +export type ThreadBackgroundTerminalsRequest = + | { method: "thread/backgroundTerminals/list"; params: ThreadBackgroundTerminalsListParams } + | { method: "thread/backgroundTerminals/terminate"; params: ThreadBackgroundTerminalsTerminateParams }; diff --git a/src/async-tasks/CodexBackgroundTerminalTasks.ts b/src/async-tasks/CodexBackgroundTerminalTasks.ts new file mode 100644 index 00000000..61ead906 --- /dev/null +++ b/src/async-tasks/CodexBackgroundTerminalTasks.ts @@ -0,0 +1,416 @@ +import type {ServerNotification} from "../app-server"; +import type {ThreadItem} from "../app-server/v2"; +import type {ACPSessionConnection} from "../ACPSessionConnection"; +import type {CodexAppServerClient} from "../CodexAppServerClient"; +import { + AIR_ASYNC_TASKS_BACKGROUNDED_KEY, + AIR_ASYNC_TASKS_KEY, + AIR_META_KEY, + JETBRAINS_META_KEY, +} from "../AirExtension"; +import {logger} from "../Logger"; +import type {ThreadBackgroundTerminal} from "./BackgroundTerminalApi"; + +type CommandExecutionItem = Extract; +type TerminalState = "completed" | "failed" | "stopped"; + +type Task = { + threadId: string; + sessionId: string; + asyncTaskId: string; + processId: string; + itemId: string; + name: string; + publication: "unpublished" | "publishing" | "published"; + announcement: Promise | null; + terminalUpdate: Promise | null; + terminalPublished: boolean; + state: "running" | "stopping" | TerminalState; +}; + +type PendingSync = { + requested: boolean; + sessionId: string; + promise: Promise; +}; + +type TerminalSnapshot = { + generation: number; + terminals: ThreadBackgroundTerminal[]; +}; + +/** Maps Codex-owned background terminals to the AIR async task extension. */ +export class CodexBackgroundTerminalTasks { + private readonly tasks = new Map(); + private readonly syncs = new Map(); + private appServerGeneration = 0; + private appServerQueriesEnabled = true; + private disposed = false; + + constructor( + readonly enabled: boolean, + private readonly rootSessionId: string, + private appServer: CodexAppServerClient, + private readonly session: ACPSessionConnection, + ) {} + + async handleNotification( + notification: ServerNotification, + sessionId: string, + commandTitle?: string, + ): Promise { + if (!this.isActive()) return; + const threadId = notificationThreadId(notification); + if (threadId === null) return; + + if (notification.method === "item/started" && notification.params.item.type === "commandExecution") { + this.observeCommandStarted(notification.params.item, threadId, sessionId, commandTitle); + return; + } + if (notification.method === "item/completed" && notification.params.item.type === "commandExecution") { + await this.observeCommandCompleted(notification.params.item, threadId); + return; + } + if (notification.method === "turn/completed") { + await this.reconcile(threadId, sessionId); + return; + } + if (notification.method === "item/started") { + this.refresh(threadId, sessionId); + } + } + + private observeCommandStarted( + item: CommandExecutionItem, + threadId: string, + sessionId: string, + commandTitle?: string, + ): void { + if (!this.isActive() || item.processId === null) return; + this.remember(threadId, sessionId, { + itemId: item.id, + processId: item.processId, + command: item.command, + }, commandTitle); + } + + private async observeCommandCompleted( + item: CommandExecutionItem, + threadId: string, + ): Promise { + if (!this.isActive()) return; + const task = this.tasks.get(wireTaskId(this.rootSessionId, threadId, item.id)); + if (task) await this.finish(task, item.status === "completed" ? "completed" : "failed"); + } + + refresh(threadId: string = this.rootSessionId, sessionId: string = this.rootSessionId): void { + void this.reconcile(threadId, sessionId); + } + + async reconcile(threadId: string = this.rootSessionId, sessionId: string = this.rootSessionId): Promise { + try { + await this.sync(threadId, sessionId); + } catch (error) { + if (this.isActive()) logger.error(`Failed to list background terminals for ${threadId}`, error); + } + } + + setAppServer(appServer: CodexAppServerClient): void { + this.appServer = appServer; + this.appServerGeneration += 1; + this.appServerQueriesEnabled = true; + } + + prepareForAppServerReplacement(): void { + this.appServerGeneration += 1; + this.appServerQueriesEnabled = false; + } + + async recover(threadId: string, sessionId: string, itemIds: ReadonlySet): Promise { + if (!this.canQueryAppServer() || itemIds.size === 0) return; + const snapshot = await this.listAll(threadId); + if (snapshot === null) return; + for (const terminal of snapshot.terminals) { + if (!this.queryIsCurrent(snapshot.generation)) return; + if (!itemIds.has(terminal.itemId)) continue; + const task = this.remember(threadId, sessionId, terminal); + if (task.publication === "unpublished" && task.state === "running") await this.announce(task); + } + } + + async finishAll(state: TerminalState): Promise { + if (!this.isActive()) return; + const errors: unknown[] = []; + for (const task of this.tasks.values()) { + try { + await this.finish(task, state); + } catch (error) { + errors.push(error); + } + } + if (errors.length > 0) { + throw new AggregateError(errors, `Failed to finish ${errors.length} background terminal task(s)`); + } + } + + async sync( + threadId: string = this.rootSessionId, + sessionId: string = this.rootSessionId, + ): Promise { + if (!this.canQueryAppServer()) return; + const current = this.syncs.get(threadId); + if (current) { + current.requested = true; + current.sessionId = sessionId; + return await current.promise; + } + + const pending: PendingSync = { + requested: false, + sessionId, + promise: Promise.resolve(), + }; + pending.promise = this.syncUntilCurrent(threadId, pending).finally(() => { + if (this.syncs.get(threadId) === pending) this.syncs.delete(threadId); + }); + this.syncs.set(threadId, pending); + await pending.promise; + } + + async stop(taskId: string): Promise { + if (!this.isActive()) return false; + const task = this.tasks.get(taskId); + if (!task || task.publication !== "published") return false; + if (task.state === "stopped" && !task.terminalPublished) { + await this.publishTerminalState(task); + return true; + } + if (task.state !== "running") return false; + task.state = "stopping"; + try { + const response = await this.appServer.threadBackgroundTerminalsTerminate({ + threadId: task.threadId, + processId: task.processId, + }); + if (!response.terminated) { + if (task.state === "stopping") task.state = "running"; + return false; + } + await this.finish(task, "stopped"); + return true; + } catch (error) { + if (task.state === "stopping") task.state = "running"; + throw error; + } + } + + clear(): void { + this.disposed = true; + this.appServerGeneration += 1; + this.appServerQueriesEnabled = false; + this.tasks.clear(); + this.syncs.clear(); + } + + private async syncThread(threadId: string, sessionId: string): Promise { + const snapshot = await this.listAll(threadId); + if (snapshot === null || !this.queryIsCurrent(snapshot.generation)) return; + + const liveTaskIds = new Set(); + for (const terminal of snapshot.terminals) { + if (!this.queryIsCurrent(snapshot.generation)) return; + liveTaskIds.add(terminal.itemId); + const task = this.remember(threadId, sessionId, terminal); + if (task.publication === "unpublished" && task.state === "running") await this.announce(task); + } + + for (const task of this.tasks.values()) { + if (!this.queryIsCurrent(snapshot.generation)) return; + if (task.threadId === threadId + && task.publication === "published" + && (task.state === "running" || task.state === "stopping") + && !liveTaskIds.has(task.itemId)) { + await this.finish(task, "stopped"); + } + } + } + + private async syncUntilCurrent(threadId: string, pending: PendingSync): Promise { + let hasFailure = false; + let failure: unknown; + do { + pending.requested = false; + try { + await this.syncThread(threadId, pending.sessionId); + hasFailure = false; + } catch (error) { + hasFailure = true; + failure = error; + } + } while (pending.requested && this.canQueryAppServer()); + if (hasFailure) throw failure; + } + + private async announce(task: Task): Promise { + if (task.publication === "published") return; + if (task.announcement !== null) { + await task.announcement; + return; + } + + task.publication = "publishing"; + const announcement = this.publishAnnouncement(task); + task.announcement = announcement; + try { + await announcement; + } finally { + if (task.announcement === announcement) task.announcement = null; + } + if (isTerminalState(task.state)) await this.publishTerminalState(task); + } + + private async publishAnnouncement(task: Task): Promise { + try { + await this.publishSpawn(task); + task.publication = "published"; + } catch (error) { + task.publication = "unpublished"; + throw error; + } + } + + private async publishSpawn(task: Task): Promise { + await this.session.update({ + sessionUpdate: "tool_call_update", + toolCallId: task.itemId, + _meta: { + [JETBRAINS_META_KEY]: { + [AIR_META_KEY]: { + [AIR_ASYNC_TASKS_KEY]: { + [AIR_ASYNC_TASKS_BACKGROUNDED_KEY]: true, + }, + }, + }, + }, + }, task.sessionId); + await this.session.update({ + sessionUpdate: "async_task_spawned", + asyncTaskId: task.asyncTaskId, + name: task.name, + taskType: "shell", + showInTranscript: false, + canStop: true, + toolCallId: task.itemId, + }, task.sessionId); + } + + private remember( + threadId: string, + sessionId: string, + terminal: ThreadBackgroundTerminal, + commandTitle?: string, + ): Task { + const asyncTaskId = wireTaskId(this.rootSessionId, threadId, terminal.itemId); + const existing = this.tasks.get(asyncTaskId); + if (existing) { + existing.processId = terminal.processId; + if (commandTitle !== undefined) existing.name = commandTitle; + return existing; + } + const task: Task = { + threadId, + sessionId, + asyncTaskId, + processId: terminal.processId, + itemId: terminal.itemId, + name: commandTitle ?? terminal.command, + publication: "unpublished", + announcement: null, + terminalUpdate: null, + terminalPublished: false, + state: "running", + }; + this.tasks.set(asyncTaskId, task); + return task; + } + + private async finish(task: Task, state: TerminalState): Promise { + if (task.state === "running" || task.state === "stopping") { + task.state = state; + } else if (task.state !== state) { + return; + } + if (task.announcement !== null) await task.announcement; + if (task.publication === "published") await this.publishTerminalState(task); + } + + private async publishTerminalState(task: Task): Promise { + if (!isTerminalState(task.state) || task.terminalPublished) return; + if (task.terminalUpdate !== null) { + await task.terminalUpdate; + return; + } + const terminalUpdate = this.session.update({ + sessionUpdate: "async_task_state_update", + asyncTaskId: task.asyncTaskId, + state: task.state, + toolCallId: task.itemId, + }, task.sessionId); + task.terminalUpdate = terminalUpdate; + try { + await terminalUpdate; + task.terminalPublished = true; + } finally { + if (task.terminalUpdate === terminalUpdate) task.terminalUpdate = null; + } + } + + private async listAll(threadId: string): Promise { + const appServer = this.appServer; + const generation = this.appServerGeneration; + const terminals: ThreadBackgroundTerminal[] = []; + const seenCursors = new Set(); + let cursor: string | null = null; + do { + const response = await appServer.threadBackgroundTerminalsList({ + threadId, + cursor, + }); + if (!this.queryIsCurrent(generation)) return null; + terminals.push(...response.data); + cursor = response.nextCursor; + if (cursor !== null) { + if (seenCursors.has(cursor)) { + throw new Error("Codex returned a repeated background terminal cursor"); + } + seenCursors.add(cursor); + } + } while (cursor !== null); + return {generation, terminals}; + } + + private isActive(): boolean { + return this.enabled && !this.disposed; + } + + private canQueryAppServer(): boolean { + return this.isActive() && this.appServerQueriesEnabled; + } + + private queryIsCurrent(generation: number): boolean { + return this.canQueryAppServer() && generation === this.appServerGeneration; + } +} + +function isTerminalState(state: Task["state"]): state is TerminalState { + return state === "completed" || state === "failed" || state === "stopped"; +} + +function wireTaskId(rootSessionId: string, threadId: string, itemId: string): string { + return threadId === rootSessionId ? itemId : `${threadId}:${itemId}`; +} + +function notificationThreadId(notification: ServerNotification): string | null { + const threadId = (notification.params as {threadId?: unknown}).threadId; + return typeof threadId === "string" ? threadId : null; +} diff --git a/src/index.ts b/src/index.ts index 68df2ccd..19759300 100644 --- a/src/index.ts +++ b/src/index.ts @@ -16,6 +16,7 @@ import { GOAL_CONTROL_METHOD, LEGACY_SET_SESSION_MODEL_METHOD, SESSION_STEERING_METHOD, } from "./AcpExtensions"; +import {ASYNC_TASK_STOP_METHOD} from "./async-tasks/AsyncTaskExtension"; const emptyExtensionParamsParser = z.preprocess( (params) => params ?? {}, @@ -44,6 +45,11 @@ const goalControlParamsParser = z.discriminatedUnion("action", [ }).passthrough(), ]); +const asyncTaskStopParamsParser = z.object({ + sessionId: z.string().trim().min(1), + asyncTaskId: z.string().trim().min(1), +}).passthrough(); + if (process.argv.includes("--version")) { console.log(`${packageJson.name} ${packageJson.version}`); process.exit(0); @@ -161,6 +167,7 @@ function startAcpServer() { .onRequest("authentication/logout", emptyExtensionParamsParser, (ctx) => getAgent().extMethod("authentication/logout", ctx.params)) .onRequest(LEGACY_SET_SESSION_MODEL_METHOD, legacySetSessionModelParamsParser, (ctx) => getAgent().extMethod(LEGACY_SET_SESSION_MODEL_METHOD, ctx.params)) .onRequest(SESSION_STEERING_METHOD, sessionSteerParamsParser, (ctx) => getAgent().extMethod(SESSION_STEERING_METHOD, ctx.params)) + .onRequest(ASYNC_TASK_STOP_METHOD, asyncTaskStopParamsParser, (ctx) => getAgent().extMethod(ASYNC_TASK_STOP_METHOD, ctx.params)) .onRequest(GOAL_CONTROL_METHOD, goalControlParamsParser, (ctx) => getAgent().extMethod(GOAL_CONTROL_METHOD, ctx.params)) .connect(acpJsonStream); } diff --git a/src/subagents/AcpSubagents.ts b/src/subagents/AcpSubagents.ts index c511b095..0b56196a 100644 --- a/src/subagents/AcpSubagents.ts +++ b/src/subagents/AcpSubagents.ts @@ -1,7 +1,6 @@ import type { ClientCapabilities, SessionCapabilities, - SessionNotification, } from "@agentclientprotocol/sdk"; import { AIR_NATIVE_SUBAGENT_SESSIONS_KEY, @@ -33,15 +32,6 @@ export type SubagentStateUpdate = { _meta?: Record | null; }; -export type AcpSessionUpdate = - | SessionNotification["update"] - | SubagentSpawnedUpdate - | SubagentStateUpdate; - -export type AcpSessionNotification = Omit & { - update: AcpSessionUpdate; -}; - export type SubagentAwareSessionCapabilities = SessionCapabilities & { subagents?: Record; }; @@ -58,10 +48,3 @@ export function clientSupportsSubagents( return clientSupportsAirCapability(capabilities, AIR_NATIVE_SUBAGENT_SESSIONS_KEY); } - -/** The only cast needed until the TypeScript SDK publishes PR #1992. */ -export function asSdkSessionNotification( - notification: AcpSessionNotification, -): SessionNotification { - return notification as SessionNotification; -} diff --git a/src/subagents/CodexSubagentEventRouter.ts b/src/subagents/CodexSubagentEventRouter.ts index 6dc32640..8d232f12 100644 --- a/src/subagents/CodexSubagentEventRouter.ts +++ b/src/subagents/CodexSubagentEventRouter.ts @@ -29,6 +29,11 @@ type PendingSubagent = { droppedBufferedNotifications: number; }; +export type ClosingChildSession = { + threadId: string; + sessionId: string; +}; + /** Owns native lifecycle, child routing, waiting, and legacy activity deduplication. */ export class CodexSubagentEventRouter { private static readonly DEFAULT_WAIT_TIMEOUT_MS = 10 * 60 * 1000; @@ -180,6 +185,27 @@ export class CodexSubagentEventRouter { : this.rootSessionId; } + closingChildSessions(notification: ServerNotification): ClosingChildSession[] { + if (!this.supported) return []; + if (notification.method === "turn/completed") { + if (terminalStateFromTurn(notification.params.turn.status) === undefined) return []; + return this.closingChildSession(notification.params.threadId); + } + if (notification.method !== "item/started" && notification.method !== "item/completed") return []; + const item = notification.params.item; + if (item.type === "subAgentActivity") { + return item.kind === "interrupted" ? this.closingChildSession(item.agentThreadId) : []; + } + if (item.type !== "collabAgentToolCall") return []; + + const closing = new Map(); + for (const [threadId, state] of Object.entries(item.agentsStates)) { + if (!state || terminalStateOf(state.status) === undefined) continue; + for (const child of this.closingChildSession(threadId)) closing.set(threadId, child); + } + return [...closing.values()]; + } + takeBufferedNotifications(): ServerNotification[] { return this.replayQueue.splice(0); } @@ -272,6 +298,13 @@ export class CodexSubagentEventRouter { || this.terminalPendingSpawns.has(threadId)); } + private closingChildSession(threadId: string): ClosingChildSession[] { + const child = this.children.get(threadId); + return child && child.terminalState === undefined + ? [{threadId, sessionId: child.sessionId}] + : []; + } + private async materialize(childSessionId: string, path: string): Promise { if (this.children.has(childSessionId)) return; const pending = this.pendingSpawns.get(childSessionId);