monoloop_loop/transaction/mcp/
gateway.rs

1//! Loopback MCP Streamable HTTP gateway with per-capability routing.
2
3use super::binding::{CapabilityToken, McpInstallError, McpRouteTable, PendingMcpBinding};
4use super::handler::TransactionMcpHandler;
5use crate::transaction::dispatcher::TransactionToolDispatcher;
6use crate::transaction::resolved_tools::ResolvedToolSet;
7use axum::body::Body;
8use axum::extract::{Path, Request, State};
9use axum::http::{Response, StatusCode};
10use axum::routing::any;
11use axum::Router;
12use monoloop_contracts::{ExchangeId, TransactionId};
13use rmcp::transport::streamable_http_server::{
14    session::local::LocalSessionManager, StreamableHttpServerConfig, StreamableHttpService,
15};
16use std::future::Future;
17use std::net::SocketAddr;
18use std::pin::Pin;
19use std::sync::Arc;
20use tokio_util::sync::CancellationToken;
21
22/// Runs one MCP HTTP request under TaskSupervisor as `TaskClass::McpRequest` (§17).
23///
24/// Injected by RuntimeOwner; standalone prepare+serve tests leave this unset
25/// and execute request work inline.
26pub trait McpRequestOwner: Send + Sync {
27    /// Own `work` for `transaction_id` and return its response.
28    fn run_owned(
29        &self,
30        transaction_id: TransactionId,
31        work: Pin<Box<dyn Future<Output = Response<Body>> + Send>>,
32    ) -> Pin<Box<dyn Future<Output = Response<Body>> + Send>>;
33}
34
35/// Bounded MCP request concurrency and duration (D-034 / Law 22).
36///
37/// Production defaults match the historical constants. Tests inject smaller
38/// budgets for exact-limit and plus-one proofs without multi-second waits.
39#[derive(Clone, Debug, PartialEq, Eq)]
40pub struct McpGatewayLimits {
41    /// Gateway-wide concurrent in-flight MCP HTTP requests.
42    pub max_global_requests: usize,
43    /// Per-capability concurrent in-flight MCP HTTP requests.
44    pub max_per_capability_requests: usize,
45    /// Absolute wall budget for body read + Streamable HTTP handle.
46    pub request_duration: std::time::Duration,
47}
48
49impl Default for McpGatewayLimits {
50    fn default() -> Self {
51        Self {
52            max_global_requests: DEFAULT_MAX_GLOBAL_MCP_REQUESTS,
53            max_per_capability_requests: DEFAULT_MAX_PER_CAPABILITY_MCP_REQUESTS,
54            request_duration: DEFAULT_MCP_REQUEST_DURATION,
55        }
56    }
57}
58
59impl McpGatewayLimits {
60    fn validated(self) -> Result<Self, McpInstallError> {
61        if self.max_global_requests == 0
62            || self.max_per_capability_requests == 0
63            || self.request_duration.is_zero()
64        {
65            return Err(McpInstallError::InvalidDescriptor);
66        }
67        Ok(self)
68    }
69}
70
71/// Axum state: routes + gateway-owned capability services (not process-global — §17).
72#[derive(Clone)]
73struct GatewayState {
74    routes: Arc<McpRouteTable>,
75    services: Arc<std::sync::Mutex<std::collections::HashMap<String, Arc<CapabilityHttpService>>>>,
76    request_permits: Arc<tokio::sync::Semaphore>,
77    request_owner: Option<Arc<dyn McpRequestOwner>>,
78    max_per_capability_requests: usize,
79    request_duration: std::time::Duration,
80}
81
82/// Cloneable handle for install/activate/revoke without owning the listener.
83#[derive(Clone)]
84pub struct McpGatewayHandle {
85    routes: Arc<McpRouteTable>,
86    services: Arc<std::sync::Mutex<std::collections::HashMap<String, Arc<CapabilityHttpService>>>>,
87    /// Retained so Clone keeps the same gateway-scoped concurrency budget.
88    #[allow(dead_code)]
89    request_permits: Arc<tokio::sync::Semaphore>,
90    base_url: String,
91    local_addr: SocketAddr,
92}
93
94impl McpGatewayHandle {
95    /// Bound loopback address.
96    pub fn local_addr(&self) -> SocketAddr {
97        self.local_addr
98    }
99
100    /// Base URL `http://127.0.0.1:port` (no path).
101    pub fn base_url(&self) -> &str {
102        &self.base_url
103    }
104
105    /// Shared route table.
106    pub fn routes(&self) -> &Arc<McpRouteTable> {
107        &self.routes
108    }
109
110    /// Install a pending capability for a transaction.
111    pub fn install_pending(
112        &self,
113        transaction_id: TransactionId,
114        tools: ResolvedToolSet,
115        dispatcher: Arc<TransactionToolDispatcher>,
116        exchange_id: ExchangeId,
117    ) -> Result<PendingMcpBinding, McpInstallError> {
118        self.install_pending_with_deadline(
119            transaction_id,
120            tools,
121            dispatcher,
122            exchange_id,
123            std::time::Instant::now() + std::time::Duration::from_secs(365 * 24 * 3600),
124        )
125    }
126
127    /// Install with the live transaction absolute Instant.
128    pub fn install_pending_with_deadline(
129        &self,
130        transaction_id: TransactionId,
131        tools: ResolvedToolSet,
132        dispatcher: Arc<TransactionToolDispatcher>,
133        exchange_id: ExchangeId,
134        transaction_deadline: std::time::Instant,
135    ) -> Result<PendingMcpBinding, McpInstallError> {
136        self.routes.install_pending_with_deadline(
137            transaction_id,
138            tools,
139            dispatcher,
140            exchange_id,
141            &self.base_url,
142            transaction_deadline,
143        )
144    }
145
146    /// Activate a pending capability.
147    pub fn activate(&self, token: &CapabilityToken) -> Result<(), McpInstallError> {
148        self.routes.activate(token)
149    }
150
151    /// Revoke one capability (idempotent).
152    pub fn revoke(&self, token: &CapabilityToken) -> bool {
153        let removed = self.routes.revoke(token);
154        if removed {
155            drop_capability_service(&self.services, &token.to_hex());
156        }
157        removed
158    }
159
160    /// Revoke every route and cancel per-capability services (shutdown / quiesce).
161    pub fn revoke_all_services(&self) {
162        let tokens = self.routes.revoke_all();
163        for hex in tokens {
164            drop_capability_service(&self.services, &hex);
165        }
166    }
167}
168
169/// Listener + router prepared without spawning (TaskSupervisor / RuntimeService).
170pub struct PreparedMcpGateway {
171    handle: McpGatewayHandle,
172    cancel: CancellationToken,
173    listener: tokio::net::TcpListener,
174    app: Router,
175}
176
177impl PreparedMcpGateway {
178    /// Cloneable install/activate handle.
179    pub fn handle(&self) -> McpGatewayHandle {
180        self.handle.clone()
181    }
182
183    /// Cancellation token that stops [`Self::serve`].
184    pub fn cancel_token(&self) -> CancellationToken {
185        self.cancel.clone()
186    }
187
188    /// Bound loopback address.
189    pub fn local_addr(&self) -> SocketAddr {
190        self.handle.local_addr()
191    }
192
193    /// Serve until [`Self::cancel_token`] is cancelled. Revokes routes on exit.
194    pub async fn serve(self) {
195        let cancel_serve = self.cancel.clone();
196        let handle = self.handle;
197        let _ = axum::serve(self.listener, self.app)
198            .with_graceful_shutdown(async move {
199                cancel_serve.cancelled().await;
200            })
201            .await;
202        handle.revoke_all_services();
203    }
204}
205
206/// MCP gateway constructors (no ambient spawn — Law 23).
207///
208/// Production RuntimeOwner uses [`PreparedMcpGateway`] under
209/// `TaskClass::RuntimeService`. Standalone tests prepare + spawn explicitly.
210pub struct McpGateway;
211
212impl McpGateway {
213    /// Build from a pre-bound non-blocking std listener (fail-closed startup bind).
214    pub fn prepare_from_std_listener(
215        std_listener: std::net::TcpListener,
216        max_routes: usize,
217        request_owner: Option<Arc<dyn McpRequestOwner>>,
218    ) -> Result<PreparedMcpGateway, McpInstallError> {
219        Self::prepare_from_std_listener_with_limits(
220            std_listener,
221            max_routes,
222            request_owner,
223            McpGatewayLimits::default(),
224        )
225    }
226
227    /// [`Self::prepare_from_std_listener`] with explicit concurrency/duration limits.
228    pub fn prepare_from_std_listener_with_limits(
229        std_listener: std::net::TcpListener,
230        max_routes: usize,
231        request_owner: Option<Arc<dyn McpRequestOwner>>,
232        limits: McpGatewayLimits,
233    ) -> Result<PreparedMcpGateway, McpInstallError> {
234        let listener = tokio::net::TcpListener::from_std(std_listener)
235            .map_err(|_| McpInstallError::InvalidDescriptor)?;
236        Self::prepare_from_tokio_listener_with_limits(listener, max_routes, request_owner, limits)
237    }
238
239    /// Build from an already-bound Tokio loopback listener (no spawn).
240    pub fn prepare_from_tokio_listener(
241        listener: tokio::net::TcpListener,
242        max_routes: usize,
243        request_owner: Option<Arc<dyn McpRequestOwner>>,
244    ) -> Result<PreparedMcpGateway, McpInstallError> {
245        Self::prepare_from_tokio_listener_with_limits(
246            listener,
247            max_routes,
248            request_owner,
249            McpGatewayLimits::default(),
250        )
251    }
252
253    /// [`Self::prepare_from_tokio_listener`] with explicit concurrency/duration limits.
254    pub fn prepare_from_tokio_listener_with_limits(
255        listener: tokio::net::TcpListener,
256        max_routes: usize,
257        request_owner: Option<Arc<dyn McpRequestOwner>>,
258        limits: McpGatewayLimits,
259    ) -> Result<PreparedMcpGateway, McpInstallError> {
260        let limits = limits.validated()?;
261        let local_addr = listener
262            .local_addr()
263            .map_err(|_| McpInstallError::InvalidDescriptor)?;
264        if !local_addr.ip().is_loopback() {
265            return Err(McpInstallError::InvalidDescriptor);
266        }
267
268        let routes = McpRouteTable::new(max_routes);
269        let services = Arc::new(std::sync::Mutex::new(std::collections::HashMap::new()));
270        let request_permits = Arc::new(tokio::sync::Semaphore::new(limits.max_global_requests));
271        let base_url = format!("http://{}", local_addr);
272        let cancel = CancellationToken::new();
273        let state = GatewayState {
274            routes: Arc::clone(&routes),
275            services: Arc::clone(&services),
276            request_permits: Arc::clone(&request_permits),
277            request_owner,
278            max_per_capability_requests: limits.max_per_capability_requests,
279            request_duration: limits.request_duration,
280        };
281
282        let app = Router::new()
283            .route("/mcp/{token}", any(mcp_dispatch))
284            .route("/mcp/{token}/{*rest}", any(mcp_dispatch_rest))
285            .with_state(state);
286
287        Ok(PreparedMcpGateway {
288            handle: McpGatewayHandle {
289                routes,
290                services,
291                request_permits,
292                base_url,
293                local_addr,
294            },
295            cancel,
296            listener,
297            app,
298        })
299    }
300}
301
302async fn mcp_dispatch(
303    State(state): State<GatewayState>,
304    Path(token): Path<String>,
305    req: Request,
306) -> Response<Body> {
307    forward_mcp(state, &token, req).await
308}
309
310async fn mcp_dispatch_rest(
311    State(state): State<GatewayState>,
312    Path((token, _rest)): Path<(String, String)>,
313    req: Request,
314) -> Response<Body> {
315    forward_mcp(state, &token, req).await
316}
317
318/// Per-capability Streamable HTTP service (shared across requests for one token).
319struct CapabilityHttpService {
320    service: StreamableHttpService<TransactionMcpHandler, LocalSessionManager>,
321    cancel: CancellationToken,
322    /// Per-capability concurrent request bound (D-034).
323    permits: Arc<tokio::sync::Semaphore>,
324}
325
326const DEFAULT_MAX_GLOBAL_MCP_REQUESTS: usize = 64;
327const DEFAULT_MAX_PER_CAPABILITY_MCP_REQUESTS: usize = 8;
328const DEFAULT_MCP_REQUEST_DURATION: std::time::Duration = std::time::Duration::from_secs(30);
329
330/// Drop and cancel the per-token Streamable HTTP service for this gateway (D-018).
331fn drop_capability_service(
332    services: &std::sync::Mutex<std::collections::HashMap<String, Arc<CapabilityHttpService>>>,
333    token_hex: &str,
334) {
335    let key = CapabilityToken::from_hex(token_hex)
336        .map(|t| t.to_hex())
337        .unwrap_or_else(|| token_hex.to_ascii_lowercase());
338    if let Ok(mut map) = services.lock() {
339        if let Some(svc) = map.remove(&key) {
340            svc.cancel.cancel();
341        }
342    }
343}
344
345async fn forward_mcp(state: GatewayState, token_hex: &str, req: Request) -> Response<Body> {
346    // Cheap fail-closed route lookup only — concurrency budget + body buffering
347    // run inside the owned McpRequest task so permits cannot outlive the handler
348    // if the axum task is dropped (Law 22 / §17).
349    let Some(canonical) = CapabilityToken::from_hex(token_hex).map(|t| t.to_hex()) else {
350        return Response::builder()
351            .status(StatusCode::NOT_FOUND)
352            .body(Body::from("unknown capability"))
353            .unwrap_or_else(|_| Response::new(Body::empty()));
354    };
355    let Some(binding) = state.routes.get_by_hex(&canonical) else {
356        drop_capability_service(&state.services, &canonical);
357        return Response::builder()
358            .status(StatusCode::NOT_FOUND)
359            .body(Body::from("unknown capability"))
360            .unwrap_or_else(|_| Response::new(Body::empty()));
361    };
362
363    let transaction_id = binding.transaction_id;
364    let state_work = state.clone();
365    let work = async move { execute_mcp_request(state_work, canonical, binding, req).await };
366
367    // RuntimeOwner path: each active request is TaskClass::McpRequest (§17).
368    if let Some(owner) = state.request_owner.as_ref() {
369        owner.run_owned(transaction_id, Box::pin(work)).await
370    } else {
371        work.await
372    }
373}
374
375/// Permit acquire + body buffer + Streamable HTTP handle (owned-task body).
376async fn execute_mcp_request(
377    state: GatewayState,
378    canonical: String,
379    binding: Arc<super::binding::McpBinding>,
380    req: Request,
381) -> Response<Body> {
382    // Re-check route after spawn (may have been revoked).
383    if state.routes.get_by_hex(&canonical).is_none() {
384        drop_capability_service(&state.services, &canonical);
385        return Response::builder()
386            .status(StatusCode::NOT_FOUND)
387            .body(Body::from("unknown capability"))
388            .unwrap_or_else(|_| Response::new(Body::empty()));
389    }
390
391    let Ok(_global) = state.request_permits.clone().try_acquire_owned() else {
392        return Response::builder()
393            .status(StatusCode::TOO_MANY_REQUESTS)
394            .body(Body::from("mcp gateway concurrency exceeded"))
395            .unwrap_or_else(|_| Response::new(Body::empty()));
396    };
397
398    let max_per_capability = state.max_per_capability_requests;
399    let service = {
400        let mut map = state.services.lock().unwrap_or_else(|e| e.into_inner());
401        map.entry(canonical.clone())
402            .or_insert_with(|| {
403                let handler = binding.handler.clone();
404                let cancel = CancellationToken::new();
405                let mut config = StreamableHttpServerConfig::default();
406                config.cancellation_token = cancel.clone();
407                config.sse_keep_alive = None;
408                config.sse_retry = None;
409                config.json_response = true;
410                Arc::new(CapabilityHttpService {
411                    service: StreamableHttpService::new(
412                        move || Ok(handler.clone()),
413                        Arc::new(LocalSessionManager::default()),
414                        config,
415                    ),
416                    cancel,
417                    permits: Arc::new(tokio::sync::Semaphore::new(max_per_capability)),
418                })
419            })
420            .clone()
421    };
422
423    let Ok(_local) = service.permits.clone().try_acquire_owned() else {
424        return Response::builder()
425            .status(StatusCode::TOO_MANY_REQUESTS)
426            .body(Body::from("mcp capability concurrency exceeded"))
427            .unwrap_or_else(|_| Response::new(Body::empty()));
428    };
429
430    let deadline_at = tokio::time::Instant::now() + state.request_duration;
431    let (parts, body) = req.into_parts();
432    let body_budget = deadline_at.saturating_duration_since(tokio::time::Instant::now());
433    let collected =
434        match tokio::time::timeout(body_budget, axum::body::to_bytes(body, 1024 * 1024)).await {
435            Ok(Ok(b)) => b,
436            Ok(Err(_)) => {
437                return Response::builder()
438                    .status(StatusCode::PAYLOAD_TOO_LARGE)
439                    .body(Body::from("request body exceeds bound"))
440                    .unwrap_or_else(|_| Response::new(Body::empty()));
441            }
442            Err(_) => {
443                return Response::builder()
444                    .status(StatusCode::GATEWAY_TIMEOUT)
445                    .body(Body::from("mcp request deadline exceeded"))
446                    .unwrap_or_else(|_| Response::new(Body::empty()));
447            }
448        };
449    let req = rewrite_path(
450        Request::from_parts(parts, Body::from(collected)),
451        &canonical,
452    );
453    let handle_budget = deadline_at.saturating_duration_since(tokio::time::Instant::now());
454    if handle_budget.is_zero() {
455        return Response::builder()
456            .status(StatusCode::GATEWAY_TIMEOUT)
457            .body(Body::from("mcp request deadline exceeded"))
458            .unwrap_or_else(|_| Response::new(Body::empty()));
459    }
460    match tokio::time::timeout(handle_budget, service.service.handle(req)).await {
461        Ok(response) => response.map(Body::new),
462        Err(_) => Response::builder()
463            .status(StatusCode::GATEWAY_TIMEOUT)
464            .body(Body::from("mcp request deadline exceeded"))
465            .unwrap_or_else(|_| Response::new(Body::empty())),
466    }
467}
468
469fn rewrite_path(req: Request, token_hex: &str) -> Request {
470    let (mut parts, body) = req.into_parts();
471    let path = parts.uri.path().to_string();
472    let query = parts.uri.query().map(|q| q.to_string());
473    let prefix = format!("/mcp/{token_hex}");
474    let new_path = if let Some(rest) = path.strip_prefix(&prefix) {
475        if rest.is_empty() {
476            "/".to_string()
477        } else {
478            rest.to_string()
479        }
480    } else {
481        path
482    };
483    let pq = match query {
484        Some(q) => format!("{new_path}?{q}"),
485        None => new_path,
486    };
487    if let Ok(uri) = pq.parse::<axum::http::Uri>() {
488        parts.uri = uri;
489    }
490    Request::from_parts(parts, body)
491}
monoloop_loop/transaction/mcp/gateway.rs

monoloop_loop/transaction/mcp/
gateway.rs
