diff --git a/apps/code/src/main/services/agent/service.ts b/apps/code/src/main/services/agent/service.ts index e151c2c29..67e17dcbc 100644 --- a/apps/code/src/main/services/agent/service.ts +++ b/apps/code/src/main/services/agent/service.ts @@ -551,11 +551,18 @@ When creating pull requests, add the following footer at the end of the PR descr }); try { + const systemPrompt = this.buildSystemPrompt( + credentials, + taskId, + customInstructions, + ); + const acpConnection = await agent.run(taskId, taskRunId, { adapter, gatewayUrl: proxyUrl, codexBinaryPath: adapter === "codex" ? getCodexBinaryPath() : undefined, model, + instructions: adapter === "codex" ? systemPrompt.append : undefined, processCallbacks: { onProcessSpawned: (info) => { this.processTracking.register( @@ -656,12 +663,6 @@ When creating pull requests, add the following footer at the end of the PR descr } } - const systemPrompt = this.buildSystemPrompt( - credentials, - taskId, - customInstructions, - ); - // Both adapters implement unstable_resumeSession: // - Claude: delegates to SDK's resumeSession with JSONL hydration // - Codex: delegates to codex-acp's loadSession internally @@ -698,11 +699,6 @@ When creating pull requests, add the following footer at the end of the PR descr taskRunId, }); } - const systemPrompt = this.buildSystemPrompt( - credentials, - taskId, - customInstructions, - ); const newSessionResponse = await connection.newSession({ cwd: repoPath, mcpServers, diff --git a/packages/agent/src/adapters/codex/codex-agent.ts b/packages/agent/src/adapters/codex/codex-agent.ts index b939e52d3..5d02be28e 100644 --- a/packages/agent/src/adapters/codex/codex-agent.ts +++ b/packages/agent/src/adapters/codex/codex-agent.ts @@ -282,18 +282,49 @@ export class CodexAcpAgent extends BaseAcpAgent { const response = await this.codexConnection.prompt(params); - // Emit PostHog usage notification - if (this.sessionState?.taskRunId && response.usage) { - await this.client.extNotification("_posthog/usage_update", { + 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; + } + + if (this.sessionState?.taskRunId) { + const { accumulatedUsage } = this.sessionState; + + await this.client.extNotification(POSTHOG_NOTIFICATIONS.TURN_COMPLETE, { sessionId: params.sessionId, - used: { - inputTokens: response.usage.inputTokens ?? 0, - outputTokens: response.usage.outputTokens ?? 0, - cachedReadTokens: response.usage.cachedReadTokens ?? 0, - cachedWriteTokens: response.usage.cachedWriteTokens ?? 0, + stopReason: response.stopReason ?? "end_turn", + usage: { + inputTokens: accumulatedUsage.inputTokens, + outputTokens: accumulatedUsage.outputTokens, + cachedReadTokens: accumulatedUsage.cachedReadTokens, + cachedWriteTokens: accumulatedUsage.cachedWriteTokens, + totalTokens: + accumulatedUsage.inputTokens + + accumulatedUsage.outputTokens + + accumulatedUsage.cachedReadTokens + + accumulatedUsage.cachedWriteTokens, }, - cost: null, }); + + if (response.usage) { + await this.client.extNotification("_posthog/usage_update", { + sessionId: params.sessionId, + used: { + inputTokens: response.usage.inputTokens ?? 0, + outputTokens: response.usage.outputTokens ?? 0, + cachedReadTokens: response.usage.cachedReadTokens ?? 0, + cachedWriteTokens: response.usage.cachedWriteTokens ?? 0, + }, + cost: null, + }); + } } return response; diff --git a/packages/agent/src/adapters/codex/codex-client.ts b/packages/agent/src/adapters/codex/codex-client.ts index f2966fe89..6306cf742 100644 --- a/packages/agent/src/adapters/codex/codex-client.ts +++ b/packages/agent/src/adapters/codex/codex-client.ts @@ -60,17 +60,34 @@ export function createCodexClient( }, async sessionUpdate(params: SessionNotification): Promise { - // Parse usage data from session updates const update = params.update as Record | undefined; if (update?.sessionUpdate === "usage_update") { const used = update.used as number | undefined; const size = update.size as number | undefined; if (used !== undefined) sessionState.contextUsed = used; if (size !== undefined) sessionState.contextSize = size; + + // Accumulate per-message token usage when available + const inputTokens = update.inputTokens as number | undefined; + const outputTokens = update.outputTokens as number | undefined; + if (inputTokens !== undefined) { + sessionState.accumulatedUsage.inputTokens += inputTokens; + } + if (outputTokens !== undefined) { + sessionState.accumulatedUsage.outputTokens += outputTokens; + } + const cachedRead = update.cachedReadTokens as number | undefined; + const cachedWrite = update.cachedWriteTokens as number | undefined; + if (cachedRead !== undefined) { + sessionState.accumulatedUsage.cachedReadTokens += cachedRead; + } + if (cachedWrite !== undefined) { + sessionState.accumulatedUsage.cachedWriteTokens += cachedWrite; + } + callbacks?.onUsageUpdate?.(update); } - // Forward to upstream client await upstreamClient.sessionUpdate(params); }, diff --git a/packages/agent/src/adapters/codex/spawn.ts b/packages/agent/src/adapters/codex/spawn.ts index 1cbb0aa03..1053d4c71 100644 --- a/packages/agent/src/adapters/codex/spawn.ts +++ b/packages/agent/src/adapters/codex/spawn.ts @@ -10,6 +10,7 @@ export interface CodexProcessOptions { apiBaseUrl?: string; apiKey?: string; model?: string; + instructions?: string; binaryPath?: string; logger?: Logger; processCallbacks?: ProcessSpawnedCallback; @@ -42,6 +43,13 @@ function buildConfigArgs(options: CodexProcessOptions): string[] { args.push("-c", `model="${options.model}"`); } + if (options.instructions) { + const escaped = options.instructions + .replace(/\\/g, "\\\\") + .replace(/"/g, '\\"'); + args.push("-c", `instructions="${escaped}"`); + } + return args; } diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index b2129e630..2ce00ae38 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -131,6 +131,7 @@ export class Agent { apiKey: gatewayConfig.apiKey, binaryPath: options.codexBinaryPath, model: sanitizedModel, + instructions: options.instructions, } : undefined, }); diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index b463189e7..3cc039d6a 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -112,6 +112,7 @@ export interface TaskExecutionOptions { model?: string; gatewayUrl?: string; codexBinaryPath?: string; + instructions?: string; processCallbacks?: ProcessSpawnedCallback; }