agent_sdk_providers/provider.rs
1//! LLM provider trait and streaming helpers.
2//!
3//! This module defines the [`LlmProvider`] trait that all LLM backends implement,
4//! as well as the [`collect_stream`] helper for consuming a streaming response.
5
6#[cfg(feature = "anthropic")]
7use agent_sdk_foundation::llm::ToolChoice;
8use agent_sdk_foundation::llm::{
9 ChatOutcome, ChatRequest, ChatResponse, ContentBlock, SpeedTier, ThinkingConfig, ThinkingMode,
10 Usage,
11};
12use anyhow::Result;
13use async_trait::async_trait;
14use futures::StreamExt;
15use serde::{Deserialize, Serialize};
16
17use crate::model_capabilities::{
18 ModelCapabilities, default_max_output_tokens, get_model_capabilities,
19};
20use crate::model_features::{ModelFeatures, get_model_features};
21use crate::streaming::{StreamAccumulator, StreamBox, StreamDelta, StreamErrorKind};
22
23/// How a provider satisfies a [`ResponseFormat`](agent_sdk_foundation::llm::ResponseFormat)
24/// structured-output request.
25///
26/// The structured-output runner consults this to decide how to shape the
27/// request and where to read the final structured value from the response.
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub enum StructuredOutputSupport {
30 /// The provider applies the schema natively (JSON-mode /
31 /// structured-outputs) when it sees `request.response_format`. The final
32 /// structured value is the JSON in the assistant's text output.
33 Native,
34 /// The provider has no native JSON-schema mode. The runner injects a
35 /// single forced "respond" tool whose `input_schema` is the output schema,
36 /// and reads the structured value from that tool call's input.
37 ToolForcing,
38}
39
40/// A single model entry returned by a provider's live model-listing endpoint.
41///
42/// This is the *dynamic* counterpart to the static
43/// [`ModelCapabilities`] table: it
44/// is populated from the provider's own `/models` API at runtime, so newly
45/// shipped models appear without an SDK code change. Fields beyond `id` are
46/// optional because not every provider's listing endpoint reports them.
47#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
48pub struct ModelInfo {
49 /// The model identifier as the provider's chat endpoint expects it.
50 pub id: String,
51 /// Human-friendly display name, when the listing endpoint provides one.
52 pub display_name: Option<String>,
53 /// Maximum total context window in tokens, when reported.
54 pub context_window: Option<u32>,
55 /// Maximum output tokens per response, when reported.
56 pub max_output_tokens: Option<u32>,
57}
58
59#[async_trait]
60pub trait LlmProvider: Send + Sync {
61 /// Non-streaming chat completion.
62 async fn chat(&self, request: ChatRequest) -> Result<ChatOutcome>;
63
64 /// List the models the provider currently exposes, queried live from the
65 /// provider's own model-listing endpoint.
66 ///
67 /// The default implementation returns an error: model discovery is an
68 /// additive capability, so a provider that has not implemented it stays
69 /// source-compatible while reporting that the operation is unsupported.
70 ///
71 /// # Errors
72 ///
73 /// Returns an error when the provider does not support live model listing,
74 /// or when the underlying HTTP request or response parsing fails.
75 async fn list_models(&self) -> Result<Vec<ModelInfo>> {
76 Err(anyhow::anyhow!(
77 "list_models is not supported for provider {}",
78 self.provider()
79 ))
80 }
81
82 /// Cheaply check whether the provider's endpoint is currently reachable,
83 /// without dispatching a billable request.
84 ///
85 /// Durable runtimes poll this to wait out connectivity loss: after a
86 /// transport-classified stream failure
87 /// ([`StreamErrorKind::is_connectivity`]) they probe until reachability
88 /// returns, then re-dispatch the real call. Only an unreachable answer
89 /// keeps that wait free of charge, so implementations should answer from
90 /// a cheap endpoint — [`probe_http_reachability`] against the API host —
91 /// never a model call.
92 ///
93 /// The default returns `true` ("reachable, or unknown"): a provider that
94 /// cannot observe reachability keeps its callers on the bounded retry
95 /// path instead of parking them in an offline wait it can never end.
96 async fn probe_connectivity(&self) -> bool {
97 true
98 }
99
100 /// Streaming chat completion.
101 ///
102 /// Returns a stream of [`StreamDelta`] events. The default implementation
103 /// calls [`chat()`](Self::chat) and converts the result to a single-chunk stream.
104 ///
105 /// Providers should override this method to provide true streaming support.
106 fn chat_stream(&self, request: ChatRequest) -> StreamBox<'_> {
107 Box::pin(async_stream::stream! {
108 match self.chat(request).await {
109 Ok(outcome) => match outcome {
110 ChatOutcome::Success(response) => {
111 // Emit content as deltas
112 for (idx, block) in response.content.iter().enumerate() {
113 match block {
114 ContentBlock::Text { text } => {
115 yield Ok(StreamDelta::TextDelta {
116 delta: text.clone(),
117 block_index: idx,
118 });
119 }
120 ContentBlock::Thinking { thinking, .. } => {
121 yield Ok(StreamDelta::ThinkingDelta {
122 delta: thinking.clone(),
123 block_index: idx,
124 });
125 }
126 ContentBlock::RedactedThinking { .. }
127 | ContentBlock::ToolResult { .. }
128 | ContentBlock::Image { .. }
129 | ContentBlock::Document { .. } => {
130 // Not streamed in the default implementation
131 }
132 ContentBlock::OpaqueReasoning { provider, data } => {
133 yield Ok(StreamDelta::OpaqueReasoning {
134 provider: provider.clone(),
135 data: data.clone(),
136 block_index: idx,
137 });
138 }
139 ContentBlock::ToolUse { id, name, input, thought_signature } => {
140 yield Ok(StreamDelta::ToolUseStart {
141 id: id.clone(),
142 name: name.clone(),
143 block_index: idx,
144 thought_signature: thought_signature.clone(),
145 });
146 yield Ok(StreamDelta::ToolInputDelta {
147 id: id.clone(),
148 delta: serde_json::to_string(input).unwrap_or_default(),
149 block_index: idx,
150 });
151 }
152 // `ContentBlock` is `#[non_exhaustive]`; a future
153 // block kind we cannot stream is skipped rather than
154 // panicking the default fallback.
155 _ => {
156 log::warn!(
157 "chat_stream fallback skipping unrecognized content block at index {idx}"
158 );
159 }
160 }
161 }
162 yield Ok(StreamDelta::Usage(response.usage));
163 yield Ok(StreamDelta::Done {
164 stop_reason: response.stop_reason,
165 served_route: Some(self.route().to_owned()),
166 });
167 }
168 ChatOutcome::RateLimited(retry_after) => {
169 yield Ok(StreamDelta::Error {
170 message: "Rate limited".to_string(),
171 kind: StreamErrorKind::RateLimited(retry_after),
172 });
173 }
174 ChatOutcome::InvalidRequest(msg) => {
175 yield Ok(StreamDelta::Error {
176 message: msg,
177 kind: StreamErrorKind::InvalidRequest,
178 });
179 }
180 ChatOutcome::ServerError(msg) => {
181 yield Ok(StreamDelta::Error {
182 message: msg,
183 kind: StreamErrorKind::ServerError,
184 });
185 }
186 // `ChatOutcome` is `#[non_exhaustive]`; an outcome this SDK
187 // version does not model is surfaced as an unclassified
188 // (non-recoverable) stream error rather than dropped.
189 _ => {
190 yield Ok(StreamDelta::Error {
191 message: "Unrecognized chat outcome".to_string(),
192 kind: StreamErrorKind::Unknown,
193 });
194 }
195 },
196 Err(e) => yield Err(e),
197 }
198 })
199 }
200
201 fn model(&self) -> &str;
202 fn provider(&self) -> &'static str;
203
204 /// The endpoint that actually serves this provider's calls.
205 ///
206 /// [`Self::provider`] names the *API shape* a request is built for, and
207 /// several distinct endpoints share one shape: `OpenRouter`, `Baseten` and
208 /// `Fireworks` are all reached through [`OpenAIProvider`] pointed at a
209 /// different `base_url`, so all three answer `"openai"`. Durable audit rows
210 /// need to tell *model-X via a gateway* from *model-X native*, so this
211 /// reports the serving route instead — the two coincide only for providers
212 /// that front a single endpoint, which is why the default delegates.
213 ///
214 /// This is *pre-dispatch* identity: for wrappers (fallback / router /
215 /// refresh) it names the configured representative, not necessarily the
216 /// backend a given call reached. Post-dispatch attribution rides the
217 /// stream instead — the serving provider stamps its route on
218 /// [`StreamDelta::Done`](crate::streaming::StreamDelta), and wrappers
219 /// forward it untouched.
220 ///
221 /// [`OpenAIProvider`]: crate::impls::openai::OpenAIProvider
222 fn route(&self) -> &str {
223 self.provider()
224 }
225
226 /// Provider-owned thinking configuration, if any.
227 fn configured_thinking(&self) -> Option<&ThinkingConfig> {
228 None
229 }
230
231 /// Canonical capability metadata for this provider/model, if known.
232 fn capabilities(&self) -> Option<&'static ModelCapabilities> {
233 get_model_capabilities(self.provider(), self.model()).or_else(|| match self.provider() {
234 "openai-responses" | "openai-codex" => get_model_capabilities("openai", self.model()),
235 "vertex" if self.model().starts_with("claude-") => {
236 get_model_capabilities("anthropic", self.model())
237 }
238 "vertex" => get_model_capabilities("gemini", self.model()),
239 _ => None,
240 })
241 }
242
243 /// API-shape feature metadata for this provider/model, when known.
244 ///
245 /// This complements [`Self::capabilities`] with endpoint, reasoning,
246 /// caching, and tool-control information that does not fit the legacy
247 /// context/pricing record.
248 fn model_features(&self) -> Option<&'static ModelFeatures> {
249 match self.provider() {
250 "openai" | "openai-responses" | "openai-codex" => get_model_features(self.model()),
251 _ => None,
252 }
253 }
254
255 /// Validate a thinking configuration against the provider/model capabilities.
256 ///
257 /// # Errors
258 ///
259 /// Returns an error when the requested thinking mode is not supported by
260 /// the active provider/model capability set.
261 fn validate_thinking_config(&self, thinking: Option<&ThinkingConfig>) -> Result<()> {
262 let Some(thinking) = thinking else {
263 return Ok(());
264 };
265
266 if self
267 .capabilities()
268 .is_some_and(|caps| !caps.supports_thinking)
269 {
270 return Err(anyhow::anyhow!(
271 "thinking is not supported for provider={} model={}",
272 self.provider(),
273 self.model()
274 ));
275 }
276
277 if matches!(thinking.mode, ThinkingMode::Adaptive)
278 && !self
279 .capabilities()
280 .is_some_and(|caps| caps.supports_adaptive_thinking)
281 {
282 return Err(anyhow::anyhow!(
283 "adaptive thinking is not supported for provider={} model={}",
284 self.provider(),
285 self.model()
286 ));
287 }
288
289 Ok(())
290 }
291
292 /// Provider-owned speed tier, when the caller configured one.
293 ///
294 /// `None` means the provider never opted in and every request runs on the
295 /// standard path.
296 fn configured_speed(&self) -> Option<SpeedTier> {
297 None
298 }
299
300 /// Validate a speed tier against the provider/model.
301 ///
302 /// The default implementation rejects every premium tier, so a provider
303 /// only serves one after opting in explicitly. Refusing is the safe
304 /// default: a premium tier the provider cannot honour would otherwise be
305 /// dropped at the wire boundary, and the caller would silently get
306 /// standard latency from a request they believe is expedited.
307 ///
308 /// # Errors
309 ///
310 /// Returns an error when the requested tier is not available for the
311 /// active provider/model.
312 fn validate_speed_tier(&self, speed: Option<SpeedTier>) -> Result<()> {
313 match speed {
314 None => Ok(()),
315 Some(tier) if !tier.is_premium() => Ok(()),
316 Some(_) => Err(anyhow::anyhow!(
317 "a premium speed tier is not supported for provider={} model={}",
318 self.provider(),
319 self.model()
320 )),
321 }
322 }
323
324 /// Resolve the effective speed tier for a request.
325 ///
326 /// # Errors
327 ///
328 /// Returns an error when the resolved tier is not available for the active
329 /// provider/model.
330 fn resolve_speed_tier(&self) -> Result<Option<SpeedTier>> {
331 let speed = self.configured_speed();
332 self.validate_speed_tier(speed)?;
333 Ok(speed)
334 }
335
336 /// Resolve the effective thinking configuration for a request.
337 ///
338 /// Request-level thinking overrides provider-owned defaults when present.
339 ///
340 /// # Errors
341 ///
342 /// Returns an error when the resolved thinking configuration is not
343 /// supported by the active provider/model capability set.
344 fn resolve_thinking_config(
345 &self,
346 request_thinking: Option<&ThinkingConfig>,
347 ) -> Result<Option<ThinkingConfig>> {
348 let thinking = request_thinking.or_else(|| self.configured_thinking());
349 self.validate_thinking_config(thinking)?;
350 Ok(thinking.cloned())
351 }
352
353 /// Default maximum output tokens for this provider/model when the caller
354 /// does not explicitly override `AgentConfig.max_tokens`.
355 fn default_max_tokens(&self) -> u32 {
356 self.capabilities()
357 .and_then(|caps| caps.max_output_tokens)
358 .or_else(|| default_max_output_tokens(self.provider(), self.model()))
359 .unwrap_or(4096)
360 }
361
362 /// How this provider satisfies a structured-output
363 /// ([`ResponseFormat`](agent_sdk_foundation::llm::ResponseFormat)) request.
364 ///
365 /// Providers with a native JSON-schema / JSON-mode wire field
366 /// (OpenAI-family, Gemini, Vertex) report
367 /// [`StructuredOutputSupport::Native`] and consume
368 /// `request.response_format` directly. Providers without one (Anthropic)
369 /// report [`StructuredOutputSupport::ToolForcing`] so the runner forces a
370 /// single "respond" tool whose schema is the output schema. The default
371 /// is the conservative [`StructuredOutputSupport::ToolForcing`], which
372 /// works for any tool-capable provider.
373 fn structured_output_support(&self) -> StructuredOutputSupport {
374 match self.provider() {
375 "openai" | "openai-responses" | "openai-codex" | "gemini" => {
376 StructuredOutputSupport::Native
377 }
378 // Vertex multiplexes Anthropic and Gemini models. Only the Gemini
379 // side has a native structured-output field; Claude-on-Vertex uses
380 // the Messages API shape, which has no `response_format`.
381 "vertex" if !self.model().starts_with("claude-") => StructuredOutputSupport::Native,
382 _ => StructuredOutputSupport::ToolForcing,
383 }
384 }
385}
386
387/// Timeout for [`probe_http_reachability`]: long enough for a slow TLS
388/// handshake, short enough that a half-open network cannot stall the
389/// connectivity-wait loop it drives.
390const PROBE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
391
392/// `true` when an HTTP response arrives from `url` over a fully-established
393/// (and, for HTTPS, certificate-validated) connection.
394///
395/// The status code is deliberately ignored — a 4xx from the provider's edge
396/// proves the network path exactly as well as a 200 does, and needs no
397/// credentials. Every transport failure (DNS, connect, TLS, timeout) reports
398/// unreachable: this feeds offline-wait loops, where the fail-safe direction
399/// is to keep waiting. A captive portal that hijacks the hostname fails
400/// certificate validation and therefore stays unreachable until the real
401/// endpoint is back.
402pub async fn probe_http_reachability(client: &reqwest::Client, url: &str) -> bool {
403 client.head(url).timeout(PROBE_TIMEOUT).send().await.is_ok()
404}
405
406/// Drop extended thinking when the request forces a specific tool.
407///
408/// Anthropic-family models (native Anthropic and Claude-on-Vertex) reject a
409/// request that pairs extended `thinking` with a `tool_choice` that names a
410/// tool (HTTP 400). The forced-tool constraint is orthogonal to *where* the
411/// thinking came from: clearing `ChatRequest.thinking` upstream is not enough,
412/// because [`resolve_thinking_config`](LlmProvider::resolve_thinking_config)
413/// falls back to the provider-configured default when the request field is
414/// `None` — resurrecting thinking on the wire. The Claude request builders
415/// therefore funnel their already-resolved thinking through this at the wire
416/// boundary so a forced-tool request (e.g. structured-output's `respond` tool)
417/// stays well-formed regardless of the provider's configured thinking default.
418///
419/// `ToolChoice::Auto` (and the absence of any `tool_choice`) keeps thinking,
420/// matching the API's rule that thinking is only incompatible with a
421/// tool-forcing choice.
422#[must_use]
423#[cfg(feature = "anthropic")]
424pub(crate) const fn thinking_for_forced_tool(
425 thinking: Option<ThinkingConfig>,
426 tool_choice: Option<&ToolChoice>,
427) -> Option<ThinkingConfig> {
428 match tool_choice {
429 Some(ToolChoice::Tool(_)) => None,
430 Some(ToolChoice::Auto) | None => thinking,
431 }
432}
433
434/// Helper function to consume a stream and collect it into a `ChatResponse`.
435///
436/// This is useful for providers that want to test their streaming implementation
437/// or for cases where you need the full response after streaming.
438///
439/// # Errors
440///
441/// Returns an error if the stream yields an error result.
442pub async fn collect_stream(mut stream: StreamBox<'_>, model: String) -> Result<ChatOutcome> {
443 let mut accumulator = StreamAccumulator::new();
444 let mut last_error: Option<(String, StreamErrorKind)> = None;
445
446 while let Some(result) = stream.next().await {
447 match result {
448 Ok(delta) => {
449 if let StreamDelta::Error { message, kind } = &delta {
450 last_error = Some((message.clone(), *kind));
451 }
452 accumulator.apply(&delta);
453 }
454 Err(e) => return Err(e),
455 }
456 }
457
458 // If we encountered an error during streaming, map kind directly
459 // to the corresponding `ChatOutcome` variant. No string-matching
460 // heuristic is needed because the kind already records the
461 // category at the construction site.
462 if let Some((message, kind)) = last_error {
463 return Ok(match kind {
464 // The rate-limit kind carries whatever retry hint the provider gave
465 // (header or error body), so the reconstructed outcome honours the
466 // server exactly like the non-streaming `chat()` path does.
467 StreamErrorKind::RateLimited(retry_after) => ChatOutcome::RateLimited(retry_after),
468 StreamErrorKind::InvalidRequest => ChatOutcome::InvalidRequest(message),
469 // `StreamErrorKind::ServerError`, plus the `#[non_exhaustive]`
470 // catch-all (`Unknown` / future kinds): an unclassified error is
471 // treated as a (non-recoverable) server error so the caller still
472 // surfaces the failure rather than silently succeeding.
473 _ => ChatOutcome::ServerError(message),
474 });
475 }
476
477 // Extract usage and stop_reason before consuming the accumulator
478 let usage = accumulator.take_usage().unwrap_or(Usage {
479 served_speed: None,
480 input_tokens: 0,
481 output_tokens: 0,
482 cached_input_tokens: 0,
483 cache_creation_input_tokens: 0,
484 });
485 let stop_reason = accumulator.take_stop_reason();
486 let content = accumulator.into_content_blocks();
487
488 // Log accumulated response for debugging
489 log::debug!(
490 "Collected stream response: model={} stop_reason={:?} usage={{input_tokens={}, output_tokens={}}} content_blocks={}",
491 model,
492 stop_reason,
493 usage.input_tokens,
494 usage.output_tokens,
495 content.len()
496 );
497 for (i, block) in content.iter().enumerate() {
498 match block {
499 ContentBlock::Text { text } => {
500 log::debug!(" content_block[{}]: Text (len={})", i, text.len());
501 }
502 ContentBlock::Thinking { thinking, .. } => {
503 log::debug!(" content_block[{}]: Thinking (len={})", i, thinking.len());
504 }
505 ContentBlock::RedactedThinking { .. } => {
506 log::debug!(" content_block[{i}]: RedactedThinking");
507 }
508 ContentBlock::OpaqueReasoning { provider, .. } => {
509 log::debug!(
510 " content_block[{i}]: OpaqueReasoning provider={provider} payload=<redacted>"
511 );
512 }
513 ContentBlock::ToolUse {
514 id, name, input, ..
515 } => {
516 log::debug!(" content_block[{i}]: ToolUse id={id} name={name} input={input}");
517 }
518 ContentBlock::ToolResult {
519 tool_use_id,
520 content: result_content,
521 is_error,
522 } => {
523 log::debug!(
524 " content_block[{}]: ToolResult tool_use_id={} is_error={:?} content_len={}",
525 i,
526 tool_use_id,
527 is_error,
528 result_content.len()
529 );
530 }
531 ContentBlock::Image { source } => {
532 log::debug!(
533 " content_block[{i}]: Image media_type={}",
534 source.media_type
535 );
536 }
537 ContentBlock::Document { source } => {
538 log::debug!(
539 " content_block[{i}]: Document media_type={}",
540 source.media_type
541 );
542 }
543 // `ContentBlock` is `#[non_exhaustive]`; log unknown future block
544 // kinds generically so the debug dump stays exhaustive.
545 _ => {
546 log::debug!(" content_block[{i}]: <unrecognized block kind>");
547 }
548 }
549 }
550
551 Ok(ChatOutcome::Success(ChatResponse {
552 id: String::new(),
553 content,
554 model,
555 stop_reason,
556 usage,
557 }))
558}
559
560#[cfg(test)]
561mod tests {
562 use super::*;
563 use anyhow::Result;
564 use async_trait::async_trait;
565
566 struct Stub {
567 provider: &'static str,
568 model: &'static str,
569 }
570
571 #[async_trait]
572 impl LlmProvider for Stub {
573 async fn chat(&self, _request: ChatRequest) -> Result<ChatOutcome> {
574 Ok(ChatOutcome::ServerError("unused".to_owned()))
575 }
576
577 fn model(&self) -> &str {
578 self.model
579 }
580
581 fn provider(&self) -> &'static str {
582 self.provider
583 }
584 }
585
586 fn support_for(provider: &'static str, model: &'static str) -> StructuredOutputSupport {
587 Stub { provider, model }.structured_output_support()
588 }
589
590 #[test]
591 fn native_providers_report_native_support() {
592 for provider in ["openai", "openai-responses", "openai-codex", "gemini"] {
593 assert_eq!(
594 support_for(provider, "any-model"),
595 StructuredOutputSupport::Native,
596 "{provider} should be native"
597 );
598 }
599 }
600
601 #[test]
602 fn anthropic_reports_tool_forcing() {
603 assert_eq!(
604 support_for("anthropic", "claude-sonnet-4-5"),
605 StructuredOutputSupport::ToolForcing
606 );
607 }
608
609 #[test]
610 fn vertex_is_native_for_gemini_models_and_tool_forcing_for_claude() {
611 assert_eq!(
612 support_for("vertex", "gemini-3-flash-preview"),
613 StructuredOutputSupport::Native
614 );
615 assert_eq!(
616 support_for("vertex", "claude-sonnet-4-5"),
617 StructuredOutputSupport::ToolForcing
618 );
619 }
620
621 #[test]
622 fn unknown_provider_defaults_to_tool_forcing() {
623 assert_eq!(
624 support_for("some-new-provider", "x"),
625 StructuredOutputSupport::ToolForcing
626 );
627 }
628
629 #[test]
630 #[cfg(feature = "anthropic")]
631 fn thinking_for_forced_tool_drops_thinking_when_a_tool_is_forced() {
632 let cfg = ThinkingConfig::new(10_000);
633 let forced = ToolChoice::Tool("respond".to_owned());
634 assert!(
635 thinking_for_forced_tool(Some(cfg), Some(&forced)).is_none(),
636 "thinking must be dropped when tool_choice names a tool"
637 );
638 }
639
640 #[test]
641 #[cfg(feature = "anthropic")]
642 fn thinking_for_forced_tool_keeps_thinking_for_auto_and_none() {
643 let auto = ToolChoice::Auto;
644 assert!(
645 thinking_for_forced_tool(Some(ThinkingConfig::new(10_000)), Some(&auto)).is_some(),
646 "thinking must survive with ToolChoice::Auto"
647 );
648 assert!(
649 thinking_for_forced_tool(Some(ThinkingConfig::new(10_000)), None).is_some(),
650 "thinking must survive with no tool_choice"
651 );
652 }
653
654 #[test]
655 #[cfg(feature = "anthropic")]
656 fn thinking_for_forced_tool_is_a_noop_when_thinking_absent() {
657 let forced = ToolChoice::Tool("respond".to_owned());
658 assert!(thinking_for_forced_tool(None, Some(&forced)).is_none());
659 }
660
661 fn stream_of(deltas: Vec<StreamDelta>) -> StreamBox<'static> {
662 Box::pin(futures::stream::iter(deltas.into_iter().map(Ok)))
663 }
664
665 #[tokio::test]
666 async fn collect_stream_preserves_the_rate_limit_retry_hint() -> Result<()> {
667 let hint = std::time::Duration::from_secs(30);
668 let outcome = collect_stream(
669 stream_of(vec![StreamDelta::Error {
670 message: "Rate limited".to_owned(),
671 kind: StreamErrorKind::RateLimited(Some(hint)),
672 }]),
673 "test-model".to_owned(),
674 )
675 .await?;
676
677 assert!(
678 matches!(outcome, ChatOutcome::RateLimited(Some(delay)) if delay == hint),
679 "streamed Retry-After must survive reconstruction, got {outcome:?}"
680 );
681 Ok(())
682 }
683
684 #[tokio::test]
685 async fn collect_stream_reports_no_hint_when_the_provider_gave_none() -> Result<()> {
686 let outcome = collect_stream(
687 stream_of(vec![StreamDelta::Error {
688 message: "Rate limited".to_owned(),
689 kind: StreamErrorKind::RateLimited(None),
690 }]),
691 "test-model".to_owned(),
692 )
693 .await?;
694
695 assert!(
696 matches!(outcome, ChatOutcome::RateLimited(None)),
697 "a hintless rate limit must not invent a delay, got {outcome:?}"
698 );
699 Ok(())
700 }
701
702 #[tokio::test]
703 async fn default_chat_stream_carries_the_chat_retry_hint_into_the_error_delta() -> Result<()> {
704 // A provider that only implements `chat()` gets its streaming from the
705 // trait's default adapter; a rate limit's hint must not be lost there.
706 struct RateLimitedStub {
707 model: &'static str,
708 }
709
710 #[async_trait]
711 impl LlmProvider for RateLimitedStub {
712 async fn chat(&self, _request: ChatRequest) -> Result<ChatOutcome> {
713 Ok(ChatOutcome::RateLimited(Some(
714 std::time::Duration::from_secs(12),
715 )))
716 }
717 fn model(&self) -> &str {
718 self.model
719 }
720 fn provider(&self) -> &'static str {
721 "stub"
722 }
723 }
724
725 let stub = RateLimitedStub {
726 model: "stub-model",
727 };
728 let request = ChatRequest::new("", vec![agent_sdk_foundation::llm::Message::user("hi")]);
729 let outcome =
730 collect_stream(Box::pin(stub.chat_stream(request)), "stub-model".to_owned()).await?;
731
732 assert!(
733 matches!(outcome, ChatOutcome::RateLimited(Some(delay))
734 if delay == std::time::Duration::from_secs(12)),
735 "default chat_stream must forward chat()'s Retry-After, got {outcome:?}"
736 );
737 Ok(())
738 }
739
740 #[tokio::test]
741 async fn collect_stream_maps_other_kinds_without_a_hint() -> Result<()> {
742 let outcome = collect_stream(
743 stream_of(vec![StreamDelta::Error {
744 message: "boom".to_owned(),
745 kind: StreamErrorKind::ServerError,
746 }]),
747 "test-model".to_owned(),
748 )
749 .await?;
750
751 assert!(matches!(outcome, ChatOutcome::ServerError(msg) if msg == "boom"));
752 Ok(())
753 }
754}