diff --git a/README.md b/README.md index 2bda2a26..efbf06c4 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ Use [OpenAI Codex](https://github.com/openai/codex) from [Agent Client Protocol] - Model, reasoning effort, fast mode, approval, and sandbox mode configuration. - 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. -- Subagent launches as standard ACP tool calls, with Codex thread identity and activity details in namespaced `_meta.codex.subagent` metadata. +- Native ACP subagent sessions with separate child histories and root-routed permissions. - Session-scoped long-running goals through the provider-neutral [goal extension](docs/goal-extension.md). - Client-provided MCP servers over command-based stdio config and HTTP transport. - Slash commands: `/status`, `/mcp`, `/skills`, `/goal`, `/review`, `/review-branch`, `/review-commit`, `/compact`, and `/logout`, as well as configured skills. @@ -75,6 +75,16 @@ npm run bundle:all See [readme-dev.md](readme-dev.md) for local client configuration, binary packaging, and Codex type regeneration. +### Subagent sessions + +Subagents are exposed only after bilateral capability negotiation. Until the released ACP SDKs +preserve the draft `clientCapabilities.subagents` field, a supporting client may advertise +`nativeSubagentSessions` in `_meta.jetbrains.air.capabilities`; the adapter mirrors the capability +in its initialize response. The canonical field remains supported and takes precedence once it is +available. Without either client signal, subagent lifecycle retains its legacy ordinary ACP +tool-call representation, while child permission and elicitation requests are handled on the root +session. + ## License By contributing, you agree that your contributions will be licensed under the Apache 2.0 License. diff --git a/src/ACPSessionConnection.ts b/src/ACPSessionConnection.ts index 286630ac..e29a6514 100644 --- a/src/ACPSessionConnection.ts +++ b/src/ACPSessionConnection.ts @@ -1,5 +1,8 @@ import * as acp from "@agentclientprotocol/sdk"; -import type {SessionNotification} from "@agentclientprotocol/sdk"; +import { + type AcpSessionUpdate, + asSdkSessionNotification, +} from "./subagents/AcpSubagents"; export type AcpClientConnection = Pick; @@ -12,12 +15,12 @@ export class ACPSessionConnection { this.sessionId = sessionId; } - async update(update: UpdateSessionEvent) { - await this.connection.notify(acp.methods.client.session.update, { - sessionId: this.sessionId, + async update(update: UpdateSessionEvent, sessionId: string = this.sessionId) { + await this.connection.notify(acp.methods.client.session.update, asSdkSessionNotification({ + sessionId, update: update - }); + })); } } -export type UpdateSessionEvent = SessionNotification["update"]; +export type UpdateSessionEvent = AcpSessionUpdate; diff --git a/src/AirExtension.ts b/src/AirExtension.ts index cf97b512..9e39a18d 100644 --- a/src/AirExtension.ts +++ b/src/AirExtension.ts @@ -12,5 +12,22 @@ export const AIR_EXTENSION_VERSION_KEY = "version"; 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_AGENT_FILE_CHANGE_REPORT_REQUEST_KEY = "agentFileChangeReportRequest"; export const AIR_EXTENSION_VERSION = 1; + +export function clientSupportsAirCapability( + capabilities: ClientCapabilities | null | undefined, + capability: string, +): boolean { + const jetbrains = capabilities?._meta?.[JETBRAINS_META_KEY] as Record | undefined; + const air = jetbrains?.[AIR_META_KEY] as Record | undefined; + const version = air?.[AIR_EXTENSION_VERSION_KEY]; + const supported = air?.[AIR_EXTENSION_CAPABILITIES_KEY]; + return typeof version === "number" + && Number.isInteger(version) + && version >= AIR_EXTENSION_VERSION + && Array.isArray(supported) + && supported.includes(capability); +} +import type {ClientCapabilities} from "@agentclientprotocol/sdk"; diff --git a/src/CodexAcpClient.ts b/src/CodexAcpClient.ts index bec75265..92892c3f 100644 --- a/src/CodexAcpClient.ts +++ b/src/CodexAcpClient.ts @@ -63,6 +63,7 @@ import { createReportedAgentFileChangeReport, createUnavailableAgentFileChangeReport, } from "./AgentFileChangeReport"; +import {CodexSubagentSubscriptions} from "./subagents/CodexSubagentSubscriptions"; /** * Well-known provider id for the client-configurable custom LLM gateway. @@ -109,6 +110,7 @@ export class CodexAcpClient { private pendingLoginCompleted: Promise | null = null; private pendingAccountUpdated: Promise | null = null; private readonly sessionNotificationQueues = new Map>(); + private readonly subagents: CodexSubagentSubscriptions; private skillExtraRoots: string[] = []; private configPath: string | null = null; @@ -118,6 +120,7 @@ export class CodexAcpClient { this.config = codexConfig ?? {}; this.modelProvider = modelProvider ?? null; this.gatewayConfig = null; + this.subagents = new CodexSubagentSubscriptions(codexClient); } private readonly defaultClientInfo: ClientInfo = { @@ -544,6 +547,7 @@ export class CodexAcpClient { await this.codexClient.threadUnsubscribe({threadId: sessionId}); } finally { this.codexClient.clearThreadHandlers(sessionId); + this.subagents.clear(sessionId); } } @@ -782,34 +786,19 @@ export class CodexAcpClient { sessionId: string, eventHandler: (result: ServerNotification) => void | Promise, approvalHandler: ApprovalHandler, - elicitationHandler: ElicitationHandler + elicitationHandler: ElicitationHandler, + supportsSubagents: boolean, ) { - this.codexClient.onServerNotification(sessionId, (event) => { + const dispatch = (event: ServerNotification) => { this.enqueueSessionNotification(sessionId, () => eventHandler(event)); - }); - this.codexClient.onApprovalRequest(sessionId, { - handleCommandExecution: async (params) => { - await this.waitForSessionNotifications(sessionId); - return await approvalHandler.handleCommandExecution(params); - }, - handleFileChange: async (params) => { - await this.waitForSessionNotifications(sessionId); - return await approvalHandler.handleFileChange(params); - }, - handlePermissionsRequest: async (params) => { - await this.waitForSessionNotifications(sessionId); - return await approvalHandler.handlePermissionsRequest(params); - }, - }); - this.codexClient.onElicitationRequest(sessionId, { - handleElicitation: async (params) => { - await this.waitForSessionNotifications(sessionId); - return await elicitationHandler.handleElicitation(params); - }, - handleUserInput: async (params) => { - await this.waitForSessionNotifications(sessionId); - return await elicitationHandler.handleUserInput(params); - }, + }; + this.subagents.subscribe({ + rootSessionId: sessionId, + supportsSubagents, + dispatch, + approvalHandler, + elicitationHandler, + waitForRootNotifications: () => this.waitForSessionNotifications(sessionId), }); } diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index f3ebb373..55be7be7 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -100,15 +100,22 @@ import { createUserMessageChunk, } from "./ContentChunks"; import {sameThreadGoalSnapshot, type ThreadGoalSnapshot, toThreadGoalSnapshot,} from "./ThreadGoalSnapshot"; +import { + clientSupportsSubagents, + type SubagentAwareSessionCapabilities, +} from "./subagents/AcpSubagents"; +import {CodexSubagentEventRouter} from "./subagents/CodexSubagentEventRouter"; import {randomUUID} from "node:crypto"; import {once} from "node:events"; import { AIR_AGENT_FILE_CHANGE_REPORT_KEY, + AIR_NATIVE_SUBAGENT_SESSIONS_KEY, AIR_EXTENSION_CAPABILITIES_KEY, AIR_EXTENSION_VERSION, AIR_EXTENSION_VERSION_KEY, AIR_META_KEY, AIR_SESSION_FAILURE_KEY, + clientSupportsAirCapability, JETBRAINS_META_KEY, } from "./AirExtension"; import { @@ -148,6 +155,7 @@ export interface SessionState { sessionTitle: string | null; sessionTitleSource: "unset" | "fallback" | "explicit" | "unknown"; sessionFailure?: SessionFailure; + subagents: CodexSubagentEventRouter; } export type SessionFailureCategory = @@ -173,21 +181,6 @@ export interface SessionFailure { const CODEX_PROCESS_EXITED_ERROR_CODE = 1001; -function clientSupportsAirCapability( - capabilities: acp.ClientCapabilities | null, - capability: string, -): boolean { - const jetbrains = capabilities?._meta?.[JETBRAINS_META_KEY] as Record | undefined; - const air = jetbrains?.[AIR_META_KEY] as Record | undefined; - const version = air?.[AIR_EXTENSION_VERSION_KEY]; - const supported = air?.[AIR_EXTENSION_CAPABILITIES_KEY]; - return typeof version === "number" - && Number.isInteger(version) - && version >= AIR_EXTENSION_VERSION - && Array.isArray(supported) - && supported.includes(capability); -} - function clientSupportsTypedSessionFailures(capabilities: acp.ClientCapabilities | null): boolean { return clientSupportsAirCapability(capabilities, AIR_SESSION_FAILURE_KEY); } @@ -316,6 +309,14 @@ export class CodexAcpServer { this.terminalOutputMode = resolveTerminalOutputMode(_params.clientCapabilities); this.booleanConfigOptionsSupported = clientSupportsBooleanConfigOptions(_params.clientCapabilities); await this.runWithProcessCheck(() => this.codexAcpClient.initialize(_params)); + const sessionCapabilities: SubagentAwareSessionCapabilities = { + resume: { }, + list: { }, + close: { }, + delete: { }, + additionalDirectories: {}, + ...(clientSupportsSubagents(_params.clientCapabilities) ? {subagents: {}} : {}), + }; return { protocolVersion: acp.PROTOCOL_VERSION, agentInfo: { @@ -333,13 +334,7 @@ export class CodexAcpServer { embeddedContext: true, image: true }, - sessionCapabilities: { - resume: { }, - list: { }, - close: { }, - delete: { }, - additionalDirectories: {}, - }, + sessionCapabilities, mcpCapabilities: { acp: false, http: true, @@ -362,6 +357,7 @@ export class CodexAcpServer { [AIR_EXTENSION_CAPABILITIES_KEY]: [ AIR_SESSION_FAILURE_KEY, AIR_AGENT_FILE_CHANGE_REPORT_KEY, + AIR_NATIVE_SUBAGENT_SESSIONS_KEY, ], }, }, @@ -632,6 +628,11 @@ export class CodexAcpServer { goalRevision: 0, sessionTitle: null, sessionTitleSource: "sessionId" in request ? "unknown" : "unset", + subagents: new CodexSubagentEventRouter( + sessionId, + clientSupportsSubagents(this.clientCapabilities), + new ACPSessionConnection(this.connection, sessionId), + ), }; this.sessions.set(sessionId, sessionState); resumeSubscribed = false; @@ -1649,6 +1650,11 @@ export class CodexAcpServer { goalRevision: 0, sessionTitle: null, sessionTitleSource: "unset", + subagents: new CodexSubagentEventRouter( + sessionId, + clientSupportsSubagents(this.clientCapabilities), + new ACPSessionConnection(this.connection, sessionId), + ), }; this.sessions.set(sessionId, sessionState); subscribed = false; @@ -2273,6 +2279,7 @@ export class CodexAcpServer { : null; let agentFileChangeReportTurnId: string | null = null; let agentFileChangeReportUnavailableReason: AgentFileChangeReportUnavailableReason = "providerError"; + let promptWasCancelled = false; let recoverableSessionFailure = sessionState.sessionFailure; sessionState.currentTurnId = null; sessionState.lastTokenUsage = null; @@ -2299,6 +2306,7 @@ export class CodexAcpServer { } }; const cancelledPromptResponse = (): acp.PromptResponse => { + promptWasCancelled = true; agentFileChangeReportTurnId = null; agentFileChangeReportUnavailableReason = "cancelled"; return this.cancelledPromptResponse(sessionState); @@ -2311,33 +2319,33 @@ export class CodexAcpServer { clientSupportsPlanUpdates(this.clientCapabilities), clientSupportsTypedSessionFailures(this.clientCapabilities), this.sessionFailureEpoch, + sessionState.subagents, ); eventHandler = promptEventHandler; const permissionLifecycle = this.permissionLifecycleContext(sessionState); const permissionContext = permissionLifecycle.beginPrompt(); const approvalHandler = new CodexApprovalHandler( this.connection, - sessionState, permissionContext, activePrompt.signal, ); const elicitationHandler = new CodexElicitationHandler( this.connection, - sessionState, permissionContext, this.clientCapabilities, activePrompt.signal, ); await this.codexAcpClient.subscribeToSessionEvents(params.sessionId, async (event) => { + permissionContext.handleNotification(event); + await elicitationHandler.handleNotification(event); if (!promptNotificationsActive) { await promptEventHandler.handleSessionScopedNotification(event); return; } const completesActiveTurn = event.method === "turn/completed" + && event.params.threadId === sessionState.sessionId && event.params.turn.id === sessionState.currentTurnId; - permissionContext.handleNotification(event); - await elicitationHandler.handleNotification(event); await promptEventHandler.handleNotification(event); if (completesActiveTurn) { // The prompt may remain open for plan approval after its turn has ended. Switch at @@ -2346,7 +2354,8 @@ export class CodexAcpServer { } }, approvalHandler, - elicitationHandler); + elicitationHandler, + clientSupportsSubagents(this.clientCapabilities)); if (activePrompt.signal.aborted) { return cancelledPromptResponse(); @@ -2497,6 +2506,8 @@ export class CodexAcpServer { return cancelledPromptResponse(); } + await this.codexAcpClient.waitForSessionNotifications(params.sessionId); + await eventHandler.waitForNativeSubagents(activePrompt.signal); await this.codexAcpClient.waitForSessionNotifications(params.sessionId); await eventHandler.flushPendingErrors(); await eventHandler.handleFailedTurn(turnCompleted.turn); @@ -2590,6 +2601,8 @@ export class CodexAcpServer { return cancelledPromptResponse(); } + await this.codexAcpClient.waitForSessionNotifications(params.sessionId); + await eventHandler.waitForNativeSubagents(activePrompt.signal); await this.codexAcpClient.waitForSessionNotifications(params.sessionId); await eventHandler.flushPendingErrors(); await eventHandler.handleFailedTurn(turnCompleted.turn); @@ -2661,6 +2674,15 @@ export class CodexAcpServer { // The app-server subscription is session-scoped and outlives this prompt. Flip routing before // awaiting disposal so queued late notifications cannot enter prompt-local buffers. promptNotificationsActive = false; + try { + await eventHandler?.finishOutstandingNativeSubagents( + promptWasCancelled || activePrompt.signal.aborted || this.sessionIsClosing(params.sessionId) + ? "cancelled" + : "failed", + ); + } catch (error) { + logger.error("Failed to publish terminal subagent state during prompt cleanup", error); + } if (agentFileChangeReportRequest !== null) { await this.publishAgentFileChangeReport( sessionState, diff --git a/src/CodexElicitationHandler.ts b/src/CodexElicitationHandler.ts index f9e6c6c9..f7f24f39 100644 --- a/src/CodexElicitationHandler.ts +++ b/src/CodexElicitationHandler.ts @@ -1,5 +1,4 @@ import * as acp from "@agentclientprotocol/sdk"; -import type { SessionState } from "./CodexAcpServer"; import type { ElicitationHandler } from "./CodexAppServerClient"; import type { ServerNotification } from "./app-server"; import type {JsonValue} from "./app-server/serde_json/JsonValue"; @@ -139,7 +138,6 @@ function userInputResponseValue( export class CodexElicitationHandler implements ElicitationHandler { private readonly connection: AcpClientConnection; - private readonly sessionState: SessionState; private readonly permissionContext: PermissionPromptContext; private readonly clientCapabilities: acp.ClientCapabilities | null; private readonly cancellationSignal: AbortSignal | undefined; @@ -161,13 +159,11 @@ export class CodexElicitationHandler implements ElicitationHandler { constructor( connection: AcpClientConnection, - sessionState: SessionState, permissionContext: PermissionPromptContext, clientCapabilities: acp.ClientCapabilities | null = null, cancellationSignal?: AbortSignal ) { this.connection = connection; - this.sessionState = sessionState; this.permissionContext = permissionContext; this.clientCapabilities = clientCapabilities; this.cancellationSignal = cancellationSignal; @@ -198,7 +194,7 @@ export class CodexElicitationHandler implements ElicitationHandler { if (params.mode === "url" && result.action === "accept") { this.trackUrlElicitation(params.threadId, params.elicitationId); } - await this.publishAcceptedMcpToolApproval(context, result.action === "accept"); + await this.publishAcceptedMcpToolApproval(params.threadId, context, result.action === "accept"); return result; } if (!this.canUsePermissionFallback(params)) { @@ -206,7 +202,7 @@ export class CodexElicitationHandler implements ElicitationHandler { } const {request, correlatedCallId} = buildMcpPermissionRequest( - this.sessionState.sessionId, + params.threadId, params, context, () => this.permissionContext.nextStandaloneMcpToolCallId(params.serverName), @@ -223,7 +219,7 @@ export class CodexElicitationHandler implements ElicitationHandler { ); if (correlatedCallId !== undefined && result.action === "accept") { await this.connection.notify(acp.methods.client.session.update, { - sessionId: this.sessionState.sessionId, + sessionId: params.threadId, update: { sessionUpdate: "tool_call_update", toolCallId: correlatedCallId, status: "in_progress" }, }); } @@ -351,7 +347,7 @@ export class CodexElicitationHandler implements ElicitationHandler { context: McpElicitationContext ): acp.CreateElicitationRequest { const base = { - sessionId: this.sessionState.sessionId, + sessionId: params.threadId, ...(context.correlatedCallId ? { toolCallId: context.correlatedCallId } : {}), message: params.message, _meta: recordOrNull(params._meta), @@ -429,7 +425,7 @@ export class CodexElicitationHandler implements ElicitationHandler { const firstQuestion = params.questions[0]; return { - sessionId: this.sessionState.sessionId, + sessionId: params.threadId, toolCallId: params.itemId, mode: "form", message: params.questions.length === 1 && firstQuestion @@ -525,6 +521,7 @@ export class CodexElicitationHandler implements ElicitationHandler { } private async publishAcceptedMcpToolApproval( + sessionId: string, context: McpElicitationContext, accepted: boolean ): Promise { @@ -532,7 +529,7 @@ export class CodexElicitationHandler implements ElicitationHandler { return; } await this.connection.notify(acp.methods.client.session.update, { - sessionId: this.sessionState.sessionId, + sessionId, update: { sessionUpdate: "tool_call_update", toolCallId: context.correlatedCallId, status: "in_progress" }, }); } diff --git a/src/CodexEventHandler.ts b/src/CodexEventHandler.ts index 7360b6f9..62499587 100644 --- a/src/CodexEventHandler.ts +++ b/src/CodexEventHandler.ts @@ -41,8 +41,6 @@ import type { McpStartupCompleteEvent } from "./app-server/McpStartupCompleteEve import {toTokenCount} from "./TokenCount"; import { commandExecutionUsesTerminalOutput, - createCollabAgentToolCallCompleteUpdate, - createCollabAgentToolCallUpdate, createCommandExecutionUpdate, createContextCompactionCompleteUpdate, createContextCompactionStartUpdate, @@ -59,7 +57,6 @@ import { createFuzzyFileSearchComplete, createFuzzyFileSearchStartOrUpdate, createMcpToolCallUpdate, - createSubAgentActivityUpdate, createWebSearchCompleteUpdate, createWebSearchStartUpdate, fuzzyFileSearchToolCallId, @@ -81,6 +78,8 @@ import { AIR_SESSION_FAILURE_KEY, JETBRAINS_META_KEY, } from "./AirExtension"; +import {CodexSubagentEventRouter} from "./subagents/CodexSubagentEventRouter"; +import type {SubagentState} from "./subagents/AcpSubagents"; export { stripShellPrefix }; @@ -224,7 +223,7 @@ export class CodexEventHandler { private readonly terminalCommandIds = new Set(); private readonly terminalCommandOutputIds = new Set(); private readonly agentMessagePhases = new Map(); - private readonly activeSubAgentActivities = new Set(); + private readonly subagents: CodexSubagentEventRouter; constructor( connection: AcpClientConnection, @@ -232,12 +231,18 @@ export class CodexEventHandler { supportsPlanUpdates = false, supportsTypedSessionFailures = false, sessionFailureEpoch: string = randomUUID(), + subagents: CodexSubagentEventRouter = new CodexSubagentEventRouter( + sessionState.sessionId, + false, + new ACPSessionConnection(connection, sessionState.sessionId), + ), ) { this.sessionState = sessionState; this.supportsPlanUpdates = supportsPlanUpdates; this.supportsTypedSessionFailures = supportsTypedSessionFailures; this.sessionFailureEpoch = sessionFailureEpoch; this.session = new ACPSessionConnection(connection, sessionState.sessionId); + this.subagents = subagents; if (sessionState.sessionFailure !== undefined) { this.failuresById.set(sessionState.sessionFailure.id, sessionState.sessionFailure); } @@ -361,12 +366,26 @@ export class CodexEventHandler { async handleNotification(notification: ServerNotification) { await this.flushPendingErrors(); + if (await this.subagents.handle(notification)) { + return; + } + if (this.subagents.shouldIgnore(notification)) { + return; + } const updateEvent = await this.createUpdateEvent(notification); if (updateEvent) { - await this.session.update(updateEvent); + await this.session.update(updateEvent, this.subagents.notificationSessionId(notification)); } } + async waitForNativeSubagents(signal: AbortSignal): Promise { + await this.subagents.wait(signal); + } + + async finishOutstandingNativeSubagents(state: SubagentState): Promise { + await this.subagents.finishOutstanding(state); + } + async flushPendingPlanUpdates(): Promise { this.cancelPlanUpdateTimer(); do { @@ -681,15 +700,14 @@ export class CodexEventHandler { this.activeImageGenerationItems.add(event.item.id); return createImageGenerationStartUpdate(event.item); case "collabAgentToolCall": - return createCollabAgentToolCallUpdate(event.item); + return this.subagents.legacyCollaborationStarted(event.item); case "agentMessage": this.rememberAgentMessagePhase(event.item); return null; case "contextCompaction": return createContextCompactionStartUpdate(event.item); case "subAgentActivity": - this.activeSubAgentActivities.add(event.item.id); - return createSubAgentActivityUpdate(event.item, "in_progress", "tool_call"); + return this.subagents.legacyActivityStarted(event.item); case "sleep": case "userMessage": case "hookPrompt": @@ -738,7 +756,7 @@ export class CodexEventHandler { case "webSearch": return createWebSearchCompleteUpdate(event.item); case "collabAgentToolCall": - return createCollabAgentToolCallCompleteUpdate(event.item); + return this.subagents.legacyCollaborationCompleted(event.item); case "agentMessage": this.rememberAgentMessagePhase(event.item); return null; @@ -751,12 +769,8 @@ export class CodexEventHandler { case "contextCompaction": return createContextCompactionCompleteUpdate(event.item); //ignored types - case "subAgentActivity": { - const sessionUpdate = this.activeSubAgentActivities.delete(event.item.id) - ? "tool_call_update" - : "tool_call"; - return createSubAgentActivityUpdate(event.item, "completed", sessionUpdate); - } + case "subAgentActivity": + return this.subagents.legacyActivityCompleted(event.item); case "sleep": case "userMessage": case "hookPrompt": diff --git a/src/__tests__/CodexACPAgent/collab-agent-events.test.ts b/src/__tests__/CodexACPAgent/collab-agent-events.test.ts index 5ad2c72a..b8776381 100644 --- a/src/__tests__/CodexACPAgent/collab-agent-events.test.ts +++ b/src/__tests__/CodexACPAgent/collab-agent-events.test.ts @@ -2,6 +2,8 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import type { ServerNotification } from "../../app-server"; import type { SessionState } from "../../CodexAcpServer"; import { AgentMode } from "../../AgentMode"; +import {ACPSessionConnection} from "../../ACPSessionConnection"; +import {CodexSubagentEventRouter} from "../../subagents/CodexSubagentEventRouter"; import { createCodexMockTestFixture, createTestSessionState, @@ -11,20 +13,44 @@ import { describe("CodexEventHandler - collab agent tool call events", () => { let mockFixture: CodexMockTestFixture; + let sessionState: SessionState; const sessionId = "test-session-id"; beforeEach(() => { mockFixture = createCodexMockTestFixture(); + sessionState = createTestSessionState({ + sessionId, + currentModelId: "model-id[effort]", + agentMode: AgentMode.DEFAULT_AGENT_MODE, + }); vi.clearAllMocks(); }); - const sessionState: SessionState = createTestSessionState({ - sessionId, - currentModelId: "model-id[effort]", - agentMode: AgentMode.DEFAULT_AGENT_MODE, - }); + async function initializeNativeSubagents() { + const response = await mockFixture.getCodexAcpAgent().initialize({ + protocolVersion: 1, + clientCapabilities: { + elicitation: {url: {}}, + _meta: { + jetbrains: { + air: {version: 1, capabilities: ["nativeSubagentSessions"]}, + }, + }, + }, + }); + sessionState.subagents = new CodexSubagentEventRouter( + sessionId, + true, + new ACPSessionConnection(mockFixture.getAcpConnection(), sessionId), + ); + return response; + } - it("maps live collab agent tool calls to ACP tool call updates", async () => { + it("keeps the legacy tool-call lifecycle without subagent capability and root-routes permissions", async () => { + await mockFixture.getCodexAcpAgent().initialize({ + protocolVersion: 1, + clientCapabilities: {elicitation: {form: {}}}, + }); const notifications: ServerNotification[] = [ { method: "item/started", @@ -76,16 +102,62 @@ describe("CodexEventHandler - collab agent tool call events", () => { }, }, }, + { + method: "item/agentMessage/delta", + params: { + threadId: sessionId, + turnId: "turn-1", + itemId: "parent-message", + delta: "Visible parent output", + }, + }, ]; await setupPromptAndSendNotifications(mockFixture, sessionId, sessionState, notifications); - await expect(`${mockFixture.getAcpConnectionDump([])}\n`).toMatchFileSnapshot( - "data/collab-agent-tool-call-flow.json" - ); + const collaborationUpdates = mockFixture.getAcpConnectionEvents([]) + .filter(event => event.method === "sessionUpdate") + .map(event => event.args[0].update) + .filter(update => update.toolCallId === "call-spawn-weather"); + expect(collaborationUpdates).toMatchObject([ + {sessionUpdate: "tool_call", title: "spawnAgent", status: "in_progress"}, + {sessionUpdate: "tool_call_update", title: "spawnAgent", status: "completed"}, + ]); + + mockFixture.setPermissionResponse({outcome: {outcome: "selected", optionId: "allow_once"}}); + await mockFixture.sendServerRequest("item/commandExecution/requestApproval", { + threadId: "thread-paris", + turnId: "turn-child", + itemId: "child-command", + reason: "Check the weather service", + startedAtMs: 0, + environmentId: null, + proposedExecpolicyAmendment: null, + }); + const permissionRequest = mockFixture.getAcpConnectionEvents([]) + .find(event => event.method === "requestPermission" && event.args[0].toolCall.toolCallId === "child-command"); + expect(permissionRequest?.args[0].sessionId).toBe(sessionId); + + mockFixture.setElicitationResponse({action: "accept", content: {answer: "yes"}}); + await mockFixture.sendServerRequest("mcpServer/elicitation/request", { + threadId: "thread-paris", + turnId: "turn-child", + serverName: "child-server", + mode: "form", + _meta: null, + message: "Continue?", + requestedSchema: { + type: "object", + properties: {answer: {type: "string"}}, + required: ["answer"], + }, + }); + const elicitationRequest = mockFixture.getAcpConnectionEvents([]) + .find(event => event.method === "createElicitation" && event.args[0].message === "Continue?"); + expect(elicitationRequest?.args[0].sessionId).toBe(sessionId); }); - it("maps live subagent activity to an ACP tool call", async () => { + it("keeps legacy subagent activity as a tool call without subagent capability", async () => { const notifications: ServerNotification[] = [ { method: "item/completed", @@ -102,12 +174,988 @@ describe("CodexEventHandler - collab agent tool call events", () => { }, }, }, + { + method: "item/agentMessage/delta", + params: { + threadId: sessionId, + turnId: "turn-1", + itemId: "parent-message", + delta: "Visible parent output", + }, + }, ]; await setupPromptAndSendNotifications(mockFixture, sessionId, sessionState, notifications); - await expect(`${mockFixture.getAcpConnectionDump([])}\n`).toMatchFileSnapshot( - "data/subagent-activity-flow.json" - ); + const activity = mockFixture.getAcpConnectionEvents([]) + .filter(event => event.method === "sessionUpdate") + .map(event => event.args[0].update) + .find(update => update.toolCallId === "call-spawn-weather"); + expect(activity).toMatchObject({ + sessionUpdate: "tool_call", + title: "Start subagent weather_research", + status: "completed", + }); + }); + + it("promotes subagent activity to native lifecycle when collaboration items are absent", async () => { + await initializeNativeSubagents(); + await setupPromptAndSendNotifications(mockFixture, sessionId, sessionState, [ + { + method: "item/started", + params: { + threadId: sessionId, + turnId: "turn-1", + startedAtMs: 0, + item: { + type: "subAgentActivity", + id: "activity-started", + kind: "started", + agentThreadId: "child-1", + agentPath: "/root/air_architecture", + }, + }, + }, + { + method: "item/completed", + params: { + threadId: sessionId, + turnId: "turn-1", + completedAtMs: 0, + item: { + type: "subAgentActivity", + id: "activity-started", + kind: "started", + agentThreadId: "child-1", + agentPath: "/root/air_architecture", + }, + }, + }, + { + method: "item/started", + params: { + threadId: "child-1", + turnId: "turn-child", + startedAtMs: 0, + item: { + type: "subAgentActivity", + id: "nested-started", + kind: "started", + agentThreadId: "grandchild-1", + agentPath: "/root/air_architecture/tests", + }, + }, + }, + { + method: "item/agentMessage/delta", + params: { + threadId: "grandchild-1", + turnId: "turn-grandchild", + itemId: "grandchild-message", + delta: "Nested result", + }, + }, + { + method: "item/completed", + params: { + threadId: "child-1", + turnId: "turn-child", + completedAtMs: 0, + item: { + type: "subAgentActivity", + id: "nested-interrupted", + kind: "interrupted", + agentThreadId: "grandchild-1", + agentPath: "/root/air_architecture/tests", + }, + }, + }, + { + method: "turn/completed", + params: { + threadId: "child-1", + turn: { + id: "turn-child", + items: [], + itemsView: "notLoaded", + status: "completed", + error: null, + startedAt: null, + completedAt: null, + durationMs: null, + }, + }, + }, + { + method: "turn/completed", + params: { + threadId: sessionId, + turn: { + id: "turn-1", + items: [], + itemsView: "notLoaded", + status: "completed", + error: null, + startedAt: null, + completedAt: null, + durationMs: null, + }, + }, + }, + ]); + + const updates = mockFixture.getAcpConnectionEvents([]) + .filter(event => event.method === "sessionUpdate") + .map(event => event.args[0]); + expect(updates).toEqual([ + { + sessionId, + update: { + sessionUpdate: "subagent_spawned", + subagentSessionId: "child-1", + name: "Air architecture", + description: "Delegated task for Air architecture", + capabilities: {}, + }, + }, + { + sessionId: "child-1", + update: { + sessionUpdate: "subagent_spawned", + subagentSessionId: "grandchild-1", + name: "Tests", + description: "Delegated task for Tests", + capabilities: {}, + }, + }, + { + sessionId: "grandchild-1", + update: { + sessionUpdate: "agent_message_chunk", + content: {type: "text", text: "Nested result"}, + messageId: "grandchild-message", + }, + }, + { + sessionId: "child-1", + update: { + sessionUpdate: "subagent_state_update", + subagentSessionId: "grandchild-1", + state: "cancelled", + }, + }, + { + sessionId, + update: { + sessionUpdate: "subagent_state_update", + subagentSessionId: "child-1", + state: "completed", + }, + }, + ]); + }); + + it("does not represent the root activity as a subagent", async () => { + await initializeNativeSubagents(); + await setupPromptAndSendNotifications(mockFixture, sessionId, sessionState, [ + { + method: "item/started", + params: { + threadId: sessionId, + turnId: "turn-1", + startedAtMs: 0, + item: { + type: "subAgentActivity", + id: "root-activity", + kind: "started", + agentThreadId: "root-activity-thread", + agentPath: "/root", + }, + }, + }, + { + method: "item/completed", + params: { + threadId: sessionId, + turnId: "turn-1", + completedAtMs: 0, + item: { + type: "subAgentActivity", + id: "root-activity", + kind: "started", + agentThreadId: "root-activity-thread", + agentPath: "/root/", + }, + }, + }, + { + method: "item/agentMessage/delta", + params: { + threadId: sessionId, + turnId: "turn-1", + itemId: "parent-message", + delta: "Visible root output", + }, + }, + ]); + + expect(mockFixture.getAcpConnectionEvents([]) + .filter(event => event.method === "sessionUpdate") + .map(event => event.args[0].update)) + .not.toContainEqual(expect.objectContaining({sessionUpdate: "subagent_spawned"})); + }); + + it("emits native lifecycle and routes child output after capability negotiation", async () => { + const initializeResponse = await initializeNativeSubagents(); + expect( + (initializeResponse.agentCapabilities?.sessionCapabilities as {subagents?: unknown}).subagents + ).toEqual({}); + const notifications: ServerNotification[] = [ + { + method: "item/started", + params: { + threadId: sessionId, + turnId: "turn-1", + startedAtMs: 0, + item: { + type: "collabAgentToolCall", + id: "call-spawn-weather", + tool: "spawnAgent", + status: "inProgress", + senderThreadId: "thread-main", + receiverThreadIds: ["thread-paris"], + prompt: "Find the current weather in Paris.", + model: null, + reasoningEffort: null, + agentsStates: { + "thread-paris": {status: "running", message: "Checking weather"}, + }, + }, + }, + }, + { + method: "item/started", + params: { + threadId: sessionId, + turnId: "turn-1", + startedAtMs: 0, + item: { + type: "subAgentActivity", + id: "activity-weather", + kind: "started", + agentThreadId: "thread-paris", + agentPath: "/root/weather_research", + }, + }, + }, + { + method: "item/agentMessage/delta", + params: { + threadId: "thread-paris", + turnId: "turn-child", + itemId: "child-message", + delta: "Weather found", + }, + }, + { + method: "item/completed", + params: { + threadId: sessionId, + turnId: "turn-1", + completedAtMs: 0, + item: { + type: "collabAgentToolCall", + id: "call-spawn-weather", + tool: "spawnAgent", + status: "completed", + senderThreadId: "thread-main", + receiverThreadIds: ["thread-paris"], + prompt: "Find the current weather in Paris.", + model: null, + reasoningEffort: null, + agentsStates: { + "thread-paris": {status: "completed", message: null}, + }, + }, + }, + }, + ]; + + await setupPromptAndSendNotifications(mockFixture, sessionId, sessionState, notifications); + + const updates = mockFixture.getAcpConnectionEvents([]) + .filter(event => event.method === "sessionUpdate") + .map(event => event.args[0]); + expect(updates).toEqual([ + { + sessionId, + update: { + sessionUpdate: "subagent_spawned", + subagentSessionId: "thread-paris", + name: "Weather research", + description: "Find the current weather in Paris.", + capabilities: {}, + }, + }, + { + sessionId: "thread-paris", + update: { + sessionUpdate: "agent_message_chunk", + content: {type: "text", text: "Weather found"}, + messageId: "child-message", + }, + }, + { + sessionId, + update: { + sessionUpdate: "subagent_state_update", + subagentSessionId: "thread-paris", + state: "completed", + }, + }, + ]); + + mockFixture.setPermissionResponse({outcome: {outcome: "selected", optionId: "allow_once"}}); + await mockFixture.sendServerRequest("item/commandExecution/requestApproval", { + threadId: "thread-paris", + turnId: "turn-child", + itemId: "child-command", + reason: "Check the weather service", + startedAtMs: 0, + environmentId: null, + proposedExecpolicyAmendment: null, + }); + const permissionRequest = mockFixture.getAcpConnectionEvents([]) + .find(event => event.method === "requestPermission" && event.args[0].toolCall.toolCallId === "child-command"); + expect(permissionRequest?.args[0].sessionId).toBe("thread-paris"); + }); + + it("routes nested agents through their immediate parent sessions", async () => { + await initializeNativeSubagents(); + const collabItem = ( + threadId: string, + senderThreadId: string, + receiverThreadId: string, + id: string, + status: "running" | "completed", + ): ServerNotification => ({ + method: status === "running" ? "item/started" : "item/completed", + params: { + threadId, + turnId: `turn-${threadId}`, + ...(status === "running" ? {startedAtMs: 0} : {completedAtMs: 0}), + item: { + type: "collabAgentToolCall", + id, + tool: "spawnAgent", + status: status === "running" ? "inProgress" : "completed", + senderThreadId, + receiverThreadIds: [receiverThreadId], + prompt: `Task for ${receiverThreadId}`, + model: null, + reasoningEffort: null, + agentsStates: {[receiverThreadId]: {status, message: null}}, + }, + }, + } as ServerNotification); + await setupPromptAndSendNotifications(mockFixture, sessionId, sessionState, [ + collabItem(sessionId, sessionId, "child-1", "spawn-1", "running"), + { + method: "item/started", + params: { + threadId: sessionId, + turnId: "turn-root", + startedAtMs: 0, + item: { + type: "subAgentActivity", + id: "activity-child", + kind: "started", + agentThreadId: "child-1", + agentPath: "/root/researcher", + }, + }, + }, + collabItem("child-1", "child-1", "grandchild-1", "spawn-2", "running"), + { + method: "item/started", + params: { + threadId: "child-1", + turnId: "turn-child-1", + startedAtMs: 0, + item: { + type: "subAgentActivity", + id: "activity-grandchild", + kind: "started", + agentThreadId: "grandchild-1", + agentPath: "/root/researcher/tester", + }, + }, + }, + { + method: "item/agentMessage/delta", + params: { + threadId: "grandchild-1", + turnId: "turn-grandchild", + itemId: "grandchild-message", + delta: "Nested result", + }, + }, + collabItem("child-1", "child-1", "grandchild-1", "spawn-2", "completed"), + collabItem(sessionId, sessionId, "child-1", "spawn-1", "completed"), + ]); + + const updates = mockFixture.getAcpConnectionEvents([]) + .filter(event => event.method === "sessionUpdate") + .map(event => event.args[0]); + expect(updates.map(({sessionId: target, update}) => [target, update.sessionUpdate])).toEqual([ + [sessionId, "subagent_spawned"], + ["child-1", "subagent_spawned"], + ["grandchild-1", "agent_message_chunk"], + ["child-1", "subagent_state_update"], + [sessionId, "subagent_state_update"], + ]); + }); + + it("deduplicates lifecycle, rejects blank IDs, and ignores late child output", async () => { + await initializeNativeSubagents(); + const spawn = (method: "item/started" | "item/completed"): ServerNotification => ({ + method, + params: { + threadId: sessionId, + turnId: "turn-1", + ...(method === "item/started" ? {startedAtMs: 0} : {completedAtMs: 0}), + item: { + type: "collabAgentToolCall", + id: "spawn", + tool: "spawnAgent", + status: method === "item/started" ? "inProgress" : "completed", + senderThreadId: sessionId, + receiverThreadIds: ["", "child-1", "child-1"], + prompt: "Task", + model: null, + reasoningEffort: null, + agentsStates: {"child-1": {status: method === "item/started" ? "running" : "completed", message: null}}, + }, + }, + } as ServerNotification); + await setupPromptAndSendNotifications(mockFixture, sessionId, sessionState, [ + spawn("item/started"), + { + method: "item/started", + params: { + threadId: sessionId, + turnId: "turn-1", + startedAtMs: 0, + item: { + type: "subAgentActivity", + id: "activity-child", + kind: "started", + agentThreadId: "child-1", + agentPath: "/root/researcher", + }, + }, + }, + spawn("item/completed"), + spawn("item/completed"), + { + method: "item/agentMessage/delta", + params: { + threadId: "child-1", + turnId: "turn-child", + itemId: "late-message", + delta: "Too late", + }, + }, + ]); + + const updates = mockFixture.getAcpConnectionEvents([]) + .filter(event => event.method === "sessionUpdate") + .map(event => event.args[0].update); + expect(updates).toHaveLength(2); + expect(updates.map(update => update.sessionUpdate)).toEqual([ + "subagent_spawned", + "subagent_state_update", + ]); + }); + + it("keeps unsupported collaboration controls visible in native mode", async () => { + await initializeNativeSubagents(); + const collab = ( + method: "item/started" | "item/completed", + tool: "spawnAgent" | "sendInput", + id: string, + status: "running" | "completed", + ): ServerNotification => ({ + method, + params: { + threadId: sessionId, + turnId: "turn-1", + ...(method === "item/started" ? {startedAtMs: 0} : {completedAtMs: 0}), + item: { + type: "collabAgentToolCall", + id, + tool, + status: method === "item/started" ? "inProgress" : "completed", + senderThreadId: sessionId, + receiverThreadIds: ["child-1"], + prompt: tool === "spawnAgent" ? "Child task" : "Additional direction", + model: null, + reasoningEffort: null, + agentsStates: {"child-1": {status, message: null}}, + }, + }, + } as ServerNotification); + + await setupPromptAndSendNotifications(mockFixture, sessionId, sessionState, [ + collab("item/started", "spawnAgent", "spawn", "running"), + { + method: "item/started", + params: { + threadId: sessionId, + turnId: "turn-1", + startedAtMs: 0, + item: { + type: "subAgentActivity", + id: "activity-child", + kind: "started", + agentThreadId: "child-1", + agentPath: "/root/researcher", + }, + }, + }, + collab("item/started", "sendInput", "send-input", "running"), + collab("item/completed", "sendInput", "send-input", "running"), + collab("item/completed", "spawnAgent", "spawn", "completed"), + ]); + + const updates = mockFixture.getAcpConnectionEvents([]) + .filter(event => event.method === "sessionUpdate") + .map(event => event.args[0].update); + expect(updates.map(update => [update.sessionUpdate, update.toolCallId, update.title])).toEqual([ + ["subagent_spawned", undefined, undefined], + ["tool_call", "send-input", "sendInput"], + ["tool_call_update", "send-input", "sendInput"], + ["subagent_state_update", undefined, undefined], + ]); + }); + + it("falls back to tool representation when a native spawn cannot be represented", async () => { + await initializeNativeSubagents(); + await setupPromptAndSendNotifications(mockFixture, sessionId, sessionState, [{ + method: "item/completed", + params: { + threadId: sessionId, + turnId: "turn-1", + completedAtMs: 0, + item: { + type: "collabAgentToolCall", + id: "self-spawn", + tool: "spawnAgent", + status: "failed", + senderThreadId: sessionId, + receiverThreadIds: [sessionId], + prompt: "Invalid task", + model: null, + reasoningEffort: null, + agentsStates: {}, + }, + }, + }]); + + const updates = mockFixture.getAcpConnectionEvents([]) + .filter(event => event.method === "sessionUpdate") + .map(event => event.args[0].update); + expect(updates).toHaveLength(1); + expect(updates[0]).toMatchObject({ + sessionUpdate: "tool_call_update", + toolCallId: "self-spawn", + title: "spawnAgent", + status: "failed", + }); + }); + + it("does not duplicate global notifications after subscribing to a child", async () => { + await initializeNativeSubagents(); + await setupPromptAndSendNotifications(mockFixture, sessionId, sessionState, [ + { + method: "item/started", + params: { + threadId: sessionId, + turnId: "turn-1", + startedAtMs: 0, + item: { + type: "collabAgentToolCall", + id: "spawn", + tool: "spawnAgent", + status: "inProgress", + senderThreadId: sessionId, + receiverThreadIds: ["child-1"], + prompt: "Child task", + model: null, + reasoningEffort: null, + agentsStates: {"child-1": {status: "running", message: null}}, + }, + }, + }, + {method: "warning", params: {threadId: null, message: "Global warning"}}, + { + method: "item/completed", + params: { + threadId: sessionId, + turnId: "turn-1", + completedAtMs: 0, + item: { + type: "collabAgentToolCall", + id: "spawn", + tool: "spawnAgent", + status: "completed", + senderThreadId: sessionId, + receiverThreadIds: ["child-1"], + prompt: "Child task", + model: null, + reasoningEffort: null, + agentsStates: {"child-1": {status: "completed", message: null}}, + }, + }, + }, + ]); + + const warningUpdates = mockFixture.getAcpConnectionEvents([]) + .filter(event => event.method === "sessionUpdate") + .map(event => event.args[0].update) + .filter(update => update.sessionUpdate === "agent_message_chunk" + && update.content?.text.includes("Global warning")); + expect(warningUpdates).toHaveLength(1); + }); + + it("keeps the parent prompt open until every announced child is terminal", async () => { + await initializeNativeSubagents(); + const appServer = mockFixture.getCodexAppServerClient(); + const turn = {id: "turn-1", items: [], status: "inProgress" as const, error: null}; + const completedTurn = {...turn, status: "completed" as const}; + let completeTurn!: () => void; + const completed = new Promise<{threadId: string; turn: typeof completedTurn}>(resolve => { + completeTurn = () => resolve({threadId: sessionId, turn: completedTurn}); + }); + appServer.turnStart = vi.fn().mockResolvedValue({turn}); + appServer.awaitTurnCompleted = vi.fn().mockReturnValue(completed); + vi.spyOn(mockFixture.getCodexAcpAgent(), "getSessionState").mockReturnValue(sessionState); + + const prompt = mockFixture.getCodexAcpAgent().prompt({ + sessionId, + prompt: [{type: "text", text: "Delegate work"}], + }); + await vi.waitFor(() => expect(appServer.turnStart).toHaveBeenCalled()); + const spawn = (status: "running" | "completed") => mockFixture.sendServerNotification({ + method: status === "running" ? "item/started" : "item/completed", + params: { + threadId: sessionId, + turnId: "turn-1", + ...(status === "running" ? {startedAtMs: 0} : {completedAtMs: 0}), + item: { + type: "collabAgentToolCall", + id: "spawn", + tool: "spawnAgent", + status: status === "running" ? "inProgress" : "completed", + senderThreadId: sessionId, + receiverThreadIds: ["child-1"], + prompt: "Child task", + model: null, + reasoningEffort: null, + agentsStates: {"child-1": {status, message: null}}, + }, + }, + }); + spawn("running"); + await mockFixture.sendServerNotification({ + method: "item/started", + params: { + threadId: sessionId, + turnId: "turn-1", + startedAtMs: 0, + item: { + type: "subAgentActivity", + id: "activity-child", + kind: "started", + agentThreadId: "child-1", + agentPath: "/root/researcher", + }, + }, + }); + await mockFixture.getCodexAcpClient().waitForSessionNotifications(sessionId); + await mockFixture.sendServerNotification({ + method: "turn/completed", + params: { + threadId: sessionId, + turn: { + ...completedTurn, + itemsView: "notLoaded", + startedAt: null, + completedAt: null, + durationMs: null, + }, + }, + }); + completeTurn(); + await mockFixture.getCodexAcpClient().waitForSessionNotifications(sessionId); + + let promptSettled = false; + void prompt.finally(() => { promptSettled = true; }); + await new Promise(resolve => setTimeout(resolve, 0)); + expect(promptSettled).toBe(false); + + await mockFixture.sendServerNotification({ + method: "item/started", + params: { + threadId: "child-1", + turnId: "turn-child", + startedAtMs: 0, + item: { + type: "mcpToolCall", + id: "child-mcp-call", + server: "child-server", + tool: "child-tool", + status: "inProgress", + arguments: {}, + appContext: null, + readOnlyHint: null, + pluginId: null, + result: null, + error: null, + durationMs: null, + }, + }, + }); + await mockFixture.getCodexAcpClient().waitForSessionNotifications(sessionId); + mockFixture.setPermissionResponse({outcome: {outcome: "selected", optionId: "allow_once"}}); + await mockFixture.sendServerRequest("mcpServer/elicitation/request", { + threadId: "child-1", + turnId: "turn-child", + serverName: "child-server", + mode: "form", + _meta: {codex_approval_kind: "mcp_tool_call"}, + message: "Allow child tool?", + requestedSchema: {type: "object", properties: {}}, + }); + const childPermission = mockFixture.getAcpConnectionEvents([]) + .find(event => event.method === "requestPermission" + && event.args[0].toolCall.toolCallId === "child-mcp-call"); + expect(childPermission?.args[0].sessionId).toBe("child-1"); + + mockFixture.setElicitationResponse({action: "accept"}); + await mockFixture.sendServerRequest("mcpServer/elicitation/request", { + threadId: "child-1", + turnId: "turn-child", + serverName: "child-server", + mode: "url", + _meta: null, + message: "Authorize child", + url: "https://example.com/child", + elicitationId: "child-url", + }); + await mockFixture.sendServerNotification({ + method: "serverRequest/resolved", + params: {threadId: "child-1", requestId: 1}, + }); + await mockFixture.getCodexAcpClient().waitForSessionNotifications(sessionId); + const completedElicitation = mockFixture.getAcpConnectionEvents([]) + .find(event => event.method === "completeElicitation" + && event.args[0].elicitationId === "child-url"); + expect(completedElicitation).toBeDefined(); + + spawn("completed"); + await expect(prompt).resolves.toMatchObject({stopReason: "end_turn"}); + }); + + it("waits for a pending spawn without publishing fallback identity and suppresses late activity", async () => { + await initializeNativeSubagents(); + const appServer = mockFixture.getCodexAppServerClient(); + const turn = {id: "turn-1", items: [], status: "inProgress" as const, error: null}; + const completedTurn = {...turn, status: "completed" as const}; + let completeTurn!: () => void; + appServer.turnStart = vi.fn().mockResolvedValue({turn}); + appServer.awaitTurnCompleted = vi.fn().mockReturnValue(new Promise(resolve => { + completeTurn = () => resolve({threadId: sessionId, turn: completedTurn}); + })); + vi.spyOn(mockFixture.getCodexAcpAgent(), "getSessionState").mockReturnValue(sessionState); + + const prompt = mockFixture.getCodexAcpAgent().prompt({ + sessionId, + prompt: [{type: "text", text: "Delegate work"}], + }); + await vi.waitFor(() => expect(appServer.turnStart).toHaveBeenCalled()); + const spawn = (status: "running" | "completed") => mockFixture.sendServerNotification({ + method: status === "running" ? "item/started" : "item/completed", + params: { + threadId: sessionId, + turnId: "turn-1", + ...(status === "running" ? {startedAtMs: 0} : {completedAtMs: 0}), + item: { + type: "collabAgentToolCall", + id: "spawn-without-activity", + tool: "spawnAgent", + status: status === "running" ? "inProgress" : "completed", + senderThreadId: sessionId, + receiverThreadIds: ["child-without-activity"], + prompt: "Child task", + model: null, + reasoningEffort: null, + agentsStates: { + "child-without-activity": {status, message: null}, + }, + }, + }, + }); + spawn("running"); + await mockFixture.getCodexAcpClient().waitForSessionNotifications(sessionId); + completeTurn(); + + let promptSettled = false; + void prompt.finally(() => { promptSettled = true; }); + await new Promise(resolve => setTimeout(resolve, 0)); + expect(promptSettled).toBe(false); + + spawn("completed"); + await expect(prompt).resolves.toMatchObject({stopReason: "end_turn"}); + await mockFixture.sendServerNotification({ + method: "item/started", + params: { + threadId: sessionId, + turnId: "turn-1", + startedAtMs: 0, + item: { + type: "subAgentActivity", + id: "late-activity", + kind: "started", + agentThreadId: "child-without-activity", + agentPath: "/root/late_identity", + }, + }, + }); + await mockFixture.getCodexAcpClient().waitForSessionNotifications(sessionId); + const lifecycle = mockFixture.getAcpConnectionEvents([]) + .filter(event => event.method === "sessionUpdate") + .map(event => event.args[0].update) + .filter(update => update.subagentSessionId === "child-without-activity"); + expect(lifecycle).toEqual([]); + }); + + it("finishes only the child whose turn completed", async () => { + await initializeNativeSubagents(); + const activity = (child: string): ServerNotification => ({ + method: "item/started", + params: { + threadId: sessionId, + turnId: "turn-root", + startedAtMs: 0, + item: { + type: "subAgentActivity", + id: `activity-${child}`, + kind: "started", + agentThreadId: child, + agentPath: `/root/${child}`, + }, + }, + }); + await setupPromptAndSendNotifications(mockFixture, sessionId, sessionState, [ + activity("child-a"), + activity("child-b"), + { + method: "turn/completed", + params: { + threadId: "child-a", + turn: { + id: "turn-child-a", + items: [], + itemsView: "notLoaded", + status: "completed", + error: null, + startedAt: null, + completedAt: null, + durationMs: null, + }, + }, + }, + ]); + + const terminalUpdates = mockFixture.getAcpConnectionEvents([]) + .filter(event => event.method === "sessionUpdate") + .map(event => event.args[0].update) + .filter(update => update.sessionUpdate === "subagent_state_update"); + expect(terminalUpdates).toMatchObject([ + {subagentSessionId: "child-a", state: "completed"}, + ]); + }); + + it("keeps child turn boundaries out of the root event handler", async () => { + await initializeNativeSubagents(); + const turn = (id: string, status: "inProgress" | "completed") => ({ + id, + items: [], + itemsView: "notLoaded" as const, + status, + error: null, + startedAt: null, + completedAt: null, + durationMs: null, + }); + await setupPromptAndSendNotifications(mockFixture, sessionId, sessionState, [ + { + method: "turn/started", + params: {threadId: sessionId, turn: turn("root-turn", "inProgress")}, + }, + { + method: "item/agentMessage/delta", + params: { + threadId: sessionId, + turnId: "root-turn", + itemId: "root-message", + delta: "Root output", + }, + }, + ]); + expect(sessionState.currentTurnId).toBe("root-turn"); + + await mockFixture.sendServerNotification({ + method: "item/started", + params: { + threadId: sessionId, + turnId: "root-turn", + startedAtMs: 0, + item: { + type: "subAgentActivity", + id: "child-activity", + kind: "started", + agentThreadId: "child-turn-thread", + agentPath: "/root/child_turn", + }, + }, + }); + await mockFixture.sendServerNotification({ + method: "turn/started", + params: { + threadId: "child-turn-thread", + turn: turn("child-turn", "inProgress"), + }, + }); + await mockFixture.getCodexAcpClient().waitForSessionNotifications(sessionId); + expect(sessionState.currentTurnId).toBe("root-turn"); + + await mockFixture.sendServerNotification({ + method: "turn/completed", + params: { + threadId: "child-turn-thread", + turn: turn("child-turn", "completed"), + }, + }); + await mockFixture.getCodexAcpClient().waitForSessionNotifications(sessionId); + expect(sessionState.currentTurnId).toBe("root-turn"); + + await mockFixture.sendServerNotification({ + method: "turn/completed", + params: {threadId: sessionId, turn: turn("root-turn", "completed")}, + }); + await mockFixture.getCodexAcpClient().waitForSessionNotifications(sessionId); + expect(sessionState.currentTurnId).toBeNull(); }); }); diff --git a/src/__tests__/CodexACPAgent/data/collab-agent-tool-call-flow.json b/src/__tests__/CodexACPAgent/data/collab-agent-tool-call-flow.json deleted file mode 100644 index 4391e5a9..00000000 --- a/src/__tests__/CodexACPAgent/data/collab-agent-tool-call-flow.json +++ /dev/null @@ -1,83 +0,0 @@ -{ - "method": "sessionUpdate", - "args": [ - { - "sessionId": "test-session-id", - "update": { - "sessionUpdate": "tool_call", - "toolCallId": "call-spawn-weather", - "kind": "other", - "title": "spawnAgent", - "status": "in_progress", - "rawInput": { - "prompt": "Find the current weather in Paris.", - "senderThreadId": "thread-main", - "receiverThreadIds": [ - "thread-paris" - ], - "agentsStates": { - "thread-paris": { - "status": "running", - "message": "Checking weather" - } - }, - "model": null, - "reasoningEffort": null, - "status": "inProgress" - }, - "_meta": { - "codex": { - "collaboration": { - "tool": "spawnAgent", - "senderThreadId": "thread-main", - "receiverThreadIds": [ - "thread-paris" - ] - } - } - } - } - } - ] -} -{ - "method": "sessionUpdate", - "args": [ - { - "sessionId": "test-session-id", - "update": { - "sessionUpdate": "tool_call_update", - "toolCallId": "call-spawn-weather", - "title": "spawnAgent", - "status": "completed", - "rawInput": { - "prompt": "Find the current weather in Paris.", - "senderThreadId": "thread-main", - "receiverThreadIds": [ - "thread-paris" - ], - "agentsStates": { - "thread-paris": { - "status": "completed", - "message": null - } - }, - "model": null, - "reasoningEffort": null, - "status": "completed" - }, - "_meta": { - "codex": { - "collaboration": { - "tool": "spawnAgent", - "senderThreadId": "thread-main", - "receiverThreadIds": [ - "thread-paris" - ] - } - } - } - } - } - ] -} diff --git a/src/__tests__/CodexACPAgent/data/subagent-activity-flow.json b/src/__tests__/CodexACPAgent/data/subagent-activity-flow.json deleted file mode 100644 index e1b6749c..00000000 --- a/src/__tests__/CodexACPAgent/data/subagent-activity-flow.json +++ /dev/null @@ -1,29 +0,0 @@ -{ - "method": "sessionUpdate", - "args": [ - { - "sessionId": "test-session-id", - "update": { - "sessionUpdate": "tool_call", - "title": "Start subagent weather_research", - "kind": "other", - "toolCallId": "call-spawn-weather", - "status": "completed", - "rawInput": { - "agentThreadId": "thread-paris", - "agentPath": "/root/weather_research", - "activityKind": "started" - }, - "_meta": { - "codex": { - "subagent": { - "threadId": "thread-paris", - "path": "/root/weather_research", - "activity": "started" - } - } - } - } - } - ] -} diff --git a/src/__tests__/CodexACPAgent/initialize.test.ts b/src/__tests__/CodexACPAgent/initialize.test.ts index 739c4028..57887361 100644 --- a/src/__tests__/CodexACPAgent/initialize.test.ts +++ b/src/__tests__/CodexACPAgent/initialize.test.ts @@ -73,7 +73,7 @@ describe('CodexACPAgent - initialize', () => { jetbrains: { air: { version: 1, - capabilities: ["sessionFailure", "agentFileChangeReport"], + capabilities: ["sessionFailure", "agentFileChangeReport", "nativeSubagentSessions"], }, }, }, diff --git a/src/__tests__/CodexACPAgent/typed-session-failure-wire.test.ts b/src/__tests__/CodexACPAgent/typed-session-failure-wire.test.ts index 790c2614..115ed58d 100644 --- a/src/__tests__/CodexACPAgent/typed-session-failure-wire.test.ts +++ b/src/__tests__/CodexACPAgent/typed-session-failure-wire.test.ts @@ -11,6 +11,31 @@ const typedFailureCapabilities: acp.ClientCapabilities = { }; describe("typed session failures over ACP transport", () => { + it("negotiates native subagents through AIR metadata across the SDK boundary", async () => { + const fixture = createWireFixture(); + const response = await fixture.client.initialize({ + protocolVersion: acp.PROTOCOL_VERSION, + clientCapabilities: { + _meta: { + jetbrains: { + air: {version: 1, capabilities: ["nativeSubagentSessions"]}, + }, + }, + }, + }); + + expect((response.agentCapabilities!.sessionCapabilities as {subagents?: unknown}).subagents) + .toEqual({}); + expect(response._meta).toMatchObject({ + jetbrains: { + air: { + version: 1, + capabilities: expect.arrayContaining(["nativeSubagentSessions"]), + }, + }, + }); + }); + it("returns a sanitized process-exit failure in the decoded prompt response", async () => { const fixture = createWireFixture({ exitCode: 1, diff --git a/src/__tests__/PermissionLifecycleContext.test.ts b/src/__tests__/PermissionLifecycleContext.test.ts index ece63118..ad90788d 100644 --- a/src/__tests__/PermissionLifecycleContext.test.ts +++ b/src/__tests__/PermissionLifecycleContext.test.ts @@ -12,11 +12,11 @@ function sessionState(): SessionState { } as SessionState; } -function mcpStarted(id: string, turnId: string): ServerNotification { +function mcpStarted(id: string, turnId: string, threadId = "thread"): ServerNotification { return { method: "item/started", params: { - threadId: "thread", + threadId, turnId, startedAtMs: 0, item: { @@ -37,6 +37,42 @@ function mcpStarted(id: string, turnId: string): ServerNotification { }; } +function fileChangeStarted(id: string, threadId: string): ServerNotification { + return { + method: "item/started", + params: { + threadId, + turnId: `turn-${threadId}`, + startedAtMs: 0, + item: { + type: "fileChange", + id, + changes: [{path: `/${threadId}.txt`, kind: {type: "add"}, diff: "+content"}], + status: "inProgress", + }, + }, + }; +} + +function turnCompleted(threadId: string): ServerNotification { + return { + method: "turn/completed", + params: { + threadId, + turn: { + id: `turn-${threadId}`, + items: [], + itemsView: "full", + status: "completed", + error: null, + startedAt: 0, + completedAt: 1, + durationMs: 1_000, + }, + }, + }; +} + describe("PermissionLifecycleContext", () => { it("clears MCP correlation at the turn boundary", () => { const lifecycle = new PermissionLifecycleContext(sessionState()); @@ -81,6 +117,21 @@ describe("PermissionLifecycleContext", () => { expect(currentPrompt.popPendingMcpApproval("thread", "server")).toBe("current-call"); }); + it("clears only the completed thread's permission correlation", () => { + const prompt = new PermissionLifecycleContext(sessionState()).beginPrompt(); + prompt.handleNotification(mcpStarted("call-a", "turn-a", "child-a")); + prompt.handleNotification(mcpStarted("call-b", "turn-b", "child-b")); + prompt.handleNotification(fileChangeStarted("shared-file-change", "child-a")); + prompt.handleNotification(fileChangeStarted("shared-file-change", "child-b")); + + prompt.handleNotification(turnCompleted("child-b")); + + expect(prompt.popPendingMcpApproval("child-a", "server")).toBe("call-a"); + expect(prompt.popPendingMcpApproval("child-b", "server")).toBeUndefined(); + expect(prompt.fileChange("child-a", "shared-file-change")?.changes[0]?.path).toBe("/child-a.txt"); + expect(prompt.fileChange("child-b", "shared-file-change")).toBeUndefined(); + }); + it("does not allocate a synthetic ID for native ACP elicitation", async () => { const state = sessionState(); const lifecycle = new PermissionLifecycleContext(state); @@ -90,7 +141,6 @@ describe("PermissionLifecycleContext", () => { } as unknown as AcpClientConnection; const handler = new CodexElicitationHandler( connection, - state, prompt, {elicitation: {form: {}}}, ); @@ -119,7 +169,7 @@ describe("PermissionLifecycleContext", () => { }), notify: vi.fn(), } as unknown as AcpClientConnection; - const handler = new CodexElicitationHandler(connection, state, prompt); + const handler = new CodexElicitationHandler(connection, prompt); const approval = { threadId: "thread", turnId: "turn-1", diff --git a/src/__tests__/acp-test-utils.ts b/src/__tests__/acp-test-utils.ts index 6b9e20c0..5cf73d5d 100644 --- a/src/__tests__/acp-test-utils.ts +++ b/src/__tests__/acp-test-utils.ts @@ -4,7 +4,7 @@ import {CodexAcpClient} from '../CodexAcpClient'; import {CodexAppServerClient, type CodexConnectionEvent} from '../CodexAppServerClient'; import {startCodexConnection} from "../CodexJsonRpcConnection"; import {CodexAcpServer, type SessionState} from "../CodexAcpServer"; -import type {AcpClientConnection} from "../ACPSessionConnection"; +import {ACPSessionConnection, type AcpClientConnection} from "../ACPSessionConnection"; import type {ServerNotification} from "../app-server"; import type {MessageConnection} from "vscode-jsonrpc/node"; import path from "node:path"; @@ -14,6 +14,7 @@ import {AgentMode} from "../AgentMode"; import {DEFAULT_COLLABORATION_MODE} from "../CollaborationModeConfig"; import {expect, vi} from "vitest"; import type {Model, ReasoningEffortOption} from "../app-server/v2"; +import {CodexSubagentEventRouter} from "../subagents/CodexSubagentEventRouter"; export type MethodCallEvent = { method: string; args: any[] }; @@ -69,6 +70,7 @@ export interface TestFixture { getAcpConnectionEvents(ignoredFields: string[]): MethodCallEvent[], getAcpConnectionDump(ignoredFields: string[]): string, clearAcpConnectionDump(): void, + getAcpConnection(): AcpClientConnection, } export interface CodexConnectionDumpOptions { @@ -167,6 +169,9 @@ export function createBaseTestFixture(config: ConnectionConfig): TestFixture { }, clearAcpConnectionDump() { acpConnectionEvents.splice(0, acpConnectionEvents.length); + }, + getAcpConnection(): AcpClientConnection { + return acpConnection; } }; } @@ -379,6 +384,7 @@ function anonymizeValue(value: any, path: string[], fieldsToAnonymize: Set): SessionState { + const sessionId = overrides?.sessionId ?? "session-id"; return { currentTurnId: null, lastTokenUsage: null, @@ -390,7 +396,7 @@ export function createTestSessionState(overrides?: Partial): Sessi authProvider: null, cwd: "/test/cwd", additionalDirectories: [], - sessionId: "session-id", + sessionId, currentModelId: "model-id[effort]", availableModels: [], supportedReasoningEfforts: [], @@ -403,6 +409,11 @@ export function createTestSessionState(overrides?: Partial): Sessi goalRevision: 0, sessionTitle: null, sessionTitleSource: "unknown", + subagents: new CodexSubagentEventRouter( + sessionId, + false, + new ACPSessionConnection({notify: vi.fn(), request: vi.fn()} as AcpClientConnection, sessionId), + ), ...overrides, }; } diff --git a/src/permissions/CodexApprovalHandler.ts b/src/permissions/CodexApprovalHandler.ts index 3dab58f5..38a639ee 100644 --- a/src/permissions/CodexApprovalHandler.ts +++ b/src/permissions/CodexApprovalHandler.ts @@ -1,5 +1,4 @@ import * as acp from "@agentclientprotocol/sdk"; -import type {SessionState} from "../CodexAcpServer"; import type {ApprovalHandler} from "../CodexAppServerClient"; import type { CommandExecutionRequestApprovalParams, @@ -34,7 +33,6 @@ import type {PermissionPromptContext} from "./lifecycle"; export class CodexApprovalHandler implements ApprovalHandler { constructor( private readonly connection: AcpClientConnection, - private readonly sessionState: SessionState, private readonly permissionContext: PermissionPromptContext, private readonly cancellationSignal?: AbortSignal, ) {} @@ -51,7 +49,7 @@ export class CodexApprovalHandler implements ApprovalHandler { try { const response = await this.requestPermission({ - sessionId: this.sessionState.sessionId, + sessionId: params.threadId, toolCall: commandToolCall(authoritativeParams), options: decisions.map(({option}) => option), _meta: requestPermissionMeta( @@ -70,7 +68,7 @@ export class CodexApprovalHandler implements ApprovalHandler { const decisions = fileChangeDecisionOptions(); try { const response = await this.requestPermission({ - sessionId: this.sessionState.sessionId, + sessionId: params.threadId, toolCall: fileChangeToolCall(params, this.permissionContext), options: decisions.map(({option}) => option), _meta: requestPermissionMeta(CODEX_FILE_CHANGE_PERMISSION_TITLE, params.reason), @@ -87,7 +85,7 @@ export class CodexApprovalHandler implements ApprovalHandler { ): Promise { try { const response = await this.requestPermission({ - sessionId: this.sessionState.sessionId, + sessionId: params.threadId, toolCall: additionalPermissionsToolCall( params.itemId, params.cwd, diff --git a/src/permissions/lifecycle.ts b/src/permissions/lifecycle.ts index 18060764..e2dd7bbb 100644 --- a/src/permissions/lifecycle.ts +++ b/src/permissions/lifecycle.ts @@ -22,7 +22,7 @@ export class PermissionLifecycleContext { /** Prompt-scoped permission presentation and MCP correlation state. */ export class PermissionPromptContext { - private readonly fileChanges = new Map(); + private readonly fileChanges = new Map>(); private readonly pendingMcpApprovals = new Map>(); constructor(private readonly nextStandaloneId: (serverName: string) => string) {} @@ -36,7 +36,7 @@ export class PermissionPromptContext { this.handleItemCompleted(notification.params.threadId, notification.params.item); return; case "turn/completed": - this.clearTransientState(); + this.clearTransientState(notification.params.threadId); return; case "serverRequest/resolved": this.pendingMcpApprovals.delete(notification.params.threadId); @@ -46,8 +46,8 @@ export class PermissionPromptContext { } } - fileChange(itemId: string): FileChangeItem | undefined { - return this.fileChanges.get(itemId); + fileChange(threadId: string, itemId: string): FileChangeItem | undefined { + return this.fileChanges.get(threadId)?.get(itemId); } popPendingMcpApproval(threadId: string, serverName: string): string | undefined { @@ -67,7 +67,9 @@ export class PermissionPromptContext { private handleItemStarted(threadId: string, item: ThreadItem): void { if (item.type === "fileChange") { - this.fileChanges.set(item.id, item); + const byItem = this.fileChanges.get(threadId) ?? new Map(); + byItem.set(item.id, item); + this.fileChanges.set(threadId, byItem); return; } if (item.type !== "mcpToolCall") return; @@ -80,7 +82,9 @@ export class PermissionPromptContext { private handleItemCompleted(threadId: string, item: ThreadItem): void { if (item.type === "fileChange") { - this.fileChanges.delete(item.id); + const byItem = this.fileChanges.get(threadId); + byItem?.delete(item.id); + if (byItem?.size === 0) this.fileChanges.delete(threadId); return; } if (item.type !== "mcpToolCall") return; @@ -94,8 +98,8 @@ export class PermissionPromptContext { if (byServer.size === 0) this.pendingMcpApprovals.delete(threadId); } - private clearTransientState(): void { - this.fileChanges.clear(); - this.pendingMcpApprovals.clear(); + private clearTransientState(threadId: string): void { + this.fileChanges.delete(threadId); + this.pendingMcpApprovals.delete(threadId); } } diff --git a/src/permissions/presentation.ts b/src/permissions/presentation.ts index 67704692..d2b9cede 100644 --- a/src/permissions/presentation.ts +++ b/src/permissions/presentation.ts @@ -52,7 +52,7 @@ export function fileChangeToolCall( params: FileChangeRequestApprovalParams, permissionContext: PermissionPromptContext, ): acp.ToolCallUpdate { - const item = permissionContext.fileChange(params.itemId); + const item = permissionContext.fileChange(params.threadId, params.itemId); return { toolCallId: params.itemId, kind: "edit", diff --git a/src/subagents/AcpSubagents.ts b/src/subagents/AcpSubagents.ts new file mode 100644 index 00000000..962b0cb5 --- /dev/null +++ b/src/subagents/AcpSubagents.ts @@ -0,0 +1,67 @@ +import type { + ClientCapabilities, + SessionCapabilities, + SessionNotification, +} from "@agentclientprotocol/sdk"; +import { + AIR_NATIVE_SUBAGENT_SESSIONS_KEY, + clientSupportsAirCapability, +} from "../AirExtension"; + +/** Temporary typed surface for agentclientprotocol/agent-client-protocol#1992. */ +export type SubagentSessionCapabilities = { + cancel?: boolean; + close?: boolean; + _meta?: Record | null; +}; + +export type SubagentSpawnedUpdate = { + sessionUpdate: "subagent_spawned"; + subagentSessionId: string; + name: string; + description: string; + capabilities: SubagentSessionCapabilities; + _meta?: Record | null; +}; + +export type SubagentState = "completed" | "failed" | "cancelled"; + +export type SubagentStateUpdate = { + sessionUpdate: "subagent_state_update"; + subagentSessionId: string; + state: SubagentState; + _meta?: Record | null; +}; + +export type AcpSessionUpdate = + | SessionNotification["update"] + | SubagentSpawnedUpdate + | SubagentStateUpdate; + +export type AcpSessionNotification = Omit & { + update: AcpSessionUpdate; +}; + +export type SubagentAwareSessionCapabilities = SessionCapabilities & { + subagents?: Record; +}; + +export function clientSupportsSubagents( + capabilities?: ClientCapabilities | null, +): boolean { + const subagents = ( + capabilities as (ClientCapabilities & { subagents?: unknown }) | null | undefined + )?.subagents; + if (typeof subagents === "object" && subagents !== null && !Array.isArray(subagents)) { + return true; + } + + 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/CodexAgentPath.ts b/src/subagents/CodexAgentPath.ts new file mode 100644 index 00000000..01198b5b --- /dev/null +++ b/src/subagents/CodexAgentPath.ts @@ -0,0 +1,17 @@ +export function normalizeAgentPath(path: string): string { + const normalized = path.trim().replace(/\/+$/, ""); + return normalized || "/root"; +} + +export function isRootAgentPath(path: string): boolean { + const normalized = normalizeAgentPath(path); + return normalized === "/root" || normalized === "root"; +} + +export function nameFromAgentPath(path: string, fallback: string): string { + const normalized = normalizeAgentPath(path); + const name = normalized.slice(normalized.lastIndexOf("/") + 1).trim(); + if (!name) return fallback; + const words = name.replace(/[_-]+/g, " ").replace(/\s+/g, " ").trim(); + return words ? words.charAt(0).toUpperCase() + words.slice(1) : fallback; +} diff --git a/src/subagents/CodexSubagentEventRouter.ts b/src/subagents/CodexSubagentEventRouter.ts new file mode 100644 index 00000000..2f295587 --- /dev/null +++ b/src/subagents/CodexSubagentEventRouter.ts @@ -0,0 +1,321 @@ +import type {ServerNotification} from "../app-server"; +import type {ThreadItem} from "../app-server/v2"; +import {ACPSessionConnection, type UpdateSessionEvent} from "../ACPSessionConnection"; +import {logger} from "../Logger"; +import { + createCollabAgentToolCallCompleteUpdate, + createCollabAgentToolCallUpdate, + createSubAgentActivityUpdate, +} from "../CodexToolCallMapper"; +import type {SubagentState} from "./AcpSubagents"; +import {isRootAgentPath, nameFromAgentPath, normalizeAgentPath} from "./CodexAgentPath"; + +type NativeSubagent = { + parentSessionId: string; + name: string; + description: string; + path?: string; + terminalState?: SubagentState; +}; + +type PendingSubagent = { + parentSessionId: string; + description: string; +}; + +/** Owns native lifecycle, child routing, waiting, and legacy activity deduplication. */ +export class CodexSubagentEventRouter { + private static readonly DEFAULT_WAIT_TIMEOUT_MS = 10 * 60 * 1000; + + private readonly children = new Map(); + private readonly pendingSpawns = new Map(); + private readonly terminalPendingSpawns = new Set(); + private readonly waiters = new Set<() => void>(); + private readonly activeLegacyActivities = new Set(); + + constructor( + private readonly rootSessionId: string, + private readonly supported: boolean, + private readonly session: ACPSessionConnection, + ) {} + + async handle(notification: ServerNotification): Promise { + if (notification.method === "turn/started") { + return this.isKnownChild(notification.params.threadId); + } + if (notification.method === "turn/completed") { + const childTurn = this.isKnownChild(notification.params.threadId); + const state = terminalStateFromTurn(notification.params.turn.status); + if (!state) return childTurn; + if (notification.params.threadId === this.rootSessionId) { + if (state !== "completed") await this.finishOutstanding(state); + } + else { + if (this.pendingSpawns.has(notification.params.threadId)) { + this.finishPending(notification.params.threadId); + } + else { + await this.finish(notification.params.threadId, state); + } + } + return childTurn; + } + if (notification.method !== "item/started" && notification.method !== "item/completed") { + return false; + } + const item = notification.params.item; + if (!this.supported) { + // Preserve the pre-native protocol representation for clients that + // did not negotiate child sessions. The normal event mapper renders + // collaboration lifecycle as ordinary ACP tool calls. + return false; + } + if (item.type === "subAgentActivity") { + // Codex reports the root participant through the same activity item + // shape as children. It is the parent conversation, not a subagent. + if (isRootAgentPath(item.agentPath)) return true; + if (this.terminalPendingSpawns.has(item.agentThreadId)) return true; + let hasNativeRepresentation = this.children.has(item.agentThreadId); + if (!hasNativeRepresentation) { + await this.materialize(item.agentThreadId, item.agentPath); + hasNativeRepresentation = this.children.has(item.agentThreadId); + } + if (hasNativeRepresentation && item.kind === "interrupted") { + await this.finish(item.agentThreadId, "cancelled"); + } + return hasNativeRepresentation; + } + if (item.type !== "collabAgentToolCall") return false; + + let representedSpawn = false; + if (item.tool === "spawnAgent") { + const parentSessionId = this.children.has(item.senderThreadId) + ? item.senderThreadId + : this.rootSessionId; + for (const childSessionId of item.receiverThreadIds) { + if (childSessionId.trim().length === 0) { + logger.log("Ignoring spawned subagent with an empty thread id"); + continue; + } + if (childSessionId === parentSessionId || childSessionId === this.rootSessionId) { + logger.log(`Ignoring self-referential spawned subagent ${childSessionId}`); + continue; + } + if (this.children.has(childSessionId) + || this.pendingSpawns.has(childSessionId) + || this.terminalPendingSpawns.has(childSessionId)) { + representedSpawn = true; + continue; + } + this.pendingSpawns.set(childSessionId, { + parentSessionId, + description: item.prompt?.trim() || "Delegated task", + }); + representedSpawn = true; + } + } + + for (const [childSessionId, state] of Object.entries(item.agentsStates)) { + const terminalState = state && terminalStateOf(state.status); + if (!terminalState) continue; + if (this.children.has(childSessionId)) await this.finish(childSessionId, terminalState); + else if (this.pendingSpawns.has(childSessionId)) this.finishPending(childSessionId); + } + if (item.tool === "spawnAgent" && item.status === "failed") { + for (const childSessionId of item.receiverThreadIds) { + if (this.pendingSpawns.has(childSessionId)) this.finishPending(childSessionId); + } + } + // `updated` is intentionally not synthesized: the portable protocol + // currently defines only spawn and terminal lifecycle. + return item.tool === "spawnAgent" && representedSpawn; + } + + shouldIgnore(notification: ServerNotification): boolean { + const threadId = (notification.params as {threadId?: unknown}).threadId; + const ignored = typeof threadId === "string" + && (this.children.get(threadId)?.terminalState !== undefined + || this.terminalPendingSpawns.has(threadId)); + if (ignored) logger.log(`Ignoring update for terminal subagent ${threadId}`); + return ignored; + } + + notificationSessionId(notification: ServerNotification): string { + const threadId = (notification.params as {threadId?: unknown}).threadId; + return typeof threadId === "string" && this.children.has(threadId) + ? threadId + : this.rootSessionId; + } + + legacyActivityStarted(item: ThreadItem & {type: "subAgentActivity"}): UpdateSessionEvent { + this.activeLegacyActivities.add(item.id); + return createSubAgentActivityUpdate(item, "in_progress", "tool_call"); + } + + legacyCollaborationStarted(item: ThreadItem & {type: "collabAgentToolCall"}): UpdateSessionEvent { + return createCollabAgentToolCallUpdate(item); + } + + legacyCollaborationCompleted(item: ThreadItem & {type: "collabAgentToolCall"}): UpdateSessionEvent { + return createCollabAgentToolCallCompleteUpdate(item); + } + + legacyActivityCompleted(item: ThreadItem & {type: "subAgentActivity"}): UpdateSessionEvent { + const sessionUpdate = this.activeLegacyActivities.delete(item.id) + ? "tool_call_update" + : "tool_call"; + return createSubAgentActivityUpdate(item, "completed", sessionUpdate); + } + + async wait( + signal: AbortSignal, + timeoutMs = CodexSubagentEventRouter.DEFAULT_WAIT_TIMEOUT_MS, + ): Promise { + const deadline = Date.now() + timeoutMs; + while (this.hasOutstanding()) { + if (signal.aborted) return; + const remainingMs = deadline - Date.now(); + if (remainingMs <= 0) { + logger.log(`Timed out waiting for subagents in session ${this.rootSessionId}; marking them failed`); + await this.finishOutstanding("failed"); + return; + } + const changed = await new Promise((resolve) => { + const timeout = setTimeout(() => { + this.waiters.delete(onChange); + signal.removeEventListener("abort", onAbort); + resolve(false); + }, remainingMs); + const onAbort = () => { + clearTimeout(timeout); + this.waiters.delete(onChange); + resolve(true); + }; + const onChange = () => { + clearTimeout(timeout); + signal.removeEventListener("abort", onAbort); + resolve(true); + }; + this.waiters.add(onChange); + signal.addEventListener("abort", onAbort, {once: true}); + }); + if (!changed) { + logger.log(`Timed out waiting for subagents in session ${this.rootSessionId}; marking them failed`); + await this.finishOutstanding("failed"); + return; + } + } + } + + async finishOutstanding(state: SubagentState): Promise { + for (const childSessionId of [...this.pendingSpawns.keys()]) { + this.finishPending(childSessionId); + } + for (const childSessionId of [...this.children.keys()].reverse()) { + await this.finish(childSessionId, state); + } + } + + private hasOutstanding(): boolean { + return this.pendingSpawns.size > 0 + || [...this.children.values()].some(child => child.terminalState === undefined); + } + + private isKnownChild(threadId: string): boolean { + return threadId !== this.rootSessionId + && (this.children.has(threadId) + || this.pendingSpawns.has(threadId) + || this.terminalPendingSpawns.has(threadId)); + } + + private async materialize(childSessionId: string, path: string): Promise { + if (this.children.has(childSessionId)) return; + const pending = this.pendingSpawns.get(childSessionId); + const name = nameFromAgentPath(path, fallbackName(childSessionId)); + const parentSessionId = pending?.parentSessionId + ?? this.parentSessionIdForPath(path); + const description = pending?.description ?? `Delegated task for ${name}`; + await this.session.update({ + sessionUpdate: "subagent_spawned", + subagentSessionId: childSessionId, + name, + description, + capabilities: {}, + }, parentSessionId); + this.children.set(childSessionId, { + parentSessionId, + name, + description, + path: normalizeAgentPath(path), + }); + this.pendingSpawns.delete(childSessionId); + } + + private finishPending(childSessionId: string): void { + if (!this.pendingSpawns.delete(childSessionId)) return; + this.terminalPendingSpawns.add(childSessionId); + for (const waiter of this.waiters) waiter(); + this.waiters.clear(); + } + + private async finish(childSessionId: string, state: SubagentState): Promise { + const child = this.children.get(childSessionId); + if (!child || child.terminalState !== undefined) return; + await this.session.update({ + sessionUpdate: "subagent_state_update", + subagentSessionId: childSessionId, + state, + }, child.parentSessionId); + child.terminalState = state; + for (const waiter of this.waiters) waiter(); + this.waiters.clear(); + } + + private parentSessionIdForPath(path: string): string { + const normalized = normalizeAgentPath(path); + const separator = normalized.lastIndexOf("/"); + if (separator <= 0) return this.rootSessionId; + const parentPath = normalized.slice(0, separator); + return [...this.children.entries()] + .find(([, child]) => child.path === parentPath)?.[0] + ?? this.rootSessionId; + } +} + +function terminalStateOf( + status: "pendingInit" | "running" | "completed" | "errored" | "shutdown" | "notFound" | "interrupted", +): SubagentState | undefined { + switch (status) { + case "completed": + return "completed"; + case "interrupted": + return "cancelled"; + case "errored": + case "shutdown": + case "notFound": + return "failed"; + case "pendingInit": + case "running": + return undefined; + } +} + +function terminalStateFromTurn( + status: "inProgress" | "completed" | "interrupted" | "failed", +): SubagentState | undefined { + switch (status) { + case "completed": + return "completed"; + case "interrupted": + return "cancelled"; + case "failed": + return "failed"; + case "inProgress": + return undefined; + } +} + +function fallbackName(sessionId: string): string { + const suffix = sessionId.length > 8 ? sessionId.slice(-8) : sessionId; + return `Agent ${suffix}`; +} diff --git a/src/subagents/CodexSubagentSubscriptions.ts b/src/subagents/CodexSubagentSubscriptions.ts new file mode 100644 index 00000000..cec46ebe --- /dev/null +++ b/src/subagents/CodexSubagentSubscriptions.ts @@ -0,0 +1,134 @@ +import type { + ApprovalHandler, + CodexAppServerClient, + ElicitationHandler, +} from "../CodexAppServerClient"; +import type {ServerNotification} from "../app-server"; +import {isRootAgentPath} from "./CodexAgentPath"; + +type Subscription = { + rootSessionId: string; + supportsSubagents: boolean; + dispatch(event: ServerNotification): void; + approvalHandler: ApprovalHandler; + elicitationHandler: ElicitationHandler; + waitForRootNotifications(): Promise; +}; + +type SessionSubscription = { + current: Subscription; + children: Set; +}; + +/** Discovers child threads and keeps their output/interaction boundary negotiated. */ +export class CodexSubagentSubscriptions { + private readonly sessions = new Map(); + + constructor(private readonly client: CodexAppServerClient) {} + + subscribe(subscription: Subscription): void { + const existing = this.sessions.get(subscription.rootSessionId); + if (existing) { + existing.current = subscription; + return; + } + + const session = {current: subscription, children: new Set()}; + this.sessions.set(subscription.rootSessionId, session); + this.client.onServerNotification(subscription.rootSessionId, (event) => { + // Register synchronously: app-server may emit child output directly + // after the spawning collaboration item. + this.discover(session, event); + session.current.dispatch(event); + }); + this.registerInteractiveHandlers(session, subscription.rootSessionId); + } + + clear(rootSessionId: string): void { + for (const childSessionId of this.sessions.get(rootSessionId)?.children ?? []) { + this.client.clearThreadHandlers(childSessionId); + } + this.sessions.delete(rootSessionId); + } + + private discover(session: SessionSubscription, event: ServerNotification): void { + if (event.method !== "item/started" && event.method !== "item/completed") { + return; + } + const item = event.params.item; + const childSessionIds = item.type === "collabAgentToolCall" && item.tool === "spawnAgent" + ? item.receiverThreadIds + : item.type === "subAgentActivity" && item.kind !== "interrupted" && !isRootAgentPath(item.agentPath) + ? [item.agentThreadId] + : []; + for (const childSessionId of childSessionIds) { + if (childSessionId.trim() === "") continue; + if (childSessionId === session.current.rootSessionId + || childSessionId === event.params.threadId + || session.children.has(childSessionId)) { + continue; + } + session.children.add(childSessionId); + this.client.onServerNotification(childSessionId, (childEvent) => { + const eventThreadId = (childEvent.params as {threadId?: unknown}).threadId; + if (eventThreadId !== childSessionId) return; + this.discover(session, childEvent); + if (session.current.supportsSubagents) session.current.dispatch(childEvent); + }); + // Hidden children keep only root-attributed permission requests. + this.registerInteractiveHandlers(session, childSessionId); + } + } + + private registerInteractiveHandlers(session: SessionSubscription, targetSessionId: string): void { + this.client.onApprovalRequest(targetSessionId, { + handleCommandExecution: async (params) => { + const current = session.current; + await current.waitForRootNotifications(); + return await current.approvalHandler.handleCommandExecution( + this.rootInteractionParams(current, targetSessionId, params), + ); + }, + handleFileChange: async (params) => { + const current = session.current; + await current.waitForRootNotifications(); + return await current.approvalHandler.handleFileChange( + this.rootInteractionParams(current, targetSessionId, params), + ); + }, + handlePermissionsRequest: async (params) => { + const current = session.current; + await current.waitForRootNotifications(); + return await current.approvalHandler.handlePermissionsRequest( + this.rootInteractionParams(current, targetSessionId, params), + ); + }, + }); + this.client.onElicitationRequest(targetSessionId, { + handleElicitation: async (params) => { + const current = session.current; + await current.waitForRootNotifications(); + return await current.elicitationHandler.handleElicitation( + this.rootInteractionParams(current, targetSessionId, params), + ); + }, + handleUserInput: async (params) => { + const current = session.current; + await current.waitForRootNotifications(); + return await current.elicitationHandler.handleUserInput( + this.rootInteractionParams(current, targetSessionId, params), + ); + }, + }); + } + + private rootInteractionParams( + subscription: Subscription, + targetSessionId: string, + params: T, + ): T { + return !subscription.supportsSubagents && targetSessionId !== subscription.rootSessionId + ? {...params, threadId: subscription.rootSessionId} + : params; + } +}