use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use futures::{Stream, StreamExt};
use rig::client::{CompletionClient, ModelListingClient};
use rig::completion::{
CompletionError, CompletionModel, CompletionRequest, FinishReason, ToolDefinition,
};
use rig::http_client::{HeaderMap, HeaderValue};
use rig::message::{
AssistantContent, Message, Reasoning, ReasoningContent, ToolCall, ToolCallId, ToolFunction,
UserContent,
};
use rig::providers::openai;
use rig::streaming::{StreamedAssistantContent, StreamingCompletionResponse};
use crate::event::LlmStreamEvent;
use crate::message::{AgentMessage, ContentBlock, StopReason};
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum LlmEndpoint {
#[default]
Chat,
Responses,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum EndpointPref {
#[default]
Auto,
Pinned(LlmEndpoint),
}
impl EndpointPref {
pub fn resolve(self, base_url: Option<&str>, model: &str) -> LlmEndpoint {
match self {
Self::Pinned(e) => e,
Self::Auto => {
crate::limits::model_protocol(base_url, model).unwrap_or(LlmEndpoint::Chat)
}
}
}
pub fn wire(self) -> &'static str {
match self {
Self::Auto => "auto",
Self::Pinned(LlmEndpoint::Chat) => "chat",
Self::Pinned(LlmEndpoint::Responses) => "responses",
}
}
}
impl From<Option<LlmEndpoint>> for EndpointPref {
fn from(value: Option<LlmEndpoint>) -> Self {
match value {
Some(e) => Self::Pinned(e),
None => Self::Auto,
}
}
}
#[derive(Clone, Debug, Default)]
pub struct LlmOpts {
pub model: String,
pub base_url: Option<String>,
pub api_key: Option<String>,
pub temperature: Option<f64>,
pub endpoint: LlmEndpoint,
pub session_id: Option<String>,
pub effort: Option<String>,
pub retry: RetryPolicy,
pub model_profiles: ModelProfiles,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct LlmProfile {
pub endpoint: Option<LlmEndpoint>,
pub base_url: Option<String>,
pub api_key: Option<String>,
}
impl LlmProfile {
fn apply(&self, llm: &mut LlmOpts) {
if let Some(endpoint) = self.endpoint {
llm.endpoint = endpoint;
}
if let Some(base_url) = &self.base_url {
llm.base_url = Some(base_url.clone());
}
if let Some(api_key) = &self.api_key {
llm.api_key = Some(api_key.clone());
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct LlmProvider {
pub endpoint: EndpointPref,
pub base_url: Option<String>,
pub api_key: Option<String>,
}
impl LlmProvider {
pub fn of(llm: &LlmOpts, endpoint: EndpointPref) -> Self {
Self {
endpoint,
base_url: llm.base_url.clone(),
api_key: llm.api_key.clone(),
}
}
fn apply(&self, llm: &mut LlmOpts) {
llm.base_url = self.base_url.clone();
llm.api_key = self.api_key.clone();
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct ModelProfiles {
base: Option<LlmProvider>,
profiles: std::collections::BTreeMap<String, LlmProfile>,
}
impl ModelProfiles {
pub fn new(base: LlmProvider) -> Self {
Self {
base: Some(base),
profiles: std::collections::BTreeMap::new(),
}
}
fn resolve(&self, model: &str) -> Option<&LlmProfile> {
self.profiles.get(model)
}
pub fn settle(&self, llm: &mut LlmOpts) {
debug_assert!(
self.base.is_some() || self.profiles.is_empty(),
"ModelProfiles has profiles but no base (was `insert` called on the Default?)"
);
let Some(base) = &self.base else {
return;
};
base.apply(llm);
let profile = self.resolve(&llm.model);
if let Some(profile) = profile {
profile.apply(llm);
}
let pinned_by_profile = profile.is_some_and(|p| p.endpoint.is_some());
if !pinned_by_profile {
llm.endpoint = base.endpoint.resolve(llm.base_url.as_deref(), &llm.model);
}
}
pub fn base(&self) -> Option<&LlmProvider> {
self.base.as_ref()
}
pub fn insert(&mut self, model: impl Into<String>, profile: LlmProfile) {
self.profiles.insert(model.into(), profile);
debug_assert!(
self.base.is_some() || self.profiles.is_empty(),
"insert on a ModelProfiles with no base: the profile is dead (use `new`)"
);
}
}
impl LlmOpts {
pub fn settle_provider(&mut self) {
let profiles = self.model_profiles.clone();
profiles.settle(self);
}
}
pub type LlmStream = Pin<Box<dyn Stream<Item = LlmStreamEvent> + Send>>;
pub type StreamFn =
Arc<dyn Fn(&[AgentMessage], &str, &[ToolDefinition], &LlmOpts) -> LlmStream + Send + Sync>;
#[derive(Clone, PartialEq, Eq)]
struct ClientKey {
base_url: Option<String>,
api_key: Option<String>,
session_id: Option<String>,
}
impl ClientKey {
fn of(opts: &LlmOpts) -> Self {
Self {
base_url: opts.base_url.clone(),
api_key: opts.api_key.clone(),
session_id: opts.session_id.clone(),
}
}
}
#[derive(Default)]
struct ClientCache {
slot: Mutex<Option<(ClientKey, openai::Client)>>,
#[cfg(test)]
builds: std::sync::atomic::AtomicUsize,
}
impl ClientCache {
fn get(&self, opts: &LlmOpts) -> Result<openai::Client, String> {
let key = ClientKey::of(opts);
let mut slot = self.slot.lock().unwrap_or_else(|e| e.into_inner());
if let Some((cached_key, client)) = slot.as_ref()
&& *cached_key == key
{
return Ok(client.clone());
}
let client = openai_client(opts)?;
#[cfg(test)]
self.builds
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
*slot = Some((key, client.clone()));
Ok(client)
}
}
pub fn rig_stream_fn() -> StreamFn {
let cache = ClientCache::default();
Arc::new(move |messages, system, tools, opts| adapt(&cache, opts, messages, system, tools))
}
fn openai_client(opts: &LlmOpts) -> Result<rig::providers::openai::Client, String> {
let key = opts
.api_key
.clone()
.or_else(|| std::env::var("OPENAI_API_KEY").ok());
let key = match key {
Some(k) => k,
None => {
let local = opts.base_url.as_deref().is_some_and(|url| {
let u = url.to_ascii_lowercase();
["http://localhost", "http://127.0.0.1", "http://[::1]"]
.iter()
.any(|prefix| u.starts_with(prefix))
});
if local {
"wcode-local".to_string()
} else {
return Err(
"no API key: pass LlmOpts.api_key or set OPENAI_API_KEY (a localhost base_url works keyless)"
.to_string(),
);
}
}
};
let mut builder = openai::Client::builder()
.api_key::<rig::client::BearerAuth>(key)
.http_headers(session_headers(opts));
if let Some(base_url) = &opts.base_url {
builder = builder.base_url(base_url);
}
builder.build().map_err(|e| format!("openai client: {e}"))
}
fn session_headers(opts: &LlmOpts) -> HeaderMap {
let mut headers = HeaderMap::new();
headers.insert(
"user-agent",
HeaderValue::from_static(concat!("wcode/", env!("CARGO_PKG_VERSION"))),
);
if let Some(id) = opts.session_id.as_deref()
&& let Ok(value) = HeaderValue::from_str(id)
{
headers.insert("x-opencode-session", value);
}
headers
}
pub async fn probe_endpoint(base_url: &str, timeout: std::time::Duration) -> bool {
let opts = LlmOpts {
base_url: Some(base_url.to_string()),
..LlmOpts::default()
};
matches!(
tokio::time::timeout(timeout, list_models(&opts)).await,
Ok(Ok(_))
)
}
pub async fn list_models(opts: &LlmOpts) -> Result<Vec<String>, String> {
let client = openai_client(opts)?;
let list = client.list_models().await.map_err(|e| e.to_string())?;
let mut ids: Vec<String> = list.iter().map(|m| m.id.clone()).collect();
ids.sort();
Ok(ids)
}
fn adapt(
cache: &ClientCache,
opts: &LlmOpts,
messages: &[AgentMessage],
system: &str,
tools: &[ToolDefinition],
) -> LlmStream {
if let Some(message) = invalid_effort(opts) {
return error_stream(message);
}
let client = match cache.get(opts) {
Ok(client) => client,
Err(message) => return error_stream(message),
};
let request = build_request(messages, system, tools, opts);
let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<LlmStreamEvent>();
let Ok(runtime) = tokio::runtime::Handle::try_current() else {
return error_stream("StreamFn requires a tokio runtime".to_string());
};
let policy = opts.retry;
let endpoint = opts.endpoint;
let model = opts.model.clone();
runtime.spawn(async move {
let open = || open_stream(client.clone(), endpoint, &model, request.clone());
let mut state = RetryState::new(policy.max);
let (stream, pending) = match connect_with_retry(&policy, &mut state, &open, &tx).await {
PreContent::Ready { stream, first } => (stream, first),
PreContent::Cancelled => return,
PreContent::Failed(error) => {
let _ = tx.send(LlmStreamEvent::Error {
message: error.to_string(),
fatal: fatal_class(&error),
});
return;
}
};
forward_stream(stream, pending, policy, &open, tx, state).await;
});
Box::pin(futures::stream::unfold(rx, |mut rx| async move {
rx.recv().await.map(|ev| (ev, rx))
}))
}
async fn forward_stream<S, F>(
mut stream: S,
mut pending: Option<Result<StreamedAssistantContent, CompletionError>>,
policy: RetryPolicy,
open: F,
tx: tokio::sync::mpsc::UnboundedSender<LlmStreamEvent>,
mut state: RetryState,
) where
F: Fn() -> Pin<Box<dyn std::future::Future<Output = Result<S, CompletionError>> + Send>>,
S: futures::Stream<Item = Result<StreamedAssistantContent, CompletionError>> + Unpin + Send,
{
let mut errored = false;
let mut content_forwarded = false;
let mut done_seen = false;
loop {
let item = match pending.take() {
Some(item) => item,
None => tokio::select! {
biased;
result = next_or_stall(&mut stream, policy.idle) => match result {
None => {
if done_seen || errored {
break;
}
let error = truncation_error();
if !content_forwarded && state.budget > 0 {
match pre_content_retry(&policy, &mut state, &open, &tx, &error).await {
PreContent::Ready { stream: s, first } => {
stream = s;
pending = first;
continue;
}
PreContent::Cancelled => return,
PreContent::Failed(e) => {
if tx
.send(LlmStreamEvent::Error {
message: e.to_string(),
fatal: fatal_class(&e),
})
.is_err()
{
return;
}
break;
}
}
} else {
if tx
.send(LlmStreamEvent::Error {
message: error.to_string(),
fatal: false,
})
.is_err()
{
return;
}
break;
}
}
Some(Stall::Item(item)) => item,
Some(Stall::Idle) => {
let error = stall_error(policy.idle);
if !content_forwarded && state.budget > 0 {
match pre_content_retry(&policy, &mut state, &open, &tx, &error).await {
PreContent::Ready { stream: s, first } => {
stream = s;
pending = first;
continue;
}
PreContent::Cancelled => return,
PreContent::Failed(e) => {
let fatal = fatal_class(&e);
if tx
.send(LlmStreamEvent::Error {
message: e.to_string(),
fatal,
})
.is_err()
{
return;
}
break;
}
}
} else {
if tx
.send(LlmStreamEvent::Error {
message: error.to_string(),
fatal: false,
})
.is_err()
{
return;
}
break;
}
}
},
_ = tx.closed() => break,
},
};
let events = match item {
Ok(content) => {
let is_final = matches!(content, StreamedAssistantContent::Final(_));
if errored && is_final {
continue;
}
let events = map_item(content);
content_forwarded |= !events.is_empty();
done_seen |= is_final;
events
}
Err(error) => {
if !content_forwarded && state.budget > 0 && retryable(&error) {
state.attempts += 1;
state.budget -= 1;
if !retry_wait(&tx, &policy, state.attempts, &error).await {
return;
}
match connect_with_retry(&policy, &mut state, &open, &tx).await {
PreContent::Ready { stream: s, first } => {
stream = s;
pending = first;
continue;
}
PreContent::Cancelled => return,
PreContent::Failed(e) => {
errored = true;
vec![LlmStreamEvent::Error {
message: e.to_string(),
fatal: fatal_class(&e),
}]
}
}
} else {
errored = true;
vec![LlmStreamEvent::Error {
message: error.to_string(),
fatal: fatal_class(&error),
}]
}
}
};
for ev in events {
if tx.send(ev).is_err() {
return;
}
}
}
}
fn error_stream(message: String) -> LlmStream {
Box::pin(futures::stream::iter(vec![LlmStreamEvent::Error {
message,
fatal: true,
}]))
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct RetryPolicy {
pub max: u32,
pub base: Duration,
pub cap: Duration,
pub ttft: Duration,
pub idle: Duration,
}
impl Default for RetryPolicy {
fn default() -> Self {
Self {
max: 3,
base: Duration::from_millis(500),
cap: Duration::from_secs(8),
ttft: Duration::from_secs(60),
idle: Duration::from_secs(120),
}
}
}
enum Stall<T> {
Item(T),
Idle,
}
async fn next_or_stall<S: Stream + Unpin>(
stream: &mut S,
deadline: Duration,
) -> Option<Stall<S::Item>> {
if deadline.is_zero() {
return stream.next().await.map(Stall::Item);
}
tokio::select! {
biased;
item = stream.next() => item.map(Stall::Item),
_ = tokio::time::sleep(deadline) => Some(Stall::Idle),
}
}
fn is_transient_status(code: u16) -> bool {
matches!(code, 408 | 409 | 429 | 500 | 502 | 503 | 504)
}
fn retryable(error: &CompletionError) -> bool {
if let Some(status) = error.provider_response_status() {
return is_transient_status(status.as_u16());
}
match error {
CompletionError::HttpError(_) => true,
CompletionError::ProviderError(message) => is_transient_message(message),
_ => false,
}
}
fn is_transient_message(message: &str) -> bool {
const NEEDLES: [&str; 10] = [
"error sending request",
"error trying to connect",
"connection refused",
"connection reset",
"connection closed",
"timed out",
"timeout",
"unreachable",
"dns error",
"temporary failure",
];
let m = message.to_ascii_lowercase();
NEEDLES.iter().any(|needle| m.contains(needle))
}
fn backoff(attempt: u32, base: Duration, cap: Duration) -> Duration {
let shift = attempt.saturating_sub(1).min(16);
let capped = base.saturating_mul(1u32 << shift).min(cap);
let millis = capped.as_millis() as u64;
if millis <= 1 {
return capped;
}
let jitter = clock_nanos() % (millis / 4 + 1);
Duration::from_millis(millis - jitter)
}
fn clock_nanos() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| u64::from(d.subsec_nanos()))
.unwrap_or(0)
}
fn open_stream(
client: rig::providers::openai::Client,
endpoint: LlmEndpoint,
model: &str,
request: CompletionRequest,
) -> Pin<Box<dyn std::future::Future<Output = Result<StreamingCompletionResponse, CompletionError>> + Send>> {
match endpoint {
LlmEndpoint::Chat => {
let model = client.completions_api().completion_model(model.to_string());
Box::pin(async move { model.stream(request).await })
}
LlmEndpoint::Responses => {
let model = client.completion_model(model.to_string());
Box::pin(async move { model.stream(request).await })
}
}
}
enum PreContent<S> {
Ready {
stream: S,
first: Option<Result<StreamedAssistantContent, CompletionError>>,
},
Failed(CompletionError),
Cancelled,
}
#[derive(Clone, Copy)]
struct RetryState {
budget: u32,
attempts: u32,
}
impl RetryState {
fn new(max: u32) -> Self {
RetryState {
budget: max,
attempts: 0,
}
}
}
fn fatal_class(error: &CompletionError) -> bool {
!retryable(error)
}
fn stall_error(d: Duration) -> CompletionError {
CompletionError::ProviderError(format!("stream stalled: no item within {d:?} (timed out)"))
}
fn truncation_error() -> CompletionError {
CompletionError::ProviderError(
"stream ended before its terminal record (truncated)".to_string(),
)
}
async fn connect_with_retry<S, F>(
policy: &RetryPolicy,
state: &mut RetryState,
mut open: F,
tx: &tokio::sync::mpsc::UnboundedSender<LlmStreamEvent>,
) -> PreContent<S>
where
F: FnMut() -> Pin<Box<dyn std::future::Future<Output = Result<S, CompletionError>> + Send>>,
S: futures::Stream<Item = Result<StreamedAssistantContent, CompletionError>> + Unpin + Send,
{
loop {
state.attempts += 1;
let opened = if policy.ttft.is_zero() {
Ok(open().await)
} else {
tokio::select! {
biased;
r = open() => Ok(r),
_ = tokio::time::sleep(policy.ttft) => Err(()),
}
};
let mut stream = match opened {
Ok(Ok(stream)) => stream,
Ok(Err(error)) => {
if state.budget > 0 && retryable(&error) {
state.budget -= 1;
if !retry_wait(tx, policy, state.attempts, &error).await {
return PreContent::Cancelled;
}
continue;
}
return PreContent::Failed(error);
}
Err(()) => {
let error = stall_error(policy.ttft);
if state.budget > 0 {
state.budget -= 1;
if !retry_wait(tx, policy, state.attempts, &error).await {
return PreContent::Cancelled;
}
continue;
}
return PreContent::Failed(error);
}
};
match next_or_stall(&mut stream, policy.ttft).await {
None => return PreContent::Ready { stream, first: None },
Some(Stall::Item(Err(error))) => {
if state.budget > 0 && retryable(&error) {
state.budget -= 1;
if !retry_wait(tx, policy, state.attempts, &error).await {
return PreContent::Cancelled;
}
continue;
}
return PreContent::Ready {
stream,
first: Some(Err(error)),
};
}
Some(Stall::Item(Ok(content))) => {
return PreContent::Ready {
stream,
first: Some(Ok(content)),
};
}
Some(Stall::Idle) => {
let error = stall_error(policy.ttft);
if state.budget > 0 {
state.budget -= 1;
if !retry_wait(tx, policy, state.attempts, &error).await {
return PreContent::Cancelled;
}
continue;
}
return PreContent::Failed(error);
}
}
}
}
async fn pre_content_retry<S, F>(
policy: &RetryPolicy,
state: &mut RetryState,
open: &F,
tx: &tokio::sync::mpsc::UnboundedSender<LlmStreamEvent>,
error: &CompletionError,
) -> PreContent<S>
where
F: Fn() -> Pin<Box<dyn std::future::Future<Output = Result<S, CompletionError>> + Send>>,
S: futures::Stream<Item = Result<StreamedAssistantContent, CompletionError>> + Unpin + Send,
{
state.attempts += 1;
state.budget -= 1;
if !retry_wait(tx, policy, state.attempts, error).await {
return PreContent::Cancelled;
}
connect_with_retry(policy, state, open, tx).await
}
fn retry_after(error: &CompletionError) -> Option<Duration> {
let raw = error.provider_response_headers()?.get("retry-after")?;
let seconds: u64 = raw.to_str().ok()?.trim().parse().ok()?;
Some(Duration::from_secs(seconds))
}
async fn retry_wait(
tx: &tokio::sync::mpsc::UnboundedSender<LlmStreamEvent>,
policy: &RetryPolicy,
attempt: u32,
error: &CompletionError,
) -> bool {
let delay = retry_after(error)
.map(|d| d.min(policy.cap))
.unwrap_or_else(|| backoff(attempt, policy.base, policy.cap));
let _ = tx.send(LlmStreamEvent::Retrying {
attempt,
max: policy.max,
reason: error.to_string(),
});
tokio::select! {
biased;
_ = tx.closed() => false,
_ = tokio::time::sleep(delay) => true,
}
}
fn build_request(
messages: &[AgentMessage],
system: &str,
tools: &[ToolDefinition],
opts: &LlmOpts,
) -> CompletionRequest {
let mut chat_history = Vec::with_capacity(messages.len() + 1);
if !system.is_empty() {
chat_history.push(Message::system(system));
}
chat_history.extend(messages.iter().filter_map(to_rig_message));
CompletionRequest {
model: None,
preamble: None,
chat_history,
documents: Vec::new(),
tools: tools.to_vec(),
temperature: opts.temperature,
max_tokens: None,
tool_choice: None,
additional_params: reasoning_params(opts),
output_schema: None,
record_telemetry_content: false,
}
}
const RESPONSES_EFFORT_LEVELS: &[&str] =
&["none", "minimal", "low", "medium", "high", "xhigh", "max"];
fn invalid_effort(opts: &LlmOpts) -> Option<String> {
if opts.endpoint == LlmEndpoint::Responses
&& let Some(effort) = opts.effort.as_deref()
&& !RESPONSES_EFFORT_LEVELS.contains(&effort)
{
return Some(format!(
"effort `{effort}` is not valid for the Responses endpoint; use one of: {}",
RESPONSES_EFFORT_LEVELS.join(", ")
));
}
None
}
fn reasoning_params(opts: &LlmOpts) -> Option<serde_json::Value> {
match opts.endpoint {
LlmEndpoint::Chat => opts
.effort
.as_ref()
.map(|effort| serde_json::json!({ "reasoning_effort": effort })),
LlmEndpoint::Responses => {
let mut reasoning = serde_json::Map::new();
reasoning.insert("summary".into(), serde_json::json!("auto"));
if let Some(effort) = &opts.effort {
reasoning.insert("effort".into(), serde_json::json!(effort));
}
Some(serde_json::json!({ "reasoning": reasoning }))
}
}
}
const TOOL_OUTPUT_ERROR_PREFIX: &str = "ERROR: ";
fn to_rig_message(m: &AgentMessage) -> Option<Message> {
match m {
AgentMessage::User { content } => {
let content: Vec<UserContent> = content
.iter()
.filter_map(|b| match b {
ContentBlock::Text { text } => Some(UserContent::text(text.clone())),
_ => None,
})
.collect();
(!content.is_empty()).then_some(Message::User { content })
}
AgentMessage::Assistant { content, .. } => {
let content: Vec<AssistantContent> = content
.iter()
.map(|b| match b {
ContentBlock::Text { text } => AssistantContent::text(text.clone()),
ContentBlock::Thinking { text } => {
AssistantContent::Reasoning(Reasoning::new(text))
}
ContentBlock::ToolCall {
id,
name,
arguments,
} => AssistantContent::ToolCall(ToolCall::new(
ToolCallId::new_or_mint(id.clone()),
ToolFunction::new(name.clone(), arguments.clone()),
)),
})
.collect();
(!content.is_empty()).then_some(Message::Assistant { id: None, content })
}
AgentMessage::ToolResult {
tool_call_id,
name,
output,
is_error,
} => {
let output = if *is_error {
format!("{TOOL_OUTPUT_ERROR_PREFIX}{output}")
} else {
output.clone()
};
Some(Message::tool_result(
tool_call_id.clone(),
name.clone(),
output,
))
}
}
}
fn map_item(item: StreamedAssistantContent) -> Vec<LlmStreamEvent> {
match item {
StreamedAssistantContent::Text(text) => vec![LlmStreamEvent::TextDelta(text.text)],
StreamedAssistantContent::ReasoningDelta { reasoning, .. } => {
vec![LlmStreamEvent::ThinkingDelta(reasoning)]
}
StreamedAssistantContent::Reasoning { reasoning, .. } => {
let text = reasoning_text(&reasoning);
if text.is_empty() {
vec![]
} else {
vec![LlmStreamEvent::ThinkingReplace(text)]
}
}
StreamedAssistantContent::ToolCallDelta { .. } => vec![],
StreamedAssistantContent::ToolCall { tool_call, .. } => vec![LlmStreamEvent::ToolCall {
id: tool_call.id.as_str().to_string(),
name: tool_call.function.name,
arguments: tool_call.function.arguments,
}],
StreamedAssistantContent::Final(final_record) => {
vec![LlmStreamEvent::Done {
stop_reason: final_record
.finish_reason
.map(finish_reason)
.unwrap_or(StopReason::Stop),
usage: Some(convert_usage(final_record.usage)),
}]
}
StreamedAssistantContent::Unknown(_) => vec![],
}
}
fn finish_reason(reason: FinishReason) -> StopReason {
match reason {
FinishReason::Stop => StopReason::Stop,
FinishReason::Length => StopReason::Length,
FinishReason::ToolCalls => StopReason::ToolUse,
FinishReason::ContentFilter | FinishReason::Other(_) => StopReason::Stop,
}
}
fn convert_usage(usage: rig::completion::Usage) -> crate::message::Usage {
crate::message::Usage {
input_tokens: usage.input_tokens,
output_tokens: usage.output_tokens,
cache_read_tokens: (usage.cached_input_tokens > 0).then_some(usage.cached_input_tokens),
cache_write_tokens: (usage.cache_creation_input_tokens > 0)
.then_some(usage.cache_creation_input_tokens),
}
}
fn reasoning_text(reasoning: &Reasoning) -> String {
reasoning
.content
.iter()
.filter_map(|block| match block {
ReasoningContent::Text { text, .. } => Some(text.as_str()),
ReasoningContent::Summary(summary) => Some(summary.as_str()),
ReasoningContent::Redacted { data } => Some(data.as_str()),
ReasoningContent::Encrypted(_) => None,
})
.collect::<Vec<_>>()
.join("\n")
}
#[cfg(test)]
mod tests {
use super::*;
use rig::completion::Usage as RigUsage;
use rig::message::{ToolResultContent, UserContent};
use rig::streaming::{RawStreamingToolCall, StreamFinal, ToolCallDeltaContent};
use serde_json::json;
fn tool_call_item() -> StreamedAssistantContent {
let raw =
RawStreamingToolCall::new("call_1".to_string(), "run".to_string(), json!({"x": 1}));
StreamedAssistantContent::ToolCall {
tool_call: raw.into(),
internal_call_id: "ic1".to_string(),
}
}
fn final_item(finish_reason: Option<FinishReason>) -> StreamedAssistantContent {
let usage = RigUsage {
input_tokens: 10,
output_tokens: 5,
total_tokens: 15,
cached_input_tokens: 3,
cache_creation_input_tokens: 2,
..RigUsage::default()
};
let final_record = StreamFinal::new("test", usage);
StreamedAssistantContent::Final(match finish_reason {
Some(reason) => final_record.with_finish_reason(reason),
None => final_record,
})
}
#[test]
fn text_item_maps_to_text_delta() {
let events = map_item(StreamedAssistantContent::Text(rig::message::Text::new(
"hi",
)));
assert_eq!(events, vec![LlmStreamEvent::TextDelta("hi".to_string())]);
}
#[test]
fn reasoning_delta_maps_to_thinking_delta() {
let events = map_item(StreamedAssistantContent::ReasoningDelta {
id: "r1".to_string(),
provider_id: None,
reasoning: "hmm".to_string(),
});
assert_eq!(
events,
vec![LlmStreamEvent::ThinkingDelta("hmm".to_string())]
);
}
#[test]
fn reasoning_block_replaces_with_full_text() {
let events = map_item(StreamedAssistantContent::Reasoning {
reasoning: Reasoning::new("full thought"),
id: "r1".to_string(),
});
assert_eq!(
events,
vec![LlmStreamEvent::ThinkingReplace("full thought".to_string())]
);
}
#[test]
fn reasoning_text_joins_summary_text_and_redacted_drops_encrypted() {
let reasoning = Reasoning {
id: None,
content: vec![
ReasoningContent::Summary("a".to_string()),
ReasoningContent::Text {
text: "b".to_string(),
signature: None,
},
ReasoningContent::Redacted {
data: "c".to_string(),
},
ReasoningContent::Encrypted("opaque".to_string()),
],
};
assert_eq!(reasoning_text(&reasoning), "a\nb\nc");
}
#[test]
fn reasoning_text_summary_only_is_the_summary() {
let reasoning = Reasoning::summaries(vec!["just a summary".to_string()]);
assert_eq!(reasoning_text(&reasoning), "just a summary");
}
#[test]
fn reasoning_text_encrypted_only_is_empty() {
let reasoning = Reasoning::encrypted("opaque");
assert_eq!(reasoning_text(&reasoning), "");
}
#[test]
fn summary_reasoning_block_replaces_with_full_text() {
let events = map_item(StreamedAssistantContent::Reasoning {
reasoning: Reasoning::summaries(vec!["a summary".to_string()]),
id: "rs_1".to_string(),
});
assert_eq!(
events,
vec![LlmStreamEvent::ThinkingReplace("a summary".to_string())]
);
}
#[test]
fn encrypted_only_reasoning_block_emits_no_event() {
let events = map_item(StreamedAssistantContent::Reasoning {
reasoning: Reasoning::encrypted("opaque"),
id: "rs_1".to_string(),
});
assert!(events.is_empty(), "got: {events:?}");
}
#[test]
fn tool_call_delta_ignored() {
for content in [
ToolCallDeltaContent::Name("run".to_string()),
ToolCallDeltaContent::Delta("{\"x".to_string()),
] {
let events = map_item(StreamedAssistantContent::ToolCallDelta {
internal_call_id: "ic1".to_string(),
content,
});
assert!(events.is_empty());
}
}
#[test]
fn tool_call_maps_to_call_event() {
let events = map_item(tool_call_item());
assert_eq!(
events,
vec![LlmStreamEvent::ToolCall {
id: "call_1".to_string(),
name: "run".to_string(),
arguments: json!({"x": 1}),
}]
);
}
#[test]
fn final_maps_finish_reason_and_usage() {
let cases = [
(FinishReason::Stop, StopReason::Stop),
(FinishReason::Length, StopReason::Length),
(FinishReason::ToolCalls, StopReason::ToolUse),
(FinishReason::ContentFilter, StopReason::Stop),
(FinishReason::Other("weird".to_string()), StopReason::Stop),
];
for (reason, expected) in cases {
let events = map_item(final_item(Some(reason)));
assert_eq!(
events,
vec![LlmStreamEvent::Done {
stop_reason: expected,
usage: Some(crate::message::Usage {
input_tokens: 10,
output_tokens: 5,
cache_read_tokens: Some(3),
cache_write_tokens: Some(2),
}),
}],
"wrong mapping for {expected:?}"
);
}
}
#[test]
fn final_without_finish_reason_defaults_to_stop() {
let events = map_item(final_item(None));
match &events[..] {
[LlmStreamEvent::Done { stop_reason, .. }] => {
assert_eq!(*stop_reason, StopReason::Stop)
}
other => panic!("unexpected events: {other:?}"),
}
}
#[test]
fn unknown_item_ignored() {
let events = map_item(StreamedAssistantContent::Unknown(
rig::streaming::UnknownPayload::new(json!({"mystery": 1})),
));
assert!(events.is_empty());
}
#[test]
fn build_request_round_trips_context() {
let opts = LlmOpts {
model: "m1".to_string(),
base_url: None,
api_key: None,
temperature: Some(0.5),
..LlmOpts::default()
};
let tools = vec![ToolDefinition {
name: "get_weather".to_string(),
description: "Weather lookup".to_string(),
parameters: json!({"type": "object"}),
}];
let messages = vec![
AgentMessage::user_text("What is the weather?"),
AgentMessage::Assistant {
content: vec![
ContentBlock::Thinking { text: "hmm".into() },
ContentBlock::ToolCall {
id: "t1".into(),
name: "get_weather".into(),
arguments: json!({"city": "Tokyo"}),
},
ContentBlock::Text {
text: "Let me check.".into(),
},
],
stop_reason: StopReason::ToolUse,
usage: None,
model: None,
},
AgentMessage::ToolResult {
tool_call_id: "t1".into(),
name: "get_weather".into(),
output: "sunny".into(),
is_error: false,
},
];
let request = build_request(&messages, "You are terse.", &tools, &opts);
assert!(request.preamble.is_none());
assert_eq!(request.temperature, Some(0.5));
assert_eq!(request.tools.len(), 1);
assert_eq!(request.tools[0].name, "get_weather");
let [
Message::System { content: system },
Message::User { content: user },
Message::Assistant {
content: assistant, ..
},
Message::User {
content: tool_result,
},
] = &request.chat_history[..]
else {
panic!("unexpected chat_history shape: {:?}", request.chat_history);
};
assert_eq!(system, "You are terse.");
assert!(
matches!(user.as_slice(), [UserContent::Text(t)] if t.text == "What is the weather?")
);
assert_eq!(assistant.len(), 3);
assert!(
matches!(&assistant[0], AssistantContent::Reasoning(r) if reasoning_text(r) == "hmm")
);
assert!(matches!(
&assistant[1],
AssistantContent::ToolCall(call)
if call.id.as_str() == "t1"
&& call.function.name == "get_weather"
&& call.function.arguments == json!({"city": "Tokyo"})
));
assert!(matches!(&assistant[2], AssistantContent::Text(t) if t.text == "Let me check."));
assert!(matches!(
tool_result.as_slice(),
[UserContent::ToolResult(r)] if r.call.as_str() == "t1" && r.name == "get_weather"
));
}
#[test]
fn tool_result_error_is_annotated_on_the_wire_only() {
let err = AgentMessage::ToolResult {
tool_call_id: "t1".into(),
name: "echo".into(),
output: "boom".into(),
is_error: true,
};
let ok = AgentMessage::ToolResult {
tool_call_id: "t1".into(),
name: "echo".into(),
output: "boom".into(),
is_error: false,
};
fn wire_text(m: &AgentMessage) -> String {
match to_rig_message(m) {
Some(Message::User { content }) => match &content[..] {
[UserContent::ToolResult(r)] => match &r.content[..] {
[ToolResultContent::Text(t)] => t.text.clone(),
other => panic!("unexpected tool result content: {other:?}"),
},
other => panic!("unexpected user content: {other:?}"),
},
other => panic!("unexpected wire message: {other:?}"),
}
}
assert_eq!(wire_text(&ok), "boom");
assert_eq!(wire_text(&err), "ERROR: boom");
match err {
AgentMessage::ToolResult {
output, is_error, ..
} => {
assert_eq!(output, "boom");
assert!(is_error);
}
_ => unreachable!(),
}
}
#[test]
fn endpoint_defaults_to_chat() {
let opts = LlmOpts {
model: "m1".to_string(),
base_url: None,
api_key: None,
temperature: None,
endpoint: LlmEndpoint::Chat,
effort: None,
session_id: None,
retry: RetryPolicy::default(),
model_profiles: ModelProfiles::default(),
};
assert_eq!(opts.endpoint, LlmEndpoint::Chat);
assert_eq!(LlmOpts::default().endpoint, LlmEndpoint::Chat);
}
#[test]
fn no_effort_sends_no_additional_params() {
let opts = LlmOpts::default();
let request = build_request(&[], "sys", &[], &opts);
assert_eq!(request.additional_params, None);
}
#[test]
fn chat_effort_maps_to_reasoning_effort_param() {
let opts = LlmOpts {
effort: Some("high".into()),
endpoint: LlmEndpoint::Chat,
..LlmOpts::default()
};
let request = build_request(&[], "sys", &[], &opts);
assert_eq!(
request.additional_params,
Some(json!({ "reasoning_effort": "high" }))
);
}
#[test]
fn responses_effort_maps_to_reasoning_object() {
let opts = LlmOpts {
effort: Some("xhigh".into()),
endpoint: LlmEndpoint::Responses,
..LlmOpts::default()
};
let request = build_request(&[], "sys", &[], &opts);
assert_eq!(
request.additional_params,
Some(json!({ "reasoning": { "effort": "xhigh", "summary": "auto" } }))
);
}
#[test]
fn responses_always_requests_a_reasoning_summary() {
let opts = LlmOpts {
endpoint: LlmEndpoint::Responses,
effort: None,
..LlmOpts::default()
};
let request = build_request(&[], "sys", &[], &opts);
assert_eq!(
request.additional_params,
Some(json!({ "reasoning": { "summary": "auto" } }))
);
}
#[test]
fn chat_effort_is_verbatim_passthrough() {
for level in ["low", "medium", "my-custom-level", ""] {
let opts = LlmOpts {
effort: Some(level.into()),
endpoint: LlmEndpoint::Chat,
..LlmOpts::default()
};
let request = build_request(&[], "sys", &[], &opts);
assert_eq!(
request.additional_params,
Some(json!({ "reasoning_effort": level })),
"level {level:?} must pass through verbatim on Chat"
);
}
}
#[test]
fn responses_effort_validation() {
for level in ["none", "minimal", "low", "medium", "high", "xhigh", "max"] {
let opts = LlmOpts {
effort: Some(level.into()),
endpoint: LlmEndpoint::Responses,
..LlmOpts::default()
};
assert_eq!(invalid_effort(&opts), None, "level {level} must be valid");
}
let opts = LlmOpts {
effort: Some("deep".into()),
endpoint: LlmEndpoint::Responses,
..LlmOpts::default()
};
let msg = invalid_effort(&opts).expect("unknown level must be rejected");
assert!(msg.contains("not valid"), "got: {msg}");
assert!(
msg.contains("xhigh"),
"message must list valid levels: {msg}"
);
let opts = LlmOpts {
effort: Some("deep".into()),
endpoint: LlmEndpoint::Chat,
..LlmOpts::default()
};
assert_eq!(invalid_effort(&opts), None);
}
#[test]
fn client_sends_session_header_and_user_agent() {
let opts = LlmOpts {
api_key: Some("test-key".to_string()),
session_id: Some("sess-123".to_string()),
..LlmOpts::default()
};
let headers = openai_client(&opts)
.expect("client builds")
.headers()
.clone();
assert_eq!(
headers.get("x-opencode-session").expect("session header"),
"sess-123"
);
assert!(headers.contains_key("user-agent"), "got: {headers:?}");
}
#[test]
fn client_omits_session_header_when_unset() {
let opts = LlmOpts {
api_key: Some("test-key".to_string()),
..LlmOpts::default()
};
let headers = openai_client(&opts)
.expect("client builds")
.headers()
.clone();
assert!(!headers.contains_key("x-opencode-session"));
assert!(headers.contains_key("user-agent"), "got: {headers:?}");
}
#[test]
fn localhost_base_url_builds_without_key() {
for url in [
"http://localhost:11434/v1",
"http://127.0.0.1:8080/v1",
"http://[::1]:8080/v1",
] {
let opts = LlmOpts {
base_url: Some(url.to_string()),
api_key: None,
..LlmOpts::default()
};
assert!(
openai_client(&opts).is_ok(),
"keyless {url} must build a client"
);
}
}
#[test]
fn remote_base_url_without_key_is_error() {
let opts = LlmOpts {
base_url: Some("https://api.example.com/v1".to_string()),
api_key: None,
..LlmOpts::default()
};
let err = openai_client(&opts).unwrap_err();
assert!(err.contains("no API key"), "got: {err}");
}
#[tokio::test]
async fn list_models_unreachable_is_error() {
let opts = LlmOpts {
base_url: Some("http://127.0.0.1:9/v1".to_string()),
api_key: Some("test-key".to_string()),
..LlmOpts::default()
};
let err = tokio::time::timeout(std::time::Duration::from_secs(10), list_models(&opts))
.await
.expect("list terminates")
.expect_err("unreachable host must error");
assert!(!err.is_empty());
}
#[tokio::test]
async fn list_models_without_key_is_error() {
let opts = LlmOpts::default();
let err = list_models(&opts).await.expect_err("no key must error");
assert!(err.contains("no API key"), "got: {err}");
}
#[test]
fn client_cache_reuses_until_key_changes() {
let opts = LlmOpts {
base_url: Some("http://127.0.0.1:9/v1".to_string()),
api_key: Some("k".to_string()),
session_id: Some("s1".to_string()),
..LlmOpts::default()
};
let cache = ClientCache::default();
let builds = || cache.builds.load(std::sync::atomic::Ordering::SeqCst);
cache.get(&opts).expect("builds");
cache.get(&opts).expect("builds");
assert_eq!(builds(), 1, "identical opts must reuse the client");
let model_swap = LlmOpts {
model: "other".to_string(),
effort: Some("high".to_string()),
..opts.clone()
};
cache.get(&model_swap).expect("builds");
assert_eq!(builds(), 1, "model/effort swap must reuse the client");
for changed in [
LlmOpts {
base_url: Some("http://127.0.0.1:8/v1".to_string()),
..opts.clone()
},
LlmOpts {
api_key: Some("k2".to_string()),
..opts.clone()
},
LlmOpts {
session_id: Some("s2".to_string()),
..opts.clone()
},
] {
let before = builds();
cache.get(&changed).expect("builds");
assert_eq!(builds(), before + 1, "a key change must rebuild the client");
}
}
#[test]
fn settle_determines_the_client_key() {
let base = LlmProvider {
endpoint: EndpointPref::Pinned(LlmEndpoint::Chat),
base_url: Some("http://127.0.0.1:9/v1".to_string()),
api_key: Some("k".to_string()),
};
let mut profiles = ModelProfiles::new(base.clone());
profiles.insert(
"ep-swap",
LlmProfile {
endpoint: Some(LlmEndpoint::Responses),
..LlmProfile::default()
},
);
profiles.insert(
"moved",
LlmProfile {
base_url: Some("http://127.0.0.1:8/v1".to_string()),
..LlmProfile::default()
},
);
let settled = |model: &str| {
let mut llm = LlmOpts {
model: model.to_string(),
endpoint: base.endpoint.resolve(base.base_url.as_deref(), model),
base_url: base.base_url.clone(),
api_key: base.api_key.clone(),
model_profiles: profiles.clone(),
..LlmOpts::default()
};
llm.settle_provider();
llm
};
let cache = ClientCache::default();
let builds = || cache.builds.load(std::sync::atomic::Ordering::SeqCst);
let plain = settled("m0");
cache.get(&plain).expect("builds");
assert_eq!(builds(), 1);
let ep_swap = settled("ep-swap");
assert_eq!(ep_swap.endpoint, LlmEndpoint::Responses);
assert_eq!(ep_swap.base_url, base.base_url, "the base_url is inherited");
cache.get(&ep_swap).expect("reuses");
assert_eq!(builds(), 1, "an endpoint-only swap must reuse the client");
let moved = settled("moved");
cache.get(&moved).expect("rebuilds");
assert_eq!(builds(), 2, "a base_url change must rebuild the client");
}
const ZEN: &str = "https://opencode.ai/zen/go/v1";
#[test]
fn endpoint_pref_resolve_matrix() {
assert_eq!(
EndpointPref::Pinned(LlmEndpoint::Chat).resolve(Some(ZEN), "grok-4.7"),
LlmEndpoint::Chat
);
assert_eq!(
EndpointPref::Auto.resolve(Some(ZEN), "grok-4.7"),
LlmEndpoint::Responses
);
assert_eq!(
EndpointPref::Auto.resolve(Some(ZEN), "deepseek-v4-flash"),
LlmEndpoint::Chat
);
assert_eq!(
EndpointPref::Auto.resolve(Some("https://api.openai.com/v1"), "grok-4.7"),
LlmEndpoint::Chat
);
assert_eq!(
EndpointPref::Auto.resolve(None, "grok-4.7"),
LlmEndpoint::Chat
);
}
#[test]
fn endpoint_pref_wire_and_from_option() {
assert_eq!(EndpointPref::Auto.wire(), "auto");
assert_eq!(EndpointPref::Pinned(LlmEndpoint::Chat).wire(), "chat");
assert_eq!(EndpointPref::Pinned(LlmEndpoint::Responses).wire(), "responses");
assert_eq!(EndpointPref::from(None), EndpointPref::Auto);
assert_eq!(
EndpointPref::from(Some(LlmEndpoint::Responses)),
EndpointPref::Pinned(LlmEndpoint::Responses)
);
}
#[test]
fn settle_auto_derives_from_the_catalog() {
let base = LlmProvider {
endpoint: EndpointPref::Auto,
base_url: Some(ZEN.to_string()),
api_key: None,
};
let profiles = ModelProfiles::new(base.clone());
let settled = |model: &str| {
let mut llm = LlmOpts {
model: model.to_string(),
endpoint: LlmEndpoint::Chat,
base_url: base.base_url.clone(),
model_profiles: profiles.clone(),
..LlmOpts::default()
};
llm.settle_provider();
llm
};
assert_eq!(settled("grok-4.7").endpoint, LlmEndpoint::Responses);
assert_eq!(settled("deepseek-v4-flash").endpoint, LlmEndpoint::Chat);
}
#[test]
fn settle_profile_pin_beats_auto_and_base_pin() {
let base = LlmProvider {
endpoint: EndpointPref::Auto,
base_url: Some(ZEN.to_string()),
api_key: None,
};
let mut profiles = ModelProfiles::new(base.clone());
profiles.insert(
"grok-4.7",
LlmProfile {
endpoint: Some(LlmEndpoint::Chat),
..LlmProfile::default()
},
);
let mut llm = LlmOpts {
model: "grok-4.7".into(),
base_url: base.base_url.clone(),
model_profiles: profiles,
..LlmOpts::default()
};
llm.settle_provider();
assert_eq!(
llm.endpoint,
LlmEndpoint::Chat,
"a profile pin must beat the catalog"
);
}
#[test]
fn settle_profile_base_url_only_re_resolves_the_endpoint() {
let base = LlmProvider {
endpoint: EndpointPref::Auto,
base_url: Some("http://plain/v1".into()),
api_key: None,
};
let mut profiles = ModelProfiles::new(base.clone());
profiles.insert(
"grok-4.7",
LlmProfile {
base_url: Some(ZEN.to_string()),
..LlmProfile::default()
},
);
let mut llm = LlmOpts {
model: "grok-4.7".into(),
base_url: base.base_url.clone(),
model_profiles: profiles,
..LlmOpts::default()
};
llm.settle_provider();
assert_eq!(llm.base_url.as_deref(), Some(ZEN));
assert_eq!(
llm.endpoint,
LlmEndpoint::Responses,
"a base_url-only profile must re-derive against the new base"
);
}
#[test]
fn set_model_flips_chat_to_responses_and_back() {
let base = LlmProvider {
endpoint: EndpointPref::Auto,
base_url: Some(ZEN.to_string()),
api_key: None,
};
let profiles = ModelProfiles::new(base.clone());
let mut llm = LlmOpts {
model: "deepseek-v4-flash".into(),
base_url: base.base_url.clone(),
model_profiles: profiles,
..LlmOpts::default()
};
llm.settle_provider();
assert_eq!(llm.endpoint, LlmEndpoint::Chat);
llm.model = "grok-4.7".into();
llm.settle_provider();
assert_eq!(llm.endpoint, LlmEndpoint::Responses);
llm.model = "deepseek-v4-flash".into();
llm.settle_provider();
assert_eq!(llm.endpoint, LlmEndpoint::Chat);
}
#[test]
fn responses_request_builds_against_default_client_type() {
use rig::client::CompletionClient;
let client = openai::Client::builder()
.api_key("test-key")
.base_url("http://127.0.0.1:9/v1")
.build()
.expect("client builds without network");
let model = client.completion_model("some-model");
let request = build_request(&[], "sys", &[], &LlmOpts::default());
let fut = rig::completion::CompletionModel::stream(&model, request);
drop(fut);
}
#[tokio::test]
async fn adapter_surfaces_error_without_done() {
let opts = LlmOpts {
model: "m1".to_string(),
base_url: Some("http://127.0.0.1:9/v1".to_string()),
api_key: Some("test-key".to_string()),
temperature: None,
..LlmOpts::default()
};
let stream_fn = rig_stream_fn();
let stream = stream_fn(&[], "sys", &[], &opts);
let events: Vec<LlmStreamEvent> =
tokio::time::timeout(std::time::Duration::from_secs(10), stream.collect())
.await
.expect("stream terminates");
assert!(
events
.iter()
.any(|e| matches!(e, LlmStreamEvent::Error { .. })),
"expected an Error event, got {events:?}"
);
assert!(
!events
.iter()
.any(|e| matches!(e, LlmStreamEvent::Done { .. })),
"no Done after error, got {events:?}"
);
}
#[tokio::test]
async fn responses_adapter_surfaces_error_without_done() {
let opts = LlmOpts {
model: "m1".to_string(),
base_url: Some("http://127.0.0.1:9/v1".to_string()),
api_key: Some("test-key".to_string()),
temperature: None,
endpoint: LlmEndpoint::Responses,
effort: None,
session_id: None,
retry: RetryPolicy::default(),
model_profiles: ModelProfiles::default(),
};
let stream_fn = rig_stream_fn();
let stream = stream_fn(&[], "sys", &[], &opts);
let events: Vec<LlmStreamEvent> =
tokio::time::timeout(std::time::Duration::from_secs(10), stream.collect())
.await
.expect("stream terminates");
assert!(
events
.iter()
.any(|e| matches!(e, LlmStreamEvent::Error { .. })),
"expected an Error event, got {events:?}"
);
assert!(
!events
.iter()
.any(|e| matches!(e, LlmStreamEvent::Done { .. })),
"no Done after error, got {events:?}"
);
}
#[tokio::test]
#[ignore = "live test; run with WCODE_LIVE=1 WCODE_BASE_URL=... WCODE_API_KEY=... WCODE_MODEL=... cargo test -p wcode-harness --lib -- --ignored"]
async fn live_rig_stream_fn() {
let (Ok(base_url), Ok(api_key), Ok(model)) = (
std::env::var("WCODE_BASE_URL"),
std::env::var("WCODE_API_KEY"),
std::env::var("WCODE_MODEL"),
) else {
eprintln!("skipping: WCODE_BASE_URL / WCODE_API_KEY / WCODE_MODEL not set");
return;
};
if std::env::var("WCODE_LIVE").ok().as_deref() != Some("1") {
eprintln!("skipping: WCODE_LIVE != 1");
return;
}
let opts = LlmOpts {
model,
base_url: Some(base_url),
api_key: Some(api_key),
temperature: None,
..LlmOpts::default()
};
let stream_fn = rig_stream_fn();
let stream = stream_fn(
&[AgentMessage::user_text("Reply with exactly: pong")],
"You are a terse echo assistant.",
&[],
&opts,
);
let events: Vec<LlmStreamEvent> = stream.collect().await;
let text = events
.iter()
.filter_map(|e| match e {
LlmStreamEvent::TextDelta(t) => Some(t.as_str()),
_ => None,
})
.collect::<String>();
println!("text: {text:?}");
println!("events: {events:?}");
assert!(
events
.iter()
.any(|e| matches!(e, LlmStreamEvent::Done { .. }))
);
assert!(
!events
.iter()
.any(|e| matches!(e, LlmStreamEvent::Error { .. }))
);
}
}
#[cfg(test)]
mod retry_tests {
use super::*;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
type Chunks =
futures::stream::Iter<std::vec::IntoIter<Result<StreamedAssistantContent, CompletionError>>>;
fn transient() -> CompletionError {
CompletionError::ProviderError(
"Http client error: error sending request for url (http://x)".into(),
)
}
fn fatal() -> CompletionError {
CompletionError::ProviderError("invalid_request_error: bad tool schema".into())
}
fn empty() -> Chunks {
futures::stream::iter(Vec::new())
}
fn done() -> StreamedAssistantContent {
StreamedAssistantContent::Final(rig::streaming::StreamFinal::new(
"test",
rig::completion::Usage::default(),
))
}
fn first_item_err(e: CompletionError) -> Chunks {
futures::stream::iter(vec![Err(e)])
}
#[test]
fn transient_statuses_retry_client_errors_do_not() {
for code in [408, 409, 429, 500, 502, 503, 504] {
assert!(is_transient_status(code), "{code} should retry");
}
for code in [200, 400, 401, 403, 404, 422] {
assert!(!is_transient_status(code), "{code} must not retry");
}
}
#[test]
fn retryable_classifies_transport_vs_rejections() {
assert!(retryable(&CompletionError::HttpError(
rig::http_client::Error::StreamEnded
)));
assert!(retryable(&transient()));
assert!(!retryable(&fatal()));
assert!(!retryable(&CompletionError::ResponseError("bad".into())));
}
#[test]
fn backoff_grows_and_never_exceeds_the_cap() {
let base = Duration::from_millis(100);
let cap = Duration::from_millis(400);
let d1 = backoff(1, base, cap);
assert!(d1 <= base && d1 >= base * 3 / 4, "attempt 1 = {d1:?}");
for n in 1..12 {
assert!(backoff(n, base, cap) <= cap, "attempt {n} exceeds cap");
}
}
#[tokio::test(start_paused = true)]
async fn connect_retries_transient_failures_then_succeeds() {
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
let calls = Arc::new(AtomicUsize::new(0));
let policy = RetryPolicy {
max: 3,
base: Duration::from_millis(10),
cap: Duration::from_millis(50),
..RetryPolicy::default()
};
let c = calls.clone();
let mut state = RetryState::new(policy.max);
let out = connect_with_retry(
&policy,
&mut state,
move || {
let n = c.fetch_add(1, Ordering::SeqCst) + 1;
Box::pin(async move {
if n <= 2 {
Err(transient())
} else {
Ok(empty())
}
})
},
&tx,
)
.await;
assert!(matches!(out, PreContent::Ready { .. }));
assert_eq!(calls.load(Ordering::SeqCst), 3);
assert!(matches!(
rx.try_recv(),
Ok(LlmStreamEvent::Retrying { attempt: 1, .. })
));
assert!(matches!(
rx.try_recv(),
Ok(LlmStreamEvent::Retrying { attempt: 2, .. })
));
}
#[tokio::test(start_paused = true)]
async fn connect_retries_a_first_item_transport_error() {
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
let calls = Arc::new(AtomicUsize::new(0));
let policy = RetryPolicy {
max: 2,
base: Duration::from_millis(5),
cap: Duration::from_millis(5),
..RetryPolicy::default()
};
let c = calls.clone();
let mut state = RetryState::new(policy.max);
let out = connect_with_retry(
&policy,
&mut state,
move || {
let n = c.fetch_add(1, Ordering::SeqCst) + 1;
Box::pin(async move {
if n == 1 {
Ok(first_item_err(transient()))
} else {
Ok(empty())
}
})
},
&tx,
)
.await;
assert!(matches!(out, PreContent::Ready { .. }));
assert_eq!(calls.load(Ordering::SeqCst), 2, "first-item error retried");
}
#[tokio::test(start_paused = true)]
async fn connect_gives_up_after_max_and_skips_non_retryable() {
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
let policy = RetryPolicy {
max: 1,
base: Duration::from_millis(1),
cap: Duration::from_millis(1),
..RetryPolicy::default()
};
let calls = Arc::new(AtomicUsize::new(0));
let c = calls.clone();
let mut state = RetryState::new(policy.max);
let out = connect_with_retry(
&policy,
&mut state,
move || {
c.fetch_add(1, Ordering::SeqCst);
Box::pin(async { Err::<Chunks, CompletionError>(transient()) })
},
&tx,
)
.await;
assert!(matches!(out, PreContent::Failed(_)));
assert_eq!(calls.load(Ordering::SeqCst), 2, "one attempt + one retry");
let calls = Arc::new(AtomicUsize::new(0));
let c = calls.clone();
let mut state = RetryState::new(policy.max);
let out = connect_with_retry(
&policy,
&mut state,
move || {
c.fetch_add(1, Ordering::SeqCst);
Box::pin(async { Err::<Chunks, CompletionError>(fatal()) })
},
&tx,
)
.await;
assert!(matches!(out, PreContent::Failed(_)));
assert_eq!(calls.load(Ordering::SeqCst), 1, "non-retryable returns at once");
}
#[test]
fn fatal_class_flags_hard_errors_but_not_transient_transport() {
assert!(!fatal_class(&transient()));
assert!(!fatal_class(&CompletionError::HttpError(
rig::http_client::Error::StreamEnded
)));
assert!(fatal_class(&fatal()));
assert!(fatal_class(&CompletionError::ResponseError("bad".into())));
assert!(fatal_class(&CompletionError::RequestError("bad".into())));
let json_err = serde_json::from_str::<u8>("not a number").unwrap_err();
assert!(fatal_class(&CompletionError::JsonError(json_err)));
}
#[tokio::test]
async fn refused_port_exhausts_retries_to_non_fatal_error() {
let opts = LlmOpts {
model: "m1".to_string(),
base_url: Some("http://127.0.0.1:9/v1".to_string()),
api_key: Some("test-key".to_string()),
retry: RetryPolicy {
max: 2,
base: Duration::from_millis(10),
cap: Duration::from_millis(50),
..RetryPolicy::default()
},
..LlmOpts::default()
};
let stream_fn = rig_stream_fn();
let stream = stream_fn(&[], "sys", &[], &opts);
let events: Vec<LlmStreamEvent> =
tokio::time::timeout(std::time::Duration::from_secs(10), stream.collect())
.await
.expect("stream terminates");
let retries = events
.iter()
.filter(|e| matches!(e, LlmStreamEvent::Retrying { .. }))
.count();
assert_eq!(retries, 2, "two backoff notices before giving up: {events:?}");
assert!(
matches!(events.last(), Some(LlmStreamEvent::Error { fatal: false, .. })),
"refused port is transient-class, not fatal: {events:?}"
);
}
fn err_with_retry_after(value: &str) -> CompletionError {
let mut headers = HeaderMap::new();
headers.insert("retry-after", HeaderValue::from_str(value).unwrap());
CompletionError::ProviderResponse(
rig::ProviderResponseError::without_status("429")
.with_headers(Some(Box::new(headers))),
)
}
#[test]
fn retry_after_parses_integer_seconds() {
assert_eq!(
retry_after(&err_with_retry_after("7")),
Some(Duration::from_secs(7))
);
}
#[test]
fn retry_after_is_none_when_absent() {
assert_eq!(retry_after(&transient()), None);
}
#[test]
fn retry_after_is_none_on_non_numeric() {
assert_eq!(
retry_after(&err_with_retry_after("Wed, 21 Oct 2015 07:28:00 GMT")),
None
);
assert_eq!(retry_after(&err_with_retry_after("soon")), None);
}
#[tokio::test(start_paused = true)]
async fn retry_wait_honors_retry_after_over_backoff() {
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
let policy = RetryPolicy {
base: Duration::from_millis(10),
..RetryPolicy::default()
};
let start = tokio::time::Instant::now();
assert!(retry_wait(&tx, &policy, 1, &err_with_retry_after("3")).await);
assert_eq!(start.elapsed(), Duration::from_secs(3));
}
#[tokio::test(start_paused = true)]
async fn retry_wait_clamps_retry_after_to_the_cap() {
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
let policy = RetryPolicy {
cap: Duration::from_secs(8),
..RetryPolicy::default()
};
let start = tokio::time::Instant::now();
assert!(retry_wait(&tx, &policy, 1, &err_with_retry_after("100")).await);
assert_eq!(start.elapsed(), Duration::from_secs(8));
}
type BoxStream =
Pin<Box<dyn futures::Stream<Item = Result<StreamedAssistantContent, CompletionError>> + Send>>;
fn boxed<S>(s: S) -> BoxStream
where
S: futures::Stream<Item = Result<StreamedAssistantContent, CompletionError>>
+ Send
+ 'static,
{
Box::pin(s)
}
fn silent_boxed() -> BoxStream {
boxed(futures::stream::pending::<Result<StreamedAssistantContent, CompletionError>>())
}
fn open_silent()
-> Pin<Box<dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send>> {
Box::pin(async { Ok::<_, CompletionError>(silent_boxed()) })
}
fn open_slow_ok()
-> Pin<Box<dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send>> {
Box::pin(async {
tokio::time::sleep(Duration::from_millis(50)).await;
Ok::<_, CompletionError>(boxed(futures::stream::iter(vec![Ok::<_, CompletionError>(
StreamedAssistantContent::Text(rig::message::Text::new("ok")),
)])))
})
}
#[test]
fn a_stall_error_is_transient_and_not_hard_fatal() {
let d = Duration::from_secs(1);
assert!(retryable(&stall_error(d)), "a stall must be transient (B1)");
assert!(!fatal_class(&stall_error(d)), "a stall must not be hard-fatal");
assert!(
stall_error(d).to_string().contains("timed out"),
"the (timed out) needle drives retryable()"
);
}
#[tokio::test(start_paused = true)]
async fn a_silent_stream_trips_the_ttft_retries_then_fails() {
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
let calls = Arc::new(AtomicUsize::new(0));
let policy = RetryPolicy {
max: 2,
base: Duration::from_millis(10),
cap: Duration::from_millis(10),
ttft: Duration::from_secs(1),
..RetryPolicy::default()
};
let c = calls.clone();
let open = move || {
c.fetch_add(1, Ordering::SeqCst);
open_silent()
};
let mut state = RetryState::new(policy.max);
let out = connect_with_retry(&policy, &mut state, open, &tx).await;
let PreContent::Failed(error) = out else {
panic!("expected Failed from an exhausted stall")
};
assert!(retryable(&error), "a stall is transient");
assert!(!fatal_class(&error), "a stall is not hard-fatal");
assert_eq!(calls.load(Ordering::SeqCst), 3, "one attempt + two retries");
let retries = std::iter::from_fn(|| rx.try_recv().ok())
.filter(|e| matches!(e, LlmStreamEvent::Retrying { .. }))
.count();
assert_eq!(retries, 2, "a Retrying notice per retry");
}
#[tokio::test(start_paused = true)]
async fn deltas_then_a_stall_is_a_nonfatal_error() {
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
let policy = RetryPolicy {
max: 3,
base: Duration::from_millis(10),
cap: Duration::from_millis(10),
idle: Duration::from_secs(1),
..RetryPolicy::default()
};
let stream = boxed(
futures::stream::iter(vec![Ok::<_, CompletionError>(
StreamedAssistantContent::Text(rig::message::Text::new("hi")),
)])
.chain(futures::stream::pending::<Result<StreamedAssistantContent, CompletionError>>()),
);
let state = RetryState::new(policy.max);
forward_stream(stream, None, policy, open_silent, tx, state).await;
let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
assert!(
events
.iter()
.any(|e| matches!(e, LlmStreamEvent::TextDelta(t) if t == "hi")),
"{events:?}"
);
let error = events
.iter()
.find_map(|e| match e {
LlmStreamEvent::Error { message, fatal } => Some((message.clone(), *fatal)),
_ => None,
})
.expect("the idle deadline surfaced an error");
assert!(!error.1, "a post-content stall is non-fatal: {error:?}");
assert!(
!events.iter().any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
"no retry once content flowed: {events:?}"
);
}
#[tokio::test(start_paused = true)]
async fn a_clean_stream_is_unaffected() {
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
let policy = RetryPolicy {
idle: Duration::from_secs(1),
..RetryPolicy::default()
};
let stream = boxed(futures::stream::iter(vec![
Ok::<_, CompletionError>(StreamedAssistantContent::Text(rig::message::Text::new("hi"))),
Ok::<_, CompletionError>(done()),
]));
let state = RetryState::new(policy.max);
forward_stream(stream, None, policy, open_silent, tx, state).await;
let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
assert!(
events
.iter()
.any(|e| matches!(e, LlmStreamEvent::TextDelta(t) if t == "hi")),
"{events:?}"
);
assert!(
events.iter().any(|e| matches!(e, LlmStreamEvent::Done { .. })),
"{events:?}"
);
assert!(
!events.iter().any(|e| matches!(e, LlmStreamEvent::Error { .. })),
"{events:?}"
);
assert!(
!events.iter().any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
"{events:?}"
);
}
#[tokio::test(start_paused = true)]
async fn zero_deadlines_disable_the_timeout() {
let mut silent =
futures::stream::pending::<Result<StreamedAssistantContent, CompletionError>>();
let stalled = tokio::time::timeout(
Duration::from_secs(1),
next_or_stall(&mut silent, Duration::ZERO),
)
.await;
assert!(stalled.is_err(), "ZERO must not trip the deadline");
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
let policy = RetryPolicy {
ttft: Duration::ZERO,
..RetryPolicy::default()
};
let mut state = RetryState::new(policy.max);
let out = connect_with_retry(&policy, &mut state, open_slow_ok, &tx).await;
assert!(
matches!(out, PreContent::Ready { .. }),
"a slow connect must not be timed out under ZERO"
);
}
#[tokio::test(start_paused = true)]
async fn a_stalled_connect_trips_the_ttft_retries_then_fails() {
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
let calls = Arc::new(AtomicUsize::new(0));
let policy = RetryPolicy {
max: 2,
base: Duration::from_millis(10),
cap: Duration::from_millis(10),
ttft: Duration::from_secs(1),
..RetryPolicy::default()
};
let c = calls.clone();
let open = move || {
c.fetch_add(1, Ordering::SeqCst);
Box::pin(std::future::pending::<Result<BoxStream, CompletionError>>())
as Pin<
Box<
dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send,
>,
>
};
let mut state = RetryState::new(policy.max);
let out = connect_with_retry(&policy, &mut state, open, &tx).await;
let PreContent::Failed(error) = out else {
panic!("expected Failed from an exhausted connect stall")
};
assert!(retryable(&error), "a connect stall is transient");
assert!(!fatal_class(&error), "a connect stall is not hard-fatal");
assert_eq!(calls.load(Ordering::SeqCst), 3, "one attempt + two retries");
let retries = std::iter::from_fn(|| rx.try_recv().ok())
.filter(|e| matches!(e, LlmStreamEvent::Retrying { .. }))
.count();
assert_eq!(retries, 2, "a Retrying notice per retry");
}
#[tokio::test(start_paused = true)]
async fn a_pre_content_idle_stall_retries_and_reopens() {
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
let policy = RetryPolicy {
max: 2,
base: Duration::from_millis(10),
cap: Duration::from_millis(10),
idle: Duration::from_secs(1),
..RetryPolicy::default()
};
let stream = silent_boxed();
fn open_done()
-> Pin<Box<dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send>> {
Box::pin(async {
Ok::<_, CompletionError>(boxed(futures::stream::iter(vec![Ok::<_, CompletionError>(
done(),
)])))
})
}
let state = RetryState::new(policy.max);
forward_stream(stream, None, policy, open_done, tx, state).await;
let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
assert!(
events
.iter()
.any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
"a pre-content stall backs off: {events:?}"
);
assert!(
events.iter().any(|e| matches!(e, LlmStreamEvent::Done { .. })),
"the re-opened stream completed: {events:?}"
);
assert!(
!events.iter().any(|e| matches!(e, LlmStreamEvent::Error { .. })),
"no error once the retry succeeded: {events:?}"
);
}
fn open_empty()
-> Pin<Box<dyn std::future::Future<Output = Result<BoxStream, CompletionError>> + Send>> {
Box::pin(async { Ok::<_, CompletionError>(boxed(empty())) })
}
#[tokio::test(start_paused = true)]
async fn a_delta_then_a_truncated_end_is_a_nonfatal_error() {
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
let policy = RetryPolicy {
max: 3,
base: Duration::from_millis(10),
cap: Duration::from_millis(10),
idle: Duration::ZERO,
..RetryPolicy::default()
};
let stream = boxed(futures::stream::iter(vec![Ok::<_, CompletionError>(
StreamedAssistantContent::Text(rig::message::Text::new("hi")),
)]));
let state = RetryState::new(policy.max);
forward_stream(stream, None, policy, open_silent, tx, state).await;
let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
assert!(
events
.iter()
.any(|e| matches!(e, LlmStreamEvent::TextDelta(t) if t == "hi")),
"{events:?}"
);
let (message, fatal) = events
.iter()
.find_map(|e| match e {
LlmStreamEvent::Error { message, fatal } => Some((message.clone(), *fatal)),
_ => None,
})
.expect("a truncated end surfaced an error");
assert!(!fatal, "a post-content truncation is non-fatal: {message}");
assert!(
!events
.iter()
.any(|e| matches!(e, LlmStreamEvent::Done { .. })),
"no Done on a truncated stream: {events:?}"
);
assert!(
!events
.iter()
.any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
"no retry once content flowed: {events:?}"
);
}
#[tokio::test(start_paused = true)]
async fn a_final_terminated_stream_is_unchanged() {
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
let policy = RetryPolicy {
idle: Duration::ZERO,
..RetryPolicy::default()
};
let stream = boxed(futures::stream::iter(vec![
Ok::<_, CompletionError>(StreamedAssistantContent::Text(rig::message::Text::new(
"hi",
))),
Ok::<_, CompletionError>(done()),
]));
let state = RetryState::new(policy.max);
forward_stream(stream, None, policy, open_silent, tx, state).await;
let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
assert!(
events
.iter()
.any(|e| matches!(e, LlmStreamEvent::Done { .. })),
"{events:?}"
);
assert!(
!events
.iter()
.any(|e| matches!(e, LlmStreamEvent::Error { .. })),
"a clean turn surfaces no error: {events:?}"
);
assert!(
!events
.iter()
.any(|e| matches!(e, LlmStreamEvent::Retrying { .. })),
"{events:?}"
);
}
#[tokio::test(start_paused = true)]
async fn a_pre_content_no_item_stream_retries_then_errors() {
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
let policy = RetryPolicy {
max: 2,
base: Duration::from_millis(10),
cap: Duration::from_millis(10),
idle: Duration::ZERO,
..RetryPolicy::default()
};
let state = RetryState::new(policy.max);
forward_stream(boxed(empty()), None, policy, open_empty, tx, state).await;
let events: Vec<LlmStreamEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
let retries = events
.iter()
.filter(|e| matches!(e, LlmStreamEvent::Retrying { .. }))
.count();
assert_eq!(
retries, policy.max as usize,
"one Retrying notice per retry: {events:?}"
);
let (message, fatal) = events
.iter()
.find_map(|e| match e {
LlmStreamEvent::Error { message, fatal } => Some((message.clone(), *fatal)),
_ => None,
})
.expect("an exhausted truncation budget surfaced an error");
assert!(!fatal, "an exhausted truncation is non-fatal: {message}");
assert!(
!events
.iter()
.any(|e| matches!(e, LlmStreamEvent::Done { .. })),
"no Done on a truncated stream: {events:?}"
);
}
fn spawn_models_server(delay: std::time::Duration) -> String {
use std::io::{Read as _, Write as _};
use std::net::TcpListener;
let listener = TcpListener::bind("127.0.0.1:0").expect("bind ephemeral port");
let addr = listener.local_addr().expect("local addr");
std::thread::spawn(move || {
for stream in listener.incoming() {
let Ok(mut stream) = stream else { continue };
let mut buf = Vec::new();
let mut chunk = [0u8; 1024];
loop {
match stream.read(&mut chunk) {
Ok(0) | Err(_) => break,
Ok(n) => {
buf.extend_from_slice(&chunk[..n]);
if buf.windows(4).any(|w| w == b"\r\n\r\n") {
break;
}
}
}
}
std::thread::sleep(delay);
let body = "{\"object\":\"list\",\"data\":[{\"id\":\"m1\",\"object\":\"model\"}]}";
let response = format!(
"HTTP/1.1 200 OK\r\n\
Content-Type: application/json\r\n\
Connection: close\r\n\
Content-Length: {}\r\n\r\n{}",
body.len(),
body
);
let _ = stream.write_all(response.as_bytes());
let _ = stream.flush();
}
});
format!("http://{addr}/v1")
}
#[tokio::test]
async fn a_refused_port_probes_false() {
let up = probe_endpoint(
"http://127.0.0.1:9/v1",
std::time::Duration::from_millis(500),
)
.await;
assert!(!up, "a refused port is not an endpoint");
}
#[tokio::test]
async fn a_live_models_endpoint_probes_true() {
let _ = probe_endpoint("http://127.0.0.1:9/v1", std::time::Duration::from_secs(60)).await;
let base = spawn_models_server(std::time::Duration::ZERO);
let up = probe_endpoint(&base, std::time::Duration::from_secs(5)).await;
assert!(up, "a live /models endpoint probes true");
}
#[tokio::test]
async fn a_slow_endpoint_times_out_to_false() {
let base = spawn_models_server(std::time::Duration::from_millis(500));
let up = probe_endpoint(&base, std::time::Duration::from_millis(50)).await;
assert!(!up, "a slow endpoint is not treated as up");
}
}