|
@@ -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
|
|
|
|
|
+ }
|
|
|
|
|
+ });
|
|
|
|
|
+}
|