Skip to content

Commit 4630cd7

Browse files
NikolaiSviridovNikolaiSviridov
committed
feat: LLM-26878 Add “Fast” mode for Codex-agent via session config option
1 parent bf71f6a commit 4630cd7

6 files changed

Lines changed: 322 additions & 5 deletions

File tree

‎src/CodexAcpClient.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import type {Disposable} from "vscode-jsonrpc";
1313
import type {
1414
ClientInfo,
1515
ReasoningEffort,
16+
ServiceTier,
1617
ServerNotification
1718
} from "./app-server";
1819
import type {JsonValue} from "./app-server/serde_json/JsonValue";
@@ -220,6 +221,7 @@ export class CodexAcpClient {
220221
sessionId: request.sessionId,
221222
currentModelId: currentModelId,
222223
models: codexModels,
224+
currentServiceTier: response.serviceTier ?? null,
223225
}
224226
}
225227

@@ -242,6 +244,7 @@ export class CodexAcpClient {
242244
sessionId: request.sessionId,
243245
currentModelId: currentModelId,
244246
models: codexModels,
247+
currentServiceTier: response.serviceTier ?? null,
245248
thread: response.thread,
246249
};
247250
}
@@ -271,6 +274,7 @@ export class CodexAcpClient {
271274
sessionId: response.thread.id,
272275
currentModelId: currentModelId,
273276
models: codexModels,
277+
currentServiceTier: response.serviceTier ?? null,
274278
};
275279
}
276280

@@ -394,6 +398,7 @@ export class CodexAcpClient {
394398
request: acp.PromptRequest,
395399
agentMode: AgentMode,
396400
modelId: ModelId,
401+
serviceTier: ServiceTier | null,
397402
disableSummary: boolean,
398403
cwd: string,
399404
): Promise<TurnCompletedNotification> {
@@ -412,6 +417,7 @@ export class CodexAcpClient {
412417
cwd: null,
413418
effort: effort,
414419
model: modelId.model,
420+
serviceTier: serviceTier,
415421
});
416422
}
417423

@@ -605,6 +611,7 @@ export type SessionMetadata = {
605611
sessionId: string,
606612
currentModelId: string,
607613
models: Model[],
614+
currentServiceTier?: ServiceTier | null,
608615
}
609616

610617
export type SessionMetadataWithThread = SessionMetadata & {

‎src/CodexAcpServer.ts‎

Lines changed: 56 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,14 @@ import {
3232
createFileChangeUpdate,
3333
createMcpToolCallUpdate,
3434
} from "./CodexToolCallMapper";
35+
import {
36+
createFastModeConfigOption,
37+
FAST_MODE_CONFIG_ID,
38+
FAST_MODE_OFF,
39+
FAST_MODE_ON,
40+
modelSupportsFast,
41+
resolveFastServiceTier,
42+
} from "./FastModeConfig";
3543

3644
export interface SessionState {
3745
sessionId: string,
@@ -46,6 +54,8 @@ export interface SessionState {
4654
rateLimits: RateLimitsMap | null;
4755
account: Account | null;
4856
cwd: string;
57+
fastModeEnabled: boolean;
58+
currentModelSupportsFast: boolean;
4959
sessionMcpServers?: Array<string>;
5060
}
5161

@@ -163,6 +173,7 @@ export class CodexAcpServer implements acp.Agent {
163173
const {sessionId, currentModelId, models} = sessionMetadata;
164174
const sessionMcpServers = this.resolveSessionMcpServers(requestedMcpServers, "sessionId" in request);
165175
const currentModel = this.findCurrentModel(models, currentModelId);
176+
const currentModelSupportsFast = modelSupportsFast(currentModel);
166177
const sessionState: SessionState = {
167178
sessionId: sessionId,
168179
currentModelId: currentModelId,
@@ -176,6 +187,8 @@ export class CodexAcpServer implements acp.Agent {
176187
rateLimits: null,
177188
account: account,
178189
cwd: request.cwd,
190+
fastModeEnabled: sessionMetadata.currentServiceTier === "fast",
191+
currentModelSupportsFast: currentModelSupportsFast,
179192
sessionMcpServers: sessionMcpServers,
180193
}
181194
this.sessions.set(sessionId, sessionState);
@@ -221,7 +234,8 @@ export class CodexAcpServer implements acp.Agent {
221234
});
222235
return {
223236
models: modelState,
224-
modes: modeState
237+
modes: modeState,
238+
configOptions: this.createSessionConfigOptions(this.getSessionState(sessionId)),
225239
};
226240
}
227241

@@ -236,7 +250,8 @@ export class CodexAcpServer implements acp.Agent {
236250
});
237251
return {
238252
models: modelState,
239-
modes: modeState
253+
modes: modeState,
254+
configOptions: this.createSessionConfigOptions(this.getSessionState(sessionId)),
240255
};
241256
}
242257

@@ -261,7 +276,8 @@ export class CodexAcpServer implements acp.Agent {
261276
return {
262277
sessionId: sessionId,
263278
models: modelState,
264-
modes: modeState
279+
modes: modeState,
280+
configOptions: this.createSessionConfigOptions(this.getSessionState(sessionId)),
265281
};
266282
}
267283

@@ -302,6 +318,28 @@ export class CodexAcpServer implements acp.Agent {
302318
return {};
303319
}
304320

