perf: speed up chat loading and resume

src/server/sessions/piSessionService.lifecycle.test.ts

@@ -2,8 +2,18 @@ import { mkdtemp, rm, writeFile } from "node:fs/promises";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { describe, expect, it, vi } from "vitest";
-import { PiSessionService, type PiAgentSession } from "./piSessionService.js";
-import { CapturingSessionEventHub, fakeRuntime, fakeSessionManager, runtimeCreator, sessionGateway, sessionRecord, sessionRef, type RuntimeCreator } from "./piSessionService.testSupport.js";
+import { PiSessionService, type PiAgentSession, type PiSessionRuntime } from "./piSessionService.js";
+import { CapturingSessionEventHub, emptyArchiveStore, fakeRuntime, fakeSessionManager, runtimeCreator, sessionGateway, sessionRecord, sessionRef, type RuntimeCreator } from "./piSessionService.testSupport.js";
+
+function deferred<T = void>() {
+  let resolve!: (value: T | PromiseLike<T>) => void;
+  let reject!: (reason?: unknown) => void;
+  const promise = new Promise<T>((promiseResolve, promiseReject) => {
+    resolve = promiseResolve;
+    reject = promiseReject;
+  });
+  return { promise, resolve, reject };
+}
 
 describe("PiSessionService lifecycle, listing, and reload", () => {
   it("starts sessions through an injected runtime creator", async () => {
@@ -85,6 +95,156 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
     await service.dispose();
   });
 
+  it("shares one runtime when concurrent cold lookups resolve to the same session", async () => {
+    const sessionId = "single-flight-session";
+    const createStarted = deferred();
+    const releaseCreate = deferred();
+    const winnerUnsubscribe = vi.fn();
+    const loserUnsubscribe = vi.fn();
+    const winnerSubscribe = vi.fn(() => winnerUnsubscribe);
+    const loserSubscribe = vi.fn(() => loserUnsubscribe);
+    const winner = fakeRuntime(sessionId, {
+      sessionManager: fakeSessionManager("/workspace", {
+        getSessionId: () => sessionId,
+        getBranch: () => [{ type: "message", message: { role: "user", content: "shared runtime" } }],
+      }),
+      subscribe: winnerSubscribe,
+    });
+    const loser = fakeRuntime(sessionId, {
+      sessionManager: fakeSessionManager("/workspace", { getSessionId: () => sessionId }),
+      subscribe: loserSubscribe,
+    });
+    const runtimes = [winner.runtime, loser.runtime];
+    let createCalls = 0;
+    const createAgentRuntime: RuntimeCreator = async () => {
+      const runtime = runtimes[createCalls];
+      createCalls += 1;
+      createStarted.resolve();
+      await releaseCreate.promise;
+      if (runtime === undefined) throw new Error("unexpected runtime creation");
+      return runtime;
+    };
+    const gateway = sessionGateway([sessionRecord(sessionId)]);
+    const open = vi.spyOn(gateway, "open");
+    const service = new PiSessionService(new CapturingSessionEventHub(), {
+      archiveStore: emptyArchiveStore(),
+      createAgentRuntime,
+      sessionManager: gateway,
+      heartbeatIntervalMs: 60_000,
+    });
+
+    const messagesPromise = service.messages(sessionRef(sessionId));
+    await createStarted.promise;
+    const statusPromise = service.status(sessionRef("single-flight"));
+    await new Promise<void>((resolve) => setImmediate(resolve));
+    const callsWhileOpening = createCalls;
+    releaseCreate.resolve();
+
+    const [messages, status] = await Promise.all([messagesPromise, statusPromise]);
+    const activeCount = service.activeCount();
+    await service.dispose();
+
+    expect(callsWhileOpening).toBe(1);
+    expect(createCalls).toBe(1);
+    expect(open).toHaveBeenCalledOnce();
+    expect(activeCount).toBe(1);
+    expect(messages).toEqual([{ role: "user", content: "shared runtime" }]);
+    expect(status).toMatchObject({ sessionId });
+    expect(winnerSubscribe).toHaveBeenCalledOnce();
+    expect(winnerUnsubscribe).toHaveBeenCalledOnce();
+    expect(winner.calls.dispose).toBe(1);
+    expect(loserSubscribe).not.toHaveBeenCalled();
+    expect(loserUnsubscribe).not.toHaveBeenCalled();
+    expect(loser.calls.dispose).toBe(0);
+  });
+
+  it("clears a failed pending open so the session can be retried", async () => {
+    const sessionId = "retry-open-session";
+    const bindStarted = deferred();
+    const bindResult = deferred();
+    const openingError = new Error("extension binding failed");
+    const failed = fakeRuntime(sessionId, {
+      bindExtensions: () => {
+        bindStarted.resolve();
+        return bindResult.promise;
+      },
+    });
+    const retried = fakeRuntime(sessionId);
+    const runtimes = [failed.runtime, retried.runtime];
+    let createCalls = 0;
+    const createAgentRuntime: RuntimeCreator = () => {
+      const runtime = runtimes[createCalls];
+      createCalls += 1;
+      return runtime === undefined
+        ? Promise.reject(new Error("unexpected runtime creation"))
+        : Promise.resolve(runtime);
+    };
+    const service = new PiSessionService(new CapturingSessionEventHub(), {
+      archiveStore: emptyArchiveStore(),
+      createAgentRuntime,
+      sessionManager: sessionGateway([sessionRecord(sessionId)]),
+      heartbeatIntervalMs: 60_000,
+    });
+
+    const messagesPromise = service.messages(sessionRef(sessionId));
+    await bindStarted.promise;
+    const statusPromise = service.status(sessionRef("retry-open"));
+    await new Promise<void>((resolve) => setImmediate(resolve));
+    const callsWhileOpening = createCalls;
+    const failedLookups = Promise.allSettled([messagesPromise, statusPromise]);
+    bindResult.reject(openingError);
+
+    const outcomes = await failedLookups;
+    expect(callsWhileOpening).toBe(1);
+    expect(outcomes).toHaveLength(2);
+    for (const outcome of outcomes) {
+      expect(outcome.status).toBe("rejected");
+      if (outcome.status === "rejected") expect(outcome.reason).toBe(openingError);
+    }
+    expect(service.activeCount()).toBe(0);
+    expect(failed.calls.abort).toBe(1);
+    expect(failed.calls.dispose).toBe(1);
+
+    await expect(service.status(sessionRef(sessionId))).resolves.toMatchObject({ sessionId });
+    expect(createCalls).toBe(2);
+    expect(service.activeCount()).toBe(1);
+
+    await service.dispose();
+    expect(retried.calls.dispose).toBe(1);
+  });
+
+  it("waits for an in-flight open before disposing the service", async () => {
+    const sessionId = "dispose-opening-session";
+    const createStarted = deferred();
+    const runtimeResult = deferred<PiSessionRuntime>();
+    const fake = fakeRuntime(sessionId);
+    const service = new PiSessionService(new CapturingSessionEventHub(), {
+      archiveStore: emptyArchiveStore(),
+      createAgentRuntime: () => {
+        createStarted.resolve();
+        return runtimeResult.promise;
+      },
+      sessionManager: sessionGateway([sessionRecord(sessionId)]),
+      heartbeatIntervalMs: 60_000,
+    });
+
+    const statusPromise = service.status(sessionRef(sessionId));
+    await createStarted.promise;
+    let disposeSettled = false;
+    const disposePromise = service.dispose().then(() => { disposeSettled = true; });
+    await new Promise<void>((resolve) => setImmediate(resolve));
+    const settledWhileOpening = disposeSettled;
+    runtimeResult.resolve(fake.runtime);
+
+    await expect(statusPromise).resolves.toMatchObject({ sessionId });
+    await disposePromise;
+
+    expect(settledWhileOpening).toBe(false);
+    expect(service.activeCount()).toBe(0);
+    expect(fake.calls.abort).toBe(1);
+    expect(fake.calls.dispose).toBe(1);
+  });
+
   it("binds extensions again when the SDK runtime replaces the active session", async () => {
     const hub = new CapturingSessionEventHub();
     const fake = fakeRuntime("session-1");

src/server/sessions/piSessionService.ts

@@ -31,7 +31,7 @@ import { attachmentsToInlineImages, saveAttachmentsToWorkspace } from "./attachm
 import { parsePromptAttachments } from "../../shared/promptAttachments.js";
 import type { SavedPromptAttachment, SessionBulkArchiveResponse, SessionBulkDeleteArchivedResponse, SessionBulkFailure, SessionBulkMutationRef } from "../../shared/apiTypes.js";
 
-import { cwdPathsEqual } from "../workingDirectory.js";
+import { canonicalizeStoredCwd, cwdPathsEqual } from "../workingDirectory.js";
 import type { WorkspaceActivityService } from "../activity/workspaceActivityService.js";
 import { createSpawnSessionToolDefinition, type SpawnSessionInvocation, type SpawnSessionResult } from "./spawnSessionTool.js";
 import { createSubsessionToolDefinitions, type SpawnSubsessionInvocation, type SpawnSubsessionResult, type SubsessionCheckResult, type SubsessionReadQuery, type SubsessionReadResult, type SubsessionStatus, type SubsessionSummary, type SubsessionToolDeps } from "./spawnSubsessionTool.js";
@@ -261,6 +261,11 @@ export interface PiSessionRuntime {
   dispose(): Promise<void>;
 }
 
+interface PendingSessionOpen {
+  sessionId: string;
+  promise: Promise<ActiveSession<PiSessionRuntime>>;
+}
+
 interface CreateAgentRuntimeOptions {
   cwd: string;
   agentDir: string;
@@ -404,6 +409,7 @@ export interface PiSessionServiceDependencies {
 
 export class PiSessionService {
   private readonly active = new Map<string, ActiveSession<PiSessionRuntime>>();
+  private readonly pendingSessionOpens = new Map<string, PendingSessionOpen>();
   private readonly activities = new Map<string, { phase: "active" | "idle" | "error"; label: string; detail?: string; at: string }>();
   private readonly heartbeat: NodeJS.Timeout;
   private readonly commandService: SessionCommandService<PiAgentSession>;
@@ -533,8 +539,11 @@ export class PiSessionService {
   async dispose(): Promise<void> {
     clearInterval(this.heartbeat);
     this.clearCompactionDrainTimers();
+    const pendingOpens = this.pendingSessionOpenPromises();
+    if (pendingOpens.length > 0) await Promise.allSettled(pendingOpens);
     const activeSessions = Array.from(new Set(this.active.values()));
     this.active.clear();
+    this.pendingSessionOpens.clear();
     this.activities.clear();
     this.compactionPromptQueues.clear();
     this.authLossWarnings.clear();
@@ -546,8 +555,11 @@ export class PiSessionService {
     await Promise.all(activeSessions.map(async (active) => {
       active.unsubscribe();
       this.workspaceActivity?.removeSession(active.runtime.session.sessionId, active.runtime.session.sessionManager.getCwd());
-      await active.runtime.session.abort();
-      await active.runtime.dispose();
+      try {
+        await active.runtime.session.abort();
+      } finally {
+        await active.runtime.dispose();
+      }
     }));
   }
 
@@ -1540,6 +1552,8 @@ export class PiSessionService {
   }
 
   private async closeActive(sessionId: string): Promise<void> {
+    const pendingOpens = this.pendingSessionOpenPromises(sessionId);
+    if (pendingOpens.length > 0) await Promise.allSettled(pendingOpens);
     const active = this.active.get(sessionId);
     if (!active) return;
     this.active.delete(sessionId);
@@ -1573,13 +1587,49 @@ export class PiSessionService {
     if (active !== undefined) return active;
 
     const archived = await this.getArchived(ref);
-    if (archived?.archivePath !== undefined) return this.create(this.sessionManager.open(archived.archivePath), archived.cwd);
+    if (archived?.archivePath !== undefined) {
+      const { archivePath } = archived;
+      return this.openExistingSession(
+        archived.sessionId,
+        archived.cwd,
+        () => this.sessionManager.open(archivePath),
+      );
+    }
 
     const match = isPiSessionRef(ref)
       ? (await this.sessionManager.list(ref.cwd)).find((s) => s.id === ref.id || s.id.startsWith(ref.id))
       : (await this.sessionManager.listAll?.() ?? []).find((s) => s.id === ref || s.id.startsWith(ref));
     if (!match) throw new Error("Session not found");
-    return this.create(this.sessionManager.open(match.path), match.cwd);
+    return this.openExistingSession(match.id, match.cwd, () => this.sessionManager.open(match.path));
+  }
+
+  private openExistingSession(
+    sessionId: string,
+    cwd: string,
+    openSessionManager: () => PiSessionManager,
+  ): Promise<ActiveSession<PiSessionRuntime>> {
+    const active = this.activeForLookup({ id: sessionId, cwd });
+    if (active !== undefined) return Promise.resolve(active);
+
+    const key = JSON.stringify([canonicalizeStoredCwd(cwd), sessionId]);
+    const existing = this.pendingSessionOpens.get(key);
+    if (existing !== undefined) return existing.promise;
+
+    const pending: PendingSessionOpen = {
+      sessionId,
+      promise: this.create(openSessionManager(), cwd),
+    };
+    pending.promise = pending.promise.finally(() => {
+      if (this.pendingSessionOpens.get(key) === pending) this.pendingSessionOpens.delete(key);
+    });
+    this.pendingSessionOpens.set(key, pending);
+    return pending.promise;
+  }
+
+  private pendingSessionOpenPromises(sessionId?: string): Promise<ActiveSession<PiSessionRuntime>>[] {
+    return [...this.pendingSessionOpens.values()]
+      .filter((pending) => sessionId === undefined || pending.sessionId === sessionId)
+      .map((pending) => pending.promise);
   }
 
   private async getArchived(ref: PiSessionLookup): Promise<ArchivedSessionRecord | undefined> {
@@ -1613,18 +1663,40 @@ export class PiSessionService {
       delegationToolsEnabled,
       ...(options.initialModel === undefined ? {} : { initialModel: options.initialModel }),
     });
-    await this.bindSessionExtensions(runtime.session);
     const active: ActiveSession<PiSessionRuntime> = { runtime, unsubscribe: noop };
-    this.bindRuntime(active);
-    runtime.setRebindSession(async (session) => {
-      await this.bindSessionExtensions(session);
+    try {
+      await this.bindSessionExtensions(runtime.session);
       this.bindRuntime(active);
-      await this.recoverSubsessionTrackingForOpenedSession(session);
-    });
-    this.active.set(runtime.session.sessionId, active);
-    await this.recoverSubsessionTrackingForOpenedSession(runtime.session);
-    this.publishStatus(runtime.session);
-    return active;
+      runtime.setRebindSession(async (session) => {
+        await this.bindSessionExtensions(session);
+        this.bindRuntime(active);
+        await this.recoverSubsessionTrackingForOpenedSession(session);
+      });
+      this.active.set(runtime.session.sessionId, active);
+      await this.recoverSubsessionTrackingForOpenedSession(runtime.session);
+      this.publishStatus(runtime.session);
+      return active;
+    } catch (error: unknown) {
+      active.unsubscribe();
+      let removedActive = false;
+      for (const [sessionId, candidate] of this.active.entries()) {
+        if (candidate !== active) continue;
+        this.active.delete(sessionId);
+        this.activities.delete(sessionId);
+        this.clearAuthLossWarningsForSession(sessionId);
+        this.clearCompactionPromptQueue(sessionId);
+        removedActive = true;
+      }
+      if (removedActive) {
+        this.workspaceActivity?.removeSession(runtime.session.sessionId, runtime.session.sessionManager.getCwd());
+      }
+      try {
+        await runtime.session.abort();
+      } finally {
+        await runtime.dispose();
+      }
+      throw error;
+    }
   }
 
   private async bindSessionExtensions(session: PiAgentSession): Promise<void> {

src/server/sessions/sessionRoutes.test.ts

@@ -2,7 +2,7 @@ import { resolve } from "node:path";
 import Fastify, { type FastifyInstance } from "fastify";
 import fastifyWebsocket from "@fastify/websocket";
 import { afterEach, beforeEach, describe, expect, it } from "vitest";
-import type { SessionBulkArchiveResponse, SessionBulkDeleteArchivedResponse, SessionBulkMutationRef, SessionCleanupExecuteResponse, SessionCleanupPreviewResponse } from "../../shared/apiTypes.js";
+import type { MessagePage, SessionBulkArchiveResponse, SessionBulkDeleteArchivedResponse, SessionBulkMutationRef, SessionCleanupExecuteResponse, SessionCleanupPreviewResponse } from "../../shared/apiTypes.js";
 import { SessionEventHub } from "../realtime/sessionEventHub.js";
 import { PiSessionService, type PiSessionManagerGateway, type PiSessionRef } from "./piSessionService.js";
 import { registerSessionRoutes } from "./sessionRoutes.js";
@@ -55,6 +55,32 @@ describe("session routes", () => {
     }
   });
 
+  it("omits thinking signatures from browser history without mutating service messages", async () => {
+    const routeApp = Fastify({ logger: false });
+    await routeApp.register(fastifyWebsocket);
+    const eventHub = new SessionEventHub();
+    const routeService = new CapturingRouteSessionService(eventHub);
+    const thinkingBlock = { type: "thinking", thinking: "private chain", thinkingSignature: "opaque-provider-payload", redacted: true };
+    const message = { role: "assistant", content: [thinkingBlock, { type: "text", text: "visible answer" }] };
+    routeService.messagesResponse = { messages: [message], start: 0, total: 1 };
+    registerSessionRoutes(routeApp, routeService, eventHub);
+
+    try {
+      const response = await routeApp.inject({ method: "GET", url: "/sessions/session-1/messages?limit=20" });
+
+      expect(response.statusCode).toBe(200);
+      expect(response.json()).toEqual({
+        messages: [{ role: "assistant", content: [{ type: "thinking", thinking: "private chain", redacted: true }, { type: "text", text: "visible answer" }] }],
+        start: 0,
+        total: 1,
+      });
+      expect(thinkingBlock.thinkingSignature).toBe("opaque-provider-payload");
+    } finally {
+      await routeService.dispose();
+      await routeApp.close();
+    }
+  });
+
   it("forwards prompt attachments and supports the save-attachments route", async () => {
     const routeApp = Fastify({ logger: false });
     await routeApp.register(fastifyWebsocket);
@@ -226,6 +252,7 @@ describe("session routes", () => {
 class CapturingRouteSessionService extends PiSessionService {
   readonly calls: unknown[] = [];
   readonly reloadCalls: (string | PiSessionRef)[] = [];
+  messagesResponse: unknown[] | MessagePage = [];
   readonly cleanupPreviewCalls: NormalizedSessionCleanupRequest[] = [];
   readonly cleanupCalls: NormalizedSessionCleanupRequest[] = [];
   readonly bulkArchiveCalls: SessionBulkMutationRef[][] = [];
@@ -262,6 +289,10 @@ class CapturingRouteSessionService extends PiSessionService {
     return Promise.resolve();
   }
 
+  override messages(): Promise<unknown[] | MessagePage> {
+    return Promise.resolve(this.messagesResponse);
+  }
+
   override status(lookup: string | PiSessionRef) {
     this.calls.push(lookup);
     return Promise.resolve({

src/server/sessions/sessionRoutes.ts

@@ -1,5 +1,6 @@
 import type { FastifyInstance } from "fastify";
 import type { SessionBulkMutationRequest, SessionBulkMutationRef, SessionCleanupRequest } from "../../shared/apiTypes.js";
+import { projectBrowserMessageResponse } from "../browserMessageProjection.js";
 import { normalizeRequestCwd } from "../workingDirectory.js";
 import type { SessionEventHub } from "../realtime/sessionEventHub.js";
 import type { PiSessionRef, PiSessionService } from "./piSessionService.js";
@@ -83,7 +84,8 @@ export function registerSessionRoutes(app: FastifyInstance, sessions: PiSessionS
   app.get<{ Params: { sessionId: string }; Querystring: MessageQuery }>(`${prefix}/sessions/:sessionId/messages`, async (request, reply) => {
     try {
       const page = { ...optionalField("before", optionalNumber(request.query.before)), ...optionalField("limit", optionalNumber(request.query.limit)) };
-      return await sessions.messages(sessionLookupFromQuery(request.params.sessionId, request.query), page);
+      const messages = await sessions.messages(sessionLookupFromQuery(request.params.sessionId, request.query), page);
+      return projectBrowserMessageResponse(messages);
     } catch (error) {
       return reply.code(404).send({ error: errorMessage(error) });
     }