Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

### Breaking

- Errors keep their cause, as the AI SDK's `cause`. `JsonParse`,
`NoOutputGenerated`, `InvalidResponseData`, `InvalidArgument`,
`InvalidPrompt` and `Other` are `{ message, cause }` (build them with
`AiMuxError::json_parse(..)`, `invalid_argument(..)`, ...), `ApiCallError`
has `cause` (`.caused_by(..)`), and `InvalidToolInput.cause` is the parse or
validation error itself instead of its text. A cause is a nested
`AiMuxError` or a `SourceError` chain of an OS, TLS or parser error; it is
what `source()` returns and the payload's `cause` in every binding (Node
`error.cause`, Python `__cause__`, Java/Kotlin `getCause()`, Go `Unwrap()`,
Swift/Flutter `cause`, C++ `aimux::NativeError::json()`). `FetchError`
carries a `SourceError` without the request URL. Messages are unchanged.

- aimux's public JSON equals the AI SDK's JSON. Struct fields serialize
camelCase (`toolCallId`, `providerOptions`, `inputTokens`); unions carry the
AI SDK's tag and strings (`{"type":"text-delta"}`, `"tool-call"`,
Expand Down
4 changes: 2 additions & 2 deletions aimux-core/src/decision_model.rs
Original file line number Diff line number Diff line change
Expand Up @@ -424,7 +424,7 @@ pub struct DecideResult {
}

