Add MCP config to the .jupyter directory (#1385)

* Adding basic MCP config loader.

* Add caching to mcp config loading.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update MCP loader to use .jupyter.

* Add MCP config loader tests.

* Precommit fixes.

* Cleanup bad rebase.

* Precommit fixes.

* Fix handling of None return value in get_mcp_config.

* Precommit fixes.

* Move MCP config to persona manager and address review comments.

* Precommit fixes.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: ellisonbg

packages/jupyter-ai/jupyter_ai/mcp/__init__.py

@@ -0,0 +1,108 @@
+import json
+from pathlib import Path
+from typing import Any
+
+from jsonschema import ValidationError, validate
+from jsonschema.exceptions import SchemaError
+
+SCHEMA_FILE = Path(__file__).parent / "schema.json"
+
+
+class MCPConfigLoader:
+    """Loader for MCP server configuration files with JSON schema validation."""
+
+    def __init__(self):
+        # Load the schema from the schema.json file
+        with open(SCHEMA_FILE) as f:
+            self.schema = json.load(f)
+
+        # Cache for storing configurations and their modification times
+        # Key: config_path (str), Value: (config_dict, last_modified_time)
+        self._cache: dict[str, tuple[dict[str, Any], float]] = {}
+
+    def get_config(self, jupyter_dir: str) -> dict[str, Any]:
+        """
+        Read and validate an MCP server configuration file from .jupyter/mcp/config.json.
+
+        Args:
+            jupyter_dir (str): Path to the .jupyter directory
+
+        Returns:
+            Dict[str, Any]: The validated configuration object
+
+        Raises:
+            FileNotFoundError: If the config file doesn't exist
+            json.JSONDecodeError: If the JSON is malformed
+            ValidationError: If the config doesn't match the schema
+            SchemaError: If there's an issue with the schema itself
+        """
+        config_path = Path(jupyter_dir) / "mcp" / "config.json"
+        config_path_str = str(config_path)
+
+        # Check if file exists
+        if not config_path.exists():
+            raise FileNotFoundError(f"Configuration file not found: {config_path}")
+
+        # Get current file modification time
+        current_mtime = config_path.stat().st_mtime
+
+        # Check cache first
+        if config_path_str in self._cache:
+            cached_config, cached_mtime = self._cache[config_path_str]
+
+            # If file hasn't been modified, return cached version
+            if current_mtime == cached_mtime:
+                return cached_config
+
+        # File is new or has been modified, read and validate it
+        try:
+            with open(config_path) as f:
+                config = json.load(f)
+        except json.JSONDecodeError as e:
+            raise json.JSONDecodeError(
+                f"Invalid JSON in config file {config_path}: {e.msg}", e.doc, e.pos
+            )
+
+        # Validate against schema
+        try:
+            validate(instance=config, schema=self.schema)
+        except ValidationError as e:
+            raise ValidationError(f"Configuration validation failed: {e.message}")
+        except SchemaError as e:
+            raise SchemaError(f"Schema error: {e.message}")
+
+        # Cache the validated configuration and its modification time
+        self._cache[config_path_str] = (config, current_mtime)
+
+        return config
+
+    def validate_config(self, config: dict[str, Any]) -> bool:
+        """
+        Validate a configuration object against the schema.
+
+        Args:
+            config (Dict[str, Any]): Configuration object to validate
+
+        Returns:
+            bool: True if valid, raises exception if invalid
+        """
+        validate(instance=config, schema=self.schema)
+        return True
+
+    def clear_cache(self) -> None:
+        """
+        Clear the configuration cache.
+        """
+        self._cache.clear()
+
+    def get_cache_info(self) -> dict[str, Any]:
+        """
+        Get information about the cache.
+
+        Returns:
+            Dict[str, Any]: Dictionary with cache statistics
+        """
+        return {
+            "cached_files": len(self._cache),
+            "cache_keys": list(self._cache.keys()),
+        }

packages/jupyter-ai/jupyter_ai/mcp/mcp_config_loader.py

@@ -0,0 +1,53 @@
+{
+    "$schema": "http://json-schema.org/draft-07/schema#",
+    "title": "Extended MCP Server Configuration Schema",
+    "type": "object",
+    "properties": {
+      "mcpServers": {
+        "type": "object",
+        "patternProperties": {
+          "^[a-zA-Z0-9_-]+$": {
+            "type": "object",
+            "properties": {
+              "command": {
+                "type": "string",
+                "description": "Executable command for local stdio-based MCP server"
+              },
+              "args": {
+                "type": "array",
+                "description": "Command-line arguments passed to the executable",
+                "items": { "type": "string" }
+              },
+              "url": {
+                "type": "string",
+                "format": "uri",
+                "description": "Remote server URL (e.g. 'https://my-server.com/mcp')"
+              },
+              "transport": {
+                "type": "string",
+                "enum": ["stdio", "http", "streamable_http", "sse"],
+                "description": "Communication protocol between client and server"
+              },
+              "env": {
+                "type": "object",
+                "description": "Optional environment variables passed to the server",
+                "additionalProperties": { "type": "string" }
+              },
+              "disabled": {
+                "type": "boolean",
+                "description": "Optional flag to disable this server"
+              }
+            },
+            "anyOf": [
+              { "required": ["command", "args"] },
+              { "required": ["url"] }
+            ],
+            "additionalProperties": false
+          }
+        },
+        "additionalProperties": false
+      }
+    },
+    "required": ["mcpServers"],
+    "additionalProperties": false
+  }

packages/jupyter-ai/jupyter_ai/mcp/schema.json

@@ -341,6 +341,12 @@ def get_workspace_dir(self) -> str:
         """
         return self.parent.get_workspace_dir()
 
+    def get_mcp_config(self) -> dict[str, Any]:
+        """
+        Returns the MCP config for the current chat.
+        """
+        return self.parent.get_mcp_config()
+
 
 class GenerationInterrupted(asyncio.CancelledError):
     """Exception raised when streaming is cancelled by the user"""

packages/jupyter-ai/jupyter_ai/personas/base_persona.py

@@ -4,14 +4,15 @@
 import os
 from logging import Logger
 from time import time_ns
-from typing import TYPE_CHECKING, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from importlib_metadata import entry_points
 from jupyterlab_chat.models import Message
 from jupyterlab_chat.ychat import YChat
 from traitlets.config import LoggingConfigurable
 
 from ..config_manager import ConfigManager
+from ..mcp.mcp_config_loader import MCPConfigLoader
 from .base_persona import BasePersona
 from .directories import find_dot_dir, find_workspace_dir
 
@@ -37,6 +38,7 @@ class PersonaManager(LoggingConfigurable):
     fileid_manager: BaseFileIdManager
     root_dir: str
     event_loop: AbstractEventLoop
+    _mcp_config_loader: MCPConfigLoader
 
     log: Logger  # type: ignore
     """
@@ -78,6 +80,9 @@ def __init__(
         # Store file ID
         self.file_id = room_id.split(":")[2]
 
+        # Initialize MCP config loader
+        self._mcp_config_loader = MCPConfigLoader()
+
         # Load persona classes from entry points.
         # This is stored in a class attribute (global to all instances) because
         # the entry points are immutable after the server starts, so they only
@@ -273,3 +278,13 @@ def get_workspace_dir(self) -> str:
         Returns the path to the workspace directory for the current chat.
         """
         return find_workspace_dir(self.get_chat_dir(), root_dir=self.root_dir)
+
+    def get_mcp_config(self) -> dict[str, Any]:
+        """
+        Returns the MCP config for the current chat.
+        """
+        jdir = self.get_dotjupyter_dir()
+        if jdir is None:
+            return {}
+        else:
+            return self._mcp_config_loader.get_config(jdir)