diff --git a/app/api/agent/[id]/events/route.ts b/app/api/agent/[id]/events/route.ts index 1a6f95710..5a4ab124d 100644 --- a/app/api/agent/[id]/events/route.ts +++ b/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 = { + 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) | undefined; + const fields = delta && typeof delta.type === "string" ? DELTA_FIELDS[delta.type] : undefined; + if (delta && fields) { + // 增量事件：只转发轻量 delta 字段，浏览器端据此拼接流式消息。 + const slimDelta: Record = { 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); diff --git a/app/api/agent/events-route.test.mjs b/app/api/agent/events-route.test.mjs index adc6b875b..1d3f6f674 100644 --- a/app/api/agent/events-route.test.mjs +++ b/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/); + 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); diff --git a/hooks/useAgentSession.ts b/hooks/useAgentSession.ts index a482f7120..24470c16c 100644 --- a/hooks/useAgentSession.ts +++ b/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 | null>(null); const promptRunIdRef = useRef(0); const optimisticUserMessageKeyRef = useRef(null); + // 增量拼接中的流式消息（SDK 结构）。message_start 重置、message_delta 追加、 + // message_update（服务端降级快照）覆盖校准、message_end 清空。 + const streamingMessageRef = useRef | 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; + 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 diff --git a/lib/pi-types.ts b/lib/pi-types.ts index 1af97949c..a484ac2a5 100644 --- a/lib/pi-types.ts +++ b/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; diff --git a/lib/rpc-manager.ts b/lib/rpc-manager.ts index cfb611002..74d3d602b 100644 --- a/lib/rpc-manager.ts +++ b/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; } diff --git a/lib/stream-delta.test.mjs b/lib/stream-delta.test.mjs new file mode 100644 index 000000000..c69c149a2 --- /dev/null +++ b/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"); +}); diff --git a/lib/stream-delta.ts b/lib/stream-delta.ts new file mode 100644 index 000000000..84e25957a --- /dev/null +++ b/lib/stream-delta.ts @@ -0,0 +1,93 @@ +// 流式消息增量拼接的纯函数模块。 +// +// 服务端把 SDK 的 `message_update` 事件投影为轻量 `message_delta`（剥离 +// `partial`——每个 delta 事件都携带的完整累积消息），浏览器端用本模块把 +// delta 追加到本地流式消息上。这样传输量 = 内容本身（O(n)），而不是每个 +// 更新块都重发此前生成的全部内容（O(n²)）。 + +/** 服务端转发的轻量 delta（`assistantMessageEvent` 剥离 `partial` 后）。结构 + * 对齐 `@earendil-works/pi-ai` 的 `AssistantMessageEvent`。 */ +export interface AssistantStreamDelta { + type: string; + contentIndex?: number; + delta?: string; + content?: string; + /** toolcall_start 时服务端从 partial 提取注入（SDK 事件本身不带 id/name）。 */ + id?: string; + name?: string; + toolCall?: Record; +} + +/** 流式拼接过程中的 content block。拼接发生在 SDK 结构上（toolCall 扁平： + * `id`/`name`/`arguments`），最终经 `normalizeToolCalls` 转换为渲染格式。 */ +interface StreamContentBlock { + type: string; + text?: string; + thinking?: string; + id?: string; + name?: string; + arguments?: unknown; +} + +/** 流式消息（宽松结构：不绑定具体消息类型）。 */ +export interface StreamMessage { + role?: string; + content?: unknown; + [key: string]: unknown; +} + +/** + * 把一条 `assistantMessageEvent` delta 应用到当前流式消息上（追加式拼接）。 + * 返回新消息对象；若该 delta 不改变消息内容（start/done/error/未知类型） + * 则返回原对象（引用相等，调用方可跳过渲染）。 + */ +export function applyAssistantDelta(prev: StreamMessage | null, delta: AssistantStreamDelta): StreamMessage { + const msg: StreamMessage = prev ?? { role: "assistant", content: [] }; + const content = Array.isArray(msg.content) ? [...msg.content] : []; + const idx = delta.contentIndex ?? 0; + switch (delta.type) { + case "text_start": + content[idx] = { type: "text", text: "" }; + break; + case "text_delta": { + const block = content[idx] as StreamContentBlock | undefined; + const text = block?.type === "text" && typeof block.text === "string" ? block.text : ""; + content[idx] = { type: "text", text: text + (delta.delta ?? "") }; + break; + } + case "text_end": + content[idx] = { type: "text", text: delta.content ?? "" }; + break; + case "thinking_start": + content[idx] = { type: "thinking", thinking: "" }; + break; + case "thinking_delta": { + const block = content[idx] as StreamContentBlock | undefined; + const thinking = block?.type === "thinking" && typeof block.thinking === "string" ? block.thinking : ""; + content[idx] = { type: "thinking", thinking: thinking + (delta.delta ?? "") }; + break; + } + case "thinking_end": + content[idx] = { type: "thinking", thinking: delta.content ?? "" }; + break; + case "toolcall_start": + content[idx] = { type: "toolCall", id: delta.id ?? "", name: delta.name ?? "", arguments: "" }; + break; + case "toolcall_delta": { + const block = content[idx] as StreamContentBlock | undefined; + const id = block?.type === "toolCall" && typeof block.id === "string" ? block.id : ""; + const name = block?.type === "toolCall" && typeof block.name === "string" ? block.name : ""; + const args = block?.type === "toolCall" && typeof block.arguments === "string" ? block.arguments : ""; + content[idx] = { type: "toolCall", id, name, arguments: args + (delta.delta ?? "") }; + break; + } + case "toolcall_end": + // toolCall 自带 type: "toolCall"，可直接作为 content block；arguments 为完整对象。 + content[idx] = (delta.toolCall as StreamContentBlock | undefined) ?? { type: "toolCall", id: "", name: "", arguments: {} }; + break; + default: + // start / done / error / 未知类型：内容无变化，保留当前状态。 + return msg; + } + return { ...msg, content }; +}