diff --git a/api/session_ops.py b/api/session_ops.py index d1a3241d..22ae9363 100644 --- a/api/session_ops.py +++ b/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:] diff --git a/tests/test_issue_branch_context_at_fork.py b/tests/test_issue_branch_context_at_fork.py index 0cbfbeb6..691aaefe 100644 --- a/tests/test_issue_branch_context_at_fork.py +++ b/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" \ No newline at end of file + 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]