ebdd8795e9
- capture the header from SSE/WS handshakes, store it per ModelClientSession using `Oncelock`, echo it on turn-scoped requests, and add SSE+WS integration tests for within-turn persistence + cross-turn reset. - keep `x-codex-turn-state` sticky within a user turn to maintain routing continuity for retries/tool follow-ups.
95 lines
2.7 KiB
Rust
95 lines
2.7 KiB
Rust
use crate::auth::AuthProvider;
|
|
use crate::auth::add_auth_headers;
|
|
use crate::common::ResponseStream;
|
|
use crate::error::ApiError;
|
|
use crate::provider::Provider;
|
|
use crate::telemetry::SseTelemetry;
|
|
use crate::telemetry::run_with_request_telemetry;
|
|
use codex_client::HttpTransport;
|
|
use codex_client::RequestCompression;
|
|
use codex_client::RequestTelemetry;
|
|
use codex_client::StreamResponse;
|
|
use http::HeaderMap;
|
|
use http::Method;
|
|
use serde_json::Value;
|
|
use std::sync::Arc;
|
|
use std::sync::OnceLock;
|
|
use std::time::Duration;
|
|
|
|
pub(crate) struct StreamingClient<T: HttpTransport, A: AuthProvider> {
|
|
transport: T,
|
|
provider: Provider,
|
|
auth: A,
|
|
request_telemetry: Option<Arc<dyn RequestTelemetry>>,
|
|
sse_telemetry: Option<Arc<dyn SseTelemetry>>,
|
|
}
|
|
|
|
type StreamSpawner = fn(
|
|
StreamResponse,
|
|
Duration,
|
|
Option<Arc<dyn SseTelemetry>>,
|
|
Option<Arc<OnceLock<String>>>,
|
|
) -> ResponseStream;
|
|
|
|
impl<T: HttpTransport, A: AuthProvider> StreamingClient<T, A> {
|
|
pub(crate) fn new(transport: T, provider: Provider, auth: A) -> Self {
|
|
Self {
|
|
transport,
|
|
provider,
|
|
auth,
|
|
request_telemetry: None,
|
|
sse_telemetry: None,
|
|
}
|
|
}
|
|
|
|
pub(crate) fn with_telemetry(
|
|
mut self,
|
|
request: Option<Arc<dyn RequestTelemetry>>,
|
|
sse: Option<Arc<dyn SseTelemetry>>,
|
|
) -> Self {
|
|
self.request_telemetry = request;
|
|
self.sse_telemetry = sse;
|
|
self
|
|
}
|
|
|
|
pub(crate) fn provider(&self) -> &Provider {
|
|
&self.provider
|
|
}
|
|
|
|
pub(crate) async fn stream(
|
|
&self,
|
|
path: &str,
|
|
body: Value,
|
|
extra_headers: HeaderMap,
|
|
compression: RequestCompression,
|
|
spawner: StreamSpawner,
|
|
turn_state: Option<Arc<OnceLock<String>>>,
|
|
) -> Result<ResponseStream, ApiError> {
|
|
let builder = || {
|
|
let mut req = self.provider.build_request(Method::POST, path);
|
|
req.headers.extend(extra_headers.clone());
|
|
req.headers.insert(
|
|
http::header::ACCEPT,
|
|
http::HeaderValue::from_static("text/event-stream"),
|
|
);
|
|
req.body = Some(body.clone());
|
|
req.compression = compression;
|
|
add_auth_headers(&self.auth, req)
|
|
};
|
|
|
|
let stream_response = run_with_request_telemetry(
|
|
self.provider.retry.to_policy(),
|
|
self.request_telemetry.clone(),
|
|
builder,
|
|
|req| self.transport.stream(req),
|
|
)
|
|
.await?;
|
|
|
|
Ok(spawner(
|
|
stream_response,
|
|
self.provider.stream_idle_timeout,
|
|
self.sse_telemetry.clone(),
|
|
turn_state,
|
|
))
|
|
}
|
|
}
|