Skip to content
Draft
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
32 changes: 32 additions & 0 deletions packages/agent/src/adapters/codex/codex-agent.refresh.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,7 @@ type PrivateAgent = {
promptRunning: boolean;
};
sessionId: string;
currentMcpServers: McpServer[];
sessionState: {
sessionId: string;
cwd: string;
Expand All @@ -120,6 +121,7 @@ type PrivateAgent = {
cachedWriteTokens: number;
};
configOptions: unknown[];
serviceTier?: string;
taskRunId?: string;
};
codexProcess: SpawnHandle;
Expand Down Expand Up @@ -286,4 +288,34 @@ describe("CodexAcpAgent.extMethod refresh_session", () => {
expect(spawnedProcesses).toHaveLength(2);
expect(createdConnections[1]?.loadSession).toHaveBeenCalled();
});

it("applies service tier changes by respawning with the current MCP servers", async () => {
const agent = makeAgent();
const { priv } = primeSession(agent, "s-4");
const mcpServers: McpServer[] = [
{ name: "posthog", type: "http", url: "https://mcp", headers: [] },
];
priv.currentMcpServers = mcpServers;

const response = await agent.setSessionConfigOption({
sessionId: "s-4",
configId: "service_tier",
value: "fast",
} as never);

expect(hoisted.spawnCodexProcessMock).toHaveBeenLastCalledWith(
expect.objectContaining({ serviceTier: "fast" }),
);
expect(createdConnections[1]?.loadSession).toHaveBeenCalledWith({
sessionId: "s-4",
cwd: "/tmp/repo",
mcpServers,
});
expect(response.configOptions).toContainEqual(
expect.objectContaining({
id: "service_tier",
currentValue: "fast",
}),
);
});
});
27 changes: 27 additions & 0 deletions packages/agent/src/adapters/codex/codex-agent.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import type {
NewSessionResponse,
} from "@agentclientprotocol/sdk";
import { beforeEach, describe, expect, it, vi } from "vitest";
import type { CodexProcessOptions } from "./spawn";

