From 1d5777824c617450ac2ce685d15b52c99ef69db3 Mon Sep 17 00:00:00 2001 From: jr42 Date: Thu, 26 Mar 2026 22:47:31 +0100 Subject: [PATCH 1/3] fix(mcp): handle 202 Accepted and wire session manager for Streamable HTTP (#1437) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(mcp): handle 202 Accepted for Streamable HTTP notifications The MCP Streamable HTTP spec requires servers to respond with 202 Accepted (empty body) for JSON-RPC notifications like `notifications/initialized`. The HTTP transport tried to parse this empty body as JSON, which failed and broke the session handshake — subsequent requests like `tools/list` were rejected because the server considered the session uninitialized. Add an early return for 202 responses that produces an empty McpResponse without attempting body parsing. Fixes #1436 Co-Authored-By: Claude Opus 4.6 (1M context) * fix(mcp): wire session manager into transport for non-OAuth HTTP clients The factory used McpClient::new_with_config().with_session_manager() which only set the session manager on the client, not on the HttpMcpTransport. The transport never captured Mcp-Session-Id from responses, so subsequent requests lacked the header and the server rejected them as uninitialized. Fix by constructing the HttpMcpTransport with the session manager before wrapping it in Arc, matching the pattern already used by new_authenticated(). Co-Authored-By: Claude Opus 4.6 (1M context) * refactor(mcp): deduplicate factory HTTP path, gate dead-code methods as test-only - Collapse the two identical non-OAuth HTTP branches in `create_client_from_config()` into one (early-return for the authenticated path, fall through for the common case). - Gate `McpClient::new_with_config()` and `McpClient::with_session_manager()` as `#[cfg(test)]` — the factory was their only production caller and no longer uses them. Both methods silently skip wiring the session manager into the transport, which was the root cause of #1436. - Add doc warnings on both methods explaining the footgun. Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: ilblackdragon@gmail.com --- src/tools/mcp/client.rs | 20 +++++- src/tools/mcp/factory.rs | 117 +++++++++++++++++++++++++++----- src/tools/mcp/http_transport.rs | 28 ++++++++ 3 files changed, 148 insertions(+), 17 deletions(-) diff --git a/src/tools/mcp/client.rs b/src/tools/mcp/client.rs index 148f5a86..32c5767d 100644 --- a/src/tools/mcp/client.rs +++ b/src/tools/mcp/client.rs @@ -117,6 +117,11 @@ impl McpClient { /// The config must use HTTP transport (the default); for stdio/UDS use `new_with_transport`. /// /// Returns an error if the config uses a non-HTTP transport. + /// + /// **Note:** The session manager is NOT wired into the transport. For + /// production use, prefer `create_client_from_config()` which constructs + /// the transport with session tracking. + #[cfg(test)] pub fn new_with_config(config: McpServerConfig) -> Result { if !matches!( config.effective_transport(), @@ -214,7 +219,14 @@ impl McpClient { } } - /// Attach a session manager for Streamable HTTP session tracking. + /// Attach a session manager to the **client** only. + /// + /// **Warning:** This does NOT wire the session manager into the underlying + /// `HttpMcpTransport`, so the transport will not capture `Mcp-Session-Id` + /// from responses. For production use, construct the transport with + /// `HttpMcpTransport::with_session_manager()` and pass it to + /// `new_with_transport()` instead. See `create_client_from_config()`. + #[cfg(test)] pub fn with_session_manager(mut self, session_manager: Arc) -> Self { self.session_manager = Some(session_manager); self @@ -235,6 +247,12 @@ impl McpClient { self.session_manager.is_some() } + /// Get the underlying transport (test-only). + #[cfg(test)] + pub(crate) fn transport(&self) -> &Arc { + &self.transport + } + /// Get the next request ID. fn next_request_id(&self) -> u64 { self.next_id.fetch_add(1, Ordering::SeqCst) diff --git a/src/tools/mcp/factory.rs b/src/tools/mcp/factory.rs index c31c5051..bbb09256 100644 --- a/src/tools/mcp/factory.rs +++ b/src/tools/mcp/factory.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use crate::secrets::SecretsStore; use crate::tools::mcp::config::{EffectiveTransport, McpServerConfig}; +use crate::tools::mcp::http_transport::HttpMcpTransport; use crate::tools::mcp::{McpClient, McpProcessManager, McpSessionManager, McpTransport}; /// Error returned when MCP client creation fails. @@ -78,33 +79,37 @@ pub async fn create_client_from_config( Err(McpFactoryError::UnixNotSupported { name: server_name }) } EffectiveTransport::Http => { + // Authenticated (OAuth) path: tokens exist or server requires auth. if let Some(ref secrets) = secrets { let has_tokens = crate::tools::mcp::is_authenticated(&server, secrets, user_id).await; if has_tokens || server.requires_auth() { - Ok(McpClient::new_authenticated( + return Ok(McpClient::new_authenticated( server, Arc::clone(session_manager), Arc::clone(secrets), user_id, - )) - } else { - Ok(McpClient::new_with_config(server) - .map_err(|e| McpFactoryError::InvalidConfig { - name: server_name.clone(), - reason: e.to_string(), - })? - .with_session_manager(Arc::clone(session_manager))) + )); } - } else { - Ok(McpClient::new_with_config(server) - .map_err(|e| McpFactoryError::InvalidConfig { - name: server_name, - reason: e.to_string(), - })? - .with_session_manager(Arc::clone(session_manager))) } + + // Non-OAuth HTTP: wire the session manager into the *transport* so + // it captures `Mcp-Session-Id` from responses. Passing it only to + // the client (via `with_session_manager`) is not enough — the + // transport must know about it to read/write the header. + let transport = Arc::new( + HttpMcpTransport::new(server.url.clone(), server.name.clone()) + .with_session_manager(Arc::clone(session_manager)), + ); + Ok(McpClient::new_with_transport( + server.name.clone(), + transport, + Some(Arc::clone(session_manager)), + secrets, + user_id, + Some(server), + )) } } } @@ -134,4 +139,84 @@ mod tests { "non-OAuth HTTP clients must carry a session manager" ); } + + /// Regression test: the factory must wire the session manager into the + /// *transport*, not just the client. Otherwise the transport never + /// captures `Mcp-Session-Id` from responses and subsequent requests + /// lack the header, causing the server to reject them. + #[tokio::test] + async fn test_factory_non_oauth_http_transport_captures_session_id() { + use axum::http::header::HeaderName; + use axum::{Router, http::StatusCode, response::IntoResponse, routing::post}; + use tokio::net::TcpListener; + + const SESSION_ID: &str = "test-session-abc123"; + + async fn session_echo() -> impl IntoResponse { + let body = serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "result": {} + }) + .to_string(); + ( + StatusCode::OK, + [( + HeaderName::from_static("mcp-session-id"), + SESSION_ID.to_string(), + )], + body, + ) + } + + let app = Router::new().route("/", post(session_echo)); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let url = format!("http://127.0.0.1:{}", addr.port()); + + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + + let server = McpServerConfig::new("session-test", &url); + let session_manager = Arc::new(McpSessionManager::new()); + let process_manager = Arc::new(McpProcessManager::new()); + + let client = create_client_from_config( + server, + &session_manager, + &process_manager, + None, + "test-user", + ) + .await + .expect("factory should succeed for HTTP config"); + + // Pre-create a session entry so that update_session_id has something to update. + // In production, the MCP initialize handshake calls get_or_create before responses arrive. + session_manager.get_or_create("session-test", &url).await; + + // Send a request through the client's transport to trigger session capture. + use crate::tools::mcp::protocol::McpRequest; + let request = McpRequest { + jsonrpc: "2.0".to_string(), + id: Some(1), + method: "test".to_string(), + params: Some(serde_json::json!({})), + }; + let headers = std::collections::HashMap::new(); + client + .transport() + .send(&request, &headers) + .await + .expect("request should succeed"); + + // Verify the session manager captured the session ID from the response. + let captured = session_manager.get_session_id("session-test").await; + assert_eq!( + captured.as_deref(), + Some(SESSION_ID), + "transport must capture Mcp-Session-Id into session manager" + ); + } } diff --git a/src/tools/mcp/http_transport.rs b/src/tools/mcp/http_transport.rs index 59873ce4..ea3e1c03 100644 --- a/src/tools/mcp/http_transport.rs +++ b/src/tools/mcp/http_transport.rs @@ -494,6 +494,34 @@ mod tests { assert_eq!(echoed["authorization"], "Bearer oauth-token"); } + /// Regression test for #1436: 202 Accepted responses for notifications + /// were parsed as JSON, causing "Failed to parse MCP response" errors + /// that broke the MCP session handshake. + #[tokio::test] + async fn test_wire_202_accepted_for_notification() { + use axum::{Router, http::StatusCode, routing::post}; + use tokio::net::TcpListener; + + async fn accept_notification() -> StatusCode { + StatusCode::ACCEPTED + } + + let app = Router::new().route("/", post(accept_notification)); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let url = format!("http://127.0.0.1:{}", addr.port()); + + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + + let transport = HttpMcpTransport::new(&url, "test-202"); + let request = McpRequest::initialized_notification(); + let response = transport.send(&request, &HashMap::new()).await.unwrap(); + assert!(response.result.is_none()); + assert!(response.error.is_none()); + } + #[tokio::test] async fn test_wire_custom_auth_preserved_when_no_per_request_auth() { let (url, _handle) = spawn_echo_server().await; From dd0a0e10abebcd7c161c6e86fb89b8bd06e38592 Mon Sep 17 00:00:00 2001 From: Nige Date: Thu, 26 Mar 2026 23:20:01 +0000 Subject: [PATCH 2/3] fix(routines): recover delete name after failed update fallback (#1108) Co-authored-by: ilblackdragon@gmail.com Co-authored-by: Claude Opus 4.6 (1M context) --- src/context/state.rs | 3 + src/tools/builtin/routine.rs | 38 +++++++++- tests/e2e_builtin_tool_coverage.rs | 43 +++++++++++- .../routine_update_fail_delete_fallback.json | 70 +++++++++++++++++++ 4 files changed, 148 insertions(+), 6 deletions(-) create mode 100644 tests/fixtures/llm_traces/tools/routine_update_fail_delete_fallback.json diff --git a/src/context/state.rs b/src/context/state.rs index f5307947..0bb1f29a 100644 --- a/src/context/state.rs +++ b/src/context/state.rs @@ -192,6 +192,9 @@ pub struct JobContext { /// but subsequent tools (e.g., `json`) may need the full output. This /// stash stores the complete, unsanitized output so tools can reference /// previous results by ID via `$tool_call_id` parameter syntax. + /// + /// Also used for cross-tool implicit state (keys prefixed with `__`) such + /// as `__routine_last_name` for fallback recovery in routine tool chains. #[serde(skip)] pub tool_output_stash: Arc>>, /// User's preferred timezone (IANA name, e.g. "America/New_York"). Defaults to "UTC". diff --git a/src/tools/builtin/routine.rs b/src/tools/builtin/routine.rs index bbc24139..76f6e38b 100644 --- a/src/tools/builtin/routine.rs +++ b/src/tools/builtin/routine.rs @@ -650,6 +650,23 @@ pub(crate) fn routine_update_parameters_schema() -> Value { }) } +const ROUTINE_LAST_NAME_STASH_KEY: &str = "__routine_last_name"; + +async fn stash_last_routine_name(ctx: &JobContext, name: &str) { + ctx.tool_output_stash + .write() + .await + .insert(ROUTINE_LAST_NAME_STASH_KEY.to_string(), name.to_string()); +} + +async fn restore_last_routine_name(ctx: &JobContext) -> Option { + ctx.tool_output_stash + .read() + .await + .get(ROUTINE_LAST_NAME_STASH_KEY) + .cloned() +} + fn nested_object<'a>(params: &'a Value, field: &str) -> Option<&'a Map> { params.get(field).and_then(Value::as_object) } @@ -1093,6 +1110,7 @@ impl Tool for RoutineCreateTool { ) -> Result { let start = std::time::Instant::now(); let normalized = parse_routine_create_request(¶ms)?; + stash_last_routine_name(ctx, &normalized.name).await; let trigger = build_routine_trigger(&normalized.trigger); let action = build_routine_action(&normalized.name, &normalized.prompt, &normalized.execution); @@ -1274,6 +1292,7 @@ impl Tool for RoutineUpdateTool { let start = std::time::Instant::now(); let name = require_str(¶ms, "name")?; + stash_last_routine_name(ctx, name).await; let mut routine = self .store @@ -1411,11 +1430,24 @@ impl Tool for RoutineDeleteTool { ) -> Result { let start = std::time::Instant::now(); - let name = require_str(¶ms, "name")?; + let name = if let Some(name) = params.get("name").and_then(|v| v.as_str()) { + if name.trim().is_empty() { + return Err(ToolError::InvalidParameters( + "'name' parameter cannot be empty".to_string(), + )); + } + name.to_string() + } else { + restore_last_routine_name(ctx).await.ok_or_else(|| { + ToolError::InvalidParameters( + "missing 'name' parameter and no previous routine target to infer".to_string(), + ) + })? + }; let routine = self .store - .get_routine_by_name(&ctx.user_id, name) + .get_routine_by_name(&ctx.user_id, &name) .await .map_err(|e| ToolError::ExecutionFailed(format!("DB error: {e}")))? .ok_or_else(|| ToolError::ExecutionFailed(format!("routine '{}' not found", name)))?; @@ -1430,7 +1462,7 @@ impl Tool for RoutineDeleteTool { self.engine.refresh_event_cache().await; let result = serde_json::json!({ - "name": name, + "name": &name, "deleted": deleted, }); diff --git a/tests/e2e_builtin_tool_coverage.rs b/tests/e2e_builtin_tool_coverage.rs index 1c3cc6a2..7c0c7bc7 100644 --- a/tests/e2e_builtin_tool_coverage.rs +++ b/tests/e2e_builtin_tool_coverage.rs @@ -205,7 +205,44 @@ mod tests { } // ----------------------------------------------------------------------- - // Test 5: routine_manual_create_defaults_to_tools_enabled + // Test 5: routine_update_fail_delete_fallback + // ----------------------------------------------------------------------- + + #[tokio::test] + async fn routine_update_fail_delete_fallback() { + let trace = LlmTrace::from_file(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/llm_traces/tools/routine_update_fail_delete_fallback.json" + )) + .expect("failed to load routine_update_fail_delete_fallback.json"); + + let rig = TestRigBuilder::new() + .with_trace(trace.clone()) + .with_auto_approve_tools(true) + .build() + .await; + + rig.send_message("Try converting a routine trigger, then recover by deleting it") + .await; + let responses = rig.wait_for_responses(1, Duration::from_secs(15)).await; + + rig.verify_trace_expects(&trace, &responses); + + let completed = rig.tool_calls_completed(); + assert!( + completed.iter().any(|(n, ok)| n == "routine_update" && !ok), + "routine_update should fail in this regression path: {completed:?}" + ); + assert!( + completed.iter().any(|(n, ok)| n == "routine_delete" && *ok), + "routine_delete should recover successfully via preserved routine identity: {completed:?}" + ); + + rig.shutdown(); + } + + // ----------------------------------------------------------------------- + // Test 6: routine_manual_create_defaults_to_tools_enabled // ----------------------------------------------------------------------- #[tokio::test] @@ -246,7 +283,7 @@ mod tests { } // ----------------------------------------------------------------------- - // Test 6: routine_manual_create_explicit_no_tools + // Test 7: routine_manual_create_explicit_no_tools // ----------------------------------------------------------------------- #[tokio::test] @@ -287,7 +324,7 @@ mod tests { } // ----------------------------------------------------------------------- - // Test 7: routine_history + // Test 8: routine_history // ----------------------------------------------------------------------- #[tokio::test] diff --git a/tests/fixtures/llm_traces/tools/routine_update_fail_delete_fallback.json b/tests/fixtures/llm_traces/tools/routine_update_fail_delete_fallback.json new file mode 100644 index 00000000..5c76dbb5 --- /dev/null +++ b/tests/fixtures/llm_traces/tools/routine_update_fail_delete_fallback.json @@ -0,0 +1,70 @@ +{ + "model_name": "test-routine-update-fail-delete-fallback", + "expects": { + "tools_used": ["routine_create", "routine_update", "routine_delete"], + "tool_results_contain": { + "routine_update": "Cannot update schedule or timezone on a non-cron routine.", + "routine_delete": "temp-routine" + }, + "min_responses": 1 + }, + "steps": [ + { + "response": { + "type": "tool_calls", + "tool_calls": [ + { + "id": "call_rc_fallback", + "name": "routine_create", + "arguments": { + "name": "temp-routine", + "trigger_type": "manual", + "prompt": "Temporary routine for fallback test." + } + } + ], + "input_tokens": 120, + "output_tokens": 40 + } + }, + { + "response": { + "type": "tool_calls", + "tool_calls": [ + { + "id": "call_ru_fallback", + "name": "routine_update", + "arguments": { + "name": "temp-routine", + "schedule": "0 */10 * * * *" + } + } + ], + "input_tokens": 200, + "output_tokens": 30 + } + }, + { + "response": { + "type": "tool_calls", + "tool_calls": [ + { + "id": "call_rd_fallback", + "name": "routine_delete", + "arguments": {} + } + ], + "input_tokens": 300, + "output_tokens": 20 + } + }, + { + "response": { + "type": "text", + "content": "I recovered from the failed update and cleaned up the original routine.", + "input_tokens": 380, + "output_tokens": 25 + } + } + ] +} From 5b95d222186f9ee8f89edf69480000ca42f0d7d0 Mon Sep 17 00:00:00 2001 From: Henry Park Date: Thu, 26 Mar 2026 16:45:31 -0700 Subject: [PATCH 3/3] Support direct hosted OAuth callbacks with proxy auth token (#1684) * Support direct hosted OAuth callbacks with proxy auth token * Make OAuth env tests panic-safe * Preserve public OAuth field compatibility * Fix OAuth proxy token whitespace fallback --- src/channels/web/server.rs | 436 ++++++++++++++++++++++++++++++++++++- src/cli/oauth_defaults.rs | 189 +++++++++++++++- src/extensions/manager.rs | 13 +- src/tools/wasm/loader.rs | 55 ++++- src/tools/wasm/wrapper.rs | 18 +- 5 files changed, 689 insertions(+), 22 deletions(-) diff --git a/src/channels/web/server.rs b/src/channels/web/server.rs index 26c005d4..06870ace 100644 --- a/src/channels/web/server.rs +++ b/src/channels/web/server.rs @@ -836,10 +836,10 @@ async fn oauth_callback_handler( let result: Result<(), String> = async { let token_response = if let Some(proxy_url) = &exchange_proxy_url { - let gateway_token = flow.gateway_token.as_deref().unwrap_or_default(); + let oauth_proxy_auth_token = flow.oauth_proxy_auth_token().unwrap_or_default(); oauth_defaults::exchange_via_proxy(oauth_defaults::ProxyTokenExchangeRequest { proxy_url, - gateway_token, + gateway_token: oauth_proxy_auth_token, token_url: &flow.token_url, client_id: &flow.client_id, client_secret: flow.client_secret.as_deref(), @@ -3057,6 +3057,160 @@ mod tests { .with_state(state) } + #[derive(Clone, Debug)] + struct RecordedOauthProxyRequest { + authorization: Option, + form: std::collections::HashMap, + } + + #[derive(Clone)] + struct MockOauthProxyState { + requests: Arc>>, + } + + struct MockOauthProxyServer { + addr: std::net::SocketAddr, + requests: Arc>>, + shutdown_tx: Option>, + server_task: Option>, + } + + impl MockOauthProxyServer { + async fn start() -> Self { + async fn exchange_handler( + State(state): State, + headers: axum::http::HeaderMap, + axum::Form(form): axum::Form>, + ) -> Json { + state.requests.lock().await.push(RecordedOauthProxyRequest { + authorization: headers + .get(axum::http::header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .map(str::to_string), + form, + }); + Json(serde_json::json!({ + "access_token": "proxy-access-token", + "refresh_token": "proxy-refresh-token", + "expires_in": 7200 + })) + } + + let requests = Arc::new(tokio::sync::Mutex::new(Vec::new())); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind mock oauth proxy"); + let addr = listener.local_addr().expect("mock oauth proxy addr"); + let app = Router::new() + .route("/oauth/exchange", post(exchange_handler)) + .with_state(MockOauthProxyState { + requests: Arc::clone(&requests), + }); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>(); + let server_task = tokio::spawn(async move { + let _ = axum::serve(listener, app) + .with_graceful_shutdown(async { + let _ = shutdown_rx.await; + }) + .await; + }); + + Self { + addr, + requests, + shutdown_tx: Some(shutdown_tx), + server_task: Some(server_task), + } + } + + fn base_url(&self) -> String { + format!("http://{}", self.addr) + } + + async fn requests(&self) -> Vec { + self.requests.lock().await.clone() + } + + async fn shutdown(mut self) { + if let Some(tx) = self.shutdown_tx.take() { + let _ = tx.send(()); + } + if let Some(task) = self.server_task.take() { + let _ = task.await; + } + } + } + + impl Drop for MockOauthProxyServer { + fn drop(&mut self) { + if let Some(tx) = self.shutdown_tx.take() { + let _ = tx.send(()); + } + if let Some(task) = self.server_task.take() { + task.abort(); + } + } + } + + struct EnvVarGuard { + key: &'static str, + original: Option, + } + + impl Drop for EnvVarGuard { + fn drop(&mut self) { + // SAFETY: Tests use lock_env() to serialize environment access. + unsafe { + if let Some(ref value) = self.original { + std::env::set_var(self.key, value); + } else { + std::env::remove_var(self.key); + } + } + } + } + + fn set_env_var(key: &'static str, value: Option<&str>) -> EnvVarGuard { + let original = std::env::var(key).ok(); + // SAFETY: Tests use lock_env() to serialize environment access. + unsafe { + if let Some(value) = value { + std::env::set_var(key, value); + } else { + std::env::remove_var(key); + } + } + EnvVarGuard { key, original } + } + + fn fresh_pending_oauth_flow( + secrets: Arc, + sse_manager: Option>, + oauth_proxy_auth_token: Option, + ) -> crate::cli::oauth_defaults::PendingOAuthFlow { + crate::cli::oauth_defaults::PendingOAuthFlow { + extension_name: "test_tool".to_string(), + display_name: "Test Tool".to_string(), + token_url: "https://example.com/token".to_string(), + client_id: "client123".to_string(), + client_secret: None, + redirect_uri: "https://example.com/oauth/callback".to_string(), + code_verifier: Some("test-code-verifier".to_string()), + access_token_field: "access_token".to_string(), + secret_name: "test_token".to_string(), + provider: Some("google".to_string()), + validation_endpoint: None, + scopes: vec!["email".to_string()], + user_id: "test".to_string(), + secrets, + sse_manager, + gateway_token: oauth_proxy_auth_token, + token_exchange_extra_params: std::collections::HashMap::new(), + client_id_secret_name: None, + created_at: std::time::Instant::now(), + } + } + #[tokio::test] async fn test_extensions_setup_submit_returns_failure_when_not_activated() { use axum::body::Body; @@ -3714,6 +3868,284 @@ mod tests { ); } + #[tokio::test] + async fn test_oauth_callback_accepts_versioned_hosted_state_without_instance_name() { + use axum::body::Body; + use tower::ServiceExt; + + let secrets: Arc = + Arc::new(crate::secrets::InMemorySecretsStore::new(Arc::new( + crate::secrets::SecretsCrypto::new(secrecy::SecretString::from( + TEST_GATEWAY_CRYPTO_KEY.to_string(), + )) + .expect("crypto"), + ))); + let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(secrets.clone()); + + let Some(created_at) = expired_flow_created_at() else { + eprintln!( + "Skipping versioned OAuth state without instance test: monotonic uptime below expiry window" + ); + return; + }; + let flow = crate::cli::oauth_defaults::PendingOAuthFlow { + extension_name: "test_tool".to_string(), + display_name: "Test Tool".to_string(), + token_url: "https://example.com/token".to_string(), + client_id: "client123".to_string(), + client_secret: None, + redirect_uri: "https://example.com/oauth/callback".to_string(), + code_verifier: None, + access_token_field: "access_token".to_string(), + secret_name: "test_token".to_string(), + provider: None, + validation_endpoint: None, + scopes: vec![], + user_id: "test".to_string(), + secrets, + sse_manager: None, + gateway_token: None, + token_exchange_extra_params: std::collections::HashMap::new(), + client_id_secret_name: None, + created_at, + }; + + ext_mgr + .pending_oauth_flows() + .write() + .await + .insert("test_nonce".to_string(), flow); + + let state = test_gateway_state(Some(ext_mgr.clone())); + let app = test_oauth_router(state); + let versioned_state = + crate::cli::oauth_defaults::encode_hosted_oauth_state("test_nonce", None); + + let req = axum::http::Request::builder() + .uri(format!( + "/oauth/callback?code=fake_code&state={}", + urlencoding::encode(&versioned_state) + )) + .body(Body::empty()) + .expect("request"); + + let resp = ServiceExt::>::oneshot(app, req) + .await + .expect("response"); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), 1024 * 64) + .await + .expect("body"); + let html = String::from_utf8_lossy(&body); + assert!(html.contains("Authorization Failed")); + assert!( + ext_mgr + .pending_oauth_flows() + .read() + .await + .get("test_nonce") + .is_none() + ); + } + + #[allow(clippy::await_holding_lock)] + #[tokio::test] + async fn test_oauth_callback_happy_path_with_gateway_token_fallback() { + use axum::body::Body; + use tower::ServiceExt; + + let proxy = MockOauthProxyServer::start().await; + // Keep the process-wide env locked for the full callback so the handler + // sees a stable proxy URL/token configuration throughout the test. + let _env_guard = crate::config::helpers::lock_env(); + let _exchange_url_guard = + set_env_var("IRONCLAW_OAUTH_EXCHANGE_URL", Some(&proxy.base_url())); + let _proxy_auth_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None); + let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-test-token")); + + let secrets = test_secrets_store(); + let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(Arc::clone(&secrets)); + let sse_mgr = Arc::new(SseManager::new()); + let mut receiver = sse_mgr.sender().subscribe(); + let flow = fresh_pending_oauth_flow( + Arc::clone(&secrets), + Some(Arc::clone(&sse_mgr)), + crate::cli::oauth_defaults::oauth_proxy_auth_token(), + ); + + ext_mgr + .pending_oauth_flows() + .write() + .await + .insert("test_nonce".to_string(), flow); + + let state = test_gateway_state(Some(ext_mgr.clone())); + let app = test_oauth_router(state); + let versioned_state = + crate::cli::oauth_defaults::encode_hosted_oauth_state("test_nonce", Some("myinstance")); + + let req = axum::http::Request::builder() + .uri(format!( + "/oauth/callback?code=fake_code&state={}", + urlencoding::encode(&versioned_state) + )) + .body(Body::empty()) + .expect("request"); + + let resp = ServiceExt::>::oneshot(app, req) + .await + .expect("response"); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), 1024 * 64) + .await + .expect("body"); + let html = String::from_utf8_lossy(&body); + assert!(html.contains("Test Tool Connected")); + + let requests = proxy.requests().await; + assert_eq!(requests.len(), 1); + assert_eq!( + requests[0].authorization.as_deref(), + Some("Bearer gateway-test-token") + ); + assert_eq!( + requests[0].form.get("code").map(String::as_str), + Some("fake_code") + ); + assert_eq!( + requests[0].form.get("code_verifier").map(String::as_str), + Some("test-code-verifier") + ); + + let access_token = secrets + .get_decrypted("test", "test_token") + .await + .expect("access token stored"); + assert_eq!(access_token.expose(), "proxy-access-token"); + + let refresh_token = secrets + .get_decrypted("test", "test_token_refresh_token") + .await + .expect("refresh token stored"); + assert_eq!(refresh_token.expose(), "proxy-refresh-token"); + + match receiver.recv().await.expect("auth_completed event").event { + crate::channels::web::types::AppEvent::AuthCompleted { + extension_name, + success, + .. + } => { + assert_eq!(extension_name, "test_tool"); + assert!(success, "OAuth callback should broadcast success"); + } + event => panic!("expected AuthCompleted event, got {event:?}"), + } + + proxy.shutdown().await; + } + + #[allow(clippy::await_holding_lock)] + #[tokio::test] + async fn test_oauth_callback_happy_path_with_dedicated_proxy_auth_token() { + use axum::body::Body; + use tower::ServiceExt; + + let proxy = MockOauthProxyServer::start().await; + // Keep the process-wide env locked for the full callback so the handler + // sees a stable proxy URL/token configuration throughout the test. + let _env_guard = crate::config::helpers::lock_env(); + let _exchange_url_guard = + set_env_var("IRONCLAW_OAUTH_EXCHANGE_URL", Some(&proxy.base_url())); + let _proxy_auth_guard = set_env_var( + "IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", + Some("shared-oauth-proxy-secret"), + ); + let _gateway_token_guard = set_env_var("GATEWAY_AUTH_TOKEN", None); + + let secrets = test_secrets_store(); + let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(Arc::clone(&secrets)); + let sse_mgr = Arc::new(SseManager::new()); + let mut receiver = sse_mgr.sender().subscribe(); + let flow = fresh_pending_oauth_flow( + Arc::clone(&secrets), + Some(Arc::clone(&sse_mgr)), + crate::cli::oauth_defaults::oauth_proxy_auth_token(), + ); + + ext_mgr + .pending_oauth_flows() + .write() + .await + .insert("test_nonce".to_string(), flow); + + let state = test_gateway_state(Some(ext_mgr.clone())); + let app = test_oauth_router(state); + let versioned_state = + crate::cli::oauth_defaults::encode_hosted_oauth_state("test_nonce", None); + + let req = axum::http::Request::builder() + .uri(format!( + "/oauth/callback?code=fake_code&state={}", + urlencoding::encode(&versioned_state) + )) + .body(Body::empty()) + .expect("request"); + + let resp = ServiceExt::>::oneshot(app, req) + .await + .expect("response"); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), 1024 * 64) + .await + .expect("body"); + let html = String::from_utf8_lossy(&body); + assert!(html.contains("Test Tool Connected")); + + let requests = proxy.requests().await; + assert_eq!(requests.len(), 1); + assert_eq!( + requests[0].authorization.as_deref(), + Some("Bearer shared-oauth-proxy-secret") + ); + assert_eq!( + requests[0].form.get("code").map(String::as_str), + Some("fake_code") + ); + assert_eq!( + requests[0].form.get("code_verifier").map(String::as_str), + Some("test-code-verifier") + ); + + let access_token = secrets + .get_decrypted("test", "test_token") + .await + .expect("access token stored"); + assert_eq!(access_token.expose(), "proxy-access-token"); + + let refresh_token = secrets + .get_decrypted("test", "test_token_refresh_token") + .await + .expect("refresh token stored"); + assert_eq!(refresh_token.expose(), "proxy-refresh-token"); + + match receiver.recv().await.expect("auth_completed event").event { + crate::channels::web::types::AppEvent::AuthCompleted { + extension_name, + success, + .. + } => { + assert_eq!(extension_name, "test_tool"); + assert!(success, "OAuth callback should broadcast success"); + } + event => panic!("expected AuthCompleted event, got {event:?}"), + } + + proxy.shutdown().await; + } + // --- Slack relay OAuth CSRF tests --- fn test_relay_oauth_router(state: Arc) -> Router { diff --git a/src/cli/oauth_defaults.rs b/src/cli/oauth_defaults.rs index e9001909..384d5833 100644 --- a/src/cli/oauth_defaults.rs +++ b/src/cli/oauth_defaults.rs @@ -473,7 +473,8 @@ pub struct PendingOAuthFlow { pub secrets: Arc, /// SSE broadcast manager for notifying the web UI. pub sse_manager: Option>, - /// Gateway auth token for authenticating with the platform token exchange proxy. + /// OAuth proxy auth token for authenticating with the hosted token exchange proxy. + /// Kept as `gateway_token` for public API compatibility. pub gateway_token: Option, /// Additional form params for the token exchange request. /// Used for provider-specific requirements such as RFC 8707 `resource`. @@ -496,6 +497,12 @@ impl std::fmt::Debug for PendingOAuthFlow { } } +impl PendingOAuthFlow { + pub fn oauth_proxy_auth_token(&self) -> Option<&str> { + self.gateway_token.as_deref() + } +} + /// Thread-safe registry of pending OAuth flows, keyed by CSRF `state` parameter. pub type PendingOAuthRegistry = Arc>>; @@ -529,6 +536,22 @@ pub fn exchange_proxy_url() -> Option { .filter(|url| !url.is_empty()) } +/// Returns the configured OAuth proxy auth token, if any. +/// +/// New hosted infra can inject a dedicated shared proxy secret via +/// `IRONCLAW_OAUTH_PROXY_AUTH_TOKEN`. Existing hosted instances continue to +/// work by falling back to `GATEWAY_AUTH_TOKEN`. +pub fn oauth_proxy_auth_token() -> Option { + fn normalized_env_value(key: &str) -> Option { + crate::config::helpers::env_or_override(key) + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) + } + + normalized_env_value("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN") + .or_else(|| normalized_env_value("GATEWAY_AUTH_TOKEN")) +} + /// Maximum age for pending OAuth flows (5 minutes, matching TCP listener timeout). pub const OAUTH_FLOW_EXPIRY: Duration = Duration::from_secs(300); @@ -674,6 +697,8 @@ pub fn strip_instance_prefix(state: &str) -> &str { pub struct ProxyTokenExchangeRequest<'a> { pub proxy_url: &'a str, + /// OAuth proxy auth token. + /// Kept as `gateway_token` for public API compatibility. pub gateway_token: &'a str, pub token_url: &'a str, pub client_id: &'a str, @@ -687,6 +712,8 @@ pub struct ProxyTokenExchangeRequest<'a> { pub struct ProxyRefreshTokenRequest<'a> { pub proxy_url: &'a str, + /// OAuth proxy auth token. + /// Kept as `gateway_token` for public API compatibility. pub gateway_token: &'a str, pub token_url: &'a str, pub client_id: &'a str, @@ -729,7 +756,7 @@ fn oauth_token_response_from_json( /// Exchange an OAuth authorization code via the platform's token exchange proxy. /// -/// Authenticated via the gateway auth token (Bearer header). The caller may +/// Authenticated via an OAuth proxy auth token (Bearer header). The caller may /// either rely on proxy-side secret lookup or forward a `client_secret` when /// the provider requires it. /// @@ -741,7 +768,7 @@ pub async fn exchange_via_proxy( ) -> Result { if request.gateway_token.is_empty() { return Err(OAuthCallbackError::Io( - "Gateway auth token is required for proxy token exchange".to_string(), + "OAuth proxy auth token is required for proxy token exchange".to_string(), )); } let exchange_url = format!("{}/oauth/exchange", request.proxy_url.trim_end_matches('/')); @@ -796,7 +823,7 @@ pub async fn exchange_via_proxy( /// Refresh an OAuth access token via the platform's token refresh proxy. /// -/// Authenticated via the gateway auth token (Bearer header). The caller may +/// Authenticated via an OAuth proxy auth token (Bearer header). The caller may /// either rely on proxy-side secret lookup or forward a `client_secret` when /// the provider requires it. pub async fn refresh_token_via_proxy( @@ -804,7 +831,7 @@ pub async fn refresh_token_via_proxy( ) -> Result { if request.gateway_token.is_empty() { return Err(OAuthCallbackError::Io( - "Gateway auth token is required for proxy token refresh".to_string(), + "OAuth proxy auth token is required for proxy token refresh".to_string(), )); } @@ -1010,6 +1037,37 @@ mod tests { } } + struct EnvVarGuard { + key: &'static str, + original: Option, + } + + impl Drop for EnvVarGuard { + fn drop(&mut self) { + // SAFETY: Under ENV_MUTEX, no concurrent env access. + unsafe { + if let Some(ref value) = self.original { + std::env::set_var(self.key, value); + } else { + std::env::remove_var(self.key); + } + } + } + } + + fn set_env_var(key: &'static str, value: Option<&str>) -> EnvVarGuard { + let original = std::env::var(key).ok(); + // SAFETY: Under ENV_MUTEX, no concurrent env access. + unsafe { + if let Some(value) = value { + std::env::set_var(key, value); + } else { + std::env::remove_var(key); + } + } + EnvVarGuard { key, original } + } + #[test] fn test_hosted_proxy_client_secret_suppresses_builtin_secret() { let builtin = builtin_credentials("google_oauth_token").expect("google builtin creds"); @@ -1030,6 +1088,79 @@ mod tests { assert_eq!(result, client_secret); } + #[tokio::test] + async fn test_exchange_via_proxy_sends_auth_and_form() { + let server = MockProxyServer::start().await; + let mut extra_token_params = HashMap::new(); + extra_token_params.insert("resource".to_string(), "https://mcp.notion.com".to_string()); + + let response = super::exchange_via_proxy(super::ProxyTokenExchangeRequest { + proxy_url: &server.base_url(), + gateway_token: "shared-oauth-proxy-secret", + code: "auth-code-123", + redirect_uri: "https://oauth.example.com/oauth/callback", + token_url: "https://oauth2.googleapis.com/token", + client_id: TEST_OAUTH_CLIENT_ID, + client_secret: Some(TEST_OAUTH_CLIENT_SECRET), + access_token_field: "access_token", + code_verifier: Some("code-verifier-123"), + extra_token_params: &extra_token_params, + }) + .await + .expect("proxy exchange succeeds"); + + assert_eq!(response.access_token, "proxy-access-token"); + assert_eq!( + response.refresh_token.as_deref(), + Some("proxy-refresh-token") + ); + assert_eq!(response.expires_in, Some(7200)); + + let requests = server.requests().await; + assert_eq!(requests.len(), 1); + assert_eq!( + requests[0].authorization.as_deref(), + Some("Bearer shared-oauth-proxy-secret") + ); + assert_eq!( + requests[0].form.get("code").map(String::as_str), + Some("auth-code-123") + ); + assert_eq!( + requests[0].form.get("redirect_uri").map(String::as_str), + Some("https://oauth.example.com/oauth/callback") + ); + assert_eq!( + requests[0].form.get("token_url").map(String::as_str), + Some("https://oauth2.googleapis.com/token") + ); + assert_eq!( + requests[0].form.get("client_id").map(String::as_str), + Some(TEST_OAUTH_CLIENT_ID) + ); + assert_eq!( + requests[0].form.get("client_secret").map(String::as_str), + Some(TEST_OAUTH_CLIENT_SECRET) + ); + assert_eq!( + requests[0] + .form + .get("access_token_field") + .map(String::as_str), + Some("access_token") + ); + assert_eq!( + requests[0].form.get("code_verifier").map(String::as_str), + Some("code-verifier-123") + ); + assert_eq!( + requests[0].form.get("resource").map(String::as_str), + Some("https://mcp.notion.com") + ); + + server.shutdown().await; + } + #[tokio::test] async fn test_refresh_token_via_proxy_sends_auth_and_form() { let server = MockProxyServer::start().await; @@ -1535,6 +1666,54 @@ mod tests { } } + #[test] + fn test_oauth_proxy_auth_token_prefers_dedicated_env() { + let _guard = lock_env(); + let _proxy_guard = set_env_var( + "IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", + Some("shared-proxy-secret"), + ); + let _gateway_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-token")); + + assert_eq!( + crate::cli::oauth_defaults::oauth_proxy_auth_token().as_deref(), + Some("shared-proxy-secret") + ); + } + + #[test] + fn test_oauth_proxy_auth_token_falls_back_to_gateway_token() { + let _guard = lock_env(); + let _proxy_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None); + let _gateway_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-token")); + + assert_eq!( + crate::cli::oauth_defaults::oauth_proxy_auth_token().as_deref(), + Some("gateway-token") + ); + } + + #[test] + fn test_oauth_proxy_auth_token_whitespace_dedicated_env_falls_back_to_gateway_token() { + let _guard = lock_env(); + let _proxy_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", Some(" ")); + let _gateway_guard = set_env_var("GATEWAY_AUTH_TOKEN", Some("gateway-token")); + + assert_eq!( + crate::cli::oauth_defaults::oauth_proxy_auth_token().as_deref(), + Some("gateway-token") + ); + } + + #[test] + fn test_oauth_proxy_auth_token_returns_none_when_unset() { + let _guard = lock_env(); + let _proxy_guard = set_env_var("IRONCLAW_OAUTH_PROXY_AUTH_TOKEN", None); + let _gateway_guard = set_env_var("GATEWAY_AUTH_TOKEN", None); + + assert_eq!(crate::cli::oauth_defaults::oauth_proxy_auth_token(), None); + } + #[test] fn test_strip_instance_prefix_with_colon() { use crate::cli::oauth_defaults::strip_instance_prefix; diff --git a/src/extensions/manager.rs b/src/extensions/manager.rs index 47b45a0f..55b1e96d 100644 --- a/src/extensions/manager.rs +++ b/src/extensions/manager.rs @@ -403,9 +403,10 @@ pub struct ExtensionManager { /// when running in gateway mode, consumed by the web gateway's /// `/oauth/callback` handler. pending_oauth_flows: crate::cli::oauth_defaults::PendingOAuthRegistry, - /// Gateway auth token for authenticating with the platform token exchange proxy. - /// Read once at construction from `GATEWAY_AUTH_TOKEN` env var. - gateway_token: Option, + /// OAuth proxy auth token for authenticating with the hosted token exchange proxy. + /// Resolved once at construction from `IRONCLAW_OAUTH_PROXY_AUTH_TOKEN`, + /// then `GATEWAY_AUTH_TOKEN` as a backward-compatible fallback. + oauth_proxy_auth_token: Option, /// Relay config captured at startup. Used by `auth_channel_relay` and /// `activate_channel_relay` instead of re-reading env vars. relay_config: Option, @@ -535,7 +536,7 @@ impl ExtensionManager { activation_errors: RwLock::new(HashMap::new()), sse_manager: RwLock::new(None), pending_oauth_flows: crate::cli::oauth_defaults::new_pending_oauth_registry(), - gateway_token: std::env::var("GATEWAY_AUTH_TOKEN").ok(), + oauth_proxy_auth_token: crate::cli::oauth_defaults::oauth_proxy_auth_token(), relay_config: crate::config::RelayConfig::from_env(), relay_event_tx: Arc::new(tokio::sync::Mutex::new(None)), relay_signing_secret_cache: Arc::new(std::sync::Mutex::new(None)), @@ -2788,7 +2789,7 @@ impl ExtensionManager { user_id: user_id.to_string(), secrets: Arc::clone(&self.secrets), sse_manager: self.sse_manager.read().await.clone(), - gateway_token: self.gateway_token.clone(), + gateway_token: self.oauth_proxy_auth_token.clone(), token_exchange_extra_params, client_id_secret_name: if server.oauth.is_none() { Some(server.client_id_secret_name()) @@ -3305,7 +3306,7 @@ impl ExtensionManager { user_id: user_id.to_string(), secrets: Arc::clone(&self.secrets), sse_manager: self.sse_manager.read().await.clone(), - gateway_token: self.gateway_token.clone(), + gateway_token: self.oauth_proxy_auth_token.clone(), token_exchange_extra_params: std::collections::HashMap::new(), client_id_secret_name: None, created_at: std::time::Instant::now(), diff --git a/src/tools/wasm/loader.rs b/src/tools/wasm/loader.rs index 2a7ed040..4876dc1b 100644 --- a/src/tools/wasm/loader.rs +++ b/src/tools/wasm/loader.rs @@ -446,16 +446,14 @@ fn resolve_oauth_refresh_config(cap_file: &CapabilitiesFile) -> Option, /// Hosted OAuth proxy base URL (e.g., "http://host.docker.internal:8080"). pub exchange_proxy_url: Option, - /// Gateway auth token for authenticating with the hosted OAuth proxy. + /// OAuth proxy auth token for authenticating with the hosted OAuth proxy. + /// Kept as `gateway_token` for public API compatibility. pub gateway_token: Option, /// Secret name of the access token (e.g., "google_oauth_token"). /// The refresh token lives at `{secret_name}_refresh_token`. @@ -71,6 +72,12 @@ pub struct OAuthRefreshConfig { pub provider: Option, } +impl OAuthRefreshConfig { + fn oauth_proxy_auth_token(&self) -> Option<&str> { + self.gateway_token.as_deref() + } +} + /// Pre-resolved credential for host-based injection. /// /// Built before each WASM execution by decrypting secrets from the store. @@ -1218,9 +1225,9 @@ async fn refresh_oauth_token( let refresh_name = format!("{}_refresh_token", config.secret_name); if let Some(proxy_url) = config.exchange_proxy_url.as_deref() { - let Some(gateway_token) = config.gateway_token.as_deref() else { + let Some(oauth_proxy_auth_token) = config.oauth_proxy_auth_token() else { tracing::warn!( - "OAuth refresh proxy is configured, but no gateway auth token is available" + "OAuth refresh proxy is configured, but no OAuth proxy auth token is available" ); return false; }; @@ -1235,7 +1242,7 @@ async fn refresh_oauth_token( let token_response = match oauth_defaults::refresh_token_via_proxy( oauth_defaults::ProxyRefreshTokenRequest { proxy_url, - gateway_token, + gateway_token: oauth_proxy_auth_token, token_url: &config.token_url, client_id: &config.client_id, client_secret: config.client_secret.as_deref(), @@ -2704,7 +2711,8 @@ mod tests { } #[tokio::test] - async fn test_resolve_host_credentials_skips_refresh_token_lookup_without_gateway_token() { + async fn test_resolve_host_credentials_skips_refresh_token_lookup_without_oauth_proxy_auth_token() + { use crate::secrets::{ CreateSecretParams, CredentialLocation, CredentialMapping, SecretsStore, };