fix(session): drop interleaved tool rows after mid-history truncate (#5134)

api/session_ops.py

@@ -111,9 +111,63 @@ def truncate_context_for_display_keep(
         return []
     if len(ctx) == len(msgs):
         return ctx[:keep]
+    if len(ctx) < len(msgs):
+        return ctx[:keep]
     if len(msgs) == 0:
         return []
-    # Context tail aligns with full display transcript; preserve leading-only rows.
+
+    def _row_matches(context_row: Any, msg_row: Any) -> bool:
+        """Match rows by stable id, then (role, content, timestamp), then
+        (role, content)."""
+        if not isinstance(context_row, dict) or not isinstance(msg_row, dict):
+            return False
+
+        context_id = context_row.get('id')
+        msg_id = msg_row.get('id')
+        if context_id is not None and msg_id is not None and context_id == msg_id:
+            return True
+
+        context_ts = context_row.get('timestamp')
+        msg_ts = msg_row.get('timestamp')
+        if context_ts is not None and msg_ts is not None:
+            if (
+                context_row.get('role') == msg_row.get('role')
+                and context_row.get('content') == msg_row.get('content')
+                and context_ts == msg_ts
+            ):
+                return True
+
+        return (
+            context_row.get('role') == msg_row.get('role')
+            and context_row.get('content') == msg_row.get('content')
+        )
+
+    def _first_match_from(message: Any, start_idx: int) -> int | None:
+        for idx in range(start_idx, len(ctx)):
+            if _row_matches(ctx[idx], message):
+                return idx
+        return None
+
+    matches = [None] * len(msgs)
+    next_ctx_idx = 0
+    for msg_idx, message in enumerate(msgs):
+        match_idx = _first_match_from(message, next_ctx_idx)
+        matches[msg_idx] = match_idx
+        if match_idx is not None:
+            next_ctx_idx = match_idx + 1
+
+    # Cut at the first unkept display turn, or fallback to the last kept turn
+    # if the boundary is not directly alignable.
+    if keep < len(msgs):
+        first_unkept = matches[keep]
+        if first_unkept is not None:
+            return ctx[:first_unkept]
+        if keep > 0:
+            last_kept = matches[keep - 1]
+            if last_kept is not None:
+                return ctx[:last_kept + 1]
+
+    # Final fallback preserves #5096 behavior when alignment is unreliable.
     prefix_len = max(0, len(ctx) - len(msgs))
     prefix = ctx[:prefix_len]
     suffix = ctx[prefix_len:]

tests/test_issue_branch_context_at_fork.py

@@ -31,4 +31,69 @@ def test_truncate_context_preserves_leading_compaction_row():
     ]
     out = truncate_context_for_display_keep(ctx, msgs, 2)
     assert len(out) == 3
-    assert out[0]["content"] == "compaction-ref-only"
+    assert out[0]["content"] == "compaction-ref-only"
+
+
+def test_truncate_context_for_display_keep_drops_interleaved_tool_tail():
+    msgs = [
+        {"role": "user", "content": "u1", "id": "u1", "timestamp": 1.0},
+        {"role": "assistant", "content": "a1", "id": "a1", "timestamp": 2.0},
+        {"role": "user", "content": "u2", "id": "u2", "timestamp": 3.0},
+        {"role": "assistant", "content": "a2", "id": "a2", "timestamp": 4.0},
+    ]
+    ctx = [
+        {"role": "user", "content": "u1", "id": "u1", "timestamp": 1.0},
+        {"role": "assistant", "content": "a1", "id": "a1", "timestamp": 2.0},
+        {"role": "user", "content": "u2", "id": "u2", "timestamp": 3.0},
+        {"role": "assistant", "content": "tool-call-after-cut", "id": "tool-call-after-cut"},
+        {"role": "tool", "content": "tool-result-after-cut", "id": "tool-result-after-cut"},
+        {"role": "assistant", "content": "a2", "id": "a2", "timestamp": 4.0},
+    ]
+    out = truncate_context_for_display_keep(ctx, msgs, 2)
+    contents = [m["content"] for m in out]
+    assert contents == ["u1", "a1"]
+    assert "u2" not in contents
+    assert "tool-call-after-cut" not in contents
+    assert "tool-result-after-cut" not in contents
+
+
+def test_truncate_context_for_display_keep_keeps_tool_rows_before_next_display_anchor():
+    msgs = [
+        {"role": "user", "content": "u1", "id": "u1", "timestamp": 1.0},
+        {"role": "assistant", "content": "a1", "id": "a1", "timestamp": 2.0},
+        {"role": "user", "content": "u2", "id": "u2", "timestamp": 3.0},
+        {"role": "assistant", "content": "a2", "id": "a2", "timestamp": 4.0},
+    ]
+    ctx = [
+        {"role": "user", "content": "u1", "id": "u1", "timestamp": 1.0},
+        {"role": "assistant", "content": "a1", "id": "a1", "timestamp": 2.0},
+        {"role": "assistant", "content": "assistant-think", "id": "a1-think"},
+        {"role": "assistant", "content": "assistant-tool-result", "id": "a1-tool", "timestamp": 2.5},
+        {"role": "user", "content": "u2", "id": "u2", "timestamp": 3.0},
+        {"role": "assistant", "content": "a2", "id": "a2", "timestamp": 4.0},
+    ]
+    out = truncate_context_for_display_keep(ctx, msgs, 2)
+    contents = [m["content"] for m in out]
+    assert contents == [
+        "u1",
+        "a1",
+        "assistant-think",
+        "assistant-tool-result",
+    ]
+    assert contents[-1] != "u2"
+
+
+def test_truncate_context_for_display_keep_prefers_compact_summary_fallback():
+    msgs = [
+        {"role": "user", "content": "u1", "id": "u1", "timestamp": 1.0},
+        {"role": "assistant", "content": "a1", "id": "a1", "timestamp": 2.0},
+        {"role": "user", "content": "u2", "id": "u2", "timestamp": 3.0},
+        {"role": "assistant", "content": "a2", "id": "a2", "timestamp": 4.0},
+    ]
+    ctx = [
+        {"role": "user", "content": "compact", "id": "summary"},
+        {"role": "user", "content": "u1", "id": "u1", "timestamp": 1.0},
+        {"role": "assistant", "content": "a1", "id": "a1", "timestamp": 2.0},
+    ]
+    out = truncate_context_for_display_keep(ctx, msgs, 2)
+    assert out == ctx[:2]