const mockCodexConnection = {
initialize: vi.fn(),
Expand Down Expand Up @@ -69,6 +70,7 @@ describe("CodexAcpAgent", () => {
overrides: Partial<AgentSideConnection> = {},
agentOptions?: {
onStructuredOutput?: (output: Record<string, unknown>) => Promise<void>;
codexProcessOptions?: Partial<CodexProcessOptions>;
},
): {
agent: CodexAcpAgent;
Expand All @@ -89,6 +91,7 @@ describe("CodexAcpAgent", () => {
const agent = new CodexAcpAgent(client, {
codexProcessOptions: {
cwd: process.cwd(),
...agentOptions?.codexProcessOptions,
},
onStructuredOutput: agentOptions?.onStructuredOutput,
});
Expand Down Expand Up @@ -118,6 +121,30 @@ describe("CodexAcpAgent", () => {
).toBe("read-only");
});

it("adds the Codex service tier option to new sessions", async () => {
const { agent } = createAgent(
{},
{ codexProcessOptions: { serviceTier: "fast" } },
);
mockCodexConnection.newSession.mockResolvedValue({
sessionId: "session-1",
modes: { currentModeId: "auto", availableModes: [] },
configOptions: [],
} satisfies Partial<NewSessionResponse>);

const response = await agent.newSession({
cwd: process.cwd(),
} as never);

expect(response.configOptions).toContainEqual(
expect.objectContaining({
id: "service_tier",
category: "service_tier",
currentValue: "fast",
}),
);
});

it("propagates taskRunId and fires SDK_SESSION when loading a cloud session", async () => {
const { agent, client } = createAgent();
mockCodexConnection.loadSession.mockResolvedValue({
Expand Down
148 changes: 143 additions & 5 deletions packages/agent/src/adapters/codex/codex-agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ import {
RequestError,
type ResumeSessionRequest,
type ResumeSessionResponse,
type SessionConfigOption,
type SetSessionConfigOptionRequest,
type SetSessionConfigOptionResponse,
type SetSessionModeRequest,
Expand Down Expand Up @@ -123,6 +124,31 @@ interface NewSessionMeta {
jsonSchema?: Record<string, unknown> | null;
}

type CodexServiceTier = "standard" | "fast" | "flex";

const CODEX_SERVICE_TIER_CONFIG_ID = "service_tier";
const CODEX_SERVICE_TIER_OPTIONS: Array<{
value: CodexServiceTier;
name: string;
description: string;
}> = [
{
value: "standard",
name: "Standard",
description: "Default Codex service tier",
},
{
value: "fast",
name: "Fast",
description: "Request Codex fast mode for lower latency",
},
{
value: "flex",
name: "Flex",
description: "Request Codex flex mode",
},
];

export interface CodexAcpAgentOptions {
codexProcessOptions: CodexProcessOptions;
processCallbacks?: ProcessSpawnedCallback;
Expand All @@ -142,6 +168,38 @@ function toCodexPermissionMode(mode?: string): PermissionMode {
return "auto";
}

function toCodexServiceTier(serviceTier?: string | null): CodexServiceTier {
switch (serviceTier?.toLowerCase()) {
case "fast":
case "priority":
return "fast";
case "flex":
return "flex";
default:
return "standard";
}
}

function codexServiceTierConfigValue(
serviceTier: CodexServiceTier,
): string | undefined {
return serviceTier === "standard" ? undefined : serviceTier;
}

function codexServiceTierConfigOption(
currentValue: CodexServiceTier,
): SessionConfigOption {
return {
id: CODEX_SERVICE_TIER_CONFIG_ID,
name: "Speed",
type: "select",
category: "service_tier",
currentValue,
options: CODEX_SERVICE_TIER_OPTIONS,
description: "Choose the Codex service tier for new turns",
} as SessionConfigOption;
}

/**
* Prepend `_meta.prContext` (set by the agent-server on Slack-originated
* follow-up runs) to the prompt as a text block, mirroring Claude's
Expand Down Expand Up @@ -313,6 +371,7 @@ export class CodexAcpAgent extends BaseAcpAgent {
private readonly onStructuredOutput?: (
output: Record<string, unknown>,
) => Promise<void>;
private currentMcpServers: McpServer[] = [];
// Snapshot of the initialize() request so refreshSession can replay the
// same handshake against a respawned codex-acp subprocess.
private lastInitRequest?: InitializeRequest;
Expand Down Expand Up @@ -413,9 +472,10 @@ export class CodexAcpAgent extends BaseAcpAgent {
meta,
);
const response = await this.codexConnection.newSession(injectedParams);
response.configOptions = normalizeCodexConfigOptions(
response.configOptions = this.normalizeConfigOptions(
response.configOptions,
);
this.currentMcpServers = injectedParams.mcpServers ?? [];

// Initialize session state. Mutate in place — codex-client closure-
// captured this object in the constructor and writes contextUsed/
Expand All @@ -426,6 +486,7 @@ export class CodexAcpAgent extends BaseAcpAgent {
modeId: response.modes?.currentModeId ?? "auto",
modelId: modelIdFromConfigOptions(response.configOptions),
permissionMode: requestedPermissionMode,
serviceTier: this.currentServiceTier(),
});
this.sessionId = response.sessionId;
this.sessionState.configOptions = response.configOptions ?? [];
Expand Down Expand Up @@ -461,9 +522,10 @@ export class CodexAcpAgent extends BaseAcpAgent {
meta,
);
const response = await this.codexConnection.loadSession(injectedParams);
response.configOptions = normalizeCodexConfigOptions(
response.configOptions = this.normalizeConfigOptions(
response.configOptions,
);
this.currentMcpServers = injectedParams.mcpServers ?? [];
const currentPermissionMode = getCurrentPermissionMode(
response.modes?.currentModeId,
meta?.permissionMode,
Expand All @@ -478,6 +540,7 @@ export class CodexAcpAgent extends BaseAcpAgent {
taskId: resolveTaskId(meta),
modeId: response.modes?.currentModeId ?? "auto",
permissionMode: currentPermissionMode,
serviceTier: this.currentServiceTier(),
});
this.sessionId = params.sessionId;
this.sessionState.configOptions = response.configOptions ?? [];
Expand Down Expand Up @@ -513,9 +576,10 @@ export class CodexAcpAgent extends BaseAcpAgent {

// codex-acp doesn't support resume natively, use loadSession instead
const loadResponse = await this.codexConnection.loadSession(injectedParams);
loadResponse.configOptions = normalizeCodexConfigOptions(
loadResponse.configOptions = this.normalizeConfigOptions(
loadResponse.configOptions,
);
this.currentMcpServers = injectedParams.mcpServers ?? [];
const currentPermissionMode = getCurrentPermissionMode(
loadResponse.modes?.currentModeId,
meta?.permissionMode,
Expand All @@ -525,6 +589,7 @@ export class CodexAcpAgent extends BaseAcpAgent {
taskId: resolveTaskId(meta),
modeId: loadResponse.modes?.currentModeId ?? "auto",
permissionMode: currentPermissionMode,
serviceTier: this.currentServiceTier(),
});
this.sessionId = params.sessionId;
this.sessionState.configOptions = loadResponse.configOptions ?? [];
Expand Down Expand Up @@ -562,16 +627,18 @@ export class CodexAcpAgent extends BaseAcpAgent {

// Create a new session via codex-acp (fork isn't natively supported)
const newResponse = await this.codexConnection.newSession(injectedParams);
newResponse.configOptions = normalizeCodexConfigOptions(
newResponse.configOptions = this.normalizeConfigOptions(
newResponse.configOptions,
);
this.currentMcpServers = injectedParams.mcpServers ?? [];

const requestedPermissionMode = toCodexPermissionMode(meta?.permissionMode);
resetSessionState(this.sessionState, newResponse.sessionId, params.cwd, {
taskRunId: meta?.taskRunId,
taskId: resolveTaskId(meta),
modeId: newResponse.modes?.currentModeId ?? "auto",
permissionMode: requestedPermissionMode,
serviceTier: this.currentServiceTier(),
});
this.sessionId = newResponse.sessionId;
this.sessionState.configOptions = newResponse.configOptions ?? [];
Expand All @@ -586,6 +653,38 @@ export class CodexAcpAgent extends BaseAcpAgent {
return newResponse;
}

private currentServiceTier(): CodexServiceTier {
return toCodexServiceTier(this.codexProcessOptions.serviceTier);
}

private withCodexServiceTierOption(
configOptions: SessionConfigOption[] | null | undefined,
): SessionConfigOption[] {
const serviceTierOption = codexServiceTierConfigOption(
this.currentServiceTier(),
);
const options = configOptions ?? [];
const existingIndex = options.findIndex(
(option) => option.id === CODEX_SERVICE_TIER_CONFIG_ID,
);

if (existingIndex === -1) {
return [...options, serviceTierOption];
}

return options.map((option, index) =>
index === existingIndex ? serviceTierOption : option,
);
}

private normalizeConfigOptions(
configOptions: SessionConfigOption[] | null | undefined,
): SessionConfigOption[] {
return this.withCodexServiceTierOption(
normalizeCodexConfigOptions(configOptions) ?? [],
);
}

/**
* When the caller wires up `onStructuredOutput` and provides a JSON schema
* via `_meta.jsonSchema`, inject the stdio MCP server that exposes
Expand Down Expand Up @@ -909,6 +1008,7 @@ export class CodexAcpAgent extends BaseAcpAgent {
cwd: this.sessionState.cwd,
mcpServers,
});
this.currentMcpServers = mcpServers;

// Swap everything at once so closeSession/prompt/cancel target the new
// subprocess going forward. Preserve sessionState (accumulatedUsage,
Expand Down Expand Up @@ -936,12 +1036,50 @@ export class CodexAcpAgent extends BaseAcpAgent {
return response ?? {};
}

private async setServiceTier(value: string): Promise<SessionConfigOption[]> {
const previousTier = this.currentServiceTier();
const nextTier = toCodexServiceTier(value);
if (previousTier === nextTier) {
this.sessionState.serviceTier = nextTier;
this.sessionState.configOptions = this.withCodexServiceTierOption(
this.sessionState.configOptions,
);
return this.sessionState.configOptions;
}

const previousConfigValue = this.codexProcessOptions.serviceTier;
this.codexProcessOptions.serviceTier =
codexServiceTierConfigValue(nextTier);

try {
await this.refreshSession(this.currentMcpServers);
} catch (error) {
this.codexProcessOptions.serviceTier = previousConfigValue;
this.sessionState.serviceTier = previousTier;
this.sessionState.configOptions = this.withCodexServiceTierOption(
this.sessionState.configOptions,
);
throw error;
}

this.sessionState.serviceTier = nextTier;
this.sessionState.configOptions = this.withCodexServiceTierOption(
this.sessionState.configOptions,
Comment on lines +1057 to +1067

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Duplicate "priority" alias in two normalizers

toCodexServiceTier here and normalizeServiceTier in spawn.ts both independently map "priority""fast". If a new alias or tier is added, both functions must be updated in sync. The canonical mapping belongs in one place — the spawn.ts normalizer that actually writes the flag to the process args is the logical owner; toCodexServiceTier could delegate to it (or to a shared helper) rather than re-implementing the same switch.

Prompt To Fix With AI
This is a comment left during a code review.
Path: packages/agent/src/adapters/codex/codex-agent.ts
Line: 899-909

Comment:
**Duplicate "priority" alias in two normalizers**

`toCodexServiceTier` here and `normalizeServiceTier` in `spawn.ts` both independently map `"priority"``"fast"`. If a new alias or tier is added, both functions must be updated in sync. The canonical mapping belongs in one place — the `spawn.ts` normalizer that actually writes the flag to the process args is the logical owner; `toCodexServiceTier` could delegate to it (or to a shared helper) rather than re-implementing the same switch.

How can I resolve this? If you propose a fix, please make it concise.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

);
return this.sessionState.configOptions;
}

async setSessionConfigOption(
params: SetSessionConfigOptionRequest,
): Promise<SetSessionConfigOptionResponse> {
if (params.configId === CODEX_SERVICE_TIER_CONFIG_ID) {
const configOptions = await this.setServiceTier(String(params.value));
return { configOptions };
}

const response = await this.codexConnection.setSessionConfigOption(params);
if (response.configOptions) {
response.configOptions = normalizeCodexConfigOptions(
response.configOptions = this.normalizeConfigOptions(
response.configOptions,
) as typeof response.configOptions;
this.sessionState.configOptions = response.configOptions;
Expand Down
Loading
Loading