From ad054973195328d401e0ed27c09a3ef801b7d764 Mon Sep 17 00:00:00 2001
From: Federico Jaramillo Martinez <federico@klingit.com>
Date: Thu, 7 May 2026 13:00:10 +0200
Subject: [PATCH] Improve session reconnect behavior

---
 src/client/src/api.ts                         | 10 +++++-
 src/client/src/appState.ts                    |  2 ++
 src/client/src/components/PiWebApp.ts         |  6 +++-
 src/client/src/components/PromptEditor.ts     |  9 ++++--
 src/client/src/components/SessionList.ts      | 15 +++++++--
 .../src/controllers/sessionController.ts      | 28 +++++++++++++---
 src/client/src/sessionSocket.ts               | 32 +++++++++++++++++--
 src/server/index.ts                           |  4 +++
 src/server/realtime/sessionEventHub.ts        | 13 ++++++++
 src/server/sessions/piSessionService.ts       | 10 ++++--
 10 files changed, 114 insertions(+), 15 deletions(-)

diff --git a/src/client/src/api.ts b/src/client/src/api.ts
index 72f6e7d..3be94c2 100644
--- a/src/client/src/api.ts
+++ b/src/client/src/api.ts
@@ -90,6 +90,14 @@ export const api = {
 };
 
 export function sessionEvents(sessionId: string): WebSocket {
+  return new WebSocket(`${webSocketBaseUrl()}/api/sessions/${sessionId}/events`);
+}
+
+export function globalSessionEvents(): WebSocket {
+  return new WebSocket(`${webSocketBaseUrl()}/api/sessions/events`);
+}
+
+function webSocketBaseUrl(): string {
   const protocol = location.protocol === "https:" ? "wss:" : "ws:";
-  return new WebSocket(`${protocol}//${location.host}/api/sessions/${sessionId}/events`);
+  return `${protocol}//${location.host}`;
 }
diff --git a/src/client/src/appState.ts b/src/client/src/appState.ts
index 7b8e9fc..fb7deb1 100644
--- a/src/client/src/appState.ts
+++ b/src/client/src/appState.ts
@@ -10,6 +10,7 @@ export interface AppState {
   selectedWorkspace?: Workspace;
   selectedSession?: SessionInfo;
   status?: SessionStatus;
+  sessionStatuses: Record<string, SessionStatus>;
   commandDialog?: Extract<CommandResult, { type: "select" }>;
   error: string;
 }
@@ -20,6 +21,7 @@ export function initialAppState(): AppState {
     workspaces: [],
     sessions: [],
     messages: [],
+    sessionStatuses: {},
     error: "",
   };
 }
diff --git a/src/client/src/components/PiWebApp.ts b/src/client/src/components/PiWebApp.ts
index 00c1fd3..548b6e1 100644
--- a/src/client/src/components/PiWebApp.ts
+++ b/src/client/src/components/PiWebApp.ts
@@ -12,6 +12,7 @@ import "./SessionList";
 import "./ChatView";
 import type { ChatView } from "./ChatView";
 import "./PromptEditor";
+import type { PromptEditor } from "./PromptEditor";
 import "./StatusBar";
 import "./CommandPicker";
 import { appStyles } from "./shared";
