gitops/optimclaw
1
0
Fork 0
mirror of https://github.com/outbackdingo/optimclaw.git synced 2026-08-25 14:53:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add hosted MCP server support with OAuth 2.1 and token refresh

Browse Source 
Enables connecting to official MCP servers (like Notion) instead of
building custom WASM tools. Uses OAuth 2.1 with PKCE and supports
Dynamic Client Registration for zero-config authentication.

Key features:
- OAuth 2.1 flow with PKCE for secure browser-based auth
- Dynamic Client Registration (DCR) for servers without pre-configured clients
- Automatic token refresh on 401 responses
- Session management with Mcp-Session-Id headers
- SSE streaming response handling

New CLI commands:
- `mcp add <name> <url>` - Add an MCP server
- `mcp remove <name>` - Remove an MCP server
- `mcp list` - List configured servers
- `mcp auth <name>` - Authenticate with a server
- `mcp test <name>` - Test connection

Also removes the Notion WASM tool example since it's superseded by the
Notion MCP server which provides 13 official tools.

Co-Authored-By: Claude Opus 4.5 <[email protected]>
This commit is contained in:
Illia Polosukhin
2026-02-05 17:43:26 -08:00
co-authored by Claude Opus 4.5
parent 5992e27507
commit 974bc8d407
23 changed files with 3674 additions and 1838 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CLAUDE.md
+94
View File
@@ -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/<name>/`
2. Implement the WIT interface (`wit/tool.wit`)
3. Create `<name>.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 `<tool>.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 <tool>`:
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`
Cargo.lock
Generated
+17
View File
@@ -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"
Cargo.toml
+2 -1
View File
@@ -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]
examples/wasm-tools/notion/.gitignore
-1
View File
@@ -1 +0,0 @@
/target
examples/wasm-tools/notion/Cargo.lock
Generated
-401
View File
@@ -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"
examples/wasm-tools/notion/Cargo.toml
-21
View File
@@ -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
examples/wasm-tools/notion/README.md
-186
View File
@@ -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/)
examples/wasm-tools/notion/notion_tool.capabilities.json
-26
View File
@@ -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"]
}
}
examples/wasm-tools/notion/src/api.rs
-445
View File
@@ -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<serde_json::Value, String> {
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<serde_json::Value, String> {
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<serde_json::Value, String> {
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<serde_json::Value, String> {
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<serde_json::Value, String> {
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<serde_json::Value, String> {
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<serde_json::Value>>,
icon: Option<&serde_json::Value>,
cover: Option<&serde_json::Value>,
) -> Result<serde_json::Value, String> {
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<serde_json::Value, String> {
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<serde_json::Value, String> {
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<serde_json::Value, String> {
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<serde_json::Value, String> {
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<serde_json::Value, String> {
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<serde_json::Value, String> {
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<serde_json::Value, String> {
let block_id = normalize_uuid(block_id);
notion_patch(&format!("/blocks/{}", block_id), content)
}
pub fn delete_block(block_id: &str) -> Result<serde_json::Value, String> {
let block_id = normalize_uuid(block_id);
notion_delete(&format!("/blocks/{}", block_id))
}
// ==================== Databases ====================
pub fn get_database(database_id: &str) -> Result<serde_json::Value, String> {
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<serde_json::Value>>,
page_size: u32,
start_cursor: Option<&str>,
) -> Result<serde_json::Value, String> {
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<serde_json::Value, String> {
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<serde_json::Value>>,
properties: Option<&serde_json::Value>,
) -> Result<serde_json::Value, String> {
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<serde_json::Value, String> {
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<serde_json::Value, String> {
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<serde_json::Value, String> {
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<serde_json::Value, String> {
let user_id = normalize_uuid(user_id);
notion_get(&format!("/users/{}", user_id))
}
pub fn get_me() -> Result<serde_json::Value, String> {
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()
}
examples/wasm-tools/notion/src/lib.rs
-451
View File
@@ -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<String, String> {
// 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"]
}
]
}"#;
examples/wasm-tools/notion/src/types.rs
-267
View File
@@ -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<SearchFilter>,
/// Max results (default: 10, max: 100).
#[serde(default = "default_page_size")]
page_size: u32,
/// Pagination cursor from previous response.
#[serde(default)]
start_cursor: Option<String>,
},
// ==================== 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<Vec<serde_json::Value>>,
/// Optional icon: { "emoji": "..." } or { "external": { "url": "..." } }.
#[serde(default)]
icon: Option<serde_json::Value>,
/// Optional cover: { "external": { "url": "..." } }.
#[serde(default)]
cover: Option<serde_json::Value>,
},
/// Update page properties.
UpdatePage {
/// Page ID.
page_id: String,
/// Properties to update.
properties: serde_json::Value,
/// Optional new icon.
#[serde(default)]
icon: Option<serde_json::Value>,
/// Optional new cover.
#[serde(default)]
cover: Option<serde_json::Value>,
},
/// 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<String>,
},
/// 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<serde_json::Value>,
/// Append after this block ID (for ordering).
#[serde(default)]
after: Option<String>,
},
/// 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<serde_json::Value>,
/// Sort configuration array.
#[serde(default)]
sorts: Option<Vec<serde_json::Value>>,
/// Max results (default: 10, max: 100).
#[serde(default = "default_page_size")]
page_size: u32,
/// Pagination cursor.
#[serde(default)]
start_cursor: Option<String>,
},
/// 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<serde_json::Value>,
/// Property schema: { "Name": { "title": {} }, "Status": { "select": { "options": [...] } } }.
properties: serde_json::Value,
/// Optional icon.
#[serde(default)]
icon: Option<serde_json::Value>,
/// Optional cover.
#[serde(default)]
cover: Option<serde_json::Value>,
/// 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<Vec<serde_json::Value>>,
/// Properties to add or update (optional).
#[serde(default)]
properties: Option<serde_json::Value>,
},
// ==================== 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<String>,
},
/// 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<serde_json::Value>,
},
// ==================== 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<String>,
},
/// 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_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub next_cursor: Option<String>,
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,
}
src/cli/mcp.rs
+484
View File
@@ -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<String>,
/// OAuth authorization URL (optional, can be discovered)
#[arg(long)]
auth_url: Option<String>,
/// OAuth token URL (optional, can be discovered)
#[arg(long)]
token_url: Option<String>,
/// Scopes to request (comma-separated)
#[arg(long)]
scopes: Option<String>,
/// Server description
#[arg(long)]
description: Option<String>,
},
/// 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<String>,
auth_url: Option<String>,
token_url: Option<String>,
scopes: Option<String>,
description: Option<String>,
) -> 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<String> = 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 <name> <url> [--client-id <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<Arc<dyn SecretsStore + Send + Sync>> {
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();
}
}
src/cli/mod.rs
+7
View File
@@ -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 {
src/cli/tool.rs
+586 -1
View File
@@ -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<PathBuf>,
},
/// 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<PathBuf>,
/// 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<PathBuf>, 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\
<!DOCTYPE html><html><body style=\"font-family: sans-serif; \
display: flex; justify-content: center; align-items: center; \
height: 100vh; margin: 0; background: #191919; color: white;\">\
<div style=\"text-align: center;\">\
<h1>✓ {} Connected!</h1>\
<p>You can close this window.</p>\
</div></body></html>",
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<String> {
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::*;
src/main.rs
+116 -18
View File
@@ -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<Arc<dyn SecretsStore + Send + Sync>> =
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<Arc<dyn SecretsStore + Send + Sync>> =
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()) {
src/tools/mcp/auth.rs
+904
View File
@@ -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<String>,
/// Scopes supported by this resource.
#[serde(default)]
pub scopes_supported: Vec<String>,
}
/// 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<String>,
/// Supported response types.
#[serde(default)]
pub response_types_supported: Vec<String>,
/// Supported grant types.
#[serde(default)]
pub grant_types_supported: Vec<String>,
/// Supported code challenge methods.
#[serde(default)]
pub code_challenge_methods_supported: Vec<String>,
/// Scopes supported by this server.
#[serde(default)]
pub scopes_supported: Vec<String>,
}
/// 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<String>,
/// Grant types the client will use.
pub grant_types: Vec<String>,
/// Response types the client will use.
pub response_types: Vec<String>,
/// 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<String>,
/// When the client secret expires (if applicable).
#[serde(default)]
pub client_secret_expires_at: Option<u64>,
/// Registration access token for managing the registration.
#[serde(default)]
pub registration_access_token: Option<String>,
/// Registration client URI for managing the registration.
#[serde(default)]
pub registration_client_uri: Option<String>,
}
/// 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<u64>,
/// Refresh token for obtaining new access tokens.
pub refresh_token: Option<String>,
/// Scopes granted.
pub scope: Option<String>,
}
/// Token response from the authorization server.
#[derive(Debug, Deserialize)]
struct TokenResponse {
access_token: String,
token_type: String,
expires_in: Option<u64>,
refresh_token: Option<String>,
scope: Option<String>,
}
/// 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<ProtectedResourceMetadata, AuthError> {
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<AuthorizationServerMetadata, AuthError> {
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<AuthorizationServerMetadata, AuthError> {
// 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<ClientRegistrationResponse, AuthError> {
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<dyn SecretsStore + Send + Sync>,
user_id: &str,
) -> Result<AccessToken, AuthError> {
// 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(&registration_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, String>,
) -> 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<String, AuthError> {
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\
<!DOCTYPE html><html><body style=\"font-family: sans-serif; \
display: flex; justify-content: center; align-items: center; \
height: 100vh; margin: 0; background: #191919; color: white;\">\
<div style=\"text-align: center;\">\
<h1>✓ {} Connected!</h1>\
<p>You can close this window.</p>\
</div></body></html>",
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<AccessToken, AuthError> {
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(&params)
.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<dyn SecretsStore + Send + Sync>,
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<dyn SecretsStore + Send + Sync>,
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<dyn SecretsStore + Send + Sync>,
user_id: &str,
) -> Result<String, AuthError> {
// 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<dyn SecretsStore + Send + Sync>,
user_id: &str,
) -> Result<Option<String>, 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<dyn SecretsStore + Send + Sync>,
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<dyn SecretsStore + Send + Sync>,
user_id: &str,
) -> Result<AccessToken, AuthError> {
// 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(&params)
.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"));
}
}
src/tools/mcp/client.rs
+365 -14
View File
@@ -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<Option<Vec<McpTool>>>,
/// Session manager (shared across clients).
session_manager: Option<Arc<McpSessionManager>>,
/// Secrets store for retrieving access tokens.
secrets: Option<Arc<dyn SecretsStore + Send + Sync>>,
/// User ID for secrets lookup.
user_id: String,
/// Server configuration (for token secret name lookup).
server_config: Option<McpServerConfig>,
}
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<String>) -> 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<String>, server_url: impl Into<String>) -> 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<McpSessionManager>,
secrets: Arc<dyn SecretsStore + Send + Sync>,
user_id: impl Into<String>,
) -> 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.
/// 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<Option<String>, 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
{
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<McpResponse, ToolError> {
let response = self
// 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)
.json(&request)
.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)))?;
if !response.status().is_success() {
// 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 returned status: {}",
response.status()
"MCP server '{}' requires authentication. Run: ironclaw mcp auth {}",
self.server_name, self.server_name
)));
}
response
.json()
.await
.map_err(|e| ToolError::ExternalService(format!("Failed to parse MCP response: {}", e)))
// 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<McpResponse, ToolError> {
// 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: {} - {}",
status, body
)));
}
// 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::<McpResponse>(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<InitializeResult, ToolError> {
// 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<CallToolResult, ToolError> {
// 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<dyn Tool>
})
.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<McpClient>,
}
#[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<ToolOutput, ToolError> {
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());
}
}
src/tools/mcp/config.rs
+442
View File
@@ -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<OAuthConfig>,
/// 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<String>,
}
fn default_true() -> bool {
true
}
impl McpServerConfig {
/// Create a new MCP server configuration.
pub fn new(name: impl Into<String>, url: impl Into<String>) -> 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<String>) -> 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<String>,
/// 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<String>,
/// Scopes to request.
#[serde(default)]
pub scopes: Vec<String>,
/// 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<String, String>,
}
impl OAuthConfig {
/// Create a new OAuth configuration with just a client ID.
pub fn new(client_id: impl Into<String>) -> 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<String>,
token_url: impl Into<String>,
) -> 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<String>) -> 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<McpServerConfig>,
/// 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<Item = &McpServerConfig> {
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<ConfigError> 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<McpServersFile, ConfigError> {
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<Path>) -> Result<McpServersFile, ConfigError> {
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<Path>,
) -> 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<McpServerConfig, ConfigError> {
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"
);
}
}
src/tools/mcp/mod.rs
+30 -1
View File
@@ -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;
src/tools/mcp/protocol.rs
+117
View File
@@ -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<McpToolAnnotations>,
}
/// 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<serde_json::Value>,
}
/// 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<String>,
/// Server capabilities.
#[serde(default)]
pub capabilities: ServerCapabilities,
/// Server information.
#[serde(rename = "serverInfo")]
pub server_info: Option<ServerInfo>,
/// Instructions for using this server.
pub instructions: Option<String>,
}
/// Server capabilities advertised during initialization.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ServerCapabilities {
/// Tool capabilities.
#[serde(default)]
pub tools: Option<ToolsCapability>,
/// Resource capabilities.
#[serde(default)]
pub resources: Option<ResourcesCapability>,
/// Prompt capabilities.
#[serde(default)]
pub prompts: Option<PromptsCapability>,
/// Logging capabilities.
#[serde(default)]
pub logging: Option<serde_json::Value>,
}
/// 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<String>,
}
/// Result of listing tools.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ListToolsResult {
src/tools/mcp/session.rs
+286
View File
@@ -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<String>,
/// 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<String>) -> 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<String>) {
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<HashMap<String, McpSession>>,
/// 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<String> {
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<String>) {
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<String> {
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()));
}
}
src/tools/wasm/capabilities_schema.rs
+216
View File
@@ -56,6 +56,11 @@ pub struct CapabilitiesFile {
/// Workspace file read access.
#[serde(default)]
pub workspace: Option<WorkspaceCapabilitySchema>,
/// Authentication setup instructions.
/// Used by `ironclaw config` to guide users through auth setup.
#[serde(default)]
pub auth: Option<AuthCapabilitySchema>,
}
impl CapabilitiesFile {
@@ -315,6 +320,170 @@ pub struct WorkspaceCapabilitySchema {
pub allowed_prefixes: Vec<String>,
}
/// 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<String>,
/// OAuth configuration for browser-based login.
/// If present, OAuth flow is used instead of manual token entry.
#[serde(default)]
pub oauth: Option<OAuthConfigSchema>,
/// Instructions shown to the user for obtaining credentials (manual flow).
/// Can include markdown formatting.
#[serde(default)]
pub instructions: Option<String>,
/// URL to open for setting up credentials (manual flow).
#[serde(default)]
pub setup_url: Option<String>,
/// Hint about expected token format (e.g., "Starts with 'sk-'").
/// Used for validation feedback.
#[serde(default)]
pub token_hint: Option<String>,
/// Environment variable to check before prompting.
/// If this env var is set, its value is used automatically.
#[serde(default)]
pub env_var: Option<String>,
/// Provider hint for organizing secrets (e.g., "notion", "openai").
#[serde(default)]
pub provider: Option<String>,
/// Validation endpoint to check if the token works.
/// Tool can specify an endpoint to call for validation.
#[serde(default)]
pub validation_endpoint: Option<ValidationEndpointSchema>,
}
/// 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<String>,
/// Environment variable containing the client ID.
/// Checked if client_id is not set directly.
#[serde(default)]
pub client_id_env: Option<String>,
/// 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<String>,
/// Environment variable containing the client secret.
/// Checked if client_secret is not set directly.
#[serde(default)]
pub client_secret_env: Option<String>,
/// OAuth scopes to request.
#[serde(default)]
pub scopes: Vec<String>,
/// 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<String, String>,
/// 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());
}
}
src/tools/wasm/mod.rs
+4 -1
View File
@@ -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,
};