Preserve exact chat outcomes across stream closure and reattachment

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;

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

crates/utopia-server/src/live.rs

@@ -39,6 +39,7 @@ pub struct Snapshot {
     pub content: String,
     pub steps: Vec<serde_json::Value>,
     pub sources: Vec<serde_json::Value>,
+    terminal: Option<Frame>,
 }
 
 impl Snapshot {
@@ -68,6 +69,10 @@ impl Snapshot {
         }
     }
 
+    pub(crate) fn terminal(&self) -> Option<Frame> {
+        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);
     }

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<Response>,
   handlers: ChatHandlers,
+  allowIdle = false,
 ): () => void {
   const controller = new AbortController();
+  let reader: ReadableStreamDefaultReader<Uint8Array> | 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(() => {});
+  };
 }

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<Uint8Array>({ 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<Uint8Array>({
+      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<Uint8Array>({ 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));
+  });
+});

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",

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:

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<Uint8Array>;
-    const body = new ReadableStream<Uint8Array>({ start(c) { wire = c; } });
+    let finishCleanup!: () => void;
+    const cancel = vi.fn(() => new Promise<void>((resolve) => { finishCleanup = resolve; }));
+    const body = new ReadableStream<Uint8Array>({ 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();
   });