fn invalid_input(parameter: &str, message: &str) -> AiMuxError {
AiMuxError::InvalidArgument(format!(
AiMuxError::invalid_argument(format!(
"Invalid argument for parameter {parameter}: {message}"
))
}
Expand Down Expand Up @@ -471,7 +471,7 @@ pub fn validate_decision_input(questions: &DecisionQuestions) -> Result<(), AiMu
const TOLERANCE: f64 = 1e-6;

fn invalid_answer(message: String) -> AiMuxError {
AiMuxError::InvalidResponseData(message)
AiMuxError::invalid_response_data(message)
}

fn is_probability(value: f64) -> bool {
Expand Down
57 changes: 36 additions & 21 deletions aimux-core/src/decision_replay.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ use crate::decision_model::{
DecideOptions, DecideResult, DecisionCallOptions, DecisionModel, DecisionQuestionType,
DecisionResult, DecisionUsage, decide,
};
use crate::error::ErrorCause;
use crate::recording::{OutcomeStatus, Recording, RecordingOperation};

/// Appended to a replay error when recordings written by aimux 0.5 or
Expand Down Expand Up @@ -52,8 +53,8 @@ impl MockDecisionReplayModel {
&& rec.provider.provider == provider
&& rec.provider.model_id == model_id
}) {
return Err(AiMuxError::InvalidArgument(
"decision replay: no matching provider/model recording".into(),
return Err(AiMuxError::invalid_argument(
"decision replay: no matching provider/model recording",
));
}
Ok(Self {
Expand All @@ -67,21 +68,24 @@ impl MockDecisionReplayModel {
/// # Errors
/// Returns read/parse errors or an error if the file contains no decisions.
pub fn from_jsonl(path: impl AsRef<std::path::Path>) -> Result<Self, AiMuxError> {
let content = std::fs::read_to_string(path)
.map_err(|e| AiMuxError::InvalidArgument(format!("decision replay: {e}")))?;
let content = std::fs::read_to_string(path).map_err(|e| AiMuxError::InvalidArgument {
message: format!("decision replay: {e}"),
cause: Some(ErrorCause::source(&e)),
})?;
let recordings: Vec<Recording> = content
.lines()
.filter(|line| !line.trim().is_empty())
.map(|line| {
serde_json::from_str(line).map_err(|e| AiMuxError::JsonParse(e.to_string()))
serde_json::from_str(line).map_err(|e| AiMuxError::JsonParse {
message: e.to_string(),
cause: Some(ErrorCause::source(&e)),
})
})
.collect::<Result<_, _>>()?;
let first = recordings
.iter()
.find(|rec| rec.input.operation == RecordingOperation::Decision)
.ok_or_else(|| {
AiMuxError::InvalidArgument("decision replay: no decisions in file".into())
})?;
.ok_or_else(|| AiMuxError::invalid_argument("decision replay: no decisions in file"))?;
Self::new(
first.provider.provider.clone(),
first.provider.model_id.clone(),
Expand Down Expand Up @@ -114,8 +118,12 @@ impl DecisionModel for MockDecisionReplayModel {
DecisionQuestionType::ALL
}
async fn do_decide(&self, options: &DecisionCallOptions) -> Result<DecisionResult, AiMuxError> {
let requested =
key(&serde_json::to_value(options).map_err(|e| AiMuxError::JsonParse(e.to_string()))?);
let requested = key(
&serde_json::to_value(options).map_err(|e| AiMuxError::JsonParse {
message: e.to_string(),
cause: Some(ErrorCause::source(&e)),
})?,
);
let rec = self
.recordings
.iter()
Expand Down Expand Up @@ -143,7 +151,7 @@ impl DecisionModel for MockDecisionReplayModel {
"; {older} recording(s) of this model were {OLDER_RECORDING} and never match"
));
}
AiMuxError::InvalidArgument(message)
AiMuxError::invalid_argument(message)
})?;
if rec.outcome.status != OutcomeStatus::Success {
return Err(rec
Expand All @@ -152,7 +160,7 @@ impl DecisionModel for MockDecisionReplayModel {
.clone()
.and_then(|v| serde_json::from_value(v).ok())
.unwrap_or_else(|| {
AiMuxError::InvalidResponseData(
AiMuxError::invalid_response_data(
rec.outcome
.error
.clone()
Expand All @@ -162,9 +170,12 @@ impl DecisionModel for MockDecisionReplayModel {
}
let recorded: DecideResult =
serde_json::from_value(rec.outcome.decision_result.clone().ok_or_else(|| {
AiMuxError::InvalidResponseData("decision replay: missing recorded result".into())
AiMuxError::invalid_response_data("decision replay: missing recorded result")
})?)
.map_err(|e| AiMuxError::InvalidResponseData(format!("decision replay result: {e}")))?;
.map_err(|e| AiMuxError::InvalidResponseData {
message: format!("decision replay result: {e}"),
cause: Some(ErrorCause::source(&e)),
})?;
Ok(DecisionResult {
answers: recorded.answers,
rounding: recorded.rounding,
Expand All @@ -188,8 +199,8 @@ pub async fn replay_decision_with_model(
model: &dyn DecisionModel,
) -> Result<DecideResult, AiMuxError> {
if recording.input.operation != RecordingOperation::Decision {
return Err(AiMuxError::InvalidArgument(
"decision replay requires a decision recording".into(),
return Err(AiMuxError::invalid_argument(
"decision replay requires a decision recording",
));
}
let unread = crate::replay::snake_case_option_keys(&recording.input.options);
Expand All @@ -199,7 +210,7 @@ pub async fn replay_decision_with_model(
} else {
String::new()
};
return Err(AiMuxError::JsonParse(format!(
return Err(AiMuxError::json_parse(format!(
"decision replay: {older}this version reads the options camelCase and would \
ignore {}",
unread.join(", ")
Expand All @@ -208,11 +219,15 @@ pub async fn replay_decision_with_model(
let options = crate::replay::readable_options(&recording.input.options);
let options: DecideOptions = serde_json::from_value(options).map_err(|e| {
if recording.provider.provider_id.is_empty() {
AiMuxError::JsonParse(format!(
"decision replay: the recording was {OLDER_RECORDING}: {e}"
))
AiMuxError::JsonParse {
message: format!("decision replay: the recording was {OLDER_RECORDING}: {e}"),
cause: Some(ErrorCause::source(&e)),
}
} else {
AiMuxError::JsonParse(e.to_string())
AiMuxError::JsonParse {
message: e.to_string(),
cause: Some(ErrorCause::source(&e)),
}
}
})?;
decide(model, options).await
Expand Down
24 changes: 14 additions & 10 deletions aimux-core/src/download.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ use futures::future::BoxFuture;
use crate::abort_signal::AbortSignal;
use crate::download_guard::{validate_download_target, validate_download_url};
use crate::error::AiMuxError;
use crate::error::ErrorCause;
use crate::shared::{FileBytes, GeneratedFileData};

const DEFAULT_MAX_DOWNLOAD_SIZE: usize = 2 * 1024 * 1024 * 1024;
Expand Down Expand Up @@ -161,7 +162,7 @@ impl DownloadFunction {
let requested = downloads.len();
let files = (self.0)(downloads).await?;
if files.len() != requested {
return Err(AiMuxError::InvalidArgument(format!(
return Err(AiMuxError::invalid_argument(format!(
"the download function returned {} results for {requested} URLs",
files.len()
)));
Expand Down Expand Up @@ -246,10 +247,10 @@ async fn download_inner(
let addresses = validate_download_target(&current, None).await?;
let host = parsed
.host_str()
.ok_or_else(|| AiMuxError::InvalidArgument("download URL has no host".into()))?;
.ok_or_else(|| AiMuxError::invalid_argument("download URL has no host"))?;
let port = parsed
.port_or_known_default()
.ok_or_else(|| AiMuxError::InvalidArgument("download URL has no known port".into()))?;
.ok_or_else(|| AiMuxError::invalid_argument("download URL has no known port"))?;
let addresses: Vec<_> = addresses
.into_iter()
.map(|address| SocketAddr::new(address, port))
Expand All @@ -276,15 +277,15 @@ async fn download_inner(
&& let Some(location) = response.headers().get(reqwest::header::LOCATION)
{
if redirects == MAX_DOWNLOAD_REDIRECTS {
return Err(AiMuxError::Other("download exceeded 10 redirects".into()));
return Err(AiMuxError::other("download exceeded 10 redirects"));
}
let location = location.to_str().map_err(download_error)?;
current = parsed.join(location).map_err(download_error)?.to_string();
drop(response);
continue;
}
if !status.is_success() {
return Err(AiMuxError::Other(format!(
return Err(AiMuxError::other(format!(
"download failed with HTTP status {status}"
)));
}
Expand Down Expand Up @@ -313,12 +314,15 @@ async fn download_inner(
unreachable!("download loop returns a response or redirect error")
}

fn download_error(error: impl std::fmt::Display) -> AiMuxError {
AiMuxError::Other(format!("download failed: {error}"))
fn download_error(error: impl std::error::Error + 'static) -> AiMuxError {
AiMuxError::Other {
message: format!("download failed: {error}"),
cause: Some(ErrorCause::source(&error)),
}
}

fn size_error(max_bytes: usize) -> AiMuxError {
AiMuxError::Other(format!(
AiMuxError::other(format!(
"download exceeded maximum size of {max_bytes} bytes"
))
}
Expand Down Expand Up @@ -392,7 +396,7 @@ mod tests {
.await
.unwrap_err();
assert!(
matches!(&error, AiMuxError::InvalidArgument(message)
matches!(&error, AiMuxError::InvalidArgument { message, .. }
if message == "the download function returned 1 results for 2 URLs"),
"{error:?}"
);
Expand All @@ -407,7 +411,7 @@ mod tests {
.await
.unwrap_err();
assert!(
matches!(&error, AiMuxError::InvalidArgument(message)
matches!(&error, AiMuxError::InvalidArgument { message, .. }
if message == "download URL targets a non-public address: 127.0.0.1"),
"{error:?}"
);
Expand Down
32 changes: 18 additions & 14 deletions aimux-core/src/download_guard.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
use std::net::IpAddr;

use crate::error::AiMuxError;
use crate::error::ErrorCause;

/// Compare scheme, host, and effective port. Unparseable inputs are never
/// same-origin. Public so providers can build credential allowlists on top
Expand Down Expand Up @@ -61,31 +62,33 @@ fn host_ip_literal(host: &str) -> Option<IpAddr> {
/// Callers that resolve DNS afterwards rely on this having rejected every
/// non-public literal, so keep this the single literal-check site.
pub fn validate_download_url(url: &str) -> Result<url::Url, AiMuxError> {
let parsed = url::Url::parse(url)
.map_err(|error| AiMuxError::InvalidArgument(format!("invalid download URL: {error}")))?;
let parsed = url::Url::parse(url).map_err(|error| AiMuxError::InvalidArgument {
message: format!("invalid download URL: {error}"),
cause: Some(ErrorCause::source(&error)),
})?;
if !matches!(parsed.scheme(), "http" | "https") {
return Err(AiMuxError::InvalidArgument(format!(
return Err(AiMuxError::invalid_argument(format!(
"download URL must use HTTP or HTTPS: {}",
redacted(&parsed)
)));
}
let host = parsed
.host_str()
.ok_or_else(|| AiMuxError::InvalidArgument("download URL has no host".into()))?;
.ok_or_else(|| AiMuxError::invalid_argument("download URL has no host"))?;
let normalized_host = host.to_ascii_lowercase();
let normalized_host = normalized_host.trim_end_matches('.');
if normalized_host == "localhost"
|| normalized_host.ends_with(".local")
|| normalized_host.ends_with(".localhost")
{
return Err(AiMuxError::InvalidArgument(format!(
return Err(AiMuxError::invalid_argument(format!(
"download URL targets a disallowed hostname: {normalized_host}"
)));
}
if let Some(address) = host_ip_literal(normalized_host)
&& !is_public_download_address(address)
{
return Err(AiMuxError::InvalidArgument(format!(
return Err(AiMuxError::invalid_argument(format!(
"download URL targets a non-public address: {address}"
)));
}
Expand All @@ -108,18 +111,19 @@ pub async fn validate_download_target(
let parsed = validate_download_url(url)?;
let host = parsed
.host_str()
.ok_or_else(|| AiMuxError::InvalidArgument("download URL has no host".into()))?;
.ok_or_else(|| AiMuxError::invalid_argument("download URL has no host"))?;
if let Some(address) = host_ip_literal(host) {
// validate_download_url above already rejected non-public literals.
return Ok(vec![address]);
}
let port = parsed
.port_or_known_default()
.ok_or_else(|| AiMuxError::InvalidArgument("download URL has no known port".into()))?;
.ok_or_else(|| AiMuxError::invalid_argument("download URL has no known port"))?;
let addresses: Vec<_> = tokio::net::lookup_host((host, port))
.await
.map_err(|error| {
AiMuxError::InvalidArgument(format!("download URL host could not be resolved: {error}"))
.map_err(|error| AiMuxError::InvalidArgument {
message: format!("download URL host could not be resolved: {error}"),
cause: Some(ErrorCause::source(&error)),
})?
.map(|address| address.ip())
.collect();
Expand All @@ -131,14 +135,14 @@ fn validate_resolved_download_addresses(
addresses: Vec<IpAddr>,
) -> Result<Vec<IpAddr>, AiMuxError> {
if addresses.is_empty() {
return Err(AiMuxError::InvalidArgument(format!(
return Err(AiMuxError::invalid_argument(format!(
"download URL host did not resolve to an address: {host}"
)));
}
let mut validated = Vec::with_capacity(addresses.len());
for address in addresses {
if !is_public_download_address(address) {
return Err(AiMuxError::InvalidArgument(format!(
return Err(AiMuxError::invalid_argument(format!(
"download URL resolves to a non-public address: {address}"
)));
}
Expand Down Expand Up @@ -369,7 +373,7 @@ mod tests {
let empty = validate_resolved_download_addresses("empty.example", Vec::new())
.expect_err("an empty DNS answer must fail closed");
assert!(
matches!(empty, AiMuxError::InvalidArgument(ref message) if message.contains("did not resolve"))
matches!(empty, AiMuxError::InvalidArgument { ref message, .. } if message.contains("did not resolve"))
);

let mixed = validate_resolved_download_addresses(
Expand All @@ -378,7 +382,7 @@ mod tests {
)
.expect_err("one private answer must reject the entire DNS result");
assert!(
matches!(mixed, AiMuxError::InvalidArgument(ref message) if message.contains("non-public"))
matches!(mixed, AiMuxError::InvalidArgument { ref message, .. } if message.contains("non-public"))
);

let all = validate_resolved_download_addresses(
Expand Down
Loading
Loading