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}