api: propagate request ids through admin test runs

This commit is contained in:
a.tolmachev
2026-04-19 22:10:29 +00:00
parent f3f6a3b702
commit 88033a3b3e
5 changed files with 188 additions and 2 deletions
+109
View File
@@ -5,6 +5,7 @@ use axum::{
use crate::{
auth::{require_session, require_workspace_session},
request_context::apply_request_context,
routes::{
access::{
create_invitation, create_platform_api_key, delete_invitation, delete_membership,
@@ -205,6 +206,7 @@ pub fn build_app(state: AppState) -> Router {
.merge(protected_auth_router),
)
.nest("/api/admin", admin_router)
.layer(middleware::from_fn(apply_request_context))
.with_state(state)
}
@@ -1306,6 +1308,113 @@ mod tests {
assert_eq!(operation_usage["rollup"]["calls_total"], 1);
}
#[tokio::test(flavor = "multi_thread")]
#[serial]
async fn preserves_request_id_for_test_run_invocations() {
let registry = test_registry().await;
let storage_root = test_storage_root("observability_request_id");
let upstream_base_url = spawn_upstream_server().await;
let base_url = spawn_admin_api(build_test_app(registry, storage_root)).await;
let client = authorized_client(&base_url).await;
let created = client
.post(format!("{base_url}/operations"))
.json(&test_operation_payload(
&upstream_base_url,
"crm_observability_request_id",
))
.send()
.await
.unwrap()
.json::<Value>()
.await
.unwrap();
let operation_id = created["operation_id"].as_str().unwrap().to_owned();
let response = client
.post(format!("{base_url}/operations/{operation_id}/test-runs"))
.header("x-request-id", "req_test_123")
.json(&json!({
"version": 1,
"input": { "email": "user@example.com" }
}))
.send()
.await
.unwrap();
assert_eq!(
response.headers()["x-request-id"].to_str().unwrap(),
"req_test_123"
);
response.error_for_status().unwrap();
let logs = client
.get(format!("{base_url}/logs?period=7d"))
.send()
.await
.unwrap()
.json::<Value>()
.await
.unwrap();
assert_eq!(logs["items"][0]["log"]["request_id"], "req_test_123");
}
#[tokio::test(flavor = "multi_thread")]
#[serial]
async fn generates_request_id_for_test_run_invocations() {
let registry = test_registry().await;
let storage_root = test_storage_root("observability_generated_request_id");
let upstream_base_url = spawn_upstream_server().await;
let base_url = spawn_admin_api(build_test_app(registry, storage_root)).await;
let client = authorized_client(&base_url).await;
let created = client
.post(format!("{base_url}/operations"))
.json(&test_operation_payload(
&upstream_base_url,
"crm_observability_generated_request_id",
))
.send()
.await
.unwrap()
.json::<Value>()
.await
.unwrap();
let operation_id = created["operation_id"].as_str().unwrap().to_owned();
let response = client
.post(format!("{base_url}/operations/{operation_id}/test-runs"))
.json(&json!({
"version": 1,
"input": { "email": "user@example.com" }
}))
.send()
.await
.unwrap();
let request_id = response
.headers()
.get("x-request-id")
.unwrap()
.to_str()
.unwrap()
.to_owned();
assert!(!request_id.is_empty());
response.error_for_status().unwrap();
let logs = client
.get(format!("{base_url}/logs?period=7d"))
.send()
.await
.unwrap()
.json::<Value>()
.await
.unwrap();
assert_eq!(logs["items"][0]["log"]["request_id"], request_id);
}
#[tokio::test(flavor = "multi_thread")]
#[serial]
async fn creates_publishes_and_tests_graphql_operation() {
+1
View File
@@ -1,6 +1,7 @@
mod app;
mod auth;
mod error;
mod request_context;
mod routes;
mod service;
mod state;
+64
View File
@@ -0,0 +1,64 @@
use axum::{
extract::Request,
http::{HeaderMap, HeaderName, HeaderValue},
middleware::Next,
response::Response,
};
use uuid::Uuid;
pub const REQUEST_ID_HEADER: HeaderName = HeaderName::from_static("x-request-id");
const MAX_REQUEST_ID_LEN: usize = 128;
#[derive(Clone, Debug)]
pub struct RequestContext {
pub request_id: String,
}
pub async fn apply_request_context(mut request: Request, next: Next) -> Response {
let context = RequestContext {
request_id: resolve_request_id(request.headers()),
};
request.extensions_mut().insert(context.clone());
let mut response = next.run(request).await;
if let Ok(value) = HeaderValue::from_str(&context.request_id) {
response.headers_mut().insert(REQUEST_ID_HEADER, value);
}
response
}
fn resolve_request_id(headers: &HeaderMap) -> String {
headers
.get(&REQUEST_ID_HEADER)
.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::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';')
}
#[cfg(test)]
mod tests {
use super::is_valid_request_id;
#[test]
fn accepts_visible_ascii_request_ids() {
assert!(is_valid_request_id("req_test_123"));
assert!(is_valid_request_id("trace-123/abc"));
}
#[test]
fn rejects_empty_or_control_request_ids() {
assert!(!is_valid_request_id(""));
assert!(!is_valid_request_id("bad value"));
assert!(!is_valid_request_id("bad\nvalue"));
}
}
+4 -1
View File
@@ -1,7 +1,7 @@
use axum::{
Json,
body::Bytes,
extract::{Path, Query, State},
extract::{Extension, Path, Query, State},
http::{HeaderMap, StatusCode, header},
response::IntoResponse,
};
@@ -10,6 +10,7 @@ use serde_json::{Value, json};
use crate::{
error::ApiError,
request_context::RequestContext,
service::{
ExportQuery, GenerateDraftPayload, ImportQuery, NewVersionPayload, OperationPayload,
PublishPayload, TestRunPayload, UpdateOperationPayload,
@@ -172,6 +173,7 @@ pub async fn archive_operation(
pub async fn run_test(
Path(path): Path<WorkspaceOperationPath>,
State(state): State<AppState>,
Extension(request_context): Extension<RequestContext>,
Json(payload): Json<TestRunPayload>,
) -> Result<Json<Value>, ApiError> {
let result = state
@@ -180,6 +182,7 @@ pub async fn run_test(
&path.workspace_id.as_str().into(),
&path.operation_id.as_str().into(),
payload,
&request_context.request_id,
)
.await?;
Ok(Json(json!(result)))
+10 -1
View File
@@ -594,6 +594,7 @@ struct InvocationRecordRequest<'a> {
workspace_id: &'a WorkspaceId,
agent_id: Option<&'a AgentId>,
operation: &'a RegistryOperation,
request_id: Option<&'a str>,
source: InvocationSource,
level: InvocationLevel,
status: InvocationStatus,
@@ -2093,6 +2094,7 @@ impl AdminService {
workspace_id: &WorkspaceId,
operation_id: &OperationId,
payload: TestRunPayload,
request_id: &str,
) -> Result<TestRunResult, ApiError> {
let record = self
.get_operation_version(workspace_id, operation_id, payload.version)
@@ -2112,6 +2114,7 @@ impl AdminService {
workspace_id,
agent_id: None,
operation: &record.snapshot,
request_id: Some(request_id),
source: InvocationSource::AdminTestRun,
level: InvocationLevel::Error,
status: InvocationStatus::Error,
@@ -2180,6 +2183,7 @@ impl AdminService {
workspace_id,
agent_id: None,
operation: &record.snapshot,
request_id: Some(request_id),
source: InvocationSource::AdminTestRun,
level: InvocationLevel::Info,
status: InvocationStatus::Ok,
@@ -2209,6 +2213,7 @@ impl AdminService {
workspace_id,
agent_id: None,
operation: &record.snapshot,
request_id: Some(request_id),
source: InvocationSource::AdminTestRun,
level: InvocationLevel::Error,
status: InvocationStatus::Error,
@@ -4194,6 +4199,7 @@ impl AdminService {
workspace_id,
agent_id: Some(revops_agent_id),
operation: &rest_operation.snapshot,
request_id: None,
source: InvocationSource::AgentToolCall,
level: InvocationLevel::Info,
status: InvocationStatus::Ok,
@@ -4216,6 +4222,7 @@ impl AdminService {
workspace_id,
agent_id: Some(revops_agent_id),
operation: &graphql_operation.snapshot,
request_id: None,
source: InvocationSource::AgentToolCall,
level: InvocationLevel::Info,
status: InvocationStatus::Ok,
@@ -4238,6 +4245,7 @@ impl AdminService {
workspace_id,
agent_id: Some(revops_agent_id),
operation: &graphql_operation.snapshot,
request_id: None,
source: InvocationSource::AgentToolCall,
level: InvocationLevel::Error,
status: InvocationStatus::Error,
@@ -4262,6 +4270,7 @@ impl AdminService {
workspace_id,
agent_id: None,
operation: &grpc_operation.snapshot,
request_id: None,
source: InvocationSource::AdminTestRun,
level: InvocationLevel::Info,
status: InvocationStatus::Ok,
@@ -4298,7 +4307,7 @@ impl AdminService {
status: request.status,
tool_name: request.operation.name.clone(),
message: request.message,
request_id: None,
request_id: request.request_id.map(ToOwned::to_owned),
status_code: request.status_code,
duration_ms: request.duration_ms,
error_kind: request.error_kind,