diff --git a/Cargo.lock b/Cargo.lock index 36280b5..81c96ec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1880,9 +1880,13 @@ version = "0.1.1" dependencies = [ "anyhow", "axum", + "base64", "clap", + "form_urlencoded", "futures-util", + "getrandom 0.4.3", "jsonschema", + "libc", "qmacp", "reqwest", "rmcp", @@ -1893,6 +1897,7 @@ dependencies = [ "subtle", "tokio", "tokio-tungstenite", + "tower", "tracing", "tracing-subscriber", "uuid", diff --git a/Cargo.toml b/Cargo.toml index 36380e0..3dd82c3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,10 @@ path = "src/main.rs" [dependencies] anyhow = "1" axum = "0.8" +base64 = "0.23" +form_urlencoded = "1.2" +getrandom = "0.4" +libc = "0.2" clap = { version = "4", features = ["derive", "env"] } qmacp = { git = "https://github.com/querymt/qmacp", tag = "v0.1.2" } reqwest = { version = "0.13", default-features = false, features = ["json", "form", "rustls"] } @@ -32,3 +36,4 @@ uuid = { version = "1", features = ["v4"] } futures-util = "0.3" jsonschema = { version = "0.60", default-features = false } tokio-tungstenite = "0.29" +tower = { version = "0.5", features = ["util"] } diff --git a/README.md b/README.md index 0d466d3..1ff27ce 100644 --- a/README.md +++ b/README.md @@ -20,10 +20,20 @@ cargo install --locked --path . ## Connect +For ChatGPT, qmcp provides embedded single-owner OAuth with explicit consent, +PKCE, dynamic client registration, and rotating refresh tokens. Run it on the +Mac behind an HTTPS reverse proxy; no external login service is needed: -See [How to connect](docs/how-to-connect.md) for local stdio, Bearer over HTTPS, -private tunnels, and OAuth/ChatGPT setup. The guide includes client examples, -authentication requirements, and remote-access precautions. +```sh +qmcp --transport http --permission cancel \ + --oauth-local-resource https://mcp.example.com/mcp \ + --oauth-owner-secret-file /absolute/private/path/to/owner-secret +``` + +See the [connection guide](docs/how-to-connect.md) for credential generation, +proxy configuration, stdio, static Bearer, tunnels, or external OAuth. +OAuth state is volatile: restart invalidates tokens and registered client IDs. +Approval grants all backend tools, not multi-user isolation. ## Tools @@ -33,7 +43,7 @@ Start with `caps`, `profiles`, and `models` to discover available capabilities. | Tools | Purpose | | --- | --- | | `caps`, `profiles`, `models` | Discover protocol features, profiles, and models | -| `new`, `sessions`, `find` | Create, list, and search sessions | +| `workspace`, `new`, `sessions`, `find` | Select a workspace without prompting; create, list, and search sessions | | `status`, `inspect`, `runtime` | Read configuration, history, and live runtime | | `get_lifecycle`, `set_lifecycle` | Read or explicitly patch resolved/archived state | | `pin`, `unpin` | Set session pin state | @@ -43,10 +53,22 @@ Start with `caps`, `profiles`, and `models` to discover available capabilities. | `cancel`, `close` | Request turn cancellation or explicitly close a session | | `exec` | Run a batch sequentially on one ACP connection | -Use `prompt` with `new: true` to start a session, then its returned `sessionId` -as `session_id` to continue. Pass an explicit `cwd` for workspace operations; -paths refer to the QueryMT host. `profiles`, `models`, lifecycle, and pin tools -return provider data under `result`, separate from bridge status and event metadata. +Choose a workspace once without prompting, then reuse the returned ID: + +```text +workspace({"cwd":"/srv/projects/app","profile":"my-profile"}) +prompt({"session_id":"","text":"Explain this project"}) +workspace({"session_id":"","cwd":"../other"}) +``` + +A directory switch creates a **new session**: use its returned `sessionId`. +The old ID remains unchanged; nothing implicitly closes or cancels its work. +`new` and `prompt(new=true)` also bind cwd; continuing `prompt` can omit cwd. +`status`, `inspect`, and `follow` reject directory changes. Unknown IDs require +backend `session/list` metadata; missing or incomplete metadata fails closed. +Paths are lexical POSIX paths on the QueryMT host, not a sandbox. Initial relative +paths use qmcp's startup cwd; later ones use the session cwd. `~` is rejected; +supply the backend's absolute home path. See [workspace details](docs/how-to-connect.md#select-a-session-workspace). ## Important Behavior @@ -68,6 +90,10 @@ return provider data under `result`, separate from bridge status and event metad - `find` scans at most 20 pages. Check `truncated` and `scanLimitReached` before treating its results as exhaustive. +`inspect` still uses ACP `session/load`, which may restore runtime state; it is +not a guaranteed read-only history API. A backend replay failure is a tool error, +not empty history. No history endpoint or backend fix is supplied here. + ## Development ```sh diff --git a/docs/how-to-connect.md b/docs/how-to-connect.md index 5091a22..3165af6 100644 --- a/docs/how-to-connect.md +++ b/docs/how-to-connect.md @@ -5,9 +5,10 @@ | Setup | Best for | Authentication | | --- | --- | --- | | [Local stdio](#local-stdio) | Desktop and CLI MCP clients | Local process permissions | +| [Embedded OAuth on macOS](#embedded-oauth-on-macos-chatgpt) | Single-owner public ChatGPT connections | Built-in owner login + consent; no external provider | | [Bearer over HTTPS](#bearer-over-https) | Clients with custom authorization headers | Shared token | | [Private tunnel](#private-tunnel-chatgpt) | ChatGPT without a public endpoint | OpenAI tunnel access | -| [OAuth over HTTPS](#oauth-over-https-chatgpt) | Public authenticated ChatGPT connections | External OAuth provider | +| [External OAuth provider](#external-oauth-provider-optional) | Deployments with an existing identity provider | Optional external introspection | Bearer and OAuth are HTTP authentication modes. The tunnel privately connects ChatGPT to qmcp over stdio; it is not another MCP transport. @@ -24,6 +25,30 @@ qmcp connects to `ws://127.0.0.1:3030/acp/ws` by default. Change it with `--url` or `QMCP_ACP_URL`. Prefer `wss://` for remote ACP connections; plaintext non-loopback ACP requires `--allow-insecure`. qmcp does not launch QueryMT. +## Select a Session Workspace + +Call `workspace({"cwd":"/srv/projects/app","profile":""})` +without sending a prompt. Reuse its returned `sessionId` in later calls; omit cwd +until you explicitly change it. There is no global active workspace. +`workspace({"session_id":"A","cwd":"../other"})` creates session B, leaving A +unchanged. Use B's returned ID. Continuing `prompt` with a different cwd behaves +likewise; `status`, `inspect`, and `follow` reject a different cwd. + +Paths refer to the QueryMT host and normalize `.`/`..` lexically, without local +filesystem or symlink checks. Initial relative paths use qmcp's startup cwd; +subsequent paths use the bound session cwd. Supply absolute remote paths; +`~`, Windows paths, empty paths, and control characters are rejected. +Known profiles are inherited on switches; unknown ones use the backend default +and report `profileKnown:false`. Same-directory profile changes are rejected. + +Bindings are immutable per ID, process-local, shared across HTTP calls, retained +on close, and capped at 4096 without eviction. Unknown IDs (including after +restart) require unfiltered paginated `session/list` metadata. Missing cwd, +repeated cursors, unavailable metadata, or the 20-page scan limit fail closed; +caller cwd cannot override them. Check `recovery` before retrying failed creation: +an acknowledged ID is retained, but an unacknowledged creation may still exist +remotely. Timeouts never implicitly cancel remote work. + ## Local stdio Configure your MCP client to launch qmcp: @@ -92,9 +117,73 @@ to keep qmcp private. Your account/workspace must allow tunnels and custom MCP s 4. Select the plugin with `@` in a chat and ask: **"Use qmcp to list QueryMT profiles."** Keep QueryMT and `tunnel-client` running. No public qmcp endpoint is needed. -## OAuth over HTTPS (ChatGPT) +## Embedded OAuth on macOS (ChatGPT) + +qmcp runs the single-owner authorization server on the Mac. TLS terminates at +your trusted HTTPS proxy; the canonical public resource must end in `/mcp` with +no userinfo, query, or fragment. Its origin is the issuer. + +Generate a credential only when preparing deployment (these commands do not +print it and refuse to overwrite an existing file): + +```sh +umask 077 +mkdir -p "$HOME/.config/qmcp" +chmod 700 "$HOME/.config/qmcp" +(set -C; openssl rand -hex 32 > "$HOME/.config/qmcp/owner-secret") +chmod 600 "$HOME/.config/qmcp/owner-secret" +``` + +Use only 32..64 cryptographically random bytes encoded as lowercase hex, never +a password. Keep the value in a trusted password manager for browser login, not +URLs, arguments, environment, logs, or source control. The absolute secret path +must be an owned Unix regular file, mode `0600`, without symlinks or hard links; +protect its parent directory too. qmcp neither generates nor prints the secret. + +Leave `QMCP_BEARER_TOKEN` and `QMCP_OAUTH_CLIENT_SECRET` unset, then run: + +```sh +qmcp --transport http --permission cancel \ + --oauth-local-resource https://mcp.example.com/mcp \ + --oauth-owner-secret-file "$HOME/.config/qmcp/owner-secret" +``` -Use an external [OAuth provider](https://developers.openai.com/plugins/build/auth) +The default listener is `127.0.0.1:3031`. For a separate proxy, use the +[Nginx/VPN template](#separate-nginx-server-to-a-mac-over-vpn). Forward `/mcp`, +`/.well-known/oauth-authorization-server`, both protected-resource discovery +paths, and `/oauth/register`, `/oauth/authorize`, `/oauth/token`. Preserve Host, +Authorization, Origin, and queries. The public resource Host is automatically +allowed; local/external OAuth and static Bearer modes are mutually exclusive. + +In ChatGPT, add the public `/mcp` URL, choose **OAuth**, and use DCR. Do not put +the owner credential in client ID/secret fields. During sign-in verify the domain, +unverified client name, and exact callback. Enter the credential on qmcp's owner +page and explicitly approve or deny, then click **Continue to ChatGPT**. +Registration alone grants nothing. Custom MCP servers must be enabled in your +workspace; live sign-in still needs deployment validation. + +Security/lifetime constraints: + +- Public clients use authorization-code + PKCE `S256`; DCR may additionally request + `refresh_token`. Exact client, redirect, resource, and scope are bound to grants. + Metadata URLs are never fetched. Login requires an issuer-Origin form POST, + a browser cookie, CSRF token, and explicit owner consent. Other routes reject Origin. +- Approval grants all tools on this backend (`qmcp:access`), not per-user isolation. + `--permission cancel` controls agent approvals, not MCP authentication. +- Access lasts up to 30 minutes; rotating one-use refresh tokens have an absolute + 24-hour family deadline. Replay revokes the entire family. Serialize refreshes: + concurrent refreshes or retrying a lost successful response can revoke the winner. +- Credentials are stored as digests; state is bounded and volatile. Restart invalidates client IDs and tokens; + ChatGPT may require removing/re-adding the server. Approved clients survive until + restart; unapproved registrations expire after one hour. No persistence is supplied. +- Capacity/rate errors do not evict replay evidence or consume valid credentials. + A family allows at most 127 rotations; exhaustion requires owner authorization + again. Keep proxy connection/body/rate limits and private logs: application bounds + do not prevent all resource exhaustion. Never log OAuth queries, bodies, or tokens. + +## External OAuth Provider (Optional) + +For an existing identity provider, use an external [OAuth provider](https://developers.openai.com/plugins/build/auth) with authorization-code + PKCE (`S256`) and ChatGPT client registration (pre-registered, CIMD, or DCR). qmcp provides discovery and token introspection; the provider handles login, consent, token issuance, and revocation. @@ -136,13 +225,59 @@ per-user workspaces. Tokens are checked on every request without caching; provider outages deny access. Live ChatGPT sign-in depends on your provider configuration and has not been verified by the project's mock-provider tests. +## Separate Nginx server to a Mac over VPN + +Topology: ChatGPT -> public HTTPS Nginx -> encrypted VPN -> qmcp on Mac:3031 +-> local QueryMT. Replace the example domain, VPN address, and certificate paths. +Keep QueryMT's ACP port private. On the Mac, use the protected owner file above: + +```sh +qmcp --transport http --bind 10.0.0.42:3031 --permission cancel \ + --oauth-local-resource https://mcp.example.com/mcp \ + --oauth-owner-secret-file /absolute/private/path/to/owner-secret +``` + +On the separate proxy host (not deployment instructions to execute during review): + +```nginx +server { + listen 443 ssl; + server_name mcp.example.com; + ssl_certificate /path/to/fullchain.pem; + ssl_certificate_key /path/to/privkey.pem; + location / { + proxy_pass http://10.0.0.42:3031; + proxy_http_version 1.1; + proxy_set_header Host $http_host; + proxy_set_header Authorization $http_authorization; + proxy_set_header Origin $http_origin; + proxy_set_header Connection ""; + access_log off; # OAuth queries contain codes/state. + proxy_buffering off; + proxy_request_buffering off; + proxy_cache off; + proxy_read_timeout 3600s; + proxy_send_timeout 3600s; + } +} +``` + +The upstream has no URI suffix, so discovery and authorization paths reach qmcp +unchanged. Never strip Origin to bypass validation. Firewall Mac port 3031 to the +proxy's VPN source only; do not expose it or ACP publicly. Plain HTTP assumes the +encrypted VPN hop, not an untrusted LAN. Keep the Mac awake and VPN reachable. +Validate adapted Nginx configuration before deployment. Source tests do not prove +real ChatGPT sign-in and must not restart installed services. + ## Remote Access Checklist - Use HTTPS: qmcp itself does not terminate TLS. Never send tokens over public plaintext HTTP. - Keep qmcp's backend port private to the proxy with a firewall/private network. -- Forward Authorization and an allowed Host. Browser Origin headers are rejected. -- Non-loopback binds require Bearer or OAuth auth. Wildcard binds also require - explicit `--allowed-host` values (repeatable, optionally `host:port`). +- Forward Authorization, Host and Origin unchanged. MCP rejects browser Origin; + only embedded OAuth consent POSTs accept the exact public issuer Origin. +- Non-loopback binds require Bearer or OAuth auth. Wildcard binds require an + explicit Host allowlist (`--allowed-host`, repeatable, optionally `host:port`); + embedded OAuth supplies its canonical resource Host automatically. - Use `--bind` only when the proxy needs another interface. Keep the default loopback listener when the proxy runs on the same host. - `--permission` controls QueryMT tool approval, not MCP client authentication. diff --git a/src/commands.rs b/src/commands.rs index 130e243..e221574 100644 --- a/src/commands.rs +++ b/src/commands.rs @@ -10,6 +10,7 @@ use serde_json::{Map, Value, json}; use qmacp::{client::AcpClient, error::CliError, session}; use crate::output::{self, BatchState, OperationState}; +use crate::workspace::{Binding, Workspaces}; const MIN_SECONDS: u64 = 1; const MAX_SECONDS: u64 = 86_400; @@ -86,7 +87,13 @@ pub enum Command { query: Option, provider: Option, }, - /// Create a session without sending a prompt. + /// Select a sticky per-session workspace without sending a prompt. With no session_id, an explicit cwd is required and creates a NEW QueryMT session. With an existing ID, omitted/same cwd selects that ID; different cwd creates a NEW ID and returns previousSessionId. Use the returned sessionId in later calls; the old session is untouched. Known profiles are inherited on switches, otherwise the backend default is used (unknown). Profile changes on the same workspace are not supported. Paths are POSIX paths on the QueryMT host; no ~ expansion. + Workspace { + session_id: Option, + cwd: Option, + profile: Option, + }, + /// Create a session without sending a prompt. cwd is remembered immediately; later session calls can omit it. Relative cwd uses qmcp startup cwd; paths refer to the QueryMT host, with no ~ expansion. New { cwd: Option, profile: Option, @@ -130,7 +137,7 @@ pub enum Command { Pin { session_id: String }, /// Unpin a session. Unpin { session_id: String }, - /// Read compact session history, or the full raw snapshot when full is true. + /// Read compact session history, or the full raw snapshot when full is true. Omit cwd to use the session's remembered or authoritative metadata workspace. Different cwd is rejected; use workspace to create a NEW session. Never creates or rebinds a session. Inspect { session_id: String, cwd: Option, @@ -140,7 +147,7 @@ pub enum Command { #[serde(default = "default_true")] tools: bool, }, - /// Current mode/model/profile/effort for a session. + /// Read current mode/model/profile/effort for a session. Omit cwd to use its remembered or authoritative metadata workspace. Different cwd is rejected; use workspace to create a NEW session. Never creates or rebinds a session. Status { session_id: String, cwd: Option, @@ -151,7 +158,7 @@ pub enum Command { SetModel { session_id: String, model: String }, /// Set reasoning effort (auto, low, medium, high, or max). SetEffort { session_id: String, effort: String }, - /// Send text to a session, or create one with new=true. Auto delivery prompts when idle, steers when possible, otherwise queues. + /// Send text to a session, or create one with new=true. cwd is sticky per session: omit it on later calls. Explicit different cwd creates a NEW QueryMT session, returns previousSessionId and the NEW sessionId to use later, and leaves the old session untouched. Relative cwd uses the selected session cwd (or qmcp startup cwd for new=true). No ~ expansion. Known profiles are inherited on switches; unknown uses the backend default. Same-workspace profile changes are rejected. Auto delivery prompts when idle, steers when possible, otherwise queues. Prompt { #[serde(default)] new: bool, @@ -174,7 +181,7 @@ pub enum Command { }, /// Current QueryMT runtime state for a session. Runtime { session_id: String }, - /// Collect compact session events until idle. Supply cwd for the session workspace. + /// Collect compact session events until idle. Omit cwd to use its remembered or authoritative metadata workspace. Different cwd is rejected; use workspace to create a NEW session. Never creates or rebinds a session. Follow { session_id: String, cwd: Option, @@ -246,6 +253,11 @@ pub fn decode(name: &str, mut arguments: Map) -> Result let command: Command = serde_json::from_value(Value::Object(arguments)).context("invalid tool arguments")?; match &command { + Command::Workspace { + session_id: None, + cwd: None, + .. + } => bail!("workspace requires an explicit initial cwd when session_id is omitted"), Command::Prompt { new, session_id, @@ -385,7 +397,15 @@ fn constrain_input_schema(name: &str, input: &mut Map) { if matches!(name, "prompt" | "steer" | "queue") { properties["text"]["pattern"] = json!(NON_BLANK_TEXT_PATTERN); } - if name == "prompt" { + if name == "workspace" { + input.insert( + "anyOf".into(), + json!([ + {"required": ["session_id"], "properties": {"session_id": {"type": "string"}}}, + {"required": ["cwd"], "properties": {"cwd": {"type": "string"}}}, + ]), + ); + } else if name == "prompt" { input.insert("oneOf".into(), json!([ {"required": ["new"], "properties": {"new": {"const": true}, "session_id": {"type": "null"}}}, {"required": ["session_id"], "properties": {"new": {"const": false}, "session_id": {"type": "string"}}}, @@ -417,6 +437,7 @@ pub async fn run( command: Command, state: &mut OperationState, deadline: Option, + workspaces: &Workspaces, ) -> Value { let deadline = deadline.or_else(|| { command @@ -427,7 +448,7 @@ pub async fn run( state.recovery.clear(); let operation = async { client.clear_assistant_text().await; - run_inner(client, command, state).await + run_inner(client, command, state, workspaces).await }; let outcome = match deadline { Some(deadline) => tokio::time::timeout_at(deadline, operation) @@ -454,6 +475,7 @@ async fn run_inner( client: &AcpClient, command: Command, state: &mut OperationState, + workspaces: &Workspaces, ) -> Result { let result = match command { Command::Caps {} => caps(client).await?, @@ -463,13 +485,31 @@ async fn run_inner( query, provider, } => models(client, refresh, query, provider).await?, + Command::Workspace { + session_id, + cwd, + profile, + } => workspaces + .select(client, state, session_id, cwd, profile) + .await? + .result(), Command::New { cwd, profile, mode, model, effort, - } => new_session(client, state, cwd, profile, mode, model, effort).await?, + } => { + new_session( + client, + state, + workspaces, + cwd, + profile, + (mode, model, effort), + ) + .await? + } Command::Sessions { cwd, cursor, @@ -481,7 +521,8 @@ async fn run_inner( } => { sessions( client, - cwd, + cwd.map(|cwd| workspaces.initial_cwd(Some(cwd))) + .transpose()?, cursor, all, limit, @@ -504,7 +545,8 @@ async fn run_inner( } => { find( client, - cwd, + cwd.map(|cwd| workspaces.initial_cwd(Some(cwd))) + .transpose()?, query, phase, limit, @@ -530,8 +572,14 @@ async fn run_inner( full, messages, tools, - } => inspect(client, session_id, cwd, full, messages, tools).await?, - Command::Status { session_id, cwd } => status(client, session_id, cwd).await?, + } => { + let cwd = workspaces.read_cwd(client, &session_id, cwd).await?; + inspect(client, session_id, cwd, full, messages, tools).await? + } + Command::Status { session_id, cwd } => { + let cwd = workspaces.read_cwd(client, &session_id, cwd).await?; + status(client, session_id, cwd).await? + } Command::SetMode { session_id, mode } => set_mode(client, session_id, mode).await?, Command::SetModel { session_id, model } => set_model(client, session_id, model).await?, Command::SetEffort { session_id, effort } => set_effort(client, session_id, effort).await?, @@ -560,7 +608,7 @@ async fn run_inner( delivery, until_idle, }; - prompt(client, state, args).await? + prompt(client, state, workspaces, args).await? } Command::Runtime { session_id } => runtime(client, session_id).await?, Command::Follow { @@ -568,7 +616,10 @@ async fn run_inner( cwd, interval, timeout, - } => follow(client, session_id, cwd, interval, timeout).await?, + } => { + let cwd = workspaces.read_cwd(client, &session_id, cwd).await?; + follow(client, session_id, cwd, interval, timeout).await? + } Command::Watch { session_id, interval, @@ -605,7 +656,7 @@ async fn run_inner( session_id, input_id, } => discard_queued(client, session_id, input_id).await?, - Command::Exec { calls } => exec(client, calls, state).await?, + Command::Exec { calls } => exec(client, calls, state, workspaces).await?, Command::Cancel { session_id } => cancel(client, session_id).await?, Command::Close { session_id } => close(client, session_id).await?, }; @@ -677,30 +728,33 @@ async fn models( async fn new_session( client: &AcpClient, state: &mut OperationState, + workspaces: &Workspaces, cwd: Option, profile: Option, - mode: Option, - model: Option, - effort: Option, + config: (Option, Option, Option), ) -> Result { - client.initialize().await.map_err(CliError::from_request)?; - let cwd = session::resolve_cwd(cwd).map_err(CliError::rpc)?; - let (created, created_raw) = client - .new_session(cwd.clone(), profile.as_deref()) - .await - .map_err(CliError::from_request)?; - let session_id = created.session_id.to_string(); - state.session(&session_id, &cwd, true); - apply_config(client, &session_id, mode, model, effort).await?; + let binding = Binding { + cwd: workspaces.initial_cwd(cwd)?, + profile, + }; + let source = if binding.profile.is_some() { + "explicit" + } else { + "backend-default-unknown" + }; + let (selected, created, created_raw) = workspaces + .create(client, state, binding, None, source) + .await?; + let session_id = &selected.session_id; + let cwd = &selected.binding.cwd; + apply_config(client, session_id, config.0, config.1, config.2).await?; let status = match open_session(client, session_id.clone(), cwd.clone()).await { Ok(opened) => opened.status, - Err(_) => session::config_status(&session_id, created.modes.as_ref(), &created_raw), + Err(_) => session::config_status(session_id, created.modes.as_ref(), &created_raw), }; - Ok(json!({ - "sessionId": session_id, - "cwd": cwd, - "status": status, - })) + let mut result = selected.result(); + result["status"] = status; + Ok(result) } async fn sessions( @@ -1000,7 +1054,7 @@ async fn set_pinned( async fn inspect( client: &AcpClient, session_id: String, - cwd: Option, + cwd: PathBuf, full: bool, messages: Option, tools: bool, @@ -1009,11 +1063,16 @@ async fn inspect( if !initialized.agent_capabilities.load_session { return Err(CliError::rpc("agent does not support session/load")); } - let cwd = session::resolve_cwd(cwd).map_err(CliError::rpc)?; let (loaded, load_value) = client .load_session(session_id.clone(), cwd.clone()) .await - .map_err(CliError::from_request)?; + .map_err(|error| { + if error.to_string().contains("Loaded session has no runtime ref for ACP replay") { + CliError::rpc("QueryMT backend cannot replay this session's history: Loaded session has no runtime ref for ACP replay. History is unavailable, not empty; qmcp cannot repair the owning runtime. No fallback was attempted.") + } else { + CliError::from_request(error) + } + })?; if full { return Ok(json!({ "sessionId": session_id, @@ -1036,14 +1095,10 @@ async fn inspect( )) } -async fn status( - client: &AcpClient, - session_id: String, - cwd: Option, -) -> Result { +async fn status(client: &AcpClient, session_id: String, cwd: PathBuf) -> Result { client.initialize().await.map_err(CliError::from_request)?; - let cwd = session::resolve_cwd(cwd).map_err(CliError::rpc)?; - let opened = open_session(client, session_id.clone(), cwd).await?; + let mut opened = open_session(client, session_id.clone(), cwd.clone()).await?; + opened.status["cwd"] = json!(cwd); Ok(opened.status) } @@ -1115,6 +1170,7 @@ struct PromptArgs { async fn prompt( client: &AcpClient, state: &mut OperationState, + workspaces: &Workspaces, args: PromptArgs, ) -> Result { let PromptArgs { @@ -1130,21 +1186,30 @@ async fn prompt( until_idle, } = args; let initialized = client.initialize().await.map_err(CliError::from_request)?; - let cwd = session::resolve_cwd(cwd).map_err(CliError::rpc)?; - let session_id = if new { - let (created, _) = client - .new_session(cwd.clone(), profile.as_deref()) - .await - .map_err(CliError::from_request)?; - let session_id = created.session_id.to_string(); - state.session(&session_id, &cwd, true); - session_id + let selected = if new { + let binding = Binding { + cwd: workspaces.initial_cwd(cwd)?, + profile, + }; + let source = if binding.profile.is_some() { + "explicit" + } else { + "backend-default-unknown" + }; + workspaces + .create(client, state, binding, None, source) + .await? + .0 } else { let session_id = session_id .ok_or_else(|| CliError::Usage("prompt requires session_id or new=true".into()))?; - state.session(&session_id, &cwd, false); - session_id + workspaces + .select(client, state, Some(session_id), cwd, profile) + .await? }; + selected.checkpoint(state); + let session_id = selected.session_id.clone(); + let cwd = selected.binding.cwd.clone(); // Keep the session id observable even if configuration or the prompt later fails. client .emit(json!({ "type": "session", "sessionId": session_id, "cwd": cwd })) @@ -1153,12 +1218,12 @@ async fn prompt( state .recovery .insert("deliveryState".into(), json!("not_submitted")); - let chosen = if new { + let chosen = if selected.created { Delivery::Prompt } else { choose_delivery(client, &session_id, delivery).await? }; - if !new && chosen == Delivery::Prompt { + if !selected.created && chosen == Delivery::Prompt { open_session(client, session_id.clone(), cwd.clone()).await?; } apply_config(client, &session_id, mode, model, effort).await?; @@ -1200,17 +1265,24 @@ async fn prompt( wait_until_idle(client, &session_id, Duration::from_secs(u32::MAX.into())).await?; } let text = client.take_assistant_text().await; - Ok(json!({ - "type": "done", - "sessionId": session_id, - "delivery": delivery_name(chosen), - "stopReason": stop_reason, - "text": text, - "submitted": submitted, - "deliveryState": "accepted", - "observationState": if until_idle { "idle" } else if chosen == Delivery::Prompt { "prompt_returned" } else { "not_waited" }, - "loadSession": initialized.agent_capabilities.load_session, - })) + let mut result = selected.result(); + result.as_object_mut().unwrap().extend( + json!({ + "type": "done", + "sessionId": session_id, + "delivery": delivery_name(chosen), + "stopReason": stop_reason, + "text": text, + "submitted": submitted, + "deliveryState": "accepted", + "observationState": if until_idle { "idle" } else if chosen == Delivery::Prompt { "prompt_returned" } else { "not_waited" }, + "loadSession": initialized.agent_capabilities.load_session, + }) + .as_object() + .unwrap() + .clone(), + ); + Ok(result) } fn delivery_name(delivery: Delivery) -> &'static str { @@ -1294,12 +1366,11 @@ async fn wait_until_idle_every( async fn follow( client: &AcpClient, session_id: String, - cwd: Option, + cwd: PathBuf, interval: u64, timeout: u64, ) -> Result { client.initialize().await.map_err(CliError::from_request)?; - let cwd = session::resolve_cwd(cwd).map_err(CliError::rpc)?; open_session(client, session_id.clone(), cwd.clone()).await?; client.clear_assistant_text().await; wait_until_idle_every( @@ -1310,15 +1381,22 @@ async fn follow( ) .await?; let mut text = client.take_assistant_text().await; - if text.is_empty() - && let Ok(snapshot) = - inspect(client, session_id.clone(), Some(cwd), false, Some(1), false).await - { + if text.is_empty() { + let snapshot = inspect( + client, + session_id.clone(), + cwd.clone(), + false, + Some(1), + false, + ) + .await?; text = session::last_assistant_text(&snapshot).unwrap_or_default(); } Ok(json!({ "type": "done", "sessionId": session_id, + "cwd": cwd, "stopReason": "idle", "text": text, })) @@ -1448,12 +1526,13 @@ async fn exec( client: &AcpClient, calls: Vec, state: &mut OperationState, + workspaces: &Workspaces, ) -> Result { state.batch = Some(BatchState::default()); for call in calls { state.batch.as_mut().expect("batch initialized").active_call = Some(call.name.clone()); let command = decode(&call.name, call.arguments).map_err(CliError::rpc)?; - let result = Box::pin(run(client, command, state, None)).await; + let result = Box::pin(run(client, command, state, None, workspaces)).await; let batch = state.batch.as_mut().expect("batch initialized"); batch.text_truncated |= result["textTruncated"] == true; if output::failed(&result) { diff --git a/src/lib.rs b/src/lib.rs index 39a52a5..61372fe 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,9 +2,11 @@ mod commands; mod output; +mod workspace; pub use qmacp::policy; +pub mod local_oauth; pub mod oauth; pub mod server; pub mod transport; diff --git a/src/local_oauth.rs b/src/local_oauth.rs new file mode 100644 index 0000000..407ac95 --- /dev/null +++ b/src/local_oauth.rs @@ -0,0 +1,1235 @@ +//! Embedded, volatile, single-owner authorization server. External OAuth stays separate. + +use std::{ + collections::{HashMap, HashSet}, + io::Read, + path::Path, + sync::{Arc, Mutex}, + time::{Duration, Instant}, +}; + +use anyhow::{Context, Result, bail}; +use axum::{ + Json, Router, + body::to_bytes, + extract::{Request, State}, + http::{HeaderMap, HeaderValue, Method, StatusCode, header}, + middleware::{self, Next}, + response::{Html, IntoResponse, Response}, + routing::{get, post}, +}; +use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; +use reqwest::Url; +use serde::Deserialize; +use serde_json::json; +use sha2::{Digest, Sha256}; +use subtle::ConstantTimeEq; + +use crate::transport::bearer_credential; + +const SCOPE: &str = "qmcp:access"; +const CLIENT_TTL: Duration = Duration::from_secs(3600); +const TRANSACTION_TTL: Duration = Duration::from_secs(300); +const CODE_TTL: Duration = Duration::from_secs(90); +const TOKEN_TTL: Duration = Duration::from_secs(1800); +const FAMILY_TTL: Duration = Duration::from_secs(24 * 3600); +const MAX_FAMILIES: usize = 128; +const MAX_CLIENT_FAMILIES: usize = 8; +const MAX_REFRESH_RECORDS: usize = 8192; +const MAX_FAMILY_RECORDS: usize = 128; +const MAX_CLIENTS: usize = 128; +const MAX_TRANSACTIONS: usize = 128; +const MAX_CODES: usize = 128; +const MAX_TOKENS: usize = 256; +const MAX_UNAPPROVED_CLIENTS: usize = 96; +const MAX_CLIENT_TRANSACTIONS: usize = 8; +const MAX_CLIENT_CODES: usize = 8; +const MAX_CLIENT_TOKENS: usize = 16; +const AGGREGATE_RATE: usize = 4096; +const REGISTRATION_RATE: usize = 1024; +const MAX_REGISTRATION_ORIGINS: usize = 256; +const MAX_BODY: usize = 8192; +const COOKIE: &str = "__Secure-qmcp_oauth"; +type Digest32 = [u8; 32]; +type Fields = HashMap; + +/// Deliberately no Debug: owner secrets, forms, codes and tokens must not be logged. +#[derive(Clone)] +pub struct LocalOAuth(Arc); + +struct Server { + resource: String, + issuer: String, + authority: String, + owner: Digest32, + store: Mutex, + body_slots: tokio::sync::Semaphore, +} + +#[derive(Default)] +struct Store { + clients: HashMap, + transactions: HashMap, + codes: HashMap, + tokens: HashMap, + families: HashMap, + refresh: HashMap, + registration_rate: Rate, + authorization_rate: Rate, + exchange_rate: Rate, + registration_origins: HashMap, +} + +struct Client { + redirects: Vec, + name: String, + expires: Instant, + approved: bool, + refresh_allowed: bool, + authorization_rate: Rate, + login_rate: Rate, + exchange_rate: Rate, +} + +struct Grant { + client: String, + redirect: String, + resource: String, + scope: String, + challenge: Digest32, +} + +struct Transaction { + grant: Grant, + state: Option, + csrf: Digest32, + expires: Instant, +} + +struct Code { + grant: Grant, + expires: Instant, +} + +struct Family { + client: String, + resource: String, + scope: String, + expires: Instant, + records: usize, + revoked: bool, +} + +struct RefreshRecord { + family: Digest32, + consumed: bool, +} + +struct Token { + family: Option, + client: String, + resource: String, + scope: String, + expires: Instant, +} + +/// Fixed-memory windows deliberately ignore spoofable proxy/IP headers. +#[derive(Default)] +struct Rate { + start: Option, + used: usize, +} + +impl Rate { + fn allow(&mut self, now: Instant, limit: usize) -> bool { + if self + .start + .is_none_or(|start| now.duration_since(start) >= Duration::from_secs(60)) + { + self.start = Some(now); + self.used = 0; + } + if self.used >= limit { + return false; + } + self.used += 1; + true + } +} + +impl Store { + fn prune(&mut self, now: Instant) { + self.clients.retain(|_, c| c.approved || c.expires > now); + self.registration_origins.retain(|_, rate| { + rate.start + .is_some_and(|start| now.duration_since(start) < Duration::from_secs(60)) + }); + self.transactions.retain(|_, t| t.expires > now); + self.codes.retain(|_, c| c.expires > now); + self.families.retain(|_, family| family.expires > now); + self.refresh + .retain(|_, record| self.families.contains_key(&record.family)); + self.tokens.retain(|_, t| { + t.expires > now + && t.family + .is_none_or(|id| self.families.get(&id).is_some_and(|family| !family.revoked)) + }); + } + + fn access_capacity(&self, client: &str) -> bool { + self.tokens.len() < MAX_TOKENS + && self.tokens.values().filter(|t| t.client == client).count() < MAX_CLIENT_TOKENS + } + + fn family_capacity(&self, client: &str) -> bool { + self.families.len() < MAX_FAMILIES + && self + .families + .values() + .filter(|f| f.client == client) + .count() + < MAX_CLIENT_FAMILIES + && self.refresh.len() < MAX_REFRESH_RECORDS + } + + fn revoke_family(&mut self, id: Digest32) { + if let Some(family) = self.families.get_mut(&id) { + family.revoked = true; + } + self.tokens.retain(|_, token| token.family != Some(id)); + // Keep all refresh digests, including the active one, until absolute expiry. + } +} + +impl LocalOAuth { + /// Read a generated hex credential (32..64 random bytes) from a protected file. + pub fn from_secret_file(resource: &str, path: &Path) -> Result { + if !path.is_absolute() { + bail!("owner secret file must be an absolute path"); + } + let mut options = std::fs::OpenOptions::new(); + options.read(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + // Reject final-component symlinks and avoid blocking on a FIFO/device. + options.custom_flags(libc::O_NOFOLLOW | libc::O_NONBLOCK); + } + let file = options + .open(path) + .context("open protected owner secret file")?; + let metadata = file.metadata().context("stat owner secret file")?; + if !metadata.is_file() || metadata.len() > 130 { + bail!("owner secret must be a small regular file"); + } + #[cfg(unix)] + { + use std::os::unix::fs::MetadataExt; + // SAFETY: geteuid has no arguments or memory-safety preconditions. + let uid = unsafe { libc::geteuid() }; + if metadata.mode() & 0o7777 != 0o600 || metadata.uid() != uid || metadata.nlink() != 1 { + bail!("owner secret file must be owned by this user, mode 0600, with one link"); + } + } + #[cfg(not(unix))] + bail!("embedded OAuth secret-file protection currently requires Unix"); + let mut secret = String::new(); + file.take(131) + .read_to_string(&mut secret) + .context("read owner secret file")?; + Self::new(resource, secret.trim_end_matches(['\r', '\n'])) + } + + fn new(resource: &str, secret: &str) -> Result { + let url = Url::parse(resource).context("invalid local OAuth resource URL")?; + if resource.len() > 2048 + || url.scheme() != "https" + || url.host_str().is_none() + || url.path() != "/mcp" + || !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + || url.as_str() != resource + { + bail!( + "local OAuth resource must be a canonical HTTPS /mcp URL without userinfo, query or fragment" + ); + } + // Encoding/length and obvious-weakness checks cannot prove random generation. + // Only use a CSPRNG-generated credential, never a password or a derived phrase. + if !(64..=128).contains(&secret.len()) + || !secret.len().is_multiple_of(2) + || !secret + .bytes() + .all(|b| b.is_ascii_hexdigit() && !b.is_ascii_uppercase()) + { + bail!("owner secret must be 32..64 random bytes encoded as lowercase hex"); + } + let decoded: Vec = secret + .as_bytes() + .as_chunks::<2>() + .0 + .iter() + .map(|pair| u8::from_str_radix(std::str::from_utf8(pair).unwrap(), 16).unwrap()) + .collect(); + if decoded.iter().collect::>().len() < 16 { + bail!("owner secret is obviously weak; generate it using a secure random generator"); + } + let issuer = url.origin().ascii_serialization(); + let authority = issuer.strip_prefix("https://").unwrap().to_owned(); + Ok(Self(Arc::new(Server { + resource: resource.into(), + issuer, + authority, + owner: digest(secret), + store: Mutex::new(Store::default()), + body_slots: tokio::sync::Semaphore::new(16), + }))) + } + + pub(crate) fn authority(&self) -> &str { + &self.0.authority + } + + pub(crate) fn routes(&self) -> Router { + Router::new() + .route("/.well-known/oauth-authorization-server", get(as_metadata)) + .route( + "/.well-known/oauth-protected-resource", + get(resource_metadata), + ) + .route( + "/.well-known/oauth-protected-resource/mcp", + get(resource_metadata), + ) + .route("/oauth/register", post(register)) + .route("/oauth/authorize", get(authorize).post(consent)) + .route("/oauth/token", post(exchange)) + .with_state(self.clone()) + } + + pub(crate) fn guard(&self, router: Router) -> Router { + router.layer(middleware::from_fn_with_state(self.clone(), guard)) + } + + fn accepts(&self, credential: &str) -> bool { + let now = Instant::now(); + let mut store = self.0.store.lock().unwrap(); + store.prune(now); + store.tokens.get(&digest(credential)).is_some_and(|token| { + token.resource == self.0.resource && token.scope == SCOPE && !token.client.is_empty() + }) + } +} + +fn digest(value: &str) -> Digest32 { + Sha256::digest(value.as_bytes()).into() +} + +struct Failure(StatusCode, &'static str); + +impl IntoResponse for Failure { + fn into_response(self) -> Response { + error(self.0, self.1) + } +} + +fn random() -> Result { + let mut bytes = [0u8; 32]; + getrandom::fill(&mut bytes) + .map_err(|_| Failure(StatusCode::SERVICE_UNAVAILABLE, "temporarily_unavailable"))?; + Ok(URL_SAFE_NO_PAD.encode(bytes)) +} + +fn secure(mut response: Response) -> Response { + for (key, value) in [ + ("cache-control", "no-store"), + ("pragma", "no-cache"), + ( + "content-security-policy", + "default-src 'none'; base-uri 'none'; form-action 'self'; frame-ancestors 'none'", + ), + ("x-frame-options", "DENY"), + ("referrer-policy", "no-referrer"), + ("x-content-type-options", "nosniff"), + ] { + response + .headers_mut() + .insert(key, HeaderValue::from_static(value)); + } + response +} + +fn error(status: StatusCode, name: &'static str) -> Response { + secure((status, Json(json!({"error": name}))).into_response()) +} + +fn one_header<'a>(headers: &'a HeaderMap, name: &str) -> Option<&'a str> { + let mut values = headers.get_all(name).iter(); + let value = values.next()?.to_str().ok()?; + values.next().is_none().then_some(value) +} + +async fn guard(State(auth): State, request: Request, next: Next) -> Response { + let host_ok = one_header(request.headers(), "host") + .is_some_and(|host| host.eq_ignore_ascii_case(&auth.0.authority)); + let browser_post = + request.method() == Method::POST && request.uri().path() == "/oauth/authorize"; + let origin = one_header(request.headers(), "origin"); + let origin_ok = if browser_post { + origin == Some(auth.0.issuer.as_str()) + } else { + !request.headers().contains_key("origin") + }; + if !host_ok || !origin_ok { + return error(StatusCode::FORBIDDEN, "access_denied"); + } + if request.method() == Method::POST && request.uri().query().is_some() { + return error(StatusCode::BAD_REQUEST, "invalid_request"); + } + if request.uri().to_string().len() > MAX_BODY { + return error(StatusCode::URI_TOO_LONG, "invalid_request"); + } + let mut response = secure(next.run(request).await); + if response + .headers() + .get(header::CONTENT_TYPE) + .is_some_and(|value| { + value + .to_str() + .is_ok_and(|value| value.starts_with("text/html")) + }) + { + // A no-referrer form document makes browser POST Origin opaque (null). + response + .headers_mut() + .insert("referrer-policy", HeaderValue::from_static("same-origin")); + } + response +} + +pub(crate) async fn require_token( + State(auth): State, + mut request: Request, + next: Next, +) -> Response { + if !bearer_credential(request.headers()).is_some_and(|token| auth.accepts(token)) { + let mut response = error(StatusCode::UNAUTHORIZED, "invalid_token"); + let challenge = format!( + "Bearer resource_metadata=\"{}/.well-known/oauth-protected-resource/mcp\", scope=\"{SCOPE}\"", + auth.0.issuer + ); + response + .headers_mut() + .insert(header::WWW_AUTHENTICATE, challenge.parse().unwrap()); + return response; + } + request.headers_mut().remove(header::AUTHORIZATION); + next.run(request).await +} + +async fn as_metadata(State(auth): State) -> Response { + let issuer = &auth.0.issuer; + Json(json!({ + "issuer": issuer, + "authorization_endpoint": format!("{issuer}/oauth/authorize"), + "token_endpoint": format!("{issuer}/oauth/token"), + "registration_endpoint": format!("{issuer}/oauth/register"), + "response_types_supported": ["code"], + "grant_types_supported": ["authorization_code", "refresh_token"], + "token_endpoint_auth_methods_supported": ["none"], + "code_challenge_methods_supported": ["S256"], + "scopes_supported": [SCOPE], + "authorization_response_iss_parameter_supported": true, + })) + .into_response() +} + +async fn resource_metadata(State(auth): State) -> Response { + Json(json!({ + "resource": auth.0.resource, + "authorization_servers": [auth.0.issuer], + "scopes_supported": [SCOPE], + "bearer_methods_supported": ["header"], + })) + .into_response() +} + +async fn body(auth: &LocalOAuth, request: Request, content_type: &str) -> Result, Failure> { + if !one_header(request.headers(), "content-type").is_some_and(|v| { + v.split(';') + .next() + .is_some_and(|v| v.trim() == content_type) + }) { + return Err(Failure( + StatusCode::UNSUPPORTED_MEDIA_TYPE, + "invalid_request", + )); + } + // Bound concurrent AS buffers and slow uploads without timing out MCP work. + let _permit = auth + .0 + .body_slots + .try_acquire() + .map_err(|_| Failure(StatusCode::TOO_MANY_REQUESTS, "slow_down"))?; + let bytes = tokio::time::timeout( + Duration::from_secs(10), + to_bytes(request.into_body(), MAX_BODY), + ) + .await + .map_err(|_| Failure(StatusCode::REQUEST_TIMEOUT, "invalid_request"))? + .map_err(|_| Failure(StatusCode::PAYLOAD_TOO_LARGE, "invalid_request"))?; + Ok(bytes.to_vec()) +} + +fn fields(encoded: &[u8]) -> Result { + if encoded.len() > MAX_BODY { + return Err(Failure(StatusCode::BAD_REQUEST, "invalid_request")); + } + // Reject malformed escapes/UTF-8 rather than silently changing the client's state. + let mut decoded = Vec::with_capacity(encoded.len()); + let mut index = 0; + while index < encoded.len() { + if encoded[index] == b'%' { + let byte = encoded + .get(index + 1..index + 3) + .filter(|pair| pair.iter().all(u8::is_ascii_hexdigit)) + .and_then(|pair| std::str::from_utf8(pair).ok()) + .and_then(|pair| u8::from_str_radix(pair, 16).ok()) + .ok_or(Failure(StatusCode::BAD_REQUEST, "invalid_request"))?; + decoded.push(byte); + index += 3; + } else { + decoded.push(encoded[index]); + index += 1; + } + } + std::str::from_utf8(&decoded) + .map_err(|_| Failure(StatusCode::BAD_REQUEST, "invalid_request"))?; + let mut result = Fields::new(); + for (key, value) in form_urlencoded::parse(encoded) { + let recognized = matches!( + key.as_ref(), + "client_id" + | "redirect_uri" + | "response_type" + | "response_mode" + | "scope" + | "resource" + | "state" + | "code_challenge" + | "code_challenge_method" + | "grant_type" + | "code" + | "code_verifier" + | "csrf" + | "decision" + | "owner_secret" + | "client_secret" + | "client_assertion" + | "client_assertion_type" + | "access_token" + | "refresh_token" + | "password" + | "token" + | "request" + | "request_uri" + ); + if key.len() > 64 || value.len() > 2048 || result.len() >= 32 { + return Err(Failure(StatusCode::BAD_REQUEST, "invalid_request")); + } + if result + .insert(key.into_owned(), value.into_owned()) + .is_some() + && recognized + { + return Err(Failure(StatusCode::BAD_REQUEST, "invalid_request")); + } + } + Ok(result) +} + +fn param<'a>(fields: &'a Fields, name: &str) -> &'a str { + fields.get(name).map_or("", String::as_str) +} + +#[derive(Deserialize)] +struct Registration { + redirect_uris: Vec, + client_name: Option, + token_endpoint_auth_method: Option, + grant_types: Option>, + response_types: Option>, + scope: Option, +} + +fn valid_redirect(value: &str) -> bool { + let Ok(url) = Url::parse(value) else { + return false; + }; + let loopback = url.host_str().is_some_and(|host| { + host == "localhost" + || host + .trim_matches(['[', ']']) + .parse::() + .is_ok_and(|ip| ip.is_loopback()) + }); + value.len() <= 2048 + && url.host_str().is_some() + && (url.scheme() == "https" || (url.scheme() == "http" && loopback)) + && url.username().is_empty() + && url.password().is_none() + && url.fragment().is_none() + && url.as_str() == value + && !url.query_pairs().any(|(key, _)| { + matches!( + key.as_ref(), + "code" | "state" | "iss" | "error" | "error_description" + ) + }) +} + +async fn register(State(auth): State, request: Request) -> Response { + let now = Instant::now(); + { + let mut store = auth.0.store.lock().unwrap(); + store.prune(now); + if !store.registration_rate.allow(now, REGISTRATION_RATE) { + return error(StatusCode::TOO_MANY_REQUESTS, "slow_down"); + } + } + let bytes = match body(&auth, request, "application/json").await { + Ok(bytes) => bytes, + Err(error) => return error.into_response(), + }; + let registration: Registration = match serde_json::from_slice(&bytes) { + Ok(registration) => registration, + Err(_) => return error(StatusCode::BAD_REQUEST, "invalid_client_metadata"), + }; + if registration.redirect_uris.is_empty() + || registration.redirect_uris.len() > 8 + || registration + .redirect_uris + .iter() + .any(|uri| !valid_redirect(uri)) + || registration + .client_name + .as_ref() + .is_some_and(|name| name.len() > 128 || name.chars().any(char::is_control)) + || registration + .token_endpoint_auth_method + .as_deref() + .is_some_and(|method| method != "none") + || registration.grant_types.as_ref().is_some_and(|grants| { + !grants.iter().any(|grant| grant == "authorization_code") + || grants + .iter() + .any(|grant| grant != "authorization_code" && grant != "refresh_token") + || grants.iter().collect::>().len() != grants.len() + }) + || registration + .response_types + .as_ref() + .is_some_and(|types| types != &["code"]) + || registration + .scope + .as_deref() + .is_some_and(|scope| scope != SCOPE) + { + return error(StatusCode::BAD_REQUEST, "invalid_client_metadata"); + } + let id = match random() { + Ok(id) => id, + Err(error) => return error.into_response(), + }; + let mut store = auth.0.store.lock().unwrap(); + // Bound registrations by validated callback origin and aggregate work. + // Descriptive metadata URLs are never fetched. + let origins: HashSet<_> = registration + .redirect_uris + .iter() + .map(|uri| Url::parse(uri).unwrap().origin().ascii_serialization()) + .collect(); + let new_origins = origins + .iter() + .filter(|origin| !store.registration_origins.contains_key(*origin)) + .count(); + if store.registration_origins.len() + new_origins > MAX_REGISTRATION_ORIGINS { + return error(StatusCode::TOO_MANY_REQUESTS, "slow_down"); + } + for origin in origins { + if !store + .registration_origins + .entry(origin) + .or_default() + .allow(now, 16) + { + return error(StatusCode::TOO_MANY_REQUESTS, "slow_down"); + } + } + if store.clients.len() >= MAX_CLIENTS + || store + .clients + .values() + .filter(|client| !client.approved) + .count() + >= MAX_UNAPPROVED_CLIENTS + { + return error(StatusCode::SERVICE_UNAVAILABLE, "temporarily_unavailable"); + } + let refresh_allowed = registration + .grant_types + .as_ref() + .is_some_and(|grants| grants.iter().any(|grant| grant == "refresh_token")); + let grants = if refresh_allowed { + vec!["authorization_code", "refresh_token"] + } else { + vec!["authorization_code"] + }; + let response = ( + StatusCode::CREATED, + Json(json!({ + "client_id": id, + "redirect_uris": registration.redirect_uris, "client_name": registration.client_name, + "token_endpoint_auth_method": "none", "grant_types": grants, + "response_types": ["code"], "scope": SCOPE, + })), + ) + .into_response(); + store.clients.insert( + id, + Client { + redirects: registration.redirect_uris, + name: registration + .client_name + .unwrap_or_else(|| "Unnamed public client".into()), + expires: now + CLIENT_TTL, + approved: false, + refresh_allowed, + authorization_rate: Rate::default(), + login_rate: Rate::default(), + exchange_rate: Rate::default(), + }, + ); + response +} + +fn completion( + auth: &LocalOAuth, + target: &str, + state: Option<&str>, + key: &str, + value: &str, +) -> Response { + let mut url = Url::parse(target).expect("validated registered redirect"); + { + let mut pairs = url.query_pairs_mut(); + pairs + .append_pair(key, value) + .append_pair("iss", &auth.0.issuer); + if let Some(state) = state { + pairs.append_pair("state", state); + } + } + // An ordinary user-clicked link is not a form-action redirect. Keep the + // callback/code private and do not navigate externally without another click. + Html(format!( + "Complete qmcp authorization

Complete authorization

Continue to the callback you reviewed, including after denial.

Continue to ChatGPT", + escape(url.as_str()) + )).into_response() +} + +fn escape(text: &str) -> String { + text.replace('&', "&") + .replace('<', "<") + .replace('>', ">") + .replace('"', """) + .replace('\'', "'") +} + +fn cookie(value: &str, max_age: u64) -> HeaderValue { + format!( + "{COOKIE}={value}; Max-Age={max_age}; Path=/oauth/authorize; Secure; HttpOnly; SameSite=Lax" + ) + .parse() + .unwrap() +} + +async fn authorize(State(auth): State, request: Request) -> Response { + let fields = match fields(request.uri().query().unwrap_or("").as_bytes()) { + Ok(f) => f, + Err(e) => return e.into_response(), + }; + // Never accept a credential in a URL, even as an unknown extension parameter. + if fields.keys().any(|k| { + matches!( + k.as_str(), + "owner_secret" + | "client_secret" + | "client_assertion" + | "client_assertion_type" + | "access_token" + | "refresh_token" + | "code" + | "code_verifier" + | "password" + | "token" + ) + }) || fields.get("state").is_some_and(|s| s.len() > 1024) + || fields + .get("response_mode") + .is_some_and(|mode| mode != "query") + || fields.contains_key("request") + || fields.contains_key("request_uri") + { + return error(StatusCode::BAD_REQUEST, "invalid_request"); + } + let now = Instant::now(); + let mut store = auth.0.store.lock().unwrap(); + store.prune(now); + if !store.authorization_rate.allow(now, AGGREGATE_RATE) { + return error(StatusCode::TOO_MANY_REQUESTS, "slow_down"); + } + let Some(client) = store.clients.get_mut(param(&fields, "client_id")) else { + return error(StatusCode::BAD_REQUEST, "invalid_client"); + }; + if !client.authorization_rate.allow(now, 30) { + return error(StatusCode::TOO_MANY_REQUESTS, "slow_down"); + } + let name = escape(&client.name); + let refresh_allowed = client.refresh_allowed; + let target = param(&fields, "redirect_uri"); + if !client.redirects.iter().any(|r| r == target) { + return error(StatusCode::BAD_REQUEST, "invalid_request"); + } + let state = fields.get("state").cloned(); + let fail = |name| completion(&auth, target, state.as_deref(), "error", name); + if param(&fields, "response_type") != "code" { + return fail("unsupported_response_type"); + } + if param(&fields, "resource") != auth.0.resource { + return fail("invalid_target"); + } + if param(&fields, "scope") != SCOPE { + return fail("invalid_scope"); + } + let challenge = param(&fields, "code_challenge"); + let challenge_bytes = URL_SAFE_NO_PAD.decode(challenge).ok(); + if param(&fields, "code_challenge_method") != "S256" + || challenge.len() != 43 + || !challenge_bytes + .as_ref() + .is_some_and(|b| b.len() == 32 && URL_SAFE_NO_PAD.encode(b) == challenge) + { + return fail("invalid_request"); + } + if store.transactions.len() >= MAX_TRANSACTIONS + || store + .transactions + .values() + .filter(|t| t.grant.client == param(&fields, "client_id")) + .count() + >= MAX_CLIENT_TRANSACTIONS + { + return fail("temporarily_unavailable"); + } + let transaction = match random() { + Ok(v) => v, + Err(e) => return e.into_response(), + }; + let csrf = match random() { + Ok(v) => v, + Err(e) => return e.into_response(), + }; + let access_minutes = TOKEN_TTL.as_secs() / 60; + let duration = if refresh_allowed { + format!( + "Each access token lasts up to {access_minutes} minutes. Access can be renewed without a new login for up to {} hours from the initial code exchange; renewal does not extend this deadline.", + FAMILY_TTL.as_secs() / 3600 + ) + } else { + format!("Access lasts up to {access_minutes} minutes.") + }; + let html = format!( + "qmcp owner login and authorization

Owner login + authorize

Unverified public client: {name}

Redirect: {}

Approve grants all QueryMT tools on {} (scope {SCOPE}). {duration} Only approve a client and callback you trust.

", + escape(target), + escape(&auth.0.resource) + ); + let grant = Grant { + client: param(&fields, "client_id").into(), + redirect: target.into(), + resource: auth.0.resource.clone(), + scope: SCOPE.into(), + challenge: challenge_bytes.unwrap().try_into().unwrap(), + }; + store.transactions.insert( + digest(&transaction), + Transaction { + grant, + state, + csrf: digest(&csrf), + expires: now + TRANSACTION_TTL, + }, + ); + let mut response = Html(html).into_response(); + response.headers_mut().insert( + header::SET_COOKIE, + cookie(&transaction, TRANSACTION_TTL.as_secs()), + ); + response +} + +fn browser_cookie(headers: &HeaderMap) -> Option<&str> { + let mut found = None; + for header in headers.get_all(header::COOKIE) { + for item in header.to_str().ok()?.split(';') { + let Some((key, value)) = item.trim().split_once('=') else { + continue; + }; + if key == COOKIE { + if found.is_some() || value.len() != 43 { + return None; + } + found = Some(value); + } + } + } + found +} + +async fn consent(State(auth): State, request: Request) -> Response { + let transaction = browser_cookie(request.headers()).map(digest); + let bytes = match body(&auth, request, "application/x-www-form-urlencoded").await { + Ok(b) => b, + Err(e) => return e.into_response(), + }; + let fields = match fields(&bytes) { + Ok(f) => f, + Err(e) => return e.into_response(), + }; + let now = Instant::now(); + let mut store = auth.0.store.lock().unwrap(); + store.prune(now); + // Exact Origin is checked by guard before any body or owner-attempt work. + // Invalid cookies/CSRF neither consume an attempt nor destroy a valid flow. + let Some(id) = transaction else { + return error(StatusCode::BAD_REQUEST, "invalid_request"); + }; + let Some(transaction) = store.transactions.get(&id) else { + return error(StatusCode::BAD_REQUEST, "invalid_request"); + }; + if !bool::from(transaction.csrf.ct_eq(&digest(param(&fields, "csrf")))) { + return error(StatusCode::FORBIDDEN, "access_denied"); + } + let client_id = transaction.grant.client.clone(); + let Some(client) = store.clients.get_mut(&client_id) else { + return error(StatusCode::BAD_REQUEST, "invalid_client"); + }; + if !client.login_rate.allow(now, 8) { + return error(StatusCode::TOO_MANY_REQUESTS, "slow_down"); + } + // One owner attempt per CSRF-validated transaction, bounded per client. + let transaction = store.transactions.remove(&id).unwrap(); + let grant = transaction.grant; + let approved = param(&fields, "decision") == "approve" + && !fields + .keys() + .any(|k| !matches!(k.as_str(), "csrf" | "decision" | "owner_secret")) + && bool::from(auth.0.owner.ct_eq(&digest(param(&fields, "owner_secret")))); + if approved { + store.clients.get_mut(&grant.client).unwrap().approved = true; + } + let result = if !approved { + completion( + &auth, + &grant.redirect, + transaction.state.as_deref(), + "error", + "access_denied", + ) + } else if store.codes.len() >= MAX_CODES + || store + .codes + .values() + .filter(|code| code.grant.client == grant.client) + .count() + >= MAX_CLIENT_CODES + { + completion( + &auth, + &grant.redirect, + transaction.state.as_deref(), + "error", + "temporarily_unavailable", + ) + } else { + let code = match random() { + Ok(v) => v, + Err(e) => return e.into_response(), + }; + let result = completion( + &auth, + &grant.redirect, + transaction.state.as_deref(), + "code", + &code, + ); + store.codes.insert( + digest(&code), + Code { + grant, + expires: now + CODE_TTL, + }, + ); + result + }; + let mut result = result; + result + .headers_mut() + .insert(header::SET_COOKIE, cookie("", 0)); + result +} + +async fn exchange(State(auth): State, request: Request) -> Response { + if request.headers().contains_key(header::AUTHORIZATION) { + return error(StatusCode::BAD_REQUEST, "invalid_client"); + } + { + let now = Instant::now(); + let mut store = auth.0.store.lock().unwrap(); + store.prune(now); + if !store.exchange_rate.allow(now, AGGREGATE_RATE) { + return error(StatusCode::TOO_MANY_REQUESTS, "slow_down"); + } + } + let bytes = match body(&auth, request, "application/x-www-form-urlencoded").await { + Ok(b) => b, + Err(e) => return e.into_response(), + }; + let fields = match fields(&bytes) { + Ok(f) => f, + Err(e) => return e.into_response(), + }; + let mut store = auth.0.store.lock().unwrap(); + let now = Instant::now(); + store.prune(now); + match param(&fields, "grant_type") { + "authorization_code" => exchange_code(&mut store, &fields, now), + "refresh_token" => exchange_refresh(&mut store, &fields, now), + _ => error(StatusCode::BAD_REQUEST, "unsupported_grant_type"), + } +} + +fn exchange_code(store: &mut Store, fields: &Fields, now: Instant) -> Response { + let id = digest(param(fields, "code")); + // Unknown codes cannot consume a named client's budget. + if let Some(code) = store.codes.get(&id) { + let actual_client = code.grant.client.clone(); + if let Some(client) = store.clients.get_mut(&actual_client) + && !client.exchange_rate.allow(now, 60) + { + return error(StatusCode::TOO_MANY_REQUESTS, "slow_down"); + } + } + let Some(code) = store.codes.get(&id) else { + return error(StatusCode::BAD_REQUEST, "invalid_grant"); + }; + let verifier = param(fields, "code_verifier"); + let verifier_ok = (43..=128).contains(&verifier.len()) + && verifier + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b"-._~".contains(&b)); + if code.expires <= now + || fields.contains_key("client_secret") + || fields.contains_key("client_assertion") + || fields.contains_key("client_assertion_type") + || param(fields, "client_id") != code.grant.client + || param(fields, "redirect_uri") != code.grant.redirect + || param(fields, "resource") != code.grant.resource + || fields.get("scope").is_some_and(|s| s != &code.grant.scope) + || !verifier_ok + || !bool::from(code.grant.challenge.ct_eq(&digest(verifier))) + { + // Failed PKCE/bindings consume the code; transient capacity failures do not. + store.codes.remove(&id); + return error(StatusCode::BAD_REQUEST, "invalid_grant"); + } + let refresh_allowed = store + .clients + .get(&code.grant.client) + .is_some_and(|client| client.approved && client.refresh_allowed); + if !store.access_capacity(&code.grant.client) + || (refresh_allowed && !store.family_capacity(&code.grant.client)) + { + return error(StatusCode::SERVICE_UNAVAILABLE, "temporarily_unavailable"); + } + let token = match random() { + Ok(v) => v, + Err(e) => return e.into_response(), + }; + let refresh = if refresh_allowed { + match random() { + Ok(v) => Some(v), + Err(e) => return e.into_response(), + } + } else { + None + }; + // Validation, capacity checks, removal and all issuance share the same lock. + let grant = store.codes.remove(&id).unwrap().grant; + let family_id = refresh.as_deref().map(digest); + if let Some(family) = family_id { + store.families.insert( + family, + Family { + client: grant.client.clone(), + resource: grant.resource.clone(), + scope: grant.scope.clone(), + expires: now + FAMILY_TTL, + records: 1, + revoked: false, + }, + ); + store.refresh.insert( + family, + RefreshRecord { + family, + consumed: false, + }, + ); + } + store.tokens.insert( + digest(&token), + Token { + family: family_id, + client: grant.client, + resource: grant.resource, + scope: grant.scope, + expires: now + TOKEN_TTL, + }, + ); + token_response(token, refresh, TOKEN_TTL) +} + +fn exchange_refresh(store: &mut Store, fields: &Fields, now: Instant) -> Response { + let raw = param(fields, "refresh_token"); + // Only canonical 32-byte opaque credentials, never mixed authentication/grant material. + if raw.len() != 43 + || !URL_SAFE_NO_PAD + .decode(raw) + .is_ok_and(|bytes| bytes.len() == 32 && URL_SAFE_NO_PAD.encode(bytes) == raw) + || param(fields, "client_id").is_empty() + || [ + "client_secret", + "client_assertion", + "client_assertion_type", + "code", + "code_verifier", + "redirect_uri", + "owner_secret", + "password", + "access_token", + "token", + ] + .iter() + .any(|key| fields.contains_key(*key)) + { + return error(StatusCode::BAD_REQUEST, "invalid_grant"); + } + let id = digest(raw); + let Some(record) = store.refresh.get(&id) else { + return error(StatusCode::BAD_REQUEST, "invalid_grant"); + }; + let family_id = record.family; + let family = &store.families[&family_id]; + if family.revoked + || family.expires <= now + || param(fields, "client_id") != family.client + || !store + .clients + .get(&family.client) + .is_some_and(|c| c.approved && c.refresh_allowed) + { + return error(StatusCode::BAD_REQUEST, "invalid_grant"); + } + // Validate bindings before replay handling: malformed/wrong-target attempts do not revoke. + // RFC 8707 permits omission to reuse the original (single) bound resource. + if fields + .get("resource") + .is_some_and(|resource| resource != &family.resource) + { + return error(StatusCode::BAD_REQUEST, "invalid_target"); + } + if fields + .get("scope") + .is_some_and(|scope| scope != &family.scope) + { + return error(StatusCode::BAD_REQUEST, "invalid_scope"); + } + if record.consumed { + store.revoke_family(family_id); + return error(StatusCode::BAD_REQUEST, "invalid_grant"); + } + let client_id = family.client.clone(); + let resource = family.resource.clone(); + let scope = family.scope.clone(); + let expires = family.expires; + let records = family.records; + if !store + .clients + .get_mut(&client_id) + .unwrap() + .exchange_rate + .allow(now, 60) + { + return error(StatusCode::TOO_MANY_REQUESTS, "slow_down"); + } + if !store.access_capacity(&client_id) + || store.refresh.len() >= MAX_REFRESH_RECORDS + || records >= MAX_FAMILY_RECORDS + { + return error(StatusCode::SERVICE_UNAVAILABLE, "temporarily_unavailable"); + } + let token = match random() { + Ok(v) => v, + Err(e) => return e.into_response(), + }; + let refresh = match random() { + Ok(v) => v, + Err(e) => return e.into_response(), + }; + let ttl = TOKEN_TTL.min(expires.duration_since(now)); + store.refresh.get_mut(&id).unwrap().consumed = true; + store.families.get_mut(&family_id).unwrap().records += 1; + store.refresh.insert( + digest(&refresh), + RefreshRecord { + family: family_id, + consumed: false, + }, + ); + store.tokens.insert( + digest(&token), + Token { + family: Some(family_id), + client: client_id, + resource, + scope, + expires: now + ttl, + }, + ); + token_response(token, Some(refresh), ttl) +} + +fn token_response(access: String, refresh: Option, ttl: Duration) -> Response { + let mut response = json!({ "access_token": access, "token_type": "Bearer", + "expires_in": ttl.as_secs(), "scope": SCOPE }); + if let Some(refresh) = refresh { + response["refresh_token"] = refresh.into(); + } + Json(response).into_response() +} + +#[cfg(test)] +mod tests; diff --git a/src/local_oauth/tests.rs b/src/local_oauth/tests.rs new file mode 100644 index 0000000..8a0b26a --- /dev/null +++ b/src/local_oauth/tests.rs @@ -0,0 +1,1426 @@ +use super::*; +use crate::{BridgeConfig, QmcpServer, policy::PermissionPolicy, transport}; +use axum::body::Body; +use serde_json::Value; +use tower::ServiceExt; + +const RESOURCE: &str = "https://mcp.example.com/mcp"; +const ISSUER: &str = "https://mcp.example.com"; +const CALLBACK: &str = "https://client.example/callback?existing=1"; +// Public deterministic test fixture, not a runtime credential. +const SECRET: &str = "000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f"; +const VERIFIER: &str = "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFG._~-"; +const STATE: &str = "a&b=+ /?#%\"<>\0"; + +struct Harness { + auth: LocalOAuth, + app: Router, +} +struct Reply { + status: StatusCode, + headers: HeaderMap, + text: String, +} + +impl Reply { + fn json(&self) -> Value { + serde_json::from_str(&self.text).unwrap() + } + fn location(&self) -> Url { + assert_eq!(self.status, StatusCode::OK); + assert!(!self.headers.contains_key(header::LOCATION)); + let link = self + .text + .split("") + .replace("&", "&"); + Url::parse(&link).unwrap() + } + fn parameter(&self, key: &str) -> String { + self.location() + .query_pairs() + .find(|(k, _)| k == key) + .unwrap() + .1 + .into_owned() + } +} + +fn encoded(pairs: &[(&str, &str)]) -> String { + form_urlencoded::Serializer::new(String::new()) + .extend_pairs(pairs.iter().copied()) + .finish() +} + +impl Harness { + fn new() -> Self { + let auth = LocalOAuth::new(RESOURCE, SECRET).unwrap(); + let server = QmcpServer::new( + BridgeConfig::new("127.0.0.1:1", false, PermissionPolicy::Cancel).unwrap(), + ) + .unwrap(); + let app = transport::http_router_with_local_oauth( + server, + "127.0.0.1:3031".parse().unwrap(), + auth.clone(), + &[], + ) + .unwrap(); + Self { auth, app } + } + + async fn request( + &self, + method: &str, + uri: &str, + body: String, + headers: &[(&str, &str)], + ) -> Reply { + let mut request = Request::builder() + .method(method) + .uri(uri) + .header("host", "mcp.example.com"); + for (key, value) in headers { + request = request.header(*key, *value); + } + let response = self + .app + .clone() + .oneshot(request.body(Body::from(body)).unwrap()) + .await + .unwrap(); + assert_eq!(response.headers()[header::CACHE_CONTROL], "no-store"); + assert_eq!(response.headers()["x-frame-options"], "DENY"); + let status = response.status(); + let headers = response.headers().clone(); + let text = String::from_utf8( + to_bytes(response.into_body(), 65536) + .await + .unwrap() + .to_vec(), + ) + .unwrap(); + Reply { + status, + headers, + text, + } + } + + async fn register(&self, value: Value) -> Reply { + self.request( + "POST", + "/oauth/register", + value.to_string(), + &[("content-type", "application/json")], + ) + .await + } + + async fn client(&self) -> String { + let response = self.register(json!({ "redirect_uris": [CALLBACK], "client_name": "" })).await; + assert_eq!(response.status, StatusCode::CREATED, "{}", response.text); + response.json()["client_id"].as_str().unwrap().into() + } + + async fn authorize(&self, client: &str, changes: &[(&str, &str)]) -> Reply { + let challenge = URL_SAFE_NO_PAD.encode(digest(VERIFIER)); + let mut pairs = vec![ + ("client_id", client), + ("redirect_uri", CALLBACK), + ("resource", RESOURCE), + ("scope", SCOPE), + ("response_type", "code"), + ("code_challenge_method", "S256"), + ("code_challenge", challenge.as_str()), + ("state", STATE), + ]; + for (key, value) in changes { + pairs.retain(|(k, _)| k != key); + pairs.push((key, value)); + } + self.request( + "GET", + &format!("/oauth/authorize?{}", encoded(&pairs)), + String::new(), + &[], + ) + .await + } + + fn browser(reply: &Reply) -> (String, String) { + assert_eq!(reply.status, StatusCode::OK, "{}", reply.text); + let cookie = reply.headers[header::SET_COOKIE] + .to_str() + .unwrap() + .split(';') + .next() + .unwrap() + .to_owned(); + let csrf = reply + .text + .split("name=\"csrf\" value=\"") + .nth(1) + .unwrap() + .split('"') + .next() + .unwrap() + .to_owned(); + (cookie, csrf) + } + + async fn consent( + &self, + cookie: &str, + csrf: &str, + secret: &str, + decision: &str, + origin: Option<&str>, + ) -> Reply { + let mut headers = vec![ + ("content-type", "application/x-www-form-urlencoded"), + ("cookie", cookie), + ]; + if let Some(origin) = origin { + headers.push(("origin", origin)); + } + self.request( + "POST", + "/oauth/authorize", + encoded(&[ + ("csrf", csrf), + ("owner_secret", secret), + ("decision", decision), + ]), + &headers, + ) + .await + } + + async fn code(&self, client: &str) -> String { + let page = self.authorize(client, &[]).await; + let (cookie, csrf) = Self::browser(&page); + let response = self + .consent(&cookie, &csrf, SECRET, "approve", Some(ISSUER)) + .await; + assert_eq!(response.status, StatusCode::OK); + assert_eq!(response.parameter("iss"), ISSUER); + assert_eq!(response.parameter("state"), STATE); + assert_eq!(response.parameter("existing"), "1"); + response.parameter("code") + } + + async fn refresh_client(&self) -> String { + let reply = self + .register(json!({ + "redirect_uris": [CALLBACK], "client_name": "ChatGPT", "scope": SCOPE, + "token_endpoint_auth_method": "none", + "grant_types": ["authorization_code", "refresh_token"], "response_types": ["code"], + "client_uri": "https://client.example/", "logo_uri": "https://client.example/icon", + "application_type": "web", "unknown_metadata": {"safe": [1, 2]} + })) + .await; + assert_eq!(reply.status, StatusCode::CREATED); + assert_eq!( + reply.json()["grant_types"], + json!(["authorization_code", "refresh_token"]) + ); + assert_eq!(reply.json()["token_endpoint_auth_method"], "none"); + reply.json()["client_id"].as_str().unwrap().to_owned() + } + + async fn refresh(&self, client: &str, token: &str, changes: &[(&str, &str)]) -> Reply { + let mut pairs = vec![ + ("grant_type", "refresh_token"), + ("client_id", client), + ("refresh_token", token), + ]; + for (key, value) in changes { + pairs.retain(|(k, _)| k != key); + pairs.push((key, value)); + } + self.request( + "POST", + "/oauth/token", + encoded(&pairs), + &[("content-type", "application/x-www-form-urlencoded")], + ) + .await + } + + async fn initial_tokens(&self, client: &str) -> Value { + let code = self.code(client).await; + let reply = self.exchange(client, &code, &[]).await; + assert_eq!(reply.status, StatusCode::OK); + reply.json() + } + + async fn mcp_status(&self, access: &str) -> StatusCode { + let bearer = format!("Bearer {access}"); + let initialize = json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{ + "protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test","version":"1"}}}); + self.request( + "POST", + "/mcp", + initialize.to_string(), + &[ + ("authorization", &bearer), + ("content-type", "application/json"), + ("accept", "application/json, text/event-stream"), + ], + ) + .await + .status + } + + async fn exchange(&self, client: &str, code: &str, changes: &[(&str, &str)]) -> Reply { + let mut pairs = vec![ + ("grant_type", "authorization_code"), + ("client_id", client), + ("code", code), + ("redirect_uri", CALLBACK), + ("resource", RESOURCE), + ("code_verifier", VERIFIER), + ]; + for (key, value) in changes { + pairs.retain(|(k, _)| k != key); + pairs.push((key, value)); + } + self.request( + "POST", + "/oauth/token", + encoded(&pairs), + &[("content-type", "application/x-www-form-urlencoded")], + ) + .await + } +} + +#[tokio::test] +async fn token_protects_real_mcp_transport_without_origin_stripping() { + let h = Harness::new(); + let unauthenticated = h.request("POST", "/mcp", String::new(), &[]).await; + assert_eq!(unauthenticated.status, StatusCode::UNAUTHORIZED); + assert_eq!( + unauthenticated.headers[header::WWW_AUTHENTICATE], + format!( + "Bearer resource_metadata=\"{ISSUER}/.well-known/oauth-protected-resource/mcp\", scope=\"{SCOPE}\"" + ) + ); + let client = h.client().await; + let tokens = h.initial_tokens(&client).await; + let access = tokens["access_token"].as_str().unwrap(); + let bearer = format!("Bearer {access}"); + let forbidden = h + .request( + "POST", + "/mcp", + json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{ + "protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test","version":"1"}}}).to_string(), + &[ + ("authorization", &bearer), + ("origin", ISSUER), + ("content-type", "application/json"), + ("accept", "application/json, text/event-stream"), + ], + ) + .await; + assert_eq!(forbidden.status, StatusCode::FORBIDDEN); + assert_eq!(h.mcp_status(access).await, StatusCode::OK); + for reply in [unauthenticated, forbidden] { + let output = format!("{:?}{}", reply.headers, reply.text); + for credential in [access, SECRET, VERIFIER] { + assert!(!output.contains(credential)); + } + } +} + +#[tokio::test] +async fn discovery_and_all_route_host_origin_guards() { + let h = Harness::new(); + for path in [ + "/.well-known/oauth-protected-resource", + "/.well-known/oauth-protected-resource/mcp", + ] { + let r = h.request("GET", path, String::new(), &[]).await; + assert_eq!(r.status, StatusCode::OK); + assert_eq!(r.json()["resource"], RESOURCE); + assert_eq!(r.json()["authorization_servers"], json!([ISSUER])); + } + let r = h + .request( + "GET", + "/.well-known/oauth-authorization-server", + String::new(), + &[], + ) + .await; + let metadata = r.json(); + assert_eq!(metadata["issuer"], ISSUER); + assert_eq!( + metadata["code_challenge_methods_supported"], + json!(["S256"]) + ); + assert_eq!( + metadata["grant_types_supported"], + json!(["authorization_code", "refresh_token"]) + ); + assert_eq!( + metadata["token_endpoint_auth_methods_supported"], + json!(["none"]) + ); + assert_eq!( + metadata["authorization_response_iss_parameter_supported"], + true + ); + assert!(metadata.get("refresh_token").is_none()); + for (key, path) in [ + ("authorization_endpoint", "/oauth/authorize"), + ("token_endpoint", "/oauth/token"), + ("registration_endpoint", "/oauth/register"), + ] { + assert_eq!(metadata[key], format!("{ISSUER}{path}")); + } + for path in [ + "/mcp", + "/unknown", + "/oauth/authorize", + "/oauth/register", + "/oauth/token", + "/.well-known/oauth-authorization-server", + "/.well-known/oauth-protected-resource", + ] { + let r = h + .request("GET", path, String::new(), &[("host", "evil.example")]) + .await; + assert_eq!(r.status, StatusCode::FORBIDDEN, "{path}"); + let r = h + .request("GET", path, String::new(), &[("origin", ISSUER)]) + .await; + assert_eq!(r.status, StatusCode::FORBIDDEN, "{path}"); + } + let r = h + .request( + "POST", + "/oauth/token?owner_secret=not-allowed", + String::new(), + &[], + ) + .await; + assert_eq!(r.status, StatusCode::BAD_REQUEST); +} + +#[tokio::test] +async fn dcr_validation_body_bounds_and_rate_limit() { + for value in [ + json!({"redirect_uris": []}), + json!({"redirect_uris": ["http://remote.example/cb"]}), + json!({"redirect_uris": ["https://u:p@client.example/cb"]}), + json!({"redirect_uris": ["https://client.example/cb#x"]}), + json!({"redirect_uris": ["https://client.example/cb?code=x"]}), + json!({"redirect_uris": [CALLBACK], "token_endpoint_auth_method": "client_secret_post"}), + json!({"redirect_uris": [CALLBACK], "grant_types": ["refresh_token"]}), + json!({"redirect_uris": [CALLBACK], "grant_types": []}), + json!({"redirect_uris": [CALLBACK], "grant_types": ["authorization_code", "authorization_code"]}), + json!({"redirect_uris": [CALLBACK], "grant_types": ["authorization_code", "refresh_token", "refresh_token"]}), + json!({"redirect_uris": [CALLBACK], "grant_types": ["authorization_code", "client_credentials"]}), + json!({"redirect_uris": [CALLBACK], "scope": "other"}), + json!({"redirect_uris": [CALLBACK], "response_types": ["token"]}), + json!({"redirect_uris": [CALLBACK], "client_name": "x".repeat(129)}), + json!({"redirect_uris": vec![CALLBACK; 9]}), + ] { + let h = Harness::new(); + let r = h.register(value).await; + assert_eq!(r.status, StatusCode::BAD_REQUEST, "{}", r.text); + assert!(h.auth.0.store.lock().unwrap().clients.is_empty()); + } + for uri in [ + "http://localhost:3210/callback", + "http://127.0.0.1:3210/callback", + "http://[::1]:3210/callback", + CALLBACK, + ] { + assert!(valid_redirect(uri), "{uri}"); + } + for uri in [ + "http://localhost.evil/cb", + "http://192.168.0.1/cb", + "javascript:alert(1)", + "https://client.example", + "https://client.example/cb?iss=x", + ] { + assert!(!valid_redirect(uri), "{uri}"); + } + let h = Harness::new(); + let r = h + .request( + "POST", + "/oauth/register", + "x".repeat(MAX_BODY + 1), + &[("content-type", "application/json")], + ) + .await; + assert_eq!(r.status, StatusCode::PAYLOAD_TOO_LARGE); + for _ in 0..16 { + h.client().await; + } + assert_eq!( + h.register(json!({"redirect_uris": [CALLBACK]})) + .await + .status, + StatusCode::TOO_MANY_REQUESTS + ); +} + +#[tokio::test] +async fn owner_login_consent_csrf_and_origin_are_mandatory() { + let h = Harness::new(); + let client = h.client().await; + let page = h.authorize(&client, &[]).await; + assert!( + page.text + .contains("<script>untrusted & client</script>") + ); + assert!(!page.text.contains(""), + ], + ) + .await; + Harness::browser(&page); + assert!( + page.text + .contains("Only approve a client and callback you trust.") + ); + assert!(page.text.contains(&format!( + "Approve grants all QueryMT tools on {RESOURCE} (scope {SCOPE})." + ))); + assert_eq!( + page.text + .contains("24 hours from the initial code exchange"), + refresh + ); + assert!(page.text.contains("30 minutes")); + assert!(!page.text.contains("", + listener.local_addr().unwrap() + ); + let reply = h + .register(json!({"redirect_uris":[CALLBACK],"logo_uri":url,"client_uri":url})) + .await; + let client = reply.json()["client_id"].as_str().unwrap().to_owned(); + let page = h.authorize(&client, &[("future_extension", &url)]).await; + assert!(!page.text.contains("ignored")); + let code = h.code(&client).await; + assert_eq!( + h.exchange(&client, &code, &[("future_extension", &url)]) + .await + .status, + StatusCode::OK + ); + assert!( + tokio::time::timeout(Duration::from_millis(50), listener.accept()) + .await + .is_err() + ); + for key in ["client_secret", "client_assertion", "client_assertion_type"] { + assert_eq!( + h.authorize(&client, &[(key, "unsupported")]).await.status, + StatusCode::BAD_REQUEST + ); + let code = h.code(&client).await; + assert_eq!( + h.exchange(&client, &code, &[(key, "unsupported")]) + .await + .status, + StatusCode::BAD_REQUEST + ); + } +} + +#[tokio::test] +async fn invalid_refresh_bindings_and_malformed_requests_preserve_current_token() { + let h = Harness::new(); + let client = h.refresh_client().await; + let other = h.refresh_client().await; + let tokens = h.initial_tokens(&client).await; + let raw = tokens["refresh_token"].as_str().unwrap(); + for (key, value, expected) in [ + ("client_id", other.as_str(), "invalid_grant"), + ("resource", "https://other.example/mcp", "invalid_target"), + ("scope", "qmcp:access other", "invalid_scope"), + ("refresh_token", "unknown", "invalid_grant"), + ("client_secret", "", "invalid_grant"), + ("code", "mixed-grants", "invalid_grant"), + ] { + assert_eq!( + h.refresh(&client, raw, &[(key, value)]).await.json()["error"], + expected + ); + let store = h.auth.0.store.lock().unwrap(); + assert!(!store.refresh[&digest(raw)].consumed && !store.families[&digest(raw)].revoked); + } + let body = encoded(&[ + ("grant_type", "refresh_token"), + ("client_id", &client), + ("refresh_token", raw), + ]); + for malformed in [ + format!("{body}&refresh_token={raw}"), + format!("{body}&broken=%GG"), + ] { + assert_eq!( + h.request( + "POST", + "/oauth/token", + malformed, + &[("content-type", "application/x-www-form-urlencoded")] + ) + .await + .status, + StatusCode::BAD_REQUEST + ); + } + let rotated = h + .refresh(&client, raw, &[("resource", RESOURCE), ("scope", SCOPE)]) + .await + .json(); + for changes in [ + vec![("client_id", other.as_str())], + vec![("resource", "wrong")], + vec![("scope", "wrong")], + ] { + assert_eq!( + h.refresh(&client, raw, &changes).await.status, + StatusCode::BAD_REQUEST + ); + assert!(!h.auth.0.store.lock().unwrap().families[&digest(raw)].revoked); + } + assert_eq!( + h.refresh(&client, rotated["refresh_token"].as_str().unwrap(), &[]) + .await + .status, + StatusCode::OK + ); +} + +#[tokio::test] +async fn expiry_and_restart_invalidate_grants_and_mcp_access() { + let h = Harness::new(); + let client = h.refresh_client().await; + let (cookie, csrf) = Harness::browser(&h.authorize(&client, &[]).await); + for transaction in h.auth.0.store.lock().unwrap().transactions.values_mut() { + transaction.expires = Instant::now(); + } + assert_eq!( + h.consent(&cookie, &csrf, SECRET, "approve", Some(ISSUER)) + .await + .status, + StatusCode::BAD_REQUEST + ); + let code = h.code(&client).await; + h.auth + .0 + .store + .lock() + .unwrap() + .codes + .get_mut(&digest(&code)) + .unwrap() + .expires = Instant::now(); + assert_eq!( + h.exchange(&client, &code, &[]).await.status, + StatusCode::BAD_REQUEST + ); + let tokens = h.initial_tokens(&client).await; + let access = tokens["access_token"].as_str().unwrap(); + let refresh = tokens["refresh_token"].as_str().unwrap(); + assert_eq!(h.mcp_status("bad").await, StatusCode::UNAUTHORIZED); + assert_eq!(h.mcp_status(access).await, StatusCode::OK); + h.auth + .0 + .store + .lock() + .unwrap() + .tokens + .get_mut(&digest(access)) + .unwrap() + .expires = Instant::now(); + assert_eq!(h.mcp_status(access).await, StatusCode::UNAUTHORIZED); + let deadline = Instant::now() + Duration::from_secs(10); + h.auth + .0 + .store + .lock() + .unwrap() + .families + .get_mut(&digest(refresh)) + .unwrap() + .expires = deadline; + let rotated = h.refresh(&client, refresh, &[]).await.json(); + assert!(rotated["expires_in"].as_u64().unwrap() <= 10); + { + let mut store = h.auth.0.store.lock().unwrap(); + assert!(store.tokens.values().all(|token| token.expires <= deadline)); + store.prune(deadline); + assert!(store.families.is_empty() && store.refresh.is_empty() && store.tokens.is_empty()); + } + assert_eq!( + h.refresh(&client, refresh, &[]).await.status, + StatusCode::BAD_REQUEST + ); + let restarted = Harness::new(); + assert_eq!( + restarted + .refresh(&client, rotated["refresh_token"].as_str().unwrap(), &[]) + .await + .status, + StatusCode::BAD_REQUEST + ); + assert_eq!( + restarted + .mcp_status(rotated["access_token"].as_str().unwrap()) + .await, + StatusCode::UNAUTHORIZED + ); + assert_eq!( + restarted.authorize(&client, &[]).await.status, + StatusCode::BAD_REQUEST + ); +} + +#[tokio::test] +async fn refresh_capacity_failures_preserve_credentials_and_replay_evidence() { + let h = Harness::new(); + let client = h.refresh_client().await; + // Exercise the per-client family bound with real grants, not synthetic matrices. + let mut issued = Vec::new(); + for _ in 0..MAX_CLIENT_FAMILIES { + issued.push(h.initial_tokens(&client).await); + } + // Start another owner-attempt window without waiting a minute in the test. + h.auth + .0 + .store + .lock() + .unwrap() + .clients + .get_mut(&client) + .unwrap() + .login_rate = Rate::default(); + let code = h.code(&client).await; + assert_eq!( + h.exchange(&client, &code, &[]).await.status, + StatusCode::SERVICE_UNAVAILABLE + ); + assert!( + h.auth + .0 + .store + .lock() + .unwrap() + .codes + .contains_key(&digest(&code)) + ); + let old = issued[0]["refresh_token"].as_str().unwrap(); + let current = h.refresh(&client, old, &[]).await.json(); + let raw = current["refresh_token"].as_str().unwrap(); + { + let mut store = h.auth.0.store.lock().unwrap(); + store.clients.get_mut(&client).unwrap().exchange_rate.used = 60; + } + assert_eq!( + h.refresh(&client, raw, &[]).await.status, + StatusCode::TOO_MANY_REQUESTS + ); + { + let mut store = h.auth.0.store.lock().unwrap(); + store.clients.get_mut(&client).unwrap().exchange_rate = Rate::default(); + store.families.get_mut(&digest(old)).unwrap().records = MAX_FAMILY_RECORDS; + } + assert_eq!( + h.refresh(&client, raw, &[]).await.status, + StatusCode::SERVICE_UNAVAILABLE + ); + { + let store = h.auth.0.store.lock().unwrap(); + assert!(!store.refresh[&digest(raw)].consumed); + assert!(store.refresh[&digest(old)].consumed); + } + assert_eq!( + h.refresh(&client, old, &[]).await.json()["error"], + "invalid_grant" + ); + assert_eq!( + h.mcp_status(current["access_token"].as_str().unwrap()) + .await, + StatusCode::UNAUTHORIZED + ); +} diff --git a/src/main.rs b/src/main.rs index 94f386e..bd8c8c0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -47,12 +47,16 @@ struct Cli { long, env = "QMCP_BEARER_TOKEN", hide_env_values = true, - conflicts_with = "oauth_config" + conflicts_with_all = ["oauth_config", "oauth_local_resource"] )] bearer_token: Option, /// OAuth resource/provider JSON configuration (HTTP only; see README). - #[arg(long, requires = "oauth_client_secret")] + #[arg( + long, + requires = "oauth_client_secret", + conflicts_with = "oauth_local_resource" + )] oauth_config: Option, /// Introspection client secret. Prefer the environment variable. @@ -64,14 +68,38 @@ struct Cli { )] oauth_client_secret: Option, + /// Embedded single-owner OAuth resource: canonical public HTTPS /mcp URL (HTTP only). + #[arg(long, requires = "oauth_owner_secret_file", conflicts_with_all = ["oauth_config", "bearer_token"])] + oauth_local_resource: Option, + + /// Absolute path to a generated owner credential, private regular file mode 0600. + #[arg(long, requires = "oauth_local_resource")] + oauth_owner_secret_file: Option, + /// Additional allowed HTTP hostname or host:port (repeatable; requires auth). #[arg(long = "allowed-host", action = clap::ArgAction::Append)] allowed_hosts: Vec, } +impl Cli { + fn validate(&self) -> Result<()> { + if matches!(self.transport, Transport::Stdio) + && (self.bearer_token.is_some() + || self.oauth_config.is_some() + || self.oauth_local_resource.is_some() + || self.oauth_owner_secret_file.is_some() + || !self.allowed_hosts.is_empty()) + { + bail!("HTTP authentication options require --transport http"); + } + Ok(()) + } +} + #[tokio::main] async fn main() -> Result<()> { let cli = Cli::parse(); + cli.validate()?; tracing_subscriber::fmt() .with_env_filter( EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")), @@ -85,21 +113,19 @@ async fn main() -> Result<()> { cli.permission, )?)?; match cli.transport { - Transport::Stdio => { - if cli.bearer_token.is_some() - || cli.oauth_config.is_some() - || !cli.allowed_hosts.is_empty() - { - bail!("HTTP authentication options require --transport http"); - } - transport::serve_stdio(server).await - } + Transport::Stdio => transport::serve_stdio(server).await, Transport::Http => { let token = cli .bearer_token .map(|value| transport::BearerToken::new(&value)) .transpose()?; - let router = if let Some(path) = cli.oauth_config { + let router = if let Some(resource) = cli.oauth_local_resource { + let path = cli + .oauth_owner_secret_file + .context("owner secret file is required")?; + let auth = qmcp::local_oauth::LocalOAuth::from_secret_file(&resource, &path)?; + transport::http_router_with_local_oauth(server, cli.bind, auth, &cli.allowed_hosts)? + } else if let Some(path) = cli.oauth_config { let bytes = tokio::fs::read(path) .await .context("read OAuth configuration")?; @@ -161,6 +187,43 @@ async fn shutdown_signal() { mod tests { use super::*; + #[test] + fn embedded_oauth_rejects_orphan_conflicting_and_stdio_options() { + let local = [ + "--oauth-local-resource", + "https://mcp.example.com/mcp", + "--oauth-owner-secret-file", + "/example/owner-secret", + ]; + for suffix in [ + vec!["--bearer-token", "test-static-token"], + vec![ + "--oauth-config", + "oauth.json", + "--oauth-client-secret", + "test-provider-secret", + ], + ] { + let args: Vec<_> = + [vec!["qmcp", "--transport", "http"], local.to_vec(), suffix].concat(); + assert!(Cli::try_parse_from(args).is_err()); + } + for args in [ + vec!["qmcp", "--oauth-owner-secret-file", "/example/owner-secret"], + vec![ + "qmcp", + "--oauth-local-resource", + "https://mcp.example.com/mcp", + ], + ] { + assert!(Cli::try_parse_from(args).is_err()); + } + let args = [vec!["qmcp"], local.to_vec()].concat(); + assert!(Cli::try_parse_from(args).unwrap().validate().is_err()); + let args = [vec!["qmcp", "--transport", "http"], local.to_vec()].concat(); + assert!(Cli::try_parse_from(args).unwrap().validate().is_ok()); + } + #[test] fn oauth_configuration_requires_a_secret_and_cannot_fall_back_to_a_static_token() { for args in [ diff --git a/src/oauth.rs b/src/oauth.rs index 7bb3a89..c42ecda 100644 --- a/src/oauth.rs +++ b/src/oauth.rs @@ -496,7 +496,7 @@ mod tests { .as_array() .unwrap() .len(), - 25 + 26 ); current.lock().unwrap()["active"] = json!(false); let response = call(Some("active-token")).await.unwrap(); diff --git a/src/server.rs b/src/server.rs index 605d8e5..4265efe 100644 --- a/src/server.rs +++ b/src/server.rs @@ -45,6 +45,7 @@ impl BridgeConfig { pub struct QmcpServer { config: Arc, tools: Arc>, + workspaces: Arc, } impl QmcpServer { @@ -53,6 +54,7 @@ impl QmcpServer { Ok(Self { config: Arc::new(config), tools: Arc::new(tool_catalog()?), + workspaces: Arc::new(crate::workspace::Workspaces::new()?), }) } @@ -92,7 +94,7 @@ impl QmcpServer { let mut log = EventLog::default(); let mut state = OperationState::default(); let result = { - let operation = commands::run(&client, command, &mut state, deadline); + let operation = commands::run(&client, command, &mut state, deadline, &self.workspaces); tokio::pin!(operation); loop { tokio::select! { @@ -117,7 +119,7 @@ impl ServerHandler for QmcpServer { fn get_info(&self) -> ServerConfig { ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) .with_server_info(Implementation::new("qmcp", env!("CARGO_PKG_VERSION"))) - .with_instructions("Control an existing QueryMT ACP WebSocket server. Discover capabilities, profiles, models, and session ids rather than inventing values. Keep returned sessionId values. Pass cwd for session operations. Timeouts and MCP cancellation do not cancel remote work: check runtime before retrying. Only close sessions when explicitly requested. Streaming ACP events are collected in each tool result.") + .with_instructions("Control an existing QueryMT ACP WebSocket server. Discover capabilities, profiles, models, and session ids rather than inventing values. Keep returned sessionId values and pass them as session_id; there is no global active workspace. Use workspace with an explicit initial cwd to select without prompting. cwd is sticky per session across calls. workspace or prompt with a different cwd creates a NEW QueryMT session; use its new sessionId, not previousSessionId. Read-only status/inspect/follow never rebind or create sessions. Timeouts and MCP cancellation do not cancel remote work: check runtime before retrying. Only close sessions when explicitly requested. Streaming ACP events are collected in each tool result.") } async fn list_tools( @@ -196,10 +198,13 @@ fn tool_catalog() -> Result> { name.as_str(), "set_lifecycle" | "pin" | "unpin" | "set_mode" | "set_model" | "set_effort" ); - let annotations = ToolAnnotations::new() + let mut annotations = ToolAnnotations::new() .read_only(read_only) .idempotent(idempotent) .open_world(true); + if name == "workspace" { + annotations = annotations.destructive(false); + } tools.push(Tool::new(name, description, input).with_annotations(annotations)); } if tools.is_empty() { @@ -274,6 +279,9 @@ mod tests { json!({"protocolVersion": 1, "agentCapabilities": {"loadSession": true, "sessionCapabilities": {"resume": {}}}}) } "session/resume" | "session/load" => json!({}), + "session/list" => { + json!({"sessions": [{"sessionId": "fixture-session", "cwd": "/workspace"}]}) + } "_querymt/session/runtimeState" if input_id.is_null() => { json!({"phase": "tools", "steerable": true}) } @@ -336,6 +344,100 @@ mod tests { ); } + #[tokio::test] + async fn cancelled_creation_keeps_workspace_binding_across_server_clones() { + use futures_util::{SinkExt, StreamExt}; + use std::time::Duration; + use tokio_tungstenite::{accept_async, tungstenite::Message}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("ws://{}/acp/ws", listener.local_addr().unwrap()); + let (ready_tx, ready_rx) = tokio::sync::oneshot::channel(); + let peer = tokio::spawn(async move { + let mut ready = Some(ready_tx); + for _ in 0..2 { + let (socket, _) = listener.accept().await.unwrap(); + let mut socket = accept_async(socket).await.unwrap(); + while let Some(Ok(message)) = socket.next().await { + let Message::Text(text) = message else { + continue; + }; + let request: Value = serde_json::from_str(&text).unwrap(); + let result = match request["method"].as_str().unwrap() { + "initialize" => { + json!({"protocolVersion": 1, "agentCapabilities": {"loadSession": true, "sessionCapabilities": {"resume": {}}}}) + } + "session/new" => { + assert_eq!(request["params"]["cwd"], "/remote/cancelled-creation"); + json!({"sessionId": "created-before-cancellation"}) + } + "session/set_mode" => { + // Creation was acknowledged and configuration now stalls. + ready.take().unwrap().send(()).unwrap(); + continue; + } + "session/resume" | "session/load" => { + assert_eq!(request["params"]["cwd"], "/remote/cancelled-creation"); + assert_eq!( + request["params"]["sessionId"], + "created-before-cancellation" + ); + json!({}) + } + _ => panic!( + "no rediscovery, implicit cancel, close, or extra creation: {request}" + ), + }; + socket + .send(Message::Text( + json!({"jsonrpc": "2.0", "id": request["id"], "result": result}) + .to_string() + .into(), + )) + .await + .unwrap(); + } + } + }); + let server = + QmcpServer::new(BridgeConfig::new(&url, false, PermissionPolicy::Cancel).unwrap()) + .unwrap(); + let clone = server.clone(); + assert!(Arc::ptr_eq(&server.workspaces, &clone.workspaces)); + let command = commands::decode("prompt", json!({"new": true, "cwd": "/remote/cancelled-creation", "text": "not sent", "mode": "stall"}).as_object().unwrap().clone()).unwrap(); + let result = tokio::time::timeout( + Duration::from_secs(2), + server.execute_with_cancellation(command, async { + ready_rx.await.unwrap(); + }), + ) + .await + .unwrap(); + assert_eq!(result["ok"], false, "{result}"); + assert_eq!( + result["recovery"]["sessionId"], + "created-before-cancellation" + ); + assert_eq!(result["recovery"]["sessionCreated"], true); + let command = commands::decode( + "status", + json!({"session_id": "created-before-cancellation"}) + .as_object() + .unwrap() + .clone(), + ) + .unwrap(); + let result = clone + .execute_with_cancellation(command, std::future::pending()) + .await; + assert_eq!(result["ok"], true, "{result}"); + assert_eq!(result["cwd"], "/remote/cancelled-creation"); + tokio::time::timeout(Duration::from_secs(2), peer) + .await + .unwrap() + .unwrap(); + } + #[test] fn advertised_schemas_match_decoder_boundaries() { let tools = tool_catalog().unwrap(); @@ -454,6 +556,23 @@ mod tests { ] { check("exec", json!({"calls": [entry]})); } + for arguments in [ + json!({}), + json!({"session_id": null}), + json!({"cwd": null}), + json!({"cwd": "/remote"}), + json!({"session_id": "s"}), + json!({"session_id": "s", "cwd": null}), + json!({"session_id": "s", "cwd": "/remote", "profile": "p"}), + json!({"session_id": "s", "extra": true}), + ] { + check("workspace", arguments); + } + check( + "exec", + json!({"calls": [{"name": "workspace", "arguments": {"cwd": "/remote"}}]}), + ); + check("exec", json!({"calls": [{"name": "workspace"}]})); check("prompt", json!({"new": true, "text": "x", "extra": true})); check("caps", json!({"extra": true})); } @@ -461,7 +580,7 @@ mod tests { #[test] fn catalog_contains_self_contained_object_schemas_and_native_prompt_fields() { let tools = tool_catalog().unwrap(); - assert_eq!(tools.len(), 25); + assert_eq!(tools.len(), 26); for tool in &tools { assert_eq!(tool.input_schema["type"], "object"); assert!( @@ -476,6 +595,11 @@ mod tests { .contains("\"$ref\"") ); } + let workspace = tools.iter().find(|t| t.name == "workspace").unwrap(); + let annotations = json!(workspace.annotations); + assert_eq!(annotations["readOnlyHint"], false); + assert_eq!(annotations["destructiveHint"], false); + assert_eq!(annotations["idempotentHint"], false); let prompt = tools.iter().find(|t| t.name == "prompt").unwrap(); assert_eq!(prompt.input_schema["properties"]["text"]["type"], "string"); assert_eq!( diff --git a/src/transport.rs b/src/transport.rs index 5e428a6..f706e7f 100644 --- a/src/transport.rs +++ b/src/transport.rs @@ -159,6 +159,24 @@ pub fn http_router_with_oauth( .route("/.well-known/oauth-protected-resource/mcp", discovery)) } +/// Embedded single-owner OAuth. Its outer guard covers AS, MCP and fallback routes; +/// only the browser consent POST accepts Origin, never the MCP transport. +pub fn http_router_with_local_oauth( + server: QmcpServer, + bind: SocketAddr, + auth: crate::local_oauth::LocalOAuth, + allowed_hosts: &[String], +) -> Result { + let mut hosts = allowed_hosts.to_vec(); + hosts.push(auth.authority().into()); + let config = http_config(bind, true, &hosts)?; + let mcp = build_http_router(server, config).layer(middleware::from_fn_with_state( + auth.clone(), + crate::local_oauth::require_token, + )); + Ok(auth.guard(mcp.merge(auth.routes()))) +} + async fn guard_discovery( State(hosts): State>, request: Request, diff --git a/src/workspace.rs b/src/workspace.rs new file mode 100644 index 0000000..12ae971 --- /dev/null +++ b/src/workspace.rs @@ -0,0 +1,665 @@ +//! Process-local, immutable session workspace bindings. No lock spans an ACP request. +use std::{ + collections::{HashMap, HashSet}, + path::{Component, Path, PathBuf}, + sync::Mutex, +}; + +use qmacp::{client::AcpClient, error::CliError, protocol::schema::v1 as acp, session}; +use serde_json::{Value, json}; + +use crate::output::OperationState; + +const CAPACITY: usize = 4096; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct Binding { + pub cwd: PathBuf, + pub profile: Option, +} + +struct CachedBinding { + binding: Binding, + // Entries are never removed: insertion order is bounded by store capacity, + // not by the (unbounded) number of operations or cancelled reservations. + generation: usize, + creation_claimed: bool, +} + +#[derive(Default)] +struct Store { + bindings: HashMap, + reserved: usize, +} + +pub(crate) struct Workspaces { + startup_cwd: PathBuf, + store: Mutex, + capacity: usize, +} + +pub(crate) struct Selection { + pub session_id: String, + pub binding: Binding, + pub created: bool, + pub previous: Option, + profile_source: &'static str, +} + +impl Selection { + pub fn result(&self) -> Value { + let mut value = json!({ + "sessionId": self.session_id, "cwd": self.binding.cwd, + "created": self.created, "switched": self.previous.is_some(), + "profile": self.binding.profile, "profileKnown": self.binding.profile.is_some(), + "profileSource": self.profile_source, + }); + if let Some(previous) = &self.previous { + value["previousSessionId"] = json!(previous); + } + value + } + + pub fn checkpoint(&self, state: &mut OperationState) { + state.session(&self.session_id, &self.binding.cwd, self.created); + state + .recovery + .extend(self.result().as_object().unwrap().clone()); + } +} + +// Capacity is reserved before session/new, released on rejection or cancellation, +// and committed synchronously on acknowledgement, before the next await point. +struct Reservation<'a> { + workspaces: &'a Workspaces, + generation: usize, +} + +impl Drop for Reservation<'_> { + fn drop(&mut self) { + self.workspaces.store.lock().unwrap().reserved -= 1; + } +} + +impl Workspaces { + pub fn new() -> anyhow::Result { + let startup = std::env::current_dir()?; + Ok(Self { + startup_cwd: normalize(&startup, &startup).map_err(anyhow::Error::msg)?, + store: Mutex::new(Store::default()), + capacity: CAPACITY, + }) + } + + pub fn initial_cwd(&self, cwd: Option) -> Result { + normalize( + cwd.as_deref().unwrap_or(&self.startup_cwd), + &self.startup_cwd, + ) + .map_err(CliError::Usage) + } + + fn reserve(&self) -> Result, CliError> { + let mut store = self.store.lock().unwrap(); + if store.bindings.len() + store.reserved >= self.capacity { + return Err(CliError::rpc( + "workspace binding capacity reached; no session was created", + )); + } + store.reserved += 1; + Ok(Reservation { + workspaces: self, + generation: store.bindings.len(), + }) + } + + fn insert( + &self, + id: &str, + binding: Binding, + reservation: Option<&Reservation<'_>>, + ) -> Result { + let mut store = self.store.lock().unwrap(); + if let Some(known) = store.bindings.get_mut(id) { + if reservation + .is_some_and(|r| known.generation < r.generation || known.creation_claimed) + { + return Err(CliError::rpc( + "session/new returned an already bound session id; refusing to alias or overwrite the existing workspace", + )); + } + if known.binding.cwd != binding.cwd + || matches!((&known.binding.profile, &binding.profile), (Some(a), Some(b)) if a != b) + { + return Err(CliError::rpc( + "backend returned conflicting workspace metadata for an immutable session id; binding was not changed", + )); + } + if known.binding.profile.is_none() { + known.binding.profile = binding.profile; + } + // Only matching authoritative metadata first cached after this + // reservation may be claimed. A second ACK cannot claim it again. + known.creation_claimed |= reservation.is_some(); + return Ok(known.binding.clone()); + } + if reservation.is_none() && store.bindings.len() + store.reserved >= self.capacity { + return Err(CliError::rpc("workspace binding capacity reached")); + } + let generation = store.bindings.len(); + store.bindings.insert( + id.to_owned(), + CachedBinding { + binding: binding.clone(), + generation, + creation_claimed: reservation.is_some(), + }, + ); + Ok(binding) + } + + pub async fn binding(&self, client: &AcpClient, id: &str) -> Result { + if let Some(known) = self.store.lock().unwrap().bindings.get(id) { + return Ok(known.binding.clone()); + } + client.initialize().await.map_err(CliError::from_request)?; + let mut cursor = None; + let mut seen = HashSet::new(); + for _ in 0..session::find_page_cap() { + // Equivalent to list_sessions(None, cursor, None): no cwd or metadata + // filters. Retain raw JSON because ACP's typed list decoder can omit + // malformed summaries (including a target without cwd). + let mut request = acp::ListSessionsRequest::new(); + if let Some(cursor) = cursor { + request = request.cursor(cursor); + } + let (_page, raw) = client.request_json(request).await.map_err(|error| { + CliError::rpc(format!("authoritative workspace metadata unavailable from session/list: {error}; refusing to guess cwd or create a session")) + })?; + let next_cursor = match raw.get("nextCursor") { + None | Some(Value::Null) => None, + Some(Value::String(next)) => Some(next.clone()), + Some(_) => { + return Err(CliError::rpc( + "authoritative workspace metadata unavailable: session/list nextCursor must be a string or null; discovery is incomplete, not an exhaustive not-found result", + )); + } + }; + let summaries = raw.get("sessions").and_then(Value::as_array).ok_or_else(|| CliError::rpc("authoritative workspace metadata unavailable: session/list has no sessions array"))?; + for item in summaries { + if item.get("sessionId").and_then(Value::as_str) != Some(id) { + continue; + } + let cwd = item.get("cwd").and_then(Value::as_str).map(PathBuf::from).ok_or_else(|| CliError::rpc("authoritative workspace metadata unavailable: target session cwd is absent; refusing to guess cwd (even from an explicit caller path)"))?; + if !cwd.is_absolute() { + return Err(CliError::rpc( + "authoritative session cwd is missing or not absolute; refusing to guess cwd (even from an explicit caller path)", + )); + } + let cwd = normalize(&cwd, &self.startup_cwd).map_err(|error| { + CliError::rpc(format!("authoritative session cwd is invalid: {error}")) + })?; + let profile = item + .pointer("/_meta/querymt/profile_id") + .and_then(Value::as_str) + .map(str::to_owned); + return self.insert(id, Binding { cwd, profile }, None); + } + match next_cursor { + None => { + return Err(CliError::rpc(format!( + "session {id} not found in authoritative session/list; no workspace was selected or created" + ))); + } + Some(next) if !seen.insert(next.clone()) => { + return Err(CliError::rpc( + "workspace metadata discovery incomplete: session/list repeated a cursor; no workspace was selected or created", + )); + } + Some(next) => cursor = Some(next), + } + } + Err(CliError::rpc( + "workspace metadata discovery incomplete: session/list page cap reached; session may exist beyond the scan; no workspace was selected or created", + )) + } + + pub async fn read_cwd( + &self, + client: &AcpClient, + id: &str, + cwd: Option, + ) -> Result { + let binding = self.binding(client, id).await?; + if let Some(cwd) = cwd { + let requested = normalize(&cwd, &binding.cwd).map_err(CliError::Usage)?; + if requested != binding.cwd { + return Err(CliError::Usage("read-only session operations cannot change cwd; use workspace with this session_id and cwd to create a NEW session, then use its returned sessionId".into())); + } + } + Ok(binding.cwd) + } + + pub async fn create( + &self, + client: &AcpClient, + state: &mut OperationState, + binding: Binding, + previous: Option, + profile_source: &'static str, + ) -> Result<(Selection, acp::NewSessionResponse, Value), CliError> { + let reservation = self.reserve()?; + client.initialize().await.map_err(CliError::from_request)?; + let (created, raw) = client + .new_session(binding.cwd.clone(), binding.profile.as_deref()) + .await + .map_err(CliError::from_request)?; + let session_id = created.session_id.to_string(); + let binding = match self.insert(&session_id, binding, Some(&reservation)) { + Ok(binding) => binding, + Err(error) => { + state.recovery.insert( + "workspaceConflict".into(), + json!({ + "acknowledgedSessionId": session_id, + "selectionCommitted": false, + }), + ); + return Err(error); + } + }; + let profile_source = + if binding.profile.is_some() && profile_source == "backend-default-unknown" { + "authoritative-metadata" + } else { + profile_source + }; + let selection = Selection { + session_id, + binding, + created: true, + previous, + profile_source, + }; + selection.checkpoint(state); + drop(reservation); + Ok((selection, created, raw)) + } + + pub async fn select( + &self, + client: &AcpClient, + state: &mut OperationState, + id: Option, + cwd: Option, + profile: Option, + ) -> Result { + let Some(id) = id else { + let cwd = cwd.ok_or_else(|| { + CliError::Usage( + "workspace requires an explicit initial cwd when session_id is omitted".into(), + ) + })?; + let binding = Binding { + cwd: self.initial_cwd(Some(cwd))?, + profile, + }; + let source = if binding.profile.is_some() { + "explicit" + } else { + "backend-default-unknown" + }; + return self + .create(client, state, binding, None, source) + .await + .map(|created| created.0); + }; + let binding = self.binding(client, &id).await?; + let requested = cwd + .map(|cwd| normalize(&cwd, &binding.cwd).map_err(CliError::Usage)) + .transpose()? + .unwrap_or(binding.cwd.clone()); + if requested == binding.cwd { + if profile + .as_ref() + .is_some_and(|profile| Some(profile) != binding.profile.as_ref()) + { + return Err(CliError::Usage("profile changes on an existing session are not supported by workspace/prompt; the requested profile differs or the existing profile is unknown; create a new workspace without session_id to select a profile".into())); + } + return Ok(Selection { + session_id: id, + binding, + created: false, + previous: None, + profile_source: "remembered-or-metadata", + }); + } + let source = if profile.is_some() { + "explicit" + } else if binding.profile.is_some() { + "inherited" + } else { + "backend-default-unknown" + }; + let binding = Binding { + cwd: requested, + profile: profile.or(binding.profile), + }; + self.create(client, state, binding, Some(id), source) + .await + .map(|created| created.0) + } +} + +/// POSIX lexical paths on the ACP host, not filesystem canonicalization on qmcp. +fn normalize(path: &Path, base: &Path) -> Result { + let text = path.to_str().ok_or("cwd must be valid UTF-8")?; + if text.is_empty() || text.chars().any(char::is_control) { + return Err("cwd must not be empty or contain control characters".into()); + } + if text.starts_with('~') { + return Err("cwd cannot use ~ expansion on a remote host; provide an explicit absolute path on the QueryMT host".into()); + } + if text.contains('\\') || text.as_bytes().get(1) == Some(&b':') || text.starts_with("//") { + return Err("cwd must be an unambiguous POSIX path, not a Windows drive/UNC path".into()); + } + let joined = if path.is_absolute() { + path.to_path_buf() + } else { + base.join(path) + }; + let mut result = PathBuf::new(); + for component in joined.components() { + match component { + Component::RootDir => result.push("/"), + Component::CurDir => {} + Component::ParentDir => { + result.pop(); + } + Component::Normal(name) => result.push(name), + Component::Prefix(_) => return Err("cwd must be a POSIX path".into()), + } + } + if !result.is_absolute() { + return Err("cwd must resolve to an absolute POSIX path".into()); + } + Ok(result) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn paths_are_lexical_and_do_not_require_a_local_filesystem() { + let base = Path::new("/remote/nonexistent/project"); + for (input, expected) in [ + (".", "/remote/nonexistent/project"), + ("a/../b", "/remote/nonexistent/project/b"), + ("../other", "/remote/nonexistent/other"), + ("/a/./b/../../..", "/"), + ] { + assert_eq!( + normalize(Path::new(input), base).unwrap(), + Path::new(expected) + ); + } + for input in [ + "", + "~", + "~/project", + "~user/project", + "C:", + "C:/work", + "a\nb", + "a\0b", + "//host/work", + "C:\\work", + ] { + assert!(normalize(Path::new(input), base).is_err(), "{input:?}"); + } + } + + #[test] + fn startup_relative_paths_are_stable() { + let workspaces = Workspaces::new().unwrap(); + let expected = + normalize(Path::new("relative/../project"), &workspaces.startup_cwd).unwrap(); + assert_eq!( + workspaces + .initial_cwd(Some("relative/../project".into())) + .unwrap(), + expected + ); + assert_eq!( + workspaces.initial_cwd(None).unwrap(), + workspaces.startup_cwd + ); + } + + #[tokio::test] + async fn full_store_fails_before_any_new_session_request() { + use futures_util::StreamExt; + use qmacp::policy::PermissionPolicy; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("ws://{}/acp/ws", listener.local_addr().unwrap()); + let peer = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut socket = tokio_tungstenite::accept_async(stream).await.unwrap(); + while let Some(Ok(message)) = socket.next().await { + assert!( + !message.is_text(), + "a full workspace store must not send any ACP request" + ); + } + }); + let (events, _receiver) = tokio::sync::mpsc::channel(16); + let client = AcpClient::connect_with_events(&url, &url, PermissionPolicy::Cancel, events) + .await + .unwrap(); + let mut workspaces = Workspaces::new().unwrap(); + workspaces.capacity = 0; + let mut state = OperationState::default(); + let result = workspaces + .create( + &client, + &mut state, + Binding { + cwd: "/remote".into(), + profile: None, + }, + None, + "backend-default-unknown", + ) + .await; + assert!(matches!(result, Err(error) if error.to_string().contains("capacity"))); + assert!(state.recovery.is_empty()); + drop(client); + tokio::time::timeout(std::time::Duration::from_secs(2), peer) + .await + .unwrap() + .unwrap(); + } + + #[test] + fn reservation_generations_reconcile_once_and_refund_at_full_capacity() { + let mut workspaces = Workspaces::new().unwrap(); + workspaces.capacity = 3; + let binding = Binding { + cwd: "/remote".into(), + profile: Some("inherited".into()), + }; + let first = workspaces.reserve().unwrap(); + let second = workspaces.reserve().unwrap(); + workspaces + .insert( + "fresh", + Binding { + profile: None, + ..binding.clone() + }, + None, + ) + .unwrap(); + assert!(workspaces.reserve().is_err()); + assert_eq!( + workspaces + .insert("fresh", binding.clone(), Some(&first)) + .unwrap(), + binding + ); + drop(first); + assert!( + workspaces + .insert("fresh", binding.clone(), Some(&second)) + .is_err() + ); + drop(second); + { + let store = workspaces.store.lock().unwrap(); + assert_eq!(store.reserved, 0); + assert_eq!(store.bindings.len(), 1); + } + let reuse = workspaces.reserve().unwrap(); + let spare = workspaces.reserve().unwrap(); + assert!(workspaces.reserve().is_err()); + drop(reuse); + drop(spare); + // An ID discovered before the reservation is never a fresh creation, + // even when cwd and profile are identical and it is metadata-only. + workspaces + .insert("prebound", binding.clone(), None) + .unwrap(); + let reservation = workspaces.reserve().unwrap(); + assert!( + workspaces + .insert("prebound", binding, Some(&reservation)) + .is_err() + ); + drop(reservation); + assert_eq!(workspaces.store.lock().unwrap().reserved, 0); + } + + #[test] + fn conflicting_metadata_does_not_mutate_bindings_or_consume_creation_claim() { + let workspaces = Workspaces::new().unwrap(); + let reservation = workspaces.reserve().unwrap(); + let binding = Binding { + cwd: "/remote".into(), + profile: Some("authoritative".into()), + }; + workspaces.insert("fresh", binding.clone(), None).unwrap(); + for conflict in [ + Binding { + cwd: "/wrong".into(), + ..binding.clone() + }, + Binding { + profile: Some("wrong".into()), + ..binding.clone() + }, + ] { + assert!( + workspaces + .insert("fresh", conflict, Some(&reservation)) + .is_err() + ); + let store = workspaces.store.lock().unwrap(); + assert_eq!(store.bindings["fresh"].binding, binding); + assert!(!store.bindings["fresh"].creation_claimed); + } + assert_eq!( + workspaces + .insert( + "fresh", + Binding { + profile: None, + ..binding.clone() + }, + Some(&reservation) + ) + .unwrap(), + binding + ); + drop(reservation); + assert_eq!(workspaces.store.lock().unwrap().reserved, 0); + } + + #[tokio::test] + async fn cancellation_while_awaiting_creation_ack_releases_capacity_without_inventing_an_id() { + use futures_util::{SinkExt, StreamExt}; + use qmacp::policy::PermissionPolicy; + use tokio_tungstenite::tungstenite::Message; + tokio::time::timeout(std::time::Duration::from_secs(5), async { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("ws://{}/acp/ws", listener.local_addr().unwrap()); + let (ready, observed) = tokio::sync::oneshot::channel(); + let peer = tokio::spawn(async move { + let mut ready = Some(ready); + for attempt in 0..2 { + let (socket, _) = listener.accept().await.unwrap(); + let mut socket = tokio_tungstenite::accept_async(socket).await.unwrap(); + while let Some(Ok(message)) = socket.next().await { + let Message::Text(text) = message else { continue }; + let request: Value = serde_json::from_str(&text).unwrap(); + let result = match request["method"].as_str().unwrap() { + "initialize" => json!({"protocolVersion": 1, "agentCapabilities": {}}), + "session/new" if attempt == 0 => { + // The backend may have created an orphan, but no ID is acknowledged. + ready.take().unwrap().send(()).unwrap(); + continue; + } + "session/new" => json!({"sessionId": "acknowledged"}), + _ => panic!("no implicit remote cancellation, close, load or prompt: {request}"), + }; + socket.send(Message::Text(json!({"jsonrpc": "2.0", "id": request["id"], "result": result}).to_string().into())).await.unwrap(); + } + } + }); + let mut workspaces = Workspaces::new().unwrap(); + workspaces.capacity = 1; + let mut state = OperationState::default(); + let binding = Binding { cwd: "/remote".into(), profile: None }; + let (events, _receiver) = tokio::sync::mpsc::channel(16); + let client = AcpClient::connect_with_events(&url, &url, PermissionPolicy::Cancel, events).await.unwrap(); + { + let operation = workspaces.create(&client, &mut state, binding.clone(), None, "backend-default-unknown"); + tokio::pin!(operation); + tokio::select! { + result = &mut operation => panic!("unexpected creation ACK: {}", result.is_ok()), + _ = observed => {}, + } + } + assert!(state.recovery.is_empty()); + { + let store = workspaces.store.lock().unwrap(); + assert!(store.bindings.is_empty()); + assert_eq!(store.reserved, 0); + } + drop(client); + let (events, _receiver) = tokio::sync::mpsc::channel(16); + let client = AcpClient::connect_with_events(&url, &url, PermissionPolicy::Cancel, events).await.unwrap(); + workspaces.create(&client, &mut state, binding, Some("previous-valid".into()), "backend-default-unknown").await.unwrap(); + assert_eq!(state.recovery["sessionId"], "acknowledged"); + assert_eq!(state.recovery["sessionCreated"], true); + assert_eq!(workspaces.store.lock().unwrap().reserved, 0); + // A subsequent invalid ACK must not overwrite any prior valid checkpoint. + workspaces.capacity = 2; + let prior = state.recovery.clone(); + let error = workspaces.create(&client, &mut state, + Binding { cwd: "/fake-new".into(), profile: None }, + Some("fake-previous".into()), "backend-default-unknown").await.err().unwrap(); + let failed = state.finish(crate::output::failure(error)); + for (key, value) in prior { + assert_eq!(failed["recovery"][&key], value); + } + assert_eq!(failed["recovery"]["previousSessionId"], "previous-valid"); + assert_ne!(failed["recovery"]["previousSessionId"], "fake-previous"); + assert_eq!(failed["recovery"]["cwd"], "/remote"); + assert_eq!(failed["recovery"]["workspaceConflict"]["acknowledgedSessionId"], "acknowledged"); + assert_eq!(workspaces.store.lock().unwrap().reserved, 0); + drop(client); + peer.await.unwrap(); + }).await.expect("cancellation and capacity reuse must complete without waiting for an ACK"); + } +} diff --git a/tests/bridge.rs b/tests/bridge.rs index c3eb81a..12d10b9 100644 --- a/tests/bridge.rs +++ b/tests/bridge.rs @@ -1,4 +1,5 @@ use std::{ + collections::VecDeque, path::PathBuf, process::Stdio, sync::{ @@ -16,6 +17,7 @@ use tokio::{ io::{AsyncBufReadExt, AsyncWriteExt, BufReader, Lines}, net::{TcpListener, TcpStream}, process::{Child, ChildStdin, ChildStdout}, + sync::oneshot, task::{JoinHandle, JoinSet}, }; use tokio_tungstenite::{WebSocketStream, accept_async, tungstenite::Message}; @@ -25,6 +27,22 @@ struct MockAcp { requests: Arc>>, active: Arc, task: JoinHandle<()>, + workspace: Arc>, +} + +#[derive(Default)] +struct MockWorkspace { + created: Vec, + listing: Option>, + reject_creation: bool, + replay_error: bool, + forced_id: Option, + creation_gates: VecDeque, +} + +struct CreationGate { + observed: oneshot::Sender<()>, + release: oneshot::Receiver<()>, } #[derive(Clone, Copy, Default)] @@ -77,6 +95,8 @@ impl MockAcp { let requests = Arc::new(Mutex::new(Vec::new())); let active = Arc::new(AtomicUsize::new(0)); let recorded = requests.clone(); + let workspace = Arc::new(Mutex::new(MockWorkspace::default())); + let workspaces = workspace.clone(); let connections = active.clone(); let task = tokio::spawn(async move { let mut tasks = JoinSet::new(); @@ -87,12 +107,13 @@ impl MockAcp { let recorded = recorded.clone(); let active = connections.clone(); let payload = extension_payload.clone(); + let workspace = workspaces.clone(); tasks.spawn(async move { active.fetch_add(1, Ordering::SeqCst); let _guard = ConnectionGuard(active); tokio::time::sleep(delay).await; let socket = accept_async(stream).await.unwrap(); - serve_acp(socket, recorded, behavior, listing_pages, payload).await; + serve_acp(socket, recorded, behavior, listing_pages, payload, workspace).await; }); } Some(result) = tasks.join_next(), if !tasks.is_empty() => { result.unwrap(); } @@ -104,6 +125,7 @@ impl MockAcp { requests, active, task, + workspace, } } @@ -147,6 +169,7 @@ async fn serve_acp( behavior: SubmissionBehavior, listing_pages: Option, extension_payload: Option, + workspace: Arc>, ) { let mut submitted = false; while let Some(Ok(message)) = socket.next().await { @@ -163,6 +186,10 @@ async fn serve_acp( continue; } let params = &request["params"]; + if method == "session/load" && workspace.lock().unwrap().replay_error { + send(&mut socket, json!({"jsonrpc":"2.0","id":id,"error":{"code":-32000,"message":"Loaded session has no runtime ref for ACP replay"}})).await; + continue; + } if matches!( method, "_querymt/profiles" @@ -181,6 +208,38 @@ async fn serve_acp( .await; continue; } + if method == "session/new" && workspace.lock().unwrap().reject_creation { + send(&mut socket, json!({"jsonrpc": "2.0", "id": id, "error": {"code": -32000, "message": "creation rejected"}})).await; + continue; + } + if method == "session/list" { + let listing = workspace.lock().unwrap().listing.clone(); + if let Some(pages) = listing { + assert!( + params.get("cwd").is_none(), + "workspace discovery must be unfiltered" + ); + assert!( + params.get("_meta").is_none(), + "workspace discovery must include all lifecycles" + ); + let page = params["cursor"] + .as_str() + .and_then(|cursor| cursor.strip_prefix("page-")) + .and_then(|n| n.parse::().ok()) + .unwrap_or(1) + - 1; + let reply = &pages[page.min(pages.len() - 1)]; + let mut response = json!({"jsonrpc": "2.0", "id": id}); + response[if reply.get("error").is_some() { + "error" + } else { + "result" + }] = reply.get("error").unwrap_or(reply).clone(); + send(&mut socket, response).await; + continue; + } + } let result = match method { "initialize" => { json!({"protocolVersion": 1, "agentInfo": {"name": "fixture", "version": "1"}, "agentCapabilities": {"loadSession": true, "sessionCapabilities": {"list": {}, "resume": {}, "close": {}}}}) @@ -189,7 +248,20 @@ async fn serve_acp( json!({"features": {"session_lifecycle": true, "session_pins": true}, "methods": ["querymt/session/getLifecycle", "querymt/session/setLifecycle", "querymt/session/setPinned"]}) } "_querymt/profiles" => json!({"profiles": [{"id": "fixture-profile"}]}), - "session/new" => json!({"sessionId": "fixture-session"}), + "session/new" => { + let mut workspace = workspace.lock().unwrap(); + let number = workspace.created.len() + 1; + let session_id = workspace.forced_id.clone().unwrap_or_else(|| { + if number == 1 { + "fixture-session".to_owned() + } else { + format!("fixture-session-{number}") + } + }); + workspace.created.push(json!({"sessionId": session_id, "cwd": params["cwd"], "_meta": params["_meta"]})); + json!({"sessionId": session_id}) + } + "session/close" => json!({}), "session/resume" | "session/load" => { if matches!(behavior, SubmissionBehavior::HistoryUnavailable) { send(&mut socket, json!({"jsonrpc": "2.0", "id": id, "error": {"code": -32000, "message": "history replay unavailable"}})).await; @@ -218,7 +290,8 @@ async fn serve_acp( } else if params["cursor"] == "page-2" { json!({"sessions": [{"sessionId": "second-session", "cwd": "/workspace", "title": "Second"}]}) } else { - json!({"sessions": [{"sessionId": "fixture-session", "cwd": "/workspace", "title": "Fixture"}], "nextCursor": "page-2"}) + let created = workspace.lock().unwrap().created.clone(); + json!({"sessions": if created.is_empty() { vec![json!({"sessionId": "fixture-session", "cwd": "/workspace", "title": "Fixture"})] } else { created }, "nextCursor": "page-2"}) } } "_querymt/session/runtimeState" => { @@ -267,7 +340,7 @@ async fn serve_acp( continue; } let permission = acp::RequestPermissionRequest::new( - "fixture-session", + params["sessionId"].as_str().unwrap().to_owned(), acp::ToolCallUpdate::new("tool-1", acp::ToolCallUpdateFields::new()), vec![acp::PermissionOption::new( "allow", @@ -287,7 +360,7 @@ async fn serve_acp( } else { format!("reply:{text}") }; - send(&mut socket, json!({"jsonrpc": "2.0", "method": "session/update", "params": {"sessionId": "fixture-session", "update": {"sessionUpdate": "agent_message_chunk", "content": {"type": "text", "text": reply}}}})).await; + send(&mut socket, json!({"jsonrpc": "2.0", "method": "session/update", "params": {"sessionId": params["sessionId"], "update": {"sessionUpdate": "agent_message_chunk", "content": {"type": "text", "text": reply}}}})).await; json!({"stopReason": if text == "cancelled" { "cancelled" } else { "end_turn" }}) } _ => { @@ -295,6 +368,13 @@ async fn serve_acp( continue; } }; + if method == "session/new" { + let gate = workspace.lock().unwrap().creation_gates.pop_front(); + if let Some(gate) = gate { + gate.observed.send(()).unwrap(); + gate.release.await.unwrap(); + } + } send( &mut socket, json!({"jsonrpc": "2.0", "id": id, "result": result}), @@ -303,6 +383,67 @@ async fn serve_acp( } } +#[tokio::test] +async fn inspect_reports_backend_replay_limitation_without_empty_history_or_fallback() { + for http in [false, true] { + let backend = MockAcp::start().await; + backend.workspace.lock().unwrap().replay_error = true; + let mut mcp = Mcp::start(&backend.url, http).await; + for full in [false, true] { + let result = mcp + .tool( + "inspect", + json!({"session_id":"fixture-session","full":full}), + ) + .await; + assert_eq!(result["isError"], true, "{result}"); + let content = &result["structuredContent"]; + let error = content["error"].as_str().unwrap(); + assert!(error.contains("QueryMT backend cannot replay"), "{result}"); + assert!(error.contains("History is unavailable, not empty")); + assert!(error.contains("Loaded session has no runtime ref for ACP replay")); + assert!(content.get("messages").is_none() && content.get("raw").is_none()); + } + let requests = backend.recorded(); + assert_eq!( + requests + .iter() + .filter(|r| r["method"] == "session/load") + .count(), + 2 + ); + assert!(requests.iter().all(|r| matches!( + r["method"].as_str(), + Some("initialize" | "session/list" | "session/load") + ))); + // A failed history fallback after successful resume/poll is not a done/empty answer. + let follow = mcp + .tool( + "follow", + json!({"session_id":"fixture-session","timeout":1,"interval":1}), + ) + .await; + assert_eq!(follow["isError"], true, "{follow}"); + assert!( + follow["structuredContent"]["error"] + .as_str() + .unwrap() + .contains("History is unavailable, not empty") + ); + assert!(follow["structuredContent"].get("text").is_none()); + assert!(backend.recorded().iter().all(|r| matches!( + r["method"].as_str(), + Some( + "initialize" + | "session/list" + | "session/load" + | "session/resume" + | "_querymt/session/runtimeState" + ) + ))); + } +} + enum Mcp { Stdio { child: Box, @@ -447,11 +588,592 @@ impl Drop for Mcp { } } +#[tokio::test] +async fn workspace_is_sticky_across_calls_and_switches_without_touching_old_sessions() { + for http in [false, true] { + let backend = MockAcp::start().await; + let mut mcp = Mcp::start(&backend.url, http).await; + assert_eq!(mcp.tool("workspace", json!({})).await["isError"], true); + assert!(backend.recorded().is_empty()); + let first = mcp + .tool( + "workspace", + json!({"cwd":"/remote/project/./sub/..", "profile":"fixture-profile"}), + ) + .await; + assert_eq!(first["isError"], false, "{first}"); + assert_eq!(first["structuredContent"]["cwd"], "/remote/project"); + assert_eq!(first["structuredContent"]["created"], true); + assert!( + !backend + .recorded() + .iter() + .any(|r| r["method"] == "session/prompt") + ); + let first = first["structuredContent"]["sessionId"].as_str().unwrap(); + let same = mcp + .tool("workspace", json!({"session_id":first,"cwd":"sub/.."})) + .await; + assert_eq!(same["structuredContent"]["sessionId"], first); + assert_eq!(same["structuredContent"]["created"], false); + assert_eq!( + mcp.tool("workspace", json!({"session_id":first,"profile":"other"})) + .await["isError"], + true + ); + let prompted = mcp + .tool("prompt", json!({"session_id":first,"text":"continue"})) + .await; + assert_eq!(prompted["structuredContent"]["cwd"], "/remote/project"); + for name in ["status", "inspect", "follow"] { + let result = mcp.tool(name, json!({"session_id":first})).await; + assert_eq!(result["isError"], false, "{result}"); + assert_eq!(result["structuredContent"]["cwd"], "/remote/project"); + let before = backend.recorded().len(); + assert_eq!( + mcp.tool(name, json!({"session_id":first,"cwd":"/wrong"})) + .await["isError"], + true + ); + assert_eq!(backend.recorded().len(), before); + } + let switched = mcp + .tool("workspace", json!({"session_id":first,"cwd":"../other"})) + .await; + assert_eq!(switched["isError"], false, "{switched}"); + assert_eq!(switched["structuredContent"]["previousSessionId"], first); + assert_eq!(switched["structuredContent"]["cwd"], "/remote/other"); + assert_eq!(switched["structuredContent"]["profileSource"], "inherited"); + let second = switched["structuredContent"]["sessionId"].as_str().unwrap(); + assert_ne!(first, second); + let old = mcp.tool("status", json!({"session_id":first})).await; + assert_eq!(old["structuredContent"]["cwd"], "/remote/project"); + let requests = backend.recorded(); + let creations: Vec<_> = requests + .iter() + .filter(|r| r["method"] == "session/new") + .collect(); + assert_eq!(creations.len(), 2); + assert!( + creations + .iter() + .all(|r| r["params"]["_meta"]["querymt"]["profile_id"] == "fixture-profile") + ); + assert!(!requests.iter().any(|r| r["method"] == "session/list")); + assert_no_implicit_cancel(&backend); + } +} + +#[tokio::test] +async fn workspace_metadata_recovery_is_unfiltered_paginated_and_never_trusts_caller_cwd() { + let backend = MockAcp::start().await; + backend.workspace.lock().unwrap().listing = Some(vec![ + json!({"sessions": [{"sessionId": "unrelated", "cwd": "/unrelated"}], "nextCursor": "page-2"}), + json!({"sessions": [{"sessionId": "old", "cwd": "/remote/original/./project", "_meta": {"querymt": {"profile_id": "listed-profile"}}}]}), + ]); + // Separate qmcp instances simulate losing all process-local bindings. + for http in [false, true] { + let mut mcp = Mcp::start(&backend.url, http).await; + let rejected = mcp + .tool( + "status", + json!({"session_id": "old", "cwd": "/requested/not-authoritative"}), + ) + .await; + assert_eq!(rejected["isError"], true, "{rejected}"); + let selected = mcp.tool("workspace", json!({"session_id": "old"})).await; + assert_eq!(selected["isError"], false, "{selected}"); + assert_eq!( + selected["structuredContent"]["cwd"], + "/remote/original/project" + ); + assert_eq!(selected["structuredContent"]["profile"], "listed-profile"); + let switched = mcp + .tool( + "prompt", + json!({"session_id": "old", "cwd": "../next", "text": "new workspace"}), + ) + .await; + assert_eq!(switched["isError"], false, "{switched}"); + assert_eq!(switched["structuredContent"]["previousSessionId"], "old"); + assert_eq!( + switched["structuredContent"]["cwd"], + "/remote/original/next" + ); + assert_eq!(switched["structuredContent"]["profileSource"], "inherited"); + let old = mcp.tool("status", json!({"session_id": "old"})).await; + assert_eq!(old["structuredContent"]["cwd"], "/remote/original/project"); + } + let requests = backend.recorded(); + let listing: Vec<_> = requests + .iter() + .filter(|r| r["method"] == "session/list") + .collect(); + assert_eq!(listing.len(), 4); + for page in listing { + assert!(page["params"].get("cwd").is_none()); + assert!(page["params"].get("_meta").is_none()); + } + for request in requests.iter().filter(|r| { + matches!( + r["method"].as_str(), + Some("session/load" | "session/resume") + ) + }) { + assert_eq!(request["params"]["sessionId"], "old"); + assert_eq!(request["params"]["cwd"], "/remote/original/project"); + } + assert_no_implicit_cancel(&backend); +} + +#[tokio::test] +async fn workspace_discovery_fails_closed_for_incomplete_or_unavailable_metadata() { + let cases = [ + (vec![json!({"sessions":[]})], "not found"), + ( + vec![json!({"sessions":[],"nextCursor":17})], + "metadata unavailable", + ), + ( + vec![json!({"sessions":[],"nextCursor":"page-1"})], + "repeated a cursor", + ), + ( + vec![json!({"sessions":[{"sessionId":"unknown","cwd":"relative"}]})], + "missing or not absolute", + ), + ( + vec![json!({"sessions":[{"sessionId":"unknown"}]})], + "metadata unavailable", + ), + ( + vec![json!({"error":{"code":-32601,"message":"list unsupported"}})], + "metadata unavailable", + ), + ( + (1..=20) + .map(|page| json!({"sessions":[],"nextCursor":format!("page-{}",page+1)})) + .collect(), + "page cap reached", + ), + ]; + for (pages, expected) in cases { + let backend = MockAcp::start().await; + backend.workspace.lock().unwrap().listing = Some(pages); + let mut mcp = Mcp::start(&backend.url, true).await; + let result = mcp + .tool( + "workspace", + json!({"session_id":"unknown","cwd":"/caller/cannot/prove/this"}), + ) + .await; + assert_eq!(result["isError"], true, "{result}"); + assert!( + result["structuredContent"]["error"] + .as_str() + .unwrap() + .contains(expected), + "{result}" + ); + assert!( + backend + .recorded() + .iter() + .all(|r| matches!(r["method"].as_str(), Some("initialize" | "session/list"))) + ); + } + // Each entry point must recover metadata before doing any session work. + for name in ["prompt", "status", "inspect", "follow"] { + let backend = MockAcp::start().await; + backend.workspace.lock().unwrap().listing = + Some(vec![json!({"sessions":[{"sessionId":"unknown"}]})]); + let mut mcp = Mcp::start(&backend.url, true).await; + let mut args = json!({"session_id":"unknown","cwd":"/caller/path"}); + if name == "prompt" { + args["text"] = json!("must not send"); + } + let result = mcp.tool(name, args).await; + assert_eq!(result["isError"], true, "{result}"); + assert!( + result["structuredContent"]["error"] + .as_str() + .unwrap() + .contains("metadata unavailable") + ); + assert!( + backend + .recorded() + .iter() + .all(|r| matches!(r["method"].as_str(), Some("initialize" | "session/list"))) + ); + } +} + +#[tokio::test] +async fn workspace_profile_unknown_is_explicit_and_paths_are_remote_lexical() { + let backend = MockAcp::start().await; + let mut mcp = Mcp::start(&backend.url, true).await; + let selected = mcp + .tool("workspace", json!({"session_id":"fixture-session"})) + .await; + assert_eq!(selected["structuredContent"]["profileKnown"], false); + assert_eq!( + mcp.tool( + "workspace", + json!({"session_id":"fixture-session","profile":"cannot-assume"}) + ) + .await["isError"], + true + ); + backend + .workspace + .lock() + .unwrap() + .created + .push(json!({"sessionId":"fixture-session","cwd":"/workspace"})); + let switched = mcp + .tool( + "workspace", + json!({"session_id":"fixture-session","cwd":"/remote/nonexistent"}), + ) + .await; + assert_eq!(switched["isError"], false, "{switched}"); + assert_eq!( + switched["structuredContent"]["profileSource"], + "backend-default-unknown" + ); + assert_eq!(switched["structuredContent"]["profileKnown"], false); + let explicit = mcp + .tool( + "workspace", + json!({"session_id":"fixture-session-2","cwd":"/other","profile":"explicit-profile"}), + ) + .await; + assert_eq!(explicit["structuredContent"]["profile"], "explicit-profile"); + let requests = backend.recorded(); + let creations: Vec<_> = requests + .iter() + .filter(|r| r["method"] == "session/new") + .collect(); + assert!(creations[0]["params"].get("_meta").is_none()); + assert_eq!( + creations[1]["params"]["_meta"]["querymt"]["profile_id"], + "explicit-profile" + ); + assert_no_implicit_cancel(&backend); +} + +#[tokio::test] +async fn workspace_batches_and_close_keep_bindings_and_new_commands_register_immediately() { + for http in [false, true] { + for name in ["new", "prompt"] { + let backend = MockAcp::start().await; + let mut mcp = Mcp::start(&backend.url, http).await; + let mut args = json!({"cwd": "/remote/partial", "mode": "invalid"}); + if name == "prompt" { + args["new"] = json!(true); + args["text"] = json!("not sent"); + } + let failed = mcp.tool(name, args).await; + assert_eq!(failed["isError"], true, "{failed}"); + assert_eq!( + failed["structuredContent"]["recovery"]["sessionCreated"], + true + ); + let status = mcp + .tool("status", json!({"session_id": "fixture-session"})) + .await; + assert_eq!(status["isError"], false, "{status}"); + assert_eq!(status["structuredContent"]["cwd"], "/remote/partial"); + let closed = mcp + .tool("close", json!({"session_id": "fixture-session"})) + .await; + assert_eq!(closed["isError"], false, "{closed}"); + let batch = mcp.tool("exec", json!({"calls": [ + {"name": "workspace", "arguments": {"session_id": "fixture-session"}}, + {"name": "status", "arguments": {"session_id": "fixture-session"}}, + {"name": "prompt", "arguments": {"session_id": "fixture-session", "cwd": "../switched", "text": "batch"}}, + {"name": "status", "arguments": {"session_id": "fixture-session-2"}}, + {"name": "inspect", "arguments": {"session_id": "fixture-session"}}, + ]})).await; + assert_eq!(batch["isError"], false, "{batch}"); + let steps = batch["structuredContent"]["results"].as_array().unwrap(); + assert_eq!(steps[0]["result"]["cwd"], "/remote/partial"); + assert_eq!(steps[1]["result"]["cwd"], "/remote/partial"); + assert_eq!(steps[2]["result"]["previousSessionId"], "fixture-session"); + assert_eq!(steps[3]["result"]["cwd"], "/remote/switched"); + assert_eq!(steps[4]["result"]["cwd"], "/remote/partial"); + assert!( + !backend + .recorded() + .iter() + .any(|r| r["method"] == "session/list") + ); + } + } +} + +#[tokio::test] +async fn workspace_rejected_creation_and_failed_switch_do_not_rebind_old_id() { + let backend = MockAcp::start().await; + let mut mcp = Mcp::start(&backend.url, true).await; + backend.workspace.lock().unwrap().reject_creation = true; + let rejected = mcp + .tool("workspace", json!({"cwd": "/remote/rejected"})) + .await; + assert_eq!(rejected["isError"], true, "{rejected}"); + assert!(rejected["structuredContent"].get("recovery").is_none()); + backend.workspace.lock().unwrap().reject_creation = false; + let created = mcp.tool("new", json!({"cwd": "/remote/old"})).await; + assert_eq!(created["isError"], false, "{created}"); + backend.workspace.lock().unwrap().reject_creation = true; + let rejected = mcp + .tool( + "workspace", + json!({"session_id": "fixture-session", "cwd": "../rejected"}), + ) + .await; + assert_eq!(rejected["isError"], true); + backend.workspace.lock().unwrap().reject_creation = false; + let failed = mcp.tool("prompt", json!({"session_id": "fixture-session", "cwd": "../new", "text": "not sent", "mode": "invalid"})).await; + assert_eq!(failed["isError"], true, "{failed}"); + let recovery = &failed["structuredContent"]["recovery"]; + assert_eq!(recovery["previousSessionId"], "fixture-session"); + assert_eq!(recovery["sessionId"], "fixture-session-2"); + assert_eq!(recovery["sessionCreated"], true); + assert_eq!(recovery["switched"], true); + for (id, cwd) in [ + ("fixture-session", "/remote/old"), + ("fixture-session-2", "/remote/new"), + ] { + let status = mcp.tool("status", json!({"session_id": id})).await; + assert_eq!(status["isError"], false, "{status}"); + assert_eq!(status["structuredContent"]["cwd"], cwd); + } + assert_no_implicit_cancel(&backend); +} + +#[tokio::test] +async fn concurrent_unknown_workspace_switches_never_bind_caller_paths_to_the_old_id() { + let backend = MockAcp::start().await; + backend.workspace.lock().unwrap().listing = Some(vec![ + json!({"sessions": [{"sessionId": "old", "cwd": "/authoritative/old"}]}), + ]); + let mut mcp = Mcp::start(&backend.url, true).await; + let Mcp::Http { client, url, .. } = &mcp else { + unreachable!() + }; + let request = |cwd| { + client.post(url).header("accept", "application/json, text/event-stream") + .header("mcp-protocol-version", "2025-11-25") + .json(&json!({"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "workspace", "arguments": {"session_id": "old", "cwd": cwd}}})) + }; + let (first, second) = + tokio::join!(request("/remote/one").send(), request("/remote/two").send()); + let first: Value = first.unwrap().json().await.unwrap(); + let second: Value = second.unwrap().json().await.unwrap(); + let first = &first["result"]["structuredContent"]; + let second = &second["result"]["structuredContent"]; + assert_eq!(first["ok"], true, "{first}"); + assert_eq!(second["ok"], true, "{second}"); + assert_ne!(first["sessionId"], second["sessionId"]); + assert_eq!(first["previousSessionId"], "old"); + assert_eq!(second["previousSessionId"], "old"); + for (id, cwd) in [ + (first["sessionId"].as_str().unwrap(), "/remote/one"), + (second["sessionId"].as_str().unwrap(), "/remote/two"), + ("old", "/authoritative/old"), + ] { + let status = mcp.tool("status", json!({"session_id": id})).await; + assert_eq!(status["isError"], false, "{status}"); + assert_eq!(status["structuredContent"]["cwd"], cwd); + } + assert_no_implicit_cancel(&backend); +} + +fn gated_creation(backend: &MockAcp) -> (oneshot::Receiver<()>, oneshot::Sender<()>) { + let (observed, ready) = oneshot::channel(); + let (release, wait) = oneshot::channel(); + backend + .workspace + .lock() + .unwrap() + .creation_gates + .push_back(CreationGate { + observed, + release: wait, + }); + (ready, release) +} + +fn concurrent_tool(mcp: &Mcp, name: &str, arguments: Value) -> JoinHandle { + let Mcp::Http { client, url, .. } = mcp else { + unreachable!() + }; + let request = client.post(url).header("accept", "application/json, text/event-stream") + .header("mcp-protocol-version", "2025-11-25") + .json(&json!({"jsonrpc": "2.0", "id": 91, "method": "tools/call", "params": {"name": name, "arguments": arguments}})); + tokio::spawn(async move { + let response: Value = request.send().await.unwrap().json().await.unwrap(); + response["result"].clone() + }) +} + +fn assert_creation_conflict(result: &Value, acknowledged_id: &str) { + assert_eq!(result["isError"], true, "{result}"); + let content = &result["structuredContent"]; + let recovery = &content["recovery"]; + assert_eq!( + recovery["workspaceConflict"]["acknowledgedSessionId"], acknowledged_id, + "{result}" + ); + assert_eq!(recovery["workspaceConflict"]["selectionCommitted"], false); + for key in [ + "sessionId", + "cwd", + "created", + "sessionCreated", + "switched", + "previousSessionId", + ] { + assert!( + recovery.get(key).is_none(), + "invalid selection in recovery: {result}" + ); + assert!( + content.get(key).is_none(), + "invalid usable selection: {result}" + ); + } +} + +#[tokio::test] +async fn workspace_metadata_before_creation_ack_reconciles_or_reports_a_conflict() { + tokio::time::timeout(Duration::from_secs(10), async { + for (metadata_cwd, success) in [("/remote/new", true), ("/contradiction", false)] { + let backend = MockAcp::start().await; + let mut mcp = Mcp::start(&backend.url, true).await; + mcp.tool( + "workspace", + json!({"cwd":"/remote/old","profile":"inherited"}), + ) + .await; + let (ready, release) = gated_creation(&backend); + let pending = concurrent_tool( + &mcp, + "prompt", + json!({"session_id":"fixture-session","cwd":"../new","text":"switch prompt"}), + ); + ready.await.unwrap(); + // The backend has created B, but another call discovers B before its ACK. + backend.workspace.lock().unwrap().listing = Some(vec![ + json!({"sessions":[{"sessionId":"fixture-session-2","cwd":metadata_cwd}]}), + ]); + let discovered = mcp + .tool("workspace", json!({"session_id":"fixture-session-2"})) + .await; + assert_eq!(discovered["structuredContent"]["created"], false); + assert!( + !backend + .recorded() + .iter() + .any(|r| r["method"] == "session/prompt") + ); + release.send(()).unwrap(); + let result = pending.await.unwrap(); + if success { + assert_eq!(result["isError"], false, "{result}"); + assert_eq!( + result["structuredContent"]["sessionId"], + "fixture-session-2" + ); + assert_eq!(result["structuredContent"]["profile"], "inherited"); + assert_eq!( + result["structuredContent"]["previousSessionId"], + "fixture-session" + ); + } else { + assert_creation_conflict(&result, "fixture-session-2"); + assert!( + !backend + .recorded() + .iter() + .any(|r| r["method"] == "session/prompt") + ); + } + let remembered = mcp + .tool("workspace", json!({"session_id":"fixture-session-2"})) + .await; + assert_eq!(remembered["structuredContent"]["cwd"], metadata_cwd); + let old = mcp + .tool("workspace", json!({"session_id":"fixture-session"})) + .await; + assert_eq!(old["structuredContent"]["cwd"], "/remote/old"); + assert_no_implicit_cancel(&backend); + } + }) + .await + .expect("barrier-controlled ACK test must complete"); +} + +#[tokio::test] +async fn workspace_prebound_duplicate_ack_never_publishes_an_invalid_selection() { + for from_metadata in [false, true] { + for bound_cwd in ["/remote/new", "/remote/different"] { + let backend = MockAcp::start().await; + let mut mcp = Mcp::start(&backend.url, true).await; + let old = mcp.tool("workspace", json!({"cwd": "/remote/old"})).await; + assert_eq!(old["isError"], false); + let duplicate_id = if from_metadata { + backend.workspace.lock().unwrap().listing = Some(vec![ + json!({"sessions": [{"sessionId": "prebound", "cwd": bound_cwd}]}), + ]); + "prebound" + } else { + let prebound = mcp.tool("workspace", json!({"cwd": bound_cwd})).await; + assert_eq!(prebound["isError"], false); + "fixture-session-2" + }; + let known = mcp + .tool("workspace", json!({"session_id": duplicate_id})) + .await; + assert_eq!(known["isError"], false); + backend.workspace.lock().unwrap().forced_id = Some(duplicate_id.into()); + for name in ["workspace", "prompt"] { + let mut args = json!({"session_id": "fixture-session", "cwd": "/remote/new"}); + if name == "prompt" { + args["text"] = json!("must not be sent"); + } + let failed = mcp.tool(name, args).await; + assert_creation_conflict(&failed, duplicate_id); + } + // Also cover the backend returning A itself for an A -> new switch. + backend.workspace.lock().unwrap().forced_id = Some("fixture-session".into()); + let failed = mcp.tool("prompt", json!({"session_id": "fixture-session", "cwd": "/remote/new", "text": "must not be sent"})).await; + assert_creation_conflict(&failed, "fixture-session"); + for (id, cwd) in [ + ("fixture-session", "/remote/old"), + (duplicate_id, bound_cwd), + ] { + let known = mcp.tool("workspace", json!({"session_id": id})).await; + assert_eq!(known["structuredContent"]["cwd"], cwd); + } + assert!( + !backend + .recorded() + .iter() + .any(|r| r["method"] == "session/prompt") + ); + assert_no_implicit_cancel(&backend); + } + } +} + async fn exercise_bridge(http: bool) { let backend = MockAcp::start().await; let mut mcp = Mcp::start(&backend.url, http).await; let catalog = mcp.rpc("tools/list", json!({})).await; - assert_eq!(catalog["tools"].as_array().unwrap().len(), 25); + assert_eq!(catalog["tools"].as_array().unwrap().len(), 26); assert!( backend.recorded().is_empty(), "discovery must not connect to ACP" @@ -619,7 +1341,7 @@ async fn failures_are_tool_errors_and_timeouts_do_not_cancel_remote_work() { } #[tokio::test] -async fn long_prompt_does_not_block_concurrent_runtime_requests() { +async fn long_prompt_does_not_block_concurrent_status_runtime_steer_or_cancel() { let backend = MockAcp::start().await; let mut mcp = Mcp::start(&backend.url, true).await; let Mcp::Http { client, url, .. } = &mcp else { @@ -651,6 +1373,25 @@ async fn long_prompt_does_not_block_concurrent_runtime_requests() { .await .unwrap(); assert_eq!(runtime["isError"], false); + for (name, args) in [ + ("status", json!({"session_id": "fixture-session"})), + ( + "steer", + json!({"session_id": "fixture-session", "run_id": "fixture-run", "text": "steer while prompt waits"}), + ), + ("cancel", json!({"session_id": "fixture-session"})), + ] { + let result = tokio::time::timeout(Duration::from_millis(500), mcp.tool(name, args)) + .await + .expect("workspace locks must not block control commands"); + assert_eq!(result["isError"], false, "{name}: {result}"); + } + assert!( + !backend + .recorded() + .iter() + .any(|r| r["method"] == "session/list") + ); let finished = waiting.await.unwrap(); assert_eq!(finished["result"]["structuredContent"]["code"], 4); backend.wait_disconnected().await; @@ -849,10 +1590,20 @@ async fn invalid_runtime_and_missing_active_run_never_submit_work() { .await; assert_eq!(response["isError"], true, "{response}"); assert_eq!(response["structuredContent"]["code"], 3); - assert!(backend.recorded().iter().all(|request| matches!( + let requests = backend.recorded(); + assert!(requests.iter().all(|request| matches!( request["method"].as_str(), - Some("initialize" | "_querymt/session/runtimeState") + Some("initialize" | "session/list" | "_querymt/session/runtimeState") ))); + assert!( + requests + .iter() + .filter(|request| request["method"] == "session/list") + .all(|request| { + request["params"].get("cwd").is_none() + && request["params"].get("_meta").is_none() + }) + ); } }