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,
apps/mcp-server/src/main.rs
+190 -8
View File
@@ -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
.json(&payload)
.send()
.await
.unwrap()
.json::<Value>()
.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(