From 50f1ddc9c514021ac02ac1e3e393b2c35eb7eea4 Mon Sep 17 00:00:00 2001 From: Federico Jaramillo Martinez Date: Tue, 26 May 2026 08:52:43 +0200 Subject: [PATCH] fix: refresh session message counts live --- .changeset/live-session-message-counts.md | 5 ++ src/client/src/api/parsers.test.ts | 2 + src/client/src/api/parsers.ts | 1 + .../src/controllers/sessionController.test.ts | 48 ++++++++++++++++++- .../src/controllers/sessionController.ts | 22 ++++++++- src/server/sessions/piSessionService.test.ts | 2 + src/server/sessions/piSessionService.ts | 1 + src/shared/apiTypes.ts | 1 + 8 files changed, 80 insertions(+), 2 deletions(-) create mode 100644 .changeset/live-session-message-counts.md diff --git a/.changeset/live-session-message-counts.md b/.changeset/live-session-message-counts.md new file mode 100644 index 0000000..dc683f9 --- /dev/null +++ b/.changeset/live-session-message-counts.md @@ -0,0 +1,5 @@ +--- +"@jmfederico/pi-web": patch +--- + +Refresh session list message counts from live session status updates. diff --git a/src/client/src/api/parsers.test.ts b/src/client/src/api/parsers.test.ts index 0988d07..9629ebd 100644 --- a/src/client/src/api/parsers.test.ts +++ b/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" } }, diff --git a/src/client/src/api/parsers.ts b/src/client/src/api/parsers.ts index 7b466dc..28538be 100644 --- a/src/client/src/api/parsers.ts +++ b/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"]), diff --git a/src/client/src/controllers/sessionController.test.ts b/src/client/src/controllers/sessionController.test.ts index 8bfc745..1d4421c 100644 --- a/src/client/src/controllers/sessionController.test.ts +++ b/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 }); diff --git a/src/client/src/controllers/sessionController.ts b/src/client/src/controllers/sessionController.ts index adf7843..b8b8918 100644 --- a/src/client/src/controllers/sessionController.ts +++ b/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, sessio return Object.fromEntries(Object.entries(activities).filter(([id]) => id !== sessionId)); } +function sessionMessageCountPatch(state: AppState, sessionId: string, messageCount: number | undefined): Pick, "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); } diff --git a/src/server/sessions/piSessionService.test.ts b/src/server/sessions/piSessionService.test.ts index ffde6fb..74e3e69 100644 --- a/src/server/sessions/piSessionService.test.ts +++ b/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(); }); diff --git a/src/server/sessions/piSessionService.ts b/src/server/sessions/piSessionService.ts index c869bc3..93e9d64 100644 --- a/src/server/sessions/piSessionService.ts +++ b/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 }), diff --git a/src/shared/apiTypes.ts b/src/shared/apiTypes.ts index 08fcb54..3d2ae87 100644 --- a/src/shared/apiTypes.ts +++ b/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 };