feat(sessions): stream in-flight partial when joining a mid-turn session

Seed the in-flight partial assistant message (text, thinking, and
in-progress tool calls) when opening or reconnecting to a session that is
mid-stream, then continue streaming live deltas on top of it. Replaces the
blocking "Catching up..." placeholder and the end-of-turn transcript reload.

Server stamps every per-session UI event with a monotonic seq at the
SessionEventHub publish choke point and exposes
GET /sessions/:sessionId/stream-snapshot returning { seq, partial }. The
client fetches the snapshot on join, seeds the normalized partial into the
in-memory transcript (never the history cache), and applies buffered/live
events using the seq watermark for exactly-once delivery.

The snapshot is a progressive enhancement: a 404 from an older remote
pi-web or a not-yet-restarted session daemon falls back to an empty seed
(seq 0, drops nothing), so sessions still open and stream normally. The
stream-snapshot route is registered in the federation allowlist for
remote-machine proxying.

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

@@ -566,3 +566,65 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
     await service.dispose();
   });
 });
+
+describe("PiSessionService.streamSnapshot", () => {
+  it("returns a null partial with the current watermark when idle", async () => {
+    const hub = new CapturingSessionEventHub();
+    const fake = fakeRuntime("snap-idle");
+    const service = new PiSessionService(hub, {
+      agentDir: TEST_AGENT_DIR,
+      createAgentRuntime: runtimeCreator(fake.runtime),
+      sessionManager: sessionGateway([]),
+      heartbeatIntervalMs: 60_000,
+    });
+    try {
+      await service.start("/workspace");
+
+      const snapshot = await service.streamSnapshot(sessionRef("snap-idle"));
+
+      expect(snapshot).toEqual({ seq: 0, partial: null });
+    } finally {
+      await service.dispose();
+    }
+  });
+
+  it("projects the in-flight partial and matches the event watermark mid-stream", async () => {
+    const hub = new CapturingSessionEventHub();
+    const streamingMessage = {
+      role: "assistant",
+      content: [
+        { type: "thinking", thinking: "weighing options", thinkingSignature: "opaque" },
+        { type: "text", text: "partial answer" },
+        { type: "toolCall", id: "call-1", name: "edit", arguments: { path: "a.ts" } },
+      ],
+    };
+    const fake = fakeRuntime("snap-live", { state: { streamingMessage } });
+    const service = new PiSessionService(hub, {
+      agentDir: TEST_AGENT_DIR,
+      createAgentRuntime: runtimeCreator(fake.runtime),
+      sessionManager: sessionGateway([]),
+      heartbeatIntervalMs: 60_000,
+    });
+    try {
+      await service.start("/workspace");
+      // Advance the per-session watermark to a known value.
+      hub.setSeq("snap-live", 5);
+
+      const snapshot = await service.streamSnapshot(sessionRef("snap-live"));
+
+      expect(snapshot.seq).toBe(5);
+      expect(snapshot.partial).toEqual({
+        role: "assistant",
+        content: [
+          { type: "thinking", thinking: "weighing options" },
+          { type: "text", text: "partial answer" },
+          { type: "toolCall", id: "call-1", name: "edit", arguments: { path: "a.ts" } },
+        ],
+      });
+      // The runtime message is not mutated by the browser projection.
+      expect(streamingMessage.content[0]).toHaveProperty("thinkingSignature", "opaque");
+    } finally {
+      await service.dispose();
+    }
+  });
+});

src/server/sessions/piSessionService.testSupport.ts

