Deploy 607a293

app/pipeline/nodes/generate.py

@@ -264,9 +264,12 @@ def make_generate_node(llm_client: LLMClient, gemini_client=None) -> Callable[[P
                 complexity="simple",
             )
             full_answer = ""
-            async for token in stream:
-                full_answer += token
-                writer({"type": "token", "text": token})
+            try:
+                async for token in stream:
+                    full_answer += token
+                    writer({"type": "token", "text": token})
+            except (Exception, asyncio.CancelledError) as exc:
+                logger.warning("Enumeration stream interrupted: %s / Returning partial answer.", exc)
             return {"answer": full_answer, "sources": source_refs, "path": "enumeration"}
 
         # ── Not-found path ────────────────────────────────────────────────────────────
@@ -337,73 +340,85 @@ def make_generate_node(llm_client: LLMClient, gemini_client=None) -> Callable[[P
         token_chunk_count = 0    # number of token chunks received
         think_first_emitted = False  # CoT first-sentence status already sent
 
-        async for token in stream:
-            raw_answer += token
-            token_chunk_count += 1
-
-            if think_done or no_cot:
-                # Phase 2: emit all but the last _TAG_BOUNDARY_WINDOW chars so that
-                # a split closing tag ("</thi" in chunk N, "nk>" in chunk N+1) is
-                # always buffered until complete before routing to the token event.
-                phase2_buf += token
-                if len(phase2_buf) > _TAG_BOUNDARY_WINDOW:
-                    emit_part = phase2_buf[:-_TAG_BOUNDARY_WINDOW]
-                    phase2_buf = phase2_buf[-_TAG_BOUNDARY_WINDOW:]
-                    writer({"type": "token", "text": emit_part})
-                continue
-
-            buf += token
-
-            if not in_think:
-                if "<think>" in buf:
-                    in_think = True
-                    pre = buf[: buf.index("<think>")]
-                    if pre.strip():
-                        # Text before the think tag is part of the answer.
-                        writer({"type": "token", "text": pre})
-                    buf = buf[buf.index("<think>") + 7:]  # 7 = len("<think>")
-                elif token_chunk_count >= _THINK_LOOKAHEAD:
-                    # No CoT block in first 50 chunks — emit buffered and go direct.
-                    no_cot = True
-                    writer({"type": "token", "text": buf})
-                    buf = ""
-            else:
-                # Phase 1: inside the <think> block; buffer until </think>.
-                if "</think>" in buf:
-                    idx = buf.index("</think>")
-                    think_txt = buf[:idx].strip()
-                    after_think = buf[idx + 9:]  # 9 = len("</think>")
-
-                    if think_txt and not think_first_emitted:
-                        # Surface the first sentence as a legible status label.
-                        for j, ch in enumerate(think_txt):
-                            if ch in ".?!\n":
-                                first_sent = think_txt[: j + 1].strip()[:120]
-                                writer({"type": "status", "label": first_sent})
-                                think_first_emitted = True
-                                break
-
-                    if think_txt:
-                        writer({"type": "thinking", "text": think_txt})
-
-                    think_done = True
-                    buf = ""
-                    if after_think.strip():
-                        writer({"type": "token", "text": after_think})
-
-        # Flush Phase 2 tail buffer — strip any orphaned closing think tags that
-        # arrived in the final window before emitting to the client.
-        if phase2_buf:
-            clean_tail = _THINK_CLOSE_RE.sub("", phase2_buf)
-            if clean_tail:
-                writer({"type": "token", "text": clean_tail})
-
-        # Flush Phase 1 / lookahead buffer if stream ended mid-detection.
-        # Apply canonical strip so an incomplete opening tag doesn't surface.
-        if buf:
-            clean_buf = _strip_think_tags(buf)
-            if clean_buf:
-                writer({"type": "token", "text": clean_buf})
+        try:
+            async for token in stream:
+                raw_answer += token
+                token_chunk_count += 1
+
+                if think_done or no_cot:
+                    # Phase 2: emit all but the last _TAG_BOUNDARY_WINDOW chars so that
+                    # a split closing tag ("</thi" in chunk N, "nk>" in chunk N+1) is
+                    # always buffered until complete before routing to the token event.
+                    phase2_buf += token
+                    if len(phase2_buf) > _TAG_BOUNDARY_WINDOW:
+                        emit_part = phase2_buf[:-_TAG_BOUNDARY_WINDOW]
+                        phase2_buf = phase2_buf[-_TAG_BOUNDARY_WINDOW:]
+                        writer({"type": "token", "text": emit_part})
+                    continue
+
+                buf += token
+
+                if not in_think:
+                    if "<think>" in buf:
+                        in_think = True
+                        pre = buf[: buf.index("<think>")]
+                        if pre.strip():
+                            # Text before the think tag is part of the answer.
+                            writer({"type": "token", "text": pre})
+                        buf = buf[buf.index("<think>") + 7:]  # 7 = len("<think>")
+                    elif token_chunk_count >= _THINK_LOOKAHEAD:
+                        # No CoT block in first 50 chunks — emit buffered and go direct.
+                        no_cot = True
+                        writer({"type": "token", "text": buf})
+                        buf = ""
+                else:
+                    # Phase 1: inside the <think> block; buffer until </think>.
+                    if "</think>" in buf:
+                        idx = buf.index("</think>")
+                        think_txt = buf[:idx].strip()
+                        after_think = buf[idx + 9:]  # 9 = len("</think>")
+
+                        if think_txt and not think_first_emitted:
+                            # Surface the first sentence as a legible status label.
+                            for j, ch in enumerate(think_txt):
+                                if ch in ".?!\n":
+                                    first_sent = think_txt[: j + 1].strip()[:120]
+                                    writer({"type": "status", "label": first_sent})
+                                    think_first_emitted = True
+                                    break
+
+                        if think_txt:
+                            writer({"type": "thinking", "text": think_txt})
+
+                        think_done = True
+                        buf = ""
+                        if after_think.strip():
+                            writer({"type": "token", "text": after_think})
+
+            # Flush Phase 2 tail buffer — strip any orphaned closing think tags that
+            # arrived in the final window before emitting to the client.
+            if phase2_buf:
+                clean_tail = _THINK_CLOSE_RE.sub("", phase2_buf)
+                if clean_tail:
+                    writer({"type": "token", "text": clean_tail})
+
+            # Flush Phase 1 / lookahead buffer if stream ended mid-detection.
+            # Apply canonical strip so an incomplete opening tag doesn't surface.
+            if buf:
+                clean_buf = _strip_think_tags(buf)
+                if clean_buf:
+                    writer({"type": "token", "text": clean_buf})
+        except (Exception, asyncio.CancelledError) as exc:
+            logger.warning("Generation stream interrupted: %s / Returning partial answer.", exc)
+            # Flush whatever we have if the stream was cut
+            if buf:
+                clean_buf = _strip_think_tags(buf)
+                if clean_buf:
+                    writer({"type": "token", "text": clean_buf})
+            if phase2_buf:
+                clean_tail = _THINK_CLOSE_RE.sub("", phase2_buf)
+                if clean_tail:
+                    writer({"type": "token", "text": clean_tail})
 
         # ── Strip CoT scratchpad — single canonical pass ──────────────────────
         # _strip_think_tags is the ONLY think-tag removal in the codebase.