diff --git a/api/streaming.py b/api/streaming.py index d097fdb9..4d96c160 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -1441,18 +1441,130 @@ def _clean_synthetic_control_messages_with_provenance(messages): return _drop_synthetic_control_messages(raw_messages), has_verification_nudge -def _prepare_marker_clean_writeback(previous_context_messages, result_messages): +def _active_turn_checkpoint_identity(session, stream_id, msg_text): + """Capture the stream-owned pending turn before settlement mutates it.""" + pending_text = getattr(session, 'pending_user_message', None) + return { + 'stream_id': stream_id, + 'text': pending_text if pending_text is not None else msg_text, + 'timestamp': getattr(session, 'pending_started_at', None), + 'source': getattr(session, 'pending_user_source', None) or 'webui', + 'attachments': copy.deepcopy(getattr(session, 'pending_attachments', None) or []), + } + + +def _active_turn_checkpoint_matches(message, identity): + if not isinstance(identity, dict) or not identity.get('stream_id'): + return False + if identity.get('timestamp') in (None, ''): + return False + if not isinstance(message, dict) or message.get('role') != 'user': + return False + try: + return ( + _normalize_user_text(message.get('content')) + == _normalize_user_text(identity.get('text')) + and int(message.get('timestamp')) == int(identity.get('timestamp')) + and (message.get('_source') or 'webui') == (identity.get('source') or 'webui') + and list(message.get('attachments') or []) == list(identity.get('attachments') or []) + ) + except (TypeError, ValueError): + return False + + +def _active_turn_has_checkpoint(messages, identity): + return any(_active_turn_checkpoint_matches(message, identity) for message in messages or []) + + +def _materialize_active_turn_user(identity, msg_text, source): + message = { + 'role': 'user', + 'content': identity.get('text') if isinstance(identity, dict) else msg_text, + } + if isinstance(identity, dict): + if identity.get('timestamp') is not None: + message['timestamp'] = identity['timestamp'] + if identity.get('attachments'): + message['attachments'] = copy.deepcopy(identity['attachments']) + message['_source'] = identity.get('source') or source or 'webui' + else: + stamp_message_source(message, source) + return message + + +def _ensure_active_turn_boundary(previous_context, result_messages, identity, msg_text, source): + """Insert the pending turn before assistant/tool output when it is absent.""" + result_messages = list(result_messages or []) + if not result_messages or not isinstance(identity, dict): + return result_messages + if _active_turn_has_checkpoint(result_messages, identity): + return result_messages + current_user_key = _message_identity({'role': 'user', 'content': msg_text}) + current_start = len(list(previous_context or [])) if _messages_have_prefix( + result_messages, list(previous_context or []) + ) else 0 + if any( + isinstance(message, dict) + and message.get('role') == 'user' + and _message_identity(message) == current_user_key + and not _active_turn_checkpoint_matches(message, identity) + for message in result_messages[current_start:] + ): + return result_messages + if _messages_have_prefix(result_messages, list(previous_context or [])): + insert_at = len(list(previous_context or [])) + elif all( + _is_context_compression_marker(message) + or (isinstance(message, dict) and message.get('role') in ('assistant', 'tool')) + for message in result_messages + ): + insert_at = 0 + else: + return result_messages + return ( + result_messages[:insert_at] + + [_materialize_active_turn_user(identity, msg_text, source)] + + result_messages[insert_at:] + ) + + +def _align_active_turn_boundaries(previous_display, previous_context, identity): + """Make a context-only exact checkpoint visible before shared settlement.""" + display = list(previous_display or []) + context = list(previous_context or []) + if not isinstance(identity, dict): + return display, context + if _active_turn_has_checkpoint(display, identity): + return display, context + checkpoint = next( + (message for message in context if _active_turn_checkpoint_matches(message, identity)), + None, + ) + if checkpoint is not None: + display.append(copy.deepcopy(checkpoint)) + return display, context + + +def _prepare_marker_clean_writeback( + previous_context_messages, + result_messages, + active_turn_identity=None, +): """Return marker-cleaned rows, next context rows, and nudge provenance.""" cleaned, has_verification_nudge = _clean_synthetic_control_messages_with_provenance( result_messages ) + provenance = { + 'verification_nudge_seen': has_verification_nudge, + 'active_turn_identity': copy.deepcopy(active_turn_identity), + } if cleaned or not result_messages: return ( cleaned, _restore_reasoning_metadata(previous_context_messages, cleaned), - has_verification_nudge, + provenance, ) - return [], list(previous_context_messages or []), has_verification_nudge + return [], list(previous_context_messages or []), provenance def _current_turn_already_has_visible_assistant_answer(messages, *, current_user_text=None): @@ -5750,7 +5862,7 @@ def _merge_display_messages_after_agent_result( result_messages, msg_text, source: str = "webui", - verification_nudge_provenance: bool = False, + verification_nudge_provenance=None, ): """Keep UI transcript durable while allowing model context to compact. @@ -5800,12 +5912,20 @@ def _merge_display_messages_after_agent_result( previous_display = _deduped previous_context = list(previous_context or []) result_messages = list(result_messages or []) - verification_nudge_provenance = verification_nudge_provenance or any( + if isinstance(verification_nudge_provenance, dict): + _writeback_provenance = verification_nudge_provenance + else: + _writeback_provenance = { + 'verification_nudge_seen': bool(verification_nudge_provenance), + 'active_turn_identity': None, + } + _verification_nudge_seen = _writeback_provenance.get('verification_nudge_seen', False) or any( isinstance(message, dict) and message.get('role') == 'user' and _is_synthetic_control_message(message) for message in result_messages ) + _active_turn_identity = _writeback_provenance.get('active_turn_identity') # Same marker filter for the model-history inputs: the synthetic verify-loop # answer/nudge live in the agent's returned messages and prior context, and # would otherwise slip into the merged transcript as a real delta. (#5334) @@ -5993,19 +6113,32 @@ def _merge_display_messages_after_agent_result( _message_identity(m) == current_user_key or _looks_like_current_user_turn(m, msg_text) for m in candidates ) - current_user_already_checkpointed = bool( - merged - and ( - _message_identity(merged[-1]) == current_user_key - or _looks_like_current_user_turn(merged[-1], msg_text) + current_user_already_checkpointed = ( + _active_turn_has_checkpoint(merged, _active_turn_identity) + if _active_turn_identity + else bool( + merged + and ( + _message_identity(merged[-1]) == current_user_key + or _looks_like_current_user_turn(merged[-1], msg_text) + ) ) ) verification_user_already_checkpointed = bool( - verification_nudge_provenance - and msg_text - and any( - _looks_like_current_user_turn(_last_user_row(history), msg_text) - for history in (previous_display, previous_context) + _verification_nudge_seen + and ( + ( + _active_turn_identity + and _active_turn_has_checkpoint(previous_display, _active_turn_identity) + ) + or ( + not _active_turn_identity + and msg_text + and any( + _looks_like_current_user_turn(_last_user_row(history), msg_text) + for history in (previous_display, previous_context) + ) + ) ) ) if ( @@ -6024,8 +6157,9 @@ def _merge_display_messages_after_agent_result( # directly would make the assistant bubble appear attached to the prior # exchange and then clear the pending prompt. Materialize the current # turn at the transcript boundary before the assistant/tool response. - current_user_msg = {'role': 'user', 'content': msg_text} - stamp_message_source(current_user_msg, source) + current_user_msg = _materialize_active_turn_user( + _active_turn_identity, msg_text, source + ) insert_at = 0 while insert_at < len(candidates) and _is_context_compression_marker(candidates[insert_at]): insert_at += 1 @@ -7704,6 +7838,7 @@ def _run_agent_streaming( _turn_session_identity_tokens = _set_turn_session_identity(session_id) s = get_session(session_id) _turn_pending_source = getattr(s, 'pending_user_source', None) or 'webui' + _active_turn_identity = _active_turn_checkpoint_identity(s, stream_id, msg_text) update_active_run(stream_id, phase="running", session_id=session_id) s.workspace = str(Path(workspace).expanduser().resolve()) _last_persisted_model = None @@ -9319,6 +9454,7 @@ def _run_agent_streaming( ) = _prepare_marker_clean_writeback( _previous_context_messages, _result_messages, + _active_turn_identity, ) if cancel_event.is_set(): _finalize_cancelled_turn(s, ephemeral=False) @@ -9352,9 +9488,29 @@ def _run_agent_streaming( _next_context_messages, msg_text, ) + _next_context_messages = _ensure_active_turn_boundary( + _previous_context_messages, + _next_context_messages, + _active_turn_identity, + msg_text, + _turn_pending_source, + ) s.context_messages = _deduplicate_context_messages(_next_context_messages) - s.messages = _merge_display_messages_after_agent_result( + if _result_messages: + s.context_messages = _ensure_active_turn_boundary( + _previous_context_messages, + s.context_messages, + _active_turn_identity, + msg_text, + _turn_pending_source, + ) + _previous_display_for_writeback, _ = _align_active_turn_boundaries( _previous_messages, + s.context_messages, + _active_turn_identity, + ) + s.messages = _merge_display_messages_after_agent_result( + _previous_display_for_writeback, _previous_context_messages, _restore_display_reasoning_metadata(_previous_messages, _result_messages), msg_text, @@ -9702,6 +9858,7 @@ def _run_agent_streaming( ) = _prepare_marker_clean_writeback( _previous_context_messages, _result_messages, + _active_turn_identity, ) if _result_messages: # Mint ids on the shared result rows BEFORE dedupe @@ -9715,9 +9872,29 @@ def _run_agent_streaming( _next_context_messages, msg_text, ) + _next_context_messages = _ensure_active_turn_boundary( + _previous_context_messages, + _next_context_messages, + _active_turn_identity, + msg_text, + _turn_pending_source, + ) s.context_messages = _deduplicate_context_messages(_next_context_messages) - s.messages = _merge_display_messages_after_agent_result( + if _result_messages: + s.context_messages = _ensure_active_turn_boundary( + _previous_context_messages, + s.context_messages, + _active_turn_identity, + msg_text, + _turn_pending_source, + ) + _previous_display_for_writeback, _ = _align_active_turn_boundaries( _previous_messages, + s.context_messages, + _active_turn_identity, + ) + s.messages = _merge_display_messages_after_agent_result( + _previous_display_for_writeback, _previous_context_messages, _restore_reasoning_metadata(_previous_messages, _result_messages), msg_text, @@ -10932,6 +11109,7 @@ def _run_agent_streaming( ) = _prepare_marker_clean_writeback( _previous_context_messages, _result_messages, + _active_turn_identity, ) if _result_messages: # Mint ids on the shared result rows BEFORE dedupe @@ -10945,9 +11123,29 @@ def _run_agent_streaming( _next_context_messages, msg_text, ) + _next_context_messages = _ensure_active_turn_boundary( + _previous_context_messages, + _next_context_messages, + _active_turn_identity, + msg_text, + _turn_pending_source, + ) s.context_messages = _deduplicate_context_messages(_next_context_messages) - s.messages = _merge_display_messages_after_agent_result( + if _result_messages: + s.context_messages = _ensure_active_turn_boundary( + _previous_context_messages, + s.context_messages, + _active_turn_identity, + msg_text, + _turn_pending_source, + ) + _previous_display_for_writeback, _ = _align_active_turn_boundaries( _previous_messages, + s.context_messages, + _active_turn_identity, + ) + s.messages = _merge_display_messages_after_agent_result( + _previous_display_for_writeback, _previous_context_messages, _restore_reasoning_metadata(_previous_messages, _result_messages), msg_text, diff --git a/tests/test_issue6481_verification_evidence_phantom.py b/tests/test_issue6481_verification_evidence_phantom.py index 349b5340..3484fa05 100644 --- a/tests/test_issue6481_verification_evidence_phantom.py +++ b/tests/test_issue6481_verification_evidence_phantom.py @@ -39,10 +39,21 @@ def _verification_followup(marker, corrective=CORRECTIVE_ANSWER): ] -def _merge(result_messages, previous=None, msg_text="Fix the failing test."): +def _merge( + result_messages, + previous=None, + previous_context=None, + msg_text="Fix the failing test.", + writeback_provenance=None, +): previous = list(previous or [{"role": "user", "content": "Fix the failing test."}]) + previous_context = list(previous_context if previous_context is not None else previous) return streaming._merge_display_messages_after_agent_result( - previous, previous, result_messages, msg_text + previous, + previous_context, + result_messages, + msg_text, + verification_nudge_provenance=writeback_provenance, ) @@ -140,3 +151,97 @@ def test_unmarked_repeated_identical_prompt_materializes_new_user_row(): ("user", "Fix the failing test."), ("assistant", "The second attempt."), ] + + +def test_marked_repeated_old_prompt_materializes_deferred_current_turn(): + prompt = "Fix the failing test." + previous = [ + {"role": "user", "content": prompt, "timestamp": 1.0}, + {"role": "assistant", "content": "The earlier attempt is complete."}, + ] + corrective = "Verification failed. I fixed the parser and reran the tests." + + for marker in ("_verification_stop_synthetic", "_pre_verify_synthetic"): + merged = _merge( + previous + + [ + {"role": "user", "content": "[System: verify the workspace]", marker: True}, + {"role": "assistant", "content": corrective}, + ], + previous=previous, + msg_text=prompt, + writeback_provenance={ + "verification_nudge_seen": True, + "active_turn_identity": { + "stream_id": "direct-stream", + "text": prompt, + "timestamp": 2.0, + "source": "webui", + "attachments": [], + }, + }, + ) + + assert _role_content_sequence(merged) == [ + ("user", prompt), + ("assistant", "The earlier attempt is complete."), + ("user", prompt), + ("assistant", corrective), + ] + assert [message.get("content") for message in merged].count(prompt) == 2 + + +def test_context_only_exact_checkpoint_does_not_suppress_display_boundary(): + prompt = "Fix the failing test." + previous_display = [ + {"role": "user", "content": prompt, "timestamp": 1.0}, + {"role": "assistant", "content": "The earlier attempt is complete."}, + ] + previous_context = previous_display + [ + { + "role": "user", + "content": prompt, + "timestamp": 2.0, + "_source": "webui", + "attachments": [], + }, + ] + corrective = "Verification failed. I fixed the parser and reran the tests." + aligned_display, _ = streaming._align_active_turn_boundaries( + previous_display, + previous_context, + { + "stream_id": "direct-stream", + "text": prompt, + "timestamp": 2.0, + "source": "webui", + "attachments": [], + }, + ) + + merged = _merge( + [ + {"role": "user", "content": "[System: verify the workspace]", "_verification_stop_synthetic": True}, + {"role": "assistant", "content": corrective}, + ], + previous=aligned_display, + previous_context=previous_context, + msg_text=prompt, + writeback_provenance={ + "verification_nudge_seen": True, + "active_turn_identity": { + "stream_id": "direct-stream", + "text": prompt, + "timestamp": 2.0, + "source": "webui", + "attachments": [], + }, + }, + ) + + assert _role_content_sequence(merged) == [ + ("user", prompt), + ("assistant", "The earlier attempt is complete."), + ("user", prompt), + ("assistant", corrective), + ] diff --git a/tests/test_tool_limit_terminal_state.py b/tests/test_tool_limit_terminal_state.py index eeda1a8d..3cd2f448 100644 --- a/tests/test_tool_limit_terminal_state.py +++ b/tests/test_tool_limit_terminal_state.py @@ -22,6 +22,7 @@ def _run_streaming_with_fake_agent( prior_messages=None, prior_context_messages=None, msg_text="Do the long task.", + pending_started_at=1.0, ): session_dir = tmp_path / "sessions" session_dir.mkdir() @@ -43,17 +44,25 @@ def _run_streaming_with_fake_agent( session_id = "tool_limit_session" stream_id = "stream-tool-limit" + session_messages = list(prior_messages or []) + session_context_messages = list(prior_context_messages or []) + for history in (session_messages, session_context_messages): + for message in history: + if isinstance(message, dict) and message.get("role") == "user": + message.setdefault("timestamp", 1.0) + message.setdefault("_source", "webui") + message.setdefault("attachments", []) session = Session( session_id=session_id, title="Tool limit test", workspace=str(tmp_path), model="gpt-4o", - messages=list(prior_messages or []), - context_messages=list(prior_context_messages or []), + messages=session_messages, + context_messages=session_context_messages, ) session.active_stream_id = stream_id session.pending_user_message = msg_text - session.pending_started_at = 1.0 + session.pending_started_at = pending_started_at session.save() models.SESSIONS[session_id] = session streaming.SESSIONS[session_id] = session @@ -317,6 +326,97 @@ def test_verification_nudge_is_removed_while_corrective_followup_persists( ) == 1 +@pytest.mark.parametrize("marker", ["_verification_stop_synthetic", "_pre_verify_synthetic"]) +def test_deferred_repeated_old_prompt_marked_followup_materializes_current_turn( + tmp_path, + monkeypatch, + marker, +): + prompt = "Fix the failing test." + old_turn = [ + {"role": "user", "content": prompt, "timestamp": 1.0}, + {"role": "assistant", "content": "The earlier attempt is complete."}, + ] + corrective = "Verification failed. I fixed the parser and reran the tests." + result = { + "messages": old_turn + [ + {"role": "user", "content": "[System: verify the workspace]", marker: True}, + {"role": "assistant", "content": corrective}, + ], + } + + _events, payload = _run_streaming_with_fake_agent( + tmp_path, + monkeypatch, + result, + prior_messages=old_turn, + prior_context_messages=old_turn, + msg_text=prompt, + pending_started_at=2.0, + ) + + expected = [ + ("user", prompt), + ("assistant", "The earlier attempt is complete."), + ("user", prompt), + ("assistant", corrective), + ] + for field in ("messages", "context_messages"): + assert [(message["role"], message.get("content")) for message in payload[field]] == expected + assert [message.get("content") for message in payload[field]].count(prompt) == 2 + assert marker not in json.dumps(payload[field]) + + +@pytest.mark.parametrize("marker", ["_verification_stop_synthetic", "_pre_verify_synthetic"]) +def test_eager_exact_checkpoint_does_not_duplicate_repeated_prompt( + tmp_path, + monkeypatch, + marker, +): + prompt = "Fix the failing test." + old_turn = [ + {"role": "user", "content": prompt, "timestamp": 1.0}, + {"role": "assistant", "content": "The earlier attempt is complete."}, + ] + current_checkpoint = [ + { + "role": "user", + "content": prompt, + "timestamp": 2.0, + "_source": "webui", + "attachments": [], + }, + ] + corrective = "Verification failed. I fixed the parser and reran the tests." + result = { + "messages": [ + {"role": "user", "content": "[System: verify the workspace]", marker: True}, + {"role": "assistant", "content": corrective}, + ] + } + + _events, payload = _run_streaming_with_fake_agent( + tmp_path, + monkeypatch, + result, + prior_messages=old_turn, + prior_context_messages=old_turn + current_checkpoint, + msg_text=prompt, + pending_started_at=2.0, + ) + + expected = [ + ("user", prompt), + ("assistant", "The earlier attempt is complete."), + ("user", prompt), + ("assistant", corrective), + ] + for field in ("messages", "context_messages"): + assert [(message["role"], message.get("content")) for message in payload[field]] == expected + assert [message.get("content") for message in payload[field]].count(prompt) == 2 + assert marker not in json.dumps(payload[field]) + + @pytest.mark.parametrize("marker", ["_verification_stop_synthetic", "_pre_verify_synthetic"]) def test_synthetic_only_verification_delta_preserves_prior_histories(tmp_path, monkeypatch, marker): prior_turn = [