mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
fix(mcp): open MCP OAuth in same browser as gateway (#951)
* fix(mcp): use gateway callback for MCP OAuth so auth opens in same browser When MCP OAuth is triggered from the web gateway, the auth URL was being opened via `open::that()` which launches the OS default browser instead of the browser already running the gateway UI. This changes the MCP OAuth flow to use the same gateway callback pattern as WASM extensions: in gateway mode, the auth URL is returned to the frontend via SSE and opened with `window.open()`, keeping the user in the same browser. Also adds RFC 8707 `resource` parameter support to the gateway token exchange path, scoping issued tokens to the correct MCP server. Closes #299 Co-Authored-By: Claude Opus 4.6 <[email protected]> * style: cargo fmt Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix(mcp): persist DCR client_id in gateway OAuth callback for token refresh The gateway callback handler stored access and refresh tokens but not the DCR client_id. When the token expired, refresh failed with "No client ID found" because get_client_id() could not find it in secrets. Adds client_id_secret_name to PendingOAuthFlow so the gateway callback handler persists the client_id alongside the tokens, matching the behavior of the CLI flow in authorize_mcp_server(). Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix(mcp): return AuthRequired on 401 so activate triggers OAuth flow activate_mcp() returned ActivationFailed for all errors including 401 auth responses, so the activate handler never triggered the OAuth flow. Now 401/auth errors return AuthRequired, which the handler detects and redirects to the OAuth flow — matching the WASM extension pattern. Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix(mcp): fix gateway OAuth flow, approval cards, and auto-activation - Add explicit gateway_mode flag on ExtensionManager (set at startup by web gateway) so MCP OAuth returns auth URLs to the frontend instead of calling open::that() on the server machine. - Auto-activate extensions after successful OAuth callback so the UI transitions from "Activate" to "Active" without a second click. - Send ApprovalNeeded status (not generic "Awaiting approval") from thread_ops.rs for all three NeedApproval paths so the web UI shows approval cards for deferred tool calls. - Remove duplicate ApprovalNeeded send from agent_loop.rs (thread_ops.rs is now the canonical sender). - Skip approval for tool_auth in gateway mode since it only returns a URL. - Revert fragile active-server detection heuristic from system prompt. Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: address PR review findings - Use Release/Acquire ordering for gateway_mode AtomicBool instead of Relaxed to ensure visibility across threads. - Report activation failure as error in OAuth callback SSE event instead of silently falling back to the success message. - Fix EnvGuard::drop to remove env var when original was unset. - Replace hardcoded /tmp/ path with std::env::temp_dir() in test helper. Co-Authored-By: Claude Opus 4.6 <[email protected]> * test(mcp): add E2E trace test for MCP extension lifecycle with mock server Add a full MCP extension lifecycle E2E test that exercises: - Turn 1: tool_search → tool_install → text (extension discovery and install) - Token injection + activate (simulating OAuth completion) - Turn 2: MCP tool calls (notion-search → notion-fetch → text) Includes a mock MCP server (tests/support/mock_mcp_server.rs) with OAuth discovery, DCR, token exchange, and JSON-RPC endpoints. The mock server validates Bearer auth and serves pre-configured tool responses. Also adds inject_registry_entry() to ExtensionManager for test use and exposes extension_manager from TestRig. Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: address PR review findings (round 2) - Only fall back to manual token entry on AuthNotSupported, propagate real errors from auth_mcp_build_url() instead of masking them - Use mcp:-prefixed provider string in PendingOAuthFlow for consistency with CLI MCP auth token storage - Only persist client_id_secret_name for DCR flows (not pre-configured OAuth) - Fix gateway_callback_redirect_uri to use /oauth/callback path - Bypass exchange proxy when flow has RFC 8707 resource parameter - Remove client_id double-prefix in oauth callback handler - Remove weak tests that didn't exercise production logic - Add clarifying comments for exchange_oauth_code delegation Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: keep OAuth success independent of activation, fix wait_for_responses scoping - OAuth success is now reported accurately even when auto-activation fails (tokens are already stored, so auth succeeded) - E2E test waits for turn1_count + 1 responses to ensure turn-2 behavior is actually observed Co-Authored-By: Claude Opus 4.6 <[email protected]> --------- Co-authored-by: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Illia Polosukhin
co-authored by
Claude Opus 4.6
Claude Opus 4.6
parent
0b81342b5c
commit
8a26cfae73
Executable
+23
@@ -0,0 +1,23 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
# Pre-push hook: run clippy and tests before pushing.
|
||||
# Install: git config core.hooksPath .githooks
|
||||
|
||||
echo "pre-push: running clippy..."
|
||||
if ! cargo clippy --all --benches --tests --examples --all-features -- -D warnings; then
|
||||
echo ""
|
||||
echo "Push blocked: clippy warnings found."
|
||||
echo "To bypass: git push --no-verify"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "pre-push: running tests..."
|
||||
if ! cargo test; then
|
||||
echo ""
|
||||
echo "Push blocked: tests failed."
|
||||
echo "To bypass: git push --no-verify"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "pre-push: all checks passed."
|
||||
+5
-24
@@ -18,7 +18,7 @@ use crate::agent::self_repair::{DefaultSelfRepair, RepairResult, SelfRepair};
|
||||
use crate::agent::session_manager::SessionManager;
|
||||
use crate::agent::submission::{Submission, SubmissionParser, SubmissionResult};
|
||||
use crate::agent::{HeartbeatConfig as AgentHeartbeatConfig, Router, Scheduler};
|
||||
use crate::channels::{ChannelManager, IncomingMessage, OutgoingResponse, StatusUpdate};
|
||||
use crate::channels::{ChannelManager, IncomingMessage, OutgoingResponse};
|
||||
use crate::config::{AgentConfig, HeartbeatConfig, RoutineConfig, SkillsConfig};
|
||||
use crate::context::ContextManager;
|
||||
use crate::db::Database;
|
||||
@@ -936,29 +936,10 @@ impl Agent {
|
||||
SubmissionResult::Ok { message } => Ok(message),
|
||||
SubmissionResult::Error { message } => Ok(Some(format!("Error: {}", message))),
|
||||
SubmissionResult::Interrupted => Ok(Some("Interrupted.".into())),
|
||||
SubmissionResult::NeedApproval {
|
||||
request_id,
|
||||
tool_name,
|
||||
description,
|
||||
parameters,
|
||||
} => {
|
||||
// Each channel renders the approval prompt via send_status.
|
||||
// Web gateway shows an inline card, REPL prints a formatted prompt, etc.
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::ApprovalNeeded {
|
||||
request_id: request_id.to_string(),
|
||||
tool_name,
|
||||
description,
|
||||
parameters,
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
|
||||
// Empty string signals the caller to skip respond() (no duplicate text)
|
||||
SubmissionResult::NeedApproval { .. } => {
|
||||
// ApprovalNeeded status was already sent by thread_ops.rs before
|
||||
// returning this result. Empty string signals the caller to skip
|
||||
// respond() (no duplicate text).
|
||||
Ok(Some(String::new()))
|
||||
}
|
||||
}
|
||||
|
||||
+18
-3
@@ -486,7 +486,12 @@ impl Agent {
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::Status("Awaiting approval".into()),
|
||||
StatusUpdate::ApprovalNeeded {
|
||||
request_id: request_id.to_string(),
|
||||
tool_name: tool_name.clone(),
|
||||
description: description.clone(),
|
||||
parameters: parameters.clone(),
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
@@ -1297,7 +1302,12 @@ impl Agent {
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::Status("Awaiting approval".into()),
|
||||
StatusUpdate::ApprovalNeeded {
|
||||
request_id: request_id.to_string(),
|
||||
tool_name: tool_name.clone(),
|
||||
description: description.clone(),
|
||||
parameters: parameters.clone(),
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
@@ -1368,7 +1378,12 @@ impl Agent {
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::Status("Awaiting approval".into()),
|
||||
StatusUpdate::ApprovalNeeded {
|
||||
request_id: request_id.to_string(),
|
||||
tool_name: tool_name.clone(),
|
||||
description: description.clone(),
|
||||
parameters: parameters.clone(),
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -581,7 +581,12 @@ async fn oauth_callback_handler(
|
||||
let exchange_proxy_url = std::env::var("IRONCLAW_OAUTH_EXCHANGE_URL").ok();
|
||||
|
||||
let result: Result<(), String> = async {
|
||||
let token_response = if let Some(ref proxy_url) = exchange_proxy_url {
|
||||
let token_response = if let (Some(proxy_url), None) = (&exchange_proxy_url, &flow.resource)
|
||||
{
|
||||
// Use the platform exchange proxy when configured and no resource
|
||||
// parameter is needed. The proxy holds client_secret server-side so
|
||||
// the container never sees it. MCP flows (resource.is_some()) bypass
|
||||
// the proxy because it doesn't forward the RFC 8707 resource param.
|
||||
let gateway_token = flow.gateway_token.as_deref().unwrap_or_default();
|
||||
oauth_defaults::exchange_via_proxy(
|
||||
proxy_url,
|
||||
@@ -594,7 +599,10 @@ async fn oauth_callback_handler(
|
||||
.await
|
||||
.map_err(|e| e.to_string())?
|
||||
} else {
|
||||
oauth_defaults::exchange_oauth_code(
|
||||
// Direct token exchange: uses exchange_oauth_code_with_resource so MCP
|
||||
// flows can include the RFC 8707 `resource` parameter to scope the
|
||||
// issued token to the specific MCP server.
|
||||
oauth_defaults::exchange_oauth_code_with_resource(
|
||||
&flow.token_url,
|
||||
&flow.client_id,
|
||||
flow.client_secret.as_deref(),
|
||||
@@ -602,6 +610,7 @@ async fn oauth_callback_handler(
|
||||
&flow.redirect_uri,
|
||||
flow.code_verifier.as_deref(),
|
||||
&flow.access_token_field,
|
||||
flow.resource.as_deref(),
|
||||
)
|
||||
.await
|
||||
.map_err(|e| e.to_string())?
|
||||
@@ -628,6 +637,19 @@ async fn oauth_callback_handler(
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
|
||||
// For MCP OAuth flows (identified by resource field), persist the
|
||||
// client_id so token refresh works without re-authentication.
|
||||
// The CLI flow stores this in authorize_mcp_server(); the gateway
|
||||
// callback must do the same.
|
||||
if let Some(ref client_id_secret) = flow.client_id_secret_name {
|
||||
let params = crate::secrets::CreateSecretParams::new(client_id_secret, &flow.client_id)
|
||||
.with_provider(flow.provider.as_ref().cloned().unwrap_or_default());
|
||||
flow.secrets
|
||||
.create(&flow.user_id, params)
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
.await;
|
||||
@@ -659,12 +681,35 @@ async fn oauth_callback_handler(
|
||||
}
|
||||
}
|
||||
|
||||
// After successful OAuth, auto-activate the extension so it moves
|
||||
// from "Installed (Authenticate)" → "Active" without a second click.
|
||||
// OAuth success is independent of activation — tokens are already stored.
|
||||
// Report auth as successful and attempt activation as a bonus step.
|
||||
let final_message = if success {
|
||||
match ext_mgr.activate(&flow.extension_name).await {
|
||||
Ok(result) => result.message,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
extension = %flow.extension_name,
|
||||
error = %e,
|
||||
"Auto-activation after OAuth failed"
|
||||
);
|
||||
format!(
|
||||
"{} authenticated successfully. Activation failed: {}. Try activating manually.",
|
||||
flow.display_name, e
|
||||
)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
message
|
||||
};
|
||||
|
||||
// Broadcast SSE event to notify the web UI
|
||||
if let Some(ref sender) = flow.sse_sender {
|
||||
let _ = sender.send(SseEvent::AuthCompleted {
|
||||
extension_name: flow.extension_name,
|
||||
success,
|
||||
message,
|
||||
message: final_message.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
@@ -2966,6 +3011,8 @@ mod tests {
|
||||
secrets,
|
||||
sse_sender: None,
|
||||
gateway_token: None,
|
||||
resource: None,
|
||||
client_id_secret_name: None,
|
||||
created_at: std::time::Instant::now()
|
||||
.checked_sub(std::time::Duration::from_secs(600))
|
||||
.expect("System uptime is too low to run expired flow test"),
|
||||
@@ -3075,6 +3122,8 @@ mod tests {
|
||||
secrets,
|
||||
sse_sender: None,
|
||||
gateway_token: None,
|
||||
resource: None,
|
||||
client_id_secret_name: None,
|
||||
// Expired — handler will reject after lookup (no network I/O)
|
||||
created_at: std::time::Instant::now()
|
||||
.checked_sub(std::time::Duration::from_secs(600))
|
||||
|
||||
@@ -172,6 +172,35 @@ pub async fn exchange_oauth_code(
|
||||
redirect_uri: &str,
|
||||
code_verifier: Option<&str>,
|
||||
access_token_field: &str,
|
||||
) -> Result<OAuthTokenResponse, OAuthCallbackError> {
|
||||
// Delegates to exchange_oauth_code_with_resource with resource=None.
|
||||
// Non-MCP OAuth flows don't need the RFC 8707 resource parameter.
|
||||
exchange_oauth_code_with_resource(
|
||||
token_url,
|
||||
client_id,
|
||||
client_secret,
|
||||
code,
|
||||
redirect_uri,
|
||||
code_verifier,
|
||||
access_token_field,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Exchange an OAuth authorization code for tokens, with optional RFC 8707 `resource` parameter.
|
||||
///
|
||||
/// The `resource` parameter scopes the issued token to a specific server (used by MCP OAuth).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn exchange_oauth_code_with_resource(
|
||||
token_url: &str,
|
||||
client_id: &str,
|
||||
client_secret: Option<&str>,
|
||||
code: &str,
|
||||
redirect_uri: &str,
|
||||
code_verifier: Option<&str>,
|
||||
access_token_field: &str,
|
||||
resource: Option<&str>,
|
||||
) -> Result<OAuthTokenResponse, OAuthCallbackError> {
|
||||
let client = reqwest::Client::new();
|
||||
let mut token_params = vec![
|
||||
@@ -184,6 +213,12 @@ pub async fn exchange_oauth_code(
|
||||
token_params.push(("code_verifier", verifier.to_string()));
|
||||
}
|
||||
|
||||
// RFC 8707: include the `resource` parameter so the authorization server
|
||||
// scopes the issued token to the specific MCP server (protected resource).
|
||||
if let Some(resource) = resource {
|
||||
token_params.push(("resource", resource.to_string()));
|
||||
}
|
||||
|
||||
let mut request = client.post(token_url);
|
||||
|
||||
if let Some(secret) = client_secret {
|
||||
@@ -388,6 +423,12 @@ pub struct PendingOAuthFlow {
|
||||
pub sse_sender: Option<tokio::sync::broadcast::Sender<crate::channels::web::types::SseEvent>>,
|
||||
/// Gateway auth token for authenticating with the platform token exchange proxy.
|
||||
pub gateway_token: Option<String>,
|
||||
/// RFC 8707 resource parameter (MCP OAuth only).
|
||||
/// Sent during token exchange to scope the token to a specific MCP server.
|
||||
pub resource: Option<String>,
|
||||
/// Secret name for persisting the client ID (MCP OAuth only).
|
||||
/// Needed so token refresh can find the client_id after the session ends.
|
||||
pub client_id_secret_name: Option<String>,
|
||||
/// When this flow was created (for expiry).
|
||||
pub created_at: std::time::Instant,
|
||||
}
|
||||
@@ -975,4 +1016,42 @@ mod tests {
|
||||
assert_eq!(strip_instance_prefix("abc123"), "abc123");
|
||||
assert_eq!(strip_instance_prefix(""), "");
|
||||
}
|
||||
|
||||
/// Verify that `build_oauth_url` includes the RFC 8707 `resource` parameter
|
||||
/// when passed through `extra_params`, which is how MCP OAuth gateway mode
|
||||
/// scopes tokens to a specific MCP server.
|
||||
#[test]
|
||||
fn test_build_oauth_url_includes_resource_via_extra_params() {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::cli::oauth_defaults::build_oauth_url;
|
||||
|
||||
let mut extra = HashMap::new();
|
||||
extra.insert(
|
||||
"resource".to_string(),
|
||||
"https://mcp.example.com".to_string(),
|
||||
);
|
||||
|
||||
let result = build_oauth_url(
|
||||
"https://auth.example.com/authorize",
|
||||
"client-123",
|
||||
"https://gateway.example.com/oauth/callback",
|
||||
&["read".to_string()],
|
||||
true,
|
||||
&extra,
|
||||
);
|
||||
|
||||
// The resource parameter should be URL-encoded in the auth URL
|
||||
assert!(
|
||||
result
|
||||
.url
|
||||
.contains("resource=https%3A%2F%2Fmcp.example.com"),
|
||||
"Expected resource param in URL: {}",
|
||||
result.url
|
||||
);
|
||||
// State and PKCE should be present
|
||||
assert!(result.url.contains("state="));
|
||||
assert!(result.url.contains("code_challenge="));
|
||||
assert!(result.code_verifier.is_some());
|
||||
}
|
||||
}
|
||||
|
||||
+432
-44
@@ -27,7 +27,7 @@ use crate::secrets::{CreateSecretParams, SecretsStore};
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::tools::mcp::McpClient;
|
||||
use crate::tools::mcp::auth::{
|
||||
PkceChallenge, authorize_mcp_server, build_authorization_url, discover_full_oauth_metadata,
|
||||
authorize_mcp_server, canonical_resource_uri, discover_full_oauth_metadata,
|
||||
find_available_port, is_authenticated, register_client,
|
||||
};
|
||||
use crate::tools::mcp::config::McpServerConfig;
|
||||
@@ -108,6 +108,13 @@ pub struct ExtensionManager {
|
||||
/// Relay config captured at startup. Used by `auth_channel_relay` and
|
||||
/// `activate_channel_relay` instead of re-reading env vars.
|
||||
relay_config: Option<crate::config::RelayConfig>,
|
||||
/// When `true`, OAuth flows always return an auth URL to the caller
|
||||
/// instead of opening a browser on the server via `open::that()`.
|
||||
/// Set by the web gateway at startup via `enable_gateway_mode()`.
|
||||
gateway_mode: std::sync::atomic::AtomicBool,
|
||||
/// The gateway's own base URL for building OAuth redirect URIs.
|
||||
/// Set by the web gateway at startup via `enable_gateway_mode()`.
|
||||
gateway_base_url: RwLock<Option<String>>,
|
||||
}
|
||||
|
||||
/// Sanitize a URL for logging by removing query parameters and credentials.
|
||||
@@ -181,9 +188,75 @@ impl ExtensionManager {
|
||||
pending_oauth_flows: crate::cli::oauth_defaults::new_pending_oauth_registry(),
|
||||
gateway_token: std::env::var("GATEWAY_AUTH_TOKEN").ok(),
|
||||
relay_config: crate::config::RelayConfig::from_env(),
|
||||
gateway_mode: std::sync::atomic::AtomicBool::new(false),
|
||||
gateway_base_url: RwLock::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// Enable gateway mode so OAuth flows return auth URLs to the frontend
|
||||
/// instead of calling `open::that()` on the server.
|
||||
///
|
||||
/// `base_url` is the gateway's own public URL (e.g. `https://my-gateway.example.com`),
|
||||
/// used to build OAuth redirect URIs when `IRONCLAW_OAUTH_CALLBACK_URL` is not set.
|
||||
pub async fn enable_gateway_mode(&self, base_url: String) {
|
||||
self.gateway_mode
|
||||
.store(true, std::sync::atomic::Ordering::Release);
|
||||
*self.gateway_base_url.write().await = Some(base_url);
|
||||
}
|
||||
|
||||
/// Returns `true` if OAuth should use gateway mode (return auth URL to
|
||||
/// frontend) rather than CLI mode (open browser on server via `open::that`).
|
||||
///
|
||||
/// Gateway mode is active when any of:
|
||||
/// - `enable_gateway_mode()` was called (web gateway is running), OR
|
||||
/// - `IRONCLAW_OAUTH_CALLBACK_URL` is set to a non-loopback URL, OR
|
||||
/// - `self.tunnel_url` is set to a non-loopback URL
|
||||
pub fn should_use_gateway_mode(&self) -> bool {
|
||||
if self.gateway_mode.load(std::sync::atomic::Ordering::Acquire) {
|
||||
return true;
|
||||
}
|
||||
if crate::cli::oauth_defaults::use_gateway_callback() {
|
||||
return true;
|
||||
}
|
||||
self.tunnel_url
|
||||
.as_ref()
|
||||
.filter(|u| !u.is_empty())
|
||||
.and_then(|raw| url::Url::parse(raw).ok())
|
||||
.and_then(|u| u.host_str().map(String::from))
|
||||
.map(|host| !crate::cli::oauth_defaults::is_loopback_host(&host))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// Returns the OAuth redirect URI for gateway mode, or `None` for local mode.
|
||||
///
|
||||
/// Priority:
|
||||
/// 1. `IRONCLAW_OAUTH_CALLBACK_URL` env var (via `callback_url()`)
|
||||
/// 2. `gateway_base_url` (set by `enable_gateway_mode()`)
|
||||
/// 3. `tunnel_url` (from config)
|
||||
/// 4. `None` (local/CLI mode)
|
||||
async fn gateway_callback_redirect_uri(&self) -> Option<String> {
|
||||
use crate::cli::oauth_defaults;
|
||||
if oauth_defaults::use_gateway_callback() {
|
||||
return Some(format!("{}/oauth/callback", oauth_defaults::callback_url()));
|
||||
}
|
||||
// Use gateway_base_url from enable_gateway_mode()
|
||||
if let Some(ref base) = *self.gateway_base_url.read().await {
|
||||
let base = base.trim_end_matches('/');
|
||||
return Some(format!("{}/oauth/callback", base));
|
||||
}
|
||||
// Fall back to tunnel_url
|
||||
self.tunnel_url
|
||||
.as_ref()
|
||||
.filter(|u| !u.is_empty())
|
||||
.and_then(|raw| url::Url::parse(raw).ok())
|
||||
.and_then(|u| u.host_str().map(String::from))
|
||||
.filter(|host| !oauth_defaults::is_loopback_host(host))
|
||||
.map(|_| {
|
||||
let base = self.tunnel_url.as_ref().unwrap().trim_end_matches('/');
|
||||
format!("{}/oauth/callback", base)
|
||||
})
|
||||
}
|
||||
|
||||
/// Get the relay config stored at startup.
|
||||
fn relay_config(&self) -> Result<&crate::config::RelayConfig, ExtensionError> {
|
||||
self.relay_config.as_ref().ok_or_else(|| {
|
||||
@@ -193,6 +266,12 @@ impl ExtensionManager {
|
||||
})
|
||||
}
|
||||
|
||||
/// Inject a registry entry for testing. The entry is added to the discovery
|
||||
/// cache so it appears in search results alongside built-in entries.
|
||||
pub async fn inject_registry_entry(&self, entry: crate::extensions::RegistryEntry) {
|
||||
self.registry.cache_discovered(vec![entry]).await;
|
||||
}
|
||||
|
||||
/// Configure the channel runtime infrastructure for hot-activating WASM channels.
|
||||
///
|
||||
/// Call after construction (and after wrapping in `Arc`) once the channel
|
||||
@@ -1684,19 +1763,13 @@ impl ExtensionManager {
|
||||
return Ok(AuthResult::authenticated(name, ExtensionKind::McpServer));
|
||||
}
|
||||
|
||||
// Run the full OAuth flow (opens browser, waits for callback)
|
||||
match authorize_mcp_server(&server, &self.secrets, &self.user_id).await {
|
||||
Ok(_token) => {
|
||||
tracing::info!("MCP server '{}' authenticated via OAuth", name);
|
||||
Ok(AuthResult::authenticated(name, ExtensionKind::McpServer))
|
||||
}
|
||||
Err(crate::tools::mcp::auth::AuthError::NotSupported) => {
|
||||
// Server doesn't support OAuth, try building a URL first
|
||||
match self.auth_mcp_build_url(name, &server).await {
|
||||
// In gateway mode, build an auth URL and return it for the frontend to
|
||||
// open in the same browser. The gateway's /oauth/callback handler will
|
||||
// complete the token exchange.
|
||||
if self.should_use_gateway_mode() {
|
||||
return match self.auth_mcp_build_url(name, &server).await {
|
||||
Ok(result) => Ok(result),
|
||||
Err(_) => {
|
||||
// No OAuth, no DCR: fall back to manual token entry
|
||||
Ok(AuthResult::awaiting_token(
|
||||
Err(ExtensionError::AuthNotSupported(_)) => Ok(AuthResult::awaiting_token(
|
||||
name,
|
||||
ExtensionKind::McpServer,
|
||||
format!(
|
||||
@@ -1705,8 +1778,31 @@ impl ExtensionManager {
|
||||
name
|
||||
),
|
||||
None,
|
||||
))
|
||||
)),
|
||||
Err(e) => Err(e),
|
||||
};
|
||||
}
|
||||
|
||||
// CLI/local mode: run the full blocking OAuth flow (opens browser, waits for callback)
|
||||
match authorize_mcp_server(&server, &self.secrets, &self.user_id).await {
|
||||
Ok(_token) => {
|
||||
tracing::info!("MCP server '{}' authenticated via OAuth", name);
|
||||
Ok(AuthResult::authenticated(name, ExtensionKind::McpServer))
|
||||
}
|
||||
Err(crate::tools::mcp::auth::AuthError::NotSupported) => {
|
||||
// Server doesn't support OAuth, try building a URL
|
||||
match self.auth_mcp_build_url(name, &server).await {
|
||||
Ok(result) => Ok(result),
|
||||
Err(_) => Ok(AuthResult::awaiting_token(
|
||||
name,
|
||||
ExtensionKind::McpServer,
|
||||
format!(
|
||||
"Server '{}' does not support OAuth. \
|
||||
Please provide an API token/key for this server.",
|
||||
name
|
||||
),
|
||||
None,
|
||||
)),
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
@@ -1725,8 +1821,12 @@ impl ExtensionManager {
|
||||
}
|
||||
}
|
||||
|
||||
/// Build an auth URL for cases where non-interactive auth is needed
|
||||
/// (e.g., running via Telegram where we can't open a browser).
|
||||
/// Build an auth URL for MCP OAuth.
|
||||
///
|
||||
/// In gateway mode, stores a `PendingOAuthFlow` so the web gateway's
|
||||
/// `/oauth/callback` handler can complete the token exchange — the auth
|
||||
/// URL is sent to the frontend which opens it in the same browser.
|
||||
/// In local/CLI mode, builds the URL for the user to open manually.
|
||||
async fn auth_mcp_build_url(
|
||||
&self,
|
||||
name: &str,
|
||||
@@ -1734,45 +1834,119 @@ impl ExtensionManager {
|
||||
) -> Result<AuthResult, ExtensionError> {
|
||||
// Try to discover OAuth metadata and build a URL the user can open manually
|
||||
let metadata = discover_full_oauth_metadata(&server.url)
|
||||
.await
|
||||
.map_err(|e| match e {
|
||||
crate::tools::mcp::auth::AuthError::NotSupported => {
|
||||
ExtensionError::AuthNotSupported(e.to_string())
|
||||
}
|
||||
_ => ExtensionError::AuthFailed(e.to_string()),
|
||||
})?;
|
||||
|
||||
use crate::cli::oauth_defaults;
|
||||
|
||||
let is_gateway = self.should_use_gateway_mode();
|
||||
|
||||
// Build redirect URI: gateway uses the public callback URL,
|
||||
// local mode binds a random port.
|
||||
let redirect_uri = if let Some(uri) = self.gateway_callback_redirect_uri().await {
|
||||
uri
|
||||
} else {
|
||||
let port = find_available_port()
|
||||
.await
|
||||
.map_err(|e| ExtensionError::AuthFailed(e.to_string()))?;
|
||||
format!("http://localhost:{}/callback", port.1)
|
||||
};
|
||||
|
||||
// Try DCR if no client_id configured
|
||||
let (client_id, redirect_uri) = if let Some(ref oauth) = server.oauth {
|
||||
let port = find_available_port()
|
||||
.await
|
||||
.map_err(|e| ExtensionError::AuthFailed(e.to_string()))?;
|
||||
let redirect = format!("http://localhost:{}/callback", port.1);
|
||||
(oauth.client_id.clone(), redirect)
|
||||
let (client_id, client_secret) = if let Some(ref oauth) = server.oauth {
|
||||
(oauth.client_id.clone(), None)
|
||||
} else if let Some(ref reg_endpoint) = metadata.registration_endpoint {
|
||||
let port = find_available_port()
|
||||
.await
|
||||
.map_err(|e| ExtensionError::AuthFailed(e.to_string()))?;
|
||||
let redirect = format!("http://localhost:{}/callback", port.1);
|
||||
|
||||
let registration = register_client(reg_endpoint, &redirect)
|
||||
let registration = register_client(reg_endpoint, &redirect_uri)
|
||||
.await
|
||||
.map_err(|e| ExtensionError::AuthFailed(e.to_string()))?;
|
||||
|
||||
(registration.client_id, redirect)
|
||||
(registration.client_id, None)
|
||||
} else {
|
||||
return Err(ExtensionError::AuthFailed(
|
||||
return Err(ExtensionError::AuthNotSupported(
|
||||
"Server doesn't support OAuth or Dynamic Client Registration".to_string(),
|
||||
));
|
||||
};
|
||||
|
||||
let pkce = PkceChallenge::generate();
|
||||
let auth_url = build_authorization_url(
|
||||
// RFC 8707: resource parameter to scope the token to this MCP server
|
||||
let resource = canonical_resource_uri(&server.url);
|
||||
|
||||
// Build authorization URL with CSRF state using the shared oauth_defaults
|
||||
// builder, which generates PKCE + state for us.
|
||||
let mut extra_params = server
|
||||
.oauth
|
||||
.as_ref()
|
||||
.map(|o| o.extra_params.clone())
|
||||
.unwrap_or_default();
|
||||
extra_params.insert("resource".to_string(), resource.clone());
|
||||
|
||||
let scopes = server
|
||||
.oauth
|
||||
.as_ref()
|
||||
.map(|o| o.scopes.clone())
|
||||
.unwrap_or_else(|| metadata.scopes_supported.clone());
|
||||
|
||||
let oauth_result = oauth_defaults::build_oauth_url(
|
||||
&metadata.authorization_endpoint,
|
||||
&client_id,
|
||||
&redirect_uri,
|
||||
&metadata.scopes_supported,
|
||||
Some(&pkce),
|
||||
&std::collections::HashMap::new(),
|
||||
None,
|
||||
&scopes,
|
||||
true, // Always use PKCE for MCP
|
||||
&extra_params,
|
||||
);
|
||||
let expected_state = oauth_result.state;
|
||||
let code_verifier = oauth_result.code_verifier;
|
||||
|
||||
if is_gateway {
|
||||
// Gateway mode: store pending flow for the /oauth/callback handler.
|
||||
oauth_defaults::sweep_expired_flows(&self.pending_oauth_flows).await;
|
||||
|
||||
// Platform routing: prepend instance name to state
|
||||
let platform_state = oauth_defaults::build_platform_state(&expected_state);
|
||||
let auth_url = if platform_state != expected_state {
|
||||
oauth_result.url.replace(
|
||||
&format!("state={}", urlencoding::encode(&expected_state)),
|
||||
&format!("state={}", urlencoding::encode(&platform_state)),
|
||||
)
|
||||
} else {
|
||||
oauth_result.url
|
||||
};
|
||||
|
||||
let flow = oauth_defaults::PendingOAuthFlow {
|
||||
extension_name: name.to_string(),
|
||||
display_name: server.name.clone(),
|
||||
token_url: metadata.token_endpoint,
|
||||
client_id,
|
||||
client_secret,
|
||||
redirect_uri,
|
||||
code_verifier,
|
||||
access_token_field: "access_token".to_string(),
|
||||
secret_name: server.token_secret_name(),
|
||||
provider: Some(format!("mcp:{}", name)),
|
||||
validation_endpoint: None,
|
||||
scopes,
|
||||
user_id: self.user_id.clone(),
|
||||
secrets: Arc::clone(&self.secrets),
|
||||
sse_sender: self.sse_sender.read().await.clone(),
|
||||
gateway_token: self.gateway_token.clone(),
|
||||
resource: Some(resource),
|
||||
client_id_secret_name: if server.oauth.is_none() {
|
||||
Some(server.client_id_secret_name())
|
||||
} else {
|
||||
None
|
||||
},
|
||||
created_at: std::time::Instant::now(),
|
||||
};
|
||||
|
||||
self.pending_oauth_flows
|
||||
.write()
|
||||
.await
|
||||
.insert(expected_state, flow);
|
||||
|
||||
// Store pending auth for later callback handling
|
||||
self.pending_auth.write().await.insert(
|
||||
name.to_string(),
|
||||
PendingAuth {
|
||||
@@ -1787,9 +1961,28 @@ impl ExtensionManager {
|
||||
name,
|
||||
ExtensionKind::McpServer,
|
||||
auth_url,
|
||||
"gateway".to_string(),
|
||||
))
|
||||
} else {
|
||||
// Local mode: return URL for manual opening
|
||||
self.pending_auth.write().await.insert(
|
||||
name.to_string(),
|
||||
PendingAuth {
|
||||
_name: name.to_string(),
|
||||
_kind: ExtensionKind::McpServer,
|
||||
created_at: std::time::Instant::now(),
|
||||
task_handle: None,
|
||||
},
|
||||
);
|
||||
|
||||
Ok(AuthResult::awaiting_authorization(
|
||||
name,
|
||||
ExtensionKind::McpServer,
|
||||
oauth_result.url,
|
||||
"local".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
async fn auth_wasm_tool(&self, name: &str) -> Result<AuthResult, ExtensionError> {
|
||||
// Read the capabilities file to get auth config
|
||||
@@ -2203,7 +2396,10 @@ impl ExtensionManager {
|
||||
flows.retain(|_, flow| flow.extension_name != name);
|
||||
}
|
||||
|
||||
let redirect_uri = format!("{}/callback", oauth_defaults::callback_url());
|
||||
let redirect_uri = self
|
||||
.gateway_callback_redirect_uri()
|
||||
.await
|
||||
.unwrap_or_else(|| format!("{}/callback", oauth_defaults::callback_url()));
|
||||
|
||||
// Merge scopes from all tools sharing this provider
|
||||
let merged_scopes = self
|
||||
@@ -2228,7 +2424,7 @@ impl ExtensionManager {
|
||||
.clone()
|
||||
.unwrap_or_else(|| name.to_string());
|
||||
|
||||
if oauth_defaults::use_gateway_callback() {
|
||||
if self.should_use_gateway_mode() {
|
||||
// Gateway mode: store pending flow state for the web gateway's
|
||||
// `/oauth/callback` handler to complete the exchange. No TCP listener
|
||||
// needed — the OAuth provider redirects to the gateway URL.
|
||||
@@ -2264,6 +2460,8 @@ impl ExtensionManager {
|
||||
secrets: Arc::clone(&self.secrets),
|
||||
sse_sender: self.sse_sender.read().await.clone(),
|
||||
gateway_token: self.gateway_token.clone(),
|
||||
resource: None,
|
||||
client_id_secret_name: None,
|
||||
created_at: std::time::Instant::now(),
|
||||
};
|
||||
|
||||
@@ -2605,11 +2803,17 @@ impl ExtensionManager {
|
||||
.await
|
||||
.map_err(|e| ExtensionError::ActivationFailed(e.to_string()))?;
|
||||
|
||||
// Try to list and create tools
|
||||
let mcp_tools = client
|
||||
.list_tools()
|
||||
.await
|
||||
.map_err(|e| ExtensionError::ActivationFailed(e.to_string()))?;
|
||||
// Try to list and create tools.
|
||||
// A 401/auth error means the server requires OAuth — surface as
|
||||
// AuthRequired so the activate handler triggers the OAuth flow.
|
||||
let mcp_tools = client.list_tools().await.map_err(|e| {
|
||||
let msg = e.to_string();
|
||||
if msg.contains("requires authentication") || msg.contains("401") {
|
||||
ExtensionError::AuthRequired
|
||||
} else {
|
||||
ExtensionError::ActivationFailed(msg)
|
||||
}
|
||||
})?;
|
||||
|
||||
let tool_impls = client
|
||||
.create_tools()
|
||||
@@ -4766,6 +4970,190 @@ mod tests {
|
||||
assert!(result.contains("/v1/users/123/profile"));
|
||||
}
|
||||
|
||||
// ---- gateway mode detection tests ----
|
||||
// Regression tests for a bug where MCP OAuth called `open::that()` on the
|
||||
// server machine instead of returning an auth URL to the gateway frontend.
|
||||
// The root cause was that `should_use_gateway_mode()` only checked the
|
||||
// `IRONCLAW_OAUTH_CALLBACK_URL` env var, ignoring `self.tunnel_url`.
|
||||
|
||||
/// Serializes env-mutating tests to prevent parallel races.
|
||||
static GATEWAY_ENV_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(());
|
||||
|
||||
/// Build a minimal ExtensionManager with a custom tunnel_url.
|
||||
fn make_manager_with_tunnel(tunnel_url: Option<String>) -> ExtensionManager {
|
||||
use crate::secrets::{InMemorySecretsStore, SecretsCrypto};
|
||||
use crate::tools::mcp::process::McpProcessManager;
|
||||
use crate::tools::mcp::session::McpSessionManager;
|
||||
|
||||
let key = secrecy::SecretString::from(crate::secrets::keychain::generate_master_key_hex());
|
||||
let crypto = Arc::new(SecretsCrypto::new(key).expect("crypto"));
|
||||
let secrets: Arc<dyn crate::secrets::SecretsStore + Send + Sync> =
|
||||
Arc::new(InMemorySecretsStore::new(crypto));
|
||||
let tools = Arc::new(crate::tools::ToolRegistry::new());
|
||||
let mcp = Arc::new(McpSessionManager::new());
|
||||
let dir = std::env::temp_dir().join("ironclaw-test-gateway-mode");
|
||||
|
||||
ExtensionManager::new(
|
||||
mcp,
|
||||
Arc::new(McpProcessManager::new()),
|
||||
secrets,
|
||||
tools,
|
||||
None,
|
||||
None,
|
||||
dir.clone(),
|
||||
dir,
|
||||
tunnel_url,
|
||||
"test".to_string(),
|
||||
None,
|
||||
vec![],
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn should_use_gateway_mode_true_for_tunnel_url() {
|
||||
let _guard = GATEWAY_ENV_MUTEX.lock().expect("env mutex poisoned");
|
||||
let original = std::env::var("IRONCLAW_OAUTH_CALLBACK_URL").ok();
|
||||
// SAFETY: Under GATEWAY_ENV_MUTEX, no concurrent env access.
|
||||
unsafe {
|
||||
std::env::remove_var("IRONCLAW_OAUTH_CALLBACK_URL");
|
||||
}
|
||||
|
||||
let mgr = make_manager_with_tunnel(Some("https://my-gateway.example.com".into()));
|
||||
assert!(
|
||||
mgr.should_use_gateway_mode(),
|
||||
"should detect gateway mode from tunnel_url"
|
||||
);
|
||||
|
||||
unsafe {
|
||||
if let Some(val) = original {
|
||||
std::env::set_var("IRONCLAW_OAUTH_CALLBACK_URL", val);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn should_use_gateway_mode_false_without_tunnel() {
|
||||
let _guard = GATEWAY_ENV_MUTEX.lock().expect("env mutex poisoned");
|
||||
let original = std::env::var("IRONCLAW_OAUTH_CALLBACK_URL").ok();
|
||||
unsafe {
|
||||
std::env::remove_var("IRONCLAW_OAUTH_CALLBACK_URL");
|
||||
}
|
||||
|
||||
let mgr = make_manager_with_tunnel(None);
|
||||
assert!(
|
||||
!mgr.should_use_gateway_mode(),
|
||||
"should not detect gateway mode without tunnel_url or env var"
|
||||
);
|
||||
|
||||
unsafe {
|
||||
if let Some(val) = original {
|
||||
std::env::set_var("IRONCLAW_OAUTH_CALLBACK_URL", val);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn should_use_gateway_mode_false_for_loopback_tunnel() {
|
||||
let _guard = GATEWAY_ENV_MUTEX.lock().expect("env mutex poisoned");
|
||||
let original = std::env::var("IRONCLAW_OAUTH_CALLBACK_URL").ok();
|
||||
unsafe {
|
||||
std::env::remove_var("IRONCLAW_OAUTH_CALLBACK_URL");
|
||||
}
|
||||
|
||||
let mgr = make_manager_with_tunnel(Some("http://127.0.0.1:3001".into()));
|
||||
assert!(
|
||||
!mgr.should_use_gateway_mode(),
|
||||
"should not detect gateway mode for loopback tunnel_url"
|
||||
);
|
||||
|
||||
unsafe {
|
||||
if let Some(val) = original {
|
||||
std::env::set_var("IRONCLAW_OAUTH_CALLBACK_URL", val);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Helper to run an async test body while holding the env mutex.
|
||||
/// Clears `IRONCLAW_OAUTH_CALLBACK_URL` for the duration, restoring on drop.
|
||||
struct EnvGuard {
|
||||
original: Option<String>,
|
||||
_mutex: std::sync::MutexGuard<'static, ()>,
|
||||
}
|
||||
|
||||
impl EnvGuard {
|
||||
fn new() -> Self {
|
||||
let guard = GATEWAY_ENV_MUTEX.lock().expect("env mutex poisoned");
|
||||
let original = std::env::var("IRONCLAW_OAUTH_CALLBACK_URL").ok();
|
||||
// SAFETY: Under GATEWAY_ENV_MUTEX, no concurrent env access.
|
||||
unsafe {
|
||||
std::env::remove_var("IRONCLAW_OAUTH_CALLBACK_URL");
|
||||
}
|
||||
Self {
|
||||
original,
|
||||
_mutex: guard,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for EnvGuard {
|
||||
fn drop(&mut self) {
|
||||
// SAFETY: Under GATEWAY_ENV_MUTEX (still held by _mutex), no concurrent env access.
|
||||
unsafe {
|
||||
if let Some(ref val) = self.original {
|
||||
std::env::set_var("IRONCLAW_OAUTH_CALLBACK_URL", val);
|
||||
} else {
|
||||
std::env::remove_var("IRONCLAW_OAUTH_CALLBACK_URL");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_callback_redirect_uri_from_tunnel_url() {
|
||||
let _env = EnvGuard::new();
|
||||
|
||||
let mgr = make_manager_with_tunnel(Some("https://my-gateway.example.com".into()));
|
||||
assert_eq!(
|
||||
mgr.gateway_callback_redirect_uri().await,
|
||||
Some("https://my-gateway.example.com/oauth/callback".to_string()),
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_callback_redirect_uri_none_without_tunnel() {
|
||||
let _env = EnvGuard::new();
|
||||
|
||||
let mgr = make_manager_with_tunnel(None);
|
||||
assert_eq!(mgr.gateway_callback_redirect_uri().await, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_callback_redirect_uri_trims_trailing_slash() {
|
||||
let _env = EnvGuard::new();
|
||||
|
||||
let mgr = make_manager_with_tunnel(Some("https://my-gateway.example.com/".into()));
|
||||
assert_eq!(
|
||||
mgr.gateway_callback_redirect_uri().await,
|
||||
Some("https://my-gateway.example.com/oauth/callback".to_string()),
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_mode_enabled_explicitly() {
|
||||
let _env = EnvGuard::new();
|
||||
|
||||
let mgr = make_manager_with_tunnel(None);
|
||||
assert!(!mgr.should_use_gateway_mode());
|
||||
|
||||
mgr.enable_gateway_mode("https://my-gateway.example.com".into())
|
||||
.await;
|
||||
assert!(mgr.should_use_gateway_mode());
|
||||
assert_eq!(
|
||||
mgr.gateway_callback_redirect_uri().await,
|
||||
Some("https://my-gateway.example.com/oauth/callback".to_string()),
|
||||
);
|
||||
}
|
||||
|
||||
// ── Regression tests for PR #677 (unify-extension-lifecycle) ─────────
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -517,6 +517,9 @@ pub enum ExtensionError {
|
||||
#[error("Authentication failed: {0}")]
|
||||
AuthFailed(String),
|
||||
|
||||
#[error("Server does not support OAuth: {0}")]
|
||||
AuthNotSupported(String),
|
||||
|
||||
#[error("Activation failed: {0}")]
|
||||
ActivationFailed(String),
|
||||
|
||||
|
||||
@@ -472,6 +472,14 @@ async fn async_main() -> anyhow::Result<()> {
|
||||
gw = gw.with_log_level_handle(Arc::clone(&log_level_handle));
|
||||
gw = gw.with_tool_registry(Arc::clone(&components.tools));
|
||||
if let Some(ref ext_mgr) = components.extension_manager {
|
||||
// Enable gateway mode so MCP OAuth returns auth URLs to the frontend
|
||||
// instead of calling open::that() on the server.
|
||||
let gw_base = config
|
||||
.tunnel
|
||||
.public_url
|
||||
.clone()
|
||||
.unwrap_or_else(|| format!("http://{}:{}", gw_config.host, gw_config.port));
|
||||
ext_mgr.enable_gateway_mode(gw_base).await;
|
||||
gw = gw.with_extension_manager(Arc::clone(ext_mgr));
|
||||
}
|
||||
if !components.catalog_entries.is_empty() {
|
||||
|
||||
@@ -256,8 +256,14 @@ impl Tool for ToolAuthTool {
|
||||
}
|
||||
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
// In gateway mode, tool_auth only returns an auth URL for the frontend
|
||||
// to open — no browser is launched server-side, so no approval needed.
|
||||
if self.manager.should_use_gateway_mode() {
|
||||
ApprovalRequirement::Never
|
||||
} else {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── tool_activate ────────────────────────────────────────────────────────
|
||||
@@ -733,6 +739,22 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_auth_no_approval_in_gateway_mode() {
|
||||
let manager = test_manager_stub();
|
||||
manager
|
||||
.enable_gateway_mode("http://localhost:3000".to_string())
|
||||
.await;
|
||||
let tool = ToolAuthTool {
|
||||
manager: manager.clone(),
|
||||
};
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({})),
|
||||
ApprovalRequirement::Never,
|
||||
"tool_auth should not require approval in gateway mode"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_upgrade_schema() {
|
||||
use crate::tools::tool::ApprovalRequirement;
|
||||
|
||||
@@ -403,4 +403,136 @@ mod advanced {
|
||||
rig.verify_trace_expects(&trace, &responses);
|
||||
rig.shutdown();
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 8. MCP extension lifecycle (search → install → activate → use)
|
||||
//
|
||||
// Exercises the MCP extension flow with a mock MCP server:
|
||||
// Turn 1: tool_search → tool_install → text
|
||||
// (inject token + activate between turns)
|
||||
// Turn 2: mock-notion_notion-search → mock-notion_notion-fetch → text
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn mcp_extension_lifecycle() {
|
||||
use crate::support::mock_mcp_server::{MockToolResponse, start_mock_mcp_server};
|
||||
use ironclaw::extensions::{AuthHint, ExtensionKind, ExtensionSource, RegistryEntry};
|
||||
|
||||
// 1. Start mock MCP server with pre-configured tool responses.
|
||||
let mock_server = start_mock_mcp_server(vec![
|
||||
MockToolResponse {
|
||||
name: "notion-search".into(),
|
||||
content: serde_json::json!({
|
||||
"results": [
|
||||
{"id": "page-001", "title": "Project Alpha", "type": "page"},
|
||||
{"id": "page-002", "title": "Sprint Planning", "type": "page"}
|
||||
]
|
||||
}),
|
||||
},
|
||||
MockToolResponse {
|
||||
name: "notion-fetch".into(),
|
||||
content: serde_json::json!({
|
||||
"id": "page-001",
|
||||
"title": "Project Alpha",
|
||||
"content": "Status: In Progress\n- Sprint planning on March 15\n- API redesign review pending"
|
||||
}),
|
||||
},
|
||||
])
|
||||
.await;
|
||||
|
||||
// 2. Load trace fixture.
|
||||
let trace =
|
||||
LlmTrace::from_file(format!("{FIXTURES}/mcp_extension_lifecycle.json")).unwrap();
|
||||
|
||||
// 3. Build rig with auto-approve (so tool_install doesn't block).
|
||||
let rig = TestRigBuilder::new()
|
||||
.with_trace(trace.clone())
|
||||
.with_auto_approve_tools(true)
|
||||
.with_max_tool_iterations(15)
|
||||
.build()
|
||||
.await;
|
||||
|
||||
// 4. Inject mock-notion registry entry pointing to the mock server.
|
||||
let ext_mgr = rig
|
||||
.extension_manager()
|
||||
.expect("test rig must expose extension manager");
|
||||
ext_mgr
|
||||
.inject_registry_entry(RegistryEntry {
|
||||
name: "mock-notion".to_string(),
|
||||
display_name: "Mock Notion".to_string(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
description: "Test MCP server for E2E lifecycle test".to_string(),
|
||||
keywords: vec!["mock-notion".into(), "notion".into()],
|
||||
source: ExtensionSource::McpUrl {
|
||||
url: mock_server.mcp_url(),
|
||||
},
|
||||
fallback_source: None,
|
||||
auth_hint: AuthHint::Dcr,
|
||||
version: None,
|
||||
})
|
||||
.await;
|
||||
|
||||
// 5. Turn 1: "setup mock-notion" → search → install → text.
|
||||
rig.send_message("setup mock-notion").await;
|
||||
let r1 = rig.wait_for_responses(1, TIMEOUT).await;
|
||||
assert!(!r1.is_empty(), "Turn 1: no response");
|
||||
|
||||
// 6. Simulate OAuth completion: inject token + activate.
|
||||
// This mirrors what the gateway's oauth_callback_handler does after
|
||||
// the user completes the OAuth flow in their browser.
|
||||
let secret_name = "mcp_mock-notion_access_token";
|
||||
ext_mgr
|
||||
.secrets()
|
||||
.create(
|
||||
"default",
|
||||
ironclaw::secrets::CreateSecretParams::new(secret_name, "mock-access-token")
|
||||
.with_provider("mcp:mock-notion".to_string()),
|
||||
)
|
||||
.await
|
||||
.expect("failed to inject test token");
|
||||
|
||||
let activate_result = ext_mgr.activate("mock-notion").await;
|
||||
assert!(
|
||||
activate_result.is_ok(),
|
||||
"activation failed: {:?}",
|
||||
activate_result.err()
|
||||
);
|
||||
|
||||
// 7. Turn 2: "check what's in my notion" → notion-search → notion-fetch → text.
|
||||
// Wait for r1.len() + 1 to ensure we observe at least one new turn-2 response.
|
||||
let turn1_count = r1.len();
|
||||
rig.send_message("it's done, check what's in my notion")
|
||||
.await;
|
||||
let r2 = rig.wait_for_responses(turn1_count + 1, TIMEOUT).await;
|
||||
assert!(
|
||||
r2.len() > turn1_count,
|
||||
"Turn 2: expected new responses beyond turn 1's {turn1_count}, got {}",
|
||||
r2.len()
|
||||
);
|
||||
|
||||
// 8. Verify tool calls across both turns.
|
||||
let started = rig.tool_calls_started();
|
||||
assert!(
|
||||
started.iter().any(|s| s == "tool_search"),
|
||||
"tool_search not called: {started:?}"
|
||||
);
|
||||
assert!(
|
||||
started.iter().any(|s| s == "tool_install"),
|
||||
"tool_install not called: {started:?}"
|
||||
);
|
||||
|
||||
// Verify MCP tools were called in turn 2.
|
||||
assert!(
|
||||
started.iter().any(|s| s.starts_with("mock-notion_")),
|
||||
"No mock-notion MCP tools called: {started:?}"
|
||||
);
|
||||
|
||||
// Verify all tools that completed did so successfully.
|
||||
let completed = rig.tool_calls_completed();
|
||||
let failed: Vec<_> = completed.iter().filter(|(_, success)| !success).collect();
|
||||
assert!(failed.is_empty(), "Tools failed: {failed:?}");
|
||||
|
||||
mock_server.shutdown().await;
|
||||
rig.shutdown();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
{
|
||||
"model_name": "advanced-mcp-extension-lifecycle",
|
||||
"expects": {
|
||||
"tools_used": ["tool_search", "tool_install"],
|
||||
"tools_order": ["tool_search", "tool_install"],
|
||||
"all_tools_succeeded": true,
|
||||
"min_responses": 2
|
||||
},
|
||||
"turns": [
|
||||
{
|
||||
"user_input": "setup mock-notion",
|
||||
"steps": [
|
||||
{
|
||||
"request_hint": { "last_user_message_contains": "setup mock-notion" },
|
||||
"response": {
|
||||
"type": "tool_calls",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_search_1",
|
||||
"name": "tool_search",
|
||||
"arguments": { "query": "mock-notion" }
|
||||
}
|
||||
],
|
||||
"input_tokens": 500,
|
||||
"output_tokens": 30
|
||||
}
|
||||
},
|
||||
{
|
||||
"request_hint": { "last_user_message_contains": "setup mock-notion", "min_message_count": 4 },
|
||||
"response": {
|
||||
"type": "tool_calls",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_install_1",
|
||||
"name": "tool_install",
|
||||
"arguments": { "name": "mock-notion" }
|
||||
}
|
||||
],
|
||||
"input_tokens": 600,
|
||||
"output_tokens": 30
|
||||
}
|
||||
},
|
||||
{
|
||||
"request_hint": { "last_user_message_contains": "setup mock-notion", "min_message_count": 6 },
|
||||
"response": {
|
||||
"type": "text",
|
||||
"content": "I've installed Mock Notion. Please authenticate to connect your account — once done, tell me and I'll load the MCP tools.",
|
||||
"input_tokens": 700,
|
||||
"output_tokens": 35
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"user_input": "it's done, check what's in my notion",
|
||||
"steps": [
|
||||
{
|
||||
"request_hint": { "last_user_message_contains": "notion" },
|
||||
"response": {
|
||||
"type": "tool_calls",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_ns_1",
|
||||
"name": "mock-notion_notion-search",
|
||||
"arguments": { "query": "recent notes" }
|
||||
}
|
||||
],
|
||||
"input_tokens": 900,
|
||||
"output_tokens": 30
|
||||
}
|
||||
},
|
||||
{
|
||||
"request_hint": { "min_message_count": 4 },
|
||||
"response": {
|
||||
"type": "tool_calls",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_nf_1",
|
||||
"name": "mock-notion_notion-fetch",
|
||||
"arguments": { "query": "page-001" }
|
||||
}
|
||||
],
|
||||
"input_tokens": 1000,
|
||||
"output_tokens": 30
|
||||
}
|
||||
},
|
||||
{
|
||||
"response": {
|
||||
"type": "text",
|
||||
"content": "Here's what I found in your Notion:\n\n**Project Alpha** — Status: In Progress\n- Sprint planning on March 15\n- API redesign review pending\n\nLet me know if you want more details on any item.",
|
||||
"input_tokens": 1100,
|
||||
"output_tokens": 50
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,340 @@
|
||||
//! Mock MCP server for E2E testing of the extension lifecycle.
|
||||
//!
|
||||
//! Provides a minimal HTTP server with:
|
||||
//! - OAuth 2.1 discovery (`.well-known/oauth-protected-resource`, `.well-known/oauth-authorization-server`)
|
||||
//! - Dynamic Client Registration (`/register`)
|
||||
//! - Token exchange (`/token`)
|
||||
//! - MCP JSON-RPC endpoint (`/mcp`) with `initialize`, `tools/list`, `tools/call`
|
||||
//!
|
||||
//! Tool call responses are pre-configured via `MockToolResponse`.
|
||||
|
||||
#![allow(dead_code)]
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::extract::State;
|
||||
use axum::http::{HeaderMap, StatusCode};
|
||||
use axum::response::IntoResponse;
|
||||
use axum::routing::{get, post};
|
||||
use axum::{Json, Router};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::sync::oneshot;
|
||||
|
||||
/// A pre-configured response for a specific MCP tool call.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct MockToolResponse {
|
||||
/// Tool name (e.g., "notion-search").
|
||||
pub name: String,
|
||||
/// JSON response content for `tools/call`.
|
||||
pub content: serde_json::Value,
|
||||
}
|
||||
|
||||
/// A running mock MCP server.
|
||||
pub struct MockMcpServer {
|
||||
/// Base URL including port (e.g., "http://127.0.0.1:12345").
|
||||
pub base_url: String,
|
||||
/// Shutdown signal sender.
|
||||
shutdown_tx: Option<oneshot::Sender<()>>,
|
||||
/// Server task handle.
|
||||
handle: Option<tokio::task::JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl MockMcpServer {
|
||||
/// The MCP endpoint URL for use in registry entries.
|
||||
pub fn mcp_url(&self) -> String {
|
||||
format!("{}/mcp", self.base_url)
|
||||
}
|
||||
|
||||
/// Shut down the server.
|
||||
pub async fn shutdown(mut self) {
|
||||
if let Some(tx) = self.shutdown_tx.take() {
|
||||
let _ = tx.send(());
|
||||
}
|
||||
if let Some(h) = self.handle.take() {
|
||||
let _ = h.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for MockMcpServer {
|
||||
fn drop(&mut self) {
|
||||
if let Some(tx) = self.shutdown_tx.take() {
|
||||
let _ = tx.send(());
|
||||
}
|
||||
if let Some(h) = self.handle.take() {
|
||||
h.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Shared state for the mock server handlers.
|
||||
struct MockState {
|
||||
/// Base URL (filled after bind).
|
||||
base_url: String,
|
||||
/// Tool definitions served by tools/list.
|
||||
tools: Vec<McpToolDef>,
|
||||
/// Pre-configured tool call responses keyed by tool name.
|
||||
/// Multiple calls to the same tool return responses in order.
|
||||
tool_responses: HashMap<String, Vec<serde_json::Value>>,
|
||||
/// Counter for tool_responses consumption (per tool name).
|
||||
tool_response_idx: std::sync::Mutex<HashMap<String, usize>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Serialize)]
|
||||
struct McpToolDef {
|
||||
name: String,
|
||||
description: String,
|
||||
#[serde(rename = "inputSchema")]
|
||||
input_schema: serde_json::Value,
|
||||
}
|
||||
|
||||
/// Start a mock MCP server on a random port.
|
||||
///
|
||||
/// `tool_responses` configures what `tools/call` returns for each tool name.
|
||||
/// Multiple responses for the same tool are returned in order.
|
||||
pub async fn start_mock_mcp_server(tool_responses: Vec<MockToolResponse>) -> MockMcpServer {
|
||||
// Build tool definitions and response map.
|
||||
let mut tools = Vec::new();
|
||||
let mut response_map: HashMap<String, Vec<serde_json::Value>> = HashMap::new();
|
||||
let mut seen_tools = std::collections::HashSet::new();
|
||||
|
||||
for tr in &tool_responses {
|
||||
if seen_tools.insert(tr.name.clone()) {
|
||||
tools.push(McpToolDef {
|
||||
name: tr.name.clone(),
|
||||
description: format!("Mock tool: {}", tr.name),
|
||||
input_schema: serde_json::json!({"type": "object", "properties": {}}),
|
||||
});
|
||||
}
|
||||
response_map
|
||||
.entry(tr.name.clone())
|
||||
.or_default()
|
||||
.push(tr.content.clone());
|
||||
}
|
||||
|
||||
// Bind to a random port.
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("failed to bind mock MCP server");
|
||||
let addr: SocketAddr = listener.local_addr().expect("no local addr");
|
||||
let base_url = format!("http://127.0.0.1:{}", addr.port());
|
||||
|
||||
let state = Arc::new(MockState {
|
||||
base_url: base_url.clone(),
|
||||
tools,
|
||||
tool_responses: response_map,
|
||||
tool_response_idx: std::sync::Mutex::new(HashMap::new()),
|
||||
});
|
||||
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/.well-known/oauth-protected-resource/mcp",
|
||||
get(handle_protected_resource),
|
||||
)
|
||||
.route(
|
||||
"/.well-known/oauth-authorization-server",
|
||||
get(handle_auth_server_metadata),
|
||||
)
|
||||
.route("/register", post(handle_register))
|
||||
.route("/authorize", get(handle_authorize))
|
||||
.route("/token", post(handle_token))
|
||||
.route("/mcp", post(handle_mcp))
|
||||
.with_state(state);
|
||||
|
||||
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
|
||||
let handle = tokio::spawn(async move {
|
||||
axum::serve(listener, app)
|
||||
.with_graceful_shutdown(async {
|
||||
let _ = shutdown_rx.await;
|
||||
})
|
||||
.await
|
||||
.expect("mock MCP server failed");
|
||||
});
|
||||
|
||||
// Wait briefly for the server to start accepting.
|
||||
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
|
||||
|
||||
MockMcpServer {
|
||||
base_url,
|
||||
shutdown_tx: Some(shutdown_tx),
|
||||
handle: Some(handle),
|
||||
}
|
||||
}
|
||||
|
||||
// ── OAuth discovery endpoints ───────────────────────────────────────────
|
||||
|
||||
async fn handle_protected_resource(State(state): State<Arc<MockState>>) -> impl IntoResponse {
|
||||
Json(serde_json::json!({
|
||||
"resource": format!("{}/mcp", state.base_url),
|
||||
"authorization_servers": [state.base_url],
|
||||
"scopes_supported": ["read", "write"]
|
||||
}))
|
||||
}
|
||||
|
||||
async fn handle_auth_server_metadata(State(state): State<Arc<MockState>>) -> impl IntoResponse {
|
||||
Json(serde_json::json!({
|
||||
"issuer": state.base_url,
|
||||
"authorization_endpoint": format!("{}/authorize", state.base_url),
|
||||
"token_endpoint": format!("{}/token", state.base_url),
|
||||
"registration_endpoint": format!("{}/register", state.base_url),
|
||||
"response_types_supported": ["code"],
|
||||
"grant_types_supported": ["authorization_code"],
|
||||
"code_challenge_methods_supported": ["S256"],
|
||||
"scopes_supported": ["read", "write"]
|
||||
}))
|
||||
}
|
||||
|
||||
// ── OAuth DCR ───────────────────────────────────────────────────────────
|
||||
|
||||
async fn handle_register() -> impl IntoResponse {
|
||||
Json(serde_json::json!({
|
||||
"client_id": "mock-client-id",
|
||||
"client_name": "ironclaw-test",
|
||||
"redirect_uris": [],
|
||||
"grant_types": ["authorization_code"],
|
||||
"response_types": ["code"],
|
||||
"token_endpoint_auth_method": "none"
|
||||
}))
|
||||
}
|
||||
|
||||
// ── OAuth authorize (auto-approve) ──────────────────────────────────────
|
||||
|
||||
/// In a real flow, this would show a consent screen. For testing, we just
|
||||
/// need the endpoint to exist. The test will bypass OAuth by injecting
|
||||
/// tokens directly.
|
||||
async fn handle_authorize() -> impl IntoResponse {
|
||||
// Return a simple HTML page; in practice the test injects tokens directly.
|
||||
axum::response::Html(
|
||||
"<html><body>Mock OAuth: authorize endpoint. Tests bypass this.</body></html>",
|
||||
)
|
||||
}
|
||||
|
||||
// ── OAuth token exchange ────────────────────────────────────────────────
|
||||
|
||||
async fn handle_token() -> impl IntoResponse {
|
||||
Json(serde_json::json!({
|
||||
"access_token": "mock-access-token",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600,
|
||||
"refresh_token": "mock-refresh-token"
|
||||
}))
|
||||
}
|
||||
|
||||
// ── MCP JSON-RPC endpoint ───────────────────────────────────────────────
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct JsonRpcRequest {
|
||||
jsonrpc: String,
|
||||
id: Option<serde_json::Value>,
|
||||
method: String,
|
||||
#[serde(default)]
|
||||
params: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
async fn handle_mcp(
|
||||
State(state): State<Arc<MockState>>,
|
||||
headers: HeaderMap,
|
||||
Json(req): Json<JsonRpcRequest>,
|
||||
) -> impl IntoResponse {
|
||||
// Check for auth header.
|
||||
let auth = headers
|
||||
.get("authorization")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("");
|
||||
|
||||
if !auth.starts_with("Bearer ") || &auth[7..] != "mock-access-token" {
|
||||
// Return 401 with WWW-Authenticate header per MCP OAuth spec.
|
||||
let www_auth = format!(
|
||||
"Bearer resource_metadata=\"{}/.well-known/oauth-protected-resource/mcp\"",
|
||||
state.base_url
|
||||
);
|
||||
return (
|
||||
StatusCode::UNAUTHORIZED,
|
||||
[("www-authenticate", www_auth.as_str())],
|
||||
Json(serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": req.id,
|
||||
"error": {"code": -32000, "message": "Unauthorized"}
|
||||
})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
// Handle notifications (no id) silently.
|
||||
if req.id.is_none() {
|
||||
return StatusCode::OK.into_response();
|
||||
}
|
||||
|
||||
let response = match req.method.as_str() {
|
||||
"initialize" => serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": req.id,
|
||||
"result": {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"serverInfo": {
|
||||
"name": "mock-mcp-server",
|
||||
"version": "1.0.0"
|
||||
},
|
||||
"capabilities": {
|
||||
"tools": {}
|
||||
}
|
||||
}
|
||||
}),
|
||||
"tools/list" => {
|
||||
let tools: Vec<serde_json::Value> = state
|
||||
.tools
|
||||
.iter()
|
||||
.map(|t| serde_json::to_value(t).unwrap())
|
||||
.collect();
|
||||
serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": req.id,
|
||||
"result": {
|
||||
"tools": tools
|
||||
}
|
||||
})
|
||||
}
|
||||
"tools/call" => {
|
||||
let tool_name = req
|
||||
.params
|
||||
.as_ref()
|
||||
.and_then(|p| p.get("name"))
|
||||
.and_then(|n| n.as_str())
|
||||
.unwrap_or("unknown");
|
||||
|
||||
let content = {
|
||||
let mut idx_map = state.tool_response_idx.lock().unwrap();
|
||||
let idx = idx_map.entry(tool_name.to_string()).or_insert(0);
|
||||
let responses = state.tool_responses.get(tool_name);
|
||||
let result = responses
|
||||
.and_then(|r| r.get(*idx))
|
||||
.cloned()
|
||||
.unwrap_or_else(|| serde_json::json!({"error": "no mock response configured"}));
|
||||
*idx += 1;
|
||||
result
|
||||
};
|
||||
|
||||
serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": req.id,
|
||||
"result": {
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": serde_json::to_string(&content).unwrap_or_default()
|
||||
}
|
||||
]
|
||||
}
|
||||
})
|
||||
}
|
||||
_ => serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": req.id,
|
||||
"error": {"code": -32601, "message": format!("Method not found: {}", req.method)}
|
||||
}),
|
||||
};
|
||||
|
||||
Json(response).into_response()
|
||||
}
|
||||
@@ -4,6 +4,7 @@ pub mod cleanup;
|
||||
pub mod gateway_workflow_harness;
|
||||
pub mod instrumented_llm;
|
||||
pub mod metrics;
|
||||
pub mod mock_mcp_server;
|
||||
pub mod mock_openai_server;
|
||||
pub mod test_channel;
|
||||
pub mod test_rig;
|
||||
|
||||
@@ -50,6 +50,9 @@ pub struct TestRig {
|
||||
/// The underlying TraceLlm for inspecting captured requests.
|
||||
#[cfg(feature = "libsql")]
|
||||
trace_llm: Option<Arc<TraceLlm>>,
|
||||
/// Extension manager for direct extension operations in tests.
|
||||
#[cfg(feature = "libsql")]
|
||||
extension_manager: Option<Arc<ironclaw::extensions::ExtensionManager>>,
|
||||
/// Temp directory guard -- keeps the libSQL database file alive.
|
||||
#[cfg(feature = "libsql")]
|
||||
_temp_dir: tempfile::TempDir,
|
||||
@@ -76,6 +79,11 @@ impl TestRig {
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Return the extension manager for direct extension operations in tests.
|
||||
pub fn extension_manager(&self) -> Option<&Arc<ironclaw::extensions::ExtensionManager>> {
|
||||
self.extension_manager.as_ref()
|
||||
}
|
||||
|
||||
/// Wait until at least `n` responses have been captured, or `timeout` elapses.
|
||||
pub async fn wait_for_responses(&self, n: usize, timeout: Duration) -> Vec<OutgoingResponse> {
|
||||
self.channel.wait_for_responses(n, timeout).await
|
||||
@@ -600,6 +608,7 @@ impl TestRigBuilder {
|
||||
// Save references for test accessors.
|
||||
let db_ref = components.db.clone().expect("test rig requires a database");
|
||||
let workspace_ref = components.workspace.clone();
|
||||
let ext_mgr_ref = components.extension_manager.clone();
|
||||
|
||||
// 7. Construct AgentDeps from AppComponents (mirrors main.rs).
|
||||
let deps = AgentDeps {
|
||||
@@ -695,6 +704,7 @@ impl TestRigBuilder {
|
||||
db: db_ref,
|
||||
workspace: workspace_ref,
|
||||
trace_llm: trace_llm_ref,
|
||||
extension_manager: ext_mgr_ref,
|
||||
_temp_dir: temp_dir,
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user