From 5934c2fe8a034a27565e90f4769e548a716320c0 Mon Sep 17 00:00:00 2001 From: Lumen Yang Date: Fri, 22 May 2026 09:48:28 +0000 Subject: [PATCH] fix: address context replay review feedback --- api/routes.py | 5 + api/streaming.py | 3 +- tests/test_issue1217_transcript_compaction.py | 121 +++++++++++++++++- 3 files changed, 127 insertions(+), 2 deletions(-) diff --git a/api/routes.py b/api/routes.py index fe789dea..97db1f2e 100644 --- a/api/routes.py +++ b/api/routes.py @@ -9079,6 +9079,7 @@ def _handle_chat_sync(handler, body): session_id=s.session_id, ) from api.streaming import ( + _dedupe_replayed_context_messages, _merge_display_messages_after_agent_result, _restore_display_reasoning_metadata, _restore_reasoning_metadata, @@ -9129,6 +9130,10 @@ def _handle_chat_sync(handler, body): _previous_context_messages, _result_messages, ) + _next_context_messages = _dedupe_replayed_context_messages( + _previous_context_messages, + _next_context_messages, + ) s.context_messages = _next_context_messages s.messages = _merge_display_messages_after_agent_result( _previous_messages, diff --git a/api/streaming.py b/api/streaming.py index 5081d70e..70a0a87d 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -2469,7 +2469,8 @@ def _dedupe_replayed_context_messages(previous_context, result_messages): return result_messages candidates = result_messages[len(previous_context):] candidates = _strip_replayed_prefix(previous_context, candidates) - candidates = _strip_replayed_context_items(previous_context, candidates) + if candidates: + candidates = _strip_replayed_context_items(previous_context, candidates) return previous_context + candidates diff --git a/tests/test_issue1217_transcript_compaction.py b/tests/test_issue1217_transcript_compaction.py index e45e9d11..7b870d61 100644 --- a/tests/test_issue1217_transcript_compaction.py +++ b/tests/test_issue1217_transcript_compaction.py @@ -1,7 +1,11 @@ -from api.models import Session, reconciled_state_db_messages_for_session import contextlib +import io +import json +import sys from types import SimpleNamespace +from api.models import Session, reconciled_state_db_messages_for_session + from api.streaming import ( _assistant_reply_added_after_current_turn, _context_messages_for_new_turn, @@ -426,6 +430,31 @@ def test_prefer_context_reconcile_keeps_state_delta_when_no_mirrored_prefix(): assert reconciled == sidecar_context + state_messages +def test_prefer_context_reconcile_strips_small_mirrored_context_prefix(): + sidecar_context = [ + {"role": "user", "content": "[Session Arc Summary] compacted"}, + {"role": "assistant", "content": "last compacted answer"}, + ] + state_messages = [ + {"role": "user", "content": "[Session Arc Summary] compacted", "timestamp": 100.0}, + {"role": "assistant", "content": "last compacted answer", "timestamp": 101.0}, + {"role": "user", "content": "fresh follow-up", "timestamp": 200.0}, + ] + session = SimpleNamespace( + session_id="small-context-prefix", + context_messages=sidecar_context, + messages=[], + ) + + reconciled = reconciled_state_db_messages_for_session( + session, prefer_context=True, state_messages=state_messages + ) + + assert reconciled == sidecar_context + [ + {"role": "user", "content": "fresh follow-up", "timestamp": 200.0}, + ] + + def test_prefer_context_reconcile_strips_mirrored_rows_without_sidecar_timestamps(): sidecar_context = [ {"role": "assistant", "content": "cron banner"}, @@ -502,6 +531,96 @@ def test_non_streaming_chat_writeback_dedupes_full_context_replay(): {"role": "assistant", "content": "short answer"}, ] +class _FakePostHandler: + def __init__(self): + self.status = None + self.headers = {} + self.body = bytearray() + self.wfile = self + + def send_response(self, status): + self.status = status + + def send_header(self, name, value): + self.headers[name] = value + + def end_headers(self): + pass + + def write(self, data): + self.body.extend(data) + + def json_body(self): + return json.loads(bytes(self.body).decode("utf-8")) + + +def test_handle_chat_sync_writeback_dedupes_full_context_replay(tmp_path, monkeypatch): + import api.config as config + import api.models as models + import api.routes as routes + + state_dir = tmp_path / "state" + session_dir = state_dir / "sessions" + session_dir.mkdir(parents=True) + monkeypatch.setattr(models, "SESSION_DIR", session_dir) + monkeypatch.setattr(models, "SESSION_INDEX_FILE", state_dir / "session_index.json") + monkeypatch.setattr(routes, "SESSION_INDEX_FILE", state_dir / "session_index.json") + monkeypatch.setattr(routes, "get_session", models.get_session) + monkeypatch.setattr(routes, "title_from", models.title_from) + monkeypatch.setattr(config, "get_config", lambda: {"model": "test-model", "provider": "test-provider"}) + monkeypatch.setattr(routes, "get_config", lambda: {"model": "test-model", "provider": "test-provider"}) + monkeypatch.setattr(routes, "resolve_trusted_workspace", lambda value: tmp_path) + monkeypatch.setattr(routes, "load_settings", lambda: {}) + monkeypatch.setattr(routes, "_resolve_cli_toolsets", lambda: []) + + previous_context = [ + {"role": "assistant", "content": "cron banner"}, + {"role": "user", "content": "[Session Arc Summary (d1, node 39)]\n" + "old context\n" * 400}, + {"role": "assistant", "content": "previous answer"}, + ] + session = Session( + session_id="sync_chat_replay", + workspace=str(tmp_path), + messages=list(previous_context), + context_messages=list(previous_context), + model="test-model", + model_provider="test-provider", + ) + session.save(touch_updated_at=False) + + replayed_result = previous_context + previous_context + [ + {"role": "user", "content": "simple follow-up"}, + {"role": "assistant", "content": "short answer"}, + ] + + class FakeAgent: + def __init__(self, **_kwargs): + pass + + def run_conversation(self, **_kwargs): + return { + "messages": replayed_result, + "final_response": "short answer", + "completed": True, + } + + monkeypatch.setitem(sys.modules, "run_agent", SimpleNamespace(AIAgent=FakeAgent)) + + handler = _FakePostHandler() + routes._handle_chat_sync( + handler, + {"session_id": session.session_id, "message": "simple follow-up", "workspace": str(tmp_path)}, + ) + + assert handler.status == 200 + reloaded = Session.load(session.session_id) + assert reloaded.context_messages == previous_context + [ + {"role": "user", "content": "simple follow-up"}, + {"role": "assistant", "content": "short answer"}, + ] + assert reloaded.context_messages.count(previous_context[0]) == 1 + + def test_session_context_falls_back_to_display_messages_for_legacy_sessions(tmp_path): messages = [ {"role": "user", "content": "legacy prompt"},