Skip to content

Commit c12fc63

Browse files
NikolaiSviridovNikolaiSviridov
committed
Fix auth
1 parent ed8057e commit c12fc63

11 files changed

Lines changed: 387 additions & 109 deletions

‎src/CodexAcpClient.ts‎

Lines changed: 26 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@ import type {AuthenticationLogoutResponse, AuthenticationStatusResponse} from ".
3434
* Converts ACP requests into corresponding app-server operations.
3535
*/
3636
export class CodexAcpClient {
37-
3837
private readonly codexClient: CodexAppServerClient;
3938
private readonly config: JsonObject;
4039
private readonly modelProvider: string | null;
@@ -69,19 +68,27 @@ export class CodexAcpClient {
6968
}
7069

7170
switch (authRequest.methodId) {
72-
case "api-key":
71+
case "api-key": {
7372
if (!authRequest._meta || !authRequest._meta["api-key"]) throw RequestError.invalidRequest();
73+
const loginCompletedVersion = this.codexClient.getAccountLoginCompletedVersion();
7474
await this.codexClient.accountLogin({
7575
type: "apiKey",
7676
apiKey: authRequest._meta["api-key"].apiKey
7777
});
78-
break;
79-
case "chat-gpt":
78+
this.gatewayConfig = null;
79+
const result = await this.codexClient.awaitLoginCompleted(loginCompletedVersion);
80+
return result.success;
81+
}
82+
case "chat-gpt": {
83+
const loginCompletedVersion = this.codexClient.getAccountLoginCompletedVersion();
8084
const loginResponse = await this.codexClient.accountLogin({type: "chatgpt"});
8185
if (loginResponse.type == "chatgpt") {
8286
await open(loginResponse.authUrl);
8387
}
84-
break;
88+
this.gatewayConfig = null;
89+
const result = await this.codexClient.awaitLoginCompleted(loginCompletedVersion);
90+
return result.success;
91+
}
8592
case "gateway":
8693
if (!authRequest._meta) throw RequestError.invalidRequest();
8794

@@ -111,9 +118,7 @@ export class CodexAcpClient {
111118

112119
// Reset the gateway config to null if another authentication method was used
113120
this.gatewayConfig = null;
114-
115-
const result = await this.codexClient.awaitLoginCompleted()
116-
return result.success;
121+
return false;
117122
}
118123

119124

@@ -154,8 +159,9 @@ export class CodexAcpClient {
154159
}
155160

156161
async logout(): Promise<AuthenticationLogoutResponse> {
162+
const accountUpdatedVersion = this.codexClient.getAccountUpdatedVersion();
157163
await this.codexClient.accountLogout();
158-
await this.codexClient.awaitAccountUpdated();
164+
await this.codexClient.awaitAccountUpdated(accountUpdatedVersion);
159165
return {};
160166
}
161167

@@ -178,6 +184,7 @@ export class CodexAcpClient {
178184

179185
async resumeSession(request: acp.ResumeSessionRequest): Promise<SessionMetadata> {
180186
await this.refreshSkills(request.cwd, request._meta);
187+
const mcpStartupVersion = this.codexClient.getMcpStartupCompleteVersion();
181188

182189
const response = await this.codexClient.threadResume({
183190
approvalPolicy: null,
@@ -200,10 +207,12 @@ export class CodexAcpClient {
200207
sessionId: request.sessionId,
201208
currentModelId: currentModelId,
202209
models: codexModels,
210+
mcpStartupVersion,
203211
}
204212
}
205213

206214
async loadSession(request: acp.LoadSessionRequest): Promise<SessionMetadataWithThread> {
215+
const mcpStartupVersion = this.codexClient.getMcpStartupCompleteVersion();
207216
const response = await this.codexClient.threadResume({
208217
approvalPolicy: null,
209218
sandbox: null,
@@ -226,11 +235,13 @@ export class CodexAcpClient {
226235
currentModelId: currentModelId,
227236
models: codexModels,
228237
thread: response.thread,
238+
mcpStartupVersion,
229239
};
230240
}
231241

232242
async newSession(request: acp.NewSessionRequest): Promise<SessionMetadata> {
233243
await this.refreshSkills(request.cwd, request._meta);
244+
const mcpStartupVersion = this.codexClient.getMcpStartupCompleteVersion();
234245

235246
const response = await this.codexClient.threadStart({
236247
config: this.createSessionConfig(request.cwd, request.mcpServers),
@@ -256,12 +267,13 @@ export class CodexAcpClient {
256267
sessionId: response.thread.id,
257268
currentModelId: currentModelId,
258269
models: codexModels,
270+
mcpStartupVersion,
259271
};
260272
}
261273

262-
async awaitMcpServers(): Promise<Array<string>>{
263-
const response = await this.codexClient.awaitMcpStartup();
264-
return response.ready;
274+
async awaitMcpStartup(mcpStartupVersion: number): Promise<Array<string>> {
275+
const startup = await this.codexClient.awaitMcpStartup(mcpStartupVersion);
276+
return startup.ready;
265277
}
266278

267279
private createSessionConfig(projectPath: string, mcpServers: Array<McpServer>): JsonObject {
@@ -390,6 +402,7 @@ export class CodexAcpClient {
390402
"vscode",
391403
"exec",
392404
"appServer",
405+
"custom",
393406
"subAgent",
394407
"subAgentReview",
395408
"subAgentCompact",
@@ -523,6 +536,7 @@ export type SessionMetadata = {
523536
sessionId: string,
524537
currentModelId: string,
525538
models: Model[],
539+
mcpStartupVersion: number,
526540
}
527541

528542
export type SessionMetadataWithThread = SessionMetadata & {

‎src/CodexAcpServer.ts‎

Lines changed: 25 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import {CodexApprovalHandler} from "./CodexApprovalHandler";
1010
import {CodexAuthMethods, type CodexAuthRequest} from "./CodexAuthMethod";
1111
import {CodexAcpClient, type SessionMetadata, type SessionMetadataWithThread} from "./CodexAcpClient";
1212
import {ACPSessionConnection, type UpdateSessionEvent} from "./ACPSessionConnection";
13-
import type {Account, CollabAgentToolCallStatus, Model, RateLimitSnapshot, Thread, ThreadItem, UserInput, ReasoningEffortOption} from "./app-server/v2";
13+
import type {Account, CollabAgentToolCallStatus, Model, Thread, ThreadItem, UserInput, ReasoningEffortOption} from "./app-server/v2";
1414
import type {RateLimitsMap} from "./RateLimitsMap";
1515
import type {InputModality, ReasoningEffort} from "./app-server";
1616
import {ModelId} from "./ModelId";
@@ -127,7 +127,6 @@ export class CodexAcpServer implements acp.Agent {
127127

128128
async getOrCreateSession(request: acp.NewSessionRequest | acp.ResumeSessionRequest): Promise<[SessionId, SessionModelState, SessionModeState]> {
129129
await this.checkAuthorization();
130-
const pendingMcpServers: Promise<Array<string>> = this.codexAcpClient.awaitMcpServers();
131130

132131
let sessionMetadata: SessionMetadata;
133132
if ("sessionId" in request) {
@@ -139,9 +138,8 @@ export class CodexAcpServer implements acp.Agent {
139138
}
140139

141140
const accountResponse = await this.runWithProcessCheck(() => this.codexAcpClient.getAccount());
142-
const {sessionId, currentModelId, models} = sessionMetadata;
143-
logger.log(`Waiting MCP servers to start...`)
144-
const sessionMcpServers = await pendingMcpServers;
141+
const {sessionId, currentModelId, models, mcpStartupVersion} = sessionMetadata;
142+
const sessionMcpServers = await this.resolveSessionMcpServers(request.mcpServers ?? [], mcpStartupVersion, "sessionId" in request);
145143
const currentModel = this.findCurrentModel(models, currentModelId);
146144
const sessionState: SessionState = {
147145
sessionId: sessionId,
@@ -332,17 +330,15 @@ export class CodexAcpServer implements acp.Agent {
332330
thread: Thread;
333331
}> {
334332
await this.checkAuthorization();
335-
const pendingMcpServers: Promise<Array<string>> = this.codexAcpClient.awaitMcpServers();
336333

337334
logger.log(`Load existing session: ${request.sessionId}...`);
338335
const sessionMetadata: SessionMetadataWithThread = await this.runWithProcessCheck(() =>
339336
this.codexAcpClient.loadSession(request)
340337
);
341338

342339
const accountResponse = await this.runWithProcessCheck(() => this.codexAcpClient.getAccount());
343-
const {sessionId, currentModelId, models, thread} = sessionMetadata;
344-
logger.log("Waiting MCP servers to start...");
345-
const sessionMcpServers = await pendingMcpServers;
340+
const {sessionId, currentModelId, models, thread, mcpStartupVersion} = sessionMetadata;
341+
const sessionMcpServers = await this.resolveSessionMcpServers(request.mcpServers ?? [], mcpStartupVersion, true);
346342
const currentModel = this.findCurrentModel(models, currentModelId);
347343
const sessionState: SessionState = {
348344
sessionId: sessionId,
@@ -601,6 +597,22 @@ export class CodexAcpServer implements acp.Agent {
601597
return sessionState;
602598
}
603599

600+
private async resolveSessionMcpServers(
601+
mcpServers: Array<acp.McpServer>,
602+
mcpStartupVersion: number,
603+
recoverFromStartup: boolean,
604+
): Promise<Array<string>> {
605+
const requestedServerNames = getRequestedMcpServerNames(mcpServers);
606+
if (requestedServerNames.length > 0) {
607+
return requestedServerNames;
608+
}
609+
if (!recoverFromStartup) {
610+
return [];
611+
}
612+
logger.log("Recovering MCP servers from startup state...");
613+
return await this.runWithProcessCheck(() => this.codexAcpClient.awaitMcpStartup(mcpStartupVersion));
614+
}
615+
604616
async prompt(params: acp.PromptRequest): Promise<acp.PromptResponse> {
605617
logger.log("Prompt received", {
606618
sessionId: params.sessionId,
@@ -747,3 +759,7 @@ export class CodexAcpServer implements acp.Agent {
747759
}
748760
}
749761
}
762+
763+
function getRequestedMcpServerNames(mcpServers: Array<acp.McpServer>): Array<string> {
764+
return Array.from(new Set(mcpServers.map(server => server.name)));
765+
}

0 commit comments

Comments
 (0)