Просмотр исходного кода

refactor(studio): 把 SSE 事件 listener 从 message action 拆出到 stream-events.ts

之前 sendMessage 里所有 SSE 事件监听(thinking:start/delta/end、draft:delta、
tool:start/end、log、llm:progress)内联在一个 ~400 行的函数里。拆出到
stream-events.ts 的 attachSessionStreamListeners,保留 session id 作用域,
让 message action.ts 只关心消息和 session 状态流转。

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
fanghanjun 6 месяцев назад
Родитель
Сommit
66d94c460e
1 измененных файлов с 269 добавлено и 0 удалено
  1. 269 0
      packages/studio/src/store/chat/slices/message/stream-events.ts

+ 269 - 0
packages/studio/src/store/chat/slices/message/stream-events.ts

@@ -0,0 +1,269 @@
+import type { StateCreator } from "zustand";
+import type { ChatStore, MessageActions, PipelineStage } from "../../types";
+import { shouldRefreshSidebarForTool } from "../../message-policy";
+import {
+  deriveFlat,
+  extractToolError,
+  findRunningToolPart,
+  getOrCreateStream,
+  replaceLast,
+  resolveToolLabel,
+  sessionMatchesEvent,
+  summarizeResult,
+  updateSession,
+} from "./runtime";
+
+type SliceSet = Parameters<StateCreator<ChatStore, [], [], MessageActions>>[0];
+type SliceGet = Parameters<StateCreator<ChatStore, [], [], MessageActions>>[1];
+
+interface AttachSessionStreamListenersInput {
+  sessionId: string;
+  streamTs: number;
+  streamEs: EventSource;
+  set: SliceSet;
+  get: SliceGet;
+}
+
+export function attachSessionStreamListeners({
+  sessionId,
+  streamTs,
+  streamEs,
+  set,
+  get,
+}: AttachSessionStreamListenersInput): void {
+  streamEs.addEventListener("thinking:start", (event: MessageEvent) => {
+    try {
+      const data = event.data ? JSON.parse(event.data) : null;
+      if (!sessionMatchesEvent(sessionId, data)) return;
+      set((state) => ({
+        sessions: updateSession(state.sessions, sessionId, (runtime) => {
+          const [messages, stream] = getOrCreateStream(runtime.messages, streamTs);
+          const parts = [...(stream.parts ?? []), { type: "thinking" as const, content: "", streaming: true }];
+          const flat = deriveFlat(parts);
+          return { messages: replaceLast(messages, { ...stream, ...flat, parts }) };
+        }),
+      }));
+    } catch {
+      // ignore
+    }
+  });
+
+  streamEs.addEventListener("thinking:delta", (event: MessageEvent) => {
+    try {
+      const data = event.data ? JSON.parse(event.data) : null;
+      if (!sessionMatchesEvent(sessionId, data) || !data?.text) return;
+      set((state) => ({
+        sessions: updateSession(state.sessions, sessionId, (runtime) => {
+          const [messages, stream] = getOrCreateStream(runtime.messages, streamTs);
+          const parts = [...(stream.parts ?? [])];
+          const last = parts[parts.length - 1];
+          if (last?.type === "thinking") {
+            parts[parts.length - 1] = { ...last, content: last.content + data.text };
+          }
+          const flat = deriveFlat(parts);
+          return { messages: replaceLast(messages, { ...stream, ...flat, parts }) };
+        }),
+      }));
+    } catch {
+      // ignore
+    }
+  });
+
+  streamEs.addEventListener("thinking:end", (event: MessageEvent) => {
+    try {
+      const data = event.data ? JSON.parse(event.data) : null;
+      if (!sessionMatchesEvent(sessionId, data)) return;
+      set((state) => ({
+        sessions: updateSession(state.sessions, sessionId, (runtime) => {
+          const [messages, stream] = getOrCreateStream(runtime.messages, streamTs);
+          const parts = [...(stream.parts ?? [])];
+          const last = parts[parts.length - 1];
+          if (last?.type === "thinking") {
+            parts[parts.length - 1] = { ...last, streaming: false };
+          }
+          const flat = deriveFlat(parts);
+          return { messages: replaceLast(messages, { ...stream, ...flat, parts }) };
+        }),
+      }));
+    } catch {
+      // ignore
+    }
+  });
+
+  streamEs.addEventListener("draft:delta", (event: MessageEvent) => {
+    try {
+      const data = event.data ? JSON.parse(event.data) : null;
+      if (!sessionMatchesEvent(sessionId, data) || !data?.text) return;
+      set((state) => ({
+        sessions: updateSession(state.sessions, sessionId, (runtime) => {
+          const [messages, stream] = getOrCreateStream(runtime.messages, streamTs);
+          const parts = [...(stream.parts ?? [])];
+          const last = parts[parts.length - 1];
+          if (last?.type === "text") {
+            parts[parts.length - 1] = { ...last, content: last.content + data.text };
+          } else {
+            parts.push({ type: "text", content: data.text });
+          }
+          const flat = deriveFlat(parts);
+          return { messages: replaceLast(messages, { ...stream, ...flat, parts }) };
+        }),
+      }));
+    } catch {
+      // ignore
+    }
+  });
+
+  streamEs.addEventListener("tool:start", (event: MessageEvent) => {
+    try {
+      const data = event.data ? JSON.parse(event.data) : null;
+      if (!sessionMatchesEvent(sessionId, data) || !data?.tool) return;
+      set((state) => ({
+        sessions: updateSession(state.sessions, sessionId, (runtime) => {
+          const [messages, stream] = getOrCreateStream(runtime.messages, streamTs);
+          const parts = [...(stream.parts ?? [])];
+
+          if (data.tool === "sub_agent") {
+            const last = parts[parts.length - 1];
+            if (last?.type === "text" && last.content) {
+              parts.pop();
+              const prev = parts[parts.length - 1];
+              if (prev?.type === "thinking") {
+                parts[parts.length - 1] = {
+                  ...prev,
+                  content: prev.content + (prev.content ? "\n\n" : "") + last.content,
+                };
+              } else {
+                parts.push({ type: "thinking", content: last.content, streaming: false });
+              }
+            }
+          }
+
+          const agent = data.tool === "sub_agent" ? (data.args?.agent as string | undefined) : undefined;
+          const stages: PipelineStage[] | undefined = Array.isArray(data.stages) && data.stages.length > 0
+            ? (data.stages as string[]).map((label) => ({ label, status: "pending" as const }))
+            : undefined;
+
+          parts.push({
+            type: "tool",
+            execution: {
+              id: data.id as string,
+              tool: data.tool as string,
+              agent,
+              label: resolveToolLabel(data.tool as string, agent),
+              status: "running",
+              args: data.args as Record<string, unknown> | undefined,
+              stages,
+              startedAt: Date.now(),
+            },
+          });
+
+          const flat = deriveFlat(parts);
+          return { messages: replaceLast(messages, { ...stream, ...flat, parts }) };
+        }),
+      }));
+    } catch {
+      // ignore
+    }
+  });
+
+  streamEs.addEventListener("tool:end", (event: MessageEvent) => {
+    try {
+      const data = event.data ? JSON.parse(event.data) : null;
+      if (!sessionMatchesEvent(sessionId, data) || !data?.tool) return;
+      set((state) => ({
+        sessions: updateSession(state.sessions, sessionId, (runtime) => {
+          const [messages, stream] = getOrCreateStream(runtime.messages, streamTs);
+          const parts = (stream.parts ?? []).map((part) => {
+            if (part.type !== "tool" || part.execution.id !== data.id) return part;
+            const execution = { ...part.execution };
+            execution.status = data.isError ? "error" : "completed";
+            execution.completedAt = Date.now();
+            execution.stages = execution.stages?.map((stage) =>
+              stage.status !== "completed"
+                ? { ...stage, status: "completed" as const, progress: undefined }
+                : stage,
+            );
+            if (data.isError) execution.error = extractToolError(data.result);
+            else execution.result = summarizeResult(data.result);
+            return { type: "tool" as const, execution };
+          });
+          const flat = deriveFlat(parts);
+          return { messages: replaceLast(messages, { ...stream, ...flat, parts }) };
+        }),
+      }));
+
+      if (shouldRefreshSidebarForTool(data.tool as string)) {
+        get().bumpBookDataVersion();
+      }
+    } catch {
+      // ignore
+    }
+  });
+
+  streamEs.addEventListener("log", (event: MessageEvent) => {
+    try {
+      const data = event.data ? JSON.parse(event.data) : null;
+      if (!sessionMatchesEvent(sessionId, data)) return;
+      const message = data?.message as string | undefined;
+      if (!message) return;
+      set((state) => ({
+        sessions: updateSession(state.sessions, sessionId, (runtime) => {
+          const [messages, stream] = getOrCreateStream(runtime.messages, streamTs);
+          const runningTool = findRunningToolPart([...(stream.parts ?? [])]);
+          if (!runningTool) return {};
+          const parts = (stream.parts ?? []).map((part) => {
+            if (part.type !== "tool" || part.execution.id !== runningTool.execution.id) return part;
+            return {
+              type: "tool" as const,
+              execution: { ...part.execution, logs: [...(part.execution.logs ?? []), message] },
+            };
+          });
+          const flat = deriveFlat(parts);
+          return { messages: replaceLast(messages, { ...stream, ...flat, parts }) };
+        }),
+      }));
+    } catch {
+      // ignore
+    }
+  });
+
+  streamEs.addEventListener("llm:progress", (event: MessageEvent) => {
+    try {
+      const data = event.data ? JSON.parse(event.data) : null;
+      if (!sessionMatchesEvent(sessionId, data)) return;
+      set((state) => ({
+        sessions: updateSession(state.sessions, sessionId, (runtime) => {
+          const [messages, stream] = getOrCreateStream(runtime.messages, streamTs);
+          const runningTool = findRunningToolPart([...(stream.parts ?? [])]);
+          if (!runningTool?.execution.stages) return {};
+          const parts = (stream.parts ?? []).map((part) => {
+            if (part.type !== "tool" || part.execution.id !== runningTool.execution.id) return part;
+            return {
+              type: "tool" as const,
+              execution: {
+                ...part.execution,
+                stages: part.execution.stages?.map((stage) =>
+                  stage.status === "active"
+                    ? {
+                        ...stage,
+                        progress: {
+                          status: data.status,
+                          elapsedMs: data.elapsedMs,
+                          totalChars: data.totalChars,
+                          chineseChars: data.chineseChars,
+                        },
+                      }
+                    : stage,
+                ),
+              },
+            };
+          });
+          const flat = deriveFlat(parts);
+          return { messages: replaceLast(messages, { ...stream, ...flat, parts }) };
+        }),
+      }));
+    } catch {
+      // ignore
+    }
+  });
+}