mcp: propagate request ids through tool invocations

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<Value>,
 ) -> 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<S>(
     status: StatusCode,
     stream: S,
@@ -2219,6 +2253,24 @@ fn tool_definitions(tool: &PublishedAgentTool) -> Vec<Value> {
     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,

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(&registry, &operation, "sales-request-id").await;
+        let api_key = create_platform_api_key(
+            &registry,
+            "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::<Value>().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(&registry, &operation, "sales-generated-request-id").await;
+        let api_key = create_platform_api_key(
+            &registry,
+            "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::<Value>().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::<Value>()
+            .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
+        if let Some(request_id) = request_id {
-            .json(&payload)
+            request = request.header("x-request-id", request_id);
-            .send()
+        }
-            .await
+
-            .unwrap()
+        request.json(&payload).send().await.unwrap()
-            .json::<Value>()
-            .await
-            .unwrap()
     }
 
     async fn create_platform_api_key(