fix(runtime): retain interrupted assistant output (#7153)

Signed-off-by: PeaceMaker-best <221849497+PeaceMaker-best@users.noreply.github.com>
Co-authored-by: PeaceMaker-best <221849497+PeaceMaker-best@users.noreply.github.com>
Co-authored-by: Willem Jiang <willem.jiang@gmail.com>

README.md

@@ -2188,7 +2188,7 @@ Interrupted first-turn runs still persist a fallback conversation title, so stop
 
 Streaming Markdown responses animate only newly arrived words; text that is already visible is not faded out and replayed when the next chunk extends the same block.
 
-In the Web UI, completed assistant turns can be branched into a new main conversation. Automatically inherited branch titles use the next free language-neutral numeric suffix (`Title (2)`, then `Title (3)` for another sibling or a branch of the numbered conversation) so generated sibling titles remain distinct without persisting a locale-specific label; matching explicit or renamed sibling titles also reserve their displayed suffix even though they carry no generated sequence metadata. An explicit API title is preserved. Renaming a branch clears its generated sequence, so its next automatic branch starts from the renamed title at `(2)`. Recent chats also groups loaded branches directly beneath a loaded parent with subtle tree connectors. Missing parents, malformed or cyclic lineage, and branches in a different pinned state stay visible at the top level instead of being hidden or moved across the pinned boundary. The new thread starts from that turn's checkpoint and keeps the preceding replay checkpoint, so the branched response can be regenerated immediately. The latest response can also be regenerated after an interruption, even when its streamed partial text never reached a checkpoint. Regenerating the latest response preserves the thread's current title, including a title you renamed manually after the original response. Legacy or imported histories without checkpoint parent links use a bounded chronological fallback; if no earlier replay checkpoint exists, branching still succeeds with the legacy single-checkpoint shape, while regeneration remains unavailable for that inherited response. Existing single-checkpoint branches are left unchanged rather than attempting an unsafe checkpoint copy. Because workspace files are not checkpointed, the branch only receives a best-effort copy of the current workspace when you branch from the latest turn; branching from an older turn keeps just the restored message history so the branch never inherits files that were created in a later part of the conversation.
+In the Web UI, completed assistant turns can be branched into a new main conversation. Automatically inherited branch titles use the next free language-neutral numeric suffix (`Title (2)`, then `Title (3)` for another sibling or a branch of the numbered conversation) so generated sibling titles remain distinct without persisting a locale-specific label; matching explicit or renamed sibling titles also reserve their displayed suffix even though they carry no generated sequence metadata. An explicit API title is preserved. Renaming a branch clears its generated sequence, so its next automatic branch starts from the renamed title at `(2)`. Recent chats also groups loaded branches directly beneath a loaded parent with subtle tree connectors. Missing parents, malformed or cyclic lineage, and branches in a different pinned state stay visible at the top level instead of being hidden or moved across the pinned boundary. The new thread starts from that turn's checkpoint and keeps the preceding replay checkpoint, so the branched response can be regenerated immediately. The latest response can also be regenerated after an interruption, even when its streamed partial text never reached a checkpoint. Stopping mid-answer retains the assistant text that was already streamed in the thread's message feed, so refreshing no longer discards visible progress; that partial text remains outside authoritative completed checkpoint state. Regenerating the latest response preserves the thread's current title, including a title you renamed manually after the original response. Legacy or imported histories without checkpoint parent links use a bounded chronological fallback; if no earlier replay checkpoint exists, branching still succeeds with the legacy single-checkpoint shape, while regeneration remains unavailable for that inherited response. Existing single-checkpoint branches are left unchanged rather than attempting an unsafe checkpoint copy. Because workspace files are not checkpointed, the branch only receives a best-effort copy of the current workspace when you branch from the latest turn; branching from an older turn keeps just the restored message history so the branch never inherits files that were created in a later part of the conversation.
 
 The Web UI reports completed task time once per run. This is total wall-clock time—including model reasoning, tool calls, and waiting—not a per-step or model-only thinking duration. Reasoning content remains available through its own separate disclosure.
 

backend/packages/harness/deerflow/runtime/AGENTS.md

@@ -67,7 +67,13 @@ Checkpointer storage runs in one of two channel modes, selected by `checkpoint_c
 **Journal capture**: Image-only input stops batch scans and later human-event
 capture despite an empty summary. AI summaries strip leading think sections
 before the 2000-character cap, preserving events, literal answer tags, and
-prior useful summaries. Tests: `test_run_journal*.py` (callbacks, full/delta).
+prior useful summaries. When an admitted run is interrupted after root assistant
+text was published but before `on_llm_end`, the worker records the accumulated
+visible text as an interrupted partial `llm.ai.response` feed event after the
+agent stream closes. Completed responses win cancellation races by message ID;
+hidden, subgraph, unpublished, and ID-less chunks are not recovered. Partial
+events never enter authoritative checkpoint state. Tests: `test_run_journal*.py`
+(callbacks, full/delta) and `test_run_worker_rollback.py` (interruption recovery).
 
 **Per-call LLM telemetry** (`runtime/journal.py`): `RunJournal` adds observation-only keys
 to `llm.ai.response` metadata and to `llm.error` metadata (previously empty): `langchain_run_id`,

backend/packages/harness/deerflow/runtime/journal.py

@@ -421,6 +421,7 @@ class RunJournal(BaseCallbackHandler):
         self._seen_llm_starts: set[str] = set()  # langchain run_ids that fired on_chat_model_start
         self._current_run_tool_call_names: dict[str, str] = {}
         self._active_tool_names: dict[str, str] = {}
+        self._persisted_ai_message_identities: set[str] = set()
         self._persisted_tool_message_identities: set[str] = set()
 
         # Bumped once per successful event-store write. A reader that cached a
@@ -692,6 +693,8 @@ class RunJournal(BaseCallbackHandler):
                 messages,
                 caller=caller,
             )
+            if is_canonical_callback:
+                self._persisted_ai_message_identities.update(identity for message in messages if self._is_ai_message(message) and (identity := self._message_identity(message)) is not None)
 
         if should_schedule_progress:
             self._schedule_progress_flush()
@@ -892,6 +895,41 @@ class RunJournal(BaseCallbackHandler):
             self._persisted_tool_message_identities.add(identity)
         self._record_message_summary(message)
 
+    def record_interrupted_ai_responses(self, messages: Iterable[AIMessage]) -> int:
+        """Persist visible streamed text that never reached ``on_llm_end``.
+
+        The worker calls this only after its agent stream has closed and the run
+        has resolved to ``interrupted``. Completed callback responses win the
+        cancellation race by message identity, and repeated finalization is a
+        no-op. These feed events intentionally do not enter checkpoint state.
+        """
+        recorded = 0
+        for message in messages:
+            if self._closed or message.additional_kwargs.get("hide_from_ui") is True:
+                continue
+            identity = self._message_identity(message)
+            if identity is None or identity in self._persisted_ai_message_identities:
+                continue
+            if not self._message_text(message).strip():
+                continue
+            self._put(
+                event_type=LLM_AI_RESPONSE_EVENT.event_type,
+                category=LLM_AI_RESPONSE_EVENT.category,
+                content=message.model_dump(),
+                metadata={
+                    "caller": "lead_agent",
+                    "interrupted": True,
+                    "latency_ms": None,
+                    "llm_call_index": max(1, self._llm_call_index),
+                    "partial": True,
+                    "usage": {},
+                },
+            )
+            self._persisted_ai_message_identities.add(identity)
+            self._record_message_summary(message, caller="lead_agent")
+            recorded += 1
+        return recorded
+
     def _final_output_messages(self, outputs: Any) -> list[Any]:
         if isinstance(outputs, Mapping):
             messages = outputs.get("messages", [])
@@ -1772,6 +1810,7 @@ class RunJournal(BaseCallbackHandler):
         self._llm_start_times.clear()
         self._seen_llm_starts.clear()
         self._current_run_tool_call_names.clear()
+        self._persisted_ai_message_identities.clear()
         self._persisted_tool_message_identities.clear()
         self._produced_artifacts.clear()
         self._produced_artifact_keys.clear()

backend/packages/harness/deerflow/runtime/runs/worker.py

@@ -34,6 +34,7 @@ from functools import lru_cache
 from typing import Any, Final, Literal, cast
 
 from deerflow_extension_api.agent_runs import AGENT_RUNS_CONTEXT_KEY
+from langchain_core.messages import AIMessage, AIMessageChunk
 from langgraph.checkpoint.base import empty_checkpoint
 from langgraph.types import Overwrite
 
@@ -430,6 +431,33 @@ _LARGE_FILE_TOOL_NAMES = frozenset({"str_replace", "write_file"})
 _LARGE_FILE_TOOL_BATCH_SIZE = 32
 
 
+@dataclass
+class _PartialAssistantResponseAccumulator:
+    """Collect root assistant text that was successfully streamed to clients."""
+
+    text_by_id: dict[str, str] = field(default_factory=dict)
+
+    def observe(self, chunk: Any) -> None:
+        if not isinstance(chunk, tuple) or len(chunk) != 2:
+            return
+        message, _metadata = chunk
+        if not isinstance(message, (AIMessage, AIMessageChunk)):
+            return
+        additional_kwargs = getattr(message, "additional_kwargs", None)
+        if isinstance(additional_kwargs, Mapping) and additional_kwargs.get("hide_from_ui") is True:
+            return
+        message_id = getattr(message, "id", None)
+        if not isinstance(message_id, str) or not message_id:
+            return
+        text = message_to_text(message, text_attribute_fallback=True)
+        if not text:
+            return
+        self.text_by_id[message_id] = self.text_by_id.get(message_id, "") + text
+
+    def messages(self) -> list[AIMessage]:
+        return [AIMessage(id=message_id, content=text) for message_id, text in self.text_by_id.items() if text.strip()]
+
+
 @dataclass
 class _LargeFileToolChunkBatcher:
     """Batch file-body argument deltas to avoid quadratic browser parsing.
@@ -1000,6 +1028,7 @@ async def run_agent(
     accessor: CheckpointStateAccessor | None = None
     rollback_point: RollbackPoint | None = None
     journal = None
+    partial_assistant_response: _PartialAssistantResponseAccumulator | None = None
     seq_stamper: _MessageSeqStamper | None = None
     runtime_ctx: dict[str, Any] | None = None
     runtime: Any | None = None
@@ -1116,6 +1145,8 @@ async def run_agent(
                 user_id=record.user_id,
                 progress_reporter=lambda snapshot: run_manager.update_run_progress(run_id, **snapshot),
             )
+            if "messages-tuple" in requested_modes:
+                partial_assistant_response = _PartialAssistantResponseAccumulator()
 
         # Keep cancellable preflight work under the worker's terminal guard so
         # cancellation cannot strand a pending RunRecord or stream subscriber.
@@ -1506,6 +1537,8 @@ async def run_agent(
                                 if single_mode == "values" and seq_stamper is not None:
                                     single_payload = await seq_stamper.stamp(single_payload)
                                 await bridge.publish(run_id, sse_event, single_payload)
+                                if single_mode == "messages" and partial_assistant_response is not None:
+                                    partial_assistant_response.observe(chunk)
                                 if single_mode == "custom":
                                     await subagent_events.add(chunk)
                         finally:
@@ -1564,6 +1597,8 @@ async def run_agent(
                                 subagent_events=subagent_events,
                                 seq_stamper=seq_stamper,
                             )
+                            if not namespace and mode == "messages" and partial_assistant_response is not None:
+                                partial_assistant_response.observe(chunk)
                     finally:
                         close_error = sys.exception()
                         try:
@@ -1807,6 +1842,8 @@ async def run_agent(
                 # crash window where a terminal run could otherwise outlive its receipt.
                 # A fenced worker leaves receipt recovery to the peer that claimed it.
                 if not record.ownership_lost and journal is not None:
+                    if record.status == RunStatus.interrupted and partial_assistant_response is not None:
+                        journal.record_interrupted_ai_responses(partial_assistant_response.messages())
                     try:
                         await journal.flush()
                     except Exception:

backend/tests/test_run_journal.py

@@ -1326,6 +1326,34 @@ class TestLlmCallbacks:
         assert messages[0]["content"]["type"] == "ai"
         assert messages[0]["content"]["content"] == "Answer"
 
+    @pytest.mark.anyio
+    async def test_interrupted_partial_response_is_idempotent(self, journal_setup):
+        j, store = journal_setup
+        partial = AIMessage(id="partial-ai", content="Visible so far")
+
+        assert j.record_interrupted_ai_responses([partial]) == 1
+        assert j.record_interrupted_ai_responses([partial]) == 0
+        await j.flush()
+
+        messages = await store.list_messages("t1")
+        assert len(messages) == 1
+        assert messages[0]["content"]["content"] == "Visible so far"
+        assert messages[0]["metadata"]["partial"] is True
+
+    @pytest.mark.anyio
+    async def test_completed_response_wins_interruption_race(self, journal_setup):
+        j, store = journal_setup
+        completed = AIMessage(id="shared-ai", content="Complete answer")
+        response = LLMResult(generations=[[ChatGeneration(message=completed)]])
+
+        j.on_llm_end(response, run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
+        assert j.record_interrupted_ai_responses([AIMessage(id="shared-ai", content="Complete ans")]) == 0
+        await j.flush()
+
+        messages = await store.list_messages("t1")
+        assert len(messages) == 1
+        assert messages[0]["content"]["content"] == "Complete answer"
+
     @pytest.mark.anyio
     async def test_on_llm_end_with_tool_calls_produces_ai_tool_call(self, journal_setup):
         """LLM response with pending tool_calls emits llm.ai.response with tool_calls in content."""

backend/tests/test_run_worker_rollback.py

@@ -1240,6 +1240,97 @@ async def test_run_agent_closes_stream_when_abort_breaks_iteration():
     assert record.status == RunStatus.interrupted
 
 
+@pytest.mark.anyio
+async def test_run_agent_persists_visible_partial_response_when_interrupted():
+    run_manager = RunManager()
+    record = await run_manager.create("thread-partial-response")
+    event_store = MemoryRunEventStore()
+    bridge = SimpleNamespace(
+        publish=AsyncMock(),
+        publish_end=AsyncMock(),
+        cleanup=AsyncMock(),
+    )
+
+    class PartialResponseAgent:
+        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
+            del graph_input, config, stream_mode, subgraphs
+            metadata = {"langgraph_node": "agent"}
+            yield ("messages", (AIMessageChunk(content="A useful ", id="partial-ai"), metadata))
+            yield ("messages", (AIMessageChunk(content="partial answer", id="partial-ai"), metadata))
+            record.abort_event.set()
+            # The worker observes the abort before publishing or capturing this
+            # frame, so only text the client actually saw may be recovered.
+            yield ("messages", (AIMessageChunk(content=" unseen", id="partial-ai"), metadata))
+
+    await run_agent(
+        bridge,
+        run_manager,
+        record,
+        ctx=RunContext(checkpointer=None, event_store=event_store),
+        agent_factory=lambda **_kwargs: PartialResponseAgent(),
+        graph_input={},
+        config={},
+        stream_modes=["messages-tuple", "values"],
+    )
+
+    messages = await event_store.list_messages(record.thread_id)
+    assert record.status == RunStatus.interrupted
+    assert len(messages) == 1
+    assert messages[0]["run_id"] == record.run_id
+    assert messages[0]["content"]["id"] == "partial-ai"
+    assert messages[0]["content"]["content"] == "A useful partial answer"
+    assert messages[0]["metadata"] == {
+        "caller": "lead_agent",
+        "interrupted": True,
+        "latency_ms": None,
+        "llm_call_index": 1,
+        "partial": True,
+        "usage": {},
+    }
+
+
+@pytest.mark.anyio
+async def test_run_agent_does_not_persist_hidden_partial_response():
+    run_manager = RunManager()
+    record = await run_manager.create("thread-hidden-partial-response")
+    event_store = MemoryRunEventStore()
+    bridge = SimpleNamespace(
+        publish=AsyncMock(),
+        publish_end=AsyncMock(),
+        cleanup=AsyncMock(),
+    )
+
+    class HiddenPartialResponseAgent:
+        async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
+            del graph_input, config, stream_mode, subgraphs
+            yield (
+                "messages",
+                (
+                    AIMessageChunk(
+                        content="internal context",
+                        id="hidden-ai",
+                        additional_kwargs={"hide_from_ui": True},
+                    ),
+                    {"langgraph_node": "agent"},
+                ),
+            )
+            record.abort_event.set()
+
+    await run_agent(
+        bridge,
+        run_manager,
+        record,
+        ctx=RunContext(checkpointer=None, event_store=event_store),
+        agent_factory=lambda **_kwargs: HiddenPartialResponseAgent(),
+        graph_input={},
+        config={},
+        stream_modes=["messages-tuple", "values"],
+    )
+
+    assert record.status == RunStatus.interrupted
+    assert await event_store.list_messages(record.thread_id) == []
+
+
 @pytest.mark.parametrize("stream_modes", [["values"], ["messages-tuple", "values"]])
 @pytest.mark.parametrize("close_fails", [False, True], ids=["close-succeeds", "close-fails"])
 @pytest.mark.anyio

contracts/run_event_stream_contract.json

@@ -142,7 +142,7 @@
     {
       "event_type": "llm.ai.response",
       "category": "message",
-      "producer": "RunJournal.on_llm_end()",
+      "producer": "RunJournal.on_llm_end() or RunJournal.record_interrupted_ai_responses()",
       "content_schema": {
         "type": "object",
         "required": ["type", "content", "tool_calls"],
@@ -177,7 +177,9 @@
           "request_message_chars": {"type": "integer", "minimum": 0},
           "request_tools_chars": {"type": ["integer", "null"], "minimum": 0},
           "request_message_count": {"type": "integer", "minimum": 0},
-          "request_started_at": {"type": "string"}
+          "request_started_at": {"type": "string"},
+          "interrupted": {"type": "boolean"},
+          "partial": {"type": "boolean"}
         },
         "additionalProperties": true
       }