diff --git a/CLAUDE.md b/CLAUDE.md index 1c5d81eb..e55bcc24 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -319,12 +319,106 @@ Key test patterns: ## Adding a New Tool +### Built-in Tools (Rust) + 1. Create `src/tools/builtin/my_tool.rs` 2. Implement the `Tool` trait 3. Add `mod my_tool;` and `pub use` in `src/tools/builtin/mod.rs` 4. Register in `ToolRegistry::register_builtin_tools()` in `registry.rs` 5. Add tests +### WASM Tools (Recommended) + +WASM tools are the preferred way to add new capabilities. They run in a sandboxed environment with explicit capabilities. + +1. Create a new crate in `examples/wasm-tools//` +2. Implement the WIT interface (`wit/tool.wit`) +3. Create `.capabilities.json` declaring required permissions +4. Build with `cargo build --target wasm32-wasip2 --release` +5. Install with `ironclaw tool install path/to/tool.wasm` + +See `examples/wasm-tools/` for examples. + +## Tool Architecture Principles + +**CRITICAL: Keep tool-specific logic out of the main agent codebase.** + +The main agent provides generic infrastructure; tools are self-contained units that declare their requirements through capabilities files. + +### What Goes in Tools (capabilities.json) + +- API endpoints the tool needs (HTTP allowlist) +- Credentials required (secret names, injection locations) +- Rate limits and timeouts +- Auth setup instructions (see below) +- Workspace paths the tool can read + +### What Does NOT Go in Main Agent + +- Service-specific auth flows (OAuth for Notion, Slack, etc.) +- Service-specific CLI commands (`auth notion`, `auth slack`) +- Service-specific configuration handling +- Hardcoded API URLs or token formats + +### Tool Authentication + +Tools declare their auth requirements in `.capabilities.json` under the `auth` section. Two methods are supported: + +#### OAuth (Browser-based login) + +For services that support OAuth, users just click through browser login: + +```json +{ + "auth": { + "secret_name": "notion_api_token", + "display_name": "Notion", + "oauth": { + "authorization_url": "https://api.notion.com/v1/oauth/authorize", + "token_url": "https://api.notion.com/v1/oauth/token", + "client_id_env": "NOTION_OAUTH_CLIENT_ID", + "client_secret_env": "NOTION_OAUTH_CLIENT_SECRET", + "scopes": [], + "use_pkce": false, + "extra_params": { "owner": "user" } + }, + "env_var": "NOTION_TOKEN" + } +} +``` + +To enable OAuth for a tool: +1. Register a public OAuth app with the service (e.g., notion.so/my-integrations) +2. Configure redirect URIs: `http://localhost:9876/callback` through `http://localhost:9886/callback` +3. Set environment variables for client_id and client_secret + +#### Manual Token Entry (Fallback) + +For services without OAuth or when OAuth isn't configured: + +```json +{ + "auth": { + "secret_name": "openai_api_key", + "display_name": "OpenAI", + "instructions": "Get your API key from platform.openai.com/api-keys", + "setup_url": "https://platform.openai.com/api-keys", + "token_hint": "Starts with 'sk-'", + "env_var": "OPENAI_API_KEY" + } +} +``` + +#### Auth Flow Priority + +When running `ironclaw tool auth `: + +1. Check `env_var` - if set in environment, use it directly +2. Check `oauth` - if configured, open browser for OAuth flow +3. Fall back to `instructions` + manual token entry + +The agent reads auth config from the tool's capabilities file and provides the appropriate flow. No service-specific code in the main agent. + ## Adding a New Channel 1. Create `src/channels/my_channel.rs` diff --git a/Cargo.lock b/Cargo.lock index d8635375..655ae87f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2118,6 +2118,7 @@ dependencies = [ "anyhow", "async-trait", "axum", + "base64 0.22.1", "blake3", "bollard", "bytes", @@ -3256,6 +3257,7 @@ dependencies = [ "base64 0.22.1", "bytes", "futures-core", + "futures-util", "http", "http-body", "http-body-util", @@ -3275,12 +3277,14 @@ dependencies = [ "sync_wrapper", "tokio", "tokio-rustls", + "tokio-util", "tower", "tower-http", "tower-service", "url", "wasm-bindgen", "wasm-bindgen-futures", + "wasm-streams", "web-sys", "webpki-roots", ] @@ -4734,6 +4738,19 @@ dependencies = [ "wasmparser 0.244.0", ] +[[package]] +name = "wasm-streams" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "wasmparser" version = "0.220.1" diff --git a/Cargo.toml b/Cargo.toml index 2dc13f28..45524ee3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,7 +13,7 @@ tokio-stream = "0.1" futures = "0.3" # HTTP client -reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } +reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "stream"] } # Serialization serde = { version = "1", features = ["derive"] } @@ -97,6 +97,7 @@ hyper = { version = "1.5", features = ["server", "http1", "http2"] } hyper-util = { version = "0.1", features = ["server", "tokio", "http1", "http2"] } http-body-util = "0.1" bytes = "1" +base64 = "0.22.1" # macOS keychain [target.'cfg(target_os = "macos")'.dependencies] diff --git a/examples/wasm-tools/notion/.gitignore b/examples/wasm-tools/notion/.gitignore deleted file mode 100644 index ea8c4bf7..00000000 --- a/examples/wasm-tools/notion/.gitignore +++ /dev/null @@ -1 +0,0 @@ -/target diff --git a/examples/wasm-tools/notion/Cargo.lock b/examples/wasm-tools/notion/Cargo.lock deleted file mode 100644 index 1e971b88..00000000 --- a/examples/wasm-tools/notion/Cargo.lock +++ /dev/null @@ -1,401 +0,0 @@ -# This file is automatically @generated by Cargo. -# It is not intended for manual editing. -version = 4 - -[[package]] -name = "ahash" -version = "0.8.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" -dependencies = [ - "cfg-if", - "once_cell", - "version_check", - "zerocopy", -] - -[[package]] -name = "anyhow" -version = "1.0.100" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a23eb6b1614318a8071c9b2521f36b424b2c83db5eb3a0fead4a6c0809af6e61" - -[[package]] -name = "bitflags" -version = "2.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "812e12b5285cc515a9c72a5c1d3b6d46a19dac5acfef5265968c166106e31dd3" - -[[package]] -name = "cfg-if" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" - -[[package]] -name = "equivalent" -version = "1.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" - -[[package]] -name = "hashbrown" -version = "0.14.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" -dependencies = [ - "ahash", -] - -[[package]] -name = "hashbrown" -version = "0.16.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" - -[[package]] -name = "heck" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" - -[[package]] -name = "id-arena" -version = "2.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" - -[[package]] -name = "indexmap" -version = "2.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7714e70437a7dc3ac8eb7e6f8df75fd8eb422675fc7678aff7364301092b1017" -dependencies = [ - "equivalent", - "hashbrown 0.16.1", - "serde", - "serde_core", -] - -[[package]] -name = "itoa" -version = "1.0.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" - -[[package]] -name = "leb128" -version = "0.2.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "884e2677b40cc8c339eaefcb701c32ef1fd2493d71118dc0ca4b6a736c93bd67" - -[[package]] -name = "log" -version = "0.4.29" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" - -[[package]] -name = "memchr" -version = "2.7.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f52b00d39961fc5b2736ea853c9cc86238e165017a493d1d5c8eac6bdc4cc273" - -[[package]] -name = "notion-tool" -version = "0.1.0" -dependencies = [ - "serde", - "serde_json", - "wit-bindgen", -] - -[[package]] -name = "once_cell" -version = "1.21.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" - -[[package]] -name = "prettyplease" -version = "0.2.37" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" -dependencies = [ - "proc-macro2", - "syn", -] - -[[package]] -name = "proc-macro2" -version = "1.0.106" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" -dependencies = [ - "unicode-ident", -] - -[[package]] -name = "quote" -version = "1.0.44" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "21b2ebcf727b7760c461f091f9f0f539b77b8e87f2fd88131e7f1b433b3cece4" -dependencies = [ - "proc-macro2", -] - -[[package]] -name = "semver" -version = "1.0.27" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2" - -[[package]] -name = "serde" -version = "1.0.228" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" -dependencies = [ - "serde_core", - "serde_derive", -] - -[[package]] -name = "serde_core" -version = "1.0.228" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" -dependencies = [ - "serde_derive", -] - -[[package]] -name = "serde_derive" -version = "1.0.228" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "serde_json" -version = "1.0.149" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" -dependencies = [ - "itoa", - "memchr", - "serde", - "serde_core", - "zmij", -] - -[[package]] -name = "smallvec" -version = "1.15.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" - -[[package]] -name = "spdx" -version = "0.10.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c3e17e880bafaeb362a7b751ec46bdc5b61445a188f80e0606e68167cd540fa3" -dependencies = [ - "smallvec", -] - -[[package]] -name = "syn" -version = "2.0.114" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4d107df263a3013ef9b1879b0df87d706ff80f65a86ea879bd9c31f9b307c2a" -dependencies = [ - "proc-macro2", - "quote", - "unicode-ident", -] - -[[package]] -name = "unicode-ident" -version = "1.0.22" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5" - -[[package]] -name = "unicode-xid" -version = "0.2.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" - -[[package]] -name = "version_check" -version = "0.9.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" - -[[package]] -name = "wasm-encoder" -version = "0.220.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e913f9242315ca39eff82aee0e19ee7a372155717ff0eb082c741e435ce25ed1" -dependencies = [ - "leb128", - "wasmparser", -] - -[[package]] -name = "wasm-metadata" -version = "0.220.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "185dfcd27fa5db2e6a23906b54c28199935f71d9a27a1a27b3a88d6fee2afae7" -dependencies = [ - "anyhow", - "indexmap", - "serde", - "serde_derive", - "serde_json", - "spdx", - "wasm-encoder", - "wasmparser", -] - -[[package]] -name = "wasmparser" -version = "0.220.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8d07b6a3b550fefa1a914b6d54fc175dd11c3392da11eee604e6ffc759805d25" -dependencies = [ - "ahash", - "bitflags", - "hashbrown 0.14.5", - "indexmap", - "semver", -] - -[[package]] -name = "wit-bindgen" -version = "0.36.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a2b3e15cd6068f233926e7d8c7c588b2ec4fb7cc7bf3824115e7c7e2a8485a3" -dependencies = [ - "wit-bindgen-rt", - "wit-bindgen-rust-macro", -] - -[[package]] -name = "wit-bindgen-core" -version = "0.36.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b632a5a0fa2409489bd49c9e6d99fcc61bb3d4ce9d1907d44662e75a28c71172" -dependencies = [ - "anyhow", - "heck", - "wit-parser", -] - -[[package]] -name = "wit-bindgen-rt" -version = "0.36.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7947d0131c7c9da3f01dfde0ab8bd4c4cf3c5bd49b6dba0ae640f1fa752572ea" -dependencies = [ - "bitflags", -] - -[[package]] -name = "wit-bindgen-rust" -version = "0.36.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4329de4186ee30e2ef30a0533f9b3c123c019a237a7c82d692807bf1b3ee2697" -dependencies = [ - "anyhow", - "heck", - "indexmap", - "prettyplease", - "syn", - "wasm-metadata", - "wit-bindgen-core", - "wit-component", -] - -[[package]] -name = "wit-bindgen-rust-macro" -version = "0.36.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "177fb7ee1484d113b4792cc480b1ba57664bbc951b42a4beebe573502135b1fc" -dependencies = [ - "anyhow", - "prettyplease", - "proc-macro2", - "quote", - "syn", - "wit-bindgen-core", - "wit-bindgen-rust", -] - -[[package]] -name = "wit-component" -version = "0.220.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b505603761ed400c90ed30261f44a768317348e49f1864e82ecdc3b2744e5627" -dependencies = [ - "anyhow", - "bitflags", - "indexmap", - "log", - "serde", - "serde_derive", - "serde_json", - "wasm-encoder", - "wasm-metadata", - "wasmparser", - "wit-parser", -] - -[[package]] -name = "wit-parser" -version = "0.220.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae2a7999ed18efe59be8de2db9cb2b7f84d88b27818c79353dfc53131840fe1a" -dependencies = [ - "anyhow", - "id-arena", - "indexmap", - "log", - "semver", - "serde", - "serde_derive", - "serde_json", - "unicode-xid", - "wasmparser", -] - -[[package]] -name = "zerocopy" -version = "0.8.38" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57cf3aa6855b23711ee9852dfc97dfaa51c45feaba5b645d0c777414d494a961" -dependencies = [ - "zerocopy-derive", -] - -[[package]] -name = "zerocopy-derive" -version = "0.8.38" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a616990af1a287837c4fe6596ad77ef57948f787e46ce28e166facc0cc1cb75" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "zmij" -version = "1.0.19" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ff05f8caa9038894637571ae6b9e29466c1f4f829d26c9b28f869a29cbe3445" diff --git a/examples/wasm-tools/notion/Cargo.toml b/examples/wasm-tools/notion/Cargo.toml deleted file mode 100644 index cc3d8764..00000000 --- a/examples/wasm-tools/notion/Cargo.toml +++ /dev/null @@ -1,21 +0,0 @@ -[package] -name = "notion-tool" -version = "0.1.0" -edition = "2021" -description = "Notion integration tool for NEAR Agent (WASM component)" -license = "MIT OR Apache-2.0" -publish = false - -[lib] -crate-type = ["cdylib"] - -[dependencies] -wit-bindgen = "0.36" -serde = { version = "1", features = ["derive"] } -serde_json = "1" - -[profile.release] -opt-level = "s" -lto = true -strip = true -codegen-units = 1 diff --git a/examples/wasm-tools/notion/README.md b/examples/wasm-tools/notion/README.md deleted file mode 100644 index 549d9b10..00000000 --- a/examples/wasm-tools/notion/README.md +++ /dev/null @@ -1,186 +0,0 @@ -# Notion WASM Tool - -A WASM-sandboxed Notion integration for NEAR Agent. Follows the MCP server pattern with domain-grouped operations. - -## Building - -```bash -cargo build --release --target wasm32-wasip1 -``` - -The compiled module will be at `target/wasm32-wasip1/release/notion_tool.wasm`. - -## Capabilities - -This tool requires: - -- **HTTP**: Access to `api.notion.com/v1/*` (GET, POST, PATCH, DELETE) -- **Secrets**: `notion_api_token` (injected as Bearer authorization) - -See `notion_tool.capabilities.json` for the full capability configuration. - -## Setup - -1. Create a Notion integration at https://www.notion.so/my-integrations -2. Copy the "Internal Integration Secret" -3. Add it to the agent's secrets store as `notion_api_token` -4. Share your Notion pages/databases with the integration - -## Supported Actions - -### Search - -| Action | Required Params | Optional Params | -|--------|-----------------|-----------------| -| `search` | - | `query`, `filter`, `page_size`, `start_cursor` | - -### Pages - -| Action | Required Params | Optional Params | -|--------|-----------------|-----------------| -| `get_page` | `page_id` | - | -| `create_page` | `parent`, `properties` | `children`, `icon`, `cover` | -| `update_page` | `page_id`, `properties` | `icon`, `cover` | -| `archive_page` | `page_id` | - | -| `restore_page` | `page_id` | - | - -### Blocks - -| Action | Required Params | Optional Params | -|--------|-----------------|-----------------| -| `get_blocks` | `block_id` | `page_size`, `start_cursor` | -| `append_blocks` | `block_id`, `children` | `after` | -| `get_block` | `block_id` | - | -| `update_block` | `block_id`, `content` | - | -| `delete_block` | `block_id` | - | - -### Databases - -| Action | Required Params | Optional Params | -|--------|-----------------|-----------------| -| `get_database` | `database_id` | - | -| `query_database` | `database_id` | `filter`, `sorts`, `page_size`, `start_cursor` | -| `create_database` | `parent`, `title`, `properties` | `icon`, `cover`, `is_inline` | -| `update_database` | `database_id` | `title`, `properties` | - -### Comments - -| Action | Required Params | Optional Params | -|--------|-----------------|-----------------| -| `get_comments` | `block_id` | `page_size`, `start_cursor` | -| `add_comment` | `parent`, `rich_text` | - | - -### Users - -| Action | Required Params | Optional Params | -|--------|-----------------|-----------------| -| `list_users` | - | `page_size`, `start_cursor` | -| `get_user` | `user_id` | - | -| `get_me` | - | - | - -## Examples - -### Search for pages - -```json -{ - "action": "search", - "query": "meeting notes", - "filter": { "property": "object", "value": "page" }, - "page_size": 10 -} -``` - -### Query a database with filters - -```json -{ - "action": "query_database", - "database_id": "abc123-def456-...", - "filter": { - "property": "Status", - "select": { "equals": "Done" } - }, - "sorts": [ - { "property": "Created", "direction": "descending" } - ], - "page_size": 20 -} -``` - -### Create a page in a database - -```json -{ - "action": "create_page", - "parent": { "database_id": "abc123-def456-..." }, - "properties": { - "Name": { - "title": [{ "text": { "content": "New Task" } }] - }, - "Status": { - "select": { "name": "To Do" } - } - }, - "children": [ - { - "object": "block", - "type": "paragraph", - "paragraph": { - "rich_text": [{ "text": { "content": "Task description here." } }] - } - } - ] -} -``` - -### Append content to a page - -```json -{ - "action": "append_blocks", - "block_id": "page-id-here", - "children": [ - { - "object": "block", - "type": "heading_2", - "heading_2": { - "rich_text": [{ "text": { "content": "New Section" } }] - } - }, - { - "object": "block", - "type": "paragraph", - "paragraph": { - "rich_text": [{ "text": { "content": "Some content..." } }] - } - } - ] -} -``` - -### Add a comment - -```json -{ - "action": "add_comment", - "parent": { "page_id": "page-id-here" }, - "rich_text": [ - { "text": { "content": "This is a comment from the agent." } } - ] -} -``` - -## Security - -- Runs in WASM sandbox with fuel metering and memory limits -- API token injected at host boundary (never visible to WASM) -- Only `api.notion.com/v1/*` endpoints allowed -- Responses scanned for secret leakage -- Rate limited: 50 req/min, 1000 req/hour - -## References - -- [Notion API Documentation](https://developers.notion.com/reference/intro) -- [Notion MCP Server](https://github.com/makenotion/notion-mcp-server) -- [NEAR Agent WASM Tool System](../../../src/tools/wasm/) diff --git a/examples/wasm-tools/notion/notion_tool.capabilities.json b/examples/wasm-tools/notion/notion_tool.capabilities.json deleted file mode 100644 index f64eadc2..00000000 --- a/examples/wasm-tools/notion/notion_tool.capabilities.json +++ /dev/null @@ -1,26 +0,0 @@ -{ - "http": { - "allowlist": [ - { - "host": "api.notion.com", - "path_prefix": "/v1/", - "methods": ["GET", "POST", "PATCH", "DELETE"] - } - ], - "credentials": { - "notion_api_token": { - "secret_name": "notion_api_token", - "location": { "type": "bearer" }, - "host_patterns": ["api.notion.com"] - } - }, - "rate_limit": { - "requests_per_minute": 50, - "requests_per_hour": 1000 - }, - "timeout_secs": 30 - }, - "secrets": { - "allowed_names": ["notion_api_token"] - } -} diff --git a/examples/wasm-tools/notion/src/api.rs b/examples/wasm-tools/notion/src/api.rs deleted file mode 100644 index fd9bc5b0..00000000 --- a/examples/wasm-tools/notion/src/api.rs +++ /dev/null @@ -1,445 +0,0 @@ -//! Notion API implementation. -//! -//! All API calls go through the host's HTTP capability, which handles -//! credential injection and rate limiting. The WASM tool never sees -//! the actual API token. - -use crate::near::agent::host; -use crate::types::*; - -const NOTION_API_BASE: &str = "https://api.notion.com/v1"; -const NOTION_VERSION: &str = "2022-06-28"; - -/// Make a Notion API GET request. -fn notion_get(path: &str) -> Result { - let url = format!("{}{}", NOTION_API_BASE, path); - - let headers = serde_json::json!({ - "Notion-Version": NOTION_VERSION - }); - let headers_str = serde_json::to_string(&headers).map_err(|e| e.to_string())?; - - host::log(host::LogLevel::Debug, &format!("Notion GET: {}", path)); - - let response = host::http_request("GET", &url, &headers_str, None)?; - - if response.status < 200 || response.status >= 300 { - let body = String::from_utf8_lossy(&response.body); - return Err(format!( - "Notion API returned status {}: {}", - response.status, body - )); - } - - let body = String::from_utf8(response.body) - .map_err(|e| format!("Invalid UTF-8 in response: {}", e))?; - serde_json::from_str(&body).map_err(|e| format!("Invalid JSON in response: {}", e)) -} - -/// Make a Notion API POST request. -fn notion_post(path: &str, body: &serde_json::Value) -> Result { - let url = format!("{}{}", NOTION_API_BASE, path); - - let headers = serde_json::json!({ - "Notion-Version": NOTION_VERSION, - "Content-Type": "application/json" - }); - let headers_str = serde_json::to_string(&headers).map_err(|e| e.to_string())?; - let body_str = serde_json::to_string(body).map_err(|e| e.to_string())?; - - host::log(host::LogLevel::Debug, &format!("Notion POST: {}", path)); - - let response = host::http_request("POST", &url, &headers_str, Some(body_str.as_bytes()))?; - - if response.status < 200 || response.status >= 300 { - let body = String::from_utf8_lossy(&response.body); - return Err(format!( - "Notion API returned status {}: {}", - response.status, body - )); - } - - let body = String::from_utf8(response.body) - .map_err(|e| format!("Invalid UTF-8 in response: {}", e))?; - serde_json::from_str(&body).map_err(|e| format!("Invalid JSON in response: {}", e)) -} - -/// Make a Notion API PATCH request. -fn notion_patch(path: &str, body: &serde_json::Value) -> Result { - let url = format!("{}{}", NOTION_API_BASE, path); - - let headers = serde_json::json!({ - "Notion-Version": NOTION_VERSION, - "Content-Type": "application/json" - }); - let headers_str = serde_json::to_string(&headers).map_err(|e| e.to_string())?; - let body_str = serde_json::to_string(body).map_err(|e| e.to_string())?; - - host::log(host::LogLevel::Debug, &format!("Notion PATCH: {}", path)); - - let response = host::http_request("PATCH", &url, &headers_str, Some(body_str.as_bytes()))?; - - if response.status < 200 || response.status >= 300 { - let body = String::from_utf8_lossy(&response.body); - return Err(format!( - "Notion API returned status {}: {}", - response.status, body - )); - } - - let body = String::from_utf8(response.body) - .map_err(|e| format!("Invalid UTF-8 in response: {}", e))?; - serde_json::from_str(&body).map_err(|e| format!("Invalid JSON in response: {}", e)) -} - -/// Make a Notion API DELETE request. -fn notion_delete(path: &str) -> Result { - let url = format!("{}{}", NOTION_API_BASE, path); - - let headers = serde_json::json!({ - "Notion-Version": NOTION_VERSION - }); - let headers_str = serde_json::to_string(&headers).map_err(|e| e.to_string())?; - - host::log(host::LogLevel::Debug, &format!("Notion DELETE: {}", path)); - - let response = host::http_request("DELETE", &url, &headers_str, None)?; - - if response.status < 200 || response.status >= 300 { - let body = String::from_utf8_lossy(&response.body); - return Err(format!( - "Notion API returned status {}: {}", - response.status, body - )); - } - - let body = String::from_utf8(response.body) - .map_err(|e| format!("Invalid UTF-8 in response: {}", e))?; - serde_json::from_str(&body).map_err(|e| format!("Invalid JSON in response: {}", e)) -} - -// ==================== Search ==================== - -pub fn search( - query: &str, - filter: Option<&SearchFilter>, - page_size: u32, - start_cursor: Option<&str>, -) -> Result { - let mut body = serde_json::json!({ - "page_size": page_size.min(100) - }); - - if !query.is_empty() { - body["query"] = serde_json::Value::String(query.to_string()); - } - - if let Some(f) = filter { - body["filter"] = serde_json::json!({ - "property": f.property, - "value": f.value - }); - } - - if let Some(cursor) = start_cursor { - body["start_cursor"] = serde_json::Value::String(cursor.to_string()); - } - - notion_post("/search", &body) -} - -// ==================== Pages ==================== - -pub fn get_page(page_id: &str) -> Result { - let page_id = normalize_uuid(page_id); - notion_get(&format!("/pages/{}", page_id)) -} - -pub fn create_page( - parent: &serde_json::Value, - properties: &serde_json::Value, - children: Option<&Vec>, - icon: Option<&serde_json::Value>, - cover: Option<&serde_json::Value>, -) -> Result { - let mut body = serde_json::json!({ - "parent": parent, - "properties": properties - }); - - if let Some(c) = children { - if !c.is_empty() { - body["children"] = serde_json::Value::Array(c.clone()); - } - } - - if let Some(i) = icon { - body["icon"] = i.clone(); - } - - if let Some(c) = cover { - body["cover"] = c.clone(); - } - - notion_post("/pages", &body) -} - -pub fn update_page( - page_id: &str, - properties: &serde_json::Value, - icon: Option<&serde_json::Value>, - cover: Option<&serde_json::Value>, -) -> Result { - let page_id = normalize_uuid(page_id); - - let mut body = serde_json::json!({ - "properties": properties - }); - - if let Some(i) = icon { - body["icon"] = i.clone(); - } - - if let Some(c) = cover { - body["cover"] = c.clone(); - } - - notion_patch(&format!("/pages/{}", page_id), &body) -} - -pub fn archive_page(page_id: &str) -> Result { - let page_id = normalize_uuid(page_id); - notion_patch( - &format!("/pages/{}", page_id), - &serde_json::json!({ "archived": true }), - ) -} - -pub fn restore_page(page_id: &str) -> Result { - let page_id = normalize_uuid(page_id); - notion_patch( - &format!("/pages/{}", page_id), - &serde_json::json!({ "archived": false }), - ) -} - -// ==================== Blocks ==================== - -pub fn get_blocks( - block_id: &str, - page_size: u32, - start_cursor: Option<&str>, -) -> Result { - let block_id = normalize_uuid(block_id); - let mut path = format!( - "/blocks/{}/children?page_size={}", - block_id, - page_size.min(100) - ); - - if let Some(cursor) = start_cursor { - path.push_str(&format!("&start_cursor={}", cursor)); - } - - notion_get(&path) -} - -pub fn append_blocks( - block_id: &str, - children: &[serde_json::Value], - after: Option<&str>, -) -> Result { - let block_id = normalize_uuid(block_id); - - let mut body = serde_json::json!({ - "children": children - }); - - if let Some(after_id) = after { - body["after"] = serde_json::Value::String(normalize_uuid(after_id)); - } - - notion_patch(&format!("/blocks/{}/children", block_id), &body) -} - -pub fn get_block(block_id: &str) -> Result { - let block_id = normalize_uuid(block_id); - notion_get(&format!("/blocks/{}", block_id)) -} - -pub fn update_block( - block_id: &str, - content: &serde_json::Value, -) -> Result { - let block_id = normalize_uuid(block_id); - notion_patch(&format!("/blocks/{}", block_id), content) -} - -pub fn delete_block(block_id: &str) -> Result { - let block_id = normalize_uuid(block_id); - notion_delete(&format!("/blocks/{}", block_id)) -} - -// ==================== Databases ==================== - -pub fn get_database(database_id: &str) -> Result { - let database_id = normalize_uuid(database_id); - notion_get(&format!("/databases/{}", database_id)) -} - -pub fn query_database( - database_id: &str, - filter: Option<&serde_json::Value>, - sorts: Option<&Vec>, - page_size: u32, - start_cursor: Option<&str>, -) -> Result { - let database_id = normalize_uuid(database_id); - - let mut body = serde_json::json!({ - "page_size": page_size.min(100) - }); - - if let Some(f) = filter { - body["filter"] = f.clone(); - } - - if let Some(s) = sorts { - if !s.is_empty() { - body["sorts"] = serde_json::Value::Array(s.clone()); - } - } - - if let Some(cursor) = start_cursor { - body["start_cursor"] = serde_json::Value::String(cursor.to_string()); - } - - notion_post(&format!("/databases/{}/query", database_id), &body) -} - -pub fn create_database( - parent: &serde_json::Value, - title: &[serde_json::Value], - properties: &serde_json::Value, - icon: Option<&serde_json::Value>, - cover: Option<&serde_json::Value>, - is_inline: bool, -) -> Result { - let mut body = serde_json::json!({ - "parent": parent, - "title": title, - "properties": properties, - "is_inline": is_inline - }); - - if let Some(i) = icon { - body["icon"] = i.clone(); - } - - if let Some(c) = cover { - body["cover"] = c.clone(); - } - - notion_post("/databases", &body) -} - -pub fn update_database( - database_id: &str, - title: Option<&Vec>, - properties: Option<&serde_json::Value>, -) -> Result { - let database_id = normalize_uuid(database_id); - - let mut body = serde_json::json!({}); - - if let Some(t) = title { - body["title"] = serde_json::Value::Array(t.clone()); - } - - if let Some(p) = properties { - body["properties"] = p.clone(); - } - - notion_patch(&format!("/databases/{}", database_id), &body) -} - -// ==================== Comments ==================== - -pub fn get_comments( - block_id: &str, - page_size: u32, - start_cursor: Option<&str>, -) -> Result { - let block_id = normalize_uuid(block_id); - let mut path = format!( - "/comments?block_id={}&page_size={}", - block_id, - page_size.min(100) - ); - - if let Some(cursor) = start_cursor { - path.push_str(&format!("&start_cursor={}", cursor)); - } - - notion_get(&path) -} - -pub fn add_comment( - parent: &serde_json::Value, - rich_text: &[serde_json::Value], -) -> Result { - let body = serde_json::json!({ - "parent": parent, - "rich_text": rich_text - }); - - notion_post("/comments", &body) -} - -// ==================== Users ==================== - -pub fn list_users(page_size: u32, start_cursor: Option<&str>) -> Result { - let mut path = format!("/users?page_size={}", page_size.min(100)); - - if let Some(cursor) = start_cursor { - path.push_str(&format!("&start_cursor={}", cursor)); - } - - notion_get(&path) -} - -pub fn get_user(user_id: &str) -> Result { - let user_id = normalize_uuid(user_id); - notion_get(&format!("/users/{}", user_id)) -} - -pub fn get_me() -> Result { - notion_get("/users/me") -} - -// ==================== Helpers ==================== - -/// Normalize a UUID by removing dashes if needed. -/// Notion accepts both formats, but we normalize for consistency. -fn normalize_uuid(id: &str) -> String { - // If it already has dashes in the right places, return as-is - if id.len() == 36 && id.chars().filter(|c| *c == '-').count() == 4 { - return id.to_string(); - } - - // If it's 32 characters without dashes, add them - let clean: String = id.chars().filter(|c| c.is_ascii_hexdigit()).collect(); - if clean.len() == 32 { - return format!( - "{}-{}-{}-{}-{}", - &clean[0..8], - &clean[8..12], - &clean[12..16], - &clean[16..20], - &clean[20..32] - ); - } - - // Otherwise return as-is (let the API handle validation) - id.to_string() -} diff --git a/examples/wasm-tools/notion/src/lib.rs b/examples/wasm-tools/notion/src/lib.rs deleted file mode 100644 index 2e6f1eae..00000000 --- a/examples/wasm-tools/notion/src/lib.rs +++ /dev/null @@ -1,451 +0,0 @@ -//! Notion WASM Tool for NEAR Agent. -//! -//! This is a standalone WASM component that provides Notion integration. -//! It follows the MCP server pattern with domain-grouped operations. -//! -//! # Capabilities Required -//! -//! - HTTP: `api.notion.com/v1/*` (GET, POST, PATCH, DELETE) -//! - Secrets: `notion_api_token` (injected automatically as Bearer token) -//! -//! # Supported Actions -//! -//! ## Search -//! - `search`: Search across pages and databases -//! -//! ## Pages -//! - `get_page`: Retrieve a page by ID -//! - `create_page`: Create a new page in a database or as child of another page -//! - `update_page`: Update page properties -//! - `archive_page`: Archive (soft-delete) a page -//! - `restore_page`: Restore an archived page -//! -//! ## Blocks -//! - `get_blocks`: Get child blocks of a page/block -//! - `append_blocks`: Append content blocks to a page/block -//! - `get_block`: Get a single block -//! - `update_block`: Update a block's content -//! - `delete_block`: Delete a block -//! -//! ## Databases -//! - `get_database`: Get database schema -//! - `query_database`: Query with filters and sorts -//! - `create_database`: Create a new database -//! - `update_database`: Update database title/properties -//! -//! ## Comments -//! - `get_comments`: Get comments on a page/block -//! - `add_comment`: Add a comment -//! -//! ## Users -//! - `list_users`: List workspace users -//! - `get_user`: Get a specific user -//! - `get_me`: Get the bot user -//! -//! # Example Usage -//! -//! ```json -//! {"action": "search", "query": "meeting notes", "page_size": 5} -//! ``` -//! -//! ```json -//! { -//! "action": "query_database", -//! "database_id": "abc123...", -//! "filter": {"property": "Status", "select": {"equals": "Done"}}, -//! "sorts": [{"property": "Created", "direction": "descending"}] -//! } -//! ``` - -mod api; -mod types; - -use types::NotionAction; - -// Generate bindings from the WIT interface. -wit_bindgen::generate!({ - world: "sandboxed-tool", - path: "../../../wit/tool.wit", -}); - -/// Implementation of the tool interface. -struct NotionTool; - -impl exports::near::agent::tool::Guest for NotionTool { - fn execute(req: exports::near::agent::tool::Request) -> exports::near::agent::tool::Response { - match execute_inner(&req.params) { - Ok(result) => exports::near::agent::tool::Response { - output: Some(result), - error: None, - }, - Err(e) => exports::near::agent::tool::Response { - output: None, - error: Some(e), - }, - } - } - - fn schema() -> String { - TOOL_SCHEMA.to_string() - } - - fn description() -> String { - "Notion workspace integration for searching, managing pages/databases/blocks, \ - comments, and users. Requires a Notion integration token with appropriate \ - capabilities (read content, update content, insert content)." - .to_string() - } -} - -/// Inner execution logic with proper error handling. -fn execute_inner(params: &str) -> Result { - // Check if the Notion token is configured - if !crate::near::agent::host::secret_exists("notion_api_token") { - return Err( - "Notion API token not configured. Please add the 'notion_api_token' secret." - .to_string(), - ); - } - - // Parse the action from JSON - let action: NotionAction = - serde_json::from_str(params).map_err(|e| format!("Invalid parameters: {}", e))?; - - crate::near::agent::host::log( - crate::near::agent::host::LogLevel::Info, - &format!("Executing Notion action: {:?}", action), - ); - - // Dispatch to the appropriate handler - let result = match action { - // Search - NotionAction::Search { - query, - filter, - page_size, - start_cursor, - } => api::search(&query, filter.as_ref(), page_size, start_cursor.as_deref())?, - - // Pages - NotionAction::GetPage { page_id } => api::get_page(&page_id)?, - - NotionAction::CreatePage { - parent, - properties, - children, - icon, - cover, - } => api::create_page( - &parent, - &properties, - children.as_ref(), - icon.as_ref(), - cover.as_ref(), - )?, - - NotionAction::UpdatePage { - page_id, - properties, - icon, - cover, - } => api::update_page(&page_id, &properties, icon.as_ref(), cover.as_ref())?, - - NotionAction::ArchivePage { page_id } => api::archive_page(&page_id)?, - - NotionAction::RestorePage { page_id } => api::restore_page(&page_id)?, - - // Blocks - NotionAction::GetBlocks { - block_id, - page_size, - start_cursor, - } => api::get_blocks(&block_id, page_size, start_cursor.as_deref())?, - - NotionAction::AppendBlocks { - block_id, - children, - after, - } => api::append_blocks(&block_id, &children, after.as_deref())?, - - NotionAction::GetBlock { block_id } => api::get_block(&block_id)?, - - NotionAction::UpdateBlock { block_id, content } => api::update_block(&block_id, &content)?, - - NotionAction::DeleteBlock { block_id } => api::delete_block(&block_id)?, - - // Databases - NotionAction::GetDatabase { database_id } => api::get_database(&database_id)?, - - NotionAction::QueryDatabase { - database_id, - filter, - sorts, - page_size, - start_cursor, - } => api::query_database( - &database_id, - filter.as_ref(), - sorts.as_ref(), - page_size, - start_cursor.as_deref(), - )?, - - NotionAction::CreateDatabase { - parent, - title, - properties, - icon, - cover, - is_inline, - } => api::create_database( - &parent, - &title, - &properties, - icon.as_ref(), - cover.as_ref(), - is_inline, - )?, - - NotionAction::UpdateDatabase { - database_id, - title, - properties, - } => api::update_database(&database_id, title.as_ref(), properties.as_ref())?, - - // Comments - NotionAction::GetComments { - block_id, - page_size, - start_cursor, - } => api::get_comments(&block_id, page_size, start_cursor.as_deref())?, - - NotionAction::AddComment { parent, rich_text } => api::add_comment(&parent, &rich_text)?, - - // Users - NotionAction::ListUsers { - page_size, - start_cursor, - } => api::list_users(page_size, start_cursor.as_deref())?, - - NotionAction::GetUser { user_id } => api::get_user(&user_id)?, - - NotionAction::GetMe => api::get_me()?, - }; - - serde_json::to_string(&result).map_err(|e| format!("Failed to serialize response: {}", e)) -} - -// Export the tool implementation. -export!(NotionTool); - -/// JSON Schema for the tool's parameters. -/// -/// This schema uses `oneOf` to describe each action's specific parameters. -const TOOL_SCHEMA: &str = r#"{ - "type": "object", - "required": ["action"], - "oneOf": [ - { - "properties": { - "action": { "const": "search" }, - "query": { "type": "string", "description": "Text to search for" }, - "filter": { - "type": "object", - "properties": { - "property": { "type": "string", "const": "object" }, - "value": { "type": "string", "enum": ["page", "database"] } - }, - "description": "Filter by object type" - }, - "page_size": { "type": "integer", "minimum": 1, "maximum": 100, "default": 10 }, - "start_cursor": { "type": "string", "description": "Pagination cursor" } - }, - "required": ["action"] - }, - { - "properties": { - "action": { "const": "get_page" }, - "page_id": { "type": "string", "description": "Page ID (UUID)" } - }, - "required": ["action", "page_id"] - }, - { - "properties": { - "action": { "const": "create_page" }, - "parent": { - "type": "object", - "description": "Parent reference: {\"database_id\": \"...\"} or {\"page_id\": \"...\"}" - }, - "properties": { "type": "object", "description": "Page properties" }, - "children": { - "type": "array", - "items": { "type": "object" }, - "description": "Block children for page content" - }, - "icon": { "type": "object", "description": "Page icon" }, - "cover": { "type": "object", "description": "Page cover image" } - }, - "required": ["action", "parent", "properties"] - }, - { - "properties": { - "action": { "const": "update_page" }, - "page_id": { "type": "string" }, - "properties": { "type": "object", "description": "Properties to update" }, - "icon": { "type": "object" }, - "cover": { "type": "object" } - }, - "required": ["action", "page_id", "properties"] - }, - { - "properties": { - "action": { "const": "archive_page" }, - "page_id": { "type": "string" } - }, - "required": ["action", "page_id"] - }, - { - "properties": { - "action": { "const": "restore_page" }, - "page_id": { "type": "string" } - }, - "required": ["action", "page_id"] - }, - { - "properties": { - "action": { "const": "get_blocks" }, - "block_id": { "type": "string", "description": "Block or page ID" }, - "page_size": { "type": "integer", "minimum": 1, "maximum": 100, "default": 50 }, - "start_cursor": { "type": "string" } - }, - "required": ["action", "block_id"] - }, - { - "properties": { - "action": { "const": "append_blocks" }, - "block_id": { "type": "string", "description": "Block or page ID to append to" }, - "children": { - "type": "array", - "items": { "type": "object" }, - "description": "Block objects to append" - }, - "after": { "type": "string", "description": "Insert after this block ID" } - }, - "required": ["action", "block_id", "children"] - }, - { - "properties": { - "action": { "const": "get_block" }, - "block_id": { "type": "string" } - }, - "required": ["action", "block_id"] - }, - { - "properties": { - "action": { "const": "update_block" }, - "block_id": { "type": "string" }, - "content": { "type": "object", "description": "Block content by type" } - }, - "required": ["action", "block_id", "content"] - }, - { - "properties": { - "action": { "const": "delete_block" }, - "block_id": { "type": "string" } - }, - "required": ["action", "block_id"] - }, - { - "properties": { - "action": { "const": "get_database" }, - "database_id": { "type": "string" } - }, - "required": ["action", "database_id"] - }, - { - "properties": { - "action": { "const": "query_database" }, - "database_id": { "type": "string" }, - "filter": { "type": "object", "description": "Notion filter object" }, - "sorts": { - "type": "array", - "items": { "type": "object" }, - "description": "Sort configuration" - }, - "page_size": { "type": "integer", "minimum": 1, "maximum": 100, "default": 10 }, - "start_cursor": { "type": "string" } - }, - "required": ["action", "database_id"] - }, - { - "properties": { - "action": { "const": "create_database" }, - "parent": { "type": "object", "description": "Parent page or workspace" }, - "title": { - "type": "array", - "items": { "type": "object" }, - "description": "Database title as rich text" - }, - "properties": { "type": "object", "description": "Property schema" }, - "icon": { "type": "object" }, - "cover": { "type": "object" }, - "is_inline": { "type": "boolean", "default": false } - }, - "required": ["action", "parent", "title", "properties"] - }, - { - "properties": { - "action": { "const": "update_database" }, - "database_id": { "type": "string" }, - "title": { "type": "array", "items": { "type": "object" } }, - "properties": { "type": "object" } - }, - "required": ["action", "database_id"] - }, - { - "properties": { - "action": { "const": "get_comments" }, - "block_id": { "type": "string", "description": "Block or page ID" }, - "page_size": { "type": "integer", "minimum": 1, "maximum": 100, "default": 50 }, - "start_cursor": { "type": "string" } - }, - "required": ["action", "block_id"] - }, - { - "properties": { - "action": { "const": "add_comment" }, - "parent": { - "type": "object", - "description": "{\"page_id\": \"...\"} or {\"discussion_id\": \"...\"}" - }, - "rich_text": { - "type": "array", - "items": { "type": "object" }, - "description": "Comment content as rich text" - } - }, - "required": ["action", "parent", "rich_text"] - }, - { - "properties": { - "action": { "const": "list_users" }, - "page_size": { "type": "integer", "minimum": 1, "maximum": 100, "default": 50 }, - "start_cursor": { "type": "string" } - }, - "required": ["action"] - }, - { - "properties": { - "action": { "const": "get_user" }, - "user_id": { "type": "string" } - }, - "required": ["action", "user_id"] - }, - { - "properties": { - "action": { "const": "get_me" } - }, - "required": ["action"] - } - ] -}"#; diff --git a/examples/wasm-tools/notion/src/types.rs b/examples/wasm-tools/notion/src/types.rs deleted file mode 100644 index 58fd408a..00000000 --- a/examples/wasm-tools/notion/src/types.rs +++ /dev/null @@ -1,267 +0,0 @@ -//! Types for Notion API requests and responses. - -use serde::{Deserialize, Serialize}; - -/// Input parameters for the Notion tool. -#[derive(Debug, Deserialize)] -#[serde(tag = "action", rename_all = "snake_case")] -pub enum NotionAction { - // ==================== Search ==================== - /// Search across all pages and databases. - Search { - /// Text query to search for. - #[serde(default)] - query: String, - /// Filter by object type: "page" or "database". - #[serde(default)] - filter: Option, - /// Max results (default: 10, max: 100). - #[serde(default = "default_page_size")] - page_size: u32, - /// Pagination cursor from previous response. - #[serde(default)] - start_cursor: Option, - }, - - // ==================== Pages ==================== - /// Retrieve a page by ID. - GetPage { - /// Page ID (UUID format, with or without dashes). - page_id: String, - }, - - /// Create a new page. - CreatePage { - /// Parent reference: { "database_id": "..." } or { "page_id": "..." }. - parent: serde_json::Value, - /// Page properties matching the parent database schema. - properties: serde_json::Value, - /// Optional page content as an array of block objects. - #[serde(default)] - children: Option>, - /// Optional icon: { "emoji": "..." } or { "external": { "url": "..." } }. - #[serde(default)] - icon: Option, - /// Optional cover: { "external": { "url": "..." } }. - #[serde(default)] - cover: Option, - }, - - /// Update page properties. - UpdatePage { - /// Page ID. - page_id: String, - /// Properties to update. - properties: serde_json::Value, - /// Optional new icon. - #[serde(default)] - icon: Option, - /// Optional new cover. - #[serde(default)] - cover: Option, - }, - - /// Archive (soft-delete) a page. - ArchivePage { - /// Page ID. - page_id: String, - }, - - /// Restore an archived page. - RestorePage { - /// Page ID. - page_id: String, - }, - - // ==================== Blocks ==================== - /// Get child blocks of a block or page. - GetBlocks { - /// Block or page ID. - block_id: String, - /// Max results (default: 50, max: 100). - #[serde(default = "default_block_page_size")] - page_size: u32, - /// Pagination cursor. - #[serde(default)] - start_cursor: Option, - }, - - /// Append blocks to a page or block. - AppendBlocks { - /// Block or page ID to append to. - block_id: String, - /// Array of block objects to append. - children: Vec, - /// Append after this block ID (for ordering). - #[serde(default)] - after: Option, - }, - - /// Retrieve a single block. - GetBlock { - /// Block ID. - block_id: String, - }, - - /// Update a block's content. - UpdateBlock { - /// Block ID. - block_id: String, - /// Block content (varies by type, e.g., { "paragraph": { "rich_text": [...] } }). - content: serde_json::Value, - }, - - /// Delete a block. - DeleteBlock { - /// Block ID. - block_id: String, - }, - - // ==================== Databases ==================== - /// Retrieve a database schema. - GetDatabase { - /// Database ID. - database_id: String, - }, - - /// Query a database with filters and sorts. - QueryDatabase { - /// Database ID. - database_id: String, - /// Filter object (Notion filter format). - #[serde(default)] - filter: Option, - /// Sort configuration array. - #[serde(default)] - sorts: Option>, - /// Max results (default: 10, max: 100). - #[serde(default = "default_page_size")] - page_size: u32, - /// Pagination cursor. - #[serde(default)] - start_cursor: Option, - }, - - /// Create a new database. - CreateDatabase { - /// Parent reference: { "page_id": "..." } or { "type": "workspace", "workspace": true }. - parent: serde_json::Value, - /// Database title as rich text array. - title: Vec, - /// Property schema: { "Name": { "title": {} }, "Status": { "select": { "options": [...] } } }. - properties: serde_json::Value, - /// Optional icon. - #[serde(default)] - icon: Option, - /// Optional cover. - #[serde(default)] - cover: Option, - /// Make database inline (default: false). - #[serde(default)] - is_inline: bool, - }, - - /// Update database title or properties schema. - UpdateDatabase { - /// Database ID. - database_id: String, - /// New title (optional). - #[serde(default)] - title: Option>, - /// Properties to add or update (optional). - #[serde(default)] - properties: Option, - }, - - // ==================== Comments ==================== - /// Get comments on a block or page. - GetComments { - /// Block or page ID. - block_id: String, - /// Max results (default: 50, max: 100). - #[serde(default = "default_block_page_size")] - page_size: u32, - /// Pagination cursor. - #[serde(default)] - start_cursor: Option, - }, - - /// Add a comment to a page or discussion thread. - AddComment { - /// Parent reference: { "page_id": "..." } or { "discussion_id": "..." }. - parent: serde_json::Value, - /// Comment text as rich text array. - rich_text: Vec, - }, - - // ==================== Users ==================== - /// List all users in the workspace. - ListUsers { - /// Max results (default: 50, max: 100). - #[serde(default = "default_block_page_size")] - page_size: u32, - /// Pagination cursor. - #[serde(default)] - start_cursor: Option, - }, - - /// Get a specific user. - GetUser { - /// User ID. - user_id: String, - }, - - /// Get the bot user (the integration itself). - GetMe, -} - -/// Search filter for limiting results to pages or databases. -#[derive(Debug, Deserialize, Serialize)] -pub struct SearchFilter { - /// "object" field. - pub property: String, - /// "page" or "database". - pub value: String, -} - -fn default_page_size() -> u32 { - 10 -} - -fn default_block_page_size() -> u32 { - 50 -} - -// ==================== Response Types ==================== - -/// Generic Notion API response with pagination. -#[derive(Debug, Serialize)] -pub struct PaginatedResponse { - pub object: String, - pub results: Vec, - #[serde(skip_serializing_if = "Option::is_none")] - pub next_cursor: Option, - pub has_more: bool, -} - -/// Single object response. -#[derive(Debug, Serialize)] -pub struct ObjectResponse { - pub object: String, - #[serde(flatten)] - pub data: serde_json::Value, -} - -/// Success response for mutations. -#[derive(Debug, Serialize)] -pub struct MutationResponse { - pub success: bool, - #[serde(flatten)] - pub data: serde_json::Value, -} - -/// Minimal success response. -#[derive(Debug, Serialize)] -pub struct SuccessResponse { - pub success: bool, -} diff --git a/src/cli/mcp.rs b/src/cli/mcp.rs new file mode 100644 index 00000000..0e0536f6 --- /dev/null +++ b/src/cli/mcp.rs @@ -0,0 +1,484 @@ +//! MCP server management CLI commands. +//! +//! Commands for adding, removing, authenticating, and testing MCP servers. + +use std::io::Write; +use std::sync::Arc; + +use clap::Subcommand; + +use crate::config::Config; +use crate::history::Store; +use crate::secrets::{PostgresSecretsStore, SecretsCrypto, SecretsStore}; +use crate::tools::mcp::{ + McpClient, McpServerConfig, McpSessionManager, OAuthConfig, + auth::{authorize_mcp_server, is_authenticated}, + config::{ + add_mcp_server, get_mcp_server, load_mcp_servers, remove_mcp_server, save_mcp_servers, + }, +}; + +#[derive(Subcommand, Debug, Clone)] +pub enum McpCommand { + /// Add an MCP server + Add { + /// Server name (e.g., "notion", "github") + name: String, + + /// Server URL (e.g., "https://mcp.notion.com") + url: String, + + /// OAuth client ID (if authentication is required) + #[arg(long)] + client_id: Option, + + /// OAuth authorization URL (optional, can be discovered) + #[arg(long)] + auth_url: Option, + + /// OAuth token URL (optional, can be discovered) + #[arg(long)] + token_url: Option, + + /// Scopes to request (comma-separated) + #[arg(long)] + scopes: Option, + + /// Server description + #[arg(long)] + description: Option, + }, + + /// Remove an MCP server + Remove { + /// Server name to remove + name: String, + }, + + /// List configured MCP servers + List { + /// Show detailed information + #[arg(short, long)] + verbose: bool, + }, + + /// Authenticate with an MCP server (OAuth flow) + Auth { + /// Server name to authenticate + name: String, + + /// User ID for storing the token (default: "default") + #[arg(short, long, default_value = "default")] + user: String, + }, + + /// Test connection to an MCP server + Test { + /// Server name to test + name: String, + + /// User ID for authentication (default: "default") + #[arg(short, long, default_value = "default")] + user: String, + }, + + /// Enable or disable an MCP server + Toggle { + /// Server name + name: String, + + /// Enable the server + #[arg(long, conflicts_with = "disable")] + enable: bool, + + /// Disable the server + #[arg(long, conflicts_with = "enable")] + disable: bool, + }, +} + +/// Run an MCP command. +pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> { + match cmd { + McpCommand::Add { + name, + url, + client_id, + auth_url, + token_url, + scopes, + description, + } => { + add_server( + name, + url, + client_id, + auth_url, + token_url, + scopes, + description, + ) + .await + } + McpCommand::Remove { name } => remove_server(name).await, + McpCommand::List { verbose } => list_servers(verbose).await, + McpCommand::Auth { name, user } => auth_server(name, user).await, + McpCommand::Test { name, user } => test_server(name, user).await, + McpCommand::Toggle { + name, + enable, + disable, + } => toggle_server(name, enable, disable).await, + } +} + +/// Add a new MCP server. +async fn add_server( + name: String, + url: String, + client_id: Option, + auth_url: Option, + token_url: Option, + scopes: Option, + description: Option, +) -> anyhow::Result<()> { + let mut config = McpServerConfig::new(&name, &url); + + if let Some(desc) = description { + config = config.with_description(desc); + } + + // Track if auth is required + let requires_auth = client_id.is_some(); + + // Set up OAuth if client_id is provided + if let Some(client_id) = client_id { + let mut oauth = OAuthConfig::new(client_id); + + if let (Some(auth), Some(token)) = (auth_url, token_url) { + oauth = oauth.with_endpoints(auth, token); + } + + if let Some(scopes_str) = scopes { + let scope_list: Vec = scopes_str + .split(',') + .map(|s| s.trim().to_string()) + .collect(); + oauth = oauth.with_scopes(scope_list); + } + + config = config.with_oauth(oauth); + } + + // Validate + config.validate()?; + + // Save + add_mcp_server(config).await?; + + println!(); + println!(" ✓ Added MCP server '{}'", name); + println!(" URL: {}", url); + + if requires_auth { + println!(); + println!(" Run 'ironclaw mcp auth {}' to authenticate.", name); + } + + println!(); + + Ok(()) +} + +/// Remove an MCP server. +async fn remove_server(name: String) -> anyhow::Result<()> { + remove_mcp_server(&name).await?; + + println!(); + println!(" ✓ Removed MCP server '{}'", name); + println!(); + + Ok(()) +} + +/// List configured MCP servers. +async fn list_servers(verbose: bool) -> anyhow::Result<()> { + let servers = load_mcp_servers().await?; + + if servers.servers.is_empty() { + println!(); + println!(" No MCP servers configured."); + println!(); + println!(" Add a server with:"); + println!(" ironclaw mcp add [--client-id ]"); + println!(); + return Ok(()); + } + + println!(); + println!(" Configured MCP servers:"); + println!(); + + for server in &servers.servers { + let status = if server.enabled { "●" } else { "○" }; + let auth_status = if server.requires_auth() { + " (auth required)" + } else { + "" + }; + + if verbose { + println!(" {} {}{}", status, server.name, auth_status); + println!(" URL: {}", server.url); + if let Some(ref desc) = server.description { + println!(" Description: {}", desc); + } + if let Some(ref oauth) = server.oauth { + println!(" OAuth Client ID: {}", oauth.client_id); + if !oauth.scopes.is_empty() { + println!(" Scopes: {}", oauth.scopes.join(", ")); + } + } + println!(); + } else { + println!( + " {} {} - {}{}", + status, server.name, server.url, auth_status + ); + } + } + + if !verbose { + println!(); + println!(" Use --verbose for more details."); + } + + println!(); + + Ok(()) +} + +/// Authenticate with an MCP server. +async fn auth_server(name: String, user_id: String) -> anyhow::Result<()> { + // Get server config + let server = get_mcp_server(&name).await?; + + // Initialize secrets store + let secrets = get_secrets_store().await?; + + // Check if already authenticated + if is_authenticated(&server, &secrets, &user_id).await { + println!(); + println!(" Server '{}' is already authenticated.", name); + println!(); + print!(" Re-authenticate? [y/N]: "); + std::io::stdout().flush()?; + + let mut input = String::new(); + std::io::stdin().read_line(&mut input)?; + + if !input.trim().eq_ignore_ascii_case("y") { + return Ok(()); + } + println!(); + } + + println!(); + println!("╔════════════════════════════════════════════════════════════════╗"); + println!( + "║ {:^62}║", + format!("{} Authentication", name.to_uppercase()) + ); + println!("╚════════════════════════════════════════════════════════════════╝"); + println!(); + + // Perform OAuth flow (supports both pre-configured OAuth and DCR) + match authorize_mcp_server(&server, &secrets, &user_id).await { + Ok(_token) => { + println!(); + println!(" ✓ Successfully authenticated with '{}'!", name); + println!(); + println!(" You can now use tools from this server."); + println!(); + } + Err(crate::tools::mcp::auth::AuthError::NotSupported) => { + println!(); + println!(" ✗ Server does not support OAuth authentication."); + println!(); + println!(" The server may require a different authentication method,"); + println!(" or you may need to configure OAuth manually:"); + println!(); + println!(" ironclaw mcp remove {}", name); + println!( + " ironclaw mcp add {} {} --client-id YOUR_CLIENT_ID", + name, server.url + ); + println!(); + } + Err(e) => { + println!(); + println!(" ✗ Authentication failed: {}", e); + println!(); + return Err(e.into()); + } + } + + Ok(()) +} + +/// Test connection to an MCP server. +async fn test_server(name: String, user_id: String) -> anyhow::Result<()> { + // Get server config + let server = get_mcp_server(&name).await?; + + println!(); + println!(" Testing connection to '{}'...", name); + + // Create client + let session_manager = Arc::new(McpSessionManager::new()); + + // Always check for stored tokens (from either pre-configured OAuth or DCR) + let secrets = get_secrets_store().await?; + let has_tokens = is_authenticated(&server, &secrets, &user_id).await; + + let client = if has_tokens { + // We have stored tokens, use authenticated client + McpClient::new_authenticated(server.clone(), session_manager, secrets, user_id) + } else if server.requires_auth() { + // OAuth configured but no tokens - need to authenticate + println!(); + println!( + " ✗ Not authenticated. Run 'ironclaw mcp auth {}' first.", + name + ); + println!(); + return Ok(()); + } else { + // No OAuth and no tokens - try unauthenticated + McpClient::new_with_name(&server.name, &server.url) + }; + + // Test connection + match client.test_connection().await { + Ok(()) => { + println!(" ✓ Connection successful!"); + println!(); + + // List tools + match client.list_tools().await { + Ok(tools) => { + println!(" Available tools ({}):", tools.len()); + for tool in tools { + let approval = if tool.requires_approval() { + " [approval required]" + } else { + "" + }; + println!(" • {}{}", tool.name, approval); + if !tool.description.is_empty() { + // Truncate long descriptions + let desc = if tool.description.len() > 60 { + format!("{}...", &tool.description[..57]) + } else { + tool.description.clone() + }; + println!(" {}", desc); + } + } + } + Err(e) => { + println!(" ✗ Failed to list tools: {}", e); + } + } + } + Err(e) => { + let err_str = e.to_string(); + // Check if server requires auth but we don't have valid tokens + if err_str.contains("401") || err_str.contains("requires authentication") { + if has_tokens { + // We had tokens but they failed - need to re-authenticate + println!( + " ✗ Authentication failed (token may be expired). Try re-authenticating:" + ); + println!(" ironclaw mcp auth {}", name); + } else { + // No tokens - server requires auth + println!(" ✗ Server requires authentication."); + println!(); + println!(" Run 'ironclaw mcp auth {}' to authenticate.", name); + } + } else { + println!(" ✗ Connection failed: {}", e); + } + } + } + + println!(); + + Ok(()) +} + +/// Toggle server enabled/disabled state. +async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Result<()> { + let mut servers = load_mcp_servers().await?; + + let server = servers + .get_mut(&name) + .ok_or_else(|| anyhow::anyhow!("Server '{}' not found", name))?; + + let new_state = if enable { + true + } else if disable { + false + } else { + !server.enabled // Toggle if neither specified + }; + + server.enabled = new_state; + save_mcp_servers(&servers).await?; + + let status = if new_state { "enabled" } else { "disabled" }; + println!(); + println!(" ✓ Server '{}' is now {}.", name, status); + println!(); + + Ok(()) +} + +/// Initialize and return the secrets store. +async fn get_secrets_store() -> anyhow::Result> { + let config = Config::from_env()?; + + let master_key = config.secrets.master_key().ok_or_else(|| { + anyhow::anyhow!("SECRETS_MASTER_KEY not set. Run 'ironclaw setup' first or set it in .env") + })?; + + let store = Store::new(&config.database).await?; + store.run_migrations().await?; + + let crypto = SecretsCrypto::new(master_key.clone())?; + Ok(Arc::new(PostgresSecretsStore::new( + store.pool(), + Arc::new(crypto), + ))) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_mcp_command_parsing() { + // Just verify the command structure is valid + use clap::CommandFactory; + + // Create a dummy parent command to test subcommand parsing + #[derive(clap::Parser)] + struct TestCli { + #[command(subcommand)] + cmd: McpCommand, + } + + TestCli::command().debug_assert(); + } +} diff --git a/src/cli/mod.rs b/src/cli/mod.rs index 631a4770..4d863059 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -5,11 +5,14 @@ //! - Interactive setup wizard (`setup`) //! - Managing configuration (`config list`, `config get`, `config set`) //! - Managing WASM tools (`tool install`, `tool list`, `tool remove`) +//! - Managing MCP servers (`mcp add`, `mcp auth`, `mcp list`, `mcp test`) mod config; +mod mcp; mod tool; pub use config::{ConfigCommand, run_config_command}; +pub use mcp::{McpCommand, run_mcp_command}; pub use tool::{ToolCommand, run_tool_command}; use clap::{Parser, Subcommand}; @@ -72,6 +75,10 @@ pub enum Command { /// Manage WASM tools #[command(subcommand)] Tool(ToolCommand), + + /// Manage MCP servers (hosted tool providers) + #[command(subcommand)] + Mcp(McpCommand), } impl Cli { diff --git a/src/cli/tool.rs b/src/cli/tool.rs index 1fa2df89..79a71a7f 100644 --- a/src/cli/tool.rs +++ b/src/cli/tool.rs @@ -1,13 +1,18 @@ //! Tool management CLI commands. //! -//! Commands for installing, listing, and removing WASM tools. +//! Commands for installing, listing, removing, and authenticating WASM tools. +use std::io::Write; use std::path::{Path, PathBuf}; use std::process::Command as ProcessCommand; +use std::sync::Arc; use clap::Subcommand; use tokio::fs; +use crate::config::Config; +use crate::history::Store; +use crate::secrets::{CreateSecretParams, PostgresSecretsStore, SecretsCrypto, SecretsStore}; use crate::tools::wasm::{CapabilitiesFile, compute_binary_hash}; /// Default tools directory. @@ -79,6 +84,20 @@ pub enum ToolCommand { #[arg(short, long)] dir: Option, }, + + /// Configure authentication for a tool + Auth { + /// Name of the tool + name: String, + + /// Directory to look for tool (default: ~/.ironclaw/tools/) + #[arg(short, long)] + dir: Option, + + /// User ID for storing the secret (default: "default") + #[arg(short, long, default_value = "default")] + user: String, + }, } /// Run a tool command. @@ -96,6 +115,7 @@ pub async fn run_tool_command(cmd: ToolCommand) -> anyhow::Result<()> { ToolCommand::List { dir, verbose } => list_tools(dir, verbose).await, ToolCommand::Remove { name, dir } => remove_tool(name, dir).await, ToolCommand::Info { name_or_path, dir } => show_tool_info(name_or_path, dir).await, + ToolCommand::Auth { name, dir, user } => auth_tool(name, dir, user).await, } } @@ -658,6 +678,571 @@ fn print_capabilities_detail(caps: &CapabilitiesFile) { } } +/// Configure authentication for a tool. +async fn auth_tool(name: String, dir: Option, user_id: String) -> anyhow::Result<()> { + let tools_dir = dir.unwrap_or_else(default_tools_dir); + let caps_path = tools_dir.join(format!("{}.capabilities.json", name)); + + if !caps_path.exists() { + anyhow::bail!( + "Tool '{}' not found or has no capabilities file at {}", + name, + caps_path.display() + ); + } + + // Parse capabilities + let content = fs::read_to_string(&caps_path).await?; + let caps = CapabilitiesFile::from_json(&content) + .map_err(|e| anyhow::anyhow!("Invalid capabilities file: {}", e))?; + + // Check for auth section + let auth = caps.auth.ok_or_else(|| { + anyhow::anyhow!( + "Tool '{}' has no auth configuration.\n\ + The tool may not require authentication, or auth setup is not defined.", + name + ) + })?; + + let display_name = auth.display_name.as_deref().unwrap_or(&name); + + let header = format!("{} Authentication", display_name); + println!(); + println!("╔════════════════════════════════════════════════════════════════╗"); + println!("║ {:^62}║", header); + println!("╚════════════════════════════════════════════════════════════════╝"); + println!(); + + // Initialize secrets store + let config = Config::from_env()?; + let master_key = config.secrets.master_key().ok_or_else(|| { + anyhow::anyhow!("SECRETS_MASTER_KEY not set. Run 'ironclaw setup' first or set it in .env") + })?; + + let store = Store::new(&config.database).await?; + store.run_migrations().await?; + + let crypto = SecretsCrypto::new(master_key.clone())?; + let secrets_store = Arc::new(PostgresSecretsStore::new(store.pool(), Arc::new(crypto))); + + // Check if already configured + let already_configured = secrets_store + .exists(&user_id, &auth.secret_name) + .await + .unwrap_or(false); + + if already_configured { + println!(" {} is already configured.", display_name); + println!(); + print!(" Replace existing credentials? [y/N]: "); + std::io::stdout().flush()?; + + let mut input = String::new(); + std::io::stdin().read_line(&mut input)?; + + if !input.trim().eq_ignore_ascii_case("y") { + println!(); + println!(" Keeping existing credentials."); + return Ok(()); + } + println!(); + } + + // Check for environment variable + if let Some(ref env_var) = auth.env_var { + if let Ok(token) = std::env::var(env_var) { + if !token.is_empty() { + println!(" Found {} in environment.", env_var); + println!(); + + // Validate if endpoint is provided + if let Some(ref validation) = auth.validation_endpoint { + print!(" Validating token..."); + std::io::stdout().flush()?; + + match validate_token(&token, validation, &auth.secret_name).await { + Ok(()) => { + println!(" ✓"); + } + Err(e) => { + println!(" ✗"); + println!(" Validation failed: {}", e); + println!(); + println!(" Falling back to manual entry..."); + return auth_tool_manual(&secrets_store, &user_id, &auth).await; + } + } + } + + // Save the token + save_token(&secrets_store, &user_id, &auth, &token).await?; + print_success(display_name); + return Ok(()); + } + } + } + + // Check for OAuth configuration + if let Some(ref oauth) = auth.oauth { + return auth_tool_oauth(&secrets_store, &user_id, &auth, oauth).await; + } + + // Fall back to manual entry + auth_tool_manual(&secrets_store, &user_id, &auth).await +} + +/// OAuth browser-based login flow. +async fn auth_tool_oauth( + store: &PostgresSecretsStore, + user_id: &str, + auth: &crate::tools::wasm::AuthCapabilitySchema, + oauth: &crate::tools::wasm::OAuthConfigSchema, +) -> anyhow::Result<()> { + use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; + use rand::RngCore; + use sha2::{Digest, Sha256}; + use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; + use tokio::net::TcpListener; + + let display_name = auth.display_name.as_deref().unwrap_or(&auth.secret_name); + + // Get client_id from config or env + let client_id = oauth + .client_id + .clone() + .or_else(|| { + oauth + .client_id_env + .as_ref() + .and_then(|env| std::env::var(env).ok()) + }) + .ok_or_else(|| { + anyhow::anyhow!( + "OAuth client_id not configured.\n\ + Set it in the capabilities file or via environment variable." + ) + })?; + + // Get client_secret if provided + let client_secret = oauth.client_secret.clone().or_else(|| { + oauth + .client_secret_env + .as_ref() + .and_then(|env| std::env::var(env).ok()) + }); + + println!(" Starting OAuth authentication..."); + println!(); + + // Find an available port for the callback + let mut listener = None; + let mut port = 0; + + for p in 9876..=9886 { + match TcpListener::bind(format!("127.0.0.1:{}", p)).await { + Ok(l) => { + listener = Some(l); + port = p; + break; + } + Err(_) => continue, + } + } + + let listener = listener.ok_or_else(|| anyhow::anyhow!("Could not find available port"))?; + let redirect_uri = format!("http://localhost:{}/callback", port); + + // Generate PKCE verifier and challenge + let (code_verifier, code_challenge) = if oauth.use_pkce { + let mut verifier_bytes = [0u8; 32]; + rand::thread_rng().fill_bytes(&mut verifier_bytes); + let verifier = URL_SAFE_NO_PAD.encode(verifier_bytes); + + let mut hasher = Sha256::new(); + hasher.update(verifier.as_bytes()); + let challenge = URL_SAFE_NO_PAD.encode(hasher.finalize()); + + (Some(verifier), Some(challenge)) + } else { + (None, None) + }; + + // Build authorization URL + let mut auth_url = format!( + "{}?client_id={}&response_type=code&redirect_uri={}", + oauth.authorization_url, + urlencoding::encode(&client_id), + urlencoding::encode(&redirect_uri) + ); + + if !oauth.scopes.is_empty() { + auth_url.push_str(&format!( + "&scope={}", + urlencoding::encode(&oauth.scopes.join(" ")) + )); + } + + if let Some(ref challenge) = code_challenge { + auth_url.push_str(&format!( + "&code_challenge={}&code_challenge_method=S256", + challenge + )); + } + + // Add extra params + for (key, value) in &oauth.extra_params { + auth_url.push_str(&format!( + "&{}={}", + urlencoding::encode(key), + urlencoding::encode(value) + )); + } + + println!(" Opening browser for {} login...", display_name); + println!(); + + if let Err(e) = open::that(&auth_url) { + println!(" Could not open browser: {}", e); + println!(" Please open this URL manually:"); + println!(" {}", auth_url); + } + + println!(" Waiting for authorization..."); + + // Wait for callback with timeout + let timeout = std::time::Duration::from_secs(300); + let code = tokio::time::timeout(timeout, async { + loop { + let (mut socket, _) = listener.accept().await?; + + let mut reader = BufReader::new(&mut socket); + let mut request_line = String::new(); + reader.read_line(&mut request_line).await?; + + // Parse GET /callback?code=xxx HTTP/1.1 + if let Some(path) = request_line.split_whitespace().nth(1) { + if path.starts_with("/callback") { + if let Some(query) = path.split('?').nth(1) { + for param in query.split('&') { + let parts: Vec<&str> = param.splitn(2, '=').collect(); + if parts.len() == 2 && parts[0] == "code" { + let code = urlencoding::decode(parts[1]) + .unwrap_or_else(|_| parts[1].into()) + .into_owned(); + + // Send success response + let response = format!( + "HTTP/1.1 200 OK\r\n\ + Content-Type: text/html\r\n\ + \r\n\ + \ +
\ +

