Notify connected agents of server restarts and hot-reloads so they can reload

Author: KoolADE85

dash/mcp/_server.py

@@ -8,6 +8,7 @@
 
 import json
 import logging
+import uuid
 from typing import TYPE_CHECKING, Any
 
 from flask import Response, request
@@ -51,6 +52,24 @@ def enable_mcp_server(app: Dash, mcp_path: str) -> None:
     app.mcp_decorated_functions = dict(MCP_DECORATED_FUNCTIONS)
     MCP_DECORATED_FUNCTIONS.clear()
 
+    _session_id: str | None = None
+
+    def _get_or_create_session_id() -> str:
+        """Read the hot-reload hash or generate a stable fallback."""
+        # pylint: disable=protected-access
+        reload_hash = app._hot_reload.hash
+        return reload_hash if reload_hash is not None else uuid.uuid4().hex
+
+    def _is_session_stale(client_session_id: str | None) -> bool:
+        """True when the client's session doesn't match or the hash changed."""
+        if client_session_id != _session_id:
+            return True
+        # pylint: disable=protected-access
+        reload_hash = app._hot_reload.hash
+        if reload_hash is None:
+            return False
+        return reload_hash != _session_id
+
     # -- Streamable HTTP endpoint --------------------------------------------
 
     def mcp_handler() -> Response:
@@ -75,6 +94,42 @@ def _handle_get() -> Response:
             status=405,
         )
 
+    def _check_session(method: str) -> bool:
+        """Validate the session header.
+
+        Raises ``ValueError`` when the header is missing.
+        Returns ``True`` when the session was stale and transparently
+        recovered, or ``False`` when the session is valid.
+        """
+        nonlocal _session_id
+        if method == "initialize":
+            _session_id = _get_or_create_session_id()
+            return False
+        client_session_id = request.headers.get("Mcp-Session-Id")
+        if _session_id is not None and not client_session_id:
+            raise ValueError("Missing Mcp-Session-Id header")
+        if _is_session_stale(client_session_id):
+            _session_id = _get_or_create_session_id()
+            logger.debug("MCP session recovered: %s", _session_id)
+            return True
+        return False
+
+    def _json_response(*messages: dict) -> Response:
+        """Wrap one or more JSON-RPC messages in a Flask Response.
+
+        A single message is serialised as a JSON object; multiple
+        messages are serialised as a JSON array.
+        """
+        body = messages[0] if len(messages) == 1 else list(messages)
+        resp = Response(
+            json.dumps(body),
+            content_type="application/json",
+            status=200,
+        )
+        if _session_id is not None:
+            resp.headers["Mcp-Session-Id"] = _session_id
+        return resp
+
     def _handle_post() -> Response:
         content_type = request.content_type or ""
         if "application/json" not in content_type:
@@ -92,16 +147,33 @@ def _handle_post() -> Response:
                 status=400,
             )
 
+        method = data.get("method", "")
+
+        try:
+            is_stale_session = _check_session(method)
+        except ValueError as err:
+            return Response(
+                json.dumps({"error": str(err)}),
+                content_type="application/json",
+                status=400,
+            )
+
         response_data = _process_mcp_message(data)
 
         if response_data is None:
             return Response("", status=202)
 
-        return Response(
-            json.dumps(response_data),
-            content_type="application/json",
-            status=200,
-        )
+        if is_stale_session:
+            return _json_response(
+                {"jsonrpc": "2.0", "method": "notifications/tools/list_changed"},
+                {
+                    "jsonrpc": "2.0",
+                    "method": "notifications/resources/list_changed",
+                },
+                response_data,
+            )
+
+        return _json_response(response_data)
 
     def _handle_delete() -> Response:
         # No sessions to terminate — server is stateless.
