Skip to content

refactor(narratives): type the MCP tool definitions and schema helpers - #544

Merged
nick-nlb merged 2 commits into
datacommonsorg:mainfrom
nick-nlb:narratives-mcp-tool-definitions
Oct 22, 2026
Merged

nick-nlb merged 2 commits into
datacommonsorg:mainfrom
nick-nlb:narratives-mcp-tool-definitions

Conversation

@nick-nlb

Copy link
Copy Markdown
Collaborator

Overview

The MCP tool definitions returned by tools/list were passed through the agent as dict[str, Any]: from the tool cache in mcp/client.py, into mcp/capabilities.py and /agent/health, through workflows/mcp_loop.py where they were converted into declaration dictionaries, and into gemini/client.py, which validated those dictionaries again on every request. The schema helpers in mcp/schema.py also took and returned dict[str, Any].

This change parses each tool into a frozen McpTool in the MCP client, builds types.FunctionDeclaration objects in the discovery loop, passes them to the Gemini client without re-validating them, and types the schema helpers with JsonObject and JsonValue. Behavior is unchanged: the declarations sent to Gemini and the /agent/health response are identical to main.

Related Issues

Changes Made

Changes:

  • MCP tool definitions (mcp/client.py, mcp/capabilities.py): Replaces the tool dictionaries with a frozen McpTool(name, description, input_schema) dataclass across the tool cache, async_get_tools, cached_tools, and capabilities.from_tools, and returns a JsonObject from Capabilities.describe().
  • Schema helpers (mcp/schema.py): Types transform_schema_for_gemini, _describe_value_shape, and fix_tool_arguments with JsonObject and JsonValue, narrowing nested schema values with isinstance.
  • Gemini function declarations (workflows/mcp_loop.py, gemini/client.py): Builds types.FunctionDeclaration objects in _gemini_tools and passes Sequence[types.FunctionDeclaration] directly to the Gemini client.
  • Unit tests (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 to McpTool and adds tests for the Gemini declaration conversion, the schema narrowing branches, and the /agent/health response.

Testing Done

  • pnpm test:agent (uv run pytest): 886 passed.
  • pnpm test:ui (vitest run): 26 files, 286 passed.
  • pnpm lint:ui (tsc --noEmit and biome check): passed.
  • pnpm lint:agent (ruff format --check, ruff check, and mypy --strict): passed across all 67 source files.
  • pnpm build:ui: passed.
  • pnpm lint:deploy: terraform fmt -check, pnpm run lint:shell, and pnpm run lint:docker passed (terraform validate fails with a pre-existing GCS backend configuration error on main; no deploy/ files were modified).

Checklist:

  • Unit tests passed
  • Integration tests passed (N/A: the agent has no integration suite; the live checks are under Manual)
  • Manual verification
  • Any updated goldens or fixtures were reviewed and are intentional (no goldens changed; the test fixtures were updated to McpTool)

Manual

Verified against public Data Commons (MCP_SERVER_URL=https://api.datacommons.org/mcp), comparing main and this branch:

  • Declarations and schemas: Fetched the live tools/list response (six tools) and ran it through both main and this branch. Confirmed that the transform_schema_for_gemini output for each tool and the model_dump(mode="json") of the Gemini request's tools config are identical.
  • Health endpoint: Fetched /agent/health before and after a chat turn on both branches and confirmed the responses are identical.
  • Chat turn: Sent "What is the population of France?" to POST /agent/chat/stream and confirmed that the turn called search_indicators and get_observations, streamed data, citations, a chart config, and follow-up questions, and completed normally.

Risk & Rollback

Agent, UI, and /agent/health behavior were not changed. The Gemini declarations are covered by test_each_tool_is_declared_with_its_transformed_schema in workflows/mcp_loop_test.py, and the /agent/health response is covered by test_health_reports_status_mcp_endpoint_and_tool_surface in server/routes/system_test.py. Rollback is reverting the PR.

Follow-ups

  • Agent loop rework

Checklist

  • I have read AGENTS.md and followed CODING_GUIDELINES.md, plus FRONTEND.md for UI changes.
  • I have run the app's lint, test, and build commands, as documented in that application's guide.
  • I have commented my code, particularly in hard-to-understand areas.
  • My changes generate no new warnings.

Note: Only Maintainers can approve and merge PRs. Expected initial review time: 3 business days.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread narratives/agent/src/narratives_agent/mcp/schema.py
@nick-nlb
nick-nlb marked this pull request as ready for review October 22, 2026 22:24

@juliawu juliawu left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you! Loving how readable all the typing makes the code. LGTM, with a few nits from gemini.

Comment thread narratives/agent/src/narratives_agent/workflows/mcp_loop.py Outdated
Comment thread narratives/agent/src/narratives_agent/workflows/mcp_loop_test.py
@nick-nlb
nick-nlb merged commit 5936d12 into datacommonsorg:main Oct 22, 2026
12 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants