fix(mcp): handle 400 auth errors, clear auth mode after OAuth, trim tokens (#1158)

* fix(mcp): handle 400 auth errors, clear auth mode after OAuth, trim tokens

Three bugs prevented MCP server authentication (e.g. GitHub MCP) from
working correctly:

1. **400 treated as auth-required**: GitHub's MCP endpoint returns 400
   "Authorization header is badly formatted" instead of 401 when auth
   is missing. Broadened auth detection in activate_mcp, send_request,
   and discover_via_401 to also match 400+authorization errors.

2. **Auth mode not cleared after OAuth callback**: The OAuth callback
   handler and setup submit handler did not call clear_auth_mode(),
   leaving pending_auth on the thread. The next user message was
   intercepted as a token instead of triggering an LLM turn.

3. **Token trimming**: Tokens with leading/trailing whitespace or
   newlines produced malformed Authorization headers. Now trimmed
   before storage (configure) and before use (build_request_headers).

Adds E2E tests with a mock MCP server (JSON-RPC + OAuth discovery +
DCR + token exchange) covering install -> activate -> OAuth callback ->
LLM turn lifecycle, plus a GitHub-style 400 error variant.

[skip-regression-check]

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix(mcp): add TTL to PendingAuth and clear auth mode on all failure paths

Auth mode (pending_auth on a Thread) had no timeout and several code
paths that failed to clear it, causing user messages to be swallowed
indefinitely. This adds defense-in-depth:

- Add created_at + 5-minute TTL to PendingAuth; auto-clear on next
  message if expired (safety net for edge cases like user closing
  browser mid-OAuth)
- Clear auth mode on OAuth callback failure paths (unknown/consumed
  state, expired flow)
- Move clear_auth_mode before configure() match in setup_submit so
  it runs on failure too (addresses Copilot review feedback)

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix(ci): exclude test hunks from unwrap/assert pre-commit check

The pre-commit safety script only excluded files in tests/ but not
#[cfg(test)] mod tests blocks inside src/ files. Use the git diff @@
hunk header context (which includes the enclosing function name) to
detect and skip test hunks.

Also removes unnecessary // safety: comments from test assertions.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: restore formatting in test assertions

The replace_all edit that removed // safety: comments collapsed
newlines. Restore proper line breaks.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: address Copilot review - tighten pre-commit filter, document TTL sync

- pre-commit-safety.sh: only exclude `mod tests` hunks (not `fn test_*`)
  to avoid hiding unwrap/assert in production functions like test_server()
- session.rs: extract AUTH_MODE_TTL_SECS constant and add doc comment
  linking to OAUTH_FLOW_EXPIRY to prevent silent drift

[skip-regression-check]

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix(mcp): return error on expired auth input, clear auth on all OAuth paths

- When auth mode TTL expires and the user sends a message (possibly a
  pasted token), return an explicit "expired, please retry" response
  instead of forwarding the content to the LLM/history
- Add clear_auth_mode() to all early-return paths in oauth_callback_handler
  (provider error, missing state/code, no extension manager)

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
This commit is contained in:
Illia Polosukhin
2026-03-15 05:42:49 +00:00
committed by GitHub
co-authored by Claude Opus 4.6
parent 27e21fdabe
commit 62d16e69ac
9 changed files with 760 additions and 29 deletions
+8
View File
@@ -136,6 +136,14 @@ fi
PROD_DIFF="$DIFF_OUTPUT" PROD_DIFF="$DIFF_OUTPUT"
# Strip hunks from test-only files (tests/ directory, *_test.rs, test_*.rs) # Strip hunks from test-only files (tests/ directory, *_test.rs, test_*.rs)
PROD_DIFF=$(echo "$PROD_DIFF" | grep -v '^+++ b/tests/' || true) PROD_DIFF=$(echo "$PROD_DIFF" | grep -v '^+++ b/tests/' || true)
# Strip hunks whose @@ context line indicates a test module.
# git diff includes the enclosing function/module name after @@.
# Only match `mod tests` (the conventional #[cfg(test)] module) — do NOT
# match `fn test_*` because production code can have functions named test_*.
PROD_DIFF=$(echo "$PROD_DIFF" | awk '
/^@@ / { in_test = ($0 ~ /mod tests/) }
!in_test { print }
' || true)
if echo "$PROD_DIFF" | grep -nE '^\+' \ if echo "$PROD_DIFF" | grep -nE '^\+' \
| grep -E '\.(unwrap|expect)\(|[^_]assert(_eq|_ne)?!' \ | grep -E '\.(unwrap|expect)\(|[^_]assert(_eq|_ne)?!' \
| grep -vE 'debug_assert|// safety:|#\[cfg\(test\)\]|#\[test\]|mod tests' \ | grep -vE 'debug_assert|// safety:|#\[cfg\(test\)\]|#\[test\]|mod tests' \
+23
View File
@@ -838,6 +838,28 @@ impl Agent {
}; };
if let Some(pending) = pending_auth { if let Some(pending) = pending_auth {
if pending.is_expired() {
// TTL exceeded — clear stale auth mode
tracing::warn!(
extension = %pending.extension_name,
"Auth mode expired after TTL, clearing"
);
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
thread.pending_auth = None;
}
}
// If this was a user message (possibly a pasted token), return an
// explicit error instead of forwarding it to the LLM/history.
if matches!(submission, Submission::UserInput { .. }) {
return Ok(Some(format!(
"Authentication for **{}** expired. Please try again.",
pending.extension_name
)));
}
// Control submissions (interrupt, undo, etc.) fall through to normal handling
} else {
match &submission { match &submission {
Submission::UserInput { content } => { Submission::UserInput { content } => {
return self return self
@@ -854,6 +876,7 @@ impl Agent {
} }
} }
} }
}
tracing::trace!( tracing::trace!(
"Received message from {} on {} ({} chars)", "Received message from {} on {} ({} chars)",
+45 -9
View File
@@ -12,7 +12,7 @@
use std::collections::{HashMap, HashSet}; use std::collections::{HashMap, HashSet};
use chrono::{DateTime, Utc}; use chrono::{DateTime, TimeDelta, Utc};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use uuid::Uuid; use uuid::Uuid;
@@ -135,6 +135,12 @@ pub enum ThreadState {
/// Pending auth token request. /// Pending auth token request.
/// ///
/// Auth mode TTL — must stay in sync with
/// `crate::cli::oauth_defaults::OAUTH_FLOW_EXPIRY` (5 minutes / 300 s).
/// Defined separately to avoid a session→cli module dependency.
const AUTH_MODE_TTL_SECS: i64 = 300;
const AUTH_MODE_TTL: TimeDelta = TimeDelta::seconds(AUTH_MODE_TTL_SECS);
/// When `tool_auth` returns `awaiting_token`, the thread enters auth mode. /// When `tool_auth` returns `awaiting_token`, the thread enters auth mode.
/// The next user message is intercepted before entering the normal pipeline /// The next user message is intercepted before entering the normal pipeline
/// (no logging, no turn creation, no history) and routed directly to the /// (no logging, no turn creation, no history) and routed directly to the
@@ -143,6 +149,16 @@ pub enum ThreadState {
pub struct PendingAuth { pub struct PendingAuth {
/// Extension name to authenticate. /// Extension name to authenticate.
pub extension_name: String, pub extension_name: String,
/// When this auth mode was entered. Used for TTL expiry.
#[serde(default = "Utc::now")]
pub created_at: DateTime<Utc>,
}
impl PendingAuth {
/// Returns `true` if this auth mode has exceeded the TTL.
pub fn is_expired(&self) -> bool {
Utc::now() - self.created_at > AUTH_MODE_TTL
}
} }
/// Pending tool approval request stored on a thread. /// Pending tool approval request stored on a thread.
@@ -298,7 +314,10 @@ impl Thread {
/// Enter auth mode: next user message will be routed directly to /// Enter auth mode: next user message will be routed directly to
/// the credential store, bypassing the normal pipeline entirely. /// the credential store, bypassing the normal pipeline entirely.
pub fn enter_auth_mode(&mut self, extension_name: String) { pub fn enter_auth_mode(&mut self, extension_name: String) {
self.pending_auth = Some(PendingAuth { extension_name }); self.pending_auth = Some(PendingAuth {
extension_name,
created_at: Utc::now(),
});
self.updated_at = Utc::now(); self.updated_at = Utc::now();
} }
@@ -687,15 +706,16 @@ mod tests {
#[test] #[test]
fn test_enter_auth_mode() { fn test_enter_auth_mode() {
let before = Utc::now();
let mut thread = Thread::new(Uuid::new_v4()); let mut thread = Thread::new(Uuid::new_v4());
assert!(thread.pending_auth.is_none()); assert!(thread.pending_auth.is_none());
thread.enter_auth_mode("telegram".to_string()); thread.enter_auth_mode("telegram".to_string());
assert!(thread.pending_auth.is_some()); assert!(thread.pending_auth.is_some());
assert_eq!( let pending = thread.pending_auth.as_ref().unwrap();
thread.pending_auth.as_ref().unwrap().extension_name, assert_eq!(pending.extension_name, "telegram");
"telegram" assert!(pending.created_at >= before);
); assert!(!pending.is_expired());
} }
#[test] #[test]
@@ -705,8 +725,9 @@ mod tests {
let pending = thread.take_pending_auth(); let pending = thread.take_pending_auth();
assert!(pending.is_some()); assert!(pending.is_some());
assert_eq!(pending.unwrap().extension_name, "notion"); let pending = pending.unwrap();
assert_eq!(pending.extension_name, "notion");
assert!(!pending.is_expired());
// Should be cleared after take // Should be cleared after take
assert!(thread.pending_auth.is_none()); assert!(thread.pending_auth.is_none());
assert!(thread.take_pending_auth().is_none()); assert!(thread.take_pending_auth().is_none());
@@ -720,10 +741,25 @@ mod tests {
let json = serde_json::to_string(&thread).expect("should serialize"); let json = serde_json::to_string(&thread).expect("should serialize");
assert!(json.contains("pending_auth")); assert!(json.contains("pending_auth"));
assert!(json.contains("openai")); assert!(json.contains("openai"));
assert!(json.contains("created_at"));
let restored: Thread = serde_json::from_str(&json).expect("should deserialize"); let restored: Thread = serde_json::from_str(&json).expect("should deserialize");
assert!(restored.pending_auth.is_some()); assert!(restored.pending_auth.is_some());
assert_eq!(restored.pending_auth.unwrap().extension_name, "openai"); let pending = restored.pending_auth.unwrap();
assert_eq!(pending.extension_name, "openai");
assert!(!pending.is_expired());
}
#[test]
fn test_pending_auth_expiry() {
let mut pending = PendingAuth {
extension_name: "test".to_string(),
created_at: Utc::now(),
};
assert!(!pending.is_expired());
// Backdate beyond the TTL
pending.created_at = Utc::now() - AUTH_MODE_TTL - TimeDelta::seconds(1);
assert!(pending.is_expired());
} }
#[test] #[test]
+23 -3
View File
@@ -526,23 +526,33 @@ async fn oauth_callback_handler(
.get("error_description") .get("error_description")
.cloned() .cloned()
.unwrap_or_else(|| error.clone()); .unwrap_or_else(|| error.clone());
clear_auth_mode(&state).await;
return oauth_error_page(&description); return oauth_error_page(&description);
} }
let state_param = match params.get("state") { let state_param = match params.get("state") {
Some(s) if !s.is_empty() => s.clone(), Some(s) if !s.is_empty() => s.clone(),
_ => return oauth_error_page("IronClaw"), _ => {
clear_auth_mode(&state).await;
return oauth_error_page("IronClaw");
}
}; };
let code = match params.get("code") { let code = match params.get("code") {
Some(c) if !c.is_empty() => c.clone(), Some(c) if !c.is_empty() => c.clone(),
_ => return oauth_error_page("IronClaw"), _ => {
clear_auth_mode(&state).await;
return oauth_error_page("IronClaw");
}
}; };
// Look up the pending flow by CSRF state (atomic remove prevents replay) // Look up the pending flow by CSRF state (atomic remove prevents replay)
let ext_mgr = match state.extension_manager.as_ref() { let ext_mgr = match state.extension_manager.as_ref() {
Some(mgr) => mgr, Some(mgr) => mgr,
None => return oauth_error_page("IronClaw"), None => {
clear_auth_mode(&state).await;
return oauth_error_page("IronClaw");
}
}; };
// Strip instance prefix from state for registry lookup. // Strip instance prefix from state for registry lookup.
@@ -563,6 +573,7 @@ async fn oauth_callback_handler(
lookup_key = %lookup_key, lookup_key = %lookup_key,
"OAuth callback received with unknown or expired state" "OAuth callback received with unknown or expired state"
); );
clear_auth_mode(&state).await;
return oauth_error_page("IronClaw"); return oauth_error_page("IronClaw");
} }
}; };
@@ -581,6 +592,7 @@ async fn oauth_callback_handler(
message: "OAuth flow expired. Please try again.".to_string(), message: "OAuth flow expired. Please try again.".to_string(),
}); });
} }
clear_auth_mode(&state).await;
return oauth_error_page(&flow.display_name); return oauth_error_page(&flow.display_name);
} }
@@ -690,6 +702,10 @@ async fn oauth_callback_handler(
} }
} }
// Clear auth mode regardless of outcome so the next user message goes
// through to the LLM instead of being intercepted as a token.
clear_auth_mode(&state).await;
// After successful OAuth, auto-activate the extension so it moves // After successful OAuth, auto-activate the extension so it moves
// from "Installed (Authenticate)" → "Active" without a second click. // from "Installed (Authenticate)" → "Active" without a second click.
// OAuth success is independent of activation — tokens are already stored. // OAuth success is independent of activation — tokens are already stored.
@@ -2182,6 +2198,10 @@ async fn extensions_setup_submit_handler(
"Extension manager not available (secrets store required)".to_string(), "Extension manager not available (secrets store required)".to_string(),
))?; ))?;
// Clear auth mode regardless of outcome so the next user message goes
// through to the LLM instead of being intercepted as a token.
clear_auth_mode(&state).await;
match ext_mgr.configure(&name, &req.secrets).await { match ext_mgr.configure(&name, &req.secrets).await {
Ok(result) => { Ok(result) => {
// Broadcast auth_completed so the chat UI can dismiss any in-progress // Broadcast auth_completed so the chat UI can dismiss any in-progress
+11 -3
View File
@@ -2864,9 +2864,16 @@ impl ExtensionManager {
// Try to list and create tools. // Try to list and create tools.
// A 401/auth error means the server requires OAuth — surface as // A 401/auth error means the server requires OAuth — surface as
// AuthRequired so the activate handler triggers the OAuth flow. // AuthRequired so the activate handler triggers the OAuth flow.
// Some servers (e.g. GitHub MCP) return 400 with "Authorization header
// is badly formatted" instead of 401 when auth is missing or invalid.
let mcp_tools = client.list_tools().await.map_err(|e| { let mcp_tools = client.list_tools().await.map_err(|e| {
let msg = e.to_string(); let msg = e.to_string();
if msg.contains("requires authentication") || msg.contains("401") { let msg_lower = msg.to_ascii_lowercase();
if msg_lower.contains("requires authentication")
|| msg.contains("401")
|| (msg.contains("400")
&& (msg_lower.contains("authorization") || msg_lower.contains("authenticate")))
{
ExtensionError::AuthRequired ExtensionError::AuthRequired
} else { } else {
ExtensionError::ActivationFailed(msg) ExtensionError::ActivationFailed(msg)
@@ -3843,11 +3850,12 @@ impl ExtensionManager {
secret_name, name secret_name, name
))); )));
} }
if secret_value.trim().is_empty() { let trimmed_value = secret_value.trim();
if trimmed_value.is_empty() {
continue; continue;
} }
let params = let params =
CreateSecretParams::new(secret_name, secret_value).with_provider(name.to_string()); CreateSecretParams::new(secret_name, trimmed_value).with_provider(name.to_string());
self.secrets self.secrets
.create(&self.user_id, params) .create(&self.user_id, params)
.await .await
+12 -3
View File
@@ -443,6 +443,11 @@ async fn fetch_resource_metadata(url: &str) -> Result<ProtectedResourceMetadata,
} }
/// Try to discover OAuth metadata via 401 challenge response. /// Try to discover OAuth metadata via 401 challenge response.
///
/// Also accepts 400 responses, since some servers return 400 for
/// unauthenticated requests. In practice the 400 path rarely yields a
/// `WWW-Authenticate` header (GitHub's MCP does not), so discovery
/// typically falls through to strategy 2 (RFC 9728) or 3 (direct).
async fn discover_via_401(server_url: &str) -> Result<AuthorizationServerMetadata, AuthError> { async fn discover_via_401(server_url: &str) -> Result<AuthorizationServerMetadata, AuthError> {
validate_url_safe(server_url).await?; validate_url_safe(server_url).await?;
@@ -459,9 +464,13 @@ async fn discover_via_401(server_url: &str) -> Result<AuthorizationServerMetadat
log_redirect_if_applicable(server_url, &response); log_redirect_if_applicable(server_url, &response);
if response.status().as_u16() != 401 { let status = response.status().as_u16();
// Accept 401 (standard) and 400 (some servers like GitHub MCP use this).
// In both cases, look for WWW-Authenticate header with discovery metadata.
if status != 401 && status != 400 {
return Err(AuthError::DiscoveryFailed(format!( return Err(AuthError::DiscoveryFailed(format!(
"Expected 401, got {}", "Expected 401 or 400, got {}",
response.status() response.status()
))); )));
} }
@@ -471,7 +480,7 @@ async fn discover_via_401(server_url: &str) -> Result<AuthorizationServerMetadat
.get("WWW-Authenticate") .get("WWW-Authenticate")
.and_then(|v| v.to_str().ok()) .and_then(|v| v.to_str().ok())
.ok_or_else(|| { .ok_or_else(|| {
AuthError::DiscoveryFailed("No WWW-Authenticate header in 401 response".to_string()) AuthError::DiscoveryFailed(format!("No WWW-Authenticate header in {} response", status))
})?; })?;
let resource_metadata_url = parse_resource_metadata_url(www_auth).ok_or_else(|| { let resource_metadata_url = parse_resource_metadata_url(www_auth).ok_or_else(|| {
+142 -2
View File
@@ -275,7 +275,10 @@ impl McpClient {
.keys() .keys()
.any(|k| k.eq_ignore_ascii_case("authorization")); .any(|k| k.eq_ignore_ascii_case("authorization"));
if !has_custom_auth && let Some(token) = self.get_access_token().await? { if !has_custom_auth && let Some(token) = self.get_access_token().await? {
headers.insert("Authorization".to_string(), format!("Bearer {}", token)); let trimmed = token.trim();
if !trimmed.is_empty() {
headers.insert("Authorization".to_string(), format!("Bearer {}", trimmed));
}
} }
if let Some(ref session_manager) = self.session_manager if let Some(ref session_manager) = self.session_manager
&& let Some(session_id) = session_manager.get_session_id(&self.server_name).await && let Some(session_id) = session_manager.get_session_id(&self.server_name).await
@@ -302,7 +305,12 @@ impl McpClient {
match result { match result {
Ok(response) => return Ok(response), Ok(response) => return Ok(response),
Err(ToolError::ExternalService(ref msg)) Err(ToolError::ExternalService(ref msg))
if msg.contains("401") || msg.contains("Unauthorized") => if msg.contains("401")
|| msg.contains("Unauthorized")
|| (msg.contains("400") && {
let lower = msg.to_ascii_lowercase();
lower.contains("authorization") || lower.contains("authenticate")
}) =>
{ {
if attempt == 0 if attempt == 0
&& let Some(ref secrets) = self.secrets && let Some(ref secrets) = self.secrets
@@ -1113,4 +1121,136 @@ mod tests {
let approval = wrapper.requires_approval(&serde_json::json!({})); let approval = wrapper.requires_approval(&serde_json::json!({}));
assert_eq!(approval, ApprovalRequirement::Never); assert_eq!(approval, ApprovalRequirement::Never);
} }
// Regression test: empty/whitespace-only tokens must not produce a
// malformed `Authorization: Bearer ` header (GitHub MCP returns 400
// "Authorization header is badly formatted" in this case).
#[tokio::test]
async fn test_build_headers_skips_empty_token() {
use crate::secrets::{CreateSecretParams, DecryptedSecret, Secret, SecretError, SecretRef};
use uuid::Uuid;
// In-memory secrets store that returns a whitespace-only string for the token.
struct EmptyTokenStore;
#[async_trait]
impl crate::secrets::SecretsStore for EmptyTokenStore {
async fn create(
&self,
_user_id: &str,
_params: CreateSecretParams,
) -> Result<Secret, SecretError> {
unimplemented!()
}
async fn get(&self, _user_id: &str, _name: &str) -> Result<Secret, SecretError> {
unimplemented!()
}
async fn get_decrypted(
&self,
_user_id: &str,
_name: &str,
) -> Result<DecryptedSecret, SecretError> {
DecryptedSecret::from_bytes(b" ".to_vec())
}
async fn exists(&self, _user_id: &str, _name: &str) -> Result<bool, SecretError> {
Ok(true)
}
async fn delete(&self, _user_id: &str, _name: &str) -> Result<bool, SecretError> {
Ok(true)
}
async fn list(&self, _user_id: &str) -> Result<Vec<SecretRef>, SecretError> {
Ok(Vec::new())
}
async fn record_usage(&self, _secret_id: Uuid) -> Result<(), SecretError> {
Ok(())
}
async fn is_accessible(
&self,
_user_id: &str,
_secret_name: &str,
_allowed_secrets: &[String],
) -> Result<bool, SecretError> {
Ok(true)
}
}
let config = McpServerConfig::new("github", "https://api.githubcopilot.com/mcp/");
let session_manager = Arc::new(McpSessionManager::new());
let secrets: Arc<dyn crate::secrets::SecretsStore + Send + Sync> =
Arc::new(EmptyTokenStore);
let client = McpClient::new_authenticated(config, session_manager, secrets, "test-user");
let headers = client.build_request_headers().await.unwrap(); // safety: test
assert!(
// safety: test
!headers.contains_key("Authorization"),
"Empty/whitespace token must not produce an Authorization header, got: {:?}",
headers.get("Authorization")
);
}
// Regression test: tokens with leading/trailing whitespace must be trimmed
// before being used in the Authorization header.
#[tokio::test]
async fn test_build_headers_trims_token() {
use crate::secrets::{CreateSecretParams, DecryptedSecret, Secret, SecretError, SecretRef};
use uuid::Uuid;
struct PaddedTokenStore;
#[async_trait]
impl crate::secrets::SecretsStore for PaddedTokenStore {
async fn create(
&self,
_user_id: &str,
_params: CreateSecretParams,
) -> Result<Secret, SecretError> {
unimplemented!()
}
async fn get(&self, _user_id: &str, _name: &str) -> Result<Secret, SecretError> {
unimplemented!()
}
async fn get_decrypted(
&self,
_user_id: &str,
_name: &str,
) -> Result<DecryptedSecret, SecretError> {
DecryptedSecret::from_bytes(b" gho_abc123 \n".to_vec())
}
async fn exists(&self, _user_id: &str, _name: &str) -> Result<bool, SecretError> {
Ok(true)
}
async fn delete(&self, _user_id: &str, _name: &str) -> Result<bool, SecretError> {
Ok(true)
}
async fn list(&self, _user_id: &str) -> Result<Vec<SecretRef>, SecretError> {
Ok(Vec::new())
}
async fn record_usage(&self, _secret_id: Uuid) -> Result<(), SecretError> {
Ok(())
}
async fn is_accessible(
&self,
_user_id: &str,
_secret_name: &str,
_allowed_secrets: &[String],
) -> Result<bool, SecretError> {
Ok(true)
}
}
let config = McpServerConfig::new("github", "https://api.githubcopilot.com/mcp/");
let session_manager = Arc::new(McpSessionManager::new());
let secrets: Arc<dyn crate::secrets::SecretsStore + Send + Sync> =
Arc::new(PaddedTokenStore);
let client = McpClient::new_authenticated(config, session_manager, secrets, "test-user");
let headers = client.build_request_headers().await.unwrap(); // safety: test
assert_eq!(
// safety: test
headers.get("Authorization").unwrap(), // safety: test
"Bearer gho_abc123",
"Token must be trimmed before use in Authorization header"
);
}
} }
+132
View File
@@ -225,6 +225,128 @@ async def models(_request: web.Request) -> web.Response:
}) })
# ── Mock MCP Server ──────────────────────────────────────────────────────────
#
# Simulates an MCP server that requires OAuth. Unauthenticated requests get
# 401 + WWW-Authenticate (standard MCP flow) or 400 "Authorization header is
# badly formatted" (GitHub-style). Authenticated requests return valid
# JSON-RPC responses for initialize and tools/list.
async def mcp_endpoint(request: web.Request) -> web.Response:
"""Handle POST /mcp — JSON-RPC MCP endpoint requiring Bearer auth."""
auth = request.headers.get("Authorization", "")
if not auth.startswith("Bearer ") or len(auth.split(" ", 1)[1].strip()) == 0:
# Return 401 with WWW-Authenticate header for OAuth discovery
resource_meta_url = f"http://127.0.0.1:{request.app['port']}/.well-known/oauth-protected-resource"
return web.Response(
status=401,
headers={"WWW-Authenticate": f'Bearer resource_metadata="{resource_meta_url}"'},
text="Unauthorized",
)
return await _mcp_handle_authed(request)
async def mcp_endpoint_400(request: web.Request) -> web.Response:
"""Handle POST /mcp-400 — MCP endpoint that returns 400 (GitHub-style).
Simulates GitHub's MCP server which returns 400 "Authorization header
is badly formatted" instead of 401 when auth is missing or invalid.
"""
auth = request.headers.get("Authorization", "")
if not auth.startswith("Bearer ") or len(auth.split(" ", 1)[1].strip()) == 0:
return web.Response(
status=400,
text="bad request: Authorization header is badly formatted",
)
return await _mcp_handle_authed(request)
async def _mcp_handle_authed(request: web.Request) -> web.Response:
"""Handle an authenticated MCP JSON-RPC request."""
body = await request.json()
method = body.get("method", "")
req_id = body.get("id")
if method == "initialize":
return web.json_response({
"jsonrpc": "2.0", "id": req_id,
"result": {
"protocolVersion": "2024-11-05",
"capabilities": {"tools": {}},
"serverInfo": {"name": "mock-mcp", "version": "1.0.0"},
},
})
if method == "notifications/initialized":
return web.json_response({"jsonrpc": "2.0", "id": req_id, "result": {}})
if method == "tools/list":
return web.json_response({
"jsonrpc": "2.0", "id": req_id,
"result": {"tools": [{
"name": "mock_search",
"description": "A mock search tool for testing",
"inputSchema": {"type": "object", "properties": {
"query": {"type": "string"},
}},
}]},
})
return web.json_response({"jsonrpc": "2.0", "id": req_id, "error": {
"code": -32601, "message": f"Method not found: {method}",
}})
async def mcp_protected_resource(request: web.Request) -> web.Response:
"""GET /.well-known/oauth-protected-resource[/{path}] — RFC 9728 discovery.
Production code appends the MCP server path after the well-known suffix
(e.g. /.well-known/oauth-protected-resource/mcp-400), so this handler
accepts an optional tail and returns a resource matching the request.
"""
port = request.app["port"]
tail = request.match_info.get("tail", "mcp")
return web.json_response({
"resource": f"http://127.0.0.1:{port}/{tail}",
"authorization_servers": [f"http://127.0.0.1:{port}"],
})
async def mcp_auth_server_metadata(request: web.Request) -> web.Response:
"""GET /.well-known/oauth-authorization-server[/{path}] — OAuth metadata."""
port = request.app["port"]
base = f"http://127.0.0.1:{port}"
return web.json_response({
"issuer": base,
"authorization_endpoint": f"{base}/oauth/authorize",
"token_endpoint": f"{base}/oauth/token",
"registration_endpoint": f"{base}/oauth/register",
"scopes_supported": ["read", "write"],
"response_types_supported": ["code"],
"grant_types_supported": ["authorization_code", "refresh_token"],
"code_challenge_methods_supported": ["S256"],
})
async def mcp_oauth_register(request: web.Request) -> web.Response:
"""POST /oauth/register — Dynamic Client Registration."""
body = await request.json()
return web.json_response({
"client_id": "mock-mcp-client-id",
"client_name": body.get("client_name", "IronClaw"),
"redirect_uris": body.get("redirect_uris", []),
})
async def mcp_oauth_token(request: web.Request) -> web.Response:
"""POST /oauth/token — Token endpoint for MCP OAuth."""
data = await request.post()
code = data.get("code", "")
return web.json_response({
"access_token": f"mcp-token-{code}",
"token_type": "Bearer",
"expires_in": 3600,
})
def main(): def main():
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument("--port", type=int, default=0) parser.add_argument("--port", type=int, default=0)
@@ -236,6 +358,15 @@ def main():
app.router.add_get("/v1/models", models) app.router.add_get("/v1/models", models)
app.router.add_get("/models", models) app.router.add_get("/models", models)
app.router.add_post("/oauth/exchange", oauth_exchange) app.router.add_post("/oauth/exchange", oauth_exchange)
# Mock MCP server endpoints
app.router.add_post("/mcp", mcp_endpoint)
app.router.add_post("/mcp-400", mcp_endpoint_400)
app.router.add_get("/.well-known/oauth-protected-resource", mcp_protected_resource)
app.router.add_get("/.well-known/oauth-protected-resource/{tail:.*}", mcp_protected_resource)
app.router.add_get("/.well-known/oauth-authorization-server", mcp_auth_server_metadata)
app.router.add_get("/.well-known/oauth-authorization-server/{tail:.*}", mcp_auth_server_metadata)
app.router.add_post("/oauth/register", mcp_oauth_register)
app.router.add_post("/oauth/token", mcp_oauth_token)
async def start(): async def start():
runner = web.AppRunner(app) runner = web.AppRunner(app)
@@ -243,6 +374,7 @@ def main():
site = web.TCPSite(runner, "127.0.0.1", args.port) site = web.TCPSite(runner, "127.0.0.1", args.port)
await site.start() await site.start()
port = site._server.sockets[0].getsockname()[1] port = site._server.sockets[0].getsockname()[1]
app["port"] = port # used by MCP handlers
print(f"MOCK_LLM_PORT={port}", flush=True) print(f"MOCK_LLM_PORT={port}", flush=True)
await asyncio.Event().wait() await asyncio.Event().wait()
+355
View File
@@ -0,0 +1,355 @@
"""MCP server auth flow E2E tests.
Tests the full MCP server lifecycle: install MCP server (pointing at mock) ->
activate triggers auth (401/400 -> AuthRequired -> OAuth URL) -> OAuth callback
completes -> auth mode cleared (next message triggers LLM turn) -> MCP tools
available.
Regression coverage for:
- 400 "Authorization header is badly formatted" treated as auth-required
- OAuth discovery via 401 + WWW-Authenticate header
- clear_auth_mode after OAuth callback (user message not swallowed)
- Token trimming (whitespace/newline in stored tokens)
The mock_llm.py serves a mock MCP server at /mcp with full OAuth discovery
endpoints (.well-known/oauth-protected-resource, DCR, token exchange).
"""
from urllib.parse import parse_qs, urlparse
import httpx
import pytest
from helpers import SEL, api_get, api_post
def _extract_state(auth_url: str) -> str:
"""Extract the CSRF state parameter from an OAuth authorization URL."""
parsed = urlparse(auth_url)
qs = parse_qs(parsed.query)
assert "state" in qs, f"auth_url should contain state param: {auth_url}"
return qs["state"][0]
async def _get_extension(base_url, name):
"""Get a specific extension from the extensions list, or None."""
r = await api_get(base_url, "/api/extensions")
for ext in r.json().get("extensions", []):
if ext["name"] == name:
return ext
return None
async def _ensure_removed(base_url, name):
"""Remove extension if already installed."""
ext = await _get_extension(base_url, name)
if ext:
await api_post(base_url, f"/api/extensions/{name}/remove", timeout=30)
# ── Section A: Install MCP Server ────────────────────────────────────────
async def test_mcp_install(ironclaw_server, mock_llm_server):
"""Install a mock MCP server pointing at mock_llm.py's /mcp endpoint."""
await _ensure_removed(ironclaw_server, "mock-mcp")
mcp_url = f"{mock_llm_server}/mcp"
r = await api_post(
ironclaw_server,
"/api/extensions/install",
json={"name": "mock-mcp", "url": mcp_url, "kind": "mcp_server"},
timeout=30,
)
assert r.status_code == 200
data = r.json()
assert data.get("success") is True, f"Install failed: {data}"
ext = await _get_extension(ironclaw_server, "mock-mcp")
assert ext is not None, "mock-mcp should appear in extensions list"
assert ext["kind"] == "mcp_server"
# ── Section B: Activate Triggers Auth ────────────────────────────────────
async def test_mcp_activate_triggers_auth(ironclaw_server):
"""Activating an unauthenticated MCP server triggers the OAuth flow.
The mock MCP returns 401 with WWW-Authenticate when no Bearer token
is present. The activate handler should detect this as auth-required
and return an auth_url.
"""
ext = await _get_extension(ironclaw_server, "mock-mcp")
if ext is None:
pytest.skip("mock-mcp not installed")
r = await api_post(
ironclaw_server,
"/api/extensions/mock-mcp/activate",
timeout=30,
)
assert r.status_code == 200
data = r.json()
# Activation should fail with an auth_url (OAuth needed)
# OR it should return awaiting_token (manual token prompt)
auth_url = data.get("auth_url")
awaiting_token = data.get("awaiting_token")
assert auth_url is not None or awaiting_token, (
f"Activate should require auth, got: {data}"
)
# ── Section C: OAuth Round-Trip ──────────────────────────────────────────
async def test_mcp_oauth_callback(ironclaw_server):
"""Complete the OAuth flow via setup + callback for the MCP server."""
ext = await _get_extension(ironclaw_server, "mock-mcp")
if ext is None:
pytest.skip("mock-mcp not installed")
# Configure with empty secrets to trigger OAuth
r = await api_post(
ironclaw_server,
"/api/extensions/mock-mcp/setup",
json={"secrets": {}},
timeout=30,
)
assert r.status_code == 200
data = r.json()
# If no auth_url, try activate to trigger it
auth_url = data.get("auth_url")
if auth_url is None:
r = await api_post(
ironclaw_server,
"/api/extensions/mock-mcp/activate",
timeout=30,
)
data = r.json()
auth_url = data.get("auth_url")
if auth_url is None:
# Server might have been auto-authenticated via DCR; check if active
ext = await _get_extension(ironclaw_server, "mock-mcp")
if ext and ext.get("authenticated"):
return # Already authenticated, skip callback test
pytest.skip("Could not obtain auth_url for mock-mcp")
csrf_state = _extract_state(auth_url)
# Hit the OAuth callback endpoint
async with httpx.AsyncClient() as client:
r = await client.get(
f"{ironclaw_server}/oauth/callback",
params={"code": "mock_mcp_code", "state": csrf_state},
timeout=30,
follow_redirects=True,
)
assert r.status_code == 200, f"Callback returned {r.status_code}: {r.text[:300]}"
body = r.text.lower()
assert "connected" in body or "success" in body, (
f"Callback should indicate success: {r.text[:500]}"
)
async def test_mcp_authenticated_after_oauth(ironclaw_server):
"""After OAuth callback, MCP server shows authenticated=True."""
ext = await _get_extension(ironclaw_server, "mock-mcp")
if ext is None:
pytest.skip("mock-mcp not installed")
assert ext["authenticated"] is True, (
f"mock-mcp should be authenticated after OAuth: {ext}"
)
async def test_mcp_tools_registered(ironclaw_server):
"""After authentication, MCP tools appear in the extension."""
ext = await _get_extension(ironclaw_server, "mock-mcp")
if ext is None:
pytest.skip("mock-mcp not installed")
tools = ext.get("tools", [])
assert len(tools) > 0, f"mock-mcp should have tools after auth: {ext}"
# The mock MCP serves a tool named "mock_search", prefixed with server name
tool_names = [t for t in tools if "mock_search" in t]
assert len(tool_names) > 0, f"Expected mock_search tool, got: {tools}"
# ── Section D: Auth Mode Cleared — LLM Turn Fires ───────────────────────
async def test_mcp_auth_mode_cleared_llm_turn_fires(ironclaw_server, page):
"""After OAuth completes, the next user message triggers an LLM turn.
Regression test: previously, pending_auth was not cleared by the OAuth
callback handler, so the next user message was consumed as a token and
the LLM turn never fired.
"""
chat_input = page.locator(SEL["chat_input"])
await chat_input.wait_for(state="visible", timeout=5000)
assistant_sel = SEL["message_assistant"]
before_count = await page.locator(assistant_sel).count()
# Send a normal message — should trigger LLM, not be swallowed by auth
await chat_input.fill("hello")
await chat_input.press("Enter")
# Wait for assistant response
expected = before_count + 1
await page.wait_for_function(
"""({ assistantSelector, expectedCount }) => {
const messages = document.querySelectorAll(assistantSelector);
return messages.length >= expectedCount;
}""",
arg={"assistantSelector": assistant_sel, "expectedCount": expected},
timeout=15000,
)
text = await page.locator(assistant_sel).last.inner_text()
assert len(text.strip()) > 0, "Assistant should have responded"
# ── Section E: GitHub-style 400 Error ─────────────────────────────────────
async def test_mcp_400_activate_triggers_auth(ironclaw_server, mock_llm_server):
"""MCP server returning 400 "Authorization header is badly formatted"
is treated as auth-required (regression for GitHub MCP).
Previously, only 401 triggered the auth flow. GitHub's MCP returns 400
with "Authorization header is badly formatted" instead.
"""
await _ensure_removed(ironclaw_server, "mock-mcp-400")
mcp_url = f"{mock_llm_server}/mcp-400"
r = await api_post(
ironclaw_server,
"/api/extensions/install",
json={"name": "mock-mcp-400", "url": mcp_url, "kind": "mcp_server"},
timeout=30,
)
assert r.status_code == 200
assert r.json().get("success") is True, f"Install failed: {r.json()}"
# Activate should detect 400 + "authorization" as auth-required
r = await api_post(
ironclaw_server,
"/api/extensions/mock-mcp-400/activate",
timeout=30,
)
assert r.status_code == 200, f"Activate returned {r.status_code}: {r.text[:300]}"
data = r.json()
# The 400 should be treated as auth-required, returning an auth_url
# or awaiting_token — not a raw "400 Bad Request" activation error.
auth_url = data.get("auth_url")
awaiting_token = data.get("awaiting_token")
assert auth_url is not None or awaiting_token, (
f"400 auth error should trigger auth flow (auth_url or awaiting_token), got: {data}"
)
async def test_mcp_400_oauth_discovery_returns_auth_url(ironclaw_server):
"""OAuth discovery succeeds for the 400-variant via RFC 9728 (strategy 2).
Strategy 1 (discover_via_401) fails because /mcp-400 returns 400 without
a WWW-Authenticate header. Strategy 2 queries
/.well-known/oauth-protected-resource/mcp-400 (path-suffixed) and must
find the mock's wildcard route. Without that route, discovery fails
entirely and only awaiting_token (manual) is returned no auth_url.
This test would have failed before the wildcard .well-known routes were
added to mock_llm.py.
"""
ext = await _get_extension(ironclaw_server, "mock-mcp-400")
if ext is None:
pytest.skip("mock-mcp-400 not installed")
# Re-activate to get a fresh auth response
r = await api_post(
ironclaw_server,
"/api/extensions/mock-mcp-400/activate",
timeout=30,
)
assert r.status_code == 200, f"Activate returned {r.status_code}: {r.text[:300]}"
data = r.json()
auth_url = data.get("auth_url")
assert auth_url is not None, (
f"OAuth discovery must produce an auth_url (not just awaiting_token). "
f"Strategy 2 (RFC 9728) likely failed — check .well-known wildcard routes. "
f"Got: {data}"
)
async def test_mcp_400_full_oauth_roundtrip(ironclaw_server):
"""Complete OAuth round-trip for the 400-variant MCP server.
Exercises the full path: activate 400 detected as auth-required
OAuth discovery via strategy 2 (path-suffixed .well-known) DCR
auth_url returned callback completes token exchange extension
authenticated with tools.
Without the wildcard .well-known routes, OAuth discovery fails and
no auth_url is produced, so this test would fail at the csrf_state
extraction step.
"""
ext = await _get_extension(ironclaw_server, "mock-mcp-400")
if ext is None:
pytest.skip("mock-mcp-400 not installed")
# Get a fresh auth_url via activate
r = await api_post(
ironclaw_server,
"/api/extensions/mock-mcp-400/activate",
timeout=30,
)
data = r.json()
auth_url = data.get("auth_url")
if auth_url is None:
pytest.skip("No auth_url from activate (discovery may not have succeeded)")
csrf_state = _extract_state(auth_url)
# Complete OAuth callback
async with httpx.AsyncClient() as client:
r = await client.get(
f"{ironclaw_server}/oauth/callback",
params={"code": "mock_400_code", "state": csrf_state},
timeout=30,
follow_redirects=True,
)
assert r.status_code == 200, f"Callback returned {r.status_code}: {r.text[:300]}"
body = r.text.lower()
assert "connected" in body or "success" in body, (
f"400-variant OAuth callback should succeed: {r.text[:500]}"
)
# Verify authenticated + tools loaded
ext = await _get_extension(ironclaw_server, "mock-mcp-400")
assert ext is not None, "mock-mcp-400 should still be installed"
assert ext["authenticated"] is True, (
f"mock-mcp-400 should be authenticated after OAuth: {ext}"
)
tools = ext.get("tools", [])
assert len(tools) > 0, f"mock-mcp-400 should have tools after auth: {ext}"
async def test_mcp_400_cleanup(ironclaw_server):
"""Clean up the 400-variant MCP server."""
await _ensure_removed(ironclaw_server, "mock-mcp-400")
ext = await _get_extension(ironclaw_server, "mock-mcp-400")
assert ext is None, "mock-mcp-400 should be removed"
# ── Section F: Cleanup ───────────────────────────────────────────────────
async def test_mcp_cleanup(ironclaw_server):
"""Remove mock-mcp (cleanup for other test files)."""
await _ensure_removed(ironclaw_server, "mock-mcp")
ext = await _get_extension(ironclaw_server, "mock-mcp")
assert ext is None, "mock-mcp should be removed"