Update codebase_search to return XML string instead of dictionary.

Author: IvanBiruk

src/tests/test_search_tool.py

@@ -8,8 +8,8 @@
 
 @pytest.mark.asyncio
 @patch('tools.search.get_api_key_from_context')
-async def test_codebase_search_returns_dict(mock_get_api_key):
-    """Test that codebase_search returns a dictionary with structured_content."""
+async def test_codebase_search_returns_xml_string(mock_get_api_key):
+    """Test that codebase_search returns an XML string directly."""
     # Mock the API key function
     mock_get_api_key.return_value = "test_key"
 
@@ -56,15 +56,9 @@ async def test_codebase_search_returns_dict(mock_get_api_key):
         include_content=False
     )
 
-    # Verify result is a dictionary
-    assert isinstance(result, dict), "codebase_search should return a dictionary"
-
-    # Verify it has structured_content field
-    assert "structured_content" in result, "Result should have structured_content field"
-
-    # Verify the structured_content is a string (XML)
-    assert isinstance(result["structured_content"], str), "structured_content should be a string"
+    # Verify result is a string (XML)
+    assert isinstance(result, str), "codebase_search should return an XML string"
 
     # Verify it contains expected XML structure
-    assert "<results>" in result["structured_content"], "Should contain results tag"
-    assert "<search_result" in result["structured_content"], "Should contain search_result tag"
+    assert "<results>" in result, "Should contain results tag"
+    assert "<search_result" in result, "Should contain search_result tag"

src/tools/search.py

@@ -1,6 +1,6 @@
 """Search tool implementation."""
 
-from typing import Dict, List, Optional, Union
+from typing import Any, Dict, 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
-) -> Dict:
+) -> Union[str, Dict[str, Any]]:
     """
     Use `codebase_search` tool to search for code in the codebase.
 
@@ -154,8 +154,8 @@ async def codebase_search(
         # Transform the response to XML format for better LLM processing
         xml_content = transform_search_response_to_xml(search_results, include_content)
 
-        # Return as a dictionary with structured_content field for MCP compatibility
-        return {"structured_content": xml_content}
+        # Return XML string directly
+        return xml_content
 
     except (httpx.HTTPStatusError, Exception) as e:
         error_msg = await handle_api_error(ctx, e, "code search")