Codex Handbook
codex-api/src/telemetry.rs 98 lines
use crate::error::ApiError;use codex_client::Request;use codex_client::RequestTelemetry;use codex_client::Response;use codex_client::RetryPolicy;use codex_client::StreamResponse;use codex_client::TransportError;use codex_client::run_with_retry;use http::StatusCode;use std::future::Future;use std::sync::Arc;use std::time::Duration;use tokio::time::Instant;use tokio_tungstenite::tungstenite::Error;use tokio_tungstenite::tungstenite::Message;/// Generic telemetry.pub trait SseTelemetry: Send + Sync {    fn on_sse_poll(        &self,        result: &Result<            Option<                Result<                    eventsource_stream::Event,                    eventsource_stream::EventStreamError<TransportError>,                >,            >,            tokio::time::error::Elapsed,        >,        duration: Duration,    );}/// Telemetry for Responses WebSocket transport.pub trait WebsocketTelemetry: Send + Sync {    fn on_ws_request(&self, duration: Duration, error: Option<&ApiError>, connection_reused: bool);    fn on_ws_event(        &self,        result: &Result<Option<Result<Message, Error>>, ApiError>,        duration: Duration,    );}pub(crate) trait WithStatus {    fn status(&self) -> StatusCode;}fn http_status(err: &TransportError) -> Option<StatusCode> {    match err {        TransportError::Http { status, .. } => Some(*status),        _ => None,    }}impl WithStatus for Response {    fn status(&self) -> StatusCode {        self.status    }}impl WithStatus for StreamResponse {    fn status(&self) -> StatusCode {        self.status    }}pub(crate) async fn run_with_request_telemetry<T, F, Fut>(    policy: RetryPolicy,    telemetry: Option<Arc<dyn RequestTelemetry>>,    make_request: impl FnMut() -> Request,    send: F,) -> Result<T, TransportError>where    T: WithStatus,    F: Clone + Fn(Request) -> Fut,    Fut: Future<Output = Result<T, TransportError>>,{    // Wraps `run_with_retry` to attach per-attempt request telemetry for both    // unary and streaming HTTP calls.    run_with_retry(policy, make_request, move |req, attempt| {        let telemetry = telemetry.clone();        let send = send.clone();        async move {            let start = Instant::now();            let result = send(req).await;            if let Some(t) = telemetry.as_ref() {                let (status, err) = match &result {                    Ok(resp) => (Some(resp.status()), None),                    Err(err) => (http_status(err), Some(err)),                };                t.on_request(attempt, status, err, start.elapsed());            }            result        }    })    .await}