From 59b6022dd2185203f4eb79c8debbbcf1a9cbf5d0 Mon Sep 17 00:00:00 2001 From: Frank Song Date: Fri, 17 Jul 2026 23:46:37 +0800 Subject: [PATCH] fix(streaming): harden recovered reasoning context ownership --- api/models.py | 65 +++++++++-- tests/test_issue3929_partial_work_recovery.py | 109 +++++++++++++++++- 2 files changed, 162 insertions(+), 12 deletions(-) diff --git a/api/models.py b/api/models.py index 0ad32368..ceb662d2 100644 --- a/api/models.py +++ b/api/models.py @@ -819,13 +819,25 @@ def _active_stream_ids(): return active_ids -def _append_recovered_turn_to_context(session, recovered: dict) -> None: - context_messages = getattr(session, 'context_messages', None) - if not isinstance(context_messages, list) or not context_messages: - return +def _recovered_model_context_projection(message: dict) -> dict | None: + if not isinstance(message, dict): + return None + projected = dict(message) + projected.pop('reasoning', None) + if projected.get('_error'): + return None + projected_text = _normalize_journal_recovery_text(projected.get('content')) + if not projected_text and not projected.get('tool_call_id') and not projected.get('tool_calls'): + return None + return projected + + +def _append_recovered_context_projection( + session, + context_messages: list, + recovered: dict, +) -> None: recovered_text = _normalize_journal_recovery_text(recovered.get('content')) - if not recovered_text and not recovered.get('tool_call_id') and not recovered.get('tool_calls'): - return if recovered_text: if recovered.get('role') == 'user': if _message_matches_pending_checkpoint( @@ -845,6 +857,27 @@ def _append_recovered_turn_to_context(session, recovered: dict) -> None: context_messages.append(dict(recovered)) +def _seed_recovered_context_from_messages(session, context_messages: list) -> None: + for message in getattr(session, 'messages', None) or []: + projected = _recovered_model_context_projection(message) + if projected is None: + continue + _append_recovered_context_projection(session, context_messages, projected) + + +def _append_recovered_turn_to_context(session, recovered: dict) -> None: + context_messages = getattr(session, 'context_messages', None) + if not isinstance(context_messages, list): + context_messages = [] + session.context_messages = context_messages + if not context_messages: + _seed_recovered_context_from_messages(session, context_messages) + projected = _recovered_model_context_projection(recovered) + if projected is None: + return + _append_recovered_context_projection(session, context_messages, projected) + + def _append_recovered_pending_turn(session, *, timestamp: int | None = None) -> dict | None: pending_text = str(session.pending_user_message or '') if not pending_text: @@ -2517,8 +2550,13 @@ def _append_journaled_partial_output( 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: + if pending_text and not _message_matches_pending_checkpoint( + messages[owner_idx], + session.pending_user_message, + session.pending_started_at, + session.pending_user_source, + session.pending_attachments, + ): return False for candidate_idx in range(existing_idx + 1, initial_message_count): @@ -2526,8 +2564,17 @@ def _append_journaled_partial_output( 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: + candidate_matches_checkpoint = pending_text and _message_matches_pending_checkpoint( + candidate, + session.pending_user_message, + session.pending_started_at, + session.pending_user_source, + session.pending_attachments, + ) + if candidate_matches_checkpoint and candidate.get('_recovered'): continue + if pending_text and candidate_text == pending_text: + return False return False return True diff --git a/tests/test_issue3929_partial_work_recovery.py b/tests/test_issue3929_partial_work_recovery.py index b69e369f..cef22048 100644 --- a/tests/test_issue3929_partial_work_recovery.py +++ b/tests/test_issue3929_partial_work_recovery.py @@ -12,6 +12,7 @@ from api.models import ( _apply_core_sync_or_error_marker, ) from api.run_journal import append_run_event +import api.streaming as streaming from api.streaming import _sanitize_messages_for_api @@ -276,25 +277,127 @@ def test_reasoning_backfill_does_not_claim_matching_content_from_prior_turn(): ) +def test_repeated_pending_prompt_uses_checkpoint_owner_not_prior_same_text(): + session_id = "issue3929_repeated_pending_prompt" + stream_id = "stream_repeated_pending_prompt" + repeated_text = "The same sufficiently long recovered answer appears." + prior_assistant = {"role": "assistant", "content": repeated_text} + pending_started_at = 2_000 + session = Session( + session_id=session_id, + title="Repeated pending prompt", + messages=[ + {"role": "user", "content": "Continue", "timestamp": 1_000}, + prior_assistant, + { + "role": "user", + "content": "Continue", + "timestamp": pending_started_at, + "_recovered": True, + }, + ], + context_messages=[ + {"role": "user", "content": "Continue", "timestamp": 1_000}, + {"role": "assistant", "content": repeated_text}, + { + "role": "user", + "content": "Continue", + "timestamp": pending_started_at, + "_recovered": True, + }, + ], + pending_user_message="Continue", + pending_started_at=pending_started_at, + ) + append_run_event( + session_id, + stream_id, + "reasoning", + {"text": "This Thinking belongs to the second Continue turn."}, + ) + 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 to the second Continue turn." + ) + assert "second Continue turn" not in json.dumps( + session.context_messages, ensure_ascii=False, + ) + + +def test_empty_context_recovery_seeds_reasoning_free_model_context(): + session_id = "issue3929_empty_context_reasoning" + stream_id = "stream_empty_context_reasoning" + session = Session( + session_id=session_id, + title="Empty context reasoning", + messages=[ + {"role": "user", "content": "Continue", "timestamp": 1234}, + ], + context_messages=[], + ) + append_run_event( + session_id, + stream_id, + "reasoning", + {"text": "Display-only Thinking must not become model context."}, + ) + + assert _append_journaled_partial_output( + session, stream_id, dedupe_existing=True, + ) is True + + assert [message.get("content") for message in session.context_messages] == [ + "Continue", + ] + serialized_context = json.dumps( + session.context_messages, ensure_ascii=False, + ) + assert "Display-only Thinking" not in serialized_context + assert all("reasoning" not in message for message in session.context_messages) + + next_turn_context = streaming._context_messages_for_new_turn( + session, "Now continue from there", + ) + assert next_turn_context == session.context_messages + + 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" + pending_started_at = 3_000 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}, + {"role": "user", "content": pending_text, "timestamp": pending_started_at}, core_assistant, - {"role": "user", "content": pending_text, "_recovered": True}, + { + "role": "user", + "content": pending_text, + "timestamp": pending_started_at, + "_recovered": True, + }, ], context_messages=[ - {"role": "user", "content": pending_text}, + {"role": "user", "content": pending_text, "timestamp": pending_started_at}, {"role": "assistant", "content": recovered_text}, ], pending_user_message=pending_text, + pending_started_at=pending_started_at, ) append_run_event( session_id,