✓ {} Connected!

\ +

You can close this window.

\ +
", + display_name + ); + let _ = socket.write_all(response.as_bytes()).await; + let _ = socket.shutdown().await; + + return Ok::<_, anyhow::Error>(code); + } + } + + // Check for error + if query.contains("error=") { + let response = + "HTTP/1.1 400 Bad Request\r\n\r\nAuthorization denied"; + let _ = socket.write_all(response.as_bytes()).await; + return Err(anyhow::anyhow!("Authorization denied by user")); + } + } + } + } + + let response = "HTTP/1.1 404 Not Found\r\n\r\n"; + let _ = socket.write_all(response.as_bytes()).await; + } + }) + .await + .map_err(|_| anyhow::anyhow!("Timed out waiting for authorization"))??; + + println!(); + println!(" Exchanging code for token..."); + + // Exchange code for token + let client = reqwest::Client::new(); + let mut token_params = vec![ + ("grant_type", "authorization_code".to_string()), + ("code", code), + ("redirect_uri", redirect_uri), + ]; + + if let Some(ref verifier) = code_verifier { + token_params.push(("code_verifier", verifier.to_string())); + } + + // Build token request + let mut request = client.post(&oauth.token_url); + + // Use Basic auth if client_secret is provided, otherwise include client_id in body + if let Some(ref secret) = client_secret { + request = request.basic_auth(&client_id, Some(secret)); + } else { + token_params.push(("client_id", client_id)); + } + + let token_response = request.form(&token_params).send().await?; + + if !token_response.status().is_success() { + let status = token_response.status(); + let body = token_response.text().await.unwrap_or_default(); + return Err(anyhow::anyhow!( + "Token exchange failed: {} - {}", + status, + body + )); + } + + let token_data: serde_json::Value = token_response.json().await?; + let access_token = token_data + .get(&oauth.access_token_field) + .and_then(|v| v.as_str()) + .ok_or_else(|| { + anyhow::anyhow!( + "No {} in token response: {:?}", + oauth.access_token_field, + token_data + ) + })?; + + // Save the token + save_token(store, user_id, auth, access_token).await?; + + // Extract any additional info for display + let workspace_name = token_data + .get("workspace_name") + .and_then(|v| v.as_str()) + .or_else(|| token_data.get("team_name").and_then(|v| v.as_str())); + + println!(); + println!(" ✓ {} connected!", display_name); + if let Some(workspace) = workspace_name { + println!(" Workspace: {}", workspace); + } + println!(); + println!(" The tool can now access the API."); + println!(); + + Ok(()) +} + +/// Manual token entry flow. +async fn auth_tool_manual( + store: &PostgresSecretsStore, + user_id: &str, + auth: &crate::tools::wasm::AuthCapabilitySchema, +) -> anyhow::Result<()> { + let display_name = auth.display_name.as_deref().unwrap_or(&auth.secret_name); + + // Show instructions + if let Some(ref instructions) = auth.instructions { + println!(" Setup instructions:"); + println!(); + for line in instructions.lines() { + println!(" {}", line); + } + println!(); + } + + // Offer to open setup URL + if let Some(ref url) = auth.setup_url { + print!(" Press Enter to open setup page (or 's' to skip): "); + std::io::stdout().flush()?; + + let mut input = String::new(); + std::io::stdin().read_line(&mut input)?; + + if !input.trim().eq_ignore_ascii_case("s") { + if let Err(e) = open::that(url) { + println!(" Could not open browser: {}", e); + println!(" Please open manually: {}", url); + } else { + println!(" Opening browser..."); + } + } + println!(); + } + + // Show token hint + if let Some(ref hint) = auth.token_hint { + println!(" Token format: {}", hint); + println!(); + } + + // Prompt for token + print!(" Paste your token: "); + std::io::stdout().flush()?; + + let token = read_hidden_input()?; + println!(); + + if token.is_empty() { + println!(" No token provided. Aborting."); + return Ok(()); + } + + // Validate if endpoint is provided + if let Some(ref validation) = auth.validation_endpoint { + print!(" Validating token..."); + std::io::stdout().flush()?; + + match validate_token(&token, validation, &auth.secret_name).await { + Ok(()) => { + println!(" ✓"); + } + Err(e) => { + println!(" ✗"); + println!(" Validation failed: {}", e); + println!(); + print!(" Save anyway? [y/N]: "); + std::io::stdout().flush()?; + + let mut confirm = String::new(); + std::io::stdin().read_line(&mut confirm)?; + + if !confirm.trim().eq_ignore_ascii_case("y") { + println!(" Aborting."); + return Ok(()); + } + } + } + } + + // Save the token + save_token(store, user_id, auth, &token).await?; + print_success(display_name); + Ok(()) +} + +/// Read input with hidden characters. +fn read_hidden_input() -> anyhow::Result { + use crossterm::{ + event::{self, Event, KeyCode, KeyModifiers}, + terminal, + }; + + let mut input = String::new(); + + terminal::enable_raw_mode()?; + + loop { + if let Event::Key(key_event) = event::read()? { + match key_event.code { + KeyCode::Enter => { + break; + } + KeyCode::Backspace => { + if !input.is_empty() { + input.pop(); + print!("\x08 \x08"); + std::io::stdout().flush()?; + } + } + KeyCode::Char('c') if key_event.modifiers.contains(KeyModifiers::CONTROL) => { + terminal::disable_raw_mode()?; + return Err(anyhow::anyhow!("Interrupted")); + } + KeyCode::Char(c) => { + input.push(c); + print!("*"); + std::io::stdout().flush()?; + } + _ => {} + } + } + } + + terminal::disable_raw_mode()?; + + Ok(input) +} + +/// Validate a token against the validation endpoint. +async fn validate_token( + token: &str, + validation: &crate::tools::wasm::ValidationEndpointSchema, + _secret_name: &str, +) -> anyhow::Result<()> { + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(10)) + .build()?; + + // Build request based on method + let request = match validation.method.to_uppercase().as_str() { + "GET" => client.get(&validation.url), + "POST" => client.post(&validation.url), + _ => client.get(&validation.url), + }; + + // Add authorization header (assume Bearer for now, could be extended) + let response = request + .header("Authorization", format!("Bearer {}", token)) + .header("Notion-Version", "2022-06-28") // Notion-specific, but harmless for others + .send() + .await?; + + if response.status().as_u16() == validation.success_status { + Ok(()) + } else { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + Err(anyhow::anyhow!( + "HTTP {} (expected {}): {}", + status, + validation.success_status, + if body.len() > 100 { + format!("{}...", &body[..100]) + } else { + body + } + )) + } +} + +/// Save token to secrets store. +async fn save_token( + store: &PostgresSecretsStore, + user_id: &str, + auth: &crate::tools::wasm::AuthCapabilitySchema, + token: &str, +) -> anyhow::Result<()> { + let mut params = CreateSecretParams::new(&auth.secret_name, token); + + if let Some(ref provider) = auth.provider { + params = params.with_provider(provider); + } + + store + .create(user_id, params) + .await + .map_err(|e| anyhow::anyhow!("Failed to save token: {}", e))?; + + Ok(()) +} + +/// Print success message. +fn print_success(display_name: &str) { + println!(); + println!(" ✓ {} connected!", display_name); + println!(); + println!(" The tool can now access the API."); + println!(); +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/main.rs b/src/main.rs index 38b2da19..0261fad2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -14,7 +14,7 @@ use ironclaw::{ WasmChannelRuntime, WasmChannelRuntimeConfig, WasmChannelServer, }, }, - cli::{Cli, Command, run_tool_command}, + cli::{Cli, Command, run_mcp_command, run_tool_command}, config::Config, context::ContextManager, history::Store, @@ -25,6 +25,7 @@ use ironclaw::{ setup::{SetupConfig, SetupWizard}, tools::{ ToolRegistry, + mcp::{McpClient, McpSessionManager, config::load_mcp_servers, is_authenticated}, wasm::{WasmToolLoader, WasmToolRuntime}, }, workspace::{EmbeddingProvider, NearAiEmbeddings, OpenAiEmbeddings, Workspace}, @@ -51,6 +52,16 @@ async fn main() -> anyhow::Result<()> { return ironclaw::cli::run_config_command(config_cmd.clone()) .map_err(|e| anyhow::anyhow!("{}", e)); } + Some(Command::Mcp(mcp_cmd)) => { + // Simple logging for MCP commands + tracing_subscriber::fmt() + .with_env_filter( + EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")), + ) + .init(); + + return run_mcp_command(mcp_cmd.clone()).await; + } Some(Command::Setup { skip_auth, channels_only, @@ -310,6 +321,110 @@ async fn main() -> anyhow::Result<()> { } } } + + // Create secrets store if master key is configured (needed for MCP auth and WASM channels) + let secrets_store: Option> = + if let (Some(store), Some(master_key)) = (&store, config.secrets.master_key()) { + match SecretsCrypto::new(master_key.clone()) { + Ok(crypto) => Some(Arc::new(PostgresSecretsStore::new( + store.pool(), + Arc::new(crypto), + ))), + Err(e) => { + tracing::warn!("Failed to initialize secrets crypto: {}", e); + None + } + } + } else { + None + }; + + // Load configured MCP servers + let mcp_session_manager = Arc::new(McpSessionManager::new()); + if let Some(ref secrets) = secrets_store { + match load_mcp_servers().await { + Ok(servers) => { + let enabled_count = servers.servers.iter().filter(|s| s.enabled).count(); + if enabled_count > 0 { + tracing::info!("Loading {} configured MCP server(s)...", enabled_count); + } + + for server in servers.enabled_servers() { + tracing::debug!( + "Checking authentication for MCP server '{}'...", + server.name + ); + // Check for stored tokens (from either pre-configured OAuth or DCR) + let has_tokens = is_authenticated(server, secrets, "default").await; + tracing::debug!("MCP server '{}' has_tokens={}", server.name, has_tokens); + + let client = if has_tokens || server.requires_auth() { + // Use authenticated client if we have tokens or OAuth is configured + McpClient::new_authenticated( + server.clone(), + Arc::clone(&mcp_session_manager), + Arc::clone(secrets), + "default", + ) + } else { + // No tokens and no OAuth - try unauthenticated + McpClient::new_with_name(&server.name, &server.url) + }; + + tracing::debug!("Fetching tools from MCP server '{}'...", server.name); + match client.list_tools().await { + Ok(mcp_tools) => { + tracing::debug!( + "Got {} tools from MCP server '{}'", + mcp_tools.len(), + server.name + ); + match client.create_tools().await { + Ok(tool_impls) => { + for tool in tool_impls { + tools.register(tool).await; + } + tracing::info!( + "Loaded {} tools from MCP server '{}'", + mcp_tools.len(), + server.name + ); + } + Err(e) => { + tracing::warn!( + "Failed to create tools from MCP server '{}': {}", + server.name, + e + ); + } + } + } + Err(e) => { + // Check if it's an auth error + let err_str = e.to_string(); + if err_str.contains("401") || err_str.contains("authentication") { + tracing::warn!( + "MCP server '{}' requires authentication. Run: ironclaw mcp auth {}", + server.name, + server.name + ); + } else { + tracing::warn!( + "Failed to connect to MCP server '{}': {}", + server.name, + e + ); + } + } + } + } + } + Err(e) => { + tracing::debug!("No MCP servers configured ({})", e); + } + } + } + tracing::info!( "Tool registry initialized with {} total tools", tools.count() @@ -345,23 +460,6 @@ async fn main() -> anyhow::Result<()> { } } - // Create secrets store if master key is configured (needed for Telegram webhook registration) - let secrets_store: Option> = - if let (Some(store), Some(master_key)) = (&store, config.secrets.master_key()) { - match SecretsCrypto::new(master_key.clone()) { - Ok(crypto) => Some(Arc::new(PostgresSecretsStore::new( - store.pool(), - Arc::new(crypto), - ))), - Err(e) => { - tracing::warn!("Failed to initialize secrets crypto: {}", e); - None - } - } - } else { - None - }; - // Load WASM channels if enabled if config.channels.wasm_channels_enabled && config.channels.wasm_channels_dir.exists() { match WasmChannelRuntime::new(WasmChannelRuntimeConfig::default()) { diff --git a/src/tools/mcp/auth.rs b/src/tools/mcp/auth.rs new file mode 100644 index 00000000..98b61afb --- /dev/null +++ b/src/tools/mcp/auth.rs @@ -0,0 +1,904 @@ +//! OAuth 2.1 authentication for MCP servers. +//! +//! Implements the MCP Authorization specification using OAuth 2.1 with PKCE. +//! See: https://spec.modelcontextprotocol.io/specification/2025-03-26/basic/authorization/ + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; +use rand::RngCore; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; +use tokio::net::TcpListener; + +use crate::secrets::{CreateSecretParams, SecretsStore}; +use crate::tools::mcp::config::McpServerConfig; + +/// OAuth authorization error. +#[derive(Debug, thiserror::Error)] +pub enum AuthError { + #[error("Server does not support OAuth authorization")] + NotSupported, + + #[error("Failed to discover authorization endpoints: {0}")] + DiscoveryFailed(String), + + #[error("Authorization denied by user")] + AuthorizationDenied, + + #[error("Token exchange failed: {0}")] + TokenExchangeFailed(String), + + #[error("Token expired and refresh failed: {0}")] + RefreshFailed(String), + + #[error("No access token available")] + NoToken, + + #[error("Timeout waiting for authorization callback")] + Timeout, + + #[error("Could not bind to callback port")] + PortUnavailable, + + #[error("HTTP error: {0}")] + Http(String), + + #[error("Secrets error: {0}")] + Secrets(String), +} + +/// OAuth protected resource metadata. +/// Discovered from /.well-known/oauth-protected-resource. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ProtectedResourceMetadata { + /// The protected resource identifier. + pub resource: String, + + /// Authorization servers that can issue tokens for this resource. + #[serde(default)] + pub authorization_servers: Vec, + + /// Scopes supported by this resource. + #[serde(default)] + pub scopes_supported: Vec, +} + +/// OAuth authorization server metadata. +/// Discovered from /.well-known/oauth-authorization-server. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AuthorizationServerMetadata { + /// Authorization server issuer. + pub issuer: String, + + /// Authorization endpoint URL. + pub authorization_endpoint: String, + + /// Token endpoint URL. + pub token_endpoint: String, + + /// Dynamic client registration endpoint (if DCR is supported). + #[serde(default)] + pub registration_endpoint: Option, + + /// Supported response types. + #[serde(default)] + pub response_types_supported: Vec, + + /// Supported grant types. + #[serde(default)] + pub grant_types_supported: Vec, + + /// Supported code challenge methods. + #[serde(default)] + pub code_challenge_methods_supported: Vec, + + /// Scopes supported by this server. + #[serde(default)] + pub scopes_supported: Vec, +} + +/// Dynamic Client Registration request. +#[derive(Debug, Clone, Serialize)] +pub struct ClientRegistrationRequest { + /// Human-readable client name. + pub client_name: String, + + /// Redirect URIs for OAuth callbacks. + pub redirect_uris: Vec, + + /// Grant types the client will use. + pub grant_types: Vec, + + /// Response types the client will use. + pub response_types: Vec, + + /// Token endpoint authentication method. + pub token_endpoint_auth_method: String, +} + +/// Dynamic Client Registration response. +#[derive(Debug, Clone, Deserialize)] +pub struct ClientRegistrationResponse { + /// The assigned client ID. + pub client_id: String, + + /// Client secret (if issued). + #[serde(default)] + pub client_secret: Option, + + /// When the client secret expires (if applicable). + #[serde(default)] + pub client_secret_expires_at: Option, + + /// Registration access token for managing the registration. + #[serde(default)] + pub registration_access_token: Option, + + /// Registration client URI for managing the registration. + #[serde(default)] + pub registration_client_uri: Option, +} + +/// Access token with optional refresh token and expiry. +#[derive(Debug, Clone)] +pub struct AccessToken { + /// The access token value. + pub access_token: String, + + /// Token type (usually "Bearer"). + pub token_type: String, + + /// Seconds until expiration (if provided). + pub expires_in: Option, + + /// Refresh token for obtaining new access tokens. + pub refresh_token: Option, + + /// Scopes granted. + pub scope: Option, +} + +/// Token response from the authorization server. +#[derive(Debug, Deserialize)] +struct TokenResponse { + access_token: String, + token_type: String, + expires_in: Option, + refresh_token: Option, + scope: Option, +} + +/// PKCE verifier and challenge pair. +#[derive(Debug, Clone)] +pub struct PkceChallenge { + /// Code verifier (high-entropy random string). + pub verifier: String, + /// Code challenge (S256 hash of verifier). + pub challenge: String, +} + +impl PkceChallenge { + /// Generate a new PKCE challenge pair. + pub fn generate() -> Self { + let mut verifier_bytes = [0u8; 32]; + rand::thread_rng().fill_bytes(&mut verifier_bytes); + let verifier = URL_SAFE_NO_PAD.encode(verifier_bytes); + + let mut hasher = Sha256::new(); + hasher.update(verifier.as_bytes()); + let challenge = URL_SAFE_NO_PAD.encode(hasher.finalize()); + + Self { + verifier, + challenge, + } + } +} + +/// Discover protected resource metadata from an MCP server. +pub async fn discover_protected_resource( + server_url: &str, +) -> Result { + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(10)) + .build() + .map_err(|e| AuthError::Http(e.to_string()))?; + + // Parse the server URL to extract the origin (scheme + host + port) + // The .well-known endpoints are always at the root of the origin, not under any path + let parsed = reqwest::Url::parse(server_url) + .map_err(|e| AuthError::DiscoveryFailed(format!("Invalid server URL: {}", e)))?; + let origin = parsed.origin().ascii_serialization(); + + // Try the well-known endpoint at the origin root + let well_known_url = format!("{}/.well-known/oauth-protected-resource", origin); + + let response = client + .get(&well_known_url) + .send() + .await + .map_err(|e| AuthError::DiscoveryFailed(e.to_string()))?; + + if !response.status().is_success() { + return Err(AuthError::NotSupported); + } + + response + .json() + .await + .map_err(|e| AuthError::DiscoveryFailed(format!("Invalid metadata: {}", e))) +} + +/// Discover authorization server metadata. +pub async fn discover_authorization_server( + auth_server_url: &str, +) -> Result { + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(10)) + .build() + .map_err(|e| AuthError::Http(e.to_string()))?; + + let base_url = auth_server_url.trim_end_matches('/'); + let well_known_url = format!("{}/.well-known/oauth-authorization-server", base_url); + + let response = client + .get(&well_known_url) + .send() + .await + .map_err(|e| AuthError::DiscoveryFailed(e.to_string()))?; + + if !response.status().is_success() { + return Err(AuthError::DiscoveryFailed(format!( + "HTTP {}", + response.status() + ))); + } + + response + .json() + .await + .map_err(|e| AuthError::DiscoveryFailed(format!("Invalid metadata: {}", e))) +} + +/// Discover OAuth endpoints for an MCP server. +/// +/// First checks if endpoints are explicitly configured, then falls back to discovery. +pub async fn discover_oauth_endpoints( + server_config: &McpServerConfig, +) -> Result<(String, String), AuthError> { + let oauth = server_config + .oauth + .as_ref() + .ok_or(AuthError::NotSupported)?; + + // If endpoints are explicitly configured, use them + if let (Some(auth_url), Some(token_url)) = (&oauth.authorization_url, &oauth.token_url) { + return Ok((auth_url.clone(), token_url.clone())); + } + + // Try to discover from the server + let resource_meta = discover_protected_resource(&server_config.url).await?; + + // Get the first authorization server + let auth_server_url = resource_meta + .authorization_servers + .first() + .ok_or_else(|| AuthError::DiscoveryFailed("No authorization servers listed".to_string()))?; + + // Discover the authorization server metadata + let auth_meta = discover_authorization_server(auth_server_url).await?; + + Ok((auth_meta.authorization_endpoint, auth_meta.token_endpoint)) +} + +/// Discover full OAuth metadata including DCR support. +/// +/// Returns authorization server metadata which includes registration_endpoint if DCR is supported. +pub async fn discover_full_oauth_metadata( + server_url: &str, +) -> Result { + // Try to discover from the server + let resource_meta = discover_protected_resource(server_url).await?; + + // Get the first authorization server + let auth_server_url = resource_meta + .authorization_servers + .first() + .ok_or_else(|| AuthError::DiscoveryFailed("No authorization servers listed".to_string()))?; + + // Discover the authorization server metadata + discover_authorization_server(auth_server_url).await +} + +/// Perform Dynamic Client Registration with an authorization server. +/// +/// This allows clients to register themselves at runtime without pre-configured credentials. +pub async fn register_client( + registration_endpoint: &str, + redirect_uri: &str, +) -> Result { + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(30)) + .build() + .map_err(|e| AuthError::Http(e.to_string()))?; + + let request = ClientRegistrationRequest { + client_name: "IronClaw".to_string(), + redirect_uris: vec![redirect_uri.to_string()], + grant_types: vec![ + "authorization_code".to_string(), + "refresh_token".to_string(), + ], + response_types: vec!["code".to_string()], + token_endpoint_auth_method: "none".to_string(), // Public client (no secret) + }; + + let response = client + .post(registration_endpoint) + .json(&request) + .send() + .await + .map_err(|e| AuthError::DiscoveryFailed(format!("DCR request failed: {}", e)))?; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + return Err(AuthError::DiscoveryFailed(format!( + "DCR failed: HTTP {} - {}", + status, body + ))); + } + + response + .json() + .await + .map_err(|e| AuthError::DiscoveryFailed(format!("Invalid DCR response: {}", e))) +} + +/// Perform the OAuth 2.1 authorization flow for an MCP server. +/// +/// Supports two modes: +/// 1. Pre-configured OAuth: Uses the client_id from server config +/// 2. Dynamic Client Registration: Discovers and registers with the server automatically +/// +/// Flow: +/// 1. Discovers authorization endpoints from the server +/// 2. If no client_id configured, attempts Dynamic Client Registration (DCR) +/// 3. Generates PKCE challenge +/// 4. Opens browser for user authorization +/// 5. Receives callback with authorization code +/// 6. Exchanges code for access token +/// 7. Stores token securely +pub async fn authorize_mcp_server( + server_config: &McpServerConfig, + secrets: &Arc, + user_id: &str, +) -> Result { + // Find an available port for the callback first (needed for DCR) + let (listener, port) = find_available_port().await?; + let redirect_uri = format!("http://localhost:{}/callback", port); + + // Determine client_id and endpoints + let (client_id, authorization_url, token_url, use_pkce, scopes, extra_params) = + if let Some(oauth) = &server_config.oauth { + // Pre-configured OAuth + let (auth_url, tok_url) = discover_oauth_endpoints(server_config).await?; + ( + oauth.client_id.clone(), + auth_url, + tok_url, + oauth.use_pkce, + oauth.scopes.clone(), + oauth.extra_params.clone(), + ) + } else { + // Try Dynamic Client Registration + println!(" Discovering OAuth endpoints..."); + let auth_meta = discover_full_oauth_metadata(&server_config.url).await?; + + let registration_endpoint = auth_meta + .registration_endpoint + .ok_or(AuthError::NotSupported)?; + + println!(" Registering client dynamically..."); + let registration = register_client(®istration_endpoint, &redirect_uri).await?; + println!(" ✓ Client registered: {}", registration.client_id); + + ( + registration.client_id, + auth_meta.authorization_endpoint, + auth_meta.token_endpoint, + true, // Always use PKCE for DCR clients + auth_meta.scopes_supported, + HashMap::new(), + ) + }; + + // Generate PKCE challenge + let pkce = if use_pkce { + Some(PkceChallenge::generate()) + } else { + None + }; + + // Build authorization URL + let auth_url = build_authorization_url( + &authorization_url, + &client_id, + &redirect_uri, + &scopes, + pkce.as_ref(), + &extra_params, + ); + + // Open browser + println!(" Opening browser for {} login...", server_config.name); + if let Err(e) = open::that(&auth_url) { + println!(" Could not open browser: {}", e); + println!(" Please open this URL manually:"); + println!(" {}", auth_url); + } + + println!(" Waiting for authorization..."); + + // Wait for callback + let code = wait_for_authorization_callback(listener, &server_config.name).await?; + + println!(" Exchanging code for token..."); + + // Exchange code for token + let token = + exchange_code_for_token(&token_url, &client_id, &code, &redirect_uri, pkce.as_ref()) + .await?; + + // Store the tokens + store_tokens(secrets, user_id, server_config, &token).await?; + + // Store the client_id for DCR (needed for token refresh) + if server_config.oauth.is_none() { + store_client_id(secrets, user_id, server_config, &client_id).await?; + } + + Ok(token) +} + +/// Find an available port for the OAuth callback. +async fn find_available_port() -> Result<(TcpListener, u16), AuthError> { + for port in 9876..=9886 { + if let Ok(listener) = TcpListener::bind(format!("127.0.0.1:{}", port)).await { + return Ok((listener, port)); + } + } + Err(AuthError::PortUnavailable) +} + +/// Build the authorization URL with all required parameters. +fn build_authorization_url( + base_url: &str, + client_id: &str, + redirect_uri: &str, + scopes: &[String], + pkce: Option<&PkceChallenge>, + extra_params: &HashMap, +) -> String { + let mut url = format!( + "{}?client_id={}&response_type=code&redirect_uri={}", + base_url, + urlencoding::encode(client_id), + urlencoding::encode(redirect_uri) + ); + + if !scopes.is_empty() { + url.push_str(&format!( + "&scope={}", + urlencoding::encode(&scopes.join(" ")) + )); + } + + if let Some(pkce) = pkce { + url.push_str(&format!( + "&code_challenge={}&code_challenge_method=S256", + pkce.challenge + )); + } + + for (key, value) in extra_params { + url.push_str(&format!( + "&{}={}", + urlencoding::encode(key), + urlencoding::encode(value) + )); + } + + url +} + +/// Wait for the authorization callback and extract the code. +async fn wait_for_authorization_callback( + listener: TcpListener, + server_name: &str, +) -> Result { + let timeout = Duration::from_secs(300); + + tokio::time::timeout(timeout, async { + loop { + let (mut socket, _) = listener + .accept() + .await + .map_err(|e| AuthError::Http(e.to_string()))?; + + let mut reader = BufReader::new(&mut socket); + let mut request_line = String::new(); + reader + .read_line(&mut request_line) + .await + .map_err(|e| AuthError::Http(e.to_string()))?; + + // Parse GET /callback?code=xxx HTTP/1.1 + if let Some(path) = request_line.split_whitespace().nth(1) { + if path.starts_with("/callback") { + if let Some(query) = path.split('?').nth(1) { + // Check for error first + if query.contains("error=") { + let response = "HTTP/1.1 400 Bad Request\r\n\r\nAuthorization denied"; + let _ = socket.write_all(response.as_bytes()).await; + return Err(AuthError::AuthorizationDenied); + } + + // Look for code + for param in query.split('&') { + let parts: Vec<&str> = param.splitn(2, '=').collect(); + if parts.len() == 2 && parts[0] == "code" { + let code = urlencoding::decode(parts[1]) + .unwrap_or_else(|_| parts[1].into()) + .into_owned(); + + // Send success response + let response = format!( + "HTTP/1.1 200 OK\r\n\ + Content-Type: text/html\r\n\ + \r\n\ + \ +
\ +

