From f24094c0fab381991a2ec080dcf615d147cd9d3a Mon Sep 17 00:00:00 2001
From: Federico Jaramillo Martinez <federico@klingit.com>
Date: Mon, 11 May 2026 21:33:36 +0200
Subject: [PATCH] Refresh chat history after reconnect

---
 src/client/src/components/PiWebApp.ts         |  8 ++++++
 .../src/controllers/sessionController.ts      | 28 ++++++++++++++++++-
 src/client/src/sessionSocket.ts               |  9 +++++-
 3 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/src/client/src/components/PiWebApp.ts b/src/client/src/components/PiWebApp.ts
index fdf223b..68a3d97 100644
--- a/src/client/src/components/PiWebApp.ts
+++ b/src/client/src/components/PiWebApp.ts
@@ -68,6 +68,10 @@ export class PiWebApp extends LitElement {
   private terminalAutoStartWorkspaceId: string | undefined;
   private readonly plugins = createPluginRegistry();
   private readonly onPopState = () => void this.withChatScrollTransition(() => this.restoreRoute(false));
+  private readonly onFocus = () => { void this.sessions.refreshSelectedSession(); };
+  private readonly onVisibilityChange = () => {
+    if (document.visibilityState === "visible") void this.sessions.refreshSelectedSession();
+  };
   private readonly onKeyDown = (event: KeyboardEvent) => {
     if (this.keyboard.handle(event, this.getActions())) {
       event.preventDefault();
@@ -78,6 +82,8 @@ export class PiWebApp extends LitElement {
   override connectedCallback(): void {
     super.connectedCallback();
     window.addEventListener("popstate", this.onPopState);
+    window.addEventListener("focus", this.onFocus);
+    document.addEventListener("visibilitychange", this.onVisibilityChange);
     window.addEventListener("keydown", this.onKeyDown);
     this.connectRealtime();
     void this.loadExternalPlugins();
@@ -86,6 +92,8 @@ export class PiWebApp extends LitElement {
 
   override disconnectedCallback(): void {
     window.removeEventListener("popstate", this.onPopState);
+    window.removeEventListener("focus", this.onFocus);
+    document.removeEventListener("visibilitychange", this.onVisibilityChange);
     window.removeEventListener("keydown", this.onKeyDown);
     this.keyboard.reset();
     this.sessions.dispose();
diff --git a/src/client/src/controllers/sessionController.ts b/src/client/src/controllers/sessionController.ts
index 6e3bca9..c53a976 100644
--- a/src/client/src/controllers/sessionController.ts
+++ b/src/client/src/controllers/sessionController.ts
@@ -74,7 +74,11 @@ export class SessionController {
         return;
       }
       const buffered: SessionUiEvent[] = [];
-      this.socket.connect(session.id, (event) => buffered.push(event));
+      this.socket.connect(
+        session.id,
+        (event) => buffered.push(event),
+        () => { void this.refreshSelectedSession(session.id); },
+      );
       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);
@@ -278,6 +282,28 @@ export class SessionController {
     }
   }
 
+  async refreshSelectedSession(sessionId = this.getState().selectedSession?.id): Promise<void> {
+    const session = this.getState().selectedSession;
+    if (sessionId === undefined || session?.id !== sessionId || session.archived === true) return;
+    try {
+      this.flushPendingTranscriptEvents();
+      const [page, status] = await Promise.all([api.messages(sessionId, { limit: MESSAGE_PAGE_SIZE }), api.status(sessionId)]);
+      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,
+        status,
+        activity: this.getState().sessionActivities[sessionId],
+        isReceivingPartialStream: status.isStreaming,
+      });
+      this.applyStatus(status);
+    } catch (error) {
+      if (this.getState().selectedSession?.id === sessionId) this.setState({ error: String(error) });
+    }
+  }
+
   private replaceSession(session: SessionInfo) {
     const current = this.getState().selectedSession;
     this.setState({
diff --git a/src/client/src/sessionSocket.ts b/src/client/src/sessionSocket.ts
index df526da..f4fdaf2 100644
--- a/src/client/src/sessionSocket.ts
+++ b/src/client/src/sessionSocket.ts
@@ -10,11 +10,14 @@ export class SessionSocket {
   private reconnectTimer?: number;
   private reconnectDelay = 500;
   private shouldReconnect = false;
+  private hasOpened = false;
+  private onReconnect: (() => void) | undefined;
 
-  connect(sessionId: string, onEvent: (event: SessionUiEvent) => void): void {
+  connect(sessionId: string, onEvent: (event: SessionUiEvent) => void, onReconnect?: () => void): void {
     this.close();
     this.sessionId = sessionId;
     this.onEvent = onEvent;
+    this.onReconnect = onReconnect;
     this.shouldReconnect = true;
     this.open();
   }
@@ -30,6 +33,8 @@ export class SessionSocket {
     this.socket = undefined;
     this.sessionId = undefined;
     this.onEvent = undefined;
+    this.onReconnect = undefined;
+    this.hasOpened = false;
   }
 
   private open(): void {
@@ -38,6 +43,8 @@ export class SessionSocket {
     this.socket = socket;
     socket.onopen = () => {
       this.reconnectDelay = 500;
+      if (this.hasOpened) this.onReconnect?.();
+      this.hasOpened = true;
     };
     socket.onmessage = (message) => void this.handleMessage(message.data);
     socket.onerror = () => { socket.close(); };