perf: stream agent message deltas instead of full snapshots (#375)

app/api/agent/[id]/events/route.ts

@@ -5,9 +5,52 @@ export const dynamic = "force-dynamic";
 
 const OMITTED_EVENT_TYPES = new Set(["turn_start", "turn_end", "tool_execution_update"]);
 
+// assistantMessageEvent 中可安全转发给浏览器的轻量字段。`partial`（完整累积
+// 消息）被有意剥离——转发它会重新引入 O(n²) 传输（每个 delta 都重发此前
+// 生成的全部内容）。`done`/`error` 不在此列：它们的 reason 已包含在
+// message_end 的完整消息中，前端从 message_end 获取即可。
+const DELTA_FIELDS: Record<string, readonly string[]> = {
+  start: [],
+  text_start: ["contentIndex"],
+  text_delta: ["contentIndex", "delta"],
+  text_end: ["contentIndex", "content"],
+  thinking_start: ["contentIndex"],
+  thinking_delta: ["contentIndex", "delta"],
+  thinking_end: ["contentIndex", "content"],
+  toolcall_start: ["contentIndex"],
+  toolcall_delta: ["contentIndex", "delta"],
+  toolcall_end: ["contentIndex", "toolCall"],
+};
+
+// done/error 不转发：reason 已随 message_end 的完整消息下发，前端无需重复接收。
+const OMITTED_DELTA_TYPES = new Set(["done", "error"]);
+
 function toClientEvent(event: AgentEvent): AgentEvent | null {
   if (OMITTED_EVENT_TYPES.has(event.type)) return null;
   if (event.type === "message_update") {
+    const delta = event.assistantMessageEvent as ({ type: string } & Record<string, unknown>) | undefined;
+    const fields = delta && typeof delta.type === "string" ? DELTA_FIELDS[delta.type] : undefined;
+    if (delta && fields) {
+      // 增量事件：只转发轻量 delta 字段，浏览器端据此拼接流式消息。
+      const slimDelta: Record<string, unknown> = { type: delta.type };
+      for (const field of fields) slimDelta[field] = delta[field];
+      // toolcall_start 只携带 contentIndex；id/name 仅存在于 partial（完整累积
+      // 消息）中。只提取这两个字段注入，不转发整个 partial，避免重新引入
+      // O(n²) 传输——否则流式期间工具卡片无名无 id。
+      if (delta.type === "toolcall_start") {
+        const partial = delta.partial as
+          | { content?: Array<{ type: string; id?: string; name?: string }> }
+          | undefined;
+        const block = partial?.content?.[Number(delta.contentIndex)];
+        if (block?.type === "toolCall") {
+          slimDelta.id = block.id;
+          slimDelta.name = block.name;
+        }
+      }
+      return { type: "message_delta", assistantMessageEvent: slimDelta } as unknown as AgentEvent;
+    }
+    if (delta && OMITTED_DELTA_TYPES.has(delta.type)) return null;
+    // 无 delta 或未知 delta 类型：降级为完整快照（罕见兜底，浏览器覆盖校准）。
     const clientEvent = { ...event };
     delete clientEvent.assistantMessageEvent;
     return clientEvent;
@@ -48,6 +91,16 @@ export async function GET(
       // Send initial connected event
       encode({ type: "connected", sessionId: id });
 
+      // 重连恢复：若 session 正在流式（页面关闭后重新打开、网络闪断重连），
+      // 先把当前部分消息作为完整快照发给浏览器，让它重建增量拼接基座——
+      // 否则订阅只收到之后的 delta，关闭期间生成的内容要等 message_end
+      // 才显示。SDK 保证 streamingMessage 仅在消息流式期间非空，故此处不会
+      // 注入已完成的旧消息。
+      const streamingMessage = session.streamingMessage;
+      if (streamingMessage) {
+        encode({ type: "message_update", message: streamingMessage });
+      }
+
       const unsubscribe = session.onEvent((event) => {
         const clientEvent = toClientEvent(event);
         if (clientEvent) encode(clientEvent);

app/api/agent/events-route.test.mjs

@@ -7,11 +7,35 @@ const runningEventsSource = await readFile(new URL("./running/events/route.ts",
 
 test("agent SSE projects SDK events onto the fields consumed by the web client", () => {
   assert.match(agentEventsSource, /OMITTED_EVENT_TYPES = new Set\(\["turn_start", "turn_end", "tool_execution_update"\]\)/);
-  assert.match(agentEventsSource, /delete clientEvent\.assistantMessageEvent/);
   assert.match(agentEventsSource, /event\.type === "agent_end"\) return \{ type: "agent_end" \}/);
   assert.match(agentEventsSource, /const clientEvent = toClientEvent\(event\)/);
 });
 
+test("message_update is projected to a light delta event with partial stripped", () => {
+  // 轻量字段白名单存在，且覆盖 SDK 的全部 delta 类型
+  assert.match(agentEventsSource, /DELTA_FIELDS: Record<string, readonly string\[\]>/);
+  assert.match(agentEventsSource, /text_delta: \["contentIndex", "delta"\]/);
+  assert.match(agentEventsSource, /thinking_delta: \["contentIndex", "delta"\]/);
+  assert.match(agentEventsSource, /toolcall_delta: \["contentIndex", "delta"\]/);
+  assert.match(agentEventsSource, /toolcall_end: \["contentIndex", "toolCall"\]/);
+  // 已知 delta 类型 → 剥离 partial 后转发为 message_delta（浏览器增量拼接）
+  assert.match(agentEventsSource, /type: "message_delta", assistantMessageEvent: slimDelta/);
+  assert.doesNotMatch(agentEventsSource, /slimDelta\["partial"\]/);
+  // done/error 不转发（reason 已随 message_end 的完整消息下发）
+  assert.match(agentEventsSource, /OMITTED_DELTA_TYPES = new Set\(\["done", "error"\]\)/);
+  // toolcall_start 从 partial 提取 id/name（避免流式期间工具无名）
+  assert.match(agentEventsSource, /slimDelta\.id = block\.id/);
+  // 无 delta / 未知 delta 类型 → 降级为完整快照（罕见兜底），仍删除 assistantMessageEvent
+  assert.match(agentEventsSource, /delete clientEvent\.assistantMessageEvent/);
+  assert.match(agentEventsSource, /降级为完整快照/);
+});
+
+test("SSE reconnect injects a full-message snapshot when the session is mid-stream", () => {
+  assert.match(agentEventsSource, /const streamingMessage = session\.streamingMessage/);
+  assert.match(agentEventsSource, /\{ type: "message_update", message: streamingMessage \}/);
+  assert.match(agentEventsSource, /重建增量拼接基座/);
+});
+
 test("SSE routes reuse one TextEncoder per stream", () => {
   for (const source of [agentEventsSource, runningEventsSource]) {
     assert.equal((source.match(/new TextEncoder\(\)/g) ?? []).length, 1);

hooks/useAgentSession.ts

@@ -10,6 +10,7 @@ import type {
   SessionTreeNode,
 } from "@/lib/types";
 import { normalizeToolCalls } from "@/lib/normalize";
+import { applyAssistantDelta, type AssistantStreamDelta } from "@/lib/stream-delta";
 import { sendAgentCommand } from "@/lib/agent-client";
 import { getToolNamesForPreset, type ToolEntry } from "@/lib/tool-presets";
 import type { SessionStatsInfo } from "@/lib/pi-types";
@@ -407,6 +408,9 @@ export function useAgentSession(opts: UseAgentSessionOptions) {
   const thinkingLevelOverrideRef = useRef<Exclude<ThinkingLevelOption, "auto"> | null>(null);
   const promptRunIdRef = useRef(0);
   const optimisticUserMessageKeyRef = useRef<string | null>(null);
+  // 增量拼接中的流式消息（SDK 结构）。message_start 重置、message_delta 追加、
+  // message_update（服务端降级快照）覆盖校准、message_end 清空。
+  const streamingMessageRef = useRef<Partial<AgentMessage> | null>(null);
 
   const setToolPresetState = opts.setToolPreset ?? setToolPreset;
 
@@ -1052,6 +1056,7 @@ export function useAgentSession(opts: UseAgentSessionOptions) {
         // compacting, or continuing messages queued by extension handlers.
         // Keep the stream open until prompt_done/agent_settled and the idle grace.
         if (!agentRunningRef.current) break;
+        streamingMessageRef.current = null;
         setAgentPhase(null);
         setRetryInfo(null);
         dispatch({ type: "end" });
@@ -1126,16 +1131,35 @@ export function useAgentSession(opts: UseAgentSessionOptions) {
           break;
         }
         if (msg) {
+          // 完整快照（message_start 的初始消息，或服务端在 delta 缺失/未知时
+          // 降级发送的快照）：重置增量拼接基座并覆盖渲染。
+          streamingMessageRef.current = msg;
           dispatch({ type: "update", message: normalizeToolCalls(msg as AgentMessage) });
         }
         setAgentPhase(null);
         break;
       }
+      case "message_delta": {
+        // 轻量增量事件（服务端剥离 partial 后转发）：追加到本地流式消息。
+        // Same late-event guard as message_start/message_update.
+        if (!agentRunningRef.current) break;
+        const delta = event.assistantMessageEvent as AssistantStreamDelta | undefined;
+        if (!delta || typeof delta.type !== "string") break;
+        const prev = streamingMessageRef.current;
+        const next = applyAssistantDelta(prev, delta);
+        if (next === prev) break; // start/done/error/未知类型：内容无变化
+        streamingMessageRef.current = next as unknown as Partial<AgentMessage>;
+        dispatch({ type: "update", message: normalizeToolCalls(next as unknown as AgentMessage) });
+        setAgentPhase(null);
+        break;
+      }
       case "message_end": {
         // Same late-event guard: after reconcile finished this run,
         // loadSession already loaded this message from the session file —
         // appending it again would duplicate it.
         if (!agentRunningRef.current) break;
+        // 消息完成：丢弃增量拼接状态（完整消息由下方 append）。
+        streamingMessageRef.current = null;
         const completed = event.message as AgentMessage | undefined;
         if (completed && completed.role === "user") {
           // Delivered steering/follow-up messages surface here as user

lib/pi-types.ts

@@ -128,7 +128,7 @@ export interface AgentSessionLike {
   };
   readonly sessionManager: SessionManager;
   readonly settingsManager: SettingsManager;
-  readonly agent: { state?: { systemPrompt?: string; thinkingLevel?: string } };
+  readonly agent: { state?: { systemPrompt?: string; thinkingLevel?: string; streamingMessage?: unknown } };
   readonly extensionRunner: ExtensionRunnerLike;
   readonly promptTemplates: readonly PromptTemplateLike[];
   readonly resourceLoader: ResourceLoaderLike;

lib/rpc-manager.ts

@@ -12,7 +12,7 @@ import { getProjectTrustStatus, projectTrustReloadOptions } from "./project-trus
 import { persistExplicitStartupPreferences } from "./startup-preferences";
 import type { SlashCommandInfo } from "@earendil-works/pi-coding-agent";
 import type { AgentSessionLike, ExtensionUiContextLike, ToolInfo } from "./pi-types";
-import type { ExtensionUiRequest, ExtensionUiResponse, ExtensionWidgetItem } from "./types";
+import type { AgentMessage, ExtensionUiRequest, ExtensionUiResponse, ExtensionWidgetItem } from "./types";
 import { createHeadlessCustomUiTui, DEFAULT_CUSTOM_UI_COLUMNS } from "./custom-ui-terminal";
 
 // ============================================================================
@@ -166,6 +166,14 @@ export class AgentSessionWrapper {
     return this.inner.sessionManager.getCwd();
   }
 
+  /**
+   * 当前流式中的部分消息（若有）。仅在消息流式期间非空（SDK 在
+   * message_end/agent_end 时清空）。SSE 重连时用它重建前端拼接基座。
+   */
+  get streamingMessage(): AgentMessage | undefined {
+    return this.inner.agent.state?.streamingMessage as AgentMessage | undefined;
+  }
+
   isAlive(): boolean {
     return this._alive;
   }

lib/stream-delta.test.mjs

@@ -0,0 +1,95 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+import { createJiti } from "jiti";
+
+const jiti = createJiti(import.meta.url, { tsconfigPaths: true });
+const { applyAssistantDelta } = await jiti.import("./stream-delta.ts");
+
+test("appends text deltas to the current text block", () => {
+  let msg = applyAssistantDelta(null, { type: "text_start", contentIndex: 0 });
+  msg = applyAssistantDelta(msg, { type: "text_delta", contentIndex: 0, delta: "Hello" });
+  msg = applyAssistantDelta(msg, { type: "text_delta", contentIndex: 0, delta: " world" });
+  assert.deepEqual(msg.content, [{ type: "text", text: "Hello world" }]);
+});
+
+test("text_end replaces the block with the full content", () => {
+  let msg = applyAssistantDelta(null, { type: "text_start", contentIndex: 0 });
+  msg = applyAssistantDelta(msg, { type: "text_delta", contentIndex: 0, delta: "partial" });
+  msg = applyAssistantDelta(msg, { type: "text_end", contentIndex: 0, content: "full content" });
+  assert.deepEqual(msg.content, [{ type: "text", text: "full content" }]);
+});
+
+test("interleaved thinking and text blocks keep their contentIndex", () => {
+  let msg = applyAssistantDelta(null, { type: "thinking_start", contentIndex: 0 });
+  msg = applyAssistantDelta(msg, { type: "thinking_delta", contentIndex: 0, delta: "think" });
+  msg = applyAssistantDelta(msg, { type: "text_start", contentIndex: 1 });
+  msg = applyAssistantDelta(msg, { type: "text_delta", contentIndex: 1, delta: "answer" });
+  msg = applyAssistantDelta(msg, { type: "thinking_delta", contentIndex: 0, delta: "ing" });
+  assert.deepEqual(msg.content, [
+    { type: "thinking", thinking: "thinking" },
+    { type: "text", text: "answer" },
+  ]);
+});
+
+test("toolcall deltas accumulate JSON arguments as a string until toolcall_end", () => {
+  let msg = applyAssistantDelta(null, { type: "toolcall_start", contentIndex: 0, id: "call_1", name: "search" });
+  // 服务端从 partial 提取的 id/name 在流式期间即可用
+  assert.equal(msg.content[0].id, "call_1");
+  assert.equal(msg.content[0].name, "search");
+  msg = applyAssistantDelta(msg, { type: "toolcall_delta", contentIndex: 0, delta: '{"q' });
+  msg = applyAssistantDelta(msg, { type: "toolcall_delta", contentIndex: 0, delta: 'uery":"x"}' });
+  // 流式中 arguments 是字符串累积（normalize 后显示为空对象），
+  assert.equal(msg.content[0].type, "toolCall");
+  assert.equal(msg.content[0].arguments, '{"query":"x"}');
+  // toolcall_end 用完整对象覆盖
+  const toolCall = { type: "toolCall", id: "call_1", name: "search", arguments: { query: "x" } };
+  msg = applyAssistantDelta(msg, { type: "toolcall_end", contentIndex: 0, toolCall });
+  assert.deepEqual(msg.content[0], toolCall);
+});
+
+test("toolcall_start without injected id/name still creates a placeholder", () => {
+  const msg = applyAssistantDelta(null, { type: "toolcall_start", contentIndex: 0 });
+  assert.equal(msg.content[0].type, "toolCall");
+  assert.equal(msg.content[0].id, "");
+  assert.equal(msg.content[0].name, "");
+});
+
+test("deltas that do not change content return the same reference", () => {
+  const base = applyAssistantDelta(null, { type: "text_start", contentIndex: 0 });
+  for (const type of ["start", "done", "error", "unknown_future_type"]) {
+    const next = applyAssistantDelta(base, { type });
+    assert.equal(next, base, `${type} must not rebuild the message`);
+  }
+});
+
+test("applying a delta never mutates the previous message", () => {
+  const prev = { role: "assistant", content: [{ type: "text", text: "Hi" }] };
+  const next = applyAssistantDelta(prev, { type: "text_delta", contentIndex: 0, delta: "!" });
+  assert.equal(prev.content[0].text, "Hi"); // 原对象不变
+  assert.equal(next.content[0].text, "Hi!");
+  assert.notEqual(next, prev);
+});
+
+test("delta without a prior block creates the block at contentIndex", () => {
+  // 重连等场景下可能缺失 message_start/text_start，直接拼接应能自举。
+  const msg = applyAssistantDelta(null, { type: "text_delta", contentIndex: 2, delta: "late" });
+  assert.equal(msg.content.length, 3);
+  assert.deepEqual(msg.content[2], { type: "text", text: "late" });
+});
+
+test("keeps message metadata from the base message", () => {
+  const prev = { role: "assistant", model: "claude-x", provider: "acme" };
+  const next = applyAssistantDelta(prev, { type: "text_delta", contentIndex: 0, delta: "hi" });
+  assert.equal(next.model, "claude-x");
+  assert.equal(next.provider, "acme");
+});
+
+test("reconnect: snapshot snapshot then deltas assembles the full message", () => {
+  // 模拟页面重开后 SSE 重连：错过 message_start 和早前的 delta，服务端注入
+  // 当前部分消息快照（含关闭期间生成的全部内容），之后续收 delta 拼接。
+  const snapshot = { role: "assistant", model: "m", provider: "p", content: [{ type: "text", text: "Hello wor" }] };
+  let msg = applyAssistantDelta(snapshot, { type: "text_delta", contentIndex: 0, delta: "ld" });
+  msg = applyAssistantDelta(msg, { type: "text_delta", contentIndex: 0, delta: "!" });
+  assert.equal(msg.content[0].text, "Hello world!");
+  assert.equal(msg.model, "m");
+});