Repository navigation
refactor(narratives): type the MCP tool definitions and schema helpers - #544
Conversation
There was a problem hiding this comment.
Code Review
This pull request refactors the MCP (Model Context Protocol) tool handling to use strongly-typed domain objects (McpTool and Gemini's types.FunctionDeclaration) instead of raw untrusted dictionaries across module boundaries. It also enhances the robustness of schema transformations by skipping malformed keywords and handling validation errors. Feedback was provided regarding a potential ValidationError in transform_schema_for_gemini when a non-string description is copied from a non-map schema, which could fail the entire tool declaration loop; it is recommended to ensure description is only copied if it is a string.
juliawu
left a comment
There was a problem hiding this comment.
Thank you! Loving how readable all the typing makes the code. LGTM, with a few nits from gemini.
Overview
The MCP tool definitions returned by
tools/listwere passed through the agent asdict[str, Any]: from the tool cache inmcp/client.py, intomcp/capabilities.pyand/agent/health, throughworkflows/mcp_loop.pywhere they were converted into declaration dictionaries, and intogemini/client.py, which validated those dictionaries again on every request. The schema helpers inmcp/schema.pyalso took and returneddict[str, Any].This change parses each tool into a frozen
McpToolin the MCP client, buildstypes.FunctionDeclarationobjects in the discovery loop, passes them to the Gemini client without re-validating them, and types the schema helpers withJsonObjectandJsonValue. Behavior is unchanged: the declarations sent to Gemini and the/agent/healthresponse are identical tomain.Related Issues
Changes Made
Changes:
mcp/client.py,mcp/capabilities.py): Replaces the tool dictionaries with a frozenMcpTool(name, description, input_schema)dataclass across the tool cache,async_get_tools,cached_tools, andcapabilities.from_tools, and returns aJsonObjectfromCapabilities.describe().mcp/schema.py): Typestransform_schema_for_gemini,_describe_value_shape, andfix_tool_argumentswithJsonObjectandJsonValue, narrowing nested schema values withisinstance.workflows/mcp_loop.py,gemini/client.py): Buildstypes.FunctionDeclarationobjects in_gemini_toolsand passesSequence[types.FunctionDeclaration]directly to the Gemini client.mcp/client_test.py,mcp/capabilities_test.py,mcp/schema_test.py,server/routes/system_test.py,workflows/mcp_loop_test.py,workflows/chat_pipeline_test.py,gemini/client_test.py): Updates the tool fixtures toMcpTooland adds tests for the Gemini declaration conversion, the schema narrowing branches, and the/agent/healthresponse.Testing Done
pnpm test:agent(uv run pytest): 886 passed.pnpm test:ui(vitest run): 26 files, 286 passed.pnpm lint:ui(tsc --noEmitandbiome check): passed.pnpm lint:agent(ruff format --check,ruff check, andmypy --strict): passed across all 67 source files.pnpm build:ui: passed.pnpm lint:deploy:terraform fmt -check,pnpm run lint:shell, andpnpm run lint:dockerpassed (terraform validatefails with a pre-existing GCS backend configuration error onmain; nodeploy/files were modified).Checklist:
McpTool)Manual
Verified against public Data Commons (
MCP_SERVER_URL=https://api.datacommons.org/mcp), comparingmainand this branch:tools/listresponse (six tools) and ran it through bothmainand this branch. Confirmed that thetransform_schema_for_geminioutput for each tool and themodel_dump(mode="json")of the Gemini request'stoolsconfig are identical./agent/healthbefore and after a chat turn on both branches and confirmed the responses are identical.POST /agent/chat/streamand confirmed that the turn calledsearch_indicatorsandget_observations, streamed data, citations, a chart config, and follow-up questions, and completed normally.Risk & Rollback
Agent, UI, and
/agent/healthbehavior were not changed. The Gemini declarations are covered bytest_each_tool_is_declared_with_its_transformed_schemainworkflows/mcp_loop_test.py, and the/agent/healthresponse is covered bytest_health_reports_status_mcp_endpoint_and_tool_surfaceinserver/routes/system_test.py. Rollback is reverting the PR.Follow-ups
Checklist
AGENTS.mdand followedCODING_GUIDELINES.md, plusFRONTEND.mdfor UI changes.Note: Only Maintainers can approve and merge PRs. Expected initial review time: 3 business days.