mcp: propagate request ids through tool invocations
This commit is contained in:
a.tolmachev
+190
-8
@@ -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::<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(®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::<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(
|
||||
|
||||
Reference in New Issue
Block a user