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,