diff --git a/src/clean_agent_preview.py b/src/clean_agent_preview.py index eeaac1239..d168657c1 100644 --- a/src/clean_agent_preview.py +++ b/src/clean_agent_preview.py @@ -4313,8 +4313,17 @@ async def stream_preview(*, endpoint_url, model, messages, headers, turn_contrac history[-1]['content'] = content yield event({'delta': content}) break + research_expansion_due = ( + broad_current_web_request(direct_user_text) + and successful_web_searches == 1 + and web_search_attempts < 2 + and not breadth_recovery_attempted + and not search_completion_attempted + and round_number < round_limit + ) if ( - requested_web_source_links(direct_user_text) + not research_expansion_due + and 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 @@ -4341,7 +4350,8 @@ async def stream_preview(*, endpoint_url, model, messages, headers, turn_contrac yield event({'type': 'completion_recovery', 'reason': 'requested_source_link_missing'}) continue if ( - contentless_final_response(content) + not research_expansion_due + and contentless_final_response(content) and answer_recovery_attempts == 0 and round_number < round_limit ): @@ -4359,14 +4369,7 @@ async def stream_preview(*, endpoint_url, model, messages, headers, turn_contrac }) yield event({'type': 'completion_recovery', 'reason': 'contentless_answer'}) continue - if ( - broad_current_web_request(direct_user_text) - and successful_web_searches == 1 - and web_search_attempts < 2 - and not breadth_recovery_attempted - and not search_completion_attempted - and round_number < round_limit - ): + if research_expansion_due: breadth_recovery_attempted = True force_web_search_next_round = True replace_streamed_draft_on_finish = True diff --git a/tests/test_clean_agent_preview.py b/tests/test_clean_agent_preview.py index e91ccf2b0..28d7e4b8d 100644 --- a/tests/test_clean_agent_preview.py +++ b/tests/test_clean_agent_preview.py @@ -1199,7 +1199,8 @@ async def test_stream_bounds_research_to_two_searches_fetch_then_synthesis(monke @pytest.mark.asyncio @pytest.mark.parametrize('embedded_article', [False, True]) @pytest.mark.parametrize('empty_second_search', [False, True]) -async def test_stream_retries_an_obviously_truncated_broad_web_answer(monkeypatch, embedded_article, empty_second_search): +@pytest.mark.parametrize('sources_requested', [False, True]) +async def test_stream_retries_an_obviously_truncated_broad_web_answer(monkeypatch, embedded_article, empty_second_search, sources_requested): """Broad current research expands, retrieves evidence, then synthesizes.""" import src.clean_agent_preview as module @@ -1286,7 +1287,7 @@ async def test_stream_retries_an_obviously_truncated_broad_web_answer(monkeypatc contract = resolve_full_inventory_contract(schemas=schemas, policy=ToolPolicy()) raw = [chunk async for chunk in stream_preview( endpoint_url='http://test', model='test', - messages=[{'role': 'user', 'content': 'Latest news in AI?'}], + messages=[{'role': 'user', 'content': 'Latest news in AI?' + (' Include source links.' if sources_requested else '')}], headers={}, turn_contract=contract, session_id='test', owner='test', disabled_tools=set(), tool_policy=ToolPolicy(), max_rounds=5, )]