diff --git a/src/clean_agent_preview.py b/src/clean_agent_preview.py index 01ec94713..17f8dde78 100644 --- a/src/clean_agent_preview.py +++ b/src/clean_agent_preview.py @@ -3249,6 +3249,8 @@ def requested_web_source_links(user_text): r'\b(?:return|give|show|include|provide|cite|find)\b.{0,35}\b(?:source\s+)?links?\b' r'|\b(?:\d+|one|two|three|four|five)\s+(?:official\s+)?(?:source\s+)?links?\b' r'|\bofficial\s+source\b' + r'|\b(?:with|include|provide|cite|show|give|find)\s+(?:the\s+)?(?:official\s+)?(?:sources|citations)\b' + r'|\b(?:find|locate|get|download)\b.{0,60}\bofficial\b.{0,60}\b(?:manual|guide|handbook|pdf|documentation)\b' r'|\b(?:find|locate|get|download)\b.{0,80}\b(?:manual|guide|handbook|pdf)\b.{0,40}\b(?:online|official)\b', str(user_text or ''), re.IGNORECASE, @@ -3959,6 +3961,7 @@ async def stream_preview(*, endpoint_url, model, messages, headers, turn_contrac search_completion_attempted = False budget_completion_attempted = False answer_recovery_attempts = 0 + citation_recovery_attempted = False force_no_tools_next_round = False force_web_search_next_round = ( broad_current_web_request(direct_user_text) and not native_workspace_enabled @@ -3970,7 +3973,7 @@ async def stream_preview(*, endpoint_url, model, messages, headers, turn_contrac official_source_retry_attempted = False note_search_recovery_attempted = False replace_streamed_draft_on_finish = False - buffer_completion_drafts = broad_current_web_request(direct_user_text) + buffer_completion_drafts = broad_current_web_request(direct_user_text) or requested_web_source_links(direct_user_text) usage_in = usage_out = 0 has_real_usage = False first_request_tokens = last_request_tokens = 0 @@ -4222,6 +4225,33 @@ async def stream_preview(*, endpoint_url, model, messages, headers, turn_contrac history[-1]['content'] = content yield event({'delta': content}) break + if ( + requested_web_source_links(direct_user_text) + and successful_web_searches + and not re.search(r'https?://\S+', content or '') + and not citation_recovery_attempted + and answer_recovery_attempts < 2 + and round_number < round_limit + ): + citation_recovery_attempted = True + answer_recovery_attempts += 1 + force_no_tools_next_round = True + replace_streamed_draft_on_finish = True + history.pop() + history.append({ + 'role': 'user', '_harness_control': True, + 'content': ( + 'The user explicitly requested a source or document link, but the ' + 'draft omitted it. Complete the answer using exact URLs already ' + 'present in the tool evidence. Choose only a URL that supports the ' + 'associated claim or requested document; do not invent a URL or ' + 'choose the first result merely because it is first. If the ' + 'requested source was not found, state that limitation plainly. ' + 'No additional tool call is needed for this completion check.' + ), + }) + yield event({'type': 'completion_recovery', 'reason': 'requested_source_link_missing'}) + continue if ( contentless_final_response(content) and answer_recovery_attempts == 0 diff --git a/tests/test_search_observation_budget.py b/tests/test_search_observation_budget.py index 2c497aab8..ff76d7a59 100644 --- a/tests/test_search_observation_budget.py +++ b/tests/test_search_observation_budget.py @@ -32,18 +32,22 @@ def test_external_edits_are_not_mistaken_for_inline_text(prompt): @pytest.mark.asyncio -async def test_runtime_does_not_append_unverified_search_result_as_citation(monkeypatch): +@pytest.mark.parametrize('repair_missing_link', [False, True]) +async def test_runtime_does_not_append_unverified_search_result_as_citation(monkeypatch, repair_missing_link): import src.clean_agent_preview as runtime from src.tool_schemas import FUNCTION_TOOL_SCHEMAS from src.tool_policy import ToolPolicy from src.turn_contract import resolve_full_inventory_contract answer = 'The retrieved page describes an older version; it does not establish the latest release.' - packets = iter([ + packets_list = [ {'choices': [{'delta': {'tool_calls': [{'index': 0, 'id': 'lookup', 'function': { 'name': 'web_search', 'arguments': '{"query":"latest Python official source"}', }}]}}]}, {'choices': [{'delta': {'content': answer}}]}, - ]) + ] + if repair_missing_link: + packets_list.append({'choices': [{'delta': {'content': answer + ' See the older release: https://python.org/old-release/'}}]}) + packets = iter(packets_list) class Response: def __init__(self, payload): self.payload = payload async def __aenter__(self): return self @@ -68,12 +72,17 @@ async def test_runtime_does_not_append_unverified_search_result_as_citation(monk endpoint_url='http://test', model='test', messages=[{'role': 'user', 'content': 'latest Python version? official source please'}], headers={}, turn_contract=contract, session_id='test', owner='test', - disabled_tools=set(), tool_policy=ToolPolicy(), max_rounds=3, + disabled_tools=set(), tool_policy=ToolPolicy(), max_rounds=3 if repair_missing_link else 2, )] events = [json.loads(chunk[6:]) for chunk in raw if '[DONE]' not in chunk] - final = ''.join(event.get('delta', '') for event in events) + finals = [event['content'] for event in events if event.get('type') == 'final_response'] + final = finals[-1] if finals else ''.join(event.get('delta', '') for event in events) assert answer in final - assert 'old-release' not in final + if repair_missing_link: + assert 'https://python.org/old-release/' in final + assert sum(event.get('reason') == 'requested_source_link_missing' for event in events) == 1 + else: + assert 'old-release' not in final assert '[Source:' not in final assert not any(event.get('type') == 'error' for event in events)