fix(session): keep truncate rewind on complete turns (#5134)

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))

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