@@ -129,8 +201,8 @@ def _handle_initialize() -> InitializeResult:
     return InitializeResult(
         protocolVersion=LATEST_PROTOCOL_VERSION,
         capabilities=ServerCapabilities(
-            tools=ToolsCapability(listChanged=False),
-            resources=ResourcesCapability(),
+            tools=ToolsCapability(listChanged=True),
+            resources=ResourcesCapability(listChanged=True),
         ),
         serverInfo=Implementation(name="Plotly Dash", version=__version__),
         instructions=(

dash/mcp/primitives/tools/results/result_dataframe.py

@@ -47,7 +47,7 @@ class DataFrameResult(ResultFormatter):
     """Produce a markdown table for tabular component output values."""
 
     @classmethod
-    def format(cls, output: MCPOutput, returned_output_value: Any) -> list[TextContent]:
+    def format(cls, output: MCPOutput, returned_output_value: Any) -> list[TextContent]:  # type: ignore[override]
         if not TABULAR.matches(output.get("component_type"), output["property"]):
             return []
         if (

tests/integration/mcp/test_mcp_session.py

@@ -0,0 +1,128 @@
+"""MCP session lifecycle — end-to-end over a real Dash server.
+
+Exercises the full MCP session flow (initialize → operate → hot-reload
+recovery) against a live ``dash_duo`` server using real HTTP requests.
+Unit-level checks (status codes, header mechanics) live in
+``tests/unit/mcp/test_mcp_session.py``; these tests verify the broader
+behavioral contract.
+"""
+
+import requests
+
+from dash import Dash, Input, Output, html
+
+from tests.integration.mcp.conftest import _mcp_post
+
+
+def _mcp_post_with_session(
+    server_url, method, params=None, request_id=1, session_id=None
+):
+    """Like ``_mcp_post`` but forwards an ``Mcp-Session-Id`` header."""
+    headers = {"Content-Type": "application/json"}
+    if session_id is not None:
+        headers["Mcp-Session-Id"] = session_id
+    return requests.post(
+        f"{server_url}/_mcp",
+        json={
+            "jsonrpc": "2.0",
+            "method": method,
+            "id": request_id,
+            "params": params or {},
+        },
+        headers=headers,
+        timeout=5,
+    )
+
+
+def test_mcpse_e2e001_full_session_lifecycle(dash_duo):
+    """Initialize → tools/list → tools/call with session headers throughout."""
+    app = Dash(__name__)
+    app.layout = html.Div([html.Div(id="inp"), html.Div(id="out")])
+
+    @app.callback(Output("out", "children"), Input("inp", "children"))
+    def echo(v):
+        return f"echo: {v}"
+
+    dash_duo.start_server(app)
+    url = dash_duo.server.url
+
+    init = _mcp_post_with_session(url, "initialize")
+    assert init.status_code == 200
+    sid = init.headers.get("Mcp-Session-Id")
+    assert sid
+
+    notif = _mcp_post_with_session(
+        url, "notifications/initialized", session_id=sid, request_id=None
+    )
+    assert notif.status_code == 202
+
+    tools_resp = _mcp_post_with_session(url, "tools/list", session_id=sid, request_id=2)
+    assert tools_resp.status_code == 200
+    tools = tools_resp.json()["result"]["tools"]
+    assert any("echo" in t["name"] for t in tools)
+
+    tool_name = next(t["name"] for t in tools if "echo" in t["name"])
+    call_resp = _mcp_post_with_session(
+        url,
+        "tools/call",
+        params={"name": tool_name, "arguments": {"v": "hello"}},
+        session_id=sid,
+        request_id=3,
+    )
+    assert call_resp.status_code == 200
+    assert call_resp.headers.get("Mcp-Session-Id") == sid
+
+
+def test_mcpse_e2e002_stale_session_recovers_with_notifications(dash_duo):
+    """Simulate a hot-reload hash change and verify transparent recovery."""
+    app = Dash(__name__)
+    app.layout = html.Div([html.Div(id="inp"), html.Div(id="out")])
+
+    @app.callback(Output("out", "children"), Input("inp", "children"))
+    def echo(v):
+        return f"echo: {v}"
+
+    dash_duo.start_server(app)
+    url = dash_duo.server.url
+
+    app._hot_reload.hash = "original_hash"
+
+    init = _mcp_post_with_session(url, "initialize")
+    sid = init.headers["Mcp-Session-Id"]
+    assert sid == "original_hash"
+
+    resp = _mcp_post_with_session(url, "tools/list", session_id=sid, request_id=2)
+    assert resp.status_code == 200
+
+    app._hot_reload.hash = "new_hash"
+
+    resp = _mcp_post_with_session(url, "tools/list", session_id=sid, request_id=3)
+    assert resp.status_code == 200
+    new_sid = resp.headers["Mcp-Session-Id"]
+    assert new_sid == "new_hash"
+
+    data = resp.json()
+    assert isinstance(data, list)
+    assert len(data) == 3
+    assert data[0]["method"] == "notifications/tools/list_changed"
+    assert data[1]["method"] == "notifications/resources/list_changed"
+    assert "result" in data[2]
+    assert "tools" in data[2]["result"]
+
+    resp = _mcp_post_with_session(url, "tools/list", session_id=new_sid, request_id=4)
+    assert resp.status_code == 200
+    data = resp.json()
+    assert isinstance(data, dict)
+    assert "result" in data
+
+
+def test_mcpse_e2e003_capabilities_advertise_list_changed(dash_duo):
+    """Server capabilities include listChanged for tools and resources."""
+    app = Dash(__name__)
+    app.layout = html.Div(id="root")
+    dash_duo.start_server(app)
+
+    resp = _mcp_post(dash_duo.server.url, "initialize")
+    caps = resp.json()["result"]["capabilities"]
+    assert caps["tools"]["listChanged"] is True
+    assert caps["resources"]["listChanged"] is True