improvement: mcp server ui, auth, and rules

marimo/_ai/_tools/base.py

@@ -1,7 +1,6 @@
 # Copyright 2025 Marimo. All rights reserved.
 from __future__ import annotations
 
-import dataclasses
 import inspect
 import re
 from abc import ABC, abstractmethod
@@ -237,7 +236,7 @@ def as_backend_tool(
     # helpers
     def _coerce_args(self, args: Any) -> ArgsT:  # type: ignore[override]
         """If Args is a dataclass and args is a dict, construct it; else pass through."""
-        if dataclasses.is_dataclass(args):
+        if is_dataclass(args):
             # Already parsed
             return args  # type: ignore[return-value]
         return parse_raw(args, self.Args)

marimo/_ai/_tools/tools/rules.py

@@ -0,0 +1,63 @@
+# Copyright 2025 Marimo. All rights reserved.
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import marimo._utils.requests as requests
+from marimo import _loggers
+from marimo._ai._tools.base import ToolBase
+from marimo._ai._tools.types import EmptyArgs, SuccessResult
+
+LOGGER = _loggers.marimo_logger()
+
+# We load the rules remotely, so we can update these without requiring a new release.
+# If requested, we can bundle this into the library instead.
+MARIMO_RULES_URL = "https://docs.marimo.io/CLAUDE.md"
+
+
+@dataclass
+class GetMarimoRulesOutput(SuccessResult):
+    rules_content: Optional[str] = None
+    source_url: str = MARIMO_RULES_URL
+
+
+class GetMarimoRules(ToolBase[EmptyArgs, GetMarimoRulesOutput]):
+    """Get the official marimo rules and guidelines for AI assistants.
+
+    Returns:
+        The content of the rules file.
+    """
+
+    def handle(self, args: EmptyArgs) -> GetMarimoRulesOutput:
+        del args
+
+        try:
+            response = requests.get(MARIMO_RULES_URL, timeout=10)
+            response.raise_for_status()
+
+            return GetMarimoRulesOutput(
+                rules_content=response.text(),
+                source_url=MARIMO_RULES_URL,
+                next_steps=[
+                    "Follow the guidelines in the rules when working with marimo notebooks",
+                ],
+            )
+
+        except Exception as e:
+            LOGGER.warning(
+                "Failed to fetch marimo rules from %s: %s",
+                MARIMO_RULES_URL,
+                str(e),
+            )
+
+            return GetMarimoRulesOutput(
+                status="error",
+                message=f"Failed to fetch marimo rules: {str(e)}",
+                source_url=MARIMO_RULES_URL,
+                next_steps=[
+                    "Check internet connectivity",
+                    "Verify the rules URL is accessible",
+                    "Try again later if the service is temporarily unavailable",
+                ],
+            )

marimo/_ai/_tools/tools_registry.py

@@ -9,9 +9,11 @@
 from marimo._ai._tools.tools.datasource import GetDatabaseTables
 from marimo._ai._tools.tools.errors import GetNotebookErrors
 from marimo._ai._tools.tools.notebooks import GetActiveNotebooks
+from marimo._ai._tools.tools.rules import GetMarimoRules
 from marimo._ai._tools.tools.tables_and_variables import GetTablesAndVariables
 
 SUPPORTED_BACKEND_AND_MCP_TOOLS: list[type[ToolBase[Any, Any]]] = [
+    GetMarimoRules,
     GetActiveNotebooks,
     GetCellRuntimeData,
     GetLightweightCellMap,

marimo/_mcp/server/lifespan.py

@@ -4,7 +4,6 @@
 from typing import TYPE_CHECKING
 
 from marimo._loggers import marimo_logger
-from marimo._mcp.server.main import setup_mcp_server
 
 LOGGER = marimo_logger()
 
@@ -17,11 +16,15 @@ async def mcp_server_lifespan(app: "Starlette") -> AsyncIterator[None]:
     """Lifespan for MCP server functionality (exposing marimo as MCP server)."""
 
     try:
-        session_manager = setup_mcp_server(app)
+        mcp_app = app.state.mcp
+        if mcp_app is None:
+            LOGGER.warning("MCP server not found in app state")
+            yield
+            return
 
-        async with session_manager.run():
+        # Session manager owns request lifecycle during app run
+        async with mcp_app.session_manager.run():
             LOGGER.info("MCP server session manager started")
-            # Session manager owns request lifecycle during app run
             yield
 
     except ImportError as e:

marimo/_mcp/server/main.py

@@ -8,8 +8,6 @@
 
 from typing import TYPE_CHECKING
 
-from mcp.server.fastmcp import FastMCP
-
 from marimo._ai._tools.base import ToolContext
 from marimo._ai._tools.tools_registry import SUPPORTED_BACKEND_AND_MCP_TOOLS
 from marimo._loggers import marimo_logger
@@ -18,11 +16,10 @@
 
 
 if TYPE_CHECKING:
-    from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
     from starlette.applications import Starlette
 
 
-def setup_mcp_server(app: "Starlette") -> "StreamableHTTPSessionManager":
+def setup_mcp_server(app: "Starlette") -> None:
     """Create and configure MCP server for marimo integration.
 
     Args:
@@ -33,17 +30,20 @@ def setup_mcp_server(app: "Starlette") -> "StreamableHTTPSessionManager":
     Returns:
         StreamableHTTPSessionManager: MCP session manager
     """
+    from mcp.server.fastmcp import FastMCP
+    from starlette.middleware.base import BaseHTTPMiddleware
+    from starlette.responses import JSONResponse
     from starlette.routing import Mount
+    from starlette.types import Receive, Scope, Send
 
     mcp = FastMCP(
         "marimo-mcp-server",
         stateless_http=True,
         log_level="WARNING",
+        # Change base path from /mcp to /server
+        streamable_http_path="/server",
     )
 
-    # Change base path from /mcp to /server
-    mcp.settings.streamable_http_path = "/server"
-
     # Register all tools
     context = ToolContext(app=app)
     for tool in SUPPORTED_BACKEND_AND_MCP_TOOLS:
@@ -53,7 +53,23 @@ def setup_mcp_server(app: "Starlette") -> "StreamableHTTPSessionManager":
     # Initialize streamable HTTP app
     mcp_app = mcp.streamable_http_app()
 
+    # Middleware to require edit scope
+    class RequiresEditMiddleware(BaseHTTPMiddleware):
+        async def __call__(
+            self, scope: Scope, receive: Receive, send: Send
+        ) -> None:
+            auth = scope.get("auth")
+            if auth is None or "edit" not in auth.scopes:
+                response = JSONResponse(
+                    {"detail": "Forbidden"},
+                    status_code=403,
+                )
+                return await response(scope, receive, send)
+
+            return await self.app(scope, receive, send)
+
+    mcp_app.add_middleware(RequiresEditMiddleware)
+
     # Add to the top of the routes to avoid conflicts with other routes
     app.routes.insert(0, Mount("/mcp", mcp_app))
-
-    return mcp.session_manager
+    app.state.mcp = mcp

marimo/_server/print.py

@@ -163,14 +163,18 @@ def print_mcp_server(mcp_url: str, server_token: str | None) -> None:
     """Print MCP server configuration when MCP is enabled."""
     print_()
     print_tabbed(
-        f"{_utf8('🔗')} {green('Experimental MCP Server Configuration', bold=True)}"
+        f"{_utf8('🔗')} {green('Experimental MCP server configuration', bold=True)}"
     )
     print_tabbed(
-        f"{_utf8('➜')}  {green('MCP Server URL')}: {_colorized_url(mcp_url)}"
+        f"{_utf8('➜')}  {green('MCP server URL')}: {_colorized_url(mcp_url)}"
+    )
+    # Add to Claude code
+    print_tabbed(
+        f"{_utf8('➜')}  {green('Add to Claude Code')}: claude mcp add --transport http marimo {mcp_url}"
     )
     if server_token is not None:
         print_tabbed(
-            f"{_utf8('➜')}  {green('Add Header')}: Marimo-Server-Token: {muted(server_token)}"
+            f"{_utf8('➜')}  {green('Add header')}: Marimo-Server-Token: {muted(server_token)}"
         )
     print_()
 

marimo/_server/start.py

@@ -13,6 +13,7 @@
 from marimo._cli.print import echo
 from marimo._config.manager import get_default_config_manager
 from marimo._config.settings import GLOBAL_SETTINGS
+from marimo._mcp.server.main import setup_mcp_server
 from marimo._messaging.ops import StartupLogs
 from marimo._runtime.requests import SerializedCLIArgs
 from marimo._server.file_router import AppFileRouter
@@ -174,6 +175,16 @@ def start(
     Start the server.
     """
 
+    # Defaults when mcp is enabled
+    if mcp:
+        # Turn on watch mode
+        watch = True
+        # Turn off skew protection for MCP server
+        # since it is more convenient to connect to.
+        # Skew protection is not a security thing, but rather
+        # prevents connecting to old servers.
+        skew_protection = False
+
     # Find a free port if none is specified
     # if the user specifies a port, we don't try to find a free one
     port = port or find_free_port(DEFAULT_PORT, addr=host)
@@ -240,7 +251,9 @@ def start(
         *LIFESPAN_REGISTRY.get_all(),
     ]
 
-    if mcp and mode == SessionMode.EDIT:
+    mcp_enabled = mcp and mode == SessionMode.EDIT
+
+    if mcp_enabled:
         from marimo._mcp.server.lifespan import mcp_server_lifespan
 
         lifespans_list.append(mcp_server_lifespan)
@@ -260,6 +273,9 @@ def start(
         timeout=timeout,
     )
 
+    if mcp_enabled:
+        setup_mcp_server(app)
+
     app.state.port = external_port
     app.state.host = external_host
 

marimo/_utils/file_watcher.py

@@ -24,7 +24,7 @@ def create(path: Path, callback: Callback) -> FileWatcher:
             LOGGER.debug("Using watchdog file watcher")
             return _create_watchdog(path, callback, asyncio.get_event_loop())
         else:
-            LOGGER.warning(
+            LOGGER.info(
                 "watchdog is not installed, using polling file watcher"
             )
             return PollingFileWatcher(path, callback, asyncio.get_event_loop())

tests/_ai/tools/tools/test_rules.py

@@ -0,0 +1,79 @@
+# Copyright 2025 Marimo. All rights reserved.
+
+from __future__ import annotations
+
+from unittest.mock import Mock, patch
+
+import pytest
+
+from marimo._ai._tools.base import ToolContext
+from marimo._ai._tools.tools.rules import GetMarimoRules
+from marimo._ai._tools.types import EmptyArgs
+
+
+@pytest.fixture
+def tool() -> GetMarimoRules:
+    """Create a GetMarimoRules tool instance."""
+    return GetMarimoRules(ToolContext())
+
+
+def test_get_rules_success(tool: GetMarimoRules) -> None:
+    """Test successfully fetching marimo rules."""
+    mock_response = Mock()
+    mock_response.text.return_value = "# Marimo Rules\n\nTest content"
+    mock_response.raise_for_status = Mock()
+
+    with patch("marimo._utils.requests.get", return_value=mock_response):
+        result = tool.handle(EmptyArgs())
+
+    assert result.status == "success"
+    assert result.rules_content == "# Marimo Rules\n\nTest content"
+    assert result.source_url == "https://docs.marimo.io/CLAUDE.md"
+    assert len(result.next_steps) == 1
+    assert "Follow the guidelines" in result.next_steps[0]
+    mock_response.raise_for_status.assert_called_once()
+
+
+def test_get_rules_http_error(tool: GetMarimoRules) -> None:
+    """Test handling HTTP errors when fetching rules."""
+    mock_response = Mock()
+    mock_response.raise_for_status.side_effect = Exception("404 Not Found")
+
+    with patch("marimo._utils.requests.get", return_value=mock_response):
+        result = tool.handle(EmptyArgs())
+
+    assert result.status == "error"
+    assert result.rules_content is None
+    assert "Failed to fetch marimo rules" in result.message
+    assert "404 Not Found" in result.message
+    assert result.source_url == "https://docs.marimo.io/CLAUDE.md"
+    assert len(result.next_steps) == 3
+    assert "Check internet connectivity" in result.next_steps[0]
+
+
+def test_get_rules_network_error(tool: GetMarimoRules) -> None:
+    """Test handling network errors when fetching rules."""
+    with patch(
+        "marimo._utils.requests.get",
+        side_effect=Exception("Connection refused"),
+    ):
+        result = tool.handle(EmptyArgs())
+
+    assert result.status == "error"
+    assert result.rules_content is None
+    assert "Failed to fetch marimo rules" in result.message
+    assert "Connection refused" in result.message
+    assert len(result.next_steps) == 3
+
+
+def test_get_rules_timeout(tool: GetMarimoRules) -> None:
+    """Test handling timeout when fetching rules."""
+    with patch(
+        "marimo._utils.requests.get",
+        side_effect=Exception("Request timeout"),
+    ):
+        result = tool.handle(EmptyArgs())
+
+    assert result.status == "error"
+    assert result.rules_content is None
+    assert "Request timeout" in result.message