fix: unlock send immediately after SSE close

Normal Q&A no longer waits on GET /runs before the next message. Probe run status only on 409 or dropped streams.

Co-authored-by: Cursor <cursoragent@cursor.com>

frontend-web/src/core/threads/hooks.ts

@@ -31,6 +31,7 @@ import {
 import {
   collectHumanMessageIds,
   hasServerEchoedOptimisticHuman,
+  isStaleStreamFinish,
   mergeStreamDisplayValues,
   shouldDiscardTerminalOverlay,
 } from "./stream-display";
@@ -190,11 +191,9 @@ function getRunMetadataStorage(): {
 
 /**
  * SDK ``useStream`` stores the in-flight run id at ``lg:stream:{threadId}``
- * but **deletes that key on stream onSuccess**. A premature SSE close
- * (proxy timeout / false-positive disconnect) looks like success, so the
- * reconnect key vanishes while the backend run is still occupying the
- * thread. We keep a second key the SDK never touches and only clear it
- * after ``runs.get`` confirms a terminal status.
+ * but **deletes that key on stream onSuccess**. We keep a second key the SDK
+ * never touches so a 409 / dropped-SSE exception can still rejoin. Normal
+ * SSE close clears it immediately — do not wait for ``runs.get``.
  */
 const ACTIVE_RUN_STORAGE_PREFIX = "df:active-run:";
 
@@ -276,36 +275,18 @@ export async function stopThreadRun(
 /**
  * 协助卡提交前：若上一轮还在流式（弱模型出卡后偶发未停干净），先取消服务端 run
  * 并断开本地 SSE，避免 409「已有活跃 run」或用户点了选项却发不出去。
+ *
+ * 不在发送前 GET /runs —— 协助卡出现时图多半已经结束，为兜底去查会让每次
+ * 点选项都卡住。未在流式则直接发；真遇到 409 再走 confirmOrRejoin。
  */
 export async function stopActiveThreadBeforeClarificationReply(
   threadId: string | null | undefined,
   thread: { isLoading?: boolean; stop: () => void | Promise<void> },
 ): Promise<void> {
-  let storedRunId: string | null = null;
-  if (threadId && typeof window !== "undefined") {
-    try {
-      storedRunId =
-        normalizeStoredRunId(
-          window.sessionStorage.getItem(`lg:stream:${threadId}`),
-        ) || readActiveRunId(threadId);
-    } catch {
-      storedRunId = null;
-    }
+  if (!thread.isLoading) {
+    if (threadId) clearActiveRunId(threadId);
+    return;
   }
-  // 协助卡出现时图已经 jump_to=end，run 多半已是 terminal。再 stop()/cancel
-  // 会再触发 onFinish → getState overlay，把刚发出的用户回复冻住、整段攒着出。
-  if (storedRunId && threadId) {
-    try {
-      const run = await getAPIClient().runs.get(threadId, storedRunId);
-      if (run && !isInflightRunStatus(run.status)) {
-        clearActiveRunId(threadId);
-        return;
-      }
-    } catch {
-      /* 404 / 网络：按下面的 best-effort stop 走 */
-    }
-  }
-  if (!thread.isLoading && !storedRunId) return;
   try {
     await stopThreadRun(threadId);
   } catch {
@@ -363,12 +344,9 @@ const TERMINAL_RUN_STATUSES = new Set([
   "interrupted",
 ]);
 
-// A completed server run normally sends a terminal SSE event.  In local
-// development, however, a Vite hot update, a proxy disconnect, or a brief
-// network interruption can make the browser miss that final event while the
-// backend continues (and successfully completes) the run.  Do not probe on
-// every render: only treat the stream as suspicious after it has been quiet
-// for a while, then reconcile it from the authoritative run/state endpoints.
+// Only probe /runs after an *exception* path has already set streamHoldOpen
+// (dropped SSE, 409, React #185). Normal Q&A must not GET /runs on finish or
+// every 5s — that made every send wait for a status round-trip.
 const STREAM_RECONCILIATION_IDLE_MS = 12_000;
 const STREAM_RECONCILIATION_POLL_MS = 5_000;
 // A run can legitimately be quiet while a long tool call / slow model write is
@@ -608,14 +586,9 @@ export function useThreadStream({
       "terminal" | "rejoined" | "unknown"
     >
   >(async () => "unknown");
-  // Keep the send button in "stop" until we confirm the backend run is
-  // terminal. A premature SSE close must not look like the turn ended.
-  const [streamHoldOpen, setStreamHoldOpen] = useState(
-    () =>
-      Boolean(
-        threadId && typeof window !== "undefined" && readActiveRunId(threadId),
-      ),
-  );
+  // Held only on exception recovery (dropped SSE / 409). Normal turns rely
+  // on SDK `isLoading` so the send button unlocks as soon as SSE closes.
+  const [streamHoldOpen, setStreamHoldOpen] = useState(false);
   const sendInFlightRef = useRef(false);
   const sendEpochRef = useRef(0);
   // Filled after useStream; onFinish/onError may read it asynchronously.
@@ -707,7 +680,6 @@ export function useThreadStream({
       if (meta.run_id) {
         activeRunIdRef.current = meta.run_id;
         persistActiveRunId(meta.thread_id, meta.run_id);
-        setStreamHoldOpen(true);
       }
       const metadataAgentId = context.agent_id ?? context.agent_name;
       const extraMetadata =
@@ -944,34 +916,30 @@ export function useThreadStream({
       }
       toast.error(message);
     },
-    onFinish(state) {
+    onFinish(state, runMeta) {
       lastStreamActivityAtRef.current = Date.now();
-      // Do not trust a clean SSE close as "the run is done". streamResumable
-      // chats keep the backend run alive on disconnect; the SDK still fires
-      // onFinish and would otherwise restore the send button.
-      //
-      // If a newer turn is already streaming (common: 用户协助卡提交后，上一轮
-      // interrupt 的迟到 onFinish 才到), do not freeze/rejoin — that would
-      // steal the live SSE and dump tokens only at the end.
-      if (threadStreamRef.current?.isLoading || sendInFlightRef.current) {
+      // Trust a clean SSE close. Querying /runs here blocked every send for
+      // ~5s to paper over a rare premature disconnect; that case is handled
+      // when the next send hits 409 (or onError for a dropped stream).
+      if (
+        isStaleStreamFinish({
+          finishingRunId: runMeta?.run_id,
+          currentRunId: activeRunIdRef.current,
+        })
+      ) {
         return;
       }
-      setStreamHoldOpen(true);
-      void confirmOrRejoinRef.current("finish").then((result) => {
-        if (result === "rejoined") {
-          return;
-        }
-        if (threadStreamRef.current?.isLoading || sendInFlightRef.current) {
-          return;
-        }
-        setStreamHoldOpen(false);
-        // run 真正终结（非 rejoin）：若深度研究报告流仍 active（tool_finished
-        // 帧丢失/被取消），标记中断，避免进度条永久转圈。
-        const finThreadId = threadIdRef.current ?? onStreamThreadId ?? null;
-        if (finThreadId) finalizeResearchReportStream(finThreadId);
+      const finThreadId = threadIdRef.current ?? onStreamThreadId ?? null;
+      if (finThreadId) {
+        clearActiveRunId(finThreadId);
+        finalizeResearchReportStream(finThreadId);
+      }
+      activeRunIdRef.current = null;
+      setStreamHoldOpen(false);
+      if (state?.values) {
         listeners.current.onFinish?.(state.values);
-        void queryClient.invalidateQueries({ queryKey: ["threads", "search"] });
-      });
+      }
+      void queryClient.invalidateQueries({ queryKey: ["threads", "search"] });
     },
   });
 
@@ -1009,14 +977,6 @@ export function useThreadStream({
       if (streamReconciliationInFlightRef.current && reason === "poll") {
         return "unknown";
       }
-      // Stale finish/error from a prior interrupt must not joinStream over a
-      // turn the user already started (协助卡提交后的实时流会被打断并攒批)。
-      if (
-        (reason === "finish" || reason === "error") &&
-        (sendInFlightRef.current || threadStreamRef.current?.isLoading)
-      ) {
-        return "rejoined";
-      }
       streamReconciliationInFlightRef.current = true;
       const epochAtStart = sendEpochRef.current;
       try {
@@ -1192,11 +1152,10 @@ export function useThreadStream({
       if (dropped && idleFor < 800) {
         return;
       }
-      const stored = activeRunIdRef.current ?? readActiveRunId(activeThreadId);
-      if (!stored && !streamHoldOpen && threadStreamRef.current?.isLoading) {
-        return;
-      }
-      if (!stored && !streamHoldOpen) {
+      // Do not probe just because a run id is stored. Only the exception
+      // path sets streamHoldOpen; polling otherwise made every turn wait
+      // on GET /runs before the user could send again.
+      if (!streamHoldOpen) {
         return;
       }
       await confirmOrRejoin("poll");

frontend-web/src/core/threads/stream-display.test.ts

@@ -5,6 +5,7 @@ import type { Message } from "@langchain/langgraph-sdk";
 import {
   collectHumanMessageIds,
   hasServerEchoedOptimisticHuman,
+  isStaleStreamFinish,
   mergeStreamDisplayValues,
   shouldDiscardTerminalOverlay,
 } from "./stream-display";
@@ -28,6 +29,30 @@ test("live streaming ignores frozen checkpoint messages so tokens can paint", ()
   assert.deepEqual(afterStream?.messages, overrides.messages);
 });
 
+test("late onFinish of a prior run is stale once a newer run is live", () => {
+  assert.equal(
+    isStaleStreamFinish({
+      finishingRunId: "old-run",
+      currentRunId: "new-run",
+    }),
+    true,
+  );
+  assert.equal(
+    isStaleStreamFinish({
+      finishingRunId: "run-1",
+      currentRunId: "run-1",
+    }),
+    false,
+  );
+  assert.equal(
+    isStaleStreamFinish({
+      finishingRunId: "run-1",
+      currentRunId: null,
+    }),
+    false,
+  );
+});
+
 test("a later send discards the previous run's getState overlay", () => {
   assert.equal(
     shouldDiscardTerminalOverlay({

frontend-web/src/core/threads/stream-display.ts

@@ -51,6 +51,25 @@ export function shouldDiscardTerminalOverlay(opts: {
   return false;
 }
 
+/**
+ * True when an SSE onFinish belongs to a *previous* run, after the user
+ * already started a newer turn (协助卡提交是典型场景).
+ *
+ * Do **not** use `sendInFlight` / `isLoading` for this: `submit()` stays
+ * pending until the current stream ends, so those flags are still true on a
+ * normal same-turn finish and would skip unlock for every Q&A.
+ */
+export function isStaleStreamFinish(opts: {
+  finishingRunId?: string | null;
+  currentRunId?: string | null;
+}): boolean {
+  return Boolean(
+    opts.finishingRunId &&
+      opts.currentRunId &&
+      opts.finishingRunId !== opts.currentRunId,
+  );
+}
+
 export function collectHumanMessageIds(
   ...groups: Array<Message[] | undefined>
 ): Set<string> {