bsodfather/crank
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

mcp: propagate request ids through tool invocations

Browse Source
This commit is contained in:
a.tolmachev
2026-05-01 11:31:12 +00:00
parent 88033a3b3e
commit 3b1a7c6993
2 changed files with 252 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/mcp-server/src/app.rs
+62 -10
View File
@@ -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,