@@ -20,6 +21,7 @@ import { appStyles } from "./shared";
 export class PiWebApp extends LitElement {
   @state() private state: AppState = initialAppState();
   @query("chat-view") private chatView?: ChatView;
+  @query("prompt-editor") private promptEditor?: PromptEditor;
 
   private readonly sessions = new SessionController(
     () => this.state,
@@ -42,6 +44,7 @@ export class PiWebApp extends LitElement {
   connectedCallback(): void {
     super.connectedCallback();
     window.addEventListener("popstate", this.onPopState);
+    this.sessions.connectStatusUpdates();
     void this.loadProjectsAndRestoreRoute();
   }
 
@@ -75,6 +78,7 @@ export class PiWebApp extends LitElement {
     await this.chatView?.updateComplete;
     await nextFrame();
     this.chatView?.restoreScrollPosition();
+    this.promptEditor?.focusInput();
   }
 
   private updateUrl() {
@@ -96,7 +100,7 @@ export class PiWebApp extends LitElement {
           </header>
           <project-list .projects=${state.projects} .selected=${state.selectedProject} .onSelect=${(project: Project) => this.withChatScrollTransition(() => this.workspaces.selectProject(project))}></project-list>
           <workspace-list .workspaces=${state.workspaces} .selected=${state.selectedWorkspace} .onSelect=${(workspace: Workspace) => this.withChatScrollTransition(() => this.workspaces.selectWorkspace(workspace))}></workspace-list>
-          <session-list .sessions=${state.sessions} .selected=${state.selectedSession} .canStart=${!!state.selectedWorkspace} .onStart=${() => this.withChatScrollTransition(() => this.sessions.startSession())} .onSelect=${(session: SessionInfo) => this.withChatScrollTransition(() => this.sessions.selectSession(session))}></session-list>
+          <session-list .sessions=${state.sessions} .statuses=${state.sessionStatuses} .selected=${state.selectedSession} .canStart=${!!state.selectedWorkspace} .onStart=${() => this.withChatScrollTransition(() => this.sessions.startSession())} .onSelect=${(session: SessionInfo) => this.withChatScrollTransition(() => this.sessions.selectSession(session))}></session-list>
         </aside>
         <main>
           ${state.error ? html`<div class="error">${state.error}</div>` : null}
diff --git a/src/client/src/components/PromptEditor.ts b/src/client/src/components/PromptEditor.ts
index c250961..5299297 100644
--- a/src/client/src/components/PromptEditor.ts
+++ b/src/client/src/components/PromptEditor.ts
@@ -1,5 +1,5 @@
 import { LitElement, html } from "lit";
-import { customElement, property, state } from "lit/decorators.js";
+import { customElement, property, query, state } from "lit/decorators.js";
 import { api, type FileSuggestion, type SlashCommand } from "../api";
 import { promptEditorStyles, type CompletionItem } from "./shared";
 import "./AutocompleteMenu";
@@ -11,6 +11,7 @@ export class PromptEditor extends LitElement {
   @property() cwd?: string;
   @property({ attribute: false }) onSend?: (text: string) => void;
   @property({ attribute: false }) onCloseSession?: () => void;
+  @query("textarea") private textarea?: HTMLTextAreaElement;
   @state() private draft = "";
   @state() private completions: CompletionItem[] = [];
   @state() private selectedIndex = 0;
@@ -30,11 +31,15 @@ export class PromptEditor extends LitElement {
           <autocomplete-menu .items=${this.completions} .selectedIndex=${this.selectedIndex} .onPick=${(item: CompletionItem) => this.pick(item)}></autocomplete-menu>
         </div>
         <button ?disabled=${this.disabled} @click=${this.send}>Send</button>
-        <button ?disabled=${this.disabled} @click=${() => this.onCloseSession?.()}>Close</button>
+        <button ?disabled=${this.disabled} title="Stop this session runtime on the server" @click=${() => this.onCloseSession?.()}>Stop</button>
       </footer>
     `;
   }
 
+  focusInput() {
+    this.textarea?.focus();
+  }
+
   private updateDraft(value: string) {
     this.draft = value;
     void this.refreshCompletions();
diff --git a/src/client/src/components/SessionList.ts b/src/client/src/components/SessionList.ts
index 54b5f6f..738237d 100644
--- a/src/client/src/components/SessionList.ts
+++ b/src/client/src/components/SessionList.ts
@@ -1,11 +1,12 @@
 import { LitElement, html } from "lit";
 import { customElement, property } from "lit/decorators.js";
-import type { SessionInfo } from "../api";
+import type { SessionInfo, SessionStatus } from "../api";
 import { listStyles } from "./shared";
 
 @customElement("session-list")
 export class SessionList extends LitElement {
   @property({ attribute: false }) sessions: SessionInfo[] = [];
+  @property({ attribute: false }) statuses: Record<string, SessionStatus> = {};
   @property({ attribute: false }) selected?: SessionInfo;
   @property({ type: Boolean }) canStart = false;
   @property({ attribute: false }) onSelect?: (session: SessionInfo) => void;
@@ -17,12 +18,22 @@ export class SessionList extends LitElement {
         <h2>Sessions <button ?disabled=${!this.canStart} @click=${() => this.onStart?.()}>+</button></h2>
         ${this.sessions.map((session) => html`
           <button class=${this.selected?.id === session.id ? "selected" : ""} @click=${() => this.onSelect?.(session)}>
-            <span>${session.name || session.firstMessage || session.id.slice(0, 8)}</span><small>${session.messageCount} messages</small>
+            <span>${session.name || session.firstMessage || session.id.slice(0, 8)}</span><small>${this.renderStatus(session)}${session.messageCount} messages</small>
           </button>
         `)}
       </section>
     `;
   }
 
+  private renderStatus(session: SessionInfo) {
+    const status = this.statuses[session.id];
+    if (!status) return "";
+    if (status.isStreaming) return "● streaming · ";
+    if (status.isBashRunning) return "● bash · ";
+    if (status.isCompacting) return "● compacting · ";
+    if (status.pendingMessageCount) return `● ${status.pendingMessageCount} pending · `;
+    return "";
+  }
+
   static styles = listStyles;
 }
diff --git a/src/client/src/controllers/sessionController.ts b/src/client/src/controllers/sessionController.ts
index d9f6f26..bf230f3 100644
--- a/src/client/src/controllers/sessionController.ts
+++ b/src/client/src/controllers/sessionController.ts
@@ -1,15 +1,21 @@
-import { api, type CommandResult, type SessionInfo } from "../api";
+import { api, type CommandResult, type SessionInfo, type SessionStatus } from "../api";
 import { appendText, normalizeMessages, textMessage } from "../chatMessages";
-import { SessionSocket, type SessionUiEvent } from "../sessionSocket";
+import { GlobalSessionSocket, SessionSocket, type SessionUiEvent } from "../sessionSocket";
 import type { GetState, SetState, UpdateUrl } from "./types";
 
 export class SessionController {
   private readonly socket = new SessionSocket();
+  private readonly globalSocket = new GlobalSessionSocket();
 
   constructor(private readonly getState: GetState, private readonly setState: SetState, private readonly updateUrl: UpdateUrl) {}
 
+  connectStatusUpdates() {
+    this.globalSocket.connect((event) => this.applyStatus(event.status));
+  }
+
   dispose() {
     this.socket.close();
+    this.globalSocket.close();
   }
 
   clearActiveSession() {
@@ -32,8 +38,13 @@ export class SessionController {
   async selectSession(session: SessionInfo, options?: { updateUrl?: boolean }) {
     this.socket.close();
     try {
-      this.setState({ selectedSession: session, messages: normalizeMessages(await api.messages(session.id)), status: await api.status(session.id) });
-      this.socket.connect(session.id, (event) => this.applyEvent(event));
+      const buffered: SessionUiEvent[] = [];
+      this.socket.connect(session.id, (event) => buffered.push(event));
+      const [messages, status] = await Promise.all([api.messages(session.id), api.status(session.id)]);
+      this.setState({ selectedSession: session, messages: normalizeMessages(messages), status });
+      this.applyStatus(status);
+      for (const event of buffered) this.applyEvent(event);
+      this.socket.setHandler((event) => this.applyEvent(event));
       if (options?.updateUrl !== false) this.updateUrl();
     } catch (error) {
       this.setState({ error: String(error) });
@@ -104,6 +115,13 @@ export class SessionController {
     }
   }
 
+  private applyStatus(status: SessionStatus) {
+    this.setState({
+      sessionStatuses: { ...this.getState().sessionStatuses, [status.sessionId]: status },
+      status: this.getState().selectedSession?.id === status.sessionId ? status : this.getState().status,
+    });
+  }
+
   private applyEvent(event: SessionUiEvent) {
     const messages = this.getState().messages;
     if (event.type === "assistant.delta") {
@@ -113,7 +131,7 @@ export class SessionController {
     } else if (event.type === "tool.end") {
       this.setState({ messages: [...messages, textMessage("tool", `${event.isError ? "✖" : "✓"} ${event.toolName}`)] });
     } else if (event.type === "status.update") {
-      this.setState({ status: event.status });
+      this.applyStatus(event.status);
     } else if (event.type === "session.error") {
       this.setState({ messages: [...messages, textMessage("system", event.message)] });
     }
diff --git a/src/client/src/sessionSocket.ts b/src/client/src/sessionSocket.ts
index 720014d..e40ba8d 100644
--- a/src/client/src/sessionSocket.ts
+++ b/src/client/src/sessionSocket.ts
@@ -1,4 +1,4 @@
-import { sessionEvents, type SessionStatus } from "./api";
+import { globalSessionEvents, sessionEvents, type SessionStatus } from "./api";
 
 export type SessionUiEvent =
   | { type: "assistant.delta"; text: string }
@@ -9,13 +9,40 @@ export type SessionUiEvent =
 
 export class SessionSocket {
   private socket?: WebSocket;
+  private onEvent?: (event: SessionUiEvent) => void;
 
   connect(sessionId: string, onEvent: (event: SessionUiEvent) => void): void {
     this.close();
+    this.onEvent = onEvent;
     this.socket = sessionEvents(sessionId);
+    this.socket.onmessage = (message) => this.handleMessage(message.data);
+  }
+
+  setHandler(onEvent: (event: SessionUiEvent) => void): void {
+    this.onEvent = onEvent;
+  }
+
+  close(): void {
+    this.socket?.close();
+    this.socket = undefined;
+    this.onEvent = undefined;
+  }
+
+  private handleMessage(data: string): void {
+    const event = JSON.parse(data);
+    if (isSessionUiEvent(event)) this.onEvent?.(event);
+  }
+}
+
+export class GlobalSessionSocket {
+  private socket?: WebSocket;
+
+  connect(onEvent: (event: Extract<SessionUiEvent, { type: "status.update" }>) => void): void {
+    this.close();
+    this.socket = globalSessionEvents();
     this.socket.onmessage = (message) => {
       const event = JSON.parse(message.data);
-      if (isSessionUiEvent(event)) onEvent(event);
+      if (event?.type === "status.update") onEvent(event);
     };
   }
 
@@ -28,3 +55,4 @@ export class SessionSocket {
 function isSessionUiEvent(event: any): event is SessionUiEvent {
   return ["assistant.delta", "tool.start", "tool.end", "status.update", "session.error"].includes(event?.type);
 }
+
diff --git a/src/server/index.ts b/src/server/index.ts
index 11c5b05..667180e 100644
--- a/src/server/index.ts
+++ b/src/server/index.ts
@@ -113,6 +113,10 @@ app.get<{ Params: { sessionId: string } }>("/api/sessions/:sessionId/events", {
   eventHub.add(request.params.sessionId, socket);
 });
 
+app.get("/api/sessions/events", { websocket: true }, (socket) => {
+  eventHub.addGlobal(socket);
+});
+
 app.get<{ Querystring: { cwd?: string; q?: string; kind?: "tracked" | "untracked" | "other" } }>("/api/files", async (request, reply) => {
   if (!request.query.cwd) return reply.code(400).send({ error: "cwd query parameter is required" });
   try {
diff --git a/src/server/realtime/sessionEventHub.ts b/src/server/realtime/sessionEventHub.ts
index 58886eb..1c08db4 100644
--- a/src/server/realtime/sessionEventHub.ts
+++ b/src/server/realtime/sessionEventHub.ts
@@ -2,6 +2,7 @@ import type { WebSocket } from "ws";
 
 export class SessionEventHub {
   private readonly socketsBySession = new Map<string, Set<WebSocket>>();
+  private readonly globalSockets = new Set<WebSocket>();
 
   add(sessionId: string, socket: WebSocket): void {
     let sockets = this.socketsBySession.get(sessionId);
@@ -13,10 +14,22 @@ export class SessionEventHub {
     socket.on("close", () => sockets?.delete(socket));
   }
 
+  addGlobal(socket: WebSocket): void {
+    this.globalSockets.add(socket);
+    socket.on("close", () => this.globalSockets.delete(socket));
+  }
+
   publish(sessionId: string, event: unknown): void {
     const payload = JSON.stringify(event);
     for (const socket of this.socketsBySession.get(sessionId) ?? []) {
       if (socket.readyState === socket.OPEN) socket.send(payload);
     }
   }
+
+  publishGlobal(event: unknown): void {
+    const payload = JSON.stringify(event);
+    for (const socket of this.globalSockets) {
+      if (socket.readyState === socket.OPEN) socket.send(payload);
+    }
+  }
 }
diff --git a/src/server/sessions/piSessionService.ts b/src/server/sessions/piSessionService.ts
index 08b010b..f98bd20 100644
--- a/src/server/sessions/piSessionService.ts
+++ b/src/server/sessions/piSessionService.ts
@@ -135,7 +135,7 @@ export class PiSessionService {
     this.bindRuntime(active);
     runtime.setRebindSession(async () => this.bindRuntime(active));
     this.active.set(runtime.session.sessionId, active);
-    this.events.publish(runtime.session.sessionId, { type: "status.update", status: this.statusFromSession(runtime.session) });
+    this.publishStatus(runtime.session);
     return active;
   }
 
@@ -147,11 +147,17 @@ export class PiSessionService {
     const { session } = active.runtime;
     active.unsubscribe = session.subscribe((event) => {
       this.events.publish(session.sessionId, toClientEvent(event));
-      this.events.publish(session.sessionId, { type: "status.update", status: this.statusFromSession(session) });
+      this.publishStatus(session);
     });
     this.active.set(session.sessionId, active);
   }
 
+  private publishStatus(session: AgentSession): void {
+    const status = this.statusFromSession(session);
+    this.events.publish(session.sessionId, { type: "status.update", status });
+    this.events.publishGlobal({ type: "status.update", status });
+  }
+
   private statusFromSession(session: AgentSession): ClientSessionStatus {
     const stats = session.getSessionStats();
     return {