feat: Auto generate MCP tools through graphql introspection (#6397)

* feat: mcp autogen first pass

* fix: add medium sized tests

* fix: support nested payload fields

* fix: handle queries

* fix: collect files

* fix: further clean up

* chore: update lock

* fix: deps

* more clean up

* fix

* fix: additional logging

* Delete .claude/settings.local.json

Author: ravenac95

apps/frontend/app/api/v1/osograph/schema/graphql/system.graphql

@@ -22,38 +22,45 @@ type System {
 extend type Mutation {
   """
   System only. Mark a run as started.
+  @system-only
   """
   startRun(input: StartRunInput!): StartRunPayload!
 
   """
   System only. Update run metadata. This can be called at any time
+  @system-only
   """
   updateRunMetadata(input: UpdateRunMetadataInput!): UpdateRunMetadataPayload!
 
   """
   System only. Mark a run as finished.
+  @system-only
   """
   finishRun(input: FinishRunInput!): FinishRunPayload!
 
   """
   System only. Mark a step as started
+  @system-only
   """
   startStep(input: StartStepInput!): StartStepPayload!
 
   """
   System only. Mark a step as finished
+  @system-only
   """
   finishStep(input: FinishStepInput!): FinishStepPayload!
 
   """
   System only. Create a materialization for a step
+  @system-only
   """
   createMaterialization(
     input: CreateMaterializationInput!
   ): CreateMaterializationPayload!
 
   """
   System only. Save the generated published notebook HTML to object storage
+  @system-only
   """
   savePublishedNotebookHtml(
     input: SavePublishedNotebookHtmlInput!

uv.lock

@@ -7,7 +7,7 @@ requires-python = ">=3.10"
 authors = [{ name = "OSO Team", email = "opensource-observer@googlegroups.com" }]
 dependencies = [
     "arize-phoenix-otel>=0.10.1",
-    "arize-phoenix[evals]==11.28.0",
+    "arize-phoenix[evals]>=12.0.0",
     "discord-py>=2.5.2",
     "dotenv>=0.9.9",
     "llama-index>=0.12.29",
@@ -26,7 +26,7 @@ dependencies = [
     "pytest-asyncio>=0.26.0",
     "scikit-learn>=1.6.1",
     "sqlglot[rs]>=26.16.4",
-    "uvicorn[standard]==0.34.3",
+    "uvicorn[standard]>=0.35",
 ]
 
 [tool.uv.sources]

warehouse/oso_agent/pyproject.toml

@@ -0,0 +1,14 @@
+import logging
+import sys
+
+import pytest
+
+logger = logging.getLogger(__name__)
+
+
+@pytest.fixture(scope="session", autouse=True)
+def setup_debug_logging_for_tests() -> None:
+    root_logger = logging.getLogger(__name__.split(".")[0])
+    root_logger.setLevel(logging.DEBUG)
+
+    logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)

warehouse/oso_mcp/oso_mcp/conftest.py

@@ -0,0 +1,17 @@
+"""
+This should be considered a _test_ only module that is used to run the dev
+server for FastMCP
+"""
+
+from dotenv import load_dotenv
+from oso_core.logging import setup_module_logging
+from oso_mcp.server.app import setup_mcp_app
+from oso_mcp.server.config import MCPConfig
+
+load_dotenv()
+
+setup_module_logging("oso_mcp")
+
+config = MCPConfig()
+
+app = setup_mcp_app(config)

warehouse/oso_mcp/oso_mcp/serve.py

@@ -1,17 +1,23 @@
 import base64
 import hashlib
+import os
 import uuid
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from typing import Any, Generic, List, Optional, TypeVar, Union
 
 import requests
-from mcp.server.fastmcp import Context, FastMCP
+from fastmcp import Context, FastMCP
+from oso_mcp.server.graphql import generate_from_schema
+from oso_mcp.server.graphql.mutations import RegexMutationFilter
 from pyoso import Client, ClientConfig
 
 from .config import MCPConfig
 
+CURR_DIR = os.path.dirname(os.path.abspath(__file__))
+REPO_DIR = os.path.abspath(os.path.join(CURR_DIR, "../../../../"))
+
 MCP_SSE_PORT = 8000
 
 P = TypeVar("P")
@@ -67,10 +73,23 @@ def setup_mcp_app(config: MCPConfig):
         "OSO Data Lake Explorer",
         port=config.port,
         host=config.host,
-        dependencies=["pyoso", "python-dotenv", "requests"],
+        dependencies=["pyoso", "python-dotenv", "requests", "httpx"],
         lifespan=default_lifespan(config),
     )
 
+    # Wire autogenerated tools from graphql queries and mutations
+    generate_from_schema(
+        schema_path=os.path.join(
+            REPO_DIR, "apps/frontend/app/api/v1/osograph/schema/graphql"
+        ),
+        mcp=mcp,
+        config=config,
+        filters=[
+            RegexMutationFilter(patterns=["@mcp-ignore", "@system-only"]),
+        ],
+        client_schema_path=os.path.join(CURR_DIR, "graphql/queries"),
+    )
+
     @mcp.tool(
         description="Convert a natural language question into a SQL query using the OSO text2sql agent. Returns the generated SQL string.",
     )
@@ -118,9 +137,7 @@ async def query_text2sql_agent(nl_query: str, ctx: Context) -> McpResponse:
     @mcp.tool(
         description="Generates a deterministic OSO ID (SHA256 hash base64 encoded) from a list of input values. Use this to verify IDs in tests.",
     )
-    async def generate_oso_id(
-        args: List[Any], ctx: Context
-    ) -> McpResponse:
+    async def generate_oso_id(args: List[Any], ctx: Context) -> McpResponse:
         """
         Generates a deterministic OSO ID.
 

warehouse/oso_mcp/oso_mcp/server/app.py

@@ -1,5 +1,4 @@
 import logging
-import os
 import typing as t
 
 from pydantic import Field, SecretStr
@@ -22,24 +21,29 @@ class MCPConfig(BaseSettings):
     model_config = mcp_config_dict()
 
     oso_api_key: SecretStr = Field(
-        default_factory=lambda: SecretStr(
-            os.environ.get(
-                "COPILOT_MCP_OSO_API_KEY", os.environ.get("MCP_OSO_API_KEY", "")
-            )
-        ),
         description="API key for the OSO API",
         json_schema_extra={
             "required": True
         },  # This is the key to make the field required
     )
 
-    pyoso_base_url: str = Field(
-        default="https://www.opensource.observer",
+    oso_base_url: str = Field(
+        default="https://www.oso.xyz",
+        description="Base URL for the OSO API",
+    )
+
+    pyoso_base_path: str = Field(
+        default="/api/v1/",
         description="Base URL for the OSO pyoso client",
     )
 
-    text2sql_endpoint: str = Field(
-        default="https://www.opensource.observer/api/v1/text2sql",
+    graphql_path: str = Field(
+        default="/api/v1/osograph",
+        description="Path for the OSO GraphQL endpoint",
+    )
+
+    text2sql_path: str = Field(
+        default="/api/v1/text2sql",
         description="URL endpoint for the OSO text2sql service",
     )
 
@@ -136,3 +140,18 @@ def print_env_schema(cls):
                 print(f"{env_var}: OPTIONAL (DEFAULT={default_display})")
 
         print("=" * 50)
+
+    @property
+    def graphql_endpoint(self) -> str:
+        """Get the full GraphQL endpoint URL."""
+        return f"{self.oso_base_url}{self.graphql_path}"
+
+    @property
+    def text2sql_endpoint(self) -> str:
+        """Get the full Text2SQL endpoint URL."""
+        return f"{self.oso_base_url}{self.text2sql_path}"
+
+    @property
+    def pyoso_base_url(self) -> str:
+        """Get the full base URL for the pyoso client."""
+        return f"{self.oso_base_url}{self.pyoso_base_path}"

warehouse/oso_mcp/oso_mcp/server/config.py

@@ -0,0 +1,8 @@
+"""GraphQL tool generation for FastMCP.
+
+This package provides utilities to automatically generate FastMCP tools from GraphQL schemas.
+"""
+
+from .generator import generate_from_schema
+
+__all__ = ["generate_from_schema"]

warehouse/oso_mcp/oso_mcp/server/graphql/__init__.py

@@ -0,0 +1,141 @@
+"""Main GraphQL tool generator orchestrator."""
+
+import json
+import logging
+import typing as t
+from contextlib import asynccontextmanager
+
+import httpx
+from ariadne_codegen.schema import get_graphql_schema_from_path
+from fastmcp import FastMCP
+from oso_mcp.server.config import MCPConfig
+
+from .mutations import MutationExtractor
+from .pydantic_generator import PydanticModelGenerator
+from .queries import QueryDocumentParser, QueryExtractor
+from .tool_generator import ToolGenerator
+from .types import AsyncGraphQLClient, GraphQLClientFactory, MutationFilter
+
+logger = logging.getLogger(__name__)
+
+
+class OSOAsyncGraphQLClient(AsyncGraphQLClient):
+    """Asynchronous GraphQL client using httpx."""
+
+    def __init__(
+        self,
+        endpoint: str,
+        http_client: httpx.AsyncClient,
+        api_key: str,
+    ):
+        self.endpoint = endpoint
+        self.http_client = http_client
+        self.api_key = api_key
+
+    async def execute(
+        self,
+        query: str,
+        operation_name: str,
+        variables: dict[str, t.Any] | None = None,
+        headers: dict[str, str] | None = None,
+    ) -> t.Any:
+        """Execute a GraphQL query asynchronously.
+
+        Args:
+            query: GraphQL query string
+            variables: Optional variables for the query
+
+        Returns:
+            Parsed JSON response from the GraphQL server
+        """
+        payload: dict[str, t.Any] = {
+            "query": query,
+            "operationName": operation_name,
+        }
+        if variables:
+            payload["variables"] = variables
+        else:
+            payload["variables"] = {}
+
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+            "Agent": "oso-mcp-client/0.0",
+        }
+        logger.debug(f"Executing GraphQL request: \n\n {json.dumps(payload, indent=2)}")
+
+        # Make HTTP request
+        response = await self.http_client.post(
+            self.endpoint, json=payload, headers=headers
+        )
+        response.raise_for_status()
+        return response.json()
+
+
+def default_http_client_factory(config: MCPConfig) -> GraphQLClientFactory:
+    """Create a default HTTP client for GraphQL requests."""
+
+    @asynccontextmanager
+    async def _http_client_factory() -> t.AsyncGenerator[AsyncGraphQLClient, None]:
+        async with httpx.AsyncClient(timeout=httpx.Timeout(10.0, read=30.0)) as client:
+            yield OSOAsyncGraphQLClient(
+                endpoint=config.graphql_endpoint,
+                http_client=client,
+                api_key=config.oso_api_key.get_secret_value(),
+            )
+
+    return _http_client_factory
+
+
+def generate_from_schema(
+    schema_path: str,
+    mcp: FastMCP,
+    filters: list[MutationFilter],
+    config: MCPConfig,
+    client_schema_path: str | None = None,
+    graphql_client_factory: GraphQLClientFactory | None = None,
+) -> None:
+    """Generate and register FastMCP tools from GraphQL schema.
+
+    Args:
+        schema_path: Path to GraphQL schema directory or file
+        mcp: FastMCP instance to register tools on
+        config: Tool configuration
+        client_schema_path: Optional path to directory containing client
+                          GraphQL query files
+    """
+    # Load GraphQL schema
+    schema = get_graphql_schema_from_path(schema_path)
+
+    # Create Pydantic model generator
+    model_generator = PydanticModelGenerator()
+
+    # Extract mutations from schema
+    mutation_extractor = MutationExtractor(schema)
+    mutations = mutation_extractor.extract_mutations(model_generator, filters)
+
+    # Extract queries from client files if provided
+    queries = []
+    if client_schema_path:
+        # Parse client query files
+        parser = QueryDocumentParser(client_schema_path)
+        query_docs = parser.parse_all()
+
+        # Extract queries
+        query_extractor = QueryExtractor()
+        queries = query_extractor.extract_queries(schema, query_docs, model_generator)
+
+    if not graphql_client_factory:
+        graphql_client_factory = default_http_client_factory(config)
+
+    # Generate and register tools
+    tool_gen = ToolGenerator(
+        mcp,
+        mutations,
+        graphql_endpoint=config.graphql_endpoint,
+        graphql_client_factory=graphql_client_factory,
+        queries=queries,
+    )
+    tool_gen.generate_mutation_tools()
+    if queries:
+        tool_gen.generate_query_tools()