mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
[3/4] Add executor-backed RMCP HTTP client (#18583)
### Why The RMCP layer needs a Streamable HTTP client that can talk either directly over `reqwest` or through the executor HTTP runner without duplicating MCP session logic higher in the stack. This PR adds that client-side transport boundary so remote Streamable HTTP MCP can reuse the same RMCP flow as the local path. ### What - Add a shared `rmcp-client/src/streamable_http/` module with: - `transport_client.rs` for the local-or-remote transport enum - `local_client.rs` for the direct `reqwest` implementation - `remote_client.rs` for the executor-backed implementation - `common.rs` for the small shared Streamable HTTP helpers - Teach `RmcpClient` to build Streamable HTTP transports in either local or remote mode while keeping the existing OAuth ownership in RMCP. - Translate remote POST, GET, and DELETE session operations into executor `http/request` calls. - Preserve RMCP session expiry handling and reconnect behavior for the remote transport. - Add remote transport coverage in `rmcp-client/tests/streamable_http_remote.rs` and keep the shared test support in `rmcp-client/tests/streamable_http_test_support.rs`. ### Verification - `cargo check -p codex-rmcp-client` - online CI ### Stack 1. #18581 protocol 2. #18582 runner 3. #18583 RMCP client 4. #18584 manager wiring and local/remote coverage --------- Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
committed by
GitHub
Unverified
parent
83ec1eb5d6
commit
0e78ce80ee
@@ -0,0 +1,391 @@
|
||||
//! RMCP Streamable HTTP adapter built on top of the shared `HttpClient`
|
||||
//! capability.
|
||||
//!
|
||||
//! This module runs in the orchestrator process. It turns high-level RMCP
|
||||
//! operations like `post_message` and `get_stream` into calls on
|
||||
//! `Arc<dyn HttpClient>`, which may be:
|
||||
//! - a local HTTP client that issues requests from the orchestrator, or
|
||||
//! - a remote HTTP client that forwards requests to the remote runtime
|
||||
|
||||
use std::io;
|
||||
use std::sync::Arc;
|
||||
|
||||
use bytes::Bytes;
|
||||
use codex_exec_server::ExecServerError;
|
||||
use codex_exec_server::HttpClient;
|
||||
use codex_exec_server::HttpHeader;
|
||||
use codex_exec_server::HttpRequestParams;
|
||||
use codex_exec_server::HttpResponseBodyStream;
|
||||
use futures::StreamExt;
|
||||
use futures::stream;
|
||||
use futures::stream::BoxStream;
|
||||
use reqwest::StatusCode;
|
||||
use reqwest::header::ACCEPT;
|
||||
use reqwest::header::AUTHORIZATION;
|
||||
use reqwest::header::CONTENT_TYPE;
|
||||
use reqwest::header::HeaderMap;
|
||||
use reqwest::header::HeaderName;
|
||||
use rmcp::model::ClientJsonRpcMessage;
|
||||
use rmcp::model::ServerJsonRpcMessage;
|
||||
use rmcp::transport::streamable_http_client::AuthRequiredError;
|
||||
use rmcp::transport::streamable_http_client::StreamableHttpClient;
|
||||
use rmcp::transport::streamable_http_client::StreamableHttpError;
|
||||
use rmcp::transport::streamable_http_client::StreamableHttpPostResponse;
|
||||
use sse_stream::Sse;
|
||||
use sse_stream::SseStream;
|
||||
|
||||
const EVENT_STREAM_MIME_TYPE: &str = "text/event-stream";
|
||||
const JSON_MIME_TYPE: &str = "application/json";
|
||||
const HEADER_SESSION_ID: &str = "Mcp-Session-Id";
|
||||
const NON_JSON_RESPONSE_BODY_PREVIEW_BYTES: usize = 8_192;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct StreamableHttpClientAdapter {
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
default_headers: HeaderMap,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub(crate) enum StreamableHttpClientAdapterError {
|
||||
#[error("streamable HTTP session expired with 404 Not Found")]
|
||||
SessionExpired404,
|
||||
#[error(transparent)]
|
||||
HttpRequest(#[from] ExecServerError),
|
||||
#[error("invalid HTTP header: {0}")]
|
||||
Header(String),
|
||||
}
|
||||
|
||||
impl StreamableHttpClientAdapter {
|
||||
pub(crate) fn new(http_client: Arc<dyn HttpClient>, default_headers: HeaderMap) -> Self {
|
||||
Self {
|
||||
http_client,
|
||||
default_headers,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl StreamableHttpClient for StreamableHttpClientAdapter {
|
||||
type Error = StreamableHttpClientAdapterError;
|
||||
|
||||
async fn post_message(
|
||||
&self,
|
||||
uri: Arc<str>,
|
||||
message: ClientJsonRpcMessage,
|
||||
session_id: Option<Arc<str>>,
|
||||
auth_token: Option<String>,
|
||||
) -> std::result::Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
|
||||
let mut headers = self.default_headers.clone();
|
||||
insert_header(
|
||||
&mut headers,
|
||||
ACCEPT,
|
||||
[EVENT_STREAM_MIME_TYPE, JSON_MIME_TYPE].join(", "),
|
||||
StreamableHttpClientAdapterError::Header,
|
||||
)?;
|
||||
insert_header(
|
||||
&mut headers,
|
||||
CONTENT_TYPE,
|
||||
JSON_MIME_TYPE.to_string(),
|
||||
StreamableHttpClientAdapterError::Header,
|
||||
)?;
|
||||
if let Some(auth_token) = auth_token {
|
||||
insert_header(
|
||||
&mut headers,
|
||||
AUTHORIZATION,
|
||||
format!("Bearer {auth_token}"),
|
||||
StreamableHttpClientAdapterError::Header,
|
||||
)?;
|
||||
}
|
||||
if let Some(session_id_value) = session_id.as_ref() {
|
||||
insert_header(
|
||||
&mut headers,
|
||||
HeaderName::from_static("mcp-session-id"),
|
||||
session_id_value.to_string(),
|
||||
StreamableHttpClientAdapterError::Header,
|
||||
)?;
|
||||
}
|
||||
|
||||
let body = serde_json::to_vec(&message).map_err(StreamableHttpError::Deserialize)?;
|
||||
let (response, mut body_stream) = self
|
||||
.http_client
|
||||
.http_request_stream(HttpRequestParams {
|
||||
method: "POST".to_string(),
|
||||
url: uri.to_string(),
|
||||
headers: protocol_headers(&headers),
|
||||
body: Some(body.into()),
|
||||
timeout_ms: None,
|
||||
request_id: "buffered-request".to_string(),
|
||||
stream_response: true,
|
||||
})
|
||||
.await
|
||||
.map_err(StreamableHttpClientAdapterError::from)
|
||||
.map_err(StreamableHttpError::Client)?;
|
||||
|
||||
if response.status == StatusCode::NOT_FOUND.as_u16() && session_id.is_some() {
|
||||
return Err(StreamableHttpError::Client(
|
||||
StreamableHttpClientAdapterError::SessionExpired404,
|
||||
));
|
||||
}
|
||||
if response.status == StatusCode::UNAUTHORIZED.as_u16()
|
||||
&& let Some(header) =
|
||||
response_header(&response.headers, reqwest::header::WWW_AUTHENTICATE)
|
||||
{
|
||||
return Err(StreamableHttpError::AuthRequired(AuthRequiredError {
|
||||
www_authenticate_header: header,
|
||||
}));
|
||||
}
|
||||
if matches!(
|
||||
StatusCode::from_u16(response.status).ok(),
|
||||
Some(StatusCode::ACCEPTED | StatusCode::NO_CONTENT)
|
||||
) {
|
||||
return Ok(StreamableHttpPostResponse::Accepted);
|
||||
}
|
||||
|
||||
let content_type = response_header(&response.headers, CONTENT_TYPE);
|
||||
let session_id = response_header(&response.headers, HEADER_SESSION_ID);
|
||||
match content_type.as_deref() {
|
||||
Some(content_type) if content_type.starts_with(EVENT_STREAM_MIME_TYPE) => {
|
||||
let event_stream = sse_stream_from_body(body_stream);
|
||||
Ok(StreamableHttpPostResponse::Sse(event_stream, session_id))
|
||||
}
|
||||
Some(content_type) if content_type.starts_with(JSON_MIME_TYPE) => {
|
||||
let body = collect_body(&mut body_stream).await?;
|
||||
let message: ServerJsonRpcMessage =
|
||||
serde_json::from_slice(&body).map_err(StreamableHttpError::Deserialize)?;
|
||||
Ok(StreamableHttpPostResponse::Json(message, session_id))
|
||||
}
|
||||
_ => {
|
||||
let body = collect_body(&mut body_stream).await?;
|
||||
let content_type = content_type.unwrap_or_else(|| "missing-content-type".into());
|
||||
Err(StreamableHttpError::UnexpectedContentType(Some(format!(
|
||||
"{content_type}; body: {}",
|
||||
body_preview(String::from_utf8_lossy(&body).to_string())
|
||||
))))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_session(
|
||||
&self,
|
||||
uri: Arc<str>,
|
||||
session: Arc<str>,
|
||||
auth_token: Option<String>,
|
||||
) -> std::result::Result<(), StreamableHttpError<Self::Error>> {
|
||||
let mut headers = self.default_headers.clone();
|
||||
if let Some(auth_token) = auth_token {
|
||||
insert_header(
|
||||
&mut headers,
|
||||
AUTHORIZATION,
|
||||
format!("Bearer {auth_token}"),
|
||||
StreamableHttpClientAdapterError::Header,
|
||||
)?;
|
||||
}
|
||||
insert_header(
|
||||
&mut headers,
|
||||
HeaderName::from_static("mcp-session-id"),
|
||||
session.to_string(),
|
||||
StreamableHttpClientAdapterError::Header,
|
||||
)?;
|
||||
|
||||
let response = self
|
||||
.http_client
|
||||
.http_request(HttpRequestParams {
|
||||
method: "DELETE".to_string(),
|
||||
url: uri.to_string(),
|
||||
headers: protocol_headers(&headers),
|
||||
body: None,
|
||||
timeout_ms: None,
|
||||
request_id: "buffered-request".to_string(),
|
||||
stream_response: false,
|
||||
})
|
||||
.await
|
||||
.map_err(StreamableHttpClientAdapterError::from)
|
||||
.map_err(StreamableHttpError::Client)?;
|
||||
|
||||
if response.status == StatusCode::METHOD_NOT_ALLOWED.as_u16() {
|
||||
return Ok(());
|
||||
}
|
||||
if !status_is_success(response.status) {
|
||||
return Err(StreamableHttpError::UnexpectedServerResponse(
|
||||
format!("DELETE returned HTTP {}", response.status).into(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_stream(
|
||||
&self,
|
||||
uri: Arc<str>,
|
||||
session_id: Arc<str>,
|
||||
last_event_id: Option<String>,
|
||||
auth_token: Option<String>,
|
||||
) -> std::result::Result<
|
||||
BoxStream<'static, std::result::Result<Sse, sse_stream::Error>>,
|
||||
StreamableHttpError<Self::Error>,
|
||||
> {
|
||||
let mut headers = self.default_headers.clone();
|
||||
insert_header(
|
||||
&mut headers,
|
||||
ACCEPT,
|
||||
[EVENT_STREAM_MIME_TYPE, JSON_MIME_TYPE].join(", "),
|
||||
StreamableHttpClientAdapterError::Header,
|
||||
)?;
|
||||
insert_header(
|
||||
&mut headers,
|
||||
HeaderName::from_static("mcp-session-id"),
|
||||
session_id.to_string(),
|
||||
StreamableHttpClientAdapterError::Header,
|
||||
)?;
|
||||
if let Some(last_event_id) = last_event_id {
|
||||
insert_header(
|
||||
&mut headers,
|
||||
HeaderName::from_static("last-event-id"),
|
||||
last_event_id,
|
||||
StreamableHttpClientAdapterError::Header,
|
||||
)?;
|
||||
}
|
||||
if let Some(auth_token) = auth_token {
|
||||
insert_header(
|
||||
&mut headers,
|
||||
AUTHORIZATION,
|
||||
format!("Bearer {auth_token}"),
|
||||
StreamableHttpClientAdapterError::Header,
|
||||
)?;
|
||||
}
|
||||
|
||||
let (response, body_stream) = self
|
||||
.http_client
|
||||
.http_request_stream(HttpRequestParams {
|
||||
method: "GET".to_string(),
|
||||
url: uri.to_string(),
|
||||
headers: protocol_headers(&headers),
|
||||
body: None,
|
||||
timeout_ms: None,
|
||||
request_id: "buffered-request".to_string(),
|
||||
stream_response: true,
|
||||
})
|
||||
.await
|
||||
.map_err(StreamableHttpClientAdapterError::from)
|
||||
.map_err(StreamableHttpError::Client)?;
|
||||
|
||||
if response.status == StatusCode::METHOD_NOT_ALLOWED.as_u16() {
|
||||
return Err(StreamableHttpError::ServerDoesNotSupportSse);
|
||||
}
|
||||
if response.status == StatusCode::NOT_FOUND.as_u16() {
|
||||
return Err(StreamableHttpError::Client(
|
||||
StreamableHttpClientAdapterError::SessionExpired404,
|
||||
));
|
||||
}
|
||||
if !status_is_success(response.status) {
|
||||
return Err(StreamableHttpError::UnexpectedServerResponse(
|
||||
format!("GET returned HTTP {}", response.status).into(),
|
||||
));
|
||||
}
|
||||
|
||||
match response_header(&response.headers, CONTENT_TYPE).as_deref() {
|
||||
Some(content_type) if is_streamable_http_content_type(content_type) => {}
|
||||
Some(content_type) => {
|
||||
return Err(StreamableHttpError::UnexpectedContentType(Some(
|
||||
content_type.to_string(),
|
||||
)));
|
||||
}
|
||||
None => {
|
||||
return Err(StreamableHttpError::UnexpectedContentType(None));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(sse_stream_from_body(body_stream))
|
||||
}
|
||||
}
|
||||
|
||||
fn body_preview(body: impl Into<String>) -> String {
|
||||
let mut body_preview = body.into();
|
||||
let body_len = body_preview.len();
|
||||
if body_len > NON_JSON_RESPONSE_BODY_PREVIEW_BYTES {
|
||||
let mut boundary = NON_JSON_RESPONSE_BODY_PREVIEW_BYTES;
|
||||
while !body_preview.is_char_boundary(boundary) {
|
||||
boundary = boundary.saturating_sub(1);
|
||||
}
|
||||
body_preview.truncate(boundary);
|
||||
body_preview.push_str(&format!(
|
||||
"... (truncated {} bytes)",
|
||||
body_len.saturating_sub(boundary)
|
||||
));
|
||||
}
|
||||
body_preview
|
||||
}
|
||||
|
||||
fn insert_header<Error>(
|
||||
headers: &mut HeaderMap,
|
||||
name: HeaderName,
|
||||
value: String,
|
||||
map_error: impl FnOnce(String) -> Error,
|
||||
) -> std::result::Result<(), StreamableHttpError<Error>>
|
||||
where
|
||||
Error: std::error::Error + Send + Sync + 'static,
|
||||
{
|
||||
let value = reqwest::header::HeaderValue::from_str(&value)
|
||||
.map_err(|error| StreamableHttpError::Client(map_error(error.to_string())))?;
|
||||
headers.insert(name, value);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn is_streamable_http_content_type(content_type: &str) -> bool {
|
||||
content_type
|
||||
.as_bytes()
|
||||
.starts_with(EVENT_STREAM_MIME_TYPE.as_bytes())
|
||||
|| content_type
|
||||
.as_bytes()
|
||||
.starts_with(JSON_MIME_TYPE.as_bytes())
|
||||
}
|
||||
|
||||
fn protocol_headers(headers: &HeaderMap) -> Vec<HttpHeader> {
|
||||
headers
|
||||
.iter()
|
||||
.filter_map(|(name, value)| {
|
||||
Some(HttpHeader {
|
||||
name: name.as_str().to_string(),
|
||||
value: value.to_str().ok()?.to_string(),
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn response_header(headers: &[HttpHeader], name: impl AsRef<str>) -> Option<String> {
|
||||
let name = name.as_ref();
|
||||
headers
|
||||
.iter()
|
||||
.find(|header| header.name.eq_ignore_ascii_case(name))
|
||||
.map(|header| header.value.clone())
|
||||
}
|
||||
|
||||
fn status_is_success(status: u16) -> bool {
|
||||
StatusCode::from_u16(status).is_ok_and(|status| status.is_success())
|
||||
}
|
||||
|
||||
async fn collect_body(
|
||||
body_stream: &mut HttpResponseBodyStream,
|
||||
) -> std::result::Result<Vec<u8>, StreamableHttpError<StreamableHttpClientAdapterError>> {
|
||||
let mut body = Vec::new();
|
||||
while let Some(chunk) = body_stream
|
||||
.recv()
|
||||
.await
|
||||
.map_err(StreamableHttpClientAdapterError::from)
|
||||
.map_err(StreamableHttpError::Client)?
|
||||
{
|
||||
body.extend_from_slice(&chunk);
|
||||
}
|
||||
Ok(body)
|
||||
}
|
||||
|
||||
fn sse_stream_from_body(
|
||||
body_stream: HttpResponseBodyStream,
|
||||
) -> BoxStream<'static, std::result::Result<Sse, sse_stream::Error>> {
|
||||
SseStream::from_byte_stream(stream::unfold(body_stream, |mut body_stream| async move {
|
||||
match body_stream.recv().await {
|
||||
Ok(Some(bytes)) => Some((Ok(Bytes::from(bytes)), body_stream)),
|
||||
Ok(None) => None,
|
||||
Err(error) => Some((Err(io::Error::other(error)), body_stream)),
|
||||
}
|
||||
}))
|
||||
.boxed()
|
||||
}
|
||||
@@ -1,6 +1,7 @@
|
||||
mod auth_status;
|
||||
mod elicitation_client_service;
|
||||
mod executor_process_transport;
|
||||
mod http_client_adapter;
|
||||
mod logging_client_handler;
|
||||
mod oauth;
|
||||
mod perform_oauth_login;
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
use std::borrow::Cow;
|
||||
use std::collections::HashMap;
|
||||
use std::ffi::OsString;
|
||||
use std::future::Future;
|
||||
@@ -14,16 +13,12 @@ use anyhow::Result;
|
||||
use anyhow::anyhow;
|
||||
use codex_client::build_reqwest_client_with_custom_ca;
|
||||
use codex_config::types::McpServerEnvVar;
|
||||
use codex_exec_server::HttpClient;
|
||||
use futures::FutureExt;
|
||||
use futures::StreamExt;
|
||||
use futures::future::BoxFuture;
|
||||
use futures::stream::BoxStream;
|
||||
use oauth2::TokenResponse;
|
||||
use reqwest::header::ACCEPT;
|
||||
use reqwest::header::AUTHORIZATION;
|
||||
use reqwest::header::CONTENT_TYPE;
|
||||
use reqwest::header::HeaderMap;
|
||||
use reqwest::header::WWW_AUTHENTICATE;
|
||||
use rmcp::model::CallToolRequestParams;
|
||||
use rmcp::model::CallToolResult;
|
||||
use rmcp::model::ClientNotification;
|
||||
@@ -52,16 +47,11 @@ use rmcp::transport::StreamableHttpClientTransport;
|
||||
use rmcp::transport::auth::AuthClient;
|
||||
use rmcp::transport::auth::AuthError;
|
||||
use rmcp::transport::auth::OAuthState;
|
||||
use rmcp::transport::streamable_http_client::AuthRequiredError;
|
||||
use rmcp::transport::streamable_http_client::StreamableHttpClient;
|
||||
use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig;
|
||||
use rmcp::transport::streamable_http_client::StreamableHttpError;
|
||||
use rmcp::transport::streamable_http_client::StreamableHttpPostResponse;
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
use sse_stream::Sse;
|
||||
use sse_stream::SseStream;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::sync::Semaphore;
|
||||
use tokio::sync::watch;
|
||||
@@ -69,6 +59,8 @@ use tokio::time;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::elicitation_client_service::ElicitationClientService;
|
||||
use crate::http_client_adapter::StreamableHttpClientAdapter;
|
||||
use crate::http_client_adapter::StreamableHttpClientAdapterError;
|
||||
use crate::load_oauth_tokens;
|
||||
use crate::oauth::OAuthPersistor;
|
||||
use crate::oauth::StoredOAuthTokens;
|
||||
@@ -79,239 +71,15 @@ use crate::utils::apply_default_headers;
|
||||
use crate::utils::build_default_headers;
|
||||
use codex_config::types::OAuthCredentialsStoreMode;
|
||||
|
||||
const EVENT_STREAM_MIME_TYPE: &str = "text/event-stream";
|
||||
const JSON_MIME_TYPE: &str = "application/json";
|
||||
const HEADER_LAST_EVENT_ID: &str = "Last-Event-Id";
|
||||
const HEADER_SESSION_ID: &str = "Mcp-Session-Id";
|
||||
const NON_JSON_RESPONSE_BODY_PREVIEW_BYTES: usize = 8_192;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct StreamableHttpResponseClient {
|
||||
inner: reqwest::Client,
|
||||
}
|
||||
|
||||
impl StreamableHttpResponseClient {
|
||||
fn new(inner: reqwest::Client) -> Self {
|
||||
Self { inner }
|
||||
}
|
||||
|
||||
fn reqwest_error(
|
||||
error: reqwest::Error,
|
||||
) -> StreamableHttpError<StreamableHttpResponseClientError> {
|
||||
StreamableHttpError::Client(StreamableHttpResponseClientError::from(error))
|
||||
}
|
||||
}
|
||||
|
||||
fn build_http_client(default_headers: &HeaderMap) -> Result<reqwest::Client> {
|
||||
let builder = apply_default_headers(reqwest::Client::builder(), default_headers);
|
||||
Ok(build_reqwest_client_with_custom_ca(builder)?)
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
enum StreamableHttpResponseClientError {
|
||||
#[error("streamable HTTP session expired with 404 Not Found")]
|
||||
SessionExpired404,
|
||||
#[error(transparent)]
|
||||
Reqwest(#[from] reqwest::Error),
|
||||
}
|
||||
|
||||
impl StreamableHttpClient for StreamableHttpResponseClient {
|
||||
type Error = StreamableHttpResponseClientError;
|
||||
|
||||
async fn post_message(
|
||||
&self,
|
||||
uri: Arc<str>,
|
||||
message: rmcp::model::ClientJsonRpcMessage,
|
||||
session_id: Option<Arc<str>>,
|
||||
auth_token: Option<String>,
|
||||
) -> std::result::Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
|
||||
let mut request = self
|
||||
.inner
|
||||
.post(uri.as_ref())
|
||||
.header(ACCEPT, [EVENT_STREAM_MIME_TYPE, JSON_MIME_TYPE].join(", "));
|
||||
if let Some(auth_header) = auth_token {
|
||||
request = request.bearer_auth(auth_header);
|
||||
}
|
||||
if let Some(session_id_value) = session_id.as_ref() {
|
||||
request = request.header(HEADER_SESSION_ID, session_id_value.as_ref());
|
||||
}
|
||||
|
||||
let response = request
|
||||
.json(&message)
|
||||
.send()
|
||||
.await
|
||||
.map_err(StreamableHttpResponseClient::reqwest_error)?;
|
||||
if response.status() == reqwest::StatusCode::NOT_FOUND && session_id.is_some() {
|
||||
return Err(StreamableHttpError::Client(
|
||||
StreamableHttpResponseClientError::SessionExpired404,
|
||||
));
|
||||
}
|
||||
if response.status() == reqwest::StatusCode::UNAUTHORIZED
|
||||
&& let Some(header) = response.headers().get(WWW_AUTHENTICATE)
|
||||
{
|
||||
let header = header
|
||||
.to_str()
|
||||
.map_err(|_| {
|
||||
StreamableHttpError::UnexpectedServerResponse(Cow::Borrowed(
|
||||
"invalid www-authenticate header value",
|
||||
))
|
||||
})?
|
||||
.to_string();
|
||||
return Err(StreamableHttpError::AuthRequired(AuthRequiredError {
|
||||
www_authenticate_header: header,
|
||||
}));
|
||||
}
|
||||
|
||||
let status = response.status();
|
||||
if matches!(
|
||||
status,
|
||||
reqwest::StatusCode::ACCEPTED | reqwest::StatusCode::NO_CONTENT
|
||||
) {
|
||||
return Ok(StreamableHttpPostResponse::Accepted);
|
||||
}
|
||||
|
||||
let content_type = response
|
||||
.headers()
|
||||
.get(CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_string);
|
||||
let session_id = response
|
||||
.headers()
|
||||
.get(HEADER_SESSION_ID)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_string);
|
||||
|
||||
match content_type.as_deref() {
|
||||
Some(ct) if ct.as_bytes().starts_with(EVENT_STREAM_MIME_TYPE.as_bytes()) => {
|
||||
let event_stream = SseStream::from_byte_stream(response.bytes_stream()).boxed();
|
||||
Ok(StreamableHttpPostResponse::Sse(event_stream, session_id))
|
||||
}
|
||||
Some(ct) if ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes()) => {
|
||||
let message = response
|
||||
.json()
|
||||
.await
|
||||
.map_err(StreamableHttpResponseClient::reqwest_error)?;
|
||||
Ok(StreamableHttpPostResponse::Json(message, session_id))
|
||||
}
|
||||
_ => {
|
||||
let body = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(StreamableHttpResponseClient::reqwest_error)?;
|
||||
let mut body_preview = body;
|
||||
let body_len = body_preview.len();
|
||||
if body_len > NON_JSON_RESPONSE_BODY_PREVIEW_BYTES {
|
||||
let mut boundary = NON_JSON_RESPONSE_BODY_PREVIEW_BYTES;
|
||||
while !body_preview.is_char_boundary(boundary) {
|
||||
boundary = boundary.saturating_sub(1);
|
||||
}
|
||||
body_preview.truncate(boundary);
|
||||
body_preview.push_str(&format!(
|
||||
"... (truncated {} bytes)",
|
||||
body_len.saturating_sub(boundary)
|
||||
));
|
||||
}
|
||||
|
||||
let content_type = content_type.unwrap_or_else(|| "missing-content-type".into());
|
||||
Err(StreamableHttpError::UnexpectedContentType(Some(format!(
|
||||
"{content_type}; body: {body_preview}"
|
||||
))))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_session(
|
||||
&self,
|
||||
uri: Arc<str>,
|
||||
session: Arc<str>,
|
||||
auth_token: Option<String>,
|
||||
) -> std::result::Result<(), StreamableHttpError<Self::Error>> {
|
||||
let mut request_builder = self.inner.delete(uri.as_ref());
|
||||
if let Some(auth_header) = auth_token {
|
||||
request_builder = request_builder.bearer_auth(auth_header);
|
||||
}
|
||||
let response = request_builder
|
||||
.header(HEADER_SESSION_ID, session.as_ref())
|
||||
.send()
|
||||
.await
|
||||
.map_err(StreamableHttpResponseClient::reqwest_error)?;
|
||||
|
||||
if response.status() == reqwest::StatusCode::METHOD_NOT_ALLOWED {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
response
|
||||
.error_for_status()
|
||||
.map_err(StreamableHttpResponseClient::reqwest_error)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_stream(
|
||||
&self,
|
||||
uri: Arc<str>,
|
||||
session_id: Arc<str>,
|
||||
last_event_id: Option<String>,
|
||||
auth_token: Option<String>,
|
||||
) -> std::result::Result<
|
||||
BoxStream<'static, std::result::Result<Sse, sse_stream::Error>>,
|
||||
StreamableHttpError<Self::Error>,
|
||||
> {
|
||||
let mut request_builder = self
|
||||
.inner
|
||||
.get(uri.as_ref())
|
||||
.header(ACCEPT, [EVENT_STREAM_MIME_TYPE, JSON_MIME_TYPE].join(", "))
|
||||
.header(HEADER_SESSION_ID, session_id.as_ref());
|
||||
if let Some(last_event_id) = last_event_id {
|
||||
request_builder = request_builder.header(HEADER_LAST_EVENT_ID, last_event_id);
|
||||
}
|
||||
if let Some(auth_header) = auth_token {
|
||||
request_builder = request_builder.bearer_auth(auth_header);
|
||||
}
|
||||
|
||||
let response = request_builder
|
||||
.send()
|
||||
.await
|
||||
.map_err(StreamableHttpResponseClient::reqwest_error)?;
|
||||
if response.status() == reqwest::StatusCode::METHOD_NOT_ALLOWED {
|
||||
return Err(StreamableHttpError::ServerDoesNotSupportSse);
|
||||
}
|
||||
if response.status() == reqwest::StatusCode::NOT_FOUND {
|
||||
return Err(StreamableHttpError::Client(
|
||||
StreamableHttpResponseClientError::SessionExpired404,
|
||||
));
|
||||
}
|
||||
|
||||
let response = response
|
||||
.error_for_status()
|
||||
.map_err(StreamableHttpResponseClient::reqwest_error)?;
|
||||
match response.headers().get(CONTENT_TYPE) {
|
||||
Some(ct)
|
||||
if ct.as_bytes().starts_with(EVENT_STREAM_MIME_TYPE.as_bytes())
|
||||
|| ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes()) => {}
|
||||
Some(ct) => {
|
||||
return Err(StreamableHttpError::UnexpectedContentType(Some(
|
||||
String::from_utf8_lossy(ct.as_bytes()).to_string(),
|
||||
)));
|
||||
}
|
||||
None => {
|
||||
return Err(StreamableHttpError::UnexpectedContentType(None));
|
||||
}
|
||||
}
|
||||
|
||||
let event_stream = SseStream::from_byte_stream(response.bytes_stream()).boxed();
|
||||
Ok(event_stream)
|
||||
}
|
||||
}
|
||||
|
||||
enum PendingTransport {
|
||||
Stdio {
|
||||
transport: StdioServerTransport,
|
||||
},
|
||||
StreamableHttp {
|
||||
transport: StreamableHttpClientTransport<StreamableHttpResponseClient>,
|
||||
transport: StreamableHttpClientTransport<StreamableHttpClientAdapter>,
|
||||
},
|
||||
StreamableHttpWithOAuth {
|
||||
transport: StreamableHttpClientTransport<AuthClient<StreamableHttpResponseClient>>,
|
||||
transport: StreamableHttpClientTransport<AuthClient<StreamableHttpClientAdapter>>,
|
||||
oauth_persistor: OAuthPersistor,
|
||||
},
|
||||
}
|
||||
@@ -339,6 +107,7 @@ enum TransportRecipe {
|
||||
http_headers: Option<HashMap<String, String>>,
|
||||
env_http_headers: Option<HashMap<String, String>>,
|
||||
store_mode: OAuthCredentialsStoreMode,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -536,6 +305,7 @@ impl RmcpClient {
|
||||
http_headers: Option<HashMap<String, String>>,
|
||||
env_http_headers: Option<HashMap<String, String>>,
|
||||
store_mode: OAuthCredentialsStoreMode,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
) -> Result<Self> {
|
||||
let transport_recipe = TransportRecipe::StreamableHttp {
|
||||
server_name: server_name.to_string(),
|
||||
@@ -544,6 +314,7 @@ impl RmcpClient {
|
||||
http_headers,
|
||||
env_http_headers,
|
||||
store_mode,
|
||||
http_client,
|
||||
};
|
||||
let transport = Self::create_pending_transport(&transport_recipe).await?;
|
||||
Ok(Self {
|
||||
@@ -895,6 +666,7 @@ impl RmcpClient {
|
||||
http_headers,
|
||||
env_http_headers,
|
||||
store_mode,
|
||||
http_client,
|
||||
} => {
|
||||
let default_headers =
|
||||
build_default_headers(http_headers.clone(), env_http_headers.clone())?;
|
||||
@@ -919,6 +691,7 @@ impl RmcpClient {
|
||||
initial_tokens.clone(),
|
||||
*store_mode,
|
||||
default_headers.clone(),
|
||||
Arc::clone(http_client),
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -945,9 +718,11 @@ impl RmcpClient {
|
||||
let http_config =
|
||||
StreamableHttpClientTransportConfig::with_uri(url.clone())
|
||||
.auth_header(access_token);
|
||||
let http_client = build_http_client(&default_headers)?;
|
||||
let transport = StreamableHttpClientTransport::with_client(
|
||||
StreamableHttpResponseClient::new(http_client),
|
||||
StreamableHttpClientAdapter::new(
|
||||
Arc::clone(http_client),
|
||||
default_headers,
|
||||
),
|
||||
http_config,
|
||||
);
|
||||
Ok(PendingTransport::StreamableHttp { transport })
|
||||
@@ -961,10 +736,8 @@ impl RmcpClient {
|
||||
http_config = http_config.auth_header(bearer_token);
|
||||
}
|
||||
|
||||
let http_client = build_http_client(&default_headers)?;
|
||||
|
||||
let transport = StreamableHttpClientTransport::with_client(
|
||||
StreamableHttpResponseClient::new(http_client),
|
||||
StreamableHttpClientAdapter::new(Arc::clone(http_client), default_headers),
|
||||
http_config,
|
||||
);
|
||||
Ok(PendingTransport::StreamableHttp { transport })
|
||||
@@ -1084,12 +857,12 @@ impl RmcpClient {
|
||||
|
||||
error
|
||||
.error
|
||||
.downcast_ref::<StreamableHttpError<StreamableHttpResponseClientError>>()
|
||||
.downcast_ref::<StreamableHttpError<StreamableHttpClientAdapterError>>()
|
||||
.is_some_and(|error| {
|
||||
matches!(
|
||||
error,
|
||||
StreamableHttpError::Client(
|
||||
StreamableHttpResponseClientError::SessionExpired404
|
||||
StreamableHttpClientAdapterError::SessionExpired404
|
||||
)
|
||||
)
|
||||
})
|
||||
@@ -1156,12 +929,18 @@ async fn create_oauth_transport_and_runtime(
|
||||
initial_tokens: StoredOAuthTokens,
|
||||
credentials_store: OAuthCredentialsStoreMode,
|
||||
default_headers: HeaderMap,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
) -> Result<(
|
||||
StreamableHttpClientTransport<AuthClient<StreamableHttpResponseClient>>,
|
||||
StreamableHttpClientTransport<AuthClient<StreamableHttpClientAdapter>>,
|
||||
OAuthPersistor,
|
||||
)> {
|
||||
let http_client = build_http_client(&default_headers)?;
|
||||
let mut oauth_state = OAuthState::new(url.to_string(), Some(http_client.clone())).await?;
|
||||
let builder = apply_default_headers(reqwest::Client::builder(), &default_headers);
|
||||
let oauth_metadata_client = build_reqwest_client_with_custom_ca(builder)?;
|
||||
// TODO(aibrahim): teach OAuth bootstrap and refresh to use the same
|
||||
// shared HTTP client abstraction instead of always creating the local
|
||||
// reqwest metadata client here.
|
||||
let mut oauth_state =
|
||||
OAuthState::new(url.to_string(), Some(oauth_metadata_client.clone())).await?;
|
||||
|
||||
oauth_state
|
||||
.set_credentials(
|
||||
@@ -1178,7 +957,10 @@ async fn create_oauth_transport_and_runtime(
|
||||
}
|
||||
};
|
||||
|
||||
let auth_client = AuthClient::new(StreamableHttpResponseClient::new(http_client), manager);
|
||||
let auth_client = AuthClient::new(
|
||||
StreamableHttpClientAdapter::new(http_client, default_headers),
|
||||
manager,
|
||||
);
|
||||
let auth_manager = auth_client.auth_manager.clone();
|
||||
|
||||
let transport = StreamableHttpClientTransport::with_client(
|
||||
|
||||
Reference in New Issue
Block a user