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