diff --git a/Cargo.lock b/Cargo.lock index 364c36e..c1eece4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -177,14 +177,18 @@ version = "0.10.5" dependencies = [ "async-trait", "bytes", + "fastrand", + "futures-timer", "futures-util", "http", + "httpdate", "reqwest", "reqwest-middleware", "serde", "serde_json", "thiserror 2.0.18", "tokio", + "zeroize", ] [[package]] @@ -356,9 +360,14 @@ dependencies = [ "agentkit-http", "agentkit-loop", "async-trait", + "base64 0.22.1", + "futures-util", + "reqwest", "serde", "serde_json", "thiserror 2.0.18", + "tokio", + "zeroize", ] [[package]] diff --git a/README.md b/README.md index fe66200..f084ed1 100644 --- a/README.md +++ b/README.md @@ -122,9 +122,25 @@ Shell: The filesystem crate also supports session-scoped read-before-write enforcement through `FileSystemToolResources` and `FileSystemToolPolicy`. +## Provider authentication and resilience + +Provider configs use `agentkit_http::Authentication` as their first-class +credential type. For Baseten, Cerebras, Groq, Mistral, OpenAI, OpenRouter, and +authenticated Ollama/vLLM endpoints, passing a bare string to an authentication +argument is shorthand for bearer authentication. Custom refresh-capable +credentials can be installed with `.with_authentication_provider(...)`. +Anthropic is the exception: `AnthropicConfig::new(...)` sends `x-api-key`, while +`AnthropicConfig::with_auth_token(...)` explicitly selects a bearer auth token. +Ollama and vLLM authentication is optional. + +Provider resilience is also opt-in. Configs store +`Option` and default to `None`; call +`.with_resilience(...)` to enable retries and timeouts. Leaving it as `None` +preserves the existing single-attempt behavior. + ## Quick start -1. Set your OpenRouter API key and model — either through environment variables or directly in code via `OpenRouterConfig::new(api_key, model)`. +1. Set your OpenRouter API key and model — either through environment variables or directly in code via `OpenRouterConfig::new(authentication, model)`. 2. Run one of the examples. Example commands: diff --git a/book/src/ch01-model-adapter.md b/book/src/ch01-model-adapter.md index 1f167b2..744553f 100644 --- a/book/src/ch01-model-adapter.md +++ b/book/src/ch01-model-adapter.md @@ -254,10 +254,17 @@ pub enum ModelTurnEvent { Delta(Delta), ToolCall(ToolCallPart), Usage(Usage), + ResponseAttemptSuperseded, Finished(ModelTurnResult), } ``` +`ResponseAttemptSuperseded` is capability-gated. An adapter emits it after events +from a failed visible attempt and before any replacement-attempt events. The loop +forwards it as `AgentEvent::ResponseAttemptSuperseded`. A consumer that enables +`SessionConfig::with_response_attempt_supersession()` must discard all deltas, +tool calls, usage updates, and reconstruction state from the preceding attempt. + ## Building an adapter from scratch To see what the traits require, consider a hypothetical model provider that does not use the OpenAI format. Suppose "AcmeAI" has a proprietary REST API: @@ -287,12 +294,14 @@ No `messages` array. No `choices` wrapper. No `tool_calls`. A completely differe ```rust pub struct AcmeAdapter { client: Client, - api_key: String, + authentication: Authentication, + resilience: Option, } pub struct AcmeSession { client: Client, - api_key: String, + authentication: Authentication, + resilience: Option, } #[async_trait] @@ -302,7 +311,8 @@ impl ModelAdapter for AcmeAdapter { async fn start_session(&self, _config: SessionConfig) -> Result { Ok(AcmeSession { client: self.client.clone(), - api_key: self.api_key.clone(), + authentication: self.authentication.clone(), + resilience: self.resilience.clone(), }) } } @@ -348,10 +358,15 @@ impl ModelSession for AcmeSession { "config": { "temperature": 0.5, "max_tokens": 256 }, }); - let resp: AcmeResponse = self.client + let authentication = self.authentication.authenticate(None).await + .map_err(|e| LoopError::Provider(e.to_string()))?; + let mut request = self.client .post("https://api.acme.ai/v1/generate") - .bearer_auth(&self.api_key) - .json(&body) + .json(&body); + for (name, value) in authentication.headers() { + request = request.header(name, value); + } + let resp: AcmeResponse = request .send().await .map_err(|e| LoopError::Provider(e.to_string()))? .json().await @@ -619,9 +634,9 @@ This is the complete provider. All of the transcript conversion, tool call seria Not all OpenAI-compatible providers are identical. The three hooks exist for providers that need to customise the standard request/response flow. -OpenRouter uses all three: +OpenRouter uses all three, while exposing its first-class credential through `authentication()`: -1. **`preprocess_request`** — adds bearer auth, `X-Title`, and `HTTP-Referer` headers +1. **`preprocess_request`** — adds `X-Title` and `HTTP-Referer` headers 2. **`preprocess_response`** — the API sometimes returns HTTP 200 with an error payload instead of a proper error status; the hook parses these and converts them to errors before the adapter attempts normal deserialization 3. **`postprocess_response`** — extracts the `cost` field from the usage object (OpenRouter-specific, not part of the standard format) and adds `openrouter.model` and `openrouter.refusal` to the item metadata @@ -633,11 +648,15 @@ impl CompletionsProvider for OpenRouterProvider { fn endpoint_url(&self) -> &str { &self.base_url } fn config(&self) -> &OpenRouterRequestConfig { &self.request_config } + fn authentication(&self) -> Option { + Some(self.authentication.clone()) + } + fn preprocess_request( &self, builder: agentkit_http::HttpRequestBuilder, ) -> agentkit_http::HttpRequestBuilder { - let mut builder = builder.bearer_auth(&self.api_key); + let mut builder = builder; if let Some(app_name) = &self.app_name { builder = builder.header("X-Title", app_name); } diff --git a/book/src/ch03-transcript-model.md b/book/src/ch03-transcript-model.md index 4d50093..f32f7d7 100644 --- a/book/src/ch03-transcript-model.md +++ b/book/src/ch03-transcript-model.md @@ -49,18 +49,23 @@ pub enum ItemKind { The variants are ordered: `System < Developer < User < Assistant < Tool < Context`. This ordering is used by compaction strategies that need to sort or prioritise items by role. -Role mapping to provider wire formats: - -| agentkit `ItemKind` | OpenAI role | What it carries | -| ------------------- | ------------- | ---------------------------------- | -| `System` | `"system"` | Hardcoded application instructions | -| `Developer` | `"system"` | Developer-level instructions | -| `User` | `"user"` | End-user messages | -| `Assistant` | `"assistant"` | Model-generated text + tool calls | -| `Tool` | `"tool"` | Tool execution results | -| `Context` | `"system"` | Project context (AGENTS.md, etc.) | - -System, Developer, and Context all map to `"system"` in the OpenAI wire format, but they carry different semantic intent. The distinction matters for compaction: system items are never trimmed, context items may be refreshed, and developer items sit between the two. Collapsing them into a single kind would lose information that compaction strategies need. +Role mapping depends on the OpenAI API and profile. Chat Completions uses: + +| agentkit `ItemKind` | Chat Completions role | What it carries | +| ------------------- | --------------------- | ---------------------------------- | +| `System` | `"system"` | Hardcoded application instructions | +| `Developer` | `"developer"` | Developer-level instructions | +| `User` | `"user"` | End-user messages | +| `Assistant` | `"assistant"` | Model-generated text + tool calls | +| `Tool` | `"tool"` | Tool execution results | +| `Context` | `"system"` | Project context (AGENTS.md, etc.) | + +For the Responses API, the public profile keeps System and Context items as +`system`; the private profile sends both as `developer`. Context content is +encoded unchanged, without consumer-specific prose. These kinds retain distinct +semantic intent even when a provider maps them to the same role. The distinction +matters for compaction: system items are never trimmed, context items may be +refreshed, and developer items sit between the two. ### Why item-based, not message-based diff --git a/book/src/ch05-model-adapter.md b/book/src/ch05-model-adapter.md index 52ce6a4..6bcf4cf 100644 --- a/book/src/ch05-model-adapter.md +++ b/book/src/ch05-model-adapter.md @@ -105,6 +105,7 @@ pub enum ModelTurnEvent { Delta(Delta), ToolCall(ToolCallPart), Usage(Usage), + ResponseAttemptSuperseded, Finished(ModelTurnResult), } ``` @@ -129,11 +130,21 @@ Turn event timeline: Usage(Usage) ← token counts + ResponseAttemptSuperseded ← optional; invalidates every event above + Delta(...) ← replacement attempt starts after the marker + Finished(ModelTurnResult) ← always last ``` `Finished` always comes last. `Usage` typically comes just before `Finished` but some providers interleave it with deltas. `ToolCall` events represent fully assembled tool calls — the adapter has already accumulated the streaming chunks internally. +`ResponseAttemptSuperseded` is emitted only when the consumer enabled +`SessionConfig::with_response_attempt_supersession()`. It appears after the failed +visible attempt and before replacement output. The loop resets its attempt-local +tool-call and usage state and forwards `AgentEvent::ResponseAttemptSuperseded`; +consumers must discard every delta, tool call, usage update, and reconstruction +state from the preceding attempt. + ### ModelTurnResult ```rust diff --git a/book/src/ch15-caching.md b/book/src/ch15-caching.md index 402a70f..44e8cc6 100644 --- a/book/src/ch15-caching.md +++ b/book/src/ch15-caching.md @@ -13,6 +13,7 @@ pub struct SessionConfig { pub session_id: SessionId, pub metadata: MetadataMap, pub cache: Option, + pub consumer_capabilities: SessionConsumerCapabilities, } pub struct TurnRequest { @@ -83,16 +84,12 @@ The simplest place to configure caching is the session: ```rust let mut driver = agent - .start(SessionConfig { - session_id: SessionId::new("coding-agent"), - metadata: MetadataMap::new(), - cache: Some(PromptCacheRequest { - mode: PromptCacheMode::BestEffort, - strategy: PromptCacheStrategy::Automatic, - retention: Some(PromptCacheRetention::Short), - key: None, - }), - }) + .start(SessionConfig::new("coding-agent").with_cache(PromptCacheRequest { + mode: PromptCacheMode::BestEffort, + strategy: PromptCacheStrategy::Automatic, + retention: Some(PromptCacheRetention::Short), + key: None, + })) .await?; ``` @@ -225,16 +222,12 @@ This makes caching visible to reporters and host-side cost accounting without ex For most hosts, start here: ```rust -SessionConfig { - session_id: SessionId::new("demo"), - metadata: MetadataMap::new(), - cache: Some(PromptCacheRequest { - mode: PromptCacheMode::BestEffort, - strategy: PromptCacheStrategy::Automatic, - retention: Some(PromptCacheRetention::Short), - key: None, - }), -} +SessionConfig::new("demo").with_cache(PromptCacheRequest { + mode: PromptCacheMode::BestEffort, + strategy: PromptCacheStrategy::Automatic, + retention: Some(PromptCacheRetention::Short), + key: None, +}) ``` Then reach for explicit breakpoints only when you need to control exact cache boundaries. diff --git a/crates/agentkit-adapter-completions/src/lib.rs b/crates/agentkit-adapter-completions/src/lib.rs index cdadad0..1bd7409 100644 --- a/crates/agentkit-adapter-completions/src/lib.rs +++ b/crates/agentkit-adapter-completions/src/lib.rs @@ -25,17 +25,26 @@ mod response; mod sse; mod stream; +#[cfg(test)] +mod resilience_tests; + use std::collections::VecDeque; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::time::Duration; use agentkit_core::{MetadataMap, TurnCancellation, Usage}; -use agentkit_http::{BodyStream, Http, HttpError, HttpRequestBuilder, StatusCode}; +use agentkit_http::{ + Authentication, AuthenticationAttempt, BodyStream, Bytes, Http, HttpError, HttpRequestBuilder, + HttpResponse, LogicalDeadline, ResilienceConfig, StatusCode, TruncatedStreamDetector, + is_retryable_body_read, is_retryable_status, next_body_chunk_bounded, run_bounded, + sleep as resilience_sleep, +}; use agentkit_loop::{ LoopError, ModelAdapter, ModelSession, ModelTurn, ModelTurnEvent, SessionConfig, TurnRequest, }; use async_trait::async_trait; -use futures_util::StreamExt; -use futures_util::future::{Either, select}; +use futures_util::future::{BoxFuture, Either, select}; use serde::Serialize; use serde_json::Value; @@ -89,6 +98,18 @@ pub trait CompletionsProvider: Send + Sync + Clone { builder } + /// Optional generic authentication. Existing provider request hooks continue + /// to work when this returns `None`. + fn authentication(&self) -> Option { + None + } + + /// Optional retry and timeout policy. `None` preserves the original single- + /// attempt, no-timeout behavior. + fn resilience_config(&self) -> Option { + None + } + /// Hook to map a normalized prompt cache request into the provider's JSON /// request body. /// @@ -174,6 +195,8 @@ pub struct CompletionsAdapter { /// `gen_ai.provider.name` attribute from the OTel GenAI semantic /// conventions. provider_label: String, + authentication: Option, + resilience: Option, } impl CompletionsAdapter

{ @@ -186,10 +209,14 @@ impl CompletionsAdapter

{ .map(Http::new) .map_err(|error| CompletionsError::HttpClient(HttpError::request(error)))?; + let authentication = provider.authentication(); + let resilience = provider.resilience_config(); Ok(Self { client, provider_label: provider.provider_name().to_lowercase(), provider: Arc::new(provider), + authentication, + resilience, }) } @@ -197,12 +224,28 @@ impl CompletionsAdapter

{ /// attach auth headers via `default_headers`, supply custom TLS/proxies, /// or plug in a non-reqwest backend. pub fn with_client(provider: P, client: Http) -> Self { + let authentication = provider.authentication(); + let resilience = provider.resilience_config(); Self { client, provider_label: provider.provider_name().to_lowercase(), provider: Arc::new(provider), + authentication, + resilience, } } + + /// Overrides the provider's generic authentication configuration. + pub fn with_authentication(mut self, authentication: impl Into) -> Self { + self.authentication = Some(authentication.into()); + self + } + + /// Enables request and pre-visible-output stream retries and timeouts. + pub fn with_resilience(mut self, resilience: ResilienceConfig) -> Self { + self.resilience = Some(resilience); + self + } } /// An active session with a chat completions provider. @@ -213,6 +256,8 @@ pub struct CompletionsSession { provider: Arc

, model: Option, provider_label: String, + authentication: Option, + resilience: Option, _session_config: SessionConfig, } @@ -226,12 +271,23 @@ enum TurnInner { Streaming(Box), } +type StreamReplay = Box< + dyn FnMut() -> BoxFuture<'static, Result<(BodyStream, TruncatedStreamDetector), LoopError>> + + Send, +>; + struct StreamingState { body: BodyStream, decoder: SseDecoder, translator: EventTranslator, pending: VecDeque, eof: bool, + visible_output: bool, + detect_truncation: bool, + idle_timeout: Option, + deadline: Option, + integrity: TruncatedStreamDetector, + replay: Option, postprocess: PostprocessResponse, } @@ -242,7 +298,15 @@ impl CompletionsTurn { } } - fn streaming(body: BodyStream, postprocess: PostprocessResponse) -> Self { + fn streaming( + body: BodyStream, + postprocess: PostprocessResponse, + idle_timeout: Option, + detect_truncation: bool, + deadline: Option, + integrity: TruncatedStreamDetector, + replay: Option, + ) -> Self { Self { inner: TurnInner::Streaming(Box::new(StreamingState { body, @@ -250,12 +314,230 @@ impl CompletionsTurn { translator: EventTranslator::new(), pending: VecDeque::new(), eof: false, + visible_output: false, + detect_truncation, + idle_timeout, + deadline, + integrity, + replay, postprocess, })), } } } +#[derive(Clone)] +struct ReplayRequest { + client: Http, + provider: Arc

, + body: Bytes, + authentication: Option, + authentication_attempt: Arc>>, + reauthenticated: Arc, + resilience: Option, + deadline: Option, + retries_used: Arc, +} + +impl ReplayRequest

{ + async fn authentication_attempt(&self) -> Result, LoopError> { + let existing = self + .authentication_attempt + .lock() + .map_err(|_| LoopError::Provider("authentication state lock poisoned".into()))? + .clone(); + if existing.is_some() || self.authentication.is_none() { + return Ok(existing); + } + let authenticate = self + .authentication + .as_ref() + .expect("checked above") + .authenticate(None); + let attempt = run_bounded(authenticate, None, self.deadline.as_ref(), "authentication") + .await + .map_err(|error| LoopError::Provider(format!("authentication failed: {error}")))?; + *self + .authentication_attempt + .lock() + .map_err(|_| LoopError::Provider("authentication state lock poisoned".into()))? = + Some(attempt.clone()); + Ok(Some(attempt)) + } + + async fn reauthenticate(&self, previous: AuthenticationAttempt) -> Result<(), LoopError> { + let authenticate = self + .authentication + .as_ref() + .expect("401 refresh requires authentication") + .authenticate(Some(&previous)); + let next = run_bounded( + authenticate, + None, + self.deadline.as_ref(), + "reauthentication", + ) + .await + .map_err(|error| LoopError::Provider(format!("reauthentication failed: {error}")))?; + *self + .authentication_attempt + .lock() + .map_err(|_| LoopError::Provider("authentication state lock poisoned".into()))? = + Some(next); + Ok(()) + } + + fn reserve_retry(&self) -> Option { + let maximum = self.resilience.as_ref()?.max_retries; + self.retries_used + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |used| { + (used < maximum).then_some(used + 1) + }) + .ok() + } + + fn retry_delay( + &self, + retry_number: usize, + headers: Option<&agentkit_http::HeaderMap>, + ) -> Duration { + self.resilience + .as_ref() + .expect("a reserved retry has resilience") + .retry_delay(retry_number, headers) + } + + async fn wait_before_retry( + &self, + retry_number: usize, + headers: Option<&agentkit_http::HeaderMap>, + ) -> Result<(), LoopError> { + self.wait_for_retry_delay(self.retry_delay(retry_number, headers)) + .await + } + + async fn wait_for_retry_delay(&self, delay: Duration) -> Result<(), LoopError> { + run_bounded( + async { + resilience_sleep(delay).await; + Ok(()) + }, + None, + self.deadline.as_ref(), + "retry backoff", + ) + .await + .map_err(|error| LoopError::Provider(format!("retry backoff failed: {error}"))) + } + + async fn execute_attempt( + &self, + request: agentkit_http::HttpRequest, + ) -> Result { + let timeout = self + .resilience + .as_ref() + .and_then(|config| config.attempt_timeout); + run_bounded( + self.client.execute(request), + timeout, + self.deadline.as_ref(), + "HTTP request attempt", + ) + .await + } + + async fn read_response_text(&self, response: HttpResponse) -> Result { + let timeout = self + .resilience + .as_ref() + .and_then(|config| config.attempt_timeout); + run_bounded( + response.text(), + timeout, + self.deadline.as_ref(), + "HTTP response body", + ) + .await + } + + async fn open_response(&self) -> Result { + let provider_name = self.provider.provider_name().to_owned(); + loop { + let authentication = self.authentication_attempt().await?; + let mut builder = self + .client + .post(self.provider.endpoint_url()) + .header("Content-Type", "application/json"); + builder = self.provider.preprocess_request(builder); + if self.provider.streaming() { + builder = builder.header("Accept", "text/event-stream"); + } + if let Some(authentication) = authentication.as_ref() { + builder = builder.headers(authentication.headers().clone()); + } + let request = builder.body(self.body.clone()).build().map_err(|error| { + LoopError::Provider(format!("{provider_name} request failed: {error}")) + })?; + + match self.execute_attempt(request).await { + Ok(response) => { + if response.status() == StatusCode::UNAUTHORIZED + && let Some(authentication) = authentication.as_ref() + && !self.reauthenticated.swap(true, Ordering::SeqCst) + { + drop(response); + self.reauthenticate(authentication.clone()).await?; + continue; + } + if is_retryable_status(response.status()) + && let Some(retry_number) = self.reserve_retry() + { + let delay = self.retry_delay(retry_number, Some(response.headers())); + drop(response); + self.wait_for_retry_delay(delay).await?; + continue; + } + return Ok(response); + } + Err(error) if error.is_retryable_transport() => { + if let Some(retry_number) = self.reserve_retry() { + self.wait_before_retry(retry_number, None).await?; + continue; + } + return Err(LoopError::Provider(format!( + "{provider_name} request failed: {error}" + ))); + } + Err(error) => { + return Err(LoopError::Provider(format!( + "{provider_name} request failed: {error}" + ))); + } + } + } + } + + async fn replay_stream(&self) -> Result<(BodyStream, TruncatedStreamDetector), LoopError> { + let retry_number = self.reserve_retry().ok_or_else(|| { + LoopError::Provider( + "completions stream failed before output and retry budget is exhausted".into(), + ) + })?; + self.wait_before_retry(retry_number, None).await?; + let response = self.open_response().await?; + if !response.status().is_success() { + return Err(LoopError::Provider(format!( + "{} stream replay failed with status {}", + self.provider.provider_name(), + response.status() + ))); + } + let integrity = TruncatedStreamDetector::from_headers(response.headers()); + Ok((response.bytes_stream(), integrity)) + } +} + #[async_trait] impl ModelAdapter for CompletionsAdapter

{ type Session = CompletionsSession

; @@ -277,6 +559,8 @@ impl ModelAdapter for CompletionsAdapter

{ provider: self.provider.clone(), model, provider_label: self.provider_label.clone(), + authentication: self.authentication.clone(), + resilience: self.resilience.clone(), _session_config: config, }) } @@ -297,55 +581,92 @@ impl ModelSession for CompletionsSession

{ ) -> Result { let provider = self.provider.clone(); let provider_name = provider.provider_name().to_owned(); + let deadline = self + .resilience + .as_ref() + .map(|config| LogicalDeadline::new(config.retry_budget)); let request_future = async { let body = request::build_request_body(provider.as_ref(), &turn_request) .map_err(|e| LoopError::Provider(e.to_string()))?; + let body = serde_json::to_vec(&body) + .map(Bytes::from) + .map_err(|e| LoopError::Provider(format!("failed to serialize request: {e}")))?; + let replay_request = ReplayRequest { + client: self.client.clone(), + provider: provider.clone(), + body, + authentication: self.authentication.clone(), + authentication_attempt: Arc::new(std::sync::Mutex::new(None)), + reauthenticated: Arc::new(AtomicBool::new(false)), + resilience: self.resilience.clone(), + deadline: deadline.clone(), + retries_used: Arc::new(AtomicUsize::new(0)), + }; + + loop { + let response = replay_request.open_response().await?; + let status = response.status(); + if provider.streaming() && status.is_success() { + let provider_for_postprocess = provider.clone(); + let postprocess: PostprocessResponse = Arc::new(move |usage, metadata, raw| { + provider_for_postprocess.postprocess_response(usage, metadata, raw); + }); + let integrity = TruncatedStreamDetector::from_headers(response.headers()); + let idle_timeout = self + .resilience + .as_ref() + .and_then(|config| config.stream_idle_timeout); + let replay = self.resilience.as_ref().and_then(|config| { + (config.max_retries > 0).then(|| { + let replay_request = replay_request.clone(); + Box::new(move || { + let replay_request = replay_request.clone(); + Box::pin(async move { replay_request.replay_stream().await }) + as BoxFuture<'static, _> + }) as StreamReplay + }) + }); + return Ok(CompletionsTurn::streaming( + response.bytes_stream(), + postprocess, + idle_timeout, + self.resilience.is_some(), + deadline.clone(), + integrity, + replay, + )); + } - let http = self - .client - .post(provider.endpoint_url()) - .header("Content-Type", "application/json"); - - let mut http = provider.preprocess_request(http); - if provider.streaming() { - http = http.header("Accept", "text/event-stream"); - } - - let response = http.json(&body).send().await.map_err(|error| { - LoopError::Provider(format!("{provider_name} request failed: {error}")) - })?; - - let status = response.status(); - if provider.streaming() && status.is_success() { - let provider_for_postprocess = provider.clone(); - let postprocess: PostprocessResponse = Arc::new(move |usage, metadata, raw| { - provider_for_postprocess.postprocess_response(usage, metadata, raw); - }); - return Ok(CompletionsTurn::streaming( - response.bytes_stream(), - postprocess, - )); - } - - let body = response.text().await.map_err(|error| { - LoopError::Provider(format!( - "failed to read {provider_name} response body: {error}" - )) - })?; - - provider.preprocess_response(status, &body)?; + let body_result = replay_request.read_response_text(response).await; + let body = match body_result { + Ok(body) => body, + Err(error) if is_retryable_body_read(status, &error) => { + if let Some(retry_number) = replay_request.reserve_retry() { + replay_request.wait_before_retry(retry_number, None).await?; + continue; + } + return Err(LoopError::Provider(format!( + "failed to read {provider_name} response body: {error}" + ))); + } + Err(error) => { + return Err(LoopError::Provider(format!( + "failed to read {provider_name} response body: {error}" + ))); + } + }; - if !status.is_success() { - return Err(LoopError::Provider(format!( - "{provider_name} request failed with status {status}: {body}" - ))); + provider.preprocess_response(status, &body)?; + if !status.is_success() { + return Err(LoopError::Provider(format!( + "{provider_name} request failed with status {status}: {body}" + ))); + } + let (events, _raw) = response::build_turn_from_response(provider.as_ref(), &body) + .map_err(|e| LoopError::Provider(e.to_string()))?; + return Ok(CompletionsTurn::buffered(events)); } - - let (events, _raw) = response::build_turn_from_response(provider.as_ref(), &body) - .map_err(|e| LoopError::Provider(e.to_string()))?; - - Ok(CompletionsTurn::buffered(events)) }; if let Some(cancellation) = cancellation { @@ -391,6 +712,12 @@ impl ModelTurn for CompletionsTurn { translator, pending, eof, + visible_output, + detect_truncation, + idle_timeout, + deadline, + integrity, + replay, postprocess, } = state.as_mut(); next_streaming_event( @@ -399,6 +726,12 @@ impl ModelTurn for CompletionsTurn { translator, pending, eof, + visible_output, + *detect_truncation, + *idle_timeout, + deadline.as_ref(), + integrity, + replay, postprocess, cancellation, ) @@ -408,38 +741,56 @@ impl ModelTurn for CompletionsTurn { } } +async fn next_stream_chunk( + body: &mut BodyStream, + idle_timeout: Option, + deadline: Option<&LogicalDeadline>, + cancellation: Option<&TurnCancellation>, +) -> Result, HttpError>, LoopError> { + let next = next_body_chunk_bounded(body, idle_timeout, deadline); + futures_util::pin_mut!(next); + if let Some(cancellation) = cancellation { + let cancelled = cancellation.cancelled(); + futures_util::pin_mut!(cancelled); + match select(next, cancelled).await { + Either::Left((chunk, _)) => Ok(chunk), + Either::Right((_, _)) => Err(LoopError::Cancelled), + } + } else { + Ok(next.await) + } +} + +#[allow(clippy::too_many_arguments)] async fn next_streaming_event( body: &mut BodyStream, decoder: &mut SseDecoder, translator: &mut EventTranslator, pending: &mut VecDeque, eof: &mut bool, + visible_output: &mut bool, + detect_truncation: bool, + idle_timeout: Option, + deadline: Option<&LogicalDeadline>, + integrity: &mut TruncatedStreamDetector, + replay: &mut Option, postprocess: &PostprocessResponse, cancellation: Option, ) -> Result, LoopError> { loop { if let Some(event) = pending.pop_front() { + *visible_output = true; return Ok(Some(event)); } if *eof || translator.is_done() { return Ok(None); } - let chunk = if let Some(cancellation) = cancellation.as_ref() { - let next = body.next(); - futures_util::pin_mut!(next); - let cancelled = cancellation.cancelled(); - futures_util::pin_mut!(cancelled); - match select(next, cancelled).await { - Either::Left((chunk, _)) => chunk, - Either::Right((_, _)) => return Err(LoopError::Cancelled), - } - } else { - body.next().await - }; + let chunk = next_stream_chunk(body, idle_timeout, deadline, cancellation.as_ref()).await?; - match chunk { - Some(Ok(bytes)) => { + let failure = match chunk { + Ok(Some(bytes)) => { + integrity.observe(&bytes); let text = std::str::from_utf8(&bytes).map_err(|e| { LoopError::Provider(format!("invalid UTF-8 in completions stream: {e}")) })?; @@ -451,15 +802,54 @@ async fn next_streaming_event( pending.push_back(event); } } + None } - Some(Err(e)) => { - return Err(LoopError::Provider(format!( - "completions stream body error: {e}" - ))); + Ok(None) if !detect_truncation => { + *eof = true; + None } - None => { + Ok(None) => { *eof = true; + let message = match integrity.finish() { + Ok(()) => "completions stream ended before its terminal event".to_owned(), + Err(error) => format!("completions stream body error: {error}"), + }; + Some(LoopError::Provider(message)) + } + Err(error) if !error.is_retryable_transport() => { + return Err(LoopError::Provider(format!( + "completions stream body error: {error}" + ))); } + Err(error) => Some(LoopError::Provider(format!( + "completions stream body error: {error}" + ))), + }; + + let Some(failure) = failure else { + continue; + }; + if *visible_output || replay.is_none() { + return Err(failure); } + + let replay_future = replay.as_mut().expect("checked above")(); + futures_util::pin_mut!(replay_future); + let replayed = if let Some(cancellation) = cancellation.as_ref() { + let cancelled = cancellation.cancelled(); + futures_util::pin_mut!(cancelled); + match select(replay_future, cancelled).await { + Either::Left((result, _)) => result, + Either::Right((_, _)) => return Err(LoopError::Cancelled), + } + } else { + replay_future.await + }?; + *body = replayed.0; + *integrity = replayed.1; + *decoder = SseDecoder::new(); + *translator = EventTranslator::new(); + pending.clear(); + *eof = false; } } diff --git a/crates/agentkit-adapter-completions/src/resilience_tests.rs b/crates/agentkit-adapter-completions/src/resilience_tests.rs new file mode 100644 index 0000000..7a7211d --- /dev/null +++ b/crates/agentkit-adapter-completions/src/resilience_tests.rs @@ -0,0 +1,393 @@ +use std::collections::VecDeque; +use std::future::Future; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::task::{Context, Poll, Wake, Waker}; + +use agentkit_http::{ + Authentication, AuthenticationAttempt, AuthenticationProvider, Bytes, HeaderMap, HeaderValue, + HttpClient, HttpError, HttpRequest, HttpResponse, ResilienceConfig, StatusCode, + TruncatedStreamDetector, header, +}; +use agentkit_loop::{ModelTurn, ModelTurnEvent}; +use async_trait::async_trait; +use futures_util::stream; +use serde::Serialize; + +use super::*; + +fn block_on(future: F) -> F::Output { + struct ThreadWake(std::thread::Thread); + impl Wake for ThreadWake { + fn wake(self: Arc) { + self.0.unpark(); + } + fn wake_by_ref(self: &Arc) { + self.0.unpark(); + } + } + let waker = Waker::from(Arc::new(ThreadWake(std::thread::current()))); + let mut context = Context::from_waker(&waker); + let mut future = Box::pin(future); + loop { + match future.as_mut().poll(&mut context) { + Poll::Ready(output) => return output, + Poll::Pending => std::thread::park(), + } + } +} + +#[derive(Clone, Serialize)] +struct Config { + model: &'static str, +} + +#[derive(Clone)] +struct Provider; + +impl CompletionsProvider for Provider { + type Config = Config; + + fn provider_name(&self) -> &str { + "test" + } + fn endpoint_url(&self) -> &str { + "https://example.test/chat" + } + fn config(&self) -> &Self::Config { + static CONFIG: Config = Config { model: "test" }; + &CONFIG + } + fn streaming(&self) -> bool { + false + } +} + +struct SequenceClient { + statuses: std::sync::Mutex>, + bodies: std::sync::Mutex>, +} + +#[async_trait] +impl HttpClient for SequenceClient { + async fn execute(&self, request: HttpRequest) -> Result { + self.bodies + .lock() + .unwrap() + .push(request.body.unwrap_or_default()); + let status = self.statuses.lock().unwrap().pop_front().unwrap(); + Ok(HttpResponse::new( + status, + HeaderMap::new(), + request.url, + Box::pin(stream::empty()), + )) + } +} + +struct RefreshingAuth { + calls: Arc, +} + +struct NeverAuth; + +#[async_trait] +impl AuthenticationProvider for NeverAuth { + async fn authenticate( + &self, + _previous: Option<&AuthenticationAttempt>, + ) -> Result { + futures_util::future::pending().await + } +} + +struct HangingClient; + +#[async_trait] +impl HttpClient for HangingClient { + async fn execute(&self, _request: HttpRequest) -> Result { + futures_util::future::pending().await + } +} + +#[async_trait] +impl AuthenticationProvider for RefreshingAuth { + async fn authenticate( + &self, + previous: Option<&AuthenticationAttempt>, + ) -> Result { + let generation = self.calls.fetch_add(1, Ordering::SeqCst); + assert_eq!( + previous + .and_then(|attempt| attempt.state::()) + .copied(), + generation.checked_sub(1) + ); + let mut headers = HeaderMap::new(); + headers.insert( + header::AUTHORIZATION, + HeaderValue::from_static("Bearer redacted"), + ); + Ok(AuthenticationAttempt::new(headers, generation)) + } +} + +fn replay_request( + client: Arc, + authentication: Option, + resilience: Option, +) -> ReplayRequest { + let deadline = resilience + .as_ref() + .map(|config| LogicalDeadline::new(config.retry_budget)); + ReplayRequest { + client: Http::from_arc(client), + provider: Arc::new(Provider), + body: Bytes::from_static(b"stable request bytes"), + authentication, + authentication_attempt: Arc::new(std::sync::Mutex::new(None)), + reauthenticated: Arc::new(AtomicBool::new(false)), + resilience, + deadline, + retries_used: Arc::new(AtomicUsize::new(0)), + } +} + +#[test] +fn performs_only_one_reactive_401_reauthentication() { + block_on(async { + let client = Arc::new(SequenceClient { + statuses: std::sync::Mutex::new(VecDeque::from([ + StatusCode::UNAUTHORIZED, + StatusCode::UNAUTHORIZED, + StatusCode::OK, + ])), + bodies: std::sync::Mutex::new(Vec::new()), + }); + let auth_calls = Arc::new(AtomicUsize::new(0)); + let request = replay_request( + client.clone(), + Some(Authentication::new(RefreshingAuth { + calls: auth_calls.clone(), + })), + None, + ); + + let response = request.open_response().await.unwrap(); + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + assert_eq!(auth_calls.load(Ordering::SeqCst), 2); + assert_eq!(client.bodies.lock().unwrap().len(), 2); + }); +} + +#[test] +fn retries_retryable_status_with_identical_body_bytes() { + block_on(async { + let client = Arc::new(SequenceClient { + statuses: std::sync::Mutex::new(VecDeque::from([ + StatusCode::SERVICE_UNAVAILABLE, + StatusCode::OK, + ])), + bodies: std::sync::Mutex::new(Vec::new()), + }); + let resilience = ResilienceConfig { + max_retries: 1, + retry_budget: Duration::from_secs(1), + attempt_timeout: None, + stream_idle_timeout: None, + initial_backoff: Duration::ZERO, + max_backoff: Duration::ZERO, + }; + let request = replay_request(client.clone(), None, Some(resilience)); + + assert_eq!( + request.open_response().await.unwrap().status(), + StatusCode::OK + ); + let bodies = client.bodies.lock().unwrap(); + assert_eq!( + bodies.as_slice(), + [ + Bytes::from_static(b"stable request bytes"), + Bytes::from_static(b"stable request bytes") + ] + ); + }); +} + +#[test] +fn logical_retry_budget_bounds_authentication_attempt_body_backoff_and_stream() { + block_on(async { + let short = || ResilienceConfig { + max_retries: 1, + retry_budget: Duration::from_millis(20), + attempt_timeout: Some(Duration::from_secs(10)), + stream_idle_timeout: None, + initial_backoff: Duration::from_secs(10), + max_backoff: Duration::from_secs(10), + }; + let assert_budget = |error: LoopError| { + assert!( + error.to_string().contains("logical request retry budget"), + "unexpected error: {error}" + ); + }; + + let auth_client = Arc::new(SequenceClient { + statuses: std::sync::Mutex::new(VecDeque::from([StatusCode::OK])), + bodies: std::sync::Mutex::new(Vec::new()), + }); + let auth_request = replay_request( + auth_client, + Some(Authentication::new(NeverAuth)), + Some(short()), + ); + assert_budget(auth_request.open_response().await.unwrap_err()); + + let attempt_config = short(); + let attempt_request = ReplayRequest { + client: Http::new(HangingClient), + provider: Arc::new(Provider), + body: Bytes::from_static(b"body"), + authentication: None, + authentication_attempt: Arc::new(std::sync::Mutex::new(None)), + reauthenticated: Arc::new(AtomicBool::new(false)), + deadline: Some(LogicalDeadline::new(attempt_config.retry_budget)), + resilience: Some(attempt_config), + retries_used: Arc::new(AtomicUsize::new(0)), + }; + assert_budget(attempt_request.open_response().await.unwrap_err()); + + let body_request = replay_request( + Arc::new(SequenceClient { + statuses: std::sync::Mutex::new(VecDeque::new()), + bodies: std::sync::Mutex::new(Vec::new()), + }), + None, + Some(short()), + ); + let pending_body: BodyStream = Box::pin(stream::pending()); + let response = HttpResponse::new( + StatusCode::OK, + HeaderMap::new(), + "https://example.test".into(), + pending_body, + ); + let body_error = body_request.read_response_text(response).await.unwrap_err(); + assert!( + body_error + .to_string() + .contains("logical request retry budget") + ); + + let backoff_client = Arc::new(SequenceClient { + statuses: std::sync::Mutex::new(VecDeque::from([ + StatusCode::SERVICE_UNAVAILABLE, + StatusCode::OK, + ])), + bodies: std::sync::Mutex::new(Vec::new()), + }); + let backoff_request = replay_request(backoff_client.clone(), None, Some(short())); + assert_budget(backoff_request.open_response().await.unwrap_err()); + assert_eq!(backoff_client.bodies.lock().unwrap().len(), 1); + + let stream_deadline = LogicalDeadline::new(Duration::from_millis(20)); + let mut turn = CompletionsTurn::streaming( + Box::pin(stream::pending()), + Arc::new(|_, _, _| {}), + None, + true, + Some(stream_deadline), + TruncatedStreamDetector::default(), + None, + ); + assert_budget(turn.next_event(None).await.unwrap_err()); + }); +} + +#[test] +fn no_resilience_preserves_legacy_non_terminal_eof() { + block_on(async { + let mut turn = CompletionsTurn::streaming( + Box::pin(stream::empty()), + Arc::new(|_, _, _| {}), + None, + false, + None, + TruncatedStreamDetector::default(), + None, + ); + assert!(turn.next_event(None).await.unwrap().is_none()); + }); +} + +#[test] +fn stream_does_not_retry_non_transient_body_errors() { + block_on(async { + let replay_calls = Arc::new(AtomicUsize::new(0)); + let replay_counter = replay_calls.clone(); + let replay: StreamReplay = Box::new(move || { + replay_counter.fetch_add(1, Ordering::SeqCst); + Box::pin(async { unreachable!("non-transient errors must not replay") }) + }); + let initial = stream::iter([Err(HttpError::Other("protocol failure".into()))]); + let mut turn = CompletionsTurn::streaming( + Box::pin(initial), + Arc::new(|_, _, _| {}), + None, + true, + Some(LogicalDeadline::new(Duration::from_secs(1))), + TruncatedStreamDetector::default(), + Some(replay), + ); + + let error = turn.next_event(None).await.unwrap_err(); + assert!(error.to_string().contains("protocol failure")); + assert_eq!(replay_calls.load(Ordering::SeqCst), 0); + }); +} + +#[test] +fn stream_retries_before_output_but_never_after_output() { + block_on(async { + let replay_calls = Arc::new(AtomicUsize::new(0)); + let replay_counter = replay_calls.clone(); + let replay: StreamReplay = Box::new(move || { + replay_counter.fetch_add(1, Ordering::SeqCst); + let body: BodyStream = + Box::pin(stream::iter([Ok::<_, HttpError>(Bytes::from_static( + b"data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"}}]}\n\n", + ))])); + Box::pin(async move { Ok((body, TruncatedStreamDetector::default())) }) + }); + let initial = stream::iter([Err(HttpError::body(std::io::Error::other("disconnected")))]); + let mut turn = CompletionsTurn::streaming( + Box::pin(initial), + Arc::new(|_, _, _| {}), + None, + true, + Some(LogicalDeadline::new(Duration::from_secs(1))), + TruncatedStreamDetector::default(), + Some(replay), + ); + + assert!(matches!( + turn.next_event(None).await.unwrap(), + Some(ModelTurnEvent::Delta(_)) + )); + assert_eq!(replay_calls.load(Ordering::SeqCst), 1); + let mut failure = None; + for _ in 0..4 { + match turn.next_event(None).await { + Err(error) => { + failure = Some(error); + break; + } + Ok(Some(_)) => {} + Ok(None) => panic!("truncated stream was accepted"), + } + } + assert!(failure.unwrap().to_string().contains("terminal event")); + assert_eq!(replay_calls.load(Ordering::SeqCst), 1); + }); +} diff --git a/crates/agentkit-http/Cargo.toml b/crates/agentkit-http/Cargo.toml index eeaab8f..6b1be74 100644 --- a/crates/agentkit-http/Cargo.toml +++ b/crates/agentkit-http/Cargo.toml @@ -16,11 +16,15 @@ reqwest-middleware-client = ["dep:reqwest", "dep:reqwest-middleware"] [dependencies] async-trait.workspace = true bytes = "1" +fastrand = "2.5.0" +futures-timer.workspace = true futures-util.workspace = true http = "1" +httpdate = "1.0.3" serde.workspace = true serde_json.workspace = true thiserror.workspace = true +zeroize = "1.8.2" reqwest = { workspace = true, optional = true } reqwest-middleware = { version = "0.5", optional = true } diff --git a/crates/agentkit-http/README.md b/crates/agentkit-http/README.md index ed8ed2f..705e2cf 100644 --- a/crates/agentkit-http/README.md +++ b/crates/agentkit-http/README.md @@ -50,6 +50,24 @@ let body: serde_json::Value = response.json().await?; `request` constructors and is cheap to clone — it holds an `Arc` over the underlying client. +## Authentication and resilience + +`Authentication` is a cloneable, type-erased handle over an async +`AuthenticationProvider`. Providers return sensitive headers plus opaque attempt +state, which lets an adapter perform one reactive 401 refresh without learning +how credentials are stored. Optional non-secret credential bindings let +replay-sensitive adapters bind continuation data to an account or credential +generation. Bare `String`, `Box`, `Arc`, `Cow`, and `&str` values +convert into bearer authentication; borrowed values are copied into owned +zeroizing storage. Bearer and arbitrary-header credentials remain redacted, and +owned credential storage is zeroized when dropped. + +`ResilienceConfig` is opt-in. It combines a wall-clock retry budget, an optional +attempt cap and attempt/stream-idle timeouts with capped exponential full-jitter +backoff and validated `Retry-After`/rate-limit reset hints. Adapters retain +control of replay safety; configuring resilience on `Http` alone does not replay +requests. + ## Bring your own client Any type that implements `HttpClient` plugs in unchanged. This is the seam diff --git a/crates/agentkit-http/src/authentication.rs b/crates/agentkit-http/src/authentication.rs new file mode 100644 index 0000000..3759df3 --- /dev/null +++ b/crates/agentkit-http/src/authentication.rs @@ -0,0 +1,299 @@ +use std::any::Any; +use std::borrow::Cow; +use std::fmt; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use async_trait::async_trait; +use http::{HeaderMap, HeaderName, HeaderValue, header}; +use zeroize::Zeroizing; + +use crate::HttpError; + +/// The headers and provider-private state produced by one authentication attempt. +/// +/// The state is deliberately opaque to callers. An authentication provider can +/// recover it with [`AuthenticationAttempt::state`] when a 401 response asks it +/// to refresh credentials. Providers can additionally attach a stable, non-secret +/// credential identity or generation with [`AuthenticationAttempt::with_binding`]. +#[derive(Clone)] +pub struct AuthenticationAttempt { + headers: HeaderMap, + state: Arc, + binding: Option>, +} + +impl AuthenticationAttempt { + /// Creates an attempt with provider-private state. + pub fn new(mut headers: HeaderMap, state: T) -> Self + where + T: Any + Send + Sync, + { + mark_sensitive(&mut headers); + Self { + headers, + state: Arc::new(state), + binding: None, + } + } + + /// Creates an attempt which does not need provider-private state. + pub fn stateless(headers: HeaderMap) -> Self { + Self::new(headers, ()) + } + + /// Attaches a stable, non-secret credential identity or generation. + /// + /// The binding is intended for replay/continuation checks. It must not be a + /// raw bearer token, API key, header value, or derivative of secret material. + pub fn with_binding(mut self, binding: impl Into>) -> Self { + self.binding = Some(binding.into()); + self + } + + /// Returns the stable, non-secret credential binding, when supplied. + pub fn binding(&self) -> Option<&str> { + self.binding.as_deref() + } + + /// Returns the authentication headers. Header values are marked sensitive. + pub fn headers(&self) -> &HeaderMap { + &self.headers + } + + /// Downcasts the opaque state saved by the provider on the prior attempt. + pub fn state(&self) -> Option<&T> { + self.state.downcast_ref() + } +} + +impl fmt::Debug for AuthenticationAttempt { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("AuthenticationAttempt") + .field("header_names", &self.headers.keys().collect::>()) + .field("state", &"") + .field("binding_present", &self.binding.is_some()) + .finish() + } +} + +/// Asynchronously supplies authentication headers and optional refresh state. +#[async_trait] +pub trait AuthenticationProvider: Send + Sync + 'static { + /// Authenticates a request. `previous` is `None` initially and contains the + /// exact opaque prior attempt when called reactively after a 401 response. + async fn authenticate( + &self, + previous: Option<&AuthenticationAttempt>, + ) -> Result; +} + +/// Clone-cheap, type-erased authentication provider handle. +/// +/// Converting a bare `String`, `Box`, `Arc`, `Cow`, or `&str` +/// into `Authentication` creates bearer authentication. Borrowed strings are +/// copied into owned storage whose bytes are zeroized on drop. +#[derive(Clone)] +pub struct Authentication { + inner: Arc, +} + +impl Authentication { + pub fn new(provider: P) -> Self { + Self { + inner: Arc::new(provider), + } + } + + pub fn from_arc(provider: Arc) -> Self { + Self { inner: provider } + } + + pub async fn authenticate( + &self, + previous: Option<&AuthenticationAttempt>, + ) -> Result { + self.inner.authenticate(previous).await + } + + /// Creates bearer authentication from an owned token. The retained token + /// bytes are overwritten when the last handle is dropped. + pub fn bearer(token: impl Into) -> Self { + Self::new(StaticAuthentication::bearer(token.into().into_bytes())) + } + + /// Creates bearer authentication from a static token. + /// + /// This compatibility helper copies the token into owned storage whose bytes + /// are overwritten when the last handle is dropped. + pub fn bearer_static(token: &'static str) -> Self { + Self::bearer(token.to_owned()) + } + + /// Creates authentication using one arbitrary secret header. + pub fn header(name: HeaderName, value: impl Into) -> Self { + Self::new(StaticAuthentication::header( + name, + value.into().into_bytes(), + )) + } + + /// Creates authentication using one arbitrary static secret header. + /// + /// This compatibility helper copies the value into owned zeroizing storage. + pub fn header_static(name: HeaderName, value: &'static str) -> Self { + Self::header(name, value.to_owned()) + } +} + +impl fmt::Debug for Authentication { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Authentication").finish_non_exhaustive() + } +} + +impl From for Authentication { + fn from(token: String) -> Self { + Self::bearer(token) + } +} + +impl From> for Authentication { + fn from(token: Box) -> Self { + Self::bearer(token.into_string()) + } +} + +impl From> for Authentication { + fn from(token: Arc) -> Self { + Self::bearer(token.as_ref().to_owned()) + } +} + +impl<'a> From> for Authentication { + fn from(token: Cow<'a, str>) -> Self { + match token { + Cow::Owned(token) => Self::bearer(token), + Cow::Borrowed(token) => Self::bearer(token.to_owned()), + } + } +} + +impl From<&str> for Authentication { + fn from(token: &str) -> Self { + Self::bearer(token.to_owned()) + } +} + +impl From<&String> for Authentication { + fn from(token: &String) -> Self { + Self::bearer(token.to_owned()) + } +} + +impl From<(HeaderName, HeaderValue)> for Authentication { + fn from((name, value): (HeaderName, HeaderValue)) -> Self { + Self::new(FixedHeaders::one(name, value)) + } +} + +impl From for Authentication { + fn from(headers: HeaderMap) -> Self { + Self::new(FixedHeaders(headers)) + } +} + +struct StaticAuthentication { + name: HeaderName, + value: Zeroizing>, + bearer: bool, + binding: Arc, +} + +fn static_authentication_binding() -> Arc { + static NEXT_ID: AtomicU64 = AtomicU64::new(1); + let created_at = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + format!( + "static-auth-v1-{}-{created_at:x}-{:x}", + std::process::id(), + NEXT_ID.fetch_add(1, Ordering::Relaxed) + ) + .into() +} + +impl StaticAuthentication { + fn bearer(value: Vec) -> Self { + Self { + name: header::AUTHORIZATION, + value: Zeroizing::new(value), + bearer: true, + binding: static_authentication_binding(), + } + } + fn header(name: HeaderName, value: Vec) -> Self { + Self { + name, + value: Zeroizing::new(value), + bearer: false, + binding: static_authentication_binding(), + } + } + fn value(&self) -> Result { + let bytes = self.value.as_slice(); + let rendered = if self.bearer { + let mut rendered = Zeroizing::new(Vec::with_capacity(7 + bytes.len())); + rendered.extend_from_slice(b"Bearer "); + rendered.extend_from_slice(bytes); + rendered + } else { + Zeroizing::new(bytes.to_vec()) + }; + let result = HeaderValue::from_bytes(&rendered) + .map_err(|error| HttpError::InvalidHeader(format!("authentication value: {error}"))); + let mut value = result?; + value.set_sensitive(true); + Ok(value) + } +} + +#[async_trait] +impl AuthenticationProvider for StaticAuthentication { + async fn authenticate( + &self, + _previous: Option<&AuthenticationAttempt>, + ) -> Result { + let mut headers = HeaderMap::new(); + headers.insert(self.name.clone(), self.value()?); + Ok(AuthenticationAttempt::stateless(headers).with_binding(self.binding.clone())) + } +} + +struct FixedHeaders(HeaderMap); +impl FixedHeaders { + fn one(name: HeaderName, mut value: HeaderValue) -> Self { + value.set_sensitive(true); + let mut headers = HeaderMap::new(); + headers.insert(name, value); + Self(headers) + } +} + +#[async_trait] +impl AuthenticationProvider for FixedHeaders { + async fn authenticate( + &self, + _previous: Option<&AuthenticationAttempt>, + ) -> Result { + Ok(AuthenticationAttempt::stateless(self.0.clone())) + } +} + +fn mark_sensitive(headers: &mut HeaderMap) { + for value in headers.values_mut() { + value.set_sensitive(true); + } +} diff --git a/crates/agentkit-http/src/error.rs b/crates/agentkit-http/src/error.rs index 41e7064..3752bb4 100644 --- a/crates/agentkit-http/src/error.rs +++ b/crates/agentkit-http/src/error.rs @@ -1,4 +1,5 @@ use std::error::Error as StdError; +use std::time::Duration; use thiserror::Error; @@ -24,6 +25,15 @@ pub enum HttpError { #[error("response body read failed: {0}")] Body(#[source] BoxError), + #[error("{operation} timed out after {timeout:?}")] + Timeout { + operation: &'static str, + timeout: Duration, + }, + + #[error("response body was truncated (expected {expected} bytes, received {received})")] + TruncatedBody { expected: u64, received: u64 }, + #[error("{0}")] Other(String), } @@ -42,4 +52,13 @@ impl HttpError { { Self::Body(Box::new(err)) } + + /// Whether the failure is transient at the HTTP transport layer. + /// The adapter still decides whether replaying its request is safe. + pub fn is_retryable_transport(&self) -> bool { + matches!( + self, + Self::Request(_) | Self::Body(_) | Self::Timeout { .. } | Self::TruncatedBody { .. } + ) + } } diff --git a/crates/agentkit-http/src/lib.rs b/crates/agentkit-http/src/lib.rs index 03ba80b..f8af347 100644 --- a/crates/agentkit-http/src/lib.rs +++ b/crates/agentkit-http/src/lib.rs @@ -5,9 +5,11 @@ //! streaming). Enable the `reqwest-client` feature (default) for an //! `impl HttpClient for reqwest::Client`; disable it to compile trait-only. +mod authentication; mod client; mod error; mod request; +mod resilience; mod response; #[cfg(feature = "reqwest-client")] @@ -16,11 +18,17 @@ mod reqwest_impl; #[cfg(feature = "reqwest-middleware-client")] mod reqwest_middleware_impl; +pub use authentication::{Authentication, AuthenticationAttempt, AuthenticationProvider}; pub use client::{Http, HttpClient}; pub use error::{BoxError, HttpError}; pub use request::{HttpRequest, HttpRequestBuilder}; +pub use resilience::{ + LogicalDeadline, ResilienceConfig, TruncatedStreamDetector, is_retryable_body_read, + is_retryable_status, next_body_chunk, next_body_chunk_bounded, retry_hint, run_bounded, sleep, +}; pub use response::{BodyStream, HttpResponse}; +pub use bytes::Bytes; pub use http::{HeaderMap, HeaderName, HeaderValue, Method, StatusCode, Uri, header}; #[cfg(test)] @@ -29,6 +37,7 @@ mod tests { use async_trait::async_trait; use bytes::Bytes; use futures_util::stream; + use std::borrow::Cow; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -94,6 +103,186 @@ mod tests { assert_eq!(decoded, Resp { ok: true }); } + #[tokio::test] + async fn type_erased_authentication_preserves_opaque_prior_state_and_redacts() { + struct Refreshing(AtomicUsize); + + #[async_trait] + impl AuthenticationProvider for Refreshing { + async fn authenticate( + &self, + previous: Option<&AuthenticationAttempt>, + ) -> Result { + let generation = self.0.fetch_add(1, Ordering::SeqCst); + assert_eq!( + previous + .and_then(|attempt| attempt.state::()) + .copied(), + generation.checked_sub(1) + ); + let mut headers = HeaderMap::new(); + headers.insert( + header::AUTHORIZATION, + HeaderValue::from_str(&format!("Bearer secret-{generation}")).unwrap(), + ); + Ok(AuthenticationAttempt::new(headers, generation) + .with_binding(format!("credential-generation-{generation}"))) + } + } + + let authentication = Authentication::new(Refreshing(AtomicUsize::new(0))); + let first = authentication.authenticate(None).await.unwrap(); + let second = authentication.authenticate(Some(&first)).await.unwrap(); + assert_eq!(second.state::(), Some(&1)); + assert_eq!(first.binding(), Some("credential-generation-0")); + assert_eq!(second.binding(), Some("credential-generation-1")); + assert_eq!(second.clone().binding(), second.binding()); + let debug = format!("{authentication:?} {second:?}"); + assert!(!debug.contains("secret")); + assert!(!debug.contains("credential-generation-1")); + assert!(debug.contains("binding_present: true")); + assert!(second.headers()[header::AUTHORIZATION].is_sensitive()); + + let borrowed_source = String::from("borrowed-secret"); + let borrowed: Authentication = borrowed_source.as_str().into(); + drop(borrowed_source); + + let borrowed_string_source = String::from("borrowed-string-secret"); + let borrowed_string: Authentication = (&borrowed_string_source).into(); + drop(borrowed_string_source); + + let cow_borrowed_source = String::from("cow-borrowed-secret"); + let cow_borrowed: Authentication = Cow::Borrowed(cow_borrowed_source.as_str()).into(); + drop(cow_borrowed_source); + + let bearers = [ + (String::from("string-secret").into(), "string-secret"), + ( + String::from("box-secret").into_boxed_str().into(), + "box-secret", + ), + (Arc::::from("arc-secret").into(), "arc-secret"), + ( + Cow::<'static, str>::Owned(String::from("cow-owned-secret")).into(), + "cow-owned-secret", + ), + (borrowed, "borrowed-secret"), + (borrowed_string, "borrowed-string-secret"), + (cow_borrowed, "cow-borrowed-secret"), + ("static-secret".into(), "static-secret"), + ]; + for (bearer, secret) in bearers { + let attempt = bearer.authenticate(None).await.unwrap(); + assert_eq!( + attempt.headers()[header::AUTHORIZATION].to_str().unwrap(), + format!("Bearer {secret}") + ); + assert!(attempt.headers()[header::AUTHORIZATION].is_sensitive()); + assert!(attempt.binding().is_some()); + assert_eq!(attempt.clone().binding(), attempt.binding()); + assert!(!attempt.binding().unwrap().contains(secret)); + let debug = format!("{bearer:?} {attempt:?}"); + assert!(!debug.contains(secret)); + assert!(debug.contains("binding_present: true")); + } + + let header = Authentication::header( + HeaderName::from_static("x-api-key"), + String::from("owned-header-secret"), + ); + let header_attempt = header.authenticate(None).await.unwrap(); + assert!(header_attempt.binding().is_some()); + assert!( + !header_attempt + .binding() + .unwrap() + .contains("owned-header-secret") + ); + assert!(header_attempt.headers()["x-api-key"].is_sensitive()); + assert!(!format!("{header_attempt:?}").contains("owned-header-secret")); + } + + #[test] + fn resilience_classifies_hints_and_truncation() { + assert!(is_retryable_status(StatusCode::TOO_MANY_REQUESTS)); + assert!(is_retryable_status(StatusCode::BAD_GATEWAY)); + assert!(!is_retryable_status(StatusCode::BAD_REQUEST)); + let body_error = HttpError::body(std::io::Error::other("disconnected")); + assert!(is_retryable_body_read(StatusCode::OK, &body_error)); + assert!(!is_retryable_body_read( + StatusCode::BAD_REQUEST, + &body_error + )); + assert!(!is_retryable_body_read( + StatusCode::OK, + &HttpError::Other("protocol error".into()) + )); + + let mut headers = HeaderMap::new(); + headers.insert("retry-after", HeaderValue::from_static("3")); + assert_eq!( + retry_hint(&headers), + Some(std::time::Duration::from_secs(3)) + ); + headers.insert( + "retry-after", + HeaderValue::from_static("Wed, 21 Oct 2099 07:28:00 GMT"), + ); + assert!(retry_hint(&headers).is_some_and(|delay| !delay.is_zero())); + headers.insert( + "retry-after", + HeaderValue::from_static("Tue, 31 Feb 2099 07:28:00 GMT"), + ); + assert_eq!(retry_hint(&headers), None); + + headers.clear(); + headers.insert(header::CONTENT_LENGTH, HeaderValue::from_static("5")); + let mut detector = TruncatedStreamDetector::from_headers(&headers); + detector.observe(&Bytes::from_static(b"123")); + assert!(matches!( + detector.finish(), + Err(HttpError::TruncatedBody { + expected: 5, + received: 3 + }) + )); + } + + #[test] + fn retry_hints_handle_non_finite_and_overflow_values() { + let hint = |name: &'static str, value: &'static str| { + let mut headers = HeaderMap::new(); + headers.insert(name, HeaderValue::from_static(value)); + retry_hint(&headers) + }; + + assert_eq!(hint("retry-after", "NaN"), None); + assert_eq!(hint("retry-after", "inf"), None); + assert_eq!(hint("retry-after", "1e999"), None); + assert_eq!(hint("retry-after", "-inf"), None); + assert_eq!(hint("retry-after", "-1"), None); + assert_eq!(hint("x-ratelimit-reset", "NaN"), None); + assert_eq!(hint("x-ratelimit-reset", "inf"), None); + assert_eq!( + hint("x-ratelimit-reset-requests-day", "42s"), + Some(std::time::Duration::from_secs(42)) + ); + assert_eq!( + hint("x-ratelimit-reset-tokens-minute", "1.5s"), + Some(std::time::Duration::from_millis(1500)) + ); + + let config = ResilienceConfig { + initial_backoff: std::time::Duration::from_millis(5), + max_backoff: std::time::Duration::from_millis(5), + ..ResilienceConfig::default() + }; + let mut headers = HeaderMap::new(); + headers.insert("retry-after", HeaderValue::from_static("inf")); + assert!(config.retry_delay(0, Some(&headers)) <= std::time::Duration::from_millis(5)); + assert!(format!("{config:?}").contains("retry_budget")); + } + #[tokio::test] async fn error_for_status_flags_4xx() { let stub = StubClient { diff --git a/crates/agentkit-http/src/resilience.rs b/crates/agentkit-http/src/resilience.rs new file mode 100644 index 0000000..28e6c26 --- /dev/null +++ b/crates/agentkit-http/src/resilience.rs @@ -0,0 +1,293 @@ +use std::future::Future; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use bytes::Bytes; +use futures_timer::Delay; +use futures_util::StreamExt; +use futures_util::future::{Either, select}; +use http::{HeaderMap, StatusCode}; + +use crate::{BodyStream, HttpError}; + +/// Opt-in retry and timeout policy shared by HTTP-backed adapters. +/// `retry_budget` is the primary wall-clock budget for a logical request. +/// `max_retries` is an optional safety cap in addition to the initial try; the +/// default leaves the wall-clock budget in control. +/// `retry_budget` bounds the complete logical request, including authentication, +/// attempts, response reads, stream reads, and backoff. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ResilienceConfig { + pub max_retries: usize, + pub retry_budget: Duration, + pub attempt_timeout: Option, + pub stream_idle_timeout: Option, + pub initial_backoff: Duration, + pub max_backoff: Duration, +} + +impl Default for ResilienceConfig { + fn default() -> Self { + Self { + max_retries: usize::MAX, + retry_budget: Duration::from_secs(60), + attempt_timeout: Some(Duration::from_secs(30)), + stream_idle_timeout: Some(Duration::from_secs(30)), + initial_backoff: Duration::from_millis(200), + max_backoff: Duration::from_secs(10), + } + } +} + +impl ResilienceConfig { + pub fn no_retries() -> Self { + Self { + max_retries: 0, + ..Self::default() + } + } + + /// Computes capped exponential full-jitter backoff, honoring server hints. + /// `retry_number` starts at zero for the first retry. + pub fn retry_delay(&self, retry_number: usize, headers: Option<&HeaderMap>) -> Duration { + if let Some(hint) = headers.and_then(retry_hint) { + // Server-directed waits are not exponential backoff. The logical + // request deadline remains authoritative and bounds this delay. + return hint.min(self.retry_budget); + } + let shift = retry_number.min(31) as u32; + let cap = self + .initial_backoff + .saturating_mul(1_u32 << shift) + .min(self.max_backoff); + full_jitter(cap) + } +} + +/// HTTP statuses which are normally safe to retry when the adapter can replay. +pub fn is_retryable_status(status: StatusCode) -> bool { + matches!(status.as_u16(), 408 | 425 | 429 | 500 | 502 | 503 | 504) +} + +/// Whether a body-read failure can safely replay a successful response. +pub fn is_retryable_body_read(status: StatusCode, error: &HttpError) -> bool { + status.is_success() && error.is_retryable_transport() +} + +/// Parses Retry-After and common rate-limit reset headers into a delay. +pub fn retry_hint(headers: &HeaderMap) -> Option { + if let Some(value) = headers + .get("retry-after") + .and_then(|value| value.to_str().ok()) + { + if let Ok(seconds) = value.trim().parse::() + && let Some(delay) = duration_from_seconds(seconds) + { + return Some(delay); + } + if let Ok(timestamp) = httpdate::parse_http_date(value) { + return Some( + timestamp + .duration_since(SystemTime::now()) + .unwrap_or_default(), + ); + } + } + for name in [ + "ratelimit-reset", + "x-ratelimit-reset", + "x-rate-limit-reset", + "x-ratelimit-reset-requests-day", + "x-ratelimit-reset-tokens-minute", + ] { + if let Some(value) = header_number(headers, name) { + if value.is_nan() { + continue; + } + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs_f64(); + let seconds = if value > now - 60.0 { + value - now + } else { + value + }; + if let Some(delay) = duration_from_seconds(seconds) { + return Some(delay); + } + } + } + None +} + +fn duration_from_seconds(seconds: f64) -> Option { + if !seconds.is_finite() || seconds < 0.0 { + return None; + } + if seconds >= Duration::MAX.as_secs_f64() { + return Some(Duration::MAX); + } + Duration::try_from_secs_f64(seconds).ok() +} + +fn header_number(headers: &HeaderMap, name: &str) -> Option { + headers + .get(name)? + .to_str() + .ok()? + .trim() + .trim_end_matches('s') + .parse() + .ok() +} + +fn full_jitter(cap: Duration) -> Duration { + let cap_nanos = cap.as_nanos().min(u64::MAX as u128) as u64; + Duration::from_nanos(fastrand::u64(..=cap_nanos)) +} + +/// Wall-clock deadline shared by all work for one logical HTTP request. +#[derive(Clone, Debug)] +pub struct LogicalDeadline { + started_at: Instant, + budget: Duration, +} + +impl LogicalDeadline { + pub fn new(budget: Duration) -> Self { + Self { + started_at: Instant::now(), + budget, + } + } + + fn remaining(&self) -> Result { + let elapsed = self.started_at.elapsed(); + if elapsed >= self.budget { + Err(self.timeout_error()) + } else { + Ok(self.budget - elapsed) + } + } + + fn timeout_error(&self) -> HttpError { + HttpError::Timeout { + operation: "logical request retry budget", + timeout: self.budget, + } + } +} + +/// Runs an HTTP operation within its own timeout and the logical request deadline. +pub async fn run_bounded( + future: F, + operation_timeout: Option, + deadline: Option<&LogicalDeadline>, + operation: &'static str, +) -> Result +where + F: Future>, +{ + let remaining = deadline.map(LogicalDeadline::remaining).transpose()?; + let timeout = match (operation_timeout, remaining) { + (Some(operation), Some(remaining)) => operation.min(remaining), + (Some(operation), None) => operation, + (None, Some(remaining)) => remaining, + (None, None) => return future.await, + }; + let budget_limited = remaining + .is_some_and(|remaining| operation_timeout.is_none_or(|operation| remaining <= operation)); + + futures_util::pin_mut!(future); + let timeout_wait = sleep(timeout); + futures_util::pin_mut!(timeout_wait); + match select(future, timeout_wait).await { + Either::Left((result, _)) => result, + Either::Right((_, _)) if budget_limited => Err(deadline + .expect("budget-limited timeout has deadline") + .timeout_error()), + Either::Right((_, _)) => Err(HttpError::Timeout { operation, timeout }), + } +} + +/// Sleeps without tying the HTTP abstraction to a particular async runtime. +/// Dropping the returned future cancels the timer registration. +pub async fn sleep(duration: Duration) { + Delay::new(duration).await; +} + +/// Reads one body chunk and fails if the stream stays idle for `timeout`. +/// Dropping this future cancels the body poll and the caller's wait. +pub async fn next_body_chunk( + body: &mut BodyStream, + timeout: Option, +) -> Result, HttpError> { + next_body_chunk_bounded(body, timeout, None).await +} + +/// Reads one body chunk using a single timer for the idle and logical deadlines. +pub async fn next_body_chunk_bounded( + body: &mut BodyStream, + idle_timeout: Option, + deadline: Option<&LogicalDeadline>, +) -> Result, HttpError> { + let remaining = deadline.map(LogicalDeadline::remaining).transpose()?; + let timeout = match (idle_timeout, remaining) { + (Some(idle), Some(remaining)) => idle.min(remaining), + (Some(idle), None) => idle, + (None, Some(remaining)) => remaining, + (None, None) => return body.next().await.transpose(), + }; + let budget_limited = + remaining.is_some_and(|remaining| idle_timeout.is_none_or(|idle| remaining <= idle)); + + let next = body.next(); + futures_util::pin_mut!(next); + let timeout_wait = sleep(timeout); + futures_util::pin_mut!(timeout_wait); + match select(next, timeout_wait).await { + Either::Left((chunk, _)) => chunk.transpose(), + Either::Right((_, _)) if budget_limited => Err(deadline + .expect("budget-limited timeout has deadline") + .timeout_error()), + Either::Right((_, _)) => Err(HttpError::Timeout { + operation: "response stream idle", + timeout, + }), + } +} + +/// Byte-count hook for detecting a body shorter than Content-Length. +#[derive(Clone, Debug, Default)] +pub struct TruncatedStreamDetector { + expected: Option, + received: u64, +} + +impl TruncatedStreamDetector { + pub fn from_headers(headers: &HeaderMap) -> Self { + let expected = headers + .get(http::header::CONTENT_LENGTH) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse().ok()); + Self { + expected, + received: 0, + } + } + pub fn observe(&mut self, bytes: &Bytes) { + self.received = self.received.saturating_add(bytes.len() as u64); + } + pub fn finish(&self) -> Result<(), HttpError> { + if self + .expected + .is_some_and(|expected| self.received < expected) + { + return Err(HttpError::TruncatedBody { + expected: self.expected.unwrap_or_default(), + received: self.received, + }); + } + Ok(()) + } +} diff --git a/crates/agentkit-http/src/response.rs b/crates/agentkit-http/src/response.rs index 6790de7..b547879 100644 --- a/crates/agentkit-http/src/response.rs +++ b/crates/agentkit-http/src/response.rs @@ -14,6 +14,17 @@ pub struct HttpResponse { body: BodyStream, } +impl std::fmt::Debug for HttpResponse { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("HttpResponse") + .field("status", &self.status) + .field("header_names", &self.headers.keys().collect::>()) + .field("final_url", &self.final_url) + .field("body", &"") + .finish() + } +} + impl HttpResponse { pub fn new( status: StatusCode, diff --git a/crates/agentkit-integration-tests/src/snapshot.rs b/crates/agentkit-integration-tests/src/snapshot.rs index d3a3d07..39f140f 100644 --- a/crates/agentkit-integration-tests/src/snapshot.rs +++ b/crates/agentkit-integration-tests/src/snapshot.rs @@ -308,7 +308,9 @@ fn normalise_recording(recording: &mut SessionRecording) { result.output_items.iter_mut().for_each(normalise_item); canonicalise_metadata(&mut result.metadata); } - ModelTurnEvent::Delta(_) | ModelTurnEvent::Usage(_) => {} + ModelTurnEvent::Delta(_) + | ModelTurnEvent::Usage(_) + | ModelTurnEvent::ResponseAttemptSuperseded => {} } } } diff --git a/crates/agentkit-loop/src/lib.rs b/crates/agentkit-loop/src/lib.rs index ec9288d..04db465 100644 --- a/crates/agentkit-loop/src/lib.rs +++ b/crates/agentkit-loop/src/lib.rs @@ -243,6 +243,22 @@ impl TelemetryConfig { } } +/// Capabilities supported by the consumer of model-turn events. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct SessionConsumerCapabilities { + /// The consumer can discard all events from a superseded response attempt. + #[serde(default)] + pub response_attempt_supersession: bool, +} + +impl SessionConsumerCapabilities { + /// Enables response-attempt supersession support. + pub fn with_response_attempt_supersession(mut self) -> Self { + self.response_attempt_supersession = true; + self + } +} + /// Configuration required to start a new model session. /// /// Pass this to [`Agent::start`] to initialise the underlying [`ModelSession`] @@ -265,6 +281,9 @@ pub struct SessionConfig { pub metadata: MetadataMap, /// Default provider-side prompt caching policy for turns in this session. pub cache: Option, + /// Features that the consumer of model-turn events can safely handle. + #[serde(default)] + pub consumer_capabilities: SessionConsumerCapabilities, } impl SessionConfig { @@ -274,6 +293,7 @@ impl SessionConfig { session_id: session_id.into(), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), } } @@ -294,6 +314,14 @@ impl SessionConfig { self.cache = None; self } + + /// Declares that the event consumer can discard superseded response attempts. + pub fn with_response_attempt_supersession(mut self) -> Self { + self.consumer_capabilities = self + .consumer_capabilities + .with_response_attempt_supersession(); + self + } } /// Strength of a prompt-cache request. @@ -559,6 +587,11 @@ pub enum ModelTurnEvent { ToolCall(ToolCallPart), /// Updated token usage statistics. Usage(Usage), + /// Supersedes every previously emitted event from the current response attempt. + /// + /// This marker is ordered after the failed attempt's deltas, tool calls, and usage and + /// before replacement-attempt output. It is emitted only when the session consumer opted in. + ResponseAttemptSuperseded, /// The model has finished generating for this turn. Finished(ModelTurnResult), } @@ -934,6 +967,11 @@ pub enum AgentEvent { }, /// Updated token usage statistics. UsageUpdated(Usage), + /// All events from the preceding model response attempt are superseded. + /// + /// Consumers that opted in must discard that attempt's deltas, tool calls, usage updates, + /// and reconstruction state before handling replacement output. + ResponseAttemptSuperseded, /// Non-fatal warning (e.g. a tool failure that was recovered from). Warning { message: String }, /// The agent run has failed with an unrecoverable error. @@ -2368,6 +2406,11 @@ where saw_tool_call = true; self.emit(AgentEvent::ToolCallRequested(call.clone())); } + ModelTurnEvent::ResponseAttemptSuperseded => { + saw_tool_call = false; + latest_usage = None; + self.emit(AgentEvent::ResponseAttemptSuperseded); + } ModelTurnEvent::Finished(result) => { finished_result = Some(result); break; @@ -4311,6 +4354,7 @@ mod tests { use super::*; struct FakeAdapter; + struct SupersedingAdapter; struct SlowAdapter; struct RecordingAdapter { seen_descriptions: StdArc>>>, @@ -4320,6 +4364,9 @@ mod tests { struct DualApprovalAdapter; struct FakeSession; + struct SupersedingSession { + supersession_enabled: bool, + } struct SlowSession; struct RecordingSession { seen_descriptions: StdArc>>>, @@ -4332,6 +4379,10 @@ mod tests { events: VecDeque, } + struct SupersedingTurn { + events: VecDeque, + } + struct SlowTurn { emitted: bool, } @@ -4565,6 +4616,17 @@ mod tests { } } + #[async_trait] + impl ModelAdapter for SupersedingAdapter { + type Session = SupersedingSession; + + async fn start_session(&self, config: SessionConfig) -> Result { + Ok(SupersedingSession { + supersession_enabled: config.consumer_capabilities.response_attempt_supersession, + }) + } + } + #[async_trait] impl ModelAdapter for SlowAdapter { type Session = SlowSession; @@ -4604,6 +4666,38 @@ mod tests { } } + #[async_trait] + impl ModelSession for SupersedingSession { + type Turn = SupersedingTurn; + + async fn begin_turn( + &mut self, + _request: TurnRequest, + _cancellation: Option, + ) -> Result { + assert!(self.supersession_enabled); + Ok(SupersedingTurn { + events: VecDeque::from([ + ModelTurnEvent::Usage(Usage::default()), + ModelTurnEvent::ToolCall(ToolCallPart::new( + "discarded-call", + "discarded-tool", + json!({}), + )), + ModelTurnEvent::ResponseAttemptSuperseded, + ModelTurnEvent::Finished(ModelTurnResult { + finish_reason: FinishReason::Completed, + output_items: vec![Item::text(ItemKind::Assistant, "replacement")], + usage: None, + metadata: MetadataMap::new(), + model: None, + response_id: None, + }), + ]), + }) + } + } + #[async_trait] impl ModelSession for FakeSession { type Turn = FakeTurn; @@ -4913,6 +5007,16 @@ mod tests { } } + #[async_trait] + impl ModelTurn for SupersedingTurn { + async fn next_event( + &mut self, + _cancellation: Option, + ) -> Result, LoopError> { + Ok(self.events.pop_front()) + } + } + #[async_trait] impl ModelTurn for SlowTurn { async fn next_event( @@ -5327,6 +5431,58 @@ mod tests { } } + #[test] + fn session_consumer_capabilities_are_typed_and_serde_defaulted() { + let config = SessionConfig::new("session").with_response_attempt_supersession(); + assert!(config.consumer_capabilities.response_attempt_supersession); + + let decoded: SessionConfig = serde_json::from_value(json!({ + "session_id": "session", + "metadata": {}, + "cache": null + })) + .unwrap(); + assert_eq!( + decoded.consumer_capabilities, + SessionConsumerCapabilities::default() + ); + } + + #[tokio::test] + async fn response_attempt_supersession_is_forwarded_and_resets_attempt_state() { + let events = StdArc::new(StdMutex::new(Vec::new())); + let agent = Agent::builder() + .model(SupersedingAdapter) + .observer(RecordingObserver { + events: events.clone(), + }) + .build() + .unwrap(); + let mut driver = agent + .start(SessionConfig::new("supersession-session").with_response_attempt_supersession()) + .await + .unwrap(); + driver + .submit_input(vec![Item::text(ItemKind::User, "hello")]) + .unwrap(); + + let LoopStep::Finished(result) = run_until_finished(&mut driver).await else { + panic!("turn did not finish"); + }; + assert!(result.usage.is_none()); + + let events = events.lock().unwrap(); + let tool_call = events + .iter() + .position(|event| matches!(event, AgentEvent::ToolCallRequested(_))) + .unwrap(); + let superseded = events + .iter() + .position(|event| matches!(event, AgentEvent::ResponseAttemptSuperseded)) + .unwrap(); + assert!(tool_call < superseded); + } + fn turn_lifecycle_events( events: &[AgentEvent], ) -> Vec<(agentkit_core::TurnId, Option)> { @@ -5532,6 +5688,7 @@ mod tests { session_id: SessionId::new("session-1"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -5609,6 +5766,7 @@ mod tests { session_id: SessionId::new("session-mutation-point"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -6000,6 +6158,7 @@ mod tests { session_id: SessionId::new("session-no-valid-input"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -6043,6 +6202,7 @@ mod tests { session_id: SessionId::new("session-2"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -6101,6 +6261,7 @@ mod tests { session_id: SessionId::new("session-failing-start-event"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -6144,6 +6305,7 @@ mod tests { session_id: SessionId::new("session-run-then-deny-start-event"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -6173,10 +6335,9 @@ mod tests { .get(TOOL_RESULT_FAILURE_KIND_METADATA_KEY) .and_then(Value::as_str) == Some(TOOL_RESULT_FAILURE_KIND_PERMISSION_DENIED) - && result + && !result .metadata - .get(TOOL_RESULT_NOT_STARTED_METADATA_KEY) - .is_none() + .contains_key(TOOL_RESULT_NOT_STARTED_METADATA_KEY) ))); } @@ -6212,6 +6373,7 @@ mod tests { session_id: SessionId::new("session-background"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -6443,6 +6605,7 @@ mod tests { session_id: SessionId::new("session-detached-progress"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -6533,6 +6696,7 @@ mod tests { session_id: SessionId::new("session-cancel-delayed-background-approval"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -7133,6 +7297,7 @@ mod tests { session_id: SessionId::new("session-cancel"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -7235,6 +7400,7 @@ mod tests { session_id: SessionId::new("session-mixed-cancel"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -7324,6 +7490,7 @@ mod tests { session_id: SessionId::new("session-cancel-mid-call"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -7543,6 +7710,7 @@ mod tests { session_id: SessionId::new("session-approval"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -7600,6 +7768,7 @@ mod tests { session_id: SessionId::new("session-approval-start-event"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -7662,6 +7831,7 @@ mod tests { session_id: SessionId::new("session-cancel-pending-approval"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -7758,6 +7928,7 @@ mod tests { session_id: SessionId::new("session-resolved-approval-cancel-race"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -7801,6 +7972,7 @@ mod tests { session_id: SessionId::new("session-approval-patched"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -7852,6 +8024,7 @@ mod tests { session_id: SessionId::new("session-dual-approval"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -7986,6 +8159,7 @@ mod tests { session_id: SessionId::new("session-dual-approval-cancel"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -8083,6 +8257,7 @@ mod tests { session_id: SessionId::new("session-4"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -8205,6 +8380,7 @@ mod tests { session_id: SessionId::new("session-dynamic-tools"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -8261,6 +8437,7 @@ mod tests { session_id: SessionId::new("session-catalog-events"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -8328,6 +8505,7 @@ mod tests { session_id: SessionId::new("session-cache"), metadata: MetadataMap::new(), cache: Some(default_cache.clone()), + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -8388,6 +8566,7 @@ mod tests { session_id: SessionId::new("yield-session"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); @@ -8479,6 +8658,7 @@ mod tests { session_id: SessionId::new("observer-session"), metadata: MetadataMap::new(), cache: None, + consumer_capabilities: SessionConsumerCapabilities::default(), }) .await .unwrap(); diff --git a/crates/agentkit-provider-anthropic/README.md b/crates/agentkit-provider-anthropic/README.md index 1db8334..13aeed0 100644 --- a/crates/agentkit-provider-anthropic/README.md +++ b/crates/agentkit-provider-anthropic/README.md @@ -42,6 +42,22 @@ async fn main() -> Result<(), Box> { | `ANTHROPIC_VERSION` | no | `2023-06-01` | | `ANTHROPIC_BETA` | no | comma-separated list of anthropic-beta tags | +## Authentication and resilience + +`AnthropicConfig` stores credentials as a first-class +`agentkit_http::Authentication`. Unlike the other provider constructors, +`AnthropicConfig::new("...", model, max_tokens)` treats its string as an +Anthropic API key and sends it in `x-api-key`. To send a bearer token in +`Authorization`, use the explicit +`AnthropicConfig::with_auth_token("...", model, max_tokens)` constructor. Use +`AnthropicConfig::from_authentication(...)` or `.with_authentication(...)` for +an existing `Authentication`, and `.with_authentication_provider(...)` for a +custom refresh-capable `AuthenticationProvider`. + +Resilience is opt-in: `resilience` is an `Option` that +defaults to `None`. Calling `.with_resilience(...)` enables retries and +timeouts; leaving it as `None` preserves the existing single-attempt behavior. + ## Server tools Known server tools (`WebSearchTool`, `WebFetchTool`, `CodeExecutionTool`, etc.) diff --git a/crates/agentkit-provider-anthropic/src/config.rs b/crates/agentkit-provider-anthropic/src/config.rs index b665434..3c57944 100644 --- a/crates/agentkit-provider-anthropic/src/config.rs +++ b/crates/agentkit-provider-anthropic/src/config.rs @@ -1,3 +1,7 @@ +use std::borrow::Cow; +use std::sync::Arc; + +use agentkit_http::{Authentication, AuthenticationProvider, HeaderName, ResilienceConfig}; use serde_json::{Value, json}; use crate::error::AnthropicError; @@ -8,6 +12,115 @@ pub const DEFAULT_ENDPOINT: &str = "https://api.anthropic.com/v1/messages"; /// Default `anthropic-version` header. pub const DEFAULT_ANTHROPIC_VERSION: &str = "2023-06-01"; +/// Anthropic API-key authentication for the `x-api-key` header. +#[derive(Clone, Debug)] +pub struct AnthropicApiKey(Authentication); + +impl AnthropicApiKey { + /// Wraps an Anthropic API key without retaining it as a plain config string. + pub fn new(api_key: impl Into) -> Self { + Self(Authentication::header( + HeaderName::from_static("x-api-key"), + api_key, + )) + } +} + +impl From for AnthropicApiKey { + fn from(api_key: String) -> Self { + Self::new(api_key) + } +} + +impl From<&str> for AnthropicApiKey { + fn from(api_key: &str) -> Self { + Self::new(api_key) + } +} + +impl From<&String> for AnthropicApiKey { + fn from(api_key: &String) -> Self { + Self::new(api_key) + } +} + +impl From> for AnthropicApiKey { + fn from(api_key: Box) -> Self { + Self::new(api_key) + } +} + +impl From> for AnthropicApiKey { + fn from(api_key: Arc) -> Self { + Self::new(api_key.as_ref()) + } +} + +impl<'a> From> for AnthropicApiKey { + fn from(api_key: Cow<'a, str>) -> Self { + Self::new(api_key) + } +} + +impl From for Authentication { + fn from(api_key: AnthropicApiKey) -> Self { + api_key.0 + } +} + +/// Anthropic bearer authentication for the `Authorization` header. +#[derive(Clone, Debug)] +pub struct AnthropicAuthToken(Authentication); + +impl AnthropicAuthToken { + /// Wraps an Anthropic auth token without retaining it as a plain config string. + pub fn new(auth_token: impl Into) -> Self { + Self(Authentication::bearer(auth_token)) + } +} + +impl From for AnthropicAuthToken { + fn from(auth_token: String) -> Self { + Self::new(auth_token) + } +} + +impl From<&str> for AnthropicAuthToken { + fn from(auth_token: &str) -> Self { + Self::new(auth_token) + } +} + +impl From<&String> for AnthropicAuthToken { + fn from(auth_token: &String) -> Self { + Self::new(auth_token) + } +} + +impl From> for AnthropicAuthToken { + fn from(auth_token: Box) -> Self { + Self::new(auth_token) + } +} + +impl From> for AnthropicAuthToken { + fn from(auth_token: Arc) -> Self { + Self::new(auth_token.as_ref()) + } +} + +impl<'a> From> for AnthropicAuthToken { + fn from(auth_token: Cow<'a, str>) -> Self { + Self::new(auth_token) + } +} + +impl From for Authentication { + fn from(auth_token: AnthropicAuthToken) -> Self { + auth_token.0 + } +} + /// Extended thinking configuration. #[derive(Clone, Debug)] pub enum ThinkingConfig { @@ -143,10 +256,8 @@ pub struct AnthropicMcpServer(pub Value); /// argument — the Messages API rejects requests without it. #[derive(Clone)] pub struct AnthropicConfig { - /// Anthropic API key (`x-api-key` header). - pub api_key: Option, - /// OAuth / bearer token; if set, takes precedence over `api_key`. - pub auth_token: Option, + /// Authentication applied to each request. + pub authentication: Authentication, /// Endpoint URL. Defaults to the Anthropic production endpoint. pub base_url: String, @@ -205,47 +316,57 @@ pub struct AnthropicConfig { /// [`AnthropicConfig::with_streaming`] for debugging or when an upstream /// proxy doesn't forward SSE bodies. pub streaming: bool, + /// Optional retry and timeout policy. `None` preserves single-attempt behavior. + pub resilience: Option, +} + +impl std::fmt::Debug for AnthropicConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AnthropicConfig") + .field("authentication", &"") + .field("base_url", &self.base_url) + .field("anthropic_version", &self.anthropic_version) + .field("anthropic_beta", &self.anthropic_beta) + .field("model", &self.model) + .field("max_tokens", &self.max_tokens) + .field("temperature", &self.temperature) + .field("top_p", &self.top_p) + .field("top_k", &self.top_k) + .field("stop_sequences", &self.stop_sequences) + .field("thinking", &self.thinking) + .field("service_tier", &self.service_tier) + .field("tool_choice", &self.tool_choice) + .field("disable_parallel_tool_use", &self.disable_parallel_tool_use) + .field("output_format", &self.output_format) + .field("output_effort", &self.output_effort) + .field("streaming", &self.streaming) + .field("resilience", &self.resilience) + .finish_non_exhaustive() + } } impl AnthropicConfig { - /// Creates a new configuration using an API key. + /// Creates a new configuration using `x-api-key` authentication. pub fn new( - api_key: impl Into, + api_key: impl Into, model: impl Into, max_tokens: u32, ) -> Result { - if max_tokens == 0 { - return Err(AnthropicError::InvalidMaxTokens); - } - Ok(Self { - api_key: Some(api_key.into()), - auth_token: None, - base_url: DEFAULT_ENDPOINT.into(), - anthropic_version: DEFAULT_ANTHROPIC_VERSION.into(), - anthropic_beta: Vec::new(), - model: model.into(), - max_tokens, - temperature: None, - top_p: None, - top_k: None, - stop_sequences: None, - thinking: None, - service_tier: None, - metadata_user_id: None, - tool_choice: None, - disable_parallel_tool_use: None, - server_tools: Vec::new(), - container: None, - output_format: None, - output_effort: None, - mcp_servers: Vec::new(), - streaming: true, - }) + Self::from_authentication(api_key.into(), model, max_tokens) } /// Creates a new configuration using a bearer auth token. pub fn with_auth_token( - auth_token: impl Into, + auth_token: impl Into, + model: impl Into, + max_tokens: u32, + ) -> Result { + Self::from_authentication(auth_token.into(), model, max_tokens) + } + + /// Creates a new configuration using arbitrary authentication. + pub fn from_authentication( + authentication: impl Into, model: impl Into, max_tokens: u32, ) -> Result { @@ -253,8 +374,7 @@ impl AnthropicConfig { return Err(AnthropicError::InvalidMaxTokens); } Ok(Self { - api_key: None, - auth_token: Some(auth_token.into()), + authentication: authentication.into(), base_url: DEFAULT_ENDPOINT.into(), anthropic_version: DEFAULT_ANTHROPIC_VERSION.into(), anthropic_beta: Vec::new(), @@ -275,6 +395,7 @@ impl AnthropicConfig { output_effort: None, mcp_servers: Vec::new(), streaming: true, + resilience: None, }) } @@ -322,6 +443,23 @@ impl AnthropicConfig { // --- Builder methods --- + /// Replaces the configured authentication. + pub fn with_authentication(mut self, authentication: impl Into) -> Self { + self.authentication = authentication.into(); + self + } + + /// Uses a custom refresh-capable authentication provider. + pub fn with_authentication_provider(self, provider: P) -> Self { + self.with_authentication(Authentication::new(provider)) + } + + /// Enables request and pre-visible-output retries and timeouts. + pub fn with_resilience(mut self, resilience: ResilienceConfig) -> Self { + self.resilience = Some(resilience); + self + } + /// Overrides the endpoint URL. pub fn with_base_url(mut self, url: impl Into) -> Self { self.base_url = url.into(); diff --git a/crates/agentkit-provider-anthropic/src/lib.rs b/crates/agentkit-provider-anthropic/src/lib.rs index b7898cd..f4e4e4c 100644 --- a/crates/agentkit-provider-anthropic/src/lib.rs +++ b/crates/agentkit-provider-anthropic/src/lib.rs @@ -38,21 +38,28 @@ mod stream; use std::collections::{BTreeSet, VecDeque}; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::time::Duration; use agentkit_core::TurnCancellation; -use agentkit_http::{BodyStream, Http, HttpError, HttpRequestBuilder}; +use agentkit_http::{ + Authentication, AuthenticationAttempt, AuthenticationProvider, BodyStream, Bytes, Http, + HttpError, HttpResponse, LogicalDeadline, ResilienceConfig, StatusCode, + TruncatedStreamDetector, is_retryable_body_read, is_retryable_status, next_body_chunk_bounded, + run_bounded, sleep as resilience_sleep, +}; use agentkit_loop::{ LoopError, ModelAdapter, ModelSession, ModelTurn, ModelTurnEvent, SessionConfig, TurnRequest, }; use async_trait::async_trait; -use futures_util::StreamExt; -use futures_util::future::{Either, select}; +use futures_util::future::{BoxFuture, Either, select}; use crate::stream::{EventTranslator, SseDecoder}; pub use crate::config::{ - AnthropicConfig, AnthropicMcpServer, DEFAULT_ANTHROPIC_VERSION, DEFAULT_ENDPOINT, OutputEffort, - OutputFormat, ServiceTier, ThinkingConfig, ToolChoice, + AnthropicApiKey, AnthropicAuthToken, AnthropicConfig, AnthropicMcpServer, + DEFAULT_ANTHROPIC_VERSION, DEFAULT_ENDPOINT, OutputEffort, OutputFormat, ServiceTier, + ThinkingConfig, ToolChoice, }; pub use crate::error::AnthropicError; pub use crate::server_tool::{ @@ -78,10 +85,7 @@ impl AnthropicAdapter { .build() .map(Http::new) .map_err(|error| AnthropicError::HttpClient(HttpError::request(error)))?; - Ok(Self { - client, - config: Arc::new(config), - }) + Ok(Self::with_client(config, client)) } /// Creates a new adapter using a pre-configured [`Http`] client. @@ -91,12 +95,31 @@ impl AnthropicAdapter { config: Arc::new(config), } } + + /// Overrides the configured `x-api-key` or bearer-token authentication. + pub fn with_authentication(mut self, authentication: impl Into) -> Self { + Arc::make_mut(&mut self.config).authentication = authentication.into(); + self + } + + /// Overrides authentication with a custom refresh-capable provider. + pub fn with_authentication_provider(self, provider: P) -> Self { + self.with_authentication(Authentication::new(provider)) + } + + /// Enables request and pre-visible-output stream retries and timeouts. + pub fn with_resilience(mut self, resilience: ResilienceConfig) -> Self { + Arc::make_mut(&mut self.config).resilience = Some(resilience); + self + } } /// An active session with the Anthropic Messages API. pub struct AnthropicSession { client: Http, config: Arc, + authentication: Option, + resilience: Option, _session_config: SessionConfig, } @@ -119,12 +142,246 @@ enum TurnInner { Streaming(Box), } +type StreamReplay = Box< + dyn FnMut() -> BoxFuture<'static, Result<(BodyStream, TruncatedStreamDetector), LoopError>> + + Send, +>; + struct StreamingState { body: BodyStream, decoder: SseDecoder, translator: EventTranslator, pending: VecDeque, eof: bool, + visible_output: bool, + detect_truncation: bool, + idle_timeout: Option, + deadline: Option, + integrity: TruncatedStreamDetector, + replay: Option, +} + +#[derive(Clone)] +struct ReplayRequest { + client: Http, + config: Arc, + body: Bytes, + authentication: Option, + authentication_attempt: Arc>>, + reauthenticated: Arc, + resilience: Option, + deadline: Option, + retries_used: Arc, +} + +impl ReplayRequest { + async fn authentication_attempt(&self) -> Result { + let existing = self + .authentication_attempt + .lock() + .map_err(|_| LoopError::Provider("authentication state lock poisoned".into()))? + .clone(); + if let Some(existing) = existing { + return Ok(existing); + } + let authentication = self + .authentication + .as_ref() + .ok_or_else(|| LoopError::Provider(AnthropicError::MissingCredentials.to_string()))?; + let attempt = run_bounded( + authentication.authenticate(None), + None, + self.deadline.as_ref(), + "authentication", + ) + .await + .map_err(|error| LoopError::Provider(format!("authentication failed: {error}")))?; + *self + .authentication_attempt + .lock() + .map_err(|_| LoopError::Provider("authentication state lock poisoned".into()))? = + Some(attempt.clone()); + Ok(attempt) + } + + async fn reauthenticate(&self, previous: AuthenticationAttempt) -> Result<(), LoopError> { + let authenticate = self + .authentication + .as_ref() + .expect("401 refresh requires authentication") + .authenticate(Some(&previous)); + let next = run_bounded( + authenticate, + None, + self.deadline.as_ref(), + "reauthentication", + ) + .await + .map_err(|error| LoopError::Provider(format!("reauthentication failed: {error}")))?; + *self + .authentication_attempt + .lock() + .map_err(|_| LoopError::Provider("authentication state lock poisoned".into()))? = + Some(next); + Ok(()) + } + + fn reserve_retry(&self) -> Option { + let maximum = self.resilience.as_ref()?.max_retries; + self.retries_used + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |used| { + (used < maximum).then_some(used + 1) + }) + .ok() + } + + fn retry_delay( + &self, + retry_number: usize, + headers: Option<&agentkit_http::HeaderMap>, + ) -> Duration { + self.resilience + .as_ref() + .expect("reserved retry requires resilience") + .retry_delay(retry_number, headers) + } + + async fn wait_before_retry( + &self, + retry_number: usize, + headers: Option<&agentkit_http::HeaderMap>, + ) -> Result<(), LoopError> { + self.wait_for_retry_delay(self.retry_delay(retry_number, headers)) + .await + } + + async fn wait_for_retry_delay(&self, delay: Duration) -> Result<(), LoopError> { + run_bounded( + async { + resilience_sleep(delay).await; + Ok(()) + }, + None, + self.deadline.as_ref(), + "retry backoff", + ) + .await + .map_err(|error| LoopError::Provider(format!("retry backoff failed: {error}"))) + } + + async fn execute_attempt( + &self, + request: agentkit_http::HttpRequest, + ) -> Result { + let timeout = self + .resilience + .as_ref() + .and_then(|config| config.attempt_timeout); + run_bounded( + self.client.execute(request), + timeout, + self.deadline.as_ref(), + "HTTP request attempt", + ) + .await + } + + async fn read_response_text(&self, response: HttpResponse) -> Result { + let timeout = self + .resilience + .as_ref() + .and_then(|config| config.attempt_timeout); + run_bounded( + response.text(), + timeout, + self.deadline.as_ref(), + "HTTP response body", + ) + .await + } + + async fn open_response(&self) -> Result { + loop { + let authentication = self.authentication_attempt().await?; + let mut builder = self + .client + .post(&self.config.base_url) + .header("Content-Type", "application/json") + .header("anthropic-version", self.config.anthropic_version.as_str()); + + let betas = collect_beta_flags(&self.config); + if !betas.is_empty() { + builder = builder.header( + "anthropic-beta", + betas.into_iter().collect::>().join(","), + ); + } + builder = builder.header( + "User-Agent", + concat!("agentkit-provider-anthropic/", env!("CARGO_PKG_VERSION")), + ); + if self.config.streaming { + builder = builder.header("Accept", "text/event-stream"); + } + builder = builder.headers(authentication.headers().clone()); + let request = builder.body(self.body.clone()).build().map_err(|error| { + LoopError::Provider(format!("Anthropic request failed: {error}")) + })?; + + match self.execute_attempt(request).await { + Ok(response) => { + if response.status() == StatusCode::UNAUTHORIZED + && !self.reauthenticated.swap(true, Ordering::SeqCst) + { + drop(response); + self.reauthenticate(authentication).await?; + continue; + } + if is_retryable_status(response.status()) + && let Some(retry_number) = self.reserve_retry() + { + let delay = self.retry_delay(retry_number, Some(response.headers())); + drop(response); + self.wait_for_retry_delay(delay).await?; + continue; + } + return Ok(response); + } + Err(error) if error.is_retryable_transport() => { + if let Some(retry_number) = self.reserve_retry() { + self.wait_before_retry(retry_number, None).await?; + continue; + } + return Err(LoopError::Provider(format!( + "Anthropic request failed: {error}" + ))); + } + Err(error) => { + return Err(LoopError::Provider(format!( + "Anthropic request failed: {error}" + ))); + } + } + } + } + + async fn replay_stream(&self) -> Result<(BodyStream, TruncatedStreamDetector), LoopError> { + let retry_number = self.reserve_retry().ok_or_else(|| { + LoopError::Provider( + "Anthropic stream failed before output and retry budget is exhausted".into(), + ) + })?; + self.wait_before_retry(retry_number, None).await?; + let response = self.open_response().await?; + if !response.status().is_success() { + return Err(LoopError::Provider(format!( + "Anthropic stream replay failed with status {}", + response.status() + ))); + } + let integrity = TruncatedStreamDetector::from_headers(response.headers()); + Ok((response.bytes_stream(), integrity)) + } } #[async_trait] @@ -135,6 +392,8 @@ impl ModelAdapter for AnthropicAdapter { Ok(AnthropicSession { client: self.client.clone(), config: self.config.clone(), + authentication: Some(self.config.authentication.clone()), + resilience: self.config.resilience.clone(), _session_config: config, }) } @@ -155,69 +414,92 @@ impl ModelSession for AnthropicSession { ) -> Result { let config = self.config.clone(); - let request_future = async move { + let request_future = async { let body = request::build_request_body(&config, &turn_request) .map_err(|e| LoopError::Provider(e.to_string()))?; + let body = serde_json::to_vec(&body) + .map(Bytes::from) + .map_err(|e| LoopError::Provider(format!("failed to serialize request: {e}")))?; + let deadline = self + .resilience + .as_ref() + .map(|resilience| LogicalDeadline::new(resilience.retry_budget)); + let replay_request = ReplayRequest { + client: self.client.clone(), + config: config.clone(), + body, + authentication: self.authentication.clone(), + authentication_attempt: Arc::new(std::sync::Mutex::new(None)), + reauthenticated: Arc::new(AtomicBool::new(false)), + resilience: self.resilience.clone(), + deadline: deadline.clone(), + retries_used: Arc::new(AtomicUsize::new(0)), + }; - let betas = collect_beta_flags(&config); - - let mut http = self - .client - .post(&config.base_url) - .header("Content-Type", "application/json") - .header("anthropic-version", config.anthropic_version.as_str()); - - http = attach_auth(http, &config)?; - - if !betas.is_empty() { - let joined = betas.into_iter().collect::>().join(","); - http = http.header("anthropic-beta", joined); - } - - http = http.header( - "User-Agent", - concat!("agentkit-provider-anthropic/", env!("CARGO_PKG_VERSION")), - ); - - if config.streaming { - http = http.header("Accept", "text/event-stream"); - } - - let response = http.json(&body).send().await.map_err(|error| { - LoopError::Provider(format!("Anthropic request failed: {error}")) - })?; - - let status = response.status(); - - if !status.is_success() { - // Drain the body for the error message, regardless of mode — - // the server typically returns JSON error details here. - let body_text = response.text().await.unwrap_or_default(); - return Err(LoopError::Provider(format!( - "Anthropic request failed with status {status}: {body_text}" - ))); - } + loop { + let response = replay_request.open_response().await?; + let status = response.status(); + if config.streaming && status.is_success() { + let integrity = TruncatedStreamDetector::from_headers(response.headers()); + let idle_timeout = self + .resilience + .as_ref() + .and_then(|resilience| resilience.stream_idle_timeout); + let replay = self.resilience.as_ref().and_then(|resilience| { + (resilience.max_retries > 0).then(|| { + let replay_request = replay_request.clone(); + Box::new(move || { + let replay_request = replay_request.clone(); + Box::pin(async move { replay_request.replay_stream().await }) + as BoxFuture<'static, _> + }) as StreamReplay + }) + }); + return Ok(AnthropicTurn { + inner: TurnInner::Streaming(Box::new(StreamingState { + body: response.bytes_stream(), + decoder: SseDecoder::new(), + translator: EventTranslator::new(), + pending: VecDeque::new(), + eof: false, + visible_output: false, + detect_truncation: self.resilience.is_some(), + idle_timeout, + deadline: deadline.clone(), + integrity, + replay, + })), + }); + } - if config.streaming { - Ok(AnthropicTurn { - inner: TurnInner::Streaming(Box::new(StreamingState { - body: response.bytes_stream(), - decoder: SseDecoder::new(), - translator: EventTranslator::new(), - pending: VecDeque::new(), - eof: false, - })), - }) - } else { - let body_text = response.text().await.map_err(|error| { - LoopError::Provider(format!("failed to read Anthropic response body: {error}")) - })?; + let body_text = match replay_request.read_response_text(response).await { + Ok(body) => body, + Err(error) if is_retryable_body_read(status, &error) => { + if let Some(retry_number) = replay_request.reserve_retry() { + replay_request.wait_before_retry(retry_number, None).await?; + continue; + } + return Err(LoopError::Provider(format!( + "failed to read Anthropic response body: {error}" + ))); + } + Err(error) => { + return Err(LoopError::Provider(format!( + "failed to read Anthropic response body: {error}" + ))); + } + }; + if !status.is_success() { + return Err(LoopError::Provider(format!( + "Anthropic request failed with status {status}: {body_text}" + ))); + } let events = response::build_turn_from_response(&body_text) .map_err(|e| LoopError::Provider(e.to_string()))?; - Ok(AnthropicTurn { + return Ok(AnthropicTurn { inner: TurnInner::Buffered { events }, - }) + }); } }; @@ -264,49 +546,83 @@ impl ModelTurn for AnthropicTurn { translator, pending, eof, + visible_output, + detect_truncation, + idle_timeout, + deadline, + integrity, + replay, } = state.as_mut(); - next_streaming_event(body, decoder, translator, pending, eof, cancellation).await + next_streaming_event( + body, + decoder, + translator, + pending, + eof, + visible_output, + *detect_truncation, + *idle_timeout, + deadline.as_ref(), + integrity, + replay, + cancellation, + ) + .await } } } } -/// Pulls the next event from an active SSE stream, decoding more bytes as -/// needed. Returns `Ok(None)` once the translator has emitted `Finished` and -/// the pending queue is empty. +async fn next_stream_chunk( + body: &mut BodyStream, + idle_timeout: Option, + deadline: Option<&LogicalDeadline>, + cancellation: Option<&TurnCancellation>, +) -> Result, HttpError>, LoopError> { + let next = next_body_chunk_bounded(body, idle_timeout, deadline); + futures_util::pin_mut!(next); + if let Some(cancellation) = cancellation { + let cancelled = cancellation.cancelled(); + futures_util::pin_mut!(cancelled); + match select(next, cancelled).await { + Either::Left((chunk, _)) => Ok(chunk), + Either::Right((_, _)) => Err(LoopError::Cancelled), + } + } else { + Ok(next.await) + } +} + +/// Pulls the next event from an active SSE stream, replaying only failures +/// which occur before any event has become visible to the caller. +#[allow(clippy::too_many_arguments)] async fn next_streaming_event( body: &mut BodyStream, decoder: &mut SseDecoder, translator: &mut EventTranslator, pending: &mut VecDeque, eof: &mut bool, + visible_output: &mut bool, + detect_truncation: bool, + idle_timeout: Option, + deadline: Option<&LogicalDeadline>, + integrity: &mut TruncatedStreamDetector, + replay: &mut Option, cancellation: Option, ) -> Result, LoopError> { loop { if let Some(event) = pending.pop_front() { + *visible_output = true; return Ok(Some(event)); } if *eof || translator.is_done() { return Ok(None); } - // Await the next chunk, racing against cancellation so long-lived - // streams can be interrupted mid-response. - let chunk = if let Some(cancellation) = cancellation.as_ref() { - let next = body.next(); - futures_util::pin_mut!(next); - let cancelled = cancellation.cancelled(); - futures_util::pin_mut!(cancelled); - match select(next, cancelled).await { - Either::Left((chunk, _)) => chunk, - Either::Right((_, _)) => return Err(LoopError::Cancelled), - } - } else { - body.next().await - }; - - match chunk { - Some(Ok(bytes)) => { + let chunk = next_stream_chunk(body, idle_timeout, deadline, cancellation.as_ref()).await?; + let failure = match chunk { + Ok(Some(bytes)) => { + integrity.observe(&bytes); let text = std::str::from_utf8(&bytes).map_err(|e| { LoopError::Provider(format!("invalid UTF-8 in Anthropic stream: {e}")) })?; @@ -315,32 +631,56 @@ async fn next_streaming_event( pending.push_back(produced); } } + None } - Some(Err(e)) => { - return Err(LoopError::Provider(format!( - "Anthropic stream body error: {e}" - ))); + Ok(None) if !detect_truncation => { + *eof = true; + None } - None => { + Ok(None) => { *eof = true; + let message = match integrity.finish() { + Ok(()) => "Anthropic stream ended before its terminal event".to_owned(), + Err(error) => format!("Anthropic stream body error: {error}"), + }; + Some(LoopError::Provider(message)) + } + Err(error) if !error.is_retryable_transport() => { + return Err(LoopError::Provider(format!( + "Anthropic stream body error: {error}" + ))); } + Err(error) => Some(LoopError::Provider(format!( + "Anthropic stream body error: {error}" + ))), + }; + + let Some(failure) = failure else { + continue; + }; + if *visible_output || replay.is_none() { + return Err(failure); } - } -} -fn attach_auth( - builder: HttpRequestBuilder, - config: &AnthropicConfig, -) -> Result { - if let Some(token) = &config.auth_token { - return Ok(builder.bearer_auth(token)); - } - if let Some(key) = &config.api_key { - return Ok(builder.header("x-api-key", key.as_str())); + let replay_future = replay.as_mut().expect("checked above")(); + futures_util::pin_mut!(replay_future); + let replayed = if let Some(cancellation) = cancellation.as_ref() { + let cancelled = cancellation.cancelled(); + futures_util::pin_mut!(cancelled); + match select(replay_future, cancelled).await { + Either::Left((result, _)) => result, + Either::Right((_, _)) => return Err(LoopError::Cancelled), + } + } else { + replay_future.await + }?; + *body = replayed.0; + *integrity = replayed.1; + *decoder = SseDecoder::new(); + *translator = EventTranslator::new(); + pending.clear(); + *eof = false; } - Err(LoopError::Provider( - AnthropicError::MissingCredentials.to_string(), - )) } fn collect_beta_flags(config: &AnthropicConfig) -> BTreeSet { @@ -356,12 +696,74 @@ fn collect_beta_flags(config: &AnthropicConfig) -> BTreeSet { #[cfg(test)] mod tests { use agentkit_core::{CancellationController, FinishReason}; - use agentkit_http::HttpError; + use agentkit_http::{HeaderMap, HeaderValue, HttpClient, HttpError, HttpRequest, header}; use bytes::Bytes; use futures_util::stream; use super::*; + struct AlwaysUnauthorized { + authorizations: std::sync::Mutex>, + } + + #[async_trait] + impl HttpClient for AlwaysUnauthorized { + async fn execute(&self, request: HttpRequest) -> Result { + let authorization = request.headers[header::AUTHORIZATION] + .to_str() + .unwrap() + .to_owned(); + self.authorizations.lock().unwrap().push(authorization); + Ok(HttpResponse::new( + StatusCode::UNAUTHORIZED, + HeaderMap::new(), + request.url, + Box::pin(stream::empty()), + )) + } + } + + struct RefreshingAuthentication { + calls: Arc, + } + + #[async_trait] + impl AuthenticationProvider for RefreshingAuthentication { + async fn authenticate( + &self, + previous: Option<&AuthenticationAttempt>, + ) -> Result { + let generation = self.calls.fetch_add(1, Ordering::SeqCst); + assert_eq!( + previous + .and_then(|attempt| attempt.state::()) + .copied(), + generation.checked_sub(1) + ); + let mut headers = HeaderMap::new(); + headers.insert( + header::AUTHORIZATION, + HeaderValue::from_str(&format!("Bearer refreshed-{generation}")).unwrap(), + ); + Ok(AuthenticationAttempt::new(headers, generation)) + } + } + + #[test] + fn config_debug_redacts_api_key_and_auth_token() { + let api_key = AnthropicConfig::new("anthropic-secret", "debug-model", 1024).unwrap(); + assert_eq!(api_key.resilience, None); + let auth_token = + AnthropicConfig::with_auth_token("bearer-secret", "debug-model", 1024).unwrap(); + + for debug in [format!("{api_key:?}"), format!("{auth_token:?}")] { + assert!(!debug.contains("anthropic-secret")); + assert!(!debug.contains("bearer-secret")); + assert!(debug.contains("")); + assert!(debug.contains("debug-model")); + } + } + #[test] fn rejects_zero_max_tokens() { match AnthropicConfig::new("k", "claude-opus-4-7", 0) { @@ -370,6 +772,158 @@ mod tests { } } + #[tokio::test] + async fn replay_request_reauthenticates_exactly_once() { + let client = Arc::new(AlwaysUnauthorized { + authorizations: std::sync::Mutex::new(Vec::new()), + }); + let calls = Arc::new(AtomicUsize::new(0)); + let config = Arc::new( + AnthropicConfig::new("unused", "claude-opus-4-7", 1024) + .unwrap() + .with_streaming(false), + ); + let request = ReplayRequest { + client: Http::from_arc(client.clone()), + config, + body: Bytes::from_static(b"{}"), + authentication: Some(Authentication::new(RefreshingAuthentication { + calls: calls.clone(), + })), + authentication_attempt: Arc::new(std::sync::Mutex::new(None)), + reauthenticated: Arc::new(AtomicBool::new(false)), + resilience: None, + deadline: None, + retries_used: Arc::new(AtomicUsize::new(0)), + }; + + let response = request.open_response().await.unwrap(); + + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + assert_eq!(calls.load(Ordering::SeqCst), 2); + assert_eq!( + *client.authorizations.lock().unwrap(), + ["Bearer refreshed-0", "Bearer refreshed-1"] + ); + } + + #[tokio::test] + async fn logical_deadline_starts_before_initial_authentication() { + let client = Arc::new(AlwaysUnauthorized { + authorizations: std::sync::Mutex::new(Vec::new()), + }); + let resilience = ResilienceConfig { + retry_budget: Duration::ZERO, + ..ResilienceConfig::default() + }; + let deadline = LogicalDeadline::new(resilience.retry_budget); + let request = ReplayRequest { + client: Http::from_arc(client.clone()), + config: Arc::new(AnthropicConfig::new("unused", "claude-opus-4-7", 1024).unwrap()), + body: Bytes::from_static(b"{}"), + authentication: Some(Authentication::bearer("unused")), + authentication_attempt: Arc::new(std::sync::Mutex::new(None)), + reauthenticated: Arc::new(AtomicBool::new(false)), + resilience: Some(resilience), + deadline: Some(deadline), + retries_used: Arc::new(AtomicUsize::new(0)), + }; + + let error = request.open_response().await.unwrap_err(); + + assert!(error.to_string().contains("logical request retry budget")); + assert!(client.authorizations.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn default_authentication_preserves_anthropic_header_schemes() { + let borrowed_key = String::from("borrowed-key"); + let key_configs = [ + AnthropicConfig::new("secret-key", "claude-opus-4-7", 1024).unwrap(), + AnthropicConfig::new(&borrowed_key, "claude-opus-4-7", 1024).unwrap(), + AnthropicConfig::new(Box::::from("boxed-key"), "claude-opus-4-7", 1024).unwrap(), + AnthropicConfig::new( + std::sync::Arc::::from("arc-key"), + "claude-opus-4-7", + 1024, + ) + .unwrap(), + AnthropicConfig::new( + std::borrow::Cow::Borrowed("cow-key"), + "claude-opus-4-7", + 1024, + ) + .unwrap(), + ]; + for key_config in key_configs { + let key_attempt = key_config.authentication.authenticate(None).await.unwrap(); + assert!(key_attempt.headers().get("x-api-key").is_some()); + assert!( + key_attempt + .headers() + .get(agentkit_http::header::AUTHORIZATION) + .is_none() + ); + assert!(key_attempt.headers()["x-api-key"].is_sensitive()); + } + + let key_config = AnthropicConfig::new("secret-key", "claude-opus-4-7", 1024).unwrap(); + let key_attempt = key_config.authentication.authenticate(None).await.unwrap(); + assert_eq!(key_attempt.headers()["x-api-key"], "secret-key"); + assert!( + key_attempt + .headers() + .get(agentkit_http::header::AUTHORIZATION) + .is_none() + ); + assert!(key_attempt.headers()["x-api-key"].is_sensitive()); + + let token_config = + AnthropicConfig::with_auth_token("secret-token", "claude-opus-4-7", 1024).unwrap(); + let token_attempt = token_config + .authentication + .authenticate(None) + .await + .unwrap(); + assert_eq!( + token_attempt.headers()[agentkit_http::header::AUTHORIZATION], + "Bearer secret-token" + ); + assert!(token_attempt.headers().get("x-api-key").is_none()); + assert!(token_attempt.headers()[agentkit_http::header::AUTHORIZATION].is_sensitive()); + + let borrowed_token = String::from("borrowed-token"); + for token_config in [ + AnthropicConfig::with_auth_token(&borrowed_token, "claude-opus-4-7", 1024).unwrap(), + AnthropicConfig::with_auth_token( + Box::::from("boxed-token"), + "claude-opus-4-7", + 1024, + ) + .unwrap(), + AnthropicConfig::with_auth_token( + std::sync::Arc::::from("arc-token"), + "claude-opus-4-7", + 1024, + ) + .unwrap(), + AnthropicConfig::with_auth_token( + std::borrow::Cow::Borrowed("cow-token"), + "claude-opus-4-7", + 1024, + ) + .unwrap(), + ] { + let token_attempt = token_config + .authentication + .authenticate(None) + .await + .unwrap(); + assert!(token_attempt.headers().get("x-api-key").is_none()); + assert!(token_attempt.headers()[agentkit_http::header::AUTHORIZATION].is_sensitive()); + } + } + #[test] fn beta_flags_union_includes_server_tool_requirements() { let cfg = AnthropicConfig::new("k", "claude-opus-4-7", 1024) @@ -403,6 +957,12 @@ mod tests { translator: EventTranslator::new(), pending: VecDeque::new(), eof: false, + visible_output: false, + detect_truncation: false, + idle_timeout: None, + deadline: None, + integrity: TruncatedStreamDetector::from_headers(&agentkit_http::HeaderMap::new()), + replay: None, })), } } @@ -429,6 +989,81 @@ mod tests { assert!(seen_finished, "turn never emitted Finished"); } + #[tokio::test(flavor = "current_thread")] + async fn streaming_turn_never_replays_after_visible_output() { + let chunk = Bytes::from_static( + b"event: message_start\ndata: {\"message\":{\"id\":\"m\",\"model\":\"x\",\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\nevent: content_block_start\ndata: {\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\nevent: content_block_delta\ndata: {\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"}}\n\n", + ); + let body: BodyStream = Box::pin(stream::iter(vec![ + Ok(chunk), + Err(HttpError::Other("stream failed".into())), + ])); + let replay_count = Arc::new(AtomicUsize::new(0)); + let count = replay_count.clone(); + let replay: StreamReplay = Box::new(move || { + count.fetch_add(1, Ordering::SeqCst); + Box::pin(async { Err(LoopError::Provider("unexpected replay".into())) }) + }); + let mut turn = AnthropicTurn { + inner: TurnInner::Streaming(Box::new(StreamingState { + body, + decoder: SseDecoder::new(), + translator: EventTranslator::new(), + pending: VecDeque::new(), + eof: false, + visible_output: false, + detect_truncation: false, + idle_timeout: None, + deadline: None, + integrity: TruncatedStreamDetector::from_headers(&agentkit_http::HeaderMap::new()), + replay: Some(replay), + })), + }; + + let mut visible_events = 0; + loop { + match turn.next_event(None).await { + Ok(Some(_)) => visible_events += 1, + Err(_) => break, + Ok(None) => panic!("stream ended without surfacing the body error"), + } + } + assert!(visible_events > 0); + assert_eq!(replay_count.load(Ordering::SeqCst), 0); + } + + #[tokio::test(flavor = "current_thread")] + async fn streaming_turn_does_not_replay_non_transient_body_errors() { + let body: BodyStream = Box::pin(stream::iter([Err(HttpError::Other( + "protocol failure".into(), + ))])); + let replay_count = Arc::new(AtomicUsize::new(0)); + let count = replay_count.clone(); + let replay: StreamReplay = Box::new(move || { + count.fetch_add(1, Ordering::SeqCst); + Box::pin(async { unreachable!("non-transient errors must not replay") }) + }); + let mut turn = AnthropicTurn { + inner: TurnInner::Streaming(Box::new(StreamingState { + body, + decoder: SseDecoder::new(), + translator: EventTranslator::new(), + pending: VecDeque::new(), + eof: false, + visible_output: false, + detect_truncation: true, + idle_timeout: None, + deadline: None, + integrity: TruncatedStreamDetector::default(), + replay: Some(replay), + })), + }; + + let error = turn.next_event(None).await.unwrap_err(); + assert!(error.to_string().contains("protocol failure")); + assert_eq!(replay_count.load(Ordering::SeqCst), 0); + } + #[tokio::test(flavor = "current_thread")] async fn streaming_turn_respects_pre_fired_cancellation() { let chunks = vec![ diff --git a/crates/agentkit-provider-baseten/README.md b/crates/agentkit-provider-baseten/README.md index cdda15c..158b3a9 100644 --- a/crates/agentkit-provider-baseten/README.md +++ b/crates/agentkit-provider-baseten/README.md @@ -28,6 +28,18 @@ let adapter = BasetenAdapter::new(config)?; | `BASETEN_MODEL` | yes | -- | | `BASETEN_BASE_URL` | no | `https://inference.baseten.co/v1/chat/completions` | +## Authentication and resilience + +`BasetenConfig` stores credentials as a first-class +`agentkit_http::Authentication`. A bare string passed to `BasetenConfig::new` +or `.with_authentication(...)` is shorthand for bearer authentication. Use +`.with_authentication_provider(...)` for a custom refresh-capable +`AuthenticationProvider`. + +Resilience is opt-in: `resilience` is an `Option` that +defaults to `None`. Calling `.with_resilience(...)` enables retries and +timeouts; leaving it as `None` preserves the existing single-attempt behavior. + To use an OpenAI-compatible dedicated deployment, set `BASETEN_BASE_URL` to the full endpoint, for example: diff --git a/crates/agentkit-provider-baseten/src/lib.rs b/crates/agentkit-provider-baseten/src/lib.rs index 1a7cfac..fc72625 100644 --- a/crates/agentkit-provider-baseten/src/lib.rs +++ b/crates/agentkit-provider-baseten/src/lib.rs @@ -7,6 +7,7 @@ use agentkit_adapter_completions::{ CompletionsAdapter, CompletionsError, CompletionsProvider, CompletionsSession, CompletionsTurn, }; +use agentkit_http::{Authentication, AuthenticationProvider, ResilienceConfig}; use agentkit_loop::{LoopError, ModelAdapter, SessionConfig}; use async_trait::async_trait; use serde::Serialize; @@ -18,14 +19,16 @@ const DEFAULT_ENDPOINT: &str = "https://inference.baseten.co/v1/chat/completions /// /// Use a Baseten Model API slug such as `"openai/gpt-oss-120b"`, or the /// served model name from a dedicated deployment. -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct BasetenConfig { - /// Baseten workspace API key. - pub api_key: String, + /// Authentication used for API requests. String values become bearer authentication. + pub authentication: Authentication, /// Model API slug or dedicated deployment's served model name. pub model: String, /// Full chat completions endpoint URL. pub base_url: String, + /// Optional retry and timeout policy. `None` preserves single-attempt behavior. + pub resilience: Option, /// Sampling temperature. pub temperature: Option, /// Maximum number of generated tokens. @@ -40,13 +43,31 @@ pub struct BasetenConfig { pub streaming: bool, } +impl std::fmt::Debug for BasetenConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BasetenConfig") + .field("authentication", &"") + .field("model", &self.model) + .field("base_url", &self.base_url) + .field("resilience", &self.resilience) + .field("temperature", &self.temperature) + .field("max_tokens", &self.max_tokens) + .field("top_p", &self.top_p) + .field("top_k", &self.top_k) + .field("parallel_tool_calls", &self.parallel_tool_calls) + .field("streaming", &self.streaming) + .finish() + } +} + impl BasetenConfig { /// Creates a configuration for Baseten's shared Model API endpoint. - pub fn new(api_key: impl Into, model: impl Into) -> Self { + pub fn new(authentication: impl Into, model: impl Into) -> Self { Self { - api_key: api_key.into(), + authentication: authentication.into(), model: model.into(), base_url: DEFAULT_ENDPOINT.into(), + resilience: None, temperature: None, max_tokens: None, top_p: None, @@ -56,6 +77,29 @@ impl BasetenConfig { } } + /// Replaces request authentication. String values become bearer authentication. + pub fn with_authentication(mut self, authentication: impl Into) -> Self { + self.authentication = authentication.into(); + self + } + + /// Compatibility builder for bearer API-key authentication. + pub fn with_api_key(self, api_key: impl Into) -> Self { + self.with_authentication(api_key) + } + + /// Uses a custom refresh-capable authentication provider. + pub fn with_authentication_provider(mut self, provider: P) -> Self { + self.authentication = Authentication::new(provider); + self + } + + /// Enables request retries and timeouts. + pub fn with_resilience(mut self, resilience: ResilienceConfig) -> Self { + self.resilience = Some(resilience); + self + } + /// Overrides the full chat completions endpoint URL. /// /// For a dedicated production deployment, use a URL like @@ -141,7 +185,8 @@ pub struct BasetenRequestConfig { /// Baseten implementation of [`CompletionsProvider`]. #[derive(Clone, Debug)] pub struct BasetenProvider { - api_key: String, + authentication: Authentication, + resilience: Option, base_url: String, streaming: bool, request_config: BasetenRequestConfig, @@ -150,7 +195,8 @@ pub struct BasetenProvider { impl From for BasetenProvider { fn from(config: BasetenConfig) -> Self { Self { - api_key: config.api_key, + authentication: config.authentication, + resilience: config.resilience, base_url: config.base_url, streaming: config.streaming, request_config: BasetenRequestConfig { @@ -184,12 +230,20 @@ impl CompletionsProvider for BasetenProvider { &self, builder: agentkit_http::HttpRequestBuilder, ) -> agentkit_http::HttpRequestBuilder { - builder.bearer_auth(&self.api_key).header( + builder.header( "User-Agent", concat!("agentkit-provider-baseten/", env!("CARGO_PKG_VERSION")), ) } + fn authentication(&self) -> Option { + Some(self.authentication.clone()) + } + + fn resilience_config(&self) -> Option { + self.resilience.clone() + } + fn streaming(&self) -> bool { self.streaming } @@ -212,6 +266,21 @@ impl BasetenAdapter { config, ))?)) } + + /// Overrides the default API-key authentication. + pub fn with_authentication(self, authentication: impl Into) -> Self { + Self(self.0.with_authentication(authentication)) + } + + /// Overrides authentication with a custom refresh-capable provider. + pub fn with_authentication_provider(self, provider: P) -> Self { + self.with_authentication(Authentication::new(provider)) + } + + /// Enables retry and timeout behavior. + pub fn with_resilience(self, resilience: ResilienceConfig) -> Self { + Self(self.0.with_resilience(resilience)) + } } #[async_trait] @@ -239,6 +308,23 @@ pub enum BasetenError { mod tests { use super::*; + #[test] + fn config_authentication_and_resilience_reach_provider() { + let provider = BasetenProvider::from( + BasetenConfig::new("secret", "model").with_resilience(ResilienceConfig::default()), + ); + assert!(provider.authentication().is_some()); + assert!(provider.resilience_config().is_some()); + } + + #[test] + fn config_debug_redacts_authentication() { + let debug = format!("{:?}", BasetenConfig::new("baseten-secret", "debug-model")); + assert!(!debug.contains("baseten-secret")); + assert!(debug.contains("")); + assert!(debug.contains("debug-model")); + } + #[test] fn defaults_to_model_api_with_streaming() { let provider = BasetenProvider::from(BasetenConfig::new("key", "model")); diff --git a/crates/agentkit-provider-cerebras/README.md b/crates/agentkit-provider-cerebras/README.md index 567d91c..33eb246 100644 --- a/crates/agentkit-provider-cerebras/README.md +++ b/crates/agentkit-provider-cerebras/README.md @@ -50,11 +50,20 @@ async fn main() -> Result<(), Box> { | `batch` | Files + Batch async bulk inference API | | `experimental` | Enables all preview/private-preview features | -## Retry +## Authentication and resilience -`agentkit-http` is a trait façade without a retry layer. Callers wanting retry -build their `Http` with the `reqwest-middleware-client` feature + -`reqwest-retry` and pass it via `CerebrasAdapter::with_client`. +`CerebrasConfig` stores credentials as a first-class +`agentkit_http::Authentication`. A bare string passed to `CerebrasConfig::new` +or `.with_authentication(...)` is shorthand for bearer authentication. Use +`.with_authentication_provider(...)` for a custom refresh-capable +`AuthenticationProvider`. The adapter shares this authentication across chat, +files, batches, and models requests; clients returned by `.files()`, +`.batches()`, and `.models()` use the same config. + +Resilience is opt-in and shared across those APIs too. `resilience` is an +`Option` that defaults to `None`. Calling +`.with_resilience(...)` enables retries and timeouts; leaving it as `None` +preserves the existing single-attempt behavior. ## Metadata keys diff --git a/crates/agentkit-provider-cerebras/src/batch.rs b/crates/agentkit-provider-cerebras/src/batch.rs index 49b87f7..7415618 100644 --- a/crates/agentkit-provider-cerebras/src/batch.rs +++ b/crates/agentkit-provider-cerebras/src/batch.rs @@ -21,6 +21,7 @@ use futures_util::future::{Either, select}; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value, json}; +use crate::RequestExecutor; use crate::config::{CerebrasConfig, OutputFormat, ReasoningConfig}; use crate::error::CerebrasError; use crate::files::{FilePurpose, FilesClient}; @@ -192,6 +193,14 @@ impl<'a> BatchClient<'a> { Self { http, config } } + fn executor(&self) -> RequestExecutor { + RequestExecutor::new( + self.http, + self.config.authentication.clone(), + self.config.resilience.clone(), + ) + } + /// Assembles the JSONL input via [`crate::request::build_chat_body`], /// uploads it, and submits a batch job. Per-line overrides layer on top /// of the adapter's base config. @@ -267,11 +276,10 @@ impl<'a> BatchClient<'a> { "metadata": metadata, }); let response = self - .http - .post(&url) - .bearer_auth(&self.config.api_key) - .json(&body) - .send() + .executor() + // Creating a batch has no documented idempotency key, so apply + // authentication and timeouts but do not retry the POST. + .execute_buffered_without_retries(|| self.http.post(&url).json(&body)) .await?; if !response.status().is_success() { let status = response.status().as_u16(); @@ -285,10 +293,8 @@ impl<'a> BatchClient<'a> { pub async fn list(&self) -> Result, CerebrasError> { let url = format!("{}/batches", self.config.base_url); let response = self - .http - .get(&url) - .bearer_auth(&self.config.api_key) - .send() + .executor() + .execute_buffered(|| self.http.get(&url)) .await?; if !response.status().is_success() { let status = response.status().as_u16(); @@ -303,10 +309,8 @@ impl<'a> BatchClient<'a> { pub async fn retrieve(&self, id: &str) -> Result { let url = format!("{}/batches/{id}", self.config.base_url); let response = self - .http - .get(&url) - .bearer_auth(&self.config.api_key) - .send() + .executor() + .execute_buffered(|| self.http.get(&url)) .await?; if !response.status().is_success() { let status = response.status().as_u16(); @@ -320,10 +324,9 @@ impl<'a> BatchClient<'a> { pub async fn cancel(&self, id: &str) -> Result { let url = format!("{}/batches/{id}/cancel", self.config.base_url); let response = self - .http - .post(&url) - .bearer_auth(&self.config.api_key) - .send() + .executor() + // The API does not expose an idempotency key for cancellation. + .execute_buffered_without_retries(|| self.http.post(&url)) .await?; if !response.status().is_success() { let status = response.status().as_u16(); diff --git a/crates/agentkit-provider-cerebras/src/config.rs b/crates/agentkit-provider-cerebras/src/config.rs index f420610..0610bf5 100644 --- a/crates/agentkit-provider-cerebras/src/config.rs +++ b/crates/agentkit-provider-cerebras/src/config.rs @@ -6,6 +6,7 @@ use std::collections::BTreeMap; +use agentkit_http::{Authentication, AuthenticationProvider, ResilienceConfig}; use serde_json::{Map, Value, json}; use crate::error::BuildError; @@ -340,13 +341,16 @@ impl CompressionConfig { /// /// Build one with [`CerebrasConfig::new`] or [`CerebrasConfig::from_env`], /// then refine via the `with_*` methods. -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct CerebrasConfig { // --- auth & transport --- - /// `Authorization: Bearer `. - pub api_key: String, + /// Authentication used for every Cerebras API surface. Strings become bearer authentication. + pub authentication: Authentication, /// Endpoint base URL. pub base_url: String, + /// Optional retry and timeout policy shared by chat, files, batches, and models. + /// `None` preserves the original single-attempt, no-timeout behavior. + pub resilience: Option, /// `X-Cerebras-Version-Patch` header. `None` opts into the current API /// default. pub version_patch: Option, @@ -425,16 +429,62 @@ pub struct CerebrasConfig { pub compression: Option, } +impl std::fmt::Debug for CerebrasConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let mut debug = f.debug_struct("CerebrasConfig"); + debug + .field("authentication", &"") + .field("base_url", &self.base_url) + .field("resilience", &self.resilience) + .field("version_patch", &self.version_patch) + .field( + "extra_header_names", + &self + .extra_headers + .iter() + .map(|(name, _)| name) + .collect::>(), + ) + .field("model", &self.model) + .field("max_completion_tokens", &self.max_completion_tokens) + .field("min_tokens", &self.min_tokens) + .field("temperature", &self.temperature) + .field("top_p", &self.top_p) + .field("frequency_penalty", &self.frequency_penalty) + .field("presence_penalty", &self.presence_penalty) + .field("stop", &self.stop) + .field("seed", &self.seed) + .field("logprobs", &self.logprobs) + .field("top_logprobs", &self.top_logprobs) + .field("tool_choice", &self.tool_choice) + .field("parallel_tool_calls", &self.parallel_tool_calls) + .field("tool_strict", &self.tool_strict) + .field("output_format", &self.output_format) + .field("reasoning", &self.reasoning) + .field("streaming", &self.streaming); + #[cfg(feature = "predicted-outputs")] + debug.field("prediction", &self.prediction); + #[cfg(feature = "service-tiers")] + debug + .field("service_tier", &self.service_tier) + .field("queue_threshold_ms", &self.queue_threshold_ms); + #[cfg(feature = "compression")] + debug.field("compression", &self.compression); + debug.finish_non_exhaustive() + } +} + impl CerebrasConfig { - /// Creates a new configuration with the given API key and model. - pub fn new(api_key: impl Into, model: impl Into) -> Result { - let api_key = api_key.into(); - if api_key.is_empty() { - return Err(BuildError::MissingEnv("CEREBRAS_API_KEY")); - } + /// Creates a new configuration with the given authentication and model. + /// String authentication values become bearer tokens. + pub fn new( + authentication: impl Into, + model: impl Into, + ) -> Result { Ok(Self { - api_key, + authentication: authentication.into(), base_url: DEFAULT_BASE_URL.into(), + resilience: None, version_patch: DEFAULT_VERSION_PATCH, extra_headers: Vec::new(), extra_body: None, @@ -480,6 +530,9 @@ impl CerebrasConfig { pub fn from_env() -> Result { let api_key = std::env::var("CEREBRAS_API_KEY") .map_err(|_| BuildError::MissingEnv("CEREBRAS_API_KEY"))?; + if api_key.is_empty() { + return Err(BuildError::MissingEnv("CEREBRAS_API_KEY")); + } let model = std::env::var("CEREBRAS_MODEL") .map_err(|_| BuildError::MissingEnv("CEREBRAS_MODEL"))?; let mut config = Self::new(api_key, model)?; @@ -571,6 +624,29 @@ impl CerebrasConfig { // --- Builder methods --- + /// Replaces authentication for every Cerebras API surface. + pub fn with_authentication(mut self, authentication: impl Into) -> Self { + self.authentication = authentication.into(); + self + } + + /// Compatibility builder for bearer API-key authentication. + pub fn with_api_key(self, api_key: impl Into) -> Self { + self.with_authentication(api_key) + } + + /// Uses a custom refresh-capable authentication provider. + pub fn with_authentication_provider(mut self, provider: P) -> Self { + self.authentication = Authentication::new(provider); + self + } + + /// Enables request retries and timeouts for every Cerebras API surface. + pub fn with_resilience(mut self, resilience: ResilienceConfig) -> Self { + self.resilience = Some(resilience); + self + } + /// Overrides the endpoint base URL. pub fn with_base_url(mut self, url: impl Into) -> Self { self.base_url = url.into(); @@ -737,9 +813,26 @@ mod tests { use super::*; #[test] - fn rejects_empty_api_key() { - let err = CerebrasConfig::new("", "gpt-oss-120b").unwrap_err(); - assert!(matches!(err, BuildError::MissingEnv(_))); + fn config_debug_redacts_authentication_and_extra_header_values() { + let mut config = CerebrasConfig::new("cerebras-secret", "debug-model").unwrap(); + config + .extra_headers + .push(("x-private".into(), "header-secret".into())); + let debug = format!("{config:?}"); + assert!(!debug.contains("cerebras-secret")); + assert!(!debug.contains("header-secret")); + assert!(debug.contains("")); + assert!(debug.contains("debug-model")); + assert!(debug.contains("x-private")); + } + + #[test] + fn config_accepts_first_class_authentication_and_resilience() { + let config = CerebrasConfig::new(Authentication::bearer("secret"), "gpt-oss-120b") + .unwrap() + .with_resilience(ResilienceConfig::default()); + assert!(config.resilience.is_some()); + assert!(!format!("{config:?}").contains("secret")); } #[test] diff --git a/crates/agentkit-provider-cerebras/src/files.rs b/crates/agentkit-provider-cerebras/src/files.rs index 4871bc6..18bdba3 100644 --- a/crates/agentkit-provider-cerebras/src/files.rs +++ b/crates/agentkit-provider-cerebras/src/files.rs @@ -9,6 +9,7 @@ use agentkit_http::{BodyStream, Http}; use bytes::Bytes; use serde::{Deserialize, Serialize}; +use crate::RequestExecutor; use crate::config::CerebrasConfig; use crate::error::CerebrasError; @@ -60,6 +61,14 @@ impl<'a> FilesClient<'a> { Self { http, config } } + fn executor(&self) -> RequestExecutor { + RequestExecutor::new( + self.http, + self.config.authentication.clone(), + self.config.resilience.clone(), + ) + } + /// Uploads a JSONL batch file. Returns the created [`FileObject`]. pub async fn upload( &self, @@ -78,12 +87,15 @@ impl<'a> FilesClient<'a> { let body = build_multipart(&boundary, filename, bytes.into(), purpose); let content_type = format!("multipart/form-data; boundary={boundary}"); let response = self - .http - .post(&url) - .bearer_auth(&self.config.api_key) - .header("Content-Type", content_type) - .body(body) - .send() + .executor() + // Uploads have no stable provider idempotency contract. Bound the + // call, including auth and body reads, but never retry the POST. + .execute_buffered_without_retries(|| { + self.http + .post(&url) + .header("Content-Type", content_type.clone()) + .body(body.clone()) + }) .await?; if !response.status().is_success() { let status = response.status().as_u16(); @@ -97,10 +109,8 @@ impl<'a> FilesClient<'a> { pub async fn list(&self) -> Result, CerebrasError> { let url = format!("{}/files", self.config.base_url); let response = self - .http - .get(&url) - .bearer_auth(&self.config.api_key) - .send() + .executor() + .execute_buffered(|| self.http.get(&url)) .await?; if !response.status().is_success() { let status = response.status().as_u16(); @@ -115,10 +125,8 @@ impl<'a> FilesClient<'a> { pub async fn retrieve(&self, id: &str) -> Result { let url = format!("{}/files/{id}", self.config.base_url); let response = self - .http - .get(&url) - .bearer_auth(&self.config.api_key) - .send() + .executor() + .execute_buffered(|| self.http.get(&url)) .await?; if !response.status().is_success() { let status = response.status().as_u16(); @@ -129,37 +137,42 @@ impl<'a> FilesClient<'a> { } /// Streams a file's content (batch output or error reports). + /// + /// With resilience enabled, opening this idempotent GET may be retried. The + /// returned stream enforces the same logical deadline, configured idle + /// timeout, and content-length truncation detection. It reports stream + /// failures directly and never replays after the stream is returned. pub async fn content(&self, id: &str) -> Result { let url = format!("{}/files/{id}/content", self.config.base_url); - let response = self - .http - .get(&url) - .bearer_auth(&self.config.api_key) - .send() - .await?; + let executor = self.executor(); + let response = executor.execute(|| self.http.get(&url)).await?; if !response.status().is_success() { let status = response.status().as_u16(); - let body = response.text().await.unwrap_or_default(); + let body = executor + .read_response_text(response) + .await + .unwrap_or_default(); return Err(CerebrasError::Status { status, body }); } - Ok(response.bytes_stream()) + Ok(executor.guard_body_stream(response)) } - /// Deletes a file. + /// Deletes a file without replaying the request or buffering a successful response. pub async fn delete(&self, id: &str) -> Result<(), CerebrasError> { let url = format!("{}/files/{id}", self.config.base_url); - let response = self - .http - .delete(&url) - .bearer_auth(&self.config.api_key) - .send() + let executor = self.executor(); + let response = executor + .execute_without_retries(|| self.http.delete(&url)) .await?; - if !response.status().is_success() { - let status = response.status().as_u16(); - let body = response.text().await.unwrap_or_default(); - return Err(CerebrasError::Status { status, body }); + if response.status().is_success() { + return Ok(()); } - Ok(()) + let status = response.status().as_u16(); + let body = executor + .read_response_text(response) + .await + .unwrap_or_default(); + Err(CerebrasError::Status { status, body }) } } @@ -202,8 +215,37 @@ fn build_multipart( #[cfg(test)] mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use agentkit_http::{ + HttpClient, HttpError, HttpRequest, HttpResponse, ResilienceConfig, StatusCode, + }; + use async_trait::async_trait; + use super::*; + struct DeleteClient { + calls: AtomicUsize, + fail_first: bool, + } + + #[async_trait] + impl HttpClient for DeleteClient { + async fn execute(&self, request: HttpRequest) -> Result { + assert_eq!(request.method, agentkit_http::Method::DELETE); + let call = self.calls.fetch_add(1, Ordering::SeqCst); + if self.fail_first && call == 0 { + return Err(HttpError::request(std::io::Error::other("disconnected"))); + } + Ok(HttpResponse::new( + StatusCode::NO_CONTENT, + agentkit_http::HeaderMap::new(), + request.url, + Box::pin(futures_util::stream::pending()), + )) + } + } + #[test] fn multipart_body_has_purpose_and_file_parts() { let body = build_multipart( @@ -220,4 +262,45 @@ mod tests { assert!(text.contains("{}\n")); assert!(text.trim_end().ends_with("--BND--")); } + + #[tokio::test] + async fn delete_returns_without_buffering_a_success_body() { + let client = Arc::new(DeleteClient { + calls: AtomicUsize::new(0), + fail_first: false, + }); + let http = Http::from_arc(client.clone()); + let config = Arc::new( + CerebrasConfig::new("secret", "test-model") + .unwrap() + .with_base_url("https://example.test"), + ); + let files = FilesClient::new(&http, config); + + files.delete("file-1").await.unwrap(); + + assert_eq!(client.calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn delete_does_not_replay_transport_failures() { + let client = Arc::new(DeleteClient { + calls: AtomicUsize::new(0), + fail_first: true, + }); + let http = Http::from_arc(client.clone()); + let config = Arc::new( + CerebrasConfig::new("secret", "test-model") + .unwrap() + .with_base_url("https://example.test") + .with_resilience(ResilienceConfig { + max_retries: 3, + ..ResilienceConfig::default() + }), + ); + let files = FilesClient::new(&http, config); + + assert!(files.delete("file-1").await.is_err()); + assert_eq!(client.calls.load(Ordering::SeqCst), 1); + } } diff --git a/crates/agentkit-provider-cerebras/src/lib.rs b/crates/agentkit-provider-cerebras/src/lib.rs index 579c7fc..856ba60 100644 --- a/crates/agentkit-provider-cerebras/src/lib.rs +++ b/crates/agentkit-provider-cerebras/src/lib.rs @@ -40,16 +40,22 @@ mod sse; mod stream; use std::collections::VecDeque; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; +use std::time::Duration; use agentkit_core::TurnCancellation; -use agentkit_http::{BodyStream, Http, HttpError, HttpRequestBuilder}; +use agentkit_http::{ + Authentication, AuthenticationAttempt, AuthenticationProvider, BodyStream, Bytes, Http, + HttpError, HttpRequestBuilder, HttpResponse, LogicalDeadline, ResilienceConfig, StatusCode, + TruncatedStreamDetector, is_retryable_body_read, is_retryable_status, next_body_chunk_bounded, + run_bounded, sleep as resilience_sleep, +}; use agentkit_loop::{ LoopError, ModelAdapter, ModelSession, ModelTurn, ModelTurnEvent, SessionConfig, TurnRequest, }; use async_trait::async_trait; -use futures_util::StreamExt; -use futures_util::future::{Either, select}; +use futures_util::future::{BoxFuture, Either, select}; pub use crate::config::{ CerebrasConfig, DEFAULT_BASE_URL, DEFAULT_VERSION_PATCH, OutputFormat, PartKindName, @@ -93,11 +99,7 @@ impl CerebrasAdapter { .build() .map(Http::new) .map_err(|error| CerebrasError::Http(HttpError::request(error)))?; - Ok(Self { - client, - config: Arc::new(config), - last_rate_limit: Arc::new(Mutex::new(None)), - }) + Self::with_client(config, client) } /// Creates a new adapter using a pre-configured [`Http`] client. @@ -110,6 +112,23 @@ impl CerebrasAdapter { }) } + /// Overrides the configured authentication. + pub fn with_authentication(mut self, authentication: impl Into) -> Self { + Arc::make_mut(&mut self.config).authentication = authentication.into(); + self + } + + /// Overrides authentication with a custom refresh-capable provider. + pub fn with_authentication_provider(self, provider: P) -> Self { + self.with_authentication(Authentication::new(provider)) + } + + /// Enables request and pre-visible-output stream retries and timeouts. + pub fn with_resilience(mut self, resilience: ResilienceConfig) -> Self { + Arc::make_mut(&mut self.config).resilience = Some(resilience); + self + } + /// Reads the latest rate-limit snapshot, if any response has been received. pub fn last_rate_limit(&self) -> Option { self.last_rate_limit.lock().ok()?.clone() @@ -151,12 +170,442 @@ enum TurnInner { Streaming(Box), } +type StreamReplay = Box< + dyn FnMut() -> BoxFuture<'static, Result<(BodyStream, TruncatedStreamDetector), LoopError>> + + Send, +>; + struct StreamingState { body: BodyStream, decoder: SseDecoder, translator: EventTranslator, pending: VecDeque, eof: bool, + visible_output: bool, + detect_truncation: bool, + idle_timeout: Option, + deadline: Option, + integrity: TruncatedStreamDetector, + replay: Option, +} + +#[cfg(any(feature = "batch", test))] +struct GuardedBodyState { + body: BodyStream, + idle_timeout: Option, + deadline: Option, + integrity: TruncatedStreamDetector, + done: bool, +} + +pub(crate) struct BufferedResponse { + status: StatusCode, + body: Bytes, +} + +impl BufferedResponse { + pub(crate) fn status(&self) -> StatusCode { + self.status + } + + pub(crate) async fn text(self) -> Result { + String::from_utf8(self.body.to_vec()).map_err(|error| HttpError::Body(Box::new(error))) + } + + pub(crate) async fn json(self) -> Result { + serde_json::from_slice(&self.body).map_err(HttpError::Deserialize) + } +} + +/// Authentication and resilience state for one replayable logical request. +#[derive(Clone)] +pub(crate) struct RequestExecutor { + client: Http, + authentication: Authentication, + authentication_attempt: Arc>>, + reauthenticated: Arc, + resilience: Option, + deadline: Option, + retries_used: Arc, +} + +impl RequestExecutor { + pub(crate) fn new( + client: &Http, + authentication: Authentication, + resilience: Option, + ) -> Self { + let deadline = resilience + .as_ref() + .map(|resilience| LogicalDeadline::new(resilience.retry_budget)); + Self { + client: client.clone(), + authentication, + authentication_attempt: Arc::new(Mutex::new(None)), + reauthenticated: Arc::new(AtomicBool::new(false)), + resilience, + deadline, + retries_used: Arc::new(AtomicUsize::new(0)), + } + } + + async fn authentication_attempt(&self) -> Result { + let existing = self + .authentication_attempt + .lock() + .map_err(|_| HttpError::Other("authentication state lock poisoned".into()))? + .clone(); + if let Some(existing) = existing { + return Ok(existing); + } + let attempt = run_bounded( + self.authentication.authenticate(None), + None, + self.deadline.as_ref(), + "authentication", + ) + .await?; + *self + .authentication_attempt + .lock() + .map_err(|_| HttpError::Other("authentication state lock poisoned".into()))? = + Some(attempt.clone()); + Ok(attempt) + } + + async fn reauthenticate(&self, previous: AuthenticationAttempt) -> Result<(), HttpError> { + let next = run_bounded( + self.authentication.authenticate(Some(&previous)), + None, + self.deadline.as_ref(), + "reauthentication", + ) + .await?; + *self + .authentication_attempt + .lock() + .map_err(|_| HttpError::Other("authentication state lock poisoned".into()))? = + Some(next); + Ok(()) + } + + pub(crate) fn reserve_retry(&self) -> Option { + let maximum = self.resilience.as_ref()?.max_retries; + self.retries_used + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |used| { + (used < maximum).then_some(used + 1) + }) + .ok() + } + + fn retry_delay( + &self, + retry_number: usize, + headers: Option<&agentkit_http::HeaderMap>, + ) -> Duration { + self.resilience + .as_ref() + .expect("reserved retry requires resilience") + .retry_delay(retry_number, headers) + } + + pub(crate) async fn wait_before_retry( + &self, + retry_number: usize, + headers: Option<&agentkit_http::HeaderMap>, + ) -> Result<(), HttpError> { + self.wait_for_retry_delay(self.retry_delay(retry_number, headers)) + .await + } + + async fn wait_for_retry_delay(&self, delay: Duration) -> Result<(), HttpError> { + run_bounded( + async { + resilience_sleep(delay).await; + Ok(()) + }, + None, + self.deadline.as_ref(), + "retry backoff", + ) + .await + } + + async fn execute_attempt( + &self, + request: agentkit_http::HttpRequest, + ) -> Result { + let timeout = self + .resilience + .as_ref() + .and_then(|config| config.attempt_timeout); + run_bounded( + self.client.execute(request), + timeout, + self.deadline.as_ref(), + "HTTP request attempt", + ) + .await + } + + async fn read_response_bytes(&self, response: HttpResponse) -> Result { + let timeout = self + .resilience + .as_ref() + .and_then(|config| config.attempt_timeout); + run_bounded( + response.bytes(), + timeout, + self.deadline.as_ref(), + "HTTP response body", + ) + .await + } + + pub(crate) async fn read_response_text( + &self, + response: HttpResponse, + ) -> Result { + let timeout = self + .resilience + .as_ref() + .and_then(|config| config.attempt_timeout); + run_bounded( + response.text(), + timeout, + self.deadline.as_ref(), + "HTTP response body", + ) + .await + } + + pub(crate) async fn execute(&self, build: F) -> Result + where + F: Fn() -> HttpRequestBuilder, + { + self.execute_inner(build, true).await + } + + #[cfg(feature = "batch")] + pub(crate) async fn execute_without_retries( + &self, + build: F, + ) -> Result + where + F: Fn() -> HttpRequestBuilder, + { + self.execute_inner(build, false).await + } + + async fn execute_inner(&self, build: F, retry_safe: bool) -> Result + where + F: Fn() -> HttpRequestBuilder, + { + loop { + let authentication = self.authentication_attempt().await?; + let request = build().headers(authentication.headers().clone()).build()?; + match self.execute_attempt(request).await { + Ok(response) => { + if response.status() == StatusCode::UNAUTHORIZED + && !self.reauthenticated.swap(true, Ordering::SeqCst) + { + drop(response); + self.reauthenticate(authentication).await?; + continue; + } + if retry_safe + && is_retryable_status(response.status()) + && let Some(retry_number) = self.reserve_retry() + { + let delay = self.retry_delay(retry_number, Some(response.headers())); + drop(response); + self.wait_for_retry_delay(delay).await?; + continue; + } + return Ok(response); + } + Err(error) if retry_safe && error.is_retryable_transport() => { + if let Some(retry_number) = self.reserve_retry() { + self.wait_before_retry(retry_number, None).await?; + continue; + } + return Err(error); + } + Err(error) => return Err(error), + } + } + } + + pub(crate) async fn execute_buffered(&self, build: F) -> Result + where + F: Fn() -> HttpRequestBuilder, + { + self.execute_buffered_inner(build, true).await + } + + #[cfg(any(feature = "batch", test))] + pub(crate) async fn execute_buffered_without_retries( + &self, + build: F, + ) -> Result + where + F: Fn() -> HttpRequestBuilder, + { + self.execute_buffered_inner(build, false).await + } + + async fn execute_buffered_inner( + &self, + build: F, + retry_safe: bool, + ) -> Result + where + F: Fn() -> HttpRequestBuilder, + { + loop { + let response = self.execute_inner(&build, retry_safe).await?; + let status = response.status(); + match self.read_response_bytes(response).await { + Ok(body) => return Ok(BufferedResponse { status, body }), + Err(error) if retry_safe && is_retryable_body_read(status, &error) => { + if let Some(retry_number) = self.reserve_retry() { + self.wait_before_retry(retry_number, None).await?; + continue; + } + return Err(error); + } + Err(error) => return Err(error), + } + } + } + + /// Applies the configured logical deadline, idle timeout, and content-length + /// truncation check without changing the public `BodyStream` type. The + /// stream reports failures but never replays after it has been returned. + #[cfg(any(feature = "batch", test))] + pub(crate) fn guard_body_stream(&self, response: HttpResponse) -> BodyStream { + if self.resilience.is_none() { + return response.bytes_stream(); + } + let state = GuardedBodyState { + idle_timeout: self + .resilience + .as_ref() + .and_then(|resilience| resilience.stream_idle_timeout), + deadline: self.deadline.clone(), + integrity: TruncatedStreamDetector::from_headers(response.headers()), + body: response.bytes_stream(), + done: false, + }; + Box::pin(futures_util::stream::unfold( + state, + |mut state| async move { + if state.done { + return None; + } + let next = next_body_chunk_bounded( + &mut state.body, + state.idle_timeout, + state.deadline.as_ref(), + ) + .await; + match next { + Ok(Some(bytes)) => { + state.integrity.observe(&bytes); + Some((Ok(bytes), state)) + } + Ok(None) => match state.integrity.finish() { + Ok(()) => None, + Err(error) => { + state.done = true; + Some((Err(error), state)) + } + }, + Err(error) => { + state.done = true; + Some((Err(error), state)) + } + } + }, + )) + } +} + +#[derive(Clone)] +struct InferenceRequest { + executor: RequestExecutor, + config: Arc, + body: Bytes, + content_type: &'static str, + content_encoding: Option<&'static str>, + extra_headers: Vec<(&'static str, String)>, + rate_limit_slot: Arc>>, +} + +impl InferenceRequest { + async fn open_response(&self) -> Result { + let url = format!("{}/chat/completions", self.config.base_url); + let response = self + .executor + .execute(|| { + let mut builder = self + .executor + .client + .post(&url) + .header("Content-Type", self.content_type); + if let Some(encoding) = self.content_encoding { + builder = builder.header("Content-Encoding", encoding); + } + if let Some(patch) = self.config.version_patch { + builder = builder.header( + crate::version::VERSION_PATCH_HEADER, + crate::version::format_version_patch(patch), + ); + } + for (name, value) in &self.extra_headers { + builder = builder.header(*name, value.clone()); + } + builder = builder.header( + "User-Agent", + concat!("agentkit-provider-cerebras/", env!("CARGO_PKG_VERSION")), + ); + if self.config.streaming { + builder = builder.header("Accept", "text/event-stream"); + } + for (name, value) in &self.config.extra_headers { + builder = builder.header(name.as_str(), value.as_str()); + } + builder.body(self.body.clone()) + }) + .await + .map_err(|error| LoopError::Provider(format!("Cerebras request failed: {error}")))?; + let snapshot = RateLimitSnapshot::from_headers(response.headers()); + if let Ok(mut slot) = self.rate_limit_slot.lock() { + *slot = Some(snapshot); + } + Ok(response) + } + + async fn replay_stream(&self) -> Result<(BodyStream, TruncatedStreamDetector), LoopError> { + let retry_number = self.executor.reserve_retry().ok_or_else(|| { + LoopError::Provider( + "Cerebras stream failed before output and retry budget is exhausted".into(), + ) + })?; + self.executor + .wait_before_retry(retry_number, None) + .await + .map_err(|error| LoopError::Provider(format!("retry backoff failed: {error}")))?; + let response = self.open_response().await?; + if !response.status().is_success() { + return Err(LoopError::Provider(format!( + "Cerebras stream replay failed with status {}", + response.status() + ))); + } + let integrity = TruncatedStreamDetector::from_headers(response.headers()); + Ok((response.bytes_stream(), integrity)) + } } #[async_trait] @@ -189,26 +638,23 @@ impl ModelSession for CerebrasSession { let config = self.config.clone(); let rate_limit_slot = self.rate_limit_slot.clone(); - let request_future = async move { + let request_future = async { let built = request::build_chat_body(&config, &turn_request) .map_err(|e| LoopError::Provider(e.to_string()))?; - let url = format!("{}/chat/completions", config.base_url); - let mut http = self.client.post(&url).bearer_auth(&config.api_key); - #[cfg(feature = "compression")] - let (body_bytes, content_type, content_encoding) = match &config.compression { + let (body, content_type, content_encoding) = match &config.compression { Some(cfg) => { let encoded = crate::compression::encode_body(&built.body, cfg) .map_err(LoopError::Provider)?; ( - bytes::Bytes::from(encoded.body), + Bytes::from(encoded.body), encoded.content_type, encoded.content_encoding, ) } None => ( - bytes::Bytes::from( + Bytes::from( serde_json::to_vec(&built.body) .map_err(|e| LoopError::Provider(format!("json serialize: {e}")))?, ), @@ -217,8 +663,8 @@ impl ModelSession for CerebrasSession { ), }; #[cfg(not(feature = "compression"))] - let (body_bytes, content_type, content_encoding) = ( - bytes::Bytes::from( + let (body, content_type, content_encoding) = ( + Bytes::from( serde_json::to_vec(&built.body) .map_err(|e| LoopError::Provider(format!("json serialize: {e}")))?, ), @@ -226,70 +672,94 @@ impl ModelSession for CerebrasSession { None::<&'static str>, ); - http = http.header("Content-Type", content_type); - if let Some(enc) = content_encoding { - http = http.header("Content-Encoding", enc); - } - if let Some(patch) = config.version_patch { - http = http.header( - crate::version::VERSION_PATCH_HEADER, - crate::version::format_version_patch(patch), - ); - } - for (k, v) in &built.extra_headers { - http = http.header(*k, v.clone()); - } - http = http.header( - "User-Agent", - concat!("agentkit-provider-cerebras/", env!("CARGO_PKG_VERSION")), - ); - if config.streaming { - http = http.header("Accept", "text/event-stream"); - } - for (k, v) in &config.extra_headers { - http = http.header(k.as_str(), v.as_str()); - } - http = attach_body(http, body_bytes); - - let response = http.send().await.map_err(|error| { - LoopError::Provider(format!("Cerebras request failed: {error}")) - })?; + let inference_request = InferenceRequest { + executor: RequestExecutor::new( + &self.client, + config.authentication.clone(), + config.resilience.clone(), + ), + config: config.clone(), + body, + content_type, + content_encoding, + extra_headers: built.extra_headers, + rate_limit_slot, + }; - { - let snap = RateLimitSnapshot::from_headers(response.headers()); - if let Ok(mut slot) = rate_limit_slot.lock() { - *slot = Some(snap); + loop { + let response = inference_request.open_response().await?; + let status = response.status(); + if config.streaming && status.is_success() { + let integrity = TruncatedStreamDetector::from_headers(response.headers()); + let idle_timeout = config + .resilience + .as_ref() + .and_then(|resilience| resilience.stream_idle_timeout); + let replay = config.resilience.as_ref().and_then(|resilience| { + (resilience.max_retries > 0).then(|| { + let inference_request = inference_request.clone(); + Box::new(move || { + let inference_request = inference_request.clone(); + Box::pin(async move { inference_request.replay_stream().await }) + as BoxFuture<'static, _> + }) as StreamReplay + }) + }); + return Ok(CerebrasTurn { + inner: TurnInner::Streaming(Box::new(StreamingState { + body: response.bytes_stream(), + decoder: SseDecoder::new(), + translator: EventTranslator::new(), + pending: VecDeque::new(), + eof: false, + visible_output: false, + detect_truncation: config.resilience.is_some(), + idle_timeout, + deadline: inference_request.executor.deadline.clone(), + integrity, + replay, + })), + }); } - } - let status = response.status(); - if !status.is_success() { - let body_text = response.text().await.unwrap_or_default(); - return Err(LoopError::Provider(format!( - "Cerebras request failed with status {status}: {body_text}" - ))); - } - - if config.streaming { - Ok(CerebrasTurn { - inner: TurnInner::Streaming(Box::new(StreamingState { - body: response.bytes_stream(), - decoder: SseDecoder::new(), - translator: EventTranslator::new(), - pending: VecDeque::new(), - eof: false, - })), - }) - } else { - let body_text = response.text().await.map_err(|error| { - LoopError::Provider(format!("failed to read Cerebras response body: {error}")) - })?; + let body_text = match inference_request + .executor + .read_response_text(response) + .await + { + Ok(body) => body, + Err(error) if is_retryable_body_read(status, &error) => { + if let Some(retry_number) = inference_request.executor.reserve_retry() { + inference_request + .executor + .wait_before_retry(retry_number, None) + .await + .map_err(|error| { + LoopError::Provider(format!("retry backoff failed: {error}")) + })?; + continue; + } + return Err(LoopError::Provider(format!( + "failed to read Cerebras response body: {error}" + ))); + } + Err(error) => { + return Err(LoopError::Provider(format!( + "failed to read Cerebras response body: {error}" + ))); + } + }; + if !status.is_success() { + return Err(LoopError::Provider(format!( + "Cerebras request failed with status {status}: {body_text}" + ))); + } let events = response::build_turn_from_response(&body_text) .map_err(|e| LoopError::Provider(e.to_string()))?; - Ok(CerebrasTurn { + return Ok(CerebrasTurn { inner: TurnInner::Buffered { events }, - }) + }); } }; @@ -315,10 +785,6 @@ impl ModelSession for CerebrasSession { } } -fn attach_body(builder: HttpRequestBuilder, body: bytes::Bytes) -> HttpRequestBuilder { - builder.body(body) -} - #[async_trait] impl ModelTurn for CerebrasTurn { async fn next_event( @@ -340,67 +806,141 @@ impl ModelTurn for CerebrasTurn { translator, pending, eof, + visible_output, + detect_truncation, + idle_timeout, + deadline, + integrity, + replay, } = state.as_mut(); - next_streaming_event(body, decoder, translator, pending, eof, cancellation).await + next_streaming_event( + body, + decoder, + translator, + pending, + eof, + visible_output, + *detect_truncation, + *idle_timeout, + deadline.as_ref(), + integrity, + replay, + cancellation, + ) + .await } } } } +async fn next_stream_chunk( + body: &mut BodyStream, + idle_timeout: Option, + deadline: Option<&LogicalDeadline>, + cancellation: Option<&TurnCancellation>, +) -> Result, HttpError>, LoopError> { + let next = next_body_chunk_bounded(body, idle_timeout, deadline); + futures_util::pin_mut!(next); + if let Some(cancellation) = cancellation { + let cancelled = cancellation.cancelled(); + futures_util::pin_mut!(cancelled); + match select(next, cancelled).await { + Either::Left((chunk, _)) => Ok(chunk), + Either::Right((_, _)) => Err(LoopError::Cancelled), + } + } else { + Ok(next.await) + } +} + +#[allow(clippy::too_many_arguments)] async fn next_streaming_event( body: &mut BodyStream, decoder: &mut SseDecoder, translator: &mut EventTranslator, pending: &mut VecDeque, eof: &mut bool, + visible_output: &mut bool, + detect_truncation: bool, + idle_timeout: Option, + deadline: Option<&LogicalDeadline>, + integrity: &mut TruncatedStreamDetector, + replay: &mut Option, cancellation: Option, ) -> Result, LoopError> { loop { if let Some(event) = pending.pop_front() { + *visible_output = true; return Ok(Some(event)); } if *eof || translator.is_done() { return Ok(None); } - let chunk = if let Some(cancellation) = cancellation.as_ref() { - let next = body.next(); - futures_util::pin_mut!(next); - let cancelled = cancellation.cancelled(); - futures_util::pin_mut!(cancelled); - match select(next, cancelled).await { - Either::Left((chunk, _)) => chunk, - Either::Right((_, _)) => return Err(LoopError::Cancelled), - } - } else { - body.next().await - }; - - match chunk { - Some(Ok(bytes)) => { + let chunk = next_stream_chunk(body, idle_timeout, deadline, cancellation.as_ref()).await?; + let failure = match chunk { + Ok(Some(bytes)) => { + integrity.observe(&bytes); let text = std::str::from_utf8(&bytes).map_err(|e| { LoopError::Provider(format!("invalid UTF-8 in Cerebras stream: {e}")) })?; for sse in decoder.feed(text) { - match translator.handle(&sse) { - Ok(produced) => { - for ev in produced { - pending.push_back(ev); - } - } - Err(e) => return Err(LoopError::Provider(e.to_string())), + for event in translator + .handle(&sse) + .map_err(|e| LoopError::Provider(e.to_string()))? + { + pending.push_back(event); } } + None } - Some(Err(e)) => { - return Err(LoopError::Provider(format!( - "Cerebras stream body error: {e}" - ))); + Ok(None) if !detect_truncation => { + *eof = true; + None } - None => { + Ok(None) => { *eof = true; + let message = match integrity.finish() { + Ok(()) => "Cerebras stream ended before its terminal event".to_owned(), + Err(error) => format!("Cerebras stream body error: {error}"), + }; + Some(LoopError::Provider(message)) } + Err(error) if !error.is_retryable_transport() => { + return Err(LoopError::Provider(format!( + "Cerebras stream body error: {error}" + ))); + } + Err(error) => Some(LoopError::Provider(format!( + "Cerebras stream body error: {error}" + ))), + }; + + let Some(failure) = failure else { + continue; + }; + if *visible_output || replay.is_none() { + return Err(failure); } + + let replay_future = replay.as_mut().expect("checked above")(); + futures_util::pin_mut!(replay_future); + let replayed = if let Some(cancellation) = cancellation.as_ref() { + let cancelled = cancellation.cancelled(); + futures_util::pin_mut!(cancelled); + match select(replay_future, cancelled).await { + Either::Left((result, _)) => result, + Either::Right((_, _)) => return Err(LoopError::Cancelled), + } + } else { + replay_future.await + }?; + *body = replayed.0; + *integrity = replayed.1; + *decoder = SseDecoder::new(); + *translator = EventTranslator::new(); + pending.clear(); + *eof = false; } } @@ -408,9 +948,186 @@ async fn next_streaming_event( mod tests { use super::*; use agentkit_core::{CancellationController, FinishReason}; - use agentkit_http::HttpError; + use agentkit_http::{HeaderMap, HeaderValue, HttpClient, HttpError, HttpRequest, header}; use bytes::Bytes; - use futures_util::stream; + use futures_util::{StreamExt, stream}; + + struct AlwaysUnauthorized { + authorizations: Mutex>, + } + + #[async_trait] + impl HttpClient for AlwaysUnauthorized { + async fn execute(&self, request: HttpRequest) -> Result { + let authorization = request.headers[header::AUTHORIZATION] + .to_str() + .unwrap() + .to_owned(); + self.authorizations.lock().unwrap().push(authorization); + Ok(HttpResponse::new( + StatusCode::UNAUTHORIZED, + HeaderMap::new(), + request.url, + Box::pin(stream::empty()), + )) + } + } + + struct FixedStatus { + calls: AtomicUsize, + status: StatusCode, + } + + #[async_trait] + impl HttpClient for FixedStatus { + async fn execute(&self, request: HttpRequest) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(HttpResponse::new( + self.status, + HeaderMap::new(), + request.url, + Box::pin(stream::empty()), + )) + } + } + + struct RefreshingAuthentication { + calls: Arc, + } + + #[async_trait] + impl AuthenticationProvider for RefreshingAuthentication { + async fn authenticate( + &self, + previous: Option<&AuthenticationAttempt>, + ) -> Result { + let generation = self.calls.fetch_add(1, Ordering::SeqCst); + assert_eq!( + previous + .and_then(|attempt| attempt.state::()) + .copied(), + generation.checked_sub(1) + ); + let mut headers = HeaderMap::new(); + headers.insert( + header::AUTHORIZATION, + HeaderValue::from_str(&format!("Bearer refreshed-{generation}")).unwrap(), + ); + Ok(AuthenticationAttempt::new(headers, generation)) + } + } + + #[tokio::test] + async fn request_executor_replays_unauthorized_exactly_once() { + let client = Arc::new(AlwaysUnauthorized { + authorizations: Mutex::new(Vec::new()), + }); + let calls = Arc::new(AtomicUsize::new(0)); + let http = Http::from_arc(client.clone()); + let executor = RequestExecutor::new( + &http, + Authentication::new(RefreshingAuthentication { + calls: calls.clone(), + }), + None, + ); + + let response = executor + .execute(|| http.get("https://example.test/models")) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + assert_eq!(calls.load(Ordering::SeqCst), 2); + assert_eq!( + *client.authorizations.lock().unwrap(), + ["Bearer refreshed-0", "Bearer refreshed-1"] + ); + } + + #[tokio::test] + async fn logical_deadline_starts_before_auxiliary_authentication() { + let client = Arc::new(FixedStatus { + calls: AtomicUsize::new(0), + status: StatusCode::OK, + }); + let http = Http::from_arc(client.clone()); + let resilience = ResilienceConfig { + retry_budget: Duration::ZERO, + ..ResilienceConfig::default() + }; + let executor = + RequestExecutor::new(&http, Authentication::bearer("unused"), Some(resilience)); + + let error = executor + .execute(|| http.get("https://example.test/models")) + .await + .unwrap_err(); + + assert!(error.to_string().contains("logical request retry budget")); + assert_eq!(client.calls.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn unsafe_post_policy_does_not_retry_retryable_statuses() { + let client = Arc::new(FixedStatus { + calls: AtomicUsize::new(0), + status: StatusCode::SERVICE_UNAVAILABLE, + }); + let http = Http::from_arc(client.clone()); + let resilience = ResilienceConfig { + initial_backoff: Duration::ZERO, + max_backoff: Duration::ZERO, + ..ResilienceConfig::default() + }; + let executor = + RequestExecutor::new(&http, Authentication::bearer("unused"), Some(resilience)); + + let response = executor + .execute_buffered_without_retries(|| http.post("https://example.test/files")) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(client.calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn guarded_file_stream_reports_content_length_truncation() { + let http = Http::from_arc(Arc::new(FixedStatus { + calls: AtomicUsize::new(0), + status: StatusCode::OK, + })); + let executor = RequestExecutor::new( + &http, + Authentication::bearer("unused"), + Some(ResilienceConfig::default()), + ); + let mut headers = HeaderMap::new(); + headers.insert(header::CONTENT_LENGTH, HeaderValue::from_static("5")); + let response = HttpResponse::new( + StatusCode::OK, + headers, + "https://example.test/files/1/content".into(), + Box::pin(stream::once(async { + Ok::<_, HttpError>(Bytes::from_static(b"123")) + })), + ); + let mut body = executor.guard_body_stream(response); + + assert_eq!( + body.next().await.unwrap().unwrap(), + Bytes::from_static(b"123") + ); + assert!(matches!( + body.next().await.unwrap(), + Err(HttpError::TruncatedBody { + expected: 5, + received: 3 + }) + )); + assert!(body.next().await.is_none()); + } fn streaming_turn_from(chunks: Vec<&'static str>) -> CerebrasTurn { let body: BodyStream = Box::pin(stream::iter( @@ -425,6 +1142,12 @@ mod tests { translator: EventTranslator::new(), pending: VecDeque::new(), eof: false, + visible_output: false, + detect_truncation: false, + idle_timeout: None, + deadline: None, + integrity: TruncatedStreamDetector::from_headers(&agentkit_http::HeaderMap::new()), + replay: None, })), } } @@ -447,6 +1170,80 @@ mod tests { assert!(saw_finished); } + #[tokio::test(flavor = "current_thread")] + async fn streaming_turn_never_replays_after_visible_output() { + let body: BodyStream = Box::pin(stream::iter(vec![ + Ok(Bytes::from_static( + b"data: {\"id\":\"m\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}]}\n\n", + )), + Err(HttpError::Other("stream failed".into())), + ])); + let replay_count = Arc::new(AtomicUsize::new(0)); + let count = replay_count.clone(); + let replay: StreamReplay = Box::new(move || { + count.fetch_add(1, Ordering::SeqCst); + Box::pin(async { Err(LoopError::Provider("unexpected replay".into())) }) + }); + let mut turn = CerebrasTurn { + inner: TurnInner::Streaming(Box::new(StreamingState { + body, + decoder: SseDecoder::new(), + translator: EventTranslator::new(), + pending: VecDeque::new(), + eof: false, + visible_output: false, + detect_truncation: false, + idle_timeout: None, + deadline: None, + integrity: TruncatedStreamDetector::from_headers(&agentkit_http::HeaderMap::new()), + replay: Some(replay), + })), + }; + + let mut visible_events = 0; + loop { + match turn.next_event(None).await { + Ok(Some(_)) => visible_events += 1, + Err(_) => break, + Ok(None) => panic!("stream ended without surfacing the body error"), + } + } + assert!(visible_events > 0); + assert_eq!(replay_count.load(Ordering::SeqCst), 0); + } + + #[tokio::test(flavor = "current_thread")] + async fn streaming_turn_does_not_replay_non_transient_body_errors() { + let body: BodyStream = Box::pin(stream::iter([Err(HttpError::Other( + "protocol failure".into(), + ))])); + let replay_count = Arc::new(AtomicUsize::new(0)); + let count = replay_count.clone(); + let replay: StreamReplay = Box::new(move || { + count.fetch_add(1, Ordering::SeqCst); + Box::pin(async { unreachable!("non-transient errors must not replay") }) + }); + let mut turn = CerebrasTurn { + inner: TurnInner::Streaming(Box::new(StreamingState { + body, + decoder: SseDecoder::new(), + translator: EventTranslator::new(), + pending: VecDeque::new(), + eof: false, + visible_output: false, + detect_truncation: true, + idle_timeout: None, + deadline: None, + integrity: TruncatedStreamDetector::default(), + replay: Some(replay), + })), + }; + + let error = turn.next_event(None).await.unwrap_err(); + assert!(error.to_string().contains("protocol failure")); + assert_eq!(replay_count.load(Ordering::SeqCst), 0); + } + #[tokio::test(flavor = "current_thread")] async fn streaming_turn_respects_pre_fired_cancellation() { let chunks = vec!["data: {\"id\":\"m\",\"choices\":[]}\n\n"]; diff --git a/crates/agentkit-provider-cerebras/src/models.rs b/crates/agentkit-provider-cerebras/src/models.rs index d945b82..ee136bb 100644 --- a/crates/agentkit-provider-cerebras/src/models.rs +++ b/crates/agentkit-provider-cerebras/src/models.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use agentkit_http::Http; use serde::{Deserialize, Serialize}; +use crate::RequestExecutor; use crate::config::CerebrasConfig; use crate::error::CerebrasError; @@ -27,12 +28,12 @@ impl<'a> ModelsClient<'a> { /// `GET /v1/models` — list every model visible to the API key. pub async fn list(&self) -> Result, CerebrasError> { let url = format!("{}/models", self.config.base_url); - let response = self - .http - .get(&url) - .bearer_auth(&self.config.api_key) - .send() - .await?; + let executor = RequestExecutor::new( + self.http, + self.config.authentication.clone(), + self.config.resilience.clone(), + ); + let response = executor.execute_buffered(|| self.http.get(&url)).await?; if !response.status().is_success() { let status = response.status().as_u16(); let body = response.text().await.unwrap_or_default(); @@ -45,12 +46,12 @@ impl<'a> ModelsClient<'a> { /// `GET /v1/models/{id}` — fetch a single model object. pub async fn retrieve(&self, id: &str) -> Result { let url = format!("{}/models/{id}", self.config.base_url); - let response = self - .http - .get(&url) - .bearer_auth(&self.config.api_key) - .send() - .await?; + let executor = RequestExecutor::new( + self.http, + self.config.authentication.clone(), + self.config.resilience.clone(), + ); + let response = executor.execute_buffered(|| self.http.get(&url)).await?; if !response.status().is_success() { let status = response.status().as_u16(); let body = response.text().await.unwrap_or_default(); diff --git a/crates/agentkit-provider-groq/README.md b/crates/agentkit-provider-groq/README.md index 435c71c..3fdeef1 100644 --- a/crates/agentkit-provider-groq/README.md +++ b/crates/agentkit-provider-groq/README.md @@ -18,7 +18,7 @@ buffered response path. ## Configuration -Create a config with `GroqConfig::new(api_key, model)` and chain `.with_*()` builders for optional parameters. Alternatively, `GroqConfig::from_env()` reads from environment variables: +Create a config with `GroqConfig::new(authentication, model)` and chain `.with_*()` builders for optional parameters. Alternatively, `GroqConfig::from_env()` reads from environment variables: | Variable | Required | Default | | --------------- | -------- | ------------------------------------------------- | @@ -26,6 +26,18 @@ Create a config with `GroqConfig::new(api_key, model)` and chain `.with_*()` bui | `GROQ_MODEL` | no | `llama-3.1-8b-instant` | | `GROQ_BASE_URL` | no | `https://api.groq.com/openai/v1/chat/completions` | +## Authentication and resilience + +`GroqConfig` stores credentials as a first-class +`agentkit_http::Authentication`. A bare string passed to `GroqConfig::new` or +`.with_authentication(...)` is shorthand for bearer authentication. Use +`.with_authentication_provider(...)` for a custom refresh-capable +`AuthenticationProvider`. + +Resilience is opt-in: `resilience` is an `Option` that +defaults to `None`. Calling `.with_resilience(...)` enables retries and +timeouts; leaving it as `None` preserves the existing single-attempt behavior. + ## Examples ### Minimal chat agent diff --git a/crates/agentkit-provider-groq/src/lib.rs b/crates/agentkit-provider-groq/src/lib.rs index 1ffb0c1..5fd5664 100644 --- a/crates/agentkit-provider-groq/src/lib.rs +++ b/crates/agentkit-provider-groq/src/lib.rs @@ -30,6 +30,7 @@ use agentkit_adapter_completions::{ CompletionsAdapter, CompletionsError, CompletionsProvider, CompletionsSession, CompletionsTurn, }; +use agentkit_http::{Authentication, AuthenticationProvider, ResilienceConfig}; use agentkit_loop::{LoopError, ModelAdapter, SessionConfig}; use async_trait::async_trait; use serde::Serialize; @@ -51,14 +52,16 @@ const DEFAULT_ENDPOINT: &str = "https://api.groq.com/openai/v1/chat/completions" /// .with_temperature(0.0) /// .with_max_completion_tokens(4096); /// ``` -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct GroqConfig { - /// Groq API key (starts with `gsk_`). - pub api_key: String, + /// Authentication used for API requests. String values become bearer authentication. + pub authentication: Authentication, /// Model identifier, e.g. `"llama-3.3-70b-versatile"` or `"llama-3.1-8b-instant"`. pub model: String, /// Chat completions endpoint URL. Defaults to the Groq production URL. pub base_url: String, + /// Optional retry and timeout policy. `None` preserves single-attempt behavior. + pub resilience: Option, /// Sampling temperature (0.0 = deterministic, higher = more creative). pub temperature: Option, /// Maximum number of completion tokens the model may generate. @@ -73,13 +76,30 @@ pub struct GroqConfig { pub streaming: bool, } +impl std::fmt::Debug for GroqConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GroqConfig") + .field("authentication", &"") + .field("model", &self.model) + .field("base_url", &self.base_url) + .field("resilience", &self.resilience) + .field("temperature", &self.temperature) + .field("max_completion_tokens", &self.max_completion_tokens) + .field("top_p", &self.top_p) + .field("parallel_tool_calls", &self.parallel_tool_calls) + .field("streaming", &self.streaming) + .finish() + } +} + impl GroqConfig { - /// Creates a new configuration with the given API key and model identifier. - pub fn new(api_key: impl Into, model: impl Into) -> Self { + /// Creates a new configuration with the given authentication and model identifier. + pub fn new(authentication: impl Into, model: impl Into) -> Self { Self { - api_key: api_key.into(), + authentication: authentication.into(), model: model.into(), base_url: DEFAULT_ENDPOINT.into(), + resilience: None, temperature: None, max_completion_tokens: None, top_p: None, @@ -88,6 +108,29 @@ impl GroqConfig { } } + /// Replaces request authentication. String values become bearer authentication. + pub fn with_authentication(mut self, authentication: impl Into) -> Self { + self.authentication = authentication.into(); + self + } + + /// Compatibility builder for bearer API-key authentication. + pub fn with_api_key(self, api_key: impl Into) -> Self { + self.with_authentication(api_key) + } + + /// Uses a custom refresh-capable authentication provider. + pub fn with_authentication_provider(mut self, provider: P) -> Self { + self.authentication = Authentication::new(provider); + self + } + + /// Enables request retries and timeouts. + pub fn with_resilience(mut self, resilience: ResilienceConfig) -> Self { + self.resilience = Some(resilience); + self + } + /// Overrides the default chat completions endpoint URL. pub fn with_base_url(mut self, url: impl Into) -> Self { self.base_url = url.into(); @@ -163,7 +206,8 @@ pub struct GroqRequestConfig { /// The Groq provider, implementing [`CompletionsProvider`]. #[derive(Clone, Debug)] pub struct GroqProvider { - api_key: String, + authentication: Authentication, + resilience: Option, base_url: String, streaming: bool, request_config: GroqRequestConfig, @@ -172,7 +216,8 @@ pub struct GroqProvider { impl From for GroqProvider { fn from(config: GroqConfig) -> Self { Self { - api_key: config.api_key, + authentication: config.authentication, + resilience: config.resilience, base_url: config.base_url, streaming: config.streaming, request_config: GroqRequestConfig { @@ -203,12 +248,20 @@ impl CompletionsProvider for GroqProvider { &self, builder: agentkit_http::HttpRequestBuilder, ) -> agentkit_http::HttpRequestBuilder { - builder.bearer_auth(&self.api_key).header( + builder.header( "User-Agent", concat!("agentkit-provider-groq/", env!("CARGO_PKG_VERSION")), ) } + fn authentication(&self) -> Option { + Some(self.authentication.clone()) + } + + fn resilience_config(&self) -> Option { + self.resilience.clone() + } + fn streaming(&self) -> bool { self.streaming } @@ -246,6 +299,21 @@ impl GroqAdapter { let provider = GroqProvider::from(config); Ok(Self(CompletionsAdapter::new(provider)?)) } + + /// Overrides the default API-key authentication. + pub fn with_authentication(self, authentication: impl Into) -> Self { + Self(self.0.with_authentication(authentication)) + } + + /// Overrides authentication with a custom refresh-capable provider. + pub fn with_authentication_provider(self, provider: P) -> Self { + self.with_authentication(Authentication::new(provider)) + } + + /// Enables retry and timeout behavior. + pub fn with_resilience(self, resilience: ResilienceConfig) -> Self { + Self(self.0.with_resilience(resilience)) + } } #[async_trait] @@ -268,3 +336,25 @@ pub enum GroqError { #[error(transparent)] Completions(#[from] CompletionsError), } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn config_authentication_and_resilience_reach_provider() { + let provider = GroqProvider::from( + GroqConfig::new("secret", "model").with_resilience(ResilienceConfig::default()), + ); + assert!(provider.authentication().is_some()); + assert!(provider.resilience_config().is_some()); + } + + #[test] + fn config_debug_redacts_authentication() { + let debug = format!("{:?}", GroqConfig::new("groq-secret", "debug-model")); + assert!(!debug.contains("groq-secret")); + assert!(debug.contains("")); + assert!(debug.contains("debug-model")); + } +} diff --git a/crates/agentkit-provider-mistral/README.md b/crates/agentkit-provider-mistral/README.md index 678c6db..db4475b 100644 --- a/crates/agentkit-provider-mistral/README.md +++ b/crates/agentkit-provider-mistral/README.md @@ -20,7 +20,7 @@ used by most other OpenAI-compatible APIs. ## Configuration -Create a config with `MistralConfig::new(api_key, model)` and chain `.with_*()` builders for optional parameters. Alternatively, `MistralConfig::from_env()` reads from environment variables: +Create a config with `MistralConfig::new(authentication, model)` and chain `.with_*()` builders for optional parameters. Alternatively, `MistralConfig::from_env()` reads from environment variables: | Variable | Required | Default | | ------------------ | -------- | -------------------------------------------- | @@ -28,6 +28,18 @@ Create a config with `MistralConfig::new(api_key, model)` and chain `.with_*()` | `MISTRAL_MODEL` | no | `mistral-small-latest` | | `MISTRAL_BASE_URL` | no | `https://api.mistral.ai/v1/chat/completions` | +## Authentication and resilience + +`MistralConfig` stores credentials as a first-class +`agentkit_http::Authentication`. A bare string passed to `MistralConfig::new` +or `.with_authentication(...)` is shorthand for bearer authentication. Use +`.with_authentication_provider(...)` for a custom refresh-capable +`AuthenticationProvider`. + +Resilience is opt-in: `resilience` is an `Option` that +defaults to `None`. Calling `.with_resilience(...)` enables retries and +timeouts; leaving it as `None` preserves the existing single-attempt behavior. + ## Examples ### Minimal chat agent diff --git a/crates/agentkit-provider-mistral/src/lib.rs b/crates/agentkit-provider-mistral/src/lib.rs index 01b8ff7..4ceee7f 100644 --- a/crates/agentkit-provider-mistral/src/lib.rs +++ b/crates/agentkit-provider-mistral/src/lib.rs @@ -32,6 +32,7 @@ use agentkit_adapter_completions::{ CompletionsAdapter, CompletionsError, CompletionsProvider, CompletionsSession, CompletionsTurn, }; +use agentkit_http::{Authentication, AuthenticationProvider, ResilienceConfig}; use agentkit_loop::{LoopError, ModelAdapter, SessionConfig}; use async_trait::async_trait; use serde::Serialize; @@ -53,14 +54,16 @@ const DEFAULT_ENDPOINT: &str = "https://api.mistral.ai/v1/chat/completions"; /// .with_temperature(0.0) /// .with_max_tokens(4096); /// ``` -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct MistralConfig { - /// Mistral API key. - pub api_key: String, + /// Authentication used for API requests. String values become bearer authentication. + pub authentication: Authentication, /// Model identifier, e.g. `"mistral-large-latest"` or `"mistral-small-latest"`. pub model: String, /// Chat completions endpoint URL. Defaults to the Mistral production URL. pub base_url: String, + /// Optional retry and timeout policy. `None` preserves single-attempt behavior. + pub resilience: Option, /// Sampling temperature (0.0 = deterministic, higher = more creative). pub temperature: Option, /// Maximum number of tokens the model may generate. Mistral uses `max_tokens` @@ -83,13 +86,31 @@ pub struct MistralConfig { pub strict_alternating_roles: bool, } +impl std::fmt::Debug for MistralConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("MistralConfig") + .field("authentication", &"") + .field("model", &self.model) + .field("base_url", &self.base_url) + .field("resilience", &self.resilience) + .field("temperature", &self.temperature) + .field("max_tokens", &self.max_tokens) + .field("top_p", &self.top_p) + .field("parallel_tool_calls", &self.parallel_tool_calls) + .field("streaming", &self.streaming) + .field("strict_alternating_roles", &self.strict_alternating_roles) + .finish() + } +} + impl MistralConfig { - /// Creates a new configuration with the given API key and model identifier. - pub fn new(api_key: impl Into, model: impl Into) -> Self { + /// Creates a new configuration with the given authentication and model identifier. + pub fn new(authentication: impl Into, model: impl Into) -> Self { Self { - api_key: api_key.into(), + authentication: authentication.into(), model: model.into(), base_url: DEFAULT_ENDPOINT.into(), + resilience: None, temperature: None, max_tokens: None, top_p: None, @@ -99,6 +120,29 @@ impl MistralConfig { } } + /// Replaces request authentication. String values become bearer authentication. + pub fn with_authentication(mut self, authentication: impl Into) -> Self { + self.authentication = authentication.into(); + self + } + + /// Compatibility builder for bearer API-key authentication. + pub fn with_api_key(self, api_key: impl Into) -> Self { + self.with_authentication(api_key) + } + + /// Uses a custom refresh-capable authentication provider. + pub fn with_authentication_provider(mut self, provider: P) -> Self { + self.authentication = Authentication::new(provider); + self + } + + /// Enables request retries and timeouts. + pub fn with_resilience(mut self, resilience: ResilienceConfig) -> Self { + self.resilience = Some(resilience); + self + } + /// Overrides the default chat completions endpoint URL. pub fn with_base_url(mut self, url: impl Into) -> Self { self.base_url = url.into(); @@ -185,7 +229,8 @@ pub struct MistralRequestConfig { /// The Mistral provider, implementing [`CompletionsProvider`]. #[derive(Clone, Debug)] pub struct MistralProvider { - api_key: String, + authentication: Authentication, + resilience: Option, base_url: String, streaming: bool, strict_alternating_roles: bool, @@ -195,7 +240,8 @@ pub struct MistralProvider { impl From for MistralProvider { fn from(config: MistralConfig) -> Self { Self { - api_key: config.api_key, + authentication: config.authentication, + resilience: config.resilience, base_url: config.base_url, streaming: config.streaming, strict_alternating_roles: config.strict_alternating_roles, @@ -227,12 +273,20 @@ impl CompletionsProvider for MistralProvider { &self, builder: agentkit_http::HttpRequestBuilder, ) -> agentkit_http::HttpRequestBuilder { - builder.bearer_auth(&self.api_key).header( + builder.header( "User-Agent", concat!("agentkit-provider-mistral/", env!("CARGO_PKG_VERSION")), ) } + fn authentication(&self) -> Option { + Some(self.authentication.clone()) + } + + fn resilience_config(&self) -> Option { + self.resilience.clone() + } + fn streaming(&self) -> bool { self.streaming } @@ -274,6 +328,21 @@ impl MistralAdapter { let provider = MistralProvider::from(config); Ok(Self(CompletionsAdapter::new(provider)?)) } + + /// Overrides the default API-key authentication. + pub fn with_authentication(self, authentication: impl Into) -> Self { + Self(self.0.with_authentication(authentication)) + } + + /// Overrides authentication with a custom refresh-capable provider. + pub fn with_authentication_provider(self, provider: P) -> Self { + self.with_authentication(Authentication::new(provider)) + } + + /// Enables retry and timeout behavior. + pub fn with_resilience(self, resilience: ResilienceConfig) -> Self { + Self(self.0.with_resilience(resilience)) + } } #[async_trait] @@ -296,3 +365,25 @@ pub enum MistralError { #[error(transparent)] Completions(#[from] CompletionsError), } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn config_authentication_and_resilience_reach_provider() { + let provider = MistralProvider::from( + MistralConfig::new("secret", "model").with_resilience(ResilienceConfig::default()), + ); + assert!(provider.authentication().is_some()); + assert!(provider.resilience_config().is_some()); + } + + #[test] + fn config_debug_redacts_authentication() { + let debug = format!("{:?}", MistralConfig::new("mistral-secret", "debug-model")); + assert!(!debug.contains("mistral-secret")); + assert!(debug.contains("")); + assert!(debug.contains("debug-model")); + } +} diff --git a/crates/agentkit-provider-ollama/README.md b/crates/agentkit-provider-ollama/README.md index 871703e..26af930 100644 --- a/crates/agentkit-provider-ollama/README.md +++ b/crates/agentkit-provider-ollama/README.md @@ -15,9 +15,10 @@ OpenAI-compatible chat completions endpoint. It handles request translation and response normalization for Ollama-backed sessions. Streaming is enabled by default; use `.with_streaming(false)` to force the buffered response path. -No API key is required — Ollama runs locally and does not authenticate requests +Authentication is optional because Ollama does not authenticate local requests by default. You need a running Ollama server (e.g. `ollama serve`) with your -desired model pulled (e.g. `ollama pull llama3.1:8b`). +desired model pulled (e.g. `ollama pull llama3.1:8b`). Protected remote +endpoints can opt into authentication as described below. ## Configuration @@ -28,6 +29,18 @@ Create a config with `OllamaConfig::new(model)` and chain `.with_*()` builders f | `OLLAMA_MODEL` | yes | -- | | `OLLAMA_BASE_URL` | no | `http://localhost:11434/v1/chat/completions` | +## Authentication and resilience + +`OllamaConfig::new` leaves its first-class +`Option` as `None`. For a protected endpoint, a +bare string passed to `.with_authentication(...)` is shorthand for bearer +authentication. Use `.with_authentication_provider(...)` for a custom +refresh-capable `AuthenticationProvider`. + +Resilience is also opt-in: `resilience` is an `Option` that +defaults to `None`. Calling `.with_resilience(...)` enables retries and +timeouts; leaving it as `None` preserves the existing single-attempt behavior. + ## Examples ### Minimal chat agent diff --git a/crates/agentkit-provider-ollama/src/lib.rs b/crates/agentkit-provider-ollama/src/lib.rs index c1bc05f..e845e03 100644 --- a/crates/agentkit-provider-ollama/src/lib.rs +++ b/crates/agentkit-provider-ollama/src/lib.rs @@ -34,6 +34,7 @@ use agentkit_adapter_completions::{ CompletionsAdapter, CompletionsError, CompletionsProvider, CompletionsSession, CompletionsTurn, }; +use agentkit_http::{Authentication, AuthenticationProvider, ResilienceConfig}; use agentkit_loop::{LoopError, ModelAdapter, SessionConfig}; use async_trait::async_trait; use serde::Serialize; @@ -62,6 +63,10 @@ pub struct OllamaConfig { pub model: String, /// Chat completions endpoint URL. Defaults to `http://localhost:11434/v1/chat/completions`. pub base_url: String, + /// Optional authentication for protected endpoints. Strings become bearer authentication. + pub authentication: Option, + /// Optional retry and timeout policy. `None` preserves single-attempt behavior. + pub resilience: Option, /// Sampling temperature (0.0 = deterministic, higher = more creative). pub temperature: Option, /// Maximum number of tokens to generate (Ollama's equivalent of `max_completion_tokens`). @@ -88,6 +93,8 @@ impl OllamaConfig { Self { model: model.into(), base_url: DEFAULT_ENDPOINT.into(), + authentication: None, + resilience: None, temperature: None, num_predict: None, top_k: None, @@ -98,6 +105,24 @@ impl OllamaConfig { } } + /// Sets authentication for a protected Ollama endpoint. + pub fn with_authentication(mut self, authentication: impl Into) -> Self { + self.authentication = Some(authentication.into()); + self + } + + /// Uses a custom refresh-capable authentication provider. + pub fn with_authentication_provider(mut self, provider: P) -> Self { + self.authentication = Some(Authentication::new(provider)); + self + } + + /// Enables request retries and timeouts. + pub fn with_resilience(mut self, resilience: ResilienceConfig) -> Self { + self.resilience = Some(resilience); + self + } + /// Overrides the default chat completions endpoint URL. pub fn with_base_url(mut self, url: impl Into) -> Self { self.base_url = url.into(); @@ -190,6 +215,8 @@ pub struct OllamaRequestConfig { #[derive(Clone, Debug)] pub struct OllamaProvider { base_url: String, + authentication: Option, + resilience: Option, streaming: bool, strict_alternating_roles: bool, request_config: OllamaRequestConfig, @@ -199,6 +226,8 @@ impl From for OllamaProvider { fn from(config: OllamaConfig) -> Self { Self { base_url: config.base_url, + authentication: config.authentication, + resilience: config.resilience, streaming: config.streaming, strict_alternating_roles: config.strict_alternating_roles, request_config: OllamaRequestConfig { @@ -236,6 +265,14 @@ impl CompletionsProvider for OllamaProvider { ) } + fn authentication(&self) -> Option { + self.authentication.clone() + } + + fn resilience_config(&self) -> Option { + self.resilience.clone() + } + fn streaming(&self) -> bool { self.streaming } @@ -279,6 +316,21 @@ impl OllamaAdapter { let provider = OllamaProvider::from(config); Ok(Self(CompletionsAdapter::new(provider)?)) } + + /// Adds optional authentication for a protected Ollama endpoint. + pub fn with_authentication(self, authentication: impl Into) -> Self { + Self(self.0.with_authentication(authentication)) + } + + /// Adds refresh-capable authentication for a protected Ollama endpoint. + pub fn with_authentication_provider(self, provider: P) -> Self { + self.with_authentication(Authentication::new(provider)) + } + + /// Enables retry and timeout behavior. + pub fn with_resilience(self, resilience: ResilienceConfig) -> Self { + Self(self.0.with_resilience(resilience)) + } } #[async_trait] @@ -301,3 +353,23 @@ pub enum OllamaError { #[error(transparent)] Completions(#[from] CompletionsError), } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn authentication_and_resilience_are_optional_and_propagate() { + let default_provider = OllamaProvider::from(OllamaConfig::new("model")); + assert!(default_provider.authentication().is_none()); + assert!(default_provider.resilience_config().is_none()); + + let config = OllamaConfig::new("model") + .with_authentication("ollama-secret") + .with_resilience(ResilienceConfig::default()); + assert!(!format!("{config:?}").contains("ollama-secret")); + let provider = OllamaProvider::from(config); + assert!(provider.authentication().is_some()); + assert!(provider.resilience_config().is_some()); + } +} diff --git a/crates/agentkit-provider-openai/Cargo.toml b/crates/agentkit-provider-openai/Cargo.toml index 8d52170..287e3e0 100644 --- a/crates/agentkit-provider-openai/Cargo.toml +++ b/crates/agentkit-provider-openai/Cargo.toml @@ -15,6 +15,13 @@ agentkit-core = { version = "0.10.5", path = "../agentkit-core" } agentkit-http = { version = "0.10.5", path = "../agentkit-http" } agentkit-loop = { version = "0.10.5", path = "../agentkit-loop" } async-trait.workspace = true +base64.workspace = true +futures-util.workspace = true +reqwest.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true +zeroize = "1.8.2" + +[dev-dependencies] +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/crates/agentkit-provider-openai/README.md b/crates/agentkit-provider-openai/README.md index 787ec52..a9a1782 100644 --- a/crates/agentkit-provider-openai/README.md +++ b/crates/agentkit-provider-openai/README.md @@ -9,11 +9,16 @@ OpenAI model adapter for the agentkit agent loop. -This crate provides `OpenAIAdapter` and `OpenAIConfig` for connecting the -agent loop to the [OpenAI](https://platform.openai.com) chat completions API. -It handles request translation, response normalization, usage reporting, and -prompt cache integration for OpenAI-backed sessions. Streaming is enabled by -default; use `.with_streaming(false)` to force the buffered response path. +This crate provides two OpenAI adapters: + +- `OpenAIChatCompletionsAdapter` for `/v1/chat/completions`. The historical + `OpenAIAdapter` name remains a compatibility alias. +- `OpenAIResponsesAdapter` for the public `/v1/responses` API and configurable + private Responses-compatible deployments. + +Both adapters translate AgentKit transcripts, tools, usage, and finish reasons. +Chat-completions streaming is enabled by default; use `.with_streaming(false)` +to force its buffered response path. Applications that want an OpenAI-powered agent will usually use this crate through the umbrella `agentkit` crate's `provider-openai` feature, or depend on @@ -21,7 +26,7 @@ it directly when assembling a smaller runtime. ## Configuration -Create a config with `OpenAIConfig::new(api_key, model)` and chain `.with_*()` builders for optional parameters. Alternatively, `OpenAIConfig::from_env()` reads from environment variables: +Create a config with `OpenAIConfig::new(authentication, model)` and chain `.with_*()` builders for optional parameters. Alternatively, `OpenAIConfig::from_env()` reads from environment variables: | Variable | Required | Default | | ----------------- | -------- | -------------------------------------------- | @@ -29,6 +34,107 @@ Create a config with `OpenAIConfig::new(api_key, model)` and chain `.with_*()` b | `OPENAI_MODEL` | no | `gpt-4o` | | `OPENAI_BASE_URL` | no | `https://api.openai.com/v1/chat/completions` | +## Authentication and resilience + +`OpenAIConfig` and `OpenAIResponsesConfig` store credentials as first-class +`agentkit_http::Authentication` values. A bare string passed to either `new` +constructor or `.with_authentication(...)` is shorthand for bearer +authentication. Both configs and adapters provide +`.with_authentication_provider(...)` for a custom refresh-capable +`AuthenticationProvider`; the provider receives the opaque prior authentication +attempt during the single reactive 401 refresh. The standard static bearer and +header constructors attach an ephemeral, non-secret binding automatically. +Custom providers that need Responses continuation replay must attach a stable, +non-secret credential identity or generation with +`AuthenticationAttempt::with_binding`. If a reactive refresh changes that +binding, the adapter fails rather than sending the already encoded, binding-bound +body. Resilience is stored as `Option` and defaults to +`None`. Calling `.with_resilience(...)` opts into retries and timeouts; leaving +it as `None` preserves the existing single-attempt behavior for transient +transport/status/stream failures (the one permitted 401 refresh remains part +of authentication). + +Responses retries reuse one clone-cheap serialized request body and a +body-bound idempotency key. Events stream as soon as they are decoded. A failed +attempt is replayed automatically only before its first event becomes visible. +After visible output, replay is disabled unless the upstream consumer explicitly +calls `SessionConfig::with_response_attempt_supersession()`. On a visible-attempt +retry, the adapter emits `ModelTurnEvent::ResponseAttemptSuperseded` after the +failed attempt's events and before replacement output; the loop forwards it as +`AgentEvent::ResponseAttemptSuperseded`. Consumers must discard every delta, tool +call, usage update, and reconstruction state from that preceding attempt. +Cancellation, the logical retry deadline, stream-idle timeout, absolute +per-attempt deadline (including stream reads), auth/refresh, and backoff remain +bounded. Responses default to a 32 MiB serialized request limit, 16 MiB per +attempt, and 64 MiB aggregate wire limit across retries. Requests and responses +also default to at most 100,000 items/events and 8 MiB per text, reasoning, +ciphertext, or media field. Override these with `OpenAIResponsesLimits` and +`.with_limits(...)`; for example, a host can set `max_items` to `10_000` and +`max_text_bytes` to `1024 * 1024` while retaining the default aggregate bounds. +Limits must be non-zero; the per-field bound must fit both request and attempt +bounds, and the per-attempt bound must fit the aggregate wire bound. + +## Responses API + +`OpenAIResponsesConfig::new(authentication, model)` deliberately matches +`OpenAIConfig::new(authentication, model)`. The profile-specific constructors use +`public(model, authentication)` and `chatgpt_private(model, authentication)`; use +those named forms when selecting a profile so the different argument order is +explicit. + +```rust,no_run +use agentkit_provider_openai::{OpenAIResponsesAdapter, OpenAIResponsesConfig}; + +# fn main() -> Result<(), Box> { +let config = OpenAIResponsesConfig::new("sk-...", "gpt-5"); +let adapter = OpenAIResponsesAdapter::new(config)?; +# let _ = adapter; +# Ok(()) +# } +``` + +For the private ChatGPT Codex-shaped endpoint, inject authentication explicitly: + +```rust,no_run +use agentkit_http::Authentication; +use agentkit_provider_openai::{OpenAIResponsesAdapter, OpenAIResponsesConfig}; + +# fn main() -> Result<(), Box> { +let authentication = Authentication::bearer("short-lived-access-token"); +let config = OpenAIResponsesConfig::chatgpt_private("gpt-5-codex", authentication) + .with_endpoint("https://chatgpt.com/backend-api/codex/responses"); +let adapter = OpenAIResponsesAdapter::new(config)?; +# let _ = adapter; +# Ok(()) +# } +``` + +`with_headers` supplies ordinary non-authentication headers. +`with_user_agent` and `with_originator` provide explicit request attribution, +and `with_request_policy` can override public/private request-field differences. +The synchronous `encode_request` helper cannot resolve an authentication binding +and rejects bound continuation metadata; submit through `OpenAIResponsesAdapter` +to replay adapter-emitted function-call, reasoning, or generated-image state. +The public and private profiles request encrypted reasoning continuation data by +default. Context content is sent unchanged: the public profile keeps system and +Context items as system messages, while the private profile downgrades both to +developer messages, +defaults `parallel_tool_calls` to `true`, +omits unsupported `max_output_tokens`, sends `originator`/`session-id`, and +replays validated `x-codex-turn-state` only within one logical turn and its +retries. HTTP turn state is accepted only from a successful SSE response; the +equivalent `response.metadata` headers can update retry context, and all header, +metadata, and retry values must agree. It does not perform credential discovery +or model catalog lookups. + +Continuation metadata is versioned and bound to the authentication binding, +model, session, provider item ID, and item kind. Durable encrypted reasoning, +function-call, and generated-image continuation metadata is emitted and replayed +only when authentication supplies a binding. Valid metadata for another binding +is omitted safely; malformed metadata is a protocol error. The private profile accepts image/audio transcript inputs and +media-bearing tool outputs. The public profile continues to reject the private +audio shape. + ## Examples ### Minimal chat agent diff --git a/crates/agentkit-provider-openai/src/lib.rs b/crates/agentkit-provider-openai/src/lib.rs index 95fd48a..2847e83 100644 --- a/crates/agentkit-provider-openai/src/lib.rs +++ b/crates/agentkit-provider-openai/src/lib.rs @@ -26,10 +26,21 @@ //! } //! ``` +mod responses; + +pub use responses::{ + OpenAIResponsesAdapter, OpenAIResponsesConfig, OpenAIResponsesError, OpenAIResponsesLimits, + OpenAIResponsesProfile, OpenAIResponsesRequestPolicy, OpenAIResponsesSession, + OpenAIResponsesTurn, +}; + +use std::fmt; + use agentkit_adapter_completions::{ CompletionsAdapter, CompletionsError, CompletionsProvider, CompletionsSession, CompletionsTurn, }; use agentkit_core::{MetadataMap, Usage}; +use agentkit_http::{Authentication, AuthenticationProvider, ResilienceConfig}; use agentkit_loop::{ LoopError, ModelAdapter, PromptCacheMode, PromptCacheRetention, PromptCacheStrategy, SessionConfig, TurnRequest, @@ -56,10 +67,10 @@ const DEFAULT_ENDPOINT: &str = "https://api.openai.com/v1/chat/completions"; /// .with_temperature(0.0) /// .with_max_completion_tokens(4096); /// ``` -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct OpenAIConfig { - /// OpenAI API key (starts with `sk-`). - pub api_key: String, + /// Authentication applied to each request. Bare strings use bearer authentication. + pub authentication: Authentication, /// Model identifier, e.g. `"gpt-4o"` or `"gpt-4o-mini"`. pub model: String, /// Chat completions endpoint URL. Defaults to the OpenAI production URL. @@ -80,13 +91,36 @@ pub struct OpenAIConfig { pub parallel_tool_calls: Option, /// Request SSE streaming responses. Defaults to `true`. pub streaming: bool, + /// Optional retry and timeout policy. `None` preserves single-attempt behavior. + pub resilience: Option, +} + +impl fmt::Debug for OpenAIConfig { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("OpenAIConfig") + .field("authentication", &"") + .field("model", &self.model) + .field("base_url", &self.base_url) + .field("temperature", &self.temperature) + .field("max_completion_tokens", &self.max_completion_tokens) + .field("top_p", &self.top_p) + .field("frequency_penalty", &self.frequency_penalty) + .field("presence_penalty", &self.presence_penalty) + .field("parallel_tool_calls", &self.parallel_tool_calls) + .field("streaming", &self.streaming) + .field("resilience", &self.resilience) + .finish() + } } impl OpenAIConfig { - /// Creates a new configuration with the given API key and model identifier. - pub fn new(api_key: impl Into, model: impl Into) -> Self { + /// Creates a new configuration with the given authentication and model identifier. + /// + /// Bare strings are treated as bearer tokens. + pub fn new(authentication: impl Into, model: impl Into) -> Self { Self { - api_key: api_key.into(), + authentication: authentication.into(), model: model.into(), base_url: DEFAULT_ENDPOINT.into(), temperature: None, @@ -96,9 +130,27 @@ impl OpenAIConfig { presence_penalty: None, parallel_tool_calls: None, streaming: true, + resilience: None, } } + /// Replaces the configured authentication. + pub fn with_authentication(mut self, authentication: impl Into) -> Self { + self.authentication = authentication.into(); + self + } + + /// Uses a custom refresh-capable authentication provider. + pub fn with_authentication_provider(self, provider: P) -> Self { + self.with_authentication(Authentication::new(provider)) + } + + /// Enables request and pre-visible-output retries and timeouts. + pub fn with_resilience(mut self, resilience: ResilienceConfig) -> Self { + self.resilience = Some(resilience); + self + } + /// Overrides the default chat completions endpoint URL. pub fn with_base_url(mut self, url: impl Into) -> Self { self.base_url = url.into(); @@ -192,7 +244,8 @@ pub struct OpenAIRequestConfig { /// The OpenAI provider, implementing [`CompletionsProvider`]. #[derive(Clone, Debug)] pub struct OpenAIProvider { - api_key: String, + authentication: Authentication, + resilience: Option, base_url: String, streaming: bool, request_config: OpenAIRequestConfig, @@ -201,7 +254,8 @@ pub struct OpenAIProvider { impl From for OpenAIProvider { fn from(config: OpenAIConfig) -> Self { Self { - api_key: config.api_key, + authentication: config.authentication, + resilience: config.resilience, base_url: config.base_url, streaming: config.streaming, request_config: OpenAIRequestConfig { @@ -217,6 +271,13 @@ impl From for OpenAIProvider { } } +pub(crate) fn prompt_cache_retention_value(retention: PromptCacheRetention) -> &'static str { + match retention { + PromptCacheRetention::Default | PromptCacheRetention::Short => "in_memory", + PromptCacheRetention::Extended => "24h", + } +} + impl CompletionsProvider for OpenAIProvider { type Config = OpenAIRequestConfig; @@ -234,12 +295,20 @@ impl CompletionsProvider for OpenAIProvider { &self, builder: agentkit_http::HttpRequestBuilder, ) -> agentkit_http::HttpRequestBuilder { - builder.bearer_auth(&self.api_key).header( + builder.header( "User-Agent", concat!("agentkit-provider-openai/", env!("CARGO_PKG_VERSION")), ) } + fn authentication(&self) -> Option { + Some(self.authentication.clone()) + } + + fn resilience_config(&self) -> Option { + self.resilience.clone() + } + fn streaming(&self) -> bool { self.streaming } @@ -270,11 +339,10 @@ impl CompletionsProvider for OpenAIProvider { } if let Some(retention) = cache.retention { - let value = match retention { - PromptCacheRetention::Default | PromptCacheRetention::Short => "in_memory", - PromptCacheRetention::Extended => "24h", - }; - body.insert("prompt_cache_retention".into(), Value::String(value.into())); + body.insert( + "prompt_cache_retention".into(), + Value::String(prompt_cache_retention_value(retention).into()), + ); } if matches!(cache.strategy, PromptCacheStrategy::Explicit { .. }) @@ -327,25 +395,47 @@ impl CompletionsProvider for OpenAIProvider { /// # } /// ``` #[derive(Clone)] -pub struct OpenAIAdapter(CompletionsAdapter); +pub struct OpenAIChatCompletionsAdapter(CompletionsAdapter); -/// An active session with the OpenAI API. -pub type OpenAISession = CompletionsSession; +/// Compatibility alias for the chat-completions adapter. +pub type OpenAIAdapter = OpenAIChatCompletionsAdapter; -/// A completed turn from the OpenAI API. -pub type OpenAITurn = CompletionsTurn; +/// An active chat-completions session with the OpenAI API. +pub type OpenAIChatCompletionsSession = CompletionsSession; +/// Compatibility alias for an active chat-completions session. +pub type OpenAISession = OpenAIChatCompletionsSession; -impl OpenAIAdapter { +/// A chat-completions turn from the OpenAI API. +pub type OpenAIChatCompletionsTurn = CompletionsTurn; +/// Compatibility alias for a chat-completions turn. +pub type OpenAITurn = OpenAIChatCompletionsTurn; + +impl OpenAIChatCompletionsAdapter { /// Creates a new adapter from the given configuration. pub fn new(config: OpenAIConfig) -> Result { let provider = OpenAIProvider::from(config); Ok(Self(CompletionsAdapter::new(provider)?)) } + + /// Overrides the API-key authentication with an arbitrary authentication handle. + pub fn with_authentication(self, authentication: impl Into) -> Self { + Self(self.0.with_authentication(authentication)) + } + + /// Overrides authentication with a custom refresh-capable provider. + pub fn with_authentication_provider(self, provider: P) -> Self { + self.with_authentication(Authentication::new(provider)) + } + + /// Enables request and pre-visible-output retries and timeouts. + pub fn with_resilience(self, resilience: ResilienceConfig) -> Self { + Self(self.0.with_resilience(resilience)) + } } #[async_trait] -impl ModelAdapter for OpenAIAdapter { - type Session = OpenAISession; +impl ModelAdapter for OpenAIChatCompletionsAdapter { + type Session = OpenAIChatCompletionsSession; async fn start_session(&self, config: SessionConfig) -> Result { self.0.start_session(config).await @@ -381,6 +471,25 @@ mod tests { } } + #[test] + fn openai_config_debug_redacts_api_key() { + let debug = format!("{:?}", OpenAIConfig::new("super-secret", "gpt-test")); + assert!(!debug.contains("super-secret")); + assert!(debug.contains("")); + } + + #[test] + fn openai_config_resilience_reaches_provider() { + let default_provider = OpenAIProvider::from(OpenAIConfig::new("sk-test", "gpt-test")); + assert_eq!(default_provider.resilience_config(), None); + + let resilience = ResilienceConfig::no_retries(); + let provider = OpenAIProvider::from( + OpenAIConfig::new("sk-test", "gpt-test").with_resilience(resilience.clone()), + ); + assert_eq!(provider.resilience_config(), Some(resilience)); + } + #[test] fn openai_maps_automatic_cache_request() { let provider = OpenAIProvider::from(OpenAIConfig::new("sk-test", "gpt-5.1")); diff --git a/crates/agentkit-provider-openai/src/responses.rs b/crates/agentkit-provider-openai/src/responses.rs new file mode 100644 index 0000000..1f1cae1 --- /dev/null +++ b/crates/agentkit-provider-openai/src/responses.rs @@ -0,0 +1,4925 @@ +//! Transport-neutral OpenAI Responses API adapter. + +use std::collections::{BTreeMap, BTreeSet, VecDeque}; +use std::fmt; +use std::future::Future; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use base64::Engine as _; +use base64::engine::general_purpose::STANDARD; + +use agentkit_core::{ + DataRef, Delta, FinishReason, Item, ItemKind, MediaPart, MessageId, MetadataMap, Modality, + Part, PartId, PartKind, ReasoningPart, TextPart, TokenUsage, ToolCallPart, ToolOutput, + TurnCancellation, Usage, +}; +use agentkit_http::{ + Authentication, AuthenticationAttempt, AuthenticationProvider, HeaderMap, HeaderValue, Http, + HttpError, ResilienceConfig, StatusCode, TruncatedStreamDetector, is_retryable_status, + next_body_chunk, sleep, +}; +use agentkit_loop::{ + LoopError, ModelAdapter, ModelSession, ModelTurn, ModelTurnEvent, ModelTurnResult, + PromptCacheMode, PromptCacheStrategy, SessionConfig, TurnRequest, set_provider_finish_reasons, +}; +use async_trait::async_trait; +use futures_util::future::{Either, select}; +use serde_json::{Map, Value, json}; +use thiserror::Error; +use zeroize::{Zeroize, Zeroizing}; + +const PUBLIC_ENDPOINT: &str = "https://api.openai.com/v1/responses"; +const PRIVATE_ENDPOINT: &str = "https://chatgpt.com/backend-api/codex/responses"; +const DEFAULT_MAX_REQUEST_BYTES: usize = 32 * 1024 * 1024; +const DEFAULT_MAX_ATTEMPT_BYTES: usize = 16 * 1024 * 1024; +const DEFAULT_MAX_WIRE_BYTES: usize = 64 * 1024 * 1024; +const DEFAULT_MAX_TEXT_BYTES: usize = 8 * 1024 * 1024; +const DEFAULT_MAX_ITEMS: usize = 100_000; +const CONTINUATION_METADATA: &str = "openai.responses.continuation.v1"; +const CONTINUATION_SCHEMA_VERSION: u64 = 3; +const GENERATED_IMAGE_METADATA: &str = "openai.responses.generated_image.v1"; +const X_CODEX_TURN_STATE: &str = "x-codex-turn-state"; +const MAX_CACHE_KEY_BYTES: usize = 256; + +/// Selects the public Responses API or ChatGPT's private Codex-shaped profile. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OpenAIResponsesProfile { + /// `https://api.openai.com/v1/responses`, including public request fields. + Public, + /// `https://chatgpt.com/backend-api/codex/responses`, with its narrower field policy. + ChatGptPrivate, +} + +/// Bounds serialized requests, streamed responses, counts, and individual fields. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct OpenAIResponsesLimits { + /// Maximum serialized JSON request body size. + pub max_request_bytes: usize, + /// Maximum wire bytes accepted from one response attempt. + pub max_attempt_bytes: usize, + /// Maximum aggregate wire bytes accepted across all attempts for one logical turn. + pub max_wire_bytes: usize, + /// Maximum number of request items/tools, response indexes, or SSE events. + pub max_items: usize, + /// Maximum bytes accepted in any text, reasoning, ciphertext, or media field. + pub max_text_bytes: usize, +} + +impl Default for OpenAIResponsesLimits { + fn default() -> Self { + Self { + max_request_bytes: DEFAULT_MAX_REQUEST_BYTES, + max_attempt_bytes: DEFAULT_MAX_ATTEMPT_BYTES, + max_wire_bytes: DEFAULT_MAX_WIRE_BYTES, + max_items: DEFAULT_MAX_ITEMS, + max_text_bytes: DEFAULT_MAX_TEXT_BYTES, + } + } +} + +/// Controls fields whose support differs between Responses deployments. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct OpenAIResponsesRequestPolicy { + /// Downgrade system/context items to developer messages. Required by the private profile. + pub downgrade_system_to_developer: bool, + /// Send `store`; both built-in profiles default to `false`. + pub store: Option, + /// Ask the service to return encrypted reasoning for safe continuation. + pub include_encrypted_reasoning: bool, + /// Send `max_output_tokens` when configured. The private profile rejects it. + pub send_max_output_tokens: bool, +} + +impl OpenAIResponsesRequestPolicy { + pub fn public() -> Self { + Self { + downgrade_system_to_developer: false, + store: Some(false), + include_encrypted_reasoning: true, + send_max_output_tokens: true, + } + } + + pub fn chatgpt_private() -> Self { + Self { + downgrade_system_to_developer: true, + store: Some(false), + include_encrypted_reasoning: true, + send_max_output_tokens: false, + } + } +} + +/// Configuration shared by public and private Responses endpoints. +#[derive(Clone)] +pub struct OpenAIResponsesConfig { + pub model: String, + pub endpoint: String, + pub profile: OpenAIResponsesProfile, + pub headers: HeaderMap, + pub request_policy: OpenAIResponsesRequestPolicy, + pub reasoning_effort: Option, + pub max_output_tokens: Option, + pub parallel_tool_calls: Option, + authentication: Authentication, + resilience: Option, + limits: OpenAIResponsesLimits, + user_agent: Option, + originator: Option, +} + +impl fmt::Debug for OpenAIResponsesConfig { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("OpenAIResponsesConfig") + .field("model", &self.model) + .field("endpoint", &self.endpoint) + .field("profile", &self.profile) + .field("header_names", &self.headers.keys().collect::>()) + .field("request_policy", &self.request_policy) + .field("reasoning_effort", &self.reasoning_effort) + .field("max_output_tokens", &self.max_output_tokens) + .field("parallel_tool_calls", &self.parallel_tool_calls) + .field("authentication", &"") + .field("resilience", &self.resilience) + .field("limits", &self.limits) + .field("user_agent", &self.user_agent) + .field("originator", &self.originator) + .finish() + } +} + +impl OpenAIResponsesConfig { + /// Creates a public Responses configuration using the given authentication. + /// + /// The argument order is `(authentication, model)`, matching [`crate::OpenAIConfig::new`]. + /// Bare strings are treated as bearer tokens. + pub fn new(authentication: impl Into, model: impl Into) -> Self { + Self::public(model, authentication) + } + + /// Creates a public Responses configuration with injectable authentication. + /// + /// Unlike [`Self::new`], the profile-specific constructors take `(model, authentication)`. + pub fn public(model: impl Into, authentication: impl Into) -> Self { + Self { + model: model.into(), + endpoint: PUBLIC_ENDPOINT.into(), + profile: OpenAIResponsesProfile::Public, + headers: HeaderMap::new(), + request_policy: OpenAIResponsesRequestPolicy::public(), + reasoning_effort: None, + max_output_tokens: None, + parallel_tool_calls: None, + authentication: authentication.into(), + resilience: None, + limits: OpenAIResponsesLimits::default(), + user_agent: None, + originator: None, + } + } + + /// Creates a private ChatGPT Codex Responses configuration. + /// + /// Profile-specific constructors take `(model, authentication)`; [`Self::new`] takes + /// `(authentication, model)` for compatibility with [`crate::OpenAIConfig::new`]. + pub fn chatgpt_private( + model: impl Into, + authentication: impl Into, + ) -> Self { + Self { + model: model.into(), + endpoint: PRIVATE_ENDPOINT.into(), + profile: OpenAIResponsesProfile::ChatGptPrivate, + headers: HeaderMap::new(), + request_policy: OpenAIResponsesRequestPolicy::chatgpt_private(), + reasoning_effort: None, + max_output_tokens: None, + parallel_tool_calls: Some(true), + authentication: authentication.into(), + resilience: None, + limits: OpenAIResponsesLimits::default(), + user_agent: None, + originator: None, + } + } + + pub fn with_endpoint(mut self, endpoint: impl Into) -> Self { + self.endpoint = endpoint.into(); + self + } + + /// Replaces non-authentication request headers. Authentication headers are applied last. + pub fn with_headers(mut self, headers: HeaderMap) -> Self { + self.headers = headers; + self + } + + pub fn with_authentication(mut self, authentication: impl Into) -> Self { + self.authentication = authentication.into(); + self + } + + /// Uses a custom refresh-capable authentication provider. + pub fn with_authentication_provider(self, provider: P) -> Self { + self.with_authentication(Authentication::new(provider)) + } + + pub fn with_request_policy(mut self, policy: OpenAIResponsesRequestPolicy) -> Self { + self.request_policy = policy; + self + } + + pub fn with_reasoning_effort(mut self, effort: impl Into) -> Self { + self.reasoning_effort = Some(effort.into()); + self + } + + pub fn with_max_output_tokens(mut self, value: u32) -> Self { + self.max_output_tokens = Some(value); + self + } + + pub fn with_parallel_tool_calls(mut self, value: bool) -> Self { + self.parallel_tool_calls = Some(value); + self + } + + /// Opts into retries and stream/attempt timeouts. `None` means no transient retry. + pub fn with_resilience(mut self, resilience: ResilienceConfig) -> Self { + self.resilience = Some(resilience); + self + } + + /// Overrides request, response, item/event, and per-field limits. + pub fn with_limits(mut self, limits: OpenAIResponsesLimits) -> Self { + self.limits = limits; + self + } + + /// Overrides the adapter's default HTTP user agent. + pub fn with_user_agent(mut self, user_agent: impl Into) -> Self { + self.user_agent = Some(user_agent.into()); + self + } + + /// Sets or overrides the Responses `originator` header. + pub fn with_originator(mut self, originator: impl Into) -> Self { + self.originator = Some(originator.into()); + self + } + + /// Encodes only the transport-neutral JSON request. + /// + /// This synchronous helper does not resolve authentication and therefore has no + /// authentication binding. It rejects adapter-emitted continuation metadata instead of + /// silently dropping bound function-call IDs, reasoning, or generated-image state. Submit + /// the request through [`OpenAIResponsesAdapter`] when continuation replay is required. + pub fn encode_request(&self, request: &TurnRequest) -> Result { + encode_request(self, request) + } +} + +/// Errors produced while configuring or encoding a Responses request. +#[derive(Debug, Error)] +pub enum OpenAIResponsesError { + #[error("failed to create HTTP client: {0}")] + HttpClient(#[source] HttpError), + #[error("invalid Responses request: {0}")] + InvalidRequest(String), + #[error("Responses protocol error: {0}")] + Protocol(String), + #[error("Responses request serialization failed: {0}")] + Serialize(#[source] serde_json::Error), +} + +/// OpenAI Responses adapter with live SSE delivery and capability-gated attempt replacement. +#[derive(Clone)] +pub struct OpenAIResponsesAdapter { + client: Http, + config: Arc, +} + +impl OpenAIResponsesAdapter { + pub fn new(config: OpenAIResponsesConfig) -> Result { + validate_limits(config.limits)?; + let client = reqwest::Client::builder() + .build() + .map(Http::new) + .map_err(|error| OpenAIResponsesError::HttpClient(HttpError::request(error)))?; + Ok(Self::with_client(config, client)) + } + + /// Creates an adapter over an arbitrary AgentKit HTTP transport. + pub fn with_client(config: OpenAIResponsesConfig, client: Http) -> Self { + Self { + client, + config: Arc::new(config), + } + } + + /// Overrides authentication after constructing the adapter. + pub fn with_authentication(mut self, authentication: impl Into) -> Self { + Arc::make_mut(&mut self.config).authentication = authentication.into(); + self + } + + /// Overrides authentication with a custom refresh-capable provider. + pub fn with_authentication_provider(self, provider: P) -> Self { + self.with_authentication(Authentication::new(provider)) + } + + /// Enables retries and stream/attempt timeouts after constructing the adapter. + pub fn with_resilience(mut self, resilience: ResilienceConfig) -> Self { + Arc::make_mut(&mut self.config).resilience = Some(resilience); + self + } +} + +#[async_trait] +impl ModelAdapter for OpenAIResponsesAdapter { + type Session = OpenAIResponsesSession; + + async fn start_session(&self, config: SessionConfig) -> Result { + Ok(OpenAIResponsesSession { + client: self.client.clone(), + config: self.config.clone(), + session: config, + }) + } + + fn provider_name(&self) -> Option<&str> { + Some("openai") + } +} + +/// Active Responses session. +pub struct OpenAIResponsesSession { + client: Http, + config: Arc, + session: SessionConfig, +} + +#[async_trait] +impl ModelSession for OpenAIResponsesSession { + type Turn = OpenAIResponsesTurn; + + async fn begin_turn( + &mut self, + request: TurnRequest, + cancellation: Option, + ) -> Result { + if cancelled(cancellation.as_ref()) { + return Err(LoopError::Cancelled); + } + // One logical deadline starts before encoding and initial authentication. + let deadline = self + .config + .resilience + .as_ref() + .map(|config| LogicalDeadline::new(config.retry_budget)); + deadline_remaining(deadline.as_ref()).map_err(http_loop_error)?; + let auth_timeout = self + .config + .resilience + .as_ref() + .and_then(|config| config.attempt_timeout); + let auth = cancellable( + run_bounded_http( + self.config.authentication.authenticate(None), + auth_timeout, + deadline.as_ref(), + "OpenAI authentication", + ), + cancellation.as_ref(), + ) + .await? + .map_err(http_loop_error)?; + let mut value = encode_request_bound(&self.config, &request, auth.binding()) + .map_err(|error| LoopError::Provider(error.to_string()))?; + // Serialize once. Every status, transport, and stream retry reuses these exact bytes. + let body = agentkit_http::Bytes::from_owner(Zeroizing::new( + serde_json::to_vec(&value) + .map_err(OpenAIResponsesError::Serialize) + .map_err(|error| LoopError::Provider(error.to_string()))?, + )); + zeroize_encrypted_content(&mut value); + let idempotency_key = stable_idempotency_key( + &self.session.session_id.to_string(), + &request.turn_id.to_string(), + &body, + ); + let supersession_enabled = self + .session + .consumer_capabilities + .response_attempt_supersession; + OpenAIResponsesTurn::open( + ResponsesRequestContext { + client: self.client.clone(), + config: self.config.clone(), + body, + idempotency_key, + session_id: request.session_id.to_string(), + turn_state: Arc::new(Mutex::new(None)), + auth, + deadline, + retries: 0, + refreshed: false, + wire_bytes: 0, + }, + supersession_enabled, + cancellation.as_ref(), + ) + .await + } + + fn model_name(&self) -> Option<&str> { + Some(&self.config.model) + } + + fn provider_name(&self) -> Option<&str> { + Some("openai") + } +} + +/// Live-streaming Responses turn. +pub struct OpenAIResponsesTurn { + context: ResponsesRequestContext, + attempt: Option, + supersession_enabled: bool, + attempt_output_emitted: bool, + pending_reopen: bool, + pending_delay: Duration, + finished: bool, +} + +#[async_trait] +impl ModelTurn for OpenAIResponsesTurn { + async fn next_event( + &mut self, + cancellation: Option, + ) -> Result, LoopError> { + self.next_event_inner(cancellation.as_ref()).await + } +} + +impl OpenAIResponsesTurn { + async fn open( + context: ResponsesRequestContext, + supersession_enabled: bool, + cancellation: Option<&TurnCancellation>, + ) -> Result { + let mut turn = Self { + context, + attempt: None, + supersession_enabled, + attempt_output_emitted: false, + pending_reopen: true, + pending_delay: Duration::ZERO, + finished: false, + }; + turn.reopen(cancellation).await?; + Ok(turn) + } + + async fn reopen(&mut self, cancellation: Option<&TurnCancellation>) -> Result<(), LoopError> { + if !self.pending_delay.is_zero() { + let delay = self.pending_delay; + self.pending_delay = Duration::ZERO; + cancellable( + run_bounded_http( + async { + sleep(delay).await; + Ok(()) + }, + None, + self.context.deadline.as_ref(), + "retry backoff", + ), + cancellation, + ) + .await? + .map_err(http_loop_error)?; + } + self.attempt = Some(open_live_attempt(&mut self.context, cancellation).await?); + self.pending_reopen = false; + Ok(()) + } + + async fn next_event_inner( + &mut self, + cancellation: Option<&TurnCancellation>, + ) -> Result, LoopError> { + loop { + if cancelled(cancellation) { + return Err(LoopError::Cancelled); + } + if self.pending_reopen { + self.reopen(cancellation).await?; + } + let expired_attempt_timeout = self + .attempt + .as_ref() + .and_then(|attempt| attempt.deadline.as_ref()) + .filter(|deadline| deadline.started_at.elapsed() >= deadline.budget) + .map(|deadline| deadline.budget); + if expired_attempt_timeout.is_none() + && let Some(event) = self.attempt.as_mut().and_then(LiveAttempt::pop_event) + { + self.attempt_output_emitted = true; + if matches!(event, ModelTurnEvent::Finished(_)) { + self.finished = true; + } + return Ok(Some(event)); + } + if self.finished { + return Ok(None); + } + + let pending = if let Some(timeout) = expired_attempt_timeout { + Err(attempt_timeout_failure(timeout)) + } else { + self.attempt + .as_mut() + .expect("live attempt is open") + .decoder + .process_pending() + }; + let result = if let Err(failure) = pending { + Err(failure) + } else if self + .attempt + .as_ref() + .expect("live attempt is open") + .decoder + .peek_event() + .is_some_and(|event| !matches!(event, ModelTurnEvent::Finished(_))) + { + continue; + } else if self.attempt.as_ref().expect("live attempt is open").eof { + let attempt = self.attempt.as_mut().expect("live attempt is open"); + let finished = attempt.decoder.finish_live(); + if finished.is_ok() { + attempt.closed = true; + } + finished + } else { + let remaining = + deadline_remaining(self.context.deadline.as_ref()).map_err(http_loop_error)?; + let idle_timeout = self + .context + .config + .resilience + .as_ref() + .and_then(|config| config.stream_idle_timeout); + let attempt_remaining = self + .attempt + .as_ref() + .and_then(|attempt| attempt.deadline.as_ref()) + .map(|deadline| { + deadline + .budget + .saturating_sub(deadline.started_at.elapsed()) + }); + let timeout = [idle_timeout, remaining, attempt_remaining] + .into_iter() + .flatten() + .min(); + let chunk = { + let attempt = self.attempt.as_mut().expect("live attempt is open"); + cancellable(next_body_chunk(&mut attempt.body, timeout), cancellation).await + }; + match chunk { + Err(error) => Err(nonretryable(error)), + Ok(Err(error)) => { + if self.context.deadline.as_ref().is_some_and(|deadline| { + deadline.started_at.elapsed() >= deadline.budget + }) { + Err(nonretryable(http_loop_error(HttpError::Timeout { + operation: "logical request retry budget", + timeout: self + .context + .deadline + .as_ref() + .expect("checked logical deadline") + .budget, + }))) + } else if let Some(timeout) = self + .attempt + .as_ref() + .and_then(|attempt| attempt.deadline.as_ref()) + .filter(|deadline| deadline.started_at.elapsed() >= deadline.budget) + .map(|deadline| deadline.budget) + { + Err(attempt_timeout_failure(timeout)) + } else { + Err(transport_failure(error)) + } + } + Ok(Ok(Some(chunk))) => { + if let Some(total) = self.context.wire_bytes.checked_add(chunk.len()) { + self.context.wire_bytes = total; + if total > self.context.config.limits.max_wire_bytes { + Err(protocol_failure( + "Responses logical turn exceeds configured wire-byte limit", + )) + } else { + let attempt = self.attempt.as_mut().expect("live attempt is open"); + attempt.truncated.observe(&chunk); + attempt.decoder.push(&chunk) + } + } else { + Err(protocol_failure("Responses wire-byte count overflowed")) + } + } + Ok(Ok(None)) => { + let attempt = self.attempt.as_mut().expect("live attempt is open"); + attempt.eof = true; + attempt.truncated.finish().map_err(transport_failure) + } + } + }; + if let Err(failure) = result { + if !failure.retryable + || self.context.retries + >= self + .context + .config + .resilience + .as_ref() + .map_or(0, |config| config.max_retries) + { + return Err(*failure.error); + } + if self.attempt_output_emitted && !self.supersession_enabled { + return Err(*failure.error); + } + let delay = self + .context + .config + .resilience + .as_ref() + .expect("retry requires resilience") + .retry_delay(self.context.retries, failure.headers.as_ref()); + self.context.retries += 1; + self.attempt = None; + self.pending_reopen = true; + self.pending_delay = delay; + if self.attempt_output_emitted { + self.attempt_output_emitted = false; + return Ok(Some(ModelTurnEvent::ResponseAttemptSuperseded)); + } + } + } + } +} + +fn encode_request( + config: &OpenAIResponsesConfig, + request: &TurnRequest, +) -> Result { + if request + .transcript + .iter() + .flat_map(|item| &item.parts) + .any(part_has_continuation_metadata) + { + return Err(invalid_request( + "OpenAIResponsesConfig::encode_request cannot encode authentication-bound continuation metadata; submit through OpenAIResponsesAdapter", + )); + } + encode_request_bound(config, request, None) +} + +fn part_has_continuation_metadata(part: &Part) -> bool { + match part { + Part::ToolCall(call) => call.metadata.contains_key(CONTINUATION_METADATA), + Part::Reasoning(reasoning) => reasoning.metadata.contains_key(CONTINUATION_METADATA), + Part::Media(media) => media.metadata.contains_key(CONTINUATION_METADATA), + Part::Text(_) + | Part::Structured(_) + | Part::ToolResult(_) + | Part::File(_) + | Part::Custom(_) => false, + } +} + +fn encode_request_bound( + config: &OpenAIResponsesConfig, + request: &TurnRequest, + authentication_binding: Option<&str>, +) -> Result { + validate_limits(config.limits)?; + if request.transcript.len() > config.limits.max_items + || request.available_tools.len() > config.limits.max_items + { + return Err(OpenAIResponsesError::InvalidRequest( + "too many transcript items or tools".into(), + )); + } + let session_id = request.session_id.to_string(); + let mut input = Vec::new(); + for item in &request.transcript { + let encoded = encode_item(config, &session_id, authentication_binding, item)?; + if input.len().saturating_add(encoded.len()) > config.limits.max_items { + return Err(invalid_request( + "Responses transcript expands beyond the configured item limit", + )); + } + input.extend(encoded); + } + let tools = request + .available_tools + .iter() + .map(|tool| { + validate_tool_name(&tool.name.0)?; + validate_request_field( + &tool.description, + config.limits.max_text_bytes, + "tool description", + )?; + let schema = serde_json::to_string(&tool.input_schema) + .map_err(OpenAIResponsesError::Serialize)?; + validate_request_field(&schema, config.limits.max_text_bytes, "tool input schema")?; + Ok(json!({ + "type": "function", + "name": tool.name.0, + "description": tool.description, + "parameters": tool.input_schema, + "strict": false + })) + }) + .collect::, OpenAIResponsesError>>()?; + let mut body = Map::new(); + body.insert("model".into(), Value::String(config.model.clone())); + body.insert("input".into(), Value::Array(input)); + body.insert("tools".into(), Value::Array(tools)); + body.insert("tool_choice".into(), Value::String("auto".into())); + body.insert("stream".into(), Value::Bool(true)); + if let Some(value) = config.request_policy.store { + body.insert("store".into(), Value::Bool(value)); + } + if config.request_policy.include_encrypted_reasoning { + body.insert("include".into(), json!(["reasoning.encrypted_content"])); + } + if config.reasoning_effort.is_some() || config.request_policy.include_encrypted_reasoning { + let mut reasoning = Map::new(); + reasoning.insert("summary".into(), Value::String("auto".into())); + if let Some(effort) = &config.reasoning_effort { + reasoning.insert("effort".into(), Value::String(effort.clone())); + } + body.insert("reasoning".into(), Value::Object(reasoning)); + } + if let Some(value) = config.parallel_tool_calls { + body.insert("parallel_tool_calls".into(), Value::Bool(value)); + } + if config.request_policy.send_max_output_tokens + && let Some(value) = config.max_output_tokens + { + body.insert("max_output_tokens".into(), Value::from(value)); + } + apply_prompt_cache(&mut body, request, config.profile)?; + let value = Value::Object(body); + let encoded = + Zeroizing::new(serde_json::to_vec(&value).map_err(OpenAIResponsesError::Serialize)?); + if encoded.len() > config.limits.max_request_bytes { + return Err(invalid_request(format!( + "Responses serialized request exceeds {} bytes", + config.limits.max_request_bytes + ))); + } + Ok(value) +} + +fn encode_item( + config: &OpenAIResponsesConfig, + session_id: &str, + authentication_binding: Option<&str>, + item: &Item, +) -> Result, OpenAIResponsesError> { + if item.parts.len() > config.limits.max_items { + return Err(invalid_request( + "Responses transcript item has too many parts", + )); + } + let role = match item.kind { + ItemKind::System | ItemKind::Context + if !config.request_policy.downgrade_system_to_developer => + { + "system" + } + ItemKind::System | ItemKind::Developer | ItemKind::Context => "developer", + ItemKind::User => "user", + ItemKind::Assistant => "assistant", + ItemKind::Tool => "tool", + ItemKind::Notification => "user", + }; + if matches!( + item.kind, + ItemKind::System | ItemKind::Developer | ItemKind::Context | ItemKind::Notification + ) { + let mut text = stringify_message_parts(&item.parts, item.kind)?; + if item.kind == ItemKind::Notification { + text = wrap_notification(&text); + } + if text.is_empty() { + return Ok(Vec::new()); + } + validate_request_field(&text, config.limits.max_text_bytes, "message text")?; + return Ok(vec![json!({ + "type": "message", + "role": role, + "content": [{"type": "input_text", "text": text}] + })]); + } + let mut output = Vec::new(); + let mut content = Vec::new(); + for part in &item.parts { + validate_part_role(item.kind, part)?; + match part { + Part::Text(text) => { + let text = match item.kind { + ItemKind::Notification => { + format!("{}", text.text) + } + _ => text.text.clone(), + }; + validate_request_field(&text, config.limits.max_text_bytes, "message text")?; + content.push(json!({ + "type": if role == "assistant" { "output_text" } else { "input_text" }, + "text": text + })); + } + Part::Structured(value) => { + let text = + serde_json::to_string(&value.value).map_err(OpenAIResponsesError::Serialize)?; + validate_request_field( + &text, + config.limits.max_text_bytes, + "structured message text", + )?; + content.push(json!({ + "type": if role == "assistant" { "output_text" } else { "input_text" }, + "text": text + })); + } + Part::ToolCall(call) => { + validate_tool_name(&call.name)?; + let continuation = continuation_from_metadata( + &call.metadata, + config, + session_id, + authentication_binding, + "function_call", + false, + )?; + let arguments = + serde_json::to_string(&call.input).map_err(OpenAIResponsesError::Serialize)?; + validate_request_field( + &arguments, + config.limits.max_text_bytes, + "function arguments", + )?; + let mut value = json!({ + "type": "function_call", + "call_id": call.id.0, + "name": call.name, + "arguments": arguments + }); + if let Some(continuation) = continuation { + value["id"] = Value::String(continuation.item_id.to_owned()); + } + output.push(value); + } + Part::ToolResult(result) => output.push(json!({ + "type": "function_call_output", + "call_id": result.call_id.0, + "output": encode_tool_output(&result.output, config)? + })), + Part::Reasoning(reasoning) => { + if let Some(continuation) = continuation_from_metadata( + &reasoning.metadata, + config, + session_id, + authentication_binding, + "reasoning", + true, + )? { + output.push(json!({ + "id": continuation.item_id, + "type": "reasoning", + "summary": [], + "encrypted_content": continuation.encrypted_content.expect("validated encrypted reasoning") + })); + } + // A readable summary without encrypted continuation state is deliberately not + // converted into ordinary assistant text. + } + Part::Media(media) => { + if role == "assistant" { + let generated = + generated_image_item(media, config, session_id, authentication_binding)?; + if let Some(generated) = generated { + output.push(generated); + } else { + return Err(invalid_request( + "assistant media is not a persisted Responses generated image", + )); + } + } else if role == "tool" { + return Err(invalid_request( + "tool item media must be carried by a tool result", + )); + } else { + content.push(encode_media(media, config)?); + } + } + Part::File(_) | Part::Custom(_) => { + return Err(invalid_request( + "Responses transcript contains unsupported content", + )); + } + } + } + if !content.is_empty() && role != "tool" { + output.insert( + 0, + json!({"type": "message", "role": role, "content": content}), + ); + } + Ok(output) +} + +fn validate_tool_name(name: &str) -> Result<(), OpenAIResponsesError> { + if name.is_empty() + || name.len() > 64 + || !name.chars().all(|character| { + character.is_ascii_alphanumeric() || character == '_' || character == '-' + }) + { + return Err(invalid_request(format!( + "invalid OpenAI tool name `{name}`" + ))); + } + Ok(()) +} + +fn validate_part_role(role: ItemKind, part: &Part) -> Result<(), OpenAIResponsesError> { + let supported = match role { + ItemKind::System | ItemKind::Developer | ItemKind::Context | ItemKind::Notification => { + matches!( + part, + Part::Text(_) | Part::Structured(_) | Part::Reasoning(_) + ) + } + ItemKind::User => matches!(part, Part::Text(_) | Part::Structured(_) | Part::Media(_)), + ItemKind::Assistant => matches!( + part, + Part::Text(_) + | Part::Structured(_) + | Part::ToolCall(_) + | Part::Reasoning(_) + | Part::Media(_) + ), + ItemKind::Tool => matches!(part, Part::ToolResult(_)), + }; + if supported { + Ok(()) + } else { + Err(unsupported_part(role, part)) + } +} + +fn unsupported_part(role: ItemKind, part: &Part) -> OpenAIResponsesError { + invalid_request(format!( + "Responses does not support {} content on {role:?} items", + match part { + Part::Text(_) => "text", + Part::Media(_) => "media", + Part::File(_) => "file", + Part::Structured(_) => "structured", + Part::Reasoning(_) => "reasoning", + Part::ToolCall(_) => "tool-call", + Part::ToolResult(_) => "tool-result", + Part::Custom(_) => "custom", + } + )) +} + +fn stringify_message_parts(parts: &[Part], role: ItemKind) -> Result { + let mut segments = Vec::new(); + for part in parts { + validate_part_role(role, part)?; + match part { + Part::Text(text) => segments.push(text.text.clone()), + Part::Structured(value) => segments.push( + serde_json::to_string_pretty(&value.value) + .map_err(OpenAIResponsesError::Serialize)?, + ), + Part::Reasoning(reasoning) => { + if let Some(summary) = &reasoning.summary { + segments.push(summary.clone()); + } + } + _ => unreachable!("validated string content"), + } + } + Ok(segments.join("\n\n")) +} + +fn wrap_notification(text: &str) -> String { + format!("\n{text}\n") +} + +fn apply_prompt_cache( + body: &mut Map, + request: &TurnRequest, + profile: OpenAIResponsesProfile, +) -> Result<(), OpenAIResponsesError> { + let Some(cache) = &request.cache else { + return Ok(()); + }; + if matches!(cache.mode, PromptCacheMode::Disabled) { + return Ok(()); + } + if matches!(cache.strategy, PromptCacheStrategy::Explicit { .. }) + && matches!(cache.mode, PromptCacheMode::Required) + { + return Err(invalid_request( + "Responses does not support required explicit cache breakpoints", + )); + } + if let Some(retention) = cache.retention + && profile == OpenAIResponsesProfile::Public + { + body.insert( + "prompt_cache_retention".into(), + Value::String(crate::prompt_cache_retention_value(retention).into()), + ); + } + if let Some(key) = &cache.key { + if key.is_empty() || key.len() > MAX_CACHE_KEY_BYTES || !key.is_ascii() { + return Err(invalid_request( + "Responses prompt cache key is outside canonical bounds", + )); + } + body.insert("prompt_cache_key".into(), Value::String(key.clone())); + } + Ok(()) +} + +fn encode_media( + media: &MediaPart, + config: &OpenAIResponsesConfig, +) -> Result { + let profile = config.profile; + let expected = match media.modality { + Modality::Image => "image/", + Modality::Audio if profile == OpenAIResponsesProfile::ChatGptPrivate => "audio/", + Modality::Audio => { + return Err(invalid_request( + "public Responses audio input is not supported by this adapter", + )); + } + Modality::Video | Modality::Binary => { + return Err(invalid_request( + "Responses supports only image and private-profile audio media input", + )); + } + }; + if !media.mime_type.starts_with(expected) || media.mime_type.contains(['\r', '\n', ';', ',']) { + return Err(invalid_request("Responses media has an invalid MIME type")); + } + let url = media_data_url(media, config.limits.max_text_bytes)?; + Ok(match media.modality { + Modality::Image => { + json!({"type": "input_image", "image_url": url, "detail": "high"}) + } + Modality::Audio => json!({"type": "input_audio", "audio_url": url}), + Modality::Video | Modality::Binary => unreachable!("rejected above"), + }) +} + +fn media_data_url( + media: &MediaPart, + max_text_bytes: usize, +) -> Result { + let url = match &media.data { + DataRef::InlineBytes(bytes) => Ok(format!( + "data:{};base64,{}", + media.mime_type, + STANDARD.encode(bytes) + )), + DataRef::InlineText(text) if text.starts_with("data:") => { + validate_data_url(text, &media.mime_type)?; + Ok(text.clone()) + } + DataRef::InlineText(text) => { + validate_base64(text, "inline media")?; + Ok(format!("data:{};base64,{text}", media.mime_type)) + } + DataRef::Uri(uri) if uri.starts_with("data:") => { + validate_data_url(uri, &media.mime_type)?; + Ok(uri.clone()) + } + DataRef::Uri(uri) + if media.modality == Modality::Image + && uri.len() <= max_text_bytes + && reqwest::Url::parse(uri) + .is_ok_and(|url| matches!(url.scheme(), "http" | "https")) => + { + Ok(uri.clone()) + } + DataRef::Uri(_) => Err(invalid_request( + "Responses cannot read this media URI; use inline bytes or an HTTP(S) image URL", + )), + DataRef::Handle(_) => Err(invalid_request( + "Responses cannot resolve media handles; use inline bytes", + )), + }?; + validate_request_field(&url, max_text_bytes, "media")?; + Ok(url) +} + +fn validate_data_url(value: &str, mime_type: &str) -> Result<(), OpenAIResponsesError> { + let payload = value + .strip_prefix(&format!("data:{mime_type};base64,")) + .filter(|payload| !payload.is_empty()) + .ok_or_else(|| invalid_request("Responses media data URL is not canonical base64"))?; + validate_base64(payload, "media data URL") +} + +fn validate_base64(value: &str, field: &str) -> Result<(), OpenAIResponsesError> { + let decoded = STANDARD + .decode(value) + .map_err(|_| invalid_request(format!("Responses {field} is not valid base64")))?; + if decoded.is_empty() || STANDARD.encode(&decoded) != value { + return Err(invalid_request(format!( + "Responses {field} is not canonical base64" + ))); + } + Ok(()) +} + +fn generated_image_item( + media: &MediaPart, + config: &OpenAIResponsesConfig, + session_id: &str, + authentication_binding: Option<&str>, +) -> Result, OpenAIResponsesError> { + let Some(metadata) = media.metadata.get(GENERATED_IMAGE_METADATA) else { + return Ok(None); + }; + let metadata = metadata + .as_object() + .filter(|metadata| (2..=3).contains(&metadata.len())) + .ok_or_else(|| protocol_error("generated image metadata is malformed"))?; + let item_id = bounded_metadata_string(metadata.get("item_id"), "generated image item_id")?; + if metadata.get("status").and_then(Value::as_str) != Some("completed") + || media.modality != Modality::Image + || media.mime_type != "image/png" + { + return Err(protocol_error("generated image metadata is invalid")); + } + let Some(continuation) = continuation_from_metadata( + &media.metadata, + config, + session_id, + authentication_binding, + "image_generation_call", + false, + )? + else { + return Ok(None); + }; + if continuation.item_id != item_id { + return Err(protocol_error( + "generated image continuation item binding is invalid", + )); + } + let revised_prompt = metadata + .get("revised_prompt") + .filter(|value| !value.is_null()) + .map(|value| { + let value = value.as_str().ok_or_else(|| { + protocol_error("Responses generated image revised prompt is invalid") + })?; + validate_request_field( + value, + config.limits.max_text_bytes, + "generated image revised prompt", + )?; + Ok(value) + }) + .transpose()?; + let result = match &media.data { + DataRef::InlineBytes(bytes) if !bytes.is_empty() => STANDARD.encode(bytes), + DataRef::InlineText(text) if !text.starts_with("data:") => { + validate_base64(text, "persisted generated image result")?; + text.clone() + } + DataRef::InlineText(text) => { + validate_data_url(text, "image/png")?; + text.split_once(',') + .map(|(_, payload)| payload.to_owned()) + .ok_or_else(|| protocol_error("persisted generated image data URL is malformed"))? + } + DataRef::Uri(_) | DataRef::Handle(_) | DataRef::InlineBytes(_) => { + return Err(invalid_request( + "Responses cannot replay a generated image without inline bytes", + )); + } + }; + validate_request_field( + &result, + config.limits.max_text_bytes, + "persisted generated image result", + )?; + Ok(Some(json!({ + "id": item_id, + "type": "image_generation_call", + "status": "completed", + "revised_prompt": revised_prompt, + "result": result, + }))) +} + +struct Continuation<'a> { + item_id: &'a str, + encrypted_content: Option<&'a str>, +} + +fn continuation_from_metadata<'a>( + metadata: &'a MetadataMap, + config: &OpenAIResponsesConfig, + session_id: &str, + authentication_binding: Option<&str>, + expected_kind: &str, + encrypted_required: bool, +) -> Result>, OpenAIResponsesError> { + let Some(raw) = metadata.get(CONTINUATION_METADATA) else { + return Ok(None); + }; + let object = raw + .as_object() + .ok_or_else(|| protocol_error("Responses continuation metadata is not an object"))?; + let expected_len = if encrypted_required { 7 } else { 6 }; + if object.len() != expected_len + || object.get("schema_version").and_then(Value::as_u64) != Some(CONTINUATION_SCHEMA_VERSION) + || object.get("kind").and_then(Value::as_str) != Some(expected_kind) + || object.get("model").and_then(Value::as_str) != Some(&*config.model) + || object.get("session_id").and_then(Value::as_str) != Some(session_id) + { + return Err(protocol_error( + "Responses continuation metadata binding is invalid", + )); + } + let metadata_binding = bounded_metadata_string( + object.get("authentication_binding"), + "authentication binding", + )?; + let item_id = bounded_metadata_string(object.get("item_id"), "item_id")?; + let encrypted_content = object.get("encrypted_content").and_then(Value::as_str); + if encrypted_required + && encrypted_content + .is_none_or(|value| value.is_empty() || value.len() > config.limits.max_text_bytes) + { + return Err(protocol_error( + "Responses encrypted continuation is invalid", + )); + } + if authentication_binding != Some(metadata_binding) { + return Ok(None); + } + Ok(Some(Continuation { + item_id, + encrypted_content, + })) +} + +fn bounded_metadata_string<'a>( + value: Option<&'a Value>, + name: &str, +) -> Result<&'a str, OpenAIResponsesError> { + value + .and_then(Value::as_str) + .filter(|value| !value.is_empty() && value.len() <= 512) + .ok_or_else(|| protocol_error(format!("Responses continuation {name} is invalid"))) +} + +#[allow(clippy::too_many_arguments)] +fn continuation_metadata( + model: &str, + session_id: &str, + authentication_binding: Option<&str>, + _response_id: &str, + item_id: &str, + _output_index: u64, + kind: &str, + encrypted_content: Option<&str>, +) -> MetadataMap { + let Some(authentication_binding) = authentication_binding else { + return MetadataMap::new(); + }; + let mut value = json!({ + "schema_version": CONTINUATION_SCHEMA_VERSION, + "authentication_binding": authentication_binding, + "model": model, + "session_id": session_id, + "item_id": item_id, + "kind": kind, + }); + if let Some(encrypted) = encrypted_content { + value["encrypted_content"] = Value::String(encrypted.to_owned()); + } + MetadataMap::from([(CONTINUATION_METADATA.into(), value)]) +} + +fn validate_request_field( + value: &str, + max_text_bytes: usize, + field: &str, +) -> Result<(), OpenAIResponsesError> { + if value.len() > max_text_bytes { + return Err(invalid_request(format!( + "Responses {field} exceeds {max_text_bytes} bytes" + ))); + } + Ok(()) +} + +fn validate_request_value( + value: &Value, + max_text_bytes: usize, + field: &str, +) -> Result<(), OpenAIResponsesError> { + match value { + Value::String(value) => validate_request_field(value, max_text_bytes, field), + Value::Array(values) => { + for value in values { + validate_request_value(value, max_text_bytes, field)?; + } + Ok(()) + } + Value::Object(values) => { + for value in values.values() { + validate_request_value(value, max_text_bytes, field)?; + } + Ok(()) + } + Value::Null | Value::Bool(_) | Value::Number(_) => Ok(()), + } +} + +fn validate_limits(limits: OpenAIResponsesLimits) -> Result<(), OpenAIResponsesError> { + if limits.max_request_bytes == 0 + || limits.max_attempt_bytes == 0 + || limits.max_wire_bytes == 0 + || limits.max_items == 0 + || limits.max_text_bytes == 0 + { + return Err(invalid_request("Responses limits must be non-zero")); + } + if limits.max_text_bytes > limits.max_request_bytes + || limits.max_text_bytes > limits.max_attempt_bytes + || limits.max_attempt_bytes > limits.max_wire_bytes + { + return Err(invalid_request( + "Responses limits are inconsistent: text bytes must fit request and attempt bounds, and attempt bytes must fit the aggregate wire bound", + )); + } + Ok(()) +} + +fn invalid_request(message: impl Into) -> OpenAIResponsesError { + OpenAIResponsesError::InvalidRequest(message.into()) +} + +fn protocol_error(message: impl Into) -> OpenAIResponsesError { + OpenAIResponsesError::Protocol(message.into()) +} + +fn encode_tool_output( + output: &ToolOutput, + config: &OpenAIResponsesConfig, +) -> Result { + let value = if config.profile == OpenAIResponsesProfile::Public { + match output { + ToolOutput::Text(value) => Ok(Value::String(value.clone())), + ToolOutput::Structured(value) => serde_json::to_string(value) + .map(Value::String) + .map_err(OpenAIResponsesError::Serialize), + ToolOutput::Parts(parts) => serde_json::to_string(parts) + .map(Value::String) + .map_err(OpenAIResponsesError::Serialize), + ToolOutput::Files(files) => serde_json::to_string(files) + .map(Value::String) + .map_err(OpenAIResponsesError::Serialize), + } + } else { + match output { + ToolOutput::Text(value) => Ok(Value::String(value.clone())), + ToolOutput::Structured(value) => serde_json::to_string(value) + .map(Value::String) + .map_err(OpenAIResponsesError::Serialize), + ToolOutput::Parts(parts) if parts.iter().any(|part| matches!(part, Part::Media(_))) => { + parts + .iter() + .map(|part| match part { + Part::Text(text) => Ok(json!({"type": "input_text", "text": text.text})), + Part::Structured(value) => Ok(json!({ + "type": "input_text", + "text": serde_json::to_string(&value.value) + .map_err(OpenAIResponsesError::Serialize)?, + })), + Part::Media(media) => encode_media(media, config), + _ => Err(invalid_request( + "private Responses tool output contains unsupported content", + )), + }) + .collect::, _>>() + .map(Value::Array) + } + ToolOutput::Parts(parts) => parts + .iter() + .map(|part| match part { + Part::Text(text) => Ok(text.text.clone()), + Part::Structured(value) => { + serde_json::to_string(&value.value).map_err(OpenAIResponsesError::Serialize) + } + _ => Err(invalid_request( + "private Responses tool output contains unsupported content", + )), + }) + .collect::, _>>() + .map(|parts| Value::String(parts.join("\n"))), + ToolOutput::Files(_) => Err(invalid_request( + "private Responses file tool output is unsupported", + )), + } + }?; + validate_request_value(&value, config.limits.max_text_bytes, "tool output")?; + Ok(value) +} + +#[derive(Clone, Debug)] +struct LogicalDeadline { + started_at: Instant, + budget: Duration, +} + +impl LogicalDeadline { + fn new(budget: Duration) -> Self { + Self { + started_at: Instant::now(), + budget, + } + } + + fn remaining(&self) -> Result { + let elapsed = self.started_at.elapsed(); + if elapsed >= self.budget { + Err(HttpError::Timeout { + operation: "logical request retry budget", + timeout: self.budget, + }) + } else { + Ok(self.budget - elapsed) + } + } +} + +fn deadline_remaining(deadline: Option<&LogicalDeadline>) -> Result, HttpError> { + deadline.map(LogicalDeadline::remaining).transpose() +} + +async fn run_bounded_http( + future: F, + operation_timeout: Option, + deadline: Option<&LogicalDeadline>, + operation: &'static str, +) -> Result +where + F: Future>, +{ + let remaining = deadline_remaining(deadline)?; + let timeout = match (operation_timeout, remaining) { + (Some(operation), Some(remaining)) => operation.min(remaining), + (Some(operation), None) => operation, + (None, Some(remaining)) => remaining, + (None, None) => return future.await, + }; + let budget_limited = remaining + .is_some_and(|remaining| operation_timeout.is_none_or(|operation| remaining <= operation)); + futures_util::pin_mut!(future); + let timer = sleep(timeout); + futures_util::pin_mut!(timer); + match select(future, timer).await { + Either::Left((result, _)) => result, + Either::Right((_, _)) if budget_limited => Err(HttpError::Timeout { + operation: "logical request retry budget", + timeout: deadline + .expect("budget-limited operation has deadline") + .budget, + }), + Either::Right((_, _)) => Err(HttpError::Timeout { operation, timeout }), + } +} + +fn http_loop_error(error: HttpError) -> LoopError { + LoopError::Provider(format!("OpenAI Responses {error}")) +} + +#[derive(Debug)] +struct AttemptFailure { + error: Box, + retryable: bool, + headers: Option, +} + +struct ResponsesRequestContext { + client: Http, + config: Arc, + body: agentkit_http::Bytes, + idempotency_key: String, + session_id: String, + turn_state: Arc>>, + auth: AuthenticationAttempt, + deadline: Option, + retries: usize, + refreshed: bool, + wire_bytes: usize, +} + +struct LiveAttempt { + body: agentkit_http::BodyStream, + truncated: TruncatedStreamDetector, + decoder: ResponsesSseDecoder, + deadline: Option, + eof: bool, + closed: bool, +} + +impl LiveAttempt { + fn pop_event(&mut self) -> Option { + if !self.closed && matches!(self.decoder.peek_event(), Some(ModelTurnEvent::Finished(_))) { + return None; + } + self.decoder.pop_event() + } +} + +async fn open_live_attempt( + context: &mut ResponsesRequestContext, + cancellation: Option<&TurnCancellation>, +) -> Result { + loop { + let attempt_timeout = context + .config + .resilience + .as_ref() + .and_then(|config| config.attempt_timeout); + let attempt_deadline = attempt_timeout.map(LogicalDeadline::new); + let result = attempt_with_timeout( + send_live_attempt(context, cancellation), + attempt_timeout, + context.deadline.as_ref(), + cancellation, + ) + .await; + match result { + Ok(mut attempt) => { + attempt.deadline = attempt_deadline; + return Ok(attempt); + } + Err(failure) if is_unauthorized(&failure.error) && !context.refreshed => { + let binding = context.auth.binding().map(str::to_owned); + let refreshed = cancellable( + run_bounded_http( + context + .config + .authentication + .authenticate(Some(&context.auth)), + context + .config + .resilience + .as_ref() + .and_then(|config| config.attempt_timeout), + context.deadline.as_ref(), + "OpenAI reauthentication", + ), + cancellation, + ) + .await? + .map_err(http_loop_error)?; + if refreshed.binding() != binding.as_deref() { + return Err(LoopError::Provider( + "OpenAI authentication binding changed during reactive refresh".into(), + )); + } + context.auth = refreshed; + context.refreshed = true; + } + Err(failure) + if failure.retryable + && context.retries + < context + .config + .resilience + .as_ref() + .map_or(0, |config| config.max_retries) => + { + let delay = context + .config + .resilience + .as_ref() + .expect("retry requires resilience") + .retry_delay(context.retries, failure.headers.as_ref()); + context.retries += 1; + cancellable( + run_bounded_http( + async { + sleep(delay).await; + Ok(()) + }, + None, + context.deadline.as_ref(), + "retry backoff", + ), + cancellation, + ) + .await? + .map_err(http_loop_error)?; + } + Err(failure) => return Err(*failure.error), + } + } +} + +async fn send_live_attempt( + context: &ResponsesRequestContext, + cancellation: Option<&TurnCancellation>, +) -> Result { + deadline_remaining(context.deadline.as_ref()) + .map_err(|error| nonretryable(http_loop_error(error)))?; + let mut headers = context.config.headers.clone(); + headers.insert("accept", HeaderValue::from_static("text/event-stream")); + headers.insert("content-type", HeaderValue::from_static("application/json")); + if let Some(user_agent) = &context.config.user_agent { + headers.insert( + "user-agent", + HeaderValue::from_str(user_agent) + .map_err(|_| protocol_failure("invalid user-agent header"))?, + ); + } else { + headers + .entry("user-agent") + .or_insert(HeaderValue::from_static(concat!( + "agentkit-provider-openai/", + env!("CARGO_PKG_VERSION") + ))); + } + headers.insert( + "idempotency-key", + HeaderValue::from_str(&context.idempotency_key) + .map_err(|_| protocol_failure("invalid idempotency key"))?, + ); + if let Some(originator) = &context.config.originator { + headers.insert( + "originator", + HeaderValue::from_str(originator) + .map_err(|_| protocol_failure("invalid originator header"))?, + ); + } + let sent_turn_state = if context.config.profile == OpenAIResponsesProfile::ChatGptPrivate { + headers.remove(X_CODEX_TURN_STATE); + headers + .entry("originator") + .or_insert(HeaderValue::from_static("agentkit")); + headers.entry("session-id").or_insert( + HeaderValue::from_str(&context.session_id) + .map_err(|_| protocol_failure("invalid session ID header"))?, + ); + let state = context + .turn_state + .lock() + .map_err(|_| protocol_failure("turn-state lock poisoned"))? + .clone(); + if let Some(value) = &state { + headers.insert(X_CODEX_TURN_STATE, value.clone()); + } + state + } else { + None + }; + headers.extend(context.auth.headers().clone()); + let response = cancellable( + context + .client + .post(&context.config.endpoint) + .headers(headers) + .body(context.body.clone()) + .send(), + cancellation, + ) + .await + .map_err(nonretryable)? + .map_err(transport_failure)?; + + let status = response.status(); + if status == StatusCode::UNAUTHORIZED { + return Err(AttemptFailure { + error: Box::new(LoopError::Provider( + "OpenAI Responses returned 401 Unauthorized".into(), + )), + retryable: false, + headers: None, + }); + } + if !status.is_success() { + return Err(AttemptFailure { + error: Box::new(LoopError::Provider(format!( + "OpenAI Responses returned HTTP {status}" + ))), + retryable: is_retryable_status(status) + || (context.config.profile == OpenAIResponsesProfile::ChatGptPrivate + && status.as_u16() == 529), + headers: retry_headers(response.headers()), + }); + } + if response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| value.split(';').next().unwrap_or("").trim() != "text/event-stream") + { + return Err(nonretryable(LoopError::Provider( + "OpenAI Responses response is not text/event-stream".into(), + ))); + } + if context.config.profile == OpenAIResponsesProfile::ChatGptPrivate + && let Some(captured) = validated_turn_state_header(response.headers())? + { + if sent_turn_state + .as_ref() + .is_some_and(|expected| expected != captured) + { + return Err(protocol_failure("provider changed x-codex-turn-state")); + } + *context + .turn_state + .lock() + .map_err(|_| protocol_failure("turn-state lock poisoned"))? = Some(captured); + } + let truncated = TruncatedStreamDetector::from_headers(response.headers()); + Ok(LiveAttempt { + body: response.bytes_stream(), + truncated, + decoder: ResponsesSseDecoder::with_policy( + &context.config.model, + &context.session_id, + context.config.profile, + context.config.request_policy.include_encrypted_reasoning, + context.auth.binding(), + context.turn_state.clone(), + context.config.limits, + ), + deadline: None, + eof: false, + closed: false, + }) +} + +fn attempt_timeout_failure(timeout: Duration) -> AttemptFailure { + AttemptFailure { + error: Box::new(LoopError::Provider(format!( + "Responses attempt timed out after {timeout:?}" + ))), + retryable: true, + headers: None, + } +} + +async fn attempt_with_timeout( + future: F, + timeout: Option, + deadline: Option<&LogicalDeadline>, + cancellation: Option<&TurnCancellation>, +) -> Result +where + F: Future>, +{ + let timed = async { + let remaining = match deadline_remaining(deadline) { + Ok(value) => value, + Err(error) => return Err(nonretryable(http_loop_error(error))), + }; + let timeout = match (timeout, remaining) { + (Some(operation), Some(remaining)) => operation.min(remaining), + (Some(operation), None) => operation, + (None, Some(remaining)) => remaining, + (None, None) => return future.await, + }; + let budget_limited = remaining.is_some_and(|remaining| timeout == remaining); + futures_util::pin_mut!(future); + let timer = sleep(timeout); + futures_util::pin_mut!(timer); + match select(future, timer).await { + Either::Left((result, _)) => result, + Either::Right((_, _)) if budget_limited => { + Err(nonretryable(http_loop_error(HttpError::Timeout { + operation: "logical request retry budget", + timeout: deadline.expect("budget timeout has deadline").budget, + }))) + } + Either::Right((_, _)) => Err(attempt_timeout_failure(timeout)), + } + }; + match cancellable(timed, cancellation).await { + Ok(result) => result, + Err(error) => Err(AttemptFailure { + error: Box::new(error), + retryable: false, + headers: None, + }), + } +} + +struct ResponsesSseDecoder { + buffer: Zeroizing>, + buffer_start: usize, + received: usize, + max_attempt_bytes: usize, + state: ResponsesState, +} + +impl ResponsesSseDecoder { + #[cfg(test)] + fn new(model: &str, session_id: &str) -> Self { + Self::with_policy( + model, + session_id, + OpenAIResponsesProfile::Public, + true, + Some("test-authentication-binding"), + Arc::new(Mutex::new(None)), + OpenAIResponsesLimits::default(), + ) + } + + fn with_policy( + model: &str, + session_id: &str, + profile: OpenAIResponsesProfile, + require_encrypted_reasoning: bool, + authentication_binding: Option<&str>, + turn_state: Arc>>, + limits: OpenAIResponsesLimits, + ) -> Self { + Self { + buffer: Zeroizing::new(Vec::new()), + buffer_start: 0, + received: 0, + max_attempt_bytes: limits.max_attempt_bytes, + state: ResponsesState::new( + model, + session_id, + profile, + require_encrypted_reasoning, + authentication_binding, + turn_state, + limits, + ), + } + } + + fn push(&mut self, bytes: &[u8]) -> Result<(), AttemptFailure> { + self.received = self.received.saturating_add(bytes.len()); + if self.received > self.max_attempt_bytes { + return Err(protocol_failure( + "Responses SSE attempt exceeds configured byte limit", + )); + } + if self.buffer_start != 0 { + self.buffer[..self.buffer_start].zeroize(); + self.buffer.drain(..self.buffer_start); + self.buffer_start = 0; + } + zeroizing_extend(&mut self.buffer, bytes, self.max_attempt_bytes)?; + self.process_pending() + } + + fn process_pending(&mut self) -> Result<(), AttemptFailure> { + while self.state.events.is_empty() { + let Some((relative_end, delimiter)) = frame_end(&self.buffer[self.buffer_start..]) + else { + if self.buffer.len().saturating_sub(self.buffer_start) > self.max_attempt_bytes { + return Err(protocol_failure( + "Responses SSE event exceeds configured byte limit", + )); + } + return Ok(()); + }; + let end = self.buffer_start + relative_end; + if end - self.buffer_start > self.max_attempt_bytes { + return Err(protocol_failure( + "Responses SSE event exceeds configured byte limit", + )); + } + let mut frame = Zeroizing::new(self.buffer[self.buffer_start..end].to_vec()); + self.buffer_start = end + delimiter; + self.consume_frame(&mut frame)?; + } + Ok(()) + } + + #[cfg(test)] + fn process_all_pending(&mut self) -> Result<(), AttemptFailure> { + while let Some((relative_end, delimiter)) = frame_end(&self.buffer[self.buffer_start..]) { + let end = self.buffer_start + relative_end; + if end - self.buffer_start > self.max_attempt_bytes { + return Err(protocol_failure( + "Responses SSE event exceeds configured byte limit", + )); + } + let mut frame = Zeroizing::new(self.buffer[self.buffer_start..end].to_vec()); + self.buffer_start = end + delimiter; + self.consume_frame(&mut frame)?; + } + Ok(()) + } + + fn peek_event(&self) -> Option<&ModelTurnEvent> { + self.state.events.front() + } + + fn pop_event(&mut self) -> Option { + self.state.events.pop_front() + } + + fn finish_live(&mut self) -> Result<(), AttemptFailure> { + if self.buffer_start != self.buffer.len() { + return Err(AttemptFailure { + error: Box::new(LoopError::Provider( + "Responses SSE closed with a partial event".into(), + )), + retryable: true, + headers: None, + }); + } + self.state.validate_finished() + } + + #[cfg(test)] + fn finish(mut self) -> Result, AttemptFailure> { + self.process_all_pending()?; + if self.buffer_start != self.buffer.len() { + return Err(AttemptFailure { + error: Box::new(LoopError::Provider( + "Responses SSE closed with a partial event".into(), + )), + retryable: true, + headers: None, + }); + } + self.state.finish() + } + + fn consume_frame(&mut self, frame: &mut [u8]) -> Result<(), AttemptFailure> { + let text = std::str::from_utf8(frame) + .map_err(|_| protocol_failure("Responses SSE event is not UTF-8"))?; + let mut event = None; + let mut data = Vec::new(); + for raw in text.lines() { + let line = raw.trim_end_matches('\r'); + if line.starts_with(':') || line.is_empty() { + continue; + } + let (field, value) = line.split_once(':').unwrap_or((line, "")); + let value = value.strip_prefix(' ').unwrap_or(value); + match field { + "event" => event = Some(value), + "data" => data.push(value), + "id" | "retry" => {} + _ => {} + } + } + if data.is_empty() { + return Ok(()); + } + let data = Zeroizing::new(data.join("\n")); + if data.as_str() == "[DONE]" { + return Err(protocol_failure( + "Responses SSE used an unsupported terminal marker", + )); + } + let mut value: Value = serde_json::from_str(data.as_str()) + .map_err(|_| protocol_failure("Responses SSE data is malformed JSON"))?; + let kind = value + .get("type") + .and_then(Value::as_str) + .ok_or_else(|| protocol_failure("Responses SSE event omitted type"))? + .to_owned(); + if event.is_some_and(|name| name != kind) { + zeroize_encrypted_content(&mut value); + return Err(protocol_failure("Responses SSE event name/type mismatch")); + } + let result = self.state.consume(&kind, &value); + zeroize_encrypted_content(&mut value); + result + } +} + +struct PartAccumulator { + id: PartId, + text: String, +} + +struct ResponsesState { + requested_model: String, + profile: OpenAIResponsesProfile, + limits: OpenAIResponsesLimits, + session_id: String, + authentication_binding: Option, + turn_state: Arc>>, + require_encrypted_reasoning: bool, + sequence: Option, + created: bool, + terminal: bool, + response_id: Option, + response_model: Option, + events: VecDeque, + event_count: usize, + sse_event_count: usize, + output: BTreeMap, + item_indices: BTreeMap, + item_types: BTreeMap, + seen_ids: BTreeSet, + seen_indices: BTreeSet, + done_ids: BTreeSet, + seen_call_ids: BTreeSet, + content_added: BTreeSet<(String, u64)>, + content_done: BTreeSet<(String, u64)>, + summary_added: BTreeSet<(String, u64)>, + summary_done: BTreeSet<(String, u64)>, + argument_done: BTreeSet, + text: BTreeMap<(String, u64), PartAccumulator>, + reasoning: BTreeMap<(String, u64), PartAccumulator>, + function_arguments: BTreeMap, + usage: Option, + finish_reason: Option, + provider_finish_reason: Option, + tool_call: bool, + next_media: usize, +} + +impl ResponsesState { + fn new( + model: &str, + session_id: &str, + profile: OpenAIResponsesProfile, + require_encrypted_reasoning: bool, + authentication_binding: Option<&str>, + turn_state: Arc>>, + limits: OpenAIResponsesLimits, + ) -> Self { + Self { + requested_model: model.to_owned(), + profile, + limits, + session_id: session_id.to_owned(), + authentication_binding: authentication_binding.map(str::to_owned), + turn_state, + require_encrypted_reasoning, + sequence: None, + created: false, + terminal: false, + response_id: None, + response_model: None, + events: VecDeque::new(), + event_count: 0, + sse_event_count: 0, + output: BTreeMap::new(), + item_indices: BTreeMap::new(), + item_types: BTreeMap::new(), + seen_ids: BTreeSet::new(), + seen_indices: BTreeSet::new(), + done_ids: BTreeSet::new(), + seen_call_ids: BTreeSet::new(), + content_added: BTreeSet::new(), + content_done: BTreeSet::new(), + summary_added: BTreeSet::new(), + summary_done: BTreeSet::new(), + argument_done: BTreeSet::new(), + text: BTreeMap::new(), + reasoning: BTreeMap::new(), + function_arguments: BTreeMap::new(), + usage: None, + finish_reason: None, + provider_finish_reason: None, + tool_call: false, + next_media: 0, + } + } + + fn consume(&mut self, kind: &str, value: &Value) -> Result<(), AttemptFailure> { + if self.sse_event_count >= self.limits.max_items { + return Err(protocol_failure( + "Responses attempt produced too many SSE events", + )); + } + self.sse_event_count += 1; + if self.terminal { + return Err(protocol_failure( + "Responses event followed a terminal event", + )); + } + if let Some(sequence) = value.get("sequence_number").and_then(Value::as_u64) { + let expected = self + .sequence + .map_or(sequence, |last| last.saturating_add(1)); + if sequence != expected { + return Err(protocol_failure( + "Responses SSE sequence is duplicate or out of order", + )); + } + self.sequence = Some(sequence); + } + + match kind { + "response.created" => self.created(value), + "response.in_progress" => { + self.require_created()?; + self.observe_response(value.get("response"), "response.in_progress") + } + "response.output_item.added" => self.output_item_added(value), + "response.content_part.added" => self.section_added(value, false), + "response.content_part.done" => self.section_done(value, false), + "response.reasoning_summary_part.added" => self.section_added(value, true), + "response.reasoning_summary_part.done" => self.section_done(value, true), + "response.output_text.delta" | "response.refusal.delta" => { + self.text_delta(value, false) + } + "response.reasoning_summary_text.delta" => self.text_delta(value, true), + "response.output_text.done" | "response.refusal.done" => self.text_done(value, false), + "response.reasoning_summary_text.done" => self.text_done(value, true), + "response.function_call_arguments.delta" => self.arguments_delta(value), + "response.function_call_arguments.done" => self.arguments_done(value), + "response.output_item.done" => self.output_item_done(value), + "response.completed" => self.complete(value, FinishReason::Completed), + "response.incomplete" => self.incomplete(value), + "response.failed" | "error" => { + self.require_created_if_response_event(kind)?; + if kind == "response.failed" { + if let Some(id) = value.pointer("/response/id") { + let id = bounded_id(Some(id))?; + if self.response_id.as_deref() != Some(id) { + return Err(protocol_failure( + "response.failed changed the response ID", + )); + } + } + if let Some(model) = value.pointer("/response/model") { + let model = bounded_id(Some(model))?; + if self + .response_model + .as_deref() + .is_some_and(|expected| expected != model) + { + return Err(protocol_failure( + "response.failed changed the response model", + )); + } + } + } + let code = value + .pointer("/response/error/code") + .or_else(|| value.pointer("/error/code")) + .or_else(|| value.get("code")) + .and_then(Value::as_str) + .unwrap_or("unknown"); + let retryable = stream_failure_retryable(self.profile, value, kind); + Err(AttemptFailure { + error: Box::new(LoopError::Provider(format!( + "OpenAI Responses stream failed ({code})" + ))), + retryable, + headers: None, + }) + } + "keepalive" => Ok(()), + "response.metadata" => { + self.require_created()?; + if let Some(response) = value.get("response") { + self.observe_response(Some(response), "response.metadata")?; + } + if let Some(captured) = responses_turn_state(value)? { + let mut turn_state = self + .turn_state + .lock() + .map_err(|_| protocol_failure("turn-state lock poisoned"))?; + if turn_state + .as_ref() + .is_some_and(|expected| expected != captured) + { + return Err(protocol_failure( + "response metadata changed x-codex-turn-state", + )); + } + *turn_state = Some(captured); + } + Ok(()) + } + _ if self.profile == OpenAIResponsesProfile::ChatGptPrivate => Ok(()), + _ => Err(protocol_failure("unsupported Responses SSE event kind")), + } + } + + fn created(&mut self, value: &Value) -> Result<(), AttemptFailure> { + if self.created { + return Err(protocol_failure("duplicate response.created")); + } + self.created = true; + self.observe_response(value.get("response"), "response.created")?; + if self.response_id.is_none() { + return Err(protocol_failure("response.created omitted response ID")); + } + Ok(()) + } + + fn require_created(&self) -> Result<(), AttemptFailure> { + if self.created { + Ok(()) + } else { + Err(protocol_failure( + "Responses event preceded response.created", + )) + } + } + + fn require_created_if_response_event(&self, kind: &str) -> Result<(), AttemptFailure> { + if kind == "response.failed" { + self.require_created() + } else { + Ok(()) + } + } + + fn observe_response( + &mut self, + response: Option<&Value>, + event: &str, + ) -> Result<(), AttemptFailure> { + let response = response + .and_then(Value::as_object) + .ok_or_else(|| protocol_failure("Responses lifecycle event omitted response"))?; + let id = bounded_id(response.get("id"))?; + if self + .response_id + .as_deref() + .is_some_and(|expected| expected != id) + { + return Err(protocol_failure( + "Responses lifecycle event changed response ID", + )); + } + self.response_id.get_or_insert_with(|| id.to_owned()); + if let Some(model) = response.get("model") { + let model = bounded_id(Some(model))?; + if self + .response_model + .as_deref() + .is_some_and(|expected| expected != model) + { + return Err(protocol_failure( + "Responses lifecycle event changed response model", + )); + } + self.response_model.get_or_insert_with(|| model.to_owned()); + } + let _ = event; + Ok(()) + } + + fn output_item_added(&mut self, value: &Value) -> Result<(), AttemptFailure> { + self.require_created()?; + let item = value + .get("item") + .and_then(Value::as_object) + .ok_or_else(|| protocol_failure("output_item.added omitted item"))?; + let id = bounded_id(item.get("id"))?; + let index = nonnegative(value, "output_index", self.limits.max_items)?; + let kind = item.get("type").and_then(Value::as_str); + if kind.is_some_and(|kind| { + !matches!( + kind, + "message" | "reasoning" | "function_call" | "image_generation_call" + ) + }) { + return Err(protocol_failure("unsupported Responses output item")); + } + if !self.seen_ids.insert(id.to_owned()) { + return Err(protocol_failure("duplicate output item ID")); + } + if !self.seen_indices.insert(index) { + return Err(protocol_failure("duplicate output item index")); + } + self.item_indices.insert(id.to_owned(), index); + if let Some(kind) = kind { + self.item_types.insert(id.to_owned(), kind.to_owned()); + } + Ok(()) + } + + fn event_item( + &self, + value: &Value, + index_field: &str, + ) -> Result<(String, u64), AttemptFailure> { + self.require_created()?; + let id = bounded_id(value.get("item_id"))?; + let output_index = nonnegative(value, "output_index", self.limits.max_items)?; + if self.item_indices.get(id).copied() != Some(output_index) { + return Err(protocol_failure( + "content event refers to an unknown or inconsistent output item", + )); + } + Ok(( + id.to_owned(), + nonnegative(value, index_field, self.limits.max_items)?, + )) + } + + fn section_added(&mut self, value: &Value, reasoning: bool) -> Result<(), AttemptFailure> { + let field = if reasoning { + "summary_index" + } else { + "content_index" + }; + let key = self.event_item(value, field)?; + let expected_type = if reasoning { + "summary_text" + } else { + "output_text" + }; + let part = value + .get("part") + .and_then(Value::as_object) + .ok_or_else(|| protocol_failure("content-part add omitted part"))?; + let part_type = part.get("type").and_then(Value::as_str); + let supported = + part_type == Some(expected_type) || (!reasoning && part_type == Some("refusal")); + if !supported { + return Err(protocol_failure("unsupported Responses content part")); + } + let set = if reasoning { + &mut self.summary_added + } else { + &mut self.content_added + }; + if !set.insert(key) { + return Err(protocol_failure("duplicate Responses content-part add")); + } + Ok(()) + } + + fn section_done(&mut self, value: &Value, reasoning: bool) -> Result<(), AttemptFailure> { + let field = if reasoning { + "summary_index" + } else { + "content_index" + }; + let key = self.event_item(value, field)?; + let expected_type = if reasoning { + "summary_text" + } else { + "output_text" + }; + let part = value + .get("part") + .and_then(Value::as_object) + .ok_or_else(|| protocol_failure("content-part done omitted part"))?; + let part_type = part.get("type").and_then(Value::as_str); + let supported = + part_type == Some(expected_type) || (!reasoning && part_type == Some("refusal")); + if !supported { + return Err(protocol_failure( + "unsupported completed Responses content part", + )); + } + let added = if reasoning { + &self.summary_added + } else { + &self.content_added + }; + let done = if reasoning { + &mut self.summary_done + } else { + &mut self.content_done + }; + if !added.contains(&key) || !done.insert(key) { + return Err(protocol_failure( + "Responses content part completed without add or twice", + )); + } + Ok(()) + } + + fn text_delta(&mut self, value: &Value, reasoning: bool) -> Result<(), AttemptFailure> { + let field = if reasoning { + "summary_index" + } else { + "content_index" + }; + let key = self.event_item(value, field)?; + let added = if reasoning { + &self.summary_added + } else { + &self.content_added + }; + if !added.contains(&key) { + return Err(protocol_failure( + "Responses text delta preceded content-part add", + )); + } + let delta = bounded_field(value, "delta", self.limits.max_text_bytes)?; + let target = if reasoning { + &mut self.reasoning + } else { + &mut self.text + }; + append_part( + target, + key, + delta, + if reasoning { + PartKind::Reasoning + } else { + PartKind::Text + }, + &mut self.events, + &mut self.event_count, + self.limits, + ) + } + + fn text_done(&mut self, value: &Value, reasoning: bool) -> Result<(), AttemptFailure> { + let field = if reasoning { + "summary_index" + } else { + "content_index" + }; + let key = self.event_item(value, field)?; + let added = if reasoning { + &self.summary_added + } else { + &self.content_added + }; + if !added.contains(&key) { + return Err(protocol_failure( + "Responses text done preceded content-part add", + )); + } + if let Some(completed) = value.get("text").or_else(|| value.get("refusal")) { + let completed = completed + .as_str() + .filter(|completed| completed.len() <= self.limits.max_text_bytes) + .ok_or_else(|| { + protocol_failure("Responses completed text is outside configured bounds") + })?; + let target = if reasoning { + &self.reasoning + } else { + &self.text + }; + if target + .get(&key) + .is_some_and(|streamed| streamed.text != completed) + { + return Err(protocol_failure( + "Responses text done changed streamed content", + )); + } + } + Ok(()) + } + + fn arguments_delta(&mut self, value: &Value) -> Result<(), AttemptFailure> { + let (id, _) = self.event_item(value, "output_index")?; + if self + .item_types + .get(&id) + .is_some_and(|kind| kind != "function_call") + { + return Err(protocol_failure( + "function arguments referred to a non-call item", + )); + } + append_bounded( + &mut self.function_arguments, + &id, + bounded_field(value, "delta", self.limits.max_text_bytes)?, + self.limits.max_text_bytes, + ) + } + + fn arguments_done(&mut self, value: &Value) -> Result<(), AttemptFailure> { + let id = bounded_id(value.get("item_id"))?.to_owned(); + let output_index = nonnegative(value, "output_index", self.limits.max_items)?; + if self.item_indices.get(&id).copied() != Some(output_index) + || !self.argument_done.insert(id.clone()) + { + return Err(protocol_failure( + "function arguments done is inconsistent or duplicate", + )); + } + if let Some(arguments) = value.get("arguments") { + let arguments = arguments + .as_str() + .filter(|arguments| arguments.len() <= self.limits.max_text_bytes) + .ok_or_else(|| { + protocol_failure("Responses function arguments are outside configured bounds") + })?; + if self + .function_arguments + .get(&id) + .is_some_and(|streamed| streamed != arguments) + { + return Err(protocol_failure( + "function arguments done changed streamed arguments", + )); + } + } + Ok(()) + } + + fn output_item_done(&mut self, value: &Value) -> Result<(), AttemptFailure> { + self.require_created()?; + let item = value + .get("item") + .ok_or_else(|| protocol_failure("output item omitted item"))?; + let id = bounded_id(item.get("id"))?; + if !self.seen_ids.contains(id) || !self.done_ids.insert(id.to_owned()) { + return Err(protocol_failure( + "output item completed without add or completed twice", + )); + } + let output_index = nonnegative(value, "output_index", self.limits.max_items)?; + if self.item_indices.get(id).copied() != Some(output_index) { + return Err(protocol_failure("completed output item index changed")); + } + if self.item_types.get(id).is_some_and(|expected| { + Some(expected.as_str()) != item.get("type").and_then(Value::as_str) + }) { + return Err(protocol_failure("completed output item type changed")); + } + self.output_item(output_index, item) + } + + fn output_item(&mut self, output_index: u64, item: &Value) -> Result<(), AttemptFailure> { + let item_id = bounded_id(item.get("id"))?; + match item.get("type").and_then(Value::as_str) { + Some("message") => { + if item.get("role").and_then(Value::as_str) != Some("assistant") { + return Err(protocol_failure("output message role is not assistant")); + } + let content = item + .get("content") + .and_then(Value::as_array) + .ok_or_else(|| protocol_failure("output message content is malformed"))?; + if content.len() > self.limits.max_items { + return Err(protocol_failure( + "output message has too many content parts", + )); + } + let mut parts = Vec::new(); + for (index, raw) in content.iter().enumerate() { + let key = (item_id.to_owned(), index as u64); + let content_type = raw.get("type").and_then(Value::as_str); + if !matches!(content_type, Some("output_text" | "refusal")) + || !self.content_added.contains(&key) + || !self.content_done.contains(&key) + { + return Err(protocol_failure( + "output message content lifecycle is incomplete", + )); + } + let text = if content_type == Some("refusal") { + bounded_field(raw, "refusal", self.limits.max_text_bytes)? + } else { + bounded_field(raw, "text", self.limits.max_text_bytes)? + }; + let streamed = self.text.remove(&key); + if streamed + .as_ref() + .is_some_and(|streamed| streamed.text != text) + { + return Err(protocol_failure( + "completed output changed streamed text content", + )); + } + let part = Part::Text(TextPart::new(text)); + if streamed.is_some() { + push_event( + &mut self.events, + &mut self.event_count, + ModelTurnEvent::Delta(Delta::CommitPart { part: part.clone() }), + self.limits.max_items, + )?; + } + parts.push(part); + } + self.output + .insert(output_index, provider_item(item_id, parts)); + } + Some("function_call") => { + let call_id = bounded_nonempty_field(item, "call_id", self.limits.max_text_bytes)?; + if !self.seen_call_ids.insert(call_id.to_owned()) { + return Err(protocol_failure("duplicate function call ID")); + } + let name = bounded_nonempty_field(item, "name", self.limits.max_text_bytes)?; + validate_tool_name(name) + .map_err(|error| protocol_failure(&format!("provider returned {error}")))?; + let arguments = bounded_field(item, "arguments", self.limits.max_text_bytes)?; + if let Some(streamed) = self.function_arguments.remove(item_id) + && (streamed != arguments || !self.argument_done.contains(item_id)) + { + return Err(protocol_failure( + "completed function arguments changed streamed arguments", + )); + } + let input: Value = serde_json::from_str(arguments) + .map_err(|_| protocol_failure("function-call arguments are not JSON"))?; + if !input.is_object() { + return Err(protocol_failure( + "function-call arguments are not an object", + )); + } + let response_id = self + .response_id + .as_deref() + .expect("created response has ID"); + let call = + ToolCallPart::new(call_id, name, input).with_metadata(continuation_metadata( + &self.requested_model, + &self.session_id, + self.authentication_binding.as_deref(), + response_id, + item_id, + output_index, + "function_call", + None, + )); + self.tool_call = true; + self.output.insert( + output_index, + provider_item(item_id, vec![Part::ToolCall(call)]), + ); + } + Some("image_generation_call") => { + let status = bounded_nonempty_field(item, "status", self.limits.max_text_bytes)?; + if status != "completed" { + return Err(protocol_failure("image generation did not complete")); + } + let result = item + .get("result") + .and_then(Value::as_str) + .filter(|result| { + !result.is_empty() && result.len() <= self.limits.max_text_bytes + }) + .ok_or_else(|| { + protocol_failure("generated image result is outside configured byte bounds") + })?; + let bytes = STANDARD + .decode(result) + .map_err(|_| protocol_failure("generated image result is not valid base64"))?; + if bytes.is_empty() || STANDARD.encode(&bytes) != result { + return Err(protocol_failure( + "generated image result is not canonical base64", + )); + } + let revised_prompt = item + .get("revised_prompt") + .filter(|value| !value.is_null()) + .map(|_| bounded_field(item, "revised_prompt", self.limits.max_text_bytes)) + .transpose()?; + let response_id = self + .response_id + .as_deref() + .expect("created response has ID"); + let mut metadata = continuation_metadata( + &self.requested_model, + &self.session_id, + self.authentication_binding.as_deref(), + response_id, + item_id, + output_index, + "image_generation_call", + None, + ); + metadata.insert( + GENERATED_IMAGE_METADATA.to_owned(), + json!({ + "item_id": item_id, + "status": status, + "revised_prompt": revised_prompt, + }), + ); + let media = + MediaPart::new(Modality::Image, "image/png", DataRef::InlineBytes(bytes)) + .with_metadata(metadata); + self.next_media += 1; + let placeholder = format!("[Image #{}]", self.next_media); + let placeholder_id = PartId::new(format!("generated-image-{output_index}")); + push_event( + &mut self.events, + &mut self.event_count, + ModelTurnEvent::Delta(Delta::BeginPart { + part_id: placeholder_id.clone(), + kind: PartKind::Text, + }), + self.limits.max_items, + )?; + push_event( + &mut self.events, + &mut self.event_count, + ModelTurnEvent::Delta(Delta::AppendText { + part_id: placeholder_id, + chunk: placeholder.clone(), + }), + self.limits.max_items, + )?; + push_event( + &mut self.events, + &mut self.event_count, + ModelTurnEvent::Delta(Delta::CommitPart { + part: Part::Text(TextPart::new(placeholder)), + }), + self.limits.max_items, + )?; + self.output.insert( + output_index, + provider_item(item_id, vec![Part::Media(media)]), + ); + } + Some("reasoning") => { + let summaries = item + .get("summary") + .and_then(Value::as_array) + .ok_or_else(|| protocol_failure("reasoning summary is malformed"))?; + if summaries.len() > self.limits.max_items { + return Err(protocol_failure( + "reasoning output has too many summary parts", + )); + } + let encrypted = item + .get("encrypted_content") + .and_then(Value::as_str) + .filter(|value| !value.is_empty() && value.len() <= self.limits.max_text_bytes); + if self.require_encrypted_reasoning && encrypted.is_none() { + return Err(protocol_failure( + "encrypted reasoning is missing or outside bounds", + )); + } + let mut summary_texts = Vec::new(); + for (index, raw) in summaries.iter().enumerate() { + let key = (item_id.to_owned(), index as u64); + if raw.get("type").and_then(Value::as_str) != Some("summary_text") + || !self.summary_added.contains(&key) + || !self.summary_done.contains(&key) + { + return Err(protocol_failure( + "reasoning summary lifecycle is incomplete", + )); + } + let text = bounded_field(raw, "text", self.limits.max_text_bytes)?; + let streamed = self.reasoning.remove(&key); + if streamed + .as_ref() + .is_some_and(|streamed| streamed.text != text) + { + return Err(protocol_failure( + "completed reasoning changed streamed summary", + )); + } + if streamed.is_some() { + push_event( + &mut self.events, + &mut self.event_count, + ModelTurnEvent::Delta(Delta::CommitPart { + part: Part::Reasoning(ReasoningPart::summary(text)), + }), + self.limits.max_items, + )?; + } + summary_texts.push(text); + } + let response_id = self + .response_id + .as_deref() + .expect("created response has ID"); + let metadata = encrypted.map_or_else(MetadataMap::new, |encrypted| { + continuation_metadata( + &self.requested_model, + &self.session_id, + self.authentication_binding.as_deref(), + response_id, + item_id, + output_index, + "reasoning", + Some(encrypted), + ) + }); + let part = Part::Reasoning(ReasoningPart { + summary: (!summary_texts.is_empty()).then(|| summary_texts.join("\n\n")), + data: None, + redacted: encrypted.is_some(), + metadata, + }); + self.output + .insert(output_index, provider_item(item_id, vec![part])); + } + _ => return Err(protocol_failure("unsupported Responses output item")), + } + Ok(()) + } + + fn complete(&mut self, value: &Value, default: FinishReason) -> Result<(), AttemptFailure> { + self.require_created()?; + self.observe_response(value.get("response"), "response.completed")?; + if self.seen_ids != self.done_ids + || self.content_added != self.content_done + || self.summary_added != self.summary_done + || !self.text.is_empty() + || !self.reasoning.is_empty() + { + return Err(protocol_failure( + "response.completed preceded complete output items", + )); + } + let response = value.get("response").expect("observed response"); + if let Some(raw) = response.get("usage") { + self.usage = Some(parse_usage(raw)); + } + self.provider_finish_reason = response + .get("status") + .and_then(Value::as_str) + .filter(|reason| !reason.is_empty() && reason.len() <= 128 && reason.is_ascii()) + .map(str::to_owned) + .or_else(|| Some("completed".into())); + self.finish_reason = Some(if self.tool_call { + FinishReason::ToolCall + } else { + default + }); + self.terminal = true; + self.queue_terminal()?; + Ok(()) + } + + fn incomplete(&mut self, value: &Value) -> Result<(), AttemptFailure> { + self.require_created()?; + self.observe_response(value.get("response"), "response.incomplete")?; + let response = value.get("response").expect("observed response"); + let reason = response + .pointer("/incomplete_details/reason") + .and_then(Value::as_str) + .filter(|reason| !reason.is_empty() && reason.len() <= 128 && reason.is_ascii()) + .ok_or_else(|| protocol_failure("response.incomplete omitted a valid reason"))?; + self.provider_finish_reason = Some(reason.to_owned()); + self.finish_reason = Some(match reason { + "max_output_tokens" => FinishReason::MaxTokens, + "content_filter" => FinishReason::Blocked, + other => FinishReason::Other(other.to_owned()), + }); + if let Some(raw) = response.get("usage") { + self.usage = Some(parse_usage(raw)); + } + self.flush_partial_text()?; + self.output + .retain(|_, item| item.parts.iter().all(|part| matches!(part, Part::Text(_)))); + self.reasoning.clear(); + self.function_arguments.clear(); + self.terminal = true; + self.queue_terminal()?; + Ok(()) + } + + fn flush_partial_text(&mut self) -> Result<(), AttemptFailure> { + let partial = std::mem::take(&mut self.text); + for ((item_id, _), accumulator) in partial { + let Some(index) = self.item_indices.get(&item_id).copied() else { + continue; + }; + if accumulator.text.is_empty() { + continue; + } + let part = Part::Text(TextPart::new(accumulator.text)); + push_event( + &mut self.events, + &mut self.event_count, + ModelTurnEvent::Delta(Delta::CommitPart { part: part.clone() }), + self.limits.max_items, + )?; + self.output + .insert(index, provider_item(&item_id, vec![part])); + } + Ok(()) + } + + fn queue_terminal(&mut self) -> Result<(), AttemptFailure> { + for call in self + .output + .values() + .flat_map(|item| item.parts.iter()) + .filter_map(|part| match part { + Part::ToolCall(call) => Some(call.clone()), + _ => None, + }) + { + push_event( + &mut self.events, + &mut self.event_count, + ModelTurnEvent::Delta(Delta::CommitPart { + part: Part::ToolCall(call.clone()), + }), + self.limits.max_items, + )?; + push_event( + &mut self.events, + &mut self.event_count, + ModelTurnEvent::ToolCall(call), + self.limits.max_items, + )?; + } + if let Some(usage) = self.usage.clone() { + push_event( + &mut self.events, + &mut self.event_count, + ModelTurnEvent::Usage(usage), + self.limits.max_items, + )?; + } + let mut metadata = MetadataMap::from([( + "openai.responses.profile".into(), + Value::String("responses".into()), + )]); + set_provider_finish_reasons(&mut metadata, self.provider_finish_reason.iter().cloned()); + push_event( + &mut self.events, + &mut self.event_count, + ModelTurnEvent::Finished(ModelTurnResult { + finish_reason: self.finish_reason.clone().expect("terminal reason"), + output_items: std::mem::take(&mut self.output).into_values().collect(), + usage: self.usage.clone(), + metadata, + model: self.response_model.clone(), + response_id: self.response_id.clone(), + }), + self.limits.max_items, + )?; + Ok(()) + } + + fn validate_finished(&self) -> Result<(), AttemptFailure> { + if !self.created { + return Err(protocol_failure( + "Responses SSE closed before response.created", + )); + } + if !self.terminal { + return Err(AttemptFailure { + error: Box::new(LoopError::Provider( + "OpenAI Responses SSE stream closed before a terminal event".into(), + )), + retryable: true, + headers: None, + }); + } + if self.events.len() > self.limits.max_items { + return Err(protocol_failure( + "Responses attempt produced too many events", + )); + } + Ok(()) + } + + #[cfg(test)] + fn finish(mut self) -> Result, AttemptFailure> { + self.validate_finished()?; + Ok(std::mem::take(&mut self.events)) + } +} + +fn provider_item(item_id: &str, parts: Vec) -> Item { + let mut item = Item::new(ItemKind::Assistant, parts); + item.id = Some(MessageId::new(item_id)); + item +} + +fn push_event( + events: &mut VecDeque, + event_count: &mut usize, + event: ModelTurnEvent, + max_items: usize, +) -> Result<(), AttemptFailure> { + if *event_count >= max_items { + return Err(protocol_failure( + "Responses attempt produced too many events", + )); + } + *event_count += 1; + events.push_back(event); + Ok(()) +} + +fn append_part( + parts: &mut BTreeMap<(String, u64), PartAccumulator>, + key: (String, u64), + delta: &str, + kind: PartKind, + events: &mut VecDeque, + event_count: &mut usize, + limits: OpenAIResponsesLimits, +) -> Result<(), AttemptFailure> { + if !parts.contains_key(&key) { + let id = PartId::new(format!( + "openai-responses:{}:{}:{:?}:{}", + key.0.len(), + key.0, + kind, + key.1 + )); + push_event( + events, + event_count, + ModelTurnEvent::Delta(Delta::BeginPart { + part_id: id.clone(), + kind, + }), + limits.max_items, + )?; + parts.insert( + key.clone(), + PartAccumulator { + id, + text: String::new(), + }, + ); + } + let part = parts.get_mut(&key).expect("inserted part accumulator"); + if part.text.len().saturating_add(delta.len()) > limits.max_text_bytes { + return Err(protocol_failure( + "Responses streamed content exceeds configured text bounds", + )); + } + part.text.push_str(delta); + push_event( + events, + event_count, + ModelTurnEvent::Delta(Delta::AppendText { + part_id: part.id.clone(), + chunk: delta.to_owned(), + }), + limits.max_items, + ) +} + +fn nonnegative(value: &Value, field: &str, max_items: usize) -> Result { + value + .get(field) + .and_then(Value::as_u64) + .filter(|value| *value < max_items as u64) + .ok_or_else(|| protocol_failure("Responses index is missing or outside bounds")) +} + +fn bounded_id(value: Option<&Value>) -> Result<&str, AttemptFailure> { + value + .and_then(Value::as_str) + .filter(|value| !value.is_empty() && value.len() <= 512) + .ok_or_else(|| protocol_failure("Responses ID is missing or outside bounds")) +} + +fn bounded_nonempty_field<'a>( + value: &'a Value, + field: &str, + max_text_bytes: usize, +) -> Result<&'a str, AttemptFailure> { + bounded_field(value, field, max_text_bytes).and_then(|value| { + if value.is_empty() { + Err(protocol_failure("Responses field is empty")) + } else { + Ok(value) + } + }) +} + +fn zeroizing_extend( + buffer: &mut Zeroizing>, + bytes: &[u8], + limit: usize, +) -> Result<(), AttemptFailure> { + let new_len = buffer + .len() + .checked_add(bytes.len()) + .filter(|length| *length <= limit) + .ok_or_else(|| protocol_failure("Responses SSE buffer exceeds canonical bounds"))?; + if new_len > buffer.capacity() { + let capacity = buffer.capacity().saturating_mul(2).max(new_len).min(limit); + let mut replacement = Vec::with_capacity(capacity); + replacement.extend_from_slice(buffer); + let mut previous = std::mem::replace(&mut **buffer, replacement); + previous.zeroize(); + } + buffer.extend_from_slice(bytes); + Ok(()) +} + +fn zeroize_encrypted_content(value: &mut Value) { + match value { + Value::Object(object) => { + for (key, value) in object { + if key == "encrypted_content" { + if let Value::String(secret) = value { + secret.zeroize(); + } + } else { + zeroize_encrypted_content(value); + } + } + } + Value::Array(values) => values.iter_mut().for_each(zeroize_encrypted_content), + _ => {} + } +} + +fn responses_turn_state(value: &Value) -> Result, AttemptFailure> { + let mut state = None; + for headers in [value.pointer("/response/headers"), value.get("headers")] + .into_iter() + .flatten() + .filter_map(Value::as_object) + { + for raw in headers.iter().filter_map(|(name, value)| { + name.eq_ignore_ascii_case(X_CODEX_TURN_STATE) + .then_some(value) + }) { + let values: Vec<&str> = match raw { + Value::String(value) => vec![value], + Value::Array(values) if !values.is_empty() => values + .iter() + .map(Value::as_str) + .collect::>>() + .ok_or_else(|| protocol_failure("response metadata turn state is invalid"))?, + _ => { + return Err(protocol_failure("response metadata turn state is invalid")); + } + }; + for value in values { + let value = HeaderValue::from_str(value) + .map_err(|_| protocol_failure("response metadata turn state is invalid"))?; + if state.as_ref().is_some_and(|expected| expected != value) { + return Err(protocol_failure( + "response metadata changed x-codex-turn-state", + )); + } + state = Some(value); + } + } + } + Ok(state) +} + +fn retry_headers(headers: &HeaderMap) -> Option { + let mut retained = HeaderMap::new(); + for name in [ + "retry-after", + "ratelimit-reset", + "x-ratelimit-reset", + "x-rate-limit-reset", + ] { + if let Some(value) = headers.get(name) { + retained.insert(agentkit_http::HeaderName::from_static(name), value.clone()); + } + } + (!retained.is_empty()).then_some(retained) +} + +fn validated_turn_state_header(headers: &HeaderMap) -> Result, AttemptFailure> { + let mut state: Option = None; + for value in headers.get_all(X_CODEX_TURN_STATE) { + if state.as_ref().is_some_and(|expected| expected != value) { + return Err(protocol_failure( + "provider returned conflicting x-codex-turn-state headers", + )); + } + state = Some(value.clone()); + } + Ok(state) +} + +fn parse_usage(value: &Value) -> Usage { + let input = value + .get("input_tokens") + .and_then(Value::as_u64) + .unwrap_or(0); + let output = value + .get("output_tokens") + .and_then(Value::as_u64) + .unwrap_or(0); + let mut tokens = TokenUsage::new(input, output); + if let Some(reasoning) = value + .pointer("/output_tokens_details/reasoning_tokens") + .and_then(Value::as_u64) + { + tokens = tokens.with_reasoning_tokens(reasoning); + } + if let Some(cached) = value + .pointer("/input_tokens_details/cached_tokens") + .and_then(Value::as_u64) + { + tokens = tokens.with_cached_input_tokens(cached); + } + Usage::new(tokens) +} + +fn bounded_field<'a>( + value: &'a Value, + field: &str, + max_text_bytes: usize, +) -> Result<&'a str, AttemptFailure> { + value + .get(field) + .and_then(Value::as_str) + .filter(|value| value.len() <= max_text_bytes) + .ok_or_else(|| protocol_failure("Responses text field is missing or too large")) +} + +fn append_bounded( + target: &mut BTreeMap, + id: &str, + delta: &str, + max_text_bytes: usize, +) -> Result<(), AttemptFailure> { + let value = target.entry(id.to_owned()).or_default(); + if value.len().saturating_add(delta.len()) > max_text_bytes { + return Err(protocol_failure( + "Responses output exceeds configured text bounds", + )); + } + value.push_str(delta); + Ok(()) +} + +fn frame_end(buffer: &[u8]) -> Option<(usize, usize)> { + [ + buffer + .windows(4) + .position(|value| value == b"\r\n\r\n") + .map(|index| (index, 4)), + buffer + .windows(2) + .position(|value| value == b"\n\n") + .map(|index| (index, 2)), + buffer + .windows(2) + .position(|value| value == b"\r\r") + .map(|index| (index, 2)), + ] + .into_iter() + .flatten() + .min_by_key(|(index, _)| *index) +} + +fn stable_idempotency_key(session: &str, turn: &str, body: &[u8]) -> String { + let mut left = 0xcbf29ce484222325_u64; + let mut right = 0x9e3779b97f4a7c15_u64; + for byte in session + .bytes() + .chain([0]) + .chain(turn.bytes()) + .chain([0]) + .chain(body.iter().copied()) + { + left ^= u64::from(byte); + left = left.wrapping_mul(0x100000001b3); + right ^= left.rotate_left(17).wrapping_add(u64::from(byte)); + right = right.wrapping_mul(0xff51afd7ed558ccd); + } + format!("agentkit-{left:016x}{right:016x}") +} + +fn transport_failure(error: HttpError) -> AttemptFailure { + let retryable = error.is_retryable_transport(); + AttemptFailure { + error: Box::new(LoopError::Provider(format!( + "OpenAI Responses transport failed: {error}" + ))), + retryable, + headers: None, + } +} + +fn protocol_failure(message: &str) -> AttemptFailure { + nonretryable(LoopError::Provider(message.into())) +} + +fn stream_failure_retryable(profile: OpenAIResponsesProfile, value: &Value, kind: &str) -> bool { + let error = if kind == "response.failed" { + value.pointer("/response/error").unwrap_or(&Value::Null) + } else { + value.get("error").unwrap_or(value) + }; + let code = error + .get("code") + .or_else(|| value.get("code")) + .and_then(Value::as_str) + .unwrap_or("unknown"); + if profile == OpenAIResponsesProfile::Public { + return matches!( + code, + "server_error" | "rate_limit_exceeded" | "temporarily_unavailable" + ); + } + let error_type = error + .get("type") + .and_then(Value::as_str) + .unwrap_or("unknown"); + let status = error + .get("status") + .or_else(|| value.get("status")) + .or_else(|| value.pointer("/response/status_code")) + .and_then(Value::as_u64); + let authentication = status.is_some_and(|status| status == 401 || status == 403) + || [code, error_type].iter().any(|value| { + matches!( + *value, + "authentication_error" + | "invalid_api_key" + | "invalid_authentication" + | "unauthorized" + ) + }); + let permanent = status.is_some_and(|status| { + matches!( + status, + 400 | 402 | 403 | 404 | 405 | 406 | 410 | 413 | 415 | 422 | 501 | 505 + ) + }) || [code, error_type].iter().any(|value| { + [ + "billing", + "content_policy", + "deactivated", + "insufficient", + "invalid", + "not_found", + "not_supported", + "permission", + "quota", + "unsupported", + ] + .iter() + .any(|marker| value.contains(marker)) + }); + !authentication + && !permanent + && status + .is_none_or(|status| matches!(status, 408 | 425 | 429 | 500 | 502 | 503 | 504 | 529)) +} + +fn nonretryable(error: LoopError) -> AttemptFailure { + AttemptFailure { + error: Box::new(error), + retryable: false, + headers: None, + } +} + +fn is_unauthorized(error: &LoopError) -> bool { + matches!(error, LoopError::Provider(message) if message.contains("401 Unauthorized")) +} + +fn cancelled(cancellation: Option<&TurnCancellation>) -> bool { + cancellation.is_some_and(TurnCancellation::is_cancelled) +} + +async fn cancellable( + future: F, + cancellation: Option<&TurnCancellation>, +) -> Result +where + F: Future, +{ + let Some(cancellation) = cancellation else { + return Ok(future.await); + }; + if cancellation.is_cancelled() { + return Err(LoopError::Cancelled); + } + futures_util::pin_mut!(future); + let cancelled = cancellation.cancelled(); + futures_util::pin_mut!(cancelled); + match select(future, cancelled).await { + Either::Left((result, _)) => Ok(result), + Either::Right((_, _)) => Err(LoopError::Cancelled), + } +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use agentkit_core::{Item, SessionId, TurnId}; + use agentkit_http::{HeaderValue, HttpClient, HttpRequest, HttpResponse, StatusCode, header}; + use async_trait::async_trait; + use futures_util::stream; + + use super::*; + + #[derive(Clone)] + struct WireResponse { + status: StatusCode, + headers: HeaderMap, + body: &'static str, + } + + struct ScriptedClient { + responses: Mutex>, + requests: Mutex>, + } + + #[async_trait] + impl HttpClient for ScriptedClient { + async fn execute(&self, request: HttpRequest) -> Result { + self.requests.lock().unwrap().push(request.clone()); + let response = self.responses.lock().unwrap().pop_front().unwrap(); + let body = agentkit_http::Bytes::from_static(response.body.as_bytes()); + Ok(HttpResponse::new( + response.status, + response.headers, + request.url, + Box::pin(stream::once(async move { Ok(body) })), + )) + } + } + + struct ActiveStreamClient { + requests: AtomicUsize, + } + + #[async_trait] + impl HttpClient for ActiveStreamClient { + async fn execute(&self, request: HttpRequest) -> Result { + self.requests.fetch_add(1, Ordering::SeqCst); + let body = stream::unfold((), |_| async { + sleep(Duration::from_millis(1)).await; + Some(( + Ok(agentkit_http::Bytes::from_static(b": keepalive\n\n")), + (), + )) + }); + Ok(HttpResponse::new( + StatusCode::OK, + sse_headers(), + request.url, + Box::pin(body), + )) + } + } + + struct RefreshingAuth { + calls: Arc, + } + + #[async_trait] + impl agentkit_http::AuthenticationProvider for RefreshingAuth { + async fn authenticate( + &self, + previous: Option<&AuthenticationAttempt>, + ) -> Result { + let call = self.calls.fetch_add(1, Ordering::SeqCst); + if call == 0 { + assert!(previous.is_none()); + } else { + assert_eq!(previous.and_then(|value| value.state::()), Some(&0)); + } + let mut headers = HeaderMap::new(); + headers.insert( + header::AUTHORIZATION, + HeaderValue::from_str(if call == 0 { + "Bearer first" + } else { + "Bearer refreshed" + }) + .unwrap(), + ); + Ok(AuthenticationAttempt::new(headers, call)) + } + } + + struct BindingChangingAuth { + calls: Arc, + } + + #[async_trait] + impl AuthenticationProvider for BindingChangingAuth { + async fn authenticate( + &self, + _previous: Option<&AuthenticationAttempt>, + ) -> Result { + let call = self.calls.fetch_add(1, Ordering::SeqCst); + Ok(AuthenticationAttempt::stateless(HeaderMap::new()) + .with_binding(format!("binding-{call}"))) + } + } + + fn request() -> TurnRequest { + TurnRequest { + session_id: SessionId::new("session"), + turn_id: TurnId::new("turn"), + transcript: vec![Item::text(ItemKind::User, "hello")], + available_tools: Vec::new(), + cache: None, + metadata: MetadataMap::new(), + } + } + + fn sse_headers() -> HeaderMap { + let mut headers = HeaderMap::new(); + headers.insert( + header::CONTENT_TYPE, + HeaderValue::from_static("text/event-stream"), + ); + headers + } + + const SUCCESS: &str = r#"event: response.created +data: {"type":"response.created","sequence_number":1,"response":{"id":"resp-1","model":"gpt-test"}} + +event: response.output_item.added +data: {"type":"response.output_item.added","sequence_number":2,"output_index":0,"item":{"id":"msg-1","type":"message"}} + +event: response.content_part.added +data: {"type":"response.content_part.added","sequence_number":3,"item_id":"msg-1","output_index":0,"content_index":0,"part":{"type":"output_text"}} + +event: response.output_text.delta +data: {"type":"response.output_text.delta","sequence_number":4,"item_id":"msg-1","output_index":0,"content_index":0,"delta":"hello"} + +event: response.output_text.done +data: {"type":"response.output_text.done","sequence_number":5,"item_id":"msg-1","output_index":0,"content_index":0,"text":"hello"} + +event: response.content_part.done +data: {"type":"response.content_part.done","sequence_number":6,"item_id":"msg-1","output_index":0,"content_index":0,"part":{"type":"output_text","text":"hello"}} + +event: response.output_item.done +data: {"type":"response.output_item.done","sequence_number":7,"output_index":0,"item":{"id":"msg-1","type":"message","role":"assistant","content":[{"type":"output_text","text":"hello"}]}} + +event: response.output_item.added +data: {"type":"response.output_item.added","sequence_number":8,"output_index":1,"item":{"id":"reason-1","type":"reasoning"}} + +event: response.reasoning_summary_part.added +data: {"type":"response.reasoning_summary_part.added","sequence_number":9,"item_id":"reason-1","output_index":1,"summary_index":0,"part":{"type":"summary_text"}} + +event: response.reasoning_summary_text.delta +data: {"type":"response.reasoning_summary_text.delta","sequence_number":10,"item_id":"reason-1","output_index":1,"summary_index":0,"delta":"brief"} + +event: response.reasoning_summary_text.done +data: {"type":"response.reasoning_summary_text.done","sequence_number":11,"item_id":"reason-1","output_index":1,"summary_index":0,"text":"brief"} + +event: response.reasoning_summary_part.done +data: {"type":"response.reasoning_summary_part.done","sequence_number":12,"item_id":"reason-1","output_index":1,"summary_index":0,"part":{"type":"summary_text","text":"brief"}} + +event: response.output_item.done +data: {"type":"response.output_item.done","sequence_number":13,"output_index":1,"item":{"id":"reason-1","type":"reasoning","summary":[{"type":"summary_text","text":"brief"}],"encrypted_content":"opaque"}} + +event: response.output_item.added +data: {"type":"response.output_item.added","sequence_number":14,"output_index":2,"item":{"id":"call-item","type":"function_call"}} + +event: response.function_call_arguments.delta +data: {"type":"response.function_call_arguments.delta","sequence_number":15,"item_id":"call-item","output_index":2,"delta":"{\"q\":1}"} + +event: response.function_call_arguments.done +data: {"type":"response.function_call_arguments.done","sequence_number":16,"item_id":"call-item","output_index":2,"arguments":"{\"q\":1}"} + +event: response.output_item.done +data: {"type":"response.output_item.done","sequence_number":17,"output_index":2,"item":{"id":"call-item","type":"function_call","call_id":"call-1","name":"lookup","arguments":"{\"q\":1}"}} + +event: response.completed +data: {"type":"response.completed","sequence_number":18,"response":{"id":"resp-1","model":"gpt-test","usage":{"input_tokens":3,"output_tokens":5,"output_tokens_details":{"reasoning_tokens":2}}}} + +"#; + + #[test] + fn public_and_private_profiles_encode_distinct_fields() { + let mut request = request(); + request.transcript[0] = Item::text(ItemKind::System, "rules"); + + let public = OpenAIResponsesConfig::new("secret", "gpt-test") + .with_max_output_tokens(99) + .encode_request(&request) + .unwrap(); + assert_eq!(public.pointer("/input/0/role"), Some(&json!("system"))); + assert_eq!(public["max_output_tokens"], 99); + assert_eq!(public["include"], json!(["reasoning.encrypted_content"])); + + let private = + OpenAIResponsesConfig::chatgpt_private("gpt-test", Authentication::bearer("secret")) + .with_max_output_tokens(99) + .encode_request(&request) + .unwrap(); + assert_eq!(private.pointer("/input/0/role"), Some(&json!("developer"))); + assert!(private.get("max_output_tokens").is_none()); + assert_eq!(private["include"], json!(["reasoning.encrypted_content"])); + assert_eq!(private["parallel_tool_calls"], true); + } + + #[test] + fn context_content_is_unchanged_and_honors_role_downgrade_policy() { + let mut request = request(); + request.transcript[0] = Item::text(ItemKind::Context, "project facts"); + + let public = OpenAIResponsesConfig::new("secret", "gpt-test") + .encode_request(&request) + .unwrap(); + let private = + OpenAIResponsesConfig::chatgpt_private("gpt-test", Authentication::bearer("secret")) + .encode_request(&request) + .unwrap(); + + assert_eq!(public.pointer("/input/0/role"), Some(&json!("system"))); + assert_eq!(private.pointer("/input/0/role"), Some(&json!("developer"))); + assert_eq!( + public.pointer("/input/0/content/0/text"), + Some(&json!("project facts")) + ); + assert_eq!( + private.pointer("/input/0/content/0/text"), + Some(&json!("project facts")) + ); + } + + #[tokio::test] + async fn refreshes_once_and_replays_stable_body_and_idempotency_key_before_output() { + let calls = Arc::new(AtomicUsize::new(0)); + let client = Arc::new(ScriptedClient { + responses: Mutex::new(VecDeque::from([ + WireResponse { + status: StatusCode::UNAUTHORIZED, + headers: HeaderMap::new(), + body: "", + }, + WireResponse { + status: StatusCode::OK, + headers: sse_headers(), + body: "event: response.created\ndata: {\"type\":\"response.created\",\"sequence_number\":1,\"response\":{\"id\":\"discarded-response\",\"model\":\"gpt-test\"}}\n\n", + }, + WireResponse { + status: StatusCode::OK, + headers: sse_headers(), + body: SUCCESS, + }, + ])), + requests: Mutex::new(Vec::new()), + }); + let config = OpenAIResponsesConfig::public( + "gpt-test", + Authentication::new(RefreshingAuth { + calls: calls.clone(), + }), + ) + .with_resilience(ResilienceConfig { + max_retries: 1, + retry_budget: Duration::from_secs(1), + attempt_timeout: None, + stream_idle_timeout: None, + initial_backoff: Duration::ZERO, + max_backoff: Duration::ZERO, + }); + let adapter = OpenAIResponsesAdapter::with_client(config, Http::from_arc(client.clone())); + let mut session = adapter + .start_session(SessionConfig::new("session")) + .await + .unwrap(); + let mut turn = session.begin_turn(request(), None).await.unwrap(); + let mut events = Vec::new(); + while let Some(event) = turn.next_event(None).await.unwrap() { + events.push(event); + } + + assert_eq!(calls.load(Ordering::SeqCst), 2); + let requests = client.requests.lock().unwrap(); + assert_eq!(requests.len(), 3); + assert_eq!(requests[0].body, requests[1].body); + assert_eq!(requests[1].body, requests[2].body); + let keys: Vec<_> = requests + .iter() + .map(|value| value.headers["idempotency-key"].clone()) + .collect(); + assert!(keys.windows(2).all(|pair| pair[0] == pair[1])); + assert_eq!(requests[0].headers[header::AUTHORIZATION], "Bearer first"); + assert_eq!( + requests[1].headers[header::AUTHORIZATION], + "Bearer refreshed" + ); + assert_eq!( + requests[2].headers[header::AUTHORIZATION], + "Bearer refreshed" + ); + assert!(events.iter().any(|event| matches!( + event, + ModelTurnEvent::Delta(Delta::AppendText { chunk, .. }) if chunk == "hello" + ))); + assert!(!events.iter().any(|event| matches!( + event, + ModelTurnEvent::Delta(Delta::AppendText { chunk, .. }) if chunk == "discarded" + ))); + assert!( + events.iter().any( + |event| matches!(event, ModelTurnEvent::ToolCall(call) if call.name == "lookup") + ) + ); + assert!(events.iter().any(|event| matches!( + event, + ModelTurnEvent::Delta(Delta::CommitPart { part: Part::ToolCall(call) }) + if call.name == "lookup" + ))); + assert!(events.iter().any(|event| matches!( + event, + ModelTurnEvent::Usage(usage) + if usage.tokens.as_ref().is_some_and(|tokens| tokens.reasoning_tokens == Some(2)) + ))); + assert!( + matches!(events.last(), Some(ModelTurnEvent::Finished(result)) if result.finish_reason == FinishReason::ToolCall && result.response_id.as_deref() == Some("resp-1")) + ); + } + + #[tokio::test] + async fn reactive_refresh_rejects_changed_authentication_binding() { + let calls = Arc::new(AtomicUsize::new(0)); + let client = Arc::new(ScriptedClient { + responses: Mutex::new(VecDeque::from([WireResponse { + status: StatusCode::UNAUTHORIZED, + headers: HeaderMap::new(), + body: "", + }])), + requests: Mutex::new(Vec::new()), + }); + let config = OpenAIResponsesConfig::public( + "gpt-test", + Authentication::new(BindingChangingAuth { + calls: calls.clone(), + }), + ) + .with_endpoint("https://example.test/responses"); + let adapter = OpenAIResponsesAdapter::with_client(config, Http::from_arc(client.clone())); + let mut session = adapter + .start_session(SessionConfig::new("session")) + .await + .unwrap(); + let error = match session.begin_turn(request(), None).await { + Ok(_) => panic!("changed authentication binding unexpectedly succeeded"), + Err(error) => error, + }; + assert!(error.to_string().contains("authentication binding changed")); + assert_eq!(calls.load(Ordering::SeqCst), 2); + assert_eq!(client.requests.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn retries_retryable_status_and_honors_none_as_no_retry() { + for (resilience, expected_requests, succeeds) in [ + ( + Some(ResilienceConfig { + max_retries: 1, + retry_budget: Duration::from_secs(1), + attempt_timeout: None, + stream_idle_timeout: None, + initial_backoff: Duration::ZERO, + max_backoff: Duration::ZERO, + }), + 2, + true, + ), + (None, 1, false), + ] { + let client = Arc::new(ScriptedClient { + responses: Mutex::new(VecDeque::from([ + WireResponse { + status: StatusCode::SERVICE_UNAVAILABLE, + headers: HeaderMap::new(), + body: "", + }, + WireResponse { + status: StatusCode::OK, + headers: sse_headers(), + body: SUCCESS, + }, + ])), + requests: Mutex::new(Vec::new()), + }); + let mut config = OpenAIResponsesConfig::new("secret", "gpt-test"); + config.resilience = resilience; + let adapter = + OpenAIResponsesAdapter::with_client(config, Http::from_arc(client.clone())); + let mut session = adapter + .start_session(SessionConfig::new("session")) + .await + .unwrap(); + let result = session.begin_turn(request(), None).await; + assert_eq!(result.is_ok(), succeeds); + assert_eq!(client.requests.lock().unwrap().len(), expected_requests); + } + } + + #[test] + fn strict_lifecycle_rejects_invalid_order_sequence_and_terminal_events() { + let before_create = b"event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"sequence_number\":1,\"output_index\":0,\"item\":{\"id\":\"x\",\"type\":\"message\"}}\n\n"; + let mut decoder = ResponsesSseDecoder::new("gpt-test", "session"); + assert!(decoder.push(before_create).is_err()); + + let duplicate_sequence = concat!( + "event: response.created\ndata: {\"type\":\"response.created\",\"sequence_number\":1,\"response\":{\"id\":\"r\",\"model\":\"gpt-test\"}}\n\n", + "event: response.in_progress\ndata: {\"type\":\"response.in_progress\",\"sequence_number\":1,\"response\":{\"id\":\"r\",\"model\":\"gpt-test\"}}\n\n", + ); + let mut decoder = ResponsesSseDecoder::new("gpt-test", "session"); + assert!(decoder.push(duplicate_sequence.as_bytes()).is_err()); + + let mut decoder = ResponsesSseDecoder::new("gpt-test", "session"); + decoder.push(SUCCESS.as_bytes()).unwrap(); + decoder.process_all_pending().unwrap(); + while decoder.pop_event().is_some() {} + let trailing = b"event: response.in_progress\ndata: {\"type\":\"response.in_progress\",\"sequence_number\":19,\"response\":{\"id\":\"resp-1\",\"model\":\"gpt-test\"}}\n\n"; + assert!(decoder.push(trailing).is_err()); + } + + #[test] + fn incomplete_is_an_explicit_terminal_with_normalized_finish_reason() { + let wire = concat!( + "event: response.created\ndata: {\"type\":\"response.created\",\"sequence_number\":1,\"response\":{\"id\":\"r\",\"model\":\"gpt-test\"}}\n\n", + "event: response.incomplete\ndata: {\"type\":\"response.incomplete\",\"sequence_number\":2,\"response\":{\"id\":\"r\",\"model\":\"gpt-test\",\"incomplete_details\":{\"reason\":\"max_output_tokens\"}}}\n\n", + ); + let mut decoder = ResponsesSseDecoder::new("gpt-test", "session"); + decoder.push(wire.as_bytes()).unwrap(); + let events = decoder.finish().unwrap(); + assert!( + matches!(events.back(), Some(ModelTurnEvent::Finished(result)) if result.finish_reason == FinishReason::MaxTokens) + ); + + let unknown = concat!( + "event: response.created\ndata: {\"type\":\"response.created\",\"sequence_number\":1,\"response\":{\"id\":\"r2\",\"model\":\"gpt-test\"}}\n\n", + "event: response.incomplete\ndata: {\"type\":\"response.incomplete\",\"sequence_number\":2,\"response\":{\"id\":\"r2\",\"model\":\"gpt-test\",\"incomplete_details\":{\"reason\":\"provider_new_reason\"}}}\n\n", + ); + let mut decoder = ResponsesSseDecoder::new("gpt-test", "session"); + decoder.push(unknown.as_bytes()).unwrap(); + let events = decoder.finish().unwrap(); + let Some(ModelTurnEvent::Finished(result)) = events.back() else { + panic!("incomplete response must finish"); + }; + assert_eq!( + result.finish_reason, + FinishReason::Other("provider_new_reason".into()) + ); + assert_eq!( + result.metadata[agentkit_loop::PROVIDER_FINISH_REASONS_METADATA_KEY], + json!(["provider_new_reason"]) + ); + } + + #[test] + fn continuation_is_versioned_bound_and_preserves_function_item_id() { + let mut decoder = ResponsesSseDecoder::new("gpt-test", "session"); + decoder.push(SUCCESS.as_bytes()).unwrap(); + let events = decoder.finish().unwrap(); + let result = events + .iter() + .find_map(|event| match event { + ModelTurnEvent::Finished(result) => Some(result), + _ => None, + }) + .unwrap(); + assert!(result.output_items.iter().all(|item| item.id.is_some())); + let mut replay = request(); + replay.transcript = result.output_items.clone(); + let config = + OpenAIResponsesConfig::chatgpt_private("gpt-test", Authentication::bearer("secret")); + assert!(matches!( + config.encode_request(&replay), + Err(OpenAIResponsesError::InvalidRequest(_)) + )); + let encoded = + encode_request_bound(&config, &replay, Some("test-authentication-binding")).unwrap(); + let input = encoded["input"].as_array().unwrap(); + assert!( + input + .iter() + .any(|item| item.get("type") == Some(&json!("function_call")) + && item.get("id") == Some(&json!("call-item"))) + ); + assert!( + input + .iter() + .any(|item| item.get("type") == Some(&json!("reasoning")) + && item.get("id") == Some(&json!("reason-1")) + && item.get("encrypted_content") == Some(&json!("opaque"))) + ); + let metadata = result + .output_items + .iter() + .flat_map(|item| &item.parts) + .find_map(|part| match part { + Part::Reasoning(reasoning) => reasoning.metadata.get(CONTINUATION_METADATA), + _ => None, + }) + .unwrap(); + assert_eq!(metadata["schema_version"], CONTINUATION_SCHEMA_VERSION); + assert_eq!( + metadata["authentication_binding"], + "test-authentication-binding" + ); + assert_eq!(metadata["model"], "gpt-test"); + assert_eq!(metadata["session_id"], "session"); + assert!(metadata.get("response_id").is_none()); + assert!(metadata.get("output_index").is_none()); + + let mismatched = encode_request_bound(&config, &replay, Some("other-binding")).unwrap(); + let mismatched = mismatched["input"].as_array().unwrap(); + assert!(mismatched.iter().all(|item| { + item.get("type") != Some(&json!("reasoning")) + && item.get("id") != Some(&json!("call-item")) + })); + + let mut malformed = replay.clone(); + let reasoning = malformed + .transcript + .iter_mut() + .flat_map(|item| &mut item.parts) + .find_map(|part| match part { + Part::Reasoning(reasoning) => Some(reasoning), + _ => None, + }) + .unwrap(); + reasoning + .metadata + .insert(CONTINUATION_METADATA.into(), json!("malformed")); + assert!(matches!( + encode_request_bound(&config, &malformed, Some("test-authentication-binding")), + Err(OpenAIResponsesError::Protocol(_)) + )); + } + + #[test] + fn continuation_metadata_is_not_emitted_without_authentication_binding() { + let mut decoder = ResponsesSseDecoder::with_policy( + "gpt-test", + "session", + OpenAIResponsesProfile::Public, + true, + None, + Arc::new(Mutex::new(None)), + OpenAIResponsesLimits::default(), + ); + decoder.push(SUCCESS.as_bytes()).unwrap(); + let events = decoder.finish().unwrap(); + let result = events + .iter() + .find_map(|event| match event { + ModelTurnEvent::Finished(result) => Some(result), + _ => None, + }) + .unwrap(); + assert!( + result + .output_items + .iter() + .flat_map(|item| &item.parts) + .all(|part| { + match part { + Part::Reasoning(reasoning) => reasoning.metadata.is_empty(), + Part::ToolCall(call) => call.metadata.is_empty(), + _ => true, + } + }) + ); + } + + #[test] + fn public_profile_keeps_unencrypted_reasoning_as_reasoning() { + let wire = SUCCESS.replace(",\"encrypted_content\":\"opaque\"", ""); + let mut decoder = ResponsesSseDecoder::with_policy( + "gpt-test", + "session", + OpenAIResponsesProfile::Public, + false, + Some("test-authentication-binding"), + Arc::new(Mutex::new(None)), + OpenAIResponsesLimits::default(), + ); + decoder.push(wire.as_bytes()).unwrap(); + let events = decoder.finish().unwrap(); + let result = events + .iter() + .find_map(|event| match event { + ModelTurnEvent::Finished(result) => Some(result), + _ => None, + }) + .unwrap(); + assert!(result.output_items.iter().flat_map(|item| &item.parts).any(|part| { + matches!(part, Part::Reasoning(reasoning) if reasoning.summary.as_deref() == Some("brief") && !reasoning.redacted) + })); + } + + #[test] + fn role_validation_and_media_encoding_are_strict() { + let mut invalid = request(); + invalid.transcript = vec![Item::new( + ItemKind::User, + vec![Part::Reasoning(ReasoningPart::summary("do not expose"))], + )]; + assert!( + OpenAIResponsesConfig::new("secret", "gpt-test") + .encode_request(&invalid) + .is_err() + ); + + let mut request = request(); + request.transcript = vec![Item::new( + ItemKind::User, + vec![Part::media( + Modality::Image, + "image/png", + DataRef::InlineBytes(vec![1, 2, 3]), + )], + )]; + let encoded = OpenAIResponsesConfig::new("secret", "gpt-test") + .encode_request(&request) + .unwrap(); + let content = encoded + .pointer("/input/0/content") + .and_then(Value::as_array) + .unwrap(); + assert_eq!(content.len(), 1); + assert_eq!(content[0]["type"], "input_image"); + assert!( + content[0]["image_url"] + .as_str() + .unwrap() + .starts_with("data:image/png;base64,") + ); + } + + #[tokio::test] + async fn private_profile_scopes_turn_state_to_accepted_logical_turn() { + let mut headers = sse_headers(); + headers.insert(X_CODEX_TURN_STATE, HeaderValue::from_static("state-1")); + let client = Arc::new(ScriptedClient { + responses: Mutex::new(VecDeque::from([ + WireResponse { + status: StatusCode::SERVICE_UNAVAILABLE, + headers: headers.clone(), + body: "", + }, + WireResponse { + status: StatusCode::OK, + headers: headers.clone(), + body: SUCCESS, + }, + WireResponse { + status: StatusCode::OK, + headers, + body: SUCCESS, + }, + ])), + requests: Mutex::new(Vec::new()), + }); + let config = + OpenAIResponsesConfig::chatgpt_private("gpt-test", Authentication::bearer("secret")) + .with_resilience(ResilienceConfig { + max_retries: 1, + retry_budget: Duration::from_secs(1), + attempt_timeout: None, + stream_idle_timeout: None, + initial_backoff: Duration::ZERO, + max_backoff: Duration::ZERO, + }); + let adapter = OpenAIResponsesAdapter::with_client(config, Http::from_arc(client.clone())); + let mut session = adapter + .start_session(SessionConfig::new("session")) + .await + .unwrap(); + session.begin_turn(request(), None).await.unwrap(); + let mut second = request(); + second.turn_id = TurnId::new("turn-2"); + session.begin_turn(second, None).await.unwrap(); + let requests = client.requests.lock().unwrap(); + assert_eq!(requests[0].headers["originator"], "agentkit"); + assert_eq!(requests[0].headers["session-id"], "session"); + assert!(!requests[0].headers.contains_key(X_CODEX_TURN_STATE)); + assert!(!requests[1].headers.contains_key(X_CODEX_TURN_STATE)); + assert!(!requests[2].headers.contains_key(X_CODEX_TURN_STATE)); + assert_eq!(requests[0].body, requests[1].body); + assert_eq!( + requests[0].headers["idempotency-key"], + requests[1].headers["idempotency-key"] + ); + } + + #[tokio::test] + async fn response_metadata_turn_state_updates_same_turn_retry_context() { + const FAILED: &str = concat!( + "event: response.created\ndata: {\"type\":\"response.created\",\"sequence_number\":1,\"response\":{\"id\":\"resp-1\",\"model\":\"gpt-test\"}}\n\n", + "event: response.metadata\ndata: {\"type\":\"response.metadata\",\"sequence_number\":2,\"headers\":{\"X-Codex-Turn-State\":[\"retry-state\"]}}\n\n", + "event: response.failed\ndata: {\"type\":\"response.failed\",\"sequence_number\":3,\"response\":{\"id\":\"resp-1\",\"model\":\"gpt-test\",\"error\":{\"code\":\"server_error\"}}}\n\n", + ); + let client = Arc::new(ScriptedClient { + responses: Mutex::new(VecDeque::from([ + WireResponse { + status: StatusCode::OK, + headers: sse_headers(), + body: FAILED, + }, + WireResponse { + status: StatusCode::OK, + headers: sse_headers(), + body: SUCCESS, + }, + ])), + requests: Mutex::new(Vec::new()), + }); + let config = + OpenAIResponsesConfig::chatgpt_private("gpt-test", Authentication::bearer("secret")) + .with_resilience(ResilienceConfig { + max_retries: 1, + retry_budget: Duration::from_secs(1), + attempt_timeout: None, + stream_idle_timeout: None, + initial_backoff: Duration::ZERO, + max_backoff: Duration::ZERO, + }); + let adapter = OpenAIResponsesAdapter::with_client(config, Http::from_arc(client.clone())); + let mut session = adapter + .start_session(SessionConfig::new("session")) + .await + .unwrap(); + let mut turn = session.begin_turn(request(), None).await.unwrap(); + while turn.next_event(None).await.unwrap().is_some() {} + let requests = client.requests.lock().unwrap(); + assert!(!requests[0].headers.contains_key(X_CODEX_TURN_STATE)); + assert_eq!(requests[1].headers[X_CODEX_TURN_STATE], "retry-state"); + assert_eq!(requests[0].body, requests[1].body); + } + + struct SlowAuthentication; + + #[async_trait] + impl AuthenticationProvider for SlowAuthentication { + async fn authenticate( + &self, + _previous: Option<&AuthenticationAttempt>, + ) -> Result { + sleep(Duration::from_millis(100)).await; + Ok(AuthenticationAttempt::stateless(HeaderMap::new())) + } + } + + #[test] + fn prompt_cache_modes_keys_and_redacted_debug_are_validated() { + let config = OpenAIResponsesConfig::new("super-secret", "gpt-test"); + assert!(!format!("{config:?}").contains("super-secret")); + + let mut disabled = request(); + disabled.cache = Some(agentkit_loop::PromptCacheRequest::disabled().with_key("ignored")); + let body = config.encode_request(&disabled).unwrap(); + assert!(body.get("prompt_cache_key").is_none()); + + let mut explicit = request(); + explicit.cache = Some(agentkit_loop::PromptCacheRequest::explicit_required([])); + assert!(config.encode_request(&explicit).is_err()); + + let mut invalid_key = request(); + invalid_key.cache = + Some(agentkit_loop::PromptCacheRequest::automatic().with_key("x".repeat(257))); + assert!(config.encode_request(&invalid_key).is_err()); + } + + #[tokio::test] + async fn logical_deadline_starts_before_initial_authentication() { + let client = Arc::new(ScriptedClient { + responses: Mutex::new(VecDeque::new()), + requests: Mutex::new(Vec::new()), + }); + let config = + OpenAIResponsesConfig::public("gpt-test", Authentication::new(SlowAuthentication)) + .with_resilience(ResilienceConfig { + max_retries: 1, + retry_budget: Duration::from_millis(10), + attempt_timeout: Some(Duration::from_secs(1)), + stream_idle_timeout: Some(Duration::from_secs(1)), + initial_backoff: Duration::ZERO, + max_backoff: Duration::ZERO, + }); + let adapter = OpenAIResponsesAdapter::with_client(config, Http::from_arc(client.clone())); + let mut session = adapter + .start_session(SessionConfig::new("session")) + .await + .unwrap(); + let error = session.begin_turn(request(), None).await.err().unwrap(); + assert!(error.to_string().contains("logical request retry budget")); + assert!(client.requests.lock().unwrap().is_empty()); + } + + struct SlowRefreshAuthentication(AtomicUsize); + + #[async_trait] + impl AuthenticationProvider for SlowRefreshAuthentication { + async fn authenticate( + &self, + previous: Option<&AuthenticationAttempt>, + ) -> Result { + let call = self.0.fetch_add(1, Ordering::SeqCst); + if call == 0 { + assert!(previous.is_none()); + } else { + assert!(previous.is_some()); + sleep(Duration::from_millis(100)).await; + } + Ok(AuthenticationAttempt::new(HeaderMap::new(), call)) + } + } + + #[tokio::test] + async fn logical_deadline_also_bounds_reactive_refresh() { + let client = Arc::new(ScriptedClient { + responses: Mutex::new(VecDeque::from([WireResponse { + status: StatusCode::UNAUTHORIZED, + headers: HeaderMap::new(), + body: "", + }])), + requests: Mutex::new(Vec::new()), + }); + let config = OpenAIResponsesConfig::public( + "gpt-test", + Authentication::new(SlowRefreshAuthentication(AtomicUsize::new(0))), + ) + .with_resilience(ResilienceConfig { + max_retries: 1, + retry_budget: Duration::from_millis(10), + attempt_timeout: Some(Duration::from_secs(1)), + stream_idle_timeout: Some(Duration::from_secs(1)), + initial_backoff: Duration::ZERO, + max_backoff: Duration::ZERO, + }); + let adapter = OpenAIResponsesAdapter::with_client(config, Http::from_arc(client)); + let mut session = adapter + .start_session(SessionConfig::new("session")) + .await + .unwrap(); + let error = session.begin_turn(request(), None).await.err().unwrap(); + assert!(error.to_string().contains("logical request retry budget")); + } + + #[test] + fn request_helpers_preserve_boundaries_cache_and_idempotency() { + assert_eq!(frame_end(b"data: a\n\nlater\r\n\r\n"), Some((7, 2))); + assert_ne!( + stable_idempotency_key("session", "turn", b"one"), + stable_idempotency_key("session", "turn", b"two") + ); + + let mut request = request(); + request.cache = Some( + agentkit_loop::PromptCacheRequest::automatic() + .with_retention(agentkit_loop::PromptCacheRetention::Extended), + ); + let public = OpenAIResponsesConfig::new("secret", "gpt-test") + .encode_request(&request) + .unwrap(); + assert_eq!(public["prompt_cache_retention"], "24h"); + let private = + OpenAIResponsesConfig::chatgpt_private("gpt-test", Authentication::bearer("secret")) + .encode_request(&request) + .unwrap(); + assert!(private.get("prompt_cache_retention").is_none()); + } + + #[test] + fn notifications_are_wrapped_once_and_invalid_tools_and_audio_are_rejected() { + let mut notification = request(); + notification.transcript = vec![Item::new( + ItemKind::Notification, + vec![Part::text("notice"), Part::structured(json!({"code": 7}))], + )]; + let encoded = OpenAIResponsesConfig::new("secret", "gpt-test") + .encode_request(¬ification) + .unwrap(); + assert_eq!( + encoded + .pointer("/input/0/content/0/text") + .and_then(Value::as_str), + Some("\nnotice\n\n{\n \"code\": 7\n}\n") + ); + + let mut invalid_tool = request(); + invalid_tool.transcript = vec![Item::new( + ItemKind::Assistant, + vec![Part::ToolCall(ToolCallPart::new( + "call", + "not a valid tool", + json!({}), + ))], + )]; + assert!( + OpenAIResponsesConfig::new("secret", "gpt-test") + .encode_request(&invalid_tool) + .is_err() + ); + + let mut audio = request(); + audio.transcript = vec![Item::new( + ItemKind::User, + vec![Part::media( + Modality::Audio, + "audio/wav", + DataRef::InlineBytes(vec![1, 2, 3]), + )], + )]; + assert!( + OpenAIResponsesConfig::new("secret", "gpt-test") + .encode_request(&audio) + .is_err() + ); + } + + #[test] + fn refusal_and_keepalive_decode_without_losing_provider_item_id() { + let body = concat!( + "event: keepalive\ndata: {\"type\":\"keepalive\",\"sequence_number\":0}\n\n", + "event: response.created\ndata: {\"type\":\"response.created\",\"sequence_number\":1,\"response\":{\"id\":\"resp-refusal\",\"model\":\"gpt-test\"}}\n\n", + "event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"sequence_number\":2,\"output_index\":0,\"item\":{\"id\":\"refusal-item\",\"type\":\"message\"}}\n\n", + "event: response.content_part.added\ndata: {\"type\":\"response.content_part.added\",\"sequence_number\":3,\"item_id\":\"refusal-item\",\"output_index\":0,\"content_index\":0,\"part\":{\"type\":\"refusal\"}}\n\n", + "event: response.refusal.delta\ndata: {\"type\":\"response.refusal.delta\",\"sequence_number\":4,\"item_id\":\"refusal-item\",\"output_index\":0,\"content_index\":0,\"delta\":\"no\"}\n\n", + "event: response.refusal.done\ndata: {\"type\":\"response.refusal.done\",\"sequence_number\":5,\"item_id\":\"refusal-item\",\"output_index\":0,\"content_index\":0,\"refusal\":\"no\"}\n\n", + "event: response.content_part.done\ndata: {\"type\":\"response.content_part.done\",\"sequence_number\":6,\"item_id\":\"refusal-item\",\"output_index\":0,\"content_index\":0,\"part\":{\"type\":\"refusal\",\"refusal\":\"no\"}}\n\n", + "event: response.output_item.done\ndata: {\"type\":\"response.output_item.done\",\"sequence_number\":7,\"output_index\":0,\"item\":{\"id\":\"refusal-item\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"refusal\",\"refusal\":\"no\"}]}}\n\n", + "event: response.completed\ndata: {\"type\":\"response.completed\",\"sequence_number\":8,\"response\":{\"id\":\"resp-refusal\",\"model\":\"gpt-test\",\"status\":\"completed\"}}\n\n", + ); + let mut decoder = ResponsesSseDecoder::new("gpt-test", "session"); + decoder.push(body.as_bytes()).unwrap(); + let result = decoder + .finish() + .unwrap() + .into_iter() + .find_map(|event| match event { + ModelTurnEvent::Finished(result) => Some(result), + _ => None, + }) + .unwrap(); + assert_eq!( + result.output_items[0].id.as_ref().unwrap().0, + "refusal-item" + ); + assert!(matches!( + &result.output_items[0].parts[0], + Part::Text(text) if text.text == "no" + )); + } + + #[tokio::test] + async fn static_bearer_authentication_has_a_stable_non_secret_binding() { + let authentication = Authentication::bearer("secret"); + let first = authentication.authenticate(None).await.unwrap(); + let second = authentication.authenticate(Some(&first)).await.unwrap(); + assert!(first.binding().is_some()); + assert_eq!(first.binding(), second.binding()); + assert_ne!(first.binding(), Some("secret")); + } + + #[tokio::test] + async fn visible_stream_failure_is_fatal_without_replacement_capability() { + const PARTIAL: &str = concat!( + "event: response.created\ndata: {\"type\":\"response.created\",\"sequence_number\":1,\"response\":{\"id\":\"failed\",\"model\":\"gpt-test\"}}\n\n", + "event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"sequence_number\":2,\"output_index\":0,\"item\":{\"id\":\"failed-item\",\"type\":\"message\"}}\n\n", + "event: response.content_part.added\ndata: {\"type\":\"response.content_part.added\",\"sequence_number\":3,\"item_id\":\"failed-item\",\"output_index\":0,\"content_index\":0,\"part\":{\"type\":\"output_text\"}}\n\n", + "event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"sequence_number\":4,\"item_id\":\"failed-item\",\"output_index\":0,\"content_index\":0,\"delta\":\"visible\"}\n\n", + ); + let client = Arc::new(ScriptedClient { + responses: Mutex::new(VecDeque::from([ + WireResponse { + status: StatusCode::OK, + headers: sse_headers(), + body: PARTIAL, + }, + WireResponse { + status: StatusCode::OK, + headers: sse_headers(), + body: SUCCESS, + }, + ])), + requests: Mutex::new(Vec::new()), + }); + let config = + OpenAIResponsesConfig::new("secret", "gpt-test").with_resilience(ResilienceConfig { + max_retries: 1, + retry_budget: Duration::from_secs(1), + attempt_timeout: None, + stream_idle_timeout: None, + initial_backoff: Duration::ZERO, + max_backoff: Duration::ZERO, + }); + let adapter = OpenAIResponsesAdapter::with_client(config, Http::from_arc(client.clone())); + let mut session = adapter + .start_session(SessionConfig::new("session")) + .await + .unwrap(); + let mut turn = session.begin_turn(request(), None).await.unwrap(); + let mut saw_visible = false; + let error = loop { + match turn.next_event(None).await { + Ok(Some(ModelTurnEvent::Delta(Delta::AppendText { chunk, .. }))) => { + saw_visible |= chunk == "visible"; + } + Ok(Some(_)) => {} + Ok(None) => panic!("truncated attempt unexpectedly completed"), + Err(error) => break error, + } + }; + assert!(saw_visible); + assert!(matches!(error, LoopError::Provider(_))); + assert_eq!(client.requests.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn opted_in_visible_stream_failure_emits_supersession_then_replays() { + const PARTIAL: &str = concat!( + "event: response.created\ndata: {\"type\":\"response.created\",\"sequence_number\":1,\"response\":{\"id\":\"failed\",\"model\":\"gpt-test\"}}\n\n", + "event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"sequence_number\":2,\"output_index\":0,\"item\":{\"id\":\"failed-item\",\"type\":\"message\"}}\n\n", + "event: response.content_part.added\ndata: {\"type\":\"response.content_part.added\",\"sequence_number\":3,\"item_id\":\"failed-item\",\"output_index\":0,\"content_index\":0,\"part\":{\"type\":\"output_text\"}}\n\n", + "event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"sequence_number\":4,\"item_id\":\"failed-item\",\"output_index\":0,\"content_index\":0,\"delta\":\"replace me\"}\n\n", + ); + let client = Arc::new(ScriptedClient { + responses: Mutex::new(VecDeque::from([ + WireResponse { + status: StatusCode::OK, + headers: sse_headers(), + body: PARTIAL, + }, + WireResponse { + status: StatusCode::OK, + headers: sse_headers(), + body: SUCCESS, + }, + ])), + requests: Mutex::new(Vec::new()), + }); + let config = + OpenAIResponsesConfig::new("secret", "gpt-test").with_resilience(ResilienceConfig { + max_retries: 1, + retry_budget: Duration::from_secs(1), + attempt_timeout: None, + stream_idle_timeout: None, + initial_backoff: Duration::ZERO, + max_backoff: Duration::ZERO, + }); + let adapter = OpenAIResponsesAdapter::with_client(config, Http::from_arc(client.clone())); + let session_config = SessionConfig::new("session").with_response_attempt_supersession(); + let mut session = adapter.start_session(session_config).await.unwrap(); + let mut turn = session.begin_turn(request(), None).await.unwrap(); + let mut saw_supersession = false; + let mut saw_replacement = false; + while let Some(event) = turn.next_event(None).await.unwrap() { + match event { + ModelTurnEvent::ResponseAttemptSuperseded => { + assert!(!saw_supersession); + saw_supersession = true; + } + ModelTurnEvent::Delta(Delta::AppendText { chunk, .. }) if chunk == "hello" => { + assert!(saw_supersession); + saw_replacement = true; + } + _ => {} + } + } + assert!(saw_supersession); + assert!(saw_replacement); + assert_eq!(client.requests.lock().unwrap().len(), 2); + } + + #[tokio::test] + async fn performs_only_one_reactive_refresh() { + let calls = Arc::new(AtomicUsize::new(0)); + let client = Arc::new(ScriptedClient { + responses: Mutex::new(VecDeque::from([ + WireResponse { + status: StatusCode::UNAUTHORIZED, + headers: HeaderMap::new(), + body: "", + }, + WireResponse { + status: StatusCode::UNAUTHORIZED, + headers: HeaderMap::new(), + body: "", + }, + ])), + requests: Mutex::new(Vec::new()), + }); + let config = OpenAIResponsesConfig::public( + "gpt-test", + Authentication::new(RefreshingAuth { + calls: calls.clone(), + }), + ); + let adapter = OpenAIResponsesAdapter::with_client(config, Http::from_arc(client.clone())); + let mut session = adapter + .start_session(SessionConfig::new("session")) + .await + .unwrap(); + assert!(session.begin_turn(request(), None).await.is_err()); + assert_eq!(calls.load(Ordering::SeqCst), 2); + assert_eq!(client.requests.lock().unwrap().len(), 2); + } + + #[test] + fn private_audio_and_media_tool_outputs_preserve_multimodal_content() { + let config = + OpenAIResponsesConfig::chatgpt_private("gpt-test", Authentication::bearer("secret")); + let mut audio = request(); + audio.transcript = vec![Item::new( + ItemKind::User, + vec![Part::media( + Modality::Audio, + "audio/wav", + DataRef::InlineBytes(vec![1, 2, 3]), + )], + )]; + let encoded = config.encode_request(&audio).unwrap(); + assert_eq!(encoded["input"][0]["content"][0]["type"], "input_audio"); + assert_eq!( + encoded["input"][0]["content"][0]["audio_url"], + "data:audio/wav;base64,AQID" + ); + + let mut tool = request(); + tool.transcript = vec![Item::new( + ItemKind::Tool, + vec![Part::ToolResult(agentkit_core::ToolResultPart::success( + "call-1", + ToolOutput::parts(vec![ + Part::text("screenshot"), + Part::media( + Modality::Image, + "image/png", + DataRef::InlineBytes(vec![1, 2, 3]), + ), + ]), + ))], + )]; + let encoded = config.encode_request(&tool).unwrap(); + let output = encoded["input"][0]["output"].as_array().unwrap(); + assert_eq!(output[0]["type"], "input_text"); + assert_eq!(output[1]["type"], "input_image"); + } + + #[test] + fn media_validation_rejects_noncanonical_base64_and_unsafe_uris() { + let config = OpenAIResponsesConfig::new("secret", "gpt-test"); + for data in [ + DataRef::InlineText("%%%".into()), + DataRef::Uri("file:///tmp/private.png".into()), + DataRef::Uri("custom://artifact/image".into()), + DataRef::Uri("data:image/jpeg;base64,AQID".into()), + ] { + let mut invalid = request(); + invalid.transcript = vec![Item::new( + ItemKind::User, + vec![Part::media(Modality::Image, "image/png", data)], + )]; + assert!(config.encode_request(&invalid).is_err()); + } + let mut valid = request(); + valid.transcript = vec![Item::new( + ItemKind::User, + vec![Part::media( + Modality::Image, + "image/png", + DataRef::Uri("https://example.com/image.png".into()), + )], + )]; + assert!(config.encode_request(&valid).is_ok()); + } + + #[test] + fn generated_image_output_is_persisted_and_replayed() { + let wire = concat!( + "event: response.created\ndata: {\"type\":\"response.created\",\"sequence_number\":1,\"response\":{\"id\":\"resp-image\",\"model\":\"gpt-test\"}}\n\n", + "event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"sequence_number\":2,\"output_index\":0,\"item\":{\"id\":\"image-1\",\"type\":\"image_generation_call\"}}\n\n", + "event: response.output_item.done\ndata: {\"type\":\"response.output_item.done\",\"sequence_number\":3,\"output_index\":0,\"item\":{\"id\":\"image-1\",\"type\":\"image_generation_call\",\"status\":\"completed\",\"revised_prompt\":\"safer prompt\",\"result\":\"AQID\"}}\n\n", + "event: response.completed\ndata: {\"type\":\"response.completed\",\"sequence_number\":4,\"response\":{\"id\":\"resp-image\",\"model\":\"gpt-test\",\"status\":\"completed\"}}\n\n", + ); + let mut decoder = ResponsesSseDecoder::with_policy( + "gpt-test", + "session", + OpenAIResponsesProfile::ChatGptPrivate, + true, + Some("test-authentication-binding"), + Arc::new(Mutex::new(None)), + OpenAIResponsesLimits::default(), + ); + decoder.push(wire.as_bytes()).unwrap(); + let events = decoder.finish().unwrap(); + let deltas = events + .iter() + .filter_map(|event| match event { + ModelTurnEvent::Delta(delta) => Some(delta), + _ => None, + }) + .collect::>(); + assert!(matches!( + deltas.as_slice(), + [ + Delta::BeginPart { + part_id: begin, + kind: PartKind::Text, + }, + Delta::AppendText { + part_id: append, + chunk, + }, + Delta::CommitPart { + part: Part::Text(committed), + }, + ] if begin.0 == "generated-image-0" + && append == begin + && chunk == "[Image #1]" + && committed.text == "[Image #1]" + )); + let result = events + .iter() + .find_map(|event| match event { + ModelTurnEvent::Finished(result) => Some(result), + _ => None, + }) + .unwrap(); + let Part::Media(media) = &result.output_items[0].parts[0] else { + panic!("generated image was not persisted as media"); + }; + assert_eq!(media.data, DataRef::InlineBytes(vec![1, 2, 3])); + assert!(media.metadata.contains_key(CONTINUATION_METADATA)); + assert!(media.metadata.contains_key(GENERATED_IMAGE_METADATA)); + + let replay = TurnRequest { + transcript: result.output_items.clone(), + ..request() + }; + let config = + OpenAIResponsesConfig::chatgpt_private("gpt-test", Authentication::bearer("secret")); + let encoded = + encode_request_bound(&config, &replay, Some("test-authentication-binding")).unwrap(); + assert_eq!(encoded["input"][0]["type"], "image_generation_call"); + assert_eq!(encoded["input"][0]["result"], "AQID"); + + let bounded = config.with_limits(OpenAIResponsesLimits { + max_text_bytes: 3, + ..OpenAIResponsesLimits::default() + }); + assert!( + encode_request_bound(&bounded, &replay, Some("test-authentication-binding")).is_err() + ); + } + + #[test] + fn private_unknown_events_are_ignored_and_statusless_failures_retry_unless_permanent() { + let unknown = b"event: response.future_hint\ndata: {\"type\":\"response.future_hint\",\"sequence_number\":1}\n\n"; + let mut private = ResponsesSseDecoder::with_policy( + "gpt-test", + "session", + OpenAIResponsesProfile::ChatGptPrivate, + true, + Some("binding"), + Arc::new(Mutex::new(None)), + OpenAIResponsesLimits::default(), + ); + assert!(private.push(unknown).is_ok()); + let mut public = ResponsesSseDecoder::new("gpt-test", "session"); + assert!(public.push(unknown).is_err()); + + assert!(stream_failure_retryable( + OpenAIResponsesProfile::ChatGptPrivate, + &json!({"type": "error", "error": {"type": "brand_new_error", "code": "never_seen_before"}}), + "error", + )); + assert!(!stream_failure_retryable( + OpenAIResponsesProfile::ChatGptPrivate, + &json!({"type": "error", "error": {"type": "invalid_request_error", "code": "invalid_prompt"}}), + "error", + )); + } + + #[test] + fn serialized_request_limit_is_configurable() { + let config = + OpenAIResponsesConfig::new("secret", "gpt-test").with_limits(OpenAIResponsesLimits { + max_request_bytes: 16, + max_text_bytes: 8, + ..OpenAIResponsesLimits::default() + }); + assert!(config.encode_request(&request()).is_err()); + } + + #[test] + fn item_and_text_defaults_preserve_public_bounds() { + let limits = OpenAIResponsesLimits::default(); + assert_eq!(limits.max_items, 100_000); + assert_eq!(limits.max_text_bytes, 8 * 1024 * 1024); + } + + #[test] + fn zero_and_inconsistent_limits_are_rejected() { + for limits in [ + OpenAIResponsesLimits { + max_items: 0, + ..OpenAIResponsesLimits::default() + }, + OpenAIResponsesLimits { + max_text_bytes: DEFAULT_MAX_REQUEST_BYTES + 1, + ..OpenAIResponsesLimits::default() + }, + OpenAIResponsesLimits { + max_attempt_bytes: DEFAULT_MAX_WIRE_BYTES + 1, + ..OpenAIResponsesLimits::default() + }, + ] { + let config = OpenAIResponsesConfig::new("secret", "gpt-test").with_limits(limits); + assert!(OpenAIResponsesAdapter::new(config).is_err()); + } + } + + #[test] + fn request_item_and_text_limits_are_configurable() { + let mut too_many = request(); + too_many + .transcript + .push(Item::text(ItemKind::User, "again")); + let item_bounded = + OpenAIResponsesConfig::new("secret", "gpt-test").with_limits(OpenAIResponsesLimits { + max_items: 1, + ..OpenAIResponsesLimits::default() + }); + assert!(item_bounded.encode_request(&too_many).is_err()); + + let text_bounded = + OpenAIResponsesConfig::new("secret", "gpt-test").with_limits(OpenAIResponsesLimits { + max_text_bytes: 4, + ..OpenAIResponsesLimits::default() + }); + assert!(text_bounded.encode_request(&request()).is_err()); + } + + #[test] + fn sse_event_index_and_text_limits_are_configurable() { + let decoder = |limits| { + ResponsesSseDecoder::with_policy( + "gpt-test", + "session", + OpenAIResponsesProfile::Public, + true, + Some("test-authentication-binding"), + Arc::new(Mutex::new(None)), + limits, + ) + }; + + let mut event_bounded = decoder(OpenAIResponsesLimits { + max_items: 3, + ..OpenAIResponsesLimits::default() + }); + assert!(event_bounded.push(SUCCESS.as_bytes()).is_err()); + + let outside_index = concat!( + "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp-1\",\"model\":\"gpt-test\"}}\n\n", + "event: response.output_item.added\ndata: {\"type\":\"response.output_item.added\",\"output_index\":1,\"item\":{\"id\":\"msg-1\",\"type\":\"message\"}}\n\n", + ); + let mut index_bounded = decoder(OpenAIResponsesLimits { + max_items: 1, + ..OpenAIResponsesLimits::default() + }); + assert!(index_bounded.push(outside_index.as_bytes()).is_err()); + + let mut text_bounded = decoder(OpenAIResponsesLimits { + max_text_bytes: 4, + ..OpenAIResponsesLimits::default() + }); + assert!(text_bounded.push(SUCCESS.as_bytes()).is_err()); + } + + #[tokio::test] + async fn private_http_529_retries_and_custom_attribution_is_sent() { + let status_529 = StatusCode::from_u16(529).unwrap(); + let client = Arc::new(ScriptedClient { + responses: Mutex::new(VecDeque::from([ + WireResponse { + status: status_529, + headers: HeaderMap::new(), + body: "", + }, + WireResponse { + status: StatusCode::OK, + headers: sse_headers(), + body: SUCCESS, + }, + ])), + requests: Mutex::new(Vec::new()), + }); + let config = + OpenAIResponsesConfig::chatgpt_private("gpt-test", Authentication::bearer("secret")) + .with_user_agent("example-client/test") + .with_originator("example-client") + .with_resilience(ResilienceConfig { + max_retries: 1, + retry_budget: Duration::from_secs(1), + attempt_timeout: None, + stream_idle_timeout: None, + initial_backoff: Duration::ZERO, + max_backoff: Duration::ZERO, + }); + let adapter = OpenAIResponsesAdapter::with_client(config, Http::from_arc(client.clone())); + let mut session = adapter + .start_session(SessionConfig::new("session")) + .await + .unwrap(); + session.begin_turn(request(), None).await.unwrap(); + let requests = client.requests.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[0].headers["user-agent"], "example-client/test"); + assert_eq!(requests[0].headers["originator"], "example-client"); + } + + #[tokio::test] + async fn attempt_deadline_bounds_continuously_active_stream() { + let client = Arc::new(ActiveStreamClient { + requests: AtomicUsize::new(0), + }); + let config = + OpenAIResponsesConfig::new("secret", "gpt-test").with_resilience(ResilienceConfig { + max_retries: 0, + retry_budget: Duration::from_secs(1), + attempt_timeout: Some(Duration::from_millis(10)), + stream_idle_timeout: Some(Duration::from_millis(100)), + initial_backoff: Duration::ZERO, + max_backoff: Duration::ZERO, + }); + let adapter = OpenAIResponsesAdapter::with_client(config, Http::from_arc(client.clone())); + let mut session = adapter + .start_session(SessionConfig::new("session")) + .await + .unwrap(); + let mut turn = session.begin_turn(request(), None).await.unwrap(); + let error = turn.next_event(None).await.unwrap_err(); + assert!(error.to_string().contains("attempt timed out")); + assert_eq!(client.requests.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn aggregate_wire_limit_spans_retries() { + let client = Arc::new(ScriptedClient { + responses: Mutex::new(VecDeque::from([ + WireResponse { + status: StatusCode::OK, + headers: sse_headers(), + body: "1234567890", + }, + WireResponse { + status: StatusCode::OK, + headers: sse_headers(), + body: "1234567890", + }, + ])), + requests: Mutex::new(Vec::new()), + }); + let config = OpenAIResponsesConfig::new("secret", "gpt-test") + .with_limits(OpenAIResponsesLimits { + max_request_bytes: DEFAULT_MAX_REQUEST_BYTES, + max_attempt_bytes: 10, + max_wire_bytes: 15, + max_items: DEFAULT_MAX_ITEMS, + max_text_bytes: 8, + }) + .with_resilience(ResilienceConfig { + max_retries: 3, + retry_budget: Duration::from_secs(1), + attempt_timeout: None, + stream_idle_timeout: None, + initial_backoff: Duration::ZERO, + max_backoff: Duration::ZERO, + }); + let adapter = OpenAIResponsesAdapter::with_client(config, Http::from_arc(client.clone())); + let mut session = adapter + .start_session(SessionConfig::new("session")) + .await + .unwrap(); + let mut turn = session.begin_turn(request(), None).await.unwrap(); + let error = turn.next_event(None).await.unwrap_err(); + assert!(error.to_string().contains("wire-byte limit")); + assert_eq!(client.requests.lock().unwrap().len(), 2); + } +} diff --git a/crates/agentkit-provider-openrouter/README.md b/crates/agentkit-provider-openrouter/README.md index 637a7bc..4103c27 100644 --- a/crates/agentkit-provider-openrouter/README.md +++ b/crates/agentkit-provider-openrouter/README.md @@ -22,7 +22,7 @@ Use it when OpenRouter is the backing model provider for your agent runtime. ## Configuration -Create a config with `OpenRouterConfig::new(api_key, model)` and chain `.with_*()` builders for optional parameters. Streaming is enabled by default; use `.with_streaming(false)` to force the buffered response path. Alternatively, `OpenRouterConfig::from_env()` reads from environment variables: +Create a config with `OpenRouterConfig::new(authentication, model)` and chain `.with_*()` builders for optional parameters. Streaming is enabled by default; use `.with_streaming(false)` to force the buffered response path. Alternatively, `OpenRouterConfig::from_env()` reads from environment variables: | Variable | Required | Default | | ---------------------------------- | -------- | ----------------------------------------------- | @@ -34,6 +34,18 @@ Create a config with `OpenRouterConfig::new(api_key, model)` and chain `.with_*( | `OPENROUTER_MAX_COMPLETION_TOKENS` | no | -- | | `OPENROUTER_TEMPERATURE` | no | -- | +## Authentication and resilience + +`OpenRouterConfig` stores credentials as a first-class +`agentkit_http::Authentication`. A bare string passed to +`OpenRouterConfig::new` or `.with_authentication(...)` is shorthand for bearer +authentication. Use `.with_authentication_provider(...)` for a custom +refresh-capable `AuthenticationProvider`. + +Resilience is opt-in: `resilience` is an `Option` that +defaults to `None`. Calling `.with_resilience(...)` enables retries and +timeouts; leaving it as `None` preserves the existing single-attempt behavior. + ## Examples ### Minimal chat agent diff --git a/crates/agentkit-provider-openrouter/src/lib.rs b/crates/agentkit-provider-openrouter/src/lib.rs index 4494e57..ea16cd0 100644 --- a/crates/agentkit-provider-openrouter/src/lib.rs +++ b/crates/agentkit-provider-openrouter/src/lib.rs @@ -35,6 +35,7 @@ use agentkit_adapter_completions::{ CompletionsAdapter, CompletionsError, CompletionsProvider, CompletionsSession, CompletionsTurn, }; use agentkit_core::{CostUsage, Item, ItemKind, MetadataMap, Part, Usage}; +use agentkit_http::{Authentication, AuthenticationProvider, ResilienceConfig}; use agentkit_loop::{ LoopError, ModelAdapter, PromptCacheBreakpoint, PromptCacheMode, PromptCacheRequest, PromptCacheRetention, PromptCacheStrategy, SessionConfig, TurnRequest, @@ -62,10 +63,10 @@ const DEFAULT_BASE_URL: &str = "https://openrouter.ai/api/v1/chat/completions"; /// .with_max_completion_tokens(4096) /// .with_app_name("my-agent"); /// ``` -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct OpenRouterConfig { - /// OpenRouter API key (starts with `sk-or-`). - pub api_key: String, + /// Authentication applied to each request. Bare strings use bearer authentication. + pub authentication: Authentication, /// Model identifier, e.g. `"anthropic/claude-sonnet-4"` or `"openrouter/auto"`. pub model: String, /// Chat completions endpoint URL. Defaults to the OpenRouter production URL. @@ -90,15 +91,38 @@ pub struct OpenRouterConfig { pub reasoning_effort: Option, /// Request SSE streaming responses. Defaults to `true`. pub streaming: bool, + /// Optional retry and timeout policy. `None` preserves single-attempt behavior. + pub resilience: Option, /// Arbitrary extra fields merged into the request body. pub extra_body: MetadataMap, } +impl std::fmt::Debug for OpenRouterConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("OpenRouterConfig") + .field("authentication", &"") + .field("model", &self.model) + .field("base_url", &self.base_url) + .field("app_name", &self.app_name) + .field("site_url", &self.site_url) + .field("max_completion_tokens", &self.max_completion_tokens) + .field("temperature", &self.temperature) + .field("parallel_tool_calls", &self.parallel_tool_calls) + .field("reasoning_effort", &self.reasoning_effort) + .field("streaming", &self.streaming) + .field("resilience", &self.resilience) + .field("extra_body", &self.extra_body) + .finish() + } +} + impl OpenRouterConfig { - /// Creates a new configuration with the given API key and model identifier. - pub fn new(api_key: impl Into, model: impl Into) -> Self { + /// Creates a new configuration with the given authentication and model identifier. + /// + /// Bare strings are treated as bearer tokens. + pub fn new(authentication: impl Into, model: impl Into) -> Self { Self { - api_key: api_key.into(), + authentication: authentication.into(), model: model.into(), base_url: DEFAULT_BASE_URL.into(), app_name: None, @@ -108,10 +132,28 @@ impl OpenRouterConfig { parallel_tool_calls: None, reasoning_effort: None, streaming: true, + resilience: None, extra_body: MetadataMap::new(), } } + /// Replaces the configured authentication. + pub fn with_authentication(mut self, authentication: impl Into) -> Self { + self.authentication = authentication.into(); + self + } + + /// Uses a custom refresh-capable authentication provider. + pub fn with_authentication_provider(self, provider: P) -> Self { + self.with_authentication(Authentication::new(provider)) + } + + /// Enables request and pre-visible-output retries and timeouts. + pub fn with_resilience(mut self, resilience: ResilienceConfig) -> Self { + self.resilience = Some(resilience); + self + } + /// Requests a reasoning effort level from reasoning-capable models. pub fn with_reasoning_effort(mut self, effort: ReasoningEffort) -> Self { self.reasoning_effort = Some(effort); @@ -303,7 +345,8 @@ impl Serialize for ReasoningEffort { /// The OpenRouter provider, implementing [`CompletionsProvider`]. #[derive(Clone, Debug)] pub struct OpenRouterProvider { - api_key: String, + authentication: Authentication, + resilience: Option, base_url: String, app_name: Option, site_url: Option, @@ -314,7 +357,8 @@ pub struct OpenRouterProvider { impl From for OpenRouterProvider { fn from(config: OpenRouterConfig) -> Self { Self { - api_key: config.api_key, + authentication: config.authentication, + resilience: config.resilience, base_url: config.base_url, app_name: config.app_name, site_url: config.site_url, @@ -352,7 +396,7 @@ impl CompletionsProvider for OpenRouterProvider { &self, builder: agentkit_http::HttpRequestBuilder, ) -> agentkit_http::HttpRequestBuilder { - let mut builder = builder.bearer_auth(&self.api_key).header( + let mut builder = builder.header( "User-Agent", concat!("agentkit-provider-openrouter/", env!("CARGO_PKG_VERSION")), ); @@ -365,6 +409,14 @@ impl CompletionsProvider for OpenRouterProvider { builder } + fn authentication(&self) -> Option { + Some(self.authentication.clone()) + } + + fn resilience_config(&self) -> Option { + self.resilience.clone() + } + fn streaming(&self) -> bool { self.streaming } @@ -872,6 +924,21 @@ impl OpenRouterAdapter { let provider = OpenRouterProvider::from(config); Ok(Self(CompletionsAdapter::new(provider)?)) } + + /// Overrides the default API-key authentication. + pub fn with_authentication(self, authentication: impl Into) -> Self { + Self(self.0.with_authentication(authentication)) + } + + /// Overrides authentication with a custom refresh-capable provider. + pub fn with_authentication_provider(self, provider: P) -> Self { + self.with_authentication(Authentication::new(provider)) + } + + /// Enables retry and timeout behavior. + pub fn with_resilience(self, resilience: ResilienceConfig) -> Self { + Self(self.0.with_resilience(resilience)) + } } #[async_trait] @@ -911,6 +978,29 @@ mod tests { use super::*; + #[test] + fn config_debug_redacts_api_key() { + let debug = format!( + "{:?}", + OpenRouterConfig::new("openrouter-secret", "debug-model") + ); + assert!(!debug.contains("openrouter-secret")); + assert!(debug.contains("")); + assert!(debug.contains("debug-model")); + } + + #[test] + fn config_resilience_reaches_provider() { + let default_provider = OpenRouterProvider::from(OpenRouterConfig::new("key", "model")); + assert_eq!(default_provider.resilience_config(), None); + + let resilience = ResilienceConfig::no_retries(); + let provider = OpenRouterProvider::from( + OpenRouterConfig::new("key", "model").with_resilience(resilience.clone()), + ); + assert_eq!(provider.resilience_config(), Some(resilience)); + } + #[test] fn adapter_reports_provider_name_without_starting_a_session() { let adapter = OpenRouterAdapter::new(OpenRouterConfig::new("key", "model")).unwrap(); diff --git a/crates/agentkit-provider-vllm/README.md b/crates/agentkit-provider-vllm/README.md index 2d300d1..499aa5a 100644 --- a/crates/agentkit-provider-vllm/README.md +++ b/crates/agentkit-provider-vllm/README.md @@ -28,6 +28,19 @@ Create a config with `VllmConfig::new(model)` and chain `.with_*()` builders for | `VLLM_BASE_URL` | no | `http://localhost:8000/v1/chat/completions` | | `VLLM_API_KEY` | no | -- | +## Authentication and resilience + +`VllmConfig::new` leaves its first-class +`Option` as `None`. For a protected server, a +bare string passed to `.with_authentication(...)` is shorthand for bearer +authentication. Use `.with_authentication_provider(...)` for a custom +refresh-capable `AuthenticationProvider`. `VLLM_API_KEY` uses the same bearer +authentication when read by `from_env()`. + +Resilience is also opt-in: `resilience` is an `Option` that +defaults to `None`. Calling `.with_resilience(...)` enables retries and +timeouts; leaving it as `None` preserves the existing single-attempt behavior. + ## Examples ### Minimal chat agent @@ -62,7 +75,7 @@ use agentkit_provider_vllm::{VllmAdapter, VllmConfig}; # fn main() -> Result<(), Box> { let config = VllmConfig::new("meta-llama/Llama-3.1-8B-Instruct") - .with_api_key("my-secret-key") + .with_authentication("my-secret-key") .with_temperature(0.0) .with_max_completion_tokens(4096); diff --git a/crates/agentkit-provider-vllm/src/lib.rs b/crates/agentkit-provider-vllm/src/lib.rs index ae33ea9..c51c4ef 100644 --- a/crates/agentkit-provider-vllm/src/lib.rs +++ b/crates/agentkit-provider-vllm/src/lib.rs @@ -32,6 +32,7 @@ use agentkit_adapter_completions::{ CompletionsAdapter, CompletionsError, CompletionsProvider, CompletionsSession, CompletionsTurn, }; +use agentkit_http::{Authentication, AuthenticationProvider, ResilienceConfig}; use agentkit_loop::{LoopError, ModelAdapter, SessionConfig}; use async_trait::async_trait; use serde::Serialize; @@ -54,15 +55,17 @@ const DEFAULT_ENDPOINT: &str = "http://localhost:8000/v1/chat/completions"; /// .with_base_url("http://gpu-server:8000/v1/chat/completions") /// .with_temperature(0.0); /// ``` -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct VllmConfig { /// HuggingFace model identifier served by the vLLM instance, /// e.g. `"meta-llama/Llama-3.1-8B-Instruct"`. pub model: String, /// Chat completions endpoint URL. Defaults to `http://localhost:8000/v1/chat/completions`. pub base_url: String, - /// Optional API key, required only if the vLLM server enforces authentication. - pub api_key: Option, + /// Optional authentication for servers that enforce it. Strings become bearer authentication. + pub authentication: Option, + /// Optional retry and timeout policy. `None` preserves single-attempt behavior. + pub resilience: Option, /// Sampling temperature (0.0 = deterministic, higher = more creative). pub temperature: Option, /// Maximum number of completion tokens the model may generate. @@ -83,13 +86,31 @@ pub struct VllmConfig { pub strict_alternating_roles: bool, } +impl std::fmt::Debug for VllmConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("VllmConfig") + .field("model", &self.model) + .field("base_url", &self.base_url) + .field("authentication", &"") + .field("resilience", &self.resilience) + .field("temperature", &self.temperature) + .field("max_completion_tokens", &self.max_completion_tokens) + .field("top_p", &self.top_p) + .field("parallel_tool_calls", &self.parallel_tool_calls) + .field("streaming", &self.streaming) + .field("strict_alternating_roles", &self.strict_alternating_roles) + .finish() + } +} + impl VllmConfig { /// Creates a new configuration with the given model identifier. pub fn new(model: impl Into) -> Self { Self { model: model.into(), base_url: DEFAULT_ENDPOINT.into(), - api_key: None, + authentication: None, + resilience: None, temperature: None, max_completion_tokens: None, top_p: None, @@ -105,9 +126,26 @@ impl VllmConfig { self } - /// Sets the API key for authenticated vLLM servers. - pub fn with_api_key(mut self, key: impl Into) -> Self { - self.api_key = Some(key.into()); + /// Sets authentication for protected vLLM servers. Strings become bearer authentication. + pub fn with_authentication(mut self, authentication: impl Into) -> Self { + self.authentication = Some(authentication.into()); + self + } + + /// Uses a custom refresh-capable authentication provider. + pub fn with_authentication_provider(mut self, provider: P) -> Self { + self.authentication = Some(Authentication::new(provider)); + self + } + + /// Compatibility alias for bearer API-key authentication. + pub fn with_api_key(self, key: impl Into) -> Self { + self.with_authentication(key) + } + + /// Enables request retries and timeouts. + pub fn with_resilience(mut self, resilience: ResilienceConfig) -> Self { + self.resilience = Some(resilience); self } @@ -191,7 +229,8 @@ pub struct VllmRequestConfig { #[derive(Clone, Debug)] pub struct VllmProvider { base_url: String, - api_key: Option, + authentication: Option, + resilience: Option, streaming: bool, strict_alternating_roles: bool, request_config: VllmRequestConfig, @@ -201,7 +240,8 @@ impl From for VllmProvider { fn from(config: VllmConfig) -> Self { Self { base_url: config.base_url, - api_key: config.api_key, + authentication: config.authentication, + resilience: config.resilience, streaming: config.streaming, strict_alternating_roles: config.strict_alternating_roles, request_config: VllmRequestConfig { @@ -232,14 +272,18 @@ impl CompletionsProvider for VllmProvider { &self, builder: agentkit_http::HttpRequestBuilder, ) -> agentkit_http::HttpRequestBuilder { - let builder = builder.header( + builder.header( "User-Agent", concat!("agentkit-provider-vllm/", env!("CARGO_PKG_VERSION")), - ); - match &self.api_key { - Some(key) => builder.bearer_auth(key), - None => builder, - } + ) + } + + fn authentication(&self) -> Option { + self.authentication.clone() + } + + fn resilience_config(&self) -> Option { + self.resilience.clone() } fn streaming(&self) -> bool { @@ -285,6 +329,21 @@ impl VllmAdapter { let provider = VllmProvider::from(config); Ok(Self(CompletionsAdapter::new(provider)?)) } + + /// Adds or overrides optional server authentication. + pub fn with_authentication(self, authentication: impl Into) -> Self { + Self(self.0.with_authentication(authentication)) + } + + /// Adds refresh-capable optional server authentication. + pub fn with_authentication_provider(self, provider: P) -> Self { + self.with_authentication(Authentication::new(provider)) + } + + /// Enables retry and timeout behavior. + pub fn with_resilience(self, resilience: ResilienceConfig) -> Self { + Self(self.0.with_resilience(resilience)) + } } #[async_trait] @@ -307,3 +366,34 @@ pub enum VllmError { #[error(transparent)] Completions(#[from] CompletionsError), } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn authentication_and_resilience_are_optional_and_propagate() { + let default_provider = VllmProvider::from(VllmConfig::new("model")); + assert!(default_provider.authentication().is_none()); + assert!(default_provider.resilience_config().is_none()); + + let provider = VllmProvider::from( + VllmConfig::new("model") + .with_authentication("vllm-secret") + .with_resilience(ResilienceConfig::default()), + ); + assert!(provider.authentication().is_some()); + assert!(provider.resilience_config().is_some()); + } + + #[test] + fn config_debug_redacts_authentication() { + let debug = format!( + "{:?}", + VllmConfig::new("debug-model").with_api_key("vllm-secret") + ); + assert!(!debug.contains("vllm-secret")); + assert!(debug.contains("")); + assert!(debug.contains("debug-model")); + } +} diff --git a/docs/cerebras-provider-plan.md b/docs/cerebras-provider-plan.md index 9746261..bf33dc7 100644 --- a/docs/cerebras-provider-plan.md +++ b/docs/cerebras-provider-plan.md @@ -293,7 +293,7 @@ struct ChoiceState { #[derive(Clone)] pub struct CerebrasConfig { // auth & transport - pub api_key: String, + pub authentication: Authentication, // strings become bearer auth pub base_url: String, // default: https://api.cerebras.ai/v1 pub version_patch: Option, // X-Cerebras-Version-Patch pub extra_headers: Vec<(String, String)>, // SDK-style passthrough @@ -329,6 +329,8 @@ pub struct CerebrasConfig { // stream pub streaming: bool, // default true + pub resilience: Option, // default None (single attempt) + // preview-gated #[cfg(feature = "predicted-outputs")] pub prediction: Option, @@ -338,11 +340,9 @@ pub struct CerebrasConfig { } ``` -Builder: one `with_*` per field (mirrors anthropic). `from_env()` reads `CEREBRAS_API_KEY` (required), `CEREBRAS_MODEL` (required), `CEREBRAS_BASE_URL`, `CEREBRAS_VERSION_PATCH`, `CEREBRAS_MAX_COMPLETION_TOKENS`. +Builder: one `with_*` per field (mirrors anthropic), including `with_authentication`, `with_authentication_provider`, and opt-in `with_resilience`. `from_env()` reads `CEREBRAS_API_KEY` (required and rejected when empty before conversion to bearer authentication), `CEREBRAS_MODEL` (required), `CEREBRAS_BASE_URL`, `CEREBRAS_VERSION_PATCH`, `CEREBRAS_MAX_COMPLETION_TOKENS`. The adapter shares one `Arc` across chat, files, batch, and models so every surface uses the same authentication and resilience policy. ### 6.1 Validation (constructor / builder) - -- `api_key` non-empty. - `top_logprobs` ∈ 0..=20; requires `logprobs == Some(true)`. - `stop.len() ≤ 4`. - `temperature` ∈ 0.0..=2.0. @@ -612,7 +612,7 @@ Interactive REPL. Primary demo of the turn-loop path. Deps: `agentkit-core`, `ag | Command | Effect | | --------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | `/set ` | Live-rebuild config for the _next_ turn (covers every CLI flag) | -| `/show` | Dumps effective `CerebrasConfig` as JSON (redacts api_key) | +| `/show` | Dumps effective `CerebrasConfig` as JSON (redacts authentication) | | `/usage` | Prints last turn's `TokenUsage` + every `cerebras.*` metadata key (cached tokens, accepted/rejected prediction tokens, time_info, service_tier_used, system_fingerprint) | | `/ratelimit` | Prints `CerebrasAdapter::last_rate_limit()` (`x-ratelimit-*` snapshot) | | `/headers` | Prints request headers used on the last turn (redacted auth) — verifies `Content-Type`, `Content-Encoding`, `X-Cerebras-Version-Patch`, `queue_threshold` | @@ -625,7 +625,7 @@ Interactive REPL. Primary demo of the turn-loop path. Deps: `agentkit-core`, `ag **Banner on startup**: prints every effective knob (redacted key) so a screenshot proves the run exercised the intended config. Mirrors `anthropic-chat`'s `print_banner` shape. -**Stream decoding**: the REPL prints `delta.content` incrementally, annotates `delta.reasoning` chunks in a distinct colour, shows tool-call assembly live, and prints the terminal `Usage`/`time_info`. Covers every `ModelTurnEvent` variant. +**Stream decoding**: the REPL prints `delta.content` incrementally, annotates `delta.reasoning` chunks in a distinct colour, shows tool-call assembly live, and prints the terminal `Usage`/`time_info`. Its exhaustive `ModelTurnEvent` handling also covers `ResponseAttemptSuperseded`; if a future adapter emits that capability-gated marker, the REPL discards all deltas, tool calls, usage, and reconstruction state from the preceding attempt before rendering replacement output. #### 11.1.2 `examples/cerebras-batch/` diff --git a/docs/getting-started.md b/docs/getting-started.md index 190df17..bb7f214 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -47,10 +47,7 @@ Then: ```rust let mut driver = agent - .start(SessionConfig { - session_id: SessionId::new("demo"), - metadata: MetadataMap::new(), - }) + .start(SessionConfig::new("demo")) .await?; driver.submit_input(vec![system_item, user_item])?; diff --git a/docs/loop.md b/docs/loop.md index aaacd34..6fbf9b8 100644 --- a/docs/loop.md +++ b/docs/loop.md @@ -400,6 +400,7 @@ pub enum ModelTurnEvent { Delta(ContentDelta), ToolCall(ToolCallPart), Usage(Usage), + ResponseAttemptSuperseded, Finished(ModelTurnResult), } ``` @@ -409,8 +410,16 @@ Where: - `Delta` carries streamed content fragments - `ToolCall` surfaces a complete tool invocation request - `Usage` provides incremental or final usage updates +- `ResponseAttemptSuperseded` invalidates the preceding visible response attempt - `Finished` includes normalized finish reason and optional final assistant items +Supersession is opt-in through +`SessionConfig::with_response_attempt_supersession()`. The model event is ordered +after the failed attempt's events and before replacement output. The loop resets +its attempt-local tool-call and usage state, then emits +`AgentEvent::ResponseAttemptSuperseded`. Consumers must discard all deltas, tool +calls, usage updates, and reconstruction state from the preceding attempt. + The loop consumes these events, updates its transcript, calls observers, and decides what to do next. ## TurnRequest and TurnResult diff --git a/examples/cerebras-chat/src/main.rs b/examples/cerebras-chat/src/main.rs index e5da4a2..8e13b7b 100644 --- a/examples/cerebras-chat/src/main.rs +++ b/examples/cerebras-chat/src/main.rs @@ -338,7 +338,7 @@ fn print_show(config: &CerebrasConfig) { json.insert("model".into(), json!(config.model)); json.insert("base_url".into(), json!(config.base_url)); json.insert("streaming".into(), json!(config.streaming)); - json.insert("api_key".into(), json!("***redacted***")); + json.insert("authentication".into(), json!("***redacted***")); if let Some(v) = config.version_patch { json.insert("version_patch".into(), json!(v)); } diff --git a/projects/how/src/main.rs b/projects/how/src/main.rs index e5c459f..6e3a3f7 100644 --- a/projects/how/src/main.rs +++ b/projects/how/src/main.rs @@ -472,16 +472,12 @@ async fn run_agent( prompt: String, cancellation: CancellationController, ) -> Result, AgentError> { - let config = OpenRouterConfig::from_env() + let mut config = OpenRouterConfig::from_env() .map_err(|e| AgentError::Other(e.into()))? .with_temperature(0.0); - let config = if model != DEFAULT_MODEL { - OpenRouterConfig::new(config.api_key.clone(), &model) - .with_temperature(0.0) - .with_base_url(config.base_url.clone()) - } else { - config - }; + if model != DEFAULT_MODEL { + config.model.clone_from(&model); + } let adapter = OpenRouterAdapter::new(config).map_err(|e| AgentError::Other(e.into()))?; let tools = ToolRegistry::new().with(IsAvailableTool);