Codex Handbook
codex-api/src/endpoint/session.rs 156 lines
use crate::auth::SharedAuthProvider;use crate::error::ApiError;use crate::provider::Provider;use crate::telemetry::run_with_request_telemetry;use codex_client::EncodedJsonBody;use codex_client::HttpTransport;use codex_client::Request;use codex_client::RequestBody;use codex_client::RequestTelemetry;use codex_client::Response;use codex_client::StreamResponse;use codex_client::TransportError;use http::HeaderMap;use http::Method;use serde_json::Value;use std::sync::Arc;use tracing::instrument;pub(crate) struct EndpointSession<T: HttpTransport> {    transport: T,    provider: Provider,    auth: SharedAuthProvider,    request_telemetry: Option<Arc<dyn RequestTelemetry>>,}impl<T: HttpTransport> EndpointSession<T> {    pub(crate) fn new(transport: T, provider: Provider, auth: SharedAuthProvider) -> Self {        Self {            transport,            provider,            auth,            request_telemetry: None,        }    }    pub(crate) fn with_request_telemetry(        mut self,        request: Option<Arc<dyn RequestTelemetry>>,    ) -> Self {        self.request_telemetry = request;        self    }    pub(crate) fn provider(&self) -> &Provider {        &self.provider    }    fn make_request(        &self,        method: &Method,        path: &str,        extra_headers: &HeaderMap,        body: Option<&RequestBody>,    ) -> Request {        let mut req = self.provider.build_request(method.clone(), path);        req.headers.extend(extra_headers.clone());        if let Some(body) = body {            req.body = Some(body.clone());        }        req    }    pub(crate) async fn execute(        &self,        method: Method,        path: &str,        extra_headers: HeaderMap,        body: Option<Value>,    ) -> Result<Response, ApiError> {        self.execute_with(method, path, extra_headers, body, |_| {})            .await    }    #[instrument(        name = "endpoint_session.execute_with",        level = "info",        skip_all,        fields(http.method = %method, api.path = path)    )]    pub(crate) async fn execute_with<C>(        &self,        method: Method,        path: &str,        extra_headers: HeaderMap,        body: Option<Value>,        configure: C,    ) -> Result<Response, ApiError>    where        C: Fn(&mut Request),    {        let body = body.map(RequestBody::Json);        let make_request = || {            let mut req = self.make_request(&method, path, &extra_headers, body.as_ref());            configure(&mut req);            req        };        let response = run_with_request_telemetry(            self.provider.retry.to_policy(),            self.request_telemetry.clone(),            make_request,            |req| {                let auth = self.auth.clone();                let transport = &self.transport;                async move {                    let req = auth.apply_auth(req).await.map_err(TransportError::from)?;                    transport.execute(req).await                }            },        )        .await?;        Ok(response)    }    #[instrument(        name = "endpoint_session.stream_encoded_json_with",        level = "info",        skip_all,        fields(http.method = %method, api.path = path)    )]    pub(crate) async fn stream_encoded_json_with<C>(        &self,        method: Method,        path: &str,        extra_headers: HeaderMap,        body: Option<EncodedJsonBody>,        configure: C,    ) -> Result<StreamResponse, ApiError>    where        C: Fn(&mut Request),    {        let body = body.map(RequestBody::EncodedJson);        let mut request = self.make_request(&method, path, &extra_headers, body.as_ref());        configure(&mut request);        let request = request.into_prepared().map_err(TransportError::Build)?;        let make_request = || request.clone();        let stream = run_with_request_telemetry(            self.provider.retry.to_policy(),            self.request_telemetry.clone(),            make_request,            |req| {                let auth = self.auth.clone();                let transport = &self.transport;                async move {                    let req = auth.apply_auth(req).await.map_err(TransportError::from)?;                    transport.stream(req).await                }            },        )        .await?;        Ok(stream)    }}