From 934e720cc225c4012cf5ae295845bbbb5551e66e Mon Sep 17 00:00:00 2001 From: Nikita Ashihmin Date: Fri, 16 Oct 2026 20:19:50 +0400 Subject: [PATCH] fix: hold a fork back from the session list until its session/fork response is out `session/fork` calls `thread/fork` and then sets the session up. Meanwhile the list subscription got `thread/started` of the fork and sent `_session/list/changes` with it before the fork response. A client then met an unknown session and listed it as a session of its own. JetBrains Air refused the fork after that (IJAI-4105). The subscriptions now keep the rows of a fork held until the response is out. A row is held by `forkedFrom` and, once `thread/fork` answers, by the fork thread id. The id matters because a row of a `thread/list` scan has no `forkedFromId`. `SessionForkOrigins` also learns the parent of the fork at once, so every row of the fork carries `forkedFrom` without a rollout read. --- src/CodexAcpClient.ts | 2 + src/CodexAcpServer.ts | 10 +++- src/SessionFork.ts | 3 + src/SessionForkOrigins.ts | 6 ++ src/SessionIndex.ts | 6 ++ src/SessionIndexService.ts | 24 ++++++++ src/SessionListSubscriptions.ts | 57 ++++++++++++++++++- .../CodexACPAgent/session-fork.test.ts | 2 +- .../SessionListSubscriptions.test.ts | 32 +++++++++++ 9 files changed, 137 insertions(+), 5 deletions(-) diff --git a/src/CodexAcpClient.ts b/src/CodexAcpClient.ts index 8dfc9cea..07a4acc7 100644 --- a/src/CodexAcpClient.ts +++ b/src/CodexAcpClient.ts @@ -711,6 +711,7 @@ export class CodexAcpClient { async forkSession( request: WithAcpMcpServers, releaseFailedFork: (threadId: string) => Promise = threadId => this.closeSession(threadId), + forkCreated: (threadId: string) => void = () => {}, ): Promise { const additionalDirectories = readAdditionalDirectories(request.cwd, request.additionalDirectories, request._meta); return await runForkSession(request, additionalDirectories, { @@ -724,6 +725,7 @@ export class CodexAcpClient { this.createModelId(models, model, reasoningEffort).toString(), getCollaborationMode: sessionId => this.getCollaborationMode(sessionId), releaseFailedFork, + forkCreated, }); } diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index 9abfa4e5..d7d466ad 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -1099,7 +1099,11 @@ export class CodexAcpServer { sessionMetadata = await this.runWithProcessCheck(() => { openGeneration = this.recovery?.generation ?? 0; mcpServerStartupVersion = readMcpServerStartupVersion(); - return this.codexAcpClient.forkSession(forkRequest, threadId => this.releaseFailedFork(threadId)); + return this.codexAcpClient.forkSession( + forkRequest, + threadId => this.releaseFailedFork(threadId), + threadId => this.sessionIndex.forkCreated(forkRequest.sessionId, threadId), + ); }); } else { logger.log(`Create new session...`); @@ -1682,6 +1686,10 @@ export class CodexAcpServer { } async forkSession(params: WithAcpMcpServers): Promise { + return this.sessionIndex.forking(params.sessionId, () => this.createFork(params)); + } + + private async createFork(params: WithAcpMcpServers): Promise { if (this.providerUpdate !== null) { await this.waitForProviderUpdate(); } diff --git a/src/SessionFork.ts b/src/SessionFork.ts index bf6cadf7..2974dd98 100644 --- a/src/SessionFork.ts +++ b/src/SessionFork.ts @@ -23,6 +23,8 @@ export type SessionForkDependencies = { getCollaborationMode(sessionId: string): ModeKind; /** Gives back the subscription of a fork that failed after `thread/fork`, unless another open of it began. */ releaseFailedFork(threadId: string): Promise; + /** `thread/fork` answered with the thread of the fork. */ + forkCreated(threadId: string): void; }; export async function forkSession( @@ -45,6 +47,7 @@ export async function forkSession( modelProvider: await dependencies.getResumeModelProvider(), threadId: request.sessionId, }); + dependencies.forkCreated(response.thread.id); // `thread/fork` subscribes this connection to the new thread, and only to it: the source thread keeps the // subscription it had. The fork stays subscribed, as ACP lets the client prompt it at once. Codex keeps an // unsubscribed thread loaded and runs its turns, but sends their notifications to subscribers only. diff --git a/src/SessionForkOrigins.ts b/src/SessionForkOrigins.ts index 447d5472..71043842 100644 --- a/src/SessionForkOrigins.ts +++ b/src/SessionForkOrigins.ts @@ -54,6 +54,12 @@ export class SessionForkOrigins { return null; } + /** This adapter forked `parentId` into `threadId`: the fork parent is known without a read of the rollout. */ + remember(threadId: string, parentId: string): void { + this.pending.delete(threadId); + remember(this.origins, threadId, parentId); + } + /** Drops the pending reads. */ dispose(): void { this.disposed = true; diff --git a/src/SessionIndex.ts b/src/SessionIndex.ts index b685e50d..34be5e23 100644 --- a/src/SessionIndex.ts +++ b/src/SessionIndex.ts @@ -577,6 +577,12 @@ export function sessionIndexRowSignature(row: acp.SessionInfo): string { ]); } +/** The AIR field `forkedFrom` of a row: the thread that its thread was forked from, else `null`. */ +export function sessionIndexRowForkedFrom(row: acp.SessionInfo): string | null { + const forkedFrom = asRecord(asRecord(asRecord(row._meta)[JETBRAINS_META_KEY])[AIR_META_KEY])[AIR_FORKED_FROM_KEY]; + return typeof forkedFrom === "string" ? forkedFrom : null; +} + function asRecord(value: unknown): Record { return value !== null && typeof value === "object" && !Array.isArray(value) ? value as Record diff --git a/src/SessionIndexService.ts b/src/SessionIndexService.ts index e416f50c..f052124b 100644 --- a/src/SessionIndexService.ts +++ b/src/SessionIndexService.ts @@ -185,6 +185,30 @@ export class SessionIndexService { return {}; } + /** + * Runs `fork`, a `session/fork` request of a thread. The list subscriptions hold back the rows of its forks until + * its response is out, see {@link SessionListSubscriptions.forkStarted}. + */ + async forking(sourceId: string, fork: () => Promise): Promise { + this.subscriptions.forkStarted(sourceId); + try { + return await fork(); + } finally { + // The response goes out when the handler returns, before any timer. + setTimeout(() => this.subscriptions.forkAnswered(sourceId), 0); + } + } + + /** + * `thread/fork` of a source answered with the thread of the fork. Every row of the fork carries `forkedFrom` from + * now on, also a row of `thread/list`, which answers no `forkedFromId`. The list subscriptions hold the fork back + * by its id until the fork response is out. + */ + forkCreated(sourceId: string, forkId: string): void { + this.forkOrigins.remember(forkId, sourceId); + this.subscriptions.forkCreated(sourceId, forkId); + } + /** The subscriptions of the connection, for tests. */ subscriptionResources(): ReturnType { return this.subscriptions.resources(); diff --git a/src/SessionListSubscriptions.ts b/src/SessionListSubscriptions.ts index 325c2305..3e50e861 100644 --- a/src/SessionListSubscriptions.ts +++ b/src/SessionListSubscriptions.ts @@ -26,7 +26,7 @@ import type {ServerNotification} from "./app-server"; import type {Thread, ThreadListParams, ThreadListResponse, ThreadReadParams, ThreadReadResponse} from "./app-server/v2"; import {CodexHomeWatcher, type CodexHomeWatcherListener} from "./CodexHomeWatcher"; import {logger} from "./Logger"; -import {sessionIndexRowSignature} from "./SessionIndex"; +import {sessionIndexRowForkedFrom, sessionIndexRowSignature} from "./SessionIndex"; import {isArchivedRolloutPath, isMissingThreadError} from "./SessionIndexMutations"; import {SessionNameLog} from "./SessionNameLog"; @@ -182,6 +182,11 @@ export class SessionListSubscriptions { private readonly lastChangeAt = new Map(); /** Threads of the own app-server that the list never shows: ephemeral, subagent and other helper threads. */ private readonly ignoredThreads = new Set(); + /** + * The threads with a `session/fork` in flight: the number of their forks in flight, and the threads that + * `thread/fork` created for them. See {@link forkStarted}. + */ + private readonly forking = new Map}>(); private disposed = false; constructor(private readonly deps: SessionListSubscriptionDeps) { @@ -345,6 +350,51 @@ export class SessionListSubscriptions { } } + /** + * A `session/fork` of a thread started. The rows of its forks wait in {@link Subscription.held} until + * {@link forkAnswered}: the client gets the id of a fork from the fork response before a list change shows the + * fork. Otherwise the client sees an unknown session and lists it as a session of its own. + */ + forkStarted(sourceId: string): void { + const source = this.forking.get(sourceId); + if (source !== undefined) source.count++; + else this.forking.set(sourceId, {count: 1, forks: new Set()}); + } + + /** + * `thread/fork` of a thread answered with the thread of the fork. Its rows wait by its id too: a row of a scan + * can come before the fork parent is known. + */ + forkCreated(sourceId: string, forkId: string): void { + this.forking.get(sourceId)?.forks.add(forkId); + } + + /** The response of a `session/fork` of a thread is out: the next flush sends the held rows of its forks. */ + forkAnswered(sourceId: string): void { + const source = this.forking.get(sourceId); + if (source === undefined) return; + if (--source.count > 0) return; + this.forking.delete(sourceId); + if ([...this.subscriptions.values()].some(subscription => subscription.held !== null)) this.requestThreads(0); + } + + /** Splits held changes into the ones to send now and the rows of forks whose fork response is not out yet. */ + private withoutForksInFlight(changes: SessionListChanges): {now: SessionListChanges | null, later: SessionListChanges | null} { + if (this.forking.size === 0) return {now: changes, later: null}; + const inFlight = (row: acp.SessionInfo) => { + const forkedFrom = sessionIndexRowForkedFrom(row); + return (forkedFrom !== null && this.forking.has(forkedFrom)) + || [...this.forking.values()].some(source => source.forks.has(row.sessionId)); + }; + const later = changes.sessions.filter(inFlight); + if (later.length === 0) return {now: changes, later: null}; + const sessions = changes.sessions.filter(row => !inFlight(row)); + return { + now: sessions.length === 0 && changes.removed.length === 0 ? null : {...changes, sessions}, + later: {subscriptionId: changes.subscriptionId, sessions: later, removed: []}, + }; + } + private threadChanged(threadId: string): void { // Named by a change of its own: no longer only a thread that a scan found. this.scanRetries.delete(threadId); @@ -607,8 +657,9 @@ export class SessionListSubscriptions { const sent: SessionListChanges[] = []; for (const subscription of this.subscriptions.values()) { if (this.disposed || !subscription.ready || subscription.held === null) continue; - const changes = subscription.held; - subscription.held = null; + const {now: changes, later} = this.withoutForksInFlight(subscription.held); + subscription.held = later; + if (changes === null) continue; await this.send(changes); sent.push(changes); } diff --git a/src/__tests__/CodexACPAgent/session-fork.test.ts b/src/__tests__/CodexACPAgent/session-fork.test.ts index b800681f..7e914610 100644 --- a/src/__tests__/CodexACPAgent/session-fork.test.ts +++ b/src/__tests__/CodexACPAgent/session-fork.test.ts @@ -42,7 +42,7 @@ describe("ACP session fork", () => { sessionId: "source-id", cwd: "/workspace", mcpServers: [], - }, expect.any(Function)); + }, expect.any(Function), expect.any(Function)); // Regression pin: v1 must keep sending `modes` (v2 drops it via `forkSessionV2`). expect(response.modes).toBeDefined(); }); diff --git a/src/__tests__/SessionListSubscriptions.test.ts b/src/__tests__/SessionListSubscriptions.test.ts index 11e26aa2..a922ceea 100644 --- a/src/__tests__/SessionListSubscriptions.test.ts +++ b/src/__tests__/SessionListSubscriptions.test.ts @@ -196,6 +196,38 @@ describe("SessionListSubscriptions", () => { subscriptions.dispose(); }); + it("holds the row of a fork back until the response of its session/fork is out", async () => { + const {codex, subscriptions, sent} = setup(); + // The resolved form, so that the subscription covers the threads on every OS. + const cwd = path.resolve("/repo"); + codex.put(thread("source", {cwd})); + const subscriptionId = await subscriptions.subscribe(cwd); + await vi.advanceTimersByTimeAsync(0); + + subscriptions.forkStarted("source"); + codex.put(thread("fork", {cwd, forkedFromId: "source"})); + // A fork whose row does not name its parent, as a row of `thread/list`: it waits by its id. + codex.put(thread("unmarked", {cwd})); + subscriptions.forkCreated("source", "unmarked"); + codex.put(thread("other", {cwd})); + for (const id of ["fork", "unmarked", "other"]) { + subscriptions.observe(own("thread/started", {thread: codex.threads.get(id)!.thread})); + } + await vi.advanceTimersByTimeAsync(2_000); + // A thread that is no fork of the source goes out meanwhile. + expect(sent).toEqual([{subscriptionId, sessions: [expect.objectContaining({sessionId: "other"})], removed: []}]); + + subscriptions.forkAnswered("source"); + await vi.advanceTimersByTimeAsync(0); + expect(sent).toHaveLength(2); + expect(sent[1]).toEqual({subscriptionId, sessions: expect.arrayContaining([ + expect.objectContaining({sessionId: "fork"}), + expect.objectContaining({sessionId: "unmarked"}), + ]), removed: []}); + expect((sent[1]!.sessions.find(row => row.sessionId === "fork")!._meta as any).jetbrains.air.forkedFrom).toBe("source"); + subscriptions.dispose(); + }); + it("sends a thread that another process changed after a WAL write, read by one thread/list without cwd", async () => { const {codex, subscriptions, sent, listener} = setup(); codex.put(thread("a", {updatedAt: 1_000}));