feat(chat): stream and persist local answer UI with final report references

2026-10-12 18:54:18 +08:00
parent a31b76a793
commit f5c3cde2c2
20 changed files with 827 additions and 105 deletions
@@ -52,6 +52,7 @@ import {
  type TodoUpdatePayload,
 } from "./chatStreamEvents.js";
 import { createTokenSmoother } from "./chatTokenSmoother.js";
+import { parseAnswerUiInput } from "../chat/answer-ui.js";

 export {
  createTokenSmoother,
@@ -79,6 +80,7 @@ type StreamPromptOptions = {
  projectId?: string;
  signal?: AbortSignal;
  suppressLegacyFrontendActions?: boolean;
+  suppressAnswerUi?: boolean;
  write: (event: string, data: Record<string, unknown>) => void;
 };

@@ -142,6 +144,7 @@ export const streamPromptResponse = async ({
  projectId,
  signal,
  suppressLegacyFrontendActions = false,
+  suppressAnswerUi = false,
  write,
 }: StreamPromptOptions): Promise<{
  aborted: boolean;
@@ -177,6 +180,8 @@ export const streamPromptResponse = async ({
  let aborted = signal?.aborted ?? false;
  let failed = false;
  let finalAnswer: string | undefined;
+  const presentedAnswerIds = new Set<string>();
+  let finalAnswerUiBlockIds: string[] = [];
  const trackBusinessStage = createBusinessStageTracker((event, data) => write(event, { ...data, session_id: clientSessionId }));
  const debugContext = {
    sessionId,
@@ -722,11 +727,36 @@ export const streamPromptResponse = async ({
          });
        }
        if (part.type === "tool") {
+          if (part.tool === "present_answer") {
+            if (suppressAnswerUi) continue;
+            if (part.state.status === "completed" && !emittedToolParts.has(part.id)) {
+              emittedToolParts.add(part.id);
+              const value = parseAnswerUiInput(part.state.input);
+              if (value) {
+                presentedAnswerIds.add(value.block_id);
+                await tokenSmoother.flush();
+                write("answer_ui", { ...value, id: value.block_id, session_id: clientSessionId });
+              } else {
+                write("progress", { id: part.id, phase: "visualization", status: "error", title: "回答界面未通过校验，保留文字说明" });
+              }
+            }
+            continue;
+          }
          if (part.tool === "final_answer") {
            if (part.state.status === "completed") {
              const output = parseToolOutput(part.state.output);
              if (output && typeof output === "object" && typeof (output as Record<string, unknown>).answer === "string") {
                finalAnswer = (output as { answer: string }).answer;
+                const refs = (output as { ui_block_ids?: unknown }).ui_block_ids;
+                finalAnswerUiBlockIds = Array.isArray(refs) ? [...new Set(refs.filter((id): id is string => typeof id === "string" && presentedAnswerIds.has(id)))].slice(0, 48) : [];
+                const hasPendingTool = [...toolStatuses.values()].some((status) => status !== "completed" && status !== "error");
+                if (presentedAnswerIds.size > 0 && !suppressAnswerUi && !hasPendingTool) {
+                  // A final rich-chat answer closes this turn; native interruption stops
+                  // automatic revisions without marking a background task complete.
+                  await runtime.abortSession(sessionId);
+                  await runtime.waitForSessionIdle(sessionId);
+                  done = true;
+                }
              }
            }
            continue;
@@ -985,7 +1015,7 @@ export const streamPromptResponse = async ({
      await emitFallbackMessage(runtime, sessionId, clientSessionId, tokenSmoother.writeToken);
    }
    await tokenSmoother.flush();
-    if (finalAnswer) write("final_answer", { session_id: clientSessionId, content: finalAnswer });
+    if (finalAnswer) write("final_answer", { session_id: clientSessionId, content: finalAnswer, ui_block_ids: finalAnswerUiBlockIds });
    emitProgress({
      id: "request-received",
      phase: "start",
@@ -1,3 +1,4 @@
+import { createAnswerStreamProjection } from "../chat/answer-stream-projection.js";
 import { credentialRefreshCoordinator } from "../product/credential-refresh.js";
 import { type Router } from "express";
 import { randomUUID } from "node:crypto";
@@ -153,13 +154,16 @@ export const toCollaborationMessages = (messages: UIMessage[]) =>
    if (message.role !== "user" && message.role !== "assistant") return [];
    const content = readUiMessageText(message);
    const reference = message.parts.filter((part) => part.type === "data-task_context");
-    return content ? [{ ...message, parts: [{ type: "text" as const, text: content }, ...reference] } as UIMessage] : [];
+    const rich = message.parts.some((part) => part.type === "data-answer_ui" || part.type === "data-final_answer");
+    const parts = rich ? message.parts.filter((part) => ["text", "data-answer_ui", "data-final_answer", "data-task_context"].includes(part.type)) : [{ type: "text" as const, text: content }, ...reference];
+    return content || rich ? [{ ...message, parts } as UIMessage] : [];
  });

 const toLegacyMessages = (messages: UIMessage[]) =>
  messages.flatMap((message) => {
    const content = readUiMessageText(message);
-    if (!content || (message.role !== "user" && message.role !== "assistant")) {
+    const rich = message.parts.some((part) => part.type === "data-answer_ui" || part.type === "data-final_answer");
+    if ((!content && !rich) || (message.role !== "user" && message.role !== "assistant")) {
      return [];
    }
    return [
@@ -167,6 +171,7 @@ const toLegacyMessages = (messages: UIMessage[]) =>
        id: message.id,
        role: message.role,
        content,
+        ...(rich ? { parts: message.parts } : {}),
      },
    ];
  });
@@ -482,12 +487,6 @@ export const registerChatStreamRoute = (
      lastRunStatuses.set(clientSessionId, "running");
      const sessionUiStateContext = toSessionUiStateContext(activeSessionRecord);
      const assistantUiMessageId = `assistant-${Date.now().toString(36)}`;
-      const textPartId = "answer";
-      const assistantParts: Record<string, unknown>[] = [];
-      let responseText = "";
-      let textStarted = false;
-      let textEnded = false;
-      let uiFinished = false;
      const syncLatestUiMessages = () => {
        latestUiMessages = [
          ...requestMessages,
@@ -499,90 +498,16 @@ export const registerChatStreamRoute = (
        ];
      };
      const writeUiChunk = (chunk: UIMessageChunk) => {
-        pipeUiChunk(res, chunk);
-      };
-      const writeTextDelta = (delta: string) => {
-        if (!delta || uiFinished) {
-          return;
-        }
-        if (!textStarted) {
-          textStarted = true;
-          assistantParts.push({
-            type: "text",
-            text: "",
-            state: "streaming",
-          });
-          writeUiChunk({ type: "text-start", id: textPartId });
-        }
-        responseText += delta;
-        const textPart = assistantParts.find(
-          (part) => part.type === "text",
-        );
-        if (textPart) {
-          textPart.text = responseText;
-          textPart.state = "streaming";
-        }
-        writeUiChunk({ type: "text-delta", id: textPartId, delta });
-        writeUiChunk({
-          type: "data-stream_token",
-          data: {
-            session_id: clientSessionId,
-            message_id: assistantUiMessageId,
-            content: delta,
-          },
-          transient: true,
-        } as UIMessageChunk);
-        syncLatestUiMessages();
-      };
-      const writeDataPart = (event: string, data: Record<string, unknown>) => {
-        if (uiFinished) {
-          return;
-        }
-        const type = `data-${event}`;
-        const id =
-          typeof data.id === "string"
-            ? data.id
-            : typeof data.request_id === "string"
-              ? data.request_id
-              : typeof data.envelope_id === "string"
-                ? data.envelope_id
-                : undefined;
-        const existingIndex =
-          id === undefined
-            ? -1
-            : assistantParts.findIndex(
-                (part) => part.type === type && part.id === id,
-              );
-        const part = id ? { type, id, data } : { type, data };
-        if (existingIndex >= 0) {
-          assistantParts[existingIndex] = part;
-        } else {
-          assistantParts.push(part);
-        }
-        writeUiChunk(part as UIMessageChunk);
-        syncLatestUiMessages();
-      };
-      const finishUiMessage = (finishReason: "stop" | "error" = "stop") => {
-        if (uiFinished) {
-          return;
-        }
-        if (!textStarted && assistantParts.length === 0) {
-          writeTextDelta("Agent 已完成处理。");
-        }
-        if (textStarted && !textEnded) {
-          textEnded = true;
-          const textPart = assistantParts.find(
-            (part) => part.type === "text",
-          );
-          if (textPart) {
-            textPart.state = "done";
-          }
-          writeUiChunk({ type: "text-end", id: textPartId });
-        }
-        syncLatestUiMessages();
-        writeUiChunk({ type: "finish", finishReason });
-        uiFinished = true;
+        if (!streamClosed && !res.writableEnded && !res.destroyed) pipeUiChunk(res, chunk);
      };
+      const projection = createAnswerStreamProjection({
+        messageId: assistantUiMessageId, sessionId: clientSessionId,
+        write: writeUiChunk, onChange: syncLatestUiMessages,
+      });
+      const assistantParts = projection.parts;
+      const writeTextDelta = projection.appendText;
+      const writeDataPart = projection.appendData;
+      const finishUiMessage = projection.finish;
      writeUiChunk({ type: "start", messageId: assistantUiMessageId });
      let persistQueue = sessionUiStateStore.write(sessionUiStateContext, {
        executionPlans: executionPlanHistory,
@@ -609,9 +534,6 @@ export const registerChatStreamRoute = (
      };
      const primarySubscriber: StreamSubscriber = {
        write: (event, data) => {
-          if (streamClosed || res.writableEnded || res.destroyed) {
-            return;
-          }
          if (event === "token") {
            writeTextDelta(typeof data.content === "string" ? data.content : "");
            return;
@@ -642,7 +564,8 @@ export const registerChatStreamRoute = (
      activeRun.subscribers.add(primarySubscriber);
      const handleClientClose = () => {
        streamClosed = true;
-        activeRun.subscribers.delete(primarySubscriber);
+        // Keep projecting and persisting parts after the browser disconnects.
+        // writeUiChunk guards the closed response; replay subscribers remain independent.
      };

      req.on("close", handleClientClose);
@@ -837,6 +760,7 @@ export const registerChatStreamRoute = (
            subscriber.write(event, subscriberData);
          }
        }
+        activeRun.messages = updateLastAssistantMessage(activeRun.messages, (message) => ({ ...message, id: assistantUiMessageId, parts: assistantParts.map((part) => ({ ...part })) }));
        void queueSessionUiStatePersist().catch((error) => {
          logger.warn({ err: error, sessionId: clientSessionId }, "failed to persist chat stream state");
        });
@@ -939,6 +863,7 @@ export const registerChatStreamRoute = (
          projectId: requestContext.projectId,
          signal: abortController.signal,
          suppressLegacyFrontendActions: parsed.data.capabilities.includes("frontend-action@1"),
+          suppressAnswerUi: sessionPurpose === "task_execution",
          write: (event, data) => {
            publish(event, data);
          },
@@ -1110,7 +1035,7 @@ export const registerChatStreamRoute = (
            progress: completeBackendProgress(message.progress),
            todos: cancelBackendTodos(message.todos),
          }));
-          if (!uiFinished && !streamClosed && !res.writableEnded && !res.destroyed) {
+          if (!projection.finished) {
            writeUiChunk({
              type: "error",
              errorText: "请求已中断",
@@ -1121,10 +1046,10 @@ export const registerChatStreamRoute = (
            logger.warn({ err: error, sessionId: clientSessionId }, "failed to persist aborted chat stream state");
          });
        }
-        if (!uiFinished && !streamClosed && !res.writableEnded && !res.destroyed) {
+        if (!projection.finished) {
          finishUiMessage(activeRun.status === "error" ? "error" : "stop");
        }
-        await persistQueue.catch((error) => {
+        await queueSessionUiStatePersist().catch((error) => {
          logger.warn({ err: error, sessionId: clientSessionId }, "failed to persist chat stream state");
        });
        sessionBridge.finalizeRequest(clientSessionId);