vtcode_webmcp/
server.rs

1use crate::error::{Result, WebmcpError};
2use crate::event_hub::{EventHubConfig, EventHubSubscription, MAX_EVENT_BYTES, WebmcpEventHub};
3use crate::pairing::{PairingDisplay, PairingManager};
4use crate::protocol::{
5    BridgeEventMessage, BridgeRequest, BridgeResponse, BridgeSettings, PROTOCOL_VERSION, PairPayload, StatusPayload,
6    is_valid_request_id, response_request_id,
7};
8use crate::remote_mcp::{RemoteMcpEndpoint, RemoteMcpServerConfig};
9use crate::runtime::RuntimeAdapter;
10use axum::Router;
11use axum::extract::{State, WebSocketUpgrade, ws};
12use axum::http::{HeaderMap, StatusCode, header::ORIGIN};
13use axum::response::{IntoResponse, Response};
14use futures::{StreamExt, future::BoxFuture};
15use serde_json::Value;
16use std::net::IpAddr;
17use std::sync::Arc;
18use std::time::{Duration, Instant};
19use tokio::net::TcpListener;
20use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot};
21use vtcode_commons::validation::is_valid_origin;
22
23const MAX_PAIRED_CONNECTIONS: usize = 64;
24const MAX_CONNECTIONS: usize = 128;
25const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30);
26const EVENT_ENVELOPE_OVERHEAD: usize = 256;
27const MIN_FRAME_BYTES: usize = 1024;
28const WRITE_TIMEOUT: Duration = Duration::from_secs(10);
29const MAX_MUTATION_QUEUE: usize = 64;
30const SESSION_LEASE_CHECK_INTERVAL: Duration = Duration::from_millis(250);
31
32/// Configuration for the WebMCP listener.
33#[derive(Debug, Clone)]
34pub struct WebmcpServerConfig {
35    /// Bind host, normally `127.0.0.1`.
36    pub host: String,
37    /// Bind port, with zero selecting an available port.
38    pub port: u16,
39    /// Explicit browser origin allowlist.
40    pub allowed_origins: Vec<String>,
41    /// Pairing lifetime and authenticated-session inactivity lease in seconds.
42    pub pairing_ttl_secs: u64,
43    /// Maximum WebSocket message size.
44    pub max_frame_bytes: usize,
45    /// Maximum concurrent adapter operations.
46    pub max_in_flight_requests: usize,
47    /// Whether remote reverse-proxy mode is explicitly enabled.
48    pub allow_remote: bool,
49    /// Public WSS URL required for remote mode.
50    pub public_url: Option<String>,
51    /// Optional read-only OpenAI-compatible MCP transport.
52    pub remote_mcp: Option<RemoteMcpServerConfig>,
53    /// Event replay and subscriber queue limits.
54    pub event_hub: EventHubConfig,
55    /// Per-operation timeout.
56    pub request_timeout: Duration,
57}
58
59impl Default for WebmcpServerConfig {
60    fn default() -> Self {
61        Self {
62            host: "127.0.0.1".to_string(),
63            port: 0,
64            allowed_origins: Vec::new(),
65            pairing_ttl_secs: 300,
66            max_frame_bytes: 1024 * 1024,
67            max_in_flight_requests: 8,
68            allow_remote: false,
69            public_url: None,
70            remote_mcp: None,
71            event_hub: EventHubConfig::default(),
72            request_timeout: Duration::from_secs(30),
73        }
74    }
75}
76
77struct ServerState {
78    adapter: Arc<dyn RuntimeAdapter>,
79    pairing: PairingManager,
80    event_hub: WebmcpEventHub,
81    dispatch: Arc<DispatchState>,
82    mutation_supervisor: MutationSupervisor,
83    paired_connections: Arc<Semaphore>,
84    connections: Arc<Semaphore>,
85    max_frame_bytes: usize,
86    request_timeout: Duration,
87    remote_mcp: Option<Arc<RemoteMcpEndpoint>>,
88}
89
90struct DispatchState {
91    adapter: Arc<dyn RuntimeAdapter>,
92    pairing: PairingManager,
93    event_hub: WebmcpEventHub,
94    settings: BridgeSettings,
95    in_flight: Arc<Semaphore>,
96    request_timeout: Duration,
97}
98
99struct MutationJob {
100    dispatch: Arc<DispatchState>,
101    origin: String,
102    token: String,
103    request: BridgeRequest,
104    result: oneshot::Sender<Result<Value>>,
105}
106
107#[derive(Clone, Default)]
108struct MutationSupervisor {
109    sender: Arc<tokio::sync::Mutex<Option<mpsc::Sender<MutationJob>>>>,
110}
111
112impl MutationSupervisor {
113    async fn submit(
114        &self,
115        dispatch: Arc<DispatchState>,
116        origin: String,
117        token: String,
118        request: BridgeRequest,
119    ) -> Result<Value> {
120        let sender = {
121            let mut sender_slot = self.sender.lock().await;
122            if let Some(sender) = sender_slot.as_ref() {
123                sender.clone()
124            } else {
125                let (sender, receiver) = mpsc::channel(MAX_MUTATION_QUEUE);
126                drop(tokio::spawn(run_mutation_supervisor(receiver)));
127                *sender_slot = Some(sender.clone());
128                sender
129            }
130        };
131        let (result, receiver) = oneshot::channel();
132        sender
133            .try_send(MutationJob { dispatch, origin, token, request, result })
134            .map_err(|error| match error {
135                mpsc::error::TrySendError::Full(_) => WebmcpError::LimitExceeded,
136                mpsc::error::TrySendError::Closed(_) => {
137                    WebmcpError::Adapter("WebMCP mutation supervisor is closed".to_string())
138                }
139            })?;
140        receiver
141            .await
142            .map_err(|_error| WebmcpError::Adapter("WebMCP mutation supervisor stopped".to_string()))?
143    }
144}
145
146async fn run_mutation_supervisor(mut receiver: mpsc::Receiver<MutationJob>) {
147    while let Some(job) = receiver.recv().await {
148        let result = dispatch_request(job.dispatch, job.origin, job.token, job.request).await;
149        drop(job.result.send(result));
150    }
151}
152
153/// Authenticated WebSocket WebMCP server.
154#[derive(Clone)]
155pub struct WebmcpServer {
156    state: Arc<ServerState>,
157    config: WebmcpServerConfig,
158}
159
160impl std::fmt::Debug for WebmcpServer {
161    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
162        formatter
163            .debug_struct("WebmcpServer")
164            .field("host", &self.config.host)
165            .field("port", &self.config.port)
166            .field("allowed_origins", &self.config.allowed_origins)
167            .finish_non_exhaustive()
168    }
169}
170
171impl WebmcpServer {
172    /// Create a server around an active or headless runtime adapter.
173    pub fn new(adapter: Arc<dyn RuntimeAdapter>, config: WebmcpServerConfig) -> Result<Self> {
174        validate_config(&config)?;
175        let pairing = PairingManager::new(&config.allowed_origins, Duration::from_secs(config.pairing_ttl_secs))?;
176        let remote_mcp = config
177            .remote_mcp
178            .clone()
179            .map(|remote_config| RemoteMcpEndpoint::new(Arc::clone(&adapter), remote_config).map(Arc::new))
180            .transpose()?;
181        let event_limit = config
182            .max_frame_bytes
183            .saturating_sub(EVENT_ENVELOPE_OVERHEAD)
184            .min(MAX_EVENT_BYTES);
185        let event_hub = WebmcpEventHub::new_with_max_event_bytes(config.event_hub, event_limit)?;
186        let in_flight = Arc::new(Semaphore::new(config.max_in_flight_requests));
187        let settings = BridgeSettings {
188            host: config.host.clone(),
189            port: config.port,
190            pairing_ttl_secs: config.pairing_ttl_secs,
191            max_frame_bytes: config.max_frame_bytes,
192            max_in_flight_requests: config.max_in_flight_requests,
193            remote_enabled: config.allow_remote,
194        };
195        let dispatch = Arc::new(DispatchState {
196            adapter: Arc::clone(&adapter),
197            pairing: pairing.clone(),
198            event_hub: event_hub.clone(),
199            settings,
200            in_flight: Arc::clone(&in_flight),
201            request_timeout: config.request_timeout,
202        });
203        Ok(Self {
204            state: Arc::new(ServerState {
205                adapter,
206                pairing,
207                event_hub,
208                dispatch,
209                mutation_supervisor: MutationSupervisor::default(),
210                paired_connections: Arc::new(Semaphore::new(MAX_PAIRED_CONNECTIONS)),
211                connections: Arc::new(Semaphore::new(MAX_CONNECTIONS)),
212                max_frame_bytes: config.max_frame_bytes,
213                request_timeout: config.request_timeout,
214                remote_mcp,
215            }),
216            config,
217        })
218    }
219
220    /// Start a new one-time terminal pairing code.
221    pub fn begin_pairing(&self) -> PairingDisplay {
222        self.state.pairing.begin_pairing()
223    }
224
225    /// Start a one-time code bound to one exact allowed browser origin.
226    pub fn begin_pairing_for_origin(&self, origin: impl Into<String>) -> Result<PairingDisplay> {
227        self.state.pairing.begin_pairing_for_origin(origin)
228    }
229
230    /// Revoke browser sessions and issue a fresh code for one exact origin.
231    pub fn replace_pairing_for_origin(&self, origin: impl Into<String>) -> Result<PairingDisplay> {
232        self.state.pairing.replace_pairing_for_origin(origin)
233    }
234
235    /// Revoke all browser sessions and pending pairing codes.
236    pub fn revoke_all_pairings(&self) {
237        self.state.pairing.revoke_all();
238    }
239
240    /// Access the canonical runtime event hub used by this server.
241    pub fn event_hub(&self) -> WebmcpEventHub {
242        self.state.event_hub.clone()
243    }
244
245    /// Build the WebSocket router. No listener is opened by this method.
246    pub fn router(&self) -> Router {
247        let mut router = Router::new().route("/webmcp", axum::routing::get(websocket_handler));
248        if let Some(remote_mcp) = self.state.remote_mcp.as_ref() {
249            let remote_routes = remote_mcp
250                .routes()
251                .with_state(Arc::clone(remote_mcp))
252                .with_state(self.state.clone());
253            router = router.merge(remote_routes);
254        }
255        router.with_state(self.state.clone())
256    }
257
258    /// Bind the configured address. Call [`Self::serve_listener`] afterwards.
259    pub async fn bind(&self) -> Result<TcpListener> {
260        TcpListener::bind((self.config.host.as_str(), self.config.port))
261            .await
262            .map_err(WebmcpError::Io)
263    }
264
265    /// Serve on a caller-provided listener until it fails or is stopped.
266    pub async fn serve_listener(&self, listener: TcpListener) -> Result<()> {
267        axum::serve(listener, self.router())
268            .await
269            .map_err(|error| WebmcpError::Adapter(format!("WebMCP listener failed: {error}")))
270    }
271
272    /// Bind and serve the configured listener.
273    pub async fn serve(&self) -> Result<()> {
274        let listener = self.bind().await?;
275        self.serve_listener(listener).await
276    }
277}
278
279async fn websocket_handler(
280    State(state): State<Arc<ServerState>>,
281    headers: HeaderMap,
282    upgrade: WebSocketUpgrade,
283) -> Response {
284    let Some(origin) = origin_from_headers(&headers) else {
285        return (StatusCode::FORBIDDEN, "WebMCP requires an Origin header").into_response();
286    };
287    if !state.pairing.is_origin_allowed(origin) {
288        return (StatusCode::FORBIDDEN, "WebMCP origin is not allowed").into_response();
289    }
290    let connection_permit = match state.connections.clone().try_acquire_owned() {
291        Ok(permit) => permit,
292        Err(_error) => return (StatusCode::TOO_MANY_REQUESTS, "WebMCP connection limit reached").into_response(),
293    };
294    let origin = origin.to_string();
295    upgrade
296        .max_message_size(state.max_frame_bytes)
297        .max_frame_size(state.max_frame_bytes)
298        .on_upgrade(move |socket| run_socket(socket, state, origin, connection_permit))
299        .into_response()
300}
301
302fn origin_from_headers(headers: &HeaderMap) -> Option<&str> {
303    headers
304        .get(ORIGIN)
305        .and_then(|value| value.to_str().ok())
306        .filter(|origin| !origin.is_empty())
307}
308
309async fn run_socket(
310    mut socket: ws::WebSocket,
311    state: Arc<ServerState>,
312    origin: String,
313    _connection_permit: OwnedSemaphorePermit,
314) {
315    let handshake_deadline = Instant::now() + HANDSHAKE_TIMEOUT;
316    loop {
317        let remaining = handshake_deadline.saturating_duration_since(Instant::now());
318        if remaining.is_zero() {
319            return;
320        }
321        let message = match tokio::time::timeout(remaining, socket.next()).await {
322            Ok(Some(message)) => message,
323            Ok(None) | Err(_) => return,
324        };
325        let Ok(message) = message else { return };
326        match message {
327            ws::Message::Text(text) => {
328                if text.len() > state.max_frame_bytes {
329                    drop(
330                        send_response(
331                            &mut socket,
332                            BridgeResponse::failure("unknown", "frame_too_large", "request exceeds the frame limit"),
333                            state.max_frame_bytes,
334                        )
335                        .await,
336                    );
337                    return;
338                }
339                match serde_json::from_slice::<BridgeRequest>(text.as_bytes()) {
340                    Ok(BridgeRequest::Pair {
341                        request_id,
342                        code,
343                        resume_token,
344                        origin: claimed_origin,
345                        after_sequence,
346                    }) => {
347                        if !is_valid_request_id(&request_id) {
348                            drop(
349                                send_response(
350                                    &mut socket,
351                                    invalid_request_id_response(&request_id),
352                                    state.max_frame_bytes,
353                                )
354                                .await,
355                            );
356                            continue;
357                        }
358                        if claimed_origin.as_deref().is_some_and(|claimed| claimed != origin) {
359                            drop(
360                                send_response(
361                                    &mut socket,
362                                    BridgeResponse::failure(
363                                        &request_id,
364                                        "origin_mismatch",
365                                        "request origin does not match the WebSocket origin",
366                                    ),
367                                    state.max_frame_bytes,
368                                )
369                                .await,
370                            );
371                            return;
372                        }
373                        let mut subscription = match state.event_hub.subscribe(after_sequence) {
374                            Ok(subscription) => subscription,
375                            Err(error) => {
376                                drop(
377                                    send_response(
378                                        &mut socket,
379                                        response_for_error(&request_id, error),
380                                        state.max_frame_bytes,
381                                    )
382                                    .await,
383                                );
384                                return;
385                            }
386                        };
387                        let connection_permit = match state.paired_connections.clone().try_acquire_owned() {
388                            Ok(permit) => permit,
389                            Err(_error) => {
390                                drop(
391                                    send_response(
392                                        &mut socket,
393                                        BridgeResponse::failure(
394                                            &request_id,
395                                            "connection_limit",
396                                            "the WebMCP server has reached its paired connection limit",
397                                        ),
398                                        state.max_frame_bytes,
399                                    )
400                                    .await,
401                                );
402                                return;
403                            }
404                        };
405                        let session = match resume_token {
406                            Some(token) => state.pairing.resume(&token, &origin),
407                            None => state.pairing.pair(&code, &origin),
408                        };
409                        match session {
410                            Ok(session) => {
411                                let response = BridgeResponse::success(
412                                    request_id,
413                                    PairPayload {
414                                        token: session.token().to_string(),
415                                        protocol_version: PROTOCOL_VERSION,
416                                        expires_in_secs: session.expires_in().as_secs().max(1),
417                                    },
418                                );
419                                if send_response(&mut socket, response, state.max_frame_bytes).await.is_err() {
420                                    return;
421                                }
422                                for event in subscription.replay() {
423                                    if send_event(&mut socket, event.sequence, &event.event, state.max_frame_bytes)
424                                        .await
425                                        .is_err()
426                                    {
427                                        return;
428                                    }
429                                }
430                                run_paired_socket(
431                                    socket,
432                                    state,
433                                    origin,
434                                    session.token().to_string(),
435                                    &mut subscription,
436                                    connection_permit,
437                                    _connection_permit,
438                                )
439                                .await;
440                                return;
441                            }
442                            Err(error) => {
443                                drop(
444                                    send_response(
445                                        &mut socket,
446                                        response_for_error(&request_id, error),
447                                        state.max_frame_bytes,
448                                    )
449                                    .await,
450                                );
451                            }
452                        }
453                    }
454                    Ok(request) => {
455                        let request_id = request.request_id().to_string();
456                        let response = if is_valid_request_id(&request_id) {
457                            response_for_error(&request_id, WebmcpError::Unauthorized)
458                        } else {
459                            invalid_request_id_response(&request_id)
460                        };
461                        drop(send_response(&mut socket, response, state.max_frame_bytes).await);
462                    }
463                    Err(error) => {
464                        drop(
465                            send_response(
466                                &mut socket,
467                                BridgeResponse::failure("unknown", "malformed_request", error.to_string()),
468                                state.max_frame_bytes,
469                            )
470                            .await,
471                        );
472                    }
473                }
474            }
475            ws::Message::Binary(_) => {
476                drop(
477                    send_response(
478                        &mut socket,
479                        BridgeResponse::failure(
480                            "unknown",
481                            "binary_not_supported",
482                            "WebMCP accepts JSON text frames only",
483                        ),
484                        state.max_frame_bytes,
485                    )
486                    .await,
487                );
488            }
489            ws::Message::Ping(payload) => {
490                if send_pong(&mut socket, payload).await.is_err() {
491                    return;
492                }
493            }
494            ws::Message::Close(_) => return,
495            ws::Message::Pong(_) => {}
496        }
497    }
498}
499
500enum PairedFrameAction {
501    Request(BridgeRequest),
502    Continue,
503    Close,
504}
505
506async fn handle_paired_frame(
507    socket: &mut ws::WebSocket,
508    message: ws::Message,
509    max_frame_bytes: usize,
510) -> PairedFrameAction {
511    match message {
512        ws::Message::Text(text) => {
513            if text.len() > max_frame_bytes {
514                drop(
515                    send_response(
516                        socket,
517                        BridgeResponse::failure("unknown", "frame_too_large", "request exceeds the frame limit"),
518                        max_frame_bytes,
519                    )
520                    .await,
521                );
522                return PairedFrameAction::Close;
523            }
524            let request = match serde_json::from_slice::<BridgeRequest>(text.as_bytes()) {
525                Ok(request) => request,
526                Err(error) => {
527                    if send_response(
528                        socket,
529                        BridgeResponse::failure("unknown", "malformed_request", error.to_string()),
530                        max_frame_bytes,
531                    )
532                    .await
533                    .is_err()
534                    {
535                        return PairedFrameAction::Close;
536                    }
537                    return PairedFrameAction::Continue;
538                }
539            };
540            if !is_valid_request_id(request.request_id()) {
541                if send_response(socket, invalid_request_id_response(request.request_id()), max_frame_bytes)
542                    .await
543                    .is_err()
544                {
545                    return PairedFrameAction::Close;
546                }
547                return PairedFrameAction::Continue;
548            }
549            PairedFrameAction::Request(request)
550        }
551        ws::Message::Binary(_) => {
552            if send_response(
553                socket,
554                BridgeResponse::failure("unknown", "binary_not_supported", "WebMCP accepts JSON text frames only"),
555                max_frame_bytes,
556            )
557            .await
558            .is_err()
559            {
560                PairedFrameAction::Close
561            } else {
562                PairedFrameAction::Continue
563            }
564        }
565        ws::Message::Ping(payload) => {
566            if send_pong(socket, payload).await.is_err() {
567                PairedFrameAction::Close
568            } else {
569                PairedFrameAction::Continue
570            }
571        }
572        ws::Message::Close(_) => PairedFrameAction::Close,
573        ws::Message::Pong(_) => PairedFrameAction::Continue,
574    }
575}
576
577async fn run_paired_socket(
578    mut socket: ws::WebSocket,
579    state: Arc<ServerState>,
580    origin: String,
581    token: String,
582    subscription: &mut EventHubSubscription,
583    _connection_permit: OwnedSemaphorePermit,
584    _all_connections_permit: OwnedSemaphorePermit,
585) {
586    let mut expiry_check = tokio::time::interval(SESSION_LEASE_CHECK_INTERVAL);
587    expiry_check.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
588    loop {
589        tokio::select! {
590            message = socket.next() => {
591                let Some(Ok(message)) = message else { return };
592                match handle_paired_frame(&mut socket, message, state.max_frame_bytes).await {
593                    PairedFrameAction::Close => return,
594                    PairedFrameAction::Continue => {}
595                    PairedFrameAction::Request(request) => {
596                        let request_id = request.request_id().to_string();
597                        let mutation_request = matches!(
598                            &request,
599                            BridgeRequest::ApplyProposal { .. } | BridgeRequest::RevertLastChange { .. }
600                        );
601                        let operation: BoxFuture<'static, Result<Value>> = if mutation_request {
602                            let supervisor = state.mutation_supervisor.clone();
603                            let dispatch = Arc::clone(&state.dispatch);
604                            let origin = origin.clone();
605                            let token = token.clone();
606                            Box::pin(async move { supervisor.submit(dispatch, origin, token, request).await })
607                        } else {
608                            Box::pin(dispatch_request(
609                                Arc::clone(&state.dispatch),
610                                origin.clone(),
611                                token.clone(),
612                                request,
613                            ))
614                        };
615                        let mut operation = operation;
616                        loop {
617                            tokio::select! {
618                                result = &mut operation => {
619                                    let response = match result {
620                                        Ok(payload) => BridgeResponse::success(request_id.clone(), payload),
621                                        Err(error) => response_for_error(&request_id, error),
622                                    };
623                                    if send_response(&mut socket, response, state.max_frame_bytes).await.is_err() { return; }
624                                    break;
625                                }
626                                message = socket.next() => {
627                                    let Some(Ok(message)) = message else { return };
628                                    match handle_paired_frame(&mut socket, message, state.max_frame_bytes).await {
629                                        PairedFrameAction::Close => return,
630                                        PairedFrameAction::Continue => {}
631                                        PairedFrameAction::Request(request) => {
632                                            let request_id = request.request_id().to_string();
633                                            let response = if request.token() != Some(token.as_str()) {
634                                                response_for_error(&request_id, WebmcpError::Unauthorized)
635                                            } else if matches!(&request, BridgeRequest::Cancel { .. }) {
636                                                match dispatch_cancel_request(&state, &origin, &token, request).await {
637                                                    Ok(payload) => BridgeResponse::success(request_id.clone(), payload),
638                                                    Err(error) => response_for_error(&request_id, error),
639                                                }
640                                            } else {
641                                                if state.pairing.refresh(&token, &origin).is_err() { return; }
642                                                BridgeResponse::failure(&request_id, "request_in_progress", "wait for the active request or cancel it")
643                                            };
644                                            if send_response(&mut socket, response, state.max_frame_bytes).await.is_err() { return; }
645                                        }
646                                    }
647                                }
648                                event = subscription.recv() => {
649                                    let Some(event) = event else {
650                                        drop(send_response(&mut socket, BridgeResponse::failure("event", "slow_client", "client could not keep up with runtime events"), state.max_frame_bytes).await);
651                                        return;
652                                    };
653                                    if state.pairing.validate(&token, &origin).is_err() { return; }
654                                    if send_event(&mut socket, event.sequence, &event.event, state.max_frame_bytes).await.is_err() { return; }
655                                }
656                                _ = expiry_check.tick() => {
657                                    // Keep the lease pinned while the authenticated operation is
658                                    // still executing. Mutations intentionally remain pending
659                                    // until the supervisor knows their final result.
660                                    if state.pairing.refresh(&token, &origin).is_err() { return; }
661                                }
662                            }
663                        }
664                    }
665                }
666            }
667            event = subscription.recv() => {
668                let Some(event) = event else {
669                    drop(send_response(&mut socket, BridgeResponse::failure("event", "slow_client", "client could not keep up with runtime events"), state.max_frame_bytes).await);
670                    return;
671                };
672                if state.pairing.validate(&token, &origin).is_err() { return; }
673                if send_event(&mut socket, event.sequence, &event.event, state.max_frame_bytes).await.is_err() { return; }
674            }
675            _ = expiry_check.tick() => {
676                if state.pairing.validate(&token, &origin).is_err() { return; }
677            }
678        }
679    }
680}
681
682async fn dispatch_request(
683    dispatch: Arc<DispatchState>,
684    origin: String,
685    token: String,
686    request: BridgeRequest,
687) -> Result<Value> {
688    if request.token() != Some(token.as_str()) {
689        return Err(WebmcpError::Unauthorized);
690    }
691    dispatch.pairing.refresh(&token, &origin)?;
692    let _permit = tokio::time::timeout(dispatch.request_timeout, dispatch.in_flight.clone().acquire_owned())
693        .await
694        .map_err(|_error| WebmcpError::Timeout(dispatch.request_timeout))?
695        .map_err(|_error| WebmcpError::Adapter("WebMCP request capacity is closed".to_string()))?;
696    // A request may have waited for capacity long enough for the session to
697    // expire or be revoked. Re-check immediately before handing it to the
698    // runtime adapter.
699    dispatch.pairing.refresh(&token, &origin)?;
700    let mutation_request =
701        matches!(&request, BridgeRequest::ApplyProposal { .. } | BridgeRequest::RevertLastChange { .. });
702    let operation = async {
703        match request {
704            BridgeRequest::Pair { .. } => Err(WebmcpError::Unauthorized),
705            BridgeRequest::ExplanationGet { scope, offset, .. } => {
706                serde_json::to_value(dispatch.adapter.explanation_get(scope, offset).await?).map_err(WebmcpError::from)
707            }
708            BridgeRequest::ExplanationEvidence { reference, offset, .. } => {
709                serde_json::to_value(dispatch.adapter.explanation_evidence(reference, offset).await?)
710                    .map_err(WebmcpError::from)
711            }
712            BridgeRequest::ExplanationNavigate { reference, .. } => {
713                Ok(serde_json::json!({"focused": dispatch.adapter.explanation_navigate(reference).await?}))
714            }
715            BridgeRequest::Status { .. } => {
716                let runtime = dispatch.adapter.status().await?;
717                serde_json::to_value(StatusPayload {
718                    protocol_version: PROTOCOL_VERSION,
719                    connected: runtime.connected,
720                    runtime,
721                    authenticated_origin: origin,
722                    settings: dispatch.settings.clone(),
723                    latest_sequence: dispatch.event_hub.latest_sequence(),
724                })
725                .map_err(WebmcpError::Json)
726            }
727            BridgeRequest::ListFiles { .. } => {
728                serde_json::to_value(dispatch.adapter.list_files().await?).map_err(WebmcpError::Json)
729            }
730            BridgeRequest::ReadFile { path, .. } => {
731                serde_json::to_value(dispatch.adapter.read_file(&path).await?).map_err(WebmcpError::Json)
732            }
733            BridgeRequest::ProposeChanges { changes, .. } => {
734                serde_json::to_value(dispatch.adapter.propose_changes(changes).await?).map_err(WebmcpError::Json)
735            }
736            BridgeRequest::ApplyProposal { proposal_id, .. } => {
737                serde_json::to_value(dispatch.adapter.apply_proposal(&proposal_id).await?).map_err(WebmcpError::Json)
738            }
739            BridgeRequest::RunChecks { command, .. } => {
740                serde_json::to_value(dispatch.adapter.run_checks(&command).await?).map_err(WebmcpError::Json)
741            }
742            BridgeRequest::RevertLastChange { change_id, .. } => {
743                serde_json::to_value(dispatch.adapter.revert_last_change(&change_id).await?).map_err(WebmcpError::Json)
744            }
745            BridgeRequest::RequestTurn { prompt, proposal_id, .. } => {
746                serde_json::to_value(dispatch.adapter.request_turn(&prompt, proposal_id.as_deref()).await?)
747                    .map_err(WebmcpError::Json)
748            }
749            BridgeRequest::Cancel { target_id, .. } => {
750                let accepted = dispatch.adapter.cancel(&target_id).await?;
751                Ok(serde_json::json!({ "cancelled": target_id, "accepted": accepted }))
752            }
753        }
754    };
755    if mutation_request {
756        // Filesystem adapters keep multi-file mutations transactional after a
757        // transport disconnect. Keep the response pending until that result is
758        // known instead of returning an ambiguous timeout to the browser.
759        operation.await
760    } else {
761        tokio::time::timeout(dispatch.request_timeout, operation)
762            .await
763            .map_err(|_error| WebmcpError::Timeout(dispatch.request_timeout))?
764    }
765}
766
767async fn dispatch_cancel_request(
768    state: &Arc<ServerState>,
769    origin: &str,
770    token: &str,
771    request: BridgeRequest,
772) -> Result<Value> {
773    if request.token() != Some(token) {
774        return Err(WebmcpError::Unauthorized);
775    }
776    state.pairing.refresh(token, origin)?;
777    let BridgeRequest::Cancel { target_id, .. } = request else {
778        return Err(WebmcpError::InvalidRequest(
779            "only cancellation requests are accepted while a request is running".to_string(),
780        ));
781    };
782    let accepted = tokio::time::timeout(state.request_timeout, state.adapter.cancel(&target_id))
783        .await
784        .map_err(|_error| WebmcpError::Timeout(state.request_timeout))??;
785    Ok(serde_json::json!({ "cancelled": target_id, "accepted": accepted }))
786}
787
788async fn send_response(socket: &mut ws::WebSocket, response: BridgeResponse, max_frame_bytes: usize) -> Result<()> {
789    let serialized = serde_json::to_string(&response)?;
790    let serialized = if serialized.len() > max_frame_bytes {
791        serde_json::to_string(&BridgeResponse::failure(
792            response.request_id,
793            "limit_exceeded",
794            "WebMCP response exceeds the configured frame limit",
795        ))?
796    } else {
797        serialized
798    };
799    if serialized.len() > max_frame_bytes {
800        return Err(WebmcpError::LimitExceeded);
801    }
802    tokio::time::timeout(WRITE_TIMEOUT, socket.send(ws::Message::Text(serialized.into())))
803        .await
804        .map_err(|_error| WebmcpError::Timeout(WRITE_TIMEOUT))?
805        .map_err(|error| WebmcpError::Adapter(error.to_string()))
806}
807
808async fn send_pong(socket: &mut ws::WebSocket, payload: axum::body::Bytes) -> Result<()> {
809    tokio::time::timeout(WRITE_TIMEOUT, socket.send(ws::Message::Pong(payload)))
810        .await
811        .map_err(|_error| WebmcpError::Timeout(WRITE_TIMEOUT))?
812        .map_err(|error| WebmcpError::Adapter(error.to_string()))
813}
814
815async fn send_event(
816    socket: &mut ws::WebSocket,
817    sequence: u64,
818    event: &vtcode_exec_events::VersionedThreadEvent,
819    max_frame_bytes: usize,
820) -> Result<()> {
821    let message = BridgeEventMessage { kind: "event", sequence, event: event.clone() };
822    let serialized = serde_json::to_string(&message)?;
823    if serialized.len() > max_frame_bytes {
824        return Err(WebmcpError::LimitExceeded);
825    }
826    tokio::time::timeout(WRITE_TIMEOUT, socket.send(ws::Message::Text(serialized.into())))
827        .await
828        .map_err(|_error| WebmcpError::Timeout(WRITE_TIMEOUT))?
829        .map_err(|error| WebmcpError::Adapter(error.to_string()))
830}
831
832fn response_for_error(request_id: &str, error: WebmcpError) -> BridgeResponse {
833    let (code, message) = match &error {
834        WebmcpError::OriginRejected(_) => ("origin_rejected", "browser origin is not allowed".to_string()),
835        WebmcpError::PairingExpired => ("pairing_expired", error.to_string()),
836        WebmcpError::PairingUsed => ("pairing_used", error.to_string()),
837        WebmcpError::Unauthorized => ("unauthorized", error.to_string()),
838        WebmcpError::LimitExceeded => ("limit_exceeded", error.to_string()),
839        WebmcpError::PathRejected(_) => ("path_rejected", error.to_string()),
840        WebmcpError::Conflict { .. } => ("conflict", error.to_string()),
841        WebmcpError::ProposalNotFound => ("proposal_not_found", error.to_string()),
842        WebmcpError::ApprovalRequired => ("approval_required", error.to_string()),
843        WebmcpError::Unsupported(_) => ("unsupported", error.to_string()),
844        WebmcpError::ChangeNotFound => ("change_not_found", error.to_string()),
845        WebmcpError::PartialApply => ("partial_apply", error.to_string()),
846        WebmcpError::SequenceGap { .. } => ("sequence_gap", error.to_string()),
847        WebmcpError::SlowClient => ("slow_client", error.to_string()),
848        WebmcpError::Timeout(_) => ("timeout", error.to_string()),
849        WebmcpError::InvalidRequest(_) | WebmcpError::Json(_) => ("invalid_request", error.to_string()),
850        WebmcpError::Io(_) | WebmcpError::Adapter(_) => {
851            ("runtime_error", "WebMCP runtime operation failed".to_string())
852        }
853    };
854    BridgeResponse::failure(response_request_id(request_id), code, message)
855}
856
857fn invalid_request_id_response(request_id: &str) -> BridgeResponse {
858    BridgeResponse::failure(
859        response_request_id(request_id),
860        "invalid_request",
861        "request_id must be between 1 and 256 UTF-8 bytes",
862    )
863}
864
865fn validate_config(config: &WebmcpServerConfig) -> Result<()> {
866    if config.host.trim().is_empty() || config.max_in_flight_requests == 0 {
867        return Err(WebmcpError::InvalidRequest("WebMCP host and limits must be non-empty".to_string()));
868    }
869    if config.max_frame_bytes < MIN_FRAME_BYTES
870        || config.max_frame_bytes > 16 * 1024 * 1024
871        || config.max_in_flight_requests > 64
872    {
873        return Err(WebmcpError::LimitExceeded);
874    }
875    if (config.allowed_origins.is_empty() && config.remote_mcp.is_none())
876        || config.allowed_origins.iter().any(|origin| !is_valid_origin(origin))
877    {
878        return Err(WebmcpError::InvalidRequest("WebMCP requires an explicit origin allowlist".to_string()));
879    }
880    let address = config
881        .host
882        .parse::<IpAddr>()
883        .map_err(|_error| WebmcpError::InvalidRequest("WebMCP host must be a literal IP address".to_string()))?;
884    if !address.is_loopback() {
885        return Err(WebmcpError::InvalidRequest(
886            "WebMCP only binds loopback; place a TLS-terminating reverse proxy in front of it for remote access"
887                .to_string(),
888        ));
889    }
890    match (config.allow_remote, config.public_url.as_deref()) {
891        (false, Some(_)) => {
892            return Err(WebmcpError::InvalidRequest("--public-url requires remote WebMCP mode".to_string()));
893        }
894        (true, None) => {
895            return Err(WebmcpError::InvalidRequest("remote WebMCP mode requires a wss:// public URL".to_string()));
896        }
897        (true, Some(url)) if !is_valid_public_url(url) => {
898            return Err(WebmcpError::InvalidRequest(
899                "remote WebMCP mode requires a valid wss:// public URL".to_string(),
900            ));
901        }
902        (false, None) => {}
903        (true, Some(_)) => {}
904    }
905    if config.request_timeout.is_zero() {
906        return Err(WebmcpError::InvalidRequest("WebMCP request timeout must be positive".to_string()));
907    }
908    if let Some(remote_mcp) = config.remote_mcp.as_ref() {
909        remote_mcp.validate()?;
910    }
911    Ok(())
912}
913
914fn is_valid_public_url(url: &str) -> bool {
915    let Ok(parsed) = url::Url::parse(url) else {
916        return false;
917    };
918    url == url.trim()
919        && !url.chars().any(char::is_whitespace)
920        && parsed.scheme() == "wss"
921        && parsed.host_str().is_some_and(|host| !host.is_empty())
922        && parsed.username().is_empty()
923        && parsed.password().is_none()
924        && parsed.query().is_none()
925        && parsed.fragment().is_none()
926}
927
928#[cfg(test)]
929mod tests {
930    use super::*;
931    use crate::FilesystemWorkspace;
932    use tempfile::TempDir;
933
934    #[tokio::test]
935    async fn explanation_requests_require_origin_bound_tokens_and_report_unsupported() {
936        let temp = TempDir::new().expect("workspace");
937        let adapter = Arc::new(FilesystemWorkspace::new(temp.path(), [], false).await.expect("adapter"));
938        let origin = "https://example.test";
939        let server = WebmcpServer::new(
940            adapter,
941            WebmcpServerConfig {
942                allowed_origins: vec![origin.to_owned(), "https://other.test".to_owned()],
943                ..Default::default()
944            },
945        )
946        .expect("server");
947        let pairing = server.begin_pairing_for_origin(origin).expect("pairing");
948        let session = server
949            .state
950            .pairing
951            .pair(pairing.code(), origin)
952            .expect("authenticated session");
953        let token = session.token().to_owned();
954        let reference = vtcode_memory::explanation::EvidenceRef {
955            session_id: "session".into(),
956            offset: 0,
957            length: 1,
958            digest: "0".repeat(64),
959            item_id: None,
960        };
961        let requests = [
962            BridgeRequest::ExplanationGet {
963                request_id: "get".into(),
964                token: token.clone(),
965                scope: vtcode_memory::explanation::ExplanationScope::Task,
966                offset: 0,
967            },
968            BridgeRequest::ExplanationEvidence {
969                request_id: "evidence".into(),
970                token: token.clone(),
971                reference: reference.clone(),
972                offset: 0,
973            },
974            BridgeRequest::ExplanationNavigate {
975                request_id: "navigate".into(),
976                token: token.clone(),
977                reference,
978            },
979        ];
980        for request in &requests {
981            let dispatch = || Arc::clone(&server.state.dispatch);
982            assert!(matches!(
983                dispatch_request(dispatch(), origin.into(), "wrong-token".into(), request.clone()).await,
984                Err(WebmcpError::Unauthorized)
985            ));
986            let error = dispatch_request(dispatch(), origin.into(), token.clone(), request.clone())
987                .await
988                .expect_err("headless adapter does not provide explanations");
989            assert!(matches!(error, WebmcpError::Unsupported(_)));
990            assert_eq!(response_for_error("operation", error).error.expect("error payload").code, "unsupported");
991        }
992        assert!(matches!(
993            dispatch_request(
994                Arc::clone(&server.state.dispatch),
995                "https://other.test".into(),
996                token.clone(),
997                requests[0].clone()
998            )
999            .await,
1000            Err(WebmcpError::Unauthorized)
1001        ));
1002        for request in requests {
1003            assert!(matches!(
1004                dispatch_request(Arc::clone(&server.state.dispatch), origin.into(), token.clone(), request).await,
1005                Err(WebmcpError::Unauthorized)
1006            ));
1007        }
1008        let status = server.state.adapter.status().await.expect("status");
1009        assert!(!status.explanations_available);
1010        assert!(!status.turns_available);
1011        assert!(std::fs::read_dir(temp.path()).expect("workspace files").next().is_none());
1012    }
1013
1014    #[tokio::test]
1015    async fn server_requires_explicit_origins_and_remote_flags() {
1016        let temp = TempDir::new().expect("temp dir");
1017        let adapter = Arc::new(FilesystemWorkspace::new(temp.path(), [], false).await.expect("adapter"));
1018        assert!(matches!(
1019            WebmcpServer::new(adapter.clone(), WebmcpServerConfig::default()),
1020            Err(WebmcpError::InvalidRequest(_))
1021        ));
1022        let config = WebmcpServerConfig {
1023            host: "0.0.0.0".to_string(),
1024            allowed_origins: vec!["https://example.test".to_string()],
1025            ..Default::default()
1026        };
1027        assert!(matches!(WebmcpServer::new(adapter, config), Err(WebmcpError::InvalidRequest(_))));
1028
1029        let remote_config = WebmcpServerConfig {
1030            host: "0.0.0.0".to_string(),
1031            allowed_origins: vec!["https://example.test".to_string()],
1032            allow_remote: true,
1033            public_url: Some("wss://bridge.example.test/webmcp".to_string()),
1034            ..Default::default()
1035        };
1036        let temp = TempDir::new().expect("temp dir");
1037        let adapter = Arc::new(FilesystemWorkspace::new(temp.path(), [], false).await.expect("adapter"));
1038        assert!(matches!(WebmcpServer::new(adapter, remote_config), Err(WebmcpError::InvalidRequest(_))));
1039
1040        let temp = TempDir::new().expect("temp dir");
1041        let adapter = Arc::new(FilesystemWorkspace::new(temp.path(), [], false).await.expect("adapter"));
1042        let valid_proxy_config = WebmcpServerConfig {
1043            allowed_origins: vec!["https://example.test".to_string()],
1044            allow_remote: true,
1045            public_url: Some("wss://bridge.example.test/webmcp".to_string()),
1046            ..Default::default()
1047        };
1048        assert!(WebmcpServer::new(adapter.clone(), valid_proxy_config).is_ok());
1049        let invalid_public_url_config = WebmcpServerConfig {
1050            allowed_origins: vec!["https://example.test".to_string()],
1051            public_url: Some("ws://bridge.example.test/webmcp".to_string()),
1052            ..Default::default()
1053        };
1054        assert!(matches!(WebmcpServer::new(adapter, invalid_public_url_config), Err(WebmcpError::InvalidRequest(_))));
1055
1056        let temp = TempDir::new().expect("temp dir");
1057        let adapter = Arc::new(FilesystemWorkspace::new(temp.path(), [], false).await.expect("adapter"));
1058        let invalid_public_url_config = WebmcpServerConfig {
1059            allowed_origins: vec!["https://example.test".to_string()],
1060            allow_remote: true,
1061            public_url: Some("wss://".to_string()),
1062            ..Default::default()
1063        };
1064        assert!(matches!(WebmcpServer::new(adapter, invalid_public_url_config), Err(WebmcpError::InvalidRequest(_))));
1065    }
1066}
vtcode_webmcp/server.rs

vtcode_webmcp/
server.rs
