Export verified user and model token usage through OTLP

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"
 

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.

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"

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;

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(())
 }

src/server.rs

@@ -20,6 +20,7 @@ pub struct Gateway {
     pub concurrency: Arc<Semaphore>,
     pub timeout: Duration,
     pub transfers: Option<crate::transfers::Transfers>,
+    pub telemetry: Option<Arc<crate::telemetry::Telemetry>>,
 }
 
 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<serde_json::value::RawValue>, HeaderMap)>,
     stream: bool,
     aggregate_output: bool,
+    mut usage: crate::telemetry::UsageGuard,
 ) -> Result<Response, GatewayError> {
     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<Json<CreateResponse>, JsonRejection>,
 ) -> Result<Response, GatewayError> {
     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)