diff --git a/api/streaming.py b/api/streaming.py index aceed689..cff54d24 100644 --- a/api/streaming.py +++ b/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 - user row with the submitted current turn as ``\\n\\n``. - The literal ``\\n\\n`` boundary must survive into the comparison; a - single-newline or space-only join is not the repair shape and must not - match. Workspace sentinels may be present on either or both halves and - are stripped before comparison. + The agent's defensive repair path can concatenate prior user context with + the submitted current turn as ``\\n\\n``. The stale portion + can be either the immediate prior user tail or a replayed prefix from an + older already-polluted user row. The literal ``\\n\\n`` boundary must survive + into the comparison; a single-newline or space-only join is not the repair + 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 ``\n\n\n\n``. + 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 - # that workspace-prefixed rows on either or both halves still match. - for candidate in (merged, _strip_workspace_prefixes_for_compare(merged)): - if '\n\n' not in candidate: - continue - head, _, rest = candidate.partition('\n\n') - if _normalize_user_text(head) == tail_norm and _normalize_user_text(rest) == current_norm: - return True - return False + + if previous_context is not None and _stale_prefix_matches_prior_user_context( + stale_prefix, + stale_segments, + previous_context, + ): + return True + + return bool( + 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 diff --git a/tests/test_stale_user_context_contamination.py b/tests/test_stale_user_context_contamination.py index 6fd24fbd..494343a4 100644 --- a/tests/test_stale_user_context_contamination.py +++ b/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