fix: hold a fork back from the session list until its session/fork response is out

src/CodexAcpClient.ts

@@ -711,6 +711,7 @@ export class CodexAcpClient {
     async forkSession(
         request: WithAcpMcpServers<acp.ForkSessionRequest>,
         releaseFailedFork: (threadId: string) => Promise<void> = threadId => this.closeSession(threadId),
+        forkCreated: (threadId: string) => void = () => {},
     ): Promise<SessionMetadata> {
         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,
         });
     }
 

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<acp.ForkSessionRequest>): Promise<acp.ForkSessionResponse> {
+        return this.sessionIndex.forking(params.sessionId, () => this.createFork(params));
+    }
+
+    private async createFork(params: WithAcpMcpServers<acp.ForkSessionRequest>): Promise<acp.ForkSessionResponse> {
         if (this.providerUpdate !== null) {
             await this.waitForProviderUpdate();
         }

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<void>;
+    /** `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.

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;

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<string, unknown> {
     return value !== null && typeof value === "object" && !Array.isArray(value)
         ? value as Record<string, unknown>

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<T>(sourceId: string, fork: () => Promise<T>): Promise<T> {
+        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<SessionListSubscriptions["resources"]> {
         return this.subscriptions.resources();

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<string, number>();
     /** Threads of the own app-server that the list never shows: ephemeral, subagent and other helper threads. */
     private readonly ignoredThreads = new Set<string>();
+    /**
+     * 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<string, {count: number, forks: Set<string>}>();
     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);
         }

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();
     });

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}));