From a3871a0001d2e9253c0723391a3a46f93bf2f616 Mon Sep 17 00:00:00 2001 From: Shooksie Date: Sun, 12 Jul 2026 07:39:27 -0600 Subject: [PATCH] fix(mcp-oauth): resilient config resolution + reduce config_source spawn amplification + cache OAuth discovery (TASK-326) --- crates/animus-mcp-oauth/Cargo.toml | 1 + crates/animus-mcp-oauth/src/client.rs | 40 ++++- crates/animus-mcp-oauth/src/config.rs | 159 +++++++++++++---- crates/animus-mcp-oauth/src/flow.rs | 48 +++-- crates/animus-mcp-oauth/src/proxy.rs | 168 +++++++++++++++++- .../tests/proxy_integration.rs | 75 +++++++- .../src/bin/animus_mcp_proxy.rs | 27 ++- .../workflow_config/config_source_client.rs | 55 ++++++ .../src/workflow_config/loading.rs | 82 +++++++++ .../src/workflow_config/mod.rs | 3 +- .../src/workflow_config/tests.rs | 71 ++++++++ 11 files changed, 665 insertions(+), 64 deletions(-) diff --git a/crates/animus-mcp-oauth/Cargo.toml b/crates/animus-mcp-oauth/Cargo.toml index 41c01533..edf8c001 100644 --- a/crates/animus-mcp-oauth/Cargo.toml +++ b/crates/animus-mcp-oauth/Cargo.toml @@ -33,6 +33,7 @@ rmcp = { version = "1.7", features = [ [dev-dependencies] protocol = { workspace = true, features = ["test-utils"] } +orchestrator-config = { workspace = true, features = ["test-utils"] } tempfile = "3" axum = "0.8" http = "1" diff --git a/crates/animus-mcp-oauth/src/client.rs b/crates/animus-mcp-oauth/src/client.rs index 16cda716..aaf3255a 100644 --- a/crates/animus-mcp-oauth/src/client.rs +++ b/crates/animus-mcp-oauth/src/client.rs @@ -59,33 +59,55 @@ impl McpSession { timeout: Duration, ) -> Result { crate::ensure_crypto_provider(); + // The ONE `config_source` resolution per `mcp call`. The resolved URL + + // flow are threaded to the spawned proxy (`--url` + `--auth-code`) so the + // proxy TRUSTS them and does NOT re-resolve — collapsing the historical + // 3× config_source spawn amplification (CLI + proxy-bin + proxy-connect) + // down to this single touch. let resolution = resolve_server_url(project_root, server, url_override)?; // OAuth servers (every flow) are served through the local proxy, which // resolves + injects the live bearer itself; plain-HTTP servers are // connected directly. let uses_oauth = resolution.is_authorization_code || resolution.broker_oauth.is_some(); if uses_oauth { - Self::connect_via_proxy(project_root, server, url_override, timeout).await + Self::connect_via_proxy(project_root, server, &resolution, timeout).await } else { Self::connect_http(&resolution.url, timeout).await } } - /// Spawn `animus-mcp-proxy --server [--url ] --project-root - /// ` and connect an MCP client to its stdio. The proxy injects the - /// cached bearer; no secret is passed on argv or read here. stderr is - /// inherited so proxy diagnostics (e.g. "run `animus mcp auth`") reach the - /// user; stdout carries only JSON-RPC frames. + /// Spawn `animus-mcp-proxy --server --url + /// [--auth-code] --project-root ` and connect an MCP client to its + /// stdio. The proxy injects the cached bearer; no secret is passed on argv + /// or read here. stderr is inherited so proxy diagnostics (e.g. "run `animus + /// mcp auth`") reach the user; stdout carries only JSON-RPC frames. + /// + /// Passing the already-resolved `--url` (and, for keychain servers, + /// `--auth-code`) lets the proxy skip its own `config_source` lookup — so a + /// bulk `mcp call` burst does one source resolution per call, here, instead + /// of one per proxy spawn. async fn connect_via_proxy( project_root: &Path, server: &str, - url_override: Option<&str>, + resolution: &crate::config::ServerResolution, timeout: Duration, ) -> Result { let mut cmd = Command::new(mcp_proxy_command()); cmd.arg("--server").arg(server).arg("--project-root").arg(project_root); - if let Some(url) = url_override.filter(|u| !u.trim().is_empty()) { - cmd.arg("--url").arg(url); + // Pass the resolved `--url` for BOTH flows so the proxy binds the exact + // upstream the parent selected (honoring a `--url` override) instead of + // re-resolving the server name to a possibly-different same-named entry. + if !resolution.url.trim().is_empty() { + cmd.arg("--url").arg(&resolution.url); + } + // `--auth-code` is added ONLY for keychain (`authorization_code`) servers: + // it tells the proxy to trust `--url` and skip its own config_source + // lookup entirely (the amplification cut). Broker flows omit it because + // they still need their full oauth block from config; the proxy re-resolves + // those — and if the source is down it errors cleanly (ConfigSourceUnavailable) + // rather than misrouting the broker server to the keychain path. + if resolution.is_authorization_code { + cmd.arg("--auth-code"); } // `TokioChildProcess::new` pipes stdin/stdout and inherits stderr; the // child is killed on drop of the returned transport (held by the diff --git a/crates/animus-mcp-oauth/src/config.rs b/crates/animus-mcp-oauth/src/config.rs index b67a95ad..0c04a488 100644 --- a/crates/animus-mcp-oauth/src/config.rs +++ b/crates/animus-mcp-oauth/src/config.rs @@ -6,7 +6,7 @@ use std::path::Path; use std::sync::Arc; use orchestrator_config::workflow_config::{ - load_workflow_config_or_default, load_workflow_config_with_metadata, OauthConfig, OauthFlow, + try_load_workflow_config, OauthConfig, OauthFlow, WorkflowConfigAvailability, }; use orchestrator_core::SecretStore; use protocol::repository_scope::{repository_scope_for_path, scoped_state_root}; @@ -21,6 +21,13 @@ pub enum ServerResolutionError { UnknownServer(String), #[error("MCP server `{0}` has no `url`; an HTTP-transport URL is required for OAuth")] MissingUrl(String), + /// The `config_source` failed to load (spawn / RPC / DB overload / validation), + /// so the server's config could NOT be determined. This is distinct from + /// [`UnknownServer`], which fires only when the config loaded and the server + /// is genuinely absent. Callers should retry / report a transient source + /// error rather than "server not configured". + #[error("MCP server `{0}` could not be resolved: the workflow config source failed to load ({1}). This is a config-source error, not a missing-server error — retry shortly or check the `config_source` plugin.")] + ConfigSourceUnavailable(String, String), #[error("failed to load workflow config: {0}")] WorkflowConfig(String), #[error("failed to load project config: {0}")] @@ -99,22 +106,30 @@ pub fn resolve_server_url( ) -> Result { // Workflow config mcp_servers first (authoritative for daemon runs). // - // A malformed `.animus/workflows.yaml` must surface its YAML/validation - // error rather than silently falling back to the builtin default (which - // would mislead the user into "unknown server"). But an *absent* workflow - // config is normal (project-config-only setups), and - // `load_workflow_config_with_metadata` returns a "missing" error in that - // case. So: only propagate the load error when a workflow YAML source - // actually exists; otherwise fall back to the builtin default and continue - // to project config. - let loaded = match load_workflow_config_with_metadata(project_root, None) { - Ok(loaded) => loaded, - Err(err) if workflow_yaml_present(project_root) => { - return Err(ServerResolutionError::WorkflowConfig(err.to_string())); + // Use the NON-SWALLOWING loader so a transient `config_source` failure (DB + // overload under bulk `mcp call`) is NOT degraded into an empty config that + // then misreports the server as "not defined". `try_load_workflow_config` + // distinguishes three cases: + // * Loaded — use its `mcp_servers`. + // * NoSource — no config_source configured; benign, fall through to + // project config (project-config-only setups). + // * SourceUnavailable — the source failed; surface a DISTINCT retryable + // error. We do this EVEN when the caller supplied a + // `--url`: a `--url` overrides only the upstream + // endpoint, not the server's oauth FLOW/block, so + // synthesizing an `authorization_code` resolution from + // the URL alone would misroute a broker-flow server + // (`manual_bearer` / `client_credentials` / + // `refresh_token`) to the keychain path. Erroring lets + // the caller retry once the source recovers. + let loaded = match try_load_workflow_config(project_root, None) { + WorkflowConfigAvailability::Loaded(loaded) => Some(loaded), + WorkflowConfigAvailability::NoSource => None, + WorkflowConfigAvailability::SourceUnavailable(err) => { + return Err(ServerResolutionError::ConfigSourceUnavailable(server.to_string(), err.to_string())); } - Err(_) => load_workflow_config_or_default(project_root), }; - if let Some(def) = loaded.config.mcp_servers.get(server) { + if let Some(def) = loaded.as_ref().and_then(|loaded| loaded.config.mcp_servers.get(server)) { return finalize(server, url_override, def.url.clone(), def.oauth.clone()); } @@ -139,21 +154,6 @@ pub fn resolve_server_url( } } -/// True when a workflow YAML source exists at `.animus/workflows.yaml` or -/// any `.animus/workflows/*.yaml`. Used to decide whether a workflow-config -/// load failure is a real (malformed-config) error worth propagating vs the -/// benign "no workflow config" case. -fn workflow_yaml_present(project_root: &Path) -> bool { - let animus = project_root.join(".animus"); - if animus.join("workflows.yaml").exists() { - return true; - } - let dir = animus.join("workflows"); - std::fs::read_dir(&dir) - .map(|entries| entries.flatten().any(|e| e.path().extension().is_some_and(|ext| ext == "yaml" || ext == "yml"))) - .unwrap_or(false) -} - fn finalize( server: &str, url_override: Option<&str>, @@ -188,3 +188,102 @@ fn finalize( }), } } + +#[cfg(test)] +mod tests { + use super::*; + use orchestrator_config::workflow_config::builtin_workflow_config; + use orchestrator_config::workflow_config::config_source_client::test_seam; + use std::sync::{Mutex, MutexGuard, OnceLock}; + + // These tests mutate the process-global config_source test seam + HOME env, + // so serialize them under one lock held for the whole test body. + fn serial() -> &'static Mutex<()> { + static LOCK: OnceLock> = OnceLock::new(); + LOCK.get_or_init(|| Mutex::new(())) + } + + /// RAII guard that points `HOME` at a temp dir for the test and RESTORES the + /// previous value on drop, so the mutation never leaks to sibling tests. + struct HomeGuard { + prev: Option, + } + + impl HomeGuard { + fn set(path: &Path) -> Self { + let prev = std::env::var_os("HOME"); + std::env::set_var("HOME", path); + Self { prev } + } + } + + impl Drop for HomeGuard { + fn drop(&mut self) { + match self.prev.take() { + Some(v) => std::env::set_var("HOME", v), + None => std::env::remove_var("HOME"), + } + } + } + + /// Common fixture: hold the serial lock, isolate `HOME` to a fresh tempdir + /// (restored on drop). Bound in the caller in declaration order so the + /// tempdir drops LAST — after `HomeGuard` has restored `HOME`. + fn fixture() -> (MutexGuard<'static, ()>, tempfile::TempDir, HomeGuard) { + let guard = serial().lock().unwrap_or_else(|p| p.into_inner()); + let temp = tempfile::tempdir().expect("tempdir"); + let home = HomeGuard::set(temp.path()); + (guard, temp, home) + } + + /// TASK-326 (fix d): a transient config_source failure must surface as the + /// DISTINCT `ConfigSourceUnavailable`, NOT `UnknownServer` — otherwise the + /// caller concludes "krisp isn't configured" when the source is merely down. + #[test] + fn source_failure_yields_config_source_unavailable_not_unknown_server() { + let (_guard, temp, _home) = fixture(); + let root = temp.path(); + + let _fail = test_seam::install_failure(root, "config_source RPC timed out"); + let err = resolve_server_url(root, "krisp", None).expect_err("must fail when source is down"); + assert!( + matches!(err, ServerResolutionError::ConfigSourceUnavailable(ref s, _) if s == "krisp"), + "expected ConfigSourceUnavailable, got {err:?}" + ); + } + + /// Even with an explicit `--url`, a source outage yields + /// `ConfigSourceUnavailable` (never a synthesized `authorization_code` + /// resolution): a `--url` overrides only the upstream endpoint, not the + /// server's oauth FLOW, so guessing the flow could misroute a broker server + /// to the keychain path. The caller retries once the source recovers. + #[test] + fn explicit_url_still_errors_when_source_unavailable() { + let (_guard, temp, _home) = fixture(); + let root = temp.path(); + + let _fail = test_seam::install_failure(root, "config_source RPC timed out"); + let err = resolve_server_url(root, "krisp", Some("https://example.test/mcp")) + .expect_err("a source outage must error even with --url, to avoid misrouting a broker server"); + assert!( + matches!(err, ServerResolutionError::ConfigSourceUnavailable(ref s, _) if s == "krisp"), + "expected ConfigSourceUnavailable, got {err:?}" + ); + } + + /// When the config LOADS but the server is genuinely absent, the error is + /// `UnknownServer` (the pre-existing behavior) — never `ConfigSourceUnavailable`. + #[test] + fn loaded_config_missing_server_yields_unknown_server() { + let (_guard, temp, _home) = fixture(); + let root = temp.path(); + + // A base with no mcp_servers loads cleanly; the server is simply absent. + let _base = test_seam::install(root, builtin_workflow_config()); + let err = resolve_server_url(root, "krisp", None).expect_err("absent server must error"); + assert!( + matches!(err, ServerResolutionError::UnknownServer(ref s) if s == "krisp"), + "expected UnknownServer, got {err:?}" + ); + } +} diff --git a/crates/animus-mcp-oauth/src/flow.rs b/crates/animus-mcp-oauth/src/flow.rs index a1724bb9..5f96f050 100644 --- a/crates/animus-mcp-oauth/src/flow.rs +++ b/crates/animus-mcp-oauth/src/flow.rs @@ -46,7 +46,7 @@ enum Exchange { } use crate::callback::CallbackListener; -use crate::config::{build_secret_store, resolve_principal_id, resolve_server_url}; +use crate::config::{build_secret_store, resolve_principal_id, resolve_server_url, ServerResolutionError}; use crate::keychain_store::KeychainCredentialStore; use crate::pending::{PendingAuth, PendingStore}; use crate::state_store::PersistentStateStore; @@ -745,17 +745,26 @@ pub async fn auth_status(project_root: &Path, server: Option<&str>, url_override let servers: Vec = match server { Some(s) => vec![s.to_string()], - None => authorization_code_servers(project_root), + // Non-swallowing enumeration: a transient config_source failure surfaces + // as an Err here rather than an empty `servers:[]` that would read as + // "no OAuth servers configured". + None => authorization_code_servers(project_root)?, }; let mut out = Vec::with_capacity(servers.len()); for name in servers { // Tokens are keyed by the upstream URL too, so resolve it. A server // that can't be resolved (e.g. dropped from config and no --url) is - // reported as unauthenticated rather than failing the whole report. - let Some(url) = resolve_server_url(project_root, &name, url_override).ok().map(|r| r.url) else { - out.push(server_state_from_creds(&name, &principal, None)); - continue; + // reported as unauthenticated rather than failing the whole report — + // EXCEPT a transient config_source failure, which must surface (not be + // silently reported as an unauthenticated server) so the user retries. + let url = match resolve_server_url(project_root, &name, url_override) { + Ok(resolution) => resolution.url, + Err(err @ ServerResolutionError::ConfigSourceUnavailable(..)) => return Err(err.into()), + Err(_) => { + out.push(server_state_from_creds(&name, &principal, None)); + continue; + } }; let store = KeychainCredentialStore::new(secrets.clone(), &name, &principal, &url); let creds = @@ -814,14 +823,27 @@ fn server_state_from_creds(server: &str, principal: &str, creds: Option Vec { - use orchestrator_config::workflow_config::{load_workflow_config_or_default, OauthFlow}; +/// +/// Uses the NON-SWALLOWING loader: a transient `config_source` failure returns +/// `Err` so `auth_status` reports a transient error rather than an empty +/// `servers:[]` (which would read as "no OAuth servers configured"). A genuinely +/// absent config source (`NoSource`) is benign — enumeration proceeds from +/// project config only. +fn authorization_code_servers(project_root: &Path) -> Result> { + use orchestrator_config::workflow_config::{try_load_workflow_config, OauthFlow, WorkflowConfigAvailability}; let mut names = std::collections::BTreeSet::new(); - let loaded = load_workflow_config_or_default(project_root); - for (name, def) in &loaded.config.mcp_servers { - if def.oauth.as_ref().is_some_and(|o| o.flow == OauthFlow::AuthorizationCode) { - names.insert(name.clone()); + match try_load_workflow_config(project_root, None) { + WorkflowConfigAvailability::Loaded(loaded) => { + for (name, def) in &loaded.config.mcp_servers { + if def.oauth.as_ref().is_some_and(|o| o.flow == OauthFlow::AuthorizationCode) { + names.insert(name.clone()); + } + } + } + WorkflowConfigAvailability::NoSource => {} + WorkflowConfigAvailability::SourceUnavailable(err) => { + return Err(anyhow!("workflow config source unavailable while listing OAuth servers: {err}")); } } @@ -838,7 +860,7 @@ fn authorization_code_servers(project_root: &Path) -> Vec { } } - names.into_iter().collect() + Ok(names.into_iter().collect()) } /// Extract the `state` query parameter from an authorization URL. diff --git a/crates/animus-mcp-oauth/src/proxy.rs b/crates/animus-mcp-oauth/src/proxy.rs index 941ed938..86fc8784 100644 --- a/crates/animus-mcp-oauth/src/proxy.rs +++ b/crates/animus-mcp-oauth/src/proxy.rs @@ -32,7 +32,7 @@ use rmcp::model::{ClientNotification, ClientRequest, ErrorCode, ServerInfo, Serv use rmcp::service::{ NotificationContext, RequestContext, RoleClient, RoleServer, RunningService, Service, ServiceError, }; -use rmcp::transport::auth::{AuthError, AuthorizationManager}; +use rmcp::transport::auth::{AuthError, AuthorizationManager, CredentialStore}; use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig; use rmcp::transport::{stdio, StreamableHttpClientTransport}; use rmcp::{serve_client, ErrorData as McpError, ServiceExt}; @@ -103,14 +103,33 @@ pub struct McpProxy { } impl McpProxy { - /// Build the proxy: resolve the server URL, load the keychain-backed - /// auth manager, fetch the first access token, and open the upstream - /// connection. + /// Build the proxy: resolve the server URL from config, load the + /// keychain-backed auth manager, fetch the first access token, and open the + /// upstream connection. + /// + /// Prefer [`connect_authorization_code`](Self::connect_authorization_code) + /// when the upstream URL is ALREADY resolved (the daemon/CLI contract passes + /// it via `--url`): that path skips the `config_source` round-trip this one + /// performs, which is what saturates the source under bulk `mcp call`. pub async fn connect(project_root: &Path, server_name: &str, url_override: Option<&str>) -> Result { let resolution = resolve_server_url(project_root, server_name, url_override)?; + Self::connect_authorization_code(project_root, server_name, &resolution.url).await + } + + /// Build the proxy for an `authorization_code` (keychain) server whose + /// upstream URL is ALREADY resolved — skipping the `config_source` lookup + /// [`connect`](Self::connect) does. The keychain principal + secret store + /// are read from the OS keychain / on-disk scoped state (NOT the config + /// source), so this performs ZERO `config_source` spawns. + pub async fn connect_authorization_code( + project_root: &Path, + server_name: &str, + upstream_url: &str, + ) -> Result { let principal = resolve_principal_id(project_root); let secrets = build_secret_store(project_root)?; - Self::connect_with_store(server_name, &resolution.url, secrets, &principal).await + let cache_dir = discovery_cache::cache_dir_for(project_root); + Self::connect_with_store(server_name, upstream_url, secrets, &principal, cache_dir.as_deref()).await } /// Connect with an explicit keychain store, principal, and upstream URL. @@ -121,22 +140,51 @@ impl McpProxy { /// the OAuth `resource` indicator (RFC 8707) and the discovery seed in /// rmcp 1.7 — it must match the URL the auth flow logged in against so /// refresh issues a token for the same audience. + /// + /// `discovery_cache_dir` (when `Some`) is a directory used to cache the + /// discovered RFC 8414/9728 OAuth metadata per `(server, url)` so repeated + /// proxy spawns don't re-hit the (throttled) upstream discovery endpoint on + /// every connect. `None` disables the cache (used by tests that want a live + /// discovery each connect). pub async fn connect_with_store( server_name: &str, upstream_url: &str, secrets: Arc, principal: &str, + discovery_cache_dir: Option<&Path>, ) -> Result { crate::ensure_crypto_provider(); let cred_store = KeychainCredentialStore::new(secrets, server_name, principal, upstream_url); + // Gate discovery priming on a stored token existing. Priming (below) + // runs a live `.well-known` discovery on a cache miss; doing that for an + // unauthenticated server (no stored token) would turn the fast-fail + // startup path into a needless network round-trip against the (throttled) + // upstream. When the credential store has no token we skip priming and + // let `initialize_from_store` take the unauthenticated fast-fail path. + let has_stored_token = matches!(cred_store.load().await, Ok(Some(_))); + let mut manager = AuthorizationManager::new(upstream_url) .await .map_err(|err| anyhow!("failed to init OAuth manager for `{server_name}`: {err}"))?; manager.set_credential_store(cred_store); - // Hydrate the manager from the stored token (discovers metadata + - // configures the client id). `false` means no usable stored token. + // (a) Prime the manager with cached OAuth discovery metadata BEFORE + // `initialize_from_store` (which only discovers when metadata is unset — + // rmcp 1.7 `transport/auth.rs`). Under bulk `mcp call` the upstream + // throttles the per-spawn `.well-known` discovery, surfacing as + // `NoAuthorizationSupport` ("No authorization support detected"); a + // cache hit avoids the network call entirely. On a miss we discover once + // and persist for subsequent spawns. Best-effort: a cache/discovery + // failure here is non-fatal — `initialize_from_store` still runs. Only + // primed when a token is stored (see `has_stored_token` above). + if has_stored_token { + prime_discovery_metadata(&mut manager, server_name, upstream_url, discovery_cache_dir).await; + } + + // Hydrate the manager from the stored token (discovers metadata when not + // already primed above + configures the client id). `false` means no + // usable stored token. let hydrated = manager .initialize_from_store() .await @@ -321,12 +369,25 @@ impl McpProxy { } } -/// Drive the proxy: serve the agent over stdio until the connection closes. +/// Drive the proxy: resolve the server URL from config, then serve the agent +/// over stdio until the connection closes. +/// +/// Prefer [`run_authorization_code`] when the URL is already resolved — it +/// avoids the `config_source` round-trip this path performs. pub async fn run(project_root: &Path, server_name: &str, url_override: Option<&str>) -> Result<()> { let proxy = McpProxy::connect(project_root, server_name, url_override).await?; serve_until_closed(proxy, server_name).await } +/// Drive the proxy for an `authorization_code` (keychain) server whose upstream +/// URL is ALREADY resolved (passed by the daemon/CLI contract via `--url`). +/// Skips the `config_source` lookup [`run`] performs, so a bulk `mcp call` +/// burst doesn't re-saturate the source on every proxy spawn. +pub async fn run_authorization_code(server_name: &str, upstream_url: &str, project_root: &Path) -> Result<()> { + let proxy = McpProxy::connect_authorization_code(project_root, server_name, upstream_url).await?; + serve_until_closed(proxy, server_name).await +} + /// Drive the proxy for a broker-flow server: tokens come from `source` /// (resolved by the caller, typically the `animus-mcp-proxy` binary over the /// OAuth broker) instead of the keychain. @@ -406,6 +467,97 @@ fn is_retryable_auth_error(err: &ServiceError) -> bool { matches!(err, ServiceError::TransportSend(_) | ServiceError::TransportClosed) } +/// Prime `manager` with OAuth discovery metadata for `(server_name, +/// upstream_url)` so `initialize_from_store` skips its per-spawn `.well-known` +/// discovery. On a cache hit, load + `set_metadata`. On a miss, discover live, +/// `set_metadata`, and persist for next time. All failures are non-fatal (we +/// simply leave the manager to discover itself). +async fn prime_discovery_metadata( + manager: &mut AuthorizationManager, + server_name: &str, + upstream_url: &str, + cache_dir: Option<&Path>, +) { + let Some(cache_dir) = cache_dir else { + return; + }; + if let Some(cached) = discovery_cache::load(cache_dir, server_name, upstream_url) { + manager.set_metadata(cached); + return; + } + match manager.discover_metadata().await { + Ok(metadata) => { + discovery_cache::store(cache_dir, server_name, upstream_url, &metadata); + manager.set_metadata(metadata); + } + Err(err) => { + // Non-fatal: leave `metadata` unset so `initialize_from_store` + // attempts its own discovery (which will surface the real error if + // the upstream is genuinely unreachable). + tracing::warn!(server = server_name, error = %err, "oauth discovery failed while priming cache; will retry via initialize_from_store"); + } + } +} + +/// Best-effort on-disk cache of discovered OAuth [`AuthorizationMetadata`], +/// keyed by `(server, upstream_url)`. The metadata is the set of public +/// `.well-known` OAuth endpoints (RFC 8414/9728) — NOT a secret — so it lives +/// next to (not inside) the encrypted token store. A stale entry self-heals via +/// the TTL below; any read/parse error is treated as a miss. +mod discovery_cache { + use std::path::{Path, PathBuf}; + use std::time::{Duration, SystemTime}; + + use rmcp::transport::auth::AuthorizationMetadata; + use sha2::{Digest, Sha256}; + + /// Cache subdirectory under the scoped state root. Versioned so a metadata + /// schema change can bump the directory without colliding with old entries. + const CACHE_SUBDIR: &str = "mcp-oauth-discovery.v1"; + /// Entries older than this (24h) are ignored (and overwritten on the next + /// miss), so a genuinely rotated upstream discovery document eventually + /// takes hold. + const TTL: Duration = Duration::from_hours(24); + + /// The discovery-cache directory for `project_root` (its scoped state root), + /// or `None` when the scope can't be resolved. + pub(super) fn cache_dir_for(project_root: &Path) -> Option { + protocol::repository_scope::scoped_state_root(project_root).map(|root| root.join(CACHE_SUBDIR)) + } + + /// Stable filename for a `(server, url)` pair. + fn entry_path(cache_dir: &Path, server: &str, url: &str) -> PathBuf { + let mut hasher = Sha256::new(); + hasher.update(server.as_bytes()); + hasher.update([0x1f]); + hasher.update(url.as_bytes()); + cache_dir.join(format!("{:x}.json", hasher.finalize())) + } + + /// Load cached metadata for `(server, url)`, or `None` on miss / expiry / + /// any read or parse error. + pub(super) fn load(cache_dir: &Path, server: &str, url: &str) -> Option { + let path = entry_path(cache_dir, server, url); + let modified = std::fs::metadata(&path).ok()?.modified().ok()?; + if SystemTime::now().duration_since(modified).map(|age| age > TTL).unwrap_or(true) { + return None; + } + let bytes = std::fs::read(&path).ok()?; + serde_json::from_slice(&bytes).ok() + } + + /// Persist `metadata` for `(server, url)`. Best-effort — any error is + /// ignored (the next connect simply re-discovers). + pub(super) fn store(cache_dir: &Path, server: &str, url: &str, metadata: &AuthorizationMetadata) { + if std::fs::create_dir_all(cache_dir).is_err() { + return; + } + if let Ok(bytes) = serde_json::to_vec(metadata) { + let _ = std::fs::write(entry_path(cache_dir, server, url), bytes); + } + } +} + fn auth_error_to_anyhow(err: AuthError, server_name: &str) -> anyhow::Error { match err { AuthError::AuthorizationRequired | AuthError::TokenExpired => { diff --git a/crates/animus-mcp-oauth/tests/proxy_integration.rs b/crates/animus-mcp-oauth/tests/proxy_integration.rs index 0852d70b..e7527d25 100644 --- a/crates/animus-mcp-oauth/tests/proxy_integration.rs +++ b/crates/animus-mcp-oauth/tests/proxy_integration.rs @@ -44,9 +44,13 @@ struct MockState { /// Set once the `/token` refresh endpoint has been called. refreshed: Arc>, token_calls: Arc>, + /// Count of OAuth Authorization-Server discovery (`.well-known`) hits, used + /// to prove the discovery-metadata cache avoids re-discovery per spawn. + well_known_calls: Arc>, } async fn well_known(State(state): State) -> impl IntoResponse { + *state.well_known_calls.lock().await += 1; let base = state.base_url.lock().await.clone(); ( StatusCode::OK, @@ -169,6 +173,7 @@ async fn spawn_mock(require_refresh: bool) -> (String, MockState) { require_refresh, refreshed: Arc::new(Mutex::new(false)), token_calls: Arc::new(Mutex::new(0)), + well_known_calls: Arc::new(Mutex::new(0)), }; let app = Router::new() .route("/.well-known/oauth-authorization-server", get(well_known)) @@ -247,7 +252,7 @@ async fn tools_list_passes_bearer_through_to_upstream() { // Fresh token, long expiry → no refresh needed. let secrets = seed_token("github", "local", &url, ACCESS_TOKEN_FRESH, 3600, true); - let proxy = animus_mcp_oauth::proxy::McpProxy::connect_with_store("github", &url, secrets, "local") + let proxy = animus_mcp_oauth::proxy::McpProxy::connect_with_store("github", &url, secrets, "local", None) .await .expect("proxy connects with stored token"); @@ -280,7 +285,7 @@ async fn upstream_401_triggers_refresh_and_retry() { // upstream additionally 401s until the refresh lands. let secrets = seed_token("linear", "local", &url, "stale-token", 1, false); - let proxy = animus_mcp_oauth::proxy::McpProxy::connect_with_store("linear", &url, secrets, "local") + let proxy = animus_mcp_oauth::proxy::McpProxy::connect_with_store("linear", &url, secrets, "local", None) .await .expect("proxy connects"); @@ -299,3 +304,69 @@ async fn upstream_401_triggers_refresh_and_retry() { "a request must carry the refreshed bearer, saw: {headers:?}" ); } + +/// TASK-326 (fix a): with a discovery-metadata cache directory supplied, the +/// SECOND proxy connect for the same `(server, url)` must reuse the cached +/// `.well-known` OAuth metadata instead of re-running discovery against the +/// (throttle-prone) upstream. We assert the discovery endpoint is hit exactly +/// ONCE across two connects. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn discovery_metadata_cache_skips_rediscovery() { + let (base, state) = spawn_mock(false).await; + let url = format!("{base}/mcp"); + // Ignore discovery hits from the readiness probe in `spawn_mock`. + *state.well_known_calls.lock().await = 0; + + let cache = tempfile::tempdir().expect("cache tempdir"); + + // First connect: cache miss → one live discovery, then cached. + let secrets1 = seed_token("krisp", "local", &url, ACCESS_TOKEN_FRESH, 3600, true); + let proxy1 = + animus_mcp_oauth::proxy::McpProxy::connect_with_store("krisp", &url, secrets1, "local", Some(cache.path())) + .await + .expect("first proxy connect"); + proxy1.forward_request_for_test(tools_list_request()).await.expect("tools/list on first connect"); + assert_eq!(*state.well_known_calls.lock().await, 1, "first connect must discover exactly once"); + + // Second connect (same server+url+cache): cache hit → NO further discovery. + let secrets2 = seed_token("krisp", "local", &url, ACCESS_TOKEN_FRESH, 3600, true); + let proxy2 = + animus_mcp_oauth::proxy::McpProxy::connect_with_store("krisp", &url, secrets2, "local", Some(cache.path())) + .await + .expect("second proxy connect"); + proxy2.forward_request_for_test(tools_list_request()).await.expect("tools/list on second connect"); + + assert_eq!( + *state.well_known_calls.lock().await, + 1, + "second connect must reuse cached discovery metadata, not re-hit the upstream .well-known endpoint" + ); +} + +/// TASK-326 (fix a, unauthenticated gate): with NO stored token for the server, +/// connecting must fail fast WITHOUT running discovery priming. Priming a live +/// `.well-known` discovery for an unauthenticated server would regress the +/// fast-fail startup path into a needless network round-trip. We assert the +/// connect errors and the discovery endpoint is never hit. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unauthenticated_server_skips_discovery_priming() { + let (base, state) = spawn_mock(false).await; + let url = format!("{base}/mcp"); + // Ignore discovery hits from the readiness probe in `spawn_mock`. + *state.well_known_calls.lock().await = 0; + + let cache = tempfile::tempdir().expect("cache tempdir"); + + // Empty secret store: no token stored for this server. + let secrets: Arc = Arc::new(MockSecretStore::new()); + let result = + animus_mcp_oauth::proxy::McpProxy::connect_with_store("unauthed", &url, secrets, "local", Some(cache.path())) + .await; + + assert!(result.is_err(), "connect must fail when no token is stored"); + assert_eq!( + *state.well_known_calls.lock().await, + 0, + "an unauthenticated server (no stored token) must skip discovery priming and fast-fail" + ); +} diff --git a/crates/orchestrator-cli/src/bin/animus_mcp_proxy.rs b/crates/orchestrator-cli/src/bin/animus_mcp_proxy.rs index 95daf58a..b4c17c46 100644 --- a/crates/orchestrator-cli/src/bin/animus_mcp_proxy.rs +++ b/crates/orchestrator-cli/src/bin/animus_mcp_proxy.rs @@ -35,6 +35,14 @@ struct Args { #[arg(long)] url: Option, + /// The caller (CLI `mcp call`) already resolved this as an + /// `authorization_code` (keychain) server and passed the upstream `--url`. + /// With both set, the proxy TRUSTS the URL and skips the `config_source` + /// round-trip entirely — the reduction that keeps bulk `mcp call` from + /// re-saturating the source on every spawn. Ignored without `--url`. + #[arg(long = "auth-code")] + auth_code: bool, + /// Project root. Defaults to the current working directory. #[arg(long)] project_root: Option, @@ -73,8 +81,25 @@ async fn main() -> Result<()> { .unwrap_or_else(|| ".".to_string()); let root = std::path::Path::new(&project_root); + // Fast path: the caller already resolved an `authorization_code` (keychain) + // server and handed us the upstream `--url`. Trust it and skip the + // `config_source` lookup entirely — this is the amplification cut that keeps + // bulk `mcp call` from re-loading the source on every proxy spawn. Keychain + // principal + secret store come from the OS keychain / on-disk scoped state, + // never the config source. + if args.auth_code { + if let Some(url) = args.url.as_deref().map(str::trim).filter(|u| !u.is_empty()) { + return animus_mcp_oauth::proxy::run_authorization_code(&args.server, url, root).await; + } + } + // stdout is the MCP stdio channel and must carry only JSON-RPC frames. // The proxy never logs tokens; diagnostics go to stderr. + // + // The single `config_source` resolution needed to split broker vs keychain + // flows. The resolved URL is threaded onward (broker path passes it to + // `run_with_bearer_source`; keychain path to `run_authorization_code`) so + // NEITHER downstream re-resolves — one source touch per spawn, not two. let resolution = animus_mcp_oauth::resolve_server_url(root, &args.server, args.url.as_deref())?; match resolution.broker_oauth { Some(oauth) => { @@ -82,6 +107,6 @@ async fn main() -> Result<()> { Arc::new(BrokerBearerSource { server: args.server.clone(), project_root: project_root.clone(), oauth }); animus_mcp_oauth::proxy::run_with_bearer_source(&args.server, &resolution.url, source).await } - None => animus_mcp_oauth::proxy::run(root, &args.server, args.url.as_deref()).await, + None => animus_mcp_oauth::proxy::run_authorization_code(&args.server, &resolution.url, root).await, } } diff --git a/crates/orchestrator-config/src/workflow_config/config_source_client.rs b/crates/orchestrator-config/src/workflow_config/config_source_client.rs index d6ec00f6..ff575bca 100644 --- a/crates/orchestrator-config/src/workflow_config/config_source_client.rs +++ b/crates/orchestrator-config/src/workflow_config/config_source_client.rs @@ -154,6 +154,16 @@ pub fn config_source_installed(project_root: &Path) -> bool { /// caller surfaces an actionable "no config_source plugin installed" error. /// Returns `(base WorkflowConfig, cache_token_version)`. pub fn resolve_plugin_base(project_root: &Path, actor: Option<&Actor>) -> Result> { + // Test-only seam: simulate a config_source plugin whose load FAILS (spawn / + // handshake / RPC / DB overload) so tests can exercise the non-swallowing + // classification in `try_load_workflow_config` (a present-but-failing source + // must surface as `SourceUnavailable`, never as an empty config). Checked + // before the success seam so an installed-failure wins. + #[cfg(any(test, feature = "test-utils"))] + if let Some(message) = test_seam::failure_for(project_root) { + return Err(anyhow!(message)); + } + // Test-only seam: a synthetic base config injected via // `set_test_plugin_base` stands in for an installed config_source plugin so // unit tests can exercise the kernel's pack-merge + validate pipeline (which @@ -550,6 +560,13 @@ pub mod test_seam { REGISTRY.get_or_init(|| Mutex::new(HashMap::new())) } + // Roots for which `resolve_plugin_base` should simulate a config_source + // load FAILURE (an installed-but-failing plugin), keyed to an error message. + fn failure_registry() -> &'static Mutex> { + static REGISTRY: OnceLock>> = OnceLock::new(); + REGISTRY.get_or_init(|| Mutex::new(HashMap::new())) + } + // Normalize the key so a test that installs against `tempdir().path()` is // still found when the loader resolves the project root via git-common-root // (which canonicalizes, e.g. /var -> /private/var on macOS). Falls back to @@ -577,6 +594,23 @@ pub mod test_seam { ActorTestBaseGuard { key } } + /// Install a simulated config_source load FAILURE for `project_root`. + /// `resolve_plugin_base` returns `Err(message)` while the guard is held, so + /// tests can prove a present-but-failing source surfaces as + /// `SourceUnavailable` (never as an empty config). + #[must_use] + pub fn install_failure(project_root: &Path, message: &str) -> FailureGuard { + let key = normalize(project_root); + failure_registry().lock().unwrap_or_else(|p| p.into_inner()).insert(key.clone(), message.to_string()); + FailureGuard { key } + } + + /// The simulated failure message installed for `project_root`, if any. + pub fn failure_for(project_root: &Path) -> Option { + let root = normalize(project_root); + failure_registry().lock().unwrap_or_else(|p| p.into_inner()).get(&root).cloned() + } + /// Clone the synthetic base installed for `(project_root, actor)`. Prefers an /// actor-scoped base (when `actor` is present and one was installed), else /// falls back to the root-only base. @@ -612,6 +646,17 @@ pub mod test_seam { actor_registry().lock().unwrap_or_else(|p| p.into_inner()).remove(&self.key); } } + + /// RAII guard that clears the simulated config_source failure on drop. + pub struct FailureGuard { + key: PathBuf, + } + + impl Drop for FailureGuard { + fn drop(&mut self) { + failure_registry().lock().unwrap_or_else(|p| p.into_inner()).remove(&self.key); + } + } } /// v0.6 cross-crate test seam: stand in for an installed `config_source` plugin @@ -640,6 +685,16 @@ pub fn install_yaml_config_source_base(project_root: &Path) -> test_seam::TestBa test_seam::install(project_root, base) } +/// v0.6 cross-crate test seam: simulate an installed `config_source` plugin +/// whose base load FAILS, so dependent crates' tests can drive the +/// non-swallowing `try_load_workflow_config` classification (a present-but- +/// failing source must surface as `SourceUnavailable`, never as an empty +/// config). Returns a guard that clears the failure on drop. +#[cfg(any(test, feature = "test-utils"))] +pub fn install_config_source_failure(project_root: &Path, message: &str) -> test_seam::FailureGuard { + test_seam::install_failure(project_root, message) +} + #[cfg(test)] mod resident_cache_tests { use super::*; diff --git a/crates/orchestrator-config/src/workflow_config/loading.rs b/crates/orchestrator-config/src/workflow_config/loading.rs index 2e50754e..9f07f0f3 100644 --- a/crates/orchestrator-config/src/workflow_config/loading.rs +++ b/crates/orchestrator-config/src/workflow_config/loading.rs @@ -272,6 +272,88 @@ pub(crate) fn compile_workflow_config_onto_base( }) } +/// Three-way outcome of attempting to source the base workflow config, +/// distinguishing "no config source is configured" (benign — callers may fall +/// back to other config) from "the config source failed to load" (transient / +/// erroring — the config could NOT be determined). +/// +/// This is the non-swallowing counterpart to +/// [`load_workflow_config_or_default`]: the `_or_default` path degrades ANY +/// load failure into an empty builtin base, which is correct for system +/// reconcilers but catastrophic for credential/OAuth resolution — a transient +/// `config_source` plugin failure (DB overload under bulk `mcp call`) would be +/// misreported as "server not defined" with an empty `mcp_servers`. Credential +/// resolvers route through [`try_load_workflow_config`] instead so a source +/// outage propagates as [`WorkflowConfigAvailability::SourceUnavailable`] +/// rather than silently emptying the config. +pub enum WorkflowConfigAvailability { + /// The config was sourced and compiled successfully. Boxed because + /// `LoadedWorkflowConfig` is large relative to the other (small) variants. + Loaded(Box), + /// No `config_source` plugin is installed (and no injected test base): the + /// project simply has no workflow-config source. Benign — callers fall back + /// to other config (e.g. project `.animus/config.json`). + NoSource, + /// A `config_source` plugin IS present but base acquisition or compilation + /// failed (spawn / handshake / RPC / DB overload, or a validation error). + /// The config could not be determined; callers must NOT treat this as an + /// empty/absent config. + SourceUnavailable(anyhow::Error), +} + +/// Non-swallowing load: classify base acquisition into +/// [`WorkflowConfigAvailability`] so credential/OAuth resolvers can tell a +/// genuinely-absent config source apart from a transient source failure. +/// +/// Unlike [`load_workflow_config_with_metadata`] (which collapses "no plugin +/// installed" into an error) and [`load_workflow_config_or_default`] (which +/// collapses every error into an empty base), this inspects the +/// `config_source` base result directly: +/// - `Ok(None)` (no plugin / no injected base) => [`WorkflowConfigAvailability::NoSource`]. +/// - `Ok(Some(base))` => compile it (pack overlays + validate); a compile +/// failure surfaces as [`WorkflowConfigAvailability::SourceUnavailable`]. +/// - `Err(_)` (plugin present but its load failed) => +/// [`WorkflowConfigAvailability::SourceUnavailable`]. +pub fn try_load_workflow_config(project_root: &Path, actor: Option<&Actor>) -> WorkflowConfigAvailability { + let (base, cache_version) = match super::config_source_client::resolve_plugin_base(project_root, actor) { + Ok(None) => return WorkflowConfigAvailability::NoSource, + Ok(Some(plugin_base)) => plugin_base, + Err(err) => return WorkflowConfigAvailability::SourceUnavailable(err), + }; + + // Reuse the SAME compiled-config cache as `load_workflow_config_with_metadata` + // (keyed by source token + pack fingerprint + actor) so an unchanged source + // and pack set short-circuits the pack-overlay merge + validate below. + // `resolve_server_url` routes OAuth resolution through this path once per + // `mcp call`, so honoring the cache keeps a bulk burst from recompiling pack + // overlays on every call. + let registry = match resolve_pack_registry(project_root) { + Ok(registry) => registry, + Err(err) => return WorkflowConfigAvailability::SourceUnavailable(err), + }; + let compile_token = format!( + "{cache_version}\u{1f}{}\u{1f}{}", + pack_registry_fingerprint(®istry), + super::config_source_client::actor_cache_key(actor), + ); + if let Some(cached) = super::config_source_client::cached_compiled(project_root, actor, &compile_token) { + return WorkflowConfigAvailability::Loaded(Box::new(cached)); + } + + match compile_workflow_config_onto_base(project_root, base) { + Ok(mut loaded) => { + // The base is plugin-sourced (identical to + // `load_workflow_config_with_metadata`), so label it `Plugin` rather + // than `Yaml`; this also keeps the shared cache entry consistent + // across both loaders. + loaded.metadata.source = WorkflowConfigSource::Plugin; + super::config_source_client::store_compiled(project_root, actor, compile_token, loaded.clone()); + WorkflowConfigAvailability::Loaded(Box::new(loaded)) + } + Err(err) => WorkflowConfigAvailability::SourceUnavailable(err), + } +} + pub fn load_workflow_config_or_default(project_root: &Path) -> LoadedWorkflowConfig { // The default/fallback path is system-initiated (daemon reconcilers, // schedulers, CLI inspection) with no authenticated actor → the global diff --git a/crates/orchestrator-config/src/workflow_config/mod.rs b/crates/orchestrator-config/src/workflow_config/mod.rs index f9bd3224..92bd4cc7 100644 --- a/crates/orchestrator-config/src/workflow_config/mod.rs +++ b/crates/orchestrator-config/src/workflow_config/mod.rs @@ -68,7 +68,8 @@ pub use environment_routing::resolve_environment; pub use loading::{ ensure_workflow_config_compiled, ensure_workflow_config_file, legacy_workflow_config_paths, load_workflow_config, load_workflow_config_or_default, load_workflow_config_or_default_for_actor, load_workflow_config_with_metadata, - workflow_config_hash, workflow_config_path, write_workflow_config, + try_load_workflow_config, workflow_config_hash, workflow_config_path, write_workflow_config, + WorkflowConfigAvailability, }; pub use resolution::{ resolve_workflow_phase_plan, resolve_workflow_rework_attempts, resolve_workflow_skip_guards, diff --git a/crates/orchestrator-config/src/workflow_config/tests.rs b/crates/orchestrator-config/src/workflow_config/tests.rs index 28bffb11..1c6c9a0d 100644 --- a/crates/orchestrator-config/src/workflow_config/tests.rs +++ b/crates/orchestrator-config/src/workflow_config/tests.rs @@ -4120,6 +4120,77 @@ mod cache_token_short_circuit { } } +// TASK-326: the non-swallowing loader must distinguish a genuinely-absent +// config source (benign) from a present-but-failing source (transient), so +// credential/OAuth resolvers never degrade a config_source outage into an empty +// config that misreports servers as "not defined". +mod try_load_availability { + use super::*; + use crate::workflow_config::config_source_client::test_seam; + use crate::workflow_config::loading::try_load_workflow_config; + use crate::workflow_config::WorkflowConfigAvailability; + + #[test] + fn injected_base_classifies_as_loaded() { + let _lock = env_lock().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); + let temp = tempfile::tempdir().expect("tempdir"); + let _home_guard = EnvVarGuard::set("HOME", temp.path()); + let root = temp.path(); + + let mut base = builtin_workflow_config(); + base.tools_allowlist = vec!["marker-loaded".to_string()]; + let _guard = test_seam::install(root, base); + + match try_load_workflow_config(root, None) { + WorkflowConfigAvailability::Loaded(loaded) => { + assert!(loaded.config.tools_allowlist.contains(&"marker-loaded".to_string())); + } + other => panic!("expected Loaded, got {}", availability_label(&other)), + } + } + + #[test] + fn no_installed_source_classifies_as_no_source() { + let _lock = env_lock().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); + let temp = tempfile::tempdir().expect("tempdir"); + let _home_guard = EnvVarGuard::set("HOME", temp.path()); + let root = temp.path(); + + // No injected base, no installed plugin => resolve_plugin_base is + // Ok(None) => a benign "no source", NOT an error. + match try_load_workflow_config(root, None) { + WorkflowConfigAvailability::NoSource => {} + other => panic!("expected NoSource, got {}", availability_label(&other)), + } + } + + #[test] + fn failing_source_classifies_as_source_unavailable() { + let _lock = env_lock().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); + let temp = tempfile::tempdir().expect("tempdir"); + let _home_guard = EnvVarGuard::set("HOME", temp.path()); + let root = temp.path(); + + // A present-but-failing source (DB overload under bulk mcp call) must + // NOT be swallowed into an empty config. + let _guard = test_seam::install_failure(root, "config_source RPC timed out"); + match try_load_workflow_config(root, None) { + WorkflowConfigAvailability::SourceUnavailable(err) => { + assert!(err.to_string().contains("config_source RPC timed out"), "err: {err}"); + } + other => panic!("expected SourceUnavailable, got {}", availability_label(&other)), + } + } + + fn availability_label(a: &WorkflowConfigAvailability) -> &'static str { + match a { + WorkflowConfigAvailability::Loaded(_) => "Loaded", + WorkflowConfigAvailability::NoSource => "NoSource", + WorkflowConfigAvailability::SourceUnavailable(_) => "SourceUnavailable", + } + } +} + mod unenforced_field_warnings { use super::super::validation::{unenforced_project_yaml_warnings, unenforced_yaml_field_warnings}; use std::fs;