Handle reconnecting mid-stream

src/client/src/appState.ts

@@ -10,6 +10,7 @@ export interface AppState {
   messagePageStart: number;
   messagePageTotal: number;
   isLoadingEarlierMessages: boolean;
+  isReceivingPartialStream: boolean;
   selectedProject: Project | undefined;
   selectedWorkspace: Workspace | undefined;
   selectedSession: SessionInfo | undefined;
@@ -44,6 +45,7 @@ export function initialAppState(): AppState {
     messagePageStart: 0,
     messagePageTotal: 0,
     isLoadingEarlierMessages: false,
+    isReceivingPartialStream: false,
     selectedProject: undefined,
     selectedWorkspace: undefined,
     selectedSession: undefined,

src/client/src/chatHistoryCache.test.ts

@@ -12,9 +12,22 @@ describe("mergeChatHistory", () => {
     expect(merged).toEqual(page(0, 5, ["a", "b", "c", "d", "e"]));
   });
 
-  it("uses incoming history when totals changed", () => {
+  it("keeps cached history when new messages were appended", () => {
+    const existing = page(0, 3, ["a", "b", "c"]);
+    const incoming = page(1, 4, ["b", "c", "d"]);
+
+    expect(mergeChatHistory(existing, incoming)).toEqual(page(0, 4, ["a", "b", "c", "d"]));
+  });
+
+  it("uses incoming history when totals shrink", () => {
     const incoming = page(0, 2, ["fresh-a", "fresh-b"]);
 
-    expect(mergeChatHistory(page(0, 1, ["stale"]), incoming)).toEqual(incoming);
+    expect(mergeChatHistory(page(0, 3, ["stale-a", "stale-b", "stale-c"]), incoming)).toEqual(incoming);
+  });
+
+  it("uses incoming history instead of creating a gapped page", () => {
+    const incoming = page(8, 10, ["i", "j"]);
+
+    expect(mergeChatHistory(page(0, 10, ["a", "b"]), incoming)).toEqual(incoming);
   });
 });

src/client/src/chatHistoryCache.ts

@@ -36,7 +36,8 @@ export function writeChatHistoryCache(sessionId: string, page: RawMessagePage):
 }
 
 export function mergeChatHistory(existing: RawMessagePage | undefined, incoming: RawMessagePage): RawMessagePage {
-  if (existing?.total !== incoming.total) return incoming;
+  if (existing === undefined) return incoming;
+  if (existing.total > incoming.total) return incoming;
 
   const start = Math.min(existing.start, incoming.start);
   const end = Math.max(existing.start + existing.messages.length, incoming.start + incoming.messages.length);
@@ -44,7 +45,15 @@ export function mergeChatHistory(existing: RawMessagePage | undefined, incoming:
   copyInto(messages, start, existing);
   copyInto(messages, start, incoming);
 
-  return { start, total: incoming.total, messages: messages.filter((message) => message !== undefined) };
+  if (hasSparseEntries(messages)) return incoming;
+  return { start, total: incoming.total, messages };
+}
+
+function hasSparseEntries(messages: unknown[]): boolean {
+  for (let index = 0; index < messages.length; index += 1) {
+    if (!(index in messages) || messages[index] === undefined) return true;
+  }
+  return false;
 }
 
 function copyInto(target: unknown[], targetStart: number, page: RawMessagePage): void {

src/client/src/components/ChatView.ts

@@ -28,6 +28,7 @@ export class ChatView extends LitElement {
   @property({ type: Number }) messageTotal = 0;
   @property({ type: Boolean }) hasMore = false;
   @property({ type: Boolean }) loadingMore = false;
+  @property({ type: Boolean }) isReceivingPartialStream = false;
   @property({ type: Boolean }) isCompacting = false;
   @property({ type: Number }) pendingMessageCount = 0;
   @property({ attribute: false }) status?: SessionStatus;
@@ -85,6 +86,12 @@ export class ChatView extends LitElement {
   }
 
   private renderSessionActivity() {
+    if (this.isReceivingPartialStream) return html`
+      <aside class="session-activity receiving" aria-live="polite">
+        <strong>Receiving answer…</strong>
+        <span>This session was reconnected mid-response. The answer will appear when complete.</span>
+      </aside>
+    `;
     if (!this.isCompacting) return null;
     return html`
       <aside class="session-activity compacting" aria-live="polite">

src/client/src/components/PiWebApp.ts

@@ -251,7 +251,7 @@ export class PiWebApp extends LitElement {
           ${state.error ? html`<div class="error">${state.error}</div>` : null}
           <div class="mobile-navigation-panel">${this.renderNavigationPanel(true)}</div>
           ${state.selectedSession ? html`
-            <chat-view .sessionId=${state.selectedSession.id} .messages=${state.messages} .messageStart=${state.messagePageStart} .messageTotal=${state.messagePageTotal} .hasMore=${state.messagePageStart > 0} .loadingMore=${state.isLoadingEarlierMessages} .isCompacting=${state.status?.isCompacting === true} .pendingMessageCount=${state.status?.pendingMessageCount ?? 0} .status=${state.status} .activity=${state.activity} .onLoadMore=${() => this.withChatPrependTransition(() => this.sessions.loadEarlierMessages())}></chat-view>
+            <chat-view .sessionId=${state.selectedSession.id} .messages=${state.messages} .messageStart=${state.messagePageStart} .messageTotal=${state.messagePageTotal} .hasMore=${state.messagePageStart > 0} .loadingMore=${state.isLoadingEarlierMessages} .isReceivingPartialStream=${state.isReceivingPartialStream} .isCompacting=${state.status?.isCompacting === true} .pendingMessageCount=${state.status?.pendingMessageCount ?? 0} .status=${state.status} .activity=${state.activity} .onLoadMore=${() => this.withChatPrependTransition(() => this.sessions.loadEarlierMessages())}></chat-view>
             <prompt-editor .sessionId=${state.selectedSession.id} .cwd=${state.selectedWorkspace?.path} .disabled=${state.selectedSession.archived === true} .canSteer=${state.status?.isStreaming === true} .isCompacting=${state.status?.isCompacting === true} .canStop=${state.status?.isStreaming === true || state.status?.isBashRunning === true || state.status?.isCompacting === true} .onSend=${(text: string, streamingBehavior?: "steer" | "followUp") => this.sessions.send(text, streamingBehavior)} .onStop=${() => this.sessions.stopActiveWork()}></prompt-editor>
             <status-bar .status=${state.status} .workspace=${state.selectedWorkspace}></status-bar>
             ${state.commandDialog !== undefined ? html`<command-picker .title=${state.commandDialog.title} .options=${state.commandDialog.options} .onPick=${(value: string) => this.sessions.respondToCommand(state.commandDialog?.requestId ?? "", value)} .onCancel=${() => { this.sessions.cancelCommand(); }}></command-picker>` : null}

src/client/src/components/shared.ts

@@ -141,7 +141,9 @@ export const chatStyles = css`
   .history-boundary { display: grid; gap: 3px; margin: 0 0 14px; color: #8b949e; font-size: 12px; text-align: center; }
   .session-activity { display: grid; gap: 4px; margin: 0 0 14px; padding: 12px; border: 1px solid #30363d; border-radius: 10px; background: #161b22; color: #e6edf3; }
   .session-activity.compacting { border-color: #a371f7; background: #21132f; }
+  .session-activity.receiving { border-color: #238636; background: #0f1b12; }
   .session-activity strong { color: #d2a8ff; }
+  .session-activity.receiving strong { color: #3fb950; }
   .session-activity span, .session-activity small { color: #8b949e; }
   .history-boundary small { color: #6e7681; }
   .label { display: block; margin-bottom: 8px; color: #8b949e; font-size: 12px; text-transform: uppercase; }

src/client/src/controllers/sessionController.ts

@@ -13,6 +13,7 @@ export class SessionController {
   private readonly socket = new SessionSocket();
   private readonly globalSocket = new GlobalSessionSocket();
   private selectionSeq = 0;
+  private catchupStreamSessionId: string | undefined;
 
   constructor(private readonly getState: GetState, private readonly setState: SetState, private readonly updateUrl: UpdateUrl) {}
 
@@ -31,7 +32,8 @@ export class SessionController {
 
   clearActiveSession() {
     this.socket.close();
-    this.setState({ selectedSession: undefined, messages: [], messagePageStart: 0, messagePageTotal: 0, isLoadingEarlierMessages: false, status: undefined, activity: undefined });
+    this.catchupStreamSessionId = undefined;
+    this.setState({ selectedSession: undefined, messages: [], messagePageStart: 0, messagePageTotal: 0, isLoadingEarlierMessages: false, isReceivingPartialStream: false, status: undefined, activity: undefined });
   }
 
   async startSession() {
@@ -49,6 +51,7 @@ export class SessionController {
   async selectSession(session: SessionInfo, options?: { updateUrl?: boolean | undefined }) {
     const seq = ++this.selectionSeq;
     this.socket.close();
+    this.catchupStreamSessionId = undefined;
     const cached = readChatHistoryCache(session.id);
     this.setState({
       selectedSession: session,
@@ -56,6 +59,7 @@ export class SessionController {
       messagePageStart: cached?.start ?? 0,
       messagePageTotal: cached?.total ?? 0,
       isLoadingEarlierMessages: false,
+      isReceivingPartialStream: false,
       status: session.archived === true ? undefined : this.getState().sessionStatuses[session.id],
       activity: session.archived === true ? undefined : this.getState().sessionActivities[session.id],
     });
@@ -64,7 +68,7 @@ export class SessionController {
         const page = await api.messages(session.id, { limit: MESSAGE_PAGE_SIZE });
         if (seq !== this.selectionSeq || this.getState().selectedSession?.id !== session.id) return;
         const history = this.mergeAndCacheHistory(session.id, page);
-        this.setState({ messages: normalizeMessages(history.messages), messagePageStart: history.start, messagePageTotal: history.total, isLoadingEarlierMessages: false, status: undefined, activity: undefined });
+        this.setState({ messages: normalizeMessages(history.messages), messagePageStart: history.start, messagePageTotal: history.total, isLoadingEarlierMessages: false, isReceivingPartialStream: false, status: undefined, activity: undefined });
         if (options?.updateUrl !== false) this.updateUrl();
         return;
       }
@@ -73,7 +77,9 @@ export class SessionController {
       const [page, status] = await Promise.all([api.messages(session.id, { limit: MESSAGE_PAGE_SIZE }), api.status(session.id)]);
       if (seq !== this.selectionSeq || this.getState().selectedSession?.id !== session.id) return;
       const history = this.mergeAndCacheHistory(session.id, page);
-      this.setState({ messages: normalizeMessages(history.messages), messagePageStart: history.start, messagePageTotal: history.total, isLoadingEarlierMessages: false, status, activity: this.getState().sessionActivities[session.id] });
+      const isReceivingPartialStream = status.isStreaming === true;
+      this.catchupStreamSessionId = isReceivingPartialStream ? session.id : undefined;
+      this.setState({ messages: normalizeMessages(history.messages), messagePageStart: history.start, messagePageTotal: history.total, isLoadingEarlierMessages: false, isReceivingPartialStream, status, activity: this.getState().sessionActivities[session.id] });
       this.applyStatus(status);
       for (const event of buffered) this.applyEvent(event);
       this.socket.setHandler((event) => { this.applyEvent(event); });
@@ -238,6 +244,7 @@ export class SessionController {
       sessionStatuses: { ...this.getState().sessionStatuses, [status.sessionId]: status },
       status: this.getState().selectedSession?.id === status.sessionId ? status : this.getState().status,
     });
+    if (this.catchupStreamSessionId === status.sessionId && !status.isStreaming) this.finishStreamCatchup(status.sessionId);
   }
 
   private applySessionName(sessionId: string, name: string | undefined) {
@@ -256,6 +263,15 @@ export class SessionController {
   }
 
   private applyEvent(event: SessionUiEvent) {
+    const selectedSessionId = this.getState().selectedSession?.id;
+    if (this.catchupStreamSessionId !== undefined && this.catchupStreamSessionId === selectedSessionId) {
+      if (event.type === "message.end" || event.type === "agent.end") {
+        this.finishStreamCatchup(this.catchupStreamSessionId);
+        return;
+      }
+      if (isTranscriptEvent(event)) return;
+    }
+
     const transcript = applyTranscriptEvent(this.getState().messages, event);
     if (transcript) {
       this.setState({ messages: transcript });
@@ -267,5 +283,27 @@ export class SessionController {
       this.applySessionName(event.sessionId, event.name);
     }
   }
+
+  private finishStreamCatchup(sessionId: string) {
+    if (this.catchupStreamSessionId !== sessionId) return;
+    this.catchupStreamSessionId = undefined;
+    if (this.getState().selectedSession?.id === sessionId) this.setState({ isReceivingPartialStream: false });
+    void this.refreshMessages(sessionId);
+  }
+
+  private async refreshMessages(sessionId: string) {
+    try {
+      const page = await api.messages(sessionId, { limit: MESSAGE_PAGE_SIZE });
+      if (this.getState().selectedSession?.id !== sessionId) return;
+      const history = this.mergeAndCacheHistory(sessionId, page);
+      this.setState({ messages: normalizeMessages(history.messages), messagePageStart: history.start, messagePageTotal: history.total });
+    } catch (error) {
+      if (this.getState().selectedSession?.id === sessionId) this.setState({ error: String(error) });
+    }
+  }
+}
+
+function isTranscriptEvent(event: SessionUiEvent): boolean {
+  return ["message.append", "assistant.delta", "tool.start", "tool.end", "shell.start", "shell.chunk", "shell.end", "command.output", "session.error"].includes(event.type);
 }