321+
async setSessionConfigOption(params: acp.SetSessionConfigOptionRequest): Promise<acp.SetSessionConfigOptionResponse> {
322+
logger.log("Set session config option requested", {
323+
sessionId: params.sessionId,
324+
configId: params.configId,
325+
});
326+
const sessionState = this.sessions.get(params.sessionId);
327+
if (!sessionState) throw new Error(`Session ${params.sessionId} not found`);
328+
329+
if (params.configId !== FAST_MODE_CONFIG_ID || ("type" in params && params.type === "boolean")) {
330+
throw RequestError.invalidParams();
331+
}
332+
333+
if (params.value !== FAST_MODE_ON && params.value !== FAST_MODE_OFF) {
334+
throw RequestError.invalidParams();
335+
}
336+
337+
sessionState.fastModeEnabled = params.value === FAST_MODE_ON;
338+
return {
339+
configOptions: this.createSessionConfigOptions(sessionState),
340+
};
341+
}
342+
305343
async unstable_setSessionModel(params: acp.SetSessionModelRequest): Promise<acp.SetSessionModelResponse | void> {
306344
logger.log("Set session model requested", {
307345
sessionId: params.sessionId,
@@ -337,10 +375,17 @@ export class CodexAcpServer implements acp.Agent {
337375
sessionState.currentModelId = ModelId.fromComponents(model, reasoningEffort).toString();
338376
sessionState.supportedReasoningEfforts = model.supportedReasoningEfforts;
339377
sessionState.supportedInputModalities = model.inputModalities;
378+
sessionState.currentModelSupportsFast = modelSupportsFast(model);
340379

341380
return {};
342381
}
343382

383+
private createSessionConfigOptions(sessionState: SessionState): Array<acp.SessionConfigOption> {
384+
return [
385+
createFastModeConfigOption(sessionState.fastModeEnabled),
386+
];
387+
}
388+
344389
private publishAvailableCommandsAsync(sessionId: string) {
345390
void this.availableCommands.publish(sessionId);
346391
}
@@ -388,6 +433,7 @@ export class CodexAcpServer implements acp.Agent {
388433
const {sessionId, currentModelId, models, thread} = sessionMetadata;
389434
const sessionMcpServers = this.resolveSessionMcpServers(requestedMcpServers, true);
390435
const currentModel = this.findCurrentModel(models, currentModelId);
436+
const currentModelSupportsFast = modelSupportsFast(currentModel);
391437
const sessionState: SessionState = {
392438
sessionId: sessionId,
393439
currentModelId: currentModelId,
@@ -401,6 +447,8 @@ export class CodexAcpServer implements acp.Agent {
401447
rateLimits: null,
402448
account: account,
403449
cwd: request.cwd,
450+
fastModeEnabled: sessionMetadata.currentServiceTier === "fast",
451+
currentModelSupportsFast: currentModelSupportsFast,
404452
sessionMcpServers: sessionMcpServers,
405453
};
406454
this.sessions.set(sessionId, sessionState);
@@ -767,8 +815,12 @@ export class CodexAcpServer implements acp.Agent {
767815
throw RequestError.invalidRequest("The current model does not support image input");
768816
}
769817
const agentMode = sessionState.agentMode;
818+
const serviceTier = resolveFastServiceTier(
819+
sessionState.fastModeEnabled,
820+
sessionState.currentModelSupportsFast,
821+
);
770822
const turnCompleted = await this.runWithProcessCheck(
771-
() => this.codexAcpClient.sendPrompt(params, agentMode, modelId, disableSummary, sessionState.cwd));
823+
() => this.codexAcpClient.sendPrompt(params, agentMode, modelId, serviceTier, disableSummary, sessionState.cwd));
772824

773825
// Check if turn was interrupted (cancelled)
774826
if (turnCompleted.turn.status === "interrupted") {

‎src/FastModeConfig.ts‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
import type {SessionConfigOption} from "@agentclientprotocol/sdk";
2+
import type {ServiceTier} from "./app-server";
3+
import type {Model} from "./app-server/v2";
4+
5+
export const FAST_MODE_CONFIG_ID = "fast-mode";
6+
export const FAST_MODE_ON = "on";
7+
export const FAST_MODE_OFF = "off";
8+
9+
const FAST_MODE_DESCRIPTION = "Use the fast service tier when it is available for the selected model. This can be more expensive.";
10+
11+
export function modelSupportsFast(model: Model | undefined): boolean {
12+
return model?.additionalSpeedTiers?.includes("fast") ?? false;
13+
}
14+
15+
export function resolveFastServiceTier(fastModeEnabled: boolean, currentModelSupportsFast: boolean): ServiceTier | null {
16+
return fastModeEnabled && currentModelSupportsFast ? "fast" : null;
17+
}
18+
19+
export function createFastModeConfigOption(fastModeEnabled: boolean): SessionConfigOption {
20+
return {
21+
id: FAST_MODE_CONFIG_ID,
22+
name: "Fast mode",
23+
description: FAST_MODE_DESCRIPTION,
24+
category: FAST_MODE_CONFIG_ID,
25+
type: "select",
26+
currentValue: fastModeEnabled ? FAST_MODE_ON : FAST_MODE_OFF,
27+
options: [
28+
{
29+
value: FAST_MODE_OFF,
30+
name: "Off",
31+
description: "Use the standard service tier.",
32+
},
33+
{
34+
value: FAST_MODE_ON,
35+
name: "On",
36+
description: FAST_MODE_DESCRIPTION,
37+
},
38+
],
39+
};
40+
}

‎src/__tests__/CodexACPAgent/data/send-attachments-turn-start.json‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,8 @@
5656
"personality": null,
5757
"cwd": "cwd",
5858
"effort": "effort",
59-
"model": "model"
59+
"model": "model",
60+
"serviceTier": null
6061
}
6162
}
6263
{

0 commit comments

Comments
 (0)