diff --git a/README.md b/README.md index 039ec18c..204e217a 100644 --- a/README.md +++ b/README.md @@ -316,7 +316,7 @@ For detailed instructions and troubleshooting, see **[docs/AUTHENTICATION.md](do ## MCP Configuration -> **⚠️ Context Window Warning:** This MCP provides **39 tools**. Disable it when not using NotebookLM to preserve context. In Claude Code: `@notebooklm-mcp` to toggle. +> **⚠️ Context Window Warning:** This MCP provides **39 tools**. Disable it when not using NotebookLM to preserve context. In Claude Code: `@notebooklm-mcp` to toggle. To keep it on but expose only a subset, see [Selective tool exposure](docs/MCP_GUIDE.md#selective-tool-exposure). ### Automatic Setup (Recommended) diff --git a/docs/MCP_GUIDE.md b/docs/MCP_GUIDE.md index c8e4f98f..2d9dcd83 100644 --- a/docs/MCP_GUIDE.md +++ b/docs/MCP_GUIDE.md @@ -317,6 +317,9 @@ pipeline(action="run", notebook_id="abc", pipeline_name="ingest-and-podcast", in | `NOTEBOOKLM_HL` | Interface language and default artifact locale, including regional BCP-47 values such as `es-419` (default: en) | | `NOTEBOOKLM_QUERY_TIMEOUT` | Query timeout (seconds) | | `NOTEBOOKLM_BASE_URL` | Override base URL for Enterprise/Workspace (default: `https://notebooklm.google.com`) | +| `NOTEBOOKLM_DISABLED_GROUPS` | Comma-separated tool groups to hide (see [Selective tool exposure](#selective-tool-exposure)) | +| `NOTEBOOKLM_DISABLED_TOOLS` | Comma-separated individual tools to hide | +| `NOTEBOOKLM_ENABLED_TOOLS` | Comma-separated tools to re-enable, overriding the two above | --- @@ -325,9 +328,34 @@ pipeline(action="run", notebook_id="abc", pipeline_name="ingest-and-podcast", in This MCP has **39 tools** which consume context. Best practices: - **Disable when not using**: In Claude Code, use `@notebooklm-mcp` to toggle +- **Hide tools you don't need**: See [Selective tool exposure](#selective-tool-exposure) below to expose only a subset - **Use unified tools**: `source_add`, `studio_create`, `download_artifact` handle multiple operations each - **Poll wisely**: Use `studio_status` sparingly - artifacts take 1-5 minutes +### Selective tool exposure + +Gating is opt-in: with no configuration all tools are visible. To reduce context, +hide tools by group or by name via environment variables. Tools are hidden rather +than removed, so no code changes are needed. + +Resolution order (later wins): `NOTEBOOKLM_DISABLED_GROUPS`, then +`NOTEBOOKLM_DISABLED_TOOLS`, then `NOTEBOOKLM_ENABLED_TOOLS`. + +```bash +# Query-first setup: hide mutating groups, keep read + chat tools +export NOTEBOOKLM_DISABLED_GROUPS="notebooks_manage,sources_manage,studio,research,sharing,notes" + +# Hide one extra tool, but keep studio_status from an otherwise-hidden group +export NOTEBOOKLM_DISABLED_TOOLS="tag" +export NOTEBOOKLM_ENABLED_TOOLS="studio_status" +``` + +Available groups: `notebooks_read`, `notebooks_manage`, `sources_read`, +`sources_manage`, `chat`, `query_multi`, `organization`, `automation`, `notes`, +`auth`, `server`, `sharing`, `research`, `studio`. + +Unknown group names are ignored. Changes take effect on server restart. + --- ## IDE Configuration diff --git a/src/notebooklm_tools/mcp/server.py b/src/notebooklm_tools/mcp/server.py index ff44b684..f36cd71f 100644 --- a/src/notebooklm_tools/mcp/server.py +++ b/src/notebooklm_tools/mcp/server.py @@ -100,6 +100,11 @@ def _register_tools() -> None: # Register collected tools with mcp register_all_tools(mcp) + # Optionally hide tool groups/tools via environment variables (opt-in). + from . import tool_groups + + tool_groups.apply(mcp) + # Register tools on import _register_tools() diff --git a/src/notebooklm_tools/mcp/tool_groups.py b/src/notebooklm_tools/mcp/tool_groups.py new file mode 100644 index 00000000..427d93f5 --- /dev/null +++ b/src/notebooklm_tools/mcp/tool_groups.py @@ -0,0 +1,126 @@ +"""Optional group-based gating of MCP tools. + +The server registers a large tool set (39 tools). Clients that only need a +subset can hide the rest to save context, without editing code, by toggling +named groups or individual tools through environment variables. + +Gating is opt-in: with no configuration, every tool stays visible and behavior +is unchanged. + +Resolution order (later wins for a given tool): + 1. env NOTEBOOKLM_DISABLED_GROUPS (comma-separated group names) hides whole + groups. + 2. env NOTEBOOKLM_DISABLED_TOOLS (comma-separated tool names) hides single + tools. + 3. env NOTEBOOKLM_ENABLED_TOOLS (comma-separated tool names) re-enables single + tools, overriding the two above. + +apply(mcp) is called once from server._register_tools() after all tools are +registered. It uses FastMCP's visibility transform +(mcp.local_provider.disable(names=...)) so no tool is unregistered, only hidden. +""" + +from __future__ import annotations + +import os +from typing import Any + +# Read/manage split so a query-first client can hide mutating tools while +# keeping the read + chat core. Each tool name appears in exactly one group; +# together the groups cover every registered tool. +TOOL_GROUPS: dict[str, set[str]] = { + "notebooks_read": { + "notebook_list", + "notebook_get", + "notebook_describe", + }, + "notebooks_manage": { + "notebook_create", + "notebook_rename", + "notebook_delete", + }, + "sources_read": { + "source_list_drive", + "source_describe", + "source_get_content", + }, + "sources_manage": { + "source_add", + "source_rename", + "source_delete", + "source_sync_drive", + }, + "chat": { + "notebook_query", + "chat_configure", + "notebook_query_start", + "notebook_query_status", + }, + "query_multi": { + "cross_notebook_query", + }, + "organization": { + "label", + "tag", + }, + "automation": { + "batch", + "pipeline", + }, + "notes": { + "note", + }, + "auth": { + "refresh_auth", + "save_auth_tokens", + }, + "server": { + "server_info", + }, + "sharing": { + "notebook_share_status", + "notebook_share_public", + "notebook_share_invite", + "notebook_share_batch", + }, + "research": { + "research_start", + "research_status", + "research_import", + }, + "studio": { + "studio_create", + "studio_status", + "studio_delete", + "studio_revise", + "download_artifact", + "export_artifact", + }, +} + + +def _env_names(var: str) -> set[str]: + raw = os.environ.get(var, "") + return {part.strip() for part in raw.split(",") if part.strip()} + + +def _resolve_disabled() -> set[str]: + """Compute the final set of tool names to hide (empty unless configured).""" + names: set[str] = set() + for group in _env_names("NOTEBOOKLM_DISABLED_GROUPS"): + names |= TOOL_GROUPS.get(group, set()) + + names |= _env_names("NOTEBOOKLM_DISABLED_TOOLS") + names -= _env_names("NOTEBOOKLM_ENABLED_TOOLS") + return names + + +def apply(mcp: Any) -> set[str]: + """Hide the resolved set of tools on the given FastMCP instance. + + Returns the set of hidden tool names (empty if nothing was hidden). + """ + names = _resolve_disabled() + if names: + mcp.local_provider.disable(names=names) + return names diff --git a/tests/test_tool_groups.py b/tests/test_tool_groups.py new file mode 100644 index 00000000..0be8d514 --- /dev/null +++ b/tests/test_tool_groups.py @@ -0,0 +1,84 @@ +"""Tests for optional MCP tool-group gating.""" + +from unittest.mock import MagicMock + +from notebooklm_tools.mcp import tool_groups + +_ENV_VARS = ( + "NOTEBOOKLM_DISABLED_GROUPS", + "NOTEBOOKLM_DISABLED_TOOLS", + "NOTEBOOKLM_ENABLED_TOOLS", +) + + +def _clear_env(monkeypatch): + for var in _ENV_VARS: + monkeypatch.delenv(var, raising=False) + + +def test_no_config_hides_nothing(monkeypatch): + _clear_env(monkeypatch) + assert tool_groups._resolve_disabled() == set() + + +def test_disabled_group_hides_its_tools(monkeypatch): + _clear_env(monkeypatch) + monkeypatch.setenv("NOTEBOOKLM_DISABLED_GROUPS", "studio") + assert tool_groups._resolve_disabled() == tool_groups.TOOL_GROUPS["studio"] + + +def test_multiple_groups_are_unioned(monkeypatch): + _clear_env(monkeypatch) + monkeypatch.setenv("NOTEBOOKLM_DISABLED_GROUPS", "notes, sharing") + expected = tool_groups.TOOL_GROUPS["notes"] | tool_groups.TOOL_GROUPS["sharing"] + assert tool_groups._resolve_disabled() == expected + + +def test_unknown_group_is_ignored(monkeypatch): + _clear_env(monkeypatch) + monkeypatch.setenv("NOTEBOOKLM_DISABLED_GROUPS", "does_not_exist") + assert tool_groups._resolve_disabled() == set() + + +def test_disabled_tools_add_single_names(monkeypatch): + _clear_env(monkeypatch) + monkeypatch.setenv("NOTEBOOKLM_DISABLED_TOOLS", "notebook_query, note") + assert tool_groups._resolve_disabled() == {"notebook_query", "note"} + + +def test_enabled_tools_override_disabled(monkeypatch): + _clear_env(monkeypatch) + monkeypatch.setenv("NOTEBOOKLM_DISABLED_GROUPS", "studio") + monkeypatch.setenv("NOTEBOOKLM_ENABLED_TOOLS", "studio_status") + result = tool_groups._resolve_disabled() + assert "studio_status" not in result + assert "studio_create" in result + + +def test_apply_disables_resolved_names(monkeypatch): + _clear_env(monkeypatch) + monkeypatch.setenv("NOTEBOOKLM_DISABLED_GROUPS", "notes") + mcp = MagicMock() + + hidden = tool_groups.apply(mcp) + + assert hidden == tool_groups.TOOL_GROUPS["notes"] + mcp.local_provider.disable.assert_called_once_with(names=hidden) + + +def test_apply_noop_when_nothing_disabled(monkeypatch): + _clear_env(monkeypatch) + mcp = MagicMock() + + hidden = tool_groups.apply(mcp) + + assert hidden == set() + mcp.local_provider.disable.assert_not_called() + + +def test_groups_are_disjoint(): + seen: dict[str, str] = {} + for group, names in tool_groups.TOOL_GROUPS.items(): + for name in names: + assert name not in seen, f"{name} in both {seen.get(name)} and {group}" + seen[name] = group