@@ -6,6 +6,7 @@ import type { PiAgentSession, PiSessionManager, PiSessionRuntime, PiSessionServi
 export class CapturingSessionEventHub extends SessionEventHub {
   readonly sessionEvents: { sessionId: string; event: SessionUiEvent }[] = [];
   readonly globalEvents: GlobalSessionEvent[] = [];
+  private readonly seqBySessionOverride = new Map<string, number>();
 
   override publish(sessionId: string, event: SessionUiEvent): void {
     this.sessionEvents.push({ sessionId, event });
@@ -14,6 +15,15 @@ export class CapturingSessionEventHub extends SessionEventHub {
   override publishGlobal(event: GlobalSessionEvent): void {
     this.globalEvents.push(event);
   }
+
+  /** Test seam: set the per-session watermark returned by {@link currentSeq}. */
+  setSeq(sessionId: string, value: number): void {
+    this.seqBySessionOverride.set(sessionId, value);
+  }
+
+  override currentSeq(sessionId: string): number {
+    return this.seqBySessionOverride.get(sessionId) ?? 0;
+  }
 }
 
 export type SessionGateway = NonNullable<PiSessionServiceDependencies["sessionManager"]>;
@@ -68,6 +78,7 @@ export function fakeRuntime(sessionId = "session-1", patch: Partial<TestSession>
     sessionId,
     sessionFile: `/tmp/${sessionId}.jsonl`,
     messages: [],
+    state: {},
     sessionName: undefined,
     model: undefined,
     thinkingLevel: "off",

src/server/sessions/piSessionService.ts

@@ -14,7 +14,8 @@ import {
   type CreateAgentSessionRuntimeFactory,
   type EditToolDetails,
 } from "@earendil-works/pi-coding-agent";
-import type { ClientArchiveSessionsResponse, ClientCommand, ClientCommandResult, ClientMessagePage, ClientSession, ClientSessionCleanupExecuteResponse, ClientSessionCleanupPreviewResponse, ClientSessionModel, ClientSessionStatus, ClientThinkingLevel, SessionUiEvent } from "../types.js";
+import type { ClientArchiveSessionsResponse, ClientCommand, ClientCommandResult, ClientMessagePage, ClientSession, ClientSessionCleanupExecuteResponse, ClientSessionCleanupPreviewResponse, ClientSessionModel, ClientSessionStatus, ClientThinkingLevel, SessionStreamSnapshot, SessionUiEvent } from "../types.js";
+import { projectBrowserMessage } from "../browserMessageProjection.js";
 import { pageMessagesAtSafeBoundary } from "./messagePaging.js";
 import type { SessionEventHub } from "../realtime/sessionEventHub.js";
 import { BUILTIN_COMMANDS } from "./builtinCommands.js";
@@ -211,6 +212,14 @@ export interface PiAgentSession {
   sessionFile: string | undefined;
   sessionName: string | undefined;
   messages: readonly unknown[];
+  /**
+   * Narrow read of the SDK `AgentState`. Only the in-flight partial is consumed
+   * here: `state.streamingMessage` is the current streamed assistant message
+   * (an `AssistantMessage`) while a turn is mid-stream, and `undefined`
+   * otherwise (idle, or during post-message tool execution). Used by
+   * {@link PiSessionService.streamSnapshot} to seed a joining client.
+   */
+  readonly state: { readonly streamingMessage?: unknown };
   model: AgentModel | undefined;
   thinkingLevel: ClientThinkingLevel;
   isStreaming: boolean;
@@ -981,6 +990,26 @@ export class PiSessionService implements SessionRouteService {
     return this.statusFromSession(await this.getOrOpen(ref));
   }
 
+  /**
+   * Join-time snapshot of the in-flight assistant stream. The `seq` watermark and
+   * the partial are read together in one synchronous tick (no await between the
+   * `currentSeq` read and the `state.streamingMessage` read) so a joining client
+   * can seed the partial and then apply only buffered live events with
+   * `seq > snapshot.seq`. The partial is browser-projected to strip thinking
+   * signatures; it is `null` when no assistant message is mid-stream.
+   */
+  async streamSnapshot(ref: PiSessionLookup): Promise<SessionStreamSnapshot> {
+    const session = await this.getOrOpen(ref);
+    // Single consistent tick: capture the watermark and the partial together so
+    // the seq matches the partial the client seeds against.
+    const seq = this.events.currentSeq(session.sessionId);
+    const streamingMessage = session.state.streamingMessage;
+    const partial = streamingMessage === undefined || streamingMessage === null
+      ? null
+      : projectBrowserMessage(streamingMessage);
+    return { seq, partial };
+  }
+
   async availableModels(ref: PiSessionLookup): Promise<ClientSessionModel[]> {
     const session = await this.getOrOpen(ref);
     session.modelRegistry.refresh();

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 { MessagePage, SessionBulkArchiveResponse, SessionBulkDeleteArchivedResponse, SessionBulkMutationRef, SessionCleanupExecuteResponse, SessionCleanupPreviewResponse, SessionStatus } from "../../shared/apiTypes.js";
+import type { MessagePage, SessionBulkArchiveResponse, SessionBulkDeleteArchivedResponse, SessionBulkMutationRef, SessionCleanupExecuteResponse, SessionCleanupPreviewResponse, SessionStatus, SessionStreamSnapshot } from "../../shared/apiTypes.js";
 import { SessionEventHub } from "../realtime/sessionEventHub.js";
 import { PiSessionService, type PiSessionManagerGateway } from "./piSessionService.js";
 import type { SessionRouteLookup, SessionRouteService } from "./sessionService.js";
@@ -168,6 +168,46 @@ describe("session routes", () => {
     }
   });
 
+  it("returns the join-time stream snapshot, forwarding workspace context", async () => {
+    const routeApp = Fastify({ logger: false });
+    await routeApp.register(fastifyWebsocket);
+    const eventHub = new SessionEventHub();
+    const routeService = new CapturingRouteSessionService();
+    routeService.streamSnapshotResponse = { seq: 7, partial: { role: "assistant", content: [{ type: "text", text: "partial" }] } };
+    registerSessionRoutes(routeApp, routeService, eventHub);
+
+    try {
+      const requestCwd = resolve("/repo");
+      const response = await routeApp.inject({ method: "GET", url: `/sessions/session-1/stream-snapshot?cwd=${encodeURIComponent(requestCwd)}` });
+
+      expect(response.statusCode).toBe(200);
+      expect(response.json()).toEqual({ seq: 7, partial: { role: "assistant", content: [{ type: "text", text: "partial" }] } });
+      expect(routeService.streamSnapshotCalls).toEqual([{ id: "session-1", cwd: requestCwd }]);
+    } finally {
+      await routeService.dispose();
+      await routeApp.close();
+    }
+  });
+
+  it("maps stream-snapshot lookup failures to 404", async () => {
+    const routeApp = Fastify({ logger: false });
+    await routeApp.register(fastifyWebsocket);
+    const eventHub = new SessionEventHub();
+    const routeService = new CapturingRouteSessionService();
+    routeService.streamSnapshot = () => Promise.reject(new Error("Session not found"));
+    registerSessionRoutes(routeApp, routeService, eventHub);
+
+    try {
+      const response = await routeApp.inject({ method: "GET", url: "/sessions/missing/stream-snapshot" });
+
+      expect(response.statusCode).toBe(404);
+      expect(response.json()).toEqual({ error: "Session not found" });
+    } finally {
+      await routeService.dispose();
+      await routeApp.close();
+    }
+  });
+
   it("clears a session queue with workspace context and returns fresh status", async () => {
     const routeApp = Fastify({ logger: false });
     await routeApp.register(fastifyWebsocket);
@@ -306,6 +346,8 @@ class CapturingRouteSessionService implements SessionRouteService {
   readonly reloadCalls: SessionRouteLookup[] = [];
   readonly clearQueueCalls: SessionRouteLookup[] = [];
   messagesResponse: unknown[] | MessagePage = [];
+  streamSnapshotResponse: SessionStreamSnapshot = { seq: 0, partial: null };
+  readonly streamSnapshotCalls: SessionRouteLookup[] = [];
   readonly cleanupPreviewCalls: NormalizedSessionCleanupRequest[] = [];
   readonly cleanupCalls: NormalizedSessionCleanupRequest[] = [];
   readonly bulkArchiveCalls: SessionBulkMutationRef[][] = [];
@@ -379,6 +421,11 @@ class CapturingRouteSessionService implements SessionRouteService {
     });
   }
 
+  streamSnapshot(lookup: SessionRouteLookup): Promise<SessionStreamSnapshot> {
+    this.streamSnapshotCalls.push(lookup);
+    return Promise.resolve(this.streamSnapshotResponse);
+  }
+
   availableModels(): Promise<[]> { return Promise.resolve([]); }
   setModel(): never { throw unusedRouteMethod("setModel"); }
   cycleModel(): never { throw unusedRouteMethod("cycleModel"); }

src/server/sessions/sessionRoutes.ts

@@ -99,6 +99,14 @@ export function registerSessionRoutes(app: FastifyInstance, sessions: SessionRou
     }
   });
 
+  app.get<{ Params: { sessionId: string }; Querystring: SessionQuery }>(`${prefix}/sessions/:sessionId/stream-snapshot`, async (request, reply) => {
+    try {
+      return await sessions.streamSnapshot(sessionLookupFromQuery(request.params.sessionId, request.query));
+    } catch (error) {
+      return reply.code(404).send({ error: errorMessage(error) });
+    }
+  });
+
   app.get<{ Params: { sessionId: string }; Querystring: SessionQuery }>(`${prefix}/sessions/:sessionId/models`, async (request, reply) => {
     try {
       return { models: await sessions.availableModels(sessionLookupFromQuery(request.params.sessionId, request.query)) };

src/server/sessions/sessionService.ts

@@ -16,6 +16,7 @@ import type {
   ClientSessionRef,
   ClientSessionStatus,
   ClientThinkingLevel,
+  SessionStreamSnapshot,
 } from "../types.js";
 import type { NormalizedSessionCleanupRequest } from "./sessionCleanup.js";
 
@@ -34,6 +35,7 @@ export interface SessionRouteService {
   start(cwd: string): Promise<ClientSession>;
   messages(ref: SessionRouteLookup, page?: { before?: number; limit?: number }): Promise<unknown[] | ClientMessagePage>;
   status(ref: SessionRouteLookup): Promise<ClientSessionStatus>;
+  streamSnapshot(ref: SessionRouteLookup): Promise<SessionStreamSnapshot>;
   clearQueue(ref: SessionRouteLookup): Promise<ClientSessionStatus>;
   availableModels(ref: SessionRouteLookup): Promise<ClientSessionModel[]>;
   setModel(ref: SessionRouteLookup, provider: string, modelId: string): Promise<ClientSessionStatus>;