diff --git a/packages/datacommons-mcp/README.md b/packages/datacommons-mcp/README.md index 1701915..b49a761 100644 --- a/packages/datacommons-mcp/README.md +++ b/packages/datacommons-mcp/README.md @@ -48,6 +48,37 @@ This transport mode is intended for local integrations and is programmatically c uvx datacommons-mcp serve stdio ``` +### Optional documentation guidance + +Documentation guidance is disabled by default. HTTP clients, including clients +connecting to a local HTTP server, can opt in through their MCP connection configuration: + +```json +"headers": { + "X-DC-Enable-Documentation": "true" +} +``` + +Use `"false"` or omit the header to leave guidance disabled. +Values are case-insensitive and surrounding +whitespace is ignored; other values produce an invalid-parameters error. +Send the same preference on every request and reconnect after changing it. + +There is no server environment setting for this feature. Stdio connections do +not receive the optional documentation guidance; their tools and skills are unchanged. + +When enabled, a documentation routing hint follows the existing server +instructions and directs the client to fetch https://docs.datacommons.org/llms.txt +for relevant documentation questions. The client needs URL-reading capability +and permission to access the index and linked pages. The server does not fetch +the index or expose it as an MCP resource. When disabled, no documentation hint +is supplied; this does not prevent access to the public URL. Existing tools +and skills are unaffected. + +To customize the hint, place `doc_instructions_extension.md` in +`DC_INSTRUCTIONS_DIR`. It uses the same override mechanism as `server.md`, +which remains separate and always precedes the extension. + ## Clients You can use any MCP-enabled agent or client to connect to your running server. For example, see the [Data Commons MCP documentation](https://github.com/datacommonsorg/agent-toolkit/blob/main/docs/user_guide.md) for guides on connecting: diff --git a/packages/datacommons-mcp/datacommons_mcp/app.py b/packages/datacommons-mcp/datacommons_mcp/app.py index 6e64a20..5f2c0a8 100644 --- a/packages/datacommons-mcp/datacommons_mcp/app.py +++ b/packages/datacommons-mcp/datacommons_mcp/app.py @@ -27,6 +27,7 @@ from datacommons_mcp.client import AgentAPIClient from datacommons_mcp.data_models.settings import DCSettings +from datacommons_mcp.middleware import DocumentationMiddleware from datacommons_mcp.utils import read_external_content, read_package_content from datacommons_mcp.version import __version__ @@ -36,6 +37,7 @@ MCP_SERVER_NAME = "DC MCP Server" DEFAULT_INSTRUCTIONS_PACKAGE = "datacommons_mcp.instructions" SERVER_INSTRUCTIONS_FILE = "server.md" +DOCUMENTATION_INSTRUCTIONS_FILE = "doc_instructions_extension.md" class DCApp: @@ -67,7 +69,13 @@ def __init__(self) -> None: ) # Load Server Instructions - server_instructions = self._load_instructions(SERVER_INSTRUCTIONS_FILE) + base_instructions = self._load_instructions(SERVER_INSTRUCTIONS_FILE) + documentation_extension = self._load_instructions( + DOCUMENTATION_INSTRUCTIONS_FILE + ) + documentation_instructions = ( + f"{base_instructions.rstrip()}\n\n{documentation_extension}" + ) @asynccontextmanager async def lifespan(_server: FastMCP) -> AsyncIterator[dict[str, Any]]: @@ -79,9 +87,15 @@ async def lifespan(_server: FastMCP) -> AsyncIterator[dict[str, Any]]: self.mcp = FastMCP( MCP_SERVER_NAME, version=__version__, - instructions=server_instructions, + instructions=base_instructions, lifespan=lifespan, ) + self.mcp.add_middleware( + DocumentationMiddleware( + base_instructions=base_instructions, + documentation_instructions=documentation_instructions, + ) + ) def _load_instructions(self, filename: str) -> str: """Loads markdown content relative to the instructions directory. diff --git a/packages/datacommons-mcp/datacommons_mcp/instructions/doc_instructions_extension.md b/packages/datacommons-mcp/datacommons_mcp/instructions/doc_instructions_extension.md new file mode 100644 index 0000000..0afdb17 --- /dev/null +++ b/packages/datacommons-mcp/datacommons_mcp/instructions/doc_instructions_extension.md @@ -0,0 +1 @@ +For Data Commons API, client library, schema, dataset coverage, concept, or integration questions, fetch https://docs.datacommons.org/llms.txt and use its index to open only the documentation pages relevant to the question. For statistical data queries, use the MCP tools and skills instead. diff --git a/packages/datacommons-mcp/datacommons_mcp/middleware.py b/packages/datacommons-mcp/datacommons_mcp/middleware.py index 1dbf11a..e6b5357 100644 --- a/packages/datacommons-mcp/datacommons_mcp/middleware.py +++ b/packages/datacommons-mcp/datacommons_mcp/middleware.py @@ -1,6 +1,11 @@ import logging from collections.abc import Awaitable, Callable +from typing import Any +from fastmcp.server.dependencies import get_http_headers +from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext +from mcp import McpError +from mcp.types import INVALID_PARAMS, ErrorData from starlette.middleware.base import BaseHTTPMiddleware from starlette.requests import Request from starlette.responses import Response @@ -9,6 +14,52 @@ logger = logging.getLogger(__name__) +DOCUMENTATION_HEADER = "X-DC-Enable-Documentation" + + +class DocumentationMiddleware(Middleware): + """Enable documentation guidance only when the HTTP client opts in.""" + + def __init__( + self, *, base_instructions: str, documentation_instructions: str + ) -> None: + self._base_instructions = base_instructions + self._documentation_instructions = documentation_instructions + + def _enabled(self) -> bool: + value = get_http_headers().get(DOCUMENTATION_HEADER.lower()) + if value is None: + return False + value = value.strip().lower() + if value not in ("true", "false"): + raise McpError( + ErrorData( + code=INVALID_PARAMS, + message=f"{DOCUMENTATION_HEADER} must be true or false", + ) + ) + return value == "true" + + async def on_request(self, context: MiddlewareContext, call_next: CallNext) -> Any: # noqa: ANN401 + self._enabled() # Validate the preference on every request. + return await call_next(context) + + async def on_initialize( + self, context: MiddlewareContext, call_next: CallNext + ) -> Any: # noqa: ANN401 + instructions = ( + self._documentation_instructions + if self._enabled() + else self._base_instructions + ) + session = context.fastmcp_context.session + # FastMCP 3.4.2 sends initialization before middleware returns. Copy the + # session options before dispatch; never mutate the shared options. + session._init_options = session._init_options.model_copy( + update={"instructions": instructions} + ) + return await call_next(context) + class APIKeyMiddleware(BaseHTTPMiddleware): """Middleware to extract X-API-Key header and set it as the override API key diff --git a/packages/datacommons-mcp/datacommons_mcp/server.py b/packages/datacommons-mcp/datacommons_mcp/server.py index ad84c61..b6f41a4 100644 --- a/packages/datacommons-mcp/datacommons_mcp/server.py +++ b/packages/datacommons-mcp/datacommons_mcp/server.py @@ -30,7 +30,6 @@ # Configure logging logger = logging.getLogger(__name__) - # Expose the FastMCP instance for the CLI mcp = app.mcp diff --git a/packages/datacommons-mcp/pyproject.toml b/packages/datacommons-mcp/pyproject.toml index c8ea58f..24c9f45 100644 --- a/packages/datacommons-mcp/pyproject.toml +++ b/packages/datacommons-mcp/pyproject.toml @@ -7,6 +7,7 @@ requires-python = ">=3.11, <3.14" dependencies = [ "fastapi>=0.115.0", "uvicorn", + # Initialization middleware uses private session options; revalidate before upgrading. "fastmcp==3.4.2", "requests", "pydantic>=2.11.7", diff --git a/packages/datacommons-mcp/tests/server_test.py b/packages/datacommons-mcp/tests/server_test.py new file mode 100644 index 0000000..548196f --- /dev/null +++ b/packages/datacommons-mcp/tests/server_test.py @@ -0,0 +1,102 @@ +# Copyright 2026 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Protocol tests for client-controlled documentation guidance.""" + +import asyncio +from functools import partial + +import pytest +from datacommons_mcp.app import DCApp +from datacommons_mcp.middleware import DOCUMENTATION_HEADER, DocumentationMiddleware +from fastmcp import Client +from fastmcp.client.transports import StreamableHttpTransport +from fastmcp.utilities.tests import run_server_async +from mcp import McpError + + +@pytest.fixture +def documentation_app(monkeypatch, tmp_path, create_test_file): + create_test_file("server.md", "Custom server instructions.\n") + create_test_file("doc_instructions_extension.md", "Custom documentation hint.\n") + monkeypatch.setenv("DC_INSTRUCTIONS_DIR", str(tmp_path)) + return DCApp + + +@pytest.mark.asyncio +async def test_documentation_http_opt_in(documentation_app, monkeypatch): + app = documentation_app() + server = app.mcp + original = server.instructions + monkeypatch.setattr( + server, "run_http_async", partial(server.run_http_async, stateless_http=True) + ) + + @server.resource("test://unrelated") + def unrelated(): + return "unchanged" + + async with run_server_async(server) as url: + + async def check(header): + enabled = header is not None and header.lower() == "true" + headers = {} if header is None else {DOCUMENTATION_HEADER: header} + async with Client(StreamableHttpTransport(url, headers=headers)) as client: + result = await client.initialize() + expected = "Custom server instructions.\n" + if enabled: + expected = ( + "Custom server instructions.\n\nCustom documentation hint.\n" + ) + assert result.instructions == expected + for _ in range(2): + resources = await client.list_resources() + assert [str(r.uri) for r in resources] == ["test://unrelated"] + assert (await client.read_resource("test://unrelated"))[ + 0 + ].text == "unchanged" + + await asyncio.gather( + *(check(value) for value in [None, "true", "false", "TRUE"]) + ) + async with Client( + StreamableHttpTransport(url, headers={DOCUMENTATION_HEADER: "invalid"}), + auto_initialize=False, + ) as client: + with pytest.raises(McpError, match="must be true or false") as error: + await client.initialize() + assert error.value.error.code == -32602 + assert server.instructions == original + + +@pytest.mark.asyncio +async def test_documentation_without_http_header(documentation_app): + """Headerless transports receive only the base instructions.""" + app = documentation_app() + async with Client(app.mcp) as client: + result = await client.initialize() + expected = "Custom server instructions.\n" + assert result.instructions == expected + + +@pytest.mark.parametrize(("value", "expected"), [(" TRUE ", True), (" false ", False)]) +def test_documentation_header_whitespace(value, expected, monkeypatch): + monkeypatch.setattr( + "datacommons_mcp.middleware.get_http_headers", + lambda: {DOCUMENTATION_HEADER.lower(): value}, + ) + middleware = DocumentationMiddleware( + base_instructions="base", + documentation_instructions="base with documentation", + ) + assert middleware._enabled() is expected diff --git a/packages/datacommons-mcp/tests/test_app.py b/packages/datacommons-mcp/tests/test_app.py index f0926e8..398259c 100644 --- a/packages/datacommons-mcp/tests/test_app.py +++ b/packages/datacommons-mcp/tests/test_app.py @@ -66,6 +66,29 @@ def test_app_initialization_override( assert instructions == "Custom Server Instructions" +def test_app_prepares_documentation_instructions( + mock_settings, mock_fastmcp, tmp_path, create_test_file +): + """Prepare the packaged extension while leaving shared instructions unchanged.""" + custom_dir = tmp_path / "instructions" + create_test_file("instructions/server.md", "Custom Server Instructions") + mock_settings.return_value.instructions_dir = str(custom_dir) + + from datacommons_mcp.app import DCApp + + _ = DCApp() + + instructions = mock_fastmcp.call_args[1]["instructions"] + assert instructions == "Custom Server Instructions" + middleware = mock_fastmcp.return_value.add_middleware.call_args[0][0] + assert middleware._base_instructions == instructions + assert ( + "https://docs.datacommons.org/llms.txt" + in middleware._documentation_instructions + ) + assert middleware._documentation_instructions.startswith(f"{instructions}\n\n") + + def test_load_instruction_tool_override(mock_settings, tmp_path, create_test_file): """Test loading tool instructions with override.""" custom_dir = tmp_path / "instructions"