From 2dd4cbbe61cd6feaf73f7b690668d44a70416cc1 Mon Sep 17 00:00:00 2001 From: Brooklyn Nicholson Date: Sun, 26 Jul 2026 19:43:23 -0500 Subject: [PATCH] fix(tui_gateway): keep the original prompt when a redirect corrects a turn An accepted mid-turn redirect wrote its correction over inflight_turn["user"]. That field is the only user text session.resume can replay, so the prompt that started the turn was gone the moment the user typed again while it ran. On the next resume the client rebuilt the thread without it. Record corrections in their own list instead, alongside the prompt. Renamed _replace_inflight_user to _record_inflight_correction now that it appends. _start_inflight_turn rebuilds the dict wholesale, so corrections cannot leak into a later turn. --- tests/test_tui_gateway_queue_on_busy.py | 4 ++- tests/test_tui_gateway_server.py | 47 ++++++++++++++++++++++++- tui_gateway/server.py | 28 +++++++++++---- 3 files changed, 70 insertions(+), 9 deletions(-) diff --git a/tests/test_tui_gateway_queue_on_busy.py b/tests/test_tui_gateway_queue_on_busy.py index ea13596ef691..1fffb7db3f68 100644 --- a/tests/test_tui_gateway_queue_on_busy.py +++ b/tests/test_tui_gateway_queue_on_busy.py @@ -66,7 +66,9 @@ def test_busy_interrupt_mode_redirects_active_turn(monkeypatch): assert resp["result"]["status"] == "redirected" assert seen == ["redirect"] - assert session["inflight_turn"]["user"] == "redirect" + # Appended, not overwritten: the original prompt must stay recoverable. + assert session["inflight_turn"]["user"] == "original request" + assert session["inflight_turn"]["corrections"] == ["redirect"] assert session.get("queued_prompt") is None diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index f57373752d39..c5ca08aaf2bb 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -7531,11 +7531,56 @@ def test_session_redirect_calls_capable_core_agent(monkeypatch): "text": "use Postgres", } assert calls == ["use Postgres"] - assert session["inflight_turn"]["user"] == "use Postgres" + # The correction is recorded alongside the prompt that started the turn, + # never over it — resume must be able to rebuild both bubbles. + assert session["inflight_turn"]["user"] == "original request" + assert session["inflight_turn"]["corrections"] == ["use Postgres"] assert session.get("last_active") is not None assert before is None or session["last_active"] >= before +def test_session_redirect_records_correction_without_erasing_prompt(): + """A redirect must not overwrite the turn's original user text. + + The inflight snapshot is the only thing session.resume can replay, so + overwriting ``user`` erased the prompt that started the turn and the + client repainted the thread with the user's message missing. + """ + session = {} + server._start_inflight_turn(session, "remove the session counts") + server._append_inflight_delta(session, "Moving.") + server._record_inflight_correction(session, "hurry up") + server._record_inflight_correction(session, "and the worktree ones") + + snapshot = server._inflight_snapshot(session) + assert snapshot is not None + + assert snapshot["user"] == "remove the session counts" + assert snapshot["corrections"] == ["hurry up", "and the worktree ones"] + + +def test_inflight_snapshot_omits_corrections_when_none_recorded(): + session = {} + server._start_inflight_turn(session, "just the prompt") + + snapshot = server._inflight_snapshot(session) + assert snapshot is not None + assert "corrections" not in snapshot + + +def test_new_turn_does_not_inherit_prior_turn_corrections(): + session = {} + server._start_inflight_turn(session, "first prompt") + server._record_inflight_correction(session, "first correction") + server._start_inflight_turn(session, "second prompt") + + snapshot = server._inflight_snapshot(session) + assert snapshot is not None + + assert snapshot["user"] == "second prompt" + assert "corrections" not in snapshot + + def test_session_redirect_queues_during_agent_build_window(monkeypatch): # A fresh turn flips running=True and builds the agent asynchronously, so # session["agent"] is briefly None. A correction landing here must queue diff --git a/tui_gateway/server.py b/tui_gateway/server.py index f90350d2ec82..872bfdbbbf30 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -6378,16 +6378,25 @@ def _append_inflight_delta(session: dict, delta: Any) -> None: session["inflight_turn"] = turn -def _replace_inflight_user(session: dict, text: Any) -> None: - """Reflect an accepted correction as the live turn's current user text.""" - user = _inflight_text(text) - if not user: +def _record_inflight_correction(session: dict, text: Any) -> None: + """Record an accepted mid-turn correction on the live turn. + + The correction is appended, never written over ``user``: a resuming client + must be able to rebuild BOTH bubbles. Overwriting the slot erased the + prompt that started the turn from the only snapshot resume can read, so a + reconnect (or a dev hot-reload that wipes the renderer cache) repainted the + thread with the user's original message missing. + """ + correction = _inflight_text(text) + if not correction: return turn = session.get("inflight_turn") if not isinstance(turn, dict): return turn = dict(turn) - turn["user"] = user + corrections = list(turn.get("corrections") or []) + corrections.append(correction) + turn["corrections"] = corrections turn["updated_at"] = time.time() session["inflight_turn"] = turn @@ -6680,7 +6689,7 @@ def _handle_busy_submit( try: if agent.redirect(plain_text): with session["history_lock"]: - _replace_inflight_user(session, plain_text) + _record_inflight_correction(session, plain_text) session["last_active"] = time.time() return _ok(rid, {"status": "redirected"}) except Exception: @@ -6751,6 +6760,11 @@ def _inflight_snapshot(session: dict) -> dict | None: "streaming": streaming, "user": user, } + corrections = [c for c in (turn.get("corrections") or []) if str(c).strip()] + if corrections: + # Mid-turn redirects. Carried alongside the original prompt (not over + # it) so resume can rebuild every user bubble the turn produced. + snapshot["corrections"] = [str(c) for c in corrections] if error: # Retained failed turn (see _fail_inflight_turn): carry the error # semantics so a resuming client can rebuild the failed-turn bubble @@ -10744,7 +10758,7 @@ def _(rid, params: dict) -> dict: return _err(rid, 5000, f"redirect failed: {exc}") if accepted: with session["history_lock"]: - _replace_inflight_user(session, text) + _record_inflight_correction(session, text) session["last_active"] = time.time() return _ok( rid,