Skip to main content

relay_knowledge/interfaces/agent/
mcp.rs

1use std::{
2    error::Error,
3    fmt,
4    future::Future,
5    time::{Duration, Instant},
6};
7
8mod audit_bridge;
9mod code_tools;
10mod http_contract;
11mod metrics;
12mod notifications;
13mod prompts;
14mod resources;
15mod scope_authorization;
16mod state;
17mod tool_registry;
18
19use axum::{
20    Router,
21    body::Bytes,
22    extract::State,
23    http::{HeaderMap, HeaderValue, StatusCode, header},
24    response::{IntoResponse, Response},
25    routing::{get, post},
26};
27use serde::Deserialize;
28use serde_json::{Value, json};
29use tokio::sync::watch;
30use tower_http::{limit::RequestBodyLimitLayer, trace::TraceLayer};
31
32use http_contract::{
33    ensure_remote_bind_allowed, validate_http_headers, validate_origin,
34    validate_protocol_version_header,
35};
36use scope_authorization::RuntimeScopeAuthorizer;
37use state::{CancellationRegistry, SessionCreateError, SessionLookupError, SessionRegistry};
38
39use crate::{
40    api::{
41        AgentRetrievalResult, ApiError, ErrorKind, GraphInspectionRequest, HybridRetrievalRequest,
42        InterfaceKind, RequestContext, RuntimeIdentity, freshness_label,
43    },
44    application::{AgentRuntimeConfig, RelayKnowledgeService},
45    domain::FreshnessPolicy,
46    net::{
47        NetworkRuntime,
48        http::HttpServeError,
49        qos::{QosPermit, QosRuntime, RejectReason},
50    },
51    observability::AgentProtocolMetrics,
52    project::PROJECT_NAME,
53};
54
55use super::{
56    AgentAdapterError, AgentAdapterErrorKind, AgentAuditEvent, AgentAuditLog, AgentAuditSink,
57    authorize_limit, validate_query_text,
58};
59use audit_bridge::{record_mcp_qos_rejection, record_mcp_tool_audit};
60use code_tools::run_code_tool;
61use tool_registry::{
62    CODE_CONTEXT_TOOL, CODE_FEATURE_FLAGS_TOOL, CODE_IMPACT_TOOL, CODE_QUERY_TOOL,
63    CODE_REPOSITORY_SET_QUERY_TOOL, CODE_SOFTWARE_QUERY_TOOL, CODEBASE_VIEW_TOOL, HEALTH_TOOL,
64    INDEX_STATUS_TOOL, INSPECT_GRAPH_TOOL, RETRIEVE_CONTEXT_TOOL, SERVICE_STATUS_TOOL,
65    is_known_tool,
66};
67
68pub const MCP_PROTOCOL_VERSION: &str = "2025-11-25";
69const MCP_PROTOCOL_VERSION_HEADER: &str = "mcp-protocol-version";
70const MCP_SESSION_ID_HEADER: &str = "mcp-session-id";
71
72/// MCP Streamable HTTP server state shared by route handlers.
73#[derive(Clone)]
74pub struct McpServer {
75    service: RelayKnowledgeService,
76    network: NetworkRuntime,
77    agent: AgentRuntimeConfig,
78    qos: QosRuntime,
79    audit: AgentAuditLog,
80    metrics: AgentProtocolMetrics,
81    cancellations: CancellationRegistry,
82    sessions: SessionRegistry,
83    scope_authorizer: RuntimeScopeAuthorizer,
84}
85
86impl McpServer {
87    /// Creates MCP server state from validated runtime boundaries.
88    pub fn new(
89        service: RelayKnowledgeService,
90        network: NetworkRuntime,
91        agent: AgentRuntimeConfig,
92    ) -> Self {
93        let metrics = service.observability().agent_metrics();
94        let qos = network.qos_runtime();
95        let audit = if agent.audit_sink_enabled {
96            AgentAuditSink::jsonl(service.agent_audit_log_path(), agent.audit_queue_depth)
97                .map(AgentAuditLog::with_sink)
98                .unwrap_or_default()
99        } else {
100            AgentAuditLog::default()
101        };
102
103        Self {
104            service,
105            network,
106            agent,
107            qos,
108            audit,
109            metrics,
110            cancellations: CancellationRegistry::default(),
111            sessions: SessionRegistry::default(),
112            scope_authorizer: RuntimeScopeAuthorizer::default(),
113        }
114    }
115
116    /// Builds the Streamable HTTP router without opening sockets.
117    pub fn router(self) -> Router {
118        let config = self.network.current();
119        let endpoint = self.agent.mcp_endpoint.clone();
120        let metrics_endpoint = metrics::metrics_endpoint(&endpoint);
121        let body_limit = usize::try_from(config.http.max_request_body_bytes).unwrap_or(usize::MAX);
122
123        Router::new()
124            .route(&endpoint, post(handle_mcp_post))
125            .route(&endpoint, axum::routing::delete(handle_mcp_delete))
126            .route(&metrics_endpoint, get(metrics::handle_metrics_get))
127            .with_state(self)
128            .layer(TraceLayer::new_for_http())
129            .layer(RequestBodyLimitLayer::new(body_limit))
130    }
131
132    /// Starts the MCP HTTP listener through `net::http`.
133    pub async fn serve_until_shutdown(
134        self,
135        shutdown: impl Future<Output = ()> + Send + 'static,
136    ) -> Result<(), McpServeError> {
137        let network_config = self.network.current();
138        let config = network_config.http;
139        let qos_policy = network_config.qos;
140        let qos = self.qos.clone();
141        let router = self.checked_router()?;
142
143        crate::net::http::serve_router_with_qos(router, config, qos, qos_policy, shutdown)
144            .await
145            .map_err(McpServeError::Http)
146    }
147
148    /// Builds the Streamable HTTP router after validating listener policy.
149    pub fn checked_router(self) -> Result<Router, McpServeError> {
150        if !self.agent.mcp_streamable_http_enabled {
151            return Err(McpServeError::Disabled);
152        }
153        ensure_remote_bind_allowed(&self.network.current().http, &self.agent.access_policy)?;
154
155        Ok(self.router())
156    }
157
158    #[cfg(test)]
159    pub fn qos_snapshot(&self) -> crate::net::qos::QosSnapshot {
160        self.qos.snapshot()
161    }
162
163    #[cfg(test)]
164    pub fn audit_snapshot(&self) -> Vec<AgentAuditEvent> {
165        self.audit.snapshot()
166    }
167}
168
169/// MCP server startup error.
170#[derive(Debug)]
171pub enum McpServeError {
172    Disabled,
173    RemoteBindDisabled,
174    Http(HttpServeError),
175}
176
177impl fmt::Display for McpServeError {
178    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
179        match self {
180            Self::Disabled => write!(formatter, "MCP Streamable HTTP is not enabled"),
181            Self::RemoteBindDisabled => {
182                write!(
183                    formatter,
184                    "MCP remote bind requires allow_remote_clients=true"
185                )
186            }
187            Self::Http(error) => write!(formatter, "{error}"),
188        }
189    }
190}
191
192impl Error for McpServeError {}
193
194#[derive(Debug, Deserialize)]
195struct JsonRpcRequest {
196    jsonrpc: Option<String>,
197    id: Option<Value>,
198    method: Option<String>,
199    #[serde(default)]
200    params: Value,
201}
202
203#[derive(Debug, Deserialize)]
204struct InitializeParams {
205    #[serde(rename = "protocolVersion")]
206    protocol_version: String,
207    capabilities: Value,
208    #[serde(rename = "clientInfo")]
209    client_info: InitializeClientInfo,
210}
211
212#[derive(Debug, Deserialize)]
213struct InitializeClientInfo {
214    name: String,
215    version: String,
216}
217
218#[derive(Debug, Deserialize)]
219struct ToolCallParams {
220    name: String,
221    #[serde(default)]
222    arguments: Value,
223}
224
225#[derive(Debug, Deserialize)]
226struct RetrieveContextArgs {
227    query: String,
228    #[serde(default)]
229    source_scope: Option<String>,
230    #[serde(default)]
231    limit: Option<usize>,
232    #[serde(default)]
233    freshness: Option<String>,
234}
235
236#[derive(Debug, Deserialize)]
237struct InspectGraphArgs {
238    #[serde(default)]
239    source_scope: Option<String>,
240}
241
242#[derive(Debug, Deserialize)]
243pub(super) struct CancelParams {
244    #[serde(rename = "requestId")]
245    request_id: Value,
246}
247
248pub(super) struct McpMethodError {
249    code: i64,
250    kind: &'static str,
251    message: String,
252}
253
254impl McpMethodError {
255    fn invalid_params(message: impl Into<String>) -> Self {
256        Self {
257            code: -32602,
258            kind: "invalid_argument",
259            message: message.into(),
260        }
261    }
262
263    fn internal(message: impl Into<String>) -> Self {
264        Self {
265            code: -32603,
266            kind: "internal",
267            message: message.into(),
268        }
269    }
270
271    fn timeout(message: impl Into<String>) -> Self {
272        Self {
273            code: -32000,
274            kind: "timeout",
275            message: message.into(),
276        }
277    }
278
279    fn api(error: ApiError) -> Self {
280        Self {
281            code: -32000,
282            kind: match error.error_kind {
283                ErrorKind::InvalidArgument => "invalid_argument",
284                ErrorKind::StorageUnavailable => "storage_unavailable",
285                ErrorKind::QosRejected => "qos_rejected",
286                ErrorKind::Timeout => "timeout",
287                ErrorKind::Internal => "internal",
288            },
289            message: error.message,
290        }
291    }
292
293    fn adapter(error: AgentAdapterError) -> Self {
294        Self {
295            code: -32000,
296            kind: error.kind.as_str(),
297            message: error.message,
298        }
299    }
300}
301
302async fn handle_mcp_post(
303    State(server): State<McpServer>,
304    headers: HeaderMap,
305    body: Bytes,
306) -> Response {
307    if let Err(status) = validate_http_headers(&server, &headers) {
308        return status.into_response();
309    }
310    if body.len() as u64 > server.network.current().http.max_request_body_bytes {
311        return StatusCode::PAYLOAD_TOO_LARGE.into_response();
312    }
313
314    let payload = match serde_json::from_slice::<Value>(&body) {
315        Ok(payload) => payload,
316        Err(error) => {
317            return json_rpc_error(Value::Null, -32700, format!("parse error: {error}"));
318        }
319    };
320    if payload.is_array() {
321        return json_rpc_error(Value::Null, -32600, "batch requests are not supported");
322    }
323    let request = match serde_json::from_value::<JsonRpcRequest>(payload.clone()) {
324        Ok(request) => request,
325        Err(error) => {
326            return json_rpc_error(Value::Null, -32600, format!("invalid request: {error}"));
327        }
328    };
329    let id = request.id.clone().unwrap_or(Value::Null);
330    if request.jsonrpc.as_deref() != Some("2.0") {
331        return json_rpc_error(id, -32600, "jsonrpc must be 2.0");
332    }
333    let Some(method) = request.method.as_deref() else {
334        if is_valid_json_rpc_response(&payload) {
335            if let Err(status) = validate_protocol_version_header(&headers, true) {
336                return status.into_response();
337            }
338            return response_message_session_response(&server, &headers);
339        }
340        if payload
341            .as_object()
342            .is_some_and(|object| object.contains_key("result") || object.contains_key("error"))
343        {
344            return StatusCode::BAD_REQUEST.into_response();
345        }
346        return json_rpc_error(id, -32600, "method is required");
347    };
348
349    if method == "initialize" {
350        let Some(id) = request.id else {
351            return json_rpc_error(Value::Null, -32600, "requests must include an id");
352        };
353        if !is_json_rpc_id(&id) {
354            return invalid_request_id_response();
355        }
356        if let Err(message) = validate_initialize_params(request.params) {
357            return json_rpc_error(id, -32602, message);
358        }
359        let Ok(permit) = admit_mcp_request(&server) else {
360            return StatusCode::TOO_MANY_REQUESTS.into_response();
361        };
362        let session_id = match server.sessions.require_session(&headers) {
363            Ok(session) => session.session_id().to_owned(),
364            Err(SessionLookupError::Missing) => match server.sessions.create_session() {
365                Ok(session_id) => session_id,
366                Err(error) => {
367                    drop(permit);
368                    return session_create_error(id, error);
369                }
370            },
371            Err(error) => {
372                drop(permit);
373                return session_lookup_error_response(error);
374            }
375        };
376        drop(permit);
377        return json_rpc_success_with_session(id, initialize_result(), &session_id);
378    }
379
380    if let Err(status) = validate_protocol_version_header(&headers, true) {
381        return status.into_response();
382    }
383
384    let session = match server.sessions.require_session(&headers) {
385        Ok(session) => session,
386        Err(error) => return session_lookup_error_response(error),
387    };
388
389    if method == "notifications/initialized" {
390        if request.id.is_some() {
391            return json_rpc_error(id, -32600, "notifications must not include an id");
392        }
393        let Ok(permit) = admit_mcp_request(&server) else {
394            return StatusCode::TOO_MANY_REQUESTS.into_response();
395        };
396        if let Err(error) = server.sessions.mark_initialized(session.session_id()) {
397            drop(permit);
398            return session_lookup_error_response(error);
399        }
400        drop(permit);
401        return StatusCode::ACCEPTED.into_response();
402    }
403
404    if !session.initialized {
405        return uninitialized_session_response(request.id);
406    }
407
408    let namespace = session.namespace();
409    if method.starts_with("notifications/") {
410        if request.id.is_some() {
411            return json_rpc_error(id, -32600, "notifications must not include an id");
412        }
413        if method == "notifications/cancelled" {
414            notifications::handle_notification(&server, method, request.params, &namespace);
415            return StatusCode::ACCEPTED.into_response();
416        }
417        let Ok(permit) = admit_mcp_request(&server) else {
418            return StatusCode::TOO_MANY_REQUESTS.into_response();
419        };
420        notifications::handle_notification(&server, method, request.params, &namespace);
421        drop(permit);
422        return StatusCode::ACCEPTED.into_response();
423    }
424
425    let Some(id) = request.id else {
426        return json_rpc_error(Value::Null, -32600, "requests must include an id");
427    };
428    let Some(request_id) = request_id_key(&namespace, &id) else {
429        return invalid_request_id_response();
430    };
431    let permit = match admit_mcp_request(&server) {
432        Ok(permit) => permit,
433        Err(reason) => {
434            let error =
435                AgentAdapterError::new(AgentAdapterErrorKind::QosRejected, qos_message(reason));
436            record_mcp_qos_rejection(&server, method, &id, error.kind.as_str());
437            server.metrics.record_rejection("mcp", error.kind.as_str());
438            return if method == "tools/call" {
439                json_rpc_success(id, tool_error_result(error))
440            } else {
441                json_rpc_error(id, -32000, error.to_string())
442            };
443        }
444    };
445
446    let started = Instant::now();
447    let mut pending_tool_audit = None;
448    let result = match method {
449        "ping" => json!({}),
450        "tools/list" => metrics::tools_list_result(&server, &session),
451        "resources/list" => resources::list_resources(&server),
452        "resources/read" => {
453            match resources::read_resource_with_timeout(&server, request.params, &request_id).await
454            {
455                Ok(result) => result,
456                Err(error) => return json_rpc_error(id, error.code, error.message),
457            }
458        }
459        "prompts/list" => prompts::list_prompts(),
460        "prompts/get" => match prompts::get_prompt(&server, request.params, &request_id).await {
461            Ok(result) => result,
462            Err(error) => return json_rpc_error(id, error.code, error.message),
463        },
464        "tools/call" => {
465            let params = match serde_json::from_value::<ToolCallParams>(request.params) {
466                Ok(params) => params,
467                Err(error) => {
468                    return json_rpc_error(
469                        id,
470                        -32602,
471                        format!("invalid tools/call params: {error}"),
472                    );
473                }
474            };
475            if !is_known_tool(&params.name) {
476                return json_rpc_error(id, -32602, "unknown tool name");
477            }
478            let outcome = run_cancellable_tool_call(&server, params, request_id).await;
479            pending_tool_audit = Some((
480                outcome.operation,
481                outcome.request_id,
482                outcome.result.clone(),
483                outcome.duration_ms,
484            ));
485            outcome.result
486        }
487        _ => return json_rpc_error(id, -32601, "method not found"),
488    };
489
490    drop(permit);
491    if let Some((operation, request_id, result, duration_ms)) = pending_tool_audit {
492        record_mcp_tool_audit(&server, &operation, &request_id, &result, duration_ms).await;
493    } else if !matches!(method, "resources/read" | "prompts/get") {
494        server
495            .metrics
496            .record_request("mcp", method, "completed", elapsed_millis(started), false);
497    }
498    json_rpc_success(id, result)
499}
500
501async fn handle_mcp_delete(State(server): State<McpServer>, headers: HeaderMap) -> Response {
502    if let Err(status) = validate_origin(&server, &headers) {
503        return status.into_response();
504    }
505    if let Err(status) = validate_protocol_version_header(&headers, true) {
506        return status.into_response();
507    }
508    let permit = match admit_mcp_request(&server) {
509        Ok(permit) => permit,
510        Err(_) => {
511            server.metrics.record_rejection("mcp", "qos_rejected");
512            return StatusCode::TOO_MANY_REQUESTS.into_response();
513        }
514    };
515    match server.sessions.terminate_session(&headers) {
516        Ok(()) => {
517            drop(permit);
518            StatusCode::ACCEPTED.into_response()
519        }
520        Err(error) => {
521            drop(permit);
522            session_lookup_error_response(error)
523        }
524    }
525}
526
527fn is_valid_json_rpc_response(payload: &Value) -> bool {
528    let Some(object) = payload.as_object() else {
529        return false;
530    };
531    let has_result = object.contains_key("result");
532    let has_error = object.contains_key("error");
533    if has_result == has_error {
534        return false;
535    }
536
537    object.get("id").is_some_and(is_json_rpc_id)
538}
539
540fn validate_initialize_params(params: Value) -> Result<(), String> {
541    let params = serde_json::from_value::<InitializeParams>(params)
542        .map_err(|error| format!("invalid initialize params: {error}"))?;
543    if params.protocol_version != MCP_PROTOCOL_VERSION {
544        return Err(format!(
545            "unsupported MCP protocol version '{}'",
546            params.protocol_version
547        ));
548    }
549    if !params.capabilities.is_object() {
550        return Err("initialize capabilities must be an object".to_owned());
551    }
552    if params.client_info.name.trim().is_empty() || params.client_info.version.trim().is_empty() {
553        return Err("initialize clientInfo requires name and version".to_owned());
554    }
555
556    Ok(())
557}
558
559fn response_message_session_response(server: &McpServer, headers: &HeaderMap) -> Response {
560    match server.sessions.require_session(headers) {
561        Ok(session) if session.initialized => StatusCode::ACCEPTED.into_response(),
562        Ok(_) => StatusCode::BAD_REQUEST.into_response(),
563        Err(error) => session_lookup_error_response(error),
564    }
565}
566
567fn session_lookup_error_response(error: SessionLookupError) -> Response {
568    match error {
569        SessionLookupError::Missing | SessionLookupError::InvalidHeader => {
570            StatusCode::BAD_REQUEST.into_response()
571        }
572        SessionLookupError::Unknown => StatusCode::NOT_FOUND.into_response(),
573    }
574}
575
576fn uninitialized_session_response(id: Option<Value>) -> Response {
577    let Some(id) = id else {
578        return StatusCode::BAD_REQUEST.into_response();
579    };
580    if is_json_rpc_id(&id) {
581        json_rpc_error(id, -32002, "MCP session is not initialized")
582    } else {
583        invalid_request_id_response()
584    }
585}
586
587fn invalid_request_id_response() -> Response {
588    json_rpc_error(Value::Null, -32600, "request id must be a string or number")
589}
590
591fn session_create_error(id: Value, error: SessionCreateError) -> Response {
592    json_rpc_error(id, -32603, format!("failed to create MCP session: {error}"))
593}
594
595fn admit_mcp_request(server: &McpServer) -> Result<QosPermit, RejectReason> {
596    let policy = server.network.current().qos;
597    if crate::net::http::qos_request_context_active() {
598        return Ok(QosPermit::already_admitted(server.qos.clone()));
599    }
600    server.qos.admit_queued_request(&policy)
601}
602
603async fn run_cancellable_tool_call(
604    server: &McpServer,
605    params: ToolCallParams,
606    request_id: String,
607) -> ToolCallOutcome {
608    let started = Instant::now();
609    let operation = params.name.clone();
610    let (mut cancellation, _registration) = server.cancellations.register(request_id.clone());
611    let timeout = Duration::from_millis(server.agent.access_policy.max_runtime_ms);
612    let tool = run_tool_call(server, params, request_id.clone());
613
614    let result = tokio::select! {
615        result = tokio::time::timeout(timeout, tool) => match result {
616            Ok(value) => value,
617            Err(_) => {
618                server.qos.record_timed_out();
619                tool_error_result(AgentAdapterError::new(
620                    AgentAdapterErrorKind::Timeout,
621                    "MCP tool call exceeded max_runtime_ms",
622                ))
623            }
624        },
625        _ = wait_for_cancellation(&mut cancellation) => {
626            tool_error_result(AgentAdapterError::new(
627                AgentAdapterErrorKind::Cancelled,
628                "MCP tool call was cancelled",
629            ))
630        }
631    };
632
633    ToolCallOutcome {
634        operation,
635        request_id,
636        result,
637        duration_ms: elapsed_millis(started),
638    }
639}
640
641struct ToolCallOutcome {
642    operation: String,
643    request_id: String,
644    result: Value,
645    duration_ms: u64,
646}
647
648async fn wait_for_cancellation(cancellation: &mut watch::Receiver<bool>) {
649    while cancellation.changed().await.is_ok() {
650        if *cancellation.borrow() {
651            return;
652        }
653    }
654
655    std::future::pending::<()>().await;
656}
657
658async fn run_tool_call(server: &McpServer, params: ToolCallParams, request_id: String) -> Value {
659    match params.name.as_str() {
660        RETRIEVE_CONTEXT_TOOL => retrieve_context_tool(server, params.arguments, request_id).await,
661        INSPECT_GRAPH_TOOL => inspect_graph_tool(server, params.arguments, request_id).await,
662        HEALTH_TOOL => health_tool(server, request_id).await,
663        SERVICE_STATUS_TOOL => service_status_tool(server, request_id).await,
664        INDEX_STATUS_TOOL => index_status_tool(server, request_id).await,
665        CODE_QUERY_TOOL
666        | CODE_CONTEXT_TOOL
667        | CODE_FEATURE_FLAGS_TOOL
668        | CODE_IMPACT_TOOL
669        | CODE_REPOSITORY_SET_QUERY_TOOL
670        | CODE_SOFTWARE_QUERY_TOOL
671        | CODEBASE_VIEW_TOOL => {
672            run_code_tool(server, params.name.as_str(), params.arguments, request_id).await
673        }
674        _ => json!({
675            "content": [{"type": "text", "text": "unknown MCP tool"}],
676            "isError": true
677        }),
678    }
679}
680
681async fn retrieve_context_tool(server: &McpServer, arguments: Value, request_id: String) -> Value {
682    let started = Instant::now();
683    let args = match serde_json::from_value::<RetrieveContextArgs>(arguments) {
684        Ok(args) => args,
685        Err(error) => return tool_error_result(invalid_arguments(error)),
686    };
687    if let Err(error) = validate_query_text("query", &args.query) {
688        return tool_error_result(error);
689    }
690    let policy = &server.agent.access_policy;
691    let limit = match authorize_limit(args.limit, policy) {
692        Ok(limit) => limit,
693        Err(error) => return tool_error_result(error),
694    };
695    let source_scope = match server
696        .scope_authorizer
697        .authorize_scope(&server.service, policy, args.source_scope)
698        .await
699    {
700        Ok(scope) => scope,
701        Err(error) => return tool_error_result(error),
702    };
703    let freshness = match parse_freshness(args.freshness.as_deref()) {
704        Ok(freshness) => freshness,
705        Err(error) => return tool_error_result(error),
706    };
707    let context = request_context(request_id.clone());
708    let identity = RuntimeIdentity::mcp(Some(request_id));
709
710    match server
711        .service
712        .retrieve_context(
713            HybridRetrievalRequest {
714                query: args.query,
715                source_scope: source_scope.clone(),
716                limit,
717                freshness,
718            },
719            context,
720        )
721        .await
722    {
723        Ok(response) => {
724            let elapsed_ms = elapsed_millis(started);
725            let result = AgentRetrievalResult::from_retrieval(
726                response,
727                identity,
728                policy.max_context_bytes,
729                elapsed_ms,
730            );
731            tool_success_result(
732                format!(
733                    "retrieved {} result(s), graph_version={}, freshness={}",
734                    result.results.len(),
735                    result.metadata.graph_version,
736                    freshness_label(freshness)
737                ),
738                json!(result),
739            )
740        }
741        Err(error) => api_error_result(error),
742    }
743}
744
745async fn inspect_graph_tool(server: &McpServer, arguments: Value, request_id: String) -> Value {
746    let args = match serde_json::from_value::<InspectGraphArgs>(arguments) {
747        Ok(args) => args,
748        Err(error) => return tool_error_result(invalid_arguments(error)),
749    };
750    let source_scope = match server
751        .scope_authorizer
752        .authorize_scope(
753            &server.service,
754            &server.agent.access_policy,
755            args.source_scope,
756        )
757        .await
758    {
759        Ok(scope) => scope,
760        Err(error) => return tool_error_result(error),
761    };
762
763    match server
764        .service
765        .inspect_graph(
766            GraphInspectionRequest { source_scope },
767            request_context(request_id),
768        )
769        .await
770    {
771        Ok(response) => tool_success_result("graph inspection completed", json!(response)),
772        Err(error) => api_error_result(error),
773    }
774}
775
776async fn health_tool(server: &McpServer, request_id: String) -> Value {
777    match server
778        .service
779        .read_only_health(request_context(request_id))
780        .await
781    {
782        Ok(response) => tool_success_result(
783            format!(
784                "health={}",
785                if response.healthy { "ok" } else { "degraded" }
786            ),
787            json!(response),
788        ),
789        Err(error) => api_error_result(error),
790    }
791}
792
793async fn service_status_tool(server: &McpServer, request_id: String) -> Value {
794    match server
795        .service
796        .read_only_service_status(request_context(request_id))
797        .await
798    {
799        Ok(response) => tool_success_result("service status loaded", json!(response)),
800        Err(error) => api_error_result(error),
801    }
802}
803
804async fn index_status_tool(server: &McpServer, request_id: String) -> Value {
805    match server.service.health(request_context(request_id)).await {
806        Ok(response) => tool_success_result(
807            "index status loaded",
808            json!({
809                "metadata": response.metadata,
810                "indexes": response.indexes,
811            }),
812        ),
813        Err(error) => api_error_result(error),
814    }
815}
816
817fn initialize_result() -> Value {
818    json!({
819        "protocolVersion": MCP_PROTOCOL_VERSION,
820        "capabilities": {
821            "tools": {},
822            "resources": {"listChanged": false},
823            "prompts": {"listChanged": false}
824        },
825        "serverInfo": {
826            "name": PROJECT_NAME,
827            "version": env!("CARGO_PKG_VERSION")
828        },
829        "instructions": "MCP tool schemas are static and storage is opened lazily on the first storage-backed tool call. For repository exploration, prefer relay_code_query or relay_code_repository_set_query and follow the explore_budget returned in structuredContent; budget tiers are 0-499 files: 1 call/15000 chars/5 files, 500-4999: 2/30000/10, 5000-14999: 3/45000/15, 15000+: 5/75000/25. Free-text queries are capped at 10000 characters and path filters at 4096 characters."
830    })
831}
832
833fn parse_freshness(value: Option<&str>) -> Result<FreshnessPolicy, AgentAdapterError> {
834    match value.unwrap_or("allow-stale") {
835        "allow-stale" => Ok(FreshnessPolicy::AllowStale),
836        "wait-until-fresh" => Ok(FreshnessPolicy::WaitUntilFresh),
837        "graph-only" => Ok(FreshnessPolicy::GraphOnly),
838        other => Err(AgentAdapterError::new(
839            AgentAdapterErrorKind::InvalidArgument,
840            format!("invalid freshness '{other}'"),
841        )),
842    }
843}
844
845fn tool_success_result(summary: impl Into<String>, structured_content: Value) -> Value {
846    json!({
847        "content": [{"type": "text", "text": summary.into()}],
848        "structuredContent": structured_content,
849        "isError": false
850    })
851}
852
853fn api_error_result(error: ApiError) -> Value {
854    tool_error_result(AgentAdapterError::new(
855        agent_error_kind(error.error_kind),
856        error.message,
857    ))
858}
859
860fn agent_error_kind(kind: ErrorKind) -> AgentAdapterErrorKind {
861    match kind {
862        ErrorKind::InvalidArgument => AgentAdapterErrorKind::InvalidArgument,
863        ErrorKind::StorageUnavailable => AgentAdapterErrorKind::StorageUnavailable,
864        ErrorKind::QosRejected => AgentAdapterErrorKind::QosRejected,
865        ErrorKind::Timeout => AgentAdapterErrorKind::Timeout,
866        ErrorKind::Internal => AgentAdapterErrorKind::Internal,
867    }
868}
869
870fn tool_error_result(error: AgentAdapterError) -> Value {
871    json!({
872        "content": [{
873            "type": "text",
874            "text": format!("{}: {}", error.kind.as_str(), error.message)
875        }],
876        "structuredContent": {
877            "error_kind": error.kind.as_str(),
878            "message": error.message,
879        },
880        "isError": true
881    })
882}
883
884fn invalid_arguments(error: serde_json::Error) -> AgentAdapterError {
885    AgentAdapterError::new(
886        AgentAdapterErrorKind::InvalidArgument,
887        format!("invalid tool arguments: {error}"),
888    )
889}
890
891fn domain_argument_error(error: impl fmt::Display) -> AgentAdapterError {
892    AgentAdapterError::new(AgentAdapterErrorKind::InvalidArgument, error.to_string())
893}
894
895fn json_rpc_success(id: Value, result: Value) -> Response {
896    json_response(
897        StatusCode::OK,
898        json!({ "jsonrpc": "2.0", "id": id, "result": result }),
899    )
900}
901
902fn json_rpc_success_with_session(id: Value, result: Value, session_id: &str) -> Response {
903    let mut response = json_rpc_success(id, result);
904    response.headers_mut().insert(
905        MCP_SESSION_ID_HEADER,
906        HeaderValue::from_str(session_id).expect("generated MCP session id is a valid header"),
907    );
908    response
909}
910
911fn json_rpc_error(id: Value, code: i64, message: impl Into<String>) -> Response {
912    json_response(
913        StatusCode::OK,
914        json!({
915            "jsonrpc": "2.0",
916            "id": id,
917            "error": {
918                "code": code,
919                "message": message.into()
920            }
921        }),
922    )
923}
924
925fn json_response(status: StatusCode, value: Value) -> Response {
926    (
927        status,
928        [(header::CONTENT_TYPE, "application/json")],
929        value.to_string(),
930    )
931        .into_response()
932}
933
934fn request_context(request_id: String) -> RequestContext {
935    RequestContext::with_ids(
936        InterfaceKind::Mcp,
937        format!("mcp-{request_id}"),
938        format!("trace-mcp-{request_id}"),
939    )
940}
941
942fn endpoint_child(endpoint: &str, child: &str) -> String {
943    if endpoint == "/" {
944        format!("/{child}")
945    } else {
946        format!("{}/{child}", endpoint.trim_end_matches('/'))
947    }
948}
949
950fn request_id_key(namespace: &str, value: &Value) -> Option<String> {
951    match value {
952        Value::String(value) => Some(format!("{namespace}|string:{value}")),
953        Value::Number(value) if value.is_i64() || value.is_u64() => {
954            Some(format!("{namespace}|number:{value}"))
955        }
956        _ => None,
957    }
958}
959
960fn is_json_rpc_id(value: &Value) -> bool {
961    match value {
962        Value::String(_) => true,
963        Value::Number(number) => number.is_i64() || number.is_u64(),
964        _ => false,
965    }
966}
967
968fn elapsed_millis(started: Instant) -> u64 {
969    u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX)
970}
971
972fn qos_message(reason: RejectReason) -> &'static str {
973    match reason {
974        RejectReason::ConnectionBudgetExceeded => "connection budget exhausted",
975        RejectReason::RequestBudgetExceeded => "request budget exhausted",
976        RejectReason::QueueBudgetExceeded => "queue budget exhausted",
977    }
978}
979
980#[cfg(test)]
981#[path = "mcp_feature_flag_tool_tests.rs"]
982mod mcp_feature_flag_tool_tests;
983#[cfg(test)]
984#[path = "mcp_issue_283_tests.rs"]
985mod mcp_issue_283_tests;
986#[cfg(test)]
987#[path = "mcp_protocol_tests.rs"]
988mod mcp_protocol_tests;
989#[cfg(test)]
990#[path = "mcp_software_tool_tests.rs"]
991mod mcp_software_tool_tests;
992#[cfg(test)]
993#[path = "mcp_test_support.rs"]
994mod mcp_test_support;
995#[cfg(test)]
996#[path = "mcp_tests.rs"]
997mod mcp_tests;
998#[cfg(test)]
999#[path = "mcp_tool_tests.rs"]
1000mod mcp_tool_tests;