mirror of
https://github.com/openai/codex.git
synced 2026-11-03 16:23:51 +00:00
[rollout-trace] Add x-codex-inference-call-id header to inference calls. (#22311)
This allows us to attach call logs to inference requests in traces.
This commit is contained in:
1 parent
d996f5366f
commit
cb55b769d1
4 files changed
+72
-6
No files matched your search
Generated
+1
@@ -3514,6 +3514,7 @@ dependencies = [
|
|||||||
"anyhow",
|
"anyhow",
|
||||||
"codex-code-mode",
|
"codex-code-mode",
|
||||||
"codex-protocol",
|
"codex-protocol",
|
||||||
|
"http 1.4.0",
|
||||||
"pretty_assertions",
|
"pretty_assertions",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
|
|||||||
@@ -1256,7 +1256,7 @@ impl ModelClientSession {
|
|||||||
self.client.state.auth_env_telemetry.clone(),
|
self.client.state.auth_env_telemetry.clone(),
|
||||||
);
|
);
|
||||||
let compression = self.responses_request_compression(client_setup.auth.as_ref());
|
let compression = self.responses_request_compression(client_setup.auth.as_ref());
|
||||||
let options = self
|
let mut options = self
|
||||||
.build_responses_options(turn_metadata_header, compression)
|
.build_responses_options(turn_metadata_header, compression)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
@@ -1269,6 +1269,7 @@ impl ModelClientSession {
|
|||||||
service_tier.clone(),
|
service_tier.clone(),
|
||||||
)?;
|
)?;
|
||||||
let inference_trace_attempt = inference_trace.start_attempt();
|
let inference_trace_attempt = inference_trace.start_attempt();
|
||||||
|
inference_trace_attempt.add_request_headers(&mut options.extra_headers);
|
||||||
inference_trace_attempt.record_started(&request);
|
inference_trace_attempt.record_started(&request);
|
||||||
let client = ApiResponsesClient::new(
|
let client = ApiResponsesClient::new(
|
||||||
transport,
|
transport,
|
||||||
|
|||||||
@@ -16,10 +16,11 @@ workspace = true
|
|||||||
anyhow = { workspace = true }
|
anyhow = { workspace = true }
|
||||||
codex-code-mode = { workspace = true }
|
codex-code-mode = { workspace = true }
|
||||||
codex-protocol = { workspace = true }
|
codex-protocol = { workspace = true }
|
||||||
|
http = { workspace = true }
|
||||||
serde = { workspace = true, features = ["derive"] }
|
serde = { workspace = true, features = ["derive"] }
|
||||||
serde_json = { workspace = true }
|
serde_json = { workspace = true }
|
||||||
tracing = { workspace = true }
|
tracing = { workspace = true }
|
||||||
uuid = { workspace = true }
|
uuid = { workspace = true, features = ["v4"] }
|
||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
pretty_assertions = { workspace = true }
|
pretty_assertions = { workspace = true }
|
||||||
|
|||||||
@@ -7,13 +7,15 @@
|
|||||||
use std::fmt::Display;
|
use std::fmt::Display;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::sync::atomic::AtomicBool;
|
use std::sync::atomic::AtomicBool;
|
||||||
use std::sync::atomic::AtomicU64;
|
|
||||||
use std::sync::atomic::Ordering;
|
use std::sync::atomic::Ordering;
|
||||||
|
|
||||||
use codex_protocol::models::ResponseItem;
|
use codex_protocol::models::ResponseItem;
|
||||||
use codex_protocol::protocol::TokenUsage;
|
use codex_protocol::protocol::TokenUsage;
|
||||||
|
use http::HeaderMap;
|
||||||
|
use http::HeaderValue;
|
||||||
use serde::Serialize;
|
use serde::Serialize;
|
||||||
use serde_json::Value as JsonValue;
|
use serde_json::Value as JsonValue;
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::model::AgentThreadId;
|
use crate::model::AgentThreadId;
|
||||||
use crate::model::CodexTurnId;
|
use crate::model::CodexTurnId;
|
||||||
@@ -23,7 +25,7 @@ use crate::raw_event::RawTraceEventContext;
|
|||||||
use crate::raw_event::RawTraceEventPayload;
|
use crate::raw_event::RawTraceEventPayload;
|
||||||
use crate::writer::TraceWriter;
|
use crate::writer::TraceWriter;
|
||||||
|
|
||||||
static NEXT_INFERENCE_ATTEMPT: AtomicU64 = AtomicU64::new(1);
|
const INFERENCE_CALL_ID_HEADER: &str = "x-codex-inference-call-id";
|
||||||
|
|
||||||
/// Turn-local inference tracing context.
|
/// Turn-local inference tracing context.
|
||||||
///
|
///
|
||||||
@@ -140,6 +142,30 @@ impl InferenceTraceAttempt {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn inference_call_id(&self) -> Option<&str> {
|
||||||
|
match &self.state {
|
||||||
|
InferenceTraceAttemptState::Disabled => None,
|
||||||
|
InferenceTraceAttemptState::Enabled(attempt) => {
|
||||||
|
Some(attempt.inference_call_id.as_str())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Adds rollout-trace propagation headers for this attempt when tracing is enabled.
|
||||||
|
pub fn add_request_headers(&self, headers: &mut HeaderMap) {
|
||||||
|
let Some(inference_call_id) = self.inference_call_id() else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let Ok(inference_call_id) = HeaderValue::from_str(inference_call_id) else {
|
||||||
|
// These IDs are generated internally as UUID strings, so rejection
|
||||||
|
// should be impossible in practice. Tracing remains best-effort,
|
||||||
|
// though, and must never make provider requests fail.
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
|
||||||
|
headers.insert(INFERENCE_CALL_ID_HEADER, inference_call_id);
|
||||||
|
}
|
||||||
|
|
||||||
/// Records the exact request object about to be sent to the model provider.
|
/// Records the exact request object about to be sent to the model provider.
|
||||||
pub fn record_started(&self, request: &impl Serialize) {
|
pub fn record_started(&self, request: &impl Serialize) {
|
||||||
let InferenceTraceAttemptState::Enabled(attempt) = &self.state else {
|
let InferenceTraceAttemptState::Enabled(attempt) = &self.state else {
|
||||||
@@ -315,8 +341,7 @@ pub(crate) fn trace_response_item_json(item: &ResponseItem) -> JsonValue {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn next_inference_call_id() -> InferenceCallId {
|
fn next_inference_call_id() -> InferenceCallId {
|
||||||
let ordinal = NEXT_INFERENCE_ATTEMPT.fetch_add(1, Ordering::Relaxed);
|
Uuid::new_v4().to_string()
|
||||||
format!("inference:{ordinal}")
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn write_json_payload_best_effort(
|
fn write_json_payload_best_effort(
|
||||||
@@ -372,6 +397,44 @@ mod tests {
|
|||||||
use crate::model::ExecutionStatus;
|
use crate::model::ExecutionStatus;
|
||||||
use crate::replay_bundle;
|
use crate::replay_bundle;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn disabled_attempt_adds_no_request_headers() {
|
||||||
|
let mut headers = HeaderMap::new();
|
||||||
|
|
||||||
|
InferenceTraceAttempt::disabled().add_request_headers(&mut headers);
|
||||||
|
|
||||||
|
assert!(headers.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn enabled_attempt_adds_inference_request_header() -> anyhow::Result<()> {
|
||||||
|
let temp = TempDir::new()?;
|
||||||
|
let writer = Arc::new(TraceWriter::create(
|
||||||
|
temp.path(),
|
||||||
|
"trace-1".to_string(),
|
||||||
|
"rollout-1".to_string(),
|
||||||
|
"thread-root".to_string(),
|
||||||
|
)?);
|
||||||
|
let context = InferenceTraceContext::enabled(
|
||||||
|
writer,
|
||||||
|
"thread-root".to_string(),
|
||||||
|
"turn-1".to_string(),
|
||||||
|
"gpt-test".to_string(),
|
||||||
|
"test-provider".to_string(),
|
||||||
|
);
|
||||||
|
let attempt = context.start_attempt();
|
||||||
|
let mut headers = HeaderMap::new();
|
||||||
|
|
||||||
|
attempt.add_request_headers(&mut headers);
|
||||||
|
|
||||||
|
let header = headers
|
||||||
|
.get(INFERENCE_CALL_ID_HEADER)
|
||||||
|
.expect("inference header present");
|
||||||
|
assert_eq!(Some(header.to_str()?), attempt.inference_call_id());
|
||||||
|
assert!(Uuid::parse_str(header.to_str()?).is_ok());
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn enabled_context_records_replayable_inference_attempt() -> anyhow::Result<()> {
|
fn enabled_context_records_replayable_inference_attempt() -> anyhow::Result<()> {
|
||||||
let temp = TempDir::new()?;
|
let temp = TempDir::new()?;
|
||||||
|
|||||||
Reference in new issue
Block a user