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