diff --git a/apps/mcp-server/src/app.rs b/apps/mcp-server/src/app.rs index 42ae260..c985be9 100644 --- a/apps/mcp-server/src/app.rs +++ b/apps/mcp-server/src/app.rs @@ -10,7 +10,7 @@ use axum::{ extract::{Path, State}, http::{ HeaderMap, HeaderValue, StatusCode, - header::{self, ACCEPT, AUTHORIZATION}, + header::{self, ACCEPT, AUTHORIZATION, HeaderName}, }, response::{ IntoResponse, Response, @@ -48,6 +48,8 @@ use crate::{ const HEADER_MCP_SESSION_ID: &str = "MCP-Session-Id"; const HEADER_MCP_PROTOCOL_VERSION: &str = "MCP-Protocol-Version"; +const HEADER_X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id"); +const MAX_REQUEST_ID_LEN: usize = 128; #[derive(Clone, Copy)] enum ResponseMode { @@ -234,22 +236,28 @@ async fn mcp_post( headers: HeaderMap, Json(message): Json, ) -> Response { + let transport_request_id = resolve_request_id(&headers); + if let Err(status) = validate_origin(&state.allowed_origins, &headers) { - return status.into_response(); + return with_request_id_header(status.into_response(), &transport_request_id); } let response_mode = match negotiate_post_response_mode(&headers) { Ok(mode) => mode, - Err(status) => return status.into_response(), + Err(status) => { + return with_request_id_header(status.into_response(), &transport_request_id); + } }; if is_response(&message) || is_notification(&message) && method_name(&message).is_none() { - return StatusCode::ACCEPTED.into_response(); + return with_request_id_header(StatusCode::ACCEPTED.into_response(), &transport_request_id); } let protocol_version = match protocol_version_from_headers(&headers) { Ok(value) => value, - Err(status) => return status.into_response(), + Err(status) => { + return with_request_id_header(status.into_response(), &transport_request_id); + } }; if let Some(session_id) = headers.get(HEADER_MCP_SESSION_ID) { @@ -258,7 +266,7 @@ async fn mcp_post( if let Err(status) = validate_session_protocol_version(&headers, &session.protocol_version) { - return status.into_response(); + return with_request_id_header(status.into_response(), &transport_request_id); } } } @@ -269,10 +277,10 @@ async fn mcp_post( _ => PlatformApiKeyScope::Read, }; if let Err(status) = require_platform_api_key(&state, &path, &headers, required_scope).await { - return status.into_response(); + return with_request_id_header(status.into_response(), &transport_request_id); } - match method_name(&message) { + let response = match method_name(&message) { Some("initialize") if is_request(&message) => { handle_initialize(state, &path, &message, response_mode).await } @@ -362,6 +370,7 @@ async fn mcp_post( response_mode, resolved, arguments, + &transport_request_id, ) .await } @@ -402,7 +411,9 @@ async fn mcp_post( None, Some(&protocol_version), ), - } + }; + + with_request_id_header(response, &transport_request_id) } async fn handle_tool_call( @@ -412,6 +423,7 @@ async fn handle_tool_call( response_mode: ResponseMode, resolved: ResolvedToolCall, arguments: Value, + transport_request_id: &str, ) -> Response { match resolved.kind { GeneratedToolKind::Base => { @@ -422,6 +434,7 @@ async fn handle_tool_call( response_mode, resolved.tool, arguments, + transport_request_id, ) .await } @@ -433,6 +446,7 @@ async fn handle_tool_call( response_mode, resolved.tool, arguments, + transport_request_id, ) .await } @@ -444,6 +458,7 @@ async fn handle_tool_call( response_mode, resolved.tool, arguments, + transport_request_id, ) .await } @@ -466,6 +481,7 @@ async fn handle_tool_call( response_mode, resolved.tool, arguments, + transport_request_id, ) .await } @@ -609,6 +625,7 @@ async fn handle_base_tool_call( response_mode: ResponseMode, tool: PublishedAgentTool, arguments: Value, + transport_request_id: &str, ) -> Response { let operation = runtime_operation(&tool); let request_preview = build_request_preview(&state.runtime, &operation, &arguments); @@ -654,6 +671,7 @@ async fn handle_base_tool_call( &state, &tool, InvocationRecord { + request_id: Some(transport_request_id), tool_name: &tool.tool_name, status: InvocationStatus::Ok, level: InvocationLevel::Info, @@ -674,6 +692,7 @@ async fn handle_base_tool_call( &state, &tool, InvocationRecord { + request_id: Some(transport_request_id), tool_name: &tool.tool_name, status: InvocationStatus::Error, level: InvocationLevel::Error, @@ -706,6 +725,7 @@ async fn handle_session_start_call( response_mode: ResponseMode, tool: PublishedAgentTool, arguments: Value, + transport_request_id: &str, ) -> Response { let runtime_operation = runtime_operation(&tool); let request_preview = build_request_preview(&state.runtime, &runtime_operation, &arguments); @@ -796,6 +816,7 @@ async fn handle_session_start_call( &state, &tool, InvocationRecord { + request_id: Some(transport_request_id), tool_name: streaming .tool_family .start_tool_name @@ -820,6 +841,7 @@ async fn handle_session_start_call( &state, &tool, InvocationRecord { + request_id: Some(transport_request_id), tool_name: streaming .tool_family .start_tool_name @@ -856,6 +878,7 @@ async fn handle_session_poll_call( response_mode: ResponseMode, tool: PublishedAgentTool, arguments: Value, + transport_request_id: &str, ) -> Response { let control: SessionControlArgs = match serde_json::from_value(arguments.clone()) { Ok(value) => value, @@ -1009,6 +1032,7 @@ async fn handle_session_poll_call( &state, &tool, InvocationRecord { + request_id: Some(transport_request_id), tool_name: streaming .tool_family .poll_tool_name @@ -1108,6 +1132,7 @@ async fn handle_async_job_start_call( response_mode: ResponseMode, tool: PublishedAgentTool, arguments: Value, + transport_request_id: &str, ) -> Response { let operation = runtime_operation(&tool); let request_preview = build_request_preview(&state.runtime, &operation, &arguments); @@ -1224,6 +1249,7 @@ async fn handle_async_job_start_call( &state, &tool, InvocationRecord { + request_id: Some(transport_request_id), tool_name: streaming .tool_family .start_tool_name @@ -1837,6 +1863,7 @@ fn build_request_preview( } struct InvocationRecord<'a> { + request_id: Option<&'a str>, tool_name: &'a str, status: InvocationStatus, level: InvocationLevel, @@ -1865,7 +1892,7 @@ async fn persist_invocation( status: record.status, tool_name: record.tool_name.to_owned(), message: record.message.to_owned(), - request_id: None, + request_id: record.request_id.map(ToOwned::to_owned), status_code: record.status_code, duration_ms, error_kind: record.error_kind.map(ToOwned::to_owned), @@ -2093,6 +2120,13 @@ fn transport_response( json_response(status, payload, session_id, protocol_version) } +fn with_request_id_header(mut response: Response, request_id: &str) -> Response { + if let Ok(value) = HeaderValue::from_str(request_id) { + response.headers_mut().insert(HEADER_X_REQUEST_ID, value); + } + response +} + fn sse_response( status: StatusCode, stream: S, @@ -2219,6 +2253,24 @@ fn tool_definitions(tool: &PublishedAgentTool) -> Vec { definitions } +fn resolve_request_id(headers: &HeaderMap) -> String { + headers + .get(&HEADER_X_REQUEST_ID) + .and_then(|value| value.to_str().ok()) + .map(str::trim) + .filter(|value| is_valid_request_id(value)) + .map(ToOwned::to_owned) + .unwrap_or_else(|| uuid::Uuid::now_v7().to_string()) +} + +fn is_valid_request_id(value: &str) -> bool { + !value.is_empty() + && value.len() <= MAX_REQUEST_ID_LEN + && value + .bytes() + .all(|byte| matches!(byte, 0x21..=0x7e) && byte != b',' && byte != b';') +} + fn tool_definition(name: &str, title: &str, description: &str, input_schema: Value) -> Value { json!({ "name": name, diff --git a/apps/mcp-server/src/main.rs b/apps/mcp-server/src/main.rs index f0fbbfc..23a9a26 100644 --- a/apps/mcp-server/src/main.rs +++ b/apps/mcp-server/src/main.rs @@ -240,6 +240,176 @@ mod tests { assert!(keys[0].api_key.last_used_at.is_some()); } + #[tokio::test] + async fn preserves_request_id_for_tool_call_invocations() { + let registry = test_registry().await; + let upstream_base_url = spawn_upstream_server().await; + let operation = test_operation(&upstream_base_url, "crm_request_id"); + + registry + .create_operation(&test_workspace_id(), &operation, Some("alice")) + .await + .unwrap(); + registry + .publish_operation(PublishRequest { + workspace_id: &test_workspace_id(), + operation_id: &operation.id, + version: 1, + published_at: &OffsetDateTime::parse("2026-03-26T10:00:00Z", &Rfc3339).unwrap(), + published_by: Some("alice"), + }) + .await + .unwrap(); + publish_agent_for_operation(®istry, &operation, "sales-request-id").await; + let api_key = create_platform_api_key( + ®istry, + "mcp-request-id", + &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write], + ) + .await; + + let base_url = spawn_mcp_server(build_test_app( + registry.clone(), + Duration::from_millis(0), + Some("https://crank.example.com".to_owned()), + )) + .await; + let client = reqwest::Client::new(); + let mcp_url = agent_mcp_url(&base_url, "sales-request-id"); + let initialized_session = initialize_session(&client, &mcp_url, &api_key).await; + + let response = post_jsonrpc_response( + &client, + &mcp_url, + &api_key, + Some(&initialized_session), + Some("req_test_123"), + json!({ + "jsonrpc": "2.0", + "id": 3, + "method": "tools/call", + "params": { + "name": "crm_request_id", + "arguments": { + "email": "user@example.com" + } + } + }), + ) + .await; + + assert_eq!( + response.headers()["x-request-id"].to_str().unwrap(), + "req_test_123" + ); + let call_result = response.json::().await.unwrap(); + assert_eq!(call_result["result"]["isError"], false); + + let logs = registry + .list_invocation_logs(ListInvocationLogsQuery { + workspace_id: &test_workspace_id(), + level: None, + search_text: None, + source: Some(crank_core::InvocationSource::AgentToolCall), + operation_id: Some(&operation.id), + agent_id: None, + created_after: None, + limit: 10, + }) + .await + .unwrap(); + + assert_eq!(logs.len(), 1); + assert_eq!(logs[0].log.request_id.as_deref(), Some("req_test_123")); + } + + #[tokio::test] + async fn generates_request_id_for_tool_call_responses_and_logs() { + let registry = test_registry().await; + let upstream_base_url = spawn_upstream_server().await; + let operation = test_operation(&upstream_base_url, "crm_generated_request_id"); + + registry + .create_operation(&test_workspace_id(), &operation, Some("alice")) + .await + .unwrap(); + registry + .publish_operation(PublishRequest { + workspace_id: &test_workspace_id(), + operation_id: &operation.id, + version: 1, + published_at: &OffsetDateTime::parse("2026-03-26T10:00:00Z", &Rfc3339).unwrap(), + published_by: Some("alice"), + }) + .await + .unwrap(); + publish_agent_for_operation(®istry, &operation, "sales-generated-request-id").await; + let api_key = create_platform_api_key( + ®istry, + "mcp-generated-request-id", + &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write], + ) + .await; + + let base_url = spawn_mcp_server(build_test_app( + registry.clone(), + Duration::from_millis(0), + Some("https://crank.example.com".to_owned()), + )) + .await; + let client = reqwest::Client::new(); + let mcp_url = agent_mcp_url(&base_url, "sales-generated-request-id"); + let initialized_session = initialize_session(&client, &mcp_url, &api_key).await; + + let response = post_jsonrpc_response( + &client, + &mcp_url, + &api_key, + Some(&initialized_session), + None, + json!({ + "jsonrpc": "2.0", + "id": 3, + "method": "tools/call", + "params": { + "name": "crm_generated_request_id", + "arguments": { + "email": "user@example.com" + } + } + }), + ) + .await; + let request_id = response + .headers() + .get("x-request-id") + .unwrap() + .to_str() + .unwrap() + .to_owned(); + assert!(!request_id.is_empty()); + + let call_result = response.json::().await.unwrap(); + assert_eq!(call_result["result"]["isError"], false); + + let logs = registry + .list_invocation_logs(ListInvocationLogsQuery { + workspace_id: &test_workspace_id(), + level: None, + search_text: None, + source: Some(crank_core::InvocationSource::AgentToolCall), + operation_id: Some(&operation.id), + agent_id: None, + created_after: None, + limit: 10, + }) + .await + .unwrap(); + + assert_eq!(logs.len(), 1); + assert_eq!(logs[0].log.request_id.as_deref(), Some(request_id.as_str())); + } + #[tokio::test] async fn initializes_and_calls_published_graphql_tool_via_mcp() { let registry = test_registry().await; @@ -1514,6 +1684,21 @@ mod tests { session_id: Option<&str>, payload: Value, ) -> Value { + post_jsonrpc_response(client, mcp_url, api_key, session_id, None, payload) + .await + .json::() + .await + .unwrap() + } + + async fn post_jsonrpc_response( + client: &reqwest::Client, + mcp_url: &str, + api_key: &str, + session_id: Option<&str>, + request_id: Option<&str>, + payload: Value, + ) -> reqwest::Response { let mut request = client .post(mcp_url) .header(header::ACCEPT, "application/json, text/event-stream") @@ -1524,14 +1709,11 @@ mod tests { request = request.header("MCP-Session-Id", session_id); } - request - .json(&payload) - .send() - .await - .unwrap() - .json::() - .await - .unwrap() + if let Some(request_id) = request_id { + request = request.header("x-request-id", request_id); + } + + request.json(&payload).send().await.unwrap() } async fn create_platform_api_key(