✓ {} Connected!

\ +

You can close this window.

\ +
", + server_name + ); + let _ = socket.write_all(response.as_bytes()).await; + let _ = socket.shutdown().await; + + return Ok(code); + } + } + } + } + } + + let response = "HTTP/1.1 404 Not Found\r\n\r\n"; + let _ = socket.write_all(response.as_bytes()).await; + } + }) + .await + .map_err(|_| AuthError::Timeout)? +} + +/// Exchange the authorization code for an access token. +async fn exchange_code_for_token( + token_url: &str, + client_id: &str, + code: &str, + redirect_uri: &str, + pkce: Option<&PkceChallenge>, +) -> Result { + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(30)) + .build() + .map_err(|e| AuthError::Http(e.to_string()))?; + + let mut params = vec![ + ("grant_type", "authorization_code".to_string()), + ("code", code.to_string()), + ("redirect_uri", redirect_uri.to_string()), + ("client_id", client_id.to_string()), + ]; + + if let Some(pkce) = pkce { + params.push(("code_verifier", pkce.verifier.clone())); + } + + let response = client + .post(token_url) + .form(¶ms) + .send() + .await + .map_err(|e| AuthError::TokenExchangeFailed(e.to_string()))?; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + return Err(AuthError::TokenExchangeFailed(format!( + "HTTP {} - {}", + status, body + ))); + } + + let token_response: TokenResponse = response + .json() + .await + .map_err(|e| AuthError::TokenExchangeFailed(format!("Invalid response: {}", e)))?; + + Ok(AccessToken { + access_token: token_response.access_token, + token_type: token_response.token_type, + expires_in: token_response.expires_in, + refresh_token: token_response.refresh_token, + scope: token_response.scope, + }) +} + +/// Store access and refresh tokens securely. +async fn store_tokens( + secrets: &Arc, + user_id: &str, + server_config: &McpServerConfig, + token: &AccessToken, +) -> Result<(), AuthError> { + // Store access token + let params = CreateSecretParams::new(server_config.token_secret_name(), &token.access_token) + .with_provider(format!("mcp:{}", server_config.name)); + + secrets + .create(user_id, params) + .await + .map_err(|e| AuthError::Secrets(e.to_string()))?; + + // Store refresh token if present + if let Some(ref refresh_token) = token.refresh_token { + let params = + CreateSecretParams::new(server_config.refresh_token_secret_name(), refresh_token) + .with_provider(format!("mcp:{}", server_config.name)); + + secrets + .create(user_id, params) + .await + .map_err(|e| AuthError::Secrets(e.to_string()))?; + } + + Ok(()) +} + +/// Store the DCR client ID for future token refresh. +async fn store_client_id( + secrets: &Arc, + user_id: &str, + server_config: &McpServerConfig, + client_id: &str, +) -> Result<(), AuthError> { + let params = CreateSecretParams::new(server_config.client_id_secret_name(), client_id) + .with_provider(format!("mcp:{}", server_config.name)); + + secrets + .create(user_id, params) + .await + .map(|_| ()) + .map_err(|e| AuthError::Secrets(e.to_string())) +} + +/// Get the client ID for a server (from config or stored DCR). +async fn get_client_id( + server_config: &McpServerConfig, + secrets: &Arc, + user_id: &str, +) -> Result { + // First check if OAuth is configured with a client_id + if let Some(ref oauth) = server_config.oauth { + return Ok(oauth.client_id.clone()); + } + + // Otherwise try to get the DCR client_id from secrets + match secrets + .get_decrypted(user_id, &server_config.client_id_secret_name()) + .await + { + Ok(client_id) => Ok(client_id.expose().to_string()), + Err(crate::secrets::SecretError::NotFound(_)) => Err(AuthError::RefreshFailed( + "No client ID found. Please re-authenticate.".to_string(), + )), + Err(e) => Err(AuthError::Secrets(e.to_string())), + } +} + +/// Get the stored access token for an MCP server. +pub async fn get_access_token( + server_config: &McpServerConfig, + secrets: &Arc, + user_id: &str, +) -> Result, AuthError> { + match secrets + .get_decrypted(user_id, &server_config.token_secret_name()) + .await + { + Ok(token) => Ok(Some(token.expose().to_string())), + Err(crate::secrets::SecretError::NotFound(_)) => Ok(None), + Err(e) => Err(AuthError::Secrets(e.to_string())), + } +} + +/// Check if a server has valid authentication. +/// +/// Returns true if: +/// - A valid access token is stored (regardless of how it was obtained) +/// - The server doesn't require authentication at all +pub async fn is_authenticated( + server_config: &McpServerConfig, + secrets: &Arc, + user_id: &str, +) -> bool { + // Check if we have a stored token (from either pre-configured OAuth or DCR) + secrets + .exists(user_id, &server_config.token_secret_name()) + .await + .unwrap_or(false) +} + +/// Refresh an access token using the refresh token. +/// +/// Works with both pre-configured OAuth and Dynamic Client Registration (DCR). +/// For DCR, retrieves the client_id from stored secrets. +pub async fn refresh_access_token( + server_config: &McpServerConfig, + secrets: &Arc, + user_id: &str, +) -> Result { + // Get client_id (from config or stored DCR) + let client_id = get_client_id(server_config, secrets, user_id).await?; + + // Get the refresh token + let refresh_token = secrets + .get_decrypted(user_id, &server_config.refresh_token_secret_name()) + .await + .map_err(|e| AuthError::RefreshFailed(format!("No refresh token: {}", e)))?; + + // Discover the token endpoint + let token_url = if let Some(ref oauth) = server_config.oauth { + if let Some(ref url) = oauth.token_url { + url.clone() + } else { + // Discover from server + let auth_meta = discover_full_oauth_metadata(&server_config.url).await?; + auth_meta.token_endpoint + } + } else { + // DCR - always discover + let auth_meta = discover_full_oauth_metadata(&server_config.url).await?; + auth_meta.token_endpoint + }; + + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(30)) + .build() + .map_err(|e| AuthError::Http(e.to_string()))?; + + let params = vec![ + ("grant_type", "refresh_token".to_string()), + ("refresh_token", refresh_token.expose().to_string()), + ("client_id", client_id), + ]; + + let response = client + .post(&token_url) + .form(¶ms) + .send() + .await + .map_err(|e| AuthError::RefreshFailed(e.to_string()))?; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + return Err(AuthError::RefreshFailed(format!( + "HTTP {} - {}", + status, body + ))); + } + + let token_response: TokenResponse = response + .json() + .await + .map_err(|e| AuthError::RefreshFailed(format!("Invalid response: {}", e)))?; + + let token = AccessToken { + access_token: token_response.access_token, + token_type: token_response.token_type, + expires_in: token_response.expires_in, + refresh_token: token_response.refresh_token, + scope: token_response.scope, + }; + + // Store the new tokens + store_tokens(secrets, user_id, server_config, &token).await?; + + Ok(token) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_pkce_challenge_generation() { + let pkce = PkceChallenge::generate(); + + // Verifier should be base64url encoded + assert!(!pkce.verifier.is_empty()); + assert!(!pkce.verifier.contains('+')); + assert!(!pkce.verifier.contains('/')); + assert!(!pkce.verifier.contains('=')); + + // Challenge should be different from verifier + assert_ne!(pkce.verifier, pkce.challenge); + + // Two challenges should be different + let pkce2 = PkceChallenge::generate(); + assert_ne!(pkce.verifier, pkce2.verifier); + } + + #[test] + fn test_build_authorization_url() { + let url = build_authorization_url( + "https://auth.example.com/authorize", + "client-123", + "http://localhost:9876/callback", + &["read".to_string(), "write".to_string()], + None, + &HashMap::new(), + ); + + assert!(url.starts_with("https://auth.example.com/authorize?")); + assert!(url.contains("client_id=client-123")); + assert!(url.contains("response_type=code")); + assert!(url.contains("redirect_uri=")); + assert!(url.contains("scope=read%20write")); + } + + #[test] + fn test_build_authorization_url_with_pkce() { + let pkce = PkceChallenge::generate(); + let url = build_authorization_url( + "https://auth.example.com/authorize", + "client-123", + "http://localhost:9876/callback", + &[], + Some(&pkce), + &HashMap::new(), + ); + + assert!(url.contains(&format!("code_challenge={}", pkce.challenge))); + assert!(url.contains("code_challenge_method=S256")); + } + + #[test] + fn test_build_authorization_url_with_extra_params() { + let mut extra = HashMap::new(); + extra.insert("owner".to_string(), "user".to_string()); + extra.insert("state".to_string(), "abc123".to_string()); + + let url = build_authorization_url( + "https://auth.example.com/authorize", + "client-123", + "http://localhost:9876/callback", + &[], + None, + &extra, + ); + + assert!(url.contains("owner=user")); + assert!(url.contains("state=abc123")); + } +} diff --git a/src/tools/mcp/client.rs b/src/tools/mcp/client.rs index 45ea3b7b..e859f173 100644 --- a/src/tools/mcp/client.rs +++ b/src/tools/mcp/client.rs @@ -1,4 +1,7 @@ //! MCP client for connecting to MCP servers. +//! +//! Supports both local (unauthenticated) and hosted (OAuth-authenticated) servers. +//! Uses the Streamable HTTP transport with session management. use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; @@ -8,63 +11,357 @@ use async_trait::async_trait; use tokio::sync::RwLock; use crate::context::JobContext; +use crate::secrets::SecretsStore; +use crate::tools::mcp::auth::refresh_access_token; +use crate::tools::mcp::config::McpServerConfig; use crate::tools::mcp::protocol::{ - CallToolResult, ListToolsResult, McpRequest, McpResponse, McpTool, + CallToolResult, InitializeResult, ListToolsResult, McpRequest, McpResponse, McpTool, }; +use crate::tools::mcp::session::McpSessionManager; use crate::tools::tool::{Tool, ToolError, ToolOutput}; /// MCP client for communicating with MCP servers. +/// +/// Supports two modes: +/// - Simple: Just a URL, no auth or session management (for local/test servers) +/// - Authenticated: Full OAuth support with session management (for hosted servers) pub struct McpClient { /// Server URL (for HTTP transport). server_url: String, + + /// Server name (for logging and session management). + server_name: String, + /// HTTP client. http_client: reqwest::Client, + /// Request ID counter. next_id: AtomicU64, + /// Cached tools. tools_cache: RwLock>>, + + /// Session manager (shared across clients). + session_manager: Option>, + + /// Secrets store for retrieving access tokens. + secrets: Option>, + + /// User ID for secrets lookup. + user_id: String, + + /// Server configuration (for token secret name lookup). + server_config: Option, } impl McpClient { - /// Create a new MCP client. + /// Create a new simple MCP client (no authentication). + /// + /// Use this for local development servers or servers that don't require auth. pub fn new(server_url: impl Into) -> Self { + let url = server_url.into(); + let name = extract_server_name(&url); + Self { - server_url: server_url.into(), + server_url: url, + server_name: name, http_client: reqwest::Client::builder() .timeout(Duration::from_secs(30)) .build() .expect("Failed to create HTTP client"), next_id: AtomicU64::new(1), tools_cache: RwLock::new(None), + session_manager: None, + secrets: None, + user_id: "default".to_string(), + server_config: None, } } + /// Create a new simple MCP client with a specific name. + /// + /// Use this when you have a configured server name but no authentication. + pub fn new_with_name(server_name: impl Into, server_url: impl Into) -> Self { + Self { + server_url: server_url.into(), + server_name: server_name.into(), + http_client: reqwest::Client::builder() + .timeout(Duration::from_secs(30)) + .build() + .expect("Failed to create HTTP client"), + next_id: AtomicU64::new(1), + tools_cache: RwLock::new(None), + session_manager: None, + secrets: None, + user_id: "default".to_string(), + server_config: None, + } + } + + /// Create a new authenticated MCP client. + /// + /// Use this for hosted MCP servers that require OAuth authentication. + pub fn new_authenticated( + config: McpServerConfig, + session_manager: Arc, + secrets: Arc, + user_id: impl Into, + ) -> Self { + Self { + server_url: config.url.clone(), + server_name: config.name.clone(), + http_client: reqwest::Client::builder() + .timeout(Duration::from_secs(30)) + .build() + .expect("Failed to create HTTP client"), + next_id: AtomicU64::new(1), + tools_cache: RwLock::new(None), + session_manager: Some(session_manager), + secrets: Some(secrets), + user_id: user_id.into(), + server_config: Some(config), + } + } + + /// Get the server name. + pub fn server_name(&self) -> &str { + &self.server_name + } + + /// Get the server URL. + pub fn server_url(&self) -> &str { + &self.server_url + } + /// Get the next request ID. fn next_request_id(&self) -> u64 { self.next_id.fetch_add(1, Ordering::SeqCst) } - /// Send a request to the MCP server. - async fn send_request(&self, request: McpRequest) -> Result { - let response = self - .http_client - .post(&self.server_url) - .json(&request) - .send() + /// Get the access token for this server (if authenticated). + /// + /// Returns the stored token regardless of whether OAuth was pre-configured + /// or obtained via Dynamic Client Registration. + async fn get_access_token(&self) -> Result, ToolError> { + let Some(ref secrets) = self.secrets else { + return Ok(None); + }; + + let Some(ref config) = self.server_config else { + return Ok(None); + }; + + // Try to get stored token (from either pre-configured OAuth or DCR) + match secrets + .get_decrypted(&self.user_id, &config.token_secret_name()) .await - .map_err(|e| ToolError::ExternalService(format!("MCP request failed: {}", e)))?; + { + Ok(token) => Ok(Some(token.expose().to_string())), + Err(crate::secrets::SecretError::NotFound(_)) => Ok(None), + Err(e) => Err(ToolError::ExternalService(format!( + "Failed to get access token: {}", + e + ))), + } + } + + /// Send a request to the MCP server with auth and session headers. + /// Automatically attempts token refresh on 401 errors. + async fn send_request(&self, request: McpRequest) -> Result { + // Try up to 2 times: first attempt, then retry after token refresh + for attempt in 0..2 { + // Request both JSON and SSE as per MCP spec + let mut req_builder = self + .http_client + .post(&self.server_url) + .header("Accept", "application/json, text/event-stream") + .header("Content-Type", "application/json") + .json(&request); + + // Add Authorization header if we have a token + if let Some(token) = self.get_access_token().await? { + req_builder = req_builder.header("Authorization", format!("Bearer {}", token)); + } + + // Add Mcp-Session-Id header if we have a session + if let Some(ref session_manager) = self.session_manager { + if let Some(session_id) = session_manager.get_session_id(&self.server_name).await { + req_builder = req_builder.header("Mcp-Session-Id", session_id); + } + } + + let response = req_builder + .send() + .await + .map_err(|e| ToolError::ExternalService(format!("MCP request failed: {}", e)))?; + + // Check for 401 Unauthorized - try to refresh token on first attempt + if response.status() == reqwest::StatusCode::UNAUTHORIZED { + if attempt == 0 { + // Try to refresh the token + if let Some(ref secrets) = self.secrets { + if let Some(ref config) = self.server_config { + tracing::debug!( + "MCP token expired, attempting refresh for '{}'", + self.server_name + ); + match refresh_access_token(config, secrets, &self.user_id).await { + Ok(_) => { + tracing::info!( + "MCP token refreshed for '{}'", + self.server_name + ); + // Continue to next iteration to retry with new token + continue; + } + Err(e) => { + tracing::debug!( + "Token refresh failed for '{}': {}", + self.server_name, + e + ); + // Fall through to return auth error + } + } + } + } + } + return Err(ToolError::ExternalService(format!( + "MCP server '{}' requires authentication. Run: ironclaw mcp auth {}", + self.server_name, self.server_name + ))); + } + + // Success path - return the parsed response + return self.parse_response(response).await; + } + + // Should not reach here, but just in case + Err(ToolError::ExternalService( + "MCP request failed after retry".to_string(), + )) + } + + /// Parse the HTTP response into an MCP response. + async fn parse_response(&self, response: reqwest::Response) -> Result { + // Extract session ID from response header + if let Some(ref session_manager) = self.session_manager { + if let Some(session_id) = response + .headers() + .get("Mcp-Session-Id") + .and_then(|v| v.to_str().ok()) + { + session_manager + .update_session_id(&self.server_name, Some(session_id.to_string())) + .await; + } + } if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); return Err(ToolError::ExternalService(format!( - "MCP server returned status: {}", - response.status() + "MCP server returned status: {} - {}", + status, body ))); } - response - .json() - .await - .map_err(|e| ToolError::ExternalService(format!("Failed to parse MCP response: {}", e))) + // Check content type to handle SSE vs JSON responses + let content_type = response + .headers() + .get("content-type") + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .to_string(); + + if content_type.contains("text/event-stream") { + // SSE response - read chunks until we get a complete JSON message + use futures::StreamExt; + + let mut stream = response.bytes_stream(); + let mut buffer = String::new(); + + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(|e| { + ToolError::ExternalService(format!("Failed to read SSE chunk: {}", e)) + })?; + + buffer.push_str(&String::from_utf8_lossy(&chunk)); + + // Look for complete SSE data lines + for line in buffer.lines() { + if let Some(json_str) = line.strip_prefix("data: ") { + // Try to parse - if valid JSON, we're done + if let Ok(response) = serde_json::from_str::(json_str) { + return Ok(response); + } + } + } + } + + Err(ToolError::ExternalService(format!( + "No valid data in SSE response: {}", + buffer + ))) + } else { + // JSON response + response.json().await.map_err(|e| { + ToolError::ExternalService(format!("Failed to parse MCP response: {}", e)) + }) + } + } + + /// Initialize the connection to the MCP server. + /// + /// This should be called once per session to establish capabilities. + pub async fn initialize(&self) -> Result { + // Check if already initialized + if let Some(ref session_manager) = self.session_manager { + if session_manager.is_initialized(&self.server_name).await { + // Return cached/default capabilities + return Ok(InitializeResult::default()); + } + } + + // Ensure we have a session + if let Some(ref session_manager) = self.session_manager { + session_manager + .get_or_create(&self.server_name, &self.server_url) + .await; + } + + let request = McpRequest::initialize(self.next_request_id()); + let response = self.send_request(request).await?; + + if let Some(error) = response.error { + return Err(ToolError::ExternalService(format!( + "MCP initialization error: {} (code {})", + error.message, error.code + ))); + } + + let result: InitializeResult = response + .result + .ok_or_else(|| { + ToolError::ExternalService("No result in initialize response".to_string()) + }) + .and_then(|r| { + serde_json::from_value(r).map_err(|e| { + ToolError::ExternalService(format!("Invalid initialize result: {}", e)) + }) + })?; + + // Mark session as initialized + if let Some(ref session_manager) = self.session_manager { + session_manager.mark_initialized(&self.server_name).await; + } + + // Send initialized notification + let notification = McpRequest::initialized_notification(); + // Fire and forget - notifications don't have responses + let _ = self.send_request(notification).await; + + Ok(result) } /// List available tools from the MCP server. @@ -74,6 +371,11 @@ impl McpClient { return Ok(tools.clone()); } + // Ensure initialized for authenticated sessions + if self.session_manager.is_some() { + self.initialize().await?; + } + let request = McpRequest::list_tools(self.next_request_id()); let response = self.send_request(request).await?; @@ -104,6 +406,11 @@ impl McpClient { name: &str, arguments: serde_json::Value, ) -> Result { + // Ensure initialized for authenticated sessions + if self.session_manager.is_some() { + self.initialize().await?; + } + let request = McpRequest::call_tool(self.next_request_id(), name, arguments); let response = self.send_request(request).await?; @@ -136,36 +443,61 @@ impl McpClient { Ok(mcp_tools .into_iter() .map(|t| { + let prefixed_name = format!("{}_{}", self.server_name, t.name); Arc::new(McpToolWrapper { tool: t, + prefixed_name, client: client.clone(), }) as Arc }) .collect()) } + + /// Test the connection to the MCP server. + pub async fn test_connection(&self) -> Result<(), ToolError> { + self.initialize().await?; + self.list_tools().await?; + Ok(()) + } } impl Clone for McpClient { fn clone(&self) -> Self { Self { server_url: self.server_url.clone(), + server_name: self.server_name.clone(), http_client: self.http_client.clone(), next_id: AtomicU64::new(self.next_id.load(Ordering::SeqCst)), tools_cache: RwLock::new(None), + session_manager: self.session_manager.clone(), + secrets: self.secrets.clone(), + user_id: self.user_id.clone(), + server_config: self.server_config.clone(), } } } +/// Extract a server name from a URL for logging/display purposes. +fn extract_server_name(url: &str) -> String { + reqwest::Url::parse(url) + .ok() + .and_then(|u| u.host_str().map(|h| h.to_string())) + .unwrap_or_else(|| "unknown".to_string()) + .replace('.', "_") +} + /// Wrapper that implements Tool for an MCP tool. struct McpToolWrapper { tool: McpTool, + /// Prefixed name (server_name_tool_name) for unique identification. + prefixed_name: String, client: Arc, } #[async_trait] impl Tool for McpToolWrapper { fn name(&self) -> &str { - &self.tool.name + &self.prefixed_name } fn description(&self) -> &str { @@ -183,6 +515,7 @@ impl Tool for McpToolWrapper { ) -> Result { let start = std::time::Instant::now(); + // Use the original tool name (without prefix) for the actual call let result = self.client.call_tool(&self.tool.name, params).await?; // Convert content blocks to a single result @@ -227,4 +560,22 @@ mod tests { assert_eq!(req.method, "tools/call"); assert!(req.params.is_some()); } + + #[test] + fn test_extract_server_name() { + assert_eq!( + extract_server_name("https://mcp.notion.com/v1"), + "mcp_notion_com" + ); + assert_eq!(extract_server_name("http://localhost:8080"), "localhost"); + assert_eq!(extract_server_name("invalid"), "unknown"); + } + + #[test] + fn test_simple_client_creation() { + let client = McpClient::new("http://localhost:8080"); + assert_eq!(client.server_url(), "http://localhost:8080"); + assert!(client.session_manager.is_none()); + assert!(client.secrets.is_none()); + } } diff --git a/src/tools/mcp/config.rs b/src/tools/mcp/config.rs new file mode 100644 index 00000000..54893654 --- /dev/null +++ b/src/tools/mcp/config.rs @@ -0,0 +1,442 @@ +//! MCP server configuration. +//! +//! Stores configuration for connecting to hosted MCP servers. +//! Configuration is persisted at ~/.ironclaw/mcp-servers.json. + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; + +use serde::{Deserialize, Serialize}; +use tokio::fs; + +use crate::tools::tool::ToolError; + +/// Configuration for connecting to a remote MCP server. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct McpServerConfig { + /// Unique name for this server (e.g., "notion", "github"). + pub name: String, + + /// Server URL (must be HTTPS for remote servers). + pub url: String, + + /// OAuth configuration (if server requires authentication). + #[serde(skip_serializing_if = "Option::is_none")] + pub oauth: Option, + + /// Whether this server is enabled. + #[serde(default = "default_true")] + pub enabled: bool, + + /// Optional description for the server. + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, +} + +fn default_true() -> bool { + true +} + +impl McpServerConfig { + /// Create a new MCP server configuration. + pub fn new(name: impl Into, url: impl Into) -> Self { + Self { + name: name.into(), + url: url.into(), + oauth: None, + enabled: true, + description: None, + } + } + + /// Set OAuth configuration. + pub fn with_oauth(mut self, oauth: OAuthConfig) -> Self { + self.oauth = Some(oauth); + self + } + + /// Set description. + pub fn with_description(mut self, description: impl Into) -> Self { + self.description = Some(description.into()); + self + } + + /// Validate the server configuration. + pub fn validate(&self) -> Result<(), ConfigError> { + if self.name.is_empty() { + return Err(ConfigError::InvalidConfig { + reason: "Server name cannot be empty".to_string(), + }); + } + + if self.url.is_empty() { + return Err(ConfigError::InvalidConfig { + reason: "Server URL cannot be empty".to_string(), + }); + } + + // Remote servers must use HTTPS (localhost is allowed for development) + let url_lower = self.url.to_lowercase(); + let is_localhost = url_lower.contains("localhost") || url_lower.contains("127.0.0.1"); + if !is_localhost && !url_lower.starts_with("https://") { + return Err(ConfigError::InvalidConfig { + reason: "Remote MCP servers must use HTTPS".to_string(), + }); + } + + Ok(()) + } + + /// Check if this server requires authentication. + pub fn requires_auth(&self) -> bool { + self.oauth.is_some() + } + + /// Get the secret name used to store the access token. + pub fn token_secret_name(&self) -> String { + format!("mcp_{}_access_token", self.name) + } + + /// Get the secret name used to store the refresh token. + pub fn refresh_token_secret_name(&self) -> String { + format!("mcp_{}_refresh_token", self.name) + } + + /// Get the secret name used to store the DCR client ID. + pub fn client_id_secret_name(&self) -> String { + format!("mcp_{}_client_id", self.name) + } +} + +/// OAuth 2.1 configuration for an MCP server. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct OAuthConfig { + /// OAuth client ID. + pub client_id: String, + + /// Authorization endpoint URL. + /// If not provided, will be discovered from /.well-known/oauth-protected-resource. + #[serde(skip_serializing_if = "Option::is_none")] + pub authorization_url: Option, + + /// Token endpoint URL. + /// If not provided, will be discovered from /.well-known/oauth-authorization-server. + #[serde(skip_serializing_if = "Option::is_none")] + pub token_url: Option, + + /// Scopes to request. + #[serde(default)] + pub scopes: Vec, + + /// Whether to use PKCE (default: true, as required by OAuth 2.1). + #[serde(default = "default_true")] + pub use_pkce: bool, + + /// Extra parameters to include in the authorization request. + #[serde(default)] + pub extra_params: HashMap, +} + +impl OAuthConfig { + /// Create a new OAuth configuration with just a client ID. + pub fn new(client_id: impl Into) -> Self { + Self { + client_id: client_id.into(), + authorization_url: None, + token_url: None, + scopes: Vec::new(), + use_pkce: true, + extra_params: HashMap::new(), + } + } + + /// Set authorization and token URLs. + pub fn with_endpoints( + mut self, + authorization_url: impl Into, + token_url: impl Into, + ) -> Self { + self.authorization_url = Some(authorization_url.into()); + self.token_url = Some(token_url.into()); + self + } + + /// Set scopes. + pub fn with_scopes(mut self, scopes: Vec) -> Self { + self.scopes = scopes; + self + } +} + +/// Configuration file containing all MCP servers. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct McpServersFile { + /// List of configured MCP servers. + #[serde(default)] + pub servers: Vec, + + /// Schema version for future compatibility. + #[serde(default = "default_schema_version")] + pub schema_version: u32, +} + +fn default_schema_version() -> u32 { + 1 +} + +impl McpServersFile { + /// Get a server by name. + pub fn get(&self, name: &str) -> Option<&McpServerConfig> { + self.servers.iter().find(|s| s.name == name) + } + + /// Get a mutable server by name. + pub fn get_mut(&mut self, name: &str) -> Option<&mut McpServerConfig> { + self.servers.iter_mut().find(|s| s.name == name) + } + + /// Add or update a server configuration. + pub fn upsert(&mut self, config: McpServerConfig) { + if let Some(existing) = self.get_mut(&config.name) { + *existing = config; + } else { + self.servers.push(config); + } + } + + /// Remove a server by name. + pub fn remove(&mut self, name: &str) -> bool { + let len_before = self.servers.len(); + self.servers.retain(|s| s.name != name); + self.servers.len() < len_before + } + + /// Get all enabled servers. + pub fn enabled_servers(&self) -> impl Iterator { + self.servers.iter().filter(|s| s.enabled) + } +} + +/// Error type for MCP configuration operations. +#[derive(Debug, thiserror::Error)] +pub enum ConfigError { + #[error("IO error: {0}")] + Io(#[from] std::io::Error), + + #[error("JSON error: {0}")] + Json(#[from] serde_json::Error), + + #[error("Invalid configuration: {reason}")] + InvalidConfig { reason: String }, + + #[error("Server not found: {name}")] + ServerNotFound { name: String }, +} + +impl From for ToolError { + fn from(err: ConfigError) -> Self { + ToolError::ExternalService(err.to_string()) + } +} + +/// Get the default MCP servers configuration path. +pub fn default_config_path() -> PathBuf { + dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join(".ironclaw") + .join("mcp-servers.json") +} + +/// Load MCP server configurations from the default location. +pub async fn load_mcp_servers() -> Result { + load_mcp_servers_from(default_config_path()).await +} + +/// Load MCP server configurations from a specific path. +pub async fn load_mcp_servers_from(path: impl AsRef) -> Result { + let path = path.as_ref(); + + if !path.exists() { + return Ok(McpServersFile::default()); + } + + let content = fs::read_to_string(path).await?; + let config: McpServersFile = serde_json::from_str(&content)?; + + Ok(config) +} + +/// Save MCP server configurations to the default location. +pub async fn save_mcp_servers(config: &McpServersFile) -> Result<(), ConfigError> { + save_mcp_servers_to(config, default_config_path()).await +} + +/// Save MCP server configurations to a specific path. +pub async fn save_mcp_servers_to( + config: &McpServersFile, + path: impl AsRef, +) -> Result<(), ConfigError> { + let path = path.as_ref(); + + // Ensure parent directory exists + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).await?; + } + + let content = serde_json::to_string_pretty(config)?; + fs::write(path, content).await?; + + Ok(()) +} + +/// Add a new MCP server configuration. +pub async fn add_mcp_server(config: McpServerConfig) -> Result<(), ConfigError> { + config.validate()?; + + let mut servers = load_mcp_servers().await?; + servers.upsert(config); + save_mcp_servers(&servers).await?; + + Ok(()) +} + +/// Remove an MCP server by name. +pub async fn remove_mcp_server(name: &str) -> Result<(), ConfigError> { + let mut servers = load_mcp_servers().await?; + + if !servers.remove(name) { + return Err(ConfigError::ServerNotFound { + name: name.to_string(), + }); + } + + save_mcp_servers(&servers).await?; + + Ok(()) +} + +/// Get a specific MCP server configuration. +pub async fn get_mcp_server(name: &str) -> Result { + let servers = load_mcp_servers().await?; + + servers + .get(name) + .cloned() + .ok_or_else(|| ConfigError::ServerNotFound { + name: name.to_string(), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::tempdir; + + #[test] + fn test_server_config_validation() { + // Valid HTTPS server + let config = McpServerConfig::new("notion", "https://mcp.notion.com"); + assert!(config.validate().is_ok()); + + // Valid localhost (allowed for dev) + let config = McpServerConfig::new("local", "http://localhost:8080"); + assert!(config.validate().is_ok()); + + // Invalid: empty name + let config = McpServerConfig::new("", "https://example.com"); + assert!(config.validate().is_err()); + + // Invalid: HTTP for remote server + let config = McpServerConfig::new("remote", "http://mcp.example.com"); + assert!(config.validate().is_err()); + } + + #[test] + fn test_oauth_config_builder() { + let oauth = OAuthConfig::new("client-123") + .with_endpoints( + "https://auth.example.com/authorize", + "https://auth.example.com/token", + ) + .with_scopes(vec!["read".to_string(), "write".to_string()]); + + assert_eq!(oauth.client_id, "client-123"); + assert!(oauth.authorization_url.is_some()); + assert!(oauth.token_url.is_some()); + assert_eq!(oauth.scopes.len(), 2); + assert!(oauth.use_pkce); + } + + #[test] + fn test_servers_file_operations() { + let mut file = McpServersFile::default(); + + // Add a server + file.upsert(McpServerConfig::new("notion", "https://mcp.notion.com")); + assert_eq!(file.servers.len(), 1); + + // Update the server + let mut updated = McpServerConfig::new("notion", "https://mcp.notion.com/v2"); + updated.enabled = false; + file.upsert(updated); + assert_eq!(file.servers.len(), 1); + assert!(!file.get("notion").unwrap().enabled); + + // Add another server + file.upsert(McpServerConfig::new("github", "https://mcp.github.com")); + assert_eq!(file.servers.len(), 2); + + // Remove a server + assert!(file.remove("notion")); + assert_eq!(file.servers.len(), 1); + assert!(file.get("notion").is_none()); + + // Remove non-existent server + assert!(!file.remove("nonexistent")); + } + + #[tokio::test] + async fn test_load_save_config() { + let dir = tempdir().unwrap(); + let path = dir.path().join("mcp-servers.json"); + + // Save a configuration + let mut config = McpServersFile::default(); + config.upsert( + McpServerConfig::new("notion", "https://mcp.notion.com").with_oauth( + OAuthConfig::new("client-123") + .with_scopes(vec!["read".to_string(), "write".to_string()]), + ), + ); + + save_mcp_servers_to(&config, &path).await.unwrap(); + + // Load it back + let loaded = load_mcp_servers_from(&path).await.unwrap(); + assert_eq!(loaded.servers.len(), 1); + + let server = loaded.get("notion").unwrap(); + assert_eq!(server.url, "https://mcp.notion.com"); + assert!(server.oauth.is_some()); + assert_eq!(server.oauth.as_ref().unwrap().client_id, "client-123"); + } + + #[tokio::test] + async fn test_load_nonexistent_returns_empty() { + let dir = tempdir().unwrap(); + let path = dir.path().join("nonexistent.json"); + + let config = load_mcp_servers_from(&path).await.unwrap(); + assert!(config.servers.is_empty()); + } + + #[test] + fn test_token_secret_names() { + let config = McpServerConfig::new("notion", "https://mcp.notion.com"); + assert_eq!(config.token_secret_name(), "mcp_notion_access_token"); + assert_eq!( + config.refresh_token_secret_name(), + "mcp_notion_refresh_token" + ); + } +} diff --git a/src/tools/mcp/mod.rs b/src/tools/mcp/mod.rs index a77b5fea..8a8b92ea 100644 --- a/src/tools/mcp/mod.rs +++ b/src/tools/mcp/mod.rs @@ -2,9 +2,38 @@ //! //! MCP allows the agent to connect to external tool servers that provide //! additional capabilities through a standardized protocol. +//! +//! Supports both local (unauthenticated) and hosted (OAuth-authenticated) servers. +//! +//! ## Usage +//! +//! ```ignore +//! // Simple client (no auth) +//! let client = McpClient::new("http://localhost:8080"); +//! +//! // Authenticated client (for hosted servers) +//! let client = McpClient::new_authenticated( +//! config, +//! session_manager, +//! secrets, +//! "user_id", +//! ); +//! +//! // List and register tools +//! let tools = client.create_tools().await?; +//! for tool in tools { +//! registry.register(tool); +//! } +//! ``` +pub mod auth; mod client; +pub mod config; mod protocol; +pub mod session; +pub use auth::{is_authenticated, refresh_access_token}; pub use client::McpClient; -pub use protocol::{McpRequest, McpResponse, McpTool}; +pub use config::{McpServerConfig, McpServersFile, OAuthConfig}; +pub use protocol::{InitializeResult, McpRequest, McpResponse, McpTool}; +pub use session::McpSessionManager; diff --git a/src/tools/mcp/protocol.rs b/src/tools/mcp/protocol.rs index 285be30e..4f5b5d98 100644 --- a/src/tools/mcp/protocol.rs +++ b/src/tools/mcp/protocol.rs @@ -2,20 +2,31 @@ use serde::{Deserialize, Serialize}; +/// MCP protocol version. +pub const PROTOCOL_VERSION: &str = "2024-11-05"; + /// An MCP tool definition. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct McpTool { /// Tool name. pub name: String, /// Tool description. + #[serde(default)] pub description: String, /// JSON Schema for input parameters. + /// Defaults to empty object schema if not provided. + #[serde(default = "default_input_schema")] pub input_schema: serde_json::Value, /// Optional annotations from the MCP server. #[serde(default)] pub annotations: Option, } +/// Default input schema (empty object). +fn default_input_schema() -> serde_json::Value { + serde_json::json!({"type": "object", "properties": {}}) +} + /// Annotations for an MCP tool that provide hints about its behavior. #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct McpToolAnnotations { @@ -84,6 +95,35 @@ impl McpRequest { } } + /// Create an initialize request. + pub fn initialize(id: u64) -> Self { + Self::new( + id, + "initialize", + Some(serde_json::json!({ + "protocolVersion": PROTOCOL_VERSION, + "capabilities": { + "roots": { "listChanged": false }, + "sampling": {} + }, + "clientInfo": { + "name": "ironclaw", + "version": env!("CARGO_PKG_VERSION") + } + })), + ) + } + + /// Create an initialized notification (sent after initialize). + pub fn initialized_notification() -> Self { + Self { + jsonrpc: "2.0".to_string(), + id: 0, // Notifications don't have IDs, but we need one for the struct + method: "notifications/initialized".to_string(), + params: None, + } + } + /// Create a tools/list request. pub fn list_tools(id: u64) -> Self { Self::new(id, "tools/list", None) @@ -129,6 +169,83 @@ pub struct McpError { pub data: Option, } +/// Result of the initialize handshake. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct InitializeResult { + /// Protocol version supported by the server. + #[serde(rename = "protocolVersion")] + pub protocol_version: Option, + + /// Server capabilities. + #[serde(default)] + pub capabilities: ServerCapabilities, + + /// Server information. + #[serde(rename = "serverInfo")] + pub server_info: Option, + + /// Instructions for using this server. + pub instructions: Option, +} + +/// Server capabilities advertised during initialization. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct ServerCapabilities { + /// Tool capabilities. + #[serde(default)] + pub tools: Option, + + /// Resource capabilities. + #[serde(default)] + pub resources: Option, + + /// Prompt capabilities. + #[serde(default)] + pub prompts: Option, + + /// Logging capabilities. + #[serde(default)] + pub logging: Option, +} + +/// Tool-related capabilities. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct ToolsCapability { + /// Whether the tool list can change. + #[serde(rename = "listChanged", default)] + pub list_changed: bool, +} + +/// Resource-related capabilities. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct ResourcesCapability { + /// Whether subscriptions are supported. + #[serde(default)] + pub subscribe: bool, + + /// Whether the resource list can change. + #[serde(rename = "listChanged", default)] + pub list_changed: bool, +} + +/// Prompt-related capabilities. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct PromptsCapability { + /// Whether the prompt list can change. + #[serde(rename = "listChanged", default)] + pub list_changed: bool, +} + +/// Server information. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ServerInfo { + /// Server name. + pub name: String, + + /// Server version. + pub version: Option, +} + /// Result of listing tools. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ListToolsResult { diff --git a/src/tools/mcp/session.rs b/src/tools/mcp/session.rs new file mode 100644 index 00000000..6046b3c9 --- /dev/null +++ b/src/tools/mcp/session.rs @@ -0,0 +1,286 @@ +//! MCP session management. +//! +//! Manages Mcp-Session-Id headers for stateful connections to MCP servers. +//! Each server can have an active session that persists across requests. + +use std::collections::HashMap; +use std::time::Instant; + +use tokio::sync::RwLock; + +/// Session state for a single MCP server connection. +#[derive(Debug, Clone)] +pub struct McpSession { + /// Session ID returned by the server (via Mcp-Session-Id header). + pub session_id: Option, + + /// Last activity timestamp for this session. + pub last_activity: Instant, + + /// Server URL this session is connected to. + pub server_url: String, + + /// Whether initialization has completed. + pub initialized: bool, +} + +impl McpSession { + /// Create a new session for a server. + pub fn new(server_url: impl Into) -> Self { + Self { + session_id: None, + last_activity: Instant::now(), + server_url: server_url.into(), + initialized: false, + } + } + + /// Update the session ID (from server response). + pub fn update_session_id(&mut self, session_id: Option) { + if session_id.is_some() { + self.session_id = session_id; + } + self.last_activity = Instant::now(); + } + + /// Mark the session as initialized. + pub fn mark_initialized(&mut self) { + self.initialized = true; + self.last_activity = Instant::now(); + } + + /// Check if the session has been idle for too long. + pub fn is_stale(&self, max_idle_secs: u64) -> bool { + self.last_activity.elapsed().as_secs() > max_idle_secs + } + + /// Touch the session to update last activity. + pub fn touch(&mut self) { + self.last_activity = Instant::now(); + } +} + +/// Manages MCP sessions for multiple servers. +pub struct McpSessionManager { + /// Active sessions by server name. + sessions: RwLock>, + + /// Maximum idle time before a session is considered stale (in seconds). + max_idle_secs: u64, +} + +impl McpSessionManager { + /// Create a new session manager with default idle timeout (30 minutes). + pub fn new() -> Self { + Self { + sessions: RwLock::new(HashMap::new()), + max_idle_secs: 1800, // 30 minutes + } + } + + /// Create a new session manager with custom idle timeout. + pub fn with_idle_timeout(max_idle_secs: u64) -> Self { + Self { + sessions: RwLock::new(HashMap::new()), + max_idle_secs, + } + } + + /// Get or create a session for a server. + pub async fn get_or_create(&self, server_name: &str, server_url: &str) -> McpSession { + let mut sessions = self.sessions.write().await; + + if let Some(session) = sessions.get(server_name) { + // Check if session is stale + if session.is_stale(self.max_idle_secs) { + // Create a fresh session + let new_session = McpSession::new(server_url); + sessions.insert(server_name.to_string(), new_session.clone()); + return new_session; + } + return session.clone(); + } + + // Create new session + let session = McpSession::new(server_url); + sessions.insert(server_name.to_string(), session.clone()); + session + } + + /// Get the current session ID for a server (if any). + pub async fn get_session_id(&self, server_name: &str) -> Option { + let sessions = self.sessions.read().await; + sessions.get(server_name).and_then(|s| s.session_id.clone()) + } + + /// Update the session ID from a server response. + pub async fn update_session_id(&self, server_name: &str, session_id: Option) { + let mut sessions = self.sessions.write().await; + if let Some(session) = sessions.get_mut(server_name) { + session.update_session_id(session_id); + } + } + + /// Mark a session as initialized. + pub async fn mark_initialized(&self, server_name: &str) { + let mut sessions = self.sessions.write().await; + if let Some(session) = sessions.get_mut(server_name) { + session.mark_initialized(); + } + } + + /// Check if a session is initialized. + pub async fn is_initialized(&self, server_name: &str) -> bool { + let sessions = self.sessions.read().await; + sessions + .get(server_name) + .map(|s| s.initialized) + .unwrap_or(false) + } + + /// Touch a session to update its activity timestamp. + pub async fn touch(&self, server_name: &str) { + let mut sessions = self.sessions.write().await; + if let Some(session) = sessions.get_mut(server_name) { + session.touch(); + } + } + + /// Terminate a session (e.g., on error or explicit disconnect). + pub async fn terminate(&self, server_name: &str) { + let mut sessions = self.sessions.write().await; + sessions.remove(server_name); + } + + /// Get all active server names. + pub async fn active_servers(&self) -> Vec { + let sessions = self.sessions.read().await; + sessions.keys().cloned().collect() + } + + /// Clean up stale sessions. + pub async fn cleanup_stale(&self) -> usize { + let mut sessions = self.sessions.write().await; + let before_len = sessions.len(); + sessions.retain(|_, session| !session.is_stale(self.max_idle_secs)); + before_len - sessions.len() + } +} + +impl Default for McpSessionManager { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_session_creation() { + let session = McpSession::new("https://mcp.example.com"); + assert!(session.session_id.is_none()); + assert!(!session.initialized); + assert_eq!(session.server_url, "https://mcp.example.com"); + } + + #[test] + fn test_session_update() { + let mut session = McpSession::new("https://mcp.example.com"); + + session.update_session_id(Some("session-123".to_string())); + assert_eq!(session.session_id, Some("session-123".to_string())); + + session.mark_initialized(); + assert!(session.initialized); + } + + #[test] + fn test_session_staleness() { + let mut session = McpSession::new("https://mcp.example.com"); + + // Fresh session should not be stale with reasonable timeout + assert!(!session.is_stale(1800)); + + // Manually set last_activity to the past to simulate staleness + session.last_activity = std::time::Instant::now() - std::time::Duration::from_secs(10); + assert!(session.is_stale(5)); + assert!(!session.is_stale(15)); + } + + #[tokio::test] + async fn test_session_manager_get_or_create() { + let manager = McpSessionManager::new(); + + // First call creates a new session + let session1 = manager + .get_or_create("notion", "https://mcp.notion.com") + .await; + assert!(session1.session_id.is_none()); + + // Update the session ID + manager + .update_session_id("notion", Some("session-abc".to_string())) + .await; + + // Second call returns existing session with the ID + let session2 = manager + .get_or_create("notion", "https://mcp.notion.com") + .await; + assert_eq!(session2.session_id, Some("session-abc".to_string())); + } + + #[tokio::test] + async fn test_session_manager_terminate() { + let manager = McpSessionManager::new(); + + manager + .get_or_create("notion", "https://mcp.notion.com") + .await; + manager + .update_session_id("notion", Some("session-123".to_string())) + .await; + + // Terminate the session + manager.terminate("notion").await; + + // Should create a fresh session now + let session = manager + .get_or_create("notion", "https://mcp.notion.com") + .await; + assert!(session.session_id.is_none()); + } + + #[tokio::test] + async fn test_session_manager_initialization() { + let manager = McpSessionManager::new(); + + manager + .get_or_create("notion", "https://mcp.notion.com") + .await; + + assert!(!manager.is_initialized("notion").await); + + manager.mark_initialized("notion").await; + + assert!(manager.is_initialized("notion").await); + } + + #[tokio::test] + async fn test_active_servers() { + let manager = McpSessionManager::new(); + + manager + .get_or_create("notion", "https://mcp.notion.com") + .await; + manager + .get_or_create("github", "https://mcp.github.com") + .await; + + let servers = manager.active_servers().await; + assert_eq!(servers.len(), 2); + assert!(servers.contains(&"notion".to_string())); + assert!(servers.contains(&"github".to_string())); + } +} diff --git a/src/tools/wasm/capabilities_schema.rs b/src/tools/wasm/capabilities_schema.rs index ef01da2f..e2f404d9 100644 --- a/src/tools/wasm/capabilities_schema.rs +++ b/src/tools/wasm/capabilities_schema.rs @@ -56,6 +56,11 @@ pub struct CapabilitiesFile { /// Workspace file read access. #[serde(default)] pub workspace: Option, + + /// Authentication setup instructions. + /// Used by `ironclaw config` to guide users through auth setup. + #[serde(default)] + pub auth: Option, } impl CapabilitiesFile { @@ -315,6 +320,170 @@ pub struct WorkspaceCapabilitySchema { pub allowed_prefixes: Vec, } +/// Authentication setup schema. +/// +/// Tools declare their auth requirements here. The agent uses this to provide +/// generic auth flows without needing service-specific code in the main codebase. +/// +/// Supports two auth methods: +/// 1. **OAuth** - Browser-based login (preferred for user-facing services) +/// 2. **Manual** - Copy/paste token from provider's dashboard +/// +/// # Example (OAuth) +/// +/// ```json +/// { +/// "auth": { +/// "secret_name": "notion_api_token", +/// "display_name": "Notion", +/// "oauth": { +/// "authorization_url": "https://api.notion.com/v1/oauth/authorize", +/// "token_url": "https://api.notion.com/v1/oauth/token", +/// "client_id": "your-client-id", +/// "scopes": [] +/// }, +/// "env_var": "NOTION_TOKEN" +/// } +/// } +/// ``` +/// +/// # Example (Manual) +/// +/// ```json +/// { +/// "auth": { +/// "secret_name": "openai_api_key", +/// "display_name": "OpenAI", +/// "instructions": "Get your API key from platform.openai.com/api-keys", +/// "setup_url": "https://platform.openai.com/api-keys", +/// "token_hint": "Starts with 'sk-'", +/// "env_var": "OPENAI_API_KEY" +/// } +/// } +/// ``` +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct AuthCapabilitySchema { + /// Name of the secret to store (e.g., "notion_api_token"). + /// Must match the secret_name in credentials if HTTP capability is used. + pub secret_name: String, + + /// Human-readable name for the service (e.g., "Notion", "Slack"). + #[serde(default)] + pub display_name: Option, + + /// OAuth configuration for browser-based login. + /// If present, OAuth flow is used instead of manual token entry. + #[serde(default)] + pub oauth: Option, + + /// Instructions shown to the user for obtaining credentials (manual flow). + /// Can include markdown formatting. + #[serde(default)] + pub instructions: Option, + + /// URL to open for setting up credentials (manual flow). + #[serde(default)] + pub setup_url: Option, + + /// Hint about expected token format (e.g., "Starts with 'sk-'"). + /// Used for validation feedback. + #[serde(default)] + pub token_hint: Option, + + /// Environment variable to check before prompting. + /// If this env var is set, its value is used automatically. + #[serde(default)] + pub env_var: Option, + + /// Provider hint for organizing secrets (e.g., "notion", "openai"). + #[serde(default)] + pub provider: Option, + + /// Validation endpoint to check if the token works. + /// Tool can specify an endpoint to call for validation. + #[serde(default)] + pub validation_endpoint: Option, +} + +/// OAuth 2.0 configuration for browser-based login. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct OAuthConfigSchema { + /// OAuth authorization URL (e.g., "https://api.notion.com/v1/oauth/authorize"). + pub authorization_url: String, + + /// OAuth token exchange URL (e.g., "https://api.notion.com/v1/oauth/token"). + pub token_url: String, + + /// OAuth client ID. + /// Can be set here or via environment variable (see client_id_env). + #[serde(default)] + pub client_id: Option, + + /// Environment variable containing the client ID. + /// Checked if client_id is not set directly. + #[serde(default)] + pub client_id_env: Option, + + /// OAuth client secret (optional, some providers don't require it with PKCE). + /// Can be set here or via environment variable (see client_secret_env). + #[serde(default)] + pub client_secret: Option, + + /// Environment variable containing the client secret. + /// Checked if client_secret is not set directly. + #[serde(default)] + pub client_secret_env: Option, + + /// OAuth scopes to request. + #[serde(default)] + pub scopes: Vec, + + /// Use PKCE (Proof Key for Code Exchange). Defaults to true. + /// Required for public clients (CLI tools). + #[serde(default = "default_true")] + pub use_pkce: bool, + + /// Additional parameters to include in the authorization URL. + #[serde(default)] + pub extra_params: std::collections::HashMap, + + /// Field name in token response containing the access token. + /// Defaults to "access_token". + #[serde(default = "default_access_token_field")] + pub access_token_field: String, +} + +fn default_true() -> bool { + true +} + +fn default_access_token_field() -> String { + "access_token".to_string() +} + +/// Schema for token validation endpoint. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct ValidationEndpointSchema { + /// URL to call for validation (e.g., "https://api.notion.com/v1/users/me"). + pub url: String, + + /// HTTP method (defaults to GET). + #[serde(default = "default_method")] + pub method: String, + + /// Expected HTTP status code for success (defaults to 200). + #[serde(default = "default_success_status")] + pub success_status: u16, +} + +fn default_method() -> String { + "GET".to_string() +} + +fn default_success_status() -> u16 { + 200 +} + #[cfg(test)] mod tests { use crate::tools::wasm::capabilities_schema::{CapabilitiesFile, CredentialLocationSchema}; @@ -503,4 +672,51 @@ mod tests { let secrets = caps.secrets.unwrap(); assert!(secrets.is_allowed("slack_bot_token")); } + + #[test] + fn test_parse_auth_capability() { + let json = r#"{ + "auth": { + "secret_name": "notion_api_token", + "display_name": "Notion", + "instructions": "Create an integration at notion.so/my-integrations", + "setup_url": "https://www.notion.so/my-integrations", + "token_hint": "Starts with 'secret_' or 'ntn_'", + "env_var": "NOTION_TOKEN", + "provider": "notion", + "validation_endpoint": { + "url": "https://api.notion.com/v1/users/me", + "method": "GET", + "success_status": 200 + } + } + }"#; + + let caps = CapabilitiesFile::from_json(json).unwrap(); + let auth = caps.auth.unwrap(); + assert_eq!(auth.secret_name, "notion_api_token"); + assert_eq!(auth.display_name, Some("Notion".to_string())); + assert_eq!(auth.env_var, Some("NOTION_TOKEN".to_string())); + assert_eq!(auth.provider, Some("notion".to_string())); + + let validation = auth.validation_endpoint.unwrap(); + assert_eq!(validation.url, "https://api.notion.com/v1/users/me"); + assert_eq!(validation.method, "GET"); + assert_eq!(validation.success_status, 200); + } + + #[test] + fn test_parse_auth_minimal() { + let json = r#"{ + "auth": { + "secret_name": "my_api_key" + } + }"#; + + let caps = CapabilitiesFile::from_json(json).unwrap(); + let auth = caps.auth.unwrap(); + assert_eq!(auth.secret_name, "my_api_key"); + assert!(auth.display_name.is_none()); + assert!(auth.setup_url.is_none()); + } } diff --git a/src/tools/wasm/mod.rs b/src/tools/wasm/mod.rs index 378aa6ba..d898dbbe 100644 --- a/src/tools/wasm/mod.rs +++ b/src/tools/wasm/mod.rs @@ -118,4 +118,7 @@ pub use storage::{ pub use loader::{DiscoveredTool, LoadResults, WasmLoadError, WasmToolLoader, discover_tools}; // Capabilities schema (for parsing *.capabilities.json files) -pub use capabilities_schema::{CapabilitiesFile, RateLimitSchema}; +pub use capabilities_schema::{ + AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema, RateLimitSchema, + ValidationEndpointSchema, +};