Refactor codebase_search to return error strings instead of dictionaries and add corresponding async test coverage.

Author: rodion-m

src/tests/test_search_tool.py

@@ -61,4 +61,89 @@ async def test_codebase_search_returns_xml_string(mock_get_api_key):
 
     # Verify it contains expected XML structure
     assert "<results>" in result, "Should contain results tag"
-    assert "<search_result" in result, "Should contain search_result tag"
+    assert "<search_result" in result, "Should contain search_result tag"
+
+
+@pytest.mark.asyncio
+async def test_codebase_search_empty_query_returns_error_string():
+    """Test that empty query returns an error string (not dict)."""
+    # Create mock context
+    ctx = MagicMock(spec=Context)
+    ctx.info = AsyncMock()
+    ctx.warning = AsyncMock()
+    ctx.error = AsyncMock()
+
+    # Create mock context with proper structure
+    mock_codealive_context = MagicMock()
+    mock_codealive_context.base_url = "https://app.codealive.ai"
+    ctx.request_context.lifespan_context = mock_codealive_context
+
+    # Call with empty query
+    result = await codebase_search(
+        ctx=ctx,
+        query="",
+        data_source_ids=["test_id"],
+        mode="auto",
+        include_content=False
+    )
+
+    # Verify result is a string (not a dict)
+    assert isinstance(result, str), "Error should be returned as a string"
+    assert "<error>" in result, "Error string should contain <error> tag"
+    assert "Query cannot be empty" in result, "Should contain error message"
+
+
+@pytest.mark.asyncio
+@patch('tools.search.get_api_key_from_context')
+async def test_codebase_search_api_error_returns_error_string(mock_get_api_key):
+    """Test that API errors return an error string (not dict)."""
+    import httpx
+
+    # Mock the API key function
+    mock_get_api_key.return_value = "test_key"
+
+    # Create mock context
+    ctx = MagicMock(spec=Context)
+    ctx.info = AsyncMock()
+    ctx.warning = AsyncMock()
+    ctx.error = AsyncMock()
+
+    # Create mock response that raises 404
+    mock_response = MagicMock()
+    mock_response.status_code = 404
+    mock_response.text = "Not found"
+
+    def raise_404():
+        raise httpx.HTTPStatusError(
+            "Not found",
+            request=MagicMock(),
+            response=mock_response
+        )
+
+    mock_response.raise_for_status = raise_404
+
+    # Create mock client that returns the error response
+    mock_client = AsyncMock()
+    mock_client.get.return_value = mock_response
+
+    # Create mock context with proper structure
+    mock_codealive_context = MagicMock()
+    mock_codealive_context.client = mock_client
+    mock_codealive_context.base_url = "https://app.codealive.ai"
+
+    ctx.request_context.lifespan_context = mock_codealive_context
+    ctx.request_context.headers = {"authorization": "Bearer test_key"}
+
+    # Call codebase_search
+    result = await codebase_search(
+        ctx=ctx,
+        query="test query",
+        data_source_ids=["invalid_id"],
+        mode="auto",
+        include_content=False
+    )
+
+    # Verify result is a string (not a dict)
+    assert isinstance(result, str), "Error should be returned as a string"
+    assert "<error>" in result, "Error string should contain <error> tag"
+    assert "404" in result, "Should contain error details"

src/tools/search.py

@@ -1,6 +1,6 @@
 """Search tool implementation."""
 
-from typing import Any, Dict, List, Optional, Union
+from typing import List, Optional, Union
 from urllib.parse import urljoin
 
 import httpx
@@ -16,7 +16,7 @@ async def codebase_search(
     data_source_ids: Optional[Union[str, List[str]]] = None,
     mode: str = "auto",
     include_content: bool = False
-) -> Union[str, Dict[str, Any]]:
+) -> str:
     """
     Use `codebase_search` tool to search for code in the codebase.
 
@@ -99,7 +99,7 @@ async def codebase_search(
 
     # Validate inputs
     if not query or not query.strip():
-        return {"error": "Query cannot be empty. Please provide a search term, function name, or description of the code you're looking for."}
+        return "<error>Query cannot be empty. Please provide a search term, function name, or description of the code you're looking for.</error>"
 
     if not data_source_ids or len(data_source_ids) == 0:
         await ctx.info("No data source IDs provided. If the API key has exactly one assigned data source, that will be used as default.")
@@ -161,4 +161,4 @@ async def codebase_search(
         error_msg = await handle_api_error(ctx, e, "code search")
         if isinstance(e, httpx.HTTPStatusError) and e.response.status_code == 404:
             error_msg = f"Error: Not found (404): One or more data sources could not be found. Check your data_source_ids."
-        return {"error": error_msg}
+        return f"<error>{error_msg}</error>"