Skip to content

Commit 41d2c3c

Browse files
NikolaiSviridovNikolaiSviridov
committed
simplify auth and mcp startup
1 parent 2f623c1 commit 41d2c3c

5 files changed

Lines changed: 42 additions & 60 deletions

File tree

‎src/CodexAcpClient.ts‎

Lines changed: 28 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@ import {AgentMode} from "./AgentMode";
1515
import path from "node:path";
1616
import {logger} from "./Logger";
1717
import type {
18+
AccountLoginCompletedNotification,
19+
AccountUpdatedNotification,
1820
GetAccountResponse,
1921
ListMcpServerStatusParams,
2022
ListMcpServerStatusResponse,
@@ -70,23 +72,23 @@ export class CodexAcpClient {
7072
switch (authRequest.methodId) {
7173
case "api-key": {
7274
if (!authRequest._meta || !authRequest._meta["api-key"]) throw RequestError.invalidRequest();
73-
const loginCompletedVersion = this.codexClient.getAccountLoginCompletedVersion();
75+
const loginCompletedPromise = this.awaitNextLoginCompleted();
7476
await this.codexClient.accountLogin({
7577
type: "apiKey",
7678
apiKey: authRequest._meta["api-key"].apiKey
7779
});
7880
this.gatewayConfig = null;
79-
const result = await this.codexClient.awaitLoginCompleted(loginCompletedVersion);
81+
const result = await loginCompletedPromise;
8082
return result.success;
8183
}
8284
case "chat-gpt": {
83-
const loginCompletedVersion = this.codexClient.getAccountLoginCompletedVersion();
85+
const loginCompletedPromise = this.awaitNextLoginCompleted();
8486
const loginResponse = await this.codexClient.accountLogin({type: "chatgpt"});
8587
if (loginResponse.type == "chatgpt") {
8688
await open(loginResponse.authUrl);
8789
}
8890
this.gatewayConfig = null;
89-
const result = await this.codexClient.awaitLoginCompleted(loginCompletedVersion);
91+
const result = await loginCompletedPromise;
9092
return result.success;
9193
}
9294
case "gateway":
@@ -159,9 +161,9 @@ export class CodexAcpClient {
159161
}
160162

161163
async logout(): Promise<AuthenticationLogoutResponse> {
162-
const accountUpdatedVersion = this.codexClient.getAccountUpdatedVersion();
164+
const accountUpdatedPromise = this.awaitNextAccountUpdated();
163165
await this.codexClient.accountLogout();
164-
await this.codexClient.awaitAccountUpdated(accountUpdatedVersion);
166+
await accountUpdatedPromise;
165167
return {};
166168
}
167169

@@ -184,7 +186,6 @@ export class CodexAcpClient {
184186

185187
async resumeSession(request: acp.ResumeSessionRequest): Promise<SessionMetadata> {
186188
await this.refreshSkills(request.cwd, request._meta);
187-
const mcpStartupVersion = this.codexClient.getMcpStartupCompleteVersion();
188189

189190
const response = await this.codexClient.threadResume({
190191
approvalPolicy: null,
@@ -207,12 +208,10 @@ export class CodexAcpClient {
207208
sessionId: request.sessionId,
208209
currentModelId: currentModelId,
209210
models: codexModels,
210-
mcpStartupVersion,
211211
}
212212
}
213213

214214
async loadSession(request: acp.LoadSessionRequest): Promise<SessionMetadataWithThread> {
215-
const mcpStartupVersion = this.codexClient.getMcpStartupCompleteVersion();
216215
const response = await this.codexClient.threadResume({
217216
approvalPolicy: null,
218217
sandbox: null,
@@ -235,13 +234,11 @@ export class CodexAcpClient {
235234
currentModelId: currentModelId,
236235
models: codexModels,
237236
thread: response.thread,
238-
mcpStartupVersion,
239237
};
240238
}
241239

242240
async newSession(request: acp.NewSessionRequest): Promise<SessionMetadata> {
243241
await this.refreshSkills(request.cwd, request._meta);
244-
const mcpStartupVersion = this.codexClient.getMcpStartupCompleteVersion();
245242

246243
const response = await this.codexClient.threadStart({
247244
config: this.createSessionConfig(request.cwd, request.mcpServers),
@@ -267,7 +264,6 @@ export class CodexAcpClient {
267264
sessionId: response.thread.id,
268265
currentModelId: currentModelId,
269266
models: codexModels,
270-
mcpStartupVersion,
271267
};
272268
}
273269

@@ -276,6 +272,10 @@ export class CodexAcpClient {
276272
return startup.ready;
277273
}
278274

275+
getMcpStartupCompleteVersion(): number {
276+
return this.codexClient.getMcpStartupCompleteVersion();
277+
}
278+
279279
private createSessionConfig(projectPath: string, mcpServers: Array<McpServer>): JsonObject {
280280
const mergedConfig = {
281281
...mergeGatewayConfig(this.config, this.gatewayConfig),
@@ -392,6 +392,22 @@ export class CodexAcpClient {
392392
return this.codexClient.listSkills(params ?? {});
393393
}
394394

395+
private async awaitNextLoginCompleted(): Promise<AccountLoginCompletedNotification> {
396+
return await new Promise((resolve) => {
397+
this.codexClient.connection.onNotification("account/login/completed", (event: AccountLoginCompletedNotification) => {
398+
resolve(event);
399+
});
400+
});
401+
}
402+
403+
private async awaitNextAccountUpdated(): Promise<AccountUpdatedNotification> {
404+
return await new Promise((resolve) => {
405+
this.codexClient.connection.onNotification("account/updated", (event: AccountUpdatedNotification) => {
406+
resolve(event);
407+
});
408+
});
409+
}
410+
395411
async listMcpServers(params: ListMcpServerStatusParams = { cursor: null, limit: null }): Promise<ListMcpServerStatusResponse> {
396412
return this.codexClient.listMcpServerStatus(params);
397413
}
@@ -535,7 +551,6 @@ export type SessionMetadata = {
535551
sessionId: string,
536552
currentModelId: string,
537553
models: Model[],
538-
mcpStartupVersion: number,
539554
}
540555

541556
export type SessionMetadataWithThread = SessionMetadata & {

‎src/CodexAcpServer.ts‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,7 @@ 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 mcpStartupVersion = this.codexAcpClient.getMcpStartupCompleteVersion();
130131

131132
let sessionMetadata: SessionMetadata;
132133
if ("sessionId" in request) {
@@ -138,7 +139,7 @@ export class CodexAcpServer implements acp.Agent {
138139
}
139140

140141
const accountResponse = await this.runWithProcessCheck(() => this.codexAcpClient.getAccount());
141-
const {sessionId, currentModelId, models, mcpStartupVersion} = sessionMetadata;
142+
const {sessionId, currentModelId, models} = sessionMetadata;
142143
const sessionMcpServers = await this.resolveSessionMcpServers(request.mcpServers ?? [], mcpStartupVersion, "sessionId" in request);
143144
const currentModel = this.findCurrentModel(models, currentModelId);
144145
const sessionState: SessionState = {
@@ -330,14 +331,15 @@ export class CodexAcpServer implements acp.Agent {
330331
thread: Thread;
331332
}> {
332333
await this.checkAuthorization();
334+
const mcpStartupVersion = this.codexAcpClient.getMcpStartupCompleteVersion();
333335

334336
logger.log(`Load existing session: ${request.sessionId}...`);
335337
const sessionMetadata: SessionMetadataWithThread = await this.runWithProcessCheck(() =>
336338
this.codexAcpClient.loadSession(request)
337339
);
338340

339341
const accountResponse = await this.runWithProcessCheck(() => this.codexAcpClient.getAccount());
340-
const {sessionId, currentModelId, models, thread, mcpStartupVersion} = sessionMetadata;
342+
const {sessionId, currentModelId, models, thread} = sessionMetadata;
341343
const sessionMcpServers = await this.resolveSessionMcpServers(request.mcpServers ?? [], mcpStartupVersion, true);
342344
const currentModel = this.findCurrentModel(models, currentModelId);
343345
const sessionState: SessionState = {

‎src/CodexAppServerClient.ts‎

Lines changed: 0 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -60,28 +60,12 @@ const FileChangeApprovalRequest = new RequestType<
6060
export class CodexAppServerClient {
6161
readonly connection: MessageConnection;
6262
private approvalHandlers = new Map<string, ApprovalHandler>();
63-
private accountLoginCompletedVersion = 0;
64-
private lastAccountLoginCompleted: AccountLoginCompletedNotification | null = null;
65-
private readonly accountLoginCompletedResolvers: Array<SignalResolver<AccountLoginCompletedNotification>> = [];
66-
private accountUpdatedVersion = 0;
67-
private lastAccountUpdated: AccountUpdatedNotification | null = null;
68-
private readonly accountUpdatedResolvers: Array<SignalResolver<AccountUpdatedNotification>> = [];
6963
private mcpStartupCompleteVersion = 0;
7064
private lastMcpStartupComplete: McpStartupCompleteEvent | null = null;
7165
private readonly mcpStartupCompleteResolvers: Array<SignalResolver<McpStartupCompleteEvent>> = [];
7266

7367
constructor(connection: MessageConnection) {
7468
this.connection = connection;
75-
this.connection.onNotification("account/login/completed", (event: AccountLoginCompletedNotification) => {
76-
this.accountLoginCompletedVersion += 1;
77-
this.lastAccountLoginCompleted = event;
78-
this.resolveSignal(event, this.accountLoginCompletedVersion, this.accountLoginCompletedResolvers);
79-
});
80-
this.connection.onNotification("account/updated", (event: AccountUpdatedNotification) => {
81-
this.accountUpdatedVersion += 1;
82-
this.lastAccountUpdated = event;
83-
this.resolveSignal(event, this.accountUpdatedVersion, this.accountUpdatedResolvers);
84-
});
8569
this.connection.onUnhandledNotification((data) => {
8670
if (isMcpStartupCompleteNotification(data)) {
8771
this.mcpStartupCompleteVersion += 1;
@@ -169,32 +153,6 @@ export class CodexAppServerClient {
169153
return await this.sendRequest({ method: "config/read", params: params });
170154
}
171155

172-
getAccountLoginCompletedVersion(): number {
173-
return this.accountLoginCompletedVersion;
174-
}
175-
176-
async awaitLoginCompleted(afterVersion: number): Promise<AccountLoginCompletedNotification> {
177-
return await this.awaitSignal(
178-
this.lastAccountLoginCompleted,
179-
this.accountLoginCompletedVersion,
180-
afterVersion,
181-
this.accountLoginCompletedResolvers
182-
);
183-
}
184-
185-
getAccountUpdatedVersion(): number {
186-
return this.accountUpdatedVersion;
187-
}
188-
189-
async awaitAccountUpdated(afterVersion: number): Promise<AccountUpdatedNotification> {
190-
return await this.awaitSignal(
191-
this.lastAccountUpdated,
192-
this.accountUpdatedVersion,
193-
afterVersion,
194-
this.accountUpdatedResolvers
195-
);
196-
}
197-
198156
getMcpStartupCompleteVersion(): number {
199157
return this.mcpStartupCompleteVersion;
200158
}

‎src/__tests__/CodexACPAgent/model-filtering.test.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,6 @@ describe("Model filtering", () => {
8080
sessionId: "session-id",
8181
currentModelId: "gpt-5.2[medium]",
8282
models,
83-
mcpStartupVersion: 0,
8483
});
8584
vi.spyOn(codexAcpClient, "getAccount").mockResolvedValue({account: null, requiresOpenaiAuth: false});
8685

‎src/login.ts‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ async function login(options: LoginOptions): Promise<boolean> {
8989
}
9090

9191
logger.log("Starting ChatGPT login...");
92-
const loginCompletedVersion = appServerClient.getAccountLoginCompletedVersion();
92+
const loginCompletedPromise = waitForNextLoginCompleted(appServerClient);
9393
const loginResponse = await appServerClient.accountLogin({type: "chatgpt"});
9494

9595
if (loginResponse.type === "chatgpt") {
@@ -101,7 +101,7 @@ async function login(options: LoginOptions): Promise<boolean> {
101101
}
102102

103103
logger.log("Waiting for login completion...");
104-
const result = await appServerClient.awaitLoginCompleted(loginCompletedVersion);
104+
const result = await loginCompletedPromise;
105105

106106
if (result.success) {
107107
logger.log("Login successful!");
@@ -116,6 +116,14 @@ async function login(options: LoginOptions): Promise<boolean> {
116116
}
117117
}
118118

119+
function waitForNextLoginCompleted(appServerClient: CodexAppServerClient): Promise<{ success: boolean }> {
120+
return new Promise((resolve) => {
121+
appServerClient.connection.onNotification("account/login/completed", (event: { success: boolean }) => {
122+
resolve(event);
123+
});
124+
});
125+
}
126+
119127
/**
120128
* Run the login command with the given CLI arguments.
121129
* @param args CLI arguments after "login" command (e.g., ["--client-name", "AIA"])

0 commit comments

Comments
 (0)