Skip to main content

molo_core/provider/
mod.rs

1//! Provider: the interface for communicating with an LLM.
2//!
3//! This module defines the `Provider` trait and its companion data types:
4//! requests ([`ChatRequest`]), responses ([`ChatResponse`] / [`StreamEvent`]),
5//! errors ([`ProviderError`]), capabilities ([`ProviderCapabilities`]),
6//! request context ([`ProviderRequestContext`]), and usage ([`Usage`]). The
7//! trait itself is vendor-agnostic; the lightweight implementations
8//! ([`FakeProvider`] and [`RetryProvider`]) live here. Concrete network
9//! providers live in optional adapter crates such as `molo-openai`.
10//!
11//! The provider contract is intentionally explicit: successful `chat`
12//! returns exactly one assistant message, successful streams terminate with
13//! one `Done`, usage is reported only when the provider supplies it, and
14//! local decode/protocol/size-limit failures are distinct from vendor API
15//! errors. Context-aware methods are the runtime boundary: runtimes pass run
16//! ids, model request ids, deadlines, cancellation, and sanitized metadata to
17//! providers explicitly. The plain `chat` and `stream_chat` methods remain
18//! convenience entry points for direct provider use.
19
20mod fake;
21mod retry;
22
23pub use fake::{FakeProvider, FakeReply};
24pub use retry::{Backoff, RetryPolicy, RetryProvider, Retryable};
25
26use crate::message::Message;
27use crate::run::{RunContext, RunMetadata};
28use crate::tool::ToolSchema;
29use futures::stream::BoxStream;
30use serde::{Deserialize, Serialize};
31use std::collections::BTreeMap;
32use std::fmt;
33use std::time::{Duration, Instant};
34
35/// Provider capability metadata used by hosts and conformance tests.
36///
37/// Capabilities are descriptive, not a security boundary. When a provider
38/// declares support for an optional capability, it should pass the matching
39/// provider conformance cases. Unsupported direct calls should return
40/// [`ProviderError::Unsupported`] rather than panic.
41#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
42pub struct ProviderCapabilities {
43    /// Supports [`Provider::stream_chat`] / [`Provider::stream_chat_with_context`].
44    pub streaming: bool,
45    /// Maps provider reasoning/thinking output to [`Message`] or
46    /// [`StreamEvent::Reasoning`].
47    pub reasoning: bool,
48    /// Supports model tool-call requests.
49    pub tool_calls: bool,
50    /// Preserves multiple tool calls in one assistant turn.
51    pub parallel_tool_calls: bool,
52    /// Accepts [`ModelOptions::structured`] as a provider-side best-effort
53    /// constraint.
54    pub structured_output: bool,
55    /// Reports provider token usage when the backend supplies it.
56    pub usage: bool,
57    /// Cooperatively observes cancellation from [`ProviderRequestContext`].
58    pub context_cancellation: bool,
59    /// Cooperatively observes deadlines from [`ProviderRequestContext`].
60    pub context_deadline: bool,
61}
62
63impl ProviderCapabilities {
64    /// Baseline provider capabilities: non-streaming text only.
65    pub fn baseline() -> Self {
66        Self::default()
67    }
68}
69
70/// Request-scoped provider context.
71///
72/// This is the provider-boundary projection of [`RunContext`]: it carries
73/// correlation ids, cancellation/deadline controls, an optional timeout hint,
74/// and sanitized host metadata. Raw prompts, source code, auth headers, API
75/// keys, and environment values should not be placed in metadata by default.
76#[derive(Clone)]
77pub struct ProviderRequestContext {
78    /// Run id shared with run summaries, event records, and tracing spans.
79    pub run_id: String,
80    /// Model request id unique within the run.
81    pub model_request_id: String,
82    /// Cooperative cancellation source.
83    pub cancellation: tokio_util::sync::CancellationToken,
84    /// Optional absolute deadline.
85    pub deadline: Option<Instant>,
86    /// Optional per-provider-call timeout hint.
87    pub timeout: Option<Duration>,
88    /// Host/framework metadata for observability and routing.
89    pub metadata: RunMetadata,
90}
91
92impl fmt::Debug for ProviderRequestContext {
93    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
94        f.debug_struct("ProviderRequestContext")
95            .field("run_id", &self.run_id)
96            .field("model_request_id", &self.model_request_id)
97            .field("cancellation", &"CancellationToken")
98            .field("deadline", &self.deadline)
99            .field("timeout", &self.timeout)
100            .field("metadata", &self.metadata)
101            .finish()
102    }
103}
104
105impl ProviderRequestContext {
106    /// Builds provider context from a run context and model request id.
107    pub fn from_run_context(model_request_id: impl Into<String>, context: &RunContext) -> Self {
108        Self {
109            run_id: context.run_id.clone(),
110            model_request_id: model_request_id.into(),
111            cancellation: context.cancellation.clone(),
112            deadline: context.deadline,
113            timeout: context.remaining(),
114            metadata: context.metadata.clone(),
115        }
116    }
117
118    /// Sets a provider-call timeout hint.
119    pub fn with_timeout(mut self, timeout: Duration) -> Self {
120        self.timeout = Some(timeout);
121        self
122    }
123
124    /// Whether cancellation has already been requested.
125    pub fn is_cancelled(&self) -> bool {
126        self.cancellation.is_cancelled()
127    }
128
129    /// Whether the deadline has elapsed.
130    pub fn is_expired(&self) -> bool {
131        self.deadline
132            .is_some_and(|deadline| Instant::now() >= deadline)
133    }
134
135    /// Remaining time before the deadline.
136    pub fn remaining(&self) -> Option<Duration> {
137        self.deadline
138            .map(|deadline| deadline.saturating_duration_since(Instant::now()))
139    }
140}
141
142/// The interface for chatting with an LLM.
143///
144/// Implementations are responsible for communicating with a specific LLM
145/// service and mapping vendor responses back to this framework's [`Message`];
146/// [`chat`](Provider::chat) returns the full reply at once, while
147/// [`stream_chat`](Provider::stream_chat) returns the same reply incrementally
148/// as a stream of events. Both share the same semantics and differ only in
149/// delivery.
150///
151/// `Send + Sync` guarantees that `Box<dyn Provider>` can be held across
152/// awaits in Agent implementations (for the same reason as
153/// [`Tool`](crate::tool::Tool)).
154///
155/// # Examples
156///
157/// The calling convention is identical for every implementation; the example
158/// below uses [`FakeProvider`]:
159///
160/// ```rust
161/// # extern crate molo_core as molo;
162/// # #[tokio::main]
163/// # async fn main() -> Result<(), molo::ProviderError> {
164/// use molo::provider::{ChatRequest, FakeProvider, FakeReply, Provider};
165///
166/// let fake = FakeProvider::new([FakeReply::Text("hi".into())]);
167/// let response = fake.chat(ChatRequest::default()).await?;
168/// assert_eq!(response.message, molo::message::Message::assistant("hi"));
169/// # Ok(())
170/// # }
171/// ```
172#[async_trait::async_trait]
173pub trait Provider: Send + Sync {
174    /// Model identifier exposed by this provider, when known.
175    ///
176    /// Agents copy this value into run summaries for observability. Providers
177    /// that are not bound to one model can keep the default `None`.
178    fn model(&self) -> Option<&str> {
179        None
180    }
181
182    /// Capability metadata for this provider instance.
183    fn capabilities(&self) -> ProviderCapabilities {
184        ProviderCapabilities::baseline()
185    }
186
187    /// Sends one turn with request-scoped provider context.
188    ///
189    /// # Errors
190    ///
191    /// Network failures / timeouts / rate limits / vendor business errors are
192    /// all returned as [`ProviderError`]; see that type's docs for error
193    /// classification and retry guidance.
194    async fn chat_with_context(
195        &self,
196        request: ChatRequest,
197        context: &ProviderRequestContext,
198    ) -> Result<ChatResponse, ProviderError>;
199
200    /// Sends one turn of conversation and returns the model's reply (text, or
201    /// a request to call tools).
202    ///
203    /// This direct-use convenience wrapper creates a generated run context.
204    /// Runtimes that already have a [`RunContext`] should call
205    /// [`chat_with_context`](Provider::chat_with_context).
206    async fn chat(&self, request: ChatRequest) -> Result<ChatResponse, ProviderError> {
207        let run = RunContext::generated();
208        let context = ProviderRequestContext::from_run_context("direct-chat", &run);
209        self.chat_with_context(request, &context).await
210    }
211
212    /// Streams one turn of conversation for direct provider use.
213    ///
214    /// Semantically identical to [`chat`](Provider::chat), except that the
215    /// reply is returned as a stream of events: several [`StreamEvent::Delta`]
216    /// items concatenated in order form the full reply, and the stream ends
217    /// with [`StreamEvent::Done`].
218    ///
219    /// # Errors
220    ///
221    /// Failures during request setup (connection / timeout / vendor rejection)
222    /// are returned as `Err`; event errors after the stream is established are
223    /// produced as `Err` items in the stream, and no success events are
224    /// produced after an error item (see [`StreamEvent`] for termination
225    /// semantics).
226    async fn stream_chat(
227        &self,
228        request: ChatRequest,
229    ) -> Result<BoxStream<'static, Result<StreamEvent, ProviderError>>, ProviderError> {
230        let run = RunContext::generated();
231        let context = ProviderRequestContext::from_run_context("direct-stream", &run);
232        self.stream_chat_with_context(request, &context).await
233    }
234
235    /// Streams one turn with request-scoped provider context.
236    ///
237    /// Runtimes should call this method so providers can observe cancellation,
238    /// deadlines, request ids, and sanitized metadata.
239    async fn stream_chat_with_context(
240        &self,
241        request: ChatRequest,
242        context: &ProviderRequestContext,
243    ) -> Result<BoxStream<'static, Result<StreamEvent, ProviderError>>, ProviderError>;
244}
245
246/// `Box<dyn Provider>` is itself a Provider: re-exposes the trait object as a
247/// value, for assembly patterns that need to "hold an instance and create a
248/// new loop per call" (e.g., a sub-agent factory that captures a provider and
249/// constructs a fresh loop for each invocation).
250#[async_trait::async_trait]
251impl Provider for Box<dyn Provider> {
252    fn model(&self) -> Option<&str> {
253        self.as_ref().model()
254    }
255
256    fn capabilities(&self) -> ProviderCapabilities {
257        self.as_ref().capabilities()
258    }
259
260    async fn chat(&self, request: ChatRequest) -> Result<ChatResponse, ProviderError> {
261        self.as_ref().chat(request).await
262    }
263
264    async fn chat_with_context(
265        &self,
266        request: ChatRequest,
267        context: &ProviderRequestContext,
268    ) -> Result<ChatResponse, ProviderError> {
269        self.as_ref().chat_with_context(request, context).await
270    }
271
272    async fn stream_chat(
273        &self,
274        request: ChatRequest,
275    ) -> Result<BoxStream<'static, Result<StreamEvent, ProviderError>>, ProviderError> {
276        self.as_ref().stream_chat(request).await
277    }
278
279    async fn stream_chat_with_context(
280        &self,
281        request: ChatRequest,
282        context: &ProviderRequestContext,
283    ) -> Result<BoxStream<'static, Result<StreamEvent, ProviderError>>, ProviderError> {
284        self.as_ref()
285            .stream_chat_with_context(request, context)
286            .await
287    }
288}
289
290/// A single conversation request.
291///
292/// # Examples
293///
294/// ```rust
295/// # extern crate molo_core as molo;
296/// use molo::message::Message;
297/// use molo::provider::ChatRequest;
298///
299/// let request = ChatRequest {
300///     messages: vec![Message::user("hi")],
301///     ..Default::default()
302/// };
303/// # let _ = request;
304/// ```
305#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
306pub struct ChatRequest {
307    /// Conversation history (in order of occurrence, assembled by the caller).
308    pub messages: Vec<Message>,
309    /// Tool definitions offered to the model; when empty, the model sees no
310    /// tools.
311    pub tools: Vec<ToolSchema>,
312    /// Model options; `Default` means all vendor defaults.
313    pub options: ModelOptions,
314}
315
316/// Model options for one conversation.
317///
318/// Common parameters are provided as typed fields (temperature / max tokens,
319/// where `None` means vendor default); **vendor-specific or framework-unknown
320/// parameters go into [`extra`](ModelOptions::extra)** and are passed through
321/// to the vendor verbatim under their wire field names — so users can use new
322/// parameters without waiting for a framework update:
323///
324/// ```rust
325/// # extern crate molo_core as molo;
326/// use molo::ModelOptions;
327///
328/// let mut options = ModelOptions::default();
329/// options.extra.insert("top_p".into(), serde_json::json!(0.9));
330/// ```
331///
332/// Extra keys that collide with framework-managed fields are ignored in favor
333/// of the typed fields.
334#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
335pub struct ModelOptions {
336    /// Sampling temperature; `None` means vendor default.
337    pub temperature: Option<f32>,
338    /// Maximum number of tokens for the reply; `None` means vendor default.
339    pub max_tokens: Option<u32>,
340    /// Vendor extension parameters: keys are wire request field names,
341    /// serialized into the request body verbatim.
342    pub extra: BTreeMap<String, serde_json::Value>,
343    /// Structured output: the final answer must be JSON conforming to this
344    /// **JSON Schema document** (the serialized `RootSchema` produced by
345    /// schemars, or a hand-written schema).
346    ///
347    /// Two layers of semantics:
348    /// - Provider side: compatible endpoints receive it via `response_format`
349    ///   to best-effort constrain the model (the OpenAI-compatible
350    ///   `json_schema` shape; unsupported endpoints ignore it or error);
351    /// - Agent side: the final answer is **validated framework-side**, and on
352    ///   mismatch the validation error is fed back to the model for a retry
353    ///   (counted against the turn budget by the agent runtime's structured
354    ///   output support).
355    ///
356    /// `None` = free-form text reply.
357    pub structured: Option<serde_json::Value>,
358}
359
360/// Token usage for one conversation.
361///
362/// Field names match the OpenAI wire format; `total_tokens` follows the
363/// vendor's convention (not necessarily the sum of the other two). `Default`
364/// = all zeros, with nothing reported.
365///
366/// Presence is carried by the enclosing type: [`ChatResponse::usage`] /
367/// [`StreamEvent::Done::usage`] are `Option<Usage>` — `None` means the
368/// endpoint did not report usage for this turn, `Some` means the reported
369/// values. The distinction matters for observability: "not reported" is not
370/// the same as "reportedly zero" (the Agent layer also tracks it in
371/// [`RunSummary::usage_omitted`](crate::run::RunSummary)).
372///
373/// [`cached_tokens`](Self::cached_tokens) carries the same distinction inside
374/// a reported usage: the prompt-cache breakdown is optional per turn even
375/// when the totals are present.
376#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
377pub struct Usage {
378    /// Input tokens for this turn.
379    pub prompt_tokens: u32,
380    /// Output tokens for this turn.
381    pub completion_tokens: u32,
382    /// Total for this turn (vendor convention).
383    pub total_tokens: u32,
384    /// Prompt tokens served from the provider's prompt cache.
385    ///
386    /// A subset of [`prompt_tokens`](Self::prompt_tokens), not an addition to
387    /// it (vendors bill cached input at a lower rate).
388    ///
389    /// `None` = the endpoint did not report a cache breakdown for this turn;
390    /// `Some(0)` = it reported that nothing hit the cache. Endpoints differ
391    /// here — prompt-caching ones report the breakdown (OpenAI / DeepSeek as
392    /// `prompt_tokens_details.cached_tokens`), others omit it — so a reported
393    /// zero must not be inferred from an omission. When turns are summed,
394    /// only the reported parts contribute and
395    /// [`RunSummary::cache_omitted`](crate::run::RunSummary) marks whether
396    /// the sum is a lower bound.
397    pub cached_tokens: Option<u32>,
398}
399
400impl Usage {
401    /// Constructs from input / output counts; the total is summed
402    /// automatically and no cache breakdown is reported
403    /// ([`with_cached_tokens`](Self::with_cached_tokens) sets one).
404    pub fn new(prompt_tokens: u32, completion_tokens: u32) -> Self {
405        Self {
406            prompt_tokens,
407            completion_tokens,
408            total_tokens: prompt_tokens + completion_tokens,
409            cached_tokens: None,
410        }
411    }
412
413    /// Sets the prompt tokens served from the provider's cache, marking the
414    /// breakdown as reported (`Some` in [`cached_tokens`](Self::cached_tokens)).
415    pub fn with_cached_tokens(mut self, cached_tokens: u32) -> Self {
416        self.cached_tokens = Some(cached_tokens);
417        self
418    }
419}
420
421/// Usage accumulates per turn (the Agent layer sums tokens across turns).
422///
423/// `cached_tokens` sums the parts that were reported: a turn that omitted the
424/// breakdown contributes nothing rather than zero.
425impl std::ops::AddAssign for Usage {
426    fn add_assign(&mut self, rhs: Self) {
427        self.prompt_tokens += rhs.prompt_tokens;
428        self.completion_tokens += rhs.completion_tokens;
429        self.total_tokens += rhs.total_tokens;
430        self.cached_tokens = match (self.cached_tokens, rhs.cached_tokens) {
431            (Some(lhs), Some(rhs)) => Some(lhs + rhs),
432            (lhs, rhs) => lhs.or(rhs),
433        };
434    }
435}
436
437/// The reply to one conversation.
438#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
439pub struct ChatResponse {
440    /// This turn's reply: a single [`Message::Assistant`] message
441    /// (text + reasoning + any tool requests from the same turn stay together,
442    /// one-to-one with the vendor wire structure);
443    /// when the model produces nothing it is still an empty Assistant.
444    ///
445    /// The Agent loop executes tool requests when it sees them and appends the
446    /// results as [`Message::ToolResult`] before continuing the conversation.
447    pub message: Message,
448    /// Why the model ended its reply; vendor-specific reasons are surfaced via
449    /// [`FinishReason::Other`].
450    pub finish_reason: FinishReason,
451    /// Token usage for this turn; `None` when the endpoint did not return it
452    /// (compatible endpoints may omit usage; see [`Usage`] for the presence
453    /// semantics shared with [`StreamEvent::Done::usage`]).
454    pub usage: Option<Usage>,
455}
456
457/// Why the model ended its reply.
458///
459/// Common reasons are typed (Stop / Length); vendor-specific or
460/// framework-unknown reasons are surfaced via [`Other`](FinishReason::Other)
461/// carrying the vendor's raw string — users can recognize new reasons without
462/// waiting for a framework update. `#[non_exhaustive]` guarantees that adding
463/// new common categories in the future is not a breaking change.
464#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
465#[non_exhaustive]
466pub enum FinishReason {
467    /// The model ended naturally (including ending the turn by requesting
468    /// tool calls).
469    Stop,
470    /// Truncated by hitting the max_tokens limit.
471    Length,
472    /// A vendor-specific or framework-unknown reason, carrying the vendor's
473    /// raw string.
474    Other(String),
475}
476
477/// An event in a streamed conversation reply.
478///
479/// One streamed reply = several [`StreamEvent::Delta`] /
480/// [`StreamEvent::ToolCall`] increments + one closing [`StreamEvent::Done`];
481/// the caller concatenates the Deltas in order to get the full reply.
482///
483/// Stream termination semantics: on normal termination `Done` is always the
484/// last success event on the stream; errors terminate the stream with an
485/// `Err` item, and no events are produced after the error item.
486///
487/// The enum is `#[non_exhaustive]` (reserved for extension): matches must
488/// include a wildcard arm.
489#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
490#[non_exhaustive]
491pub enum StreamEvent {
492    /// An incremental fragment of the reply content.
493    Delta(String),
494    /// The model requests a tool call; argument fragments have already been
495    /// aggregated by the Provider into full JSON text, with fields matching
496    /// [`crate::ToolCall`] so the Agent can use them directly.
497    ToolCall {
498        /// Unique id of this call; the execution result is paired back to it
499        /// via this id.
500        id: String,
501        /// Tool name, corresponding to the name in
502        /// [`Tool::schema`](crate::tool::Tool::schema).
503        name: String,
504        /// Arguments generated by the model (JSON text).
505        arguments: String,
506    },
507    /// An incremental fragment of the model's reasoning (thinking); the
508    /// vendor delivers it in fragments during the stream, and the Provider
509    /// forwards each fragment as it arrives (like [`StreamEvent::Delta`]) —
510    /// consumers concatenate them in order to get the full text.
511    ///
512    /// Corresponds to the reasoning field of [`Message`]; the Agent must store
513    /// it in this turn's message and carry it back verbatim in the history
514    /// (otherwise thinking models like DeepSeek / Qwen3 reject the request).
515    Reasoning(String),
516    /// The model finished its reply; the stream produces no more events after
517    /// this.
518    Done {
519        /// Why the model ended its reply.
520        reason: FinishReason,
521        /// Token usage for this turn; `None` when the endpoint did not return
522        /// it (`include_usage` off or unsupported by the endpoint).
523        usage: Option<Usage>,
524    },
525}
526
527/// Why a Provider call failed.
528///
529/// The enum categories cover the cases that need distinguishing, with details
530/// carried by fields; vendor-specific errors are mapped into this type at the
531/// implementation boundary. `#[non_exhaustive]` guarantees that adding new
532/// categories in the future is not a breaking change.
533///
534/// Error classification is the basis for retry decisions (see the `Default`
535/// judgment of [`Retryable`]): Network / Timeout / RateLimited are worth
536/// retrying, while `Api` is judged by status (5xx retried, 4xx not — retrying
537/// would not change the outcome). Local decode/protocol errors are distinct
538/// from vendor API errors and are not retried by default.
539#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
540#[non_exhaustive]
541pub enum ProviderError {
542    /// A business/API error returned by the vendor (auth failure / invalid
543    /// arguments, quota exhaustion, server error), carrying the HTTP status
544    /// and optional vendor error code.
545    // No "provider " prefix: the type name already expresses the domain,
546    // avoiding a double prefix when wrapped by AgentError::Provider
547    // ("provider error: provider api error …").
548    #[error("api error (status {status}): {message}")]
549    Api {
550        /// HTTP status code returned by the vendor.
551        status: u16,
552        /// Optional vendor error code.
553        code: Option<String>,
554        /// Error description text.
555        message: String,
556    },
557    /// Provider response was syntactically decoded but violated either the
558    /// vendor contract or molo's provider contract.
559    #[error("provider protocol error: {message}")]
560    Protocol {
561        /// Error description text.
562        message: String,
563    },
564    /// Provider response body, SSE frame, or encoded payload could not be
565    /// decoded.
566    #[error("response decode error: {message}")]
567    Decode {
568        /// Error description text.
569        message: String,
570    },
571    /// Rate limited (HTTP 429): retrying is meaningful, and the default retry
572    /// policy waits before retrying.
573    ///
574    /// `retry_after`: the wait duration parsed from the vendor's
575    /// `Retry-After` response header (numeric seconds); `None` when absent
576    /// (HTTP date formats are not parsed).
577    #[error("rate limited")]
578    RateLimited {
579        /// The wait duration indicated by the vendor (numeric seconds);
580        /// `None` when missing / not numeric.
581        retry_after: Option<Duration>,
582    },
583    /// A network-layer failure (connection failure and other transport
584    /// errors).
585    ///
586    /// Implementations map concrete transport errors (e.g. reqwest's) into
587    /// carried text, so the error type does not depend on a concrete
588    /// implementation library's types.
589    #[error("network error: {0}")]
590    Network(String),
591    /// Request timeout, carrying the stage at which it occurred (see
592    /// [`TimeoutStage`]): distinguishes "cannot connect", "total duration
593    /// elapsed" and "event interval stalled" for easier diagnosis.
594    #[error("request timed out during {0:?}")]
595    Timeout(TimeoutStage),
596    /// Provider request was cancelled through provider context.
597    #[error("provider request cancelled")]
598    Cancelled,
599    /// Provider response exceeded a configured local size limit.
600    #[error("response exceeded configured limit ({limit_bytes} bytes)")]
601    ResponseTooLarge {
602        /// Configured limit that was exceeded.
603        limit_bytes: usize,
604    },
605    /// Capability was requested from a provider that does not support it.
606    #[error("unsupported provider capability: {capability}")]
607    Unsupported {
608        /// Unsupported capability name.
609        capability: &'static str,
610    },
611}
612
613/// The stage at which a timeout occurred: one-to-one with
614/// `OpenAiProvider`'s four timeouts (connect / non-streaming total /
615/// streaming event interval / streaming total), plus the error-response-body
616/// read timeout and a generic transport timeout.
617///
618/// The enum is `#[non_exhaustive]` (reserved for extension): matches must
619/// include a wildcard arm.
620#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
621#[non_exhaustive]
622pub enum TimeoutStage {
623    /// Connect timeout (Client-level connect timeout).
624    Connect,
625    /// Request timeout: for non-streaming, covers until the response body is
626    /// fully read; for streaming, covers the connect and response-header wait.
627    Request,
628    /// Streaming event interval exceeded: no data between two events (idle
629    /// timeout).
630    Idle,
631    /// Streaming total duration exceeded: the wall-clock deadline has elapsed
632    /// and an active but never-ending stream is terminated (stream timeout).
633    StreamTotal,
634    /// Total timeout for reading an error response body.
635    ResponseBody,
636    /// Generic transport timeout (reqwest cannot distinguish the stage, e.g.
637    /// during connect or read).
638    Transport,
639}
640
641#[cfg(test)]
642mod tests {
643    use super::Usage;
644
645    #[test]
646    fn usage_add_assign_sums_reported_cache_counts() {
647        let mut total = Usage::new(10, 2).with_cached_tokens(4);
648        total += Usage::new(20, 5).with_cached_tokens(6);
649        assert_eq!(total, Usage::new(30, 7).with_cached_tokens(10));
650    }
651
652    /// A turn that did not report the breakdown contributes nothing (not
653    /// zero); a turn that reported one still contributes, so the sum stays
654    /// usable as the lower bound it is (the Agent layer marks it via
655    /// `RunSummary::cache_omitted`).
656    #[test]
657    fn usage_add_assign_ignores_turns_without_a_cache_breakdown() {
658        let mut total = Usage::new(10, 2).with_cached_tokens(4);
659        total += Usage::new(20, 5);
660        assert_eq!(total, Usage::new(30, 7).with_cached_tokens(4));
661
662        let mut none_reported = Usage::new(10, 2);
663        none_reported += Usage::new(20, 5);
664        assert_eq!(none_reported.cached_tokens, None);
665
666        // Reported zero is not the same as not reported.
667        let mut reported_zero = Usage::default();
668        reported_zero += Usage::new(10, 2).with_cached_tokens(0);
669        assert_eq!(reported_zero.cached_tokens, Some(0));
670    }
671}