feat(sessions): stream in-flight partial when joining a mid-turn session

Seed the in-flight partial assistant message (text, thinking, and
in-progress tool calls) when opening or reconnecting to a session that is
mid-stream, then continue streaming live deltas on top of it. Replaces the
blocking "Catching up..." placeholder and the end-of-turn transcript reload.

Server stamps every per-session UI event with a monotonic seq at the
SessionEventHub publish choke point and exposes
GET /sessions/:sessionId/stream-snapshot returning { seq, partial }. The
client fetches the snapshot on join, seeds the normalized partial into the
in-memory transcript (never the history cache), and applies buffered/live
events using the seq watermark for exactly-once delivery.

The snapshot is a progressive enhancement: a 404 from an older remote
pi-web or a not-yet-restarted session daemon falls back to an empty seed
(seq 0, drops nothing), so sessions still open and stream normally. The
stream-snapshot route is registered in the federation allowlist for
remote-machine proxying.
This commit is contained in:
Federico Jaramillo Martinez
2026-07-17 14:49:56 +02:00
parent ae9eaf3082
commit 2b17145291
31 changed files with 681 additions and 108 deletions
+5
View File
@@ -0,0 +1,5 @@
---
"@jmfederico/pi-web": patch
---
Stream in-flight assistant replies immediately when opening or reconnecting to a session mid-turn. The chat now seeds the partial message (text, thinking, and in-progress tool calls) and continues streaming live updates on top of it, replacing the blocking "Catching up…" placeholder and the end-of-turn transcript reload. Sessions still open normally against remote machines or session daemons that predate this feature: the snapshot is fetched as a progressive enhancement and its absence no longer blocks the transcript.
+1 -1
View File
@@ -2,4 +2,4 @@ export { activityApi, api, configApi, filesApi, gitApi, machinesApi, piPackagesA
export { globalSessionEvents, realtimeEvents, sessionEvents, terminalSocket } from "./api/sockets"; export { globalSessionEvents, realtimeEvents, sessionEvents, terminalSocket } from "./api/sockets";
export { DEFAULT_WORKSPACE_UPLOADS_FOLDER, effectiveWorkspaceUploadFolder, uploadWorkspaceFile, uploadWorkspaceFiles, workspaceEffectiveUploadFolder, workspaceUploadPath, WorkspaceUploadBatchError, WorkspaceUploadCancelledError } from "./api/workspaceUploads"; export { DEFAULT_WORKSPACE_UPLOADS_FOLDER, effectiveWorkspaceUploadFolder, uploadWorkspaceFile, uploadWorkspaceFiles, workspaceEffectiveUploadFolder, workspaceUploadPath, WorkspaceUploadBatchError, WorkspaceUploadCancelledError } from "./api/workspaceUploads";
export type { UploadWorkspaceFileOptions, UploadWorkspaceFilesOptions, WorkspaceFileUploadProgress, WorkspaceUploadBatchFileProgress, WorkspaceUploadBatchProgress, WorkspaceUploadFileFailure, WorkspaceUploadFileInput, WorkspaceUploadFolderConfig, WorkspaceUploadTask, WorkspaceUploadXhr, WorkspaceUploadXhrFactory } from "./api/workspaceUploads"; export type { UploadWorkspaceFileOptions, UploadWorkspaceFilesOptions, WorkspaceFileUploadProgress, WorkspaceUploadBatchFileProgress, WorkspaceUploadBatchProgress, WorkspaceUploadFileFailure, WorkspaceUploadFileInput, WorkspaceUploadFolderConfig, WorkspaceUploadTask, WorkspaceUploadXhr, WorkspaceUploadXhrFactory } from "./api/workspaceUploads";
export type { ActiveAgentProfileDescriptor, ArchiveSessionsResponse, AuthProviderOption, AuthProviderStatus, AuthProvidersResponse, AuthStatusSource, AuthType, CommandOption, CommandResult, DeleteWorkspaceFileResponse, FileContentMediaType, FileContentResponse, FileSuggestion, FileTreeEntry, FileTreeResponse, GitDiffResponse, GitFileState, GitStatusFile, GitStatusResponse, Machine, MachineHealth, MachineKind, MachineRuntime, MachineStatus, MessagePage, ModelSelectionResponse, MoveWorkspaceFileOptions, MoveWorkspaceFileResponse, OAuthFlowState, PiPackageInfo, PiPackageInstallRequest, PiPackageMutationAction, PiPackageMutationResponse, PiPackageRemoveRequest, PiPackageScope, PiPackageUpdateRequest, PiPackagesResponse, PiWebAgentDirEnvSource, PiWebCapability, PiWebComponentStatus, PiWebConfigEnvOverrides, PiWebConfigResponse, PiWebConfigValues, PiWebDockerMode, PiWebInstallationInfo, PiWebInstallationKind, PiWebPluginConfig, PiWebPluginConfigMap, PiWebPluginInfo, PiWebPluginsResponse, PiWebPluginScope, PiWebPluginSettings, PiWebReleaseStatus, PiWebRuntimeComponent, PiWebRuntimeResponse, PiWebShortcutConfig, PiWebStatusMessage, PiWebStatusResponse, PiWebUploadsConfig, Project, PromptAttachment, QueuedSessionMessage, RealtimeEvent, RunTerminalCommandInput, SavedPromptAttachment, SessionActivity, SessionBulkArchiveResponse, SessionBulkDeleteArchivedResponse, SessionBulkFailure, SessionBulkMutationRef, SessionBulkMutationRequest, SessionCleanupExecuteResponse, SessionCleanupPreviewResponse, SessionCleanupProjectSummary, SessionCleanupRequest, SessionCleanupThresholds, SessionCleanupTotals, SessionInfo, SessionModel, SessionRef, SessionStatus, SlashCommand, SessionUiEvent, TerminalCommandRun, TerminalCommandRunFilter, TerminalCommandRunHandle, TerminalCommandRunStatus, TerminalInfo, TerminalUiEvent, ThinkingLevel, ThinkingLevelsResponse, WriteWorkspaceFileOptions, WriteWorkspaceFileResponse, Workspace, WorkspaceActivity, WorkspaceActivityResponse, WorkspaceActivityUiEvent } from "../../shared/apiTypes"; export type { ActiveAgentProfileDescriptor, ArchiveSessionsResponse, AuthProviderOption, AuthProviderStatus, AuthProvidersResponse, AuthStatusSource, AuthType, CommandOption, CommandResult, DeleteWorkspaceFileResponse, FileContentMediaType, FileContentResponse, FileSuggestion, FileTreeEntry, FileTreeResponse, GitDiffResponse, GitFileState, GitStatusFile, GitStatusResponse, Machine, MachineHealth, MachineKind, MachineRuntime, MachineStatus, MessagePage, ModelSelectionResponse, MoveWorkspaceFileOptions, MoveWorkspaceFileResponse, OAuthFlowState, PiPackageInfo, PiPackageInstallRequest, PiPackageMutationAction, PiPackageMutationResponse, PiPackageRemoveRequest, PiPackageScope, PiPackageUpdateRequest, PiPackagesResponse, PiWebAgentDirEnvSource, PiWebCapability, PiWebComponentStatus, PiWebConfigEnvOverrides, PiWebConfigResponse, PiWebConfigValues, PiWebDockerMode, PiWebInstallationInfo, PiWebInstallationKind, PiWebPluginConfig, PiWebPluginConfigMap, PiWebPluginInfo, PiWebPluginsResponse, PiWebPluginScope, PiWebPluginSettings, PiWebReleaseStatus, PiWebRuntimeComponent, PiWebRuntimeResponse, PiWebShortcutConfig, PiWebStatusMessage, PiWebStatusResponse, PiWebUploadsConfig, Project, PromptAttachment, QueuedSessionMessage, RealtimeEvent, RunTerminalCommandInput, SavedPromptAttachment, SessionActivity, SessionBulkArchiveResponse, SessionBulkDeleteArchivedResponse, SessionBulkFailure, SessionBulkMutationRef, SessionBulkMutationRequest, SessionCleanupExecuteResponse, SessionCleanupPreviewResponse, SessionCleanupProjectSummary, SessionCleanupRequest, SessionCleanupThresholds, SessionCleanupTotals, SessionInfo, SessionModel, SessionRef, SessionStatus, SessionStreamSnapshot, SlashCommand, SessionUiEvent, TerminalCommandRun, TerminalCommandRunFilter, TerminalCommandRunHandle, TerminalCommandRunStatus, TerminalInfo, TerminalUiEvent, ThinkingLevel, ThinkingLevelsResponse, WriteWorkspaceFileOptions, WriteWorkspaceFileResponse, Workspace, WorkspaceActivity, WorkspaceActivityResponse, WorkspaceActivityUiEvent } from "../../shared/apiTypes";
+23
View File
@@ -280,6 +280,29 @@ describe("session API compatibility", () => {
expect(init?.method).toBe("POST"); expect(init?.method).toBe("POST");
expect(JSON.parse(requestBody(init))).toEqual({ cwd: "/repo with spaces" }); expect(JSON.parse(requestBody(init))).toEqual({ cwd: "/repo with spaces" });
}); });
it("reads a session stream snapshot through an encoded machine route with cwd context", async () => {
const fetchMock = stubJsonFetch({ seq: 12, partial: { role: "assistant", content: [{ type: "text", text: "streaming" }] } });
await expect(sessionsApi.streamSnapshot({ id: "s /?", cwd: "/repo with spaces" }, "remote /?")).resolves.toEqual({
seq: 12,
partial: { role: "assistant", content: [{ type: "text", text: "streaming" }] },
});
expect(fetchMock).toHaveBeenCalledOnce();
const [url, init] = fetchCall(fetchMock, 0);
expect(url).toBe("https://pi.example.test/api/machines/remote%20%2F%3F/sessions/s%20%2F%3F/stream-snapshot?cwd=%2Frepo+with+spaces");
expect(init?.method ?? "GET").toBe("GET");
});
it("reads a session stream snapshot for a legacy session-id ref without cwd context", async () => {
const fetchMock = stubJsonFetch({ seq: 0, partial: null });
await expect(sessionsApi.streamSnapshot("s 1", "remote a")).resolves.toEqual({ seq: 0, partial: null });
expect(fetchMock).toHaveBeenCalledOnce();
expect(fetchCall(fetchMock, 0)[0]).toBe("https://pi.example.test/api/machines/remote%20a/sessions/s%201/stream-snapshot");
});
}); });
describe("machine-scoped file suggestion API", () => { describe("machine-scoped file suggestion API", () => {
+2
View File
@@ -41,6 +41,7 @@ import {
parseSessionCleanupPreviewResponse, parseSessionCleanupPreviewResponse,
parseSessionInfo, parseSessionInfo,
parseSessionStatus, parseSessionStatus,
parseSessionStreamSnapshot,
parseSlashCommand, parseSlashCommand,
parseStopped, parseStopped,
parseTerminalCommandRun, parseTerminalCommandRun,
@@ -209,6 +210,7 @@ export const sessionsApi = {
deleteArchivedMany: (sessions: readonly SessionLookup[], machineId = "local") => request(`${machinePrefix(machineId)}/sessions/bulk/delete-archived`, parseSessionBulkDeleteArchivedResponse, { method: "POST", body: sessionBulkMutationBody(sessions) }), deleteArchivedMany: (sessions: readonly SessionLookup[], machineId = "local") => request(`${machinePrefix(machineId)}/sessions/bulk/delete-archived`, parseSessionBulkDeleteArchivedResponse, { method: "POST", body: sessionBulkMutationBody(sessions) }),
messages: (session: SessionLookup, options?: { limit?: number; before?: number }, machineId = "local") => request(messagePath(session, options, machineId), parseMessagePage), messages: (session: SessionLookup, options?: { limit?: number; before?: number }, machineId = "local") => request(messagePath(session, options, machineId), parseMessagePage),
status: (session: SessionLookup, machineId = "local") => request(sessionQueryPath(session, "status", machineId), parseSessionStatus), status: (session: SessionLookup, machineId = "local") => request(sessionQueryPath(session, "status", machineId), parseSessionStatus),
streamSnapshot: (session: SessionLookup, machineId = "local") => request(sessionQueryPath(session, "stream-snapshot", machineId), parseSessionStreamSnapshot),
clearQueue: (session: SessionLookup, machineId = "local") => request(sessionPath(session, "queue/clear", machineId), parseSessionStatus, { method: "POST", body: sessionBody(session) }), clearQueue: (session: SessionLookup, machineId = "local") => request(sessionPath(session, "queue/clear", machineId), parseSessionStatus, { method: "POST", body: sessionBody(session) }),
models: (session: SessionLookup, machineId = "local") => request(sessionQueryPath(session, "models", machineId), parseModelSelectionResponse), models: (session: SessionLookup, machineId = "local") => request(sessionQueryPath(session, "models", machineId), parseModelSelectionResponse),
setModel: (session: SessionLookup, provider: string, modelId: string, machineId = "local") => request(sessionPath(session, "model", machineId), parseSessionStatus, { method: "POST", body: sessionBody(session, { provider, modelId }) }), setModel: (session: SessionLookup, provider: string, modelId: string, machineId = "local") => request(sessionPath(session, "model", machineId), parseSessionStatus, { method: "POST", body: sessionBody(session, { provider, modelId }) }),
@@ -64,6 +64,7 @@ describe("federated route contract", () => {
ignoreParseFailure(sessionsApi.deleteArchivedMany([session], machineId)), ignoreParseFailure(sessionsApi.deleteArchivedMany([session], machineId)),
ignoreParseFailure(sessionsApi.messages(session, { limit: 20, before: 10 }, machineId)), ignoreParseFailure(sessionsApi.messages(session, { limit: 20, before: 10 }, machineId)),
ignoreParseFailure(sessionsApi.status(session, machineId)), ignoreParseFailure(sessionsApi.status(session, machineId)),
ignoreParseFailure(sessionsApi.streamSnapshot(session, machineId)),
ignoreParseFailure(sessionsApi.clearQueue(session, machineId)), ignoreParseFailure(sessionsApi.clearQueue(session, machineId)),
ignoreParseFailure(sessionsApi.models(session, machineId)), ignoreParseFailure(sessionsApi.models(session, machineId)),
ignoreParseFailure(sessionsApi.setModel(session, "openai", "gpt", machineId)), ignoreParseFailure(sessionsApi.setModel(session, "openai", "gpt", machineId)),
+14 -1
View File
@@ -1,6 +1,6 @@
import { describe, expect, it } from "vitest"; import { describe, expect, it } from "vitest";
import { PI_WEB_CAPABILITIES } from "../../../shared/capabilities"; import { PI_WEB_CAPABILITIES } from "../../../shared/capabilities";
import { parseCommandResult, parseFileContentResponse, parseFileSuggestion, parseGitStatusResponse, parseMachineRuntime, parseMessagePage, parsePiPackageMutationResponse, parsePiPackagesResponse, parsePiWebConfigResponse, parsePiWebPluginsResponse, parsePiWebRuntimeResponse, parsePiWebStatusResponse, parseSessionBulkArchiveResponse, parseSessionBulkDeleteArchivedResponse, parseSessionCleanupExecuteResponse, parseSessionCleanupPreviewResponse, parseSessionInfo, parseSessionStatus, parseSlashCommand, parseTerminalCommandRun, parseTerminalInfo, parseWorkspace, parseWorkspaceActivityResponse } from "./parsers"; import { parseCommandResult, parseFileContentResponse, parseFileSuggestion, parseGitStatusResponse, parseMachineRuntime, parseMessagePage, parsePiPackageMutationResponse, parsePiPackagesResponse, parsePiWebConfigResponse, parsePiWebPluginsResponse, parsePiWebRuntimeResponse, parsePiWebStatusResponse, parseSessionBulkArchiveResponse, parseSessionBulkDeleteArchivedResponse, parseSessionCleanupExecuteResponse, parseSessionCleanupPreviewResponse, parseSessionInfo, parseSessionStatus, parseSessionStreamSnapshot, parseSlashCommand, parseTerminalCommandRun, parseTerminalInfo, parseWorkspace, parseWorkspaceActivityResponse } from "./parsers";
describe("API parsers", () => { describe("API parsers", () => {
it("parses PI WEB config responses", () => { it("parses PI WEB config responses", () => {
@@ -152,6 +152,19 @@ describe("API parsers", () => {
expect(parseMessagePage({ messages: ["c"], start: 3, total: 9 })).toEqual({ messages: ["c"], start: 3, total: 9 }); expect(parseMessagePage({ messages: ["c"], start: 3, total: 9 })).toEqual({ messages: ["c"], start: 3, total: 9 });
}); });
it("parses a session stream snapshot, defaulting a missing partial to null", () => {
expect(parseSessionStreamSnapshot({ seq: 7, partial: { role: "assistant", content: [{ type: "text", text: "hi" }] } })).toEqual({
seq: 7,
partial: { role: "assistant", content: [{ type: "text", text: "hi" }] },
});
expect(parseSessionStreamSnapshot({ seq: 0, partial: null })).toEqual({ seq: 0, partial: null });
expect(parseSessionStreamSnapshot({ seq: 3 })).toEqual({ seq: 3, partial: null });
});
it("rejects a session stream snapshot without a numeric seq", () => {
expect(() => parseSessionStreamSnapshot({ partial: null })).toThrow("Expected number field: seq");
});
it("parses session cleanup preview and execute responses", () => { it("parses session cleanup preview and execute responses", () => {
const preview = { const preview = {
generatedAt: "2026-06-25T12:00:00.000Z", generatedAt: "2026-06-25T12:00:00.000Z",
+9 -1
View File
@@ -1,4 +1,4 @@
import type { ArchiveSessionsResponse, AuthProviderOption, AuthProviderStatus, AuthProvidersResponse, AuthStatusSource, AuthType, CommandOption, CommandResult, DeleteWorkspaceFileResponse, FileContentResponse, FileSuggestion, FileTreeEntry, FileTreeResponse, GitDiffResponse, GitFileState, GitStatusFile, GitStatusResponse, Machine, MachineHealth, MachineKind, MachineRuntime, MachineStatus, MessagePage, ModelSelectionResponse, MoveWorkspaceFileResponse, OAuthFlowState, PiWebAgentDirEnvSource, PiWebCapability, PiWebComponentStatus, PiWebConfigEnvOverrides, PiWebConfigResponse, PiWebConfigValues, PiWebInstallationInfo, PiWebPluginConfigMap, PiWebPluginInfo, PiWebPluginsResponse, PiWebPluginScope, PiWebReleaseStatus, PiWebRuntimeComponent, PiWebRuntimeResponse, PiWebServiceComponent, PiWebShortcutConfig, PiWebStatusMessage, PiWebStatusResponse, PiWebStatusSeverity, Project, QueuedSessionMessage, SavedPromptAttachment, SessionBulkArchiveResponse, SessionBulkDeleteArchivedResponse, SessionBulkFailure, SessionCleanupExecuteResponse, SessionCleanupPreviewResponse, SessionCleanupProjectSummary, SessionCleanupThresholds, SessionCleanupTotals, SessionInfo, SessionModel, SessionStatus, SlashCommand, TerminalCommandRun, TerminalCommandRunStatus, TerminalInfo, ThinkingLevelsResponse, WriteWorkspaceFileResponse, Workspace, WorkspaceActivity, WorkspaceActivityResponse } from "../../../shared/apiTypes"; import type { ArchiveSessionsResponse, AuthProviderOption, AuthProviderStatus, AuthProvidersResponse, AuthStatusSource, AuthType, CommandOption, CommandResult, DeleteWorkspaceFileResponse, FileContentResponse, FileSuggestion, FileTreeEntry, FileTreeResponse, GitDiffResponse, GitFileState, GitStatusFile, GitStatusResponse, Machine, MachineHealth, MachineKind, MachineRuntime, MachineStatus, MessagePage, ModelSelectionResponse, MoveWorkspaceFileResponse, OAuthFlowState, PiWebAgentDirEnvSource, PiWebCapability, PiWebComponentStatus, PiWebConfigEnvOverrides, PiWebConfigResponse, PiWebConfigValues, PiWebInstallationInfo, PiWebPluginConfigMap, PiWebPluginInfo, PiWebPluginsResponse, PiWebPluginScope, PiWebReleaseStatus, PiWebRuntimeComponent, PiWebRuntimeResponse, PiWebServiceComponent, PiWebShortcutConfig, PiWebStatusMessage, PiWebStatusResponse, PiWebStatusSeverity, Project, QueuedSessionMessage, SavedPromptAttachment, SessionBulkArchiveResponse, SessionBulkDeleteArchivedResponse, SessionBulkFailure, SessionCleanupExecuteResponse, SessionCleanupPreviewResponse, SessionCleanupProjectSummary, SessionCleanupThresholds, SessionCleanupTotals, SessionInfo, SessionModel, SessionStatus, SessionStreamSnapshot, SlashCommand, TerminalCommandRun, TerminalCommandRunStatus, TerminalInfo, ThinkingLevelsResponse, WriteWorkspaceFileResponse, Workspace, WorkspaceActivity, WorkspaceActivityResponse } from "../../../shared/apiTypes";
import type { PiPackageInfo, PiPackageMutationAction, PiPackageMutationResponse, PiPackageScope, PiPackagesResponse } from "../../../shared/apiTypes"; import type { PiPackageInfo, PiPackageMutationAction, PiPackageMutationResponse, PiPackageScope, PiPackagesResponse } from "../../../shared/apiTypes";
import { parseActiveAgentProfileDescriptor } from "../../../shared/activeAgentProfile"; import { parseActiveAgentProfileDescriptor } from "../../../shared/activeAgentProfile";
import { parseKnownPiWebCapabilities } from "../../../shared/capabilities"; import { parseKnownPiWebCapabilities } from "../../../shared/capabilities";
@@ -196,6 +196,14 @@ export function parseSessionStatus(value: unknown): SessionStatus {
}; };
} }
export function parseSessionStreamSnapshot(value: unknown): SessionStreamSnapshot {
const record = requireRecord(value);
return {
seq: requireNumber(record, "seq"),
partial: record["partial"] ?? null,
};
}
export function parseSessionCleanupPreviewResponse(value: unknown): SessionCleanupPreviewResponse { export function parseSessionCleanupPreviewResponse(value: unknown): SessionCleanupPreviewResponse {
const record = requireRecord(value); const record = requireRecord(value);
const skippedBusySessionIds = record["skippedBusySessionIds"] === undefined ? undefined : arrayOfString(record["skippedBusySessionIds"], "skippedBusySessionIds"); const skippedBusySessionIds = record["skippedBusySessionIds"] === undefined ? undefined : arrayOfString(record["skippedBusySessionIds"], "skippedBusySessionIds");
-2
View File
@@ -17,7 +17,6 @@ export interface AppState {
messagePageEnd: number; messagePageEnd: number;
messagePageTotal: number; messagePageTotal: number;
isLoadingEarlierMessages: boolean; isLoadingEarlierMessages: boolean;
isReceivingPartialStream: boolean;
/** Sessions with a prompt upload in flight, keyed by sessionId (client-owned). */ /** Sessions with a prompt upload in flight, keyed by sessionId (client-owned). */
sendingPrompts: Record<string, true>; sendingPrompts: Record<string, true>;
/** Client-side queued sends waiting for a just-created backend session, keyed by sessionId. */ /** Client-side queued sends waiting for a just-created backend session, keyed by sessionId. */
@@ -127,7 +126,6 @@ export function initialAppState(): AppState {
messagePageEnd: 0, messagePageEnd: 0,
messagePageTotal: 0, messagePageTotal: 0,
isLoadingEarlierMessages: false, isLoadingEarlierMessages: false,
isReceivingPartialStream: false,
sendingPrompts: {}, sendingPrompts: {},
clientQueuedSessionMessages: {}, clientQueuedSessionMessages: {},
startingSessionCount: 0, startingSessionCount: 0,
+37 -1
View File
@@ -1,7 +1,7 @@
import { describe, expect, it } from "vitest"; import { describe, expect, it } from "vitest";
import { groupChatMessages } from "./chatGroups"; import { groupChatMessages } from "./chatGroups";
import { normalizeMessages, textMessage } from "./chatMessages"; import { normalizeMessages, textMessage } from "./chatMessages";
import { applyTranscriptEvent } from "./chatTranscript"; import { applyTranscriptEvent, seedStreamingPartial } from "./chatTranscript";
import type { ChatLine } from "./components/shared"; import type { ChatLine } from "./components/shared";
const finalAssistant = { const finalAssistant = {
@@ -464,6 +464,42 @@ describe("applyTranscriptEvent", () => {
]); ]);
}); });
it("seeds a null or undefined partial as a no-op", () => {
const messages = [textMessage("user", "question")];
expect(seedStreamingPartial(messages, null)).toBe(messages);
expect(seedStreamingPartial(messages, undefined)).toBe(messages);
});
it("seeds an in-flight assistant partial with text and thinking so live deltas append onto it", () => {
const seeded = seedStreamingPartial([textMessage("user", "question")], {
role: "assistant",
content: [{ type: "thinking", thinking: "plan" }, { type: "text", text: "partial" }],
});
expect(seeded).toEqual([
textMessage("user", "question"),
{ role: "assistant", parts: [{ type: "thinking", text: "plan" }, { type: "text", text: "partial" }] },
]);
// A live delta continues the seeded assistant message rather than starting a new one.
expect(applyTranscriptEvent(seeded, { type: "assistant.delta", text: " answer" })).toEqual([
textMessage("user", "question"),
{ role: "assistant", parts: [{ type: "thinking", text: "plan" }, { type: "text", text: "partial answer" }] },
]);
});
it("seeds an in-progress tool call from the partial as a tool execution line", () => {
const seeded = seedStreamingPartial([textMessage("user", "run it")], {
role: "assistant",
content: [{ type: "toolCall", id: "tool-1", name: "bash", arguments: { command: "ls" } }],
});
expect(seeded).toEqual([
textMessage("user", "run it"),
{ role: "tool", parts: [{ type: "toolExecution", toolCallId: "tool-1", toolName: "bash", summary: "ls", args: { command: "ls" }, status: "pending" }] },
]);
});
it("replaces an optimistic user message when the finalized text matches", () => { it("replaces an optimistic user message when the finalized text matches", () => {
const messages = [textMessage("user", "sent prompt")]; const messages = [textMessage("user", "sent prompt")];
+21 -1
View File
@@ -1,4 +1,4 @@
import { appendText, appendThinking, normalizeMessage, previewFromDetails, summarizeArgs, textMessage } from "./chatMessages"; import { appendText, appendThinking, normalizeMessage, normalizeMessages, previewFromDetails, summarizeArgs, textMessage } from "./chatMessages";
import type { ChatLine, ToolExecutionPart } from "./components/shared"; import type { ChatLine, ToolExecutionPart } from "./components/shared";
import { appendShellChunk, finalizeShellMessage, shellStartMessage } from "./shellMessages"; import { appendShellChunk, finalizeShellMessage, shellStartMessage } from "./shellMessages";
import type { SessionUiEvent } from "./sessionSocket"; import type { SessionUiEvent } from "./sessionSocket";
@@ -20,6 +20,26 @@ interface ToolResultUpdate {
presentation: ToolResultPresentation; presentation: ToolResultPresentation;
} }
/**
* Seed the in-flight partial assistant message on top of committed history at
* join time. `partial` is the browser-projected `AssistantMessage` from the
* stream snapshot (or `null`/`undefined` when the session is not mid
* assistant-message stream). The partial is normalized the same way a streamed
* message is (text/thinking parts plus any in-progress tool call parts kept on
* the assistant line), so live `assistant.delta`/`assistant.thinking.delta`
* events append onto it and `message.end` reconciles it into the finalized shape.
* This returns a new in-memory message list only; it never writes to the raw
* history cache.
*/
export function seedStreamingPartial(messages: ChatLine[], partial: unknown): ChatLine[] {
if (partial === null || partial === undefined) return messages;
// Normalize the same way committed history is (coalescing an in-progress tool
// call into a `toolExecution` line) so the seeded partial renders identically
// to its finalized form and live `tool.*` events can target it by call id.
const lines = normalizeMessages([partial]);
return lines.length === 0 ? messages : [...messages, ...lines];
}
export function applyTranscriptEvent(messages: ChatLine[], event: SessionUiEvent): ChatLine[] | undefined { export function applyTranscriptEvent(messages: ChatLine[], event: SessionUiEvent): ChatLine[] | undefined {
if (event.type === "message.append") return appendNewMessage(messages, event.message); if (event.type === "message.append") return appendNewMessage(messages, event.message);
if (event.type === "assistant.delta") return appendText(messages, "assistant", event.text); if (event.type === "assistant.delta") return appendText(messages, "assistant", event.text);
+10 -1
View File
@@ -1,5 +1,5 @@
import { normalizeMessages } from "./chatMessages"; import { normalizeMessages } from "./chatMessages";
import { applyTranscriptEvent } from "./chatTranscript"; import { applyTranscriptEvent, seedStreamingPartial } from "./chatTranscript";
import { mergeChatHistory, readChatHistoryCache, removeChatHistoryCache, writeChatHistoryCache, type RawMessagePage } from "./chatHistoryCache"; import { mergeChatHistory, readChatHistoryCache, removeChatHistoryCache, writeChatHistoryCache, type RawMessagePage } from "./chatHistoryCache";
import type { ChatLine } from "./components/shared"; import type { ChatLine } from "./components/shared";
import type { SessionUiEvent } from "./sessionSocket"; import type { SessionUiEvent } from "./sessionSocket";
@@ -45,6 +45,15 @@ export class ChatTranscriptStore {
return applyTranscriptEvent(messages, event); return applyTranscriptEvent(messages, event);
} }
/**
* Seed the join-time in-flight partial assistant message on top of the
* committed history view. Returns a new in-memory message list; the raw
* history cache is deliberately untouched so the partial never persists.
*/
seedStreamingPartial(messages: ChatLine[], partial: unknown): ChatLine[] {
return seedStreamingPartial(messages, partial);
}
discard(sessionId: string): void { discard(sessionId: string): void {
this.rawHistoryPages.delete(sessionId); this.rawHistoryPages.delete(sessionId);
this.cache.remove?.(sessionId); this.cache.remove?.(sessionId);
+4 -35
View File
@@ -16,19 +16,6 @@ import "./ToolExecutionView";
const messageTimestampFormatter = new Intl.DateTimeFormat(undefined, { dateStyle: "medium", timeStyle: "medium" }); const messageTimestampFormatter = new Intl.DateTimeFormat(undefined, { dateStyle: "medium", timeStyle: "medium" });
const partialStreamNoticeBodies = [
"You opened this chat while the assistant was already replying. The complete answer will appear shortly.",
"We joined mid-sentence. Holding the curtain until the full reply is ready.",
"The assistant started before this tab arrived. Well show the full answer when it lands.",
"Catching the reply in one piece — no spoilers, no half-answers.",
"The tokens are still assembling themselves. Full answer incoming.",
"We arrived fashionably late to this response. The complete version will appear soon.",
] as const;
function randomPartialStreamNoticeBody(): string {
return partialStreamNoticeBodies[Math.floor(Math.random() * partialStreamNoticeBodies.length)] ?? partialStreamNoticeBodies[0];
}
function clampPercent(value: number): number { function clampPercent(value: number): number {
return clampNumber(value, 0, 100); return clampNumber(value, 0, 100);
} }
@@ -83,7 +70,6 @@ export class ChatView extends LitElement {
@property({ type: Number }) messageTotal = 0; @property({ type: Number }) messageTotal = 0;
@property({ type: Boolean }) hasMore = false; @property({ type: Boolean }) hasMore = false;
@property({ type: Boolean }) loadingMore = false; @property({ type: Boolean }) loadingMore = false;
@property({ type: Boolean }) isReceivingPartialStream = false;
@property({ type: Boolean }) isSendingPrompt = false; @property({ type: Boolean }) isSendingPrompt = false;
@property({ type: Boolean }) isCompacting = false; @property({ type: Boolean }) isCompacting = false;
@property({ type: Number }) pendingMessageCount = 0; @property({ type: Number }) pendingMessageCount = 0;
@@ -112,7 +98,6 @@ export class ChatView extends LitElement {
private groupedMessagesCache: ChatGroup[] = []; private groupedMessagesCache: ChatGroup[] = [];
private readonly messageMetaCache = new WeakMap<ChatLine, string>(); private readonly messageMetaCache = new WeakMap<ChatLine, string>();
private readonly messageCopyTextCache = new WeakMap<ChatLine, string>(); private readonly messageCopyTextCache = new WeakMap<ChatLine, string>();
private partialStreamNoticeBody: string | undefined;
private lastScrollTop = 0; private lastScrollTop = 0;
private lastClientHeight = 0; private lastClientHeight = 0;
private touchStartY: number | undefined; private touchStartY: number | undefined;
@@ -193,7 +178,6 @@ export class ChatView extends LitElement {
this.savePreviousSessionScrollPosition(changed.get("sessionId")); this.savePreviousSessionScrollPosition(changed.get("sessionId"));
this.prepareSessionUiState(); this.prepareSessionUiState();
} }
if (changed.has("isReceivingPartialStream") || (changed.has("sessionId") && this.isReceivingPartialStream)) this.syncPartialStreamNoticeBody();
if (changed.has("messages")) this.pinnedToBottom = this.pinnedToBottom && (this.didChatHeightChange() || this.isNearBottom()); if (changed.has("messages")) this.pinnedToBottom = this.pinnedToBottom && (this.didChatHeightChange() || this.isNearBottom());
} }
@@ -326,12 +310,6 @@ export class ChatView extends LitElement {
} }
private renderSessionActivity() { private renderSessionActivity() {
if (this.isReceivingPartialStream) return html`
<aside class="session-activity receiving" aria-live="polite">
<strong>Catching up…</strong>
<span>${this.currentPartialStreamNoticeBody()}</span>
</aside>
`;
if (!this.isCompacting) return null; if (!this.isCompacting) return null;
return html` return html`
<aside class="session-activity compacting" aria-live="polite"> <aside class="session-activity compacting" aria-live="polite">
@@ -342,15 +320,6 @@ export class ChatView extends LitElement {
`; `;
} }
private syncPartialStreamNoticeBody(): void {
this.partialStreamNoticeBody = this.isReceivingPartialStream ? randomPartialStreamNoticeBody() : undefined;
}
private currentPartialStreamNoticeBody(): string {
this.partialStreamNoticeBody ??= randomPartialStreamNoticeBody();
return this.partialStreamNoticeBody;
}
private activityState(): string | undefined { private activityState(): string | undefined {
const status = this.status; const status = this.status;
if (status === undefined) return this.activity?.label; if (status === undefined) return this.activity?.label;
@@ -733,10 +702,10 @@ export class ChatView extends LitElement {
} }
private shouldFallbackToBottomForMissingAnchor(): boolean { private shouldFallbackToBottomForMissingAnchor(): boolean {
// While catching up to a stream, history can temporarily omit the in-flight // Only fall back to the bottom once the full history is loaded; while earlier
// assistant message that a previous scroll save anchored to. Keep retrying // pages can still load, a missing scroll anchor should keep retrying rather
// until the final refreshed transcript has a chance to render that anchor. // than jump the user to the bottom.
return !this.hasMore && !this.isReceivingPartialStream; return !this.hasMore;
} }
private updatePinnedToBottomAfterRestore(status: Exclude<ChatScrollRestoreResult["status"], "missing">): void { private updatePinnedToBottomAfterRestore(status: Exclude<ChatScrollRestoreResult["status"], "missing">): void {
+1 -1
View File
@@ -1884,7 +1884,7 @@ export class PiWebApp extends LitElement {
private renderChatView(state: AppState, session: SessionInfo) { private renderChatView(state: AppState, session: SessionInfo) {
return html` return html`
<chat-view .sessionId=${session.id} .messages=${state.messages} .messageStart=${state.messagePageStart} .messageEnd=${state.messagePageEnd} .messageTotal=${state.messagePageTotal} .hasMore=${state.messagePageStart > 0} .loadingMore=${state.isLoadingEarlierMessages} .isReceivingPartialStream=${state.isReceivingPartialStream} .isSendingPrompt=${state.sendingPrompts[session.id] === true} .isCompacting=${state.status?.isCompacting === true} .pendingMessageCount=${state.status?.pendingMessageCount ?? 0} .clientQueuedMessages=${state.clientQueuedSessionMessages[session.id] ?? []} .status=${state.status} .activity=${state.activity} .canClearServerQueue=${this.canClearServerQueue()} .onClearServerQueue=${this.handleClearServerQueue} .onLoadMore=${() => this.withChatPrependTransition(() => this.sessions.loadEarlierMessages())}></chat-view> <chat-view .sessionId=${session.id} .messages=${state.messages} .messageStart=${state.messagePageStart} .messageEnd=${state.messagePageEnd} .messageTotal=${state.messagePageTotal} .hasMore=${state.messagePageStart > 0} .loadingMore=${state.isLoadingEarlierMessages} .isSendingPrompt=${state.sendingPrompts[session.id] === true} .isCompacting=${state.status?.isCompacting === true} .pendingMessageCount=${state.status?.pendingMessageCount ?? 0} .clientQueuedMessages=${state.clientQueuedSessionMessages[session.id] ?? []} .status=${state.status} .activity=${state.activity} .canClearServerQueue=${this.canClearServerQueue()} .onClearServerQueue=${this.handleClearServerQueue} .onLoadMore=${() => this.withChatPrependTransition(() => this.sessions.loadEarlierMessages())}></chat-view>
`; `;
} }
-2
View File
@@ -324,9 +324,7 @@ export const chatStyles = css`
.queued-kind { color: var(--pi-muted); font-size: 12px; text-transform: uppercase; } .queued-kind { color: var(--pi-muted); font-size: 12px; text-transform: uppercase; }
.session-activity { max-width: 100%; min-width: 0; box-sizing: border-box; display: grid; gap: 4px; margin: 0 0 14px; padding: 12px; border: 1px solid var(--pi-border); border-radius: 10px; background: var(--pi-surface); color: var(--pi-text); overflow: hidden; } .session-activity { max-width: 100%; min-width: 0; box-sizing: border-box; display: grid; gap: 4px; margin: 0 0 14px; padding: 12px; border: 1px solid var(--pi-border); border-radius: 10px; background: var(--pi-surface); color: var(--pi-text); overflow: hidden; }
.session-activity.compacting { border-color: var(--pi-purple-border); background: var(--pi-purple-surface); } .session-activity.compacting { border-color: var(--pi-purple-border); background: var(--pi-purple-surface); }
.session-activity.receiving { border-color: var(--pi-success-border); background: var(--pi-success-bg); }
.session-activity strong { color: var(--pi-purple); } .session-activity strong { color: var(--pi-purple); }
.session-activity.receiving strong { color: var(--pi-success); }
.session-activity span, .session-activity small { color: var(--pi-muted); } .session-activity span, .session-activity small { color: var(--pi-muted); }
.history-boundary small { color: var(--pi-dim); } .history-boundary small { color: var(--pi-dim); }
.msg-header { display: flex; align-items: center; justify-content: space-between; gap: 10px; min-height: 22px; margin-bottom: 8px; } .msg-header { display: flex; align-items: center; justify-content: space-between; gap: 10px; min-height: 22px; margin-bottom: 8px; }
@@ -59,6 +59,7 @@ describe("SessionController live events", () => {
...defaultApi, ...defaultApi,
messages: () => Promise.resolve(emptyPage), messages: () => Promise.resolve(emptyPage),
status: () => Promise.resolve(status(oldSession.id)), status: () => Promise.resolve(status(oldSession.id)),
streamSnapshot: () => Promise.resolve({ seq: 0, partial: null }),
}; };
const controller = new SessionController( const controller = new SessionController(
() => state, () => state,
@@ -28,6 +28,7 @@ describe("SessionController selected-session refresh", () => {
statusCalls += 1; statusCalls += 1;
return statusCalls === 1 ? firstStatus.promise : trailingStatus.promise; return statusCalls === 1 ? firstStatus.promise : trailingStatus.promise;
}, },
streamSnapshot: () => Promise.resolve({ seq: 0, partial: null }),
}; };
const controller = new SessionController( const controller = new SessionController(
() => state, () => state,
@@ -72,6 +73,7 @@ describe("SessionController selected-session refresh", () => {
...defaultApi, ...defaultApi,
messages: (session) => sessionLookupId(session) === oldSession.id ? stalePage.promise : Promise.resolve(replacementPage), messages: (session) => sessionLookupId(session) === oldSession.id ? stalePage.promise : Promise.resolve(replacementPage),
status: (session) => sessionLookupId(session) === oldSession.id ? staleStatus.promise : Promise.resolve(status(replacementSession.id)), status: (session) => sessionLookupId(session) === oldSession.id ? staleStatus.promise : Promise.resolve(status(replacementSession.id)),
streamSnapshot: () => Promise.resolve({ seq: 0, partial: null }),
thinkingLevels: () => Promise.resolve({ levels: [] }), thinkingLevels: () => Promise.resolve({ levels: [] }),
}; };
const controller = new SessionController( const controller = new SessionController(
@@ -93,4 +95,32 @@ describe("SessionController selected-session refresh", () => {
expect(state.messages).toEqual([{ role: "assistant", parts: [{ type: "text", text: "replacement" }] }]); expect(state.messages).toEqual([{ role: "assistant", parts: [{ type: "text", text: "replacement" }] }]);
expect(state.status?.sessionId).toBe(replacementSession.id); expect(state.status?.sessionId).toBe(replacementSession.id);
}); });
it("fetches the join-time stream snapshot alongside messages and status on refresh", async () => {
// Leg 3 contract: the snapshot is fetched for the selected session on the join
// refresh path. Seeding/watermark application is deliberately NOT asserted here
// (that is Leg 4); this only guards that the data is fetched.
const snapshotLookups: string[] = [];
let state: AppState = { ...initialAppState(), selectedWorkspace: workspace, selectedSession: oldSession, sessions: [oldSession] };
const api: typeof defaultApi = {
...defaultApi,
messages: () => Promise.resolve(page("live", 1)),
status: () => Promise.resolve({ ...status(oldSession.id), isStreaming: true }),
streamSnapshot: (session) => {
snapshotLookups.push(sessionLookupId(session));
return Promise.resolve({ seq: 5, partial: { role: "assistant", content: [{ type: "text", text: "partial" }] } });
},
};
const controller = new SessionController(
() => state,
(patch) => { state = { ...state, ...patch }; },
() => undefined,
undefined,
{ api, socket: new FakeSocket() },
);
await controller.refreshSelectedSession();
expect(snapshotLookups).toEqual([oldSession.id]);
});
}); });
@@ -32,6 +32,7 @@ describe("SessionController reload and selection", () => {
return Promise.resolve(freshPage); return Promise.resolve(freshPage);
}, },
status: (session) => Promise.resolve(status(sessionLookupId(session))), status: (session) => Promise.resolve(status(sessionLookupId(session))),
streamSnapshot: () => Promise.resolve({ seq: 0, partial: null }),
thinkingLevels: () => Promise.resolve({ levels: [] }), thinkingLevels: () => Promise.resolve({ levels: [] }),
}; };
const controller = new SessionController( const controller = new SessionController(
@@ -0,0 +1,214 @@
import { describe, expect, it } from "vitest";
import { initialAppState } from "../appState";
import { SessionController } from "./sessionController";
import { defaultApi, deferred, EmitSocket, oldSession, runPendingAnimationFrames, status, workspace, type AppState, type MessagePage, type SessionStatus, type SessionStreamSnapshot } from "./sessionController.testSupport";
function assistantPartial(text: string): SessionStreamSnapshot["partial"] {
return { role: "assistant", content: [{ type: "text", text }] };
}
describe("SessionController stream seed + watermark reconciliation", () => {
it("seeds the in-flight partial on top of committed history at join time", async () => {
const socket = new EmitSocket();
let state: AppState = { ...initialAppState(), selectedWorkspace: workspace, sessions: [oldSession] };
const api: typeof defaultApi = {
...defaultApi,
messages: () => Promise.resolve({ messages: [{ role: "user", content: "question" }], start: 0, total: 1 }),
status: () => Promise.resolve({ ...status(oldSession.id), isStreaming: true }),
streamSnapshot: () => Promise.resolve({ seq: 4, partial: assistantPartial("streaming answer") }),
};
const controller = new SessionController(
() => state,
(patch) => { state = { ...state, ...patch }; },
() => undefined,
undefined,
{ api, socket },
);
await controller.selectSession(oldSession, { updateUrl: false });
expect(state.messages).toEqual([
{ role: "user", parts: [{ type: "text", text: "question" }] },
{ role: "assistant", parts: [{ type: "text", text: "streaming answer" }] },
]);
// The seeded partial must never be written to the raw history cache.
expect(controller).toBeDefined();
});
it("drops live events at or below the watermark and applies later events exactly once", async () => {
const socket = new EmitSocket();
let state: AppState = { ...initialAppState(), selectedWorkspace: workspace, sessions: [oldSession] };
const api: typeof defaultApi = {
...defaultApi,
messages: () => Promise.resolve({ messages: [{ role: "user", content: "question" }], start: 0, total: 1 }),
status: () => Promise.resolve({ ...status(oldSession.id), isStreaming: true }),
streamSnapshot: () => Promise.resolve({ seq: 4, partial: assistantPartial("seed") }),
};
const controller = new SessionController(
() => state,
(patch) => { state = { ...state, ...patch }; },
() => undefined,
undefined,
{ api, socket },
);
await controller.selectSession(oldSession, { updateUrl: false });
// Already reflected in the seeded partial (seq <= 4): dropped.
socket.emit({ type: "assistant.delta", text: "DUP", seq: 3 });
socket.emit({ type: "assistant.delta", text: "DUP", seq: 4 });
// Past the watermark (seq > 4): appended onto the seeded partial exactly once.
socket.emit({ type: "assistant.delta", text: " more", seq: 5 });
runPendingAnimationFrames();
expect(state.messages).toEqual([
{ role: "user", parts: [{ type: "text", text: "question" }] },
{ role: "assistant", parts: [{ type: "text", text: "seed more" }] },
]);
});
it("applies buffered events replayed after join through the same watermark", async () => {
const socket = new EmitSocket();
const page = deferred<MessagePage>();
const statusResult = deferred<SessionStatus>();
const snapshot = deferred<SessionStreamSnapshot>();
let state: AppState = { ...initialAppState(), selectedWorkspace: workspace, sessions: [oldSession] };
const api: typeof defaultApi = {
...defaultApi,
messages: () => page.promise,
status: () => statusResult.promise,
streamSnapshot: () => snapshot.promise,
thinkingLevels: () => Promise.resolve({ levels: [] }),
};
const controller = new SessionController(
() => state,
(patch) => { state = { ...state, ...patch }; },
() => undefined,
undefined,
{ api, socket },
);
const selecting = controller.selectSession(oldSession, { updateUrl: false });
// Events arriving during the join fetch are buffered by selectSession.
socket.emit({ type: "assistant.delta", text: "STALE", seq: 2 });
socket.emit({ type: "assistant.delta", text: " live", seq: 6 });
page.resolve({ messages: [{ role: "user", content: "question" }], start: 0, total: 1 });
statusResult.resolve({ ...status(oldSession.id), isStreaming: true });
snapshot.resolve({ seq: 4, partial: assistantPartial("seed") });
await selecting;
runPendingAnimationFrames();
expect(state.messages).toEqual([
{ role: "user", parts: [{ type: "text", text: "question" }] },
{ role: "assistant", parts: [{ type: "text", text: "seed live" }] },
]);
});
it("handles a mid-tool join: null partial, committed tool call in history, live tool.update filtered by seq", async () => {
const socket = new EmitSocket();
let state: AppState = { ...initialAppState(), selectedWorkspace: workspace, sessions: [oldSession] };
const api: typeof defaultApi = {
...defaultApi,
messages: () => Promise.resolve({
messages: [
{ role: "user", content: "run it" },
{ role: "assistant", content: [{ type: "toolCall", id: "tool-1", name: "bash", arguments: { command: "ls" } }] },
],
start: 0,
total: 2,
}),
status: () => Promise.resolve({ ...status(oldSession.id), isStreaming: true, isBashRunning: true }),
// Mid tool execution the assistant-message stream has ended, so the
// snapshot carries no partial; the tool call is already in history.
streamSnapshot: () => Promise.resolve({ seq: 7, partial: null }),
};
const controller = new SessionController(
() => state,
(patch) => { state = { ...state, ...patch }; },
() => undefined,
undefined,
{ api, socket },
);
await controller.selectSession(oldSession, { updateUrl: false });
const toolLine = state.messages.find((line) => line.parts.some((part) => part.type === "toolExecution"));
expect(toolLine?.parts[0]).toMatchObject({ type: "toolExecution", toolCallId: "tool-1", toolName: "bash" });
// Reflected in the snapshot watermark (seq <= 7): dropped.
socket.emit({ type: "tool.update", toolName: "bash", toolCallId: "tool-1", text: "stale", content: undefined, details: undefined, seq: 7 });
// Fresh progress past the watermark: applied.
socket.emit({ type: "tool.update", toolName: "bash", toolCallId: "tool-1", text: "fresh output", content: undefined, details: undefined, seq: 8 });
runPendingAnimationFrames();
const updatedToolLine = state.messages.find((line) => line.parts.some((part) => part.type === "toolExecution"));
expect(updatedToolLine?.parts[0]).toMatchObject({ resultText: "fresh output" });
});
it("loads the transcript and streams live even when the snapshot fetch fails (older/un-restarted peer)", async () => {
const socket = new EmitSocket();
let state: AppState = { ...initialAppState(), selectedWorkspace: workspace, sessions: [oldSession] };
const api: typeof defaultApi = {
...defaultApi,
messages: () => Promise.resolve({ messages: [{ role: "user", content: "question" }], start: 0, total: 1 }),
status: () => Promise.resolve({ ...status(oldSession.id), isStreaming: true }),
// A session daemon / remote pi-web without the stream-snapshot route 404s.
streamSnapshot: () => Promise.reject(new Error("Not Found")),
thinkingLevels: () => Promise.resolve({ levels: [] }),
};
const controller = new SessionController(
() => state,
(patch) => { state = { ...state, ...patch }; },
() => undefined,
undefined,
{ api, socket },
);
await controller.selectSession(oldSession, { updateUrl: false });
// The core transcript still loads (no error banner, no dropped history).
expect(state.error).toBeFalsy();
expect(state.messages).toEqual([{ role: "user", parts: [{ type: "text", text: "question" }] }]);
// With a fallback watermark of 0, fresh live deltas still stream in.
socket.emit({ type: "assistant.delta", text: "live answer", seq: 1 });
runPendingAnimationFrames();
expect(state.messages).toEqual([
{ role: "user", parts: [{ type: "text", text: "question" }] },
{ role: "assistant", parts: [{ type: "text", text: "live answer" }] },
]);
});
it("does not seed a partial and does not filter events for an idle join", async () => {
const socket = new EmitSocket();
let state: AppState = { ...initialAppState(), selectedWorkspace: workspace, sessions: [oldSession] };
const api: typeof defaultApi = {
...defaultApi,
messages: () => Promise.resolve({ messages: [{ role: "user", content: "question" }], start: 0, total: 1 }),
status: () => Promise.resolve(status(oldSession.id)),
streamSnapshot: () => Promise.resolve({ seq: 0, partial: null }),
};
const controller = new SessionController(
() => state,
(patch) => { state = { ...state, ...patch }; },
() => undefined,
undefined,
{ api, socket },
);
await controller.selectSession(oldSession, { updateUrl: false });
expect(state.messages).toEqual([{ role: "user", parts: [{ type: "text", text: "question" }] }]);
// A watermark of 0 must not drop a fresh streamed delta (seq >= 1).
socket.emit({ type: "assistant.delta", text: "new turn", seq: 1 });
runPendingAnimationFrames();
expect(state.messages).toEqual([
{ role: "user", parts: [{ type: "text", text: "question" }] },
{ role: "assistant", parts: [{ type: "text", text: "new turn" }] },
]);
});
});
@@ -5,7 +5,7 @@ import type { SessionUiEvent } from "../sessionSocket";
import type { SessionEventSocket } from "./sessionController"; import type { SessionEventSocket } from "./sessionController";
export { api as defaultApi } from "../api"; export { api as defaultApi } from "../api";
export type { MessagePage, PromptAttachment, SessionActivity, SessionInfo, SessionRef, SessionStatus, Workspace } from "../api"; export type { MessagePage, PromptAttachment, SessionActivity, SessionInfo, SessionRef, SessionStatus, SessionStreamSnapshot, Workspace } from "../api";
export type { AppState } from "../appState"; export type { AppState } from "../appState";
export class MemoryStorage implements Storage { export class MemoryStorage implements Storage {
+42 -55
View File
@@ -1,4 +1,4 @@
import { api as defaultApi, type CommandResult, type PromptAttachment, type QueuedSessionMessage, type SessionActivity, type SessionBulkFailure, type SessionCleanupExecuteResponse, type SessionInfo, type SessionRef, type SessionStatus, type Workspace } from "../api"; import { api as defaultApi, type CommandResult, type PromptAttachment, type QueuedSessionMessage, type SessionActivity, type SessionBulkFailure, type SessionCleanupExecuteResponse, type SessionInfo, type SessionRef, type SessionStatus, type SessionStreamSnapshot, type Workspace } from "../api";
import type { AppState } from "../appState"; import type { AppState } from "../appState";
import { forgetCachedNewSession, isCachedNewSessionInfo, markCachedNewSessionInfo, mergeCachedNewSessions, rememberCachedNewSession, stripCachedNewSessionMarker } from "../cachedNewSessions"; import { forgetCachedNewSession, isCachedNewSessionInfo, markCachedNewSessionInfo, mergeCachedNewSessions, rememberCachedNewSession, stripCachedNewSessionMarker } from "../cachedNewSessions";
import { textMessage } from "../chatMessages"; import { textMessage } from "../chatMessages";
@@ -72,7 +72,12 @@ export class SessionController {
private readonly api: typeof defaultApi; private readonly api: typeof defaultApi;
private readonly transcripts: ChatTranscriptStore; private readonly transcripts: ChatTranscriptStore;
private selectionSeq = 0; private selectionSeq = 0;
private catchupStreamSessionId: string | undefined; // Join-time stream watermark for the selected session. `seq` is the
// `SessionEventHub` sequence captured together with the seeded partial by the
// stream snapshot: buffered/live events with `seq <= seq` are already reflected
// in the committed history + seeded partial and must be dropped, so every event
// applies exactly once. Reset whenever the selection changes.
private streamWatermark: { sessionId: string; seq: number } | undefined;
private pendingTranscriptEvents: SessionUiEvent[] = []; private pendingTranscriptEvents: SessionUiEvent[] = [];
private pendingStatusBySession = new Map<string, SessionStatus>(); private pendingStatusBySession = new Map<string, SessionStatus>();
private pendingActivityBySession = new Map<string, SessionActivity>(); private pendingActivityBySession = new Map<string, SessionActivity>();
@@ -111,13 +116,13 @@ export class SessionController {
clearActiveSession() { clearActiveSession() {
this.selectionSeq += 1; this.selectionSeq += 1;
this.socket.close(); this.socket.close();
this.catchupStreamSessionId = undefined; this.streamWatermark = undefined;
this.clearPendingUpdates(); this.clearPendingUpdates();
// Note: sendingPrompts is intentionally NOT cleared here. Deselecting a // Note: sendingPrompts is intentionally NOT cleared here. Deselecting a
// session must not cancel the in-flight upload indicator of the session // session must not cancel the in-flight upload indicator of the session
// that is still sending; the per-session entry is cleared by send()'s // that is still sending; the per-session entry is cleared by send()'s
// finally block when the request settles. // finally block when the request settles.
this.setState({ selectedSession: undefined, messages: [], messagePageStart: 0, messagePageEnd: 0, messagePageTotal: 0, isLoadingEarlierMessages: false, isReceivingPartialStream: false, status: undefined, activity: undefined, availableThinkingLevels: [] }); this.setState({ selectedSession: undefined, messages: [], messagePageStart: 0, messagePageEnd: 0, messagePageTotal: 0, isLoadingEarlierMessages: false, status: undefined, activity: undefined, availableThinkingLevels: [] });
} }
deselectSession(options?: { forgetRememberedSelection?: boolean | undefined; updateUrl?: boolean | undefined }) { deselectSession(options?: { forgetRememberedSelection?: boolean | undefined; updateUrl?: boolean | undefined }) {
@@ -161,7 +166,7 @@ export class SessionController {
this.sessionSelection.rememberSession({ ...session, cwd: this.workspaceSelectionKey(session.cwd) }); this.sessionSelection.rememberSession({ ...session, cwd: this.workspaceSelectionKey(session.cwd) });
const seq = ++this.selectionSeq; const seq = ++this.selectionSeq;
this.socket.close(); this.socket.close();
this.catchupStreamSessionId = undefined; this.streamWatermark = undefined;
this.clearPendingUpdates(); this.clearPendingUpdates();
const transcriptKey = this.sessionCacheKey(session.id); const transcriptKey = this.sessionCacheKey(session.id);
const cached = this.transcripts.cachedView(transcriptKey); const cached = this.transcripts.cachedView(transcriptKey);
@@ -169,7 +174,6 @@ export class SessionController {
selectedSession: session, selectedSession: session,
...cached, ...cached,
isLoadingEarlierMessages: false, isLoadingEarlierMessages: false,
isReceivingPartialStream: false,
status: session.archived === true ? undefined : this.getState().sessionStatuses[session.id], status: session.archived === true ? undefined : this.getState().sessionStatuses[session.id],
activity: session.archived === true ? undefined : this.getState().sessionActivities[session.id], activity: session.archived === true ? undefined : this.getState().sessionActivities[session.id],
availableThinkingLevels: [], availableThinkingLevels: [],
@@ -179,7 +183,7 @@ export class SessionController {
const page = await this.api.messages(session, { limit: MESSAGE_PAGE_SIZE }, selectedMachineId(this.getState())); const page = await this.api.messages(session, { limit: MESSAGE_PAGE_SIZE }, selectedMachineId(this.getState()));
if (seq !== this.selectionSeq || this.getState().selectedSession?.id !== session.id) return; if (seq !== this.selectionSeq || this.getState().selectedSession?.id !== session.id) return;
const history = this.transcripts.mergeHistory(transcriptKey, page); const history = this.transcripts.mergeHistory(transcriptKey, page);
this.setState({ ...history, isLoadingEarlierMessages: false, isReceivingPartialStream: false, status: undefined, activity: undefined }); this.setState({ ...history, isLoadingEarlierMessages: false, status: undefined, activity: undefined });
if (options?.updateUrl !== false) this.updateUrl(); if (options?.updateUrl !== false) this.updateUrl();
return; return;
} }
@@ -765,17 +769,32 @@ export class SessionController {
return this.selectedSessionRefreshes.request(key, async () => { return this.selectedSessionRefreshes.request(key, async () => {
if (!this.isCurrentRefreshTarget(target)) return; if (!this.isCurrentRefreshTarget(target)) return;
this.flushPendingUpdates(); this.flushPendingUpdates();
const [page, status] = await Promise.all([ const [page, status, streamSnapshot] = await Promise.all([
this.api.messages(target.session, { limit: MESSAGE_PAGE_SIZE }, target.machineId), this.api.messages(target.session, { limit: MESSAGE_PAGE_SIZE }, target.machineId),
this.api.status(target.session, target.machineId), this.api.status(target.session, target.machineId),
// The stream snapshot is a progressive enhancement. An older/not-yet-
// restarted session daemon or a remote machine on an older pi-web has no
// `stream-snapshot` route and returns 404; treat any failure as "no
// partial to seed". A `seq: 0` watermark drops nothing (live events start
// at seq 1, and un-stamped events fail open), so the core transcript
// still loads and streams normally.
this.api.streamSnapshot(target.session, target.machineId).catch((): SessionStreamSnapshot => ({ seq: 0, partial: null })),
]); ]);
if (!this.isCurrentRefreshTarget(target)) return; if (!this.isCurrentRefreshTarget(target)) return;
// Seed the in-flight partial assistant message on top of committed history
// and record the snapshot's sequence as the watermark. Buffered/live events
// with `seq <= watermark` are already reflected here and are dropped by
// `applyEvent`; later events (`seq > watermark`) stream in on top. The
// partial is seeded into the in-memory transcript only, never the raw
// history cache, so it never persists.
const history = this.transcripts.mergeHistory(key, page); const history = this.transcripts.mergeHistory(key, page);
const messages = this.transcripts.seedStreamingPartial(history.messages, streamSnapshot.partial);
this.streamWatermark = { sessionId: target.session.id, seq: streamSnapshot.seq };
this.setState({ this.setState({
...history, ...history,
messages,
status, status,
activity: this.getState().sessionActivities[target.session.id], activity: this.getState().sessionActivities[target.session.id],
...this.setStreamCatchup(status.isStreaming ? target.session.id : undefined),
}); });
this.applyStatus(status); this.applyStatus(status);
}); });
@@ -859,7 +878,7 @@ export class SessionController {
this.sessionSelection.rememberSession({ ...session, cwd: this.workspaceSelectionKey(session.cwd) }); this.sessionSelection.rememberSession({ ...session, cwd: this.workspaceSelectionKey(session.cwd) });
this.selectionSeq += 1; this.selectionSeq += 1;
this.socket.close(); this.socket.close();
this.catchupStreamSessionId = undefined; this.streamWatermark = undefined;
this.clearPendingUpdates(); this.clearPendingUpdates();
const state = this.getState(); const state = this.getState();
const pendingStart = this.pendingSessionStarts.get(session.id); const pendingStart = this.pendingSessionStarts.get(session.id);
@@ -872,7 +891,6 @@ export class SessionController {
messagePageEnd: 0, messagePageEnd: 0,
messagePageTotal: 0, messagePageTotal: 0,
isLoadingEarlierMessages: false, isLoadingEarlierMessages: false,
isReceivingPartialStream: false,
status: undefined, status: undefined,
activity, activity,
availableThinkingLevels: [], availableThinkingLevels: [],
@@ -1061,7 +1079,6 @@ export class SessionController {
status: state.selectedSession?.id === status.sessionId ? status : state.status, status: state.selectedSession?.id === status.sessionId ? status : state.status,
activity: state.selectedSession?.id === status.sessionId && clearsStaleActivity ? undefined : state.activity, activity: state.selectedSession?.id === status.sessionId && clearsStaleActivity ? undefined : state.activity,
}); });
if (!status.isStreaming) this.finishStreamCatchup(status.sessionId);
} }
private applySessionName(sessionId: string, name: string | undefined) { private applySessionName(sessionId: string, name: string | undefined) {
@@ -1080,14 +1097,10 @@ export class SessionController {
} }
private applyEvent(event: SessionUiEvent) { private applyEvent(event: SessionUiEvent) {
const selectedSessionId = this.getState().selectedSession?.id; // Drop events already reflected in the seeded join snapshot (committed
if (this.catchupStreamSessionId !== undefined && this.catchupStreamSessionId === selectedSessionId) { // history + partial). Everything past the watermark applies exactly once,
if (event.type === "message.end" || event.type === "agent.end") { // so live content streams directly on top of the seeded partial.
this.finishStreamCatchup(this.catchupStreamSessionId); if (this.isStreamEventBelowWatermark(event)) return;
return;
}
if (isTranscriptEvent(event)) return;
}
// Status and activity arrive once per token (the server republishes them on // Status and activity arrive once per token (the server republishes them on
// every transcript event). Buffer them alongside high-frequency transcript // every transcript event). Buffer them alongside high-frequency transcript
@@ -1175,37 +1188,15 @@ export class SessionController {
this.pendingFrame = undefined; this.pendingFrame = undefined;
} }
// Stream catch-up is a single mode with two coupled facets that must never // Watermark filter for join-time exactly-once application. An event is below
// drift: the private `catchupStreamSessionId` guard (which suppresses live // the watermark when it belongs to the selected session's seeded snapshot
// transcript events while we lack the in-flight message prefix) and the // (`seq <= watermark.seq`); such events are already reflected in the committed
// public `isReceivingPartialStream` flag (which drives the "Catching up…" // history + seeded partial and must be dropped. Events with no `seq` (which
// badge). Route every mutation of the mode through this helper so the guard // should not occur on the per-session socket) are never dropped.
// and the badge can never disagree. Catch-up only ever applies to the private isStreamEventBelowWatermark(event: SessionUiEvent): boolean {
// selected session, so an active session id always implies the badge is on. const watermark = this.streamWatermark;
private setStreamCatchup(sessionId: string | undefined): Pick<AppState, "isReceivingPartialStream"> { if (watermark === undefined || watermark.sessionId !== this.getState().selectedSession?.id) return false;
this.catchupStreamSessionId = sessionId; return event.seq !== undefined && event.seq <= watermark.seq;
return { isReceivingPartialStream: sessionId !== undefined };
}
private finishStreamCatchup(sessionId: string) {
const isSelected = this.getState().selectedSession?.id === sessionId;
const wasCatchingUp = this.catchupStreamSessionId === sessionId || (isSelected && this.getState().isReceivingPartialStream);
if (!wasCatchingUp) return;
this.catchupStreamSessionId = undefined;
if (isSelected) this.setState({ isReceivingPartialStream: false });
void this.refreshMessages(sessionId);
}
private async refreshMessages(sessionId: string) {
try {
const session = this.getState().selectedSession;
if (session?.id !== sessionId) return;
const page = await this.api.messages(session, { limit: MESSAGE_PAGE_SIZE }, selectedMachineId(this.getState()));
if (this.getState().selectedSession?.id !== sessionId) return;
this.setState(this.transcripts.mergeHistory(this.sessionCacheKey(sessionId), page));
} catch (error) {
if (this.getState().selectedSession?.id === sessionId) this.setState({ error: String(error) });
}
} }
} }
@@ -1377,10 +1368,6 @@ function sessionMessageCountPatch(state: AppState, sessionId: string, messageCou
}; };
} }
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);
}
function isHighFrequencyTranscriptEvent(event: SessionUiEvent): boolean { function isHighFrequencyTranscriptEvent(event: SessionUiEvent): boolean {
return event.type === "assistant.delta" || event.type === "assistant.thinking.delta" || event.type === "shell.chunk"; return event.type === "assistant.delta" || event.type === "assistant.thinking.delta" || event.type === "shell.chunk";
} }
+63 -1
View File
@@ -18,7 +18,7 @@ describe("SessionEventHub", () => {
hub.publish("s1", { type: "assistant.delta", text: "hello" }); hub.publish("s1", { type: "assistant.delta", text: "hello" });
expect(sessionSocket.send).toHaveBeenCalledWith(JSON.stringify({ type: "assistant.delta", text: "hello" })); expect(sessionSocket.send).toHaveBeenCalledWith(JSON.stringify({ type: "assistant.delta", text: "hello", seq: 1 }));
expect(otherSocket.send).not.toHaveBeenCalled(); expect(otherSocket.send).not.toHaveBeenCalled();
}); });
@@ -34,6 +34,7 @@ describe("SessionEventHub", () => {
expect(socket.send).toHaveBeenCalledWith(JSON.stringify({ expect(socket.send).toHaveBeenCalledWith(JSON.stringify({
type: "message.end", type: "message.end",
message: { role: "assistant", content: [{ type: "thinking", thinking: "private chain", redacted: true }, { type: "text", text: "visible answer" }] }, message: { role: "assistant", content: [{ type: "thinking", thinking: "private chain", redacted: true }, { type: "text", text: "visible answer" }] },
seq: 1,
})); }));
expect(thinkingBlock.thinkingSignature).toBe("opaque-provider-payload"); expect(thinkingBlock.thinkingSignature).toBe("opaque-provider-payload");
}); });
@@ -76,4 +77,65 @@ describe("SessionEventHub", () => {
expect(globalSocket.send).toHaveBeenCalledWith(JSON.stringify({ type: "status.update", status })); expect(globalSocket.send).toHaveBeenCalledWith(JSON.stringify({ type: "status.update", status }));
expect(sessionSocket.send).not.toHaveBeenCalled(); expect(sessionSocket.send).not.toHaveBeenCalled();
}); });
it("stamps a monotonically increasing per-session seq on published events", () => {
const hub = new SessionEventHub();
const socket = new FakeSocket();
hub.add("s1", socket);
hub.publish("s1", { type: "assistant.delta", text: "a" });
hub.publish("s1", { type: "assistant.delta", text: "b" });
hub.publish("s1", { type: "assistant.delta", text: "c" });
expect(socket.send).toHaveBeenNthCalledWith(1, JSON.stringify({ type: "assistant.delta", text: "a", seq: 1 }));
expect(socket.send).toHaveBeenNthCalledWith(2, JSON.stringify({ type: "assistant.delta", text: "b", seq: 2 }));
expect(socket.send).toHaveBeenNthCalledWith(3, JSON.stringify({ type: "assistant.delta", text: "c", seq: 3 }));
});
it("advances seq even when no sockets are attached so the watermark stays accurate", () => {
const hub = new SessionEventHub();
hub.publish("s1", { type: "assistant.delta", text: "a" });
hub.publish("s1", { type: "assistant.delta", text: "b" });
expect(hub.currentSeq("s1")).toBe(2);
const socket = new FakeSocket();
hub.add("s1", socket);
hub.publish("s1", { type: "assistant.delta", text: "c" });
expect(socket.send).toHaveBeenCalledWith(JSON.stringify({ type: "assistant.delta", text: "c", seq: 3 }));
});
it("tracks seq independently per session", () => {
const hub = new SessionEventHub();
const s1 = new FakeSocket();
const s2 = new FakeSocket();
hub.add("s1", s1);
hub.add("s2", s2);
hub.publish("s1", { type: "assistant.delta", text: "a" });
hub.publish("s1", { type: "assistant.delta", text: "b" });
hub.publish("s2", { type: "assistant.delta", text: "x" });
expect(hub.currentSeq("s1")).toBe(2);
expect(hub.currentSeq("s2")).toBe(1);
expect(s1.send).toHaveBeenLastCalledWith(JSON.stringify({ type: "assistant.delta", text: "b", seq: 2 }));
expect(s2.send).toHaveBeenCalledWith(JSON.stringify({ type: "assistant.delta", text: "x", seq: 1 }));
});
it("reports zero seq for a session that has never published", () => {
const hub = new SessionEventHub();
expect(hub.currentSeq("never")).toBe(0);
});
it("does not stamp seq on global events", () => {
const hub = new SessionEventHub();
const globalSocket = new FakeSocket();
hub.addGlobal(globalSocket);
hub.publishGlobal({ type: "session.name", sessionId: "s1", name: "Renamed" });
expect(globalSocket.send).toHaveBeenCalledWith(JSON.stringify({ type: "session.name", sessionId: "s1", name: "Renamed" }));
});
}); });
+14 -1
View File
@@ -11,6 +11,7 @@ export interface RealtimeSocket {
export class SessionEventHub { export class SessionEventHub {
private readonly socketsBySession = new Map<string, Set<RealtimeSocket>>(); private readonly socketsBySession = new Map<string, Set<RealtimeSocket>>();
private readonly globalSockets = new Set<RealtimeSocket>(); private readonly globalSockets = new Set<RealtimeSocket>();
private readonly seqBySession = new Map<string, number>();
add(sessionId: string, socket: RealtimeSocket): void { add(sessionId: string, socket: RealtimeSocket): void {
let sockets = this.socketsBySession.get(sessionId); let sockets = this.socketsBySession.get(sessionId);
@@ -30,12 +31,24 @@ export class SessionEventHub {
} }
publish(sessionId: string, event: SessionUiEvent): void { publish(sessionId: string, event: SessionUiEvent): void {
const payload = JSON.stringify(projectBrowserSessionEvent(event)); const seq = (this.seqBySession.get(sessionId) ?? 0) + 1;
this.seqBySession.set(sessionId, seq);
const payload = JSON.stringify({ ...projectBrowserSessionEvent(event), seq });
for (const socket of this.socketsBySession.get(sessionId) ?? []) { for (const socket of this.socketsBySession.get(sessionId) ?? []) {
if (socket.readyState === socket.OPEN) socket.send(payload); if (socket.readyState === socket.OPEN) socket.send(payload);
} }
} }
/**
* Last per-session sequence number stamped by {@link publish} (0 before any
* event). Callers building a join-time stream snapshot read this as the
* watermark: buffered live events with `seq <= currentSeq` are already
* reflected in the snapshot's partial and must be dropped by the client.
*/
currentSeq(sessionId: string): number {
return this.seqBySession.get(sessionId) ?? 0;
}
publishGlobal(event: GlobalSessionEvent): void { publishGlobal(event: GlobalSessionEvent): void {
this.publishRealtime(event); this.publishRealtime(event);
} }
@@ -566,3 +566,65 @@ describe("PiSessionService lifecycle, listing, and reload", () => {
await service.dispose(); await service.dispose();
}); });
}); });
describe("PiSessionService.streamSnapshot", () => {
it("returns a null partial with the current watermark when idle", async () => {
const hub = new CapturingSessionEventHub();
const fake = fakeRuntime("snap-idle");
const service = new PiSessionService(hub, {
agentDir: TEST_AGENT_DIR,
createAgentRuntime: runtimeCreator(fake.runtime),
sessionManager: sessionGateway([]),
heartbeatIntervalMs: 60_000,
});
try {
await service.start("/workspace");
const snapshot = await service.streamSnapshot(sessionRef("snap-idle"));
expect(snapshot).toEqual({ seq: 0, partial: null });
} finally {
await service.dispose();
}
});
it("projects the in-flight partial and matches the event watermark mid-stream", async () => {
const hub = new CapturingSessionEventHub();
const streamingMessage = {
role: "assistant",
content: [
{ type: "thinking", thinking: "weighing options", thinkingSignature: "opaque" },
{ type: "text", text: "partial answer" },
{ type: "toolCall", id: "call-1", name: "edit", arguments: { path: "a.ts" } },
],
};
const fake = fakeRuntime("snap-live", { state: { streamingMessage } });
const service = new PiSessionService(hub, {
agentDir: TEST_AGENT_DIR,
createAgentRuntime: runtimeCreator(fake.runtime),
sessionManager: sessionGateway([]),
heartbeatIntervalMs: 60_000,
});
try {
await service.start("/workspace");
// Advance the per-session watermark to a known value.
hub.setSeq("snap-live", 5);
const snapshot = await service.streamSnapshot(sessionRef("snap-live"));
expect(snapshot.seq).toBe(5);
expect(snapshot.partial).toEqual({
role: "assistant",
content: [
{ type: "thinking", thinking: "weighing options" },
{ type: "text", text: "partial answer" },
{ type: "toolCall", id: "call-1", name: "edit", arguments: { path: "a.ts" } },
],
});
// The runtime message is not mutated by the browser projection.
expect(streamingMessage.content[0]).toHaveProperty("thinkingSignature", "opaque");
} finally {
await service.dispose();
}
});
});
@@ -6,6 +6,7 @@ import type { PiAgentSession, PiSessionManager, PiSessionRuntime, PiSessionServi
export class CapturingSessionEventHub extends SessionEventHub { export class CapturingSessionEventHub extends SessionEventHub {
readonly sessionEvents: { sessionId: string; event: SessionUiEvent }[] = []; readonly sessionEvents: { sessionId: string; event: SessionUiEvent }[] = [];
readonly globalEvents: GlobalSessionEvent[] = []; readonly globalEvents: GlobalSessionEvent[] = [];
private readonly seqBySessionOverride = new Map<string, number>();
override publish(sessionId: string, event: SessionUiEvent): void { override publish(sessionId: string, event: SessionUiEvent): void {
this.sessionEvents.push({ sessionId, event }); this.sessionEvents.push({ sessionId, event });
@@ -14,6 +15,15 @@ export class CapturingSessionEventHub extends SessionEventHub {
override publishGlobal(event: GlobalSessionEvent): void { override publishGlobal(event: GlobalSessionEvent): void {
this.globalEvents.push(event); this.globalEvents.push(event);
} }
/** Test seam: set the per-session watermark returned by {@link currentSeq}. */
setSeq(sessionId: string, value: number): void {
this.seqBySessionOverride.set(sessionId, value);
}
override currentSeq(sessionId: string): number {
return this.seqBySessionOverride.get(sessionId) ?? 0;
}
} }
export type SessionGateway = NonNullable<PiSessionServiceDependencies["sessionManager"]>; export type SessionGateway = NonNullable<PiSessionServiceDependencies["sessionManager"]>;
@@ -68,6 +78,7 @@ export function fakeRuntime(sessionId = "session-1", patch: Partial<TestSession>
sessionId, sessionId,
sessionFile: `/tmp/${sessionId}.jsonl`, sessionFile: `/tmp/${sessionId}.jsonl`,
messages: [], messages: [],
state: {},
sessionName: undefined, sessionName: undefined,
model: undefined, model: undefined,
thinkingLevel: "off", thinkingLevel: "off",
+30 -1
View File
@@ -14,7 +14,8 @@ import {
type CreateAgentSessionRuntimeFactory, type CreateAgentSessionRuntimeFactory,
type EditToolDetails, type EditToolDetails,
} from "@earendil-works/pi-coding-agent"; } from "@earendil-works/pi-coding-agent";
import type { ClientArchiveSessionsResponse, ClientCommand, ClientCommandResult, ClientMessagePage, ClientSession, ClientSessionCleanupExecuteResponse, ClientSessionCleanupPreviewResponse, ClientSessionModel, ClientSessionStatus, ClientThinkingLevel, SessionUiEvent } from "../types.js"; import type { ClientArchiveSessionsResponse, ClientCommand, ClientCommandResult, ClientMessagePage, ClientSession, ClientSessionCleanupExecuteResponse, ClientSessionCleanupPreviewResponse, ClientSessionModel, ClientSessionStatus, ClientThinkingLevel, SessionStreamSnapshot, SessionUiEvent } from "../types.js";
import { projectBrowserMessage } from "../browserMessageProjection.js";
import { pageMessagesAtSafeBoundary } from "./messagePaging.js"; import { pageMessagesAtSafeBoundary } from "./messagePaging.js";
import type { SessionEventHub } from "../realtime/sessionEventHub.js"; import type { SessionEventHub } from "../realtime/sessionEventHub.js";
import { BUILTIN_COMMANDS } from "./builtinCommands.js"; import { BUILTIN_COMMANDS } from "./builtinCommands.js";
@@ -211,6 +212,14 @@ export interface PiAgentSession {
sessionFile: string | undefined; sessionFile: string | undefined;
sessionName: string | undefined; sessionName: string | undefined;
messages: readonly unknown[]; messages: readonly unknown[];
/**
* Narrow read of the SDK `AgentState`. Only the in-flight partial is consumed
* here: `state.streamingMessage` is the current streamed assistant message
* (an `AssistantMessage`) while a turn is mid-stream, and `undefined`
* otherwise (idle, or during post-message tool execution). Used by
* {@link PiSessionService.streamSnapshot} to seed a joining client.
*/
readonly state: { readonly streamingMessage?: unknown };
model: AgentModel | undefined; model: AgentModel | undefined;
thinkingLevel: ClientThinkingLevel; thinkingLevel: ClientThinkingLevel;
isStreaming: boolean; isStreaming: boolean;
@@ -981,6 +990,26 @@ export class PiSessionService implements SessionRouteService {
return this.statusFromSession(await this.getOrOpen(ref)); return this.statusFromSession(await this.getOrOpen(ref));
} }
/**
* Join-time snapshot of the in-flight assistant stream. The `seq` watermark and
* the partial are read together in one synchronous tick (no await between the
* `currentSeq` read and the `state.streamingMessage` read) so a joining client
* can seed the partial and then apply only buffered live events with
* `seq > snapshot.seq`. The partial is browser-projected to strip thinking
* signatures; it is `null` when no assistant message is mid-stream.
*/
async streamSnapshot(ref: PiSessionLookup): Promise<SessionStreamSnapshot> {
const session = await this.getOrOpen(ref);
// Single consistent tick: capture the watermark and the partial together so
// the seq matches the partial the client seeds against.
const seq = this.events.currentSeq(session.sessionId);
const streamingMessage = session.state.streamingMessage;
const partial = streamingMessage === undefined || streamingMessage === null
? null
: projectBrowserMessage(streamingMessage);
return { seq, partial };
}
async availableModels(ref: PiSessionLookup): Promise<ClientSessionModel[]> { async availableModels(ref: PiSessionLookup): Promise<ClientSessionModel[]> {
const session = await this.getOrOpen(ref); const session = await this.getOrOpen(ref);
session.modelRegistry.refresh(); session.modelRegistry.refresh();
+48 -1
View File
@@ -2,7 +2,7 @@ import { resolve } from "node:path";
import Fastify, { type FastifyInstance } from "fastify"; import Fastify, { type FastifyInstance } from "fastify";
import fastifyWebsocket from "@fastify/websocket"; import fastifyWebsocket from "@fastify/websocket";
import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { afterEach, beforeEach, describe, expect, it } from "vitest";
import type { MessagePage, SessionBulkArchiveResponse, SessionBulkDeleteArchivedResponse, SessionBulkMutationRef, SessionCleanupExecuteResponse, SessionCleanupPreviewResponse, SessionStatus } from "../../shared/apiTypes.js"; import type { MessagePage, SessionBulkArchiveResponse, SessionBulkDeleteArchivedResponse, SessionBulkMutationRef, SessionCleanupExecuteResponse, SessionCleanupPreviewResponse, SessionStatus, SessionStreamSnapshot } from "../../shared/apiTypes.js";
import { SessionEventHub } from "../realtime/sessionEventHub.js"; import { SessionEventHub } from "../realtime/sessionEventHub.js";
import { PiSessionService, type PiSessionManagerGateway } from "./piSessionService.js"; import { PiSessionService, type PiSessionManagerGateway } from "./piSessionService.js";
import type { SessionRouteLookup, SessionRouteService } from "./sessionService.js"; import type { SessionRouteLookup, SessionRouteService } from "./sessionService.js";
@@ -168,6 +168,46 @@ describe("session routes", () => {
} }
}); });
it("returns the join-time stream snapshot, forwarding workspace context", async () => {
const routeApp = Fastify({ logger: false });
await routeApp.register(fastifyWebsocket);
const eventHub = new SessionEventHub();
const routeService = new CapturingRouteSessionService();
routeService.streamSnapshotResponse = { seq: 7, partial: { role: "assistant", content: [{ type: "text", text: "partial" }] } };
registerSessionRoutes(routeApp, routeService, eventHub);
try {
const requestCwd = resolve("/repo");
const response = await routeApp.inject({ method: "GET", url: `/sessions/session-1/stream-snapshot?cwd=${encodeURIComponent(requestCwd)}` });
expect(response.statusCode).toBe(200);
expect(response.json()).toEqual({ seq: 7, partial: { role: "assistant", content: [{ type: "text", text: "partial" }] } });
expect(routeService.streamSnapshotCalls).toEqual([{ id: "session-1", cwd: requestCwd }]);
} finally {
await routeService.dispose();
await routeApp.close();
}
});
it("maps stream-snapshot lookup failures to 404", async () => {
const routeApp = Fastify({ logger: false });
await routeApp.register(fastifyWebsocket);
const eventHub = new SessionEventHub();
const routeService = new CapturingRouteSessionService();
routeService.streamSnapshot = () => Promise.reject(new Error("Session not found"));
registerSessionRoutes(routeApp, routeService, eventHub);
try {
const response = await routeApp.inject({ method: "GET", url: "/sessions/missing/stream-snapshot" });
expect(response.statusCode).toBe(404);
expect(response.json()).toEqual({ error: "Session not found" });
} finally {
await routeService.dispose();
await routeApp.close();
}
});
it("clears a session queue with workspace context and returns fresh status", async () => { it("clears a session queue with workspace context and returns fresh status", async () => {
const routeApp = Fastify({ logger: false }); const routeApp = Fastify({ logger: false });
await routeApp.register(fastifyWebsocket); await routeApp.register(fastifyWebsocket);
@@ -306,6 +346,8 @@ class CapturingRouteSessionService implements SessionRouteService {
readonly reloadCalls: SessionRouteLookup[] = []; readonly reloadCalls: SessionRouteLookup[] = [];
readonly clearQueueCalls: SessionRouteLookup[] = []; readonly clearQueueCalls: SessionRouteLookup[] = [];
messagesResponse: unknown[] | MessagePage = []; messagesResponse: unknown[] | MessagePage = [];
streamSnapshotResponse: SessionStreamSnapshot = { seq: 0, partial: null };
readonly streamSnapshotCalls: SessionRouteLookup[] = [];
readonly cleanupPreviewCalls: NormalizedSessionCleanupRequest[] = []; readonly cleanupPreviewCalls: NormalizedSessionCleanupRequest[] = [];
readonly cleanupCalls: NormalizedSessionCleanupRequest[] = []; readonly cleanupCalls: NormalizedSessionCleanupRequest[] = [];
readonly bulkArchiveCalls: SessionBulkMutationRef[][] = []; readonly bulkArchiveCalls: SessionBulkMutationRef[][] = [];
@@ -379,6 +421,11 @@ class CapturingRouteSessionService implements SessionRouteService {
}); });
} }
streamSnapshot(lookup: SessionRouteLookup): Promise<SessionStreamSnapshot> {
this.streamSnapshotCalls.push(lookup);
return Promise.resolve(this.streamSnapshotResponse);
}
availableModels(): Promise<[]> { return Promise.resolve([]); } availableModels(): Promise<[]> { return Promise.resolve([]); }
setModel(): never { throw unusedRouteMethod("setModel"); } setModel(): never { throw unusedRouteMethod("setModel"); }
cycleModel(): never { throw unusedRouteMethod("cycleModel"); } cycleModel(): never { throw unusedRouteMethod("cycleModel"); }
+8
View File
@@ -99,6 +99,14 @@ export function registerSessionRoutes(app: FastifyInstance, sessions: SessionRou
} }
}); });
app.get<{ Params: { sessionId: string }; Querystring: SessionQuery }>(`${prefix}/sessions/:sessionId/stream-snapshot`, async (request, reply) => {
try {
return await sessions.streamSnapshot(sessionLookupFromQuery(request.params.sessionId, request.query));
} catch (error) {
return reply.code(404).send({ error: errorMessage(error) });
}
});
app.get<{ Params: { sessionId: string }; Querystring: SessionQuery }>(`${prefix}/sessions/:sessionId/models`, async (request, reply) => { app.get<{ Params: { sessionId: string }; Querystring: SessionQuery }>(`${prefix}/sessions/:sessionId/models`, async (request, reply) => {
try { try {
return { models: await sessions.availableModels(sessionLookupFromQuery(request.params.sessionId, request.query)) }; return { models: await sessions.availableModels(sessionLookupFromQuery(request.params.sessionId, request.query)) };
+2
View File
@@ -16,6 +16,7 @@ import type {
ClientSessionRef, ClientSessionRef,
ClientSessionStatus, ClientSessionStatus,
ClientThinkingLevel, ClientThinkingLevel,
SessionStreamSnapshot,
} from "../types.js"; } from "../types.js";
import type { NormalizedSessionCleanupRequest } from "./sessionCleanup.js"; import type { NormalizedSessionCleanupRequest } from "./sessionCleanup.js";
@@ -34,6 +35,7 @@ export interface SessionRouteService {
start(cwd: string): Promise<ClientSession>; start(cwd: string): Promise<ClientSession>;
messages(ref: SessionRouteLookup, page?: { before?: number; limit?: number }): Promise<unknown[] | ClientMessagePage>; messages(ref: SessionRouteLookup, page?: { before?: number; limit?: number }): Promise<unknown[] | ClientMessagePage>;
status(ref: SessionRouteLookup): Promise<ClientSessionStatus>; status(ref: SessionRouteLookup): Promise<ClientSessionStatus>;
streamSnapshot(ref: SessionRouteLookup): Promise<SessionStreamSnapshot>;
clearQueue(ref: SessionRouteLookup): Promise<ClientSessionStatus>; clearQueue(ref: SessionRouteLookup): Promise<ClientSessionStatus>;
availableModels(ref: SessionRouteLookup): Promise<ClientSessionModel[]>; availableModels(ref: SessionRouteLookup): Promise<ClientSessionModel[]>;
setModel(ref: SessionRouteLookup, provider: string, modelId: string): Promise<ClientSessionStatus>; setModel(ref: SessionRouteLookup, provider: string, modelId: string): Promise<ClientSessionStatus>;
+1
View File
@@ -9,6 +9,7 @@ export type {
SessionCleanupPreviewResponse as ClientSessionCleanupPreviewResponse, SessionCleanupPreviewResponse as ClientSessionCleanupPreviewResponse,
SessionCleanupExecuteResponse as ClientSessionCleanupExecuteResponse, SessionCleanupExecuteResponse as ClientSessionCleanupExecuteResponse,
MessagePage as ClientMessagePage, MessagePage as ClientMessagePage,
SessionStreamSnapshot,
SessionStatus as ClientSessionStatus, SessionStatus as ClientSessionStatus,
SessionModel as ClientSessionModel, SessionModel as ClientSessionModel,
ThinkingLevel as ClientThinkingLevel, ThinkingLevel as ClientThinkingLevel,
+24 -2
View File
@@ -689,12 +689,34 @@ export interface MessagePage {
total: number; total: number;
} }
/**
* Join-time snapshot of a session's in-flight assistant stream. `seq` is the
* `SessionEventHub` watermark captured together with `partial` in a single tick,
* so a joining client can seed `partial` and then apply only buffered live events
* with `seq > snapshot.seq` (exactly-once). `partial` is a browser-projected
* in-flight `AssistantMessage` (thinking signatures stripped), or `null` when the
* session is not mid assistant-message stream.
*/
export interface SessionStreamSnapshot {
seq: number;
/** Browser-projected in-flight `AssistantMessage`, or `null` when idle. */
partial: unknown;
}
export type CommandResult = export type CommandResult =
| { type: "done"; message?: string; session?: SessionInfo; promptDraft?: string } | { type: "done"; message?: string; session?: SessionInfo; promptDraft?: string }
| { type: "select"; requestId: string; title: string; options: CommandOption[] } | { type: "select"; requestId: string; title: string; options: CommandOption[] }
| { type: "unsupported"; message: string }; | { type: "unsupported"; message: string };
export type SessionUiEvent = /**
* Transport-level per-session sequence stamp. `SessionEventHub.publish` assigns a
* monotonic `seq` to every per-session event as it is serialized to the socket.
* Clients use it as a watermark against the join-time stream snapshot so buffered
* live events are applied exactly once. Existing consumers may ignore it.
*/
export type SessionUiEvent = SessionUiEventBody & { seq?: number };
type SessionUiEventBody =
| { type: "message.append"; message: unknown } | { type: "message.append"; message: unknown }
| { type: "assistant.delta"; text: string } | { type: "assistant.delta"; text: string }
| { type: "assistant.thinking.delta"; text: string } | { type: "assistant.thinking.delta"; text: string }
@@ -715,5 +737,5 @@ export type SessionUiEvent =
| { type: "session.created"; session: SessionInfo } | { type: "session.created"; session: SessionInfo }
| { type: "pi.event"; eventType: string }; | { type: "pi.event"; eventType: string };
export type GlobalSessionEvent = Extract<SessionUiEvent, { type: "status.update" | "activity.update" | "session.name" | "session.created" }>; export type GlobalSessionEvent = Extract<SessionUiEventBody, { type: "status.update" | "activity.update" | "session.name" | "session.created" }>;
export type RealtimeEvent = GlobalSessionEvent | TerminalUiEvent | WorkspaceActivityUiEvent; export type RealtimeEvent = GlobalSessionEvent | TerminalUiEvent | WorkspaceActivityUiEvent;
+1
View File
@@ -51,6 +51,7 @@ export const FEDERATED_HTTP_ROUTES = [
{ method: "POST", path: "/sessions/bulk/delete-archived" }, { method: "POST", path: "/sessions/bulk/delete-archived" },
{ method: "GET", path: "/sessions/:sessionId/messages" }, { method: "GET", path: "/sessions/:sessionId/messages" },
{ method: "GET", path: "/sessions/:sessionId/status" }, { method: "GET", path: "/sessions/:sessionId/status" },
{ method: "GET", path: "/sessions/:sessionId/stream-snapshot" },
{ method: "GET", path: "/sessions/:sessionId/models" }, { method: "GET", path: "/sessions/:sessionId/models" },
{ method: "POST", path: "/sessions/:sessionId/model" }, { method: "POST", path: "/sessions/:sessionId/model" },
{ method: "POST", path: "/sessions/:sessionId/model/cycle" }, { method: "POST", path: "/sessions/:sessionId/model/cycle" },