diff --git a/crates/utopia-server/src/api/chat.rs b/crates/utopia-server/src/api/chat.rs index 1b08111aa..4bfdb26f8 100644 --- a/crates/utopia-server/src/api/chat.rs +++ b/crates/utopia-server/src/api/chat.rs @@ -976,6 +976,10 @@ fn sse_from( return; }; yield to_event(&snapshot.to_frame()); + if let Some(terminal) = snapshot.terminal() { + yield to_event(&terminal); + return; + } loop { match rx.recv().await { Ok(frame) => { @@ -1338,3 +1342,7 @@ mod chat_empty_reply_tests; #[cfg(test)] #[path = "chat_terminal_tests.rs"] mod chat_terminal_tests; + +#[cfg(test)] +#[path = "chat_stream_tests.rs"] +mod stream_tests; diff --git a/crates/utopia-server/src/api/chat_stream_tests.rs b/crates/utopia-server/src/api/chat_stream_tests.rs new file mode 100644 index 000000000..6feb862b0 --- /dev/null +++ b/crates/utopia-server/src/api/chat_stream_tests.rs @@ -0,0 +1,73 @@ +use super::*; +use axum::response::IntoResponse; +use std::sync::Arc; + +#[tokio::test] +async fn reattachment_preserves_terminal_frames_on_both_sides_of_emit() { + for event in ["done", "error"] { + let registry = Arc::new(crate::live::Registry::default()); + let id = Uuid::now_v7(); + let handle = registry.begin(id).await; + handle.emit(delta_event("partial")).await; + let before = sse_from(registry.attach(id).await); + handle.emit(Frame::new(event, "safe outcome".into())).await; + // Finish closes the late receiver on the old implementation, so this + // counterexample completes without timing out even when terminal is lost. + let after = sse_from(registry.attach(id).await); + handle.finish().await; + for response in [before, after] { + let body = axum::body::to_bytes(response.into_response().into_body(), 65536) + .await + .unwrap(); + let text = String::from_utf8_lossy(&body); + assert_eq!( + text.matches(&format!("event: {event}")).count(), + 1, + "{text}" + ); + assert!(text.contains("partial") && text.contains("safe outcome")); + } + let idle = axum::body::to_bytes( + sse_from(registry.attach(id).await) + .into_response() + .into_body(), + 65536, + ) + .await + .unwrap(); + assert!(String::from_utf8_lossy(&idle).contains("event: idle")); + } +} + +#[tokio::test] +async fn lagged_subscribers_receive_an_error_not_done() { + let registry = Arc::new(crate::live::Registry::default()); + let id = Uuid::now_v7(); + let handle = registry.begin(id).await; + let stream = sse_from(registry.attach(id).await); + for _ in 0..300 { + handle.emit(delta_event("x")).await; + } + handle.finish().await; + let body = axum::body::to_bytes(stream.into_response().into_body(), 65536) + .await + .unwrap(); + let text = String::from_utf8_lossy(&body); + assert!(text.contains("event: error") && !text.contains("event: done")); +} + +#[tokio::test] +async fn first_terminal_freezes_the_snapshot_and_broadcast() { + let registry = Arc::new(crate::live::Registry::default()); + let id = Uuid::now_v7(); + let handle = registry.begin(id).await; + handle.emit(delta_event("kept")).await; + handle.emit(error_event("original error")).await; + handle.emit(delta_event("discarded")).await; + handle.emit(done_event()).await; + let (snapshot, _) = registry.attach(id).await.unwrap(); + assert_eq!(snapshot.content, "kept"); + assert_eq!(snapshot.terminal().unwrap().event, "error"); + assert!(!snapshot.to_frame().data.contains("terminal")); + handle.finish().await; +} diff --git a/crates/utopia-server/src/live.rs b/crates/utopia-server/src/live.rs index 1333942cd..3b69aadd6 100644 --- a/crates/utopia-server/src/live.rs +++ b/crates/utopia-server/src/live.rs @@ -39,6 +39,7 @@ pub struct Snapshot { pub content: String, pub steps: Vec, pub sources: Vec, + terminal: Option, } impl Snapshot { @@ -68,6 +69,10 @@ impl Snapshot { } } + pub(crate) fn terminal(&self) -> Option { + self.terminal.clone() + } + pub fn to_frame(&self) -> Frame { Frame::new( "snapshot", @@ -107,7 +112,15 @@ impl Handle { /// 于是接上的时刻要么整个在这次 emit 之前,要么整个在它之后 pub async fn emit(&self, frame: Frame) { let mut snap = self.snap.write().await; + // The snapshot and the subscription boundary must include the terminal: + // a subscriber arriving after this broadcast still needs the same outcome. + if snap.terminal.is_some() { + return; + } snap.apply(&frame); + if matches!(frame.event, "done" | "error") { + snap.terminal = Some(frame.clone()); + } // 没有订阅者是常态(人走了),不是错 let _ = self.tx.send(frame); } diff --git a/web/src/api.ts b/web/src/api.ts index d1d5c92c5..1f139e9cf 100644 --- a/web/src/api.ts +++ b/web/src/api.ts @@ -2562,6 +2562,7 @@ export function reattachChat( signal, }), handlers, + true, ); } @@ -2587,57 +2588,95 @@ export function streamChat( function consumeChatStream( open: (signal: AbortSignal) => Promise, handlers: ChatHandlers, + allowIdle = false, ): () => void { const controller = new AbortController(); + let reader: ReadableStreamDefaultReader | undefined; + let terminal = false; + const fail = (message: string) => { + if (terminal || controller.signal.aborted) return; + terminal = true; + handlers.onError(message); + }; (async () => { try { const res = await open(controller.signal); + if (controller.signal.aborted) { + await res.body?.cancel(); + return; + } if (!res.ok || !res.body) { let message = res.statusText; try { const body = (await res.json()) as { error?: string }; if (body.error) message = body.error; - } catch { - /* ignore */ - } - handlers.onError(message); + } catch { /* keep the HTTP status */ } + fail(message); return; } - const reader = res.body.getReader(); + reader = res.body.getReader(); const decoder = new TextDecoder(); - let buf = ""; - for (;;) { + let line = ""; + let skipLf = false; + let event = "message"; + let data: string[] = []; + const dispatch = () => { + if (data.length === 0) return; + const value = data.join("\n"); + if (event === "done") { terminal = true; handlers.onDone(); } + else if (event === "error") fail(value); + else if (event === "idle") { + if (allowIdle) { terminal = true; handlers.onIdle?.(); } + else fail(S.ask.streamInterrupted); + } else if (event === "conversation") handlers.onConversation(JSON.parse(value).id); + else if (event === "sources") handlers.onSources(JSON.parse(value)); + else if (event === "step") handlers.onStep(JSON.parse(value)); + else if (event === "delta") handlers.onDelta(JSON.parse(value).text); + else if (event === "snapshot") handlers.onSnapshot?.(JSON.parse(value)); + }; + const finishLine = () => { + if (line === "") { + dispatch(); + event = "message"; + data = []; + } else { + const colon = line.indexOf(":"); + const field = colon < 0 ? line : line.slice(0, colon); + let value = colon < 0 ? "" : line.slice(colon + 1); + if (value.startsWith(" ")) value = value.slice(1); + if (field === "event") event = value; + else if (field === "data") data.push(value); + } + line = ""; + }; + while (!terminal && !controller.signal.aborted) { const { done, value } = await reader.read(); - if (done) break; - buf += decoder.decode(value, { stream: true }); - let idx: number; - while ((idx = buf.indexOf("\n\n")) >= 0) { - const frame = buf.slice(0, idx); - buf = buf.slice(idx + 2); - let event = "message"; - let data = ""; - for (const line of frame.split("\n")) { - if (line.startsWith("event:")) event = line.slice(6).trim(); - else if (line.startsWith("data:")) data += line.slice(5).trim(); - } - if (event === "conversation") - handlers.onConversation((JSON.parse(data) as { id: string }).id); - else if (event === "sources") - handlers.onSources(JSON.parse(data || "[]")); - else if (event === "step") - handlers.onStep(JSON.parse(data) as ChatStep); - else if (event === "delta") - handlers.onDelta((JSON.parse(data) as { text: string }).text); - else if (event === "snapshot") handlers.onSnapshot?.(JSON.parse(data)); - else if (event === "idle") handlers.onIdle?.(); - else if (event === "done") handlers.onDone(); - else if (event === "error") handlers.onError(data); + if (done || controller.signal.aborted) break; + // CR is a complete line ending, even when its optional LF arrives in the + // next byte chunk. TextDecoder independently preserves split UTF-8. + for (const char of decoder.decode(value, { stream: true })) { + if (skipLf && char === "\n") { skipLf = false; continue; } + skipLf = false; + if (char === "\r" || char === "\n") { + finishLine(); + skipLf = char === "\r"; + } else line += char; + if (terminal || controller.signal.aborted) break; } } - handlers.onDone(); + // EOF never dispatches an incomplete frame and is not an application done. + fail(S.ask.streamInterrupted); } catch (e) { - if (!controller.signal.aborted) handlers.onError(String(e)); + fail(e instanceof SyntaxError ? S.ask.streamInterrupted : String(e)); + } finally { + if (reader) { + try { await reader.cancel(); } catch { /* terminal/abort already decided */ } + reader.releaseLock(); + } } })(); - return () => controller.abort(); + return () => { + controller.abort(); + void reader?.cancel().catch(() => {}); + }; } diff --git a/web/src/chatStream.test.ts b/web/src/chatStream.test.ts new file mode 100644 index 000000000..94c8acd03 --- /dev/null +++ b/web/src/chatStream.test.ts @@ -0,0 +1,76 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { reattachChat, streamChat, type ChatHandlers } from "./api"; +vi.mock("./i18n", () => ({ S: { ask: { streamInterrupted: "Stream interrupted" } }, lang: "en" })); +afterEach(() => vi.unstubAllGlobals()); + +const encoder = new TextEncoder(); +async function replay(text: string, attach = false, bytewise = false) { + const events: [string, unknown][] = []; + const h: ChatHandlers = { + onConversation: (v) => events.push(["conversation", v]), onSources: (v) => events.push(["sources", v]), + onStep: (v) => events.push(["step", v]), onDelta: (v) => events.push(["delta", v]), + onSnapshot: (v) => events.push(["snapshot", v]), onDone: () => events.push(["done", null]), + onError: (v) => events.push(["error", v]), onIdle: () => events.push(["idle", null]), + }; + const bytes = encoder.encode(text); + const body = new ReadableStream({ start(c) { + if (bytewise) for (const byte of bytes) c.enqueue(new Uint8Array([byte])); + else c.enqueue(bytes); + c.close(); + } }); + const fetch = vi.fn().mockResolvedValue(new Response(body)); + vi.stubGlobal("fetch", fetch); + if (attach) reattachChat("kb", "c", h); else streamChat("kb", {message:"hello"}, h); + await vi.waitFor(() => expect(events.some(([e]) => ["done","error","idle"].includes(e))).toBe(true)); + expect(fetch).toHaveBeenCalledTimes(1); + return events; +} + +describe("application chat terminal outcomes", () => { + it.each(["\n", "\r\n", "\r"])("reads %j line endings once, including split UTF-8", async (nl) => { + const text = 'event: delta\ndata: {"text":"中文🙂"}\n\nevent: done\ndata: {}\n\n'.replaceAll("\n", nl); + expect(await replay(text, false, true)).toEqual([["delta","中文🙂"],["done",null]]); + }); + it("preserves multiline error data and ignores everything after it", async () => { + expect(await replay("event: error\ndata: first\ndata: second\n\nevent: done\ndata: {}\n\nevent: delta\ndata: not-json\n\n")) + .toEqual([["error","first\n second"]]); + }); + it("does not turn partial output plus EOF into success", async () => { + expect(await replay('event: delta\ndata: {"text":"partial"}\n\n')).toEqual([["delta","partial"],["error","Stream interrupted"]]); + }); + it("idle only terminates a reattachment", async () => { + expect(await replay("event: idle\ndata: {}\n\n", true)).toEqual([["idle",null]]); + expect(await replay("event: idle\ndata: {}\n\n")).toEqual([["error","Stream interrupted"]]); + }); + it("ignores frames after the first done", async () => { + expect(await replay('event: done\ndata: {}\n\nevent: delta\ndata: {"text":"late"}\n\nevent: error\ndata: late error\n\n')).toEqual([["done",null]]); + }); + it.each(["event: done\ndata: {}\n", "event: done\ndata: {}", "", "event: delta\ndata: {broken}\n\n"])("requires a complete terminal frame: %s", async (s) => { + const result = await replay(s); + expect(result).toHaveLength(1); expect(result[0][0]).toBe("error"); + }); + it("cancels an open stream after done even if cancellation rejects", async () => { + const cancel = vi.fn(() => Promise.reject(new Error("cancel failed"))); + const body = new ReadableStream({ + start(c) { c.enqueue(encoder.encode("event: done\ndata: {}\n\n")); }, cancel, + }); + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(new Response(body))); + const h = {onConversation:vi.fn(),onSources:vi.fn(),onStep:vi.fn(),onDelta:vi.fn(),onDone:vi.fn(),onError:vi.fn()}; + streamChat("kb", {message:"hello"}, h); + await vi.waitFor(() => expect(body.locked).toBe(false)); + await vi.waitFor(() => expect(cancel).toHaveBeenCalledTimes(1)); + expect(h.onDone).toHaveBeenCalledTimes(1); expect(h.onError).not.toHaveBeenCalled(); + }); + it("active abort is silent and cancels the reader", async () => { + const cancelled = vi.fn(); + const body = new ReadableStream({ cancel: cancelled }); + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(new Response(body))); + const h = {onConversation:vi.fn(),onSources:vi.fn(),onStep:vi.fn(),onDelta:vi.fn(),onDone:vi.fn(),onError:vi.fn()}; + const abort = streamChat("kb",{message:"hello"},h); + await vi.waitFor(() => expect(body.locked).toBe(true)); + abort(); + await vi.waitFor(() => expect(cancelled).toHaveBeenCalledTimes(1)); + expect(h.onDone).not.toHaveBeenCalled(); expect(h.onError).not.toHaveBeenCalled(); + await vi.waitFor(() => expect(body.locked).toBe(false)); + }); +}); diff --git a/web/src/i18n/en.ts b/web/src/i18n/en.ts index a01a07029..669e847f0 100644 --- a/web/src/i18n/en.ts +++ b/web/src/i18n/en.ts @@ -735,6 +735,7 @@ export const en = { chunkOf: (filename: string, seq: number) => `${filename} · section ${seq}`, }, ask: { + streamInterrupted: "The answer stream was interrupted. Reopen the conversation to check its status.", /* 新对话首屏问候:碑铭衬线,品牌名入句(标题不带句号) */ greeting: "Ask Utopia what it remembers", emptyTitle: "Chat", diff --git a/web/src/i18n/zh.ts b/web/src/i18n/zh.ts index 0bf28c2fa..eb4560dcf 100644 --- a/web/src/i18n/zh.ts +++ b/web/src/i18n/zh.ts @@ -676,6 +676,7 @@ export const zh: Strings = { chunkOf: (filename: string, seq: number) => `${filename} · 第 ${seq} 段`, }, ask: { + streamInterrupted: "回答连接已中断,请重新打开会话查看状态。", greeting: "问问 Utopia 都记住了什么", emptyTitle: "对话", emptyBody: diff --git a/web/src/liveAnswer.stream.test.ts b/web/src/liveAnswer.stream.test.ts index 180e12ca9..9f733af24 100644 --- a/web/src/liveAnswer.stream.test.ts +++ b/web/src/liveAnswer.stream.test.ts @@ -12,9 +12,11 @@ afterEach(() => { }); describe("live answer generation ownership", () => { - it("keeps a follow-up streaming when the previous SSE connection closes", async () => { + it("keeps a follow-up streaming when the previous SSE cleanup finishes", async () => { let wire!: ReadableStreamDefaultController; - const body = new ReadableStream({ start(c) { wire = c; } }); + let finishCleanup!: () => void; + const cancel = vi.fn(() => new Promise((resolve) => { finishCleanup = resolve; })); + const body = new ReadableStream({ start(c) { wire = c; }, cancel }); vi.stubGlobal("fetch", vi.fn().mockResolvedValue(new Response(body))); const previous = liveAnswer.begin("kb", "conversation", turns(), () => {}); const done = vi.fn(() => previous.finish()); @@ -29,8 +31,10 @@ describe("live answer generation ownership", () => { wire.enqueue(new TextEncoder().encode('event: done\ndata: {}\n\n')); await vi.waitFor(() => expect(done).toHaveBeenCalledTimes(1)); const followUp = liveAnswer.begin("kb", "conversation", turns(), () => {}); - wire.close(); - await vi.waitFor(() => expect(done).toHaveBeenCalledTimes(2)); + await vi.waitFor(() => expect(cancel).toHaveBeenCalledTimes(1)); + finishCleanup(); + await vi.waitFor(() => expect(body.locked).toBe(false)); + expect(done).toHaveBeenCalledTimes(1); expect(liveAnswer.entry("kb", "conversation")?.streaming).toBe(true); followUp.finish(); });