fix: refresh session message counts live

.changeset/live-session-message-counts.md

@@ -0,0 +1,5 @@
+---
+"@jmfederico/pi-web": patch
+---
+
+Refresh session list message counts from live session status updates.

src/client/src/api/parsers.test.ts

@@ -15,6 +15,7 @@ describe("API parsers", () => {
       isBashRunning: false,
       pendingMessageCount: 2,
       queuedMessages: [{ kind: "steer", text: "adjust this" }, { kind: "followUp", text: "then do that" }],
+      messageCount: 7,
       tokens: { input: 1, output: 2, cacheRead: 3, cacheWrite: 4, total: 10 },
       cost: 0.12,
       model: { provider: "p", id: "m", contextWindow: 100, reasoning: { effort: "low" } },
@@ -27,6 +28,7 @@ describe("API parsers", () => {
       isBashRunning: false,
       pendingMessageCount: 2,
       queuedMessages: [{ kind: "steer", text: "adjust this" }, { kind: "followUp", text: "then do that" }],
+      messageCount: 7,
       tokens: { input: 1, output: 2, cacheRead: 3, cacheWrite: 4, total: 10 },
       cost: 0.12,
       model: { provider: "p", id: "m", contextWindow: 100, reasoning: { effort: "low" } },

src/client/src/api/parsers.ts

@@ -106,6 +106,7 @@ export function parseSessionStatus(value: unknown): SessionStatus {
     isBashRunning: requireBoolean(record, "isBashRunning"),
     pendingMessageCount: requireNumber(record, "pendingMessageCount"),
     queuedMessages: record["queuedMessages"] === undefined ? [] : arrayOf(parseQueuedSessionMessage)(record["queuedMessages"]),
+    ...optionalField("messageCount", optionalNumber(record, "messageCount")),
     tokens: parseTokens(record["tokens"]),
     cost: requireNumber(record, "cost"),
     ...optionalModel(record["model"]),

src/client/src/controllers/sessionController.test.ts

@@ -1,6 +1,6 @@
 import { afterEach, describe, expect, it } from "vitest";
 import { api as defaultApi, type MessagePage, type SessionActivity, type SessionInfo, type SessionStatus, type Workspace } from "../api";
-import { loadCachedNewSessions, markCachedNewSessionInfo, rememberCachedNewSession } from "../cachedNewSessions";
+import { isCachedNewSessionInfo, loadCachedNewSessions, markCachedNewSessionInfo, rememberCachedNewSession } from "../cachedNewSessions";
 import { initialAppState, type AppState } from "../appState";
 import { loadDraft, saveDraft } from "../promptDraftStorage";
 import { SessionController, type SessionEventSocket } from "./sessionController";
@@ -120,6 +120,52 @@ describe("SessionController", () => {
     expect(state.sessionStatuses[oldSession.id]).toMatchObject({ sessionId: oldSession.id, isStreaming: false });
   });
 
+  it("updates visible session message counts from live status events", () => {
+    let state: AppState = {
+      ...initialAppState(),
+      selectedSession: oldSession,
+      sessions: [oldSession],
+    };
+    const controller = new SessionController(
+      () => state,
+      (patch) => { state = { ...state, ...patch }; },
+      () => undefined,
+      undefined,
+      { socket: new FakeSocket() },
+    );
+
+    controller.applyGlobalEvent({ type: "status.update", status: { ...status(oldSession.id), messageCount: 3 } });
+
+    expect(state.sessions[0]?.messageCount).toBe(3);
+    expect(state.selectedSession?.messageCount).toBe(3);
+  });
+
+  it("keeps live message count updates when a cached new session becomes persisted", async () => {
+    const cachedSession = markCachedNewSessionInfo(oldSession);
+    let resolvePrompt: (() => void) | undefined;
+    let state: AppState = { ...initialAppState(), selectedWorkspace: workspace, selectedSession: cachedSession, sessions: [cachedSession] };
+    const api: typeof defaultApi = {
+      ...defaultApi,
+      prompt: () => new Promise<{ accepted: true }>((resolve) => { resolvePrompt = () => { resolve({ accepted: true }); }; }),
+    };
+    const controller = new SessionController(
+      () => state,
+      (patch) => { state = { ...state, ...patch }; },
+      () => undefined,
+      undefined,
+      { api, socket: new FakeSocket() },
+    );
+
+    const send = controller.send("hello");
+    controller.applyGlobalEvent({ type: "status.update", status: { ...status(oldSession.id), messageCount: 1 } });
+    resolvePrompt?.();
+    await send;
+
+    expect(state.sessions[0]?.messageCount).toBe(1);
+    expect(isCachedNewSessionInfo(state.sessions[0])).toBe(false);
+    expect(state.selectedSession?.messageCount).toBe(1);
+  });
+
   it("recreates missing browser-cached new sessions and moves their draft", async () => {
     const storage = new MemoryStorage();
     Object.defineProperty(globalThis, "localStorage", { value: storage, configurable: true });

src/client/src/controllers/sessionController.ts

@@ -1,4 +1,5 @@
 import { api as defaultApi, type CommandResult, type SessionActivity, type SessionInfo, type SessionStatus, type ThinkingLevel } from "../api";
+import type { AppState } from "../appState";
 import { forgetCachedNewSession, isCachedNewSessionInfo, markCachedNewSessionInfo, rememberCachedNewSession, stripCachedNewSessionMarker } from "../cachedNewSessions";
 import { textMessage } from "../chatMessages";
 import { clearDraft, moveDraft, saveDraft } from "../promptDraftStorage";
@@ -417,7 +418,8 @@ export class SessionController {
 
   private markCachedNewSessionPersisted(session: SessionInfo): void {
     if (!isCachedNewSessionInfo(session)) return;
-    this.replaceSession(stripCachedNewSessionMarker(session));
+    const latest = this.getState().sessions.find((candidate) => candidate.id === session.id) ?? session;
+    this.replaceSession(stripCachedNewSessionMarker(latest));
   }
 
   private applyCommandResult(result: CommandResult) {
@@ -448,6 +450,7 @@ export class SessionController {
     const clearsStaleActivity = state.sessionActivities[status.sessionId]?.phase === "active" && !isSessionActive(status);
     this.setState({
       sessionStatuses: { ...state.sessionStatuses, [status.sessionId]: status },
+      ...sessionMessageCountPatch(state, status.sessionId, status.messageCount),
       ...(clearsStaleActivity ? { sessionActivities: omitSessionActivity(state.sessionActivities, status.sessionId) } : {}),
       status: state.selectedSession?.id === status.sessionId ? status : state.status,
       activity: state.selectedSession?.id === status.sessionId && clearsStaleActivity ? undefined : state.activity,
@@ -545,6 +548,23 @@ function omitSessionActivity(activities: Record<string, SessionActivity>, sessio
   return Object.fromEntries(Object.entries(activities).filter(([id]) => id !== sessionId));
 }
 
+function sessionMessageCountPatch(state: AppState, sessionId: string, messageCount: number | undefined): Pick<Partial<AppState>, "sessions" | "selectedSession"> {
+  if (messageCount === undefined) return {};
+
+  const sessionsChanged = state.sessions.some((session) => session.id === sessionId && session.messageCount !== messageCount);
+  const sessions = sessionsChanged
+    ? state.sessions.map((session) => session.id === sessionId ? { ...session, messageCount } : session)
+    : undefined;
+  const selectedSession = state.selectedSession?.id === sessionId && state.selectedSession.messageCount !== messageCount
+    ? { ...state.selectedSession, messageCount }
+    : state.selectedSession;
+
+  return {
+    ...(sessions === undefined ? {} : { sessions }),
+    ...(selectedSession !== state.selectedSession ? { selectedSession } : {}),
+  };
+}
+
 function isTranscriptEvent(event: SessionUiEvent): boolean {
   return ["message.append", "assistant.delta", "assistant.thinking.delta", "tool.start", "tool.update", "tool.end", "shell.start", "shell.chunk", "shell.end", "command.output", "session.error"].includes(event.type);
 }

src/server/sessions/piSessionService.test.ts

@@ -361,6 +361,7 @@ describe("PiSessionService", () => {
 
   it("includes queued message details in session status", async () => {
     const fake = fakeRuntime("status-session", {
+      messages: [{ role: "user", content: "hello" }, { role: "assistant", content: "hi" }],
       pendingMessageCount: 2,
       getSteeringMessages: () => ["adjust this turn"],
       getFollowUpMessages: () => ["then do this"],
@@ -374,6 +375,7 @@ describe("PiSessionService", () => {
     await expect(service.status("status-session")).resolves.toMatchObject({
       pendingMessageCount: 2,
       queuedMessages: [{ kind: "steer", text: "adjust this turn" }, { kind: "followUp", text: "then do this" }],
+      messageCount: 2,
     });
     await service.dispose();
   });

src/server/sessions/piSessionService.ts

@@ -761,6 +761,7 @@ export class PiSessionService {
       isBashRunning: session.isBashRunning,
       pendingMessageCount: session.pendingMessageCount,
       queuedMessages: queuedMessagesFromSession(session),
+      messageCount: session.messages.length,
       tokens: stats.tokens,
       cost: stats.cost,
       ...(contextUsage === undefined ? {} : { contextUsage }),

src/shared/apiTypes.ts

@@ -109,6 +109,7 @@ export interface SessionStatus {
   isBashRunning: boolean;
   pendingMessageCount: number;
   queuedMessages: QueuedSessionMessage[];
+  messageCount?: number;
   tokens: { input: number; output: number; cacheRead: number; cacheWrite: number; total: number };
   cost: number;
   contextUsage?: { tokens: number | null; contextWindow: number; percent: number | null };