From ad0c6deff526b2e351e6488577d3b3b31d260569 Mon Sep 17 00:00:00 2001 From: Charles Vien Date: Fri, 3 Apr 2026 08:20:35 -0700 Subject: [PATCH] Harden CodexAcpAgent session state and permission handling --- packages/agent/src/acp-extensions.ts | 3 + .../agent/src/adapters/claude/claude-agent.ts | 22 +++-- .../agent/src/adapters/codex/codex-agent.ts | 97 +++++++------------ .../agent/src/adapters/codex/session-state.ts | 3 - packages/agent/src/adapters/codex/spawn.ts | 2 + 5 files changed, 52 insertions(+), 75 deletions(-) diff --git a/packages/agent/src/acp-extensions.ts b/packages/agent/src/acp-extensions.ts index bce33ae45..ef43cce4c 100644 --- a/packages/agent/src/acp-extensions.ts +++ b/packages/agent/src/acp-extensions.ts @@ -64,4 +64,7 @@ export const POSTHOG_NOTIFICATIONS = { /** Marks a boundary for log compaction */ COMPACT_BOUNDARY: "_posthog/compact_boundary", + + /** Token usage update for a session turn */ + USAGE_UPDATE: "_posthog/usage_update", } as const; diff --git a/packages/agent/src/adapters/claude/claude-agent.ts b/packages/agent/src/adapters/claude/claude-agent.ts index 2aafaaac4..4679ae2bd 100644 --- a/packages/agent/src/adapters/claude/claude-agent.ts +++ b/packages/agent/src/adapters/claude/claude-agent.ts @@ -44,6 +44,7 @@ import { } from "@anthropic-ai/claude-agent-sdk"; import { v7 as uuidv7 } from "uuid"; import packageJson from "../../../package.json" with { type: "json" }; +import { POSTHOG_NOTIFICATIONS } from "../../acp-extensions"; import { unreachable, withTimeout } from "../../utils/common"; import { Logger } from "../../utils/logger"; import { Pushable } from "../../utils/streams"; @@ -442,16 +443,19 @@ export class ClaudeAcpAgent extends BaseAcpAgent { }); } - await this.client.extNotification("_posthog/usage_update", { - sessionId: params.sessionId, - used: { - inputTokens: message.usage.input_tokens, - outputTokens: message.usage.output_tokens, - cachedReadTokens: message.usage.cache_read_input_tokens, - cachedWriteTokens: message.usage.cache_creation_input_tokens, + await this.client.extNotification( + POSTHOG_NOTIFICATIONS.USAGE_UPDATE, + { + sessionId: params.sessionId, + used: { + inputTokens: message.usage.input_tokens, + outputTokens: message.usage.output_tokens, + cachedReadTokens: message.usage.cache_read_input_tokens, + cachedWriteTokens: message.usage.cache_creation_input_tokens, + }, + cost: message.total_cost_usd, }, - cost: message.total_cost_usd, - }); + ); const usage: Usage = { inputTokens: this.session.accumulatedUsage.inputTokens, diff --git a/packages/agent/src/adapters/codex/codex-agent.ts b/packages/agent/src/adapters/codex/codex-agent.ts index ed7735c96..f52194b3d 100644 --- a/packages/agent/src/adapters/codex/codex-agent.ts +++ b/packages/agent/src/adapters/codex/codex-agent.ts @@ -12,7 +12,6 @@ import { type AgentSideConnection, type AuthenticateRequest, - type CancelNotification, ClientSideConnection, type ForkSessionRequest, type ForkSessionResponse, @@ -102,8 +101,8 @@ export class CodexAcpAgent extends BaseAcpAgent { readonly adapterName = "codex"; declare session: CodexSession; private codexProcess: CodexProcess; - private codexConnection!: ClientSideConnection; - private sessionState!: CodexSessionState; + private codexConnection: ClientSideConnection; + private sessionState: CodexSessionState; constructor(client: AgentSideConnection, options: CodexAcpAgentOptions) { super(client); @@ -132,28 +131,14 @@ export class CodexAcpAgent extends BaseAcpAgent { cancelled: false, }; + this.sessionState = createSessionState("", cwd); + // Create the ClientSideConnection to codex-acp. // The Client handler delegates all requests from codex-acp to the upstream // PostHog Code client via our AgentSideConnection. this.codexConnection = new ClientSideConnection( (_agent) => - createCodexClient( - this.client, - this.logger, - this.sessionState ?? { - sessionId: "", - cwd: "", - modeId: "default", - configOptions: [], - accumulatedUsage: { - inputTokens: 0, - outputTokens: 0, - cachedReadTokens: 0, - cachedWriteTokens: 0, - }, - cancelled: false, - }, - ), + createCodexClient(this.client, this.logger, this.sessionState), codexStream, ); } @@ -224,9 +209,11 @@ export class CodexAcpAgent extends BaseAcpAgent { async loadSession(params: LoadSessionRequest): Promise { const response = await this.codexConnection.loadSession(params); + const meta = params._meta as NewSessionMeta | undefined; - // Update session state - this.sessionState = createSessionState(params.sessionId, params.cwd); + this.sessionState = createSessionState(params.sessionId, params.cwd, { + permissionMode: toCodeExecutionMode(meta?.permissionMode), + }); this.sessionId = params.sessionId; this.sessionState.configOptions = response.configOptions ?? []; @@ -243,11 +230,15 @@ export class CodexAcpAgent extends BaseAcpAgent { mcpServers: params.mcpServers ?? [], }); - this.sessionState = createSessionState(params.sessionId, params.cwd); + const meta = params._meta as NewSessionMeta | undefined; + this.sessionState = createSessionState(params.sessionId, params.cwd, { + taskRunId: meta?.taskRunId, + taskId: meta?.taskId ?? meta?.persistence?.taskId, + permissionMode: toCodeExecutionMode(meta?.permissionMode), + }); this.sessionId = params.sessionId; this.sessionState.configOptions = loadResponse.configOptions ?? []; - const meta = params._meta as NewSessionMeta | undefined; if (meta?.taskRunId) { await this.client.extNotification(POSTHOG_NOTIFICATIONS.SDK_SESSION, { taskRunId: meta.taskRunId, @@ -273,7 +264,12 @@ export class CodexAcpAgent extends BaseAcpAgent { _meta: params._meta, }); - this.sessionState = createSessionState(newResponse.sessionId, params.cwd); + const meta = params._meta as NewSessionMeta | undefined; + this.sessionState = createSessionState(newResponse.sessionId, params.cwd, { + taskRunId: meta?.taskRunId, + taskId: meta?.taskId ?? meta?.persistence?.taskId, + permissionMode: toCodeExecutionMode(meta?.permissionMode), + }); this.sessionId = newResponse.sessionId; this.sessionState.configOptions = newResponse.configOptions ?? []; @@ -289,31 +285,21 @@ export class CodexAcpAgent extends BaseAcpAgent { async unstable_listSessions( params: ListSessionsRequest, ): Promise { - return this.codexConnection.listSessions(params); + return this.listSessions(params); } async prompt(params: PromptRequest): Promise { - if (this.sessionState) { - this.sessionState.cancelled = false; - this.sessionState.interruptReason = undefined; - resetUsage(this.sessionState); - } + this.session.cancelled = false; + this.session.interruptReason = undefined; + resetUsage(this.sessionState); const response = await this.codexConnection.prompt(params); - if (this.sessionState && response.usage) { - // Accumulate token usage from the prompt response - this.sessionState.accumulatedUsage.inputTokens += - response.usage.inputTokens ?? 0; - this.sessionState.accumulatedUsage.outputTokens += - response.usage.outputTokens ?? 0; - this.sessionState.accumulatedUsage.cachedReadTokens += - response.usage.cachedReadTokens ?? 0; - this.sessionState.accumulatedUsage.cachedWriteTokens += - response.usage.cachedWriteTokens ?? 0; - } + // Usage is already accumulated via sessionUpdate notifications in + // codex-client.ts. Do NOT also add response.usage here or tokens + // get double-counted. - if (this.sessionState?.taskRunId) { + if (this.sessionState.taskRunId) { const { accumulatedUsage } = this.sessionState; await this.client.extNotification(POSTHOG_NOTIFICATIONS.TURN_COMPLETE, { @@ -333,7 +319,7 @@ export class CodexAcpAgent extends BaseAcpAgent { }); if (response.usage) { - await this.client.extNotification("_posthog/usage_update", { + await this.client.extNotification(POSTHOG_NOTIFICATIONS.USAGE_UPDATE, { sessionId: params.sessionId, used: { inputTokens: response.usage.inputTokens ?? 0, @@ -350,25 +336,11 @@ export class CodexAcpAgent extends BaseAcpAgent { } protected async interrupt(): Promise { - if (this.sessionState) { - this.sessionState.cancelled = true; - } await this.codexConnection.cancel({ sessionId: this.sessionId, }); } - async cancel(params: CancelNotification): Promise { - if (this.sessionState) { - this.sessionState.cancelled = true; - const meta = params._meta as { interruptReason?: string } | undefined; - if (meta?.interruptReason) { - this.sessionState.interruptReason = meta.interruptReason; - } - } - await this.codexConnection.cancel(params); - } - async setSessionMode( params: SetSessionModeRequest, ): Promise { @@ -380,10 +352,8 @@ export class CodexAcpAgent extends BaseAcpAgent { modeId: nativeMode, }); - if (this.sessionState) { - this.sessionState.modeId = nativeMode; - this.sessionState.permissionMode = requestedMode; - } + this.sessionState.modeId = nativeMode; + this.sessionState.permissionMode = requestedMode; return response ?? {}; } @@ -391,7 +361,7 @@ export class CodexAcpAgent extends BaseAcpAgent { params: SetSessionConfigOptionRequest, ): Promise { const response = await this.codexConnection.setSessionConfigOption(params); - if (this.sessionState && response.configOptions) { + if (response.configOptions) { this.sessionState.configOptions = response.configOptions; } return response; @@ -403,6 +373,7 @@ export class CodexAcpAgent extends BaseAcpAgent { async closeSession(): Promise { this.logger.info("Closing Codex session", { sessionId: this.sessionId }); + this.session.abortController.abort(); this.session.settingsManager.dispose(); try { this.codexProcess.kill(); diff --git a/packages/agent/src/adapters/codex/session-state.ts b/packages/agent/src/adapters/codex/session-state.ts index 6ba83c8fd..dcc7203ac 100644 --- a/packages/agent/src/adapters/codex/session-state.ts +++ b/packages/agent/src/adapters/codex/session-state.ts @@ -23,8 +23,6 @@ export interface CodexSessionState { contextSize?: number; contextUsed?: number; permissionMode: CodeExecutionMode; - cancelled: boolean; - interruptReason?: string; taskRunId?: string; taskId?: string; } @@ -53,7 +51,6 @@ export function createSessionState( cachedWriteTokens: 0, }, permissionMode: opts?.permissionMode ?? "default", - cancelled: false, taskRunId: opts?.taskRunId, taskId: opts?.taskId, }; diff --git a/packages/agent/src/adapters/codex/spawn.ts b/packages/agent/src/adapters/codex/spawn.ts index 1053d4c71..3bc0b53e6 100644 --- a/packages/agent/src/adapters/codex/spawn.ts +++ b/packages/agent/src/adapters/codex/spawn.ts @@ -46,6 +46,8 @@ function buildConfigArgs(options: CodexProcessOptions): string[] { if (options.instructions) { const escaped = options.instructions .replace(/\\/g, "\\\\") + .replace(/\n/g, "\\n") + .replace(/\r/g, "\\r") .replace(/"/g, '\\"'); args.push("-c", `instructions="${escaped}"`); }