fix(session): align forked context to display for large (compacted) s…

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:]

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"},