feat: clean codex-api part 1 (#10501)

codex-rs/codex-api/src/endpoint/session.rs

@@ -0,0 +1,126 @@
+use crate::auth::AuthProvider;
+use crate::auth::add_auth_headers;
+use crate::error::ApiError;
+use crate::provider::Provider;
+use crate::telemetry::run_with_request_telemetry;
+use codex_client::HttpTransport;
+use codex_client::Request;
+use codex_client::RequestTelemetry;
+use codex_client::Response;
+use codex_client::StreamResponse;
+use http::HeaderMap;
+use http::Method;
+use serde_json::Value;
+use std::sync::Arc;
+
+pub(crate) struct EndpointSession<T: HttpTransport, A: AuthProvider> {
+    transport: T,
+    provider: Provider,
+    auth: A,
+    request_telemetry: Option<Arc<dyn RequestTelemetry>>,
+}
+
+impl<T: HttpTransport, A: AuthProvider> EndpointSession<T, A> {
+    pub(crate) fn new(transport: T, provider: Provider, auth: A) -> 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<&Value>,
+    ) -> 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());
+        }
+        add_auth_headers(&self.auth, 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
+    }
+
+    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 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| self.transport.execute(req),
+        )
+        .await?;
+
+        Ok(response)
+    }
+
+    pub(crate) async fn stream_with<C>(
+        &self,
+        method: Method,
+        path: &str,
+        extra_headers: HeaderMap,
+        body: Option<Value>,
+        configure: C,
+    ) -> Result<StreamResponse, ApiError>
+    where
+        C: Fn(&mut Request),
+    {
+        let make_request = || {
+            let mut req = self.make_request(&method, path, &extra_headers, body.as_ref());
+            configure(&mut req);
+            req
+        };
+
+        let stream = run_with_request_telemetry(
+            self.provider.retry.to_policy(),
+            self.request_telemetry.clone(),
+            make_request,
+            |req| self.transport.stream(req),
+        )
+        .await?;
+
+        Ok(stream)
+    }
+}