diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index 8031a60f..11c4acb4 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -4357,7 +4357,13 @@ challenge, whatever its body (every POST, for every server, is answered by `mcp_wire::HttpClient::post_bounded`, which decides from the status — a deliberate divergence from rmcp's own client (see below), which reads a 401 carrying a JSON-RPC error body as an ordinary error response; pinned in `an_ordinary_request_is_answered_as_rmcps_own_client_would`) — makes `ServerAuth::after_rejection` force a refresh through the same `AuthorizationManager` and `McpAuthStore` `mcp-login` uses (so the new token is persisted), and the request is retried -**once**; a second 401 is returned as the server's answer. That covers `tools/call`, `resources/*`, +**once**; a second 401 is returned as the server's answer, saying to run `agent mcp-login ` +again — on a tool call as at connect, since only a new login helps once a refreshed token is refused +too. **The server's reason survives the 401:** a 401 whose small body (≤ 64 KiB) is a JSON-RPC +error fails as rmcp's own bare-401 error does, `HTTP 401 Unauthorized: (JSON-RPC error +)` — still a 401 to the refresh, and for a server with no login (a static key it stopped +accepting, say) the tool error the model sees, with the server's "invalid API key" in it rather than a +bare "Auth required"; any other 401 is `AuthRequired` (`tests/mcp_unauthorized.rs`). That covers `tools/call`, `resources/*`, `prompts/*`, `skills/*`, MCP App view reads, the handshake, the standalone stream and `events/*`. A 403 (`InsufficientScope`) never refreshes. @@ -4366,7 +4372,7 @@ in and out, a 404 on a session as `SessionExpired` (rmcp re-initializes and retr handed back as a stream, 202/204 — or an empty 200 to a notification or reply — as accepted, a JSON 200 that is not a JSON-RPC message, for a notification or reply, as accepted, a non-JSON success as an error (rmcp's unexpected content type), and a 4xx to `server/discover` as the legacy server's cue to -`initialize`. **The differences:** any 401 is `AuthRequired` (above; rmcp reads one with a JSON-RPC body as an error response — pinned in `an_ordinary_request_is_answered_as_rmcps_own_client_would`); and a JSON 200 that is not a JSON-RPC message, answering a _request_, +`initialize`. **The differences:** any 401 is a 401 (above; rmcp reads one with a JSON-RPC body as an error response — pinned in `an_ordinary_request_is_answered_as_rmcps_own_client_would`); and a JSON 200 that is not a JSON-RPC message, answering a _request_, is an error here. rmcp calls it accepted and then waits for a response that cannot come (a request is answered on its own POST, as JSON or an SSE stream carrying it) until the request's timeout; failing at once with the body in the error is the honest answer. Pinned end to end against the OAuth fixture @@ -4666,7 +4672,7 @@ mcp_events_subscribe (any session) ──► owned by that session passes through the same `rescue`). Stateless (`2026-07-28`) streamable-HTTP servers get `events/*` directly over HTTP (`MCP-Protocol-Version`, `Mcp-Method`, per-request `_meta`, the server's resolved headers and OAuth bearer), with bodies - bounded by the same per-message cap as every MCP transport (`mcp_stdio::max_message_bytes`, `BEYOND_AI_AGENT_MCP_MAX_MESSAGE_BYTES`: a unary JSON body, a unary SSE answer's event, an `events/stream` event — an over-cap answer fails its request, never read whole). An over-cap **event notification** (on a push stream, or from a stdio server) is skipped, not reconnected into: its bounded head is read structurally (`mcp_stdio::oversized_stand_in`, the same top-level member walk as `scan_head`) for its routing, cursor and id, the rest is skipped unread, and a small `$oversized` stand-in takes its place — the subscription keeps that cursor (or the next heartbeat's), records an `oversized` gap (a `gap` frame, and a notice to the model that an event was dropped unread) and carries on (`tests/mcp_message_cap.rs`) — JWKS documents at 64 KiB — and an SSE reader whose scan and drain + bounded by the same per-message cap as every MCP transport (`mcp_stdio::max_message_bytes`, `BEYOND_AI_AGENT_MCP_MAX_MESSAGE_BYTES`: a unary JSON body, a unary SSE answer's event, an `events/stream` event — an over-cap answer fails its request, never read whole). An over-cap **event notification** (on a push stream, or from a stdio server) is skipped, not reconnected into: its bounded head is read structurally (`mcp_stdio::oversized_stand_in`, the same top-level member walk as `scan_head`) for its routing, cursor and id — **in any key order**: `method`, `id` and `params` are found wherever they fall within the window, and inside `params` every scalar and `_meta` that appears whole is kept — the rest is skipped unread, and a small `$oversized` stand-in takes its place. When the head does not reach the method (a payload-first `params` fills it), the message is still taken for an event unless the head proves it is a request (an `id`) or a response; over stdio, a stand-in whose routing lay past the head is delivered to every push stream on that connection (each records the gap, since which one it belonged to is unknowable), and on a direct-HTTP stream it is that stream's. In either case the subscription keeps whatever cursor the head showed (or else the next event's or heartbeat's — never reconnecting into the same event), records an `oversized` gap (a `gap` frame, and a notice to the model that an event was dropped unread) and carries on (`tests/mcp_message_cap.rs`, both key orders over both transports) — JWKS documents at 64 KiB — and an SSE reader whose scan and drain are both linear (events are parsed in place; the consumed prefix is dropped once per read, not once per event); an older, session-bound HTTP server goes through rmcp's own connection, which carries its `Mcp-Session-Id`. diff --git a/crates/agent/src/bin/mcp_fixture_events_server.rs b/crates/agent/src/bin/mcp_fixture_events_server.rs index e243c1b0..e0849e72 100644 --- a/crates/agent/src/bin/mcp_fixture_events_server.rs +++ b/crates/agent/src/bin/mcp_fixture_events_server.rs @@ -1496,7 +1496,7 @@ async fn handle_http(state: Shared, mut stream: TcpStream) { break false; }; let terminal = frame["method"] == "notifications/events/terminated"; - let line = format!("data: {frame}\n\n"); + let line = format!("data: {}\n\n", frame_text(&frame)); if stream.write_all(line.as_bytes()).await.is_err() || stream.flush().await.is_err() { @@ -1550,8 +1550,35 @@ async fn handle_http(state: Shared, mut stream: TcpStream) { .await; } +/// A frame's wire text. `MCP_FIXTURE_KEY_ORDER=params_first` writes an event notification the +/// other way round from the usual key order — `params` before `method`, and the payload (`data`) +/// first inside `params` — as JSON allows: a client reading a bounded head of an over-cap one then +/// sees neither its method, its routing nor its cursor. +fn frame_text(v: &Value) -> String { + let params_first = std::env::var("MCP_FIXTURE_KEY_ORDER").is_ok_and(|o| o == "params_first"); + let (Some(method), Some(Value::Object(params))) = (v["method"].as_str(), v.get("params")) + else { + return v.to_string(); + }; + if !params_first || method != "notifications/events/event" { + return v.to_string(); + } + let mut members: Vec = Vec::new(); + if let Some(data) = params.get("data") { + members.push(format!("\"data\":{data}")); + } + for (k, val) in params.iter().filter(|(k, _)| *k != "data") { + members.push(format!("{}:{val}", Value::String(k.clone()))); + } + format!( + "{{\"params\":{{{}}},\"jsonrpc\":\"2.0\",\"method\":{}}}", + members.join(","), + Value::String(method.to_owned()) + ) +} + async fn write_line(out: &Stdout, v: &Value) -> bool { - let mut line = serde_json::to_vec(v).unwrap(); + let mut line = frame_text(v).into_bytes(); line.push(b'\n'); let mut out = out.lock().await; out.write_all(&line).await.is_ok() && out.flush().await.is_ok() diff --git a/crates/agent/src/tools/mcp_events/wire.rs b/crates/agent/src/tools/mcp_events/wire.rs index c2fd5888..3e7844d8 100644 --- a/crates/agent/src/tools/mcp_events/wire.rs +++ b/crates/agent/src/tools/mcp_events/wire.rs @@ -88,7 +88,16 @@ impl NotificationRouter { self.generation.fetch_add(1, Ordering::Relaxed); return; } + let params = notification.params.unwrap_or(Value::Null); let Some(id) = subscription_id else { + // The stand-in for an over-cap event whose routing lay past the bounded head + // (`mcp_stdio::oversized_stand_in`): it belongs to one of this connection's streams, + // which is unknowable — so each is told an event may have been dropped, rather than + // none. Anything else without a subscription id is dropped, as before. + if params.get("$oversized") == Some(&Value::Bool(true)) { + self.to_every_stream(method, ¶ms); + return; + } tracing::debug!( method, "events notification without a subscriptionId; dropped" @@ -103,7 +112,7 @@ impl NotificationRouter { }; let msg = StreamMsg { method: method.to_owned(), - params: notification.params.unwrap_or(Value::Null), + params, }; if slot.tx.try_send(msg).is_err() { slot.overflowed.store(true, Ordering::Release); @@ -111,6 +120,24 @@ impl NotificationRouter { } } + /// Deliver one message to every open stream, with the same overflow rule as [`Self::route`]. + fn to_every_stream(&self, method: &str, params: &Value) { + let Ok(mut inner) = self.inner.lock() else { + return; + }; + inner.streams.retain(|_, slot| { + let msg = StreamMsg { + method: method.to_owned(), + params: params.clone(), + }; + let sent = slot.tx.try_send(msg).is_ok(); + if !sent { + slot.overflowed.store(true, Ordering::Release); + } + sent + }); + } + fn register(&self, id: RequestId, tx: mpsc::Sender, overflowed: Arc) { if let Ok(mut inner) = self.inner.lock() { inner.streams.insert(id, Slot { tx, overflowed }); diff --git a/crates/agent/src/tools/mcp_oauth.rs b/crates/agent/src/tools/mcp_oauth.rs index 03aab109..b88cb138 100644 --- a/crates/agent/src/tools/mcp_oauth.rs +++ b/crates/agent/src/tools/mcp_oauth.rs @@ -508,8 +508,27 @@ impl> OAuthHttp { match call(sent.clone()).await { Err(e) if is_unauthorized(&e) => { match auth.after_rejection(sent.as_deref()).await { - // Once: a second 401 is the server's answer, returned as is. - Ok(fresh) => call(Some(fresh)).await, + // Once: a second 401 is the server's answer — with what fixes it, since a + // refreshed token was refused too and only a new login can help. + Ok(fresh) => match call(Some(fresh)).await { + Err(e) if is_unauthorized(&e) => { + Err(StreamableHttpError::UnexpectedServerResponse( + format!( + "{}; the server refused the refreshed login for `{1}` too: run \ + `agent mcp-login {1}` again", + match &e { + StreamableHttpError::UnexpectedServerResponse(m) => { + m.to_string() + } + _ => "HTTP 401 Unauthorized".to_owned(), + }, + auth.server, + ) + .into(), + )) + } + other => other, + }, Err(message) => Err(StreamableHttpError::UnexpectedServerResponse( message.into(), )), diff --git a/crates/agent/src/tools/mcp_stdio.rs b/crates/agent/src/tools/mcp_stdio.rs index 167a75f5..d8c445a6 100644 --- a/crates/agent/src/tools/mcp_stdio.rs +++ b/crates/agent/src/tools/mcp_stdio.rs @@ -360,12 +360,21 @@ pub(crate) fn scan_head(head: &[u8]) -> Head { } /// A small stand-in for an over-cap **MCP Events notification** (`notifications/events/*`), built -/// from the bounded head of the message — the oversized one itself is never held. The stand-in is -/// the same notification with `params` cut down to the members that appear, whole, before the -/// head runs out (the event's `_meta` routing, `cursor`, `eventId`, `name`, … — whatever precedes the -/// payload) plus `"$oversized": true`, so the events client can skip that event, keep its cursor and -/// tell the model it was dropped, rather than reconnecting into the same giant event forever. -/// `None` for anything else (it is dropped as before). +/// from the bounded head of the message — the oversized one itself is never held. The head is read +/// **in any key order** (JSON fixes none): the top-level members are walked as far as the window +/// reaches, wherever `method`, `id`, `result`/`error` and `params` fall, and inside `params` every +/// member that appears whole is kept if it is a scalar (`cursor`, `eventId`, `name`, …) or the small +/// `_meta` (routing) — a nested payload is skipped. The stand-in is the notification with those +/// `params` plus `"$oversized": true`, so the events client skips that event, keeps what position it +/// can and tells the model it was dropped, rather than reconnecting into the same giant event. +/// +/// When the head does not reach the `method` (a payload-first `params` fills the window), the +/// message is still taken for an event — `notifications/events/event` — unless the head proves +/// otherwise: a response (`result`/`error`) or a request (an `id`; notifications have none). An +/// over-cap message that is something else is far likelier to be an event than not on a connection +/// carrying events; the price of guessing wrong is one spurious gap notice, where guessing the +/// other way loses an event silently (stdio) or loops on it (HTTP). `None` for anything else (it +/// is dropped as before). pub(crate) fn oversized_stand_in(head: &[u8]) -> Option { let ws = |mut i: usize| { while head.get(i).is_some_and(u8::is_ascii_whitespace) { @@ -379,68 +388,106 @@ pub(crate) fn oversized_stand_in(head: &[u8]) -> Option { } i += 1; let mut method: Option = None; + let mut params = serde_json::Map::new(); loop { i = ws(i); if head.get(i) != Some(&b'"') { - return None; + break; } - let key_end = string_end(head, i)?; + let Some(key_end) = string_end(head, i) else { + break; + }; let key = &head[i + 1..key_end - 1]; i = ws(key_end); if head.get(i) != Some(&b':') { - return None; + break; } i = ws(i + 1); - if key == b"params" { - let method = method.filter(|m| m.starts_with("notifications/events/"))?; - let mut params = serde_json::Map::new(); - if head.get(i) == Some(&b'{') { - let mut j = i + 1; - loop { - j = ws(j); - if head.get(j) != Some(&b'"') { - break; - } - let Some(k_end) = string_end(head, j) else { - break; - }; - let Ok(k) = serde_json::from_slice::(&head[j..k_end]) else { - break; - }; - j = ws(k_end); - if head.get(j) != Some(&b':') { - break; - } - j = ws(j + 1); - let Some(v_end) = value_end(head, j) else { - break; - }; - let value = &head[j..v_end]; - // Scalars, and the small `_meta` (routing); a whole nested payload is skipped. - let keep = !matches!(value.first(), Some(b'{' | b'[')) || k == "_meta"; - if keep && let Ok(v) = serde_json::from_slice::(value) { - params.insert(k, v); - } - j = ws(v_end); - if head.get(j) != Some(&b',') { - break; - } - j += 1; + match key { + // A response, or a request: not an event, whatever else the head shows. + b"result" | b"error" | b"id" => return None, + b"params" => { + let (members, end) = params_members(head, i); + params = members; + match end { + Some(end) => i = end, + None => break, } } - params.insert("$oversized".into(), Value::Bool(true)); - return Some(json!({ "jsonrpc": "2.0", "method": method, "params": params })); - } - let end = value_end(head, i)?; - if key == b"method" { - method = serde_json::from_slice::(&head[i..end]).ok(); + _ => { + let Some(end) = value_end(head, i) else { + break; + }; + if key == b"method" { + method = Some(serde_json::from_slice::(&head[i..end]).ok()?); + } + i = end; + } } - i = ws(end); + i = ws(i); if head.get(i) != Some(&b',') { - return None; + break; } i += 1; } + let method = match method { + Some(m) if m.starts_with("notifications/events/") => m, + Some(_) => return None, + None => "notifications/events/event".to_owned(), + }; + params.insert("$oversized".into(), Value::Bool(true)); + Some(json!({ "jsonrpc": "2.0", "method": method, "params": params })) +} + +/// The members of the `params` object starting at `head[at]` that appear whole within `head` and +/// are worth keeping in a stand-in — scalars, and `_meta` — in any order; and where the object +/// ends, if it ends within `head`. +fn params_members(head: &[u8], at: usize) -> (serde_json::Map, Option) { + let ws = |mut i: usize| { + while head.get(i).is_some_and(u8::is_ascii_whitespace) { + i += 1; + } + i + }; + let mut kept = serde_json::Map::new(); + if head.get(at) != Some(&b'{') { + return (kept, value_end(head, at)); + } + let mut j = at + 1; + loop { + j = ws(j); + match head.get(j) { + Some(b'}') => return (kept, Some(j + 1)), + Some(b'"') => {} + _ => return (kept, None), + } + let Some(k_end) = string_end(head, j) else { + return (kept, None); + }; + let Ok(k) = serde_json::from_slice::(&head[j..k_end]) else { + return (kept, None); + }; + j = ws(k_end); + if head.get(j) != Some(&b':') { + return (kept, None); + } + j = ws(j + 1); + let Some(v_end) = value_end(head, j) else { + return (kept, None); + }; + let value = &head[j..v_end]; + // Scalars, and the small `_meta` (routing); a whole nested payload is skipped. + let keep = !matches!(value.first(), Some(b'{' | b'[')) || k == "_meta"; + if keep && let Ok(v) = serde_json::from_slice::(value) { + kept.insert(k, v); + } + j = ws(v_end); + match head.get(j) { + Some(b',') => j += 1, + Some(b'}') => return (kept, Some(j + 1)), + _ => return (kept, None), + } + } } /// The index just past the JSON string opening at `buf[at]`, if it closes within `buf`. @@ -1184,7 +1231,6 @@ mod tests { for other in [ &br#"{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1"#[..], br#"{"jsonrpc":"2.0","id":3,"result":{"x":"yyyy"#, - br#"{"jsonrpc":"2.0","params":{"cursor":"1"},"method":"notifications/events/event"}"#, b"not json", ] { assert_eq!( @@ -1195,4 +1241,48 @@ mod tests { ); } } + + /// JSON fixes no key order: the head is read wherever `method`, `id` and `params` fall. + #[test] + fn an_oversized_events_notification_is_recognized_in_any_key_order() { + // `params` (whole) before `method`. + let stand_in = oversized_stand_in( + br#"{"params":{"cursor":"7","eventId":"e1"},"jsonrpc":"2.0","method":"notifications/events/event","x":"yyyy"#, + ) + .unwrap(); + assert_eq!(stand_in["method"], "notifications/events/event"); + assert_eq!(stand_in["params"]["cursor"], "7"); + assert_eq!(stand_in["params"]["eventId"], "e1"); + // `params` first and the payload in it: the method is past the window, the routing and + // cursor ahead of the payload are kept, and it is taken for an event. + let stand_in = oversized_stand_in( + br#"{"params":{"_meta":{"io.modelcontextprotocol/subscriptionId":"s1"},"cursor":"9","data":{"blob":"xxxx"#, + ) + .unwrap(); + assert_eq!(stand_in["method"], "notifications/events/event"); + assert_eq!(stand_in["params"]["cursor"], "9"); + assert_eq!( + stand_in["params"]["_meta"]["io.modelcontextprotocol/subscriptionId"], + "s1" + ); + // Payload first: nothing identifying is visible, and it is still skipped as an event. + let stand_in = oversized_stand_in(br#"{"params":{"data":{"blob":"xxxx"#).unwrap(); + assert_eq!(stand_in["params"]["$oversized"], true); + assert!(stand_in["params"].get("cursor").is_none()); + // What the head proves is not an event is not stood in for, in any order: a request (an + // id, which notifications never carry), a response, another method. + for other in [ + &br#"{"params":{"data":"x"},"id":4,"method":"#[..], + br#"{"id":4,"params":{"data":{"blob":"xxxx"#, + br#"{"params":{"progress":1},"method":"notifications/progress","jsonrpc":"#, + br#"{"error":{"code":1,"message":"xxxx"#, + ] { + assert_eq!( + oversized_stand_in(other), + None, + "{}", + String::from_utf8_lossy(other) + ); + } + } } diff --git a/crates/agent/src/tools/mcp_wire.rs b/crates/agent/src/tools/mcp_wire.rs index dccfcac4..44a159ef 100644 --- a/crates/agent/src/tools/mcp_wire.rs +++ b/crates/agent/src/tools/mcp_wire.rs @@ -63,6 +63,28 @@ fn ordinary_limit(message: &ClientJsonRpcMessage, max: usize) -> Limit { } } +/// The most of a 401's body read for its reason: a JSON-RPC error is a few hundred bytes. +const MAX_UNAUTHORIZED_BODY: usize = 64 * 1024; + +/// The server's reason for a 401, when its body is a JSON-RPC error: `error.message` (with the +/// code). `None` for an empty, larger or non-JSON-RPC body. +async fn unauthorized_reason(response: reqwest::Response) -> Option { + if response + .content_length() + .is_some_and(|len| len > MAX_UNAUTHORIZED_BODY as u64) + { + return None; + } + let body = capped_body(response, MAX_UNAUTHORIZED_BODY).await.ok()??; + let value: Value = serde_json::from_str(&body).ok()?; + let error = value.get("error")?; + let message = error.get("message")?.as_str()?; + Some(match error.get("code").and_then(Value::as_i64) { + Some(code) => format!("{message} (JSON-RPC error {code})"), + None => message.to_owned(), + }) +} + /// Headers a configured custom header may not override (`rmcp`'s own reserved set). const RESERVED_HEADERS: [&str; 3] = ["accept", "mcp-session-id", "last-event-id"]; @@ -232,8 +254,16 @@ impl HttpClient { .and_then(|v| v.to_str().ok()) .map(str::to_string); // Any 401 — with or without a challenge, whatever its body — is the server refusing the - // credentials, decided here from the status before a body could be read as a response. + // credentials, decided here from the status before a body could be read as a response + // (so `mcp_oauth` can refresh a login). The server's own reason survives it: a 401 whose + // (small) body is a JSON-RPC error — "invalid API key" — fails as rmcp's bare-401 error + // does, `HTTP 401 …`, carrying that message; any other 401 is `AuthRequired`. if status == reqwest::StatusCode::UNAUTHORIZED { + if let Some(message) = unauthorized_reason(response).await { + return Err(StreamableHttpError::UnexpectedServerResponse(Cow::Owned( + format!("HTTP 401 Unauthorized: {message}"), + ))); + } return Err(StreamableHttpError::AuthRequired( rmcp::transport::streamable_http_client::AuthRequiredError::new( www_authenticate.unwrap_or_default(), @@ -753,8 +783,9 @@ mod tests { } } // The deliberate divergence: a 401 with no `WWW-Authenticate` whose body is a JSON-RPC - // error. rmcp hands that back as an ordinary error *response*, losing the status; here any - // 401 is `AuthRequired`, so `mcp_oauth` can refresh (and any server's caller sees why). + // error. rmcp hands that back as an ordinary error *response*, losing the status; here it + // is the 401 it is (rmcp's own bare-401 shape, `HTTP 401 …`), so `mcp_oauth` can refresh — + // carrying the server's message, so its caller still sees why. let raw: &[u8] = b"HTTP/1.1 401 Unauthorized\r\nContent-Type: application/json\r\nContent-Length: 69\r\nConnection: close\r\n\r\n{\"jsonrpc\":\"2.0\",\"id\":7,\"error\":{\"code\":-32001,\"message\":\"no token\"}}"; let ours = HttpClient::new(reqwest::Client::new()) .post_message( @@ -774,10 +805,12 @@ mod tests { HashMap::new(), ) .await; - assert_eq!(shape(&ours), "auth-required"); - assert_ne!( - shape(&theirs), - "auth-required", + assert_eq!( + shape(&ours), + "other: unexpected server response: HTTP 401 Unauthorized: no token (JSON-RPC error -32001)" + ); + assert!( + !shape(&theirs).contains("HTTP 401"), "rmcp's own client: {}", shape(&theirs) ); diff --git a/crates/agent/tests/common/mcp_oauth_fixture.rs b/crates/agent/tests/common/mcp_oauth_fixture.rs index 10e1277d..0c88c8fe 100644 --- a/crates/agent/tests/common/mcp_oauth_fixture.rs +++ b/crates/agent/tests/common/mcp_oauth_fixture.rs @@ -169,6 +169,9 @@ pub struct OAuthFixture { /// Answer a rejected request's 401 with an `application/json` JSON-RPC error body and no /// `WWW-Authenticate` challenge (what rmcp's own client reads as an ordinary error response). pub reject_with_json_body: Arc, + /// Answer every `tools/call` 401, whatever token it carries — a login the server keeps + /// refusing (a refresh gets a new token, and that is refused too). + pub reject_calls: Arc, /// Speak sessions: `initialize` issues an `Mcp-Session-Id`, every later POST must carry a live /// one (`400` without, `404` for an expired one), `GET` opens the standalone SSE stream, and /// `DELETE` ends the session. @@ -198,6 +201,12 @@ impl OAuthFixture { self.issued.lock().unwrap().clear(); } + /// Accept `token` as if issued: for a server configured with a static `Authorization` header + /// rather than a login. + pub fn issue(&self, token: &str) { + self.issued.lock().unwrap().insert(token.to_owned()); + } + /// Start one. `expires_in_secs` is the lifetime the fixture reports for each token it issues. /// `rmcp`'s `AuthorizationManager::get_access_token` refreshes proactively whenever fewer than 30 /// seconds remain, so anything below 30 is refreshed on the very next read without any waiting; @@ -215,6 +224,7 @@ impl OAuthFixture { refresh_reply: Arc::default(), hold_rejections_until: Arc::default(), reject_with_json_body: Arc::default(), + reject_calls: Arc::default(), sessions: Arc::default(), expire_sessions_after_calls: Arc::default(), sse_responses: Arc::default(), @@ -448,7 +458,8 @@ impl Shared { let request: Value = serde_json::from_slice(&req.body).unwrap_or(Value::Null); let method = request.get("method").and_then(Value::as_str).unwrap_or(""); let id = request.get("id").cloned().unwrap_or(Value::Null); - if !self.authorized(req) { + let refused = method == "tools/call" && f.reject_calls.load(Ordering::SeqCst); + if refused || !self.authorized(req) { if method == "tools/call" { f.rejected_calls.fetch_add(1, Ordering::SeqCst); let want = f.hold_rejections_until.load(Ordering::SeqCst); diff --git a/crates/agent/tests/mcp_message_cap.rs b/crates/agent/tests/mcp_message_cap.rs index fbabe284..e1cf0545 100644 --- a/crates/agent/tests/mcp_message_cap.rs +++ b/crates/agent/tests/mcp_message_cap.rs @@ -167,12 +167,21 @@ fn an_over_cap_unary_sse_answer_on_the_events_wire_is_refused() { /// An `events/stream` event over the cap is skipped, not read and not reconnected into: its bounded /// head names it (and its cursor), it is reported as a gap — to the client and to the model — and the -/// next event on the same stream is delivered. Over streamable HTTP (SSE) and over stdio. -fn an_over_cap_push_event_is_skipped_and_reported(stdio: bool) { +/// next event on the same stream is delivered. Over streamable HTTP (SSE) and over stdio, and in +/// either key order: JSON does not fix one, so with `params_first` the giant event arrives as +/// `{"params":{"data":…,"_meta":…,"cursor":…},"method":…}` — nothing that names or routes it is +/// within the bounded head, and it must still be skipped and reported, not looped on or lost. +fn an_over_cap_push_event_is_skipped_and_reported(stdio: bool, params_first: bool) { let events = json!([{ "name": "ticket.updated", "delivery": "push", "action": "follow_up" }]); let home = tempfile::tempdir().unwrap(); let control_file = home.path().join("control"); - let http_fixture = (!stdio).then(|| spawn_http_fixture(&[("MCP_FIXTURE_HEARTBEAT_MS", "200")])); + let order = if params_first { "params_first" } else { "" }; + let http_fixture = (!stdio).then(|| { + spawn_http_fixture(&[ + ("MCP_FIXTURE_HEARTBEAT_MS", "200"), + ("MCP_FIXTURE_KEY_ORDER", order), + ]) + }); let server = match &http_fixture { Some((_fx, mcp_url, _control)) => { json!({ "name": "tickets", "transport": "http", "url": mcp_url, "events": events }) @@ -180,7 +189,7 @@ fn an_over_cap_push_event_is_skipped_and_reported(stdio: bool) { None => stdio_server( "tickets", &control_file, - json!({ "MCP_FIXTURE_HEARTBEAT_MS": "200" }), + json!({ "MCP_FIXTURE_HEARTBEAT_MS": "200", "MCP_FIXTURE_KEY_ORDER": order }), events, ), }; @@ -222,12 +231,15 @@ fn an_over_cap_push_event_is_skipped_and_reported(stdio: bool) { |f| f["type"] == "mcp_event_status" && f["kind"] == "gap", ); assert_eq!(gap["reason"], "oversized", "{gap:#}"); - // The fixture serializes `cursor` ahead of the payload (and `eventId` after it): the cursor is - // kept, the id is not known. - assert!( - gap["cursor"].is_string(), - "the skipped event's cursor is kept: {gap:#}" - ); + // In the usual order the fixture serializes `cursor` ahead of the payload (and `eventId` after + // it): the cursor is kept, the id is not known. Payload first, neither is: the position moves + // on with the next event or heartbeat instead. + if !params_first { + assert!( + gap["cursor"].is_string(), + "the skipped event's cursor is kept: {gap:#}" + ); + } let f = frames.wait(Duration::from_secs(20), "the next event", |f| { f["type"] == "mcp_event" && f["event"]["eventId"] == "small-1" }); @@ -256,10 +268,20 @@ fn an_over_cap_push_event_is_skipped_and_reported(stdio: bool) { #[test] fn an_over_cap_push_event_over_http_is_skipped_and_reported_and_the_next_is_delivered() { - an_over_cap_push_event_is_skipped_and_reported(false); + an_over_cap_push_event_is_skipped_and_reported(false, false); } #[test] fn an_over_cap_push_event_over_stdio_is_skipped_and_reported_and_the_next_is_delivered() { - an_over_cap_push_event_is_skipped_and_reported(true); + an_over_cap_push_event_is_skipped_and_reported(true, false); +} + +#[test] +fn an_over_cap_push_event_with_params_first_over_http_is_skipped_and_reported_not_looped_on() { + an_over_cap_push_event_is_skipped_and_reported(false, true); +} + +#[test] +fn an_over_cap_push_event_with_params_first_over_stdio_is_skipped_and_reported_not_lost() { + an_over_cap_push_event_is_skipped_and_reported(true, true); } diff --git a/crates/agent/tests/mcp_unauthorized.rs b/crates/agent/tests/mcp_unauthorized.rs new file mode 100644 index 00000000..6335052e --- /dev/null +++ b/crates/agent/tests/mcp_unauthorized.rs @@ -0,0 +1,94 @@ +//! What a tool call that a streamable-HTTP server answers 401 tells the model. Every POST goes +//! through `mcp_wire::HttpClient::post_bounded`, which decides a 401 from its status (so a login can +//! be refreshed) — and the server's own message must survive that: an "invalid API key" from a +//! server with a static key is the tool error, and a login the server keeps refusing says to run +//! `agent mcp-login`, on the tool call and not only at connect. +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] + +mod common; + +use std::sync::atomic::Ordering; + +use common::mcp_oauth_fixture::{ + OAuthFixture, echo, logged_in, run, tool_results, write_global_settings, +}; +use common::turn_text; +use serde_json::json; + +/// `toolu_`'s result in the model request that followed it: `(text, is_error)`. +fn result_of(bodies: &[String], request: usize, id: &str) -> (String, bool) { + let results = tool_results(&bodies[request]); + let (_, text, is_error) = results + .into_iter() + .find(|(i, ..)| i == id) + .unwrap_or_else(|| panic!("no result for {id} in request {request}")); + (text, is_error) +} + +#[test] +fn a_401_from_a_server_with_a_static_key_is_the_tool_error_with_the_servers_message() { + let fixture = OAuthFixture::spawn(3600); + fixture.issue("static-key"); + let home = tempfile::tempdir().unwrap(); + write_global_settings( + home.path(), + json!([{ "name": "protected", "transport": "http", "url": fixture.url, + "headers": { "Authorization": "Bearer static-key" } }]), + ); + // The key works for one call, then the server stops accepting it — answering with a JSON-RPC + // error that says why. + fixture.reject_with_json_body.store(true, Ordering::SeqCst); + fixture.revoke_after_calls.store(1, Ordering::SeqCst); + let bodies = run( + home.path(), + vec![ + echo("toolu_1", "one"), + echo("toolu_2", "two"), + turn_text("done"), + ], + ); + let (first, failed) = result_of(&bodies, 1, "toolu_1"); + assert!(!failed && first.contains("one"), "{first}"); + let (text, is_error) = result_of(&bodies, 2, "toolu_2"); + assert!(is_error, "{text}"); + assert!( + text.contains("unauthorized: token expired"), + "the server's own message: {text}" + ); + assert!(!text.contains("mcp-login"), "no login to run: {text}"); +} + +fn refused_login(json_body: bool) -> String { + let (home, fixture) = logged_in(); + fixture + .reject_with_json_body + .store(json_body, Ordering::SeqCst); + fixture.reject_calls.store(true, Ordering::SeqCst); + let bodies = run(home.path(), vec![echo("toolu_1", "one"), turn_text("done")]); + assert_eq!( + fixture.refresh_grants.load(Ordering::SeqCst), + 1, + "the 401 was refreshed once, and the retry refused too" + ); + let (text, is_error) = result_of(&bodies, 1, "toolu_1"); + assert!(is_error, "{text}"); + assert!( + text.contains("agent mcp-login protected"), + "the tool error says how to fix it: {text}" + ); + text +} + +#[test] +fn a_login_the_server_keeps_refusing_names_mcp_login_on_the_tool_error_and_keeps_its_message() { + let text = refused_login(true); + assert!( + text.contains("unauthorized: token expired"), + "the server's own message: {text}" + ); +} + +#[test] +fn a_login_refused_with_a_bare_challenge_names_mcp_login_on_the_tool_error() { + refused_login(false); +}