use std::net::SocketAddr;
use std::num::NonZeroUsize;
use std::sync::Arc;
use serde_json::Value;
use tokio::sync::mpsc;
use crate::adapters::primary::transport::http::router::{
EventSource, Inbound, RequestHandler, Unsessioned,
};
use crate::adapters::primary::transport::http::{
EventTail, McpRouter, Refusal, ServerConfig, SessionHandle, SessionStore,
};
use crate::application::bridge::{
BridgeClientError, DownstreamConfig, DownstreamError, ToolResultCodec, Upstream, run_client,
};
use crate::application::jsonrpc::{FrameSource, SourceError};
use crate::application::session::{FrameSink, SinkError};
use crate::domain::{DegradePolicy, McpSessionId};
const TAIL_CAPACITY: NonZeroUsize = match NonZeroUsize::new(64) {
Some(capacity) => capacity,
None => NonZeroUsize::MIN,
};
#[derive(Debug)]
pub struct BridgeSession {
tail: EventTail<Value>,
inbox: mpsc::UnboundedSender<String>,
}
impl EventSource for BridgeSession {
fn tail(&self) -> &EventTail<Value> {
&self.tail
}
}
impl BridgeSession {
fn new() -> (Self, HttpFrameSource) {
let (inbox, frames) = mpsc::unbounded_channel();
let session = Self {
tail: EventTail::new(TAIL_CAPACITY),
inbox,
};
let source = HttpFrameSource {
frames: Arc::new(tokio::sync::Mutex::new(frames)),
};
(session, source)
}
fn push(&self, frame: String) -> Result<(), ()> {
self.inbox.send(frame).map_err(|_| ())
}
fn sink(&self) -> HttpFrameSink {
HttpFrameSink {
tail: self.tail.clone(),
}
}
}
struct HttpFrameSource {
frames: Arc<tokio::sync::Mutex<mpsc::UnboundedReceiver<String>>>,
}
impl core::fmt::Debug for HttpFrameSource {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("HttpFrameSource")
.field("frames", &"<one session's inbound messages>")
.finish()
}
}
impl FrameSource for HttpFrameSource {
fn next_frame(&self) -> impl Future<Output = Result<Option<String>, SourceError>> + Send {
let frames = Arc::clone(&self.frames);
async move {
let mut receiver = frames.lock().await;
Ok(receiver.recv().await)
}
}
}
#[derive(Debug)]
struct HttpFrameSink {
tail: EventTail<Value>,
}
impl FrameSink for HttpFrameSink {
fn send(&self, frame: &str) -> impl Future<Output = Result<(), SinkError>> + Send {
let pushed = serde_json::from_str::<Value>(frame)
.map(|value| self.tail.push(value))
.map_err(|_| SinkError);
std::future::ready(pushed.map(|_| ()))
}
}
#[derive(Debug)]
struct BridgeHandler<U, C> {
sessions: Arc<SessionStore<BridgeSession>>,
upstream: Arc<U>,
config: DownstreamConfig,
policy: DegradePolicy,
codec: C,
}
impl<U, C> BridgeHandler<U, C> {
fn new(
sessions: Arc<SessionStore<BridgeSession>>,
upstream: U,
config: DownstreamConfig,
policy: DegradePolicy,
codec: C,
) -> Self {
Self {
sessions,
upstream: Arc::new(upstream),
config,
policy,
codec,
}
}
}
fn accepted() -> axum::response::Response {
axum::response::Response::builder()
.status(axum::http::StatusCode::ACCEPTED)
.body(axum::body::Body::empty())
.unwrap_or_else(|_| axum::response::Response::new(axum::body::Body::empty()))
}
fn accepted_with(id: McpSessionId) -> Result<axum::response::Response, Refusal> {
let value = axum::http::HeaderValue::from_str(&id.to_string())
.map_err(|error| Refusal::Handler(format!("the id is not a header: {error}")))?;
let mut response = accepted();
response.headers_mut().insert("mcp-session-id", value);
Ok(response)
}
fn method_of(body: &[u8]) -> Option<String> {
serde_json::from_slice::<Value>(body)
.ok()?
.get("method")?
.as_str()
.map(String::from)
}
impl<U, C> RequestHandler<BridgeSession> for BridgeHandler<U, C>
where
U: Upstream + 'static,
C: ToolResultCodec + Clone + Send + Sync + 'static,
{
fn unsessioned(&self, body: &[u8], _headers: &axum::http::HeaderMap) -> Unsessioned {
match method_of(body).as_deref() {
Some("initialize") => Unsessioned::Establishing,
_ => Unsessioned::NeedsSession,
}
}
fn dispatch(
&self,
call: Inbound<'_, BridgeSession>,
) -> impl Future<Output = Result<axum::response::Response, Refusal>> + Send {
let frame = String::from_utf8(call.body.to_vec()).map_err(|_| Refusal::MalformedBody {
reason: "a request body is not UTF-8".to_owned(),
});
let session = match &call.session {
SessionHandle::Live { session, .. } => Ok(Some(Arc::clone(session))),
SessionHandle::Establishing => Ok(None),
SessionHandle::Stateless => Err(Refusal::Handler(
"a bridge is a conversation, and the stateless lifecycle has \
none; open a session"
.to_owned(),
)),
};
let sessions = Arc::clone(&self.sessions);
let upstream = Arc::clone(&self.upstream);
let config = self.config.clone();
let policy = self.policy;
let codec = self.codec.clone();
async move {
let frame = frame?;
match session? {
Some(session) => {
session
.push(frame)
.map_err(|()| Refusal::Handler("this session has ended".to_owned()))?;
Ok(accepted())
}
None => Self::open(&sessions, upstream, config, policy, codec, frame),
}
}
}
}
impl<U, C> BridgeHandler<U, C>
where
U: Upstream + 'static,
C: ToolResultCodec + Clone + Send + Sync + 'static,
{
fn open(
sessions: &SessionStore<BridgeSession>,
upstream: Arc<U>,
config: DownstreamConfig,
policy: DegradePolicy,
codec: C,
frame: String,
) -> Result<axum::response::Response, Refusal> {
let (session, source) = BridgeSession::new();
let session = Arc::new(session);
let sink = session.sink();
let id = sessions
.open(Arc::clone(&session))
.map_err(|error| Refusal::Handler(format!("no session for the client: {error}")))?;
tokio::spawn(async move {
let _ = run_client(upstream.as_ref(), policy, config, codec, source, sink).await;
});
session
.push(frame)
.map_err(|()| Refusal::Handler("the new session has ended".to_owned()))?;
accepted_with(id)
}
}
#[derive(Debug, thiserror::Error)]
pub enum BridgeServeError {
#[error("the router could not be built: {0}")]
Router(#[from] crate::adapters::primary::transport::http::router::RouterError),
#[error(transparent)]
Downstream(#[from] DownstreamError),
#[error(transparent)]
Client(#[from] BridgeClientError),
#[error("could not serve on {addr}: {source}")]
Bind {
addr: SocketAddr,
source: std::io::Error,
},
}
#[derive(Debug)]
pub struct HttpBridge<U, C> {
config: ServerConfig,
sessions: Arc<SessionStore<BridgeSession>>,
handler: BridgeHandler<U, C>,
}
impl<U, C> HttpBridge<U, C>
where
U: Upstream + 'static,
C: ToolResultCodec + Clone + Send + Sync + 'static,
{
pub fn new(
config: ServerConfig,
sessions: Arc<SessionStore<BridgeSession>>,
upstream: U,
identity: crate::application::bridge::Identity,
policy: DegradePolicy,
codec: C,
) -> Result<Self, BridgeServeError> {
let versions = crate::domain::SupportedVersions::all()
.map_err(|_| DownstreamError::UnnegotiableVersion)?;
Ok(Self {
config,
sessions: Arc::clone(&sessions),
handler: BridgeHandler::new(
sessions,
upstream,
DownstreamConfig::with(
versions,
identity,
crate::domain::TransportKind::StreamableHttp,
),
policy,
codec,
),
})
}
pub fn into_router(
self,
) -> Result<axum::Router, crate::adapters::primary::transport::http::router::RouterError> {
McpRouter::new(self.config, self.sessions, self.handler).into_router()
}
pub async fn bind(self, addr: SocketAddr) -> Result<(Bound, SocketAddr), BridgeServeError> {
let router = self.into_router()?;
let listener = tokio::net::TcpListener::bind(addr)
.await
.map_err(|source| BridgeServeError::Bind { addr, source })?;
let bound = listener
.local_addr()
.map_err(|source| BridgeServeError::Bind { addr, source })?;
Ok((
Bound {
router,
listener,
addr: bound,
},
bound,
))
}
pub async fn serve(self, addr: SocketAddr) -> Result<SocketAddr, BridgeServeError> {
let (bound, actual) = self.bind(addr).await?;
bound.serve().await?;
Ok(actual)
}
}
#[derive(Debug)]
pub struct Bound {
router: axum::Router,
listener: tokio::net::TcpListener,
addr: SocketAddr,
}
impl Bound {
#[must_use]
pub const fn addr(&self) -> SocketAddr {
self.addr
}
pub async fn serve(self) -> Result<(), BridgeServeError> {
axum::serve(self.listener, self.router)
.await
.map_err(|source| BridgeServeError::Bind {
addr: self.addr,
source,
})
}
}
#[cfg(test)]
mod tests {
use super::{BridgeSession, HttpBridge, HttpFrameSource};
use crate::adapters::primary::transport::http::{SecurityHeaders, ServerConfig, SessionStore};
use crate::application::bridge::SendFuture;
use crate::application::bridge::{BridgeError, BridgePeer, Identity, OpenedUpstream, Upstream};
use crate::application::jsonrpc::FrameSource;
use crate::application::session::{FrameSink, IdSource};
use crate::domain::{DegradePolicy, Feature, FeatureProfile, Leg, ProtocolVersion};
use serde_json::{Map, Value, json};
use std::net::{Ipv4Addr, SocketAddr};
use std::sync::Arc;
struct Counting(std::sync::atomic::AtomicU64);
impl IdSource for Counting {
fn next_session_id(&self) -> crate::domain::session::SessionId {
crate::domain::session::SessionId::new(
self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1,
)
}
}
struct Idle;
struct IdlePeer;
impl Upstream for Idle {
fn open(
&self,
) -> impl std::future::Future<
Output = Result<OpenedUpstream, crate::application::bridge::UpstreamError>,
> + Send {
let version = ProtocolVersion::V2025_11_25;
let tools = FeatureProfile::empty().with(Feature::Tools);
std::future::ready(Ok(OpenedUpstream::new(
Leg::new(version, tools),
Arc::new(IdlePeer),
)))
}
}
impl BridgePeer for IdlePeer {
fn send(
&self,
envelope: Map<String, Value>,
) -> SendFuture<'_, Result<Map<String, Value>, BridgeError>> {
Box::pin(async move { Ok(envelope) })
}
fn notify(&self, envelope: Map<String, Value>) -> SendFuture<'_, Result<(), BridgeError>> {
Box::pin(async move {
let _ = envelope;
Ok(())
})
}
fn close(&self) -> SendFuture<'_, Result<(), BridgeError>> {
Box::pin(async { Ok(()) })
}
}
fn bridge() -> Result<
HttpBridge<Idle, crate::adapters::secondary::codec::result::WireToolResults>,
Box<dyn std::error::Error>,
> {
let ids: Arc<dyn IdSource + Send + Sync> =
Arc::new(Counting(std::sync::atomic::AtomicU64::new(0)));
let sessions: Arc<SessionStore<BridgeSession>> = Arc::new(SessionStore::new(ids));
Ok(HttpBridge::new(
ServerConfig::new()?.with_security(SecurityHeaders::default()),
sessions,
Idle,
Identity::new("unit", "0.0.0"),
DegradePolicy::default(),
crate::adapters::secondary::codec::result::WireToolResults,
)?)
}
#[tokio::test]
async fn a_bridge_binds_and_reports_the_address_it_took()
-> Result<(), Box<dyn std::error::Error>> {
let loopback = SocketAddr::from((Ipv4Addr::LOCALHOST, 0));
let (bound, actual) = bridge()?.bind(loopback).await?;
if actual.port() == 0 {
return Err("the bridge reported port 0 after binding".into());
}
if bound.addr() != actual {
return Err(format!(
"the bound socket says {} and the report said {actual}",
bound.addr()
)
.into());
}
Ok(())
}
#[tokio::test]
async fn an_agreement_the_peer_cannot_hold_still_binds()
-> Result<(), Box<dyn std::error::Error>> {
let loopback = SocketAddr::from((Ipv4Addr::LOCALHOST, 0));
let (bound, _actual) = bridge()?.bind(loopback).await?;
drop(bound);
Ok(())
}
fn session() -> (BridgeSession, HttpFrameSource) {
BridgeSession::new()
}
#[tokio::test]
async fn a_frame_pushed_comes_back_out_in_order() -> Result<(), Box<dyn std::error::Error>> {
let (session, source) = session();
session
.push("first".to_owned())
.map_err(|()| "the loop had ended")?;
session
.push("second".to_owned())
.map_err(|()| "the loop had ended")?;
let first = source.next_frame().await?;
let second = source.next_frame().await?;
if (first.as_deref(), second.as_deref()) != (Some("first"), Some("second")) {
return Err(format!("the source yielded {first:?} then {second:?}").into());
}
Ok(())
}
#[tokio::test]
async fn a_dropped_session_ends_the_source_rather_than_parking_it()
-> Result<(), Box<dyn std::error::Error>> {
let (session, source) = session();
drop(session);
if source.next_frame().await?.is_some() {
return Err("a session that was dropped still produced a frame".into());
}
Ok(())
}
#[tokio::test]
async fn a_reply_is_queued_on_the_tail_the_get_route_streams()
-> Result<(), Box<dyn std::error::Error>> {
let (session, _source) = session();
let sink = session.sink();
if !sink.tail.is_empty() {
return Err("a fresh session already had a reply queued".into());
}
sink.send(&json!({"jsonrpc": "2.0", "id": 7, "result": {"ok": true}}).to_string())
.await?;
if sink.tail.is_empty() {
return Err("the reply was sent but nothing was queued".into());
}
Ok(())
}
#[tokio::test]
async fn a_frame_that_is_not_json_is_refused_rather_than_queued()
-> Result<(), Box<dyn std::error::Error>> {
let (session, _source) = session();
let sink = session.sink();
if sink.send("not json at all").await.is_ok() {
return Err("a frame that is not JSON was accepted".into());
}
Ok(())
}
}