Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 7 additions & 11 deletions apps/code/src/main/services/agent/service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
49 changes: 40 additions & 9 deletions packages/agent/src/adapters/codex/codex-agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
21 changes: 19 additions & 2 deletions packages/agent/src/adapters/codex/codex-client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -60,17 +60,34 @@ export function createCodexClient(
},

async sessionUpdate(params: SessionNotification): Promise<void> {
// Parse usage data from session updates
const update = params.update as Record<string, unknown> | 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);
},

Expand Down
8 changes: 8 additions & 0 deletions packages/agent/src/adapters/codex/spawn.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ export interface CodexProcessOptions {
apiBaseUrl?: string;
apiKey?: string;
model?: string;
instructions?: string;
binaryPath?: string;
logger?: Logger;
processCallbacks?: ProcessSpawnedCallback;
Expand Down Expand Up @@ -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;
}

Expand Down
1 change: 1 addition & 0 deletions packages/agent/src/agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,7 @@ export class Agent {
apiKey: gatewayConfig.apiKey,
binaryPath: options.codexBinaryPath,
model: sanitizedModel,
instructions: options.instructions,
}
: undefined,
});
Expand Down
1 change: 1 addition & 0 deletions packages/agent/src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,7 @@ export interface TaskExecutionOptions {
model?: string;
gatewayUrl?: string;
codexBinaryPath?: string;
instructions?: string;
processCallbacks?: ProcessSpawnedCallback;
}

Expand Down
Loading