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.budget -= 1;
994    if !retry_wait(tx, policy, state.attempts, error).await {
995        return PreContent::Cancelled;
996    }
997    connect_with_retry(policy, state, open, tx).await
998}
999
1000/// The server's `Retry-After` header (integer delta-seconds), when the failed
1001/// response preserved one. `None` when there is no preserved response, no
1002/// header, a non-UTF-8 value, or a value that is not a non-negative integer.
1003///
1004/// SCOPE: integer delta-seconds ONLY. The HTTP-date form of `Retry-After`
1005/// (RFC 9110 `Retry-After = HTTP-date / delay-seconds`) is deliberately out of
1006/// scope: a date would need wall-clock parsing against the loop's monotonic
1007/// `tokio::time` deadline, and a subtly wrong date is worse than ignoring it —
1008/// a provider that sends one simply gets the normal backoff.
1009fn retry_after(error: &CompletionError) -> Option<Duration> {
1010    // `provider_response_headers()` already unwraps both the
1011    // `HttpError(InvalidStatusCodeWithDetails)` and the `ProviderResponse`
1012    // variants, so this reads either.
1013    let raw = error.provider_response_headers()?.get("retry-after")?;
1014    // A non-UTF-8 header value is ignored, not an error.
1015    let seconds: u64 = raw.to_str().ok()?.trim().parse().ok()?;
1016    Some(Duration::from_secs(seconds))
1017}
1018
1019/// Emit a retry notice and sleep the backoff. Returns false if the consumer
1020/// dropped `tx` (cancellation) during the wait.
1021async fn retry_wait(
1022    tx: &tokio::sync::mpsc::UnboundedSender<LlmStreamEvent>,
1023    policy: &RetryPolicy,
1024    attempt: u32,
1025    error: &CompletionError,
1026) -> bool {
1027    // A server-supplied `Retry-After` is authoritative and wins over the local
1028    // backoff, but it is clamped by the existing `policy.cap` so a hostile or
1029    // buggy `Retry-After: 86400` cannot park the turn for a day.
1030    let delay = retry_after(error)
1031        .map(|d| d.min(policy.cap))
1032        .unwrap_or_else(|| backoff(attempt, policy.base, policy.cap));
1033    let _ = tx.send(LlmStreamEvent::Retrying {
1034        attempt,
1035        max: policy.max,
1036        reason: error.to_string(),
1037    });
1038    tokio::select! {
1039        biased;
1040        _ = tx.closed() => false,
1041        _ = tokio::time::sleep(delay) => true,
1042    }
1043}
1044
1045fn build_request(
1046    messages: &[AgentMessage],
1047    system: &str,
1048    tools: &[ToolDefinition],
1049    opts: &LlmOpts,
1050) -> CompletionRequest {
1051    let mut chat_history = Vec::with_capacity(messages.len() + 1);
1052    if !system.is_empty() {
1053        chat_history.push(Message::system(system));
1054    }
1055    chat_history.extend(messages.iter().filter_map(to_rig_message));
1056    CompletionRequest {
1057        model: None,
1058        // System instructions ride as the leading Message::System, the
1059        // canonical form rig's own request builder funnels `preamble` into.
1060        preamble: None,
1061        chat_history,
1062        documents: Vec::new(),
1063        tools: tools.to_vec(),
1064        temperature: opts.temperature,
1065        max_tokens: None,
1066        tool_choice: None,
1067        additional_params: reasoning_params(opts),
1068        output_schema: None,
1069        record_telemetry_content: false,
1070    }
1071}
1072
1073/// Effort levels accepted by the Responses wire. rig deserializes
1074/// `reasoning.effort` into a snake_case enum client-side (`ReasoningEffort`),
1075/// so any other string fails with an opaque RequestError at send time;
1076/// `invalid_effort` surfaces a clear error instead. Chat forwards any string.
1077const RESPONSES_EFFORT_LEVELS: &[&str] =
1078    &["none", "minimal", "low", "medium", "high", "xhigh", "max"];
1079
1080/// `Some(message)` when `opts.effort` cannot cross the configured wire: the
1081/// Responses endpoint validates effort client-side against the enum above,
1082/// while Chat forwards any string to the provider.
1083fn invalid_effort(opts: &LlmOpts) -> Option<String> {
1084    if opts.endpoint == LlmEndpoint::Responses
1085        && let Some(effort) = opts.effort.as_deref()
1086        && !RESPONSES_EFFORT_LEVELS.contains(&effort)
1087    {
1088        return Some(format!(
1089            "effort `{effort}` is not valid for the Responses endpoint; use one of: {}",
1090            RESPONSES_EFFORT_LEVELS.join(", ")
1091        ));
1092    }
1093    None
1094}
1095
1096/// Reasoning/effort passthrough, fanned out per wire shape.
1097///
1098/// Chat: `opts.effort` alone → `{"reasoning_effort": effort}`; None sends
1099/// nothing (unchanged), so backends that reject unknown fields keep working.
1100///
1101/// Responses: ALWAYS emits `{"reasoning": {"summary": "auto"}}`, merging
1102/// `"effort"` in when `opts.effort` is `Some`. The summary is what makes
1103/// reasoning DISPLAYABLE: rig auto-injects `include:["reasoning.encrypted_content"]`
1104/// whenever `additional_params` carries `reasoning` (rig-core-0.42.0
1105/// responses_api/mod.rs:1461 — `if additional_parameters.reasoning.is_some()`),
1106/// and `ReasoningSummaryLevel::Auto` serializes to snake_case "auto"
1107/// (responses_api/mod.rs:2133). Without a summary the model emits an
1108/// encrypted-only reasoning item whose displayable text is empty (see the
1109/// `map_item` guard below).
1110///
1111/// `invalid_effort` still validates the level client-side before this runs.
1112fn reasoning_params(opts: &LlmOpts) -> Option<serde_json::Value> {
1113    match opts.endpoint {
1114        LlmEndpoint::Chat => opts
1115            .effort
1116            .as_ref()
1117            .map(|effort| serde_json::json!({ "reasoning_effort": effort })),
1118        LlmEndpoint::Responses => {
1119            let mut reasoning = serde_json::Map::new();
1120            reasoning.insert("summary".into(), serde_json::json!("auto"));
1121            if let Some(effort) = &opts.effort {
1122                reasoning.insert("effort".into(), serde_json::json!(effort));
1123            }
1124            Some(serde_json::json!({ "reasoning": reasoning }))
1125        }
1126    }
1127}
1128
1129/// Prefix applied to a failed tool's output *on the wire only*. The chat wire
1130/// has no structured error flag for tool results (rig's `ToolResult` only
1131/// carries text/content), so a failed tool's output must be annotated in the
1132/// text the model sees — otherwise a tool failure looks exactly like a
1133/// success. The UI and the JSONL session keep the raw output; only the wire
1134/// copy read by the LLM is marked.
1135const TOOL_OUTPUT_ERROR_PREFIX: &str = "ERROR: ";
1136
1137fn to_rig_message(m: &AgentMessage) -> Option<Message> {
1138    match m {
1139        AgentMessage::User { content } => {
1140            let content: Vec<UserContent> = content
1141                .iter()
1142                .filter_map(|b| match b {
1143                    ContentBlock::Text { text } => Some(UserContent::text(text.clone())),
1144                    _ => None,
1145                })
1146                .collect();
1147            (!content.is_empty()).then_some(Message::User { content })
1148        }
1149        AgentMessage::Assistant { content, .. } => {
1150            let content: Vec<AssistantContent> = content
1151                .iter()
1152                .map(|b| match b {
1153                    ContentBlock::Text { text } => AssistantContent::text(text.clone()),
1154                    ContentBlock::Thinking { text } => {
1155                        AssistantContent::Reasoning(Reasoning::new(text))
1156                    }
1157                    ContentBlock::ToolCall {
1158                        id,
1159                        name,
1160                        arguments,
1161                    } => AssistantContent::ToolCall(ToolCall::new(
1162                        ToolCallId::new_or_mint(id.clone()),
1163                        ToolFunction::new(name.clone(), arguments.clone()),
1164                    )),
1165                })
1166                .collect();
1167            (!content.is_empty()).then_some(Message::Assistant { id: None, content })
1168        }
1169        AgentMessage::ToolResult {
1170            tool_call_id,
1171            name,
1172            output,
1173            is_error,
1174        } => {
1175            // The wire has no structured error flag for tool results, so mark
1176            // failures in the text the model sees. The UI and the session keep
1177            // the raw output — only this wire copy is prefixed.
1178            let output = if *is_error {
1179                format!("{TOOL_OUTPUT_ERROR_PREFIX}{output}")
1180            } else {
1181                output.clone()
1182            };
1183            Some(Message::tool_result(
1184                tool_call_id.clone(),
1185                name.clone(),
1186                output,
1187            ))
1188        }
1189    }
1190}
1191
1192fn map_item(item: StreamedAssistantContent) -> Vec<LlmStreamEvent> {
1193    match item {
1194        StreamedAssistantContent::Text(text) => vec![LlmStreamEvent::TextDelta(text.text)],
1195        StreamedAssistantContent::ReasoningDelta { reasoning, .. } => {
1196            vec![LlmStreamEvent::ThinkingDelta(reasoning)]
1197        }
1198        // A complete reasoning block supersedes the streamed deltas: the Responses
1199        // wire restates the whole reasoning item after the deltas. Dead on
1200        // chat-completions (deltas only) but LIVE on Responses — it must not be
1201        // dropped.
1202        StreamedAssistantContent::Reasoning { reasoning, .. } => {
1203            // Replacement semantics: the complete block supersedes the
1204            // accumulated deltas (the loop drops prior thinking on this event).
1205            let text = reasoning_text(&reasoning);
1206            if text.is_empty() {
1207                // Encrypted-only terminal item: no displayable text. Emitting
1208                // ThinkingReplace("") would reach `loop_.rs:replace_thinking`,
1209                // which retains-out ALL accumulated thinking and pushes an empty
1210                // block — erasing the streamed ThinkingDeltas. Yield nothing.
1211                vec![]
1212            } else {
1213                vec![LlmStreamEvent::ThinkingReplace(text)]
1214            }
1215        }
1216        // Deltas are dropped: rig's streaming fold accumulates every tool-call
1217        // fragment and emits the complete `ToolCall` on `ToolInputEnd` (both
1218        // chat-completions and responses wires do this), so the reassembled
1219        // call always follows. A provider that ends input without closing it
1220        // loses the call inside rig itself; the loop surfaces that as an error
1221        // rather than silently treating the turn as done.
1222        StreamedAssistantContent::ToolCallDelta { .. } => vec![],
1223        StreamedAssistantContent::ToolCall { tool_call, .. } => vec![LlmStreamEvent::ToolCall {
1224            id: tool_call.id.as_str().to_string(),
1225            name: tool_call.function.name,
1226            arguments: tool_call.function.arguments,
1227        }],
1228        StreamedAssistantContent::Final(final_record) => {
1229            vec![LlmStreamEvent::Done {
1230                stop_reason: final_record
1231                    .finish_reason
1232                    .map(finish_reason)
1233                    .unwrap_or(StopReason::Stop),
1234                usage: Some(convert_usage(final_record.usage)),
1235            }]
1236        }
1237        StreamedAssistantContent::Unknown(_) => vec![],
1238    }
1239}
1240
1241fn finish_reason(reason: FinishReason) -> StopReason {
1242    match reason {
1243        FinishReason::Stop => StopReason::Stop,
1244        FinishReason::Length => StopReason::Length,
1245        FinishReason::ToolCalls => StopReason::ToolUse,
1246        FinishReason::ContentFilter | FinishReason::Other(_) => StopReason::Stop,
1247    }
1248}
1249
1250fn convert_usage(usage: rig::completion::Usage) -> crate::message::Usage {
1251    crate::message::Usage {
1252        input_tokens: usage.input_tokens,
1253        output_tokens: usage.output_tokens,
1254        // Zero is rig's documented sentinel for "not reported".
1255        cache_read_tokens: (usage.cached_input_tokens > 0).then_some(usage.cached_input_tokens),
1256        cache_write_tokens: (usage.cache_creation_input_tokens > 0)
1257            .then_some(usage.cache_creation_input_tokens),
1258    }
1259}
1260
1261fn reasoning_text(reasoning: &Reasoning) -> String {
1262    reasoning
1263        .content
1264        .iter()
1265        .filter_map(|block| match block {
1266            ReasoningContent::Text { text, .. } => Some(text.as_str()),
1267            ReasoningContent::Summary(summary) => Some(summary.as_str()),
1268            ReasoningContent::Redacted { data } => Some(data.as_str()),
1269            ReasoningContent::Encrypted(_) => None,
1270        })
1271        .collect::<Vec<_>>()
1272        .join("\n")
1273}
1274
1275#[cfg(test)]
1276mod tests {
1277    use super::*;
1278    use rig::completion::Usage as RigUsage;
1279    use rig::message::{ToolResultContent, UserContent};
1280    use rig::streaming::{RawStreamingToolCall, StreamFinal, ToolCallDeltaContent};
1281    use serde_json::json;
1282
1283    fn tool_call_item() -> StreamedAssistantContent {
1284        let raw =
1285            RawStreamingToolCall::new("call_1".to_string(), "run".to_string(), json!({"x": 1}));
1286        StreamedAssistantContent::ToolCall {
1287            tool_call: raw.into(),
1288            internal_call_id: "ic1".to_string(),
1289        }
1290    }
1291
1292    fn final_item(finish_reason: Option<FinishReason>) -> StreamedAssistantContent {
1293        let usage = RigUsage {
1294            input_tokens: 10,
1295            output_tokens: 5,
1296            total_tokens: 15,
1297            cached_input_tokens: 3,
1298            cache_creation_input_tokens: 2,
1299            ..RigUsage::default()
1300        };
1301        let final_record = StreamFinal::new("test", usage);
1302        StreamedAssistantContent::Final(match finish_reason {
1303            Some(reason) => final_record.with_finish_reason(reason),
1304            None => final_record,
1305        })
1306    }
1307
1308    #[test]
1309    fn text_item_maps_to_text_delta() {
1310        let events = map_item(StreamedAssistantContent::Text(rig::message::Text::new(
1311            "hi",
1312        )));
1313        assert_eq!(events, vec![LlmStreamEvent::TextDelta("hi".to_string())]);
1314    }
1315
1316    #[test]
1317    fn reasoning_delta_maps_to_thinking_delta() {
1318        let events = map_item(StreamedAssistantContent::ReasoningDelta {
1319            id: "r1".to_string(),
1320            provider_id: None,
1321            reasoning: "hmm".to_string(),
1322        });
1323        assert_eq!(
1324            events,
1325            vec![LlmStreamEvent::ThinkingDelta("hmm".to_string())]
1326        );
1327    }
1328
1329    #[test]
1330    fn reasoning_block_replaces_with_full_text() {
1331        let events = map_item(StreamedAssistantContent::Reasoning {
1332            reasoning: Reasoning::new("full thought"),
1333            id: "r1".to_string(),
1334        });
1335        assert_eq!(
1336            events,
1337            vec![LlmStreamEvent::ThinkingReplace("full thought".to_string())]
1338        );
1339    }
1340
1341    #[test]
1342    fn reasoning_text_joins_summary_text_and_redacted_drops_encrypted() {
1343        // Exactly rig's `Reasoning::display_text()` contract: Summary | Text |
1344        // Redacted join with "\n"; Encrypted is skipped.
1345        let reasoning = Reasoning {
1346            id: None,
1347            content: vec![
1348                ReasoningContent::Summary("a".to_string()),
1349                ReasoningContent::Text {
1350                    text: "b".to_string(),
1351                    signature: None,
1352                },
1353                ReasoningContent::Redacted {
1354                    data: "c".to_string(),
1355                },
1356                ReasoningContent::Encrypted("opaque".to_string()),
1357            ],
1358        };
1359        assert_eq!(reasoning_text(&reasoning), "a\nb\nc");
1360    }
1361
1362    #[test]
1363    fn reasoning_text_summary_only_is_the_summary() {
1364        let reasoning = Reasoning::summaries(vec!["just a summary".to_string()]);
1365        assert_eq!(reasoning_text(&reasoning), "just a summary");
1366    }
1367
1368    #[test]
1369    fn reasoning_text_encrypted_only_is_empty() {
1370        let reasoning = Reasoning::encrypted("opaque");
1371        assert_eq!(reasoning_text(&reasoning), "");
1372    }
1373
1374    #[test]
1375    fn summary_reasoning_block_replaces_with_full_text() {
1376        let events = map_item(StreamedAssistantContent::Reasoning {
1377            reasoning: Reasoning::summaries(vec!["a summary".to_string()]),
1378            id: "rs_1".to_string(),
1379        });
1380        assert_eq!(
1381            events,
1382            vec![LlmStreamEvent::ThinkingReplace("a summary".to_string())]
1383        );
1384    }
1385
1386    #[test]
1387    fn encrypted_only_reasoning_block_emits_no_event() {
1388        // An encrypted-only terminal item has no displayable text; emitting
1389        // ThinkingReplace("") would erase the streamed ThinkingDeltas downstream
1390        // (see the `map_item` guard).
1391        let events = map_item(StreamedAssistantContent::Reasoning {
1392            reasoning: Reasoning::encrypted("opaque"),
1393            id: "rs_1".to_string(),
1394        });
1395        assert!(events.is_empty(), "got: {events:?}");
1396    }
1397
1398    #[test]
1399    fn tool_call_delta_ignored() {
1400        for content in [
1401            ToolCallDeltaContent::Name("run".to_string()),
1402            ToolCallDeltaContent::Delta("{\"x".to_string()),
1403        ] {
1404            let events = map_item(StreamedAssistantContent::ToolCallDelta {
1405                internal_call_id: "ic1".to_string(),
1406                content,
1407            });
1408            assert!(events.is_empty());
1409        }
1410    }
1411
1412    #[test]
1413    fn tool_call_maps_to_call_event() {
1414        let events = map_item(tool_call_item());
1415        assert_eq!(
1416            events,
1417            vec![LlmStreamEvent::ToolCall {
1418                id: "call_1".to_string(),
1419                name: "run".to_string(),
1420                arguments: json!({"x": 1}),
1421            }]
1422        );
1423    }
1424
1425    #[test]
1426    fn final_maps_finish_reason_and_usage() {
1427        let cases = [
1428            (FinishReason::Stop, StopReason::Stop),
1429            (FinishReason::Length, StopReason::Length),
1430            (FinishReason::ToolCalls, StopReason::ToolUse),
1431            (FinishReason::ContentFilter, StopReason::Stop),
1432            (FinishReason::Other("weird".to_string()), StopReason::Stop),
1433        ];
1434        for (reason, expected) in cases {
1435            let events = map_item(final_item(Some(reason)));
1436            assert_eq!(
1437                events,
1438                vec![LlmStreamEvent::Done {
1439                    stop_reason: expected,
1440                    usage: Some(crate::message::Usage {
1441                        input_tokens: 10,
1442                        output_tokens: 5,
1443                        cache_read_tokens: Some(3),
1444                        cache_write_tokens: Some(2),
1445                    }),
1446                }],
1447                "wrong mapping for {expected:?}"
1448            );
1449        }
1450    }
1451
1452    #[test]
1453    fn final_without_finish_reason_defaults_to_stop() {
1454        let events = map_item(final_item(None));
1455        match &events[..] {
1456            [LlmStreamEvent::Done { stop_reason, .. }] => {
1457                assert_eq!(*stop_reason, StopReason::Stop)
1458            }
1459            other => panic!("unexpected events: {other:?}"),
1460        }
1461    }
1462
1463    #[test]
1464    fn unknown_item_ignored() {
1465        let events = map_item(StreamedAssistantContent::Unknown(
1466            rig::streaming::UnknownPayload::new(json!({"mystery": 1})),
1467        ));
1468        assert!(events.is_empty());
1469    }
1470
1471    #[test]
1472    fn build_request_round_trips_context() {
1473        let opts = LlmOpts {
1474            model: "m1".to_string(),
1475            base_url: None,
1476            api_key: None,
1477            temperature: Some(0.5),
1478            ..LlmOpts::default()
1479        };
1480        let tools = vec![ToolDefinition {
1481            name: "get_weather".to_string(),
1482            description: "Weather lookup".to_string(),
1483            parameters: json!({"type": "object"}),
1484        }];
1485        let messages = vec![
1486            AgentMessage::user_text("What is the weather?"),
1487            AgentMessage::Assistant {
1488                content: vec![
1489                    ContentBlock::Thinking { text: "hmm".into() },
1490                    ContentBlock::ToolCall {
1491                        id: "t1".into(),
1492                        name: "get_weather".into(),
1493                        arguments: json!({"city": "Tokyo"}),
1494                    },
1495                    ContentBlock::Text {
1496                        text: "Let me check.".into(),
1497                    },
1498                ],
1499                stop_reason: StopReason::ToolUse,
1500                usage: None,
1501                model: None,
1502            },
1503            AgentMessage::ToolResult {
1504                tool_call_id: "t1".into(),
1505                name: "get_weather".into(),
1506                output: "sunny".into(),
1507                is_error: false,
1508            },
1509        ];
1510
1511        let request = build_request(&messages, "You are terse.", &tools, &opts);
1512
1513        assert!(request.preamble.is_none());
1514        assert_eq!(request.temperature, Some(0.5));
1515        assert_eq!(request.tools.len(), 1);
1516        assert_eq!(request.tools[0].name, "get_weather");
1517
1518        let [
1519            Message::System { content: system },
1520            Message::User { content: user },
1521            Message::Assistant {
1522                content: assistant, ..
1523            },
1524            Message::User {
1525                content: tool_result,
1526            },
1527        ] = &request.chat_history[..]
1528        else {
1529            panic!("unexpected chat_history shape: {:?}", request.chat_history);
1530        };
1531        assert_eq!(system, "You are terse.");
1532        assert!(
1533            matches!(user.as_slice(), [UserContent::Text(t)] if t.text == "What is the weather?")
1534        );
1535        assert_eq!(assistant.len(), 3);
1536        assert!(
1537            matches!(&assistant[0], AssistantContent::Reasoning(r) if reasoning_text(r) == "hmm")
1538        );
1539        assert!(matches!(
1540            &assistant[1],
1541            AssistantContent::ToolCall(call)
1542                if call.id.as_str() == "t1"
1543                    && call.function.name == "get_weather"
1544                    && call.function.arguments == json!({"city": "Tokyo"})
1545        ));
1546        assert!(matches!(&assistant[2], AssistantContent::Text(t) if t.text == "Let me check."));
1547        assert!(matches!(
1548            tool_result.as_slice(),
1549            [UserContent::ToolResult(r)] if r.call.as_str() == "t1" && r.name == "get_weather"
1550        ));
1551    }
1552
1553    #[test]
1554    fn tool_result_error_is_annotated_on_the_wire_only() {
1555        let err = AgentMessage::ToolResult {
1556            tool_call_id: "t1".into(),
1557            name: "echo".into(),
1558            output: "boom".into(),
1559            is_error: true,
1560        };
1561        let ok = AgentMessage::ToolResult {
1562            tool_call_id: "t1".into(),
1563            name: "echo".into(),
1564            output: "boom".into(),
1565            is_error: false,
1566        };
1567
1568        fn wire_text(m: &AgentMessage) -> String {
1569            match to_rig_message(m) {
1570                Some(Message::User { content }) => match &content[..] {
1571                    [UserContent::ToolResult(r)] => match &r.content[..] {
1572                        [ToolResultContent::Text(t)] => t.text.clone(),
1573                        other => panic!("unexpected tool result content: {other:?}"),
1574                    },
1575                    other => panic!("unexpected user content: {other:?}"),
1576                },
1577                other => panic!("unexpected wire message: {other:?}"),
1578            }
1579        }
1580
1581        // Success passes through byte-for-byte; a failed tool is marked.
1582        assert_eq!(wire_text(&ok), "boom");
1583        assert_eq!(wire_text(&err), "ERROR: boom");
1584        // The stored message itself is untouched — the flag is only consumed
1585        // at the wire seam.
1586        match err {
1587            AgentMessage::ToolResult {
1588                output, is_error, ..
1589            } => {
1590                assert_eq!(output, "boom");
1591                assert!(is_error);
1592            }
1593            _ => unreachable!(),
1594        }
1595    }
1596
1597    #[test]
1598    fn endpoint_defaults_to_chat() {
1599        let opts = LlmOpts {
1600            model: "m1".to_string(),
1601            base_url: None,
1602            api_key: None,
1603            temperature: None,
1604            endpoint: LlmEndpoint::Chat,
1605            effort: None,
1606            session_id: None,
1607            retry: RetryPolicy::default(),
1608            model_profiles: ModelProfiles::default(),
1609        };
1610        assert_eq!(opts.endpoint, LlmEndpoint::Chat);
1611        assert_eq!(LlmOpts::default().endpoint, LlmEndpoint::Chat);
1612    }
1613
1614    #[test]
1615    fn no_effort_sends_no_additional_params() {
1616        let opts = LlmOpts::default();
1617        let request = build_request(&[], "sys", &[], &opts);
1618        assert_eq!(request.additional_params, None);
1619    }
1620
1621    #[test]
1622    fn chat_effort_maps_to_reasoning_effort_param() {
1623        let opts = LlmOpts {
1624            effort: Some("high".into()),
1625            endpoint: LlmEndpoint::Chat,
1626            ..LlmOpts::default()
1627        };
1628        let request = build_request(&[], "sys", &[], &opts);
1629        assert_eq!(
1630            request.additional_params,
1631            Some(json!({ "reasoning_effort": "high" }))
1632        );
1633    }
1634
1635    #[test]
1636    fn responses_effort_maps_to_reasoning_object() {
1637        let opts = LlmOpts {
1638            effort: Some("xhigh".into()),
1639            endpoint: LlmEndpoint::Responses,
1640            ..LlmOpts::default()
1641        };
1642        let request = build_request(&[], "sys", &[], &opts);
1643        assert_eq!(
1644            request.additional_params,
1645            Some(json!({ "reasoning": { "effort": "xhigh", "summary": "auto" } }))
1646        );
1647    }
1648
1649    #[test]
1650    fn responses_always_requests_a_reasoning_summary() {
1651        // The case that was invisible before this change: with NO effort the
1652        // Responses wire must still ask for a reasoning summary — otherwise the
1653        // model emits an encrypted-only item whose displayable text is empty.
1654        let opts = LlmOpts {
1655            endpoint: LlmEndpoint::Responses,
1656            effort: None,
1657            ..LlmOpts::default()
1658        };
1659        let request = build_request(&[], "sys", &[], &opts);
1660        assert_eq!(
1661            request.additional_params,
1662            Some(json!({ "reasoning": { "summary": "auto" } }))
1663        );
1664    }
1665
1666    #[test]
1667    fn chat_effort_is_verbatim_passthrough() {
1668        // Free-style on the Chat wire: any string (including provider
1669        // dialects) is forwarded verbatim under `reasoning_effort`.
1670        for level in ["low", "medium", "my-custom-level", ""] {
1671            let opts = LlmOpts {
1672                effort: Some(level.into()),
1673                endpoint: LlmEndpoint::Chat,
1674                ..LlmOpts::default()
1675            };
1676            let request = build_request(&[], "sys", &[], &opts);
1677            assert_eq!(
1678                request.additional_params,
1679                Some(json!({ "reasoning_effort": level })),
1680                "level {level:?} must pass through verbatim on Chat"
1681            );
1682        }
1683    }
1684
1685    #[test]
1686    fn responses_effort_validation() {
1687        // Known snake_case levels pass; anything else is rejected with a clear
1688        // message (rig would otherwise fail the request with an opaque error).
1689        for level in ["none", "minimal", "low", "medium", "high", "xhigh", "max"] {
1690            let opts = LlmOpts {
1691                effort: Some(level.into()),
1692                endpoint: LlmEndpoint::Responses,
1693                ..LlmOpts::default()
1694            };
1695            assert_eq!(invalid_effort(&opts), None, "level {level} must be valid");
1696        }
1697        let opts = LlmOpts {
1698            effort: Some("deep".into()),
1699            endpoint: LlmEndpoint::Responses,
1700            ..LlmOpts::default()
1701        };
1702        let msg = invalid_effort(&opts).expect("unknown level must be rejected");
1703        assert!(msg.contains("not valid"), "got: {msg}");
1704        assert!(
1705            msg.contains("xhigh"),
1706            "message must list valid levels: {msg}"
1707        );
1708        // Chat stays verbatim even for a value Responses would reject.
1709        let opts = LlmOpts {
1710            effort: Some("deep".into()),
1711            endpoint: LlmEndpoint::Chat,
1712            ..LlmOpts::default()
1713        };
1714        assert_eq!(invalid_effort(&opts), None);
1715    }
1716
1717    #[test]
1718    fn client_sends_session_header_and_user_agent() {
1719        let opts = LlmOpts {
1720            api_key: Some("test-key".to_string()),
1721            session_id: Some("sess-123".to_string()),
1722            ..LlmOpts::default()
1723        };
1724        let headers = openai_client(&opts)
1725            .expect("client builds")
1726            .headers()
1727            .clone();
1728        assert_eq!(
1729            headers.get("x-opencode-session").expect("session header"),
1730            "sess-123"
1731        );
1732        assert!(headers.contains_key("user-agent"), "got: {headers:?}");
1733    }
1734
1735    #[test]
1736    fn client_omits_session_header_when_unset() {
1737        let opts = LlmOpts {
1738            api_key: Some("test-key".to_string()),
1739            ..LlmOpts::default()
1740        };
1741        let headers = openai_client(&opts)
1742            .expect("client builds")
1743            .headers()
1744            .clone();
1745        assert!(!headers.contains_key("x-opencode-session"));
1746        assert!(headers.contains_key("user-agent"), "got: {headers:?}");
1747    }
1748
1749    #[test]
1750    fn localhost_base_url_builds_without_key() {
1751        // Keyless local endpoints (Ollama, llama.cpp, ...) should work with no
1752        // configured key via a placeholder bearer token.
1753        for url in [
1754            "http://localhost:11434/v1",
1755            "http://127.0.0.1:8080/v1",
1756            "http://[::1]:8080/v1",
1757        ] {
1758            let opts = LlmOpts {
1759                base_url: Some(url.to_string()),
1760                api_key: None,
1761                ..LlmOpts::default()
1762            };
1763            assert!(
1764                openai_client(&opts).is_ok(),
1765                "keyless {url} must build a client"
1766            );
1767        }
1768    }
1769
1770    #[test]
1771    fn remote_base_url_without_key_is_error() {
1772        let opts = LlmOpts {
1773            base_url: Some("https://api.example.com/v1".to_string()),
1774            api_key: None,
1775            ..LlmOpts::default()
1776        };
1777        let err = openai_client(&opts).unwrap_err();
1778        assert!(err.contains("no API key"), "got: {err}");
1779    }
1780
1781    /// Unreachable endpoint: list_models surfaces a string error, no panic.
1782    #[tokio::test]
1783    async fn list_models_unreachable_is_error() {
1784        let opts = LlmOpts {
1785            base_url: Some("http://127.0.0.1:9/v1".to_string()),
1786            api_key: Some("test-key".to_string()),
1787            ..LlmOpts::default()
1788        };
1789        let err = tokio::time::timeout(std::time::Duration::from_secs(10), list_models(&opts))
1790            .await
1791            .expect("list terminates")
1792            .expect_err("unreachable host must error");
1793        assert!(!err.is_empty());
1794    }
1795
1796    #[tokio::test]
1797    async fn list_models_without_key_is_error() {
1798        let opts = LlmOpts::default();
1799        let err = list_models(&opts).await.expect_err("no key must error");
1800        assert!(err.contains("no API key"), "got: {err}");
1801    }
1802
1803    #[test]
1804    fn client_cache_reuses_until_key_changes() {
1805        let opts = LlmOpts {
1806            base_url: Some("http://127.0.0.1:9/v1".to_string()),
1807            api_key: Some("k".to_string()),
1808            session_id: Some("s1".to_string()),
1809            ..LlmOpts::default()
1810        };
1811        let cache = ClientCache::default();
1812        let builds = || cache.builds.load(std::sync::atomic::Ordering::SeqCst);
1813
1814        // Identical opts reuse the cached client (the whole point: one pool/TLS).
1815        cache.get(&opts).expect("builds");
1816        cache.get(&opts).expect("builds");
1817        assert_eq!(builds(), 1, "identical opts must reuse the client");
1818
1819        // model/temperature/effort ride on the request, not the client.
1820        let model_swap = LlmOpts {
1821            model: "other".to_string(),
1822            effort: Some("high".to_string()),
1823            ..opts.clone()
1824        };
1825        cache.get(&model_swap).expect("builds");
1826        assert_eq!(builds(), 1, "model/effort swap must reuse the client");
1827
1828        // Endpoint, key, and session header each force a rebuild.
1829        for changed in [
1830            LlmOpts {
1831                base_url: Some("http://127.0.0.1:8/v1".to_string()),
1832                ..opts.clone()
1833            },
1834            LlmOpts {
1835                api_key: Some("k2".to_string()),
1836                ..opts.clone()
1837            },
1838            LlmOpts {
1839                session_id: Some("s2".to_string()),
1840                ..opts.clone()
1841            },
1842        ] {
1843            let before = builds();
1844            cache.get(&changed).expect("builds");
1845            assert_eq!(builds(), before + 1, "a key change must rebuild the client");
1846        }
1847    }
1848
1849    #[test]
1850    fn settle_determines_the_client_key() {
1851        // `settle` fully determines base_url/api_key; the rig client caches on
1852        // those (not the endpoint), so an endpoint-only swap reuses it and a
1853        // base_url change rebuilds it.
1854        let base = LlmProvider {
1855            endpoint: EndpointPref::Pinned(LlmEndpoint::Chat),
1856            base_url: Some("http://127.0.0.1:9/v1".to_string()),
1857            api_key: Some("k".to_string()),
1858        };
1859        let mut profiles = ModelProfiles::new(base.clone());
1860        // Same provider, different endpoint (base_url/api_key inherited).
1861        profiles.insert(
1862            "ep-swap",
1863            LlmProfile {
1864                endpoint: Some(LlmEndpoint::Responses),
1865                ..LlmProfile::default()
1866            },
1867        );
1868        // A different provider entirely.
1869        profiles.insert(
1870            "moved",
1871            LlmProfile {
1872                base_url: Some("http://127.0.0.1:8/v1".to_string()),
1873                ..LlmProfile::default()
1874            },
1875        );
1876
1877        let settled = |model: &str| {
1878            let mut llm = LlmOpts {
1879                model: model.to_string(),
1880                endpoint: base.endpoint.resolve(base.base_url.as_deref(), model),
1881                base_url: base.base_url.clone(),
1882                api_key: base.api_key.clone(),
1883                model_profiles: profiles.clone(),
1884                ..LlmOpts::default()
1885            };
1886            llm.settle_provider();
1887            llm
1888        };
1889
1890        let cache = ClientCache::default();
1891        let builds = || cache.builds.load(std::sync::atomic::Ordering::SeqCst);
1892
1893        let plain = settled("m0");
1894        cache.get(&plain).expect("builds");
1895        assert_eq!(builds(), 1);
1896
1897        let ep_swap = settled("ep-swap");
1898        assert_eq!(ep_swap.endpoint, LlmEndpoint::Responses);
1899        assert_eq!(ep_swap.base_url, base.base_url, "the base_url is inherited");
1900        cache.get(&ep_swap).expect("reuses");
1901        assert_eq!(builds(), 1, "an endpoint-only swap must reuse the client");
1902
1903        let moved = settled("moved");
1904        cache.get(&moved).expect("rebuilds");
1905        assert_eq!(builds(), 2, "a base_url change must rebuild the client");
1906    }
1907
1908    const ZEN: &str = "https://opencode.ai/zen/go/v1";
1909
1910    #[test]
1911    fn endpoint_pref_resolve_matrix() {
1912        // Pinned wins over the catalog.
1913        assert_eq!(
1914            EndpointPref::Pinned(LlmEndpoint::Chat).resolve(Some(ZEN), "grok-4.7"),
1915            LlmEndpoint::Chat
1916        );
1917        // Auto picks the catalog for a responses-only zen id.
1918        assert_eq!(
1919            EndpointPref::Auto.resolve(Some(ZEN), "grok-4.7"),
1920            LlmEndpoint::Responses
1921        );
1922        // Auto ⇒ chat for a both-works id, an unknown id, and off-zen.
1923        assert_eq!(
1924            EndpointPref::Auto.resolve(Some(ZEN), "deepseek-v4-flash"),
1925            LlmEndpoint::Chat
1926        );
1927        assert_eq!(
1928            EndpointPref::Auto.resolve(Some("https://api.openai.com/v1"), "grok-4.7"),
1929            LlmEndpoint::Chat
1930        );
1931        assert_eq!(
1932            EndpointPref::Auto.resolve(None, "grok-4.7"),
1933            LlmEndpoint::Chat
1934        );
1935    }
1936
1937    #[test]
1938    fn endpoint_pref_wire_and_from_option() {
1939        assert_eq!(EndpointPref::Auto.wire(), "auto");
1940        assert_eq!(EndpointPref::Pinned(LlmEndpoint::Chat).wire(), "chat");
1941        assert_eq!(EndpointPref::Pinned(LlmEndpoint::Responses).wire(), "responses");
1942        assert_eq!(EndpointPref::from(None), EndpointPref::Auto);
1943        assert_eq!(
1944            EndpointPref::from(Some(LlmEndpoint::Responses)),
1945            EndpointPref::Pinned(LlmEndpoint::Responses)
1946        );
1947    }
1948
1949    #[test]
1950    fn settle_auto_derives_from_the_catalog() {
1951        // Base pref = Auto, zen base: a responses-only id ⇒ Responses; a both-works
1952        // id ⇒ Chat.
1953        let base = LlmProvider {
1954            endpoint: EndpointPref::Auto,
1955            base_url: Some(ZEN.to_string()),
1956            api_key: None,
1957        };
1958        let profiles = ModelProfiles::new(base.clone());
1959        let settled = |model: &str| {
1960            let mut llm = LlmOpts {
1961                model: model.to_string(),
1962                // Seed a STALE concrete `Chat`: only `settle` can produce
1963                // `Responses`, so deleting its derive step fails this test.
1964                endpoint: LlmEndpoint::Chat,
1965                base_url: base.base_url.clone(),
1966                model_profiles: profiles.clone(),
1967                ..LlmOpts::default()
1968            };
1969            llm.settle_provider();
1970            llm
1971        };
1972        assert_eq!(settled("grok-4.7").endpoint, LlmEndpoint::Responses);
1973        assert_eq!(settled("deepseek-v4-flash").endpoint, LlmEndpoint::Chat);
1974    }
1975
1976    #[test]
1977    fn settle_profile_pin_beats_auto_and_base_pin() {
1978        // A profile `endpoint = Some(Chat)` pins even on a responses-only id —
1979        // profile pin ▸ base pin ▸ catalog ▸ chat.
1980        let base = LlmProvider {
1981            endpoint: EndpointPref::Auto,
1982            base_url: Some(ZEN.to_string()),
1983            api_key: None,
1984        };
1985        let mut profiles = ModelProfiles::new(base.clone());
1986        profiles.insert(
1987            "grok-4.7",
1988            LlmProfile {
1989                endpoint: Some(LlmEndpoint::Chat),
1990                ..LlmProfile::default()
1991            },
1992        );
1993        let mut llm = LlmOpts {
1994            model: "grok-4.7".into(),
1995            base_url: base.base_url.clone(),
1996            model_profiles: profiles,
1997            ..LlmOpts::default()
1998        };
1999        llm.settle_provider();
2000        assert_eq!(
2001            llm.endpoint,
2002            LlmEndpoint::Chat,
2003            "a profile pin must beat the catalog"
2004        );
2005    }
2006
2007    #[test]
2008    fn settle_profile_base_url_only_re_resolves_the_endpoint() {
2009        // A profile that sets ONLY base_url (endpoint None) re-derives the endpoint
2010        // against the NEW base — here it moves onto zen, so auto ⇒ Responses.
2011        let base = LlmProvider {
2012            endpoint: EndpointPref::Auto,
2013            base_url: Some("http://plain/v1".into()),
2014            api_key: None,
2015        };
2016        let mut profiles = ModelProfiles::new(base.clone());
2017        profiles.insert(
2018            "grok-4.7",
2019            LlmProfile {
2020                base_url: Some(ZEN.to_string()),
2021                ..LlmProfile::default()
2022            },
2023        );
2024        let mut llm = LlmOpts {
2025            model: "grok-4.7".into(),
2026            base_url: base.base_url.clone(),
2027            model_profiles: profiles,
2028            ..LlmOpts::default()
2029        };
2030        llm.settle_provider();
2031        assert_eq!(llm.base_url.as_deref(), Some(ZEN));
2032        assert_eq!(
2033            llm.endpoint,
2034            LlmEndpoint::Responses,
2035            "a base_url-only profile must re-derive against the new base"
2036        );
2037    }
2038
2039    #[test]
2040    fn set_model_flips_chat_to_responses_and_back() {
2041        // `settle` is the funnel `Agent::set_model` calls: boot on a both-works zen
2042        // id (Chat), switch to `grok-4.7` ⇒ Responses, then back ⇒ Chat.
2043        let base = LlmProvider {
2044            endpoint: EndpointPref::Auto,
2045            base_url: Some(ZEN.to_string()),
2046            api_key: None,
2047        };
2048        let profiles = ModelProfiles::new(base.clone());
2049        let mut llm = LlmOpts {
2050            model: "deepseek-v4-flash".into(),
2051            base_url: base.base_url.clone(),
2052            model_profiles: profiles,
2053            ..LlmOpts::default()
2054        };
2055        llm.settle_provider();
2056        assert_eq!(llm.endpoint, LlmEndpoint::Chat);
2057        llm.model = "grok-4.7".into();
2058        llm.settle_provider();
2059        assert_eq!(llm.endpoint, LlmEndpoint::Responses);
2060        llm.model = "deepseek-v4-flash".into();
2061        llm.settle_provider();
2062        assert_eq!(llm.endpoint, LlmEndpoint::Chat);
2063    }
2064
2065    #[test]
2066    fn responses_request_builds_against_default_client_type() {
2067        use rig::client::CompletionClient;
2068        let client = openai::Client::builder()
2069            .api_key("test-key")
2070            .base_url("http://127.0.0.1:9/v1")
2071            .build()
2072            .expect("client builds without network");
2073        let model = client.completion_model("some-model");
2074        let request = build_request(&[], "sys", &[], &LlmOpts::default());
2075        let fut = rig::completion::CompletionModel::stream(&model, request);
2076        drop(fut);
2077    }
2078
2079    /// Unreachable endpoint: the adapter surfaces an `Error` event and no
2080    /// `Done`, at the stream level (initial request failure).
2081    #[tokio::test]
2082    async fn adapter_surfaces_error_without_done() {
2083        let opts = LlmOpts {
2084            model: "m1".to_string(),
2085            base_url: Some("http://127.0.0.1:9/v1".to_string()),
2086            api_key: Some("test-key".to_string()),
2087            temperature: None,
2088            ..LlmOpts::default()
2089        };
2090        let stream_fn = rig_stream_fn();
2091        let stream = stream_fn(&[], "sys", &[], &opts);
2092        let events: Vec<LlmStreamEvent> =
2093            tokio::time::timeout(std::time::Duration::from_secs(10), stream.collect())
2094                .await
2095                .expect("stream terminates");
2096        assert!(
2097            events
2098                .iter()
2099                .any(|e| matches!(e, LlmStreamEvent::Error { .. })),
2100            "expected an Error event, got {events:?}"
2101        );
2102        assert!(
2103            !events
2104                .iter()
2105                .any(|e| matches!(e, LlmStreamEvent::Done { .. })),
2106            "no Done after error, got {events:?}"
2107        );
2108    }
2109
2110    /// Responses path hits the same shared forwarding loop: unreachable
2111    /// endpoint surfaces `Error` and no `Done`.
2112    #[tokio::test]
2113    async fn responses_adapter_surfaces_error_without_done() {
2114        let opts = LlmOpts {
2115            model: "m1".to_string(),
2116            base_url: Some("http://127.0.0.1:9/v1".to_string()),
2117            api_key: Some("test-key".to_string()),
2118            temperature: None,
2119            endpoint: LlmEndpoint::Responses,
2120            effort: None,
2121            session_id: None,
2122            retry: RetryPolicy::default(),
2123            model_profiles: ModelProfiles::default(),
2124        };
2125        let stream_fn = rig_stream_fn();
2126        let stream = stream_fn(&[], "sys", &[], &opts);
2127        let events: Vec<LlmStreamEvent> =
2128            tokio::time::timeout(std::time::Duration::from_secs(10), stream.collect())
2129                .await
2130                .expect("stream terminates");
2131        assert!(
2132            events
2133                .iter()
2134                .any(|e| matches!(e, LlmStreamEvent::Error { .. })),
2135            "expected an Error event, got {events:?}"
2136        );
2137        assert!(
2138            !events
2139                .iter()
2140                .any(|e| matches!(e, LlmStreamEvent::Done { .. })),
2141            "no Done after error, got {events:?}"
2142        );
2143    }
2144    /// Live smoke: real OpenAI-compatible endpoint, gated on env.
2145    #[tokio::test]
2146    #[ignore = "live test; run with WCODE_LIVE=1 WCODE_BASE_URL=... WCODE_API_KEY=... WCODE_MODEL=... cargo test -p wcode-harness --lib -- --ignored"]
2147    async fn live_rig_stream_fn() {
2148        let (Ok(base_url), Ok(api_key), Ok(model)) = (
2149            std::env::var("WCODE_BASE_URL"),
2150            std::env::var("WCODE_API_KEY"),
2151            std::env::var("WCODE_MODEL"),
2152        ) else {
2153            eprintln!("skipping: WCODE_BASE_URL / WCODE_API_KEY / WCODE_MODEL not set");
2154            return;
2155        };
2156        if std::env::var("WCODE_LIVE").ok().as_deref() != Some("1") {
2157            eprintln!("skipping: WCODE_LIVE != 1");
2158            return;
2159        }
2160        let opts = LlmOpts {
2161            model,
2162            base_url: Some(base_url),
2163            api_key: Some(api_key),
2164            temperature: None,
2165            ..LlmOpts::default()
2166        };
2167        let stream_fn = rig_stream_fn();
2168        let stream = stream_fn(
2169            &[AgentMessage::user_text("Reply with exactly: pong")],
2170            "You are a terse echo assistant.",
2171            &[],
2172            &opts,
2173        );
2174        let events: Vec<LlmStreamEvent> = stream.collect().await;
2175        let text = events
2176            .iter()
2177            .filter_map(|e| match e {
2178                LlmStreamEvent::TextDelta(t) => Some(t.as_str()),
2179                _ => None,
2180            })
2181            .collect::<String>();
2182        println!("text: {text:?}");
2183        println!("events: {events:?}");
2184        assert!(
2185            events
2186                .iter()
2187                .any(|e| matches!(e, LlmStreamEvent::Done { .. }))
2188        );
2189        assert!(
2190            !events
2191                .iter()
2192                .any(|e| matches!(e, LlmStreamEvent::Error { .. }))
2193        );
2194    }
2195}
2196
2197#[cfg(test)]
2198mod retry_tests {
2199    use super::*;
2200    use std::sync::Arc;
2201    use std::sync::atomic::{AtomicUsize, Ordering};
2202
2203    type Chunks =
2204        futures::stream::Iter<std::vec::IntoIter<Result<StreamedAssistantContent, CompletionError>>>;
2205
2206    /// A transient transport failure as rig surfaces it on the OpenAI path.
2207    fn transient() -> CompletionError {
2208        CompletionError::ProviderError(
2209            "Http client error: error sending request for url (http://x)".into(),
2210        )
2211    }
2212
2213    fn fatal() -> CompletionError {
2214        CompletionError::ProviderError("invalid_request_error: bad tool schema".into())
2215    }
2216
2217    fn empty() -> Chunks {
2218        futures::stream::iter(Vec::new())
2219    }
2220
2221    /// A terminal `Final` item (the stream's `Done`).
2222    fn done() -> StreamedAssistantContent {
2223        StreamedAssistantContent::Final(rig::streaming::StreamFinal::new(
2224            "test",
2225            rig::completion::Usage::default(),
2226        ))
2227    }
2228
2229    fn first_item_err(e: CompletionError) -> Chunks {
2230        futures::stream::iter(vec![Err(e)])
2231    }
2232
2233    #[test]
2234    fn transient_statuses_retry_client_errors_do_not() {
2235        for code in [408, 409, 429, 500, 502, 503, 504] {
2236            assert!(is_transient_status(code), "{code} should retry");
2237        }
2238        for code in [200, 400, 401, 403, 404, 422] {
2239            assert!(!is_transient_status(code), "{code} must not retry");
2240        }
2241    }
2242
2243    #[test]
2244    fn retryable_classifies_transport_vs_rejections() {
2245        assert!(retryable(&CompletionError::HttpError(
2246            rig::http_client::Error::StreamEnded
2247        )));
2248        assert!(retryable(&transient()));
2249        assert!(!retryable(&fatal()));
2250        assert!(!retryable(&CompletionError::ResponseError("bad".into())));
2251    }
2252
2253    #[test]
2254    fn backoff_grows_and_never_exceeds_the_cap() {
2255        let base = Duration::from_millis(100);
2256        let cap = Duration::from_millis(400);
2257        let d1 = backoff(1, base, cap);
2258        assert!(d1 <= base && d1 >= base * 3 / 4, "attempt 1 = {d1:?}");
2259        for n in 1..12 {
2260            assert!(backoff(n, base, cap) <= cap, "attempt {n} exceeds cap");
2261        }
2262    }
2263
2264    #[tokio::test(start_paused = true)]
2265    async fn connect_retries_transient_failures_then_succeeds() {
2266        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2267        let calls = Arc::new(AtomicUsize::new(0));
2268        let policy = RetryPolicy {
2269            max: 3,
2270            base: Duration::from_millis(10),
2271            cap: Duration::from_millis(50),
2272            ..RetryPolicy::default()
2273        };
2274        let c = calls.clone();
2275        let mut state = RetryState::new(policy.max);
2276        let out = connect_with_retry(
2277            &policy,
2278            &mut state,
2279            move || {
2280                let n = c.fetch_add(1, Ordering::SeqCst) + 1;
2281                Box::pin(async move {
2282                    if n <= 2 {
2283                        Err(transient())
2284                    } else {
2285                        Ok(empty())
2286                    }
2287                })
2288            },
2289            &tx,
2290        )
2291        .await;
2292        assert!(matches!(out, PreContent::Ready { .. }));
2293        assert_eq!(calls.load(Ordering::SeqCst), 3);
2294        assert!(matches!(
2295            rx.try_recv(),
2296            Ok(LlmStreamEvent::Retrying { attempt: 1, .. })
2297        ));
2298        assert!(matches!(
2299            rx.try_recv(),
2300            Ok(LlmStreamEvent::Retrying { attempt: 2, .. })
2301        ));
2302    }
2303
2304    #[tokio::test(start_paused = true)]
2305    async fn connect_retries_a_first_item_transport_error() {
2306        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
2307        let calls = Arc::new(AtomicUsize::new(0));
2308        let policy = RetryPolicy {
2309            max: 2,
2310            base: Duration::from_millis(5),
2311            cap: Duration::from_millis(5),
2312            ..RetryPolicy::default()
2313        };
2314        let c = calls.clone();
2315        let mut state = RetryState::new(policy.max);
2316        let out = connect_with_retry(
2317            &policy,
2318            &mut state,
2319            move || {
2320                let n = c.fetch_add(1, Ordering::SeqCst) + 1;
2321                Box::pin(async move {
2322                    if n == 1 {
2323                        Ok(first_item_err(transient()))
2324                    } else {
2325                        Ok(empty())
2326                    }
2327                })
2328            },
2329            &tx,
2330        )
2331        .await;
2332        assert!(matches!(out, PreContent::Ready { .. }));
2333        assert_eq!(calls.load(Ordering::SeqCst), 2, "first-item error retried");
2334    }
2335
2336    #[tokio::test(start_paused = true)]
2337    async fn connect_gives_up_after_max_and_skips_non_retryable() {
2338        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
2339        let policy = RetryPolicy {
2340            max: 1,
2341            base: Duration::from_millis(1),
2342            cap: Duration::from_millis(1),
2343            ..RetryPolicy::default()
2344        };
2345
2346        let calls = Arc::new(AtomicUsize::new(0));
2347        let c = calls.clone();
2348        let mut state = RetryState::new(policy.max);
2349        let out = connect_with_retry(
2350            &policy,
2351            &mut state,
2352            move || {
2353                c.fetch_add(1, Ordering::SeqCst);
2354                Box::pin(async { Err::<Chunks, CompletionError>(transient()) })
2355            },
2356            &tx,
2357        )
2358        .await;
2359        assert!(matches!(out, PreContent::Failed(_)));
2360        assert_eq!(calls.load(Ordering::SeqCst), 2, "one attempt + one retry");
2361
2362        let calls = Arc::new(AtomicUsize::new(0));
2363        let c = calls.clone();
2364        let mut state = RetryState::new(policy.max);
2365        let out = connect_with_retry(
2366            &policy,
2367            &mut state,
2368            move || {
2369                c.fetch_add(1, Ordering::SeqCst);
2370                Box::pin(async { Err::<Chunks, CompletionError>(fatal()) })
2371            },
2372            &tx,
2373        )
2374        .await;
2375        assert!(matches!(out, PreContent::Failed(_)));
2376        assert_eq!(calls.load(Ordering::SeqCst), 1, "non-retryable returns at once");
2377    }
2378
2379    #[test]
2380    fn fatal_class_flags_hard_errors_but_not_transient_transport() {
2381        // Transient-class failures stay feedable after their retries are
2382        // exhausted (the loop feeds the message back to the model); hard
2383        // client/build failures are fatal (the run stops).
2384        assert!(!fatal_class(&transient()));
2385        assert!(!fatal_class(&CompletionError::HttpError(
2386            rig::http_client::Error::StreamEnded
2387        )));
2388        assert!(fatal_class(&fatal()));
2389        assert!(fatal_class(&CompletionError::ResponseError("bad".into())));
2390        assert!(fatal_class(&CompletionError::RequestError("bad".into())));
2391        let json_err = serde_json::from_str::<u8>("not a number").unwrap_err();
2392        assert!(fatal_class(&CompletionError::JsonError(json_err)));
2393    }
2394
2395    /// Live-ish check against a refused port: a transport failure exhausts the
2396    /// retry budget (Retrying notices) and surfaces as a *non-fatal* stream
2397    /// error — which is exactly what lets the loop feed it back next turn
2398    /// instead of dying.
2399    #[tokio::test]
2400    async fn refused_port_exhausts_retries_to_non_fatal_error() {
2401        let opts = LlmOpts {
2402            model: "m1".to_string(),
2403            base_url: Some("http://127.0.0.1:9/v1".to_string()),
2404            api_key: Some("test-key".to_string()),
2405            retry: RetryPolicy {
2406                max: 2,
2407                base: Duration::from_millis(10),
2408                cap: Duration::from_millis(50),
2409                ..RetryPolicy::default()
2410            },
2411            ..LlmOpts::default()
2412        };
2413        let stream_fn = rig_stream_fn();
2414        let stream = stream_fn(&[], "sys", &[], &opts);
2415        let events: Vec<LlmStreamEvent> =
2416            tokio::time::timeout(std::time::Duration::from_secs(10), stream.collect())
2417                .await
2418                .expect("stream terminates");
2419        let retries = events
2420            .iter()
2421            .filter(|e| matches!(e, LlmStreamEvent::Retrying { .. }))
2422            .count();
2423        assert_eq!(retries, 2, "two backoff notices before giving up: {events:?}");
2424        assert!(
2425            matches!(events.last(), Some(LlmStreamEvent::Error { fatal: false, .. })),
2426            "refused port is transient-class, not fatal: {events:?}"
2427        );
2428    }
2429
2430    /// A `CompletionError` carrying a preserved response with a `Retry-After`
2431    /// header (no HTTP status needed: the header path reads
2432    /// `provider_response_headers()` regardless of status).
2433    fn err_with_retry_after(value: &str) -> CompletionError {
2434        let mut headers = HeaderMap::new();
2435        headers.insert("retry-after", HeaderValue::from_str(value).unwrap());
2436        CompletionError::ProviderResponse(
2437            rig::ProviderResponseError::without_status("429")
2438                .with_headers(Some(Box::new(headers))),
2439        )
2440    }
2441
2442    #[test]
2443    fn retry_after_parses_integer_seconds() {
2444        assert_eq!(
2445            retry_after(&err_with_retry_after("7")),
2446            Some(Duration::from_secs(7))
2447        );
2448    }
2449
2450    #[test]
2451    fn retry_after_is_none_when_absent() {
2452        // A transport failure with no preserved response/headers.
2453        assert_eq!(retry_after(&transient()), None);
2454    }
2455
2456    #[test]
2457    fn retry_after_is_none_on_non_numeric() {
2458        // The HTTP-date form and garbage both parse to `None` (dates are out of
2459        // scope — see `retry_after`).
2460        assert_eq!(
2461            retry_after(&err_with_retry_after("Wed, 21 Oct 2015 07:28:00 GMT")),
2462            None
2463        );
2464        assert_eq!(retry_after(&err_with_retry_after("soon")), None);
2465    }
2466
2467    #[tokio::test(start_paused = true)]
2468    async fn retry_wait_honors_retry_after_over_backoff() {
2469        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
2470        // base 10ms gives a ~10ms backoff for attempt 1; `Retry-After: 3` must
2471        // win, so the paused clock advances ~3s instead.
2472        let policy = RetryPolicy {
2473            base: Duration::from_millis(10),
2474            ..RetryPolicy::default()
2475        };
2476        let start = tokio::time::Instant::now();
2477        assert!(retry_wait(&tx, &policy, 1, &err_with_retry_after("3")).await);
2478        assert_eq!(start.elapsed(), Duration::from_secs(3));
2479    }
2480
2481    #[tokio::test(start_paused = true)]
2482    async fn retry_wait_clamps_retry_after_to_the_cap() {
2483        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
2484        // A hostile/buggy `Retry-After: 100` must not park the turn for 100s:
2485        // it is clamped to `policy.cap` (8s), not honored verbatim.
2486        let policy = RetryPolicy {
2487            cap: Duration::from_secs(8),
2488            ..RetryPolicy::default()
2489        };
2490        let start = tokio::time::Instant::now();
2491        assert!(retry_wait(&tx, &policy, 1, &err_with_retry_after("100")).await);
2492        assert_eq!(start.elapsed(), Duration::from_secs(8));
2493    }
2494    // ---- #23: stream-stall timeouts ----
2495
2496    type BoxStream =
2497        Pin<Box<dyn futures::Stream<Item = Result<StreamedAssistantContent, CompletionError>> + Send>>;
2498
2499    fn boxed<S>(s: S) -> BoxStream
2500    where
2501        S: futures::Stream<Item = Result<StreamedAssistantContent, CompletionError>>
2502            + Send
2503            + 'static,
2504    {
2505        Box::pin(s)
2506    }
2507
2508    fn silent_boxed() -> BoxStream {
2509        boxed(futures::stream::pending::<Result<StreamedAssistantContent, CompletionError>>())
2510    }
2511
2512    /// A fresh silent stream per call — `forward_stream`/`connect_with_retry`'s
2513    /// re-open seam, never exercised where content already flowed.
2514    fn open_silent()
2515    -> Pin<Box<dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send>> {
2516        Box::pin(async { Ok::<_, CompletionError>(silent_boxed()) })
2517    }
2518
2519    /// A slow (but successful) connect — used to prove `ttft = ZERO` awaits it
2520    /// directly instead of racing an immediate `sleep(ZERO)`.
2521    fn open_slow_ok()
2522    -> Pin<Box<dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send>> {
2523        Box::pin(async {
2524            tokio::time::sleep(Duration::from_millis(50)).await;
2525            Ok::<_, CompletionError>(boxed(futures::stream::iter(vec![Ok::<_, CompletionError>(
2526                StreamedAssistantContent::Text(rig::message::Text::new("ok")),
2527            )])))
2528        })
2529    }
2530
2531    /// A stall is transient-class, so it rides the retry / feed-back paths (B1).
2532    #[test]
2533    fn a_stall_error_is_transient_and_not_hard_fatal() {
2534        let d = Duration::from_secs(1);
2535        assert!(retryable(&stall_error(d)), "a stall must be transient (B1)");
2536        assert!(!fatal_class(&stall_error(d)), "a stall must not be hard-fatal");
2537        assert!(
2538            stall_error(d).to_string().contains("timed out"),
2539            "the (timed out) needle drives retryable()"
2540        );
2541    }
2542
2543    /// A silent stream stalls on the `ttft` peek every attempt: each retry emits
2544    /// a `Retrying` notice, and once the budget is spent the failure is a
2545    /// transient (non-fatal) `stall_error`.
2546    #[tokio::test(start_paused = true)]
2547    async fn a_silent_stream_trips_the_ttft_retries_then_fails() {
2548        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2549        let calls = Arc::new(AtomicUsize::new(0));
2550        let policy = RetryPolicy {
2551            max: 2,
2552            base: Duration::from_millis(10),
2553            cap: Duration::from_millis(10),
2554            ttft: Duration::from_secs(1),
2555            ..RetryPolicy::default()
2556        };
2557        let c = calls.clone();
2558        let open = move || {
2559            c.fetch_add(1, Ordering::SeqCst);
2560            open_silent()
2561        };
2562        let mut state = RetryState::new(policy.max);
2563        let out = connect_with_retry(&policy, &mut state, open, &tx).await;
2564
2565        let PreContent::Failed(error) = out else {
2566            panic!("expected Failed from an exhausted stall")
2567        };
2568        assert!(retryable(&error), "a stall is transient");
2569        assert!(!fatal_class(&error), "a stall is not hard-fatal");
2570        assert_eq!(calls.load(Ordering::SeqCst), 3, "one attempt + two retries");
2571        let retries = std::iter::from_fn(|| rx.try_recv().ok())
2572            .filter(|e| matches!(e, LlmStreamEvent::Retrying { .. }))
2573            .count();
2574        assert_eq!(retries, 2, "a Retrying notice per retry");
2575    }
2576
2577    /// A stream that emits a delta and then goes silent trips the idle deadline:
2578    /// because content already flowed, the failure is a NON-fatal error (the
2579    /// loop feeds it back) and is not retried (that would duplicate output).
2580    #[tokio::test(start_paused = true)]
2581    async fn deltas_then_a_stall_is_a_nonfatal_error() {
2582        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2583        let policy = RetryPolicy {
2584            max: 3,
2585            base: Duration::from_millis(10),
2586            cap: Duration::from_millis(10),
2587            idle: Duration::from_secs(1),
2588            ..RetryPolicy::default()
2589        };
2590        let stream = boxed(
2591            futures::stream::iter(vec![Ok::<_, CompletionError>(
2592                StreamedAssistantContent::Text(rig::message::Text::new("hi")),
2593            )])
2594            .chain(futures::stream::pending::<Result<StreamedAssistantContent, CompletionError>>()),
2595        );
2596        let state = RetryState::new(policy.max);
2597        forward_stream(stream, None, policy, open_silent, tx, state).await;
2598
2599        let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2600        assert!(
2601            events
2602                .iter()
2603                .any(|e| matches!(e, LlmStreamEvent::TextDelta(t) if t == "hi")),
2604            "{events:?}"
2605        );
2606        let error = events
2607            .iter()
2608            .find_map(|e| match e {
2609                LlmStreamEvent::Error { message, fatal } => Some((message.clone(), *fatal)),
2610                _ => None,
2611            })
2612            .expect("the idle deadline surfaced an error");
2613        assert!(!error.1, "a post-content stall is non-fatal: {error:?}");
2614        assert!(
2615            !events.iter().any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
2616            "no retry once content flowed: {events:?}"
2617        );
2618    }
2619
2620    /// A clean stream (deltas then `Done`) is forwarded untouched: no error, no
2621    /// retry, the idle deadline never fires.
2622    #[tokio::test(start_paused = true)]
2623    async fn a_clean_stream_is_unaffected() {
2624        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2625        let policy = RetryPolicy {
2626            idle: Duration::from_secs(1),
2627            ..RetryPolicy::default()
2628        };
2629        let stream = boxed(futures::stream::iter(vec![
2630            Ok::<_, CompletionError>(StreamedAssistantContent::Text(rig::message::Text::new("hi"))),
2631            Ok::<_, CompletionError>(done()),
2632        ]));
2633        let state = RetryState::new(policy.max);
2634        forward_stream(stream, None, policy, open_silent, tx, state).await;
2635
2636        let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2637        assert!(
2638            events
2639                .iter()
2640                .any(|e| matches!(e, LlmStreamEvent::TextDelta(t) if t == "hi")),
2641            "{events:?}"
2642        );
2643        assert!(
2644            events.iter().any(|e| matches!(e, LlmStreamEvent::Done { .. })),
2645            "{events:?}"
2646        );
2647        assert!(
2648            !events.iter().any(|e| matches!(e, LlmStreamEvent::Error { .. })),
2649            "{events:?}"
2650        );
2651        assert!(
2652            !events.iter().any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
2653            "{events:?}"
2654        );
2655    }
2656
2657    /// `ZERO` disables the deadline at BOTH seams: the peek collapses to a plain
2658    /// `stream.next()` (a silent stream stays pending), and the connect race is
2659    /// skipped so a slow connect is not spuriously timed out.
2660    #[tokio::test(start_paused = true)]
2661    async fn zero_deadlines_disable_the_timeout() {
2662        // Peek seam: no `sleep(ZERO)`, so a silent stream never yields `Idle`.
2663        let mut silent =
2664            futures::stream::pending::<Result<StreamedAssistantContent, CompletionError>>();
2665        let stalled = tokio::time::timeout(
2666            Duration::from_secs(1),
2667            next_or_stall(&mut silent, Duration::ZERO),
2668        )
2669        .await;
2670        assert!(stalled.is_err(), "ZERO must not trip the deadline");
2671
2672        // Connect seam: `ZERO` awaits the connect directly.
2673        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
2674        let policy = RetryPolicy {
2675            ttft: Duration::ZERO,
2676            ..RetryPolicy::default()
2677        };
2678        let mut state = RetryState::new(policy.max);
2679        let out = connect_with_retry(&policy, &mut state, open_slow_ok, &tx).await;
2680        assert!(
2681            matches!(out, PreContent::Ready { .. }),
2682            "a slow connect must not be timed out under ZERO"
2683        );
2684    }
2685
2686    /// The genuinely-stalled *connect* branch: `open()` never resolves, so the
2687    /// nonzero `ttft` race times out on every attempt. Each retry emits a
2688    /// `Retrying` notice; once the budget is spent the failure is a transient
2689    /// (non-fatal) `stall_error` — the same shape as a peek stall, but through
2690    /// the connect race's `Err(())` arm.
2691    #[tokio::test(start_paused = true)]
2692    async fn a_stalled_connect_trips_the_ttft_retries_then_fails() {
2693        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2694        let calls = Arc::new(AtomicUsize::new(0));
2695        let policy = RetryPolicy {
2696            max: 2,
2697            base: Duration::from_millis(10),
2698            cap: Duration::from_millis(10),
2699            ttft: Duration::from_secs(1),
2700            ..RetryPolicy::default()
2701        };
2702        let c = calls.clone();
2703        let open = move || {
2704            c.fetch_add(1, Ordering::SeqCst);
2705            Box::pin(std::future::pending::<Result<BoxStream, CompletionError>>())
2706                as Pin<
2707                    Box<
2708                        dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send,
2709                    >,
2710                >
2711        };
2712        let mut state = RetryState::new(policy.max);
2713        let out = connect_with_retry(&policy, &mut state, open, &tx).await;
2714
2715        let PreContent::Failed(error) = out else {
2716            panic!("expected Failed from an exhausted connect stall")
2717        };
2718        assert!(retryable(&error), "a connect stall is transient");
2719        assert!(!fatal_class(&error), "a connect stall is not hard-fatal");
2720        assert_eq!(calls.load(Ordering::SeqCst), 3, "one attempt + two retries");
2721        let retries = std::iter::from_fn(|| rx.try_recv().ok())
2722            .filter(|e| matches!(e, LlmStreamEvent::Retrying { .. }))
2723            .count();
2724        assert_eq!(retries, 2, "a Retrying notice per retry");
2725    }
2726
2727    /// The forward loop's *pre-content* idle branch: a stream that emits nothing
2728    /// and then stalls, with budget remaining, backs off (a `Retrying` notice)
2729    /// and re-opens — here the re-opened stream succeeds and the turn completes.
2730    #[tokio::test(start_paused = true)]
2731    async fn a_pre_content_idle_stall_retries_and_reopens() {
2732        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2733        let policy = RetryPolicy {
2734            max: 2,
2735            base: Duration::from_millis(10),
2736            cap: Duration::from_millis(10),
2737            idle: Duration::from_secs(1),
2738            ..RetryPolicy::default()
2739        };
2740        // The initial stream never yields → the idle deadline trips pre-content.
2741        let stream = silent_boxed();
2742        // The re-open yields a `Final`, so the retried turn completes.
2743        fn open_done()
2744        -> Pin<Box<dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send>> {
2745            Box::pin(async {
2746                Ok::<_, CompletionError>(boxed(futures::stream::iter(vec![Ok::<_, CompletionError>(
2747                    done(),
2748                )])))
2749            })
2750        }
2751        let state = RetryState::new(policy.max);
2752        forward_stream(stream, None, policy, open_done, tx, state).await;
2753
2754        let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2755        assert!(
2756            events
2757                .iter()
2758                .any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
2759            "a pre-content stall backs off: {events:?}"
2760        );
2761        assert!(
2762            events.iter().any(|e| matches!(e, LlmStreamEvent::Done { .. })),
2763            "the re-opened stream completed: {events:?}"
2764        );
2765        assert!(
2766            !events.iter().any(|e| matches!(e, LlmStreamEvent::Error { .. })),
2767            "no error once the retry succeeded: {events:?}"
2768        );
2769    }
2770
2771    /// A re-open that ALSO ends immediately, pre-content: paired with an empty
2772    /// initial stream, it drives the truncation arm's retries to exhaustion (each
2773    /// re-open yields no item, so nothing is ever forwarded).
2774    fn open_empty()
2775    -> Pin<Box<dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send>> {
2776        Box::pin(async { Ok::<_, CompletionError>(boxed(empty())) })
2777    }
2778
2779    /// A delta followed by an early EOF (no `Final`) is a truncation: content
2780    /// already flowed, so it is NOT retried — a single NON-fatal `Error` ends the
2781    /// turn (the loop feeds it back), and no `Done` is emitted.
2782    #[tokio::test(start_paused = true)]
2783    async fn a_delta_then_a_truncated_end_is_a_nonfatal_error() {
2784        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2785        let policy = RetryPolicy {
2786            max: 3,
2787            base: Duration::from_millis(10),
2788            cap: Duration::from_millis(10),
2789            idle: Duration::ZERO,
2790            ..RetryPolicy::default()
2791        };
2792        let stream = boxed(futures::stream::iter(vec![Ok::<_, CompletionError>(
2793            StreamedAssistantContent::Text(rig::message::Text::new("hi")),
2794        )]));
2795        let state = RetryState::new(policy.max);
2796        forward_stream(stream, None, policy, open_silent, tx, state).await;
2797
2798        let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2799        assert!(
2800            events
2801                .iter()
2802                .any(|e| matches!(e, LlmStreamEvent::TextDelta(t) if t == "hi")),
2803            "{events:?}"
2804        );
2805        let (message, fatal) = events
2806            .iter()
2807            .find_map(|e| match e {
2808                LlmStreamEvent::Error { message, fatal } => Some((message.clone(), *fatal)),
2809                _ => None,
2810            })
2811            .expect("a truncated end surfaced an error");
2812        assert!(!fatal, "a post-content truncation is non-fatal: {message}");
2813        assert!(
2814            !events
2815                .iter()
2816                .any(|e| matches!(e, LlmStreamEvent::Done { .. })),
2817            "no Done on a truncated stream: {events:?}"
2818        );
2819        assert!(
2820            !events
2821                .iter()
2822                .any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
2823            "no retry once content flowed: {events:?}"
2824        );
2825    }
2826
2827    /// A `Final`-terminated stream is unchanged: it forwards `Done` and never
2828    /// trips the truncation guard (which would surface a spurious error).
2829    #[tokio::test(start_paused = true)]
2830    async fn a_final_terminated_stream_is_unchanged() {
2831        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2832        let policy = RetryPolicy {
2833            idle: Duration::ZERO,
2834            ..RetryPolicy::default()
2835        };
2836        let stream = boxed(futures::stream::iter(vec![
2837            Ok::<_, CompletionError>(StreamedAssistantContent::Text(rig::message::Text::new(
2838                "hi",
2839            ))),
2840            Ok::<_, CompletionError>(done()),
2841        ]));
2842        let state = RetryState::new(policy.max);
2843        forward_stream(stream, None, policy, open_silent, tx, state).await;
2844
2845        let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2846        assert!(
2847            events
2848                .iter()
2849                .any(|e| matches!(e, LlmStreamEvent::Done { .. })),
2850            "{events:?}"
2851        );
2852        assert!(
2853            !events
2854                .iter()
2855                .any(|e| matches!(e, LlmStreamEvent::Error { .. })),
2856            "a clean turn surfaces no error: {events:?}"
2857        );
2858        assert!(
2859            !events
2860                .iter()
2861                .any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
2862            "{events:?}"
2863        );
2864    }
2865
2866    /// A pre-content no-item stream with budget remaining retries on the shared
2867    /// budget; each re-open also ends empty, so the whole budget is spent before
2868    /// a NON-fatal error ends the turn — one `Retrying` notice per retry, no
2869    /// `Done`.
2870    #[tokio::test(start_paused = true)]
2871    async fn a_pre_content_no_item_stream_retries_then_errors() {
2872        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2873        let policy = RetryPolicy {
2874            max: 2,
2875            base: Duration::from_millis(10),
2876            cap: Duration::from_millis(10),
2877            idle: Duration::ZERO,
2878            ..RetryPolicy::default()
2879        };
2880        let state = RetryState::new(policy.max);
2881        forward_stream(boxed(empty()), None, policy, open_empty, tx, state).await;
2882
2883        let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2884        let retries = events
2885            .iter()
2886            .filter(|e| matches!(e, LlmStreamEvent::Retrying { .. }))
2887            .count();
2888        assert_eq!(
2889            retries, policy.max as usize,
2890            "one Retrying notice per retry: {events:?}"
2891        );
2892        let (message, fatal) = events
2893            .iter()
2894            .find_map(|e| match e {
2895                LlmStreamEvent::Error { message, fatal } => Some((message.clone(), *fatal)),
2896                _ => None,
2897            })
2898            .expect("an exhausted truncation budget surfaced an error");
2899        assert!(!fatal, "an exhausted truncation is non-fatal: {message}");
2900        assert!(
2901            !events
2902                .iter()
2903                .any(|e| matches!(e, LlmStreamEvent::Done { .. })),
2904            "no Done on a truncated stream: {events:?}"
2905        );
2906    }
2907
2908    // --- probe_endpoint (opt-in local detection) ---------------------------
2909
2910    /// Stand up a one-shot HTTP server on `127.0.0.1:0` that answers `GET
2911    /// /models` with an OpenAI-shaped list after `delay`. Returns its base URL.
2912    fn spawn_models_server(delay: std::time::Duration) -> String {
2913        use std::io::{Read as _, Write as _};
2914        use std::net::TcpListener;
2915        let listener = TcpListener::bind("127.0.0.1:0").expect("bind ephemeral port");
2916        let addr = listener.local_addr().expect("local addr");
2917        std::thread::spawn(move || {
2918            // Serve connections until the process ends: the client may open a
2919            // fresh connection per request, and a timed-out probe leaves one
2920            // half-read behind.
2921            for stream in listener.incoming() {
2922                let Ok(mut stream) = stream else { continue };
2923                // Read the request head (through `\r\n\r\n`) so the client's
2924                // write completes before we answer — a single `read` can race
2925                // the request into multiple segments.
2926                let mut buf = Vec::new();
2927                let mut chunk = [0u8; 1024];
2928                loop {
2929                    match stream.read(&mut chunk) {
2930                        Ok(0) | Err(_) => break,
2931                        Ok(n) => {
2932                            buf.extend_from_slice(&chunk[..n]);
2933                            if buf.windows(4).any(|w| w == b"\r\n\r\n") {
2934                                break;
2935                            }
2936                        }
2937                    }
2938                }
2939                std::thread::sleep(delay);
2940                let body = "{\"object\":\"list\",\"data\":[{\"id\":\"m1\",\"object\":\"model\"}]}";
2941                let response = format!(
2942                    "HTTP/1.1 200 OK\r\n\
2943                     Content-Type: application/json\r\n\
2944                     Connection: close\r\n\
2945                     Content-Length: {}\r\n\r\n{}",
2946                    body.len(),
2947                    body
2948                );
2949                let _ = stream.write_all(response.as_bytes());
2950                let _ = stream.flush();
2951            }
2952        });
2953        format!("http://{addr}/v1")
2954    }
2955
2956    #[tokio::test]
2957    async fn a_refused_port_probes_false() {
2958        let up = probe_endpoint(
2959            "http://127.0.0.1:9/v1",
2960            std::time::Duration::from_millis(500),
2961        )
2962        .await;
2963        assert!(!up, "a refused port is not an endpoint");
2964    }
2965
2966    #[tokio::test]
2967    async fn a_live_models_endpoint_probes_true() {
2968        // rig's HTTP client pays a one-time, process-wide init on its very first
2969        // request (seconds on macOS, where system-proxy detection is slow) — it
2970        // happens BEFORE any connection, so a refused port triggers it cheaply.
2971        // Pre-warm so the assertion below is not racing that one-time cost.
2972        let _ = probe_endpoint("http://127.0.0.1:9/v1", std::time::Duration::from_secs(60)).await;
2973        let base = spawn_models_server(std::time::Duration::ZERO);
2974        let up = probe_endpoint(&base, std::time::Duration::from_secs(5)).await;
2975        assert!(up, "a live /models endpoint probes true");
2976    }
2977
2978    #[tokio::test]
2979    async fn a_slow_endpoint_times_out_to_false() {
2980        let base = spawn_models_server(std::time::Duration::from_millis(500));
2981        let up = probe_endpoint(&base, std::time::Duration::from_millis(50)).await;
2982        assert!(!up, "a slow endpoint is not treated as up");
2983    }
2984}