From d5b357cd1c8ecce5f343063beb9b0868960ee90f Mon Sep 17 00:00:00 2001 From: b3nw Date: Sat, 4 Jul 2026 18:22:12 +0000 Subject: [PATCH] fix(session): align forked context to display for large (compacted) sessions truncate_context_for_display_keep short-circuited on `len(ctx) <= len(msgs)` and returned ctx[:keep], slicing a trimmed (shorter) model context at the DISPLAY index. For large/compacted sessions this misaligned the forked context from the visible transcript and left a dangling assistant tool_use at the tail, producing broken/disjointed forks. Small sessions were unaffected (equal-length arrays), which masked the bug. Narrow the naive-slice guard to `len(ctx) == len(msgs)` so both divergent directions fall through to the signature matcher and the cut lands on a real turn boundary. Add an ambiguous-match-aware shorter-context fallback, gated to `len(ctx) < len(msgs)` so the context-longer (#5096 summary-prefix) path is unchanged. Upholds the model-context invariant in docs/rfcs/webui-run-state-consistency-contract.md (context must not contradict what the user can see). Live case b88a443d2f2b (618 display / 600 context) now forks to context[:425] (completed tool result) instead of context[:440] (dangling call_0366, ~15 turns out of sync). Tests: correct the compact-summary-fallback assertion (it encoded the contract-violating drop of a kept user turn's assistant reply) and add shorter-context alignment, ambiguous-boundary, and zero-match regressions. --- api/session_ops.py | 41 +++++++- tests/test_issue_branch_context_at_fork.py | 104 ++++++++++++++++++++- 2 files changed, 141 insertions(+), 4 deletions(-) diff --git a/api/session_ops.py b/api/session_ops.py index 3d47a8a7d10..29432a31664 100644 --- a/api/session_ops.py +++ b/api/session_ops.py @@ -110,10 +110,25 @@ def truncate_context_for_display_keep( msgs = full_messages if isinstance(full_messages, list) else [] if not ctx: return [] - if len(ctx) <= len(msgs): - return ctx[:keep] if len(msgs) == 0: return [] + # Only the perfectly-parallel case (display and context row-for-row) can be + # sliced at the raw display index. When the two arrays differ in length — + # in EITHER direction — they have diverged and need alignment: + # * context LONGER than display → an injected summary/system prefix, etc. + # * context SHORTER than display → large-session context trimming dropped + # turns from the model context that the display still shows. + # The shorter-context case is the one that broke forked large sessions: the + # old ``len(ctx) <= len(msgs)`` guard short-circuited to ``ctx[:keep]``, + # slicing the shorter context at the display index (landing mid-turn, e.g. + # on an assistant tool_call whose result was past the cut). Fall through to + # the signature matcher for both divergent cases so the cut lands on a real + # turn boundary. Any residual dangling tool_use in the persisted context is + # made wire-safe on the send path (streaming: ``_sanitize_messages_for_api`` + # strips unanswered tool_calls; gateway: it forwards no tool_calls/tool rows + # at all), so we do not re-do that trimming here. + if len(ctx) == len(msgs): + return ctx[:keep] def _row_signature(row: Any) -> tuple[str, ...] | None: if not isinstance(row, dict): @@ -197,7 +212,27 @@ def _first_match_from(message: Any, start_idx: int) -> tuple[int | None, int | N return ctx[:ambiguous_first_unkept] return ctx[:last_kept + 1] - # Final fallback preserves #5096 behavior when alignment is unreliable. + # Both boundary rows were ambiguous/unmatched (common in large sessions + # where context rows have lost their id/timestamp so the matcher can't + # disambiguate structurally-identical rows). Only for the shorter-context + # case: cut just past the LAST display row in the kept prefix that + # resolved to a context index — preferring an exact match but accepting + # an ambiguous (weak) one, mirroring how the sibling branches above fold + # ``ambiguous_matches`` into the boundary. Accepting the weak match keeps + # the forked boundary turn's own context (often exactly that ambiguous + # row) instead of dropping back to an earlier exact match. It still errs + # toward UNDER-keeping rather than slicing at the raw display index, + # which would over-keep and mis-attribute later context rows to the kept + # display turns. The context-longer case (injected summary prefix) is + # left to the #5096 fallback below, which preserves that prefix. + if len(ctx) < len(msgs): + for i in range(keep - 1, -1, -1): + resolved = matches[i] if matches[i] is not None else ambiguous_matches[i] + if resolved is not None: + return ctx[:resolved + 1] + + # Final fallback preserves #5096 behavior when alignment is unreliable + # (no display row resolved to a context index, or keep >= len(msgs)). prefix_len = max(0, len(ctx) - len(msgs)) prefix = ctx[:prefix_len] suffix = ctx[prefix_len:] diff --git a/tests/test_issue_branch_context_at_fork.py b/tests/test_issue_branch_context_at_fork.py index fae95e82cb1..a3d5f589ab1 100644 --- a/tests/test_issue_branch_context_at_fork.py +++ b/tests/test_issue_branch_context_at_fork.py @@ -107,6 +107,15 @@ def test_truncate_context_for_display_keep_drops_unkept_tool_rows_after_user_bou def test_truncate_context_for_display_keep_prefers_compact_summary_fallback(): + # Context is SHORTER than display (a compaction summary replaced earlier + # turns) and the display has trailing turns beyond the kept prefix. Keeping + # display [u1, a1] must map to the context that represents [u1, a1] — the + # leading compaction row PLUS both u1 and a1. The old behaviour sliced the + # shorter context at the raw display index (returning [compact, u1]), which + # dropped a1 — the assistant reply to a kept user turn — from the model + # context while the display still showed it. That mismatch is the forked- + # large-session disjointedness bug; the result must now match the sibling + # ``preserves_leading_compaction_row`` case regardless of trailing turns. msgs = [ {"role": "user", "content": "u1", "id": "u1", "timestamp": 1.0}, {"role": "assistant", "content": "a1", "id": "a1", "timestamp": 2.0}, @@ -119,7 +128,7 @@ def test_truncate_context_for_display_keep_prefers_compact_summary_fallback(): {"role": "assistant", "content": "a1", "id": "a1", "timestamp": 2.0}, ] out = truncate_context_for_display_keep(ctx, msgs, 2) - assert out == ctx[:2] + assert [row["content"] for row in out] == ["compact", "u1", "a1"] def test_truncate_context_for_display_keep_keeps_real_user_turn_when_duplicate_rows_lack_identity(): @@ -135,6 +144,99 @@ def test_truncate_context_for_display_keep_keeps_real_user_turn_when_duplicate_r out = truncate_context_for_display_keep(ctx, msgs, 1) assert [row["content"] for row in out] == ["u1", "u1"] +def test_truncate_context_shorter_than_display_aligns_to_turn_boundary(): + """Large-session fork regression. + + Mirrors the live shape that broke forks: the model context has been trimmed + so it is SHORTER than the display transcript, and its rows carry no ``id`` + (and mostly no ``timestamp``). The old ``len(ctx) <= len(msgs)`` + short-circuit returned ``ctx[:keep]`` verbatim — slicing the shorter context + at the *display* index, which lands mid-turn on an assistant ``tool_use`` + whose result is past the cut. The fix routes the shorter-context case + through the signature matcher, which cuts on the last aligned turn boundary + (a completed ``tool`` result), so the context never ends on an unmatched + ``tool_use``. (Any residual dangling ``tool_use`` on the harder no-match + paths is stripped tool-id-aware by ``_sanitize_messages_for_api`` at send.) + """ + msgs = [ + {"role": "user", "content": "u1", "timestamp": 1.0}, + {"role": "assistant", "content": "a1", "timestamp": 2.0, + "tool_calls": [{"id": "c1"}]}, + {"role": "tool", "content": "r1", "tool_call_id": "c1", "timestamp": 3.0}, + {"role": "assistant", "content": "done", "timestamp": 4.0}, # keep boundary + {"role": "user", "content": "u2", "timestamp": 5.0}, + {"role": "assistant", "content": "a2", "timestamp": 6.0, + "tool_calls": [{"id": "c2"}]}, + ] + # Context: no ids, no timestamps (matcher must fall back to signature), the + # "done" summary turn was compressed out, so ctx is shorter and diverges. + ctx = [ + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1", "tool_calls": [{"id": "c1"}]}, + {"role": "tool", "content": "r1", "tool_call_id": "c1"}, + {"role": "assistant", "content": "a2", "tool_calls": [{"id": "c2"}]}, + {"role": "tool", "content": "r2", "tool_call_id": "c2"}, + ] + out = truncate_context_for_display_keep(ctx, msgs, 4) + contents = [row["content"] for row in out] + # Old raw ctx[:4] would be ["u1","a1","r1","a2"] — "a2" is a dangling + # tool_use. Alignment instead cuts at the last resolved turn boundary (r1). + assert contents == ["u1", "a1", "r1"] + assert not (out and out[-1].get("role") == "assistant" and out[-1].get("tool_calls")) + + +def test_shorter_context_ambiguous_boundary_keeps_forked_turn_via_weak_match(): + """Shorter context, forked boundary turn matches only ambiguously. + + The kept boundary turn (a1) has no id/timestamp and appears twice in the + shorter context, so the matcher records it in ambiguous_matches (not + matches). Both boundary rows are then unmatched, so the shorter-context + fallback loop runs. It must accept the weak (ambiguous) match for the + boundary turn and KEEP a1's context, rather than dropping back to the + earlier exact match (u1) and losing the very turn the user forked at. + """ + msgs = [ + {"role": "user", "content": "u1", "id": "u1"}, + {"role": "assistant", "content": "a1"}, # keep boundary, no id (ambiguous) + {"role": "user", "content": "u2"}, + {"role": "assistant", "content": "a2"}, + ] + # Shorter context (3 < 4); a1 appears twice with no id → ambiguous match. + ctx = [ + {"role": "user", "content": "u1", "id": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "assistant", "content": "a1"}, + ] + out = truncate_context_for_display_keep(ctx, msgs, 2) + # Must include the forked boundary turn a1 (weak match at ctx[1]), + # not stop at the earlier exact match u1. + assert [row["content"] for row in out] == ["u1", "a1"] + + +def test_shorter_context_zero_match_falls_back_to_best_effort_prefix(): + """Shorter context where NO kept-prefix row aligns (unalignable last resort). + + When the kept display prefix lies entirely inside a summarized region whose + rows share no signature with the display, the matcher resolves nothing and + neither the sibling branches nor the ambiguous-aware fallback loop fire. + Control reaches the #5096 last-resort fallback, which returns a best-effort + prefix of the (shorter) context. Alignment is genuinely impossible here; + wire-safety of any dangling tool_use is handled at send time. This test + pins the documented last-resort behavior so it is not changed unknowingly. + """ + msgs = [ + {"role": "user", "content": "u1", "id": "u1"}, + {"role": "assistant", "content": "a1", "id": "a1"}, + {"role": "user", "content": "u2", "id": "u2"}, + ] + ctx = [ + {"role": "user", "content": "summary-1"}, + {"role": "assistant", "content": "summary-2"}, + ] + out = truncate_context_for_display_keep(ctx, msgs, 2) + assert [row["content"] for row in out] == ["summary-1", "summary-2"] + + def test_truncate_context_for_display_keep_keeps_tool_tail_before_ambiguous_unkept_anchor(): msgs = [ {"role": "user", "content": "u1", "id": "u1"},