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 request headers go
321/// into [`headers`](ModelOptions::headers) — both are passed through to the
322/// vendor verbatim, so users can reach new parameters and endpoint
323/// requirements without waiting for a framework update:
324///
325/// ```rust
326/// # extern crate molo_core as molo;
327/// use molo::ModelOptions;
328///
329/// let mut options = ModelOptions::default();
330/// options.extra.insert("top_p".into(), serde_json::json!(0.9));
331/// options.headers.insert("x-trace-id".into(), "abc-123".into());
332/// ```
333///
334/// # Explicit settings win
335///
336/// When an `extra` key names a parameter the framework also manages
337/// (`temperature` / `max_tokens` / `response_format`), the `extra` value is
338/// the one sent: an explicitly written wire field is a deliberate override,
339/// and the typed field is dropped for that request instead of being merged
340/// (merging would emit the JSON key twice).
341///
342/// The request envelope is not a parameter: `extra` entries named `messages` /
343/// `tools` (the conversation the framework is driving), `stream` /
344/// `stream_options` (the delivery shape) and `model` (bound at provider
345/// construction and reported in run summaries — a per-request override would
346/// make those summaries name a model that did not answer) are ignored, since
347/// they would replace the request rather than tune it.
348#[derive(Clone, Default, PartialEq, Serialize, Deserialize)]
349pub struct ModelOptions {
350    /// Sampling temperature; `None` means vendor default.
351    pub temperature: Option<f32>,
352    /// Maximum number of tokens for the reply; `None` means vendor default.
353    pub max_tokens: Option<u32>,
354    /// Vendor extension parameters: keys are wire request field names,
355    /// serialized into the request body verbatim.
356    pub extra: BTreeMap<String, serde_json::Value>,
357    /// Extra HTTP request headers, sent by HTTP providers alongside the
358    /// headers the framework sets itself.
359    ///
360    /// For endpoints that require more than the request body: gateways and
361    /// proxies commonly route or trace by header (a tenant / project id, a
362    /// gateway token, an `x-trace-id`), and some vendors authenticate with a
363    /// header name of their own (`api-key`, `x-api-key`).
364    ///
365    /// ```rust
366    /// # extern crate molo_core as molo;
367    /// use molo::ModelOptions;
368    ///
369    /// let mut options = ModelOptions::default();
370    /// options.headers.insert("api-key".into(), "secret".into());
371    /// options.headers.insert("x-trace-id".into(), "abc-123".into());
372    /// ```
373    ///
374    /// # Semantics
375    ///
376    /// - **Names / values are taken as given** (HTTP matches names
377    ///   case-insensitively); an entry that is not a valid HTTP header is
378    ///   reported by the provider as [`ProviderError::Protocol`] rather than
379    ///   dropped silently, since a header that quietly never arrives is worse
380    ///   than a failed request.
381    /// - **An explicit `Authorization` wins over the provider's api key**: the
382    ///   endpoint sees the caller's value and the provider does not add a
383    ///   bearer token of its own (two `Authorization` headers would be
384    ///   ambiguous — see the adapter's request construction). `Content-Type`
385    ///   is likewise the caller's when set.
386    /// - **Transport-framing headers stay framework-owned**:
387    ///   `Content-Length` / `Transfer-Encoding` / `Connection` / `Host` /
388    ///   `Upgrade` describe the request the framework actually sends, so
389    ///   entries with those names are ignored (a stale `Content-Length`
390    ///   corrupts the framing of the body).
391    /// - **Values are masked in `Debug` output** (names stay visible), so a
392    ///   credential carried in a header does not leak through `{:?}` logging;
393    ///   the map itself still prints plainly for deliberate inspection.
394    ///
395    /// Providers that do not speak HTTP ignore this field.
396    #[serde(default)]
397    pub headers: BTreeMap<String, String>,
398    /// Structured output: the final answer must be JSON conforming to this
399    /// **JSON Schema document** (the serialized `RootSchema` produced by
400    /// schemars, or a hand-written schema).
401    ///
402    /// Two layers of semantics:
403    /// - Provider side: compatible endpoints receive it via `response_format`
404    ///   to best-effort constrain the model (the OpenAI-compatible
405    ///   `json_schema` shape; unsupported endpoints ignore it or error);
406    /// - Agent side: the final answer is **validated framework-side**, and on
407    ///   mismatch the validation error is fed back to the model for a retry
408    ///   (counted against the turn budget by the agent runtime's structured
409    ///   output support).
410    ///
411    /// `None` = free-form text reply.
412    pub structured: Option<serde_json::Value>,
413}
414
415impl fmt::Debug for ModelOptions {
416    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
417        // Hand-written Debug that masks header values: headers routinely carry
418        // credentials (see `headers`), and Debug output ends up in logs and
419        // panic messages — the same reason `OpenAiProvider` masks its api_key.
420        // Names stay visible: which headers a request carries is the part that
421        // matters when diagnosing it.
422        let headers: BTreeMap<&str, &str> = self
423            .headers
424            .keys()
425            .map(|name| (name.as_str(), "***"))
426            .collect();
427        f.debug_struct("ModelOptions")
428            .field("temperature", &self.temperature)
429            .field("max_tokens", &self.max_tokens)
430            .field("extra", &self.extra)
431            .field("headers", &headers)
432            .field("structured", &self.structured)
433            .finish()
434    }
435}
436
437/// Token usage for one conversation.
438///
439/// Field names match the OpenAI wire format; `total_tokens` follows the
440/// vendor's convention (not necessarily the sum of the other two). `Default`
441/// = all zeros, with nothing reported.
442///
443/// Presence is carried by the enclosing type: [`ChatResponse::usage`] /
444/// [`StreamEvent::Done::usage`] are `Option<Usage>` — `None` means the
445/// endpoint did not report usage for this turn, `Some` means the reported
446/// values. The distinction matters for observability: "not reported" is not
447/// the same as "reportedly zero" (the Agent layer also tracks it in
448/// [`RunSummary::usage_omitted`](crate::run::RunSummary)).
449///
450/// [`cached_tokens`](Self::cached_tokens) carries the same distinction inside
451/// a reported usage: the prompt-cache breakdown is optional per turn even
452/// when the totals are present.
453#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
454pub struct Usage {
455    /// Input tokens for this turn.
456    pub prompt_tokens: u32,
457    /// Output tokens for this turn.
458    pub completion_tokens: u32,
459    /// Total for this turn (vendor convention).
460    pub total_tokens: u32,
461    /// Prompt tokens served from the provider's prompt cache.
462    ///
463    /// A subset of [`prompt_tokens`](Self::prompt_tokens), not an addition to
464    /// it (vendors bill cached input at a lower rate).
465    ///
466    /// `None` = the endpoint did not report a cache breakdown for this turn;
467    /// `Some(0)` = it reported that nothing hit the cache. Endpoints differ
468    /// here — prompt-caching ones report the breakdown (OpenAI / DeepSeek as
469    /// `prompt_tokens_details.cached_tokens`), others omit it — so a reported
470    /// zero must not be inferred from an omission. When turns are summed,
471    /// only the reported parts contribute and
472    /// [`RunSummary::cache_omitted`](crate::run::RunSummary) marks whether
473    /// the sum is a lower bound.
474    pub cached_tokens: Option<u32>,
475}
476
477impl Usage {
478    /// Constructs from input / output counts; the total is summed
479    /// automatically and no cache breakdown is reported
480    /// ([`with_cached_tokens`](Self::with_cached_tokens) sets one).
481    pub fn new(prompt_tokens: u32, completion_tokens: u32) -> Self {
482        Self {
483            prompt_tokens,
484            completion_tokens,
485            total_tokens: prompt_tokens + completion_tokens,
486            cached_tokens: None,
487        }
488    }
489
490    /// Sets the prompt tokens served from the provider's cache, marking the
491    /// breakdown as reported (`Some` in [`cached_tokens`](Self::cached_tokens)).
492    pub fn with_cached_tokens(mut self, cached_tokens: u32) -> Self {
493        self.cached_tokens = Some(cached_tokens);
494        self
495    }
496}
497
498/// Usage accumulates per turn (the Agent layer sums tokens across turns).
499///
500/// `cached_tokens` sums the parts that were reported: a turn that omitted the
501/// breakdown contributes nothing rather than zero.
502impl std::ops::AddAssign for Usage {
503    fn add_assign(&mut self, rhs: Self) {
504        self.prompt_tokens += rhs.prompt_tokens;
505        self.completion_tokens += rhs.completion_tokens;
506        self.total_tokens += rhs.total_tokens;
507        self.cached_tokens = match (self.cached_tokens, rhs.cached_tokens) {
508            (Some(lhs), Some(rhs)) => Some(lhs + rhs),
509            (lhs, rhs) => lhs.or(rhs),
510        };
511    }
512}
513
514/// The reply to one conversation.
515#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
516pub struct ChatResponse {
517    /// This turn's reply: a single [`Message::Assistant`] message
518    /// (text + reasoning + any tool requests from the same turn stay together,
519    /// one-to-one with the vendor wire structure);
520    /// when the model produces nothing it is still an empty Assistant.
521    ///
522    /// The Agent loop executes tool requests when it sees them and appends the
523    /// results as [`Message::ToolResult`] before continuing the conversation.
524    pub message: Message,
525    /// Why the model ended its reply; vendor-specific reasons are surfaced via
526    /// [`FinishReason::Other`].
527    pub finish_reason: FinishReason,
528    /// Token usage for this turn; `None` when the endpoint did not return it
529    /// (compatible endpoints may omit usage; see [`Usage`] for the presence
530    /// semantics shared with [`StreamEvent::Done::usage`]).
531    pub usage: Option<Usage>,
532}
533
534/// Why the model ended its reply.
535///
536/// Common reasons are typed (Stop / Length); vendor-specific or
537/// framework-unknown reasons are surfaced via [`Other`](FinishReason::Other)
538/// carrying the vendor's raw string — users can recognize new reasons without
539/// waiting for a framework update. `#[non_exhaustive]` guarantees that adding
540/// new common categories in the future is not a breaking change.
541#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
542#[non_exhaustive]
543pub enum FinishReason {
544    /// The model ended naturally (including ending the turn by requesting
545    /// tool calls).
546    Stop,
547    /// Truncated by hitting the max_tokens limit.
548    Length,
549    /// A vendor-specific or framework-unknown reason, carrying the vendor's
550    /// raw string.
551    Other(String),
552}
553
554/// An event in a streamed conversation reply.
555///
556/// One streamed reply = several [`StreamEvent::Delta`] /
557/// [`StreamEvent::ToolCall`] increments + one closing [`StreamEvent::Done`];
558/// the caller concatenates the Deltas in order to get the full reply.
559///
560/// Stream termination semantics: on normal termination `Done` is always the
561/// last success event on the stream; errors terminate the stream with an
562/// `Err` item, and no events are produced after the error item.
563///
564/// The enum is `#[non_exhaustive]` (reserved for extension): matches must
565/// include a wildcard arm.
566#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
567#[non_exhaustive]
568pub enum StreamEvent {
569    /// An incremental fragment of the reply content.
570    Delta(String),
571    /// The model requests a tool call; argument fragments have already been
572    /// aggregated by the Provider into full JSON text, with fields matching
573    /// [`crate::ToolCall`] so the Agent can use them directly.
574    ToolCall {
575        /// Unique id of this call; the execution result is paired back to it
576        /// via this id.
577        id: String,
578        /// Tool name, corresponding to the name in
579        /// [`Tool::schema`](crate::tool::Tool::schema).
580        name: String,
581        /// Arguments generated by the model (JSON text).
582        arguments: String,
583    },
584    /// An incremental fragment of the model's reasoning (thinking); the
585    /// vendor delivers it in fragments during the stream, and the Provider
586    /// forwards each fragment as it arrives (like [`StreamEvent::Delta`]) —
587    /// consumers concatenate them in order to get the full text.
588    ///
589    /// Corresponds to the reasoning field of [`Message`]; the Agent must store
590    /// it in this turn's message and carry it back verbatim in the history
591    /// (otherwise thinking models like DeepSeek / Qwen3 reject the request).
592    Reasoning(String),
593    /// The model finished its reply; the stream produces no more events after
594    /// this.
595    Done {
596        /// Why the model ended its reply.
597        reason: FinishReason,
598        /// Token usage for this turn; `None` when the endpoint did not return
599        /// it (`include_usage` off or unsupported by the endpoint).
600        usage: Option<Usage>,
601    },
602}
603
604/// Why a Provider call failed.
605///
606/// The enum categories cover the cases that need distinguishing, with details
607/// carried by fields; vendor-specific errors are mapped into this type at the
608/// implementation boundary. `#[non_exhaustive]` guarantees that adding new
609/// categories in the future is not a breaking change.
610///
611/// Error classification is the basis for retry decisions (see the `Default`
612/// judgment of [`Retryable`]): Network / Timeout / RateLimited are worth
613/// retrying, while `Api` is judged by status (5xx retried, 4xx not — retrying
614/// would not change the outcome). Local decode/protocol errors are distinct
615/// from vendor API errors and are not retried by default.
616#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
617#[non_exhaustive]
618pub enum ProviderError {
619    /// A business/API error returned by the vendor (auth failure / invalid
620    /// arguments, quota exhaustion, server error), carrying the HTTP status
621    /// and optional vendor error code.
622    // No "provider " prefix: the type name already expresses the domain,
623    // avoiding a double prefix when wrapped by AgentError::Provider
624    // ("provider error: provider api error …").
625    #[error("api error (status {status}): {message}")]
626    Api {
627        /// HTTP status code returned by the vendor.
628        status: u16,
629        /// Optional vendor error code.
630        code: Option<String>,
631        /// Error description text.
632        message: String,
633    },
634    /// Provider response was syntactically decoded but violated either the
635    /// vendor contract or molo's provider contract.
636    ///
637    /// Also raised **before sending** when a caller-supplied request part
638    /// cannot be put on the wire at all — an invalid
639    /// [`ModelOptions::headers`] entry (a name or value that is not a valid
640    /// HTTP header) fails the request here instead of being dropped silently.
641    #[error("provider protocol error: {message}")]
642    Protocol {
643        /// Error description text.
644        message: String,
645    },
646    /// Provider response body, SSE frame, or encoded payload could not be
647    /// decoded.
648    #[error("response decode error: {message}")]
649    Decode {
650        /// Error description text.
651        message: String,
652    },
653    /// Rate limited (HTTP 429): retrying is meaningful, and the default retry
654    /// policy waits before retrying.
655    ///
656    /// `retry_after`: the wait duration parsed from the vendor's
657    /// `Retry-After` response header (numeric seconds); `None` when absent
658    /// (HTTP date formats are not parsed).
659    #[error("rate limited")]
660    RateLimited {
661        /// The wait duration indicated by the vendor (numeric seconds);
662        /// `None` when missing / not numeric.
663        retry_after: Option<Duration>,
664    },
665    /// A network-layer failure (connection failure and other transport
666    /// errors).
667    ///
668    /// Implementations map concrete transport errors (e.g. reqwest's) into
669    /// carried text, so the error type does not depend on a concrete
670    /// implementation library's types.
671    #[error("network error: {0}")]
672    Network(String),
673    /// Request timeout, carrying the stage at which it occurred (see
674    /// [`TimeoutStage`]): distinguishes "cannot connect", "total duration
675    /// elapsed" and "event interval stalled" for easier diagnosis.
676    #[error("request timed out during {0:?}")]
677    Timeout(TimeoutStage),
678    /// Provider request was cancelled through provider context.
679    #[error("provider request cancelled")]
680    Cancelled,
681    /// Provider response exceeded a configured local size limit.
682    #[error("response exceeded configured limit ({limit_bytes} bytes)")]
683    ResponseTooLarge {
684        /// Configured limit that was exceeded.
685        limit_bytes: usize,
686    },
687    /// Capability was requested from a provider that does not support it.
688    #[error("unsupported provider capability: {capability}")]
689    Unsupported {
690        /// Unsupported capability name.
691        capability: &'static str,
692    },
693}
694
695/// The stage at which a timeout occurred: one-to-one with
696/// `OpenAiProvider`'s four timeouts (connect / non-streaming total /
697/// streaming event interval / streaming total), plus the error-response-body
698/// read timeout and a generic transport timeout.
699///
700/// The enum is `#[non_exhaustive]` (reserved for extension): matches must
701/// include a wildcard arm.
702#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
703#[non_exhaustive]
704pub enum TimeoutStage {
705    /// Connect timeout (Client-level connect timeout).
706    Connect,
707    /// Request timeout: for non-streaming, covers until the response body is
708    /// fully read; for streaming, covers the connect and response-header wait.
709    Request,
710    /// Streaming event interval exceeded: no data between two events (idle
711    /// timeout).
712    Idle,
713    /// Streaming total duration exceeded: the wall-clock deadline has elapsed
714    /// and an active but never-ending stream is terminated (stream timeout).
715    StreamTotal,
716    /// Total timeout for reading an error response body.
717    ResponseBody,
718    /// Generic transport timeout (reqwest cannot distinguish the stage, e.g.
719    /// during connect or read).
720    Transport,
721}
722
723#[cfg(test)]
724mod tests {
725    use super::{ModelOptions, Usage};
726
727    #[test]
728    fn usage_add_assign_sums_reported_cache_counts() {
729        let mut total = Usage::new(10, 2).with_cached_tokens(4);
730        total += Usage::new(20, 5).with_cached_tokens(6);
731        assert_eq!(total, Usage::new(30, 7).with_cached_tokens(10));
732    }
733
734    /// A turn that did not report the breakdown contributes nothing (not
735    /// zero); a turn that reported one still contributes, so the sum stays
736    /// usable as the lower bound it is (the Agent layer marks it via
737    /// `RunSummary::cache_omitted`).
738    #[test]
739    fn usage_add_assign_ignores_turns_without_a_cache_breakdown() {
740        let mut total = Usage::new(10, 2).with_cached_tokens(4);
741        total += Usage::new(20, 5);
742        assert_eq!(total, Usage::new(30, 7).with_cached_tokens(4));
743
744        let mut none_reported = Usage::new(10, 2);
745        none_reported += Usage::new(20, 5);
746        assert_eq!(none_reported.cached_tokens, None);
747
748        // Reported zero is not the same as not reported.
749        let mut reported_zero = Usage::default();
750        reported_zero += Usage::new(10, 2).with_cached_tokens(0);
751        assert_eq!(reported_zero.cached_tokens, Some(0));
752    }
753
754    /// Headers routinely carry credentials: `Debug` shows which headers a
755    /// request carries (the diagnostic part) while masking the values, so a
756    /// `{:?}` in a log line cannot leak the token.
757    #[test]
758    fn model_options_debug_masks_header_values() {
759        let mut options = ModelOptions::default();
760        options.headers.insert("x-api-key".into(), "secret".into());
761
762        let debug = format!("{options:?}");
763        assert!(debug.contains("x-api-key"), "names stay visible: {debug}");
764        assert!(!debug.contains("secret"), "values are masked: {debug}");
765    }
766
767    /// Options written before `headers` existed still deserialize: the field
768    /// defaults to empty instead of being required.
769    #[test]
770    fn model_options_deserialize_without_headers() {
771        let options: ModelOptions = serde_json::from_str(
772            r#"{"temperature":0.5,"max_tokens":null,"extra":{},"structured":null}"#,
773        )
774        .unwrap();
775
776        assert_eq!(options.temperature, Some(0.5));
777        assert!(options.headers.is_empty());
778    }
779}