fix(server,web): live in-progress output survives refresh and reconnect (#77)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

packages/server/src/api/types.ts

@@ -561,9 +561,42 @@ export interface SessionPatchRequest {
   title?: string;
 }
 
+/**
+ * Live in-progress tail of a running Session, carried by `MessagesResponse.live`.
+ *
+ * Contract (see runtime/live-tail.ts and the GET /messages route): the server captures
+ * `cursor` and `fragments` atomically — in one synchronous tick, before starting the
+ * trace read — while the Session is running/compacting.
+ *   - `cursor`: the Session channel's most recently assigned SSE event id
+ *     (`<epoch>-<seq>`); every event published up to and including this id is already
+ *     reflected in `fragments`.
+ *   - `fragments`: one synthetic `partial_* start` OmniMessage per open streaming
+ *     fragment, whose payload carries the full accumulated content so far (text/thinking
+ *     prefix, tool-call name + accumulated arguments, tool-output prefix + images), with
+ *     the original `origin` chain preserved.
+ *
+ * Client usage (the bundled Web App's connect-first flow): after applying `messages`,
+ * when the cursor's epoch matches the epoch of the SSE events seen on the current
+ * connection, drop every buffered **partial** event with seq <= cursor (its content is
+ * already inside `fragments`), feed `fragments` through the normal reducer path, then
+ * replay the rest of the buffer. Buffered **complete** messages are never dropped by the
+ * cursor — the regular overlap dedup decides for them — so nothing is lost even when a
+ * complete message's trace append is still in flight at read time.
+ */
+export interface MessagesLiveTail {
+  cursor: string;
+  fragments: OmniMessage[];
+}
+
 /** Message history: the full messages and events from concatenating all of this Session's Trace files in order (excludes partial_*). */
 export interface MessagesResponse {
   messages: OmniMessage[];
+  /**
+   * Present only while the Session is running/compacting: the in-progress stream tail
+   * (open streaming fragments + the channel cursor they cover), so a client joining
+   * mid-stream can render the currently streaming message. Omitted when idle.
+   */
+  live?: MessagesLiveTail;
 }
 
 // ---------------------------------------------------------------------------

packages/server/src/http/routes/sessions.ts

@@ -15,6 +15,7 @@ import type { OmniMessage, ThinkingLevelName } from "@prismshadow/penguin-core";
 import type {
   ApprovalMode,
   FilesStatResponse,
+  MessagesLiveTail,
   MessagesResponse,
   ServerEvent,
   SessionCategory,
@@ -337,12 +338,29 @@ export function sessionsRoutes(deps: AppDeps): Hono<AppEnv> {
 
   app.get("/:sessionId/messages", async (c) => {
     const row = resolveSession(c);
+    // Live tail (running/compacting sessions only): capture the channel cursor and the
+    // open-fragment snapshot together, synchronously — no await between the two, and both
+    // BEFORE the trace read starts. That ordering is what makes the client contract safe
+    // (see MessagesLiveTail in api/types.ts): every published event with id <= cursor is
+    // already reflected in `fragments`, and partial_* messages never reach the Trace, so
+    // the client may drop its buffered partials at/or before the cursor and seed from
+    // `fragments` without loss or duplication. Complete messages are never dropped by the
+    // cursor — the client's overlap dedup against `messages` decides for them — so a
+    // complete message whose trace append is still in flight when the read starts is not
+    // lost either.
+    let live: MessagesLiveTail | undefined;
+    if (deps.manager.statusOf(row.sessionId) !== "idle") {
+      live = {
+        cursor: deps.channels.get(row.sessionId).lastEventId,
+        fragments: deps.manager.liveFragments(row.sessionId),
+      };
+    }
     const messages = await deps.traceService.readMessages(
       row.projectId,
       row.agentId,
       row.sessionId,
     );
-    return c.json({ messages } satisfies MessagesResponse);
+    return c.json({ messages, ...(live !== undefined ? { live } : {}) } satisfies MessagesResponse);
   });
 
   app.get("/:sessionId/stream", (c) => {

packages/server/src/runtime/channel.ts

@@ -73,6 +73,17 @@ export class Channel {
     return this.listeners.size;
   }
 
+  /**
+   * Id of the most recently assigned event (`<epoch>-<seq>`; seq 0 when none was assigned
+   * yet). Unicast (sendTo) seqs count too — the value is a position marker, not a buffer
+   * lookup key. Used as the live-tail cursor on GET /messages: captured in the same
+   * synchronous tick as the fragment snapshot, it tells the client which buffered events
+   * the snapshot already covers.
+   */
+  get lastEventId(): string {
+    return `${this.epoch}-${this.nextSeq - 1}`;
+  }
+
   /** Broadcast an event: number it, buffer it (evicting the oldest), notify all subscribers. */
   publish(data: unknown, event?: string): ChannelEvent {
     const entry = this.makeEvent(data, event);

packages/server/src/runtime/live-tail.ts

@@ -0,0 +1,237 @@
+/**
+ * Live tail of a running Session's stream: the accumulated state of every OPEN
+ * streaming fragment (partial_text / partial_thinking / partial_tool_call /
+ * partial_tool_call_output), kept per origin chain.
+ *
+ * Why: `partial_*` messages never reach the Trace (core's Writer filters them), so a
+ * client that joins mid-stream — a refresh during a long tool call — cannot rebuild the
+ * in-progress message from `GET /messages` alone, and its fresh EventSource carries no
+ * `Last-Event-ID` for the channel buffer to replay. This tracker lets the messages
+ * endpoint attach, alongside history, one synthetic `partial_* start` OmniMessage per
+ * open fragment whose payload carries the full accumulated content so far (text/thinking
+ * prefix, tool-call name + accumulated arguments, tool-output prefix + images); the
+ * client seeds these into its stream model and live deltas keep appending on top. See
+ * `MessagesResponse.live` in api/types.ts for the client-facing contract.
+ *
+ * Fed by SessionManager.drive in the same synchronous tick as each channel publish, so a
+ * "channel cursor + fragments" capture between two publishes is always a consistent
+ * snapshot. Mirrors core PartialAggregator's merge semantics (fragment key = payload type
+ * + tool_call_id; start reopens, delta accumulates, stop closes) with the origin chain
+ * added to the key — the aggregator collapses fragments into complete messages, while
+ * this keeps the running prefix instead. A complete model message with the same identity
+ * also closes the fragment (covers stop-less closures, e.g. interruption cleanup); the
+ * whole session entry is dropped when the run ends (SessionManager.drive's finally).
+ */
+import { isPartialPayload } from "@prismshadow/penguin-core";
+import type {
+  CompleteModelPayload,
+  OmniMessage,
+  PartialModelPayload,
+} from "@prismshadow/penguin-core";
+
+type PartialKind = PartialModelPayload["type"];
+
+/**
+ * Per-fragment accumulation cap. Environment already front-truncates tool output online
+ * (default 16k chars), so tool fragments stay small by construction; text/thinking have
+ * no upstream cap and a fast large-code reply can reach a few hundred KB (see the channel
+ * buffer sizing note) — 512KB covers that comfortably while bounding a runaway fragment.
+ * When the cap trips, the TAIL is kept (the prefix is dropped): the seeded content then
+ * joins seamlessly with the live deltas that follow the capture, and the complete message
+ * reconciles the full content at the end anyway. Trimming happens with slack so it costs
+ * one copy per `FRAGMENT_CAP_SLACK` chars of growth, not one per delta. Tail-keeping is a
+ * deliberate one-rule-for-all: it is seamless for text/thinking/tool output, and for
+ * `partial_tool_call` arguments it yields a syntactically broken JSON prefix — accepted,
+ * since arguments approaching 512KB are not reachable in practice and the complete
+ * message reconciles either way.
+ */
+const FRAGMENT_CAP = 512 * 1024;
+const FRAGMENT_CAP_SLACK = 64 * 1024;
+
+interface OpenFragment {
+  kind: PartialKind;
+  /** Origin chain copied from the opening message (absent = main session). */
+  origin?: string[];
+  /** Timestamp of the fragment's original start (reused on the synthetic start). */
+  timestamp: string;
+  /** Accumulated text / thinking / tool-call arguments / tool output (tail-capped). */
+  buffer: string;
+  name?: string;
+  toolCallId?: string;
+  /** Tool-output images: not incremental — a single delta carries the whole set; a later one overwrites. */
+  images?: string[];
+}
+
+/** The partial kind a complete payload closes out; null when it has no streamed counterpart. */
+function partialKindFor(p: CompleteModelPayload): PartialKind | null {
+  switch (p.type) {
+    case "text":
+      // Only assistant text streams; a user text mid-run (steering) must not close the model's open fragment.
+      return p.role === "assistant" ? "partial_text" : null;
+    case "thinking":
+      return "partial_thinking";
+    case "tool_call":
+      return "partial_tool_call";
+    case "tool_call_output":
+      return "partial_tool_call_output";
+    default:
+      return null;
+  }
+}
+
+/** Fragment key: origin chain + payload type + tool_call_id (same merge rule as core's PartialAggregator, origin added). */
+function fragmentKey(origin: string[] | undefined, kind: PartialKind, toolCallId: string): string {
+  // "\0" (the escape, not a raw byte -- a raw NUL makes git classify the file as binary)
+  // separates the origin chain from the kind: session ids are [A-Za-z0-9_-], so the
+  // separator can never occur inside a chain segment and keys cannot collide.
+  return `${(origin ?? []).join("/")}\0${kind}::${toolCallId}`;
+}
+
+function toolCallIdOf(p: object): string {
+  const id = (p as { tool_call_id?: unknown }).tool_call_id;
+  return typeof id === "string" ? id : "";
+}
+
+function appendDelta(frag: OpenFragment, p: PartialModelPayload): void {
+  switch (p.type) {
+    case "partial_text":
+      frag.buffer += p.text;
+      break;
+    case "partial_thinking":
+      frag.buffer += p.thinking;
+      break;
+    case "partial_tool_call":
+      frag.buffer += p.arguments;
+      if (p.name) frag.name = p.name;
+      frag.toolCallId = p.tool_call_id;
+      break;
+    case "partial_tool_call_output":
+      frag.buffer += p.output;
+      if (p.images && p.images.length > 0) frag.images = p.images;
+      frag.toolCallId = p.tool_call_id;
+      break;
+  }
+  if (frag.buffer.length > FRAGMENT_CAP + FRAGMENT_CAP_SLACK) {
+    frag.buffer = frag.buffer.slice(frag.buffer.length - FRAGMENT_CAP);
+  }
+}
+
+/** The synthetic `partial_* start` payload carrying the fragment's full accumulated content. */
+function startPayload(frag: OpenFragment): PartialModelPayload {
+  switch (frag.kind) {
+    case "partial_text":
+      return { type: "partial_text", role: "assistant", event_type: "start", text: frag.buffer };
+    case "partial_thinking":
+      return {
+        type: "partial_thinking",
+        role: "assistant",
+        event_type: "start",
+        thinking: frag.buffer,
+      };
+    case "partial_tool_call":
+      return {
+        type: "partial_tool_call",
+        role: "assistant",
+        event_type: "start",
+        name: frag.name ?? "",
+        arguments: frag.buffer,
+        tool_call_id: frag.toolCallId ?? "",
+      };
+    case "partial_tool_call_output":
+      return {
+        type: "partial_tool_call_output",
+        role: "user",
+        event_type: "start",
+        output: frag.buffer,
+        ...(frag.images !== undefined && frag.images.length > 0 ? { images: frag.images } : {}),
+        tool_call_id: frag.toolCallId ?? "",
+      };
+  }
+}
+
+export class LiveTailTracker {
+  /** sessionId → open fragments keyed by fragmentKey, in the order they were opened. */
+  private readonly sessions = new Map<string, Map<string, OpenFragment>>();
+
+  /** Feed one published message (call in the same synchronous tick as the channel publish). */
+  observe(sessionId: string, msg: OmniMessage): void {
+    if (msg.type !== "model_msg") return;
+    const p = msg.payload;
+    if (!isPartialPayload(p)) {
+      // A complete message closes the matching open fragment (normally the stop already
+      // did; this also covers stop-less closures such as interruption cleanup).
+      const kind = partialKindFor(p as CompleteModelPayload);
+      if (kind === null) return;
+      const open = this.sessions.get(sessionId);
+      if (!open) return;
+      open.delete(fragmentKey(msg.origin, kind, toolCallIdOf(p)));
+      if (open.size === 0) this.sessions.delete(sessionId);
+      return;
+    }
+    const key = fragmentKey(msg.origin, p.type, toolCallIdOf(p));
+    let open = this.sessions.get(sessionId);
+    if (p.event_type === "start") {
+      if (!open) {
+        open = new Map();
+        this.sessions.set(sessionId, open);
+      }
+      // start reopens the fragment: a previous same-key fragment (out-of-order) is replaced.
+      const frag: OpenFragment = {
+        kind: p.type,
+        timestamp: msg.timestamp,
+        buffer: "",
+        ...(msg.origin && msg.origin.length > 0 ? { origin: [...msg.origin] } : {}),
+      };
+      open.delete(key); // re-insert so the emit order tracks the reopen
+      open.set(key, frag);
+      appendDelta(frag, p);
+      return;
+    }
+    let frag = open?.get(key);
+    if (!frag) {
+      // delta/stop without a start: lenient, same as core's PartialAggregator.
+      if (p.event_type === "stop") return; // nothing was open; nothing to keep or clear
+      if (!open) {
+        open = new Map();
+        this.sessions.set(sessionId, open);
+      }
+      frag = {
+        kind: p.type,
+        timestamp: msg.timestamp,
+        buffer: "",
+        ...(msg.origin && msg.origin.length > 0 ? { origin: [...msg.origin] } : {}),
+      };
+      open.set(key, frag);
+    }
+    appendDelta(frag, p);
+    if (p.event_type === "stop") {
+      open!.delete(key);
+      if (open!.size === 0) this.sessions.delete(sessionId);
+    }
+  }
+
+  /**
+   * Synthetic `partial_* start` messages for every open fragment (in open order), each
+   * carrying the full accumulated content, the original origin chain, and the original
+   * start timestamp. Empty when the session has no open fragments.
+   */
+  fragments(sessionId: string): OmniMessage[] {
+    const open = this.sessions.get(sessionId);
+    if (!open) return [];
+    const out: OmniMessage[] = [];
+    for (const frag of open.values()) {
+      out.push({
+        timestamp: frag.timestamp,
+        type: "model_msg",
+        payload: startPayload(frag),
+        ...(frag.origin !== undefined ? { origin: [...frag.origin] } : {}),
+      });
+    }
+    return out;
+  }
+
+  /** Drop all fragment state for a session (the run ended; nothing will continue these fragments). */
+  clear(sessionId: string): void {
+    this.sessions.delete(sessionId);
+  }
+}

packages/server/src/runtime/session-manager.ts

@@ -49,6 +49,7 @@ import { ApprovalRegistry, makeApprove } from "./approvals.js";
 import type { PendingApproval } from "./approvals.js";
 import type { ChannelHub } from "./channel.js";
 import type { ErrorSink } from "./error-recorder.js";
+import { LiveTailTracker } from "./live-tail.js";
 import { asSessionSource } from "./session-sources.js";
 import type { SessionSources } from "./session-sources.js";
 import { StreamErrorWatcher } from "./stream-error-watcher.js";
@@ -329,6 +330,8 @@ export class SessionManager {
   private readonly deletingSessions = new Set<string>();
   /** Per-Agent config generation (key = agentKey), bumped by invalidateAgentRuntimes on vault updates. */
   private readonly agentGenerations = new Map<string, number>();
+  /** Open streaming fragments of running sessions (fed by drive, served to GET /messages; see live-tail.ts). */
+  private readonly liveTail = new LiveTailTracker();
   private readonly sweepTimer: NodeJS.Timeout;
 
   constructor(private readonly deps: SessionManagerDeps) {
@@ -356,6 +359,16 @@ export class SessionManager {
     return this.entries.get(sessionId)?.followUps.length ?? 0;
   }
 
+  /**
+   * Live tail of a running session: one synthetic `partial_* start` OmniMessage per open
+   * streaming fragment, carrying the full accumulated content so far (see live-tail.ts).
+   * Empty when idle or when nothing is streaming. GET /messages attaches this (with a
+   * channel cursor) so a client joining mid-stream can render the in-progress message.
+   */
+  liveFragments(sessionId: string): OmniMessage[] {
+    return this.liveTail.fragments(sessionId);
+  }
+
   /** Number of Sessions for this Agent that are currently running / compacting. */
   activeCountForAgent(projectId: string, agentId: string): number {
     let n = 0;
@@ -944,6 +957,10 @@ export class SessionManager {
             }
           }
         }
+        // Live-tail bookkeeping in the same synchronous tick as the publish below: the
+        // messages endpoint captures "channel cursor + open fragments" between two
+        // publishes, so the pair is always a consistent snapshot (see live-tail.ts).
+        this.liveTail.observe(entry.sessionId, msg);
         // Re-fetch the channel before every publish (matches publishEvent): the channel
         // may have been recycled and recreated during a long wait on approval, and
         // holding a stale reference would send output to an orphaned, detached channel.
@@ -966,6 +983,9 @@ export class SessionManager {
     } finally {
       // Wrap-up: persist any still-pending LLM failure and clear the tool-name cache (the watcher's state doesn't carry across runs).
       watcher?.close();
+      // The run is over: no fragment will ever continue, so drop the live tail before the
+      // idle flip (GET /messages stops attaching `live` the moment status reads idle).
+      this.liveTail.clear(entry.sessionId);
       entry.approvals.denyAll();
       entry.status = "idle";
       entry.abort = null;