Skip to content

Commit 3454939

Browse files
NikolaiSviridovNikolaiSviridov
authored
Merge pull request #1 from JetBrains/nikolaisv/LLM-21851-Codex-Model-selector
LLM-21851 [Codex] Model selector
2 parents 7593f6a + e654a4a commit 3454939

9 files changed

Lines changed: 422 additions & 23 deletions

‎src/CodexAcpClient.ts‎

Lines changed: 45 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,15 @@ import * as acp from "@agentclientprotocol/sdk";
33
import type {CodexAppServerClient} from "./CodexAppServerClient";
44
import {RequestError} from "@agentclientprotocol/sdk";
55
import open from "open";
6-
import type {ClientInfo, ServerNotification} from "./app-server";
6+
import type {
7+
ClientInfo,
8+
ServerNotification,
9+
SetDefaultModelParams,
10+
SetDefaultModelResponse
11+
} from "./app-server";
712
import type {JsonValue} from "./app-server/serde_json/JsonValue";
13+
import type {Model} from "./app-server/v2";
14+
import {ModelId} from "./ModelId";
815

916
/**
1017
* API for accessing the Codex App Server using ACP requests.
@@ -13,12 +20,12 @@ import type {JsonValue} from "./app-server/serde_json/JsonValue";
1320
export class CodexAcpClient {
1421

1522
private readonly codexClient: CodexAppServerClient;
16-
private readonly config: JsonObject | null;
23+
private readonly config: JsonObject;
1724
private readonly modelProvider: string | null;
1825

1926
constructor(codexClient: CodexAppServerClient, codexConfig?: JsonObject, modelProvider?: string) {
2027
this.codexClient = codexClient;
21-
this.config = codexConfig ?? null;
28+
this.config = codexConfig ?? {};
2229
this.modelProvider = modelProvider ?? null;
2330
}
2431

@@ -71,8 +78,8 @@ export class CodexAcpClient {
7178
/**
7279
* Returns a new session ID.
7380
*/
74-
async newSession(request: acp.NewSessionRequest): Promise<string> {
75-
const response = await this.codexClient.threadStart({
81+
async newSession(request: acp.NewSessionRequest): Promise<SessionMetadata> {
82+
const threadStartResponse = await this.codexClient.threadStart({
7683
config: this.config,
7784
modelProvider: this.modelProvider,
7885
model: null,
@@ -82,7 +89,16 @@ export class CodexAcpClient {
8289
baseInstructions: null,
8390
developerInstructions: null,
8491
});
85-
return response.thread.id;
92+
const codexModels = await this.fetchAvailableModels();
93+
if (codexModels.length === 0) {
94+
throw new Error("Codex did not return any models");
95+
}
96+
const currentModelId = ModelId.fromThreadResponse(threadStartResponse).toString();
97+
return {
98+
sessionId: threadStartResponse.thread.id,
99+
currentModelId: currentModelId,
100+
models: codexModels
101+
};
86102
}
87103

88104
async sendPrompt(request: acp.PromptRequest, eventHandler: (result: ServerNotification) => void): Promise<void> {
@@ -106,6 +122,28 @@ export class CodexAcpClient {
106122
await this.codexClient.awaitTurnCompleted();
107123
}
108124

125+
async setModel(params: SetDefaultModelParams): Promise<SetDefaultModelResponse> {
126+
return this.codexClient.setModelRequest(params);
127+
}
128+
129+
private async fetchAvailableModels(): Promise<Model[]> {
130+
const models: Model[] = [];
131+
let cursor: string | null = null;
132+
133+
do {
134+
const response = await this.codexClient.listModels({ cursor, limit: null });
135+
models.push(...response.data);
136+
cursor = response.nextCursor;
137+
} while (cursor);
138+
139+
return models;
140+
}
109141
}
110142

111-
export type JsonObject = { [key in string]?: JsonValue }
143+
export type JsonObject = { [key in string]?: JsonValue }
144+
145+
export type SessionMetadata = {
146+
sessionId: string,
147+
currentModelId: string,
148+
models: Model[]
149+
}

‎src/CodexAcpServer.ts‎

Lines changed: 65 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,15 @@
11
import * as acp from "@agentclientprotocol/sdk";
22
import {CodexEventHandler} from "./CodexEventHandler";
33
import {CodexAuthMethods, type CodexAuthRequest} from "./CodexAuthMethod";
4-
import {RequestError} from "@agentclientprotocol/sdk";
5-
import {CodexAcpClient} from "./CodexAcpClient";
4+
import {type ModelInfo, RequestError, type SessionModelState} from "@agentclientprotocol/sdk";
5+
import {CodexAcpClient, type SessionMetadata} from "./CodexAcpClient";
6+
import type {Model} from "./app-server/v2";
7+
import type {ReasoningEffort} from "./app-server";
8+
import {ModelId} from "./ModelId";
69

710

811
export interface SessionState {
9-
sessionId: string,
12+
sessionMetadata: SessionMetadata;
1013
pendingPrompt: AbortController | null;
1114
}
1215

@@ -52,14 +55,21 @@ export class CodexAcpServer implements acp.Agent {
5255
}
5356
}
5457

55-
const sessionId = await this.codexAcpClient.newSession(_params);
58+
const sessionMetadata = await this.codexAcpClient.newSession(_params);
59+
const {sessionId, currentModelId, models} = sessionMetadata;
5660
this.sessions.set(sessionId, {
57-
sessionId: sessionId,
58-
pendingPrompt: null,
61+
sessionMetadata: sessionMetadata,
62+
pendingPrompt: null
5963
});
6064

65+
const availableModels = this.buildAvailableModels(models);
66+
const sessionModelState: SessionModelState = {
67+
availableModels: availableModels,
68+
currentModelId: currentModelId,
69+
}
6170
return {
6271
sessionId: sessionId,
72+
models: sessionModelState,
6373
};
6474
}
6575

@@ -80,6 +90,55 @@ export class CodexAcpServer implements acp.Agent {
8090
return {};
8191
}
8292

93+
async setSessionModel(params: acp.SetSessionModelRequest): Promise<acp.SetSessionModelResponse> {
94+
const sessionState = this.sessions.get(params.sessionId);
95+
if (!sessionState) throw new Error(`Session ${params.sessionId} not found`);
96+
97+
const requestedModelId= ModelId.fromString(params.modelId);
98+
const requestedModelName = requestedModelId.model;
99+
const requestedEffort = requestedModelId.effort;
100+
101+
const model = sessionState.sessionMetadata.models.find(m => m.id === requestedModelName);
102+
if (!model) throw new Error(`Unknown model ${params.modelId}`);
103+
104+
const requestedEffortValue = requestedEffort as ReasoningEffort | undefined;
105+
let reasoningEffort: ReasoningEffort;
106+
if (requestedEffortValue) {
107+
const matchedEffort = model.supportedReasoningEfforts.find(
108+
(option) => option.reasoningEffort === requestedEffortValue
109+
)?.reasoningEffort;
110+
111+
if (!matchedEffort) {
112+
throw new Error(`Unsupported reasoning effort ${requestedEffortValue} for model ${requestedModelName}`);
113+
}
114+
115+
reasoningEffort = matchedEffort;
116+
} else {
117+
reasoningEffort = model.defaultReasoningEffort;
118+
}
119+
120+
121+
await this.codexAcpClient.setModel({
122+
model: model.model,
123+
reasoningEffort,
124+
});
125+
sessionState.sessionMetadata.currentModelId = ModelId.fromComponents(model, reasoningEffort).toString();
126+
127+
return {};
128+
}
129+
130+
131+
132+
private buildAvailableModels(models: Model[]): ModelInfo[] {
133+
return models.flatMap((model) =>
134+
model.supportedReasoningEfforts.map((effort) => ({
135+
modelId: ModelId.fromComponents(model, effort.reasoningEffort).toString(),
136+
name: `${model.displayName} (${effort.reasoningEffort})`,
137+
description: `${model.description} ${effort.description}`,
138+
}))
139+
);
140+
}
141+
83142
getSessionState(sessionId: string): SessionState {
84143
const sessionState = this.sessions.get(sessionId);
85144
if (!sessionState) {

‎src/CodexAppServerClient.ts‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,15 @@
1-
import type {MessageConnection, NotificationMessage} from "vscode-jsonrpc/node";
1+
import type {MessageConnection} from "vscode-jsonrpc/node";
22
import type {
33
ClientRequest,
4-
EventMsg,
54
InitializeParams,
65
InitializeResponse,
7-
ServerNotification
6+
ServerNotification, SetDefaultModelParams, SetDefaultModelResponse
87
} from "./app-server";
98
import type {
109
AccountLoginCompletedNotification, AccountUpdatedNotification,
1110
GetAccountParams,
12-
GetAccountResponse, LoginAccountParams, LoginAccountResponse, LogoutAccountResponse,
11+
GetAccountResponse, LoginAccountParams, LoginAccountResponse, LogoutAccountResponse, ModelListParams,
12+
ModelListResponse,
1313
ThreadStartParams,
1414
ThreadStartResponse,
1515
TurnCompletedNotification,
@@ -86,6 +86,15 @@ export class CodexAppServerClient {
8686
});
8787
}
8888

89+
90+
async setModelRequest(params: SetDefaultModelParams): Promise<SetDefaultModelResponse> {
91+
return await this.sendRequest({ method: "setDefaultModel", params });
92+
}
93+
94+
async listModels(params: ModelListParams = {cursor: null, limit: null}): Promise<ModelListResponse> {
95+
return await this.sendRequest({ method: "model/list", params });
96+
}
97+
8998
//TODO support removal (leads to duplicated processing of follow-ups)
9099
onServerNotification(callback: (event: ServerNotification) => void){
91100
this.notificationHandlers.push(callback);

‎src/CodexEventHandler.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ export class CodexEventHandler {
2626
}
2727

2828
async handleNotification(notification: ServerNotification) {
29-
const session = new ACPSessionConnection(this.connection, this.sessionState.sessionId);
29+
const session = new ACPSessionConnection(this.connection, this.sessionState.sessionMetadata.sessionId);
3030
const updateEvent = await this.createUpdateEvent(notification);
3131
if (updateEvent) {
3232
await session.update(updateEvent);

‎src/ModelId.ts‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
import type {ReasoningEffort} from "./app-server";
2+
import type {Model, ThreadStartResponse} from "./app-server/v2";
3+
4+
export class ModelId {
5+
private constructor(
6+
public readonly model: string,
7+
public readonly effort: string | null // TODO: ThreadStartResponse
8+
) {}
9+
10+
static fromComponents(model: Model, effort: ReasoningEffort): ModelId {
11+
return new ModelId(model.id, effort);
12+
}
13+
14+
static fromThreadResponse(response: ThreadStartResponse): ModelId {
15+
return new ModelId(response.model, response.reasoningEffort);
16+
}
17+
18+
static fromString(modelId: string): ModelId {
19+
const parts = modelId.split("/");
20+
const model = parts[0];
21+
const effort = parts[1] ?? null;
22+
23+
if (!model) {
24+
throw new Error(`Invalid modelId format: ${modelId}`);
25+
}
26+
return new ModelId(model, effort);
27+
}
28+
29+
toString(): string {
30+
return this.effort ? `${this.model}/${this.effort}` : this.model;
31+
}
32+
}

‎src/__tests__/CodexACPAgent/CodexAcpClient.test.ts‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import {describe, expect, it, vi, beforeEach, afterEach} from 'vitest';
1+
import {describe, expect, it, vi, beforeEach} from 'vitest';
22
import type {CodexAuthRequest} from "../../CodexAuthMethod";
33
import {createTestFixture, type TestFixture} from "../acp-test-utils";
44
import type {ServerNotification} from "../../app-server";
@@ -12,7 +12,7 @@ describe('ACP server test', () => {
1212
vi.clearAllMocks();
1313
});
1414

15-
const ignoredFields = ["thread", "cwd", "id", "createdAt", "path", "threadId", "userAgent", "sandbox", "reasoningEffort"];
15+
const ignoredFields = ["thread", "cwd", "id", "createdAt", "path", "threadId", "userAgent", "sandbox", "reasoningEffort", "conversationId"];
1616

1717
it('should start conversation', async () => {
1818
const codexAcpAgent = fixture.getCodexAcpAgent();
@@ -79,7 +79,15 @@ describe('ACP server test', () => {
7979

8080
fixture.getCodexAppServerClient().turnStart = vi.fn().mockResolvedValue(undefined);
8181
fixture.getCodexAppServerClient().awaitTurnCompleted = vi.fn().mockResolvedValue(undefined);
82-
fixture.getCodexAcpAgent().getSessionState = vi.fn().mockResolvedValue({ pendingPrompt: null, sessionId: "id" });
82+
const sessionState: SessionState = {
83+
pendingPrompt: null,
84+
sessionMetadata: {
85+
sessionId: "id",
86+
currentModelId: "model-id",
87+
models: [],
88+
}
89+
};
90+
vi.spyOn(codexAcpAgent, "getSessionState").mockReturnValue(sessionState);
8391

8492
await codexAcpAgent.prompt({ sessionId: "id", prompt: [{type: "text", text: ""}] });
8593

0 commit comments

Comments
 (0)