diff --git a/notebook_intelligence/api.py b/notebook_intelligence/api.py index 001bd79..1009eee 100644 --- a/notebook_intelligence/api.py +++ b/notebook_intelligence/api.py @@ -389,11 +389,13 @@ def _on_ui_command_response(data: dict): response.run_ui_command_response_signal.connect(_on_ui_command_response) - while True: - if resp["result"] is not None: - response.run_ui_command_response_signal.disconnect(_on_ui_command_response) - return resp["result"] - await asyncio.sleep(0.1) + try: + while True: + if resp["result"] is not None: + return resp["result"] + await asyncio.sleep(0.1) + finally: + response.run_ui_command_response_signal.disconnect(_on_ui_command_response) @dataclass class ToolPreInvokeResponse: diff --git a/notebook_intelligence/claude.py b/notebook_intelligence/claude.py index bf3de8d..02c17e6 100644 --- a/notebook_intelligence/claude.py +++ b/notebook_intelligence/claude.py @@ -4,6 +4,7 @@ import difflib import os import sys +import secrets import asyncio from enum import Enum from pathlib import Path @@ -17,6 +18,7 @@ from notebook_intelligence.api import AskUserQuestionData, BackendMessageType, CancelToken, ChatCommand, ChatModel, ChatRequest, ChatResponse, ClaudeToolType, CompletionContext, ConfirmationData, Host, InlineCompletionModel, MarkdownData, ProgressData, SignalImpl, ToolCallData from notebook_intelligence.base_chat_participant import BaseChatParticipant from notebook_intelligence.claude_sessions import CONTROL_SLASH_COMMANDS +from notebook_intelligence.feature_flags import is_external_ui_tools_active from notebook_intelligence._version import __version__ as NBI_VERSION import base64 import logging @@ -1875,6 +1877,77 @@ async def open_file_in_jupyter_ui(args) -> str: is_error=True, ) + +# The Jupyter-UI tools above are exposed to Claude Code as an EXTERNAL stdio MCP +# server (notebook_intelligence.mcp_ui_proxy) instead of an in-process sdk server, +# so they survive a managed/enterprise MCP config that forbids dynamically +# configured servers. The proxy fetches this manifest over HTTP and forwards each +# call to invoke_ui_tool, which runs the same handlers against the live chat turn. +JUPYTER_UI_TOOLS = [ + list_available_notebook_kernels, create_new_notebook, add_markdown_cell, + add_code_cell, get_number_of_cells, get_cell_type_and_source, get_cell_output, + set_cell_type_and_source, delete_cell, insert_cell, run_cell, save_notebook, + rename_notebook, run_command_in_jupyter_terminal, open_file_in_jupyter_ui, +] + +_PY_TYPE_TO_JSON = {str: "string", int: "integer", float: "number", bool: "boolean"} + + +def _tool_json_schema(input_schema) -> dict: + """Mirror claude_agent_sdk.create_sdk_mcp_server's schema build so the manifest + is identical to the former in-process server (dict params, all required).""" + if isinstance(input_schema, dict): + if "type" in input_schema and "properties" in input_schema and isinstance(input_schema["type"], str): + return input_schema + props = {name: {"type": _PY_TYPE_TO_JSON.get(t, "string")} for name, t in input_schema.items()} + return {"type": "object", "properties": props, "required": list(props)} + return {"type": "object", "properties": {}} + + +# Per-process shared secret proving a request came from the mcp_ui_proxy NBI spawned. +# The proxy sends it to UIToolsHandler in a dedicated header (X-NBI-UI-Tools-Token), +# separate from the Jupyter identity in Authorization. A bearer secret is not +# cookie-based and is therefore XSRF-immune, so the relay uses it only to exempt the +# request from the browser XSRF check (mirroring jupyter_server's own exemption for +# token-authenticated requests) -- it is not a Jupyter identity. Generated once per +# process and handed only to the proxy (via the NBI_UI_TOOLS_SECRET env var). +_UI_TOOLS_SECRET = secrets.token_urlsafe(32) + + +def get_ui_tools_secret() -> str: + return _UI_TOOLS_SECRET + + +def get_ui_tools_manifest() -> list[dict]: + """MCP tool descriptors (name/description/inputSchema) for the Jupyter-UI tools.""" + return [ + {"name": t.name, "description": t.description, "inputSchema": _tool_json_schema(t.input_schema)} + for t in JUPYTER_UI_TOOLS + ] + + +async def invoke_ui_tool(name: str, arguments: dict, timeout: float = CLAUDE_AGENT_CLIENT_RESPONSE_TIMEOUT) -> dict: + """Run a Jupyter-UI tool by name against the active chat turn's UI bridge. + + Returns an MCP tool-result dict: {"content": [...], "is_error"?: bool}. The + timeout matches the agent response window so long-running cells/commands aren't + cut off, while still bounding the call if the frontend never replies (closed tab).""" + tool_def = next((t for t in JUPYTER_UI_TOOLS if t.name == name), None) + if tool_def is None: + return tool_text_response(f"Unknown tool: {name}", is_error=True) + if get_current_response() is None: + return tool_text_response( + "No active Notebook Intelligence chat turn; UI tools are only callable while a chat request is being handled.", + is_error=True, + ) + try: + return await asyncio.wait_for(tool_def.handler(arguments or {}), timeout) + except asyncio.TimeoutError: + return tool_text_response(f"Jupyter UI command timed out after {timeout:.0f}s.", is_error=True) + except Exception as exc: + return tool_text_response(f"Jupyter UI command failed: {exc}", is_error=True) + + async def custom_permission_handler( tool_name: str, input_data: dict, @@ -2112,15 +2185,26 @@ async def handle_inline_chat_request(self, request: ChatRequest, response: ChatR def _create_client_options(self) -> ClaudeAgentOptions: claude_settings = self._host.nbi_config.claude_settings - self._jupyter_ui_tools_mcp_server = create_sdk_mcp_server( - name="nbi", - version="1.0.0", - tools=[list_available_notebook_kernels, create_new_notebook, add_markdown_cell, add_code_cell, get_number_of_cells, get_cell_type_and_source, get_cell_output, set_cell_type_and_source, delete_cell, insert_cell, run_cell, save_notebook, rename_notebook, run_command_in_jupyter_terminal, open_file_in_jupyter_ui] - ) mcp_servers = {} jupyter_ui_tools_enabled = ClaudeToolType.JupyterUITools in claude_settings.get('tools', []) + # By default the Jupyter-UI tools run as an IN-PROCESS sdk MCP server. When + # 'jupyter_ui_tools_external' is set, they are instead served by an EXTERNAL + # MCP server (notebook_intelligence.mcp_ui_proxy) named "nbi" in the client's + # MCP config — required under a managed/enterprise MCP config that rejects + # dynamically-configured (in-process) servers. Tool names (mcp__nbi__*) and + # the system prompt are identical either way; only the transport differs. + # external_ui_tools is the single source of truth shared with the UI-tools + # relay (extension.UIToolsHandler); see feature_flags.is_external_ui_tools_active. + external_ui_tools = is_external_ui_tools_active(claude_settings) + if jupyter_ui_tools_enabled and not external_ui_tools: + mcp_servers["nbi"] = create_sdk_mcp_server( + name="nbi", version="1.0.0", tools=JUPYTER_UI_TOOLS + ) if jupyter_ui_tools_enabled: - mcp_servers["nbi"] = self._jupyter_ui_tools_mcp_server + transport = "external (via managed MCP config)" if external_ui_tools else "in-process" + log.info(f"Jupyter UI tools: {transport} MCP server 'nbi' ({len(JUPYTER_UI_TOOLS)} tools)") + else: + log.debug("Jupyter UI tools: disabled") allowed_tools = [] if jupyter_ui_tools_enabled: allowed_tools.extend(["mcp__nbi__list-available-notebook-kernels", "mcp__nbi__create-new-notebook", "mcp__nbi__add-markdown-cell", "mcp__nbi__add-code-cell", "mcp__nbi__get-number-of-cells", "mcp__nbi__get-cell-type-and-source", "mcp__nbi__get-cell-output", "mcp__nbi__set-cell-type-and-source", "mcp__nbi__insert-cell", "mcp__nbi__save-notebook", "mcp__nbi__rename-notebook", "mcp__nbi__open-file-in-jupyter-ui"]) @@ -2138,6 +2222,12 @@ def _create_client_options(self) -> ClaudeAgentOptions: env['ANTHROPIC_BASE_URL'] = base_url env["CLAUDE_CODE_ENTRYPOINT"] = "notebook-intelligence" + if external_ui_tools: + # Scope the bridge credential to the only mode that consumes it. Claude + # Code propagates its env to Bash, hooks, and MCP subprocesses, so do not + # expose the secret when tools are disabled or use the in-process transport. + # The proxy sends it in a dedicated header; see get_ui_tools_secret. + env["NBI_UI_TOOLS_SECRET"] = get_ui_tools_secret() continue_conversation = claude_settings.get('continue_conversation', False) diff --git a/notebook_intelligence/extension.py b/notebook_intelligence/extension.py index e1b6213..484ed84 100644 --- a/notebook_intelligence/extension.py +++ b/notebook_intelligence/extension.py @@ -3,6 +3,7 @@ import asyncio import atexit import base64 +import hmac from dataclasses import asdict, dataclass import json from os import path @@ -40,6 +41,7 @@ VALID_POLICIES, apply_claude_policies, apply_string_overrides, + is_external_ui_tools_active, is_force_off, is_locked, resolve_feature_flag, @@ -60,6 +62,9 @@ claude_bypass_disabled_by_managed_settings, claude_managed_default_permission_mode, fetch_claude_models, + get_ui_tools_manifest, + get_ui_tools_secret, + invoke_ui_tool, model_info_from_id, resolve_permission_mode, ) @@ -2241,6 +2246,81 @@ def cancel_request(self) -> None: self._cancellation_requested = True self._cancellation_signal.emit() + +class UIToolsHandler(APIHandler): + """Bridges the Jupyter-UI tools to an external stdio MCP server + (notebook_intelligence.mcp_ui_proxy), so they work under a managed/enterprise + MCP config that forbids dynamically configured (in-process) servers. + + GET -> the tool manifest (name/description/inputSchema). + POST {"name", "arguments"} -> run the tool against the active chat turn's UI + bridge and return its MCP tool-result ({"content": [...], "is_error"?: bool}). + """ + + def check_xsrf_cookie(self): + # mcp_ui_proxy proves it is the proxy NBI spawned with a per-process bearer + # secret (claude.get_ui_tools_secret), sent in the X-NBI-UI-Tools-Token header + # — separate from the Jupyter identity in Authorization, which @authenticated + # still enforces. A bearer secret is not cookie-based and so is XSRF-immune; + # exempt those requests from the XSRF check exactly as jupyter_server exempts + # token-authenticated ones. Every other caller still gets the normal check. + provided = self.request.headers.get("X-NBI-UI-Tools-Token", "") + expected = get_ui_tools_secret() + if expected and provided and hmac.compare_digest(provided, expected): + return + return super().check_xsrf_cookie() + + def _external_mode_enabled(self) -> bool: + # Resolved per request (not at boot): jupyter_ui_tools_external is mutable via + # ConfigHandler.post, and _create_client_options reads it per request too, so + # the relay must track the live setting to stay in lockstep with the transport. + # Shared predicate (see feature_flags.is_external_ui_tools_active) — both sides + # must resolve the same claude_settings the same way or they can drift apart. + return is_external_ui_tools_active(ai_service_manager.nbi_config.claude_settings or {}) + + @tornado.web.authenticated + async def get(self): + if not self._external_mode_enabled(): + raise tornado.web.HTTPError(404) + tools = get_ui_tools_manifest() + log.debug(f"UI tools relay: served manifest ({len(tools)} tools)") + self.finish(json.dumps({"tools": tools})) + + @tornado.web.authenticated + async def post(self): + if not self._external_mode_enabled(): + raise tornado.web.HTTPError(404) + try: + data = json.loads(self.request.body or b"{}") + except json.JSONDecodeError as exc: + self.set_status(400) + self.finish(json.dumps({"error": f"Invalid JSON: {exc}"})) + return + name = data.get("name") + if not isinstance(name, str) or not name: + self.set_status(400) + self.finish(json.dumps({"error": "name is required"})) + return + # Run as a cancellable task so a client disconnect (turn cancelled / tab + # closed) tears down the pending run_ui_command instead of polling for the + # full timeout window. + log.info(f"UI tools relay: invoking '{name}'") + self._invoke_task = asyncio.ensure_future( + invoke_ui_tool(name, data.get("arguments") or {}) + ) + try: + result = await self._invoke_task + except asyncio.CancelledError: + return + self.finish(json.dumps(result)) + + def on_connection_close(self): + task = getattr(self, "_invoke_task", None) + if task is not None and not task.done(): + task.cancel() + super().on_connection_close() + + @dataclass class MessageCallbackHandlers: response_emitter: WebsocketCopilotResponseEmitter @@ -3268,6 +3348,7 @@ def _setup_handlers(self, web_app, feature_policies: dict, string_overrides: dic base_url = web_app.settings["base_url"] route_pattern_capabilities = url_path_join(base_url, "notebook-intelligence", "capabilities") route_pattern_config = url_path_join(base_url, "notebook-intelligence", "config") + route_pattern_ui_tools = url_path_join(base_url, "notebook-intelligence", "ui-tools") route_pattern_update_provider_models = url_path_join(base_url, "notebook-intelligence", "update-provider-models") route_pattern_mcp_config_file = url_path_join(base_url, "notebook-intelligence", "mcp-config-file") route_pattern_reload_mcp_servers = url_path_join(base_url, "notebook-intelligence", "reload-mcp-servers") @@ -3411,6 +3492,10 @@ def _setup_handlers(self, web_app, feature_policies: dict, string_overrides: dic NotebookIntelligence.handlers = [ (route_pattern_capabilities, GetCapabilitiesHandler), (route_pattern_config, ConfigHandler), + # Always register the relay: jupyter_ui_tools_external is runtime-mutable. + # UIToolsHandler gates every request against the live setting, so changing + # transport after boot cannot leave route registration out of sync. + (route_pattern_ui_tools, UIToolsHandler), (route_pattern_update_provider_models, UpdateProviderModelsHandler), (route_pattern_mcp_config_file, MCPConfigFileHandler), (route_pattern_reload_mcp_servers, ReloadMCPServersHandler), diff --git a/notebook_intelligence/feature_flags.py b/notebook_intelligence/feature_flags.py index 05e1ca2..077824f 100644 --- a/notebook_intelligence/feature_flags.py +++ b/notebook_intelligence/feature_flags.py @@ -55,6 +55,22 @@ def is_force_off(policies: dict, name: str) -> bool: return policies.get(name, POLICY_USER_CHOICE) == POLICY_FORCE_OFF +def is_external_ui_tools_active(claude_settings: dict) -> bool: + """True iff the Jupyter-UI tools are enabled AND served by the external MCP + proxy rather than the in-process sdk server. + + Single source of truth for the transport decision: ``claude.py`` reads it to + choose which MCP server to register (and whether to hand the bridge secret to + the subprocess env), and ``extension.py``'s UI-tools relay reads it per request + to decide whether to serve or refuse. Both call sites must resolve the same + live ``claude_settings`` the same way, or the transport and the relay gate can + drift out of lockstep. + """ + return JUPYTER_UI_TOOLS_ID in (claude_settings.get("tools") or []) and bool( + claude_settings.get("jupyter_ui_tools_external", False) + ) + + def apply_string_overrides(target: dict, overrides: dict, mapping: tuple) -> dict: """Apply value-presence-locks per ``mapping`` to a copy of ``target``. diff --git a/notebook_intelligence/mcp_ui_proxy.py b/notebook_intelligence/mcp_ui_proxy.py new file mode 100644 index 0000000..2ce40da --- /dev/null +++ b/notebook_intelligence/mcp_ui_proxy.py @@ -0,0 +1,203 @@ +"""Generic stdio MCP proxy for a Jupyter Server tool backend. + +Bridges a stdio MCP client (e.g. Claude Code launched under a managed/enterprise +MCP config that forbids dynamically-configured servers) to an authenticated HTTP +tool endpoint served by the running Jupyter Server. It discovers the server URL +and auth token at runtime, fetches the tool manifest, and forwards each tool call. + +Nothing here is vendor-specific: configuration comes from standard Jupyter env +vars (or an explicit override), and the backend speaks a minimal JSON protocol: + GET -> {"tools": [{name, description, inputSchema}]} + POST {"name","arguments"} -> {"content": [...], "is_error"?: bool} + +Launch from an MCP config entry: + {"command": "python", "args": ["-m", "notebook_intelligence.mcp_ui_proxy"]} + +Environment (all optional; discovery falls back to the Jupyter runtime file): + NBI_UI_TOOLS_URL full endpoint URL (overrides discovery) + NBI_UI_TOOLS_TOKEN Jupyter auth token for the Authorization header + (else JUPYTER_TOKEN, else the discovered server token) + NBI_UI_TOOLS_SECRET bridge secret, sent in the X-NBI-UI-Tools-Token header; + the backend uses it only to exempt the call from XSRF + JUPYTER_SERVER_URL / JUPYTER_TOKEN standard Jupyter server coordinates + NBI_UI_TOOLS_HTTP_TIMEOUT per-request timeout in seconds (default: none) + NBI_UI_TOOLS_SERVER_NAME MCP handshake name (default "nbi"; cosmetic only) +""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys +import urllib.request + +import mcp.types as types +from mcp.server.lowlevel import NotificationOptions, Server +from mcp.server.models import InitializationOptions +from mcp.server.stdio import stdio_server + + +# server_name is only the MCP handshake identity; the CLI-visible tool prefix +# (mcp____*) comes from the client config-entry KEY, not from this value. +SERVER_NAME = os.environ.get("NBI_UI_TOOLS_SERVER_NAME", "nbi") +# Must match the route registered in extension.py (NotebookIntelligence._setup_handlers). +ENDPOINT_PATH = "notebook-intelligence/ui-tools" +# No client-side timeout by default: the backend bounds the call (NBI caps it at +# the agent response window), so long-running cells/commands aren't cut off. +_http_timeout = os.environ.get("NBI_UI_TOOLS_HTTP_TIMEOUT") +HTTP_TIMEOUT = float(_http_timeout) if _http_timeout else None + + +# Reach the (loopback) backend directly, never via an ambient HTTP(S)_PROXY: keeps +# the auth token on the local connection and works regardless of NO_PROXY. +_opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) + +# Bridge secret proving the caller is the proxy NBI spawned. Sent in a dedicated +# header (never Authorization) so it can't collide with the Jupyter identity; the +# relay uses it only to exempt the request from XSRF, not as an identity. +UI_TOOLS_SECRET_HEADER = "X-NBI-UI-Tools-Token" +_secret = os.environ.get("NBI_UI_TOOLS_SECRET", "") + +server = Server(SERVER_NAME) +_endpoint = "" +_token = "" + + +def _log(msg: str) -> None: + # stdout is the JSON-RPC channel; diagnostics must go to stderr only. + print(f"[mcp_ui_proxy] {msg}", file=sys.stderr, flush=True) + + +def _match_server_by_url(servers: list, base: str) -> dict | None: + """Return the running-server dict whose url matches ``base`` (trailing-slash + insensitive), or None if none matches. Pure and unit-testable independent of + jupyter_server's list_running_servers(), so the "never pair one server's url + with another's token" invariant has direct regression coverage.""" + base = base.rstrip("/") + return next((s for s in servers if s.get("url", "").rstrip("/") == base), None) + + +def _resolve_backend() -> tuple[str, str]: + """Return (endpoint_url, token) from env, then the Jupyter runtime file.""" + endpoint = os.environ.get("NBI_UI_TOOLS_URL") + token = os.environ.get("NBI_UI_TOOLS_TOKEN") or os.environ.get("JUPYTER_TOKEN") + base = os.environ.get("JUPYTER_SERVER_URL") + # Discover from the Jupyter runtime file when we still lack a base URL (to build + # the endpoint) or a token (empty string counts as unset). Reaching this block at + # all means token is falsy — the outer condition is only true via `not token` when + # base is already set, so every branch below can assume no token is known yet. + if (not endpoint and not base) or not token: + try: + from jupyter_server.serverapp import list_running_servers + servers = list(list_running_servers()) + except Exception as exc: + servers = [] + _log(f"Could not enumerate running Jupyter servers: {exc}") + # Never take a token from a server whose URL differs from `base`: that would mean + # authenticating server A's endpoint with server B's token (403 every call). + if base: + # A base URL is pinned: only adopt a token from the server that matches it. + match = _match_server_by_url(servers, base) + if match is not None: + token = match.get("token") + else: + _log( + f"JUPYTER_SERVER_URL={base} matched no running Jupyter server; " + "proceeding without a discovered token. Set NBI_UI_TOOLS_TOKEN / " + "JUPYTER_TOKEN explicitly." + ) + elif servers: + # No base is pinned: choose one server and keep its URL/token together. + if len(servers) > 1: + _log( + "Multiple Jupyter servers found and JUPYTER_SERVER_URL is unset; " + f"using {servers[0].get('url')}. Set NBI_UI_TOOLS_URL or " + "JUPYTER_SERVER_URL to target a specific server." + ) + chosen = servers[0] + base = chosen.get("url") + token = chosen.get("token") + + if not endpoint: + if not base: + raise RuntimeError( + "Cannot locate the Jupyter Server. Set NBI_UI_TOOLS_URL (and optionally " + "NBI_UI_TOOLS_TOKEN), or launch inside a running Jupyter Server." + ) + endpoint = base.rstrip("/") + "/" + ENDPOINT_PATH + return endpoint, token or "" + + +def _request(method: str, payload: dict | None = None) -> dict: + body = json.dumps(payload).encode() if payload is not None else None + req = urllib.request.Request(_endpoint, data=body, method=method) + req.add_header("Accept", "application/json") + if body is not None: + req.add_header("Content-Type", "application/json") + if _token: + req.add_header("Authorization", f"token {_token}") + if _secret: + req.add_header(UI_TOOLS_SECRET_HEADER, _secret) + with _opener.open(req, timeout=HTTP_TIMEOUT) as resp: + return json.loads(resp.read().decode() or "{}") + + +@server.list_tools() +async def list_tools() -> list[types.Tool]: + try: + data = await asyncio.to_thread(_request, "GET") + except Exception as exc: + _log(f"failed to fetch tool manifest from {_endpoint}: {exc}") + raise + tools = [ + types.Tool( + name=t["name"], + description=t.get("description", ""), + inputSchema=t.get("inputSchema") or {"type": "object", "properties": {}}, + ) + for t in data.get("tools", []) + ] + _log(f"advertising {len(tools)} tools") + return tools + + +@server.call_tool() +async def call_tool(name: str, arguments: dict) -> types.CallToolResult: + _log(f"tool call: {name}") + try: + result = await asyncio.to_thread( + _request, "POST", {"name": name, "arguments": arguments or {}} + ) + except Exception as exc: + _log(f"tool call '{name}' failed: {exc}") + return types.CallToolResult( + content=[types.TextContent(type="text", text=f"Tool bridge error: {exc}")], + isError=True, + ) + content = [ + types.TextContent(type="text", text=str(c.get("text", ""))) + for c in (result.get("content") or []) + if isinstance(c, dict) and c.get("type") == "text" + ] or [types.TextContent(type="text", text="")] + return types.CallToolResult(content=content, isError=bool(result.get("is_error"))) + + +async def main() -> None: + global _endpoint, _token + _endpoint, _token = _resolve_backend() + _log(f"bridging to {_endpoint} (jupyter-auth: {'token' if _token else 'none'}, bridge-secret: {'yes' if _secret else 'no'})") + async with stdio_server() as (read, write): + await server.run( + read, + write, + InitializationOptions( + server_name=SERVER_NAME, + server_version="1.0.0", + capabilities=server.get_capabilities(NotificationOptions(), {}), + ), + ) + + +if __name__ == "__main__": + asyncio.run(main())