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