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, Copy, Debug, Default, PartialEq, Eq)]
41pub enum EndpointPref {
42 #[default]
44 Auto,
45 Pinned(LlmEndpoint),
47}
48
49impl EndpointPref {
50 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 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
77impl 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 pub temperature: Option<f64>,
96 pub endpoint: LlmEndpoint,
97 pub session_id: Option<String>,
100 pub effort: Option<String>,
107 pub retry: RetryPolicy,
110 pub model_profiles: ModelProfiles,
115}
116
117#[derive(Clone, Debug, Default, PartialEq, Eq)]
122pub struct LlmProfile {
123 pub endpoint: Option<LlmEndpoint>,
127 pub base_url: Option<String>,
129 pub api_key: Option<String>,
132}
133
134impl LlmProfile {
135 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#[derive(Clone, Debug, PartialEq, Eq)]
158pub struct LlmProvider {
159 pub endpoint: EndpointPref,
161 pub base_url: Option<String>,
162 pub api_key: Option<String>,
163}
164
165impl LlmProvider {
166 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 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#[derive(Clone, Debug, Default, PartialEq, Eq)]
197pub struct ModelProfiles {
198 base: Option<LlmProvider>,
200 profiles: std::collections::BTreeMap<String, LlmProfile>,
201}
202
203impl ModelProfiles {
204 pub fn new(base: LlmProvider) -> Self {
206 Self {
207 base: Some(base),
208 profiles: std::collections::BTreeMap::new(),
209 }
210 }
211
212 fn resolve(&self, model: &str) -> Option<&LlmProfile> {
214 self.profiles.get(model)
215 }
216
217 pub fn settle(&self, llm: &mut LlmOpts) {
222 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 base.apply(llm);
234 let profile = self.resolve(&llm.model);
238 if let Some(profile) = profile {
239 profile.apply(llm);
240 }
241 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 pub fn base(&self) -> Option<&LlmProvider> {
256 self.base.as_ref()
257 }
258
259 pub fn insert(&mut self, model: impl Into<String>, profile: LlmProfile) {
261 self.profiles.insert(model.into(), profile);
262 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 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#[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#[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 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
338pub 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
347fn 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 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
385fn 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
401pub 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}
420pub 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 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 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
486async 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 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 None => {
522 if done_seen || errored {
523 break;
524 }
525 let error = truncation_error();
526 if !content_forwarded && state.budget > 0 {
527 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 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 Some(Stall::Item(item)) => item,
567 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 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 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 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 Box::pin(futures::stream::iter(vec![LlmStreamEvent::Error {
675 message,
676 fatal: true,
677 }]))
678}
679
680#[derive(Clone, Copy, Debug, PartialEq, Eq)]
683pub struct RetryPolicy {
684 pub max: u32,
686 pub base: Duration,
688 pub cap: Duration,
690 pub ttft: Duration,
692 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
708enum Stall<T> {
710 Item(T),
713 Idle,
715}
716
717async 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}
733fn is_transient_status(code: u16) -> bool {
736 matches!(code, 408 | 409 | 429 | 500 | 502 | 503 | 504)
737}
738
739fn 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 CompletionError::HttpError(_) => true,
749 CompletionError::ProviderError(message) => is_transient_message(message),
753 _ => false,
754 }
755}
756
757fn 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
776fn 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
789fn 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
797fn 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
817enum PreContent<S> {
821 Ready {
822 stream: S,
823 first: Option<Result<StreamedAssistantContent, CompletionError>>,
824 },
825 Failed(CompletionError),
827 Cancelled,
829}
830
831#[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
849fn fatal_class(error: &CompletionError) -> bool {
856 !retryable(error)
857}
858
859fn stall_error(d: Duration) -> CompletionError {
873 CompletionError::ProviderError(format!("stream stalled: no item within {d:?} (timed out)"))
874}
875fn 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 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 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 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
976async 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
1000fn retry_after(error: &CompletionError) -> Option<Duration> {
1010 let raw = error.provider_response_headers()?.get("retry-after")?;
1014 let seconds: u64 = raw.to_str().ok()?.trim().parse().ok()?;
1016 Some(Duration::from_secs(seconds))
1017}
1018
1019async fn retry_wait(
1022 tx: &tokio::sync::mpsc::UnboundedSender<LlmStreamEvent>,
1023 policy: &RetryPolicy,
1024 attempt: u32,
1025 error: &CompletionError,
1026) -> bool {
1027 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 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
1073const RESPONSES_EFFORT_LEVELS: &[&str] =
1078 &["none", "minimal", "low", "medium", "high", "xhigh", "max"];
1079
1080fn 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
1096fn 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
1129const 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 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 StreamedAssistantContent::Reasoning { reasoning, .. } => {
1203 let text = reasoning_text(&reasoning);
1206 if text.is_empty() {
1207 vec![]
1212 } else {
1213 vec![LlmStreamEvent::ThinkingReplace(text)]
1214 }
1215 }
1216 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 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 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 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 assert_eq!(wire_text(&ok), "boom");
1583 assert_eq!(wire_text(&err), "ERROR: boom");
1584 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 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 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 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 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 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 #[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 cache.get(&opts).expect("builds");
1816 cache.get(&opts).expect("builds");
1817 assert_eq!(builds(), 1, "identical opts must reuse the client");
1818
1819 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 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 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 profiles.insert(
1862 "ep-swap",
1863 LlmProfile {
1864 endpoint: Some(LlmEndpoint::Responses),
1865 ..LlmProfile::default()
1866 },
1867 );
1868 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 assert_eq!(
1914 EndpointPref::Pinned(LlmEndpoint::Chat).resolve(Some(ZEN), "grok-4.7"),
1915 LlmEndpoint::Chat
1916 );
1917 assert_eq!(
1919 EndpointPref::Auto.resolve(Some(ZEN), "grok-4.7"),
1920 LlmEndpoint::Responses
1921 );
1922 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 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 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 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 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 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 #[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 #[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 #[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 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 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 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 #[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 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 assert_eq!(retry_after(&transient()), None);
2454 }
2455
2456 #[test]
2457 fn retry_after_is_none_on_non_numeric() {
2458 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 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 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 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 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 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 #[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 #[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 #[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 #[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 #[tokio::test(start_paused = true)]
2661 async fn zero_deadlines_disable_the_timeout() {
2662 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 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 #[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 #[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 let stream = silent_boxed();
2742 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 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 #[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 #[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 #[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 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 for stream in listener.incoming() {
2922 let Ok(mut stream) = stream else { continue };
2923 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 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}