fix(streaming): recover interrupted reasoning

This commit is contained in:
Frank Song 2026-07-17 17:57:54 +08:00
parent 920cdd56c0
commit b741e17275
3 changed files with 643 additions and 25 deletions

View file

@ -2306,11 +2306,23 @@ def _collapse_adjacent_duplicate_partials(messages) -> tuple[list, bool]:
return collapsed, changed
def _find_existing_assistant_for_journal_content(session, content: str) -> int | None:
def _find_existing_assistant_for_journal_content(
session,
content: str,
*,
max_index: int | None = None,
excluded_indexes: set[int] | None = None,
) -> int | None:
candidate = _normalize_journal_recovery_text(content)
if not candidate:
return None
for idx, message in enumerate(session.messages or []):
messages = session.messages or []
stop = len(messages) if max_index is None else min(len(messages), max_index)
substring_match = None
for idx in range(stop):
if excluded_indexes and idx in excluded_indexes:
continue
message = messages[idx]
if not isinstance(message, dict) or message.get('role') != 'assistant':
continue
if message.get('_error'):
@ -2320,9 +2332,9 @@ def _find_existing_assistant_for_journal_content(session, content: str) -> int |
continue
if existing == candidate:
return idx
if len(candidate) >= 24 and candidate in existing:
return idx
return None
if substring_match is None and len(candidate) >= 24 and candidate in existing:
substring_match = idx
return substring_match
def _journal_tool_already_present(
@ -2394,6 +2406,12 @@ def _run_journal_has_visible_output(session, stream_id: str | None) -> bool:
continue
if str(payload.get('text') or '').strip():
return True
if event_name == 'reasoning':
reasoning_text = str(
payload.get('text') or payload.get('reasoning') or payload.get('thinking') or ''
)
if reasoning_text.strip():
return True
if event_name == 'tool':
return True
return False
@ -2454,9 +2472,10 @@ def _append_journaled_partial_output(
"""Recover already-emitted visible output from a dead stream journal.
This repair path is intentionally conservative: it restores user-visible
assistant text and tool-card metadata that had already been emitted over
SSE before the WebUI process died. It does not restore hidden reasoning and
it does not try to continue execution.
assistant text, display-only reasoning, and tool-card metadata that had
already been emitted over SSE before the WebUI process died. Restored
reasoning stays out of ``context_messages`` so it cannot become provider-
facing history. The repair does not try to continue execution.
"""
if not stream_id:
return False
@ -2479,24 +2498,106 @@ def _append_journaled_partial_output(
appended_any = False
assistant_parts: list[str] = []
reasoning_parts: list[str] = []
assistant_started_at: float | None = None
current_assistant_idx: int | None = None
recovered_tool_calls: list[dict] = []
initial_message_count = len(session.messages or [])
claimed_existing_assistant_indexes: set[int] = set()
def content_match_can_receive_reasoning(existing_idx: int) -> bool:
messages = session.messages or []
owner_idx = None
for candidate_idx in range(existing_idx - 1, -1, -1):
candidate = messages[candidate_idx]
if isinstance(candidate, dict) and candidate.get('role') == 'user':
owner_idx = candidate_idx
break
if owner_idx is None:
return False
pending_text = _normalize_journal_recovery_text(session.pending_user_message)
owner_text = _normalize_journal_recovery_text(messages[owner_idx].get('content'))
if pending_text and owner_text != pending_text:
return False
for candidate_idx in range(existing_idx + 1, initial_message_count):
candidate = messages[candidate_idx]
if not isinstance(candidate, dict) or candidate.get('role') != 'user':
continue
candidate_text = _normalize_journal_recovery_text(candidate.get('content'))
if pending_text and candidate.get('_recovered') and candidate_text == pending_text:
continue
return False
return True
def append_context_projection(message: dict) -> None:
context_projection = dict(message)
context_projection.pop('reasoning', None)
_append_recovered_turn_to_context(session, context_projection)
def attach_display_reasoning(message: dict, reasoning: str) -> bool:
if not reasoning:
return False
existing = str(message.get('reasoning') or '').strip()
if existing:
return False
message['reasoning'] = reasoning
return True
def flush_assistant() -> int | None:
nonlocal appended_any, assistant_parts, assistant_started_at, current_assistant_idx
nonlocal appended_any, assistant_parts, reasoning_parts
nonlocal assistant_started_at, current_assistant_idx
content = ''.join(assistant_parts).strip()
reasoning = ''.join(reasoning_parts).strip()
assistant_parts = []
if not content:
reasoning_parts = []
if not content and not reasoning:
return current_assistant_idx
if dedupe_existing:
existing_idx = _find_existing_assistant_for_journal_content(session, content)
if dedupe_existing and content:
search_excluded = set(claimed_existing_assistant_indexes)
existing_idx = None
while True:
candidate_idx = _find_existing_assistant_for_journal_content(
session,
content,
max_index=initial_message_count,
excluded_indexes=search_excluded,
)
if candidate_idx is None:
break
if not reasoning or content_match_can_receive_reasoning(candidate_idx):
existing_idx = candidate_idx
break
search_excluded.add(candidate_idx)
if existing_idx is not None:
claimed_existing_assistant_indexes.add(existing_idx)
current_assistant_idx = existing_idx
assistant_started_at = None
if 0 <= existing_idx < len(session.messages):
_append_recovered_turn_to_context(session, session.messages[existing_idx])
existing_message = session.messages[existing_idx]
append_context_projection(existing_message)
if attach_display_reasoning(existing_message, reasoning):
appended_any = True
return existing_idx
if dedupe_existing and reasoning and not content:
for existing_idx in range(initial_message_count):
if existing_idx in claimed_existing_assistant_indexes:
continue
existing_message = session.messages[existing_idx]
if not isinstance(existing_message, dict):
continue
if (
existing_message.get('_recovered_from_run_journal')
and existing_message.get('_recovered_stream_id') == stream_id
and existing_message.get('role') == 'assistant'
and not str(existing_message.get('content') or '').strip()
and str(existing_message.get('reasoning') or '').strip() == reasoning
):
claimed_existing_assistant_indexes.add(existing_idx)
current_assistant_idx = existing_idx
assistant_started_at = None
return existing_idx
timestamp = int(assistant_started_at or time.time())
recovered_assistant = {
'role': 'assistant',
@ -2505,8 +2606,9 @@ def _append_journaled_partial_output(
'_recovered_from_run_journal': True,
'_recovered_stream_id': stream_id,
}
attach_display_reasoning(recovered_assistant, reasoning)
session.messages.append(recovered_assistant)
_append_recovered_turn_to_context(session, recovered_assistant)
append_context_projection(recovered_assistant)
current_assistant_idx = len(session.messages) - 1
assistant_started_at = None
appended_any = True
@ -2559,6 +2661,16 @@ def _append_journaled_partial_output(
event_name = str(event.get('event') or event.get('type') or '')
payload = event.get('payload') if isinstance(event.get('payload'), dict) else {}
created_at = event.get('created_at') if isinstance(event.get('created_at'), (int, float)) else None
if event_name == 'reasoning':
text = str(
payload.get('text') or payload.get('reasoning') or payload.get('thinking') or ''
)
if not text:
continue
if not assistant_parts and not reasoning_parts and assistant_started_at is None:
assistant_started_at = created_at or time.time()
reasoning_parts.append(text)
continue
if event_name == 'token':
text = str(payload.get('text') or '')
if not text:

View file

@ -0,0 +1,506 @@
"""Regression coverage for #3929 restart-time partial-work recovery."""
import json
import time
import pytest
import api.models as models
from api.models import (
Session,
_append_journaled_partial_output,
_apply_core_sync_or_error_marker,
)
from api.run_journal import append_run_event
from api.streaming import _sanitize_messages_for_api
@pytest.fixture(autouse=True)
def _isolate_session_state(tmp_path, monkeypatch):
session_dir = tmp_path / "sessions"
session_dir.mkdir()
monkeypatch.setattr(models, "SESSION_DIR", session_dir)
monkeypatch.setattr(models, "SESSION_INDEX_FILE", session_dir / "_index.json")
models.SESSIONS.clear()
yield
models.SESSIONS.clear()
def test_restart_recovery_preserves_display_reasoning_without_provider_replay(
tmp_path,
):
"""A dead live turn must recover all user-visible work from its journal.
Display-only reasoning belongs in the settled Worklog after reload, but it
must not leak into the model-facing history used by the next turn.
"""
session_id = "issue3929_restart"
stream_id = "stream_partial_work"
previous_messages = [
{"role": "user", "content": "Inspect the failure"},
{"role": "assistant", "content": "I will trace it."},
]
session = Session(
session_id=session_id,
title="Partial work recovery",
messages=[dict(message) for message in previous_messages],
context_messages=[dict(message) for message in previous_messages],
pending_user_message="Continue until the root cause is clear",
pending_started_at=time.time() - 120,
active_stream_id=stream_id,
)
append_run_event(
session_id,
stream_id,
"reasoning",
{"text": "Checking the first failure boundary."},
)
append_run_event(
session_id,
stream_id,
"token",
{"text": "The first boundary is the live stream."},
)
append_run_event(
session_id,
stream_id,
"tool",
{
"name": "terminal",
"preview": "rg STREAM_PARTIAL_TEXT api/streaming.py",
"args": {"command": "rg STREAM_PARTIAL_TEXT api/streaming.py"},
},
)
append_run_event(
session_id,
stream_id,
"tool_complete",
{"name": "terminal", "duration": 0.4, "is_error": False},
)
append_run_event(
session_id,
stream_id,
"reasoning",
{"text": "Comparing the journal with the session sidecar."},
)
append_run_event(
session_id,
stream_id,
"token",
{"text": "The journal still has the emitted work."},
)
core_path = tmp_path / "missing-core-transcript.json"
assert _apply_core_sync_or_error_marker(
session,
core_path,
stream_id_for_recheck=stream_id,
) is True
models.SESSIONS.clear()
reloaded = models.get_session(session_id)
recovered = [
message for message in reloaded.messages
if message.get("_recovered_from_run_journal")
]
assert [message.get("content") for message in recovered] == [
"The first boundary is the live stream.",
"The journal still has the emitted work.",
]
assert [message.get("reasoning") for message in recovered] == [
"Checking the first failure boundary.",
"Comparing the journal with the session sidecar.",
]
assert reloaded.tool_calls[0]["name"] == "terminal"
assert reloaded.tool_calls[0]["done"] is True
assert any(
message.get("_error")
and message.get("type") == "interrupted"
and "partial output above was recovered" in message.get("content", "")
for message in reloaded.messages
)
serialized_context = json.dumps(
reloaded.context_messages, ensure_ascii=False,
)
assert "Checking the first failure boundary" not in serialized_context
assert "Comparing the journal with the session sidecar" not in serialized_context
sanitized = _sanitize_messages_for_api(reloaded.context_messages)
assert all("reasoning" not in message for message in sanitized)
def test_core_sync_keeps_pending_owner_for_reasoning_only_partial(tmp_path):
"""Reasoning-only live work still proves the pending user turn had activity."""
session_id = "issue3929_reasoning_only"
stream_id = "stream_reasoning_only"
session = Session(
session_id=session_id,
title="Reasoning-only recovery",
messages=[],
context_messages=[
{"role": "user", "content": "Earlier question"},
{"role": "assistant", "content": "Earlier answer"},
],
pending_user_message="Continue the investigation",
pending_started_at=time.time() - 90,
active_stream_id=stream_id,
)
core_path = tmp_path / "core-transcript.json"
core_path.write_text(json.dumps({
"messages": [
{"role": "user", "content": "Earlier question"},
{"role": "assistant", "content": "Earlier answer"},
],
}), encoding="utf-8")
append_run_event(
session_id,
stream_id,
"reasoning",
{"text": "The worker reached a visible Thinking step before restart."},
)
assert _apply_core_sync_or_error_marker(
session,
core_path,
stream_id_for_recheck=stream_id,
) is True
recovered_users = [message for message in session.messages if message.get("_recovered")]
assert [message.get("content") for message in recovered_users] == [
"Continue the investigation",
]
recovered_reasoning = [
message for message in session.messages
if message.get("_recovered_from_run_journal") and message.get("reasoning")
]
assert [message.get("reasoning") for message in recovered_reasoning] == [
"The worker reached a visible Thinking step before restart.",
]
assert any(
message.get("_error")
and "partial output above was recovered" in message.get("content", "")
for message in session.messages
)
assert not any(
"visible Thinking step" in json.dumps(message, ensure_ascii=False)
for message in session.context_messages
)
def test_reasoning_backfill_is_idempotent_for_existing_recovered_text():
session_id = "issue3929_reasoning_dedupe"
stream_id = "stream_reasoning_dedupe"
recovered_text = "The visible partial was already restored."
session = Session(
session_id=session_id,
title="Reasoning backfill",
messages=[
{"role": "user", "content": "Continue"},
{"role": "assistant", "content": recovered_text},
],
context_messages=[
{"role": "user", "content": "Continue"},
{"role": "assistant", "content": recovered_text},
],
)
append_run_event(
session_id,
stream_id,
"reasoning",
{"text": "Backfill this Thinking row once."},
)
append_run_event(session_id, stream_id, "token", {"text": recovered_text})
assert _append_journaled_partial_output(
session, stream_id, dedupe_existing=True,
) is True
assert _append_journaled_partial_output(
session, stream_id, dedupe_existing=True,
) is False
matching = [
message for message in session.messages
if message.get("role") == "assistant" and message.get("content") == recovered_text
]
assert len(matching) == 1
assert matching[0].get("reasoning") == "Backfill this Thinking row once."
assert "Backfill this Thinking row once" not in json.dumps(
session.context_messages, ensure_ascii=False,
)
def test_reasoning_backfill_does_not_claim_matching_content_from_prior_turn():
session_id = "issue3929_reasoning_current_turn_scope"
stream_id = "stream_reasoning_current_turn_scope"
repeated_text = "The same sufficiently long partial output appears again."
prior_assistant = {"role": "assistant", "content": repeated_text}
session = Session(
session_id=session_id,
title="Current-turn reasoning scope",
messages=[
{"role": "user", "content": "Earlier request"},
prior_assistant,
{"role": "user", "content": "Current request"},
],
context_messages=[
{"role": "user", "content": "Earlier request"},
{"role": "assistant", "content": repeated_text},
{"role": "user", "content": "Current request"},
],
)
append_run_event(
session_id,
stream_id,
"reasoning",
{"text": "This Thinking belongs only to the current request."},
)
append_run_event(session_id, stream_id, "token", {"text": repeated_text})
assert _append_journaled_partial_output(
session, stream_id, dedupe_existing=True,
) is True
assert prior_assistant.get("reasoning") is None
current_rows = [
message for message in session.messages[3:]
if message.get("content") == repeated_text
]
assert len(current_rows) == 1
assert current_rows[0].get("reasoning") == (
"This Thinking belongs only to the current request."
)
assert "This Thinking belongs only to the current request" not in json.dumps(
session.context_messages, ensure_ascii=False,
)
def test_reasoning_backfill_accepts_core_row_before_recovered_owner_echo():
session_id = "issue3929_reasoning_core_owner_echo"
stream_id = "stream_reasoning_core_owner_echo"
pending_text = "Current request"
recovered_text = "The core transcript already contains this partial output."
core_assistant = {"role": "assistant", "content": recovered_text}
session = Session(
session_id=session_id,
title="Core owner echo",
messages=[
{"role": "user", "content": pending_text},
core_assistant,
{"role": "user", "content": pending_text, "_recovered": True},
],
context_messages=[
{"role": "user", "content": pending_text},
{"role": "assistant", "content": recovered_text},
],
pending_user_message=pending_text,
)
append_run_event(
session_id,
stream_id,
"reasoning",
{"text": "Attach this Thinking to the current core row."},
)
append_run_event(session_id, stream_id, "token", {"text": recovered_text})
assert _append_journaled_partial_output(
session, stream_id, dedupe_existing=True,
) is True
assert core_assistant.get("reasoning") == (
"Attach this Thinking to the current core row."
)
matching_rows = [
message for message in session.messages
if message.get("content") == recovered_text
]
assert len(matching_rows) == 1
assert "Attach this Thinking to the current core row" not in json.dumps(
session.context_messages, ensure_ascii=False,
)
def test_reasoning_only_recovery_is_idempotent_across_replays():
session_id = "issue3929_reasoning_only_dedupe"
stream_id = "stream_reasoning_only_dedupe"
session = Session(
session_id=session_id,
title="Reasoning-only replay",
messages=[{"role": "user", "content": "Continue"}],
context_messages=[{"role": "user", "content": "Continue"}],
)
append_run_event(
session_id,
stream_id,
"reasoning",
{"text": "Restore this Thinking row once."},
)
assert _append_journaled_partial_output(
session, stream_id, dedupe_existing=True,
) is True
assert _append_journaled_partial_output(
session, stream_id, dedupe_existing=True,
) is False
reasoning_rows = [
message for message in session.messages
if message.get("reasoning") == "Restore this Thinking row once."
]
assert len(reasoning_rows) == 1
def test_retry_growth_attaches_later_tool_to_reasoning_segment():
session_id = "issue3929_reasoning_then_tool_growth"
stream_id = "stream_reasoning_then_tool_growth"
session = Session(
session_id=session_id,
title="Reasoning then tool growth",
messages=[{"role": "user", "content": "Keep checking"}],
context_messages=[{"role": "user", "content": "Keep checking"}],
)
append_run_event(
session_id,
stream_id,
"reasoning",
{"text": "Inspect the first boundary."},
)
assert _append_journaled_partial_output(
session, stream_id, dedupe_existing=True,
) is True
reasoning_idx = next(
idx for idx, message in enumerate(session.messages)
if message.get("reasoning") == "Inspect the first boundary."
)
append_run_event(
session_id,
stream_id,
"tool",
{"name": "terminal", "preview": "inspect the next boundary"},
)
assert _append_journaled_partial_output(
session, stream_id, dedupe_existing=True,
) is True
reasoning_rows = [
message for message in session.messages
if message.get("reasoning") == "Inspect the first boundary."
]
assert len(reasoning_rows) == 1
assert session.tool_calls[0]["assistant_msg_idx"] == reasoning_idx
def test_identical_reasoning_segments_around_tool_remain_distinct():
session_id = "issue3929_identical_reasoning_segments"
stream_id = "stream_identical_reasoning_segments"
repeated_reasoning = "Checking the same condition again."
session = Session(
session_id=session_id,
title="Repeated Thinking segments",
messages=[{"role": "user", "content": "Check twice"}],
context_messages=[{"role": "user", "content": "Check twice"}],
)
append_run_event(
session_id,
stream_id,
"reasoning",
{"text": repeated_reasoning},
)
append_run_event(
session_id,
stream_id,
"tool",
{"name": "terminal", "preview": "first check"},
)
append_run_event(
session_id,
stream_id,
"tool_complete",
{"name": "terminal", "is_error": False},
)
append_run_event(
session_id,
stream_id,
"reasoning",
{"text": repeated_reasoning},
)
assert _append_journaled_partial_output(
session, stream_id, dedupe_existing=True,
) is True
reasoning_rows = [
(idx, message)
for idx, message in enumerate(session.messages)
if message.get("reasoning") == repeated_reasoning
]
assert len(reasoning_rows) == 2
assert reasoning_rows[0][0] != reasoning_rows[1][0]
assert session.tool_calls[0]["assistant_msg_idx"] == reasoning_rows[0][0]
def test_identical_content_segments_claim_distinct_rows_on_replay():
session_id = "issue3929_identical_content_segments"
stream_id = "stream_identical_content_segments"
repeated_content = "The same sufficiently long process update appears twice."
session = Session(
session_id=session_id,
title="Repeated process segments",
messages=[{"role": "user", "content": "Check twice"}],
context_messages=[{"role": "user", "content": "Check twice"}],
)
append_run_event(
session_id,
stream_id,
"reasoning",
{"text": "First Thinking segment."},
)
append_run_event(
session_id,
stream_id,
"token",
{"text": repeated_content},
)
append_run_event(
session_id,
stream_id,
"tool",
{"name": "terminal", "preview": "separate the segments"},
)
append_run_event(
session_id,
stream_id,
"tool_complete",
{"name": "terminal", "is_error": False},
)
append_run_event(
session_id,
stream_id,
"reasoning",
{"text": "Second Thinking segment."},
)
append_run_event(
session_id,
stream_id,
"token",
{"text": repeated_content},
)
assert _append_journaled_partial_output(
session, stream_id, dedupe_existing=True,
) is True
assert _append_journaled_partial_output(
session, stream_id, dedupe_existing=True,
) is False
matching_rows = [
message for message in session.messages
if message.get("content") == repeated_content
]
assert [message.get("reasoning") for message in matching_rows] == [
"First Thinking segment.",
"Second Thinking segment.",
]

View file

@ -763,9 +763,11 @@ class TestNonEmptyMessagesPendingCleared:
assert "partial output above was recovered" in error_msgs[0]["content"]
assert "no agent output was recovered" not in error_msgs[0]["content"]
def test_journal_recovery_does_not_materialize_reasoning_only_events(self, hermes_home, monkeypatch):
"""Run-journal repair must not turn hidden reasoning into visible chat
transcript content."""
def test_journal_recovery_restores_reasoning_only_as_display_metadata(
self, hermes_home, monkeypatch,
):
"""Run-journal repair keeps live Thinking without making it chat prose
or provider-facing history."""
s = _make_session(messages=[{"role": "user", "content": "existing turn"}])
s.pending_user_message = "Keep going"
s.pending_started_at = time.time() - 120
@ -789,16 +791,14 @@ class TestNonEmptyMessagesPendingCleared:
assert result is True
contents = [m.get("content", "") for m in s.messages]
assert not any("private scratchpad text" in c for c in contents)
reasoning_msgs = [m for m in s.messages if m.get("reasoning")]
assert [m.get("reasoning") for m in reasoning_msgs] == ["private scratchpad text"]
error_msgs = [m for m in s.messages if m.get("_error")]
assert len(error_msgs) == 1
# Reasoning-only events do not count as visible output, so the marker
# arms the lazy-retry hook (stream id is known, journal may have more
# events appear on a later read). The legacy "no agent output was
# recovered" wording is now reserved for the no-stream-id case.
assert error_msgs[0].get("_pending_journal_recovery") is True
assert error_msgs[0].get("_journal_retry_stream_id") == "reasoning_only_stream"
assert "no agent output was recovered" not in error_msgs[0]["content"]
assert "Recovering the partial output" in error_msgs[0]["content"]
assert error_msgs[0].get("_pending_journal_recovery") is not True
assert "partial output above was recovered" in error_msgs[0]["content"]
provider_history = streaming._sanitize_messages_for_api(s.messages)
assert "private scratchpad text" not in json.dumps(provider_history)
def test_journal_recovery_keeps_consecutive_tools_on_one_anchor(self, hermes_home, monkeypatch):
"""Consecutive journaled tools without an intervening visible update