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
@@ -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