2307 lines
77 KiB
Rust
2307 lines
77 KiB
Rust
use super::*;
|
|
|
|
use super::access_token::{distrust_access_token, get_access_token};
|
|
use crate::config::{RenderMode, paths};
|
|
use crate::hooks::{self, HookEvent};
|
|
use crate::{
|
|
config::{AppConfig, Input, RequestContext},
|
|
function::{FunctionDeclaration, ToolCall, ToolResult, eval_tool_calls},
|
|
render::render_stream,
|
|
utils::*,
|
|
};
|
|
|
|
use crate::vault::Vault;
|
|
use anyhow::{Context, Result, bail};
|
|
use fancy_regex::Regex;
|
|
use indexmap::IndexMap;
|
|
use inquire::{
|
|
MultiSelect, Select, Text, list_option::ListOption, required, validator::Validation,
|
|
};
|
|
use reqwest::{Client as ReqwestClient, RequestBuilder};
|
|
use serde::{Deserialize, Serialize};
|
|
use serde_json::{Value, json};
|
|
use std::sync::LazyLock;
|
|
use std::time::{Duration, Instant};
|
|
use tokio::sync::mpsc::unbounded_channel;
|
|
|
|
pub const MODELS_YAML: &str = include_str!("../../models.yaml");
|
|
|
|
pub static ALL_PROVIDER_MODELS: LazyLock<Vec<ProviderModels>> = LazyLock::new(|| {
|
|
paths::local_models_override()
|
|
.ok()
|
|
.unwrap_or_else(|| serde_yaml::from_str(MODELS_YAML).unwrap())
|
|
});
|
|
|
|
static EMBEDDING_MODEL_RE: LazyLock<Regex> = LazyLock::new(|| {
|
|
Regex::new(r"((^|/)(bge-|e5-|uae-|gte-|text-)|embed|multilingual|minilm)").unwrap()
|
|
});
|
|
|
|
static ESCAPE_SLASH_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"(?<!\\)/").unwrap());
|
|
|
|
#[async_trait::async_trait]
|
|
pub trait Client: Sync + Send {
|
|
fn app_config(&self) -> &AppConfig;
|
|
|
|
fn extra_config(&self) -> Option<&ExtraConfig>;
|
|
|
|
fn patch_config(&self) -> Option<&RequestPatch>;
|
|
|
|
fn name(&self) -> &str;
|
|
|
|
fn model(&self) -> &Model;
|
|
|
|
fn supports_oauth(&self) -> bool {
|
|
false
|
|
}
|
|
|
|
fn build_client(&self) -> Result<ReqwestClient> {
|
|
let mut builder = ReqwestClient::builder();
|
|
let extra = self.extra_config();
|
|
let timeout = extra.and_then(|v| v.connect_timeout).unwrap_or(10);
|
|
let read_timeout = extra.and_then(|v| v.read_timeout).unwrap_or(300);
|
|
if let Some(proxy) = extra.and_then(|v| v.proxy.as_deref()) {
|
|
builder = set_proxy(builder, proxy)?;
|
|
}
|
|
if let Some(user_agent) = self.app_config().user_agent.as_ref() {
|
|
builder = builder.user_agent(user_agent);
|
|
}
|
|
if read_timeout > 0 {
|
|
builder = builder.read_timeout(Duration::from_secs(read_timeout));
|
|
}
|
|
let client = builder
|
|
.connect_timeout(Duration::from_secs(timeout))
|
|
.build()
|
|
.with_context(|| "Failed to build client")?;
|
|
Ok(client)
|
|
}
|
|
|
|
/// On a 401 the cached access token is distrusted and the call retried
|
|
/// exactly once; the retry re-runs the per-client prepare step, which
|
|
/// sees the rejection marker, force-refreshes the token, and rebuilds
|
|
/// the whole request. A second 401 propagates the original error; any
|
|
/// other retry failure propagates as-is.
|
|
async fn chat_completions(&self, input: Input) -> Result<ChatCompletionsOutput> {
|
|
if self.app_config().dry_run {
|
|
let content = input.echo_messages();
|
|
return Ok(ChatCompletionsOutput::new(&content));
|
|
}
|
|
let client = self.build_client()?;
|
|
let data = input.prepare_completion_data(self.model(), false)?;
|
|
let err = match self.chat_completions_inner(&client, data).await {
|
|
Ok(output) => return Ok(output),
|
|
Err(err) => err,
|
|
};
|
|
let ret = if should_retry_auth(&err, self.name()) {
|
|
debug!(
|
|
"provider '{}' rejected access token (401); refreshing and retrying once",
|
|
self.name()
|
|
);
|
|
let data = input.prepare_completion_data(self.model(), false)?;
|
|
match self.chat_completions_inner(&client, data).await {
|
|
Err(retry_err) if is_auth_error(&retry_err) => Err(err),
|
|
ret => ret,
|
|
}
|
|
} else {
|
|
Err(err)
|
|
};
|
|
ret.with_context(|| "Failed to call chat-completions api")
|
|
}
|
|
|
|
/// Same retry-once-on-401 semantics as [`Self::chat_completions`], but
|
|
/// only while the handler has received nothing yet: retrying after
|
|
/// partial output has streamed would render it to the user twice. The
|
|
/// retry lives inside the same `select!` arm so abort stays responsive.
|
|
async fn chat_completions_streaming(
|
|
&self,
|
|
input: &Input,
|
|
handler: &mut SseHandler,
|
|
) -> Result<()> {
|
|
let abort_signal = handler.abort();
|
|
let input = input.clone();
|
|
tokio::select! {
|
|
ret = async {
|
|
if self.app_config().dry_run {
|
|
let content = input.echo_messages();
|
|
handler.text(&content)?;
|
|
return Ok(());
|
|
}
|
|
let client = self.build_client()?;
|
|
let data = input.prepare_completion_data(self.model(), true)?;
|
|
let err = match self.chat_completions_streaming_inner(&client, handler, data).await {
|
|
Ok(()) => return Ok(()),
|
|
Err(err) => err,
|
|
};
|
|
if handler.has_received_content() || !should_retry_auth(&err, self.name()) {
|
|
return Err(err);
|
|
}
|
|
debug!(
|
|
"provider '{}' rejected access token (401); refreshing and retrying once",
|
|
self.name()
|
|
);
|
|
let data = input.prepare_completion_data(self.model(), true)?;
|
|
match self.chat_completions_streaming_inner(&client, handler, data).await {
|
|
Err(retry_err) if is_auth_error(&retry_err) => Err(err),
|
|
ret => ret,
|
|
}
|
|
} => {
|
|
handler.done();
|
|
ret.with_context(|| "Failed to call chat-completions api")
|
|
}
|
|
_ = wait_abort_signal(&abort_signal) => {
|
|
handler.done();
|
|
Ok(())
|
|
},
|
|
}
|
|
}
|
|
|
|
/// Same retry-once-on-401 semantics as [`Self::chat_completions`]
|
|
/// (gemini OAuth embeddings route here).
|
|
async fn embeddings(&self, data: &EmbeddingsData) -> Result<Vec<Vec<f32>>> {
|
|
let client = self.build_client()?;
|
|
let err = match self.embeddings_inner(&client, data).await {
|
|
Ok(output) => return Ok(output),
|
|
Err(err) => err,
|
|
};
|
|
let ret = if should_retry_auth(&err, self.name()) {
|
|
debug!(
|
|
"provider '{}' rejected access token (401); refreshing and retrying once",
|
|
self.name()
|
|
);
|
|
match self.embeddings_inner(&client, data).await {
|
|
Err(retry_err) if is_auth_error(&retry_err) => Err(err),
|
|
ret => ret,
|
|
}
|
|
} else {
|
|
Err(err)
|
|
};
|
|
ret.context("Failed to call embeddings api")
|
|
}
|
|
|
|
async fn rerank(&self, data: &RerankData) -> Result<RerankOutput> {
|
|
let client = self.build_client()?;
|
|
self.rerank_inner(&client, data)
|
|
.await
|
|
.context("Failed to call rerank api")
|
|
}
|
|
|
|
async fn chat_completions_inner(
|
|
&self,
|
|
client: &ReqwestClient,
|
|
data: ChatCompletionsData,
|
|
) -> Result<ChatCompletionsOutput>;
|
|
|
|
async fn chat_completions_streaming_inner(
|
|
&self,
|
|
client: &ReqwestClient,
|
|
handler: &mut SseHandler,
|
|
data: ChatCompletionsData,
|
|
) -> Result<()>;
|
|
|
|
async fn embeddings_inner(
|
|
&self,
|
|
_client: &ReqwestClient,
|
|
_data: &EmbeddingsData,
|
|
) -> Result<EmbeddingsOutput> {
|
|
bail!("The client doesn't support embeddings api")
|
|
}
|
|
|
|
async fn rerank_inner(
|
|
&self,
|
|
_client: &ReqwestClient,
|
|
_data: &RerankData,
|
|
) -> Result<RerankOutput> {
|
|
bail!("The client doesn't support rerank api")
|
|
}
|
|
|
|
fn request_builder(
|
|
&self,
|
|
client: &reqwest::Client,
|
|
mut request_data: RequestData,
|
|
) -> RequestBuilder {
|
|
self.patch_request_data(&mut request_data);
|
|
request_data.into_builder(client)
|
|
}
|
|
|
|
/// Patches are chat-shaped unless keyed otherwise: on the Responses wire
|
|
/// only the `responses` sub-patch of a model patch and the `responses`
|
|
/// section of a client patch apply, so chat-only keys never leak into a
|
|
/// responses body.
|
|
fn patch_request_data(&self, request_data: &mut RequestData) {
|
|
let model_type = self.model().model_type();
|
|
let wire = request_data.wire;
|
|
if let Some(patch) = self.model().patch() {
|
|
match wire {
|
|
WireApi::Chat => request_data.apply_patch(patch.clone()),
|
|
WireApi::Responses => {
|
|
if let Some(patch) = patch.get("responses") {
|
|
request_data.apply_patch(patch.clone());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
let api_name = match wire {
|
|
WireApi::Responses => "responses",
|
|
WireApi::Chat => model_type.api_name(),
|
|
};
|
|
let patch_map = std::env::var(get_env_name(&format!(
|
|
"patch_{}_{api_name}",
|
|
self.model().client_name(),
|
|
)))
|
|
.ok()
|
|
.and_then(|v| serde_json::from_str(&v).ok())
|
|
.or_else(|| {
|
|
self.patch_config().and_then(|v| match wire {
|
|
WireApi::Responses => v.responses.clone(),
|
|
WireApi::Chat => model_type.extract_patch(v).cloned(),
|
|
})
|
|
});
|
|
let patch_map = match patch_map {
|
|
Some(v) => v,
|
|
_ => return,
|
|
};
|
|
for (key, patch) in patch_map {
|
|
let key = ESCAPE_SLASH_RE.replace_all(&key, r"\/");
|
|
if let Ok(regex) = Regex::new(&format!("^({key})$"))
|
|
&& let Ok(true) = regex.is_match(self.model().name())
|
|
{
|
|
request_data.apply_patch(patch);
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Default for ClientConfig {
|
|
fn default() -> Self {
|
|
Self::OpenAIConfig(OpenAIConfig::default())
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize, Default)]
|
|
pub struct ExtraConfig {
|
|
pub proxy: Option<String>,
|
|
pub connect_timeout: Option<u64>,
|
|
pub read_timeout: Option<u64>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize, Default)]
|
|
pub struct RequestPatch {
|
|
pub chat_completions: Option<ApiPatch>,
|
|
pub responses: Option<ApiPatch>,
|
|
pub embeddings: Option<ApiPatch>,
|
|
pub rerank: Option<ApiPatch>,
|
|
}
|
|
|
|
pub type ApiPatch = IndexMap<String, Value>;
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize)]
|
|
#[serde(rename_all = "lowercase")]
|
|
pub enum WireApi {
|
|
Responses,
|
|
Chat,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct RequestData {
|
|
pub url: String,
|
|
pub headers: IndexMap<String, String>,
|
|
pub body: Value,
|
|
pub wire: WireApi,
|
|
}
|
|
|
|
impl RequestData {
|
|
pub fn new<T>(url: T, body: Value) -> Self
|
|
where
|
|
T: std::fmt::Display,
|
|
{
|
|
Self {
|
|
url: url.to_string(),
|
|
headers: Default::default(),
|
|
body,
|
|
wire: WireApi::Chat,
|
|
}
|
|
}
|
|
|
|
pub fn wire(mut self, wire: WireApi) -> Self {
|
|
self.wire = wire;
|
|
self
|
|
}
|
|
|
|
pub fn bearer_auth<T>(&mut self, auth: T)
|
|
where
|
|
T: std::fmt::Display,
|
|
{
|
|
self.headers
|
|
.insert("authorization".into(), format!("Bearer {auth}"));
|
|
}
|
|
|
|
pub fn header<K, V>(&mut self, key: K, value: V)
|
|
where
|
|
K: std::fmt::Display,
|
|
V: std::fmt::Display,
|
|
{
|
|
self.headers.insert(key.to_string(), value.to_string());
|
|
}
|
|
|
|
pub fn into_builder(self, client: &ReqwestClient) -> RequestBuilder {
|
|
let RequestData {
|
|
url, headers, body, ..
|
|
} = self;
|
|
debug!("Request {url} {body}");
|
|
|
|
let mut builder = client.post(url);
|
|
for (key, value) in headers {
|
|
builder = builder.header(key, value);
|
|
}
|
|
builder = builder.json(&body);
|
|
builder
|
|
}
|
|
|
|
pub fn apply_patch(&mut self, patch: Value) {
|
|
if let Some(patch_url) = patch["url"].as_str() {
|
|
self.url = patch_url.into();
|
|
}
|
|
if let Some(patch_body) = patch.get("body") {
|
|
json_patch::merge(&mut self.body, patch_body)
|
|
}
|
|
if let Some(patch_headers) = patch["headers"].as_object() {
|
|
for (key, value) in patch_headers {
|
|
if let Some(value) = value.as_str() {
|
|
self.header(key, value)
|
|
} else if value.is_null() {
|
|
self.headers.swap_remove(key);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct ChatCompletionsData {
|
|
pub messages: Vec<Message>,
|
|
pub temperature: Option<f64>,
|
|
pub top_p: Option<f64>,
|
|
pub reasoning_effort: Option<String>,
|
|
pub functions: Option<Vec<FunctionDeclaration>>,
|
|
pub stream: bool,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct ChatCompletionsOutput {
|
|
pub text: String,
|
|
pub tool_calls: Vec<ToolCall>,
|
|
pub thinking: Vec<ThinkingBlock>,
|
|
pub usage: Option<TokenUsage>,
|
|
}
|
|
|
|
impl ChatCompletionsOutput {
|
|
pub fn new(text: &str) -> Self {
|
|
Self {
|
|
text: text.to_string(),
|
|
..Default::default()
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
|
|
pub struct TokenUsage {
|
|
pub input_tokens: Option<u64>,
|
|
pub output_tokens: Option<u64>,
|
|
pub cache_creation_input_tokens: Option<u64>,
|
|
pub cache_read_input_tokens: Option<u64>,
|
|
}
|
|
|
|
impl TokenUsage {
|
|
pub fn is_empty(&self) -> bool {
|
|
self.input_tokens.is_none()
|
|
&& self.output_tokens.is_none()
|
|
&& self.cache_creation_input_tokens.is_none()
|
|
&& self.cache_read_input_tokens.is_none()
|
|
}
|
|
|
|
pub fn accumulate(&mut self, other: &TokenUsage) {
|
|
fn add(acc: &mut Option<u64>, value: Option<u64>) {
|
|
if let Some(value) = value {
|
|
*acc = Some(acc.unwrap_or(0).saturating_add(value));
|
|
}
|
|
}
|
|
add(&mut self.input_tokens, other.input_tokens);
|
|
add(&mut self.output_tokens, other.output_tokens);
|
|
add(
|
|
&mut self.cache_creation_input_tokens,
|
|
other.cache_creation_input_tokens,
|
|
);
|
|
add(
|
|
&mut self.cache_read_input_tokens,
|
|
other.cache_read_input_tokens,
|
|
);
|
|
}
|
|
|
|
/// USD cost of this usage at the given model's prices (USD per 1M
|
|
/// tokens). The four buckets are disjoint on both Anthropic direct and
|
|
/// Bedrock Converse (`input_tokens` excludes cache tokens), so the terms
|
|
/// sum without adjustment. Strict on missing prices: any nonzero
|
|
/// component without a price yields `None` rather than a silent
|
|
/// undercount. All-zero usage is also `None` — nothing billed, nothing
|
|
/// to show.
|
|
pub fn cost_usd(&self, model: &Model) -> Option<f64> {
|
|
fn term(tokens: Option<u64>, price: Option<f64>) -> Option<f64> {
|
|
match tokens.unwrap_or(0) {
|
|
0 => Some(0.0),
|
|
tokens => Some(tokens as f64 * price? / 1_000_000.0),
|
|
}
|
|
}
|
|
let total = term(self.input_tokens, model.input_price())?
|
|
+ term(self.output_tokens, model.output_price())?
|
|
+ term(self.cache_read_input_tokens, model.cache_read_price())?
|
|
+ term(self.cache_creation_input_tokens, model.cache_write_price())?;
|
|
(total > 0.0).then_some(total)
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct EmbeddingsData {
|
|
pub texts: Vec<String>,
|
|
pub query: bool,
|
|
}
|
|
|
|
impl EmbeddingsData {
|
|
pub fn new(texts: Vec<String>, query: bool) -> Self {
|
|
Self { texts, query }
|
|
}
|
|
}
|
|
|
|
pub type EmbeddingsOutput = Vec<Vec<f32>>;
|
|
|
|
#[derive(Debug)]
|
|
pub struct RerankData {
|
|
pub query: String,
|
|
pub documents: Vec<String>,
|
|
pub top_n: usize,
|
|
}
|
|
|
|
impl RerankData {
|
|
pub fn new(query: String, documents: Vec<String>, top_n: usize) -> Self {
|
|
Self {
|
|
query,
|
|
documents,
|
|
top_n,
|
|
}
|
|
}
|
|
}
|
|
|
|
pub type RerankOutput = Vec<RerankResult>;
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
pub struct RerankResult {
|
|
pub index: usize,
|
|
}
|
|
|
|
pub type PromptAction<'a> = (&'a str, &'a str, Option<&'a str>, bool);
|
|
|
|
pub async fn create_config(
|
|
prompts: &[PromptAction<'static>],
|
|
client: &str,
|
|
vault: &Vault,
|
|
) -> Result<(String, Value)> {
|
|
let mut config = json!({
|
|
"type": client,
|
|
});
|
|
for (key, desc, help_message, is_secret) in prompts {
|
|
let env_name = format!("{client}-{key}")
|
|
.to_ascii_uppercase()
|
|
.replace("_", "-");
|
|
let required = std::env::var(&env_name).is_err();
|
|
let value = if !is_secret {
|
|
prompt_input_string(desc, required, *help_message)?
|
|
} else {
|
|
vault.add_secret(&env_name)?;
|
|
format!("{{{{{}}}}}", env_name)
|
|
};
|
|
if !value.is_empty() {
|
|
config[key] = value.into();
|
|
}
|
|
}
|
|
let model = set_client_models_config(&mut config, client).await?;
|
|
let clients = json!(vec![config]);
|
|
Ok((model, clients))
|
|
}
|
|
|
|
pub async fn create_openai_compatible_client_config(
|
|
client: &str,
|
|
) -> Result<Option<(String, Value)>> {
|
|
let api_base = OPENAI_COMPATIBLE_PROVIDERS
|
|
.into_iter()
|
|
.find(|(name, _)| client == *name)
|
|
.map(|(_, api_base)| api_base)
|
|
.unwrap_or("http(s)://{API_ADDR}/v1");
|
|
|
|
let name = if client == OpenAICompatibleClient::NAME {
|
|
let value = prompt_input_string("Provider Name", true, None)?;
|
|
value.replace(' ', "-")
|
|
} else {
|
|
client.to_string()
|
|
};
|
|
|
|
let mut config = json!({
|
|
"type": OpenAICompatibleClient::NAME,
|
|
"name": &name,
|
|
});
|
|
|
|
let api_base = if api_base.contains('{') {
|
|
prompt_input_string("API Base", true, Some(&format!("e.g. {api_base}")))?
|
|
} else {
|
|
api_base.to_string()
|
|
};
|
|
config["api_base"] = api_base.into();
|
|
|
|
let has_bundled_oauth = ALL_PROVIDER_MODELS
|
|
.iter()
|
|
.any(|p| p.provider == client && p.oauth.is_some());
|
|
|
|
let use_oauth = if has_bundled_oauth {
|
|
let choice = Select::new("Authentication method:", vec!["API Key", "OAuth"]).prompt()?;
|
|
choice == "OAuth"
|
|
} else {
|
|
false
|
|
};
|
|
|
|
if use_oauth {
|
|
config["auth"] = "oauth".into();
|
|
} else {
|
|
let api_key = prompt_input_string("API Key", false, None)?;
|
|
if !api_key.is_empty() {
|
|
config["api_key"] = api_key.into();
|
|
}
|
|
}
|
|
|
|
let model = set_client_models_config(&mut config, &name).await?;
|
|
let clients = json!(vec![config]);
|
|
Ok(Some((model, clients)))
|
|
}
|
|
|
|
fn llm_request_extras(client: &dyn Client) -> Vec<(&'static str, String)> {
|
|
vec![
|
|
("COYOTE_LLM_PROVIDER", client.name().to_string()),
|
|
("COYOTE_LLM_MODEL", client.model().id()),
|
|
]
|
|
}
|
|
|
|
fn llm_completed_extras(client: &dyn Client, started_at: Instant) -> Vec<(&'static str, String)> {
|
|
let mut extras = llm_request_extras(client);
|
|
extras.push((
|
|
"COYOTE_LLM_DURATION_MS",
|
|
started_at.elapsed().as_millis().to_string(),
|
|
));
|
|
extras
|
|
}
|
|
|
|
fn llm_failed_extras(
|
|
client: &dyn Client,
|
|
err: Option<&anyhow::Error>,
|
|
aborted: bool,
|
|
) -> Vec<(&'static str, String)> {
|
|
let mut extras = llm_request_extras(client);
|
|
if aborted {
|
|
extras.push(("COYOTE_LLM_ERROR_KIND", "aborted".to_string()));
|
|
} else if let Some(api_err) = err.and_then(|err| err.downcast_ref::<ApiStatusError>()) {
|
|
extras.push(("COYOTE_LLM_ERROR_KIND", "api_error".to_string()));
|
|
extras.push(("COYOTE_LLM_STATUS", api_err.status.to_string()));
|
|
} else {
|
|
extras.push(("COYOTE_LLM_ERROR_KIND", "other".to_string()));
|
|
}
|
|
extras.push((
|
|
"COYOTE_LLM_ERROR",
|
|
// The immediate error only, matching COYOTE_TOOL_ERROR; the context
|
|
// chain below it stays out of the env.
|
|
err.map(|err| err.to_string())
|
|
.unwrap_or_else(|| "Aborted.".to_string()),
|
|
));
|
|
|
|
extras
|
|
}
|
|
|
|
pub async fn call_chat_completions(
|
|
input: &Input,
|
|
print: bool,
|
|
extract_code: bool,
|
|
client: &dyn Client,
|
|
ctx: &mut RequestContext,
|
|
abort_signal: AbortSignal,
|
|
) -> Result<(String, Vec<ToolResult>)> {
|
|
let is_child_agent = ctx.current_depth > 0;
|
|
let suppress_spinner = is_child_agent || ctx.render_mode == RenderMode::Silent;
|
|
let spinner_message = if suppress_spinner { "" } else { "Generating" };
|
|
hooks::fire(
|
|
HookEvent::LlmRequestStarted,
|
|
ctx,
|
|
&llm_request_extras(client),
|
|
None,
|
|
);
|
|
let request_started_at = Instant::now();
|
|
let ret = abortable_run_with_spinner(
|
|
client.chat_completions(input.clone()),
|
|
spinner_message,
|
|
abort_signal.clone(),
|
|
)
|
|
.await;
|
|
|
|
match ret {
|
|
Ok(ret) => {
|
|
hooks::fire(
|
|
HookEvent::LlmRequestCompleted,
|
|
ctx,
|
|
&llm_completed_extras(client, request_started_at),
|
|
None,
|
|
);
|
|
let ChatCompletionsOutput {
|
|
mut text,
|
|
tool_calls,
|
|
thinking,
|
|
usage,
|
|
} = ret;
|
|
if !text.is_empty() {
|
|
if extract_code {
|
|
text = extract_code_block(&strip_think_tag(&text)).to_string();
|
|
}
|
|
if print {
|
|
ctx.app.config.print_markdown(&text)?;
|
|
}
|
|
}
|
|
ctx.record_token_usage(usage, client.model());
|
|
finish_completion(ctx, text, tool_calls, thinking).await
|
|
}
|
|
Err(err) => {
|
|
hooks::fire(
|
|
HookEvent::LlmRequestFailed,
|
|
ctx,
|
|
&llm_failed_extras(client, Some(&err), abort_signal.aborted()),
|
|
None,
|
|
);
|
|
Err(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
pub async fn call_chat_completions_streaming(
|
|
input: &Input,
|
|
client: &dyn Client,
|
|
ctx: &mut RequestContext,
|
|
abort_signal: AbortSignal,
|
|
) -> Result<(String, Vec<ToolResult>)> {
|
|
let (tx, rx) = unbounded_channel();
|
|
let mut handler = SseHandler::new(tx, abort_signal.clone());
|
|
let silent = ctx.render_mode == RenderMode::Silent;
|
|
if silent {
|
|
handler.set_silent(true);
|
|
}
|
|
|
|
hooks::fire(
|
|
HookEvent::LlmRequestStarted,
|
|
ctx,
|
|
&llm_request_extras(client),
|
|
None,
|
|
);
|
|
let request_started_at = Instant::now();
|
|
let (send_ret, render_ret) = tokio::join!(
|
|
client.chat_completions_streaming(input, &mut handler),
|
|
render_stream(rx, client.app_config(), abort_signal.clone(), silent),
|
|
);
|
|
|
|
let aborted_ctrlc = handler.abort().aborted_ctrlc();
|
|
let aborted_ctrld = handler.abort().aborted_ctrld();
|
|
|
|
// Record usage before any early return: the provider bills an interrupted
|
|
// call's input tokens (reported in message_start), so abort and error
|
|
// paths must still count toward the session totals.
|
|
let (text, tool_calls, thinking, usage) = handler.take();
|
|
ctx.record_token_usage(usage, client.model());
|
|
|
|
if aborted_ctrld {
|
|
hooks::fire(
|
|
HookEvent::LlmRequestFailed,
|
|
ctx,
|
|
&llm_failed_extras(client, None, true),
|
|
None,
|
|
);
|
|
bail!("Aborted.");
|
|
}
|
|
|
|
if let Err(err) = render_ret {
|
|
hooks::fire(
|
|
HookEvent::LlmRequestFailed,
|
|
ctx,
|
|
&llm_failed_extras(client, Some(&err), false),
|
|
None,
|
|
);
|
|
return Err(err);
|
|
}
|
|
|
|
if aborted_ctrlc {
|
|
hooks::fire(
|
|
HookEvent::LlmRequestFailed,
|
|
ctx,
|
|
&llm_failed_extras(client, None, true),
|
|
None,
|
|
);
|
|
if !ctx.working_mode.is_repl() || ctx.session.is_none() {
|
|
bail!("Aborted.");
|
|
}
|
|
|
|
if text.is_empty() {
|
|
if !silent && *IS_STDOUT_TERMINAL {
|
|
println!();
|
|
eprintln!("{}", error_text("Response interrupted"));
|
|
}
|
|
|
|
return Ok(("".to_string(), vec![]));
|
|
}
|
|
|
|
if !silent && *IS_STDOUT_TERMINAL {
|
|
println!();
|
|
eprintln!("{}", error_text("Response interrupted"));
|
|
}
|
|
|
|
return Ok((text, vec![]));
|
|
}
|
|
|
|
match send_ret {
|
|
Ok(_) => {
|
|
hooks::fire(
|
|
HookEvent::LlmRequestCompleted,
|
|
ctx,
|
|
&llm_completed_extras(client, request_started_at),
|
|
None,
|
|
);
|
|
if !silent && !text.is_empty() && !text.ends_with('\n') {
|
|
println!();
|
|
}
|
|
finish_completion(ctx, text, tool_calls, thinking).await
|
|
}
|
|
Err(err) => {
|
|
if !silent && !text.is_empty() {
|
|
println!();
|
|
}
|
|
hooks::fire(
|
|
HookEvent::LlmRequestFailed,
|
|
ctx,
|
|
&llm_failed_extras(client, Some(&err), false),
|
|
None,
|
|
);
|
|
Err(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Streaming transport for graph `llm` nodes: accumulates the reply and the
|
|
/// transport itself renders nothing; no spinner, delta, or reasoning
|
|
/// (tool-call rendering inside `eval_tool_calls` is unchanged from the
|
|
/// non-streaming path). Because chunks arrive as they are generated, the
|
|
/// reqwest `read_timeout` bounds only stalls between chunks; the caller's
|
|
/// own deadline bounds the whole generation. Parity with the non-streaming
|
|
/// transport is kept deliberately: the handler's call-loop detection is off
|
|
/// (that path never had it), and a model the catalog marks `no_stream`
|
|
/// takes the non-streaming request, raced against the abort. An abort
|
|
/// observed after the provider call fails outright. Partial output is
|
|
/// never returned as success.
|
|
pub async fn call_chat_completions_streaming_quiet(
|
|
input: &Input,
|
|
client: &dyn Client,
|
|
ctx: &mut RequestContext,
|
|
abort_signal: AbortSignal,
|
|
) -> Result<(String, Vec<ToolResult>)> {
|
|
hooks::fire(
|
|
HookEvent::LlmRequestStarted,
|
|
ctx,
|
|
&llm_request_extras(client),
|
|
None,
|
|
);
|
|
let request_started_at = Instant::now();
|
|
if client.model().no_stream() {
|
|
let output = tokio::select! {
|
|
ret = client.chat_completions(input.clone()) => match ret {
|
|
Ok(output) => output,
|
|
Err(err) => {
|
|
hooks::fire(
|
|
HookEvent::LlmRequestFailed,
|
|
ctx,
|
|
&llm_failed_extras(client, Some(&err), false),
|
|
None,
|
|
);
|
|
return Err(err);
|
|
}
|
|
},
|
|
_ = wait_abort_signal(&abort_signal) => {
|
|
hooks::fire(
|
|
HookEvent::LlmRequestFailed,
|
|
ctx,
|
|
&llm_failed_extras(client, None, true),
|
|
None,
|
|
);
|
|
bail!("Aborted.")
|
|
}
|
|
};
|
|
let ChatCompletionsOutput {
|
|
text,
|
|
tool_calls,
|
|
thinking,
|
|
usage,
|
|
} = output;
|
|
ctx.record_token_usage(usage, client.model());
|
|
if abort_signal.aborted() {
|
|
hooks::fire(
|
|
HookEvent::LlmRequestFailed,
|
|
ctx,
|
|
&llm_failed_extras(client, None, true),
|
|
None,
|
|
);
|
|
bail!("Aborted.");
|
|
}
|
|
hooks::fire(
|
|
HookEvent::LlmRequestCompleted,
|
|
ctx,
|
|
&llm_completed_extras(client, request_started_at),
|
|
None,
|
|
);
|
|
return finish_completion(ctx, text, tool_calls, thinking).await;
|
|
}
|
|
|
|
// `_rx` outlives the call so `handler.done()` still finds a receiver.
|
|
let (tx, _rx) = unbounded_channel();
|
|
let mut handler = SseHandler::new(tx, abort_signal.clone());
|
|
handler.set_silent(true);
|
|
handler.set_call_loop_detection(false);
|
|
|
|
let send_ret = client.chat_completions_streaming(input, &mut handler).await;
|
|
|
|
let (text, tool_calls, thinking, usage) = handler.take();
|
|
ctx.record_token_usage(usage, client.model());
|
|
|
|
if abort_signal.aborted() {
|
|
hooks::fire(
|
|
HookEvent::LlmRequestFailed,
|
|
ctx,
|
|
&llm_failed_extras(client, None, true),
|
|
None,
|
|
);
|
|
bail!("Aborted.");
|
|
}
|
|
if let Err(err) = send_ret {
|
|
hooks::fire(
|
|
HookEvent::LlmRequestFailed,
|
|
ctx,
|
|
&llm_failed_extras(client, Some(&err), false),
|
|
None,
|
|
);
|
|
return Err(err);
|
|
}
|
|
|
|
hooks::fire(
|
|
HookEvent::LlmRequestCompleted,
|
|
ctx,
|
|
&llm_completed_extras(client, request_started_at),
|
|
None,
|
|
);
|
|
finish_completion(ctx, text, tool_calls, thinking).await
|
|
}
|
|
|
|
/// Callers record token usage themselves before calling, so that abort and
|
|
/// error paths that never reach here still count exactly once per API call.
|
|
async fn finish_completion(
|
|
ctx: &mut RequestContext,
|
|
text: String,
|
|
tool_calls: Vec<ToolCall>,
|
|
thinking: Vec<ThinkingBlock>,
|
|
) -> Result<(String, Vec<ToolResult>)> {
|
|
let mut tool_results = eval_tool_calls(ctx, tool_calls).await?;
|
|
if let Some(first) = tool_results.first_mut() {
|
|
first.thinking = thinking;
|
|
}
|
|
tool_results
|
|
.iter()
|
|
.for_each(|res| ctx.tool_scope.tool_tracker.record_call(res.call.clone()));
|
|
Ok((text, tool_results))
|
|
}
|
|
|
|
pub fn noop_prepare_rerank<T>(_client: &T, _data: &RerankData) -> Result<RequestData> {
|
|
bail!("The client doesn't support rerank api")
|
|
}
|
|
|
|
pub async fn noop_rerank(_builder: RequestBuilder, _model: &Model) -> Result<RerankOutput> {
|
|
bail!("The client doesn't support rerank api")
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct ApiStatusError {
|
|
pub status: u16,
|
|
pub message: String,
|
|
}
|
|
|
|
impl std::fmt::Display for ApiStatusError {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
write!(f, "{}", self.message)
|
|
}
|
|
}
|
|
|
|
impl std::error::Error for ApiStatusError {}
|
|
|
|
/// True when the error chain bottoms out in an [`ApiStatusError`] with
|
|
/// status 401 EXACTLY. 403 (entitlement) and 429 (rate limit) are never
|
|
/// auth failures, and message text is never inspected.
|
|
fn is_auth_error(err: &anyhow::Error) -> bool {
|
|
err.downcast_ref::<ApiStatusError>()
|
|
.is_some_and(|api_err| api_err.status == 401)
|
|
}
|
|
|
|
/// Decides whether a 401 from `client_name` warrants a single retry after a
|
|
/// forced token refresh: the error must be a 401 [`ApiStatusError`], and the
|
|
/// client must have a cached access token to distrust (API-key clients have
|
|
/// none and never retry). Distrusting marks the exact rejected token so the
|
|
/// retry's prepare step force-refreshes it. There is deliberately no backoff:
|
|
/// the blast radius is bounded at one extra request per user-visible call.
|
|
///
|
|
/// Note: vertexai shares the ACCESS_TOKENS cache, so a 401 there also
|
|
/// triggers distrust+retry — deliberate.
|
|
fn should_retry_auth(err: &anyhow::Error, client_name: &str) -> bool {
|
|
if !is_auth_error(err) {
|
|
return false;
|
|
}
|
|
let Ok(token) = get_access_token(client_name) else {
|
|
return false;
|
|
};
|
|
distrust_access_token(client_name, &token)
|
|
}
|
|
|
|
pub fn catch_error(data: &Value, status: u16) -> Result<()> {
|
|
if (200..300).contains(&status) {
|
|
return Ok(());
|
|
}
|
|
debug!("Invalid response, status: {status}, data: {data}");
|
|
let api_error = |message: String| anyhow::Error::new(ApiStatusError { status, message });
|
|
if let Some(error) = data["error"].as_object() {
|
|
if let (Some(typ), Some(message)) = (
|
|
json_str_from_map(error, "type"),
|
|
json_str_from_map(error, "message"),
|
|
) {
|
|
return Err(api_error(format!("{message} (type: {typ})")));
|
|
} else if let (Some(typ), Some(message)) = (
|
|
json_str_from_map(error, "code"),
|
|
json_str_from_map(error, "message"),
|
|
) {
|
|
return Err(api_error(format!("{message} (code: {typ})")));
|
|
}
|
|
} else if let Some(error) = data["errors"][0].as_object() {
|
|
if let (Some(code), Some(message)) = (
|
|
error.get("code").and_then(|v| v.as_u64()),
|
|
json_str_from_map(error, "message"),
|
|
) {
|
|
return Err(api_error(format!("{message} (status: {code})")));
|
|
}
|
|
} else if let Some(error) = data[0]["error"].as_object() {
|
|
if let (Some(status), Some(message)) = (
|
|
json_str_from_map(error, "status"),
|
|
json_str_from_map(error, "message"),
|
|
) {
|
|
return Err(api_error(format!("{message} (status: {status})")));
|
|
}
|
|
} else if let (Some(detail), Some(status)) = (data["detail"].as_str(), data["status"].as_i64())
|
|
{
|
|
return Err(api_error(format!("{detail} (status: {status})")));
|
|
} else if let Some(error) = data["error"].as_str() {
|
|
return Err(api_error(error.to_string()));
|
|
} else if let Some(message) = data["message"].as_str() {
|
|
return Err(api_error(message.to_string()));
|
|
}
|
|
Err(api_error(format!(
|
|
"Invalid response data: {data} (status: {status})"
|
|
)))
|
|
}
|
|
|
|
pub fn json_str_from_map<'a>(
|
|
map: &'a serde_json::Map<String, Value>,
|
|
field_name: &str,
|
|
) -> Option<&'a str> {
|
|
map.get(field_name).and_then(|v| v.as_str())
|
|
}
|
|
|
|
pub async fn set_client_models_config(client_config: &mut Value, client: &str) -> Result<String> {
|
|
if let Some(provider) = ALL_PROVIDER_MODELS.iter().find(|v| v.provider == client) {
|
|
let models: Vec<String> = provider
|
|
.models
|
|
.iter()
|
|
.filter(|v| v.model_type == "chat")
|
|
.map(|v| v.name.clone())
|
|
.collect();
|
|
let model_name = select_model(models)?;
|
|
return Ok(format!("{client}:{model_name}"));
|
|
}
|
|
let mut model_names = vec![];
|
|
if let (Some(true), Some(api_base), api_key) = (
|
|
client_config["type"]
|
|
.as_str()
|
|
.map(|v| v == OpenAICompatibleClient::NAME),
|
|
client_config["api_base"].as_str(),
|
|
client_config["api_key"]
|
|
.as_str()
|
|
.map(|v| v.to_string())
|
|
.or_else(|| {
|
|
let env_name = format!("{client}_api_key").to_ascii_uppercase();
|
|
std::env::var(&env_name).ok()
|
|
}),
|
|
) {
|
|
match abortable_run_with_spinner(
|
|
fetch_models(api_base, api_key.as_deref()),
|
|
"Fetching models",
|
|
create_abort_signal(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(fetched_models) => {
|
|
model_names = MultiSelect::new("LLMs to include (required):", fetched_models)
|
|
.with_validator(|list: &[ListOption<&String>]| {
|
|
if list.is_empty() {
|
|
Ok(Validation::Invalid(
|
|
"At least one item must be selected".into(),
|
|
))
|
|
} else {
|
|
Ok(Validation::Valid)
|
|
}
|
|
})
|
|
.prompt()?;
|
|
}
|
|
Err(err) => {
|
|
eprintln!("✗ Fetch models failed: {err}");
|
|
}
|
|
}
|
|
}
|
|
if model_names.is_empty() {
|
|
model_names = prompt_input_string(
|
|
"LLMs to add",
|
|
true,
|
|
Some("Separated by commas, e.g. llama3.3,qwen2.5"),
|
|
)?
|
|
.split(',')
|
|
.filter_map(|v| {
|
|
let v = v.trim();
|
|
if v.is_empty() {
|
|
None
|
|
} else {
|
|
Some(v.to_string())
|
|
}
|
|
})
|
|
.collect::<Vec<_>>();
|
|
}
|
|
if model_names.is_empty() {
|
|
bail!("No models");
|
|
}
|
|
let models: Vec<Value> = model_names
|
|
.iter()
|
|
.map(|v| {
|
|
let l = v.to_lowercase();
|
|
if l.contains("rank") {
|
|
json!({
|
|
"name": v,
|
|
"type": "reranker",
|
|
})
|
|
} else if let Ok(true) = EMBEDDING_MODEL_RE.is_match(&l) {
|
|
json!({
|
|
"name": v,
|
|
"type": "embedding",
|
|
"default_chunk_size": 1000,
|
|
"max_batch_size": 100
|
|
})
|
|
} else if v.contains("vision") {
|
|
json!({
|
|
"name": v,
|
|
"supports_vision": true
|
|
})
|
|
} else {
|
|
json!({
|
|
"name": v,
|
|
})
|
|
}
|
|
})
|
|
.collect();
|
|
client_config["models"] = models.into();
|
|
let model_name = select_model(model_names)?;
|
|
Ok(format!("{client}:{model_name}"))
|
|
}
|
|
|
|
fn select_model(model_names: Vec<String>) -> Result<String> {
|
|
if model_names.is_empty() {
|
|
bail!("No models");
|
|
}
|
|
let model = if model_names.len() == 1 {
|
|
model_names[0].clone()
|
|
} else {
|
|
Select::new("Default Model (required):", model_names).prompt()?
|
|
};
|
|
Ok(model)
|
|
}
|
|
|
|
fn prompt_input_string(desc: &str, required: bool, help_message: Option<&str>) -> Result<String> {
|
|
let desc = if required {
|
|
format!("{desc} (required):")
|
|
} else {
|
|
format!("{desc} (optional):")
|
|
};
|
|
let mut text = Text::new(&desc);
|
|
if required {
|
|
text = text.with_validator(required!("This field is required"))
|
|
}
|
|
if let Some(help_message) = help_message {
|
|
text = text.with_help_message(help_message);
|
|
}
|
|
let text = text.prompt()?;
|
|
Ok(text)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
use super::super::access_token::{is_rejected, set_access_token};
|
|
use crate::config::{AppState, WorkingMode};
|
|
use anyhow::Error;
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
use tokio::time::Instant;
|
|
|
|
fn catch_error_message(data: &Value, status: u16) -> String {
|
|
catch_error(data, status).unwrap_err().to_string()
|
|
}
|
|
|
|
#[test]
|
|
fn test_token_usage_is_empty_treats_some_zero_as_non_empty() {
|
|
assert!(TokenUsage::default().is_empty());
|
|
let usage = TokenUsage {
|
|
input_tokens: Some(0),
|
|
..Default::default()
|
|
};
|
|
assert!(!usage.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_token_usage_accumulate_merges_per_field() {
|
|
let mut total = TokenUsage::default();
|
|
total.accumulate(&TokenUsage {
|
|
input_tokens: Some(100),
|
|
output_tokens: None,
|
|
cache_creation_input_tokens: Some(50),
|
|
cache_read_input_tokens: None,
|
|
});
|
|
total.accumulate(&TokenUsage {
|
|
input_tokens: Some(20),
|
|
output_tokens: Some(7),
|
|
cache_creation_input_tokens: None,
|
|
cache_read_input_tokens: Some(90),
|
|
});
|
|
assert_eq!(total.input_tokens, Some(120));
|
|
assert_eq!(total.output_tokens, Some(7));
|
|
assert_eq!(total.cache_creation_input_tokens, Some(50));
|
|
assert_eq!(total.cache_read_input_tokens, Some(90));
|
|
}
|
|
|
|
#[test]
|
|
fn test_token_usage_accumulate_saturates() {
|
|
let mut total = TokenUsage {
|
|
input_tokens: Some(u64::MAX),
|
|
..Default::default()
|
|
};
|
|
total.accumulate(&TokenUsage {
|
|
input_tokens: Some(1),
|
|
..Default::default()
|
|
});
|
|
assert_eq!(total.input_tokens, Some(u64::MAX));
|
|
}
|
|
|
|
fn priced_model(
|
|
input_price: Option<f64>,
|
|
output_price: Option<f64>,
|
|
cache_read_price: Option<f64>,
|
|
cache_write_price: Option<f64>,
|
|
) -> Model {
|
|
let mut data = ModelData::new("test");
|
|
data.input_price = input_price;
|
|
data.output_price = output_price;
|
|
data.cache_read_price = cache_read_price;
|
|
data.cache_write_price = cache_write_price;
|
|
Model::from_config("provider", &[data]).remove(0)
|
|
}
|
|
|
|
#[test]
|
|
fn test_token_usage_cost_usd_sums_all_four_components() {
|
|
let usage = TokenUsage {
|
|
input_tokens: Some(1_000_000),
|
|
output_tokens: Some(500_000),
|
|
cache_creation_input_tokens: Some(200_000),
|
|
cache_read_input_tokens: Some(2_000_000),
|
|
};
|
|
let model = priced_model(Some(3.0), Some(15.0), Some(0.3), Some(3.75));
|
|
// 3.0 + 7.5 + 0.6 + 0.75
|
|
let cost = usage.cost_usd(&model).unwrap();
|
|
assert!((cost - 11.85).abs() < 1e-9);
|
|
}
|
|
|
|
#[test]
|
|
fn test_token_usage_cost_usd_is_none_when_a_billed_component_lacks_a_price() {
|
|
let usage = TokenUsage {
|
|
input_tokens: Some(100),
|
|
output_tokens: Some(10),
|
|
cache_creation_input_tokens: None,
|
|
cache_read_input_tokens: Some(50),
|
|
};
|
|
let model = priced_model(Some(3.0), Some(15.0), None, Some(3.75));
|
|
assert_eq!(usage.cost_usd(&model), None);
|
|
}
|
|
|
|
#[test]
|
|
fn test_token_usage_cost_usd_ignores_missing_prices_for_unbilled_components() {
|
|
let usage = TokenUsage {
|
|
input_tokens: Some(1_000_000),
|
|
output_tokens: Some(1_000_000),
|
|
cache_creation_input_tokens: Some(0),
|
|
cache_read_input_tokens: None,
|
|
};
|
|
let model = priced_model(Some(3.0), Some(15.0), None, None);
|
|
let cost = usage.cost_usd(&model).unwrap();
|
|
assert!((cost - 18.0).abs() < 1e-9);
|
|
}
|
|
|
|
#[test]
|
|
fn test_token_usage_cost_usd_is_none_for_empty_or_all_zero_usage() {
|
|
let model = priced_model(Some(3.0), Some(15.0), Some(0.3), Some(3.75));
|
|
assert_eq!(TokenUsage::default().cost_usd(&model), None);
|
|
let zero = TokenUsage {
|
|
input_tokens: Some(0),
|
|
output_tokens: Some(0),
|
|
cache_creation_input_tokens: Some(0),
|
|
cache_read_input_tokens: Some(0),
|
|
};
|
|
assert_eq!(zero.cost_usd(&model), None);
|
|
}
|
|
|
|
#[test]
|
|
fn test_catch_error_display_json_with_type() {
|
|
let data = json!({"error": {"type": "invalid_request_error", "message": "Bad request"}});
|
|
assert_eq!(
|
|
catch_error_message(&data, 400),
|
|
"Bad request (type: invalid_request_error)"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_catch_error_display_json_with_code() {
|
|
let data = json!({"error": {"code": "rate_limited", "message": "Too many requests"}});
|
|
assert_eq!(
|
|
catch_error_message(&data, 429),
|
|
"Too many requests (code: rate_limited)"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_catch_error_display_errors_array() {
|
|
let data = json!({"errors": [{"code": 7000, "message": "No route"}]});
|
|
assert_eq!(catch_error_message(&data, 404), "No route (status: 7000)");
|
|
}
|
|
|
|
#[test]
|
|
fn test_catch_error_display_array_error_status() {
|
|
let data = json!([{"error": {"status": "PERMISSION_DENIED", "message": "Denied"}}]);
|
|
assert_eq!(
|
|
catch_error_message(&data, 403),
|
|
"Denied (status: PERMISSION_DENIED)"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_catch_error_display_detail_status() {
|
|
let data = json!({"detail": "Not found", "status": 404});
|
|
assert_eq!(catch_error_message(&data, 404), "Not found (status: 404)");
|
|
}
|
|
|
|
#[test]
|
|
fn test_catch_error_display_error_string() {
|
|
let data = json!({"error": "Something went wrong"});
|
|
assert_eq!(catch_error_message(&data, 500), "Something went wrong");
|
|
}
|
|
|
|
#[test]
|
|
fn test_catch_error_display_message_string() {
|
|
let data = json!({"message": "Unauthorized"});
|
|
assert_eq!(catch_error_message(&data, 401), "Unauthorized");
|
|
}
|
|
|
|
#[test]
|
|
fn test_catch_error_display_fallback() {
|
|
let data = json!({"unexpected": true});
|
|
assert_eq!(
|
|
catch_error_message(&data, 500),
|
|
format!("Invalid response data: {data} (status: 500)")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_catch_error_ok_on_success_status() {
|
|
let data = json!({"error": {"type": "x", "message": "y"}});
|
|
assert!(catch_error(&data, 200).is_ok());
|
|
assert!(catch_error(&data, 299).is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_catch_error_downcast_through_context_chain() {
|
|
let data = json!({"error": {"type": "authentication_error", "message": "Invalid key"}});
|
|
let err = catch_error(&data, 401)
|
|
.context("Failed to call chat-completions api")
|
|
.unwrap_err();
|
|
let api_err = err
|
|
.downcast_ref::<ApiStatusError>()
|
|
.expect("should downcast through context chain");
|
|
assert_eq!(api_err.status, 401);
|
|
assert_eq!(api_err.message, "Invalid key (type: authentication_error)");
|
|
}
|
|
|
|
#[test]
|
|
fn test_catch_error_preserves_status() {
|
|
let data = json!({"message": "Unauthorized"});
|
|
let err = catch_error(&data, 401).unwrap_err();
|
|
assert_eq!(err.downcast_ref::<ApiStatusError>().unwrap().status, 401);
|
|
|
|
let data = json!({"detail": "Rate limited", "status": 429});
|
|
let err = catch_error(&data, 429).unwrap_err();
|
|
assert_eq!(err.downcast_ref::<ApiStatusError>().unwrap().status, 429);
|
|
|
|
// The struct carries the outer HTTP status even when the body embeds another code
|
|
let data = json!({"errors": [{"code": 7000, "message": "No route"}]});
|
|
let err = catch_error(&data, 429).unwrap_err();
|
|
assert_eq!(err.downcast_ref::<ApiStatusError>().unwrap().status, 429);
|
|
}
|
|
|
|
/// Wrapped in `.context(...)` so every test below proves the downcast
|
|
/// works through an anyhow context chain, as in the trait methods.
|
|
fn api_status_error(status: u16) -> Error {
|
|
Error::new(ApiStatusError {
|
|
status,
|
|
message: format!("error (status: {status})"),
|
|
})
|
|
.context("Failed to call chat-completions api")
|
|
}
|
|
|
|
fn cache_token(client: &str, token: &str) {
|
|
set_access_token(
|
|
client,
|
|
token.into(),
|
|
chrono::Utc::now().timestamp() + 3600,
|
|
None,
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_should_retry_auth_401_with_cached_token() {
|
|
let client = "should-retry-auth-401";
|
|
cache_token(client, "at-1");
|
|
|
|
assert!(should_retry_auth(&api_status_error(401), client));
|
|
assert!(is_rejected(client, "at-1"), "rejected marker not set");
|
|
}
|
|
|
|
#[test]
|
|
fn test_should_retry_auth_non_401_statuses() {
|
|
let client = "should-retry-auth-non-401";
|
|
cache_token(client, "at-1");
|
|
|
|
for status in [403, 429, 500] {
|
|
assert!(
|
|
!should_retry_auth(&api_status_error(status), client),
|
|
"retried on {status}"
|
|
);
|
|
}
|
|
assert_eq!(get_access_token(client).unwrap(), "at-1");
|
|
assert!(!is_rejected(client, "at-1"), "marker set without a 401");
|
|
}
|
|
|
|
#[test]
|
|
fn test_should_retry_auth_non_api_status_error() {
|
|
let client = "should-retry-auth-non-api";
|
|
cache_token(client, "at-1");
|
|
|
|
let err = anyhow::anyhow!("connection reset").context("Failed to call embeddings api");
|
|
assert!(!should_retry_auth(&err, client));
|
|
assert_eq!(get_access_token(client).unwrap(), "at-1");
|
|
assert!(!is_rejected(client, "at-1"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_should_retry_auth_401_without_cached_token() {
|
|
let client = "should-retry-auth-no-token";
|
|
|
|
assert!(!should_retry_auth(&api_status_error(401), client));
|
|
assert!(!is_rejected(client, "at-1"));
|
|
}
|
|
|
|
/// Fake provider whose streaming path emits `PACED_DELTAS` text chunks
|
|
/// 60s apart, then one thinking block and one tool call, and whose
|
|
/// non-streaming path returns that same reply in one piece when
|
|
/// `nonstream_replies` is set (and bails otherwise). Exercises the real
|
|
/// trait defaults (build_client, prepare_completion_data, abort
|
|
/// `select!`) with no network.
|
|
struct PacedStreamClient {
|
|
config: AppConfig,
|
|
model: Model,
|
|
stall_forever: bool,
|
|
nonstream_replies: bool,
|
|
streaming_calls: AtomicUsize,
|
|
}
|
|
|
|
const PACED_DELTAS: usize = 10;
|
|
const PACED_GAP: Duration = Duration::from_secs(60);
|
|
|
|
fn paced_text() -> String {
|
|
(0..PACED_DELTAS).map(|i| format!("chunk{i} ")).collect()
|
|
}
|
|
|
|
fn paced_tool_call() -> ToolCall {
|
|
ToolCall::new(
|
|
"lookup".to_string(),
|
|
json!({"q": 1}),
|
|
Some("call-1".to_string()),
|
|
)
|
|
}
|
|
|
|
fn paced_thinking() -> ThinkingBlock {
|
|
ThinkingBlock::Thinking {
|
|
thinking: "hmm".to_string(),
|
|
signature: "sig".to_string(),
|
|
}
|
|
}
|
|
|
|
#[async_trait::async_trait]
|
|
impl Client for PacedStreamClient {
|
|
fn app_config(&self) -> &AppConfig {
|
|
&self.config
|
|
}
|
|
|
|
fn extra_config(&self) -> Option<&ExtraConfig> {
|
|
None
|
|
}
|
|
|
|
fn patch_config(&self) -> Option<&RequestPatch> {
|
|
None
|
|
}
|
|
|
|
fn name(&self) -> &str {
|
|
"paced-stream"
|
|
}
|
|
|
|
fn model(&self) -> &Model {
|
|
&self.model
|
|
}
|
|
|
|
async fn chat_completions_inner(
|
|
&self,
|
|
_client: &ReqwestClient,
|
|
data: ChatCompletionsData,
|
|
) -> Result<ChatCompletionsOutput> {
|
|
assert!(!data.stream, "non-streaming path must not request a stream");
|
|
if !self.nonstream_replies {
|
|
bail!("the quiet transport must not take the non-streaming path");
|
|
}
|
|
if self.stall_forever {
|
|
std::future::pending::<()>().await;
|
|
}
|
|
Ok(ChatCompletionsOutput {
|
|
text: paced_text(),
|
|
tool_calls: vec![paced_tool_call()],
|
|
thinking: vec![paced_thinking()],
|
|
usage: None,
|
|
})
|
|
}
|
|
|
|
async fn chat_completions_streaming_inner(
|
|
&self,
|
|
_client: &ReqwestClient,
|
|
handler: &mut SseHandler,
|
|
data: ChatCompletionsData,
|
|
) -> Result<()> {
|
|
self.streaming_calls.fetch_add(1, Ordering::SeqCst);
|
|
assert!(data.stream, "trait default must request a stream");
|
|
for i in 0..PACED_DELTAS {
|
|
tokio::time::sleep(PACED_GAP).await;
|
|
handler.text(&format!("chunk{i} "))?;
|
|
}
|
|
if self.stall_forever {
|
|
std::future::pending::<()>().await;
|
|
}
|
|
handler.thinking_block(paced_thinking());
|
|
handler.tool_call(paced_tool_call())
|
|
}
|
|
}
|
|
|
|
fn paced_client(stall_forever: bool) -> PacedStreamClient {
|
|
PacedStreamClient {
|
|
config: AppConfig::default(),
|
|
model: Model::default(),
|
|
stall_forever,
|
|
nonstream_replies: false,
|
|
streaming_calls: AtomicUsize::new(0),
|
|
}
|
|
}
|
|
|
|
fn no_stream_model() -> Model {
|
|
let mut data = ModelData::new("");
|
|
data.no_stream = true;
|
|
Model::from_config("", &[data]).remove(0)
|
|
}
|
|
|
|
fn assert_paced_reply(text: &str, results: &[ToolResult]) {
|
|
assert_eq!(text, paced_text());
|
|
assert_eq!(results.len(), 1);
|
|
assert_eq!(results[0].call.name, "lookup");
|
|
assert_eq!(results[0].call.id.as_deref(), Some("call-1"));
|
|
assert!(
|
|
matches!(&results[0].thinking[..], [ThinkingBlock::Thinking { thinking, .. }] if thinking == "hmm")
|
|
);
|
|
}
|
|
|
|
fn quiet_ctx() -> RequestContext {
|
|
RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd)
|
|
}
|
|
|
|
/// 600s of virtual generation completes and yields the same shape the
|
|
/// non-streaming transport would: full text, tool results with the
|
|
/// thinking stapled onto the first, and the tracker fed each call. The
|
|
/// fake ignores the reqwest client, so reqwest's `read_timeout` is not in
|
|
/// the loop here; the real-socket tests in graph/mod.rs pin that it
|
|
/// fires on a stall (`sse_read_timeout_stall_is_transient`) and not on a
|
|
/// slow but continuous stream (`sse_slow_but_continuous_stream_outlives_read_timeout`).
|
|
#[tokio::test(start_paused = true)]
|
|
async fn quiet_streaming_accumulates_across_a_long_generation() {
|
|
let mut ctx = quiet_ctx();
|
|
let input = Input::from_str(&ctx, "hi", None).unwrap();
|
|
let client = paced_client(false);
|
|
let started = Instant::now();
|
|
|
|
let (text, results) =
|
|
call_chat_completions_streaming_quiet(&input, &client, &mut ctx, create_abort_signal())
|
|
.await
|
|
.unwrap();
|
|
|
|
assert!(started.elapsed() >= PACED_GAP * PACED_DELTAS as u32);
|
|
assert_paced_reply(&text, &results);
|
|
assert!(
|
|
results[0].output["tool_call_error"].is_string(),
|
|
"an undeclared tool evaluates to a tool_call_error, not a panic: {}",
|
|
results[0].output
|
|
);
|
|
}
|
|
|
|
/// An abort mid-stream cancels the provider call promptly and fails the
|
|
/// call outright; the ten chunks already buffered are never returned.
|
|
#[tokio::test(start_paused = true)]
|
|
async fn quiet_streaming_abort_discards_partial_output() {
|
|
let mut ctx = quiet_ctx();
|
|
let input = Input::from_str(&ctx, "hi", None).unwrap();
|
|
let client = paced_client(true);
|
|
let abort = create_abort_signal();
|
|
let trigger = abort.clone();
|
|
tokio::spawn(async move {
|
|
tokio::time::sleep(PACED_GAP * PACED_DELTAS as u32 + Duration::from_secs(30)).await;
|
|
trigger.set_ctrlc();
|
|
});
|
|
let started = Instant::now();
|
|
|
|
let err = call_chat_completions_streaming_quiet(&input, &client, &mut ctx, abort)
|
|
.await
|
|
.expect_err("an aborted stream must not succeed");
|
|
|
|
assert_eq!(err.to_string(), "Aborted.");
|
|
assert!(started.elapsed() < PACED_GAP * (PACED_DELTAS as u32 + 1));
|
|
}
|
|
|
|
/// The quiet transport and the non-streaming transport hand the node the
|
|
/// same `(text, tool_results)` for the same provider reply.
|
|
#[tokio::test(start_paused = true)]
|
|
async fn quiet_streaming_matches_the_non_streaming_reply() {
|
|
let mut ctx = quiet_ctx();
|
|
ctx.render_mode = RenderMode::Silent;
|
|
let input = Input::from_str(&ctx, "hi", None).unwrap();
|
|
let mut client = paced_client(false);
|
|
client.nonstream_replies = true;
|
|
|
|
let (streamed_text, streamed) =
|
|
call_chat_completions_streaming_quiet(&input, &client, &mut ctx, create_abort_signal())
|
|
.await
|
|
.unwrap();
|
|
let (plain_text, plain) = call_chat_completions(
|
|
&input,
|
|
false,
|
|
false,
|
|
&client,
|
|
&mut ctx,
|
|
create_abort_signal(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(streamed_text, plain_text);
|
|
assert_paced_reply(&streamed_text, &streamed);
|
|
assert_paced_reply(&plain_text, &plain);
|
|
assert_eq!(streamed[0].output, plain[0].output);
|
|
}
|
|
|
|
/// A model the catalog marks `no_stream` never sees a streaming request;
|
|
/// the quiet transport takes the non-streaming path and returns the same
|
|
/// shape.
|
|
#[tokio::test]
|
|
async fn quiet_transport_honours_no_stream_models() {
|
|
let mut ctx = quiet_ctx();
|
|
let input = Input::from_str(&ctx, "hi", None).unwrap();
|
|
let mut client = paced_client(false);
|
|
client.model = no_stream_model();
|
|
client.nonstream_replies = true;
|
|
|
|
let (text, results) =
|
|
call_chat_completions_streaming_quiet(&input, &client, &mut ctx, create_abort_signal())
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(client.streaming_calls.load(Ordering::SeqCst), 0);
|
|
assert_paced_reply(&text, &results);
|
|
}
|
|
|
|
/// The `no_stream` fallback is raced against the abort like the
|
|
/// streaming path: a hung non-streaming request is cancelled promptly.
|
|
#[tokio::test(start_paused = true)]
|
|
async fn quiet_transport_no_stream_fallback_honours_abort() {
|
|
let mut ctx = quiet_ctx();
|
|
let input = Input::from_str(&ctx, "hi", None).unwrap();
|
|
let mut client = paced_client(true);
|
|
client.model = no_stream_model();
|
|
client.nonstream_replies = true;
|
|
let abort = create_abort_signal();
|
|
let trigger = abort.clone();
|
|
tokio::spawn(async move {
|
|
tokio::time::sleep(Duration::from_secs(5)).await;
|
|
trigger.set_ctrlc();
|
|
});
|
|
let started = Instant::now();
|
|
|
|
let err = call_chat_completions_streaming_quiet(&input, &client, &mut ctx, abort)
|
|
.await
|
|
.expect_err("an aborted non-streaming fallback must not succeed");
|
|
|
|
assert_eq!(err.to_string(), "Aborted.");
|
|
assert!(started.elapsed() < Duration::from_secs(10));
|
|
}
|
|
|
|
/// The user-level `stream` toggle is a rendering preference; the quiet
|
|
/// transport still streams with it off.
|
|
#[tokio::test(start_paused = true)]
|
|
async fn quiet_transport_ignores_the_user_stream_toggle() {
|
|
let app = AppState {
|
|
config: Arc::new(AppConfig {
|
|
stream: false,
|
|
..AppConfig::default()
|
|
}),
|
|
..AppState::test_default()
|
|
};
|
|
let mut ctx = RequestContext::new(Arc::new(app), WorkingMode::Cmd);
|
|
let input = Input::from_str(&ctx, "hi", None).unwrap();
|
|
assert!(!input.stream(), "fixture must have streaming disabled");
|
|
let client = paced_client(false);
|
|
|
|
let (text, results) =
|
|
call_chat_completions_streaming_quiet(&input, &client, &mut ctx, create_abort_signal())
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(client.streaming_calls.load(Ordering::SeqCst), 1);
|
|
assert_paced_reply(&text, &results);
|
|
}
|
|
|
|
/// Fake provider whose every request fails with an [`ApiStatusError`],
|
|
/// wrapped in the same context chain the real trait defaults produce.
|
|
struct FailingStatusClient {
|
|
config: AppConfig,
|
|
model: Model,
|
|
status: u16,
|
|
}
|
|
|
|
impl FailingStatusClient {
|
|
fn new(status: u16) -> Self {
|
|
Self {
|
|
config: AppConfig::default(),
|
|
model: Model::default(),
|
|
status,
|
|
}
|
|
}
|
|
|
|
fn error(&self) -> Error {
|
|
Error::new(ApiStatusError {
|
|
status: self.status,
|
|
message: format!("error (status: {})", self.status),
|
|
})
|
|
}
|
|
}
|
|
|
|
#[async_trait::async_trait]
|
|
impl Client for FailingStatusClient {
|
|
fn app_config(&self) -> &AppConfig {
|
|
&self.config
|
|
}
|
|
|
|
fn extra_config(&self) -> Option<&ExtraConfig> {
|
|
None
|
|
}
|
|
|
|
fn patch_config(&self) -> Option<&RequestPatch> {
|
|
None
|
|
}
|
|
|
|
fn name(&self) -> &str {
|
|
"failing-status"
|
|
}
|
|
|
|
fn model(&self) -> &Model {
|
|
&self.model
|
|
}
|
|
|
|
async fn chat_completions_inner(
|
|
&self,
|
|
_client: &ReqwestClient,
|
|
_data: ChatCompletionsData,
|
|
) -> Result<ChatCompletionsOutput> {
|
|
Err(self.error())
|
|
}
|
|
|
|
async fn chat_completions_streaming_inner(
|
|
&self,
|
|
_client: &ReqwestClient,
|
|
_handler: &mut SseHandler,
|
|
_data: ChatCompletionsData,
|
|
) -> Result<()> {
|
|
Err(self.error())
|
|
}
|
|
}
|
|
|
|
fn llm_hooks_map(marker: &str) -> hooks::HooksMap {
|
|
[
|
|
"llm.request.started",
|
|
"llm.request.completed",
|
|
"llm.request.failed",
|
|
]
|
|
.into_iter()
|
|
.map(|event| {
|
|
(
|
|
event.to_string(),
|
|
vec![crate::hooks::HookDef {
|
|
name: marker.to_string(),
|
|
command: "true".to_string(),
|
|
}],
|
|
)
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn llm_hooked_ctx(marker: &str) -> RequestContext {
|
|
let mut app = AppState::test_default();
|
|
app.config = Arc::new(AppConfig {
|
|
hooks: llm_hooks_map(marker),
|
|
..Default::default()
|
|
});
|
|
let mut ctx = RequestContext::new(Arc::new(app), WorkingMode::Cmd);
|
|
ctx.render_mode = RenderMode::Silent;
|
|
ctx
|
|
}
|
|
|
|
fn llm_captures(marker: &str) -> Vec<hooks::test_sink::Capture> {
|
|
hooks::test_sink::snapshot()
|
|
.into_iter()
|
|
.filter(|capture| capture.hook_name == marker)
|
|
.collect()
|
|
}
|
|
|
|
#[tokio::test(start_paused = true)]
|
|
#[serial_test::serial]
|
|
async fn plain_transport_fires_llm_request_started_and_completed() {
|
|
let marker = "llm-hooks-plain-ok-x7x";
|
|
let mut ctx = llm_hooked_ctx(marker);
|
|
let input = Input::from_str(&ctx, "hi", None).unwrap();
|
|
let mut client = paced_client(false);
|
|
client.nonstream_replies = true;
|
|
let _guard = hooks::test_sink::install();
|
|
|
|
call_chat_completions(
|
|
&input,
|
|
false,
|
|
false,
|
|
&client,
|
|
&mut ctx,
|
|
create_abort_signal(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
let captures = llm_captures(marker);
|
|
assert_eq!(captures.len(), 2, "{captures:?}");
|
|
assert_eq!(captures[0].event, HookEvent::LlmRequestStarted);
|
|
assert_eq!(
|
|
captures[0]
|
|
.envs
|
|
.get("COYOTE_LLM_PROVIDER")
|
|
.map(String::as_str),
|
|
Some("paced-stream")
|
|
);
|
|
assert!(captures[0].envs.contains_key("COYOTE_LLM_MODEL"));
|
|
assert_eq!(captures[1].event, HookEvent::LlmRequestCompleted);
|
|
assert!(captures[1].envs.contains_key("COYOTE_LLM_DURATION_MS"));
|
|
assert!(!captures[1].envs.contains_key("COYOTE_LLM_ERROR_KIND"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[serial_test::serial]
|
|
async fn plain_transport_fires_llm_request_failed_with_api_error_kind() {
|
|
let marker = "llm-hooks-plain-api-err-x7x";
|
|
let mut ctx = llm_hooked_ctx(marker);
|
|
let input = Input::from_str(&ctx, "hi", None).unwrap();
|
|
let client = FailingStatusClient::new(429);
|
|
let _guard = hooks::test_sink::install();
|
|
|
|
let err = call_chat_completions(
|
|
&input,
|
|
false,
|
|
false,
|
|
&client,
|
|
&mut ctx,
|
|
create_abort_signal(),
|
|
)
|
|
.await
|
|
.unwrap_err();
|
|
|
|
assert!(
|
|
err.downcast_ref::<ApiStatusError>().is_some(),
|
|
"hooks must not change the error: {err:?}"
|
|
);
|
|
let captures = llm_captures(marker);
|
|
assert_eq!(captures.len(), 2, "{captures:?}");
|
|
assert_eq!(captures[0].event, HookEvent::LlmRequestStarted);
|
|
let failed = &captures[1];
|
|
assert_eq!(failed.event, HookEvent::LlmRequestFailed);
|
|
assert_eq!(
|
|
failed.envs.get("COYOTE_LLM_ERROR_KIND").map(String::as_str),
|
|
Some("api_error")
|
|
);
|
|
assert_eq!(
|
|
failed.envs.get("COYOTE_LLM_STATUS").map(String::as_str),
|
|
Some("429")
|
|
);
|
|
assert_eq!(
|
|
failed.envs.get("COYOTE_LLM_ERROR").map(String::as_str),
|
|
Some("Failed to call chat-completions api"),
|
|
"the env must carry the immediate error only, not the context chain"
|
|
);
|
|
assert!(!failed.envs.contains_key("COYOTE_LLM_DURATION_MS"));
|
|
}
|
|
|
|
#[tokio::test(start_paused = true)]
|
|
#[serial_test::serial]
|
|
async fn streaming_transport_fires_llm_request_started_and_completed() {
|
|
let marker = "llm-hooks-stream-ok-x7x";
|
|
let mut ctx = llm_hooked_ctx(marker);
|
|
let input = Input::from_str(&ctx, "hi", None).unwrap();
|
|
let client = paced_client(false);
|
|
let _guard = hooks::test_sink::install();
|
|
|
|
call_chat_completions_streaming(&input, &client, &mut ctx, create_abort_signal())
|
|
.await
|
|
.unwrap();
|
|
|
|
let captures = llm_captures(marker);
|
|
assert_eq!(captures.len(), 2, "{captures:?}");
|
|
assert_eq!(captures[0].event, HookEvent::LlmRequestStarted);
|
|
assert_eq!(captures[1].event, HookEvent::LlmRequestCompleted);
|
|
assert!(captures[1].envs.contains_key("COYOTE_LLM_DURATION_MS"));
|
|
}
|
|
|
|
#[tokio::test(start_paused = true)]
|
|
#[serial_test::serial]
|
|
async fn streaming_transport_abort_fires_llm_request_failed_aborted() {
|
|
let marker = "llm-hooks-stream-abort-x7x";
|
|
let mut ctx = llm_hooked_ctx(marker);
|
|
let input = Input::from_str(&ctx, "hi", None).unwrap();
|
|
let client = paced_client(true);
|
|
let abort = create_abort_signal();
|
|
let trigger = abort.clone();
|
|
tokio::spawn(async move {
|
|
tokio::time::sleep(PACED_GAP * PACED_DELTAS as u32 + Duration::from_secs(30)).await;
|
|
trigger.set_ctrlc();
|
|
});
|
|
let _guard = hooks::test_sink::install();
|
|
|
|
let err = call_chat_completions_streaming(&input, &client, &mut ctx, abort)
|
|
.await
|
|
.expect_err("an aborted stream must not succeed");
|
|
|
|
assert_eq!(err.to_string(), "Aborted.");
|
|
let captures = llm_captures(marker);
|
|
assert_eq!(captures.len(), 2, "{captures:?}");
|
|
let failed = &captures[1];
|
|
assert_eq!(failed.event, HookEvent::LlmRequestFailed);
|
|
assert_eq!(
|
|
failed.envs.get("COYOTE_LLM_ERROR_KIND").map(String::as_str),
|
|
Some("aborted")
|
|
);
|
|
assert!(!failed.envs.contains_key("COYOTE_LLM_STATUS"));
|
|
}
|
|
|
|
#[tokio::test(start_paused = true)]
|
|
#[serial_test::serial]
|
|
async fn quiet_transport_fires_llm_request_started_and_completed() {
|
|
let marker = "llm-hooks-quiet-ok-x7x";
|
|
let mut ctx = llm_hooked_ctx(marker);
|
|
let input = Input::from_str(&ctx, "hi", None).unwrap();
|
|
let client = paced_client(false);
|
|
let _guard = hooks::test_sink::install();
|
|
|
|
call_chat_completions_streaming_quiet(&input, &client, &mut ctx, create_abort_signal())
|
|
.await
|
|
.unwrap();
|
|
|
|
let captures = llm_captures(marker);
|
|
assert_eq!(captures.len(), 2, "{captures:?}");
|
|
assert_eq!(captures[0].event, HookEvent::LlmRequestStarted);
|
|
assert_eq!(
|
|
captures[0]
|
|
.envs
|
|
.get("COYOTE_LLM_PROVIDER")
|
|
.map(String::as_str),
|
|
Some("paced-stream")
|
|
);
|
|
assert_eq!(captures[1].event, HookEvent::LlmRequestCompleted);
|
|
assert!(captures[1].envs.contains_key("COYOTE_LLM_DURATION_MS"));
|
|
}
|
|
|
|
#[tokio::test(start_paused = true)]
|
|
#[serial_test::serial]
|
|
async fn quiet_transport_abort_fires_llm_request_failed_aborted() {
|
|
let marker = "llm-hooks-quiet-abort-x7x";
|
|
let mut ctx = llm_hooked_ctx(marker);
|
|
let input = Input::from_str(&ctx, "hi", None).unwrap();
|
|
let client = paced_client(true);
|
|
let abort = create_abort_signal();
|
|
let trigger = abort.clone();
|
|
tokio::spawn(async move {
|
|
tokio::time::sleep(PACED_GAP * PACED_DELTAS as u32 + Duration::from_secs(30)).await;
|
|
trigger.set_ctrlc();
|
|
});
|
|
let _guard = hooks::test_sink::install();
|
|
|
|
let err = call_chat_completions_streaming_quiet(&input, &client, &mut ctx, abort)
|
|
.await
|
|
.expect_err("an aborted stream must not succeed");
|
|
|
|
assert_eq!(err.to_string(), "Aborted.");
|
|
let captures = llm_captures(marker);
|
|
assert_eq!(captures.len(), 2, "{captures:?}");
|
|
let failed = &captures[1];
|
|
assert_eq!(failed.event, HookEvent::LlmRequestFailed);
|
|
assert_eq!(
|
|
failed.envs.get("COYOTE_LLM_ERROR_KIND").map(String::as_str),
|
|
Some("aborted")
|
|
);
|
|
assert_eq!(
|
|
failed.envs.get("COYOTE_LLM_ERROR").map(String::as_str),
|
|
Some("Aborted.")
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[serial_test::serial]
|
|
async fn quiet_transport_fires_llm_request_failed_with_api_error_kind() {
|
|
let marker = "llm-hooks-quiet-api-err-x7x";
|
|
let mut ctx = llm_hooked_ctx(marker);
|
|
let input = Input::from_str(&ctx, "hi", None).unwrap();
|
|
let client = FailingStatusClient::new(500);
|
|
let _guard = hooks::test_sink::install();
|
|
|
|
let err =
|
|
call_chat_completions_streaming_quiet(&input, &client, &mut ctx, create_abort_signal())
|
|
.await
|
|
.unwrap_err();
|
|
|
|
assert!(err.downcast_ref::<ApiStatusError>().is_some());
|
|
let captures = llm_captures(marker);
|
|
assert_eq!(captures.len(), 2, "{captures:?}");
|
|
let failed = &captures[1];
|
|
assert_eq!(failed.event, HookEvent::LlmRequestFailed);
|
|
assert_eq!(
|
|
failed.envs.get("COYOTE_LLM_ERROR_KIND").map(String::as_str),
|
|
Some("api_error")
|
|
);
|
|
assert_eq!(
|
|
failed.envs.get("COYOTE_LLM_STATUS").map(String::as_str),
|
|
Some("500")
|
|
);
|
|
}
|
|
|
|
/// Minimal client for exercising `patch_request_data`: only the model
|
|
/// and the client-level patch config matter.
|
|
struct PatchProbeClient {
|
|
config: AppConfig,
|
|
model: Model,
|
|
patch: Option<RequestPatch>,
|
|
}
|
|
|
|
#[async_trait::async_trait]
|
|
impl Client for PatchProbeClient {
|
|
fn app_config(&self) -> &AppConfig {
|
|
&self.config
|
|
}
|
|
|
|
fn extra_config(&self) -> Option<&ExtraConfig> {
|
|
None
|
|
}
|
|
|
|
fn patch_config(&self) -> Option<&RequestPatch> {
|
|
self.patch.as_ref()
|
|
}
|
|
|
|
fn name(&self) -> &str {
|
|
self.model.client_name()
|
|
}
|
|
|
|
fn model(&self) -> &Model {
|
|
&self.model
|
|
}
|
|
|
|
async fn chat_completions_inner(
|
|
&self,
|
|
_client: &ReqwestClient,
|
|
_data: ChatCompletionsData,
|
|
) -> Result<ChatCompletionsOutput> {
|
|
bail!("the patch probe makes no requests")
|
|
}
|
|
|
|
async fn chat_completions_streaming_inner(
|
|
&self,
|
|
_client: &ReqwestClient,
|
|
_handler: &mut SseHandler,
|
|
_data: ChatCompletionsData,
|
|
) -> Result<()> {
|
|
bail!("the patch probe makes no requests")
|
|
}
|
|
}
|
|
|
|
fn patch_probe(
|
|
client_name: &str,
|
|
model_patch: Option<Value>,
|
|
patch: Option<RequestPatch>,
|
|
) -> PatchProbeClient {
|
|
let mut data = ModelData::new("gpt-test");
|
|
data.patch = model_patch;
|
|
PatchProbeClient {
|
|
config: AppConfig::default(),
|
|
model: Model::from_config(client_name, &[data]).remove(0),
|
|
patch,
|
|
}
|
|
}
|
|
|
|
fn client_patch(json: Value) -> RequestPatch {
|
|
serde_json::from_value(json).unwrap()
|
|
}
|
|
|
|
fn wire_request(wire: WireApi) -> RequestData {
|
|
RequestData::new("https://example.test/v1", json!({"model": "gpt-test"})).wire(wire)
|
|
}
|
|
|
|
#[test]
|
|
fn request_patch_accepts_a_responses_section() {
|
|
let patch = client_patch(json!({
|
|
"responses": {".*": {"body": {"store": true}}},
|
|
"chat_completions": {".*": {"body": {"temperature": 0.5}}},
|
|
}));
|
|
|
|
assert!(patch.responses.is_some());
|
|
assert!(patch.chat_completions.is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn non_map_request_patch_responses_section_is_rejected() {
|
|
let err = serde_json::from_value::<RequestPatch>(json!({"responses": "bogus"}))
|
|
.unwrap_err()
|
|
.to_string();
|
|
|
|
assert!(err.contains("expected a map"), "unexpected error: {err}");
|
|
}
|
|
|
|
#[test]
|
|
fn chat_client_patch_never_touches_the_responses_wire() {
|
|
let client = patch_probe(
|
|
"patchprobechatiso",
|
|
None,
|
|
Some(client_patch(json!({
|
|
"chat_completions": {".*": {"body": {"temperature": 0.5}}}
|
|
}))),
|
|
);
|
|
let mut request_data = wire_request(WireApi::Responses);
|
|
let before = request_data.body.clone();
|
|
|
|
client.patch_request_data(&mut request_data);
|
|
|
|
assert_eq!(request_data.body, before);
|
|
}
|
|
|
|
#[test]
|
|
fn responses_client_patch_never_touches_the_chat_wire() {
|
|
let patch = Some(client_patch(json!({
|
|
"responses": {".*": {"body": {"store": true}}}
|
|
})));
|
|
let client = patch_probe("patchproberespiso", None, patch);
|
|
|
|
let mut chat_request = wire_request(WireApi::Chat);
|
|
let before = chat_request.body.clone();
|
|
client.patch_request_data(&mut chat_request);
|
|
assert_eq!(chat_request.body, before);
|
|
|
|
let mut responses_request = wire_request(WireApi::Responses);
|
|
client.patch_request_data(&mut responses_request);
|
|
assert_eq!(responses_request.body["store"], json!(true));
|
|
}
|
|
|
|
#[test]
|
|
fn chat_shaped_model_patch_never_touches_the_responses_wire() {
|
|
let client = patch_probe(
|
|
"patchprobemodeliso",
|
|
Some(json!({"body": {"reasoning_effort": "high"}})),
|
|
None,
|
|
);
|
|
let mut request_data = wire_request(WireApi::Responses);
|
|
let before = request_data.body.clone();
|
|
|
|
client.patch_request_data(&mut request_data);
|
|
|
|
assert_eq!(request_data.body, before);
|
|
}
|
|
|
|
#[test]
|
|
fn model_patch_responses_subkey_is_selected_per_wire() {
|
|
let model_patch = json!({
|
|
"body": {"temperature": 0.2},
|
|
"responses": {"body": {"store": true}},
|
|
});
|
|
let client = patch_probe("patchprobesubkey", Some(model_patch), None);
|
|
|
|
let mut responses_request = wire_request(WireApi::Responses);
|
|
client.patch_request_data(&mut responses_request);
|
|
assert_eq!(responses_request.body["store"], json!(true));
|
|
assert!(responses_request.body.get("temperature").is_none());
|
|
|
|
let mut chat_request = wire_request(WireApi::Chat);
|
|
client.patch_request_data(&mut chat_request);
|
|
assert_eq!(chat_request.body["temperature"], json!(0.2));
|
|
assert!(chat_request.body.get("store").is_none());
|
|
}
|
|
|
|
/// Pins the silent-ignore contract for a malformed `responses` sub-key:
|
|
/// model patches are untyped JSON maps and the whole patch layer is
|
|
/// deliberately lenient — `apply_patch` only reads object-shaped
|
|
/// `url`/`body`/`headers` keys — so a non-object `responses` sub-key
|
|
/// degrades to a no-op on the Responses wire (never a request-time
|
|
/// error, never a partial merge) and stays inert on the Chat wire,
|
|
/// where the patch's top-level keys still apply normally.
|
|
#[test]
|
|
fn malformed_model_patch_responses_subkey_is_silently_ignored() {
|
|
for malformed in [json!("oops"), json!([{"body": {"store": true}}])] {
|
|
let model_patch = json!({
|
|
"body": {"temperature": 0.2},
|
|
"responses": malformed,
|
|
});
|
|
let client = patch_probe("patchprobemalformed", Some(model_patch), None);
|
|
|
|
let mut responses_request = wire_request(WireApi::Responses);
|
|
let before_url = responses_request.url.clone();
|
|
let before_body = responses_request.body.clone();
|
|
let before_headers = responses_request.headers.clone();
|
|
client.patch_request_data(&mut responses_request);
|
|
assert_eq!(responses_request.url, before_url);
|
|
assert_eq!(responses_request.body, before_body);
|
|
assert_eq!(responses_request.headers, before_headers);
|
|
|
|
let mut chat_request = wire_request(WireApi::Chat);
|
|
client.patch_request_data(&mut chat_request);
|
|
assert_eq!(chat_request.body["temperature"], json!(0.2));
|
|
assert!(chat_request.body.get("store").is_none());
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn env_patch_override_is_selected_per_wire() {
|
|
// Unique client name so the env vars cannot race parallel tests.
|
|
let client_name = "patchprobeenvwire";
|
|
let responses_env = get_env_name(&format!("patch_{client_name}_responses"));
|
|
let chat_env = get_env_name(&format!("patch_{client_name}_chat_completions"));
|
|
unsafe {
|
|
std::env::set_var(&responses_env, r#"{".*": {"body": {"store": true}}}"#);
|
|
std::env::set_var(&chat_env, r#"{".*": {"body": {"temperature": 0.7}}}"#);
|
|
}
|
|
|
|
let client = patch_probe(client_name, None, None);
|
|
let mut responses_request = wire_request(WireApi::Responses);
|
|
client.patch_request_data(&mut responses_request);
|
|
let mut chat_request = wire_request(WireApi::Chat);
|
|
client.patch_request_data(&mut chat_request);
|
|
|
|
unsafe {
|
|
std::env::remove_var(&responses_env);
|
|
std::env::remove_var(&chat_env);
|
|
}
|
|
|
|
assert_eq!(responses_request.body["store"], json!(true));
|
|
assert!(responses_request.body.get("temperature").is_none());
|
|
assert_eq!(chat_request.body["temperature"], json!(0.7));
|
|
assert!(chat_request.body.get("store").is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn chat_wire_patching_is_unchanged_by_wire_awareness() {
|
|
let client = patch_probe(
|
|
"patchproberegress",
|
|
Some(json!({"body": {"max_tokens": null}, "headers": {"x-model": "m"}})),
|
|
Some(client_patch(json!({
|
|
"chat_completions": {".*": {
|
|
"url": "https://patched.example/v1/chat/completions",
|
|
"body": {"temperature": 0.9},
|
|
}}
|
|
}))),
|
|
);
|
|
let mut request_data = RequestData::new(
|
|
"https://example.test/v1/chat/completions",
|
|
json!({"model": "gpt-test", "max_tokens": 100}),
|
|
);
|
|
request_data.bearer_auth("sk-test");
|
|
|
|
client.patch_request_data(&mut request_data);
|
|
|
|
assert_eq!(
|
|
request_data.url,
|
|
"https://patched.example/v1/chat/completions"
|
|
);
|
|
assert_eq!(
|
|
request_data.body,
|
|
json!({"model": "gpt-test", "temperature": 0.9})
|
|
);
|
|
assert_eq!(
|
|
request_data.headers.get("x-model").map(String::as_str),
|
|
Some("m")
|
|
);
|
|
assert_eq!(
|
|
request_data
|
|
.headers
|
|
.get("authorization")
|
|
.map(String::as_str),
|
|
Some("Bearer sk-test")
|
|
);
|
|
}
|
|
}
|