From 849573ed206f7c2851d7ca8f5a50eccf0bf48011 Mon Sep 17 00:00:00 2001 From: Julian Date: Fri, 31 Jul 2026 19:44:59 -0700 Subject: [PATCH] Add ACP session fork support --- src/CodexAcpClient.ts | 32 ++++ src/CodexAcpServer.ts | 132 ++++++++++++-- src/CodexAppServerClient.ts | 6 + .../data/session-fork-model-list-failed.json | 32 ++++ .../CodexACPAgent/data/session-fork.json | 29 ++++ .../CodexACPAgent/initialize.test.ts | 1 + .../CodexACPAgent/session-fork.test.ts | 163 ++++++++++++++++++ src/index.ts | 1 + 8 files changed, 380 insertions(+), 16 deletions(-) create mode 100644 src/__tests__/CodexACPAgent/data/session-fork-model-list-failed.json create mode 100644 src/__tests__/CodexACPAgent/data/session-fork.json create mode 100644 src/__tests__/CodexACPAgent/session-fork.test.ts diff --git a/src/CodexAcpClient.ts b/src/CodexAcpClient.ts index 0c7af207..7bc60ead 100644 --- a/src/CodexAcpClient.ts +++ b/src/CodexAcpClient.ts @@ -379,6 +379,38 @@ export class CodexAcpClient { }; } + async forkSession( + request: acp.ForkSessionRequest, + onSubscribed: (sessionId: string) => void + ): Promise { + const additionalDirectories = readAdditionalDirectories(request.cwd, request.additionalDirectories, request._meta); + await this.refreshSkills(request.cwd, additionalDirectories); + + const response = await this.codexClient.threadFork({ + config: await this.createSessionConfig(request.cwd, additionalDirectories, request.mcpServers ?? []), + cwd: request.cwd, + ephemeral: false, + modelProvider: await this.getResumeModelProvider(), + threadId: request.sessionId, + }); + if (response.thread.id === request.sessionId) { + throw new Error("Codex thread/fork did not return a child session id"); + } + // Codex has subscribed to the child now, so the caller must be able to clean it up if later work fails. + onSubscribed(response.thread.id); + const codexModels = await this.fetchAvailableModels(); + const currentModelId = this.createModelId(codexModels, response.model, response.reasoningEffort).toString(); + return { + sessionId: response.thread.id, + currentModelId: currentModelId, + models: codexModels, + collaborationMode: this.getCollaborationMode(response.thread.id), + modelProvider: response.modelProvider, + currentServiceTier: response.serviceTier as ServiceTier ?? null, + additionalDirectories, + }; + } + async newSession(request: acp.NewSessionRequest): Promise { const additionalDirectories = readAdditionalDirectories(request.cwd, request.additionalDirectories, request._meta); await this.refreshSkills(request.cwd, additionalDirectories); diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index 928f2b16..4e2d88d8 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -124,6 +124,17 @@ interface ActiveAuthState { authConfigured: boolean; } +interface InstallSessionStateOptions { + cwd: string; + sessionMetadata: SessionMetadata; + authState: ActiveAuthState; + authProvider: string | null; + requestedMcpServers: Array; + mcpServerStartupVersion: number | null; + recoverMcpServers: boolean; + sessionTitleSource: SessionState["sessionTitleSource"]; +} + interface PendingMcpStartupSession { requestedServers: Set; afterVersion: number; @@ -232,6 +243,7 @@ export class CodexAcpServer { }, sessionCapabilities: { resume: { }, + fork: { }, list: { }, close: { }, delete: { }, @@ -308,8 +320,12 @@ export class CodexAcpServer { } async getOrCreateSession(request: acp.NewSessionRequest | acp.ResumeSessionRequest): Promise<[SessionId, LegacySessionModelState, SessionModeState]> { + return await this.withSessionOpenErrorHandling(() => this.tryCreateSession(request)); + } + + private async withSessionOpenErrorHandling(openSession: () => Promise): Promise { try { - return await this.tryCreateSession(request); + return await openSession(); } catch (e) { const error = e instanceof Error ? e : new Error(String(e)); await this.handleError(error); @@ -439,9 +455,41 @@ export class CodexAcpServer { resumeSubscribed = false; await this.closeStaleSessionOpen(sessionId, sessionGeneration); } - const sessionMcpServers = this.resolveSessionMcpServers(requestedMcpServers, "sessionId" in request); + const sessionState = this.installSessionState({ + cwd: request.cwd, + sessionMetadata, + authState, + authProvider, + requestedMcpServers, + mcpServerStartupVersion, + recoverMcpServers: "sessionId" in request, + sessionTitleSource: "sessionId" in request ? "unknown" : "unset", + }); + resumeSubscribed = false; + + this.publishAvailableCommandsAsync(sessionState); + if ("sessionId" in request) { + this.publishCurrentGoalAsync(sessionState, sessionGeneration); + } + const sessionModelState: LegacySessionModelState = this.createModelState(models, currentModelId); + const sessionModeState: SessionModeState = sessionState.agentMode.toSessionModeState(); + + return [sessionId, sessionModelState, sessionModeState]; + } + + private installSessionState(options: InstallSessionStateOptions): SessionState { + const { + cwd, + sessionMetadata, + authState, + authProvider, + requestedMcpServers, + mcpServerStartupVersion, + recoverMcpServers, + sessionTitleSource, + } = options; + const {sessionId, currentModelId, models} = sessionMetadata; const currentModel = this.findCurrentModel(models, currentModelId); - const currentModelSupportsFast = modelSupportsFast(currentModel); const sessionState: SessionState = { sessionId: sessionId, currentModelId: currentModelId, @@ -458,18 +506,17 @@ export class CodexAcpServer { account: authState.account, authConfigured: authState.authConfigured, authProvider: authProvider, - cwd: request.cwd, + cwd: cwd, additionalDirectories: sessionMetadata.additionalDirectories, fastModeEnabled: sessionMetadata.currentServiceTier === "fast", - currentModelSupportsFast: currentModelSupportsFast, - sessionMcpServers: sessionMcpServers, + currentModelSupportsFast: modelSupportsFast(currentModel), + sessionMcpServers: this.resolveSessionMcpServers(requestedMcpServers, recoverMcpServers), terminalOutputMode: this.terminalOutputMode, goalRevision: 0, sessionTitle: null, - sessionTitleSource: "sessionId" in request ? "unknown" : "unset", + sessionTitleSource: sessionTitleSource, }; this.sessions.set(sessionId, sessionState); - resumeSubscribed = false; if (requestedMcpServers.length > 0 && mcpServerStartupVersion !== null) { this.pendingMcpStartupSessions.set(sessionId, { @@ -479,14 +526,7 @@ export class CodexAcpServer { this.publishMcpStartupStatusAsync(sessionId); } - this.publishAvailableCommandsAsync(sessionState); - if ("sessionId" in request) { - this.publishCurrentGoalAsync(sessionState, sessionGeneration); - } - const sessionModelState: LegacySessionModelState = this.createModelState(models, currentModelId); - const sessionModeState: SessionModeState = sessionState.agentMode.toSessionModeState(); - - return [sessionId, sessionModelState, sessionModeState]; + return sessionState; } private async getAuthStateForProvider(authProvider: string | null): Promise { @@ -560,6 +600,66 @@ export class CodexAcpServer { }; } + async unstable_forkSession(params: acp.ForkSessionRequest): Promise { + logger.log("Forking session...", {sessionId: params.sessionId}); + return await this.withSessionOpenErrorHandling(() => this.tryForkSession(params)); + } + + private async tryForkSession(params: acp.ForkSessionRequest): Promise { + await this.checkAuthorization(); + const requestedMcpServers = params.mcpServers ?? []; + const mcpServerStartupVersion = requestedMcpServers.length > 0 + ? this.codexAcpClient.getMcpServerStartupVersion() + : null; + let subscribedSessionId: string | null = null; + let sessionGeneration: number | null = null; + let sessionMetadata: SessionMetadata; + let sessionState: SessionState; + try { + sessionMetadata = await this.runWithProcessCheck(() => + this.codexAcpClient.forkSession(params, (sessionId) => { + subscribedSessionId = sessionId; + sessionGeneration = this.beginSessionOpen(sessionId); + }) + ); + const {sessionId} = sessionMetadata; + if (sessionGeneration === null) { + throw new Error("Codex session/fork did not report its child subscription"); + } + const authProvider = sessionMetadata.modelProvider ?? this.codexAcpClient.getModelProvider(); + const authState = await this.getAuthStateForProvider(authProvider); + if (!this.sessionOpenCanInstall(sessionId, sessionGeneration)) { + subscribedSessionId = null; + await this.closeStaleSessionOpen(sessionId, sessionGeneration); + } + sessionState = this.installSessionState({ + cwd: params.cwd, + sessionMetadata, + authState, + authProvider, + requestedMcpServers, + mcpServerStartupVersion, + recoverMcpServers: false, + sessionTitleSource: "unknown", + }); + subscribedSessionId = null; + } catch (err) { + if (subscribedSessionId !== null && sessionGeneration !== null) { + await this.cleanupStaleSessionOpen(subscribedSessionId, sessionGeneration); + } + throw err; + } + + this.publishAvailableCommandsAsync(sessionState); + this.publishCurrentGoalAsync(sessionState, sessionGeneration); + logger.log("Session forked", {parentSessionId: params.sessionId, sessionId: sessionMetadata.sessionId}); + return { + sessionId: sessionMetadata.sessionId, + modes: sessionState.agentMode.toSessionModeState(), + ...this.createSessionConfigOptionsResponse(sessionState), + }; + } + async listSessions(params: acp.ListSessionsRequest): Promise { logger.log("Listing sessions...", {cwd: params.cwd, cursor: params.cursor}); await this.checkAuthorization(); diff --git a/src/CodexAppServerClient.ts b/src/CodexAppServerClient.ts index eb26c83f..35b5ddec 100644 --- a/src/CodexAppServerClient.ts +++ b/src/CodexAppServerClient.ts @@ -41,6 +41,8 @@ import type { ThreadGoalGetResponse, ThreadGoalSetParams, ThreadGoalSetResponse, + ThreadForkParams, + ThreadForkResponse, ThreadLoadedListParams, ThreadLoadedListResponse, ThreadListParams, @@ -528,6 +530,10 @@ export class CodexAppServerClient { return await this.sendRequest({ method: "thread/resume", params: params }); } + async threadFork(params: ThreadForkParams): Promise { + return await this.sendRequest({ method: "thread/fork", params: params }); + } + getThreadSettings(threadId: string): ThreadSettings | undefined { return this.threadSettings.get(threadId); } diff --git a/src/__tests__/CodexACPAgent/data/session-fork-model-list-failed.json b/src/__tests__/CodexACPAgent/data/session-fork-model-list-failed.json new file mode 100644 index 00000000..e2668931 --- /dev/null +++ b/src/__tests__/CodexACPAgent/data/session-fork-model-list-failed.json @@ -0,0 +1,32 @@ +{ + "eventType": "request", + "method": "thread/fork", + "params": { + "config": { + "projects": { + "/workspace": { + "trust_level": "trusted" + } + } + }, + "cwd": "/workspace", + "ephemeral": false, + "modelProvider": "openai", + "threadId": "parent-session" + } +} +{ + "eventType": "response", + "placeholder": "thread/fork" +} +{ + "eventType": "request", + "method": "thread/unsubscribe", + "params": { + "threadId": "child-session" + } +} +{ + "eventType": "response", + "status": "unsubscribed" +} \ No newline at end of file diff --git a/src/__tests__/CodexACPAgent/data/session-fork.json b/src/__tests__/CodexACPAgent/data/session-fork.json new file mode 100644 index 00000000..adb84655 --- /dev/null +++ b/src/__tests__/CodexACPAgent/data/session-fork.json @@ -0,0 +1,29 @@ +{ + "eventType": "request", + "method": "thread/fork", + "params": { + "config": { + "projects": { + "/workspace": { + "trust_level": "trusted" + }, + "/shared": { + "trust_level": "trusted" + } + }, + "sandbox_workspace_write": { + "writable_roots": [ + "/shared" + ] + } + }, + "cwd": "/workspace", + "ephemeral": false, + "modelProvider": "openai", + "threadId": "parent-session" + } +} +{ + "eventType": "response", + "placeholder": "thread/fork" +} \ No newline at end of file diff --git a/src/__tests__/CodexACPAgent/initialize.test.ts b/src/__tests__/CodexACPAgent/initialize.test.ts index 9d6dc2b8..9aaea07e 100644 --- a/src/__tests__/CodexACPAgent/initialize.test.ts +++ b/src/__tests__/CodexACPAgent/initialize.test.ts @@ -49,6 +49,7 @@ describe('CodexACPAgent - initialize', () => { }, sessionCapabilities: { resume: {}, + fork: {}, list: {}, close: {}, delete: {}, diff --git a/src/__tests__/CodexACPAgent/session-fork.test.ts b/src/__tests__/CodexACPAgent/session-fork.test.ts new file mode 100644 index 00000000..0235f575 --- /dev/null +++ b/src/__tests__/CodexACPAgent/session-fork.test.ts @@ -0,0 +1,163 @@ +import {describe, expect, it, vi} from "vitest"; +import {createCodexMockTestFixture, createTestModel} from "../acp-test-utils"; +import type {ThreadForkResponse} from "../../app-server/v2"; + +const parentSessionId = "parent-session"; +const childSessionId = "child-session"; + +function createThreadForkResponse(): ThreadForkResponse { + return { + thread: { + id: childSessionId, + sessionId: childSessionId, + forkedFromId: parentSessionId, + parentThreadId: null, + preview: "", + ephemeral: false, + modelProvider: "openai", + createdAt: 0, + updatedAt: 0, + recencyAt: null, + status: {type: "idle"}, + path: "/sessions/child-session.jsonl", + cwd: "/workspace", + cliVersion: "test", + source: "appServer", + threadSource: null, + agentNickname: null, + agentRole: null, + gitInfo: null, + name: null, + turns: [], + }, + model: "model-id", + modelProvider: "openai", + serviceTier: null, + cwd: "/workspace", + instructionSources: [], + approvalPolicy: "on-request", + approvalsReviewer: "user", + sandbox: {type: "dangerFullAccess"}, + reasoningEffort: "medium", + }; +} + +describe("ACP session fork", () => { + it("advertises session fork support", async () => { + const fixture = createCodexMockTestFixture(); + + const response = await fixture.getCodexAcpAgent().initialize({protocolVersion: 1}); + + expect(response.agentCapabilities?.sessionCapabilities?.fork).toEqual({}); + }); + + it("forks a persistent Codex thread and installs the child session", async () => { + const fixture = createCodexMockTestFixture(); + const codexAcpAgent = fixture.getCodexAcpAgent(); + const codexAcpClient = fixture.getCodexAcpClient(); + const codexAppServerClient = fixture.getCodexAppServerClient(); + const model = createTestModel(); + vi.spyOn(codexAcpClient, "authRequired").mockResolvedValue(false); + vi.spyOn(codexAcpClient, "getAccount").mockResolvedValue({ + account: null, + requiresOpenaiAuth: false, + }); + vi.spyOn(codexAcpClient, "getCurrentModelProvider").mockResolvedValue("openai"); + vi.spyOn(codexAcpClient, "getGoal").mockResolvedValue(null); + vi.spyOn(codexAppServerClient, "skillsExtraRootsSet").mockResolvedValue(); + vi.spyOn(codexAppServerClient, "listSkills").mockResolvedValue({data: []}); + vi.spyOn(codexAppServerClient.connection, "sendRequest") + .mockResolvedValue(createThreadForkResponse()); + vi.spyOn(codexAppServerClient, "listModels").mockResolvedValue({ + data: [model], + nextCursor: null, + }); + + const response = await codexAcpAgent.unstable_forkSession({ + sessionId: parentSessionId, + cwd: "/workspace", + additionalDirectories: ["/shared"], + mcpServers: [], + }); + + await expect(fixture.getCodexConnectionDump([], { + placeholderResponseMethods: ["thread/fork"], + })).toMatchFileSnapshot("data/session-fork.json"); + expect(response).toMatchObject({ + sessionId: childSessionId, + modes: {currentModeId: "agent"}, + }); + expect(response.configOptions?.length).toBeGreaterThan(0); + expect(codexAcpAgent.getSessionState(childSessionId)).toMatchObject({ + sessionId: childSessionId, + cwd: "/workspace", + additionalDirectories: ["/shared"], + currentModelId: "model-id[medium]", + }); + }); + + it("unsubscribes the child when model discovery fails", async () => { + const fixture = createCodexMockTestFixture(); + const codexAcpAgent = fixture.getCodexAcpAgent(); + const codexAcpClient = fixture.getCodexAcpClient(); + const codexAppServerClient = fixture.getCodexAppServerClient(); + vi.spyOn(codexAcpClient, "authRequired").mockResolvedValue(false); + vi.spyOn(codexAcpClient, "getCurrentModelProvider").mockResolvedValue("openai"); + vi.spyOn(codexAppServerClient, "listSkills").mockResolvedValue({data: []}); + vi.spyOn(codexAppServerClient.connection, "sendRequest") + .mockResolvedValueOnce(createThreadForkResponse()) + .mockResolvedValueOnce({status: "unsubscribed"}); + vi.spyOn(codexAppServerClient, "listModels").mockRejectedValue(new Error("model list failed")); + + await expect(codexAcpAgent.unstable_forkSession({ + sessionId: parentSessionId, + cwd: "/workspace", + mcpServers: [], + })).rejects.toThrow("model list failed"); + + await expect(fixture.getCodexConnectionDump([], { + placeholderResponseMethods: ["thread/fork"], + })).toMatchFileSnapshot("data/session-fork-model-list-failed.json"); + expect(() => codexAcpAgent.getSessionState(childSessionId)).toThrow(`Session ${childSessionId} not found`); + }); + + it("normalizes authentication errors from Codex", async () => { + const fixture = createCodexMockTestFixture(); + const codexAcpAgent = fixture.getCodexAcpAgent(); + const codexAcpClient = fixture.getCodexAcpClient(); + const codexAppServerClient = fixture.getCodexAppServerClient(); + vi.spyOn(codexAcpClient, "authRequired").mockResolvedValue(false); + vi.spyOn(codexAcpClient, "getCurrentModelProvider").mockResolvedValue("openai"); + vi.spyOn(codexAppServerClient, "listSkills").mockResolvedValue({data: []}); + const errorMessage = "failed to reload config: Please log out and sign in again."; + vi.spyOn(codexAppServerClient.connection, "sendRequest") + .mockRejectedValue(new Error(errorMessage)); + const logout = vi.spyOn(codexAcpClient, "logout").mockResolvedValue(); + + await expect(codexAcpAgent.unstable_forkSession({ + sessionId: parentSessionId, + cwd: "/workspace", + mcpServers: [], + })).rejects.toMatchObject({ + data: expect.stringContaining("You have been logged out. Please try again."), + }); + expect(logout).toHaveBeenCalledOnce(); + }); + + it("rejects a fork that reuses the parent session id", async () => { + const fixture = createCodexMockTestFixture(); + const codexAcpClient = fixture.getCodexAcpClient(); + const codexAppServerClient = fixture.getCodexAppServerClient(); + const response = createThreadForkResponse(); + response.thread.id = parentSessionId; + vi.spyOn(codexAcpClient, "getCurrentModelProvider").mockResolvedValue("openai"); + vi.spyOn(codexAppServerClient, "listSkills").mockResolvedValue({data: []}); + vi.spyOn(codexAppServerClient, "threadFork").mockResolvedValue(response); + + await expect(codexAcpClient.forkSession({ + sessionId: parentSessionId, + cwd: "/workspace", + mcpServers: [], + }, vi.fn())).rejects.toThrow("Codex thread/fork did not return a child session id"); + }); +}); diff --git a/src/index.ts b/src/index.ts index 014801ff..f1c7620f 100644 --- a/src/index.ts +++ b/src/index.ts @@ -129,6 +129,7 @@ function startAcpServer() { .onRequest(acp.methods.agent.initialize, (ctx) => getAgent().initialize(ctx.params)) .onRequest(acp.methods.agent.session.new, (ctx) => getAgent().newSession(ctx.params)) .onRequest(acp.methods.agent.session.load, (ctx) => getAgent().loadSession(ctx.params)) + .onRequest(acp.methods.agent.session.fork, (ctx) => getAgent().unstable_forkSession(ctx.params)) .onRequest(acp.methods.agent.session.list, (ctx) => getAgent().listSessions(ctx.params)) .onRequest(acp.methods.agent.session.delete, (ctx) => getAgent().deleteSession(ctx.params)) .onRequest(acp.methods.agent.session.resume, (ctx) => getAgent().resumeSession(ctx.params))