126 lines
3.2 KiB
Rust
126 lines
3.2 KiB
Rust
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)
|
|
}
|
|
}
|