Skip to main content

wcode_harness/
streamfn.rs

1//! LLM stream-function seam + default `rig` (OpenAI-compatible) adapter.
2//!
3//! This module is the only place rig types appear in the kernel. A
4//! [`StreamFn`] turns harness conversation context into a [`LlmStream`] of
5//! [`LlmStreamEvent`]s; [`rig_stream_fn`] is the default implementation built
6//! on rig's Chat Completions API client.
7
8use std::pin::Pin;
9use std::sync::{Arc, Mutex};
10use std::time::{Duration, SystemTime, UNIX_EPOCH};
11
12use futures::{Stream, StreamExt};
13use rig::client::{CompletionClient, ModelListingClient};
14use rig::completion::{
15    CompletionError, CompletionModel, CompletionRequest, FinishReason, ToolDefinition,
16};
17use rig::http_client::{HeaderMap, HeaderValue};
18use rig::message::{
19    AssistantContent, Message, Reasoning, ReasoningContent, ToolCall, ToolCallId, ToolFunction,
20    UserContent,
21};
22use rig::providers::openai;
23use rig::streaming::{StreamedAssistantContent, StreamingCompletionResponse};
24
25use crate::event::LlmStreamEvent;
26use crate::message::{AgentMessage, ContentBlock, StopReason};
27
28#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
29pub enum LlmEndpoint {
30    #[default]
31    Chat,
32    Responses,
33}
34
35#[derive(Clone, Debug, Default)]
36pub struct LlmOpts {
37    pub model: String,
38    pub base_url: Option<String>,
39    pub api_key: Option<String>,
40    /// Sampling temperature. No CLI flag surfaces it yet; a library embedder
41    /// that builds `LlmOpts` directly can set it.
42    pub temperature: Option<f64>,
43    pub endpoint: LlmEndpoint,
44    /// Stable per-conversation id, sent as `x-opencode-session` for providers
45    /// (e.g. OpenCode Go) that route/cache on it. None = header omitted.
46    pub session_id: Option<String>,
47    /// Free-style reasoning effort. None = send nothing (backends that reject
48    /// unknown fields keep working). Some(level) fans out per endpoint in
49    /// `build_request()`: Chat forwards any string verbatim under
50    /// `reasoning_effort`, while the Responses endpoint validates client-side
51    /// against `none|minimal|low|medium|high|xhigh|max` (rig's typed enum) — a
52    /// value outside that set fails with a clear error before the request.
53    pub effort: Option<String>,
54    /// Retry policy for the connect phase (transient failures). See
55    /// [`RetryPolicy`].
56    pub retry: RetryPolicy,
57    /// Per-model provider profiles (`[models.<id>]`) plus the launch provider a
58    /// switch falls back to. `Agent::set_model` (and the CLI) call
59    /// `settle_provider` to re-point endpoint/base_url/api_key for the selected
60    /// model; the `Default` map (no base) never touches the provider.
61    pub model_profiles: ModelProfiles,
62}
63
64/// A per-model provider profile: how to reach the model when it differs from
65/// the launch provider. Every field is optional — `None` means "leave the
66/// launch provider's value", so an empty profile is a no-op and a profile may
67/// override only `base_url` without resetting the endpoint.
68#[derive(Clone, Debug, Default, PartialEq, Eq)]
69pub struct LlmProfile {
70    /// Endpoint for this model. `None` = keep the launch endpoint — so a
71    /// profile may override only `base_url` without silently toggling
72    /// chat ⇄ responses.
73    pub endpoint: Option<LlmEndpoint>,
74    /// Base URL for this model. `None` = inherit.
75    pub base_url: Option<String>,
76    /// API key for this model. `None` = inherit. No CLI flag surfaces a key;
77    /// env/file still feed `to_llm_opts`, and a profile sits above both.
78    pub api_key: Option<String>,
79}
80
81impl LlmProfile {
82    /// Overlay this profile onto `llm` in place; absent fields are left
83    /// untouched, so an empty profile is a no-op and only fields the profile
84    /// SETS can win.
85    /// Overlay this profile onto `llm` in place; absent fields are left
86    /// untouched, so an empty profile is a no-op and only fields the profile
87    /// SETS can win. `ModelProfiles::settle` is the only entry point.
88    fn apply(&self, llm: &mut LlmOpts) {
89        if let Some(endpoint) = self.endpoint {
90            llm.endpoint = endpoint;
91        }
92        if let Some(base_url) = &self.base_url {
93            llm.base_url = Some(base_url.clone());
94        }
95        if let Some(api_key) = &self.api_key {
96            llm.api_key = Some(api_key.clone());
97        }
98    }
99}
100
101/// The provider a session starts on — what a model switch falls back to when the
102/// model has no `[models.<id>]` entry. Captured at boot AFTER env/flag overrides,
103/// so flags/env/config globals define the DEFAULT provider.
104#[derive(Clone, Debug, PartialEq, Eq)]
105pub struct LlmProvider {
106    pub endpoint: LlmEndpoint,
107    pub base_url: Option<String>,
108    pub api_key: Option<String>,
109}
110
111impl LlmProvider {
112    /// The provider `llm` is on right now.
113    pub fn of(llm: &LlmOpts) -> Self {
114        Self {
115            endpoint: llm.endpoint,
116            base_url: llm.base_url.clone(),
117            api_key: llm.api_key.clone(),
118        }
119    }
120
121    /// Reset `llm`'s provider fields to this one.
122    fn apply(&self, llm: &mut LlmOpts) {
123        llm.endpoint = self.endpoint;
124        llm.base_url = self.base_url.clone();
125        llm.api_key = self.api_key.clone();
126    }
127}
128
129/// `id → provider profile` resolved once at boot from `[models.<id>]`, plus the
130/// launch provider (`base`) an unmapped model falls back to. The `Default` (no
131/// base) does NOT manage providers: every operation is a no-op, so a bare
132/// kernel embedder's `LlmOpts::default()` is never touched. A `BTreeMap` (not
133/// `HashMap`) keeps debug/`--dump-*` output deterministic.
134///
135/// PRECEDENCE: `profile > flag(=base) > config/env globals`. `settle` computes
136/// `base ⊕ profile`, and `base` is the post-flag provider, so a field a profile
137/// sets necessarily wins over the flag. Flags/env/config therefore set the
138/// DEFAULT provider; a `[models.<id>]` entry overrides it for that model.
139#[derive(Clone, Debug, Default, PartialEq, Eq)]
140pub struct ModelProfiles {
141    /// The provider an unmapped model falls back to. `None` = not managing.
142    base: Option<LlmProvider>,
143    profiles: std::collections::BTreeMap<String, LlmProfile>,
144}
145
146impl ModelProfiles {
147    /// Start managing: `base` is the provider an unmapped model falls back to.
148    pub fn new(base: LlmProvider) -> Self {
149        Self {
150            base: Some(base),
151            profiles: std::collections::BTreeMap::new(),
152        }
153    }
154
155    /// The profile for `model`, if one is configured (internal to `settle`).
156    fn resolve(&self, model: &str) -> Option<&LlmProfile> {
157        self.profiles.get(model)
158    }
159
160    /// Settle `llm`'s provider for `llm.model`: reset endpoint/base_url/api_key
161    /// to `base`, then overlay that model's profile. No-op when `base` is `None`.
162    pub fn settle(&self, llm: &mut LlmOpts) {
163        // A non-empty map that is not managing providers is the footgun: those
164        // profiles would be silently dead. (The `Default` is legitimately empty.)
165        debug_assert!(
166            self.base.is_some() || self.profiles.is_empty(),
167            "ModelProfiles has profiles but no base (was `insert` called on the Default?)"
168        );
169        let Some(base) = &self.base else {
170            return;
171        };
172        base.apply(llm);
173        if let Some(profile) = self.resolve(&llm.model) {
174            profile.apply(llm);
175        }
176    }
177
178    /// The launch provider an unmapped model falls back to; `None` when the map
179    /// is not managing providers (the `Default`). Re-forwarded verbatim across a
180    /// `/reload` / resume so the DEFAULT provider survives the process boundary.
181    pub fn base(&self) -> Option<&LlmProvider> {
182        self.base.as_ref()
183    }
184
185    /// Register a profile for `model` (boot-time config build).
186    pub fn insert(&mut self, model: impl Into<String>, profile: LlmProfile) {
187        self.profiles.insert(model.into(), profile);
188        // Post-condition: inserting into a `Default` (base `None`) is the footgun
189        // this catches — such a profile could never be applied by `settle`.
190        debug_assert!(
191            self.base.is_some() || self.profiles.is_empty(),
192            "insert on a ModelProfiles with no base: the profile is dead (use `new`)"
193        );
194    }
195}
196
197impl LlmOpts {
198    /// Re-point `endpoint`/`base_url`/`api_key` for `self.model`: the launch
199    /// provider (`model_profiles.base`), then this model's `[models.<id>]`
200    /// overlay. No-op when no profiles are configured. Cloning the map avoids
201    /// borrowing `self.model_profiles` while mutably borrowing `self`.
202    pub fn settle_provider(&mut self) {
203        let profiles = self.model_profiles.clone();
204        profiles.settle(self);
205    }
206}
207
208pub type LlmStream = Pin<Box<dyn Stream<Item = LlmStreamEvent> + Send>>;
209
210pub type StreamFn =
211    Arc<dyn Fn(&[AgentMessage], &str, &[ToolDefinition], &LlmOpts) -> LlmStream + Send + Sync>;
212
213/// Everything that shapes the underlying rig client: the endpoint, the key, and
214/// the routing/session header. `model`, `temperature`, and `effort` ride on the
215/// request, not the client, so a mid-conversation model swap still reuses the
216/// pooled HTTP connections. (`OPENAI_API_KEY`, the key fallback when `api_key`
217/// is None, is read once at build time and treated as constant.)
218#[derive(Clone, PartialEq, Eq)]
219struct ClientKey {
220    base_url: Option<String>,
221    api_key: Option<String>,
222    session_id: Option<String>,
223}
224
225impl ClientKey {
226    fn of(opts: &LlmOpts) -> Self {
227        Self {
228            base_url: opts.base_url.clone(),
229            api_key: opts.api_key.clone(),
230            session_id: opts.session_id.clone(),
231        }
232    }
233}
234
235/// One-entry client cache. rig's `Client` wraps a pooled `reqwest::Client`, so
236/// rebuilding it every turn would drop the pool and re-handshake TLS; this keeps
237/// one client alive across turns and rebuilds only when [`ClientKey`] changes.
238#[derive(Default)]
239struct ClientCache {
240    slot: Mutex<Option<(ClientKey, openai::Client)>>,
241    #[cfg(test)]
242    builds: std::sync::atomic::AtomicUsize,
243}
244
245impl ClientCache {
246    /// The client for `opts`, reusing the cached one when the key is unchanged.
247    fn get(&self, opts: &LlmOpts) -> Result<openai::Client, String> {
248        let key = ClientKey::of(opts);
249        let mut slot = self.slot.lock().unwrap_or_else(|e| e.into_inner());
250        if let Some((cached_key, client)) = slot.as_ref()
251            && *cached_key == key
252        {
253            return Ok(client.clone());
254        }
255        let client = openai_client(opts)?;
256        #[cfg(test)]
257        self.builds
258            .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
259        *slot = Some((key, client.clone()));
260        Ok(client)
261    }
262}
263
264/// Default adapter: OpenAI-compatible Chat Completions endpoint via rig.
265///
266/// The provider client (and its pooled connections) is built once and reused
267/// across turns; only an endpoint/key/session change rebuilds it.
268pub fn rig_stream_fn() -> StreamFn {
269    let cache = ClientCache::default();
270    Arc::new(move |messages, system, tools, opts| adapt(&cache, opts, messages, system, tools))
271}
272
273/// Build the rig client shared by streaming and model listing, so both hit
274/// the same base_url with the same key.
275fn openai_client(opts: &LlmOpts) -> Result<rig::providers::openai::Client, String> {
276    let key = opts
277        .api_key
278        .clone()
279        .or_else(|| std::env::var("OPENAI_API_KEY").ok());
280    let key = match key {
281        Some(k) => k,
282        None => {
283            // Keyless local OpenAI-compatible endpoints (Ollama, llama.cpp,
284            // vLLM, ...) accept any bearer token, so a placeholder keeps them
285            // usable while a missing key against a real endpoint fails loudly.
286            let local = opts.base_url.as_deref().is_some_and(|url| {
287                let u = url.to_ascii_lowercase();
288                ["http://localhost", "http://127.0.0.1", "http://[::1]"]
289                    .iter()
290                    .any(|prefix| u.starts_with(prefix))
291            });
292            if local {
293                "wcode-local".to_string()
294            } else {
295                return Err(
296                    "no API key: pass LlmOpts.api_key or set OPENAI_API_KEY (a localhost base_url works keyless)"
297                        .to_string(),
298                );
299            }
300        }
301    };
302    let mut builder = openai::Client::builder()
303        .api_key::<rig::client::BearerAuth>(key)
304        .http_headers(session_headers(opts));
305    if let Some(base_url) = &opts.base_url {
306        builder = builder.base_url(base_url);
307    }
308    builder.build().map_err(|e| format!("openai client: {e}"))
309}
310
311/// Default headers for every provider request: own `User-Agent` (OpenCode Go
312/// rejects generic SDK names for routing) plus `x-opencode-session` when set.
313fn session_headers(opts: &LlmOpts) -> HeaderMap {
314    let mut headers = HeaderMap::new();
315    headers.insert(
316        "user-agent",
317        HeaderValue::from_static(concat!("wcode/", env!("CARGO_PKG_VERSION"))),
318    );
319    if let Some(id) = opts.session_id.as_deref()
320        && let Ok(value) = HeaderValue::from_str(id)
321    {
322        headers.insert("x-opencode-session", value);
323    }
324    headers
325}
326
327/// Probe `base_url` for a live OpenAI-compatible endpoint: `GET
328/// {base_url}/models` under a hard `timeout`. `true` iff the endpoint answered
329/// with a model list (an empty list still counts — the endpoint is up). Any
330/// error (refused, DNS, non-2xx, timeout) is `false`; the caller treats `false`
331/// as "not this endpoint" and moves on.
332///
333/// The keyless-placeholder rule lives in `openai_client`: only a
334/// `http://localhost|127.0.0.1|[::1]` base builds without a key, so a
335/// non-loopback `$OLLAMA_HOST` with no key probes `false`.
336pub async fn probe_endpoint(base_url: &str, timeout: std::time::Duration) -> bool {
337    let opts = LlmOpts {
338        base_url: Some(base_url.to_string()),
339        ..LlmOpts::default()
340    };
341    matches!(
342        tokio::time::timeout(timeout, list_models(&opts)).await,
343        Ok(Ok(_))
344    )
345}
346/// List models via `GET {base_url}/models` using the same client config as
347/// streaming. Ids are sorted for stable display. The endpoint only returns
348/// id/name metadata — no capability flags — so effort support stays
349/// user-managed.
350pub async fn list_models(opts: &LlmOpts) -> Result<Vec<String>, String> {
351    let client = openai_client(opts)?;
352    let list = client.list_models().await.map_err(|e| e.to_string())?;
353    let mut ids: Vec<String> = list.iter().map(|m| m.id.clone()).collect();
354    ids.sort();
355    Ok(ids)
356}
357
358fn adapt(
359    cache: &ClientCache,
360    opts: &LlmOpts,
361    messages: &[AgentMessage],
362    system: &str,
363    tools: &[ToolDefinition],
364) -> LlmStream {
365    if let Some(message) = invalid_effort(opts) {
366        return error_stream(message);
367    }
368    let client = match cache.get(opts) {
369        Ok(client) => client,
370        Err(message) => return error_stream(message),
371    };
372    let request = build_request(messages, system, tools, opts);
373
374    // Both wires yield the same normalized StreamingCompletionResponse, so
375    // only model construction branches; the forwarding loop below is shared.
376    // Retry the connect (and the first item — rig defers the request into the
377    // stream). Once output is forwarded a mid-stream error is surfaced as-is. A
378    // fresh model per attempt is cheap — it just wraps the pooled client;
379    // `request` is cloned per attempt and `policy` is `Copy`.
380    let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<LlmStreamEvent>();
381    let Ok(runtime) = tokio::runtime::Handle::try_current() else {
382        return error_stream("StreamFn requires a tokio runtime".to_string());
383    };
384    let policy = opts.retry;
385    let endpoint = opts.endpoint;
386    let model = opts.model.clone();
387    runtime.spawn(async move {
388        let open = || open_stream(client.clone(), endpoint, &model, request.clone());
389        // Retry budget for the whole pre-content window, shared across the
390        // connect, the first-item peek (rig defers the request into the stream,
391        // so a transport failure surfaces there), and any transient error that
392        // slips past the peek before the first content event is forwarded.
393        let mut state = RetryState::new(policy.max);
394        let (stream, pending) = match connect_with_retry(&policy, &mut state, &open, &tx).await {
395            PreContent::Ready { stream, first } => (stream, first),
396            PreContent::Cancelled => return,
397            PreContent::Failed(error) => {
398                let _ = tx.send(LlmStreamEvent::Error {
399                    message: error.to_string(),
400                    fatal: fatal_class(&error),
401                });
402                return;
403            }
404        };
405        forward_stream(stream, pending, policy, &open, tx, state).await;
406    });
407    Box::pin(futures::stream::unfold(rx, |mut rx| async move {
408        rx.recv().await.map(|ev| (ev, rx))
409    }))
410}
411
412/// The forwarding phase of a stream turn: pull items (the buffered `pending`
413/// first, then from `stream` under the idle deadline), map them to events, and
414/// send them on `tx`. A pre-content stall re-opens via `open` on the shared
415/// retry budget (mirroring a transient connect error); a post-content stall (or
416/// an exhausted budget) surfaces a non-fatal error and ends the loop.
417async fn forward_stream<S, F>(
418    mut stream: S,
419    mut pending: Option<Result<StreamedAssistantContent, CompletionError>>,
420    policy: RetryPolicy,
421    open: F,
422    tx: tokio::sync::mpsc::UnboundedSender<LlmStreamEvent>,
423    mut state: RetryState,
424) where
425    F: Fn() -> Pin<Box<dyn std::future::Future<Output = Result<S, CompletionError>> + Send>>,
426    S: futures::Stream<Item = Result<StreamedAssistantContent, CompletionError>> + Unpin + Send,
427{
428    let mut errored = false;
429    let mut content_forwarded = false;
430    let mut done_seen = false;
431    loop {
432        // The next item: a buffered first item, or one from the stream under the
433        // idle deadline. Cancel-safe: a dropped `next_or_stall` drops an
434        // in-flight `stream.next()` poll, which is safe to re-enter (the
435        // pre-content path re-opens a fresh stream); `biased` with the item arm
436        // first so a ready item is never lost.
437        let item = match pending.take() {
438            Some(item) => item,
439            None => tokio::select! {
440                biased;
441                result = next_or_stall(&mut stream, policy.idle) => match result {
442                    // The stream ended. A `Final` record (mapped to `Done`) is the
443                    // only clean end: rig emits neither a `Final` nor an `Err` when
444                    // the provider SSE body ends early (proxy idle-close / half-close
445                    // / EOF), so the normalized stream just ends with `None` — a
446                    // truncated turn, not a stop.
447                    None => {
448                        if done_seen || errored {
449                            break;
450                        }
451                        let error = truncation_error();
452                        if !content_forwarded && state.budget > 0 {
453                            // Pre-content: retry on the shared budget, exactly like
454                            // the idle-stall arm (one retry path, via the helper).
455                            match pre_content_retry(&policy, &mut state, &open, &tx, &error).await {
456                                PreContent::Ready { stream: s, first } => {
457                                    stream = s;
458                                    pending = first;
459                                    continue;
460                                }
461                                PreContent::Cancelled => return,
462                                PreContent::Failed(e) => {
463                                    if tx
464                                        .send(LlmStreamEvent::Error {
465                                            message: e.to_string(),
466                                            fatal: fatal_class(&e),
467                                        })
468                                        .is_err()
469                                    {
470                                        return;
471                                    }
472                                    break;
473                                }
474                            }
475                        } else {
476                            // NON-fatal: the loop feeds it back next turn and, if it
477                            // persists, ends StopReason::Error.
478                            // Chosen Error.message: `truncation_error().to_string()`.
479                            if tx
480                                .send(LlmStreamEvent::Error {
481                                    message: error.to_string(),
482                                    fatal: false,
483                                })
484                                .is_err()
485                            {
486                                return;
487                            }
488                            break;
489                        }
490                    }
491                    // One item — the stream's own `Result`.
492                    Some(Stall::Item(item)) => item,
493                    // The idle deadline fired. Pre-content → a transient retry
494                    // on the shared budget (re-open, mirroring a connect error);
495                    // post-content (or an exhausted budget) → a NON-fatal error
496                    // that ends the turn (a stalled stream will not resume).
497                    Some(Stall::Idle) => {
498                        let error = stall_error(policy.idle);
499                        if !content_forwarded && state.budget > 0 {
500                            match pre_content_retry(&policy, &mut state, &open, &tx, &error).await {
501                                PreContent::Ready { stream: s, first } => {
502                                    stream = s;
503                                    pending = first;
504                                    continue;
505                                }
506                                PreContent::Cancelled => return,
507                                PreContent::Failed(e) => {
508                                    let fatal = fatal_class(&e);
509                                    if tx
510                                        .send(LlmStreamEvent::Error {
511                                            message: e.to_string(),
512                                            fatal,
513                                        })
514                                        .is_err()
515                                    {
516                                        return;
517                                    }
518                                    break;
519                                }
520                            }
521                        } else {
522                            if tx
523                                .send(LlmStreamEvent::Error {
524                                    message: error.to_string(),
525                                    fatal: false,
526                                })
527                                .is_err()
528                            {
529                                return;
530                            }
531                            break;
532                        }
533                    }
534                },
535                _ = tx.closed() => break,
536            },
537        };
538        let events = match item {
539            Ok(content) => {
540                let is_final = matches!(content, StreamedAssistantContent::Final(_));
541                if errored && is_final {
542                    continue;
543                }
544                let events = map_item(content);
545                // Replay is only safe before the first forwarded event;
546                // once anything reached the consumer a later error cannot
547                // be retried (it would duplicate output).
548                content_forwarded |= !events.is_empty();
549                done_seen |= is_final;
550                events
551            }
552            Err(error) => {
553                if !content_forwarded && state.budget > 0 && retryable(&error) {
554                    // Nothing forwarded yet: back off and replay the request
555                    // from scratch. The re-open shares the same retry budget
556                    // and keeps backoff growing monotonically.
557                    state.attempts += 1;
558                    state.budget -= 1;
559                    if !retry_wait(&tx, &policy, state.attempts, &error).await {
560                        return;
561                    }
562                    match connect_with_retry(&policy, &mut state, &open, &tx).await {
563                        PreContent::Ready { stream: s, first } => {
564                            stream = s;
565                            pending = first;
566                            continue;
567                        }
568                        PreContent::Cancelled => return,
569                        PreContent::Failed(e) => {
570                            errored = true;
571                            vec![LlmStreamEvent::Error {
572                                message: e.to_string(),
573                                fatal: fatal_class(&e),
574                            }]
575                        }
576                    }
577                } else {
578                    // Final: a hard-fatal class, retries exhausted, or a
579                    // mid-stream error after content. The loop decides (via
580                    // `fatal`) whether to feed it back or end the run.
581                    errored = true;
582                    vec![LlmStreamEvent::Error {
583                        message: error.to_string(),
584                        fatal: fatal_class(&error),
585                    }]
586                }
587            }
588        };
589        for ev in events {
590            if tx.send(ev).is_err() {
591                return;
592            }
593        }
594    }
595}
596
597fn error_stream(message: String) -> LlmStream {
598    // Pre-request failures (invalid effort, client build, missing runtime) are
599    // hard-fatal: the request itself can never succeed, so the run must stop.
600    Box::pin(futures::stream::iter(vec![LlmStreamEvent::Error {
601        message,
602        fatal: true,
603    }]))
604}
605
606/// Retry policy for a stream call: the transient-connect retries plus the stall
607/// (idle / time-to-first-token) timeouts.
608#[derive(Clone, Copy, Debug, PartialEq, Eq)]
609pub struct RetryPolicy {
610    /// Retries after the first attempt (0 disables retrying).
611    pub max: u32,
612    /// Base backoff: attempt `n` waits `base * 2^(n-1)`, capped at `cap`.
613    pub base: Duration,
614    /// Upper bound on a single backoff wait.
615    pub cap: Duration,
616    /// Time-to-first-token: covers connect + the first stream item. `ZERO` = off.
617    pub ttft: Duration,
618    /// Inter-item idle: resets on every received item. `ZERO` = off.
619    pub idle: Duration,
620}
621
622impl Default for RetryPolicy {
623    fn default() -> Self {
624        Self {
625            max: 3,
626            base: Duration::from_millis(500),
627            cap: Duration::from_secs(8),
628            ttft: Duration::from_secs(60),
629            idle: Duration::from_secs(120),
630        }
631    }
632}
633
634/// One poll of a stream under a stall deadline (see [`next_or_stall`]).
635enum Stall<T> {
636    /// One item — the stream's OWN item type (for a rig stream, a
637    /// `Result<StreamedAssistantContent, CompletionError>`), not an `Option`.
638    Item(T),
639    /// The deadline fired before an item arrived.
640    Idle,
641}
642
643/// `stream.next()` under a `deadline`; `ZERO` means "no timeout" and collapses
644/// to a plain `stream.next().await`. `biased` toward the item so a ready item
645/// wins a tie with an expired deadline. Outer `None` = the stream ended.
646async fn next_or_stall<S: Stream + Unpin>(
647    stream: &mut S,
648    deadline: Duration,
649) -> Option<Stall<S::Item>> {
650    if deadline.is_zero() {
651        return stream.next().await.map(Stall::Item);
652    }
653    tokio::select! {
654        biased;
655        item = stream.next() => item.map(Stall::Item),
656        _ = tokio::time::sleep(deadline) => Some(Stall::Idle),
657    }
658}
659/// Transient HTTP statuses worth retrying: request timeout/conflict, rate
660/// limiting, and the 5xx family. Client errors (400/401/403/404/422) are not.
661fn is_transient_status(code: u16) -> bool {
662    matches!(code, 408 | 409 | 429 | 500 | 502 | 503 | 504)
663}
664
665/// Whether a failed connect is worth retrying. A preserved provider status is
666/// authoritative (transient statuses retry, client errors don't); with no
667/// status, a transport failure retries and a build/parse error does not.
668fn retryable(error: &CompletionError) -> bool {
669    if let Some(status) = error.provider_response_status() {
670        return is_transient_status(status.as_u16());
671    }
672    match error {
673        // Transport failure with no status (connect refused / reset / timeout).
674        CompletionError::HttpError(_) => true,
675        // rig's OpenAI path folds some transport failures into a
676        // `ProviderError(String)` (reqwest's "error sending request …"); retry
677        // those, but not provider rejections that share the variant.
678        CompletionError::ProviderError(message) => is_transient_message(message),
679        _ => false,
680    }
681}
682
683/// Heuristic for transport failures that arrive as `ProviderError` strings:
684/// connect/DNS/timeout/reset wording, not provider rejections.
685fn is_transient_message(message: &str) -> bool {
686    const NEEDLES: [&str; 10] = [
687        "error sending request",
688        "error trying to connect",
689        "connection refused",
690        "connection reset",
691        "connection closed",
692        "timed out",
693        "timeout",
694        "unreachable",
695        "dns error",
696        "temporary failure",
697    ];
698    let m = message.to_ascii_lowercase();
699    NEEDLES.iter().any(|needle| m.contains(needle))
700}
701
702/// Exponential backoff for `attempt` (1-based), capped, with jitter in
703/// [75%, 100%] so repeated retries don't align on the same instant.
704fn backoff(attempt: u32, base: Duration, cap: Duration) -> Duration {
705    let shift = attempt.saturating_sub(1).min(16);
706    let capped = base.saturating_mul(1u32 << shift).min(cap);
707    let millis = capped.as_millis() as u64;
708    if millis <= 1 {
709        return capped;
710    }
711    let jitter = clock_nanos() % (millis / 4 + 1);
712    Duration::from_millis(millis - jitter)
713}
714
715/// A cheap clock seed for backoff jitter.
716fn clock_nanos() -> u64 {
717    SystemTime::now()
718        .duration_since(UNIX_EPOCH)
719        .map(|d| u64::from(d.subsec_nanos()))
720        .unwrap_or(0)
721}
722
723/// Build one connect attempt: a fresh future per call, so the handshake can be
724/// retried. Once the response is in hand a mid-stream error is not retried.
725fn open_stream(
726    client: rig::providers::openai::Client,
727    endpoint: LlmEndpoint,
728    model: &str,
729    request: CompletionRequest,
730) -> Pin<Box<dyn std::future::Future<Output = Result<StreamingCompletionResponse, CompletionError>> + Send>> {
731    match endpoint {
732        LlmEndpoint::Chat => {
733            let model = client.completions_api().completion_model(model.to_string());
734            Box::pin(async move { model.stream(request).await })
735        }
736        LlmEndpoint::Responses => {
737            let model = client.completion_model(model.to_string());
738            Box::pin(async move { model.stream(request).await })
739        }
740    }
741}
742
743/// Outcome of the pre-content phase of a stream turn: a stream with its first
744/// (buffered) item — which may itself be an error, in which case the caller
745/// classifies it — or a final failure.
746enum PreContent<S> {
747    Ready {
748        stream: S,
749        first: Option<Result<StreamedAssistantContent, CompletionError>>,
750    },
751    /// Final failure: not in the transient class, or retries exhausted.
752    Failed(CompletionError),
753    /// The consumer dropped `tx` during a backoff wait (cancellation).
754    Cancelled,
755}
756
757/// Shared retry bookkeeping for one turn's pre-content window: how many retries
758/// remain across the connect/peek and any later pre-content errors, plus how
759/// many failed attempts have accumulated so backoff keeps growing monotonically.
760#[derive(Clone, Copy)]
761struct RetryState {
762    budget: u32,
763    attempts: u32,
764}
765
766impl RetryState {
767    fn new(max: u32) -> Self {
768        RetryState {
769            budget: max,
770            attempts: 0,
771        }
772    }
773}
774
775/// Whether a failed stream call is of a hard-fatal class that retrying or
776/// re-feeding cannot fix: hard client statuses (400/401/403/404/422) and
777/// build/parse failures (Json/Request/Response/Url). Transient-class errors
778/// (see [`retryable`]) may exhaust their retries — those surface with
779/// `fatal: false` so the loop can feed the failure back to the model instead
780/// of ending the run.
781fn fatal_class(error: &CompletionError) -> bool {
782    !retryable(error)
783}
784
785/// Open a stream, retrying transient failures on the connect and on anything
786/// the peek surfaces. rig's OpenAI path defers the HTTP request into the stream,
787/// so a connection failure surfaces on the first poll; retrying only the
788/// connect would miss it. `open` is re-invoked per attempt (a fresh future).
789/// The [`RetryState`] budget is shared with the forwarding loop, so a transient
790/// error that slips past the peek (still before any content was forwarded) is
791/// likewise retried at most `policy.max` times total per turn. Once any content
792/// has been forwarded a mid-stream error is surfaced as-is — retrying would
793/// duplicate output. Each backoff is abandoned if the consumer drops `tx`.
794/// A synthesized provider error for a stream stall. The message carries a
795/// needle `is_transient_message` already matches ("timed out"), so `retryable()`
796/// is `true` and `fatal_class()` is `false` (B1) — a stall rides the existing
797/// transient-retry / feed-back paths unchanged.
798fn stall_error(d: Duration) -> CompletionError {
799    CompletionError::ProviderError(format!("stream stalled: no item within {d:?} (timed out)"))
800}
801/// A synthesized stream-truncation failure, shaped like [`stall_error`]. The
802/// stream ended without its terminal record: rig emits neither a `Final` nor an
803/// `Err` when the provider SSE body ends early (no `[DONE]`/finish_reason), so
804/// the normalized stream simply ends with `None`. Not necessarily `retryable()`
805/// — the pre-content branch re-opens unconditionally — the message only feeds
806/// `retry_wait`'s notice text and the surfaced `Error.message`.
807fn truncation_error() -> CompletionError {
808    CompletionError::ProviderError(
809        "stream ended before its terminal record (truncated)".to_string(),
810    )
811}
812
813async fn connect_with_retry<S, F>(
814    policy: &RetryPolicy,
815    state: &mut RetryState,
816    mut open: F,
817    tx: &tokio::sync::mpsc::UnboundedSender<LlmStreamEvent>,
818) -> PreContent<S>
819where
820    F: FnMut() -> Pin<Box<dyn std::future::Future<Output = Result<S, CompletionError>> + Send>>,
821    S: futures::Stream<Item = Result<StreamedAssistantContent, CompletionError>> + Unpin + Send,
822{
823    loop {
824        state.attempts += 1;
825        // Race the connect against `ttft` — zero-aware: `ZERO` disables the
826        // deadline and awaits the connect directly (never `sleep(ZERO)`, which
827        // would fire at once). A connect stall yields a transient `stall_error`.
828        let opened = if policy.ttft.is_zero() {
829            Ok(open().await)
830        } else {
831            tokio::select! {
832                biased;
833                r = open() => Ok(r),
834                _ = tokio::time::sleep(policy.ttft) => Err(()),
835            }
836        };
837        let mut stream = match opened {
838            Ok(Ok(stream)) => stream,
839            Ok(Err(error)) => {
840                if state.budget > 0 && retryable(&error) {
841                    state.budget -= 1;
842                    if !retry_wait(tx, policy, state.attempts, &error).await {
843                        return PreContent::Cancelled;
844                    }
845                    continue;
846                }
847                return PreContent::Failed(error);
848            }
849            Err(()) => {
850                let error = stall_error(policy.ttft);
851                if state.budget > 0 {
852                    state.budget -= 1;
853                    if !retry_wait(tx, policy, state.attempts, &error).await {
854                        return PreContent::Cancelled;
855                    }
856                    continue;
857                }
858                return PreContent::Failed(error);
859            }
860        };
861        // Peek the first item under the SAME `ttft`: rig defers the request into
862        // the stream, so a transport failure (or a stall) surfaces here.
863        match next_or_stall(&mut stream, policy.ttft).await {
864            None => return PreContent::Ready { stream, first: None },
865            Some(Stall::Item(Err(error))) => {
866                if state.budget > 0 && retryable(&error) {
867                    state.budget -= 1;
868                    if !retry_wait(tx, policy, state.attempts, &error).await {
869                        return PreContent::Cancelled;
870                    }
871                    continue;
872                }
873                // A first-item error that cannot be retried is still an error
874                // item: hand it to the forwarding loop, which classifies it
875                // (fatal vs feed-back) like any other.
876                return PreContent::Ready {
877                    stream,
878                    first: Some(Err(error)),
879                };
880            }
881            Some(Stall::Item(Ok(content))) => {
882                return PreContent::Ready {
883                    stream,
884                    first: Some(Ok(content)),
885                };
886            }
887            Some(Stall::Idle) => {
888                let error = stall_error(policy.ttft);
889                if state.budget > 0 {
890                    state.budget -= 1;
891                    if !retry_wait(tx, policy, state.attempts, &error).await {
892                        return PreContent::Cancelled;
893                    }
894                    continue;
895                }
896                return PreContent::Failed(error);
897            }
898        }
899    }
900}
901
902/// One pre-content retry shared by the idle-stall arm and the truncation arm:
903/// spend one attempt + one budget unit on the shared [`RetryState`], back off
904/// (abandoning if the consumer dropped `tx`), then re-open via
905/// [`connect_with_retry`] with the SAME state so the budget stays monotone.
906/// Returns the same [`PreContent<S>`] the direct connect path returns, so both
907/// callers match identically — no new enum, no duplicated match.
908async fn pre_content_retry<S, F>(
909    policy: &RetryPolicy,
910    state: &mut RetryState,
911    open: &F,
912    tx: &tokio::sync::mpsc::UnboundedSender<LlmStreamEvent>,
913    error: &CompletionError,
914) -> PreContent<S>
915where
916    F: Fn() -> Pin<Box<dyn std::future::Future<Output = Result<S, CompletionError>> + Send>>,
917    S: futures::Stream<Item = Result<StreamedAssistantContent, CompletionError>> + Unpin + Send,
918{
919    state.attempts += 1;
920    state.budget -= 1;
921    if !retry_wait(tx, policy, state.attempts, error).await {
922        return PreContent::Cancelled;
923    }
924    connect_with_retry(policy, state, open, tx).await
925}
926
927/// The server's `Retry-After` header (integer delta-seconds), when the failed
928/// response preserved one. `None` when there is no preserved response, no
929/// header, a non-UTF-8 value, or a value that is not a non-negative integer.
930///
931/// SCOPE: integer delta-seconds ONLY. The HTTP-date form of `Retry-After`
932/// (RFC 9110 `Retry-After = HTTP-date / delay-seconds`) is deliberately out of
933/// scope: a date would need wall-clock parsing against the loop's monotonic
934/// `tokio::time` deadline, and a subtly wrong date is worse than ignoring it —
935/// a provider that sends one simply gets the normal backoff.
936fn retry_after(error: &CompletionError) -> Option<Duration> {
937    // `provider_response_headers()` already unwraps both the
938    // `HttpError(InvalidStatusCodeWithDetails)` and the `ProviderResponse`
939    // variants, so this reads either.
940    let raw = error.provider_response_headers()?.get("retry-after")?;
941    // A non-UTF-8 header value is ignored, not an error.
942    let seconds: u64 = raw.to_str().ok()?.trim().parse().ok()?;
943    Some(Duration::from_secs(seconds))
944}
945
946/// Emit a retry notice and sleep the backoff. Returns false if the consumer
947/// dropped `tx` (cancellation) during the wait.
948async fn retry_wait(
949    tx: &tokio::sync::mpsc::UnboundedSender<LlmStreamEvent>,
950    policy: &RetryPolicy,
951    attempt: u32,
952    error: &CompletionError,
953) -> bool {
954    // A server-supplied `Retry-After` is authoritative and wins over the local
955    // backoff, but it is clamped by the existing `policy.cap` so a hostile or
956    // buggy `Retry-After: 86400` cannot park the turn for a day.
957    let delay = retry_after(error)
958        .map(|d| d.min(policy.cap))
959        .unwrap_or_else(|| backoff(attempt, policy.base, policy.cap));
960    let _ = tx.send(LlmStreamEvent::Retrying {
961        attempt,
962        max: policy.max,
963        reason: error.to_string(),
964    });
965    tokio::select! {
966        biased;
967        _ = tx.closed() => false,
968        _ = tokio::time::sleep(delay) => true,
969    }
970}
971
972fn build_request(
973    messages: &[AgentMessage],
974    system: &str,
975    tools: &[ToolDefinition],
976    opts: &LlmOpts,
977) -> CompletionRequest {
978    let mut chat_history = Vec::with_capacity(messages.len() + 1);
979    if !system.is_empty() {
980        chat_history.push(Message::system(system));
981    }
982    chat_history.extend(messages.iter().filter_map(to_rig_message));
983    CompletionRequest {
984        model: None,
985        // System instructions ride as the leading Message::System, the
986        // canonical form rig's own request builder funnels `preamble` into.
987        preamble: None,
988        chat_history,
989        documents: Vec::new(),
990        tools: tools.to_vec(),
991        temperature: opts.temperature,
992        max_tokens: None,
993        tool_choice: None,
994        additional_params: reasoning_params(opts),
995        output_schema: None,
996        record_telemetry_content: false,
997    }
998}
999
1000/// Effort levels accepted by the Responses wire. rig deserializes
1001/// `reasoning.effort` into a snake_case enum client-side (`ReasoningEffort`),
1002/// so any other string fails with an opaque RequestError at send time;
1003/// `invalid_effort` surfaces a clear error instead. Chat forwards any string.
1004const RESPONSES_EFFORT_LEVELS: &[&str] =
1005    &["none", "minimal", "low", "medium", "high", "xhigh", "max"];
1006
1007/// `Some(message)` when `opts.effort` cannot cross the configured wire: the
1008/// Responses endpoint validates effort client-side against the enum above,
1009/// while Chat forwards any string to the provider.
1010fn invalid_effort(opts: &LlmOpts) -> Option<String> {
1011    if opts.endpoint == LlmEndpoint::Responses
1012        && let Some(effort) = opts.effort.as_deref()
1013        && !RESPONSES_EFFORT_LEVELS.contains(&effort)
1014    {
1015        return Some(format!(
1016            "effort `{effort}` is not valid for the Responses endpoint; use one of: {}",
1017            RESPONSES_EFFORT_LEVELS.join(", ")
1018        ));
1019    }
1020    None
1021}
1022
1023/// Reasoning/effort passthrough, fanned out per wire shape.
1024///
1025/// Chat: `opts.effort` alone → `{"reasoning_effort": effort}`; None sends
1026/// nothing (unchanged), so backends that reject unknown fields keep working.
1027///
1028/// Responses: ALWAYS emits `{"reasoning": {"summary": "auto"}}`, merging
1029/// `"effort"` in when `opts.effort` is `Some`. The summary is what makes
1030/// reasoning DISPLAYABLE: rig auto-injects `include:["reasoning.encrypted_content"]`
1031/// whenever `additional_params` carries `reasoning` (rig-core-0.42.0
1032/// responses_api/mod.rs:1461 — `if additional_parameters.reasoning.is_some()`),
1033/// and `ReasoningSummaryLevel::Auto` serializes to snake_case "auto"
1034/// (responses_api/mod.rs:2133). Without a summary the model emits an
1035/// encrypted-only reasoning item whose displayable text is empty (see the
1036/// `map_item` guard below).
1037///
1038/// `invalid_effort` still validates the level client-side before this runs.
1039fn reasoning_params(opts: &LlmOpts) -> Option<serde_json::Value> {
1040    match opts.endpoint {
1041        LlmEndpoint::Chat => opts
1042            .effort
1043            .as_ref()
1044            .map(|effort| serde_json::json!({ "reasoning_effort": effort })),
1045        LlmEndpoint::Responses => {
1046            let mut reasoning = serde_json::Map::new();
1047            reasoning.insert("summary".into(), serde_json::json!("auto"));
1048            if let Some(effort) = &opts.effort {
1049                reasoning.insert("effort".into(), serde_json::json!(effort));
1050            }
1051            Some(serde_json::json!({ "reasoning": reasoning }))
1052        }
1053    }
1054}
1055
1056/// Prefix applied to a failed tool's output *on the wire only*. The chat wire
1057/// has no structured error flag for tool results (rig's `ToolResult` only
1058/// carries text/content), so a failed tool's output must be annotated in the
1059/// text the model sees — otherwise a tool failure looks exactly like a
1060/// success. The UI and the JSONL session keep the raw output; only the wire
1061/// copy read by the LLM is marked.
1062const TOOL_OUTPUT_ERROR_PREFIX: &str = "ERROR: ";
1063
1064fn to_rig_message(m: &AgentMessage) -> Option<Message> {
1065    match m {
1066        AgentMessage::User { content } => {
1067            let content: Vec<UserContent> = content
1068                .iter()
1069                .filter_map(|b| match b {
1070                    ContentBlock::Text { text } => Some(UserContent::text(text.clone())),
1071                    _ => None,
1072                })
1073                .collect();
1074            (!content.is_empty()).then_some(Message::User { content })
1075        }
1076        AgentMessage::Assistant { content, .. } => {
1077            let content: Vec<AssistantContent> = content
1078                .iter()
1079                .map(|b| match b {
1080                    ContentBlock::Text { text } => AssistantContent::text(text.clone()),
1081                    ContentBlock::Thinking { text } => {
1082                        AssistantContent::Reasoning(Reasoning::new(text))
1083                    }
1084                    ContentBlock::ToolCall {
1085                        id,
1086                        name,
1087                        arguments,
1088                    } => AssistantContent::ToolCall(ToolCall::new(
1089                        ToolCallId::new_or_mint(id.clone()),
1090                        ToolFunction::new(name.clone(), arguments.clone()),
1091                    )),
1092                })
1093                .collect();
1094            (!content.is_empty()).then_some(Message::Assistant { id: None, content })
1095        }
1096        AgentMessage::ToolResult {
1097            tool_call_id,
1098            name,
1099            output,
1100            is_error,
1101        } => {
1102            // The wire has no structured error flag for tool results, so mark
1103            // failures in the text the model sees. The UI and the session keep
1104            // the raw output — only this wire copy is prefixed.
1105            let output = if *is_error {
1106                format!("{TOOL_OUTPUT_ERROR_PREFIX}{output}")
1107            } else {
1108                output.clone()
1109            };
1110            Some(Message::tool_result(
1111                tool_call_id.clone(),
1112                name.clone(),
1113                output,
1114            ))
1115        }
1116    }
1117}
1118
1119fn map_item(item: StreamedAssistantContent) -> Vec<LlmStreamEvent> {
1120    match item {
1121        StreamedAssistantContent::Text(text) => vec![LlmStreamEvent::TextDelta(text.text)],
1122        StreamedAssistantContent::ReasoningDelta { reasoning, .. } => {
1123            vec![LlmStreamEvent::ThinkingDelta(reasoning)]
1124        }
1125        // A complete reasoning block supersedes the streamed deltas: the Responses
1126        // wire restates the whole reasoning item after the deltas. Dead on
1127        // chat-completions (deltas only) but LIVE on Responses — it must not be
1128        // dropped.
1129        StreamedAssistantContent::Reasoning { reasoning, .. } => {
1130            // Replacement semantics: the complete block supersedes the
1131            // accumulated deltas (the loop drops prior thinking on this event).
1132            let text = reasoning_text(&reasoning);
1133            if text.is_empty() {
1134                // Encrypted-only terminal item: no displayable text. Emitting
1135                // ThinkingReplace("") would reach `loop_.rs:replace_thinking`,
1136                // which retains-out ALL accumulated thinking and pushes an empty
1137                // block — erasing the streamed ThinkingDeltas. Yield nothing.
1138                vec![]
1139            } else {
1140                vec![LlmStreamEvent::ThinkingReplace(text)]
1141            }
1142        }
1143        // Deltas are dropped: rig's streaming fold accumulates every tool-call
1144        // fragment and emits the complete `ToolCall` on `ToolInputEnd` (both
1145        // chat-completions and responses wires do this), so the reassembled
1146        // call always follows. A provider that ends input without closing it
1147        // loses the call inside rig itself; the loop surfaces that as an error
1148        // rather than silently treating the turn as done.
1149        StreamedAssistantContent::ToolCallDelta { .. } => vec![],
1150        StreamedAssistantContent::ToolCall { tool_call, .. } => vec![LlmStreamEvent::ToolCall {
1151            id: tool_call.id.as_str().to_string(),
1152            name: tool_call.function.name,
1153            arguments: tool_call.function.arguments,
1154        }],
1155        StreamedAssistantContent::Final(final_record) => {
1156            vec![LlmStreamEvent::Done {
1157                stop_reason: final_record
1158                    .finish_reason
1159                    .map(finish_reason)
1160                    .unwrap_or(StopReason::Stop),
1161                usage: Some(convert_usage(final_record.usage)),
1162            }]
1163        }
1164        StreamedAssistantContent::Unknown(_) => vec![],
1165    }
1166}
1167
1168fn finish_reason(reason: FinishReason) -> StopReason {
1169    match reason {
1170        FinishReason::Stop => StopReason::Stop,
1171        FinishReason::Length => StopReason::Length,
1172        FinishReason::ToolCalls => StopReason::ToolUse,
1173        FinishReason::ContentFilter | FinishReason::Other(_) => StopReason::Stop,
1174    }
1175}
1176
1177fn convert_usage(usage: rig::completion::Usage) -> crate::message::Usage {
1178    crate::message::Usage {
1179        input_tokens: usage.input_tokens,
1180        output_tokens: usage.output_tokens,
1181        // Zero is rig's documented sentinel for "not reported".
1182        cache_read_tokens: (usage.cached_input_tokens > 0).then_some(usage.cached_input_tokens),
1183        cache_write_tokens: (usage.cache_creation_input_tokens > 0)
1184            .then_some(usage.cache_creation_input_tokens),
1185    }
1186}
1187
1188fn reasoning_text(reasoning: &Reasoning) -> String {
1189    reasoning
1190        .content
1191        .iter()
1192        .filter_map(|block| match block {
1193            ReasoningContent::Text { text, .. } => Some(text.as_str()),
1194            ReasoningContent::Summary(summary) => Some(summary.as_str()),
1195            ReasoningContent::Redacted { data } => Some(data.as_str()),
1196            ReasoningContent::Encrypted(_) => None,
1197        })
1198        .collect::<Vec<_>>()
1199        .join("\n")
1200}
1201
1202#[cfg(test)]
1203mod tests {
1204    use super::*;
1205    use rig::completion::Usage as RigUsage;
1206    use rig::message::{ToolResultContent, UserContent};
1207    use rig::streaming::{RawStreamingToolCall, StreamFinal, ToolCallDeltaContent};
1208    use serde_json::json;
1209
1210    fn tool_call_item() -> StreamedAssistantContent {
1211        let raw =
1212            RawStreamingToolCall::new("call_1".to_string(), "run".to_string(), json!({"x": 1}));
1213        StreamedAssistantContent::ToolCall {
1214            tool_call: raw.into(),
1215            internal_call_id: "ic1".to_string(),
1216        }
1217    }
1218
1219    fn final_item(finish_reason: Option<FinishReason>) -> StreamedAssistantContent {
1220        let usage = RigUsage {
1221            input_tokens: 10,
1222            output_tokens: 5,
1223            total_tokens: 15,
1224            cached_input_tokens: 3,
1225            cache_creation_input_tokens: 2,
1226            ..RigUsage::default()
1227        };
1228        let final_record = StreamFinal::new("test", usage);
1229        StreamedAssistantContent::Final(match finish_reason {
1230            Some(reason) => final_record.with_finish_reason(reason),
1231            None => final_record,
1232        })
1233    }
1234
1235    #[test]
1236    fn text_item_maps_to_text_delta() {
1237        let events = map_item(StreamedAssistantContent::Text(rig::message::Text::new(
1238            "hi",
1239        )));
1240        assert_eq!(events, vec![LlmStreamEvent::TextDelta("hi".to_string())]);
1241    }
1242
1243    #[test]
1244    fn reasoning_delta_maps_to_thinking_delta() {
1245        let events = map_item(StreamedAssistantContent::ReasoningDelta {
1246            id: "r1".to_string(),
1247            provider_id: None,
1248            reasoning: "hmm".to_string(),
1249        });
1250        assert_eq!(
1251            events,
1252            vec![LlmStreamEvent::ThinkingDelta("hmm".to_string())]
1253        );
1254    }
1255
1256    #[test]
1257    fn reasoning_block_replaces_with_full_text() {
1258        let events = map_item(StreamedAssistantContent::Reasoning {
1259            reasoning: Reasoning::new("full thought"),
1260            id: "r1".to_string(),
1261        });
1262        assert_eq!(
1263            events,
1264            vec![LlmStreamEvent::ThinkingReplace("full thought".to_string())]
1265        );
1266    }
1267
1268    #[test]
1269    fn reasoning_text_joins_summary_text_and_redacted_drops_encrypted() {
1270        // Exactly rig's `Reasoning::display_text()` contract: Summary | Text |
1271        // Redacted join with "\n"; Encrypted is skipped.
1272        let reasoning = Reasoning {
1273            id: None,
1274            content: vec![
1275                ReasoningContent::Summary("a".to_string()),
1276                ReasoningContent::Text {
1277                    text: "b".to_string(),
1278                    signature: None,
1279                },
1280                ReasoningContent::Redacted {
1281                    data: "c".to_string(),
1282                },
1283                ReasoningContent::Encrypted("opaque".to_string()),
1284            ],
1285        };
1286        assert_eq!(reasoning_text(&reasoning), "a\nb\nc");
1287    }
1288
1289    #[test]
1290    fn reasoning_text_summary_only_is_the_summary() {
1291        let reasoning = Reasoning::summaries(vec!["just a summary".to_string()]);
1292        assert_eq!(reasoning_text(&reasoning), "just a summary");
1293    }
1294
1295    #[test]
1296    fn reasoning_text_encrypted_only_is_empty() {
1297        let reasoning = Reasoning::encrypted("opaque");
1298        assert_eq!(reasoning_text(&reasoning), "");
1299    }
1300
1301    #[test]
1302    fn summary_reasoning_block_replaces_with_full_text() {
1303        let events = map_item(StreamedAssistantContent::Reasoning {
1304            reasoning: Reasoning::summaries(vec!["a summary".to_string()]),
1305            id: "rs_1".to_string(),
1306        });
1307        assert_eq!(
1308            events,
1309            vec![LlmStreamEvent::ThinkingReplace("a summary".to_string())]
1310        );
1311    }
1312
1313    #[test]
1314    fn encrypted_only_reasoning_block_emits_no_event() {
1315        // An encrypted-only terminal item has no displayable text; emitting
1316        // ThinkingReplace("") would erase the streamed ThinkingDeltas downstream
1317        // (see the `map_item` guard).
1318        let events = map_item(StreamedAssistantContent::Reasoning {
1319            reasoning: Reasoning::encrypted("opaque"),
1320            id: "rs_1".to_string(),
1321        });
1322        assert!(events.is_empty(), "got: {events:?}");
1323    }
1324
1325    #[test]
1326    fn tool_call_delta_ignored() {
1327        for content in [
1328            ToolCallDeltaContent::Name("run".to_string()),
1329            ToolCallDeltaContent::Delta("{\"x".to_string()),
1330        ] {
1331            let events = map_item(StreamedAssistantContent::ToolCallDelta {
1332                internal_call_id: "ic1".to_string(),
1333                content,
1334            });
1335            assert!(events.is_empty());
1336        }
1337    }
1338
1339    #[test]
1340    fn tool_call_maps_to_call_event() {
1341        let events = map_item(tool_call_item());
1342        assert_eq!(
1343            events,
1344            vec![LlmStreamEvent::ToolCall {
1345                id: "call_1".to_string(),
1346                name: "run".to_string(),
1347                arguments: json!({"x": 1}),
1348            }]
1349        );
1350    }
1351
1352    #[test]
1353    fn final_maps_finish_reason_and_usage() {
1354        let cases = [
1355            (FinishReason::Stop, StopReason::Stop),
1356            (FinishReason::Length, StopReason::Length),
1357            (FinishReason::ToolCalls, StopReason::ToolUse),
1358            (FinishReason::ContentFilter, StopReason::Stop),
1359            (FinishReason::Other("weird".to_string()), StopReason::Stop),
1360        ];
1361        for (reason, expected) in cases {
1362            let events = map_item(final_item(Some(reason)));
1363            assert_eq!(
1364                events,
1365                vec![LlmStreamEvent::Done {
1366                    stop_reason: expected,
1367                    usage: Some(crate::message::Usage {
1368                        input_tokens: 10,
1369                        output_tokens: 5,
1370                        cache_read_tokens: Some(3),
1371                        cache_write_tokens: Some(2),
1372                    }),
1373                }],
1374                "wrong mapping for {expected:?}"
1375            );
1376        }
1377    }
1378
1379    #[test]
1380    fn final_without_finish_reason_defaults_to_stop() {
1381        let events = map_item(final_item(None));
1382        match &events[..] {
1383            [LlmStreamEvent::Done { stop_reason, .. }] => {
1384                assert_eq!(*stop_reason, StopReason::Stop)
1385            }
1386            other => panic!("unexpected events: {other:?}"),
1387        }
1388    }
1389
1390    #[test]
1391    fn unknown_item_ignored() {
1392        let events = map_item(StreamedAssistantContent::Unknown(
1393            rig::streaming::UnknownPayload::new(json!({"mystery": 1})),
1394        ));
1395        assert!(events.is_empty());
1396    }
1397
1398    #[test]
1399    fn build_request_round_trips_context() {
1400        let opts = LlmOpts {
1401            model: "m1".to_string(),
1402            base_url: None,
1403            api_key: None,
1404            temperature: Some(0.5),
1405            ..LlmOpts::default()
1406        };
1407        let tools = vec![ToolDefinition {
1408            name: "get_weather".to_string(),
1409            description: "Weather lookup".to_string(),
1410            parameters: json!({"type": "object"}),
1411        }];
1412        let messages = vec![
1413            AgentMessage::user_text("What is the weather?"),
1414            AgentMessage::Assistant {
1415                content: vec![
1416                    ContentBlock::Thinking { text: "hmm".into() },
1417                    ContentBlock::ToolCall {
1418                        id: "t1".into(),
1419                        name: "get_weather".into(),
1420                        arguments: json!({"city": "Tokyo"}),
1421                    },
1422                    ContentBlock::Text {
1423                        text: "Let me check.".into(),
1424                    },
1425                ],
1426                stop_reason: StopReason::ToolUse,
1427                usage: None,
1428                model: None,
1429            },
1430            AgentMessage::ToolResult {
1431                tool_call_id: "t1".into(),
1432                name: "get_weather".into(),
1433                output: "sunny".into(),
1434                is_error: false,
1435            },
1436        ];
1437
1438        let request = build_request(&messages, "You are terse.", &tools, &opts);
1439
1440        assert!(request.preamble.is_none());
1441        assert_eq!(request.temperature, Some(0.5));
1442        assert_eq!(request.tools.len(), 1);
1443        assert_eq!(request.tools[0].name, "get_weather");
1444
1445        let [
1446            Message::System { content: system },
1447            Message::User { content: user },
1448            Message::Assistant {
1449                content: assistant, ..
1450            },
1451            Message::User {
1452                content: tool_result,
1453            },
1454        ] = &request.chat_history[..]
1455        else {
1456            panic!("unexpected chat_history shape: {:?}", request.chat_history);
1457        };
1458        assert_eq!(system, "You are terse.");
1459        assert!(
1460            matches!(user.as_slice(), [UserContent::Text(t)] if t.text == "What is the weather?")
1461        );
1462        assert_eq!(assistant.len(), 3);
1463        assert!(
1464            matches!(&assistant[0], AssistantContent::Reasoning(r) if reasoning_text(r) == "hmm")
1465        );
1466        assert!(matches!(
1467            &assistant[1],
1468            AssistantContent::ToolCall(call)
1469                if call.id.as_str() == "t1"
1470                    && call.function.name == "get_weather"
1471                    && call.function.arguments == json!({"city": "Tokyo"})
1472        ));
1473        assert!(matches!(&assistant[2], AssistantContent::Text(t) if t.text == "Let me check."));
1474        assert!(matches!(
1475            tool_result.as_slice(),
1476            [UserContent::ToolResult(r)] if r.call.as_str() == "t1" && r.name == "get_weather"
1477        ));
1478    }
1479
1480    #[test]
1481    fn tool_result_error_is_annotated_on_the_wire_only() {
1482        let err = AgentMessage::ToolResult {
1483            tool_call_id: "t1".into(),
1484            name: "echo".into(),
1485            output: "boom".into(),
1486            is_error: true,
1487        };
1488        let ok = AgentMessage::ToolResult {
1489            tool_call_id: "t1".into(),
1490            name: "echo".into(),
1491            output: "boom".into(),
1492            is_error: false,
1493        };
1494
1495        fn wire_text(m: &AgentMessage) -> String {
1496            match to_rig_message(m) {
1497                Some(Message::User { content }) => match &content[..] {
1498                    [UserContent::ToolResult(r)] => match &r.content[..] {
1499                        [ToolResultContent::Text(t)] => t.text.clone(),
1500                        other => panic!("unexpected tool result content: {other:?}"),
1501                    },
1502                    other => panic!("unexpected user content: {other:?}"),
1503                },
1504                other => panic!("unexpected wire message: {other:?}"),
1505            }
1506        }
1507
1508        // Success passes through byte-for-byte; a failed tool is marked.
1509        assert_eq!(wire_text(&ok), "boom");
1510        assert_eq!(wire_text(&err), "ERROR: boom");
1511        // The stored message itself is untouched — the flag is only consumed
1512        // at the wire seam.
1513        match err {
1514            AgentMessage::ToolResult {
1515                output, is_error, ..
1516            } => {
1517                assert_eq!(output, "boom");
1518                assert!(is_error);
1519            }
1520            _ => unreachable!(),
1521        }
1522    }
1523
1524    #[test]
1525    fn endpoint_defaults_to_chat() {
1526        let opts = LlmOpts {
1527            model: "m1".to_string(),
1528            base_url: None,
1529            api_key: None,
1530            temperature: None,
1531            endpoint: LlmEndpoint::Chat,
1532            effort: None,
1533            session_id: None,
1534            retry: RetryPolicy::default(),
1535            model_profiles: ModelProfiles::default(),
1536        };
1537        assert_eq!(opts.endpoint, LlmEndpoint::Chat);
1538        assert_eq!(LlmOpts::default().endpoint, LlmEndpoint::Chat);
1539    }
1540
1541    #[test]
1542    fn no_effort_sends_no_additional_params() {
1543        let opts = LlmOpts::default();
1544        let request = build_request(&[], "sys", &[], &opts);
1545        assert_eq!(request.additional_params, None);
1546    }
1547
1548    #[test]
1549    fn chat_effort_maps_to_reasoning_effort_param() {
1550        let opts = LlmOpts {
1551            effort: Some("high".into()),
1552            endpoint: LlmEndpoint::Chat,
1553            ..LlmOpts::default()
1554        };
1555        let request = build_request(&[], "sys", &[], &opts);
1556        assert_eq!(
1557            request.additional_params,
1558            Some(json!({ "reasoning_effort": "high" }))
1559        );
1560    }
1561
1562    #[test]
1563    fn responses_effort_maps_to_reasoning_object() {
1564        let opts = LlmOpts {
1565            effort: Some("xhigh".into()),
1566            endpoint: LlmEndpoint::Responses,
1567            ..LlmOpts::default()
1568        };
1569        let request = build_request(&[], "sys", &[], &opts);
1570        assert_eq!(
1571            request.additional_params,
1572            Some(json!({ "reasoning": { "effort": "xhigh", "summary": "auto" } }))
1573        );
1574    }
1575
1576    #[test]
1577    fn responses_always_requests_a_reasoning_summary() {
1578        // The case that was invisible before this change: with NO effort the
1579        // Responses wire must still ask for a reasoning summary — otherwise the
1580        // model emits an encrypted-only item whose displayable text is empty.
1581        let opts = LlmOpts {
1582            endpoint: LlmEndpoint::Responses,
1583            effort: None,
1584            ..LlmOpts::default()
1585        };
1586        let request = build_request(&[], "sys", &[], &opts);
1587        assert_eq!(
1588            request.additional_params,
1589            Some(json!({ "reasoning": { "summary": "auto" } }))
1590        );
1591    }
1592
1593    #[test]
1594    fn chat_effort_is_verbatim_passthrough() {
1595        // Free-style on the Chat wire: any string (including provider
1596        // dialects) is forwarded verbatim under `reasoning_effort`.
1597        for level in ["low", "medium", "my-custom-level", ""] {
1598            let opts = LlmOpts {
1599                effort: Some(level.into()),
1600                endpoint: LlmEndpoint::Chat,
1601                ..LlmOpts::default()
1602            };
1603            let request = build_request(&[], "sys", &[], &opts);
1604            assert_eq!(
1605                request.additional_params,
1606                Some(json!({ "reasoning_effort": level })),
1607                "level {level:?} must pass through verbatim on Chat"
1608            );
1609        }
1610    }
1611
1612    #[test]
1613    fn responses_effort_validation() {
1614        // Known snake_case levels pass; anything else is rejected with a clear
1615        // message (rig would otherwise fail the request with an opaque error).
1616        for level in ["none", "minimal", "low", "medium", "high", "xhigh", "max"] {
1617            let opts = LlmOpts {
1618                effort: Some(level.into()),
1619                endpoint: LlmEndpoint::Responses,
1620                ..LlmOpts::default()
1621            };
1622            assert_eq!(invalid_effort(&opts), None, "level {level} must be valid");
1623        }
1624        let opts = LlmOpts {
1625            effort: Some("deep".into()),
1626            endpoint: LlmEndpoint::Responses,
1627            ..LlmOpts::default()
1628        };
1629        let msg = invalid_effort(&opts).expect("unknown level must be rejected");
1630        assert!(msg.contains("not valid"), "got: {msg}");
1631        assert!(
1632            msg.contains("xhigh"),
1633            "message must list valid levels: {msg}"
1634        );
1635        // Chat stays verbatim even for a value Responses would reject.
1636        let opts = LlmOpts {
1637            effort: Some("deep".into()),
1638            endpoint: LlmEndpoint::Chat,
1639            ..LlmOpts::default()
1640        };
1641        assert_eq!(invalid_effort(&opts), None);
1642    }
1643
1644    #[test]
1645    fn client_sends_session_header_and_user_agent() {
1646        let opts = LlmOpts {
1647            api_key: Some("test-key".to_string()),
1648            session_id: Some("sess-123".to_string()),
1649            ..LlmOpts::default()
1650        };
1651        let headers = openai_client(&opts)
1652            .expect("client builds")
1653            .headers()
1654            .clone();
1655        assert_eq!(
1656            headers.get("x-opencode-session").expect("session header"),
1657            "sess-123"
1658        );
1659        assert!(headers.contains_key("user-agent"), "got: {headers:?}");
1660    }
1661
1662    #[test]
1663    fn client_omits_session_header_when_unset() {
1664        let opts = LlmOpts {
1665            api_key: Some("test-key".to_string()),
1666            ..LlmOpts::default()
1667        };
1668        let headers = openai_client(&opts)
1669            .expect("client builds")
1670            .headers()
1671            .clone();
1672        assert!(!headers.contains_key("x-opencode-session"));
1673        assert!(headers.contains_key("user-agent"), "got: {headers:?}");
1674    }
1675
1676    #[test]
1677    fn localhost_base_url_builds_without_key() {
1678        // Keyless local endpoints (Ollama, llama.cpp, ...) should work with no
1679        // configured key via a placeholder bearer token.
1680        for url in [
1681            "http://localhost:11434/v1",
1682            "http://127.0.0.1:8080/v1",
1683            "http://[::1]:8080/v1",
1684        ] {
1685            let opts = LlmOpts {
1686                base_url: Some(url.to_string()),
1687                api_key: None,
1688                ..LlmOpts::default()
1689            };
1690            assert!(
1691                openai_client(&opts).is_ok(),
1692                "keyless {url} must build a client"
1693            );
1694        }
1695    }
1696
1697    #[test]
1698    fn remote_base_url_without_key_is_error() {
1699        let opts = LlmOpts {
1700            base_url: Some("https://api.example.com/v1".to_string()),
1701            api_key: None,
1702            ..LlmOpts::default()
1703        };
1704        let err = openai_client(&opts).unwrap_err();
1705        assert!(err.contains("no API key"), "got: {err}");
1706    }
1707
1708    /// Unreachable endpoint: list_models surfaces a string error, no panic.
1709    #[tokio::test]
1710    async fn list_models_unreachable_is_error() {
1711        let opts = LlmOpts {
1712            base_url: Some("http://127.0.0.1:9/v1".to_string()),
1713            api_key: Some("test-key".to_string()),
1714            ..LlmOpts::default()
1715        };
1716        let err = tokio::time::timeout(std::time::Duration::from_secs(10), list_models(&opts))
1717            .await
1718            .expect("list terminates")
1719            .expect_err("unreachable host must error");
1720        assert!(!err.is_empty());
1721    }
1722
1723    #[tokio::test]
1724    async fn list_models_without_key_is_error() {
1725        let opts = LlmOpts::default();
1726        let err = list_models(&opts).await.expect_err("no key must error");
1727        assert!(err.contains("no API key"), "got: {err}");
1728    }
1729
1730    #[test]
1731    fn client_cache_reuses_until_key_changes() {
1732        let opts = LlmOpts {
1733            base_url: Some("http://127.0.0.1:9/v1".to_string()),
1734            api_key: Some("k".to_string()),
1735            session_id: Some("s1".to_string()),
1736            ..LlmOpts::default()
1737        };
1738        let cache = ClientCache::default();
1739        let builds = || cache.builds.load(std::sync::atomic::Ordering::SeqCst);
1740
1741        // Identical opts reuse the cached client (the whole point: one pool/TLS).
1742        cache.get(&opts).expect("builds");
1743        cache.get(&opts).expect("builds");
1744        assert_eq!(builds(), 1, "identical opts must reuse the client");
1745
1746        // model/temperature/effort ride on the request, not the client.
1747        let model_swap = LlmOpts {
1748            model: "other".to_string(),
1749            effort: Some("high".to_string()),
1750            ..opts.clone()
1751        };
1752        cache.get(&model_swap).expect("builds");
1753        assert_eq!(builds(), 1, "model/effort swap must reuse the client");
1754
1755        // Endpoint, key, and session header each force a rebuild.
1756        for changed in [
1757            LlmOpts {
1758                base_url: Some("http://127.0.0.1:8/v1".to_string()),
1759                ..opts.clone()
1760            },
1761            LlmOpts {
1762                api_key: Some("k2".to_string()),
1763                ..opts.clone()
1764            },
1765            LlmOpts {
1766                session_id: Some("s2".to_string()),
1767                ..opts.clone()
1768            },
1769        ] {
1770            let before = builds();
1771            cache.get(&changed).expect("builds");
1772            assert_eq!(builds(), before + 1, "a key change must rebuild the client");
1773        }
1774    }
1775
1776    #[test]
1777    fn settle_determines_the_client_key() {
1778        // `settle` fully determines base_url/api_key; the rig client caches on
1779        // those (not the endpoint), so an endpoint-only swap reuses it and a
1780        // base_url change rebuilds it.
1781        let base = LlmProvider {
1782            endpoint: LlmEndpoint::Chat,
1783            base_url: Some("http://127.0.0.1:9/v1".to_string()),
1784            api_key: Some("k".to_string()),
1785        };
1786        let mut profiles = ModelProfiles::new(base.clone());
1787        // Same provider, different endpoint (base_url/api_key inherited).
1788        profiles.insert(
1789            "ep-swap",
1790            LlmProfile {
1791                endpoint: Some(LlmEndpoint::Responses),
1792                ..LlmProfile::default()
1793            },
1794        );
1795        // A different provider entirely.
1796        profiles.insert(
1797            "moved",
1798            LlmProfile {
1799                base_url: Some("http://127.0.0.1:8/v1".to_string()),
1800                ..LlmProfile::default()
1801            },
1802        );
1803
1804        let settled = |model: &str| {
1805            let mut llm = LlmOpts {
1806                model: model.to_string(),
1807                endpoint: base.endpoint,
1808                base_url: base.base_url.clone(),
1809                api_key: base.api_key.clone(),
1810                model_profiles: profiles.clone(),
1811                ..LlmOpts::default()
1812            };
1813            llm.settle_provider();
1814            llm
1815        };
1816
1817        let cache = ClientCache::default();
1818        let builds = || cache.builds.load(std::sync::atomic::Ordering::SeqCst);
1819
1820        let plain = settled("m0");
1821        cache.get(&plain).expect("builds");
1822        assert_eq!(builds(), 1);
1823
1824        let ep_swap = settled("ep-swap");
1825        assert_eq!(ep_swap.endpoint, LlmEndpoint::Responses);
1826        assert_eq!(ep_swap.base_url, base.base_url, "the base_url is inherited");
1827        cache.get(&ep_swap).expect("reuses");
1828        assert_eq!(builds(), 1, "an endpoint-only swap must reuse the client");
1829
1830        let moved = settled("moved");
1831        cache.get(&moved).expect("rebuilds");
1832        assert_eq!(builds(), 2, "a base_url change must rebuild the client");
1833    }
1834
1835    #[test]
1836    fn responses_request_builds_against_default_client_type() {
1837        use rig::client::CompletionClient;
1838        let client = openai::Client::builder()
1839            .api_key("test-key")
1840            .base_url("http://127.0.0.1:9/v1")
1841            .build()
1842            .expect("client builds without network");
1843        let model = client.completion_model("some-model");
1844        let request = build_request(&[], "sys", &[], &LlmOpts::default());
1845        let fut = rig::completion::CompletionModel::stream(&model, request);
1846        drop(fut);
1847    }
1848
1849    /// Unreachable endpoint: the adapter surfaces an `Error` event and no
1850    /// `Done`, at the stream level (initial request failure).
1851    #[tokio::test]
1852    async fn adapter_surfaces_error_without_done() {
1853        let opts = LlmOpts {
1854            model: "m1".to_string(),
1855            base_url: Some("http://127.0.0.1:9/v1".to_string()),
1856            api_key: Some("test-key".to_string()),
1857            temperature: None,
1858            ..LlmOpts::default()
1859        };
1860        let stream_fn = rig_stream_fn();
1861        let stream = stream_fn(&[], "sys", &[], &opts);
1862        let events: Vec<LlmStreamEvent> =
1863            tokio::time::timeout(std::time::Duration::from_secs(10), stream.collect())
1864                .await
1865                .expect("stream terminates");
1866        assert!(
1867            events
1868                .iter()
1869                .any(|e| matches!(e, LlmStreamEvent::Error { .. })),
1870            "expected an Error event, got {events:?}"
1871        );
1872        assert!(
1873            !events
1874                .iter()
1875                .any(|e| matches!(e, LlmStreamEvent::Done { .. })),
1876            "no Done after error, got {events:?}"
1877        );
1878    }
1879
1880    /// Responses path hits the same shared forwarding loop: unreachable
1881    /// endpoint surfaces `Error` and no `Done`.
1882    #[tokio::test]
1883    async fn responses_adapter_surfaces_error_without_done() {
1884        let opts = LlmOpts {
1885            model: "m1".to_string(),
1886            base_url: Some("http://127.0.0.1:9/v1".to_string()),
1887            api_key: Some("test-key".to_string()),
1888            temperature: None,
1889            endpoint: LlmEndpoint::Responses,
1890            effort: None,
1891            session_id: None,
1892            retry: RetryPolicy::default(),
1893            model_profiles: ModelProfiles::default(),
1894        };
1895        let stream_fn = rig_stream_fn();
1896        let stream = stream_fn(&[], "sys", &[], &opts);
1897        let events: Vec<LlmStreamEvent> =
1898            tokio::time::timeout(std::time::Duration::from_secs(10), stream.collect())
1899                .await
1900                .expect("stream terminates");
1901        assert!(
1902            events
1903                .iter()
1904                .any(|e| matches!(e, LlmStreamEvent::Error { .. })),
1905            "expected an Error event, got {events:?}"
1906        );
1907        assert!(
1908            !events
1909                .iter()
1910                .any(|e| matches!(e, LlmStreamEvent::Done { .. })),
1911            "no Done after error, got {events:?}"
1912        );
1913    }
1914    /// Live smoke: real OpenAI-compatible endpoint, gated on env.
1915    #[tokio::test]
1916    #[ignore = "live test; run with WCODE_LIVE=1 WCODE_BASE_URL=... WCODE_API_KEY=... WCODE_MODEL=... cargo test -p wcode-harness --lib -- --ignored"]
1917    async fn live_rig_stream_fn() {
1918        let (Ok(base_url), Ok(api_key), Ok(model)) = (
1919            std::env::var("WCODE_BASE_URL"),
1920            std::env::var("WCODE_API_KEY"),
1921            std::env::var("WCODE_MODEL"),
1922        ) else {
1923            eprintln!("skipping: WCODE_BASE_URL / WCODE_API_KEY / WCODE_MODEL not set");
1924            return;
1925        };
1926        if std::env::var("WCODE_LIVE").ok().as_deref() != Some("1") {
1927            eprintln!("skipping: WCODE_LIVE != 1");
1928            return;
1929        }
1930        let opts = LlmOpts {
1931            model,
1932            base_url: Some(base_url),
1933            api_key: Some(api_key),
1934            temperature: None,
1935            ..LlmOpts::default()
1936        };
1937        let stream_fn = rig_stream_fn();
1938        let stream = stream_fn(
1939            &[AgentMessage::user_text("Reply with exactly: pong")],
1940            "You are a terse echo assistant.",
1941            &[],
1942            &opts,
1943        );
1944        let events: Vec<LlmStreamEvent> = stream.collect().await;
1945        let text = events
1946            .iter()
1947            .filter_map(|e| match e {
1948                LlmStreamEvent::TextDelta(t) => Some(t.as_str()),
1949                _ => None,
1950            })
1951            .collect::<String>();
1952        println!("text: {text:?}");
1953        println!("events: {events:?}");
1954        assert!(
1955            events
1956                .iter()
1957                .any(|e| matches!(e, LlmStreamEvent::Done { .. }))
1958        );
1959        assert!(
1960            !events
1961                .iter()
1962                .any(|e| matches!(e, LlmStreamEvent::Error { .. }))
1963        );
1964    }
1965}
1966
1967#[cfg(test)]
1968mod retry_tests {
1969    use super::*;
1970    use std::sync::Arc;
1971    use std::sync::atomic::{AtomicUsize, Ordering};
1972
1973    type Chunks =
1974        futures::stream::Iter<std::vec::IntoIter<Result<StreamedAssistantContent, CompletionError>>>;
1975
1976    /// A transient transport failure as rig surfaces it on the OpenAI path.
1977    fn transient() -> CompletionError {
1978        CompletionError::ProviderError(
1979            "Http client error: error sending request for url (http://x)".into(),
1980        )
1981    }
1982
1983    fn fatal() -> CompletionError {
1984        CompletionError::ProviderError("invalid_request_error: bad tool schema".into())
1985    }
1986
1987    fn empty() -> Chunks {
1988        futures::stream::iter(Vec::new())
1989    }
1990
1991    /// A terminal `Final` item (the stream's `Done`).
1992    fn done() -> StreamedAssistantContent {
1993        StreamedAssistantContent::Final(rig::streaming::StreamFinal::new(
1994            "test",
1995            rig::completion::Usage::default(),
1996        ))
1997    }
1998
1999    fn first_item_err(e: CompletionError) -> Chunks {
2000        futures::stream::iter(vec![Err(e)])
2001    }
2002
2003    #[test]
2004    fn transient_statuses_retry_client_errors_do_not() {
2005        for code in [408, 409, 429, 500, 502, 503, 504] {
2006            assert!(is_transient_status(code), "{code} should retry");
2007        }
2008        for code in [200, 400, 401, 403, 404, 422] {
2009            assert!(!is_transient_status(code), "{code} must not retry");
2010        }
2011    }
2012
2013    #[test]
2014    fn retryable_classifies_transport_vs_rejections() {
2015        assert!(retryable(&CompletionError::HttpError(
2016            rig::http_client::Error::StreamEnded
2017        )));
2018        assert!(retryable(&transient()));
2019        assert!(!retryable(&fatal()));
2020        assert!(!retryable(&CompletionError::ResponseError("bad".into())));
2021    }
2022
2023    #[test]
2024    fn backoff_grows_and_never_exceeds_the_cap() {
2025        let base = Duration::from_millis(100);
2026        let cap = Duration::from_millis(400);
2027        let d1 = backoff(1, base, cap);
2028        assert!(d1 <= base && d1 >= base * 3 / 4, "attempt 1 = {d1:?}");
2029        for n in 1..12 {
2030            assert!(backoff(n, base, cap) <= cap, "attempt {n} exceeds cap");
2031        }
2032    }
2033
2034    #[tokio::test(start_paused = true)]
2035    async fn connect_retries_transient_failures_then_succeeds() {
2036        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2037        let calls = Arc::new(AtomicUsize::new(0));
2038        let policy = RetryPolicy {
2039            max: 3,
2040            base: Duration::from_millis(10),
2041            cap: Duration::from_millis(50),
2042            ..RetryPolicy::default()
2043        };
2044        let c = calls.clone();
2045        let mut state = RetryState::new(policy.max);
2046        let out = connect_with_retry(
2047            &policy,
2048            &mut state,
2049            move || {
2050                let n = c.fetch_add(1, Ordering::SeqCst) + 1;
2051                Box::pin(async move {
2052                    if n <= 2 {
2053                        Err(transient())
2054                    } else {
2055                        Ok(empty())
2056                    }
2057                })
2058            },
2059            &tx,
2060        )
2061        .await;
2062        assert!(matches!(out, PreContent::Ready { .. }));
2063        assert_eq!(calls.load(Ordering::SeqCst), 3);
2064        assert!(matches!(
2065            rx.try_recv(),
2066            Ok(LlmStreamEvent::Retrying { attempt: 1, .. })
2067        ));
2068        assert!(matches!(
2069            rx.try_recv(),
2070            Ok(LlmStreamEvent::Retrying { attempt: 2, .. })
2071        ));
2072    }
2073
2074    #[tokio::test(start_paused = true)]
2075    async fn connect_retries_a_first_item_transport_error() {
2076        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
2077        let calls = Arc::new(AtomicUsize::new(0));
2078        let policy = RetryPolicy {
2079            max: 2,
2080            base: Duration::from_millis(5),
2081            cap: Duration::from_millis(5),
2082            ..RetryPolicy::default()
2083        };
2084        let c = calls.clone();
2085        let mut state = RetryState::new(policy.max);
2086        let out = connect_with_retry(
2087            &policy,
2088            &mut state,
2089            move || {
2090                let n = c.fetch_add(1, Ordering::SeqCst) + 1;
2091                Box::pin(async move {
2092                    if n == 1 {
2093                        Ok(first_item_err(transient()))
2094                    } else {
2095                        Ok(empty())
2096                    }
2097                })
2098            },
2099            &tx,
2100        )
2101        .await;
2102        assert!(matches!(out, PreContent::Ready { .. }));
2103        assert_eq!(calls.load(Ordering::SeqCst), 2, "first-item error retried");
2104    }
2105
2106    #[tokio::test(start_paused = true)]
2107    async fn connect_gives_up_after_max_and_skips_non_retryable() {
2108        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
2109        let policy = RetryPolicy {
2110            max: 1,
2111            base: Duration::from_millis(1),
2112            cap: Duration::from_millis(1),
2113            ..RetryPolicy::default()
2114        };
2115
2116        let calls = Arc::new(AtomicUsize::new(0));
2117        let c = calls.clone();
2118        let mut state = RetryState::new(policy.max);
2119        let out = connect_with_retry(
2120            &policy,
2121            &mut state,
2122            move || {
2123                c.fetch_add(1, Ordering::SeqCst);
2124                Box::pin(async { Err::<Chunks, CompletionError>(transient()) })
2125            },
2126            &tx,
2127        )
2128        .await;
2129        assert!(matches!(out, PreContent::Failed(_)));
2130        assert_eq!(calls.load(Ordering::SeqCst), 2, "one attempt + one retry");
2131
2132        let calls = Arc::new(AtomicUsize::new(0));
2133        let c = calls.clone();
2134        let mut state = RetryState::new(policy.max);
2135        let out = connect_with_retry(
2136            &policy,
2137            &mut state,
2138            move || {
2139                c.fetch_add(1, Ordering::SeqCst);
2140                Box::pin(async { Err::<Chunks, CompletionError>(fatal()) })
2141            },
2142            &tx,
2143        )
2144        .await;
2145        assert!(matches!(out, PreContent::Failed(_)));
2146        assert_eq!(calls.load(Ordering::SeqCst), 1, "non-retryable returns at once");
2147    }
2148
2149    #[test]
2150    fn fatal_class_flags_hard_errors_but_not_transient_transport() {
2151        // Transient-class failures stay feedable after their retries are
2152        // exhausted (the loop feeds the message back to the model); hard
2153        // client/build failures are fatal (the run stops).
2154        assert!(!fatal_class(&transient()));
2155        assert!(!fatal_class(&CompletionError::HttpError(
2156            rig::http_client::Error::StreamEnded
2157        )));
2158        assert!(fatal_class(&fatal()));
2159        assert!(fatal_class(&CompletionError::ResponseError("bad".into())));
2160        assert!(fatal_class(&CompletionError::RequestError("bad".into())));
2161        let json_err = serde_json::from_str::<u8>("not a number").unwrap_err();
2162        assert!(fatal_class(&CompletionError::JsonError(json_err)));
2163    }
2164
2165    /// Live-ish check against a refused port: a transport failure exhausts the
2166    /// retry budget (Retrying notices) and surfaces as a *non-fatal* stream
2167    /// error — which is exactly what lets the loop feed it back next turn
2168    /// instead of dying.
2169    #[tokio::test]
2170    async fn refused_port_exhausts_retries_to_non_fatal_error() {
2171        let opts = LlmOpts {
2172            model: "m1".to_string(),
2173            base_url: Some("http://127.0.0.1:9/v1".to_string()),
2174            api_key: Some("test-key".to_string()),
2175            retry: RetryPolicy {
2176                max: 2,
2177                base: Duration::from_millis(10),
2178                cap: Duration::from_millis(50),
2179                ..RetryPolicy::default()
2180            },
2181            ..LlmOpts::default()
2182        };
2183        let stream_fn = rig_stream_fn();
2184        let stream = stream_fn(&[], "sys", &[], &opts);
2185        let events: Vec<LlmStreamEvent> =
2186            tokio::time::timeout(std::time::Duration::from_secs(10), stream.collect())
2187                .await
2188                .expect("stream terminates");
2189        let retries = events
2190            .iter()
2191            .filter(|e| matches!(e, LlmStreamEvent::Retrying { .. }))
2192            .count();
2193        assert_eq!(retries, 2, "two backoff notices before giving up: {events:?}");
2194        assert!(
2195            matches!(events.last(), Some(LlmStreamEvent::Error { fatal: false, .. })),
2196            "refused port is transient-class, not fatal: {events:?}"
2197        );
2198    }
2199
2200    /// A `CompletionError` carrying a preserved response with a `Retry-After`
2201    /// header (no HTTP status needed: the header path reads
2202    /// `provider_response_headers()` regardless of status).
2203    fn err_with_retry_after(value: &str) -> CompletionError {
2204        let mut headers = HeaderMap::new();
2205        headers.insert("retry-after", HeaderValue::from_str(value).unwrap());
2206        CompletionError::ProviderResponse(
2207            rig::ProviderResponseError::without_status("429")
2208                .with_headers(Some(Box::new(headers))),
2209        )
2210    }
2211
2212    #[test]
2213    fn retry_after_parses_integer_seconds() {
2214        assert_eq!(
2215            retry_after(&err_with_retry_after("7")),
2216            Some(Duration::from_secs(7))
2217        );
2218    }
2219
2220    #[test]
2221    fn retry_after_is_none_when_absent() {
2222        // A transport failure with no preserved response/headers.
2223        assert_eq!(retry_after(&transient()), None);
2224    }
2225
2226    #[test]
2227    fn retry_after_is_none_on_non_numeric() {
2228        // The HTTP-date form and garbage both parse to `None` (dates are out of
2229        // scope — see `retry_after`).
2230        assert_eq!(
2231            retry_after(&err_with_retry_after("Wed, 21 Oct 2015 07:28:00 GMT")),
2232            None
2233        );
2234        assert_eq!(retry_after(&err_with_retry_after("soon")), None);
2235    }
2236
2237    #[tokio::test(start_paused = true)]
2238    async fn retry_wait_honors_retry_after_over_backoff() {
2239        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
2240        // base 10ms gives a ~10ms backoff for attempt 1; `Retry-After: 3` must
2241        // win, so the paused clock advances ~3s instead.
2242        let policy = RetryPolicy {
2243            base: Duration::from_millis(10),
2244            ..RetryPolicy::default()
2245        };
2246        let start = tokio::time::Instant::now();
2247        assert!(retry_wait(&tx, &policy, 1, &err_with_retry_after("3")).await);
2248        assert_eq!(start.elapsed(), Duration::from_secs(3));
2249    }
2250
2251    #[tokio::test(start_paused = true)]
2252    async fn retry_wait_clamps_retry_after_to_the_cap() {
2253        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
2254        // A hostile/buggy `Retry-After: 100` must not park the turn for 100s:
2255        // it is clamped to `policy.cap` (8s), not honored verbatim.
2256        let policy = RetryPolicy {
2257            cap: Duration::from_secs(8),
2258            ..RetryPolicy::default()
2259        };
2260        let start = tokio::time::Instant::now();
2261        assert!(retry_wait(&tx, &policy, 1, &err_with_retry_after("100")).await);
2262        assert_eq!(start.elapsed(), Duration::from_secs(8));
2263    }
2264    // ---- #23: stream-stall timeouts ----
2265
2266    type BoxStream =
2267        Pin<Box<dyn futures::Stream<Item = Result<StreamedAssistantContent, CompletionError>> + Send>>;
2268
2269    fn boxed<S>(s: S) -> BoxStream
2270    where
2271        S: futures::Stream<Item = Result<StreamedAssistantContent, CompletionError>>
2272            + Send
2273            + 'static,
2274    {
2275        Box::pin(s)
2276    }
2277
2278    fn silent_boxed() -> BoxStream {
2279        boxed(futures::stream::pending::<Result<StreamedAssistantContent, CompletionError>>())
2280    }
2281
2282    /// A fresh silent stream per call — `forward_stream`/`connect_with_retry`'s
2283    /// re-open seam, never exercised where content already flowed.
2284    fn open_silent()
2285    -> Pin<Box<dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send>> {
2286        Box::pin(async { Ok::<_, CompletionError>(silent_boxed()) })
2287    }
2288
2289    /// A slow (but successful) connect — used to prove `ttft = ZERO` awaits it
2290    /// directly instead of racing an immediate `sleep(ZERO)`.
2291    fn open_slow_ok()
2292    -> Pin<Box<dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send>> {
2293        Box::pin(async {
2294            tokio::time::sleep(Duration::from_millis(50)).await;
2295            Ok::<_, CompletionError>(boxed(futures::stream::iter(vec![Ok::<_, CompletionError>(
2296                StreamedAssistantContent::Text(rig::message::Text::new("ok")),
2297            )])))
2298        })
2299    }
2300
2301    /// A stall is transient-class, so it rides the retry / feed-back paths (B1).
2302    #[test]
2303    fn a_stall_error_is_transient_and_not_hard_fatal() {
2304        let d = Duration::from_secs(1);
2305        assert!(retryable(&stall_error(d)), "a stall must be transient (B1)");
2306        assert!(!fatal_class(&stall_error(d)), "a stall must not be hard-fatal");
2307        assert!(
2308            stall_error(d).to_string().contains("timed out"),
2309            "the (timed out) needle drives retryable()"
2310        );
2311    }
2312
2313    /// A silent stream stalls on the `ttft` peek every attempt: each retry emits
2314    /// a `Retrying` notice, and once the budget is spent the failure is a
2315    /// transient (non-fatal) `stall_error`.
2316    #[tokio::test(start_paused = true)]
2317    async fn a_silent_stream_trips_the_ttft_retries_then_fails() {
2318        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2319        let calls = Arc::new(AtomicUsize::new(0));
2320        let policy = RetryPolicy {
2321            max: 2,
2322            base: Duration::from_millis(10),
2323            cap: Duration::from_millis(10),
2324            ttft: Duration::from_secs(1),
2325            ..RetryPolicy::default()
2326        };
2327        let c = calls.clone();
2328        let open = move || {
2329            c.fetch_add(1, Ordering::SeqCst);
2330            open_silent()
2331        };
2332        let mut state = RetryState::new(policy.max);
2333        let out = connect_with_retry(&policy, &mut state, open, &tx).await;
2334
2335        let PreContent::Failed(error) = out else {
2336            panic!("expected Failed from an exhausted stall")
2337        };
2338        assert!(retryable(&error), "a stall is transient");
2339        assert!(!fatal_class(&error), "a stall is not hard-fatal");
2340        assert_eq!(calls.load(Ordering::SeqCst), 3, "one attempt + two retries");
2341        let retries = std::iter::from_fn(|| rx.try_recv().ok())
2342            .filter(|e| matches!(e, LlmStreamEvent::Retrying { .. }))
2343            .count();
2344        assert_eq!(retries, 2, "a Retrying notice per retry");
2345    }
2346
2347    /// A stream that emits a delta and then goes silent trips the idle deadline:
2348    /// because content already flowed, the failure is a NON-fatal error (the
2349    /// loop feeds it back) and is not retried (that would duplicate output).
2350    #[tokio::test(start_paused = true)]
2351    async fn deltas_then_a_stall_is_a_nonfatal_error() {
2352        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2353        let policy = RetryPolicy {
2354            max: 3,
2355            base: Duration::from_millis(10),
2356            cap: Duration::from_millis(10),
2357            idle: Duration::from_secs(1),
2358            ..RetryPolicy::default()
2359        };
2360        let stream = boxed(
2361            futures::stream::iter(vec![Ok::<_, CompletionError>(
2362                StreamedAssistantContent::Text(rig::message::Text::new("hi")),
2363            )])
2364            .chain(futures::stream::pending::<Result<StreamedAssistantContent, CompletionError>>()),
2365        );
2366        let state = RetryState::new(policy.max);
2367        forward_stream(stream, None, policy, open_silent, tx, state).await;
2368
2369        let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2370        assert!(
2371            events
2372                .iter()
2373                .any(|e| matches!(e, LlmStreamEvent::TextDelta(t) if t == "hi")),
2374            "{events:?}"
2375        );
2376        let error = events
2377            .iter()
2378            .find_map(|e| match e {
2379                LlmStreamEvent::Error { message, fatal } => Some((message.clone(), *fatal)),
2380                _ => None,
2381            })
2382            .expect("the idle deadline surfaced an error");
2383        assert!(!error.1, "a post-content stall is non-fatal: {error:?}");
2384        assert!(
2385            !events.iter().any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
2386            "no retry once content flowed: {events:?}"
2387        );
2388    }
2389
2390    /// A clean stream (deltas then `Done`) is forwarded untouched: no error, no
2391    /// retry, the idle deadline never fires.
2392    #[tokio::test(start_paused = true)]
2393    async fn a_clean_stream_is_unaffected() {
2394        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2395        let policy = RetryPolicy {
2396            idle: Duration::from_secs(1),
2397            ..RetryPolicy::default()
2398        };
2399        let stream = boxed(futures::stream::iter(vec![
2400            Ok::<_, CompletionError>(StreamedAssistantContent::Text(rig::message::Text::new("hi"))),
2401            Ok::<_, CompletionError>(done()),
2402        ]));
2403        let state = RetryState::new(policy.max);
2404        forward_stream(stream, None, policy, open_silent, tx, state).await;
2405
2406        let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2407        assert!(
2408            events
2409                .iter()
2410                .any(|e| matches!(e, LlmStreamEvent::TextDelta(t) if t == "hi")),
2411            "{events:?}"
2412        );
2413        assert!(
2414            events.iter().any(|e| matches!(e, LlmStreamEvent::Done { .. })),
2415            "{events:?}"
2416        );
2417        assert!(
2418            !events.iter().any(|e| matches!(e, LlmStreamEvent::Error { .. })),
2419            "{events:?}"
2420        );
2421        assert!(
2422            !events.iter().any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
2423            "{events:?}"
2424        );
2425    }
2426
2427    /// `ZERO` disables the deadline at BOTH seams: the peek collapses to a plain
2428    /// `stream.next()` (a silent stream stays pending), and the connect race is
2429    /// skipped so a slow connect is not spuriously timed out.
2430    #[tokio::test(start_paused = true)]
2431    async fn zero_deadlines_disable_the_timeout() {
2432        // Peek seam: no `sleep(ZERO)`, so a silent stream never yields `Idle`.
2433        let mut silent =
2434            futures::stream::pending::<Result<StreamedAssistantContent, CompletionError>>();
2435        let stalled = tokio::time::timeout(
2436            Duration::from_secs(1),
2437            next_or_stall(&mut silent, Duration::ZERO),
2438        )
2439        .await;
2440        assert!(stalled.is_err(), "ZERO must not trip the deadline");
2441
2442        // Connect seam: `ZERO` awaits the connect directly.
2443        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
2444        let policy = RetryPolicy {
2445            ttft: Duration::ZERO,
2446            ..RetryPolicy::default()
2447        };
2448        let mut state = RetryState::new(policy.max);
2449        let out = connect_with_retry(&policy, &mut state, open_slow_ok, &tx).await;
2450        assert!(
2451            matches!(out, PreContent::Ready { .. }),
2452            "a slow connect must not be timed out under ZERO"
2453        );
2454    }
2455
2456    /// The genuinely-stalled *connect* branch: `open()` never resolves, so the
2457    /// nonzero `ttft` race times out on every attempt. Each retry emits a
2458    /// `Retrying` notice; once the budget is spent the failure is a transient
2459    /// (non-fatal) `stall_error` — the same shape as a peek stall, but through
2460    /// the connect race's `Err(())` arm.
2461    #[tokio::test(start_paused = true)]
2462    async fn a_stalled_connect_trips_the_ttft_retries_then_fails() {
2463        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2464        let calls = Arc::new(AtomicUsize::new(0));
2465        let policy = RetryPolicy {
2466            max: 2,
2467            base: Duration::from_millis(10),
2468            cap: Duration::from_millis(10),
2469            ttft: Duration::from_secs(1),
2470            ..RetryPolicy::default()
2471        };
2472        let c = calls.clone();
2473        let open = move || {
2474            c.fetch_add(1, Ordering::SeqCst);
2475            Box::pin(std::future::pending::<Result<BoxStream, CompletionError>>())
2476                as Pin<
2477                    Box<
2478                        dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send,
2479                    >,
2480                >
2481        };
2482        let mut state = RetryState::new(policy.max);
2483        let out = connect_with_retry(&policy, &mut state, open, &tx).await;
2484
2485        let PreContent::Failed(error) = out else {
2486            panic!("expected Failed from an exhausted connect stall")
2487        };
2488        assert!(retryable(&error), "a connect stall is transient");
2489        assert!(!fatal_class(&error), "a connect stall is not hard-fatal");
2490        assert_eq!(calls.load(Ordering::SeqCst), 3, "one attempt + two retries");
2491        let retries = std::iter::from_fn(|| rx.try_recv().ok())
2492            .filter(|e| matches!(e, LlmStreamEvent::Retrying { .. }))
2493            .count();
2494        assert_eq!(retries, 2, "a Retrying notice per retry");
2495    }
2496
2497    /// The forward loop's *pre-content* idle branch: a stream that emits nothing
2498    /// and then stalls, with budget remaining, backs off (a `Retrying` notice)
2499    /// and re-opens — here the re-opened stream succeeds and the turn completes.
2500    #[tokio::test(start_paused = true)]
2501    async fn a_pre_content_idle_stall_retries_and_reopens() {
2502        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2503        let policy = RetryPolicy {
2504            max: 2,
2505            base: Duration::from_millis(10),
2506            cap: Duration::from_millis(10),
2507            idle: Duration::from_secs(1),
2508            ..RetryPolicy::default()
2509        };
2510        // The initial stream never yields → the idle deadline trips pre-content.
2511        let stream = silent_boxed();
2512        // The re-open yields a `Final`, so the retried turn completes.
2513        fn open_done()
2514        -> Pin<Box<dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send>> {
2515            Box::pin(async {
2516                Ok::<_, CompletionError>(boxed(futures::stream::iter(vec![Ok::<_, CompletionError>(
2517                    done(),
2518                )])))
2519            })
2520        }
2521        let state = RetryState::new(policy.max);
2522        forward_stream(stream, None, policy, open_done, tx, state).await;
2523
2524        let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2525        assert!(
2526            events
2527                .iter()
2528                .any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
2529            "a pre-content stall backs off: {events:?}"
2530        );
2531        assert!(
2532            events.iter().any(|e| matches!(e, LlmStreamEvent::Done { .. })),
2533            "the re-opened stream completed: {events:?}"
2534        );
2535        assert!(
2536            !events.iter().any(|e| matches!(e, LlmStreamEvent::Error { .. })),
2537            "no error once the retry succeeded: {events:?}"
2538        );
2539    }
2540
2541    /// A re-open that ALSO ends immediately, pre-content: paired with an empty
2542    /// initial stream, it drives the truncation arm's retries to exhaustion (each
2543    /// re-open yields no item, so nothing is ever forwarded).
2544    fn open_empty()
2545    -> Pin<Box<dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send>> {
2546        Box::pin(async { Ok::<_, CompletionError>(boxed(empty())) })
2547    }
2548
2549    /// A delta followed by an early EOF (no `Final`) is a truncation: content
2550    /// already flowed, so it is NOT retried — a single NON-fatal `Error` ends the
2551    /// turn (the loop feeds it back), and no `Done` is emitted.
2552    #[tokio::test(start_paused = true)]
2553    async fn a_delta_then_a_truncated_end_is_a_nonfatal_error() {
2554        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2555        let policy = RetryPolicy {
2556            max: 3,
2557            base: Duration::from_millis(10),
2558            cap: Duration::from_millis(10),
2559            idle: Duration::ZERO,
2560            ..RetryPolicy::default()
2561        };
2562        let stream = boxed(futures::stream::iter(vec![Ok::<_, CompletionError>(
2563            StreamedAssistantContent::Text(rig::message::Text::new("hi")),
2564        )]));
2565        let state = RetryState::new(policy.max);
2566        forward_stream(stream, None, policy, open_silent, tx, state).await;
2567
2568        let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2569        assert!(
2570            events
2571                .iter()
2572                .any(|e| matches!(e, LlmStreamEvent::TextDelta(t) if t == "hi")),
2573            "{events:?}"
2574        );
2575        let (message, fatal) = events
2576            .iter()
2577            .find_map(|e| match e {
2578                LlmStreamEvent::Error { message, fatal } => Some((message.clone(), *fatal)),
2579                _ => None,
2580            })
2581            .expect("a truncated end surfaced an error");
2582        assert!(!fatal, "a post-content truncation is non-fatal: {message}");
2583        assert!(
2584            !events
2585                .iter()
2586                .any(|e| matches!(e, LlmStreamEvent::Done { .. })),
2587            "no Done on a truncated stream: {events:?}"
2588        );
2589        assert!(
2590            !events
2591                .iter()
2592                .any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
2593            "no retry once content flowed: {events:?}"
2594        );
2595    }
2596
2597    /// A `Final`-terminated stream is unchanged: it forwards `Done` and never
2598    /// trips the truncation guard (which would surface a spurious error).
2599    #[tokio::test(start_paused = true)]
2600    async fn a_final_terminated_stream_is_unchanged() {
2601        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2602        let policy = RetryPolicy {
2603            idle: Duration::ZERO,
2604            ..RetryPolicy::default()
2605        };
2606        let stream = boxed(futures::stream::iter(vec![
2607            Ok::<_, CompletionError>(StreamedAssistantContent::Text(rig::message::Text::new(
2608                "hi",
2609            ))),
2610            Ok::<_, CompletionError>(done()),
2611        ]));
2612        let state = RetryState::new(policy.max);
2613        forward_stream(stream, None, policy, open_silent, tx, state).await;
2614
2615        let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2616        assert!(
2617            events
2618                .iter()
2619                .any(|e| matches!(e, LlmStreamEvent::Done { .. })),
2620            "{events:?}"
2621        );
2622        assert!(
2623            !events
2624                .iter()
2625                .any(|e| matches!(e, LlmStreamEvent::Error { .. })),
2626            "a clean turn surfaces no error: {events:?}"
2627        );
2628        assert!(
2629            !events
2630                .iter()
2631                .any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
2632            "{events:?}"
2633        );
2634    }
2635
2636    /// A pre-content no-item stream with budget remaining retries on the shared
2637    /// budget; each re-open also ends empty, so the whole budget is spent before
2638    /// a NON-fatal error ends the turn — one `Retrying` notice per retry, no
2639    /// `Done`.
2640    #[tokio::test(start_paused = true)]
2641    async fn a_pre_content_no_item_stream_retries_then_errors() {
2642        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2643        let policy = RetryPolicy {
2644            max: 2,
2645            base: Duration::from_millis(10),
2646            cap: Duration::from_millis(10),
2647            idle: Duration::ZERO,
2648            ..RetryPolicy::default()
2649        };
2650        let state = RetryState::new(policy.max);
2651        forward_stream(boxed(empty()), None, policy, open_empty, tx, state).await;
2652
2653        let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2654        let retries = events
2655            .iter()
2656            .filter(|e| matches!(e, LlmStreamEvent::Retrying { .. }))
2657            .count();
2658        assert_eq!(
2659            retries, policy.max as usize,
2660            "one Retrying notice per retry: {events:?}"
2661        );
2662        let (message, fatal) = events
2663            .iter()
2664            .find_map(|e| match e {
2665                LlmStreamEvent::Error { message, fatal } => Some((message.clone(), *fatal)),
2666                _ => None,
2667            })
2668            .expect("an exhausted truncation budget surfaced an error");
2669        assert!(!fatal, "an exhausted truncation is non-fatal: {message}");
2670        assert!(
2671            !events
2672                .iter()
2673                .any(|e| matches!(e, LlmStreamEvent::Done { .. })),
2674            "no Done on a truncated stream: {events:?}"
2675        );
2676    }
2677
2678    // --- probe_endpoint (opt-in local detection) ---------------------------
2679
2680    /// Stand up a one-shot HTTP server on `127.0.0.1:0` that answers `GET
2681    /// /models` with an OpenAI-shaped list after `delay`. Returns its base URL.
2682    fn spawn_models_server(delay: std::time::Duration) -> String {
2683        use std::io::{Read as _, Write as _};
2684        use std::net::TcpListener;
2685        let listener = TcpListener::bind("127.0.0.1:0").expect("bind ephemeral port");
2686        let addr = listener.local_addr().expect("local addr");
2687        std::thread::spawn(move || {
2688            // Serve connections until the process ends: the client may open a
2689            // fresh connection per request, and a timed-out probe leaves one
2690            // half-read behind.
2691            for stream in listener.incoming() {
2692                let Ok(mut stream) = stream else { continue };
2693                // Read the request head (through `\r\n\r\n`) so the client's
2694                // write completes before we answer — a single `read` can race
2695                // the request into multiple segments.
2696                let mut buf = Vec::new();
2697                let mut chunk = [0u8; 1024];
2698                loop {
2699                    match stream.read(&mut chunk) {
2700                        Ok(0) | Err(_) => break,
2701                        Ok(n) => {
2702                            buf.extend_from_slice(&chunk[..n]);
2703                            if buf.windows(4).any(|w| w == b"\r\n\r\n") {
2704                                break;
2705                            }
2706                        }
2707                    }
2708                }
2709                std::thread::sleep(delay);
2710                let body = "{\"object\":\"list\",\"data\":[{\"id\":\"m1\",\"object\":\"model\"}]}";
2711                let response = format!(
2712                    "HTTP/1.1 200 OK\r\n\
2713                     Content-Type: application/json\r\n\
2714                     Connection: close\r\n\
2715                     Content-Length: {}\r\n\r\n{}",
2716                    body.len(),
2717                    body
2718                );
2719                let _ = stream.write_all(response.as_bytes());
2720                let _ = stream.flush();
2721            }
2722        });
2723        format!("http://{addr}/v1")
2724    }
2725
2726    #[tokio::test]
2727    async fn a_refused_port_probes_false() {
2728        let up = probe_endpoint(
2729            "http://127.0.0.1:9/v1",
2730            std::time::Duration::from_millis(500),
2731        )
2732        .await;
2733        assert!(!up, "a refused port is not an endpoint");
2734    }
2735
2736    #[tokio::test]
2737    async fn a_live_models_endpoint_probes_true() {
2738        // rig's HTTP client pays a one-time, process-wide init on its very first
2739        // request (seconds on macOS, where system-proxy detection is slow) — it
2740        // happens BEFORE any connection, so a refused port triggers it cheaply.
2741        // Pre-warm so the assertion below is not racing that one-time cost.
2742        let _ = probe_endpoint("http://127.0.0.1:9/v1", std::time::Duration::from_secs(60)).await;
2743        let base = spawn_models_server(std::time::Duration::ZERO);
2744        let up = probe_endpoint(&base, std::time::Duration::from_secs(5)).await;
2745        assert!(up, "a live /models endpoint probes true");
2746    }
2747
2748    #[tokio::test]
2749    async fn a_slow_endpoint_times_out_to_false() {
2750        let base = spawn_models_server(std::time::Duration::from_millis(500));
2751        let up = probe_endpoint(&base, std::time::Duration::from_millis(50)).await;
2752        assert!(!up, "a slow endpoint is not treated as up");
2753    }
2754}