fix(streaming): clean cancelled user context chains

api/streaming.py

@@ -3846,9 +3846,14 @@ def _dedupe_replayed_context_messages(previous_context, result_messages, msg_tex
         ):
             boundary_idx = len(previous_context) - 1
             boundary_row = result_messages[boundary_idx]
-            is_stale_merge = (
+            is_stale_merge = bool(
                 previous_user_tail
-                and _detect_stale_user_merge(boundary_row, msg_text, previous_user_tail)
+                and _detect_stale_user_merge(
+                    boundary_row,
+                    msg_text,
+                    previous_user_tail,
+                    previous_context=previous_context,
+                )
             )
             if is_stale_merge or _looks_like_current_user_turn(boundary_row, msg_text):
                 if is_stale_merge:
@@ -3873,6 +3878,7 @@ def _dedupe_replayed_context_messages(previous_context, result_messages, msg_tex
             candidates,
             msg_text,
             previous_user_tail,
+            previous_context=previous_context,
         )
     candidates = _strip_replayed_prefix(previous_context, candidates)
     if candidates:
@@ -4030,40 +4036,131 @@ def _last_user_row(messages):
     return None
 
 
-def _detect_stale_user_merge(message, msg_text, previous_user_tail):
+def _stale_prefix_matches_prior_user_context(stale_prefix, stale_segments, previous_context):
+    """Return True when a stale prefix is explainable by prior user context.
+
+    First-generation repair usually produces segments matching consecutive
+    prior user rows. Once a session is already contaminated, later repair can
+    replay a stable stale prefix from an older polluted row even after newer
+    clean user turns have moved the context tail forward. Handle both shapes
+    while still requiring all evidence to come from prior user-role rows.
+    """
+    prior_rows = [
+        _stale_user_tail_candidate(msg)
+        for msg in previous_context or []
+    ]
+    prior_rows = [row for row in prior_rows if row]
+    if not prior_rows:
+        return False
+
+    if stale_segments:
+        segment_count = len(stale_segments)
+        for start in range(0, len(prior_rows) - segment_count + 1):
+            if prior_rows[start:start + segment_count] == stale_segments:
+                return True
+
+        # Already-polluted sessions may replay paragraphs from older polluted
+        # rows after newer clean turns have advanced the context tail. In that
+        # shape the stale paragraphs are still all prior user content, but they
+        # are substrings within older merged rows rather than standalone rows.
+        row_index = 0
+        row_offset = 0
+        matched_all_segments = True
+        for segment in stale_segments:
+            matched_segment = False
+            while row_index < len(prior_rows):
+                row = prior_rows[row_index]
+                pos = row.find(segment, row_offset)
+                if pos >= 0:
+                    row_offset = pos + len(segment)
+                    matched_segment = True
+                    break
+                row_index += 1
+                row_offset = 0
+            if not matched_segment:
+                matched_all_segments = False
+                break
+        if matched_all_segments:
+            return True
+
+    prefix_norm = _normalize_user_text(stale_prefix)
+    if not prefix_norm:
+        return False
+    for row in prior_rows:
+        if row == prefix_norm or row.startswith(f'{prefix_norm} '):
+            return True
+    return False
+
+
+def _detect_stale_user_merge(message, msg_text, previous_user_tail, previous_context=None):
     """Return True if `message` is the current user turn with a stale prefix merged in.
 
-    The agent's defensive repair path can concatenate the prior context-tail
+    The agent's defensive repair path can concatenate prior user context with
-    user row with the submitted current turn as ``<stale>\\n\\n<current>``.
+    the submitted current turn as ``<stale>\\n\\n<current>``. The stale portion
-    The literal ``\\n\\n`` boundary must survive into the comparison; a
+    can be either the immediate prior user tail or a replayed prefix from an
-    single-newline or space-only join is not the repair shape and must not
+    older already-polluted user row. The literal ``\\n\\n`` boundary must survive
-    match. Workspace sentinels may be present on either or both halves and
+    into the comparison; a single-newline or space-only join is not the repair
-    are stripped before comparison.
+    shape and must not match. Workspace sentinels may be present on either or
+    both halves and are stripped before comparison.
     """
     if not isinstance(message, dict) or message.get('role') != 'user':
         return False
     current_norm = _normalize_user_text(msg_text)
     if not current_norm:
         return False
-    tail_norm = _normalize_user_text(previous_user_tail) if previous_user_tail else ""
+
-    if not tail_norm or len(tail_norm) < 2:
+    merged = _raw_message_text(message.get('content', '')).replace("\r\n", "\n")
+    if "\n\n" not in merged:
         return False
-    merged = _raw_message_text(message.get('content', ''))
+
-    if not merged:
+    # The user's current turn can itself contain paragraph breaks. Find a repair
+    # boundary whose suffix normalizes to the *entire* submitted turn, then treat
+    # only the prefix as stale context. A plain split-and-last-segment check would
+    # miss ``<stale>\n\n<current paragraph A>\n\n<current paragraph B>``.
+    stale_segments = []
+    stale_prefix = ''
+    search_end = len(merged)
+    while search_end > 0:
+        boundary_idx = merged.rfind("\n\n", 0, search_end)
+        if boundary_idx < 0:
+            break
+        suffix = merged[boundary_idx + 2:]
+        if _normalize_user_text(suffix) == current_norm:
+            prefix = merged[:boundary_idx]
+            candidate_segments = [
+                _normalize_user_text(segment)
+                for segment in prefix.split("\n\n")
+            ]
+            if candidate_segments and all(candidate_segments):
+                stale_segments = candidate_segments
+                stale_prefix = prefix
+                break
+            # The suffix matched the submitted turn, but this boundary leaves
+            # blank prefix segments; try the next candidate boundary to the left.
+        search_end = boundary_idx
+    if not stale_segments:
         return False
-    # Require the literal \n\n repair boundary to be present after sentinel
+
-    # stripping. Check both the raw text and the workspace-stripped text so
+    if previous_context is not None and _stale_prefix_matches_prior_user_context(
-    # that workspace-prefixed rows on either or both halves still match.
+        stale_prefix,
-    for candidate in (merged, _strip_workspace_prefixes_for_compare(merged)):
+        stale_segments,
-        if '\n\n' not in candidate:
+        previous_context,
-            continue
+    ):
-        head, _, rest = candidate.partition('\n\n')
+        return True
-        if _normalize_user_text(head) == tail_norm and _normalize_user_text(rest) == current_norm:
+
-            return True
+    return bool(
-    return False
+        previous_context is None
+        and len(stale_segments) == 1
+        and _normalize_user_text(previous_user_tail) == stale_segments[0]
+    )
 
 
-def _strip_stale_user_merge_from_messages(messages, msg_text, previous_user_tail):
+def _strip_stale_user_merge_from_messages(
+    messages,
+    msg_text,
+    previous_user_tail,
+    previous_context=None,
+):
     """Return messages with stale-prefixed current user turns replaced by clean ones.
 
     Both context-merge (model-facing) and display-merge (visible transcript)
@@ -4075,7 +4172,12 @@ def _strip_stale_user_merge_from_messages(messages, msg_text, previous_user_tail
         return messages
     out = []
     for msg in messages:
-        if _detect_stale_user_merge(msg, msg_text, previous_user_tail):
+        if _detect_stale_user_merge(
+            msg,
+            msg_text,
+            previous_user_tail,
+            previous_context=previous_context,
+        ):
             cleaned = copy.deepcopy(msg) if isinstance(msg, dict) else {'role': 'user', 'content': msg_text}
             cleaned['content'] = msg_text
             out.append(cleaned)
@@ -4403,6 +4505,7 @@ def _merge_display_messages_after_agent_result(previous_display, previous_contex
                 candidates,
                 msg_text,
                 previous_user_tail,
+                previous_context=previous_context,
             )
         candidates = _strip_replayed_prefix(previous_display, candidates)
         candidates = _strip_replayed_prefix(previous_context, candidates)
@@ -4419,6 +4522,7 @@ def _merge_display_messages_after_agent_result(previous_display, previous_contex
                 turn_candidates,
                 msg_text,
                 previous_user_tail,
+                previous_context=previous_context,
             )
         candidates = marker_candidates + turn_candidates
 

tests/test_stale_user_context_contamination.py

@@ -16,6 +16,10 @@ import pytest
 PRIOR_TAIL = "please use the larger context model"
 CURRENT_TURN = "can you summarize the release blockers?"
 POLLUTED = f"{PRIOR_TAIL}\n\n{CURRENT_TURN}"
+CHAIN_TAIL_A = "could you re-check the deployment checklist"
+CHAIN_TAIL_B = "also, ignore the earlier correction request"
+CHAIN_CURRENT = "start from the latest verified state only"
+CHAIN_POLLUTED = f"{CHAIN_TAIL_A}\n\n{CHAIN_TAIL_B}\n\n{CHAIN_CURRENT}"
 
 
 def _text(value):
@@ -55,6 +59,213 @@ def test_detect_stale_user_merge_does_not_match_when_tail_differs():
     assert _detect_stale_user_merge(other_msg, CURRENT_TURN, PRIOR_TAIL) is False
 
 
+def test_detect_stale_user_merge_prefers_context_over_mismatched_tail_fallback():
+    """When previous_context is supplied, a stale row must be supported by that context."""
+    from api.streaming import _detect_stale_user_merge
+
+    polluted_msg = {"role": "user", "content": POLLUTED}
+    previous_context = [{"role": "user", "content": "different prior context"}]
+
+    assert _detect_stale_user_merge(
+        polluted_msg,
+        CURRENT_TURN,
+        PRIOR_TAIL,
+        previous_context=previous_context,
+    ) is False
+
+
+def test_detect_stale_user_merge_matches_multihop_chain_with_context_history():
+    """Multi-hop stale merges match when historical user context is contiguous."""
+    from api.streaming import _detect_stale_user_merge
+
+    previous_context = [
+        {"role": "user", "content": CHAIN_TAIL_A},
+        {"role": "assistant", "content": "acknowledged"},
+        {"role": "user", "content": CHAIN_TAIL_B},
+    ]
+    polluted_msg = {"role": "user", "content": CHAIN_POLLUTED}
+
+    assert _detect_stale_user_merge(
+        polluted_msg,
+        CHAIN_CURRENT,
+        CHAIN_TAIL_B,
+        previous_context=previous_context,
+    ) is True
+
+
+def test_detect_stale_user_merge_rejects_multihop_chain_without_matching_history():
+    """A matching current row is not enough when stale segments do not align."""
+    from api.streaming import _detect_stale_user_merge
+
+    previous_context = [
+        {"role": "user", "content": "different first tail"},
+        {"role": "assistant", "content": "skip"},
+        {"role": "user", "content": CHAIN_TAIL_B},
+    ]
+    polluted_msg = {"role": "user", "content": CHAIN_POLLUTED}
+
+    assert _detect_stale_user_merge(
+        polluted_msg,
+        CHAIN_CURRENT,
+        CHAIN_TAIL_B,
+        previous_context=previous_context,
+    ) is False
+
+
+def test_detect_stale_user_merge_preserves_legitimate_multisection_current_turn():
+    """A user-authored multi-section current prompt is not the stale repair shape."""
+    from api.streaming import _detect_stale_user_merge
+
+    previous_context = [
+        {"role": "user", "content": CHAIN_TAIL_A},
+        {"role": "assistant", "content": "acknowledged"},
+        {"role": "user", "content": CHAIN_TAIL_B},
+    ]
+    current_prompt = f"{CHAIN_TAIL_A}\n\n{CHAIN_TAIL_B}\n\n{CHAIN_CURRENT}"
+    current_msg = {"role": "user", "content": current_prompt}
+
+    assert _detect_stale_user_merge(
+        current_msg,
+        current_prompt,
+        CHAIN_TAIL_B,
+        previous_context=previous_context,
+    ) is False
+
+
+def test_detect_stale_user_merge_handles_single_hop_current_turn_with_paragraphs():
+    """One-hop stale repair still matches when the submitted turn has paragraphs."""
+    from api.streaming import _detect_stale_user_merge
+
+    current_prompt = "summarize this first\n\nthen list the risks"
+    polluted_msg = {"role": "user", "content": f"{PRIOR_TAIL}\n\n{current_prompt}"}
+
+    assert _detect_stale_user_merge(
+        polluted_msg,
+        current_prompt,
+        PRIOR_TAIL,
+    ) is True
+
+
+def test_detect_stale_user_merge_handles_multihop_current_turn_with_paragraphs():
+    """Multi-hop stale repair uses the full current-turn suffix, not the last paragraph."""
+    from api.streaming import _detect_stale_user_merge
+
+    previous_context = [
+        {"role": "user", "content": CHAIN_TAIL_A},
+        {"role": "assistant", "content": "acknowledged"},
+        {"role": "user", "content": CHAIN_TAIL_B},
+    ]
+    current_prompt = "summarize this first\n\nthen list the risks"
+    polluted_msg = {
+        "role": "user",
+        "content": f"{CHAIN_TAIL_A}\n\n{CHAIN_TAIL_B}\n\n{current_prompt}",
+    }
+
+    assert _detect_stale_user_merge(
+        polluted_msg,
+        current_prompt,
+        CHAIN_TAIL_B,
+        previous_context=previous_context,
+    ) is True
+
+
+def test_detect_stale_user_merge_handles_replayed_prefix_from_old_polluted_row():
+    """Already-contaminated sessions can replay an old prefix after newer clean turns."""
+    from api.streaming import _detect_stale_user_merge
+
+    stable_prefix = f"{CHAIN_TAIL_A}\n\n{CHAIN_TAIL_B}"
+    previous_context = [
+        {"role": "user", "content": f"{stable_prefix}\n\nold follow-up"},
+        {"role": "assistant", "content": "answered the old follow-up"},
+        {"role": "user", "content": "newer clean question after the polluted row"},
+        {"role": "assistant", "content": "answered the newer clean question"},
+    ]
+    current_turn = "latest question that should stand alone"
+    polluted_msg = {"role": "user", "content": f"{stable_prefix}\n\n{current_turn}"}
+
+    assert _detect_stale_user_merge(
+        polluted_msg,
+        current_turn,
+        "newer clean question after the polluted row",
+        previous_context=previous_context,
+    ) is True
+
+
+def test_detect_stale_user_merge_handles_segments_replayed_from_multiple_old_rows():
+    """Stale paragraphs can be assembled from earlier polluted rows, not only one row."""
+    from api.streaming import _detect_stale_user_merge
+
+    old_correction = "remove the reassurance sentence"
+    attachment_request = "send that with the attached image"
+    thread_check = "confirm it was sent in the right thread"
+    salary_question = "what is the average salary there"
+    current_turn = "should we mention the prior product by name?"
+    previous_context = [
+        {"role": "user", "content": f"{old_correction}\n\n{attachment_request}\n\nold current"},
+        {"role": "assistant", "content": "answered the old current"},
+        {"role": "user", "content": f"{thread_check}\n\n{salary_question}"},
+        {"role": "assistant", "content": "answered salary"},
+        {"role": "user", "content": "newer clean question"},
+    ]
+    polluted_msg = {
+        "role": "user",
+        "content": (
+            f"{old_correction}\n\n{attachment_request}\n\n{thread_check}\n\n"
+            f"{salary_question}\n\n{current_turn}"
+        ),
+    }
+
+    assert _detect_stale_user_merge(
+        polluted_msg,
+        current_turn,
+        "newer clean question",
+        previous_context=previous_context,
+    ) is True
+
+
+def test_detect_stale_user_merge_rejects_replayed_segments_in_wrong_order():
+    """Prior user substrings must explain stale paragraphs in chronological order."""
+    from api.streaming import _detect_stale_user_merge
+
+    first = "first stale paragraph"
+    second = "second stale paragraph"
+    current_turn = "latest clean question"
+    previous_context = [
+        {"role": "user", "content": second},
+        {"role": "assistant", "content": "answered second"},
+        {"role": "user", "content": first},
+    ]
+    polluted_msg = {"role": "user", "content": f"{first}\n\n{second}\n\n{current_turn}"}
+
+    assert _detect_stale_user_merge(
+        polluted_msg,
+        current_turn,
+        first,
+        previous_context=previous_context,
+    ) is False
+
+
+def test_detect_stale_user_merge_handles_prior_row_starting_with_stale_prefix():
+    """A stable stale prefix may be a leading subset of one older polluted row."""
+    from api.streaming import _detect_stale_user_merge
+
+    stable_prefix = f"{CHAIN_TAIL_A}\n\n{CHAIN_TAIL_B}"
+    current_turn = "latest question that should stand alone"
+    previous_context = [
+        {"role": "user", "content": f"{stable_prefix}\n\nolder different follow-up"},
+        {"role": "assistant", "content": "answered older follow-up"},
+        {"role": "user", "content": "newer clean question"},
+    ]
+    polluted_msg = {"role": "user", "content": f"{stable_prefix}\n\n{current_turn}"}
+
+    assert _detect_stale_user_merge(
+        polluted_msg,
+        current_turn,
+        "newer clean question",
+        previous_context=previous_context,
+    ) is True
+
+
 def test_detect_stale_user_merge_ignores_non_user_roles():
     """Only user rows are candidates for the repair-merge pattern."""
     from api.streaming import _detect_stale_user_merge
@@ -138,6 +349,31 @@ def test_strip_stale_user_merge_handles_list_content_row():
     assert cleaned[1] == messages[1]
 
 
+def test_strip_stale_user_merge_from_messages_replaces_multihop_polluted_row():
+    """Cleaner should replace a multi-hop polluted row with the clean current text."""
+    from api.streaming import _strip_stale_user_merge_from_messages
+
+    previous_context = [
+        {"role": "user", "content": CHAIN_TAIL_A},
+        {"role": "assistant", "content": "done"},
+        {"role": "user", "content": CHAIN_TAIL_B},
+    ]
+    messages = [
+        {"role": "user", "content": CHAIN_POLLUTED},
+        {"role": "assistant", "content": "pushing now"},
+    ]
+
+    cleaned = _strip_stale_user_merge_from_messages(
+        messages,
+        CHAIN_CURRENT,
+        CHAIN_TAIL_B,
+        previous_context=previous_context,
+    )
+
+    assert cleaned[0]["content"] == CHAIN_CURRENT
+    assert cleaned[1] == messages[1]
+
+
 def test_strip_stale_user_merge_does_not_touch_clean_rows():
     """Clean rows, assistant rows, and tool rows must pass through untouched."""
     from api.streaming import _strip_stale_user_merge_from_messages
@@ -255,6 +491,41 @@ def test_dedupe_replayed_context_handles_repair_replaced_tail_user_row():
     )
 
 
+def test_dedupe_replayed_context_handles_multihop_repair_replaced_tail_row():
+    """When repair replaces last tail row with a multi-hop merge, history is preserved."""
+    from api.streaming import _dedupe_replayed_context_messages
+
+    previous_context = [
+        {"role": "user", "content": CHAIN_TAIL_A},
+        {"role": "assistant", "content": "done"},
+        {"role": "user", "content": CHAIN_TAIL_B},
+    ]
+    result_messages = [
+        previous_context[0],
+        previous_context[1],
+        {"role": "user", "content": CHAIN_POLLUTED},
+        {"role": "assistant", "content": "pushing now"},
+    ]
+
+    cleaned = _dedupe_replayed_context_messages(
+        previous_context,
+        result_messages,
+        CHAIN_CURRENT,
+    )
+
+    assert [m.get("content") for m in cleaned] == [
+        CHAIN_TAIL_A,
+        "done",
+        CHAIN_TAIL_B,
+        CHAIN_CURRENT,
+        "pushing now",
+    ]
+    assert not any(
+        isinstance(m, dict) and CHAIN_POLLUTED in _text(m.get("content", ""))
+        for m in cleaned
+    )
+
+
 def test_merge_display_drops_polluted_current_when_eager_checkpoint_clean():
     """Display merge must not append a polluted row next to a clean eager checkpoint.
 
@@ -305,6 +576,94 @@ def test_merge_display_drops_polluted_current_when_eager_checkpoint_clean():
     ), "Assistant response must still be appended"
 
 
+def test_merge_display_drops_multihop_polluted_current_when_eager_checkpoint_clean():
+    """Multi-hop stale shape is normalized before eager checkpoint dedupe logic."""
+    from api.streaming import _merge_display_messages_after_agent_result
+
+    previous_display = [
+        {"role": "user", "content": CHAIN_TAIL_A},
+        {"role": "assistant", "content": "done"},
+        {"role": "user", "content": CHAIN_CURRENT},  # eager checkpoint
+    ]
+    previous_context = [
+        {"role": "user", "content": CHAIN_TAIL_A},
+        {"role": "assistant", "content": "done"},
+        {"role": "user", "content": CHAIN_TAIL_B},
+    ]
+    result_messages = [
+        *previous_context,
+        {"role": "user", "content": CHAIN_POLLUTED},
+        {"role": "assistant", "content": "pushing now"},
+    ]
+
+    merged = _merge_display_messages_after_agent_result(
+        previous_display,
+        previous_context,
+        result_messages,
+        CHAIN_CURRENT,
+    )
+
+    user_texts = [
+        _text(m.get("content", "")).strip()
+        for m in merged
+        if isinstance(m, dict) and m.get("role") == "user"
+    ]
+    assert user_texts.count(CHAIN_CURRENT) == 1, (
+        f"Should keep exactly one clean current user row; got user rows: {user_texts}"
+    )
+    assert not any(CHAIN_POLLUTED in t for t in user_texts), (
+        f"Multi-hop polluted row must be removed from display; got: {user_texts}"
+    )
+
+
+def test_merge_display_drops_replayed_old_prefix_after_newer_clean_turn():
+    """Visible transcript drops stale prefixes replayed from older polluted rows."""
+    from api.streaming import _merge_display_messages_after_agent_result
+
+    stable_prefix = f"{CHAIN_TAIL_A}\n\n{CHAIN_TAIL_B}"
+    newer_clean = "newer clean question after the polluted row"
+    current_turn = "latest question that should stand alone"
+    polluted_current = f"{stable_prefix}\n\n{current_turn}"
+    previous_display = [
+        {"role": "user", "content": f"{stable_prefix}\n\nold follow-up"},
+        {"role": "assistant", "content": "answered old follow-up"},
+        {"role": "user", "content": newer_clean},
+        {"role": "assistant", "content": "answered newer clean question"},
+        {"role": "user", "content": current_turn},  # eager checkpoint
+    ]
+    previous_context = previous_display[:-1]
+    result_messages = [
+        *previous_context,
+        {"role": "user", "content": polluted_current},
+        {"role": "assistant", "content": "latest answer"},
+    ]
+
+    merged = _merge_display_messages_after_agent_result(
+        previous_display,
+        previous_context,
+        result_messages,
+        current_turn,
+    )
+
+    user_texts = [
+        _text(m.get("content", "")).strip()
+        for m in merged
+        if isinstance(m, dict) and m.get("role") == "user"
+    ]
+    assert user_texts.count(current_turn) == 1, (
+        f"Should keep one clean current user row; got user rows: {user_texts}"
+    )
+    assert polluted_current not in user_texts, (
+        f"Replayed stale prefix must not be displayed as a user row; got: {user_texts}"
+    )
+    assert any(
+        isinstance(m, dict)
+        and m.get("role") == "assistant"
+        and "latest answer" in _text(m.get("content", ""))
+        for m in merged
+    ), "Assistant response must still be appended"
+
+
 def test_merge_display_does_not_overstrip_when_current_already_clean():
     """A legitimately new current turn that mentions the prior phrase stays intact."""
     from api.streaming import _merge_display_messages_after_agent_result