1use std::pin::Pin;
9use std::sync::{Arc, Mutex};
10use std::time::{Duration, SystemTime, UNIX_EPOCH};
11
12use futures::{Stream, StreamExt};
13use rig::client::{CompletionClient, ModelListingClient};
14use rig::completion::{
15 CompletionError, CompletionModel, CompletionRequest, FinishReason, ToolDefinition,
16};
17use rig::http_client::{HeaderMap, HeaderValue};
18use rig::message::{
19 AssistantContent, Message, Reasoning, ReasoningContent, ToolCall, ToolCallId, ToolFunction,
20 UserContent,
21};
22use rig::providers::openai;
23use rig::streaming::{StreamedAssistantContent, StreamingCompletionResponse};
24
25use crate::event::LlmStreamEvent;
26use crate::message::{AgentMessage, ContentBlock, StopReason};
27
28#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
29pub enum LlmEndpoint {
30 #[default]
31 Chat,
32 Responses,
33}
34
35#[derive(Clone, Debug, Default)]
36pub struct LlmOpts {
37 pub model: String,
38 pub base_url: Option<String>,
39 pub api_key: Option<String>,
40 pub temperature: Option<f64>,
43 pub endpoint: LlmEndpoint,
44 pub session_id: Option<String>,
47 pub effort: Option<String>,
54 pub retry: RetryPolicy,
57 pub model_profiles: ModelProfiles,
62}
63
64#[derive(Clone, Debug, Default, PartialEq, Eq)]
69pub struct LlmProfile {
70 pub endpoint: Option<LlmEndpoint>,
74 pub base_url: Option<String>,
76 pub api_key: Option<String>,
79}
80
81impl LlmProfile {
82 fn apply(&self, llm: &mut LlmOpts) {
89 if let Some(endpoint) = self.endpoint {
90 llm.endpoint = endpoint;
91 }
92 if let Some(base_url) = &self.base_url {
93 llm.base_url = Some(base_url.clone());
94 }
95 if let Some(api_key) = &self.api_key {
96 llm.api_key = Some(api_key.clone());
97 }
98 }
99}
100
101#[derive(Clone, Debug, PartialEq, Eq)]
105pub struct LlmProvider {
106 pub endpoint: LlmEndpoint,
107 pub base_url: Option<String>,
108 pub api_key: Option<String>,
109}
110
111impl LlmProvider {
112 pub fn of(llm: &LlmOpts) -> Self {
114 Self {
115 endpoint: llm.endpoint,
116 base_url: llm.base_url.clone(),
117 api_key: llm.api_key.clone(),
118 }
119 }
120
121 fn apply(&self, llm: &mut LlmOpts) {
123 llm.endpoint = self.endpoint;
124 llm.base_url = self.base_url.clone();
125 llm.api_key = self.api_key.clone();
126 }
127}
128
129#[derive(Clone, Debug, Default, PartialEq, Eq)]
140pub struct ModelProfiles {
141 base: Option<LlmProvider>,
143 profiles: std::collections::BTreeMap<String, LlmProfile>,
144}
145
146impl ModelProfiles {
147 pub fn new(base: LlmProvider) -> Self {
149 Self {
150 base: Some(base),
151 profiles: std::collections::BTreeMap::new(),
152 }
153 }
154
155 fn resolve(&self, model: &str) -> Option<&LlmProfile> {
157 self.profiles.get(model)
158 }
159
160 pub fn settle(&self, llm: &mut LlmOpts) {
163 debug_assert!(
166 self.base.is_some() || self.profiles.is_empty(),
167 "ModelProfiles has profiles but no base (was `insert` called on the Default?)"
168 );
169 let Some(base) = &self.base else {
170 return;
171 };
172 base.apply(llm);
173 if let Some(profile) = self.resolve(&llm.model) {
174 profile.apply(llm);
175 }
176 }
177
178 pub fn base(&self) -> Option<&LlmProvider> {
182 self.base.as_ref()
183 }
184
185 pub fn insert(&mut self, model: impl Into<String>, profile: LlmProfile) {
187 self.profiles.insert(model.into(), profile);
188 debug_assert!(
191 self.base.is_some() || self.profiles.is_empty(),
192 "insert on a ModelProfiles with no base: the profile is dead (use `new`)"
193 );
194 }
195}
196
197impl LlmOpts {
198 pub fn settle_provider(&mut self) {
203 let profiles = self.model_profiles.clone();
204 profiles.settle(self);
205 }
206}
207
208pub type LlmStream = Pin<Box<dyn Stream<Item = LlmStreamEvent> + Send>>;
209
210pub type StreamFn =
211 Arc<dyn Fn(&[AgentMessage], &str, &[ToolDefinition], &LlmOpts) -> LlmStream + Send + Sync>;
212
213#[derive(Clone, PartialEq, Eq)]
219struct ClientKey {
220 base_url: Option<String>,
221 api_key: Option<String>,
222 session_id: Option<String>,
223}
224
225impl ClientKey {
226 fn of(opts: &LlmOpts) -> Self {
227 Self {
228 base_url: opts.base_url.clone(),
229 api_key: opts.api_key.clone(),
230 session_id: opts.session_id.clone(),
231 }
232 }
233}
234
235#[derive(Default)]
239struct ClientCache {
240 slot: Mutex<Option<(ClientKey, openai::Client)>>,
241 #[cfg(test)]
242 builds: std::sync::atomic::AtomicUsize,
243}
244
245impl ClientCache {
246 fn get(&self, opts: &LlmOpts) -> Result<openai::Client, String> {
248 let key = ClientKey::of(opts);
249 let mut slot = self.slot.lock().unwrap_or_else(|e| e.into_inner());
250 if let Some((cached_key, client)) = slot.as_ref()
251 && *cached_key == key
252 {
253 return Ok(client.clone());
254 }
255 let client = openai_client(opts)?;
256 #[cfg(test)]
257 self.builds
258 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
259 *slot = Some((key, client.clone()));
260 Ok(client)
261 }
262}
263
264pub fn rig_stream_fn() -> StreamFn {
269 let cache = ClientCache::default();
270 Arc::new(move |messages, system, tools, opts| adapt(&cache, opts, messages, system, tools))
271}
272
273fn openai_client(opts: &LlmOpts) -> Result<rig::providers::openai::Client, String> {
276 let key = opts
277 .api_key
278 .clone()
279 .or_else(|| std::env::var("OPENAI_API_KEY").ok());
280 let key = match key {
281 Some(k) => k,
282 None => {
283 let local = opts.base_url.as_deref().is_some_and(|url| {
287 let u = url.to_ascii_lowercase();
288 ["http://localhost", "http://127.0.0.1", "http://[::1]"]
289 .iter()
290 .any(|prefix| u.starts_with(prefix))
291 });
292 if local {
293 "wcode-local".to_string()
294 } else {
295 return Err(
296 "no API key: pass LlmOpts.api_key or set OPENAI_API_KEY (a localhost base_url works keyless)"
297 .to_string(),
298 );
299 }
300 }
301 };
302 let mut builder = openai::Client::builder()
303 .api_key::<rig::client::BearerAuth>(key)
304 .http_headers(session_headers(opts));
305 if let Some(base_url) = &opts.base_url {
306 builder = builder.base_url(base_url);
307 }
308 builder.build().map_err(|e| format!("openai client: {e}"))
309}
310
311fn session_headers(opts: &LlmOpts) -> HeaderMap {
314 let mut headers = HeaderMap::new();
315 headers.insert(
316 "user-agent",
317 HeaderValue::from_static(concat!("wcode/", env!("CARGO_PKG_VERSION"))),
318 );
319 if let Some(id) = opts.session_id.as_deref()
320 && let Ok(value) = HeaderValue::from_str(id)
321 {
322 headers.insert("x-opencode-session", value);
323 }
324 headers
325}
326
327pub async fn probe_endpoint(base_url: &str, timeout: std::time::Duration) -> bool {
337 let opts = LlmOpts {
338 base_url: Some(base_url.to_string()),
339 ..LlmOpts::default()
340 };
341 matches!(
342 tokio::time::timeout(timeout, list_models(&opts)).await,
343 Ok(Ok(_))
344 )
345}
346pub async fn list_models(opts: &LlmOpts) -> Result<Vec<String>, String> {
351 let client = openai_client(opts)?;
352 let list = client.list_models().await.map_err(|e| e.to_string())?;
353 let mut ids: Vec<String> = list.iter().map(|m| m.id.clone()).collect();
354 ids.sort();
355 Ok(ids)
356}
357
358fn adapt(
359 cache: &ClientCache,
360 opts: &LlmOpts,
361 messages: &[AgentMessage],
362 system: &str,
363 tools: &[ToolDefinition],
364) -> LlmStream {
365 if let Some(message) = invalid_effort(opts) {
366 return error_stream(message);
367 }
368 let client = match cache.get(opts) {
369 Ok(client) => client,
370 Err(message) => return error_stream(message),
371 };
372 let request = build_request(messages, system, tools, opts);
373
374 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<LlmStreamEvent>();
381 let Ok(runtime) = tokio::runtime::Handle::try_current() else {
382 return error_stream("StreamFn requires a tokio runtime".to_string());
383 };
384 let policy = opts.retry;
385 let endpoint = opts.endpoint;
386 let model = opts.model.clone();
387 runtime.spawn(async move {
388 let open = || open_stream(client.clone(), endpoint, &model, request.clone());
389 let mut state = RetryState::new(policy.max);
394 let (stream, pending) = match connect_with_retry(&policy, &mut state, &open, &tx).await {
395 PreContent::Ready { stream, first } => (stream, first),
396 PreContent::Cancelled => return,
397 PreContent::Failed(error) => {
398 let _ = tx.send(LlmStreamEvent::Error {
399 message: error.to_string(),
400 fatal: fatal_class(&error),
401 });
402 return;
403 }
404 };
405 forward_stream(stream, pending, policy, &open, tx, state).await;
406 });
407 Box::pin(futures::stream::unfold(rx, |mut rx| async move {
408 rx.recv().await.map(|ev| (ev, rx))
409 }))
410}
411
412async fn forward_stream<S, F>(
418 mut stream: S,
419 mut pending: Option<Result<StreamedAssistantContent, CompletionError>>,
420 policy: RetryPolicy,
421 open: F,
422 tx: tokio::sync::mpsc::UnboundedSender<LlmStreamEvent>,
423 mut state: RetryState,
424) where
425 F: Fn() -> Pin<Box<dyn std::future::Future<Output = Result<S, CompletionError>> + Send>>,
426 S: futures::Stream<Item = Result<StreamedAssistantContent, CompletionError>> + Unpin + Send,
427{
428 let mut errored = false;
429 let mut content_forwarded = false;
430 let mut done_seen = false;
431 loop {
432 let item = match pending.take() {
438 Some(item) => item,
439 None => tokio::select! {
440 biased;
441 result = next_or_stall(&mut stream, policy.idle) => match result {
442 None => {
448 if done_seen || errored {
449 break;
450 }
451 let error = truncation_error();
452 if !content_forwarded && state.budget > 0 {
453 match pre_content_retry(&policy, &mut state, &open, &tx, &error).await {
456 PreContent::Ready { stream: s, first } => {
457 stream = s;
458 pending = first;
459 continue;
460 }
461 PreContent::Cancelled => return,
462 PreContent::Failed(e) => {
463 if tx
464 .send(LlmStreamEvent::Error {
465 message: e.to_string(),
466 fatal: fatal_class(&e),
467 })
468 .is_err()
469 {
470 return;
471 }
472 break;
473 }
474 }
475 } else {
476 if tx
480 .send(LlmStreamEvent::Error {
481 message: error.to_string(),
482 fatal: false,
483 })
484 .is_err()
485 {
486 return;
487 }
488 break;
489 }
490 }
491 Some(Stall::Item(item)) => item,
493 Some(Stall::Idle) => {
498 let error = stall_error(policy.idle);
499 if !content_forwarded && state.budget > 0 {
500 match pre_content_retry(&policy, &mut state, &open, &tx, &error).await {
501 PreContent::Ready { stream: s, first } => {
502 stream = s;
503 pending = first;
504 continue;
505 }
506 PreContent::Cancelled => return,
507 PreContent::Failed(e) => {
508 let fatal = fatal_class(&e);
509 if tx
510 .send(LlmStreamEvent::Error {
511 message: e.to_string(),
512 fatal,
513 })
514 .is_err()
515 {
516 return;
517 }
518 break;
519 }
520 }
521 } else {
522 if tx
523 .send(LlmStreamEvent::Error {
524 message: error.to_string(),
525 fatal: false,
526 })
527 .is_err()
528 {
529 return;
530 }
531 break;
532 }
533 }
534 },
535 _ = tx.closed() => break,
536 },
537 };
538 let events = match item {
539 Ok(content) => {
540 let is_final = matches!(content, StreamedAssistantContent::Final(_));
541 if errored && is_final {
542 continue;
543 }
544 let events = map_item(content);
545 content_forwarded |= !events.is_empty();
549 done_seen |= is_final;
550 events
551 }
552 Err(error) => {
553 if !content_forwarded && state.budget > 0 && retryable(&error) {
554 state.attempts += 1;
558 state.budget -= 1;
559 if !retry_wait(&tx, &policy, state.attempts, &error).await {
560 return;
561 }
562 match connect_with_retry(&policy, &mut state, &open, &tx).await {
563 PreContent::Ready { stream: s, first } => {
564 stream = s;
565 pending = first;
566 continue;
567 }
568 PreContent::Cancelled => return,
569 PreContent::Failed(e) => {
570 errored = true;
571 vec![LlmStreamEvent::Error {
572 message: e.to_string(),
573 fatal: fatal_class(&e),
574 }]
575 }
576 }
577 } else {
578 errored = true;
582 vec![LlmStreamEvent::Error {
583 message: error.to_string(),
584 fatal: fatal_class(&error),
585 }]
586 }
587 }
588 };
589 for ev in events {
590 if tx.send(ev).is_err() {
591 return;
592 }
593 }
594 }
595}
596
597fn error_stream(message: String) -> LlmStream {
598 Box::pin(futures::stream::iter(vec![LlmStreamEvent::Error {
601 message,
602 fatal: true,
603 }]))
604}
605
606#[derive(Clone, Copy, Debug, PartialEq, Eq)]
609pub struct RetryPolicy {
610 pub max: u32,
612 pub base: Duration,
614 pub cap: Duration,
616 pub ttft: Duration,
618 pub idle: Duration,
620}
621
622impl Default for RetryPolicy {
623 fn default() -> Self {
624 Self {
625 max: 3,
626 base: Duration::from_millis(500),
627 cap: Duration::from_secs(8),
628 ttft: Duration::from_secs(60),
629 idle: Duration::from_secs(120),
630 }
631 }
632}
633
634enum Stall<T> {
636 Item(T),
639 Idle,
641}
642
643async fn next_or_stall<S: Stream + Unpin>(
647 stream: &mut S,
648 deadline: Duration,
649) -> Option<Stall<S::Item>> {
650 if deadline.is_zero() {
651 return stream.next().await.map(Stall::Item);
652 }
653 tokio::select! {
654 biased;
655 item = stream.next() => item.map(Stall::Item),
656 _ = tokio::time::sleep(deadline) => Some(Stall::Idle),
657 }
658}
659fn is_transient_status(code: u16) -> bool {
662 matches!(code, 408 | 409 | 429 | 500 | 502 | 503 | 504)
663}
664
665fn retryable(error: &CompletionError) -> bool {
669 if let Some(status) = error.provider_response_status() {
670 return is_transient_status(status.as_u16());
671 }
672 match error {
673 CompletionError::HttpError(_) => true,
675 CompletionError::ProviderError(message) => is_transient_message(message),
679 _ => false,
680 }
681}
682
683fn is_transient_message(message: &str) -> bool {
686 const NEEDLES: [&str; 10] = [
687 "error sending request",
688 "error trying to connect",
689 "connection refused",
690 "connection reset",
691 "connection closed",
692 "timed out",
693 "timeout",
694 "unreachable",
695 "dns error",
696 "temporary failure",
697 ];
698 let m = message.to_ascii_lowercase();
699 NEEDLES.iter().any(|needle| m.contains(needle))
700}
701
702fn backoff(attempt: u32, base: Duration, cap: Duration) -> Duration {
705 let shift = attempt.saturating_sub(1).min(16);
706 let capped = base.saturating_mul(1u32 << shift).min(cap);
707 let millis = capped.as_millis() as u64;
708 if millis <= 1 {
709 return capped;
710 }
711 let jitter = clock_nanos() % (millis / 4 + 1);
712 Duration::from_millis(millis - jitter)
713}
714
715fn clock_nanos() -> u64 {
717 SystemTime::now()
718 .duration_since(UNIX_EPOCH)
719 .map(|d| u64::from(d.subsec_nanos()))
720 .unwrap_or(0)
721}
722
723fn open_stream(
726 client: rig::providers::openai::Client,
727 endpoint: LlmEndpoint,
728 model: &str,
729 request: CompletionRequest,
730) -> Pin<Box<dyn std::future::Future<Output = Result<StreamingCompletionResponse, CompletionError>> + Send>> {
731 match endpoint {
732 LlmEndpoint::Chat => {
733 let model = client.completions_api().completion_model(model.to_string());
734 Box::pin(async move { model.stream(request).await })
735 }
736 LlmEndpoint::Responses => {
737 let model = client.completion_model(model.to_string());
738 Box::pin(async move { model.stream(request).await })
739 }
740 }
741}
742
743enum PreContent<S> {
747 Ready {
748 stream: S,
749 first: Option<Result<StreamedAssistantContent, CompletionError>>,
750 },
751 Failed(CompletionError),
753 Cancelled,
755}
756
757#[derive(Clone, Copy)]
761struct RetryState {
762 budget: u32,
763 attempts: u32,
764}
765
766impl RetryState {
767 fn new(max: u32) -> Self {
768 RetryState {
769 budget: max,
770 attempts: 0,
771 }
772 }
773}
774
775fn fatal_class(error: &CompletionError) -> bool {
782 !retryable(error)
783}
784
785fn stall_error(d: Duration) -> CompletionError {
799 CompletionError::ProviderError(format!("stream stalled: no item within {d:?} (timed out)"))
800}
801fn truncation_error() -> CompletionError {
808 CompletionError::ProviderError(
809 "stream ended before its terminal record (truncated)".to_string(),
810 )
811}
812
813async fn connect_with_retry<S, F>(
814 policy: &RetryPolicy,
815 state: &mut RetryState,
816 mut open: F,
817 tx: &tokio::sync::mpsc::UnboundedSender<LlmStreamEvent>,
818) -> PreContent<S>
819where
820 F: FnMut() -> Pin<Box<dyn std::future::Future<Output = Result<S, CompletionError>> + Send>>,
821 S: futures::Stream<Item = Result<StreamedAssistantContent, CompletionError>> + Unpin + Send,
822{
823 loop {
824 state.attempts += 1;
825 let opened = if policy.ttft.is_zero() {
829 Ok(open().await)
830 } else {
831 tokio::select! {
832 biased;
833 r = open() => Ok(r),
834 _ = tokio::time::sleep(policy.ttft) => Err(()),
835 }
836 };
837 let mut stream = match opened {
838 Ok(Ok(stream)) => stream,
839 Ok(Err(error)) => {
840 if state.budget > 0 && retryable(&error) {
841 state.budget -= 1;
842 if !retry_wait(tx, policy, state.attempts, &error).await {
843 return PreContent::Cancelled;
844 }
845 continue;
846 }
847 return PreContent::Failed(error);
848 }
849 Err(()) => {
850 let error = stall_error(policy.ttft);
851 if state.budget > 0 {
852 state.budget -= 1;
853 if !retry_wait(tx, policy, state.attempts, &error).await {
854 return PreContent::Cancelled;
855 }
856 continue;
857 }
858 return PreContent::Failed(error);
859 }
860 };
861 match next_or_stall(&mut stream, policy.ttft).await {
864 None => return PreContent::Ready { stream, first: None },
865 Some(Stall::Item(Err(error))) => {
866 if state.budget > 0 && retryable(&error) {
867 state.budget -= 1;
868 if !retry_wait(tx, policy, state.attempts, &error).await {
869 return PreContent::Cancelled;
870 }
871 continue;
872 }
873 return PreContent::Ready {
877 stream,
878 first: Some(Err(error)),
879 };
880 }
881 Some(Stall::Item(Ok(content))) => {
882 return PreContent::Ready {
883 stream,
884 first: Some(Ok(content)),
885 };
886 }
887 Some(Stall::Idle) => {
888 let error = stall_error(policy.ttft);
889 if state.budget > 0 {
890 state.budget -= 1;
891 if !retry_wait(tx, policy, state.attempts, &error).await {
892 return PreContent::Cancelled;
893 }
894 continue;
895 }
896 return PreContent::Failed(error);
897 }
898 }
899 }
900}
901
902async fn pre_content_retry<S, F>(
909 policy: &RetryPolicy,
910 state: &mut RetryState,
911 open: &F,
912 tx: &tokio::sync::mpsc::UnboundedSender<LlmStreamEvent>,
913 error: &CompletionError,
914) -> PreContent<S>
915where
916 F: Fn() -> Pin<Box<dyn std::future::Future<Output = Result<S, CompletionError>> + Send>>,
917 S: futures::Stream<Item = Result<StreamedAssistantContent, CompletionError>> + Unpin + Send,
918{
919 state.attempts += 1;
920 state.budget -= 1;
921 if !retry_wait(tx, policy, state.attempts, error).await {
922 return PreContent::Cancelled;
923 }
924 connect_with_retry(policy, state, open, tx).await
925}
926
927fn retry_after(error: &CompletionError) -> Option<Duration> {
937 let raw = error.provider_response_headers()?.get("retry-after")?;
941 let seconds: u64 = raw.to_str().ok()?.trim().parse().ok()?;
943 Some(Duration::from_secs(seconds))
944}
945
946async fn retry_wait(
949 tx: &tokio::sync::mpsc::UnboundedSender<LlmStreamEvent>,
950 policy: &RetryPolicy,
951 attempt: u32,
952 error: &CompletionError,
953) -> bool {
954 let delay = retry_after(error)
958 .map(|d| d.min(policy.cap))
959 .unwrap_or_else(|| backoff(attempt, policy.base, policy.cap));
960 let _ = tx.send(LlmStreamEvent::Retrying {
961 attempt,
962 max: policy.max,
963 reason: error.to_string(),
964 });
965 tokio::select! {
966 biased;
967 _ = tx.closed() => false,
968 _ = tokio::time::sleep(delay) => true,
969 }
970}
971
972fn build_request(
973 messages: &[AgentMessage],
974 system: &str,
975 tools: &[ToolDefinition],
976 opts: &LlmOpts,
977) -> CompletionRequest {
978 let mut chat_history = Vec::with_capacity(messages.len() + 1);
979 if !system.is_empty() {
980 chat_history.push(Message::system(system));
981 }
982 chat_history.extend(messages.iter().filter_map(to_rig_message));
983 CompletionRequest {
984 model: None,
985 preamble: None,
988 chat_history,
989 documents: Vec::new(),
990 tools: tools.to_vec(),
991 temperature: opts.temperature,
992 max_tokens: None,
993 tool_choice: None,
994 additional_params: reasoning_params(opts),
995 output_schema: None,
996 record_telemetry_content: false,
997 }
998}
999
1000const RESPONSES_EFFORT_LEVELS: &[&str] =
1005 &["none", "minimal", "low", "medium", "high", "xhigh", "max"];
1006
1007fn invalid_effort(opts: &LlmOpts) -> Option<String> {
1011 if opts.endpoint == LlmEndpoint::Responses
1012 && let Some(effort) = opts.effort.as_deref()
1013 && !RESPONSES_EFFORT_LEVELS.contains(&effort)
1014 {
1015 return Some(format!(
1016 "effort `{effort}` is not valid for the Responses endpoint; use one of: {}",
1017 RESPONSES_EFFORT_LEVELS.join(", ")
1018 ));
1019 }
1020 None
1021}
1022
1023fn reasoning_params(opts: &LlmOpts) -> Option<serde_json::Value> {
1040 match opts.endpoint {
1041 LlmEndpoint::Chat => opts
1042 .effort
1043 .as_ref()
1044 .map(|effort| serde_json::json!({ "reasoning_effort": effort })),
1045 LlmEndpoint::Responses => {
1046 let mut reasoning = serde_json::Map::new();
1047 reasoning.insert("summary".into(), serde_json::json!("auto"));
1048 if let Some(effort) = &opts.effort {
1049 reasoning.insert("effort".into(), serde_json::json!(effort));
1050 }
1051 Some(serde_json::json!({ "reasoning": reasoning }))
1052 }
1053 }
1054}
1055
1056const TOOL_OUTPUT_ERROR_PREFIX: &str = "ERROR: ";
1063
1064fn to_rig_message(m: &AgentMessage) -> Option<Message> {
1065 match m {
1066 AgentMessage::User { content } => {
1067 let content: Vec<UserContent> = content
1068 .iter()
1069 .filter_map(|b| match b {
1070 ContentBlock::Text { text } => Some(UserContent::text(text.clone())),
1071 _ => None,
1072 })
1073 .collect();
1074 (!content.is_empty()).then_some(Message::User { content })
1075 }
1076 AgentMessage::Assistant { content, .. } => {
1077 let content: Vec<AssistantContent> = content
1078 .iter()
1079 .map(|b| match b {
1080 ContentBlock::Text { text } => AssistantContent::text(text.clone()),
1081 ContentBlock::Thinking { text } => {
1082 AssistantContent::Reasoning(Reasoning::new(text))
1083 }
1084 ContentBlock::ToolCall {
1085 id,
1086 name,
1087 arguments,
1088 } => AssistantContent::ToolCall(ToolCall::new(
1089 ToolCallId::new_or_mint(id.clone()),
1090 ToolFunction::new(name.clone(), arguments.clone()),
1091 )),
1092 })
1093 .collect();
1094 (!content.is_empty()).then_some(Message::Assistant { id: None, content })
1095 }
1096 AgentMessage::ToolResult {
1097 tool_call_id,
1098 name,
1099 output,
1100 is_error,
1101 } => {
1102 let output = if *is_error {
1106 format!("{TOOL_OUTPUT_ERROR_PREFIX}{output}")
1107 } else {
1108 output.clone()
1109 };
1110 Some(Message::tool_result(
1111 tool_call_id.clone(),
1112 name.clone(),
1113 output,
1114 ))
1115 }
1116 }
1117}
1118
1119fn map_item(item: StreamedAssistantContent) -> Vec<LlmStreamEvent> {
1120 match item {
1121 StreamedAssistantContent::Text(text) => vec![LlmStreamEvent::TextDelta(text.text)],
1122 StreamedAssistantContent::ReasoningDelta { reasoning, .. } => {
1123 vec![LlmStreamEvent::ThinkingDelta(reasoning)]
1124 }
1125 StreamedAssistantContent::Reasoning { reasoning, .. } => {
1130 let text = reasoning_text(&reasoning);
1133 if text.is_empty() {
1134 vec![]
1139 } else {
1140 vec![LlmStreamEvent::ThinkingReplace(text)]
1141 }
1142 }
1143 StreamedAssistantContent::ToolCallDelta { .. } => vec![],
1150 StreamedAssistantContent::ToolCall { tool_call, .. } => vec![LlmStreamEvent::ToolCall {
1151 id: tool_call.id.as_str().to_string(),
1152 name: tool_call.function.name,
1153 arguments: tool_call.function.arguments,
1154 }],
1155 StreamedAssistantContent::Final(final_record) => {
1156 vec![LlmStreamEvent::Done {
1157 stop_reason: final_record
1158 .finish_reason
1159 .map(finish_reason)
1160 .unwrap_or(StopReason::Stop),
1161 usage: Some(convert_usage(final_record.usage)),
1162 }]
1163 }
1164 StreamedAssistantContent::Unknown(_) => vec![],
1165 }
1166}
1167
1168fn finish_reason(reason: FinishReason) -> StopReason {
1169 match reason {
1170 FinishReason::Stop => StopReason::Stop,
1171 FinishReason::Length => StopReason::Length,
1172 FinishReason::ToolCalls => StopReason::ToolUse,
1173 FinishReason::ContentFilter | FinishReason::Other(_) => StopReason::Stop,
1174 }
1175}
1176
1177fn convert_usage(usage: rig::completion::Usage) -> crate::message::Usage {
1178 crate::message::Usage {
1179 input_tokens: usage.input_tokens,
1180 output_tokens: usage.output_tokens,
1181 cache_read_tokens: (usage.cached_input_tokens > 0).then_some(usage.cached_input_tokens),
1183 cache_write_tokens: (usage.cache_creation_input_tokens > 0)
1184 .then_some(usage.cache_creation_input_tokens),
1185 }
1186}
1187
1188fn reasoning_text(reasoning: &Reasoning) -> String {
1189 reasoning
1190 .content
1191 .iter()
1192 .filter_map(|block| match block {
1193 ReasoningContent::Text { text, .. } => Some(text.as_str()),
1194 ReasoningContent::Summary(summary) => Some(summary.as_str()),
1195 ReasoningContent::Redacted { data } => Some(data.as_str()),
1196 ReasoningContent::Encrypted(_) => None,
1197 })
1198 .collect::<Vec<_>>()
1199 .join("\n")
1200}
1201
1202#[cfg(test)]
1203mod tests {
1204 use super::*;
1205 use rig::completion::Usage as RigUsage;
1206 use rig::message::{ToolResultContent, UserContent};
1207 use rig::streaming::{RawStreamingToolCall, StreamFinal, ToolCallDeltaContent};
1208 use serde_json::json;
1209
1210 fn tool_call_item() -> StreamedAssistantContent {
1211 let raw =
1212 RawStreamingToolCall::new("call_1".to_string(), "run".to_string(), json!({"x": 1}));
1213 StreamedAssistantContent::ToolCall {
1214 tool_call: raw.into(),
1215 internal_call_id: "ic1".to_string(),
1216 }
1217 }
1218
1219 fn final_item(finish_reason: Option<FinishReason>) -> StreamedAssistantContent {
1220 let usage = RigUsage {
1221 input_tokens: 10,
1222 output_tokens: 5,
1223 total_tokens: 15,
1224 cached_input_tokens: 3,
1225 cache_creation_input_tokens: 2,
1226 ..RigUsage::default()
1227 };
1228 let final_record = StreamFinal::new("test", usage);
1229 StreamedAssistantContent::Final(match finish_reason {
1230 Some(reason) => final_record.with_finish_reason(reason),
1231 None => final_record,
1232 })
1233 }
1234
1235 #[test]
1236 fn text_item_maps_to_text_delta() {
1237 let events = map_item(StreamedAssistantContent::Text(rig::message::Text::new(
1238 "hi",
1239 )));
1240 assert_eq!(events, vec![LlmStreamEvent::TextDelta("hi".to_string())]);
1241 }
1242
1243 #[test]
1244 fn reasoning_delta_maps_to_thinking_delta() {
1245 let events = map_item(StreamedAssistantContent::ReasoningDelta {
1246 id: "r1".to_string(),
1247 provider_id: None,
1248 reasoning: "hmm".to_string(),
1249 });
1250 assert_eq!(
1251 events,
1252 vec![LlmStreamEvent::ThinkingDelta("hmm".to_string())]
1253 );
1254 }
1255
1256 #[test]
1257 fn reasoning_block_replaces_with_full_text() {
1258 let events = map_item(StreamedAssistantContent::Reasoning {
1259 reasoning: Reasoning::new("full thought"),
1260 id: "r1".to_string(),
1261 });
1262 assert_eq!(
1263 events,
1264 vec![LlmStreamEvent::ThinkingReplace("full thought".to_string())]
1265 );
1266 }
1267
1268 #[test]
1269 fn reasoning_text_joins_summary_text_and_redacted_drops_encrypted() {
1270 let reasoning = Reasoning {
1273 id: None,
1274 content: vec![
1275 ReasoningContent::Summary("a".to_string()),
1276 ReasoningContent::Text {
1277 text: "b".to_string(),
1278 signature: None,
1279 },
1280 ReasoningContent::Redacted {
1281 data: "c".to_string(),
1282 },
1283 ReasoningContent::Encrypted("opaque".to_string()),
1284 ],
1285 };
1286 assert_eq!(reasoning_text(&reasoning), "a\nb\nc");
1287 }
1288
1289 #[test]
1290 fn reasoning_text_summary_only_is_the_summary() {
1291 let reasoning = Reasoning::summaries(vec!["just a summary".to_string()]);
1292 assert_eq!(reasoning_text(&reasoning), "just a summary");
1293 }
1294
1295 #[test]
1296 fn reasoning_text_encrypted_only_is_empty() {
1297 let reasoning = Reasoning::encrypted("opaque");
1298 assert_eq!(reasoning_text(&reasoning), "");
1299 }
1300
1301 #[test]
1302 fn summary_reasoning_block_replaces_with_full_text() {
1303 let events = map_item(StreamedAssistantContent::Reasoning {
1304 reasoning: Reasoning::summaries(vec!["a summary".to_string()]),
1305 id: "rs_1".to_string(),
1306 });
1307 assert_eq!(
1308 events,
1309 vec![LlmStreamEvent::ThinkingReplace("a summary".to_string())]
1310 );
1311 }
1312
1313 #[test]
1314 fn encrypted_only_reasoning_block_emits_no_event() {
1315 let events = map_item(StreamedAssistantContent::Reasoning {
1319 reasoning: Reasoning::encrypted("opaque"),
1320 id: "rs_1".to_string(),
1321 });
1322 assert!(events.is_empty(), "got: {events:?}");
1323 }
1324
1325 #[test]
1326 fn tool_call_delta_ignored() {
1327 for content in [
1328 ToolCallDeltaContent::Name("run".to_string()),
1329 ToolCallDeltaContent::Delta("{\"x".to_string()),
1330 ] {
1331 let events = map_item(StreamedAssistantContent::ToolCallDelta {
1332 internal_call_id: "ic1".to_string(),
1333 content,
1334 });
1335 assert!(events.is_empty());
1336 }
1337 }
1338
1339 #[test]
1340 fn tool_call_maps_to_call_event() {
1341 let events = map_item(tool_call_item());
1342 assert_eq!(
1343 events,
1344 vec![LlmStreamEvent::ToolCall {
1345 id: "call_1".to_string(),
1346 name: "run".to_string(),
1347 arguments: json!({"x": 1}),
1348 }]
1349 );
1350 }
1351
1352 #[test]
1353 fn final_maps_finish_reason_and_usage() {
1354 let cases = [
1355 (FinishReason::Stop, StopReason::Stop),
1356 (FinishReason::Length, StopReason::Length),
1357 (FinishReason::ToolCalls, StopReason::ToolUse),
1358 (FinishReason::ContentFilter, StopReason::Stop),
1359 (FinishReason::Other("weird".to_string()), StopReason::Stop),
1360 ];
1361 for (reason, expected) in cases {
1362 let events = map_item(final_item(Some(reason)));
1363 assert_eq!(
1364 events,
1365 vec![LlmStreamEvent::Done {
1366 stop_reason: expected,
1367 usage: Some(crate::message::Usage {
1368 input_tokens: 10,
1369 output_tokens: 5,
1370 cache_read_tokens: Some(3),
1371 cache_write_tokens: Some(2),
1372 }),
1373 }],
1374 "wrong mapping for {expected:?}"
1375 );
1376 }
1377 }
1378
1379 #[test]
1380 fn final_without_finish_reason_defaults_to_stop() {
1381 let events = map_item(final_item(None));
1382 match &events[..] {
1383 [LlmStreamEvent::Done { stop_reason, .. }] => {
1384 assert_eq!(*stop_reason, StopReason::Stop)
1385 }
1386 other => panic!("unexpected events: {other:?}"),
1387 }
1388 }
1389
1390 #[test]
1391 fn unknown_item_ignored() {
1392 let events = map_item(StreamedAssistantContent::Unknown(
1393 rig::streaming::UnknownPayload::new(json!({"mystery": 1})),
1394 ));
1395 assert!(events.is_empty());
1396 }
1397
1398 #[test]
1399 fn build_request_round_trips_context() {
1400 let opts = LlmOpts {
1401 model: "m1".to_string(),
1402 base_url: None,
1403 api_key: None,
1404 temperature: Some(0.5),
1405 ..LlmOpts::default()
1406 };
1407 let tools = vec![ToolDefinition {
1408 name: "get_weather".to_string(),
1409 description: "Weather lookup".to_string(),
1410 parameters: json!({"type": "object"}),
1411 }];
1412 let messages = vec![
1413 AgentMessage::user_text("What is the weather?"),
1414 AgentMessage::Assistant {
1415 content: vec![
1416 ContentBlock::Thinking { text: "hmm".into() },
1417 ContentBlock::ToolCall {
1418 id: "t1".into(),
1419 name: "get_weather".into(),
1420 arguments: json!({"city": "Tokyo"}),
1421 },
1422 ContentBlock::Text {
1423 text: "Let me check.".into(),
1424 },
1425 ],
1426 stop_reason: StopReason::ToolUse,
1427 usage: None,
1428 model: None,
1429 },
1430 AgentMessage::ToolResult {
1431 tool_call_id: "t1".into(),
1432 name: "get_weather".into(),
1433 output: "sunny".into(),
1434 is_error: false,
1435 },
1436 ];
1437
1438 let request = build_request(&messages, "You are terse.", &tools, &opts);
1439
1440 assert!(request.preamble.is_none());
1441 assert_eq!(request.temperature, Some(0.5));
1442 assert_eq!(request.tools.len(), 1);
1443 assert_eq!(request.tools[0].name, "get_weather");
1444
1445 let [
1446 Message::System { content: system },
1447 Message::User { content: user },
1448 Message::Assistant {
1449 content: assistant, ..
1450 },
1451 Message::User {
1452 content: tool_result,
1453 },
1454 ] = &request.chat_history[..]
1455 else {
1456 panic!("unexpected chat_history shape: {:?}", request.chat_history);
1457 };
1458 assert_eq!(system, "You are terse.");
1459 assert!(
1460 matches!(user.as_slice(), [UserContent::Text(t)] if t.text == "What is the weather?")
1461 );
1462 assert_eq!(assistant.len(), 3);
1463 assert!(
1464 matches!(&assistant[0], AssistantContent::Reasoning(r) if reasoning_text(r) == "hmm")
1465 );
1466 assert!(matches!(
1467 &assistant[1],
1468 AssistantContent::ToolCall(call)
1469 if call.id.as_str() == "t1"
1470 && call.function.name == "get_weather"
1471 && call.function.arguments == json!({"city": "Tokyo"})
1472 ));
1473 assert!(matches!(&assistant[2], AssistantContent::Text(t) if t.text == "Let me check."));
1474 assert!(matches!(
1475 tool_result.as_slice(),
1476 [UserContent::ToolResult(r)] if r.call.as_str() == "t1" && r.name == "get_weather"
1477 ));
1478 }
1479
1480 #[test]
1481 fn tool_result_error_is_annotated_on_the_wire_only() {
1482 let err = AgentMessage::ToolResult {
1483 tool_call_id: "t1".into(),
1484 name: "echo".into(),
1485 output: "boom".into(),
1486 is_error: true,
1487 };
1488 let ok = AgentMessage::ToolResult {
1489 tool_call_id: "t1".into(),
1490 name: "echo".into(),
1491 output: "boom".into(),
1492 is_error: false,
1493 };
1494
1495 fn wire_text(m: &AgentMessage) -> String {
1496 match to_rig_message(m) {
1497 Some(Message::User { content }) => match &content[..] {
1498 [UserContent::ToolResult(r)] => match &r.content[..] {
1499 [ToolResultContent::Text(t)] => t.text.clone(),
1500 other => panic!("unexpected tool result content: {other:?}"),
1501 },
1502 other => panic!("unexpected user content: {other:?}"),
1503 },
1504 other => panic!("unexpected wire message: {other:?}"),
1505 }
1506 }
1507
1508 assert_eq!(wire_text(&ok), "boom");
1510 assert_eq!(wire_text(&err), "ERROR: boom");
1511 match err {
1514 AgentMessage::ToolResult {
1515 output, is_error, ..
1516 } => {
1517 assert_eq!(output, "boom");
1518 assert!(is_error);
1519 }
1520 _ => unreachable!(),
1521 }
1522 }
1523
1524 #[test]
1525 fn endpoint_defaults_to_chat() {
1526 let opts = LlmOpts {
1527 model: "m1".to_string(),
1528 base_url: None,
1529 api_key: None,
1530 temperature: None,
1531 endpoint: LlmEndpoint::Chat,
1532 effort: None,
1533 session_id: None,
1534 retry: RetryPolicy::default(),
1535 model_profiles: ModelProfiles::default(),
1536 };
1537 assert_eq!(opts.endpoint, LlmEndpoint::Chat);
1538 assert_eq!(LlmOpts::default().endpoint, LlmEndpoint::Chat);
1539 }
1540
1541 #[test]
1542 fn no_effort_sends_no_additional_params() {
1543 let opts = LlmOpts::default();
1544 let request = build_request(&[], "sys", &[], &opts);
1545 assert_eq!(request.additional_params, None);
1546 }
1547
1548 #[test]
1549 fn chat_effort_maps_to_reasoning_effort_param() {
1550 let opts = LlmOpts {
1551 effort: Some("high".into()),
1552 endpoint: LlmEndpoint::Chat,
1553 ..LlmOpts::default()
1554 };
1555 let request = build_request(&[], "sys", &[], &opts);
1556 assert_eq!(
1557 request.additional_params,
1558 Some(json!({ "reasoning_effort": "high" }))
1559 );
1560 }
1561
1562 #[test]
1563 fn responses_effort_maps_to_reasoning_object() {
1564 let opts = LlmOpts {
1565 effort: Some("xhigh".into()),
1566 endpoint: LlmEndpoint::Responses,
1567 ..LlmOpts::default()
1568 };
1569 let request = build_request(&[], "sys", &[], &opts);
1570 assert_eq!(
1571 request.additional_params,
1572 Some(json!({ "reasoning": { "effort": "xhigh", "summary": "auto" } }))
1573 );
1574 }
1575
1576 #[test]
1577 fn responses_always_requests_a_reasoning_summary() {
1578 let opts = LlmOpts {
1582 endpoint: LlmEndpoint::Responses,
1583 effort: None,
1584 ..LlmOpts::default()
1585 };
1586 let request = build_request(&[], "sys", &[], &opts);
1587 assert_eq!(
1588 request.additional_params,
1589 Some(json!({ "reasoning": { "summary": "auto" } }))
1590 );
1591 }
1592
1593 #[test]
1594 fn chat_effort_is_verbatim_passthrough() {
1595 for level in ["low", "medium", "my-custom-level", ""] {
1598 let opts = LlmOpts {
1599 effort: Some(level.into()),
1600 endpoint: LlmEndpoint::Chat,
1601 ..LlmOpts::default()
1602 };
1603 let request = build_request(&[], "sys", &[], &opts);
1604 assert_eq!(
1605 request.additional_params,
1606 Some(json!({ "reasoning_effort": level })),
1607 "level {level:?} must pass through verbatim on Chat"
1608 );
1609 }
1610 }
1611
1612 #[test]
1613 fn responses_effort_validation() {
1614 for level in ["none", "minimal", "low", "medium", "high", "xhigh", "max"] {
1617 let opts = LlmOpts {
1618 effort: Some(level.into()),
1619 endpoint: LlmEndpoint::Responses,
1620 ..LlmOpts::default()
1621 };
1622 assert_eq!(invalid_effort(&opts), None, "level {level} must be valid");
1623 }
1624 let opts = LlmOpts {
1625 effort: Some("deep".into()),
1626 endpoint: LlmEndpoint::Responses,
1627 ..LlmOpts::default()
1628 };
1629 let msg = invalid_effort(&opts).expect("unknown level must be rejected");
1630 assert!(msg.contains("not valid"), "got: {msg}");
1631 assert!(
1632 msg.contains("xhigh"),
1633 "message must list valid levels: {msg}"
1634 );
1635 let opts = LlmOpts {
1637 effort: Some("deep".into()),
1638 endpoint: LlmEndpoint::Chat,
1639 ..LlmOpts::default()
1640 };
1641 assert_eq!(invalid_effort(&opts), None);
1642 }
1643
1644 #[test]
1645 fn client_sends_session_header_and_user_agent() {
1646 let opts = LlmOpts {
1647 api_key: Some("test-key".to_string()),
1648 session_id: Some("sess-123".to_string()),
1649 ..LlmOpts::default()
1650 };
1651 let headers = openai_client(&opts)
1652 .expect("client builds")
1653 .headers()
1654 .clone();
1655 assert_eq!(
1656 headers.get("x-opencode-session").expect("session header"),
1657 "sess-123"
1658 );
1659 assert!(headers.contains_key("user-agent"), "got: {headers:?}");
1660 }
1661
1662 #[test]
1663 fn client_omits_session_header_when_unset() {
1664 let opts = LlmOpts {
1665 api_key: Some("test-key".to_string()),
1666 ..LlmOpts::default()
1667 };
1668 let headers = openai_client(&opts)
1669 .expect("client builds")
1670 .headers()
1671 .clone();
1672 assert!(!headers.contains_key("x-opencode-session"));
1673 assert!(headers.contains_key("user-agent"), "got: {headers:?}");
1674 }
1675
1676 #[test]
1677 fn localhost_base_url_builds_without_key() {
1678 for url in [
1681 "http://localhost:11434/v1",
1682 "http://127.0.0.1:8080/v1",
1683 "http://[::1]:8080/v1",
1684 ] {
1685 let opts = LlmOpts {
1686 base_url: Some(url.to_string()),
1687 api_key: None,
1688 ..LlmOpts::default()
1689 };
1690 assert!(
1691 openai_client(&opts).is_ok(),
1692 "keyless {url} must build a client"
1693 );
1694 }
1695 }
1696
1697 #[test]
1698 fn remote_base_url_without_key_is_error() {
1699 let opts = LlmOpts {
1700 base_url: Some("https://api.example.com/v1".to_string()),
1701 api_key: None,
1702 ..LlmOpts::default()
1703 };
1704 let err = openai_client(&opts).unwrap_err();
1705 assert!(err.contains("no API key"), "got: {err}");
1706 }
1707
1708 #[tokio::test]
1710 async fn list_models_unreachable_is_error() {
1711 let opts = LlmOpts {
1712 base_url: Some("http://127.0.0.1:9/v1".to_string()),
1713 api_key: Some("test-key".to_string()),
1714 ..LlmOpts::default()
1715 };
1716 let err = tokio::time::timeout(std::time::Duration::from_secs(10), list_models(&opts))
1717 .await
1718 .expect("list terminates")
1719 .expect_err("unreachable host must error");
1720 assert!(!err.is_empty());
1721 }
1722
1723 #[tokio::test]
1724 async fn list_models_without_key_is_error() {
1725 let opts = LlmOpts::default();
1726 let err = list_models(&opts).await.expect_err("no key must error");
1727 assert!(err.contains("no API key"), "got: {err}");
1728 }
1729
1730 #[test]
1731 fn client_cache_reuses_until_key_changes() {
1732 let opts = LlmOpts {
1733 base_url: Some("http://127.0.0.1:9/v1".to_string()),
1734 api_key: Some("k".to_string()),
1735 session_id: Some("s1".to_string()),
1736 ..LlmOpts::default()
1737 };
1738 let cache = ClientCache::default();
1739 let builds = || cache.builds.load(std::sync::atomic::Ordering::SeqCst);
1740
1741 cache.get(&opts).expect("builds");
1743 cache.get(&opts).expect("builds");
1744 assert_eq!(builds(), 1, "identical opts must reuse the client");
1745
1746 let model_swap = LlmOpts {
1748 model: "other".to_string(),
1749 effort: Some("high".to_string()),
1750 ..opts.clone()
1751 };
1752 cache.get(&model_swap).expect("builds");
1753 assert_eq!(builds(), 1, "model/effort swap must reuse the client");
1754
1755 for changed in [
1757 LlmOpts {
1758 base_url: Some("http://127.0.0.1:8/v1".to_string()),
1759 ..opts.clone()
1760 },
1761 LlmOpts {
1762 api_key: Some("k2".to_string()),
1763 ..opts.clone()
1764 },
1765 LlmOpts {
1766 session_id: Some("s2".to_string()),
1767 ..opts.clone()
1768 },
1769 ] {
1770 let before = builds();
1771 cache.get(&changed).expect("builds");
1772 assert_eq!(builds(), before + 1, "a key change must rebuild the client");
1773 }
1774 }
1775
1776 #[test]
1777 fn settle_determines_the_client_key() {
1778 let base = LlmProvider {
1782 endpoint: LlmEndpoint::Chat,
1783 base_url: Some("http://127.0.0.1:9/v1".to_string()),
1784 api_key: Some("k".to_string()),
1785 };
1786 let mut profiles = ModelProfiles::new(base.clone());
1787 profiles.insert(
1789 "ep-swap",
1790 LlmProfile {
1791 endpoint: Some(LlmEndpoint::Responses),
1792 ..LlmProfile::default()
1793 },
1794 );
1795 profiles.insert(
1797 "moved",
1798 LlmProfile {
1799 base_url: Some("http://127.0.0.1:8/v1".to_string()),
1800 ..LlmProfile::default()
1801 },
1802 );
1803
1804 let settled = |model: &str| {
1805 let mut llm = LlmOpts {
1806 model: model.to_string(),
1807 endpoint: base.endpoint,
1808 base_url: base.base_url.clone(),
1809 api_key: base.api_key.clone(),
1810 model_profiles: profiles.clone(),
1811 ..LlmOpts::default()
1812 };
1813 llm.settle_provider();
1814 llm
1815 };
1816
1817 let cache = ClientCache::default();
1818 let builds = || cache.builds.load(std::sync::atomic::Ordering::SeqCst);
1819
1820 let plain = settled("m0");
1821 cache.get(&plain).expect("builds");
1822 assert_eq!(builds(), 1);
1823
1824 let ep_swap = settled("ep-swap");
1825 assert_eq!(ep_swap.endpoint, LlmEndpoint::Responses);
1826 assert_eq!(ep_swap.base_url, base.base_url, "the base_url is inherited");
1827 cache.get(&ep_swap).expect("reuses");
1828 assert_eq!(builds(), 1, "an endpoint-only swap must reuse the client");
1829
1830 let moved = settled("moved");
1831 cache.get(&moved).expect("rebuilds");
1832 assert_eq!(builds(), 2, "a base_url change must rebuild the client");
1833 }
1834
1835 #[test]
1836 fn responses_request_builds_against_default_client_type() {
1837 use rig::client::CompletionClient;
1838 let client = openai::Client::builder()
1839 .api_key("test-key")
1840 .base_url("http://127.0.0.1:9/v1")
1841 .build()
1842 .expect("client builds without network");
1843 let model = client.completion_model("some-model");
1844 let request = build_request(&[], "sys", &[], &LlmOpts::default());
1845 let fut = rig::completion::CompletionModel::stream(&model, request);
1846 drop(fut);
1847 }
1848
1849 #[tokio::test]
1852 async fn adapter_surfaces_error_without_done() {
1853 let opts = LlmOpts {
1854 model: "m1".to_string(),
1855 base_url: Some("http://127.0.0.1:9/v1".to_string()),
1856 api_key: Some("test-key".to_string()),
1857 temperature: None,
1858 ..LlmOpts::default()
1859 };
1860 let stream_fn = rig_stream_fn();
1861 let stream = stream_fn(&[], "sys", &[], &opts);
1862 let events: Vec<LlmStreamEvent> =
1863 tokio::time::timeout(std::time::Duration::from_secs(10), stream.collect())
1864 .await
1865 .expect("stream terminates");
1866 assert!(
1867 events
1868 .iter()
1869 .any(|e| matches!(e, LlmStreamEvent::Error { .. })),
1870 "expected an Error event, got {events:?}"
1871 );
1872 assert!(
1873 !events
1874 .iter()
1875 .any(|e| matches!(e, LlmStreamEvent::Done { .. })),
1876 "no Done after error, got {events:?}"
1877 );
1878 }
1879
1880 #[tokio::test]
1883 async fn responses_adapter_surfaces_error_without_done() {
1884 let opts = LlmOpts {
1885 model: "m1".to_string(),
1886 base_url: Some("http://127.0.0.1:9/v1".to_string()),
1887 api_key: Some("test-key".to_string()),
1888 temperature: None,
1889 endpoint: LlmEndpoint::Responses,
1890 effort: None,
1891 session_id: None,
1892 retry: RetryPolicy::default(),
1893 model_profiles: ModelProfiles::default(),
1894 };
1895 let stream_fn = rig_stream_fn();
1896 let stream = stream_fn(&[], "sys", &[], &opts);
1897 let events: Vec<LlmStreamEvent> =
1898 tokio::time::timeout(std::time::Duration::from_secs(10), stream.collect())
1899 .await
1900 .expect("stream terminates");
1901 assert!(
1902 events
1903 .iter()
1904 .any(|e| matches!(e, LlmStreamEvent::Error { .. })),
1905 "expected an Error event, got {events:?}"
1906 );
1907 assert!(
1908 !events
1909 .iter()
1910 .any(|e| matches!(e, LlmStreamEvent::Done { .. })),
1911 "no Done after error, got {events:?}"
1912 );
1913 }
1914 #[tokio::test]
1916 #[ignore = "live test; run with WCODE_LIVE=1 WCODE_BASE_URL=... WCODE_API_KEY=... WCODE_MODEL=... cargo test -p wcode-harness --lib -- --ignored"]
1917 async fn live_rig_stream_fn() {
1918 let (Ok(base_url), Ok(api_key), Ok(model)) = (
1919 std::env::var("WCODE_BASE_URL"),
1920 std::env::var("WCODE_API_KEY"),
1921 std::env::var("WCODE_MODEL"),
1922 ) else {
1923 eprintln!("skipping: WCODE_BASE_URL / WCODE_API_KEY / WCODE_MODEL not set");
1924 return;
1925 };
1926 if std::env::var("WCODE_LIVE").ok().as_deref() != Some("1") {
1927 eprintln!("skipping: WCODE_LIVE != 1");
1928 return;
1929 }
1930 let opts = LlmOpts {
1931 model,
1932 base_url: Some(base_url),
1933 api_key: Some(api_key),
1934 temperature: None,
1935 ..LlmOpts::default()
1936 };
1937 let stream_fn = rig_stream_fn();
1938 let stream = stream_fn(
1939 &[AgentMessage::user_text("Reply with exactly: pong")],
1940 "You are a terse echo assistant.",
1941 &[],
1942 &opts,
1943 );
1944 let events: Vec<LlmStreamEvent> = stream.collect().await;
1945 let text = events
1946 .iter()
1947 .filter_map(|e| match e {
1948 LlmStreamEvent::TextDelta(t) => Some(t.as_str()),
1949 _ => None,
1950 })
1951 .collect::<String>();
1952 println!("text: {text:?}");
1953 println!("events: {events:?}");
1954 assert!(
1955 events
1956 .iter()
1957 .any(|e| matches!(e, LlmStreamEvent::Done { .. }))
1958 );
1959 assert!(
1960 !events
1961 .iter()
1962 .any(|e| matches!(e, LlmStreamEvent::Error { .. }))
1963 );
1964 }
1965}
1966
1967#[cfg(test)]
1968mod retry_tests {
1969 use super::*;
1970 use std::sync::Arc;
1971 use std::sync::atomic::{AtomicUsize, Ordering};
1972
1973 type Chunks =
1974 futures::stream::Iter<std::vec::IntoIter<Result<StreamedAssistantContent, CompletionError>>>;
1975
1976 fn transient() -> CompletionError {
1978 CompletionError::ProviderError(
1979 "Http client error: error sending request for url (http://x)".into(),
1980 )
1981 }
1982
1983 fn fatal() -> CompletionError {
1984 CompletionError::ProviderError("invalid_request_error: bad tool schema".into())
1985 }
1986
1987 fn empty() -> Chunks {
1988 futures::stream::iter(Vec::new())
1989 }
1990
1991 fn done() -> StreamedAssistantContent {
1993 StreamedAssistantContent::Final(rig::streaming::StreamFinal::new(
1994 "test",
1995 rig::completion::Usage::default(),
1996 ))
1997 }
1998
1999 fn first_item_err(e: CompletionError) -> Chunks {
2000 futures::stream::iter(vec![Err(e)])
2001 }
2002
2003 #[test]
2004 fn transient_statuses_retry_client_errors_do_not() {
2005 for code in [408, 409, 429, 500, 502, 503, 504] {
2006 assert!(is_transient_status(code), "{code} should retry");
2007 }
2008 for code in [200, 400, 401, 403, 404, 422] {
2009 assert!(!is_transient_status(code), "{code} must not retry");
2010 }
2011 }
2012
2013 #[test]
2014 fn retryable_classifies_transport_vs_rejections() {
2015 assert!(retryable(&CompletionError::HttpError(
2016 rig::http_client::Error::StreamEnded
2017 )));
2018 assert!(retryable(&transient()));
2019 assert!(!retryable(&fatal()));
2020 assert!(!retryable(&CompletionError::ResponseError("bad".into())));
2021 }
2022
2023 #[test]
2024 fn backoff_grows_and_never_exceeds_the_cap() {
2025 let base = Duration::from_millis(100);
2026 let cap = Duration::from_millis(400);
2027 let d1 = backoff(1, base, cap);
2028 assert!(d1 <= base && d1 >= base * 3 / 4, "attempt 1 = {d1:?}");
2029 for n in 1..12 {
2030 assert!(backoff(n, base, cap) <= cap, "attempt {n} exceeds cap");
2031 }
2032 }
2033
2034 #[tokio::test(start_paused = true)]
2035 async fn connect_retries_transient_failures_then_succeeds() {
2036 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2037 let calls = Arc::new(AtomicUsize::new(0));
2038 let policy = RetryPolicy {
2039 max: 3,
2040 base: Duration::from_millis(10),
2041 cap: Duration::from_millis(50),
2042 ..RetryPolicy::default()
2043 };
2044 let c = calls.clone();
2045 let mut state = RetryState::new(policy.max);
2046 let out = connect_with_retry(
2047 &policy,
2048 &mut state,
2049 move || {
2050 let n = c.fetch_add(1, Ordering::SeqCst) + 1;
2051 Box::pin(async move {
2052 if n <= 2 {
2053 Err(transient())
2054 } else {
2055 Ok(empty())
2056 }
2057 })
2058 },
2059 &tx,
2060 )
2061 .await;
2062 assert!(matches!(out, PreContent::Ready { .. }));
2063 assert_eq!(calls.load(Ordering::SeqCst), 3);
2064 assert!(matches!(
2065 rx.try_recv(),
2066 Ok(LlmStreamEvent::Retrying { attempt: 1, .. })
2067 ));
2068 assert!(matches!(
2069 rx.try_recv(),
2070 Ok(LlmStreamEvent::Retrying { attempt: 2, .. })
2071 ));
2072 }
2073
2074 #[tokio::test(start_paused = true)]
2075 async fn connect_retries_a_first_item_transport_error() {
2076 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
2077 let calls = Arc::new(AtomicUsize::new(0));
2078 let policy = RetryPolicy {
2079 max: 2,
2080 base: Duration::from_millis(5),
2081 cap: Duration::from_millis(5),
2082 ..RetryPolicy::default()
2083 };
2084 let c = calls.clone();
2085 let mut state = RetryState::new(policy.max);
2086 let out = connect_with_retry(
2087 &policy,
2088 &mut state,
2089 move || {
2090 let n = c.fetch_add(1, Ordering::SeqCst) + 1;
2091 Box::pin(async move {
2092 if n == 1 {
2093 Ok(first_item_err(transient()))
2094 } else {
2095 Ok(empty())
2096 }
2097 })
2098 },
2099 &tx,
2100 )
2101 .await;
2102 assert!(matches!(out, PreContent::Ready { .. }));
2103 assert_eq!(calls.load(Ordering::SeqCst), 2, "first-item error retried");
2104 }
2105
2106 #[tokio::test(start_paused = true)]
2107 async fn connect_gives_up_after_max_and_skips_non_retryable() {
2108 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
2109 let policy = RetryPolicy {
2110 max: 1,
2111 base: Duration::from_millis(1),
2112 cap: Duration::from_millis(1),
2113 ..RetryPolicy::default()
2114 };
2115
2116 let calls = Arc::new(AtomicUsize::new(0));
2117 let c = calls.clone();
2118 let mut state = RetryState::new(policy.max);
2119 let out = connect_with_retry(
2120 &policy,
2121 &mut state,
2122 move || {
2123 c.fetch_add(1, Ordering::SeqCst);
2124 Box::pin(async { Err::<Chunks, CompletionError>(transient()) })
2125 },
2126 &tx,
2127 )
2128 .await;
2129 assert!(matches!(out, PreContent::Failed(_)));
2130 assert_eq!(calls.load(Ordering::SeqCst), 2, "one attempt + one retry");
2131
2132 let calls = Arc::new(AtomicUsize::new(0));
2133 let c = calls.clone();
2134 let mut state = RetryState::new(policy.max);
2135 let out = connect_with_retry(
2136 &policy,
2137 &mut state,
2138 move || {
2139 c.fetch_add(1, Ordering::SeqCst);
2140 Box::pin(async { Err::<Chunks, CompletionError>(fatal()) })
2141 },
2142 &tx,
2143 )
2144 .await;
2145 assert!(matches!(out, PreContent::Failed(_)));
2146 assert_eq!(calls.load(Ordering::SeqCst), 1, "non-retryable returns at once");
2147 }
2148
2149 #[test]
2150 fn fatal_class_flags_hard_errors_but_not_transient_transport() {
2151 assert!(!fatal_class(&transient()));
2155 assert!(!fatal_class(&CompletionError::HttpError(
2156 rig::http_client::Error::StreamEnded
2157 )));
2158 assert!(fatal_class(&fatal()));
2159 assert!(fatal_class(&CompletionError::ResponseError("bad".into())));
2160 assert!(fatal_class(&CompletionError::RequestError("bad".into())));
2161 let json_err = serde_json::from_str::<u8>("not a number").unwrap_err();
2162 assert!(fatal_class(&CompletionError::JsonError(json_err)));
2163 }
2164
2165 #[tokio::test]
2170 async fn refused_port_exhausts_retries_to_non_fatal_error() {
2171 let opts = LlmOpts {
2172 model: "m1".to_string(),
2173 base_url: Some("http://127.0.0.1:9/v1".to_string()),
2174 api_key: Some("test-key".to_string()),
2175 retry: RetryPolicy {
2176 max: 2,
2177 base: Duration::from_millis(10),
2178 cap: Duration::from_millis(50),
2179 ..RetryPolicy::default()
2180 },
2181 ..LlmOpts::default()
2182 };
2183 let stream_fn = rig_stream_fn();
2184 let stream = stream_fn(&[], "sys", &[], &opts);
2185 let events: Vec<LlmStreamEvent> =
2186 tokio::time::timeout(std::time::Duration::from_secs(10), stream.collect())
2187 .await
2188 .expect("stream terminates");
2189 let retries = events
2190 .iter()
2191 .filter(|e| matches!(e, LlmStreamEvent::Retrying { .. }))
2192 .count();
2193 assert_eq!(retries, 2, "two backoff notices before giving up: {events:?}");
2194 assert!(
2195 matches!(events.last(), Some(LlmStreamEvent::Error { fatal: false, .. })),
2196 "refused port is transient-class, not fatal: {events:?}"
2197 );
2198 }
2199
2200 fn err_with_retry_after(value: &str) -> CompletionError {
2204 let mut headers = HeaderMap::new();
2205 headers.insert("retry-after", HeaderValue::from_str(value).unwrap());
2206 CompletionError::ProviderResponse(
2207 rig::ProviderResponseError::without_status("429")
2208 .with_headers(Some(Box::new(headers))),
2209 )
2210 }
2211
2212 #[test]
2213 fn retry_after_parses_integer_seconds() {
2214 assert_eq!(
2215 retry_after(&err_with_retry_after("7")),
2216 Some(Duration::from_secs(7))
2217 );
2218 }
2219
2220 #[test]
2221 fn retry_after_is_none_when_absent() {
2222 assert_eq!(retry_after(&transient()), None);
2224 }
2225
2226 #[test]
2227 fn retry_after_is_none_on_non_numeric() {
2228 assert_eq!(
2231 retry_after(&err_with_retry_after("Wed, 21 Oct 2015 07:28:00 GMT")),
2232 None
2233 );
2234 assert_eq!(retry_after(&err_with_retry_after("soon")), None);
2235 }
2236
2237 #[tokio::test(start_paused = true)]
2238 async fn retry_wait_honors_retry_after_over_backoff() {
2239 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
2240 let policy = RetryPolicy {
2243 base: Duration::from_millis(10),
2244 ..RetryPolicy::default()
2245 };
2246 let start = tokio::time::Instant::now();
2247 assert!(retry_wait(&tx, &policy, 1, &err_with_retry_after("3")).await);
2248 assert_eq!(start.elapsed(), Duration::from_secs(3));
2249 }
2250
2251 #[tokio::test(start_paused = true)]
2252 async fn retry_wait_clamps_retry_after_to_the_cap() {
2253 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
2254 let policy = RetryPolicy {
2257 cap: Duration::from_secs(8),
2258 ..RetryPolicy::default()
2259 };
2260 let start = tokio::time::Instant::now();
2261 assert!(retry_wait(&tx, &policy, 1, &err_with_retry_after("100")).await);
2262 assert_eq!(start.elapsed(), Duration::from_secs(8));
2263 }
2264 type BoxStream =
2267 Pin<Box<dyn futures::Stream<Item = Result<StreamedAssistantContent, CompletionError>> + Send>>;
2268
2269 fn boxed<S>(s: S) -> BoxStream
2270 where
2271 S: futures::Stream<Item = Result<StreamedAssistantContent, CompletionError>>
2272 + Send
2273 + 'static,
2274 {
2275 Box::pin(s)
2276 }
2277
2278 fn silent_boxed() -> BoxStream {
2279 boxed(futures::stream::pending::<Result<StreamedAssistantContent, CompletionError>>())
2280 }
2281
2282 fn open_silent()
2285 -> Pin<Box<dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send>> {
2286 Box::pin(async { Ok::<_, CompletionError>(silent_boxed()) })
2287 }
2288
2289 fn open_slow_ok()
2292 -> Pin<Box<dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send>> {
2293 Box::pin(async {
2294 tokio::time::sleep(Duration::from_millis(50)).await;
2295 Ok::<_, CompletionError>(boxed(futures::stream::iter(vec![Ok::<_, CompletionError>(
2296 StreamedAssistantContent::Text(rig::message::Text::new("ok")),
2297 )])))
2298 })
2299 }
2300
2301 #[test]
2303 fn a_stall_error_is_transient_and_not_hard_fatal() {
2304 let d = Duration::from_secs(1);
2305 assert!(retryable(&stall_error(d)), "a stall must be transient (B1)");
2306 assert!(!fatal_class(&stall_error(d)), "a stall must not be hard-fatal");
2307 assert!(
2308 stall_error(d).to_string().contains("timed out"),
2309 "the (timed out) needle drives retryable()"
2310 );
2311 }
2312
2313 #[tokio::test(start_paused = true)]
2317 async fn a_silent_stream_trips_the_ttft_retries_then_fails() {
2318 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2319 let calls = Arc::new(AtomicUsize::new(0));
2320 let policy = RetryPolicy {
2321 max: 2,
2322 base: Duration::from_millis(10),
2323 cap: Duration::from_millis(10),
2324 ttft: Duration::from_secs(1),
2325 ..RetryPolicy::default()
2326 };
2327 let c = calls.clone();
2328 let open = move || {
2329 c.fetch_add(1, Ordering::SeqCst);
2330 open_silent()
2331 };
2332 let mut state = RetryState::new(policy.max);
2333 let out = connect_with_retry(&policy, &mut state, open, &tx).await;
2334
2335 let PreContent::Failed(error) = out else {
2336 panic!("expected Failed from an exhausted stall")
2337 };
2338 assert!(retryable(&error), "a stall is transient");
2339 assert!(!fatal_class(&error), "a stall is not hard-fatal");
2340 assert_eq!(calls.load(Ordering::SeqCst), 3, "one attempt + two retries");
2341 let retries = std::iter::from_fn(|| rx.try_recv().ok())
2342 .filter(|e| matches!(e, LlmStreamEvent::Retrying { .. }))
2343 .count();
2344 assert_eq!(retries, 2, "a Retrying notice per retry");
2345 }
2346
2347 #[tokio::test(start_paused = true)]
2351 async fn deltas_then_a_stall_is_a_nonfatal_error() {
2352 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2353 let policy = RetryPolicy {
2354 max: 3,
2355 base: Duration::from_millis(10),
2356 cap: Duration::from_millis(10),
2357 idle: Duration::from_secs(1),
2358 ..RetryPolicy::default()
2359 };
2360 let stream = boxed(
2361 futures::stream::iter(vec![Ok::<_, CompletionError>(
2362 StreamedAssistantContent::Text(rig::message::Text::new("hi")),
2363 )])
2364 .chain(futures::stream::pending::<Result<StreamedAssistantContent, CompletionError>>()),
2365 );
2366 let state = RetryState::new(policy.max);
2367 forward_stream(stream, None, policy, open_silent, tx, state).await;
2368
2369 let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2370 assert!(
2371 events
2372 .iter()
2373 .any(|e| matches!(e, LlmStreamEvent::TextDelta(t) if t == "hi")),
2374 "{events:?}"
2375 );
2376 let error = events
2377 .iter()
2378 .find_map(|e| match e {
2379 LlmStreamEvent::Error { message, fatal } => Some((message.clone(), *fatal)),
2380 _ => None,
2381 })
2382 .expect("the idle deadline surfaced an error");
2383 assert!(!error.1, "a post-content stall is non-fatal: {error:?}");
2384 assert!(
2385 !events.iter().any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
2386 "no retry once content flowed: {events:?}"
2387 );
2388 }
2389
2390 #[tokio::test(start_paused = true)]
2393 async fn a_clean_stream_is_unaffected() {
2394 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2395 let policy = RetryPolicy {
2396 idle: Duration::from_secs(1),
2397 ..RetryPolicy::default()
2398 };
2399 let stream = boxed(futures::stream::iter(vec![
2400 Ok::<_, CompletionError>(StreamedAssistantContent::Text(rig::message::Text::new("hi"))),
2401 Ok::<_, CompletionError>(done()),
2402 ]));
2403 let state = RetryState::new(policy.max);
2404 forward_stream(stream, None, policy, open_silent, tx, state).await;
2405
2406 let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2407 assert!(
2408 events
2409 .iter()
2410 .any(|e| matches!(e, LlmStreamEvent::TextDelta(t) if t == "hi")),
2411 "{events:?}"
2412 );
2413 assert!(
2414 events.iter().any(|e| matches!(e, LlmStreamEvent::Done { .. })),
2415 "{events:?}"
2416 );
2417 assert!(
2418 !events.iter().any(|e| matches!(e, LlmStreamEvent::Error { .. })),
2419 "{events:?}"
2420 );
2421 assert!(
2422 !events.iter().any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
2423 "{events:?}"
2424 );
2425 }
2426
2427 #[tokio::test(start_paused = true)]
2431 async fn zero_deadlines_disable_the_timeout() {
2432 let mut silent =
2434 futures::stream::pending::<Result<StreamedAssistantContent, CompletionError>>();
2435 let stalled = tokio::time::timeout(
2436 Duration::from_secs(1),
2437 next_or_stall(&mut silent, Duration::ZERO),
2438 )
2439 .await;
2440 assert!(stalled.is_err(), "ZERO must not trip the deadline");
2441
2442 let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
2444 let policy = RetryPolicy {
2445 ttft: Duration::ZERO,
2446 ..RetryPolicy::default()
2447 };
2448 let mut state = RetryState::new(policy.max);
2449 let out = connect_with_retry(&policy, &mut state, open_slow_ok, &tx).await;
2450 assert!(
2451 matches!(out, PreContent::Ready { .. }),
2452 "a slow connect must not be timed out under ZERO"
2453 );
2454 }
2455
2456 #[tokio::test(start_paused = true)]
2462 async fn a_stalled_connect_trips_the_ttft_retries_then_fails() {
2463 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2464 let calls = Arc::new(AtomicUsize::new(0));
2465 let policy = RetryPolicy {
2466 max: 2,
2467 base: Duration::from_millis(10),
2468 cap: Duration::from_millis(10),
2469 ttft: Duration::from_secs(1),
2470 ..RetryPolicy::default()
2471 };
2472 let c = calls.clone();
2473 let open = move || {
2474 c.fetch_add(1, Ordering::SeqCst);
2475 Box::pin(std::future::pending::<Result<BoxStream, CompletionError>>())
2476 as Pin<
2477 Box<
2478 dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send,
2479 >,
2480 >
2481 };
2482 let mut state = RetryState::new(policy.max);
2483 let out = connect_with_retry(&policy, &mut state, open, &tx).await;
2484
2485 let PreContent::Failed(error) = out else {
2486 panic!("expected Failed from an exhausted connect stall")
2487 };
2488 assert!(retryable(&error), "a connect stall is transient");
2489 assert!(!fatal_class(&error), "a connect stall is not hard-fatal");
2490 assert_eq!(calls.load(Ordering::SeqCst), 3, "one attempt + two retries");
2491 let retries = std::iter::from_fn(|| rx.try_recv().ok())
2492 .filter(|e| matches!(e, LlmStreamEvent::Retrying { .. }))
2493 .count();
2494 assert_eq!(retries, 2, "a Retrying notice per retry");
2495 }
2496
2497 #[tokio::test(start_paused = true)]
2501 async fn a_pre_content_idle_stall_retries_and_reopens() {
2502 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2503 let policy = RetryPolicy {
2504 max: 2,
2505 base: Duration::from_millis(10),
2506 cap: Duration::from_millis(10),
2507 idle: Duration::from_secs(1),
2508 ..RetryPolicy::default()
2509 };
2510 let stream = silent_boxed();
2512 fn open_done()
2514 -> Pin<Box<dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send>> {
2515 Box::pin(async {
2516 Ok::<_, CompletionError>(boxed(futures::stream::iter(vec![Ok::<_, CompletionError>(
2517 done(),
2518 )])))
2519 })
2520 }
2521 let state = RetryState::new(policy.max);
2522 forward_stream(stream, None, policy, open_done, tx, state).await;
2523
2524 let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2525 assert!(
2526 events
2527 .iter()
2528 .any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
2529 "a pre-content stall backs off: {events:?}"
2530 );
2531 assert!(
2532 events.iter().any(|e| matches!(e, LlmStreamEvent::Done { .. })),
2533 "the re-opened stream completed: {events:?}"
2534 );
2535 assert!(
2536 !events.iter().any(|e| matches!(e, LlmStreamEvent::Error { .. })),
2537 "no error once the retry succeeded: {events:?}"
2538 );
2539 }
2540
2541 fn open_empty()
2545 -> Pin<Box<dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send>> {
2546 Box::pin(async { Ok::<_, CompletionError>(boxed(empty())) })
2547 }
2548
2549 #[tokio::test(start_paused = true)]
2553 async fn a_delta_then_a_truncated_end_is_a_nonfatal_error() {
2554 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2555 let policy = RetryPolicy {
2556 max: 3,
2557 base: Duration::from_millis(10),
2558 cap: Duration::from_millis(10),
2559 idle: Duration::ZERO,
2560 ..RetryPolicy::default()
2561 };
2562 let stream = boxed(futures::stream::iter(vec![Ok::<_, CompletionError>(
2563 StreamedAssistantContent::Text(rig::message::Text::new("hi")),
2564 )]));
2565 let state = RetryState::new(policy.max);
2566 forward_stream(stream, None, policy, open_silent, tx, state).await;
2567
2568 let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2569 assert!(
2570 events
2571 .iter()
2572 .any(|e| matches!(e, LlmStreamEvent::TextDelta(t) if t == "hi")),
2573 "{events:?}"
2574 );
2575 let (message, fatal) = events
2576 .iter()
2577 .find_map(|e| match e {
2578 LlmStreamEvent::Error { message, fatal } => Some((message.clone(), *fatal)),
2579 _ => None,
2580 })
2581 .expect("a truncated end surfaced an error");
2582 assert!(!fatal, "a post-content truncation is non-fatal: {message}");
2583 assert!(
2584 !events
2585 .iter()
2586 .any(|e| matches!(e, LlmStreamEvent::Done { .. })),
2587 "no Done on a truncated stream: {events:?}"
2588 );
2589 assert!(
2590 !events
2591 .iter()
2592 .any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
2593 "no retry once content flowed: {events:?}"
2594 );
2595 }
2596
2597 #[tokio::test(start_paused = true)]
2600 async fn a_final_terminated_stream_is_unchanged() {
2601 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2602 let policy = RetryPolicy {
2603 idle: Duration::ZERO,
2604 ..RetryPolicy::default()
2605 };
2606 let stream = boxed(futures::stream::iter(vec![
2607 Ok::<_, CompletionError>(StreamedAssistantContent::Text(rig::message::Text::new(
2608 "hi",
2609 ))),
2610 Ok::<_, CompletionError>(done()),
2611 ]));
2612 let state = RetryState::new(policy.max);
2613 forward_stream(stream, None, policy, open_silent, tx, state).await;
2614
2615 let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2616 assert!(
2617 events
2618 .iter()
2619 .any(|e| matches!(e, LlmStreamEvent::Done { .. })),
2620 "{events:?}"
2621 );
2622 assert!(
2623 !events
2624 .iter()
2625 .any(|e| matches!(e, LlmStreamEvent::Error { .. })),
2626 "a clean turn surfaces no error: {events:?}"
2627 );
2628 assert!(
2629 !events
2630 .iter()
2631 .any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
2632 "{events:?}"
2633 );
2634 }
2635
2636 #[tokio::test(start_paused = true)]
2641 async fn a_pre_content_no_item_stream_retries_then_errors() {
2642 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
2643 let policy = RetryPolicy {
2644 max: 2,
2645 base: Duration::from_millis(10),
2646 cap: Duration::from_millis(10),
2647 idle: Duration::ZERO,
2648 ..RetryPolicy::default()
2649 };
2650 let state = RetryState::new(policy.max);
2651 forward_stream(boxed(empty()), None, policy, open_empty, tx, state).await;
2652
2653 let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
2654 let retries = events
2655 .iter()
2656 .filter(|e| matches!(e, LlmStreamEvent::Retrying { .. }))
2657 .count();
2658 assert_eq!(
2659 retries, policy.max as usize,
2660 "one Retrying notice per retry: {events:?}"
2661 );
2662 let (message, fatal) = events
2663 .iter()
2664 .find_map(|e| match e {
2665 LlmStreamEvent::Error { message, fatal } => Some((message.clone(), *fatal)),
2666 _ => None,
2667 })
2668 .expect("an exhausted truncation budget surfaced an error");
2669 assert!(!fatal, "an exhausted truncation is non-fatal: {message}");
2670 assert!(
2671 !events
2672 .iter()
2673 .any(|e| matches!(e, LlmStreamEvent::Done { .. })),
2674 "no Done on a truncated stream: {events:?}"
2675 );
2676 }
2677
2678 fn spawn_models_server(delay: std::time::Duration) -> String {
2683 use std::io::{Read as _, Write as _};
2684 use std::net::TcpListener;
2685 let listener = TcpListener::bind("127.0.0.1:0").expect("bind ephemeral port");
2686 let addr = listener.local_addr().expect("local addr");
2687 std::thread::spawn(move || {
2688 for stream in listener.incoming() {
2692 let Ok(mut stream) = stream else { continue };
2693 let mut buf = Vec::new();
2697 let mut chunk = [0u8; 1024];
2698 loop {
2699 match stream.read(&mut chunk) {
2700 Ok(0) | Err(_) => break,
2701 Ok(n) => {
2702 buf.extend_from_slice(&chunk[..n]);
2703 if buf.windows(4).any(|w| w == b"\r\n\r\n") {
2704 break;
2705 }
2706 }
2707 }
2708 }
2709 std::thread::sleep(delay);
2710 let body = "{\"object\":\"list\",\"data\":[{\"id\":\"m1\",\"object\":\"model\"}]}";
2711 let response = format!(
2712 "HTTP/1.1 200 OK\r\n\
2713 Content-Type: application/json\r\n\
2714 Connection: close\r\n\
2715 Content-Length: {}\r\n\r\n{}",
2716 body.len(),
2717 body
2718 );
2719 let _ = stream.write_all(response.as_bytes());
2720 let _ = stream.flush();
2721 }
2722 });
2723 format!("http://{addr}/v1")
2724 }
2725
2726 #[tokio::test]
2727 async fn a_refused_port_probes_false() {
2728 let up = probe_endpoint(
2729 "http://127.0.0.1:9/v1",
2730 std::time::Duration::from_millis(500),
2731 )
2732 .await;
2733 assert!(!up, "a refused port is not an endpoint");
2734 }
2735
2736 #[tokio::test]
2737 async fn a_live_models_endpoint_probes_true() {
2738 let _ = probe_endpoint("http://127.0.0.1:9/v1", std::time::Duration::from_secs(60)).await;
2743 let base = spawn_models_server(std::time::Duration::ZERO);
2744 let up = probe_endpoint(&base, std::time::Duration::from_secs(5)).await;
2745 assert!(up, "a live /models endpoint probes true");
2746 }
2747
2748 #[tokio::test]
2749 async fn a_slow_endpoint_times_out_to_false() {
2750 let base = spawn_models_server(std::time::Duration::from_millis(500));
2751 let up = probe_endpoint(&base, std::time::Duration::from_millis(50)).await;
2752 assert!(!up, "a slow endpoint is not treated as up");
2753 }
2754}