diff --git a/docs/architecture.md b/docs/architecture.md index 8e871ede..39720f34 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -3,7 +3,7 @@ Turnstone is an AI orchestration platform with tool use, parallel workstreams, and persistent memory. It connects to any OpenAI-compatible API (local vLLM, OpenAI, etc.) or Anthropic's native Messages API via pluggable provider adapters, and gives the -model 14 built-in tools plus external tools via MCP (Model Context Protocol) for +model 18 built-in tools plus external tools via MCP (Model Context Protocol) for reading, writing, searching, planning, and executing code. The core design principle is a **UI-agnostic engine with pluggable frontends**. diff --git a/docs/diagrams/02-package-structure.puml b/docs/diagrams/02-package-structure.puml index ab2aa8ee..6d07a491 100644 --- a/docs/diagrams/02-package-structure.puml +++ b/docs/diagrams/02-package-structure.puml @@ -96,7 +96,7 @@ package "turnstone/sdk/" <> { ' Tool schemas package "turnstone/tools/" <> { - component [*.json\n15 tool schemas] as schemas <> + component [*.json\n18 tool schemas] as schemas <> } ' Entry point dependencies diff --git a/docs/diagrams/03-core-engine-classes.puml b/docs/diagrams/03-core-engine-classes.puml index f32f4034..a336d502 100644 --- a/docs/diagrams/03-core-engine-classes.puml +++ b/docs/diagrams/03-core-engine-classes.puml @@ -211,15 +211,23 @@ enum "WorkstreamState" as WsState { class "MCPClientManager" as MCPMgr { - _sessions: dict[str, ClientSession] - _per_server_tools: dict[str, list[dict]] + - _per_server_resources: dict[str, list[dict]] + - _per_server_prompts: dict[str, list[dict]] - _tools: list[dict] - _tool_map: dict[str, tuple] + - _resource_map: dict[str, tuple] + - _prompt_map: dict[str, tuple] - _supports_list_changed: dict[str, bool] - _listeners: list[Callable] -- + start() + get_tools() → list[dict] + + get_resources() → list[dict] + + get_prompts() → list[dict] + is_mcp_tool(name) → bool + call_tool_sync(name, args) → str + + read_resource_sync(uri) → str + + get_prompt_sync(name, args?) → list[dict] + refresh_sync(server?) → dict + add_listener(callback) + remove_listener(callback) @@ -230,6 +238,8 @@ class "MCPClientManager" as MCPMgr { bridges async MCP SDK to sync ChatSession dispatch. Push + periodic + manual refresh. + Resources + prompts discovered + alongside tools at startup. -- core/mcp_client.py } diff --git a/docs/diagrams/05-tool-pipeline.puml b/docs/diagrams/05-tool-pipeline.puml index ab565e07..1d8aba85 100644 --- a/docs/diagrams/05-tool-pipeline.puml +++ b/docs/diagrams/05-tool-pipeline.puml @@ -24,29 +24,31 @@ partition "Phase 1: Prepare" #E8F5E9 { :Dispatch to _prepare_{func_name}(); note right - **Dispatch table (16 tools):** - ┌──────────────┬──────────────────┐ - │ Tool │ Needs Approval? │ - ├──────────────┼──────────────────┤ - │ bash │ ✓ Yes │ - │ read_file │ ✗ Auto-approve │ - │ write_file │ ✓ Yes │ - │ edit_file │ ✓ Yes │ - │ search │ ✗ Auto-approve │ - │ math │ ✓ Yes │ - │ man │ ✗ Auto-approve │ - │ web_fetch │ ✓ Yes │ - │ web_search │ ✓ Yes │ - │ tool_search │ ✗ Auto-approve │ - │ task │ ✓ Yes │ - │ plan │ ✓ Yes │ - │ remember │ ✗ Auto-approve │ - │ recall │ ✗ Auto-approve │ - │ forget │ ✗ Auto-approve │ - │ notify │ ✗ Auto-approve │ - ├──────────────┼──────────────────┤ - │ mcp__* │ ✓ Yes (external) │ - └──────────────┴──────────────────┘ + **Dispatch table (18 tools):** + ┌───────────────┬──────────────────┐ + │ Tool │ Needs Approval? │ + ├───────────────┼──────────────────┤ + │ bash │ ✓ Yes │ + │ read_file │ ✗ Auto-approve │ + │ write_file │ ✓ Yes │ + │ edit_file │ ✓ Yes │ + │ search │ ✗ Auto-approve │ + │ math │ ✗ Auto-approve │ + │ man │ ✗ Auto-approve │ + │ web_fetch │ ✗ Auto-approve │ + │ web_search │ ✗ Auto-approve │ + │ tool_search │ ✗ Auto-approve │ + │ task │ ✓ Yes │ + │ plan │ ✓ Yes │ + │ remember │ ✗ Auto-approve │ + │ recall │ ✗ Auto-approve │ + │ forget │ ✗ Auto-approve │ + │ notify │ ✗ Auto-approve │ + │ read_resource │ ✓ Yes │ + │ use_prompt │ ✓ Yes │ + ├───────────────┼──────────────────┤ + │ mcp__* │ ✓ Yes (external) │ + └───────────────┴──────────────────┘ end note :Build item dict: @@ -117,6 +119,8 @@ partition "Phase 3: Execute" #E3F2FD { ├─ _exec_remember: SQLite INSERT OR REPLACE ├─ _exec_recall: SQLite FTS5/LIKE search ├─ _exec_forget: SQLite DELETE + ├─ _exec_read_resource: MCPClientManager.read_resource_sync() + ├─ _exec_use_prompt: MCPClientManager.get_prompt_sync() └─ _exec_mcp_tool: MCPClientManager.call_tool_sync() end note diff --git a/docs/diagrams/20-mcp-architecture.puml b/docs/diagrams/20-mcp-architecture.puml new file mode 100644 index 00000000..9cf1591d --- /dev/null +++ b/docs/diagrams/20-mcp-architecture.puml @@ -0,0 +1,157 @@ +@startuml +!theme plain +title Turnstone — MCP Architecture (Resources, Prompts, Tools) + +skinparam participant { + BackgroundColor<> #E1BEE7 + BackgroundColor<> #C8E6C9 + BackgroundColor<> #B3E5FC + BackgroundColor<> #FFE0B2 + BackgroundColor<> #E8EAF6 +} + +participant "MCP Server\n(external)" as MCPSrv <> +participant "MCPClientManager\n(mcp_client.py)" as MCPMgr <> +participant "ChatSession\n(session.py)" as Session <> +participant "StorageBackend\n(governance)" as Storage <> +participant "Server / Console\n(health + UI)" as UI <> + +== Startup: Connection & Discovery == + +MCPMgr -> MCPSrv : initialize (stdio or HTTP) +MCPSrv --> MCPMgr : capabilities\n(tools, resources, prompts) + +MCPMgr -> MCPSrv : tools/list +MCPSrv --> MCPMgr : Tool[] + +opt resources capability + MCPMgr -> MCPSrv : resources/list + MCPSrv --> MCPMgr : Resource[] + MCPMgr -> MCPSrv : resources/templates/list + MCPSrv --> MCPMgr : ResourceTemplate[] +end + +opt prompts capability + MCPMgr -> MCPSrv : prompts/list + MCPSrv --> MCPMgr : Prompt[] +end + +note over MCPMgr + Per-server storage: + _per_server_tools, _per_server_resources, _per_server_prompts + Copy-on-write rebuild into _tools, _resources, _prompts + Prefix: mcp__{server}__{name} +end note + +MCPMgr -> Session : notify tool listeners +MCPMgr -> Session : notify resource listeners + +== Governance Sync (on connect & refresh) == + +MCPMgr -> Storage : sync_prompts_to_storage() +note right + For each MCP prompt: + - Manual template exists? → skip + - MCP template exists? → update + - New? → create (origin="mcp", + readonly=True) + Removed prompts → delete + Protected by _sync_lock +end note + +== set_storage() from entry point == + +UI -> MCPMgr : set_storage(backend) +note right + If servers already connected, + triggers immediate sync +end note + +== Runtime: Tool Execution == + +Session -> Session : _prepare_mcp_tool(func_name, args) +note right + approval_label = func_name + (e.g. mcp__github__search) + needs_approval = True +end note +Session -> MCPMgr : call_tool_sync(name, args) +MCPMgr -> MCPSrv : tools/call +MCPSrv --> MCPMgr : ToolResult +MCPMgr --> Session : output (text) + +== Runtime: Resource Read == + +Session -> Session : _prepare_read_resource(uri) +note right + approval_label = mcp_resource__{normalized_uri} + URI normalized (.. resolved) + needs_approval = True +end note +Session -> MCPMgr : read_resource_sync(uri) +MCPMgr -> MCPSrv : resources/read +MCPSrv --> MCPMgr : ReadResourceResult +MCPMgr --> Session : content (text/blob) + +== Runtime: Prompt Invocation == + +Session -> Session : _prepare_use_prompt(name, arguments) +note right + approval_label = mcp__srv__prompt + Validated via is_mcp_prompt() + needs_approval = True +end note +Session -> MCPMgr : get_prompt_sync(name, args) +MCPMgr -> MCPSrv : prompts/get +MCPSrv --> MCPMgr : GetPromptResult +MCPMgr --> Session : messages [{role, content}] + +== Three-Tier Refresh == + +group Push Notifications + MCPSrv -> MCPMgr : ToolListChangedNotification + MCPMgr -> MCPMgr : _refresh_server_tools() + + MCPSrv -> MCPMgr : ResourceListChangedNotification + MCPMgr -> MCPMgr : _refresh_server_resources() + + MCPSrv -> MCPMgr : PromptListChangedNotification + MCPMgr -> MCPMgr : _refresh_server_prompts() + MCPMgr -> Storage : sync_prompts_to_storage() +end + +group Periodic Polling (default 4h) + MCPMgr -> MCPMgr : _periodic_refresh() + note right + Only polls capabilities + without push support. + Staggered per-server. + end note +end + +group Manual Refresh + Session -> MCPMgr : refresh_sync() + note right: /mcp refresh [server] +end + +== Policy Evaluation == + +note over Session + Tool policies use fnmatch on approval_label: + - mcp__github__* → allow (all GitHub tools/prompts) + - mcp_resource__file:///docs/* → allow + - mcp_resource__* → deny (block all resource reads) + - mcp__untrusted__* → ask +end note + +== UI Visibility == + +UI -> MCPMgr : server_count, get_resources(), get_prompts() +note over UI + /health → mcp.servers, mcp.resources, mcp.prompts + Server UI: magenta status badge + Console: cluster status bar + node detail + System message: + catalogs +end note + +@enduml diff --git a/docs/diagrams/png/02-package-structure.png b/docs/diagrams/png/02-package-structure.png index 07b44751..6ce82b08 100644 --- a/docs/diagrams/png/02-package-structure.png +++ b/docs/diagrams/png/02-package-structure.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:0ee0a9391bd19d92e9271bf6bd531e9c2e18baf8c5a11ead49b3c10db4d8939b -size 329625 +oid sha256:c9daca81971ba7a8ed6736d23d5373c69435158fa6240b9880d14fc4759ab580 +size 329673 diff --git a/docs/diagrams/png/03-core-engine-classes.png b/docs/diagrams/png/03-core-engine-classes.png index 065fbb4e..33647da3 100644 --- a/docs/diagrams/png/03-core-engine-classes.png +++ b/docs/diagrams/png/03-core-engine-classes.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:760c37e67736588dadee21d500419a48e9fc50f8bdc5667e686c580022bd40e2 -size 554869 +oid sha256:01fbb3338df6426cefc2811541a865f268673b4febf32f524c264d120bc068fa +size 589546 diff --git a/docs/diagrams/png/05-tool-pipeline.png b/docs/diagrams/png/05-tool-pipeline.png index dc644407..881072b2 100644 --- a/docs/diagrams/png/05-tool-pipeline.png +++ b/docs/diagrams/png/05-tool-pipeline.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:a3ffd93ccb634f76560f1dd65242b89cd34443b37e355f25f29a1c63a22001be -size 265259 +oid sha256:6dd3c923d1e1c49b5f91d8d342fb4b0d49a46d432460379ad146a9e3b075a05a +size 277234 diff --git a/docs/diagrams/png/20-mcp-architecture.png b/docs/diagrams/png/20-mcp-architecture.png new file mode 100644 index 00000000..53942a0f --- /dev/null +++ b/docs/diagrams/png/20-mcp-architecture.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2c19dfae7606de8277d44d11226bdfba608d830382e92ebf0bec284b901fb807 +size 248194 diff --git a/docs/governance.md b/docs/governance.md index 5a84f140..7c0a8fe8 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -43,15 +43,26 @@ Admin-defined rules that control tool execution: - **Priority**: Higher priority evaluated first, first match wins - **Enforcement**: `evaluate_tool_policies_batch()` called in `WebUI.approve_tools()` before the `auto_approve` check +- **MCP granular policies**: MCP resources and prompts are evaluated using their + `approval_label` for fine-grained control: + - Resource reads: `mcp_resource__{uri}` (e.g., `mcp_resource__file:///docs/*` to allow, + `mcp_resource__*` to deny all) + - Prompt invocations: `mcp__{server}__{prompt}` (e.g., `mcp__trusted__*` to allow, + `mcp__*` to require approval for all) + - Built-in tools continue to use `func_name` for backward compatibility ### Prompt Templates Reusable system message templates with variable substitution: - **Variables**: `{{variable_name}}` placeholders in content -- **Categories**: general, engineering, support, custom +- **Categories**: general, engineering, support, custom, mcp - **Default flag**: `is_default=true` templates intended for new workstreams - **Storage**: `prompt_templates` table with JSON `variables` array +- **MCP sync**: MCP server prompts are auto-synced into prompt_templates with + `origin="mcp"`, `mcp_server` set, and `readonly=True`. Manual templates take + precedence on name collision. Admin UI shows origin badge and disables + edit/delete for MCP-sourced templates. See `docs/tools.md` MCP Prompts section ### Usage Tracking diff --git a/docs/tools.md b/docs/tools.md index d2fb20d4..836bdef4 100644 --- a/docs/tools.md +++ b/docs/tools.md @@ -1,6 +1,6 @@ # Tools Reference -turnstone exposes 16 built-in tools plus any number of external MCP tools to the +turnstone exposes 18 built-in tools plus any number of external MCP tools to the LLM via the OpenAI function-calling interface. Built-in tools are defined as JSON files under `turnstone/tools/` and loaded at startup by `turnstone/core/tools.py`. MCP tools are discovered from configured MCP servers at startup by @@ -46,12 +46,12 @@ schema plus turnstone-specific metadata keys: | Name | Description | |---------------------|-------------| -| `TOOLS` | All 16 tool definitions (sent to the model). | +| `TOOLS` | All 18 tool definitions (sent to the model). | | `AGENT_TOOLS` | Tools with `agent: true` -- available to plan sub-agents. Read-only tools. | | `TASK_AGENT_TOOLS` | Tools with `task_agent: true` -- available to task sub-agents. Includes write operations. | | `AGENT_AUTO_TOOLS` | Set of tool names with `auto_approve: true` -- no user confirmation needed. | | `TASK_AUTO_TOOLS` | Same as `AGENT_AUTO_TOOLS` (identical filter). | -| `BUILTIN_TOOL_NAMES`| Frozenset of all 16 built-in tool names. Used by tool search to distinguish always-on tools from deferrable MCP tools. | +| `BUILTIN_TOOL_NAMES`| Frozenset of all 18 built-in tool names. Used by tool search to distinguish always-on tools from deferrable MCP tools. | | `PRIMARY_KEY_MAP` | Dict mapping tool name to its `primary_key` parameter name. | --- @@ -69,7 +69,7 @@ Tool execution follows a three-phase pipeline inside `ChatSession._execute_tools - Parses the JSON arguments (with fallback for malformed JSON). - If JSON parsing fails entirely, uses `PRIMARY_KEY_MAP` to map a bare string to the correct parameter. -- Dispatches to the matching `_prepare_{func_name}()` handler. There are 15 +- Dispatches to the matching `_prepare_{func_name}()` handler. There are 18 built-in tools plus `tool_search` (synthetic, client-side BM25 fallback) and the generic `_prepare_mcp_tool()` handler for MCP tools. - Validates arguments and builds a preview dict containing: @@ -168,6 +168,8 @@ Every tool defines a `primary_key`. The mapping is: | `recall` | `query` | | `forget` | `key` | | `notify` | `message` | +| `read_resource` | `uri` | +| `use_prompt` | `name` | --- @@ -517,6 +519,8 @@ data.get("mergedAt") is not None | `forget` | Memory | Yes | No | No | `key` | | `notify` | Notify | Yes | Yes | Yes | `message` | | `watch` | Monitor | No (create) | No | No | `command` | +| `read_resource`| MCP | No | Yes | Yes | `uri` | +| `use_prompt` | MCP | No | Yes | Yes | `name` | | `tool_search`| Search | Yes | No | No | `query` | --- @@ -569,7 +573,7 @@ CLI flags override the config file: search stays off and all tools are sent to the model directly. 2. **Partitioning**: When active, tools are split into two sets: - - **Always-on** -- the 15 built-in tools (members of `BUILTIN_TOOL_NAMES`). + - **Always-on** -- the 18 built-in tools (members of `BUILTIN_TOOL_NAMES`). These are always visible to the model. - **Deferred** -- all MCP tools. These are not sent in the tool list unless the model searches for them. @@ -593,6 +597,8 @@ where the model can interactively search for tools it needs. ## MCP Tools (External) +> See also: [MCP Architecture diagram](diagrams/png/20-mcp-architecture.png) + Turnstone supports the [Model Context Protocol](https://modelcontextprotocol.io/) (MCP) for connecting external tool servers — GitHub, databases, filesystems, or any MCP-compatible service. @@ -610,7 +616,7 @@ MCP-compatible service. 3. **Schema conversion**: Each MCP tool's `inputSchema` is converted to OpenAI function-calling format. The tool name is prefixed: `mcp__{server}__{tool}`. -4. **Merging**: MCP tools are appended after the 15 built-in tools via +4. **Merging**: MCP tools are appended after the 18 built-in tools via `merge_mcp_tools()`. Built-in tools appear first, giving them natural LLM priority. When dynamic tool search is active, MCP tools are deferred rather than directly visible -- the model discovers them via search as needed (see @@ -729,3 +735,140 @@ MCP refresh complete: MCP refresh complete: github: no changes ``` + +--- + +## MCP Resources + +MCP servers can expose **resources** -- named data items (files, database rows, +API responses) addressable by URI. turnstone discovers resources at startup and +makes them available to the model via the `read_resource` built-in tool. + +### Discovery + +During the MCP `initialize` handshake, `MCPClientManager` checks each server's +capabilities for the `resources` capability. For servers that declare it: + +1. `list_resources` fetches static resources (fixed URIs). +2. `list_resource_templates` fetches URI templates (parameterized patterns like + `db://tables/{table}/rows/{id}`). + +Both are stored as `{uri, name, description, mimeType, server}` dicts and +merged into a unified catalog. + +### Resource catalog in system message + +The first 50 resources are injected into the system message as an XML-delimited +block so the model knows what URIs are available: + +```xml + + file:///project/README.md Project readme + db://users/schema User table schema + +Use read_resource(uri='...') to access the resources listed above. +``` + +### read_resource tool + +| Parameter | Type | Required | Description | +|-----------|--------|----------|-------------| +| `uri` | string | yes | The resource URI to read. | + +- **What it does**: Reads the resource from its MCP server via `MCPClientManager.read_resource_sync()`. Returns text content for text resources or base64-encoded data for binary resources. Output is truncated by the standard tool output limiter. +- **Auto-approve**: No -- requires user confirmation (reads external data). +- **Agent availability**: `agent` and `task_agent`. + +### Capability guards + +The `read_resource` tool schema is always loaded (it is a built-in JSON schema), +but resource discovery only runs for servers that declare the `resources` +capability. Servers without the capability contribute zero resources to the +catalog. + +### Refresh + +Resource lists stay current through the same three-tier mechanism as tool lists: + +1. **Push** -- Servers declaring `resources.listChanged: true` send + `notifications/resources/list_changed`, triggering an immediate refresh. +2. **Periodic** -- Servers without push are polled on the configured refresh + interval (default 4 hours, same timer as tools). +3. **Manual** -- `/mcp refresh` re-fetches resources alongside tools. + +--- + +## MCP Prompts + +MCP servers can also expose **prompts** -- reusable message templates with +optional arguments. turnstone discovers prompts at startup for servers that +declare the `prompts` capability. + +### Discovery + +Prompt discovery mirrors resource discovery: `list_prompts` is called during +the `initialize` handshake. Each prompt is stored with its prefixed name +(`mcp__{server}__{prompt}`), description, and argument schema. + +### use_prompt tool + +| Parameter | Type | Required | Description | +|-------------|--------|----------|-------------| +| `name` | string | yes | The prompt name (e.g. `mcp__server__prompt_name`). | +| `arguments` | object | no | Key-value argument pairs for the prompt. Values must be strings. | + +- **What it does**: Invokes an MCP prompt template by name via `MCPClientManager.get_prompt_sync()`, expanding it into messages. Returns the expanded prompt content formatted as `[role]: content` blocks joined with blank lines. The prompt catalog is listed in the system message so the model knows which prompts are available. Output is truncated by the standard tool output limiter. +- **Auto-approve**: No -- requires user confirmation (invokes external prompt servers). +- **Agent availability**: `agent` and `task_agent`. + +### Invocation + +`MCPClientManager.get_prompt_sync()` calls the server's `get_prompt` method +with the provided arguments and returns the expanded messages. The `use_prompt` +built-in tool exposes this to the model as a function call. + +### Governance Sync + +Discovered MCP prompts are automatically synced into the `prompt_templates` +governance table as first-class governed templates: + +- **Origin tracking**: MCP-sourced templates have `origin="mcp"` and + `mcp_server` set to the server name. Manual templates have + `origin="manual"`. +- **Read-only**: MCP-sourced templates are `readonly=True`. The admin API + returns 403 on update/delete attempts. The admin UI disables edit/delete + buttons and shows an origin badge. +- **Precedence**: If a manual template and MCP prompt share the same name, + the manual template wins and the MCP prompt is skipped (with a log + warning). +- **Lifecycle**: Templates are created on connect, updated on prompt list + refresh, and removed when the MCP server no longer exposes the prompt. + The sync runs automatically on connect, on `PromptListChangedNotification`, + and on manual `/mcp refresh`. +- **Schema**: Migration 009 adds `origin`, `mcp_server`, and `readonly` + columns to the `prompt_templates` table. + +The `use_prompt` tool allows the model to invoke any discovered MCP prompt at +runtime. A catalog of up to 30 prompts is injected into the system message +inside `` XML tags so the model can discover available prompts. + +--- + +## MCP UI Visibility + +MCP server, resource, and prompt counts are surfaced across the UI: + +- **Server `/health` endpoint**: Returns `mcp.servers`, `mcp.resources`, + `mcp.prompts` when MCP is configured +- **Server UI**: Magenta status badge in the header showing server count, + with resource/prompt counts in tooltip +- **Console cluster status bar**: MCP metrics (servers/resources/prompts) + with magenta LED dot indicator, shown after a divider from workstream + metrics +- **Console node detail**: Per-node MCP summary showing server, resource, + and prompt counts +- **Console collector**: Aggregates MCP counts across all nodes in the + cluster overview + +MCP indicators use the `--magenta` design token for consistent theming +across light and dark modes. diff --git a/sdk/typescript/openapi-server.json b/sdk/typescript/openapi-server.json index 666a9b26..354da163 100644 --- a/sdk/typescript/openapi-server.json +++ b/sdk/typescript/openapi-server.json @@ -1177,12 +1177,44 @@ } ], "default": null + }, + "mcp": { + "anyOf": [ + { + "$ref": "#/components/schemas/McpStatus" + }, + { + "type": "null" + } + ], + "default": null } }, "required": ["status"], "title": "HealthResponse", "type": "object" }, + "McpStatus": { + "properties": { + "servers": { + "default": 0, + "title": "Servers", + "type": "integer" + }, + "resources": { + "default": 0, + "title": "Resources", + "type": "integer" + }, + "prompts": { + "default": 0, + "title": "Prompts", + "type": "integer" + } + }, + "title": "McpStatus", + "type": "object" + }, "BackendStatus": { "properties": { "status": { diff --git a/sdk/typescript/src/index.ts b/sdk/typescript/src/index.ts index c79de98f..7fa83c73 100644 --- a/sdk/typescript/src/index.ts +++ b/sdk/typescript/src/index.ts @@ -91,6 +91,7 @@ export type { SavedWorkstreamInfo, ListSavedWorkstreamsResponse, BackendStatus, + McpStatus, WorkstreamCounts, HealthResponse, AuthLoginRequest, diff --git a/sdk/typescript/src/types.ts b/sdk/typescript/src/types.ts index e7d2093d..077628d5 100644 --- a/sdk/typescript/src/types.ts +++ b/sdk/typescript/src/types.ts @@ -159,6 +159,12 @@ export interface WorkstreamCounts { error?: number; } +export interface McpStatus { + servers: number; + resources: number; + prompts: number; +} + export interface HealthResponse { status: string; version?: string; @@ -166,6 +172,7 @@ export interface HealthResponse { model?: string; workstreams?: WorkstreamCounts; backend?: BackendStatus | null; + mcp?: McpStatus | null; } // --------------------------------------------------------------------------- @@ -452,6 +459,9 @@ export interface PromptTemplateInfo { created_by: string; created: string; updated: string; + origin: string; + mcp_server: string; + readonly: boolean; } export interface CreateTemplateOptions { diff --git a/tests/test_governance_storage.py b/tests/test_governance_storage.py index fa8b2ccf..e845c0be 100644 --- a/tests/test_governance_storage.py +++ b/tests/test_governance_storage.py @@ -375,6 +375,66 @@ class TestPromptTemplateCRUD: assert t2["is_default"] is False assert isinstance(t2["is_default"], bool) + def test_create_with_mcp_origin(self, db): + db.create_prompt_template( + "t1", + "mcp__srv__prompt", + "mcp", + "content", + variables="[]", + is_default=False, + org_id="", + created_by="", + origin="mcp", + mcp_server="srv", + readonly=True, + ) + tpl = db.get_prompt_template("t1") + assert tpl is not None + assert tpl["origin"] == "mcp" + assert tpl["mcp_server"] == "srv" + assert tpl["readonly"] is True + assert isinstance(tpl["readonly"], bool) + + def test_default_origin_values(self, db): + db.create_prompt_template("t1", "basic", "general", "Hello") + tpl = db.get_prompt_template("t1") + assert tpl is not None + assert tpl["origin"] == "manual" + assert tpl["mcp_server"] == "" + assert tpl["readonly"] is False + + def test_get_prompt_template_by_name(self, db): + db.create_prompt_template("t1", "greeting", "general", "Hello!") + tpl = db.get_prompt_template_by_name("greeting") + assert tpl is not None + assert tpl["template_id"] == "t1" + assert tpl["name"] == "greeting" + + def test_get_prompt_template_by_name_nonexistent(self, db): + assert db.get_prompt_template_by_name("nope") is None + + def test_list_prompt_templates_by_origin(self, db): + db.create_prompt_template("t1", "manual_one", "general", "A", origin="manual") + db.create_prompt_template("t2", "mcp_one", "mcp", "B", origin="mcp", mcp_server="srv1") + db.create_prompt_template("t3", "mcp_two", "mcp", "C", origin="mcp", mcp_server="srv2") + result = db.list_prompt_templates_by_origin("mcp") + assert len(result) == 2 + names = [r["name"] for r in result] + assert "mcp_one" in names + assert "mcp_two" in names + + def test_delete_prompt_templates_by_server(self, db): + db.create_prompt_template("t1", "a", "mcp", "A", origin="mcp", mcp_server="srv1") + db.create_prompt_template("t2", "b", "mcp", "B", origin="mcp", mcp_server="srv1") + db.create_prompt_template("t3", "c", "mcp", "C", origin="mcp", mcp_server="srv2") + deleted = db.delete_prompt_templates_by_server("srv1") + assert deleted == 2 + # Only srv2 template remains. + remaining = db.list_prompt_templates() + assert len(remaining) == 1 + assert remaining[0]["mcp_server"] == "srv2" + # --------------------------------------------------------------------------- # Usage Events diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py index 5369d673..ff3b3bdb 100644 --- a/tests/test_mcp_client.py +++ b/tests/test_mcp_client.py @@ -51,6 +51,83 @@ def _fake_openai_tool(name: str = "mcp__test__search") -> dict[str, Any]: } +def _fake_mcp_resource( + uri: str = "file:///README.md", + name: str = "readme", + description: str = "Project readme", + mime_type: str = "text/plain", +) -> MagicMock: + """Create a mock MCP Resource object matching the SDK's Resource type.""" + res = MagicMock() + res.uri = uri + res.name = name + res.description = description + res.mimeType = mime_type + return res + + +def _fake_resource_dict( + uri: str = "file:///README.md", + name: str = "readme", + description: str = "Project readme", + mime_type: str = "text/plain", + server: str = "test", +) -> dict[str, Any]: + """Create a fake resource dict as stored in per-server state.""" + return { + "uri": uri, + "name": name, + "description": description, + "mimeType": mime_type, + "server": server, + } + + +def _fake_mcp_prompt( + name: str = "code_review", + description: str = "Generate a code review", + arguments: list[dict[str, Any]] | None = None, +) -> MagicMock: + """Create a mock MCP Prompt object matching the SDK's Prompt type.""" + prompt = MagicMock() + prompt.name = name + prompt.description = description + if arguments is None: + arg = MagicMock() + arg.name = "language" + arg.description = "Programming language" + arg.required = True + prompt.arguments = [arg] + else: + mock_args = [] + for a in arguments: + arg = MagicMock() + arg.name = a["name"] + arg.description = a.get("description", "") + arg.required = a.get("required", False) + mock_args.append(arg) + prompt.arguments = mock_args + return prompt + + +def _fake_prompt_dict( + name: str = "mcp__test__code_review", + original_name: str = "code_review", + server: str = "test", + description: str = "Generate a code review", +) -> dict[str, Any]: + """Create a fake prompt dict as stored in per-server state.""" + return { + "name": name, + "original_name": original_name, + "server": server, + "description": description, + "arguments": [ + {"name": "language", "description": "Programming language", "required": True} + ], + } + + # --------------------------------------------------------------------------- # Schema conversion # --------------------------------------------------------------------------- @@ -453,6 +530,23 @@ class TestRebuildTools: class TestRefreshServer: + @staticmethod + def _add_empty_resource_prompt_mocks( + mgr: MCPClientManager, server_name: str, mock_session: MagicMock + ) -> None: + """Add empty list_resources/list_prompts mocks so _refresh_server works.""" + mgr._supports_resources[server_name] = True + mgr._supports_prompts[server_name] = True + empty_res = MagicMock() + empty_res.resources = [] + mock_session.list_resources = AsyncMock(return_value=empty_res) + empty_tmpl = MagicMock() + empty_tmpl.resourceTemplates = [] + mock_session.list_resource_templates = AsyncMock(return_value=empty_tmpl) + empty_prompts = MagicMock() + empty_prompts.prompts = [] + mock_session.list_prompts = AsyncMock(return_value=empty_prompts) + def test_refresh_detects_added_tools(self): async def _run() -> None: mgr = MCPClientManager({}) @@ -463,6 +557,7 @@ class TestRefreshServer: _fake_mcp_tool("create"), # new tool ] mock_session.list_tools = AsyncMock(return_value=mock_result) + self._add_empty_resource_prompt_mocks(mgr, "github", mock_session) mgr._sessions["github"] = mock_session mgr._per_server_tools["github"] = [_fake_openai_tool("mcp__github__search")] mgr._rebuild_tools() @@ -481,6 +576,7 @@ class TestRefreshServer: mock_result = MagicMock() mock_result.tools = [] # all tools removed mock_session.list_tools = AsyncMock(return_value=mock_result) + self._add_empty_resource_prompt_mocks(mgr, "github", mock_session) mgr._sessions["github"] = mock_session mgr._per_server_tools["github"] = [_fake_openai_tool("mcp__github__search")] mgr._rebuild_tools() @@ -499,6 +595,7 @@ class TestRefreshServer: mock_result = MagicMock() mock_result.tools = [_fake_mcp_tool("search")] mock_session.list_tools = AsyncMock(return_value=mock_result) + self._add_empty_resource_prompt_mocks(mgr, "github", mock_session) mgr._sessions["github"] = mock_session mgr._per_server_tools["github"] = [_fake_openai_tool("mcp__github__search")] mgr._rebuild_tools() @@ -513,7 +610,7 @@ class TestRefreshServer: async def _run() -> None: mgr = MCPClientManager({}) with pytest.raises(RuntimeError, match="not connected"): - await mgr._refresh_server("ghost") + await mgr._refresh_server_tools("ghost") asyncio.run(_run()) @@ -709,3 +806,440 @@ class TestSessionRefresh: session.handle_command("/mcp refresh") session.ui.on_error.assert_called_once() assert "MCP refresh failed" in session.ui.on_error.call_args[0][0] + + +# --------------------------------------------------------------------------- +# MCP Resources +# --------------------------------------------------------------------------- + + +class TestMCPResources: + def test_resource_discovery(self): + """Mock list_resources() returning 2 resources, verify get_resources().""" + mgr = MCPClientManager({}) + mgr._per_server_resources = { + "fs": [ + _fake_resource_dict("file:///a.txt", "a", "File A", "text/plain", "fs"), + _fake_resource_dict("file:///b.txt", "b", "File B", "text/plain", "fs"), + ], + } + mgr._rebuild_resources() + resources = mgr.get_resources() + assert len(resources) == 2 + uris = {r["uri"] for r in resources} + assert uris == {"file:///a.txt", "file:///b.txt"} + assert all(r["server"] == "fs" for r in resources) + + def test_rebuild_resources_copy_on_write(self): + """Verify mutation safety — get_resources() returns independent copy.""" + mgr = MCPClientManager({}) + mgr._per_server_resources = { + "a": [_fake_resource_dict("file:///x", "x", "", "", "a")], + } + mgr._rebuild_resources() + old_resources = mgr._resources + old_map = mgr._resource_map + mgr._per_server_resources["b"] = [_fake_resource_dict("file:///y", "y", "", "", "b")] + mgr._rebuild_resources() + assert mgr._resources is not old_resources + assert mgr._resource_map is not old_map + + def test_get_resources_returns_copy(self): + mgr = MCPClientManager({}) + mgr._per_server_resources = { + "a": [_fake_resource_dict("file:///x", "x", "", "", "a")], + } + mgr._rebuild_resources() + resources = mgr.get_resources() + assert len(resources) == 1 + resources.clear() + assert len(mgr.get_resources()) == 1 + + def test_read_resource_sync(self): + """Mock session.read_resource(), verify text extraction.""" + mgr = MCPClientManager({}) + mgr._resource_map = {"file:///readme": ("fs", "file:///readme")} + mock_session = MagicMock() + mgr._sessions["fs"] = mock_session + mgr._loop = asyncio.new_event_loop() + + # Mock the read_resource result + text_content = MagicMock(spec=["text"]) + text_content.text = "Hello, world!" + mock_result = MagicMock() + mock_result.contents = [text_content] + mock_session.read_resource = AsyncMock(return_value=mock_result) + + thread = None + try: + thread = __import__("threading").Thread(target=mgr._loop.run_forever, daemon=True) + thread.start() + output = mgr.read_resource_sync("file:///readme", timeout=5) + assert output == "Hello, world!" + mock_session.read_resource.assert_awaited_once_with("file:///readme") + finally: + mgr._loop.call_soon_threadsafe(mgr._loop.stop) + if thread: + thread.join(timeout=5) + mgr._loop.close() + + def test_read_resource_sync_blob(self): + """Verify base64 blob extraction.""" + mgr = MCPClientManager({}) + mgr._resource_map = {"file:///img.png": ("fs", "file:///img.png")} + mock_session = MagicMock() + mgr._sessions["fs"] = mock_session + mgr._loop = asyncio.new_event_loop() + + blob_content = MagicMock(spec=["blob"]) + blob_content.blob = "aGVsbG8=" + mock_result = MagicMock() + mock_result.contents = [blob_content] + mock_session.read_resource = AsyncMock(return_value=mock_result) + + thread = None + try: + thread = __import__("threading").Thread(target=mgr._loop.run_forever, daemon=True) + thread.start() + output = mgr.read_resource_sync("file:///img.png", timeout=5) + assert output == "aGVsbG8=" + finally: + mgr._loop.call_soon_threadsafe(mgr._loop.stop) + if thread: + thread.join(timeout=5) + mgr._loop.close() + + def test_read_resource_sync_unknown_uri(self): + mgr = MCPClientManager({}) + with pytest.raises(ValueError, match="Unknown MCP resource"): + mgr.read_resource_sync("file:///nonexistent") + + def test_read_resource_sync_disconnected(self): + mgr = MCPClientManager({}) + mgr._resource_map = {"file:///x": ("dead", "file:///x")} + with pytest.raises(RuntimeError, match="not connected"): + mgr.read_resource_sync("file:///x") + + def test_read_resource_sync_timeout(self): + """Verify timeout handling.""" + mgr = MCPClientManager({}) + mgr._resource_map = {"file:///x": ("fs", "file:///x")} + mock_session = MagicMock() + mgr._sessions["fs"] = mock_session + mgr._loop = asyncio.new_event_loop() + + async def _slow_read(_uri: str) -> None: + await asyncio.sleep(10) + + mock_session.read_resource = _slow_read + + thread = None + try: + thread = __import__("threading").Thread(target=mgr._loop.run_forever, daemon=True) + thread.start() + with pytest.raises(TimeoutError): + mgr.read_resource_sync("file:///x", timeout=1) + finally: + mgr._loop.call_soon_threadsafe(mgr._loop.stop) + if thread: + thread.join(timeout=5) + mgr._loop.close() + + def test_resource_listener_notification(self): + """Verify callback fires on rebuild.""" + mgr = MCPClientManager({}) + calls: list[int] = [] + mgr.add_resource_listener(lambda: calls.append(1)) + mgr._per_server_resources = {"a": [_fake_resource_dict()]} + mgr._rebuild_resources() + assert len(calls) == 1 + + def test_resource_listener_remove(self): + mgr = MCPClientManager({}) + calls: list[int] = [] + cb = lambda: calls.append(1) # noqa: E731 + mgr.add_resource_listener(cb) + mgr.remove_resource_listener(cb) + mgr._rebuild_resources() + assert calls == [] + + def test_resource_listener_error_does_not_propagate(self): + mgr = MCPClientManager({}) + mgr.add_resource_listener(lambda: 1 / 0) + mgr._rebuild_resources() # should not raise + + def test_resource_refresh_on_notification(self): + """Mock notification, verify re-fetch of resources.""" + + async def _run() -> None: + mgr = MCPClientManager({}) + mock_session = MagicMock() + mgr._sessions["fs"] = mock_session + mgr._supports_resources["fs"] = True + + # Initial state + mgr._per_server_resources["fs"] = [ + _fake_resource_dict("file:///old", server="fs"), + ] + mgr._rebuild_resources() + assert len(mgr.get_resources()) == 1 + + # Mock the re-fetch returning a new resource + new_res = _fake_mcp_resource("file:///new", "new") + mock_res_result = MagicMock() + mock_res_result.resources = [new_res] + mock_session.list_resources = AsyncMock(return_value=mock_res_result) + mock_tmpl_result = MagicMock() + mock_tmpl_result.resourceTemplates = [] + mock_session.list_resource_templates = AsyncMock(return_value=mock_tmpl_result) + + await mgr._refresh_server_resources("fs") + resources = mgr.get_resources() + assert len(resources) == 1 + assert resources[0]["uri"] == "file:///new" + + asyncio.run(_run()) + + def test_rebuild_resources_empty(self): + mgr = MCPClientManager({}) + mgr._per_server_resources = {} + mgr._rebuild_resources() + assert mgr._resources == [] + assert mgr._resource_map == {} + + def test_rebuild_resources_multi_server(self): + mgr = MCPClientManager({}) + mgr._per_server_resources = { + "fs": [_fake_resource_dict("file:///a", server="fs")], + "db": [_fake_resource_dict("db://table", name="table", server="db")], + } + mgr._rebuild_resources() + assert len(mgr._resources) == 2 + assert mgr._resource_map["file:///a"] == ("fs", "file:///a") + assert mgr._resource_map["db://table"] == ("db", "db://table") + + +# --------------------------------------------------------------------------- +# MCP Prompts +# --------------------------------------------------------------------------- + + +class TestMCPPrompts: + def test_prompt_discovery(self): + """Mock list_prompts(), verify get_prompts() with correct prefixed names.""" + mgr = MCPClientManager({}) + mgr._per_server_prompts = { + "tmpl": [ + _fake_prompt_dict("mcp__tmpl__code_review", "code_review", "tmpl"), + _fake_prompt_dict("mcp__tmpl__summarize", "summarize", "tmpl"), + ], + } + mgr._rebuild_prompts() + prompts = mgr.get_prompts() + assert len(prompts) == 2 + names = {p["name"] for p in prompts} + assert names == {"mcp__tmpl__code_review", "mcp__tmpl__summarize"} + # Verify map entries + assert mgr._prompt_map["mcp__tmpl__code_review"] == ("tmpl", "code_review") + assert mgr._prompt_map["mcp__tmpl__summarize"] == ("tmpl", "summarize") + + def test_rebuild_prompts_copy_on_write(self): + """Verify mutation safety.""" + mgr = MCPClientManager({}) + mgr._per_server_prompts = { + "a": [_fake_prompt_dict("mcp__a__p1", "p1", "a")], + } + mgr._rebuild_prompts() + old_prompts = mgr._prompts + old_map = mgr._prompt_map + mgr._per_server_prompts["b"] = [_fake_prompt_dict("mcp__b__p2", "p2", "b")] + mgr._rebuild_prompts() + assert mgr._prompts is not old_prompts + assert mgr._prompt_map is not old_map + + def test_get_prompts_returns_copy(self): + mgr = MCPClientManager({}) + mgr._per_server_prompts = { + "a": [_fake_prompt_dict("mcp__a__p1", "p1", "a")], + } + mgr._rebuild_prompts() + prompts = mgr.get_prompts() + assert len(prompts) == 1 + prompts.clear() + assert len(mgr.get_prompts()) == 1 + + def test_get_prompt_sync(self): + """Mock session.get_prompt(), verify message conversion.""" + mgr = MCPClientManager({}) + mgr._prompt_map = {"mcp__tmpl__review": ("tmpl", "review")} + mock_session = MagicMock() + mgr._sessions["tmpl"] = mock_session + mgr._loop = asyncio.new_event_loop() + + # Build mock PromptMessage + msg1 = MagicMock() + msg1.role = "user" + msg1.content = MagicMock() + msg1.content.text = "Review this code" + msg2 = MagicMock() + msg2.role = "assistant" + msg2.content = MagicMock() + msg2.content.text = "Looks good!" + mock_result = MagicMock() + mock_result.messages = [msg1, msg2] + mock_session.get_prompt = AsyncMock(return_value=mock_result) + + thread = None + try: + thread = __import__("threading").Thread(target=mgr._loop.run_forever, daemon=True) + thread.start() + messages = mgr.get_prompt_sync( + "mcp__tmpl__review", arguments={"language": "python"}, timeout=5 + ) + assert len(messages) == 2 + assert messages[0] == {"role": "user", "content": "Review this code"} + assert messages[1] == {"role": "assistant", "content": "Looks good!"} + mock_session.get_prompt.assert_awaited_once_with( + "review", arguments={"language": "python"} + ) + finally: + mgr._loop.call_soon_threadsafe(mgr._loop.stop) + if thread: + thread.join(timeout=5) + mgr._loop.close() + + def test_get_prompt_sync_unknown(self): + mgr = MCPClientManager({}) + with pytest.raises(ValueError, match="Unknown MCP prompt"): + mgr.get_prompt_sync("mcp__no__such") + + def test_get_prompt_sync_disconnected(self): + mgr = MCPClientManager({}) + mgr._prompt_map = {"mcp__dead__p": ("dead", "p")} + with pytest.raises(RuntimeError, match="not connected"): + mgr.get_prompt_sync("mcp__dead__p") + + def test_get_prompt_sync_timeout(self): + """Verify timeout handling.""" + mgr = MCPClientManager({}) + mgr._prompt_map = {"mcp__tmpl__slow": ("tmpl", "slow")} + mock_session = MagicMock() + mgr._sessions["tmpl"] = mock_session + mgr._loop = asyncio.new_event_loop() + + async def _slow_prompt(_name: str, *, arguments: dict[str, str] | None = None) -> None: + await asyncio.sleep(10) + + mock_session.get_prompt = _slow_prompt + + thread = None + try: + thread = __import__("threading").Thread(target=mgr._loop.run_forever, daemon=True) + thread.start() + with pytest.raises(TimeoutError): + mgr.get_prompt_sync("mcp__tmpl__slow", timeout=1) + finally: + mgr._loop.call_soon_threadsafe(mgr._loop.stop) + if thread: + thread.join(timeout=5) + mgr._loop.close() + + def test_prompt_listener_notification(self): + """Verify callback fires on rebuild.""" + mgr = MCPClientManager({}) + calls: list[int] = [] + mgr.add_prompt_listener(lambda: calls.append(1)) + mgr._per_server_prompts = {"a": [_fake_prompt_dict()]} + mgr._rebuild_prompts() + assert len(calls) == 1 + + def test_prompt_listener_remove(self): + mgr = MCPClientManager({}) + calls: list[int] = [] + cb = lambda: calls.append(1) # noqa: E731 + mgr.add_prompt_listener(cb) + mgr.remove_prompt_listener(cb) + mgr._rebuild_prompts() + assert calls == [] + + def test_prompt_listener_error_does_not_propagate(self): + mgr = MCPClientManager({}) + mgr.add_prompt_listener(lambda: 1 / 0) + mgr._rebuild_prompts() # should not raise + + def test_is_mcp_prompt(self): + """Verify name lookup.""" + mgr = MCPClientManager({}) + mgr._prompt_map["mcp__tmpl__review"] = ("tmpl", "review") + assert mgr.is_mcp_prompt("mcp__tmpl__review") is True + assert mgr.is_mcp_prompt("nonexistent") is False + + def test_prompt_refresh_on_notification(self): + """Mock notification, verify re-fetch of prompts.""" + + async def _run() -> None: + mgr = MCPClientManager({}) + mock_session = MagicMock() + mgr._sessions["tmpl"] = mock_session + mgr._supports_prompts["tmpl"] = True + + # Initial state + mgr._per_server_prompts["tmpl"] = [ + _fake_prompt_dict("mcp__tmpl__old", "old", "tmpl"), + ] + mgr._rebuild_prompts() + assert len(mgr.get_prompts()) == 1 + + # Mock re-fetch returning a new prompt + new_prompt = _fake_mcp_prompt("new_prompt", "A new prompt") + mock_prompt_result = MagicMock() + mock_prompt_result.prompts = [new_prompt] + mock_session.list_prompts = AsyncMock(return_value=mock_prompt_result) + + await mgr._refresh_server_prompts("tmpl") + prompts = mgr.get_prompts() + assert len(prompts) == 1 + assert prompts[0]["name"] == "mcp__tmpl__new_prompt" + assert prompts[0]["original_name"] == "new_prompt" + + asyncio.run(_run()) + + def test_rebuild_prompts_empty(self): + mgr = MCPClientManager({}) + mgr._per_server_prompts = {} + mgr._rebuild_prompts() + assert mgr._prompts == [] + assert mgr._prompt_map == {} + + def test_rebuild_prompts_multi_server(self): + mgr = MCPClientManager({}) + mgr._per_server_prompts = { + "a": [_fake_prompt_dict("mcp__a__p1", "p1", "a")], + "b": [_fake_prompt_dict("mcp__b__p2", "p2", "b")], + } + mgr._rebuild_prompts() + assert len(mgr._prompts) == 2 + assert mgr._prompt_map["mcp__a__p1"] == ("a", "p1") + assert mgr._prompt_map["mcp__b__p2"] == ("b", "p2") + + +# --------------------------------------------------------------------------- +# Shutdown cleans up new state +# --------------------------------------------------------------------------- + + +class TestShutdownCleanup: + def test_shutdown_clears_resources_and_prompts(self): + mgr = MCPClientManager({}) + mgr._per_server_resources = {"a": [_fake_resource_dict()]} + mgr._rebuild_resources() + mgr._per_server_prompts = {"a": [_fake_prompt_dict()]} + mgr._rebuild_prompts() + assert mgr.get_resources() != [] + assert mgr.get_prompts() != [] + + mgr.shutdown() + assert mgr.get_resources() == [] + assert mgr.get_prompts() == [] + assert mgr._resource_map == {} + assert mgr._prompt_map == {} diff --git a/tests/test_mcp_prompt_sync.py b/tests/test_mcp_prompt_sync.py new file mode 100644 index 00000000..600bb06c --- /dev/null +++ b/tests/test_mcp_prompt_sync.py @@ -0,0 +1,241 @@ +"""Tests for MCP prompt → governance template sync and readonly API guards.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +from turnstone.core.mcp_client import MCPClientManager + + +@pytest.fixture() +def mgr() -> MCPClientManager: + """Create an MCPClientManager with no real servers (no start()).""" + return MCPClientManager({}) + + +def _make_storage() -> MagicMock: + """Create a mock storage backend with prompt template methods.""" + storage = MagicMock() + storage.get_prompt_template_by_name.return_value = None + storage.list_prompt_templates_by_origin.return_value = [] + storage.create_prompt_template.return_value = None + storage.update_prompt_template.return_value = True + storage.delete_prompt_template.return_value = True + return storage + + +class TestSyncPromptsToStorage: + def test_sync_no_storage(self, mgr: MCPClientManager) -> None: + """Without storage set, sync returns empty stats.""" + result = mgr.sync_prompts_to_storage() + assert result == {"added": [], "removed": [], "skipped": []} + + def test_sync_creates_mcp_templates(self, mgr: MCPClientManager) -> None: + """New MCP prompts are created as templates.""" + storage = _make_storage() + mgr.set_storage(storage) + + # Populate internal prompts list directly + mgr._prompts = [ + { + "name": "mcp__test__greeting", + "original_name": "greeting", + "server": "test", + "description": "Say hello", + "arguments": [ + {"name": "name", "description": "Who to greet", "required": True}, + ], + }, + ] + + result = mgr.sync_prompts_to_storage() + + assert result["added"] == ["mcp__test__greeting"] + assert result["removed"] == [] + assert result["skipped"] == [] + storage.create_prompt_template.assert_called_once() + call_kwargs = storage.create_prompt_template.call_args + assert call_kwargs[1]["name"] == "mcp__test__greeting" + assert call_kwargs[1]["origin"] == "mcp" + assert call_kwargs[1]["mcp_server"] == "test" + assert call_kwargs[1]["readonly"] is True + assert call_kwargs[1]["category"] == "mcp" + assert '"name"' in call_kwargs[1]["variables"] + + def test_sync_skips_manual_overrides(self, mgr: MCPClientManager) -> None: + """A manual template with the same name is not overwritten.""" + storage = _make_storage() + storage.get_prompt_template_by_name.return_value = { + "template_id": "existing-id", + "name": "mcp__test__greeting", + "origin": "manual", + "readonly": False, + } + mgr.set_storage(storage) + + mgr._prompts = [ + { + "name": "mcp__test__greeting", + "original_name": "greeting", + "server": "test", + "description": "Say hello", + "arguments": [], + }, + ] + + result = mgr.sync_prompts_to_storage() + + assert result["skipped"] == ["mcp__test__greeting"] + assert result["added"] == [] + storage.create_prompt_template.assert_not_called() + storage.update_prompt_template.assert_not_called() + + def test_sync_updates_existing_mcp_template(self, mgr: MCPClientManager) -> None: + """An existing MCP template gets its content/variables updated.""" + storage = _make_storage() + storage.get_prompt_template_by_name.return_value = { + "template_id": "existing-id", + "name": "mcp__test__greeting", + "origin": "mcp", + "mcp_server": "test", + "readonly": True, + } + mgr.set_storage(storage) + + mgr._prompts = [ + { + "name": "mcp__test__greeting", + "original_name": "greeting", + "server": "test", + "description": "Updated description", + "arguments": [ + {"name": "user", "description": "The user", "required": False}, + ], + }, + ] + + result = mgr.sync_prompts_to_storage() + + assert result["added"] == [] + assert result["skipped"] == [] + storage.create_prompt_template.assert_not_called() + storage.update_prompt_template.assert_called_once() + call_args = storage.update_prompt_template.call_args + assert call_args[0][0] == "existing-id" + assert "Updated description" in call_args[1]["content"] + assert "user" in call_args[1]["variables"] + + def test_sync_removes_deleted_prompts(self, mgr: MCPClientManager) -> None: + """MCP templates in storage with no matching prompt are deleted.""" + storage = _make_storage() + storage.list_prompt_templates_by_origin.return_value = [ + { + "template_id": "old-id", + "name": "mcp__test__old_prompt", + "origin": "mcp", + "mcp_server": "test", + }, + ] + mgr.set_storage(storage) + mgr._prompts = [] # No prompts at all + + result = mgr.sync_prompts_to_storage() + + assert result["removed"] == ["mcp__test__old_prompt"] + storage.delete_prompt_template.assert_called_once_with("old-id") + + +class TestSetStorageAutoSync: + """set_storage() triggers an immediate sync when servers are already connected.""" + + def test_set_storage_syncs_when_connected(self, mgr) -> None: + storage = _make_storage() + mgr._prompts = [ + { + "name": "mcp__srv__p1", + "original_name": "p1", + "server": "srv", + "description": "A prompt", + "arguments": [], + } + ] + mgr._connected.set() + + mgr.set_storage(storage) + + # Should have called create_prompt_template for the discovered prompt + storage.create_prompt_template.assert_called_once() + call_kwargs = storage.create_prompt_template.call_args + assert call_kwargs[1]["name"] == "mcp__srv__p1" + assert call_kwargs[1]["origin"] == "mcp" + + def test_set_storage_no_sync_when_not_connected(self, mgr) -> None: + storage = _make_storage() + mgr._prompts = [ + { + "name": "mcp__srv__p1", + "original_name": "p1", + "server": "srv", + "description": "A prompt", + "arguments": [], + } + ] + # _connected is NOT set + + mgr.set_storage(storage) + + # Should not have synced + storage.create_prompt_template.assert_not_called() + + +class TestReadonlyAPIGuards: + """Test that the console server API guards reject edits to readonly templates.""" + + @pytest.fixture() + def db(self, tmp_path): + """Create a fresh SQLite backend for each test.""" + from turnstone.core.storage._sqlite import SQLiteBackend + + return SQLiteBackend(str(tmp_path / "test.db")) + + def test_readonly_guard_update(self, db) -> None: + """Readonly templates cannot be updated via storage guard logic.""" + db.create_prompt_template( + "t1", + "mcp__srv__prompt", + "mcp", + "content", + variables="[]", + is_default=False, + org_id="", + created_by="", + origin="mcp", + mcp_server="srv", + readonly=True, + ) + tpl = db.get_prompt_template("t1") + assert tpl is not None + assert tpl["readonly"] is True + # Simulate API guard check + assert tpl.get("readonly") is True + + def test_readonly_guard_delete(self, db) -> None: + """Readonly templates are flagged for API-level rejection.""" + db.create_prompt_template( + "t1", + "mcp__srv__prompt", + "mcp", + "content", + variables="[]", + is_default=False, + org_id="", + created_by="", + origin="mcp", + mcp_server="srv", + readonly=True, + ) + existing = db.get_prompt_template("t1") + assert existing is not None + assert existing.get("readonly") is True diff --git a/tests/test_tool_policy.py b/tests/test_tool_policy.py index 85d398ce..d5fde9b5 100644 --- a/tests/test_tool_policy.py +++ b/tests/test_tool_policy.py @@ -84,3 +84,84 @@ def test_first_match_wins(storage): storage.create_tool_policy("p1", "deny-bash", "bash*", "deny", 100) storage.create_tool_policy("p2", "allow-bash", "bash*", "allow", 50) assert evaluate_tool_policy(storage, "bash_exec") == "deny" + + +# --------------------------------------------------------------------------- +# MCP resource and prompt policy patterns +# --------------------------------------------------------------------------- + + +def test_mcp_resource_wildcard_deny(storage): + """Deny all MCP resource reads via glob pattern.""" + storage.create_tool_policy("p1", "block-resources", "mcp_resource__*", "deny", 100) + assert evaluate_tool_policy(storage, "mcp_resource__file:///secret.txt") == "deny" + assert evaluate_tool_policy(storage, "mcp_resource__db://users") == "deny" + assert evaluate_tool_policy(storage, "read_file") is None # unrelated tool + + +def test_mcp_resource_per_server_pattern(storage): + """Allow resources from a specific server, deny others.""" + storage.create_tool_policy("p1", "block-all-resources", "mcp_resource__*", "deny", 50) + storage.create_tool_policy("p2", "allow-docs", "mcp_resource__file:///docs/*", "allow", 100) + assert evaluate_tool_policy(storage, "mcp_resource__file:///docs/readme.md") == "allow" + assert evaluate_tool_policy(storage, "mcp_resource__file:///etc/passwd") == "deny" + + +def test_mcp_prompt_wildcard_ask(storage): + """Require approval for all MCP prompt invocations.""" + storage.create_tool_policy("p1", "ask-prompts", "mcp__*", "ask", 100) + assert evaluate_tool_policy(storage, "mcp__github__code_review") == "ask" + assert evaluate_tool_policy(storage, "mcp__templates__greeting") == "ask" + assert evaluate_tool_policy(storage, "bash") is None + + +def test_mcp_prompt_per_server_allow(storage): + """Auto-approve prompts from a trusted server.""" + storage.create_tool_policy("p1", "ask-all-mcp", "mcp__*", "ask", 50) + storage.create_tool_policy("p2", "allow-trusted", "mcp__trusted__*", "allow", 100) + assert evaluate_tool_policy(storage, "mcp__trusted__greeting") == "allow" + assert evaluate_tool_policy(storage, "mcp__untrusted__evil") == "ask" + + +def test_mcp_batch_mixed(storage): + """Batch evaluation with mixed MCP and built-in tools.""" + storage.create_tool_policy("p1", "block-resources", "mcp_resource__*", "deny", 100) + storage.create_tool_policy("p2", "allow-prompts", "mcp__trusted__*", "allow", 100) + results = evaluate_tool_policies_batch( + storage, + ["mcp_resource__file:///x", "mcp__trusted__greeting", "bash", "mcp__other__y"], + ) + assert results["mcp_resource__file:///x"] == "deny" + assert results["mcp__trusted__greeting"] == "allow" + assert results["bash"] is None + assert results["mcp__other__y"] is None + + +def test_normalize_resource_uri_prevents_traversal(): + """URI normalization resolves .. segments to prevent policy traversal bypass.""" + from turnstone.core.session import ChatSession + + # Normal URI unchanged + assert ChatSession._normalize_resource_uri("file:///docs/readme.md") == "file:///docs/readme.md" + # Traversal resolved + assert ChatSession._normalize_resource_uri("file:///docs/../etc/passwd") == "file:///etc/passwd" + # Double traversal + assert ChatSession._normalize_resource_uri("file:///a/b/../../c") == "file:///c" + # Non-file scheme (netloc preserved, path normalized) + assert ChatSession._normalize_resource_uri("db://host/tables/../secrets") == "db://host/secrets" + # Percent-encoded traversal decoded before normalization + assert ( + ChatSession._normalize_resource_uri("file:///docs/%2e%2e/etc/passwd") + == "file:///etc/passwd" + ) + # Mixed percent-encoded and literal traversal + assert ChatSession._normalize_resource_uri("file:///a/%2e%2e/b/../c") == "file:///c" + + +def test_mcp_tool_granular_policy(storage): + """MCP tool calls use their prefixed func_name for granular policy matching.""" + storage.create_tool_policy("p1", "ask-all-mcp", "mcp__*", "ask", 50) + storage.create_tool_policy("p2", "allow-github", "mcp__github__*", "allow", 100) + # MCP tools now use func_name as approval_label + assert evaluate_tool_policy(storage, "mcp__github__search") == "allow" + assert evaluate_tool_policy(storage, "mcp__untrusted__exec") == "ask" diff --git a/tests/test_tools_schema.py b/tests/test_tools_schema.py index 8787b745..03a03087 100644 --- a/tests/test_tools_schema.py +++ b/tests/test_tools_schema.py @@ -72,16 +72,24 @@ class TestToolsMetadata: """Validate the metadata extracted from JSON files.""" def test_tool_count(self): - assert len(TOOLS) == 16 + assert len(TOOLS) == 18 def test_agent_tools_count(self): - assert len(AGENT_TOOLS) == 7 + assert len(AGENT_TOOLS) == 9 def test_task_agent_tools_count(self): - assert len(TASK_AGENT_TOOLS) == 10 + assert len(TASK_AGENT_TOOLS) == 12 def test_auto_approve_sets_match(self): - expected = {"read_file", "search", "math", "man", "web_fetch", "web_search", "notify"} + expected = { + "read_file", + "search", + "math", + "man", + "web_fetch", + "web_search", + "notify", + } assert expected == AGENT_AUTO_TOOLS assert expected == TASK_AUTO_TOOLS @@ -103,6 +111,8 @@ class TestToolsMetadata: "forget": "key", "notify": "message", "watch": "command", + "read_resource": "uri", + "use_prompt": "name", } assert expected == PRIMARY_KEY_MAP diff --git a/turnstone/api/console_schemas.py b/turnstone/api/console_schemas.py index 28742970..50434598 100644 --- a/turnstone/api/console_schemas.py +++ b/turnstone/api/console_schemas.py @@ -286,6 +286,9 @@ class PromptTemplateInfo(BaseModel): is_default: bool org_id: str created_by: str + origin: str = "manual" + mcp_server: str = "" + readonly: bool = False created: str updated: str diff --git a/turnstone/api/server_schemas.py b/turnstone/api/server_schemas.py index 35aacc31..b2136d51 100644 --- a/turnstone/api/server_schemas.py +++ b/turnstone/api/server_schemas.py @@ -143,6 +143,12 @@ class WorkstreamCounts(BaseModel): error: int = 0 +class McpStatus(BaseModel): + servers: int = 0 + resources: int = 0 + prompts: int = 0 + + class HealthResponse(BaseModel): status: str = Field(examples=["ok", "degraded"]) version: str = "" @@ -150,3 +156,4 @@ class HealthResponse(BaseModel): model: str = "" workstreams: WorkstreamCounts = WorkstreamCounts() backend: BackendStatus | None = None + mcp: McpStatus | None = None diff --git a/turnstone/cli.py b/turnstone/cli.py index 036d3dd4..546fb6d2 100644 --- a/turnstone/cli.py +++ b/turnstone/cli.py @@ -1002,6 +1002,9 @@ def main() -> None: mcp_tools = mcp_client.get_tools() if mcp_tools: print(f"MCP tools: {len(mcp_tools)} from {mcp_client.server_count} server(s)") + from turnstone.core.storage import get_storage as _cli_get_storage + + mcp_client.set_storage(_cli_get_storage()) print("Type /help for commands, /ws for workstreams, /exit or Ctrl+D to quit.\n") # Prompt string -- use a short display name diff --git a/turnstone/console/collector.py b/turnstone/console/collector.py index 1e203405..96559532 100644 --- a/turnstone/console/collector.py +++ b/turnstone/console/collector.py @@ -320,6 +320,9 @@ class ClusterCollector: total_tokens = 0 total_tool_calls = 0 total_ws = 0 + mcp_servers = 0 + mcp_resources = 0 + mcp_prompts = 0 versions: set[str] = set() with self._lock: for node in self._nodes.values(): @@ -332,8 +335,12 @@ class ClusterCollector: ver = node.health.get("version", "") if ver: versions.add(ver) + mcp = node.health.get("mcp", {}) + mcp_servers += mcp.get("servers", 0) + mcp_resources += mcp.get("resources", 0) + mcp_prompts += mcp.get("prompts", 0) node_count = len(self._nodes) - return { + result: dict[str, Any] = { "nodes": node_count, "workstreams": total_ws, "states": states, @@ -344,6 +351,11 @@ class ClusterCollector: "version_drift": len(versions) > 1, "versions": sorted(versions), } + if mcp_servers: + result["mcp_servers"] = mcp_servers + result["mcp_resources"] = mcp_resources + result["mcp_prompts"] = mcp_prompts + return result def get_version_info(self) -> dict[str, Any]: """Return per-node version map and drift flag.""" @@ -515,6 +527,9 @@ class ClusterCollector: total_tokens = 0 total_tool_calls = 0 total_ws = 0 + mcp_servers = 0 + mcp_resources = 0 + mcp_prompts = 0 versions: set[str] = set() for node in self._nodes.values(): @@ -530,6 +545,10 @@ class ClusterCollector: ver = node.health.get("version", "") if ver: versions.add(ver) + mcp = node.health.get("mcp", {}) + mcp_servers += mcp.get("servers", 0) + mcp_resources += mcp.get("resources", 0) + mcp_prompts += mcp.get("prompts", 0) nodes_out.append( { @@ -546,19 +565,25 @@ class ClusterCollector: node_count = len(self._nodes) + overview: dict[str, Any] = { + "nodes": node_count, + "workstreams": total_ws, + "states": states, + "aggregate": { + "total_tokens": total_tokens, + "total_tool_calls": total_tool_calls, + }, + "version_drift": len(versions) > 1, + "versions": sorted(versions), + } + if mcp_servers: + overview["mcp_servers"] = mcp_servers + overview["mcp_resources"] = mcp_resources + overview["mcp_prompts"] = mcp_prompts + return { "nodes": nodes_out, - "overview": { - "nodes": node_count, - "workstreams": total_ws, - "states": states, - "aggregate": { - "total_tokens": total_tokens, - "total_tool_calls": total_tool_calls, - }, - "version_drift": len(versions) > 1, - "versions": sorted(versions), - }, + "overview": overview, "timestamp": time.time(), } diff --git a/turnstone/console/server.py b/turnstone/console/server.py index 2442821d..285abddc 100644 --- a/turnstone/console/server.py +++ b/turnstone/console/server.py @@ -2074,6 +2074,8 @@ async def admin_update_template(request: Request) -> JSONResponse: existing = storage.get_prompt_template(template_id) if existing is None: return JSONResponse({"error": "Template not found"}, status_code=404) + if existing.get("readonly"): + return JSONResponse({"error": "MCP-sourced templates are read-only"}, status_code=403) body = await read_json_or_400(request) if isinstance(body, JSONResponse): @@ -2130,6 +2132,8 @@ async def admin_delete_template(request: Request) -> JSONResponse: existing = storage.get_prompt_template(template_id) if existing is None: return JSONResponse({"error": "Template not found"}, status_code=404) + if existing.get("readonly"): + return JSONResponse({"error": "MCP-sourced templates are read-only"}, status_code=403) storage.delete_prompt_template(template_id) diff --git a/turnstone/console/static/app.js b/turnstone/console/static/app.js index 1e5bfb50..7fbdb0fc 100644 --- a/turnstone/console/static/app.js +++ b/turnstone/console/static/app.js @@ -140,6 +140,9 @@ function recomputeOverview() { var totalTokens = 0, totalToolCalls = 0, totalWs = 0; + var mcpServers = 0, + mcpResources = 0, + mcpPrompts = 0; var versions = {}; Object.keys(clusterState.nodes).forEach(function (nid) { var node = clusterState.nodes[nid]; @@ -154,6 +157,10 @@ function recomputeOverview() { totalTokens += aggTokens || nodeWsTokens; totalToolCalls += (node.aggregate || {}).total_tool_calls || 0; if (node.version) versions[node.version] = true; + var mcp = (node.health || {}).mcp || {}; + mcpServers += mcp.servers || 0; + mcpResources += mcp.resources || 0; + mcpPrompts += mcp.prompts || 0; }); var versionList = Object.keys(versions).sort(); clusterState.overview = { @@ -167,6 +174,11 @@ function recomputeOverview() { version_drift: versionList.length > 1, versions: versionList, }; + if (mcpServers > 0) { + clusterState.overview.mcp_servers = mcpServers; + clusterState.overview.mcp_resources = mcpResources; + clusterState.overview.mcp_prompts = mcpPrompts; + } } function buildNodeInfoFromSnapshot(node) { @@ -227,6 +239,23 @@ function renderFromState() { }).length; document.getElementById("node-ws-summary").textContent = active + " active \u00b7 " + wsList.length + " total"; + var mcpSumEl = document.getElementById("node-mcp-summary"); + if (mcpSumEl) { + var mcpInfo = snapNode.health && snapNode.health.mcp; + if (mcpInfo && mcpInfo.servers > 0) { + mcpSumEl.textContent = + mcpInfo.servers + + " MCP server" + + (mcpInfo.servers !== 1 ? "s" : "") + + " \u00b7 " + + mcpInfo.resources + + " resources \u00b7 " + + mcpInfo.prompts + + " prompts"; + } else { + mcpSumEl.textContent = ""; + } + } renderWsTable(document.getElementById("node-ws-table"), wsList); } } else if (currentView === "filtered") { @@ -470,6 +499,43 @@ function renderStatusBar(overview) { verEl.appendChild(verLbl); metricsContainer.appendChild(verEl); } + // MCP aggregate metrics + if (overview.mcp_servers && overview.mcp_servers > 0) { + var mcpDivider = document.createElement("span"); + mcpDivider.className = "csb-divider"; + mcpDivider.setAttribute("aria-hidden", "true"); + metricsContainer.appendChild(mcpDivider); + var mcpTitles = { + mcp: "MCP servers", + rsrc: "MCP resources", + pmpt: "MCP prompts", + }; + var mcpMetrics = [ + { value: overview.mcp_servers, label: "mcp" }, + { value: overview.mcp_resources, label: "rsrc" }, + { value: overview.mcp_prompts, label: "pmpt" }, + ]; + mcpMetrics.forEach(function (m) { + var el = document.createElement("span"); + el.className = "csb-metric"; + el.title = mcpTitles[m.label] || ""; + if (m.label === "mcp") { + var dot = document.createElement("span"); + dot.className = "csb-mcp-dot"; + dot.setAttribute("aria-hidden", "true"); + el.appendChild(dot); + } + var valSpan = document.createElement("span"); + valSpan.className = "csb-metric-value"; + valSpan.textContent = formatCount(m.value); + var labelSpan = document.createElement("span"); + labelSpan.className = "csb-metric-label"; + labelSpan.textContent = m.label; + el.appendChild(valSpan); + el.appendChild(labelSpan); + metricsContainer.appendChild(el); + }); + } } // --- Node Grouping --- diff --git a/turnstone/console/static/governance.js b/turnstone/console/static/governance.js index 3c01bbed..7f5c7b0a 100644 --- a/turnstone/console/static/governance.js +++ b/turnstone/console/static/governance.js @@ -689,14 +689,23 @@ function _renderGovTemplates(items) { var defBadge = t.is_default ? 'default' : ""; + var originBadge = + t.origin === "mcp" + ? ' mcp:' + + escapeHtml(t.mcp_server) + + "" + : ""; var catBadge = '' + escapeHtml(t.category) + ""; + var editDisabled = t.readonly ? " disabled" : ""; + var deleteDisabled = t.readonly ? " disabled" : ""; html += '
' + '' + escapeHtml(t.name) + " " + defBadge + + originBadge + "" + '' + catBadge + @@ -707,12 +716,16 @@ function _renderGovTemplates(items) { '' + '' + + '"' + + editDisabled + + ">edit" + '' + + '"' + + deleteDisabled + + ">delete" + "
"; } el.innerHTML = html; diff --git a/turnstone/console/static/index.html b/turnstone/console/static/index.html index 17cad5d6..c2dd7ac9 100644 --- a/turnstone/console/static/index.html +++ b/turnstone/console/static/index.html @@ -41,6 +41,7 @@
WORKSTREAMS +

turnstone

+ diff --git a/turnstone/ui/static/style.css b/turnstone/ui/static/style.css index ce66979c..0996a89b 100644 --- a/turnstone/ui/static/style.css +++ b/turnstone/ui/static/style.css @@ -12,6 +12,14 @@ font-family: var(--font-display); letter-spacing: 0.02em; } +#mcp-status { + color: var(--magenta); + font-size: 11px; + font-family: var(--font-mono); + opacity: 0; + transition: opacity 0.3s; + cursor: default; +} #health-indicator { font-size: 11px; padding: 2px 8px; @@ -699,5 +707,6 @@ .approval-btn, .approval-feedback-input, #plan-buttons button, #input-area button, .dashboard-new-btn, .dashboard-input, - #health-indicator, #hamburger-btn { transition: none; } + #health-indicator, #hamburger-btn, + #mcp-status { transition: none; } }