diff --git a/api/session_ops.py b/api/session_ops.py index 22ae9363..b4a3d166 100644 --- a/api/session_ops.py +++ b/api/session_ops.py @@ -109,9 +109,7 @@ 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(ctx) < len(msgs): + if len(ctx) <= len(msgs): return ctx[:keep] if len(msgs) == 0: return [] @@ -159,13 +157,20 @@ def truncate_context_for_display_keep( # 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] + last_kept = None if keep > 0: last_kept = matches[keep - 1] - if last_kept is not None: + first_unkept = matches[keep] + if first_unkept is not None: + if ( + last_kept is not None + and isinstance(msgs[keep - 1], dict) + and msgs[keep - 1].get('role') == 'user' + ): return ctx[:last_kept + 1] + return ctx[:first_unkept] + 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)) diff --git a/tests/test_issue_branch_context_at_fork.py b/tests/test_issue_branch_context_at_fork.py index 691aaefe..5c69faaf 100644 --- a/tests/test_issue_branch_context_at_fork.py +++ b/tests/test_issue_branch_context_at_fork.py @@ -83,6 +83,29 @@ def test_truncate_context_for_display_keep_keeps_tool_rows_before_next_display_a assert contents[-1] != "u2" +def test_truncate_context_for_display_keep_drops_unkept_tool_rows_after_user_boundary(): + 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-user-keep", "id": "tool-call-after-user-keep"}, + {"role": "tool", "content": "tool-result-after-user-keep", "id": "tool-result-after-user-keep"}, + {"role": "assistant", "content": "a2", "id": "a2", "timestamp": 4.0}, + ] + out = truncate_context_for_display_keep(ctx, msgs, 3) + contents = [m["content"] for m in out] + assert contents == ["u1", "a1", "u2"] + assert "tool-call-after-user-keep" not in contents + assert "tool-result-after-user-keep" not in contents + assert "a2" not in contents + + def test_truncate_context_for_display_keep_prefers_compact_summary_fallback(): msgs = [ {"role": "user", "content": "u1", "id": "u1", "timestamp": 1.0},