diff --git a/src/AirExtension.ts b/src/AirExtension.ts index c16df898..b01d3dec 100644 --- a/src/AirExtension.ts +++ b/src/AirExtension.ts @@ -11,4 +11,5 @@ export const AIR_META_KEY = "air"; export const AIR_EXTENSION_VERSION_KEY = "version"; export const AIR_EXTENSION_CAPABILITIES_KEY = "capabilities"; export const AIR_SESSION_FAILURE_KEY = "sessionFailure"; +export const AIR_COMPLETION_DETAILS_KEY = "completionDetails"; export const AIR_EXTENSION_VERSION = 1; diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index dc21a60a..2f57bd5d 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -14,7 +14,7 @@ import { import type {McpStartupResult} from "./CodexAppServerClient"; import {type AcpClientConnection, ACPSessionConnection, type UpdateSessionEvent} from "./ACPSessionConnection"; import type {InputModality, ReasoningEffort} from "./app-server"; -import type {Account, Model, ReasoningEffortOption, Thread, ThreadGoal, ThreadItem, UserInput} from "./app-server/v2"; +import type {Account, Model, ReasoningEffortOption, Thread, ThreadGoal, ThreadItem, Turn, UserInput} from "./app-server/v2"; import type {RateLimitsMap} from "./RateLimitsMap"; import {ModelId} from "./ModelId"; import {AgentMode, MODE_CONFIG_ID} from "./AgentMode"; @@ -93,6 +93,7 @@ import { import {sameThreadGoalSnapshot, type ThreadGoalSnapshot, toThreadGoalSnapshot,} from "./ThreadGoalSnapshot"; import {randomUUID} from "node:crypto"; import { + AIR_COMPLETION_DETAILS_KEY, AIR_EXTENSION_CAPABILITIES_KEY, AIR_EXTENSION_VERSION, AIR_EXTENSION_VERSION_KEY, @@ -152,9 +153,18 @@ export interface SessionFailure { turnId?: string; } +interface CompletionDetails { + turnId: string; + partial: boolean; + retryable: boolean; +} + const CODEX_PROCESS_EXITED_ERROR_CODE = 1001; -function clientSupportsTypedSessionFailures(capabilities: acp.ClientCapabilities | null): boolean { +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]; @@ -163,7 +173,15 @@ function clientSupportsTypedSessionFailures(capabilities: acp.ClientCapabilities && Number.isInteger(version) && version >= AIR_EXTENSION_VERSION && Array.isArray(supported) - && supported.includes(AIR_SESSION_FAILURE_KEY); + && supported.includes(capability); +} + +function clientSupportsTypedSessionFailures(capabilities: acp.ClientCapabilities | null): boolean { + return clientSupportsAirCapability(capabilities, AIR_SESSION_FAILURE_KEY); +} + +function clientSupportsCompletionDetails(capabilities: acp.ClientCapabilities | null): boolean { + return clientSupportsAirCapability(capabilities, AIR_COMPLETION_DETAILS_KEY); } interface ActiveAuthState { @@ -307,7 +325,10 @@ export class CodexAcpServer { [JETBRAINS_META_KEY]: { [AIR_META_KEY]: { [AIR_EXTENSION_VERSION_KEY]: AIR_EXTENSION_VERSION, - [AIR_EXTENSION_CAPABILITIES_KEY]: [AIR_SESSION_FAILURE_KEY], + [AIR_EXTENSION_CAPABILITIES_KEY]: [ + AIR_SESSION_FAILURE_KEY, + AIR_COMPLETION_DETAILS_KEY, + ], }, }, }, @@ -2154,6 +2175,8 @@ export class CodexAcpServer { sessionState, eventHandler, commandResult.turnCompleted?.turn.id ?? sessionState.currentTurnId, + false, + commandResult.turnCompleted?.turn, ); if (terminalFailure) { return terminalFailure; @@ -2251,6 +2274,8 @@ export class CodexAcpServer { sessionState, eventHandler, turnCompleted.turn.id, + false, + turnCompleted.turn, ); if (terminalFailure) { return terminalFailure; @@ -2342,6 +2367,8 @@ export class CodexAcpServer { sessionState, eventHandler, turnCompleted.turn.id, + false, + turnCompleted.turn, ); if (implementationFailure) { return implementationFailure; @@ -2474,17 +2501,41 @@ export class CodexAcpServer { eventHandler: CodexEventHandler, turnId: string | null, allowUnattributed = false, + turn?: Turn, ): acp.PromptResponse | null { const failureMeta = eventHandler.getTerminalSessionFailureMeta(turnId, allowUnattributed); if (failureMeta === null) { return null; } + const completionDetails = turnId !== null + && clientSupportsCompletionDetails(this.clientCapabilities) + && (eventHandler.hasEmittedAssistantText() + || turn?.items.some(item => item.type === "agentMessage" && item.text.length > 0)) + ? { + turnId, + partial: true, + retryable: false, + } satisfies CompletionDetails + : null; + const jetbrainsMeta = failureMeta[JETBRAINS_META_KEY] as Record | undefined; + const airMeta = jetbrainsMeta?.[AIR_META_KEY] as Record | undefined; return { stopReason: "end_turn", usage: this.buildPromptUsage(sessionState.lastTokenUsage), _meta: { ...this.buildQuotaMeta(sessionState), ...failureMeta, + ...(completionDetails && jetbrainsMeta && airMeta + ? { + [JETBRAINS_META_KEY]: { + ...jetbrainsMeta, + [AIR_META_KEY]: { + ...airMeta, + [AIR_COMPLETION_DETAILS_KEY]: completionDetails, + }, + }, + } + : {}), }, }; } diff --git a/src/CodexEventHandler.ts b/src/CodexEventHandler.ts index df2ebf95..988af035 100644 --- a/src/CodexEventHandler.ts +++ b/src/CodexEventHandler.ts @@ -164,6 +164,7 @@ export class CodexEventHandler { private readonly terminalCommandOutputIds = new Set(); private readonly agentMessagePhases = new Map(); private readonly activeSubAgentActivities = new Set(); + private emittedAssistantText = false; constructor( connection: AcpClientConnection, @@ -183,6 +184,10 @@ export class CodexEventHandler { return this.failure; } + hasEmittedAssistantText(): boolean { + return this.emittedAssistantText; + } + getTerminalSessionFailureMeta( turnId: string | null, allowUnattributed = false, @@ -482,6 +487,9 @@ export class CodexEventHandler { } private async createTextEvent(event: AgentMessageDeltaNotification): Promise { + if (event.delta.length > 0) { + this.emittedAssistantText = true; + } const phase = this.agentMessagePhases.get(event.itemId) ?? null; return createAgentTextMessageChunk(event.delta, event.itemId, createCodexMessagePhaseMeta(phase)); } @@ -748,6 +756,9 @@ export class CodexEventHandler { } private createPlanTextEvent(text: string, messageId: string): UpdateSessionEvent { + if (text.length > 0) { + this.emittedAssistantText = true; + } return createAgentTextMessageChunk( text, messageId, diff --git a/src/__tests__/CodexACPAgent/auth-error-events.test.ts b/src/__tests__/CodexACPAgent/auth-error-events.test.ts index e49fdafa..a74f3767 100644 --- a/src/__tests__/CodexACPAgent/auth-error-events.test.ts +++ b/src/__tests__/CodexACPAgent/auth-error-events.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it, vi } from "vitest"; import * as acp from "@agentclientprotocol/sdk"; -import type { ErrorNotification, TurnCompletedNotification } from "../../app-server/v2"; +import type { ErrorNotification, Turn, TurnCompletedNotification } from "../../app-server/v2"; import type { SessionState } from "../../CodexAcpServer"; import { createCodexMockTestFixture, @@ -69,7 +69,7 @@ const configuredAuthFailureCases: Array<{ ]; const typedFailureCapabilities: acp.ClientCapabilities = { - _meta: {jetbrains: {air: {version: 1, capabilities: ["sessionFailure"]}}}, + _meta: {jetbrains: {air: {version: 1, capabilities: ["sessionFailure", "completionDetails"]}}}, }; describe("CodexEventHandler - auth error events", () => { @@ -425,6 +425,71 @@ describe("CodexEventHandler - auth error events", () => { expect(JSON.stringify(result)).not.toContain("secret completion details"); }); + it("marks assistant output partial when a typed terminal failure ends the same turn", async () => { + const sessionState = createTestSessionState({ + sessionId: "partial-failed-completion-session", + account: {type: "apiKey"}, + }); + const {result, updates} = await runPromptWithCompletedTurn( + sessionState, + typedFailureCapabilities, + createTurn("failed", "partial-failed-turn", { + message: "usage exhausted", + codexErrorInfo: "usageLimitExceeded", + additionalDetails: null, + }, [{ + type: "agentMessage", + id: "partial-message", + text: "I updated the mapper, but", + phase: "final_answer", + memoryCitation: null, + }]), + ); + + expect(result).toMatchObject({ + stopReason: "end_turn", + _meta: {jetbrains: {air: { + sessionFailure: { + category: "quota_exhausted", + turnId: "partial-failed-turn", + }, + completionDetails: { + turnId: "partial-failed-turn", + partial: true, + retryable: false, + }, + }}}, + }); + expect(updates).toEqual([]); + }); + + it("omits completion details when only session failures were negotiated", async () => { + const sessionState = createTestSessionState({ + sessionId: "session-failure-only-session", + account: {type: "apiKey"}, + }); + const {result} = await runPromptWithCompletedTurn( + sessionState, + {_meta: {jetbrains: {air: {version: 1, capabilities: ["sessionFailure"]}}}}, + createTurn("failed", "session-failure-only-turn", { + message: "usage exhausted", + codexErrorInfo: "usageLimitExceeded", + additionalDetails: null, + }, [{ + type: "agentMessage", + id: "partial-message", + text: "Partial answer", + phase: "final_answer", + memoryCitation: null, + }]), + ); + + expect(result).toMatchObject({ + _meta: {jetbrains: {air: {sessionFailure: {category: "quota_exhausted"}}}}, + }); + expect(JSON.stringify(result)).not.toContain("completionDetails"); + }); + it("publishes a late idle error through the session-scoped subscription", async () => { const mockFixture = createCodexMockTestFixture(); const codexAcpAgent = mockFixture.getCodexAcpAgent(); @@ -889,10 +954,11 @@ function createTurn( status: "inProgress" | "completed" | "failed", id = "turn-id", error: ErrorNotification["error"] | null = null, + items: Turn["items"] = [], ) { return { id, - items: [], + items, itemsView: "notLoaded" as const, status, error, diff --git a/src/__tests__/CodexACPAgent/initialize.test.ts b/src/__tests__/CodexACPAgent/initialize.test.ts index 811cef28..5c4fc8e5 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"], + capabilities: ["sessionFailure", "completionDetails"], }, }, },