From acfdcf346c5adccb46f77740554f6984a663e7b1 Mon Sep 17 00:00:00 2001 From: Tatlatat Date: Tue, 2 Jun 2026 10:57:45 +0700 Subject: [PATCH] fix(agent): map native google_search and surface empty rounds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Models (notably Gemini) emit a native 'google_search' function call, but the agent loop had no mapping for it, so the call failed to convert, the round produced 0 chars and 0 tool blocks, and generation died silently — the web client hung on 'waiting for first token' with no error (also #443). - Map google_search / google_search_retrieval / google_search_grounding to the web_search tool, and read Gemini's 'queries' array (falling back to 'query'). - In stream_agent_loop, when a round yields no response text and no tool events, emit a visible fallback message instead of leaving the user hanging. - Give the unknown-tool execution branch an explicit exit_code=1 so the failure is logged as an error rather than 'n/a'. Unknown/unconvertible tool names still return None (unchanged) so they are dropped safely rather than executed. Added tests covering the google_search mapping, the queries array, and unknown/invalid-JSON returning None. --- src/agent_loop.py | 7 ++++ src/tool_execution.py | 2 +- src/tool_parsing.py | 3 ++ src/tool_schemas.py | 9 ++++- tests/test_unknown_tool_calls.py | 63 ++++++++++++++++++++++++++++++++ 5 files changed, 82 insertions(+), 2 deletions(-) create mode 100644 tests/test_unknown_tool_calls.py diff --git a/src/agent_loop.py b/src/agent_loop.py index 1b62f82..ad0ad00 100644 --- a/src/agent_loop.py +++ b/src/agent_loop.py @@ -2161,6 +2161,13 @@ async def stream_agent_loop( # Separator in accumulated response full_response += "\n\n" + # If the response is completely empty and no tools were executed, + # yield a fallback message so the user is not left hanging. + if not full_response.strip() and not tool_events: + _error_msg = "The model returned an empty response. Please try again or switch to a different model." + yield f'data: {json.dumps({"delta": _error_msg})}\n\n' + full_response = _error_msg + # --- Final metrics --- total_duration = time.time() - total_start metrics = _compute_final_metrics( diff --git a/src/tool_execution.py b/src/tool_execution.py index 33402ed..3f239b6 100644 --- a/src/tool_execution.py +++ b/src/tool_execution.py @@ -788,7 +788,7 @@ async def execute_tool_block( result = {"error": "MCP manager not available", "exit_code": 1} else: desc = f"unknown: {tool}" - result = {"error": f"Unknown tool type: {tool}"} + result = {"error": f"Unknown tool type: {tool}", "exit_code": 1} logger.info(f"Tool executed: {desc} -> exit_code={result.get('exit_code', 'n/a')}") return desc, result diff --git a/src/tool_parsing.py b/src/tool_parsing.py index 6d7aae3..cd9ecd1 100644 --- a/src/tool_parsing.py +++ b/src/tool_parsing.py @@ -95,6 +95,9 @@ _TOOL_NAME_MAP = { "search": "web_search", "web_search": "web_search", "websearch": "web_search", + "google_search": "web_search", + "google_search_retrieval": "web_search", + "google_search_grounding": "web_search", "web_fetch": "web_fetch", "webfetch": "web_fetch", "fetch_url": "web_fetch", diff --git a/src/tool_schemas.py b/src/tool_schemas.py index 8e10122..394e138 100644 --- a/src/tool_schemas.py +++ b/src/tool_schemas.py @@ -1074,6 +1074,7 @@ def function_call_to_tool_block(name: str, arguments: str) -> Optional[ToolBlock return None tool_type = _TOOL_NAME_MAP.get(name, name) + # Allow MCP tools through (namespaced as mcp__serverid__toolname) if tool_type.startswith("mcp__"): content = json.dumps(args) if args else "{}" @@ -1093,7 +1094,13 @@ def function_call_to_tool_block(name: str, arguments: str) -> Optional[ToolBlock elif tool_type == "python": content = args.get("code", "") elif tool_type == "web_search": - content = args.get("query", "") + queries = args.get("queries") + if isinstance(queries, list) and queries: + content = str(queries[0]) + elif queries: + content = str(queries) + else: + content = args.get("query", "") elif tool_type == "read_file": content = args.get("path", "") elif tool_type == "write_file": diff --git a/tests/test_unknown_tool_calls.py b/tests/test_unknown_tool_calls.py new file mode 100644 index 0000000..bf6e4b6 --- /dev/null +++ b/tests/test_unknown_tool_calls.py @@ -0,0 +1,63 @@ +import sys +from unittest.mock import MagicMock + +# Clean up any mocks from previous tests to ensure we load real modules +for mod in ['src.agent_tools', 'src.tool_parsing', 'src.tool_schemas', 'src.tool_execution']: + sys.modules.pop(mod, None) + +# Mock heavy database/model dependencies before importing +for mod in [ + 'sqlalchemy', 'sqlalchemy.orm', 'sqlalchemy.ext', 'sqlalchemy.ext.declarative', + 'sqlalchemy.ext.hybrid', 'sqlalchemy.sql', 'sqlalchemy.sql.expression', + 'src.database', 'core.models', 'core.database', 'core.auth' +]: + if mod not in sys.modules: + sys.modules[mod] = MagicMock() + +import pytest +import src.agent_tools +from src.tool_parsing import parse_tool_blocks +from src.tool_schemas import function_call_to_tool_block +from src.tool_execution import execute_tool_block +from types import SimpleNamespace + + +def test_parse_xml_unknown_tool_returns_none(): + """XML-style tags with truly unknown tools should be filtered out (return None).""" + text = 'value1' + blocks = parse_tool_blocks(text) + assert len(blocks) == 0 + + +def test_parse_tool_call_unknown_tool_returns_none(): + """[TOOL_CALL] blocks with truly unknown tools should be filtered out (return None).""" + text = '[TOOL_CALL] {tool => "mega_blast", command => "run energy"} [/TOOL_CALL]' + blocks = parse_tool_blocks(text) + assert len(blocks) == 0 + + +def test_function_call_to_tool_block_unknown_tool_returns_none(): + """Native function calls of truly unknown tools should return None.""" + block = function_call_to_tool_block("ultra_zap", '{"power": 9000}') + assert block is None + + +def test_function_call_to_tool_block_invalid_json_returns_none(): + """Unparseable JSON arguments should result in returning None.""" + block = function_call_to_tool_block("web_search", '{"query": "valid json') # invalid JSON + assert block is None + + +def test_google_search_mapping(): + """google_search should map to web_search and extract the first query from queries list or string.""" + # List of queries case + block = function_call_to_tool_block("google_search", '{"queries": ["testing google search"]}') + assert block is not None + assert block.tool_type == "web_search" + assert block.content == "testing google search" + + # Single string query case + block = function_call_to_tool_block("google_search_retrieval", '{"queries": "testing google search string"}') + assert block is not None + assert block.tool_type == "web_search" + assert block.content == "testing google search string"