From beca1a733e3225ee491f534c2b4f4a406ceccaba Mon Sep 17 00:00:00 2001 From: reonokiy Date: Sat, 10 Oct 2026 12:20:17 +0000 Subject: [PATCH] feat: export per-request usage through OpenTelemetry Why: Attribute model token usage to verified users and expose usage details in Grafana. How: Export usage logs, spans and GenAI metrics over OTLP; consume trusted generic identity headers and track HTTP and WebSocket inference outcomes. --- Cargo.lock | 3 + Cargo.toml | 4 + docs/usage-telemetry.md | 40 +++ src/headers.rs | 6 + src/lib.rs | 1 + src/main.rs | 14 + src/server.rs | 71 +++- src/telemetry.rs | 589 ++++++++++++++++++++++++++++++++++ src/websocket.rs | 38 ++- tests/catalog.rs | 1 + tests/gateway.rs | 1 + tests/support/live_gateway.rs | 1 + 12 files changed, 756 insertions(+), 13 deletions(-) create mode 100644 docs/usage-telemetry.md create mode 100644 src/telemetry.rs diff --git a/Cargo.lock b/Cargo.lock index 744e420..fd3e6d0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1546,6 +1546,9 @@ dependencies = [ "eventsource-stream", "futures", "http 1.4.0", + "opentelemetry", + "opentelemetry-otlp", + "opentelemetry_sdk", "reqwest", "serde", "serde_ignored", diff --git a/Cargo.toml b/Cargo.toml index 07a03c4..c23498c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,6 +26,9 @@ tokio = { version = "1", features = ["full"] } toml = "0.9.5" tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } +opentelemetry = "0.31" +opentelemetry_sdk = "0.31" +opentelemetry-otlp = { version = "0.31", default-features = false, features = ["trace", "metrics", "logs", "http-proto", "reqwest-blocking-client", "reqwest-rustls"] } uuid = { version = "1", features = ["v4", "v5"] } codex-api = { git = "https://github.com/openai/codex.git", rev = "092d3acd6bec3e3a14bdc7e7a2810ab628ab759d" } codex-client = { git = "https://github.com/openai/codex.git", rev = "092d3acd6bec3e3a14bdc7e7a2810ab628ab759d" } @@ -45,6 +48,7 @@ codex-tools = { git = "https://github.com/openai/codex.git", rev = "092d3acd6bec codex-exec-server = { git = "https://github.com/openai/codex.git", rev = "092d3acd6bec3e3a14bdc7e7a2810ab628ab759d" } [dev-dependencies] +opentelemetry_sdk = { version = "0.31", features = ["testing"] } reqwest = { version = "0.12", features = ["json", "stream", "multipart"] } tempfile = "3" diff --git a/docs/usage-telemetry.md b/docs/usage-telemetry.md new file mode 100644 index 0000000..4660fbc --- /dev/null +++ b/docs/usage-telemetry.md @@ -0,0 +1,40 @@ +# Gateway usage telemetry + +Set `OTEL_EXPORTER_OTLP_ENDPOINT=http://alloy.observability.svc.cluster.local:4318` +and `OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf` to enable outbound OTLP logs, +metrics and traces. `OTEL_SERVICE_NAME` defaults to `codex-api`; standard OTLP +endpoint, headers and timeout variables are supported. `OTEL_SDK_DISABLED=true` +disables collection. This is independent of the existing `/telemetry/metrics` +proxy route. Queued telemetry is best effort, not a billing ledger. + +Set `CODEX_GATEWAY_TRUST_IDENTITY_HEADERS=true` only behind a trusted ingress. +Envoy maps verified `X-Keygate-User-Id` and `X-Keygate-Key-Id` authorization +response headers to `X-User-Id` and `X-Api-Key-Id`, overwriting client values. +These are conventional custom headers, not IETF standardized identity headers. +Internal routes overwrite the user with +`internal` and remove the key ID. Direct service access must be restricted; +without trusted identity a request is attributed to `unknown`. These headers +are never forwarded to the model provider. + +Each HTTP inference (including streaming and compact requests) and each +WebSocket response.create generates an OTLP usage log and correlated span. +Logs include user/key IDs, request/response IDs, requested and returned model, +start time, duration, outcome, input/output/total tokens and cache/reasoning +details when supplied by the upstream. Missing usage remains missing, rather +than being reported as zero. Cache and reasoning counts are subsets of input +and output respectively; do not add them again to total tokens. Prompts, +responses, tool arguments and credentials are not collected. + +Metrics use the development GenAI inference conventions: +`gen_ai.client.inference.usage.input_tokens`, `output_tokens`, +`cache_read.input_tokens`, `reasoning.output_tokens`, and +`gen_ai.client.inference.duration`. User and request IDs are excluded from +metric dimensions. `codex_api.usage.missing` counts requests without usage. +See the [upstream convention](https://github.com/open-telemetry/semantic-conventions-genai/blob/main/docs/gen-ai/client-inference.md). + +The Talos configuration routes OTLP through Alloy to Mimir, Tempo and Loki. +The provisioned `codex-api-usage` Grafana dashboard uses Loki usage events for +user/model filters, token totals and per-request details, and Mimir for rates +and latency. Request details are limited by Loki retention (currently 14 days) +and the query row limit. Deploy the gateway image containing this change +before enabling its telemetry environment variables. diff --git a/src/headers.rs b/src/headers.rs index 08865f6..7d8a11d 100644 --- a/src/headers.rs +++ b/src/headers.rs @@ -53,6 +53,12 @@ pub(crate) fn client_header(name: &str) -> bool { | "x-api-key" | "api-key" | "x-codex-gateway-authorization" + | "x-keygate-user-id" + | "x-keygate-key-id" + | "x-user-id" + | "x-api-key-id" + | "x-keygate-app" + | "x-auth-request-user" | "sec-websocket-key" | "sec-websocket-version" | "sec-websocket-accept" diff --git a/src/lib.rs b/src/lib.rs index acbe5ce..d93f864 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -13,6 +13,7 @@ mod search; pub mod server; pub mod settings; pub mod standalone; +pub mod telemetry; mod tools; pub mod transfers; pub mod transport; diff --git a/src/main.rs b/src/main.rs index 38928de..2466ce9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -57,7 +57,10 @@ async fn main() -> anyhow::Result<()> { codex_api_gateway::transfers::Transfers::new(backend.factory.clone(), &url, timeout) }) .transpose()?; + let telemetry = + tokio::task::spawn_blocking(codex_api_gateway::telemetry::Telemetry::from_env).await??; let gateway = Gateway { + telemetry: telemetry.clone(), backend, models, key, @@ -71,8 +74,19 @@ async fn main() -> anyhow::Result<()> { tracing::info!(listen=%listener.local_addr()?, "Codex Responses gateway ready"); axum::serve(listener, router(gateway)) .with_graceful_shutdown(async { + #[cfg(unix)] + { + let mut terminate = + tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + .expect("install SIGTERM handler"); + tokio::select! { _ = tokio::signal::ctrl_c() => {}, _ = terminate.recv() => {} } + } + #[cfg(not(unix))] let _ = tokio::signal::ctrl_c().await; }) .await?; + if let Some(telemetry) = telemetry { + tokio::task::spawn_blocking(move || telemetry.shutdown()).await?; + } Ok(()) } diff --git a/src/server.rs b/src/server.rs index 10b10e4..5d6e8c2 100644 --- a/src/server.rs +++ b/src/server.rs @@ -20,6 +20,7 @@ pub struct Gateway { pub concurrency: Arc, pub timeout: Duration, pub transfers: Option, + pub telemetry: Option>, } pub fn router(gateway: Gateway) -> Router { @@ -184,6 +185,7 @@ async fn create( ) })?; let stream = input.stream; + let telemetry = crate::telemetry::RequestTelemetry::new(gateway.telemetry.clone(), &headers); let headers = crate::headers::request_headers(&headers); let session_id = headers .get("thread-id") @@ -193,6 +195,7 @@ async fn create( .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); let lite = input.uses_responses_lite(model); let request = input.into_codex(model, &session_id)?; + let usage = telemetry.start(&request.model, "responses", "http"); if lite { let mut headers = forward_request_headers(&headers); headers.insert( @@ -212,10 +215,19 @@ async fn create( )), stream, true, + usage, ) .await } else { - respond(gateway, Some((request, session_id)), None, stream, true).await + respond( + gateway, + Some((request, session_id)), + None, + stream, + true, + usage, + ) + .await } } @@ -307,7 +319,12 @@ async fn native_create( )); } let forwarded = forward_request_headers(&headers); - respond(gateway, None, Some((raw, forwarded)), true, false).await + let usage = crate::telemetry::RequestTelemetry::new(gateway.telemetry.clone(), &headers).start( + body["model"].as_str().unwrap_or("unknown"), + "responses", + "http", + ); + respond(gateway, None, Some((raw, forwarded)), true, false, usage).await } pub(crate) fn forward_request_headers(headers: &HeaderMap) -> HeaderMap { @@ -349,16 +366,20 @@ async fn respond( native: Option<(Box, HeaderMap)>, stream: bool, aggregate_output: bool, + mut usage: crate::telemetry::UsageGuard, ) -> Result { let permit = gateway .concurrency .clone() .try_acquire_owned() - .map_err(|_| GatewayError { - upstream_response: None, - status: StatusCode::TOO_MANY_REQUESTS, - code: "gateway_busy", - message: "gateway concurrency limit reached".into(), + .map_err(|_| { + usage.failure("gateway_busy"); + GatewayError { + upstream_response: None, + status: StatusCode::TOO_MANY_REQUESTS, + code: "gateway_busy", + message: "gateway concurrency limit reached".into(), + } })?; let deadline = tokio::time::Instant::now() + gateway.timeout; let start = async { @@ -371,7 +392,14 @@ async fn respond( }; let mut running = tokio::time::timeout_at(deadline, start) .await - .map_err(|_| timeout_error())??; + .map_err(|_| { + usage.failure("upstream_timeout"); + timeout_error() + })? + .map_err(|e| { + usage.failure(e.code); + e + })?; let upstream_headers = forward_response_headers(&running.headers); let mut accumulator = crate::output::OutputAccumulator::default(); if stream { @@ -383,6 +411,7 @@ async fn respond( loop { match tokio::time::timeout_at(deadline, running.events.recv()).await { Ok(Some(mut event)) => { + usage.observe(&event.value); if aggregate_output && accumulator.observe(&mut event.value) { event.bytes = Bytes::from(format!("event: {}\ndata: {}\n\n",event.value["type"].as_str().unwrap_or("message"),event.value)); } @@ -404,6 +433,7 @@ async fn respond( match terminal { Some(event) if !event.successful() || parsed.is_ok() => { yield Ok(event.bytes); } _ => { + usage.failure(if timed_out { "upstream_timeout" } else { "upstream_error" }); let event = json!({"type":"error", "sequence_number":sequence, "code":"upstream_error", "message":parsed.err().unwrap_or_else(|| "upstream closed without a terminal response".into()), "param":null}); yield Ok(Bytes::from(format!("event: error\ndata: {event}\n\n"))); } @@ -426,10 +456,15 @@ async fn respond( let terminal = loop { let mut event = tokio::time::timeout_at(deadline, running.events.recv()) .await - .map_err(|_| timeout_error())? + .map_err(|_| { + usage.failure("upstream_timeout"); + timeout_error() + })? .ok_or_else(|| { + usage.failure("upstream_error"); GatewayError::internal("upstream closed without terminal response") })?; + usage.observe(&event.value); if aggregate_output { accumulator.observe(&mut event.value); } @@ -439,10 +474,19 @@ async fn respond( }; let parsed = tokio::time::timeout_at(deadline, &mut running.finished) .await - .map_err(|_| timeout_error())? - .map_err(GatewayError::internal)?; + .map_err(|_| { + usage.failure("upstream_timeout"); + timeout_error() + })? + .map_err(|e| { + usage.failure("upstream_error"); + GatewayError::internal(e) + })?; if terminal.successful() { - parsed.map_err(GatewayError::internal)?; + parsed.map_err(|e| { + usage.failure("upstream_error"); + GatewayError::internal(e) + })?; } let response = terminal .value @@ -489,6 +533,7 @@ async fn compact( body: Result, JsonRejection>, ) -> Result { authorize(&gateway, &headers)?; + let telemetry = crate::telemetry::RequestTelemetry::new(gateway.telemetry.clone(), &headers); let headers = crate::headers::request_headers(&headers); let Json(input) = body.map_err(|e| GatewayError::invalid(e.body_text()))?; let model = gateway @@ -504,6 +549,7 @@ async fn compact( .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); let lite = input.uses_responses_lite(model); let mut request = input.into_codex(model, &session_id)?; + let usage = telemetry.start(&request.model, "compact", "http"); request .input .push(codex_protocol::models::ResponseItem::CompactionTrigger {}); @@ -527,6 +573,7 @@ async fn compact( )), false, true, + usage, ) .await?; let bytes = axum::body::to_bytes(result.into_body(), 64 * 1024 * 1024) diff --git a/src/telemetry.rs b/src/telemetry.rs new file mode 100644 index 0000000..ad8a750 --- /dev/null +++ b/src/telemetry.rs @@ -0,0 +1,589 @@ +//! Usage-only OpenTelemetry. Never collect prompts, tool arguments, or credentials. +use http::HeaderMap; +use opentelemetry::{ + Context, KeyValue, + logs::{AnyValue, LogRecord, Logger, LoggerProvider, Severity}, + metrics::{Counter, Histogram, MeterProvider}, + propagation::{Extractor, TextMapPropagator}, + trace::{Span, SpanKind, Status, Tracer, TracerProvider}, +}; +use opentelemetry_sdk::{ + Resource, logs::SdkLoggerProvider, metrics::SdkMeterProvider, + propagation::TraceContextPropagator, trace::SdkTracerProvider, +}; +use serde_json::{Value, json}; +use std::{ + collections::{HashMap, HashSet, VecDeque}, + sync::Arc, + time::{Instant, SystemTime, UNIX_EPOCH}, +}; + +pub struct Telemetry { + traces: SdkTracerProvider, + logs: SdkLoggerProvider, + metrics: SdkMeterProvider, + input: Counter, + output: Counter, + cached: Counter, + reasoning: Counter, + duration: Histogram, + missing: Counter, + trust_identity_headers: bool, +} + +impl Telemetry { + /// Enable only when an OTLP endpoint is explicitly configured. Exporters use + /// standard OTEL_EXPORTER_OTLP_* endpoint, header, and timeout variables. + pub fn from_env() -> anyhow::Result>> { + if std::env::var("OTEL_SDK_DISABLED").is_ok_and(|v| v.eq_ignore_ascii_case("true")) + || ![ + "OTEL_EXPORTER_OTLP_ENDPOINT", + "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", + "OTEL_EXPORTER_OTLP_METRICS_ENDPOINT", + "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT", + ] + .iter() + .any(|name| std::env::var(name).is_ok_and(|v| !v.is_empty())) + { + return Ok(None); + } + for name in [ + "OTEL_EXPORTER_OTLP_PROTOCOL", + "OTEL_EXPORTER_OTLP_TRACES_PROTOCOL", + "OTEL_EXPORTER_OTLP_METRICS_PROTOCOL", + "OTEL_EXPORTER_OTLP_LOGS_PROTOCOL", + ] { + if std::env::var(name).is_ok_and(|v| v != "http/protobuf") { + anyhow::bail!("{name} must be http/protobuf"); + } + } + let resource = Resource::builder() + .with_service_name( + std::env::var("OTEL_SERVICE_NAME").unwrap_or_else(|_| "codex-api".into()), + ) + .with_attributes([KeyValue::new("service.version", crate::CODEX_RELEASE)]) + .build(); + let traces = SdkTracerProvider::builder() + .with_resource(resource.clone()) + .with_batch_exporter( + opentelemetry_otlp::SpanExporter::builder() + .with_http() + .build()?, + ) + .build(); + let logs = SdkLoggerProvider::builder() + .with_resource(resource.clone()) + .with_batch_exporter( + opentelemetry_otlp::LogExporter::builder() + .with_http() + .build()?, + ) + .build(); + let metrics = SdkMeterProvider::builder() + .with_resource(resource) + .with_periodic_exporter( + opentelemetry_otlp::MetricExporter::builder() + .with_http() + .build()?, + ) + .build(); + let trust = std::env::var("CODEX_GATEWAY_TRUST_IDENTITY_HEADERS") + .is_ok_and(|v| v.eq_ignore_ascii_case("true")); + Ok(Some(Arc::new(Self::new(traces, logs, metrics, trust)))) + } + + /// Providers are explicit so applications and tests do not share global state. + pub fn new( + traces: SdkTracerProvider, + logs: SdkLoggerProvider, + metrics: SdkMeterProvider, + trust_identity_headers: bool, + ) -> Self { + let meter = metrics.meter("codex-api.usage"); + Self { + input: meter + .u64_counter("gen_ai.client.inference.usage.input_tokens") + .with_unit("{token}") + .build(), + output: meter + .u64_counter("gen_ai.client.inference.usage.output_tokens") + .with_unit("{token}") + .build(), + cached: meter + .u64_counter("gen_ai.client.inference.usage.cache_read.input_tokens") + .with_unit("{token}") + .build(), + reasoning: meter + .u64_counter("gen_ai.client.inference.usage.reasoning.output_tokens") + .with_unit("{token}") + .build(), + duration: meter + .f64_histogram("gen_ai.client.inference.duration") + .with_unit("s") + .with_boundaries(vec![ + 0.01, 0.02, 0.04, 0.08, 0.16, 0.32, 0.64, 1.28, 2.56, 5.12, 10.24, 20.48, + 40.96, 81.92, + ]) + .build(), + missing: meter + .u64_counter("codex_api.usage.missing") + .with_unit("{request}") + .build(), + traces, + logs, + metrics, + trust_identity_headers, + } + } + + pub fn shutdown(&self) { + // Workers have bounded queues; an unavailable collector must not stop inference. + let _ = self.logs.shutdown(); + let _ = self.metrics.shutdown(); + let _ = self.traces.shutdown(); + } +} + +struct HeaderExtractor<'a>(&'a HeaderMap); +impl Extractor for HeaderExtractor<'_> { + fn get(&self, key: &str) -> Option<&str> { + self.0.get(key)?.to_str().ok() + } + fn keys(&self) -> Vec<&str> { + self.0.keys().map(|k| k.as_str()).collect() + } +} + +fn label(value: Option<&str>) -> Option { + value + .filter(|v| { + !v.is_empty() + && v.len() <= 256 + && v.bytes() + .all(|c| c.is_ascii_alphanumeric() || b"-_.:/".contains(&c)) + }) + .map(str::to_owned) +} + +#[derive(Clone)] +pub struct RequestTelemetry { + telemetry: Option>, + user: String, + key: Option, + parent: Context, +} +impl RequestTelemetry { + pub fn new(telemetry: Option>, headers: &HeaderMap) -> Self { + let trust = telemetry.as_ref().is_some_and(|t| t.trust_identity_headers); + let get = |name| headers.get(name).and_then(|v| v.to_str().ok()); + Self { + telemetry, + user: if trust { + label(get("x-user-id")).unwrap_or_else(|| "unknown".into()) + } else { + "unknown".into() + }, + key: if trust { + label(get("x-api-key-id")) + } else { + None + }, + parent: TraceContextPropagator::new().extract(&HeaderExtractor(headers)), + } + } + + pub fn start( + &self, + model: &str, + operation: &'static str, + transport: &'static str, + ) -> UsageGuard { + UsageGuard { + request: self.clone(), + request_id: uuid::Uuid::new_v4().to_string(), + model: label(Some(model)).unwrap_or_else(|| "unknown".into()), + response_model: None, + response_id: None, + operation, + transport, + started: SystemTime::now(), + clock: Instant::now(), + status: "cancelled", + error: Some("client_disconnect"), + usage: None, + } + } +} + +#[derive(Default, Debug)] +struct Usage { + input: Option, + output: Option, + cached: Option, + reasoning: Option, + total: Option, +} +impl Usage { + fn from_value(value: &Value) -> Option { + if !value.is_object() { + return None; + } + let input = value["input_tokens"].as_u64(); + let output = value["output_tokens"].as_u64(); + if input.is_none() && output.is_none() { + return None; + } + Some(Self { + input, + output, + cached: value["input_tokens_details"]["cached_tokens"].as_u64(), + reasoning: value["output_tokens_details"]["reasoning_tokens"].as_u64(), + total: value["total_tokens"] + .as_u64() + .or_else(|| input?.checked_add(output?)), + }) + } +} + +pub struct UsageGuard { + request: RequestTelemetry, + request_id: String, + model: String, + response_model: Option, + response_id: Option, + operation: &'static str, + transport: &'static str, + started: SystemTime, + clock: Instant, + status: &'static str, + error: Option<&'static str>, + usage: Option, +} +impl UsageGuard { + pub fn failure(&mut self, code: &'static str) { + self.status = "error"; + self.error = Some(code); + } + + pub fn observe(&mut self, event: &Value) { + let response = &event["response"]; + if let Some(id) = label(response["id"].as_str()) { + self.response_id = Some(id); + } + if let Some(model) = label(response["model"].as_str()) { + self.response_model = Some(model); + } + match event["type"].as_str() { + Some("response.completed") => { + self.status = "completed"; + self.error = None; + } + Some("response.failed") => { + self.status = "failed"; + self.error = Some("response_failed"); + } + Some("response.incomplete") => { + self.status = "incomplete"; + self.error = Some("response_incomplete"); + } + _ => return, + } + self.usage = Usage::from_value(&response["usage"]); + } +} + +impl Drop for UsageGuard { + fn drop(&mut self) { + let Some(telemetry) = &self.request.telemetry else { + return; + }; + let elapsed = self.clock.elapsed(); + let mut attributes = vec![ + KeyValue::new("gen_ai.operation.name", "chat"), + KeyValue::new("gen_ai.provider.name", "openai"), + KeyValue::new("gen_ai.request.model", self.model.clone()), + KeyValue::new("codex_api.operation", self.operation), + KeyValue::new("codex_api.transport", self.transport), + KeyValue::new("codex_api.status", self.status), + ]; + if let Some(model) = &self.response_model { + attributes.push(KeyValue::new("gen_ai.response.model", model.clone())); + } + if let Some(error) = self.error { + attributes.push(KeyValue::new("error.type", error)); + } + // User and request identifiers never become metric dimensions. + telemetry + .duration + .record(elapsed.as_secs_f64(), &attributes); + if let Some(usage) = &self.usage { + for (counter, count) in [ + (&telemetry.input, usage.input), + (&telemetry.output, usage.output), + (&telemetry.cached, usage.cached), + (&telemetry.reasoning, usage.reasoning), + ] { + if let Some(count) = count { + counter.add(count, &attributes); + } + } + } else { + telemetry.missing.add(1, &attributes); + } + attributes.push(KeyValue::new("user.id", self.request.user.clone())); + attributes.push(KeyValue::new( + "codex_api.request.id", + self.request_id.clone(), + )); + attributes.push(KeyValue::new( + "codex_api.usage.available", + self.usage.is_some(), + )); + if let Some(key) = &self.request.key { + attributes.push(KeyValue::new("codex_api.key.id", key.clone())); + } + if let Some(id) = &self.response_id { + attributes.push(KeyValue::new("gen_ai.response.id", id.clone())); + } + let mut body = json!({ + "event": "codex_api.usage", "request_id": self.request_id, + "user_id": self.request.user, "key_id": self.request.key, + "request_model": self.model, "model": self.response_model.as_ref().unwrap_or(&self.model), + "response_id": self.response_id, "operation": self.operation, "transport": self.transport, + "status": self.status, "error_type": self.error, + "started_at_unix_ms": self.started.duration_since(UNIX_EPOCH).unwrap_or_default().as_millis() as u64, + "duration_ms": elapsed.as_millis() as u64, "usage_available": self.usage.is_some(), + }); + if let Some(usage) = &self.usage { + for (field, key, count) in [ + ("input_tokens", "gen_ai.usage.input_tokens", usage.input), + ("output_tokens", "gen_ai.usage.output_tokens", usage.output), + ( + "cached_tokens", + "gen_ai.usage.cache_read.input_tokens", + usage.cached, + ), + ( + "reasoning_tokens", + "gen_ai.usage.reasoning.output_tokens", + usage.reasoning, + ), + ("total_tokens", "codex_api.usage.total_tokens", usage.total), + ] { + if let Some(count) = count { + body[field] = json!(count); + if let Ok(count) = i64::try_from(count) { + attributes.push(KeyValue::new(key, count)); + } + } + } + } + let tracer = telemetry.traces.tracer("codex-api.usage"); + let mut span = tracer + .span_builder(format!("chat {}", self.model)) + .with_kind(SpanKind::Client) + .with_start_time(self.started) + .with_attributes(attributes.clone()) + .start_with_context(&tracer, &self.request.parent); + if let Some(error) = self.error { + span.set_status(Status::error(error)); + } + let logger = telemetry.logs.logger("codex-api.usage"); + let mut log = logger.create_log_record(); + log.set_event_name("gen_ai.client.inference.operation.details"); + log.set_timestamp(SystemTime::now()); + log.set_severity_number(Severity::Info); + log.set_body(AnyValue::from(body.to_string())); + for attr in attributes { + let value = match attr.value { + opentelemetry::Value::String(v) => AnyValue::String(v), + opentelemetry::Value::I64(v) => AnyValue::Int(v), + opentelemetry::Value::Bool(v) => AnyValue::Boolean(v), + _ => continue, + }; + log.add_attribute(attr.key, value); + } + let context = span.span_context(); + log.set_trace_context( + context.trace_id(), + context.span_id(), + Some(context.trace_flags()), + ); + logger.emit(log); + span.end(); + } +} + +/// One record per inference, not per WebSocket connection. Match response IDs +/// and lanes so terminal duplicates and multiplexed turns cannot double-count. +pub struct WebSocketUsage { + request: RequestTelemetry, + pending: HashMap>, + active: HashMap, + completed: HashSet, + recent: VecDeque, +} +impl WebSocketUsage { + pub fn new(request: RequestTelemetry) -> Self { + Self { + request, + pending: HashMap::new(), + active: HashMap::new(), + completed: HashSet::new(), + recent: VecDeque::new(), + } + } + pub fn request(&mut self, text: &str) { + let Ok(value) = serde_json::from_str::(text) else { + return; + }; + if value["type"] != "response.create" || value["generate"] == false { + return; + } + let mut guard = self.request.start( + value["model"].as_str().unwrap_or("unknown"), + "responses", + "websocket", + ); + if self.active.len() + self.pending.values().map(VecDeque::len).sum::() >= 128 { + guard.failure("telemetry_capacity"); + return; + } + let lane = label(value["stream_id"].as_str()).unwrap_or_default(); + self.pending.entry(lane).or_default().push_back(guard); + } + fn take_pending(&mut self, lane: &str) -> Option { + let queue = self.pending.get_mut(lane)?; + let guard = queue.pop_front(); + if queue.is_empty() { + self.pending.remove(lane); + } + guard + } + pub fn observe(&mut self, event: &Value) { + let lane = label(event["stream_id"].as_str()).unwrap_or_default(); + let id = label(event["response"]["id"].as_str()); + match event["type"].as_str() { + Some("response.created") => { + if let Some(id) = id + && !self.active.contains_key(&id) + && let Some(mut guard) = self.take_pending(&lane) + { + guard.observe(event); + self.active.insert(id, guard); + } + } + Some("response.completed" | "response.failed" | "response.incomplete") => { + if let Some(id) = &id { + if !self.completed.insert(id.clone()) { + return; + } + self.recent.push_back(id.clone()); + if self.recent.len() > 128 + && let Some(old) = self.recent.pop_front() + { + self.completed.remove(&old); + } + } + if let Some(mut guard) = id + .as_ref() + .and_then(|id| self.active.remove(id)) + .or_else(|| self.take_pending(&lane)) + { + guard.observe(event); + } + } + _ => {} + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use opentelemetry_sdk::logs::InMemoryLogExporter; + + fn setup(trust: bool) -> (RequestTelemetry, InMemoryLogExporter) { + let exporter = InMemoryLogExporter::default(); + let telemetry = Arc::new(Telemetry::new( + SdkTracerProvider::builder().build(), + SdkLoggerProvider::builder() + .with_simple_exporter(exporter.clone()) + .build(), + SdkMeterProvider::builder().build(), + trust, + )); + let mut headers = HeaderMap::new(); + headers.insert("x-user-id", "user-1".parse().unwrap()); + headers.insert("x-api-key-id", "key-1".parse().unwrap()); + (RequestTelemetry::new(Some(telemetry), &headers), exporter) + } + + fn bodies(exporter: &InMemoryLogExporter) -> Vec { + exporter + .get_emitted_logs() + .unwrap() + .iter() + .map(|log| { + let Some(AnyValue::String(body)) = log.record.body() else { + panic!("JSON body expected") + }; + serde_json::from_str(body.as_str()).unwrap() + }) + .collect() + } + + fn completed(id: &str) -> Value { + json!({"type":"response.completed", "response":{ + "id":id, "model":"gpt-5", "usage":{ + "input_tokens":12,"output_tokens":3, + "input_tokens_details":{"cached_tokens":2}, + "output_tokens_details":{"reasoning_tokens":1} + }}}) + } + + #[test] + fn identity_usage_and_subset_totals_are_exported() { + let (request, exporter) = setup(true); + let mut guard = request.start("gpt-5", "responses", "http"); + guard.observe(&completed("resp-1")); + drop(guard); + let body = &bodies(&exporter)[0]; + assert_eq!(body["user_id"], "user-1"); + assert_eq!(body["key_id"], "key-1"); + assert_eq!(body["total_tokens"], 15); + assert_eq!(body["cached_tokens"], 2); + assert_eq!(body["reasoning_tokens"], 1); + assert_eq!(body["status"], "completed"); + } + + #[test] + fn untrusted_identity_and_missing_usage_remain_explicit() { + let (request, exporter) = setup(false); + drop(request.start("gpt-5", "responses", "http")); + let body = &bodies(&exporter)[0]; + assert_eq!(body["user_id"], "unknown"); + assert!(body["key_id"].is_null()); + assert_eq!(body["usage_available"], false); + assert!(body.get("total_tokens").is_none()); + assert_eq!(body["status"], "cancelled"); + } + + #[test] + fn websocket_turns_count_once_and_skip_non_generation() { + let (request, exporter) = setup(true); + let mut socket = WebSocketUsage::new(request); + socket.request(r#"{"type":"response.create","model":"gpt-5","generate":false}"#); + for id in ["resp-1", "resp-2"] { + socket.request(r#"{"type":"response.create","model":"gpt-5"}"#); + socket.observe(&json!({"type":"response.created","response":{"id":id}})); + socket.observe(&completed(id)); + socket.observe(&completed(id)); + } + drop(socket); + let records = bodies(&exporter); + assert_eq!(records.len(), 2); + assert!(records.iter().all(|body| body["total_tokens"] == 15)); + assert_ne!(records[0]["request_id"], records[1]["request_id"]); + } +} diff --git a/src/websocket.rs b/src/websocket.rs index 669ea72..1372e8f 100644 --- a/src/websocket.rs +++ b/src/websocket.rs @@ -24,6 +24,7 @@ pub async fn upgrade( ws: WebSocketUpgrade, ) -> Result { authorize(&gateway, &headers)?; + let telemetry = crate::telemetry::RequestTelemetry::new(gateway.telemetry.clone(), &headers); let headers = crate::headers::request_headers(&headers); let permit = gateway .concurrency @@ -56,7 +57,16 @@ pub async fn upgrade( .max_frame_size(16 * 1024 * 1024) .on_upgrade(move |client| async move { let _permit = permit; - relay(client, upstream, idle, public, &gateway.models, &session).await; + relay_with_telemetry( + client, + upstream, + idle, + public, + &gateway.models, + &session, + telemetry, + ) + .await; }); response .headers_mut() @@ -65,15 +75,39 @@ pub async fn upgrade( } pub(crate) async fn relay( + client: WebSocket, + upstream: S, + idle: std::time::Duration, + public: bool, + models: &[codex_protocol::openai_models::ModelInfo], + session: &str, +) where + S: futures::Stream> + futures::Sink + Unpin, +{ + relay_with_telemetry( + client, + upstream, + idle, + public, + models, + session, + crate::telemetry::RequestTelemetry::new(None, &HeaderMap::new()), + ) + .await; +} + +pub(crate) async fn relay_with_telemetry( mut client: WebSocket, mut upstream: S, idle: Duration, public: bool, models: &[codex_protocol::openai_models::ModelInfo], session: &str, + telemetry: crate::telemetry::RequestTelemetry, ) where S: futures::Stream> + futures::Sink + Unpin, { + let mut usage = crate::telemetry::WebSocketUsage::new(telemetry); let mut accumulator = crate::output::OutputAccumulator::default(); loop { let step = async { @@ -92,6 +126,7 @@ pub(crate) async fn relay( } } } else { v.as_str().to_owned() }; + usage.request(&text); Upstream::Text(text.into()) }, Message::Binary(v) => Upstream::Binary(v), @@ -107,6 +142,7 @@ pub(crate) async fn relay( let message = match message { Upstream::Text(v) => { let mut text = v.as_str().to_owned(); + if let Ok(event) = serde_json::from_str::(&text) { usage.observe(&event); } if public && let Ok(mut event) = serde_json::from_str::(&text) && accumulator.observe(&mut event) { text = event.to_string(); } Message::Text(text.into()) }, diff --git a/tests/catalog.rs b/tests/catalog.rs index 0119641..55ca0fb 100644 --- a/tests/catalog.rs +++ b/tests/catalog.rs @@ -73,6 +73,7 @@ async fn startup_catalog_exposes_new_models_and_uses_their_capabilities() { ); } let app = router(Gateway { + telemetry: None, backend, models, key: None, diff --git a/tests/gateway.rs b/tests/gateway.rs index d82b6c2..32f9132 100644 --- a/tests/gateway.rs +++ b/tests/gateway.rs @@ -305,6 +305,7 @@ impl Harness { .unwrap() }); let app = router(Gateway { + telemetry: None, backend, models, key: Some("client-key".into()), diff --git a/tests/support/live_gateway.rs b/tests/support/live_gateway.rs index 0c12a38..4afea95 100644 --- a/tests/support/live_gateway.rs +++ b/tests/support/live_gateway.rs @@ -53,6 +53,7 @@ impl LiveGateway { let model_catalog = Some(models.clone()); let observed = transports.clone(); let app = router(Gateway { + telemetry: None, backend, models, key: Some(key.clone()),