From cc04825c5dd49afe9dd503f4115d1e9b37e64505 Mon Sep 17 00:00:00 2001 From: kshitij Date: Mon, 3 Aug 2026 10:22:53 +0530 Subject: [PATCH] fix: content-based diff for model-switch marker merge (#76870) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The original PR #77274 used positional slicing (current_history[len(history):]) to detect the model-switch-only mutation. But _append_model_switch_marker strips prior markers in-place before appending the new one, so when a prior marker existed (every switch after the first in a session), the net length delta is zero and the slice produces an empty list — the merge path is dead code for the common case. Replace with a content-based diff: strip markers from both the turn-start snapshot and the current history, then check that the non-marker content is identical. This correctly handles the strip-and-replace behavior. Also guard against auto-compression making result["messages"] shorter than the turn-start history — use the full result as the base when that happens. Added test covering both no-prior-marker and prior-marker cases. --- tests/test_tui_gateway_server.py | 115 +++++++++++++++++++++++++++++++ tui_gateway/server.py | 30 ++++++-- 2 files changed, 141 insertions(+), 4 deletions(-) diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index f60d66f8f2507..00021195a4343 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -8907,6 +8907,121 @@ def test_prompt_submit_history_version_mismatch_surfaces_warning(monkeypatch): server._sessions.pop("sid", None) +def test_prompt_submit_merges_on_model_switch_marker(monkeypatch): + """#76870: when a model-switch marker is the only history mutation during + a turn, the agent's output must be merged into the current history (which + now contains the marker) instead of being discarded. + + This test covers BOTH cases: + - No prior marker in turn-start history (first switch in a session) + - Prior marker existed (every subsequent switch — the original PR #77274 + fix was dead code here because _append_model_switch_marker strips the + old marker before appending the new one, producing a net-zero length + delta that the positional slice missed). + """ + from tui_gateway.server import _MODEL_SWITCH_MARKER_PREFIX + + session_ref = {"s": None} + + def _make_marker(model: str) -> dict: + return { + "role": "user", + "content": f"{_MODEL_SWITCH_MARKER_PREFIX}{model}.]", + "display_kind": "model_switch", + } + + class _MarkerAgent: + def __init__(self, new_history_state: list): + self._new_history_state = new_history_state + + def run_conversation(self, prompt, conversation_history=None, stream_callback=None, **_kwargs): + # Simulate _append_model_switch_marker: strip prior markers, append new one. + with session_ref["s"]["history_lock"]: + hist = session_ref["s"]["history"] + hist[:] = [h for h in hist if not _is_marker(h)] + hist.append(_make_marker("new-model")) + session_ref["s"]["history_version"] += 1 + # result["messages"] = conversation_history + user msg + assistant reply + return { + "final_response": "agent reply", + "messages": list(conversation_history) + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "agent reply"}, + ], + } + + def _is_marker(entry) -> bool: + from tui_gateway.server import _is_model_switch_marker + return _is_model_switch_marker(entry) + + class _ImmediateThread: + def __init__(self, target=None, daemon=None): + self._target = target + + def start(self): + self._target() + + # Test both: no prior marker, and prior marker present + for label, prior_history in [ + ("no prior marker", [{"role": "user", "content": "hello"}]), + ("with prior marker", [ + {"role": "user", "content": "hello"}, + _make_marker("old-model"), + {"role": "assistant", "content": "hi there"}, + ]), + ]: + server._sessions["sid"] = _session( + agent=_MarkerAgent([]), + history=list(prior_history), + ) + session_ref["s"] = server._sessions["sid"] + emits: list[tuple] = [] + try: + monkeypatch.setattr(server.threading, "Thread", _ImmediateThread) + monkeypatch.setattr(server, "_get_usage", lambda _a: {}) + monkeypatch.setattr(server, "render_message", lambda _t, _c: "") + monkeypatch.setattr(server, "_emit", lambda *a: emits.append(a)) + + resp = server.handle_request( + { + "id": "1", + "method": "prompt.submit", + "params": {"session_id": "sid", "text": "hi"}, + } + ) + assert resp.get("result"), f"[{label}] got error: {resp.get('error')}" + + final_history = server._sessions["sid"]["history"] + + # The agent's new messages must be present in the persisted history. + assistant_msgs = [ + e for e in final_history + if isinstance(e, dict) and e.get("role") == "assistant" + and e.get("content") == "agent reply" + ] + assert len(assistant_msgs) == 1, ( + f"[{label}] agent output was not merged into history " + f"(got {len(assistant_msgs)} assistant 'agent reply' messages)" + ) + + # The model-switch marker must be present. + markers = [e for e in final_history if _is_marker(e)] + assert len(markers) == 1, ( + f"[{label}] expected exactly 1 model-switch marker, got {len(markers)}" + ) + assert "new-model" in markers[0]["content"] + + # No warning should be surfaced — the merge succeeded. + complete_calls = [a for a in emits if a[0] == "message.complete"] + assert len(complete_calls) == 1 + _, _, payload = complete_calls[0] + assert "warning" not in payload, ( + f"[{label}] merge path should not surface a warning" + ) + finally: + server._sessions.pop("sid", None) + + def test_prompt_submit_sanitizes_bracketed_paste_before_agent(monkeypatch): """prompt.submit must sanitize corrupted user text before run_conversation.""" captured: dict[str, str] = {} diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 4662a1ed72b50..c5ed0080f75be 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -9694,14 +9694,36 @@ def _run_prompt_submit( # marker inserted mid-turn (#76870). If so the # agent output is still valid — merge it into the # current history that now contains the marker. + # + # _append_model_switch_marker strips prior markers + # in-place then appends a new one, so the delta + # is NOT a simple tail-slice — we must compare + # content, not indices. current_history = list(session["history"]) - added = current_history[len(history):] + history_no_markers = [ + e for e in history if not _is_model_switch_marker(e) + ] + current_no_markers = [ + e for e in current_history if not _is_model_switch_marker(e) + ] model_switch_only = ( - len(added) >= 1 - and all(_is_model_switch_marker(e) for e in added) + current_no_markers == history_no_markers + and any( + _is_model_switch_marker(e) + for e in current_history + ) ) if model_switch_only: - new_messages = result["messages"][len(history):] + # The agent's new messages start after the + # turn-start history. Guard against + # auto-compression making result["messages"] + # shorter than history (#77274 review). + if len(result["messages"]) > len(history): + new_messages = result["messages"][len(history):] + else: + # Compression rebound the messages list — + # use the full result as the base. + new_messages = list(result["messages"]) session["history"] = current_history + new_messages session["history_version"] = current_version + 1 else: