diff --git a/crates/aion-agent/src/commands/clear.rs b/crates/aion-agent/src/commands/clear.rs index 053c74a8..b69ff281 100644 --- a/crates/aion-agent/src/commands/clear.rs +++ b/crates/aion-agent/src/commands/clear.rs @@ -31,8 +31,8 @@ impl SlashCommand for ClearCommand { mod tests { use std::sync::Arc; - use aion_providers::{LlmProvider, ProviderError}; - use aion_types::llm::{LlmEvent, LlmRequest}; + use aion_providers::{LlmProvider, ProviderFailure, ProviderStreamReceiver}; + use aion_types::llm::LlmRequest; use aion_types::message::{ContentBlock, Message, Role}; use super::*; @@ -42,10 +42,7 @@ mod tests { struct NullProvider; #[async_trait::async_trait] impl LlmProvider for NullProvider { - async fn stream( - &self, - _: &LlmRequest, - ) -> Result, ProviderError> { + async fn stream(&self, _: &LlmRequest) -> Result { let (_tx, rx) = tokio::sync::mpsc::channel(1); Ok(rx) } diff --git a/crates/aion-agent/src/commands/compact.rs b/crates/aion-agent/src/commands/compact.rs index d705e618..ffb95b97 100644 --- a/crates/aion-agent/src/commands/compact.rs +++ b/crates/aion-agent/src/commands/compact.rs @@ -83,8 +83,8 @@ impl SlashCommand for CompactCommand { mod tests { use std::sync::Arc; - use aion_providers::{LlmProvider, ProviderError}; - use aion_types::llm::{LlmEvent, LlmRequest}; + use aion_providers::{LlmProvider, ProviderFailure, ProviderStreamReceiver}; + use aion_types::llm::LlmRequest; use aion_types::message::{ContentBlock, Message, Role}; use super::*; @@ -95,10 +95,7 @@ mod tests { struct NullProvider; #[async_trait::async_trait] impl LlmProvider for NullProvider { - async fn stream( - &self, - _: &LlmRequest, - ) -> Result, ProviderError> { + async fn stream(&self, _: &LlmRequest) -> Result { let (_tx, rx) = tokio::sync::mpsc::channel(1); Ok(rx) } diff --git a/crates/aion-agent/src/commands/help.rs b/crates/aion-agent/src/commands/help.rs index 26175d24..4fbb5119 100644 --- a/crates/aion-agent/src/commands/help.rs +++ b/crates/aion-agent/src/commands/help.rs @@ -40,8 +40,8 @@ impl SlashCommand for HelpCommand { mod tests { use std::sync::{Arc, Mutex}; - use aion_providers::{LlmProvider, ProviderError}; - use aion_types::llm::{LlmEvent, LlmRequest}; + use aion_providers::{LlmProvider, ProviderFailure, ProviderStreamReceiver}; + use aion_types::llm::LlmRequest; use aion_types::message::Message; use super::*; @@ -81,10 +81,7 @@ mod tests { struct NullProvider; #[async_trait::async_trait] impl LlmProvider for NullProvider { - async fn stream( - &self, - _: &LlmRequest, - ) -> Result, ProviderError> { + async fn stream(&self, _: &LlmRequest) -> Result { let (_tx, rx) = tokio::sync::mpsc::channel(1); Ok(rx) } diff --git a/crates/aion-agent/src/compact/auto.rs b/crates/aion-agent/src/compact/auto.rs index 0bdc3d18..7f9bac0f 100644 --- a/crates/aion-agent/src/compact/auto.rs +++ b/crates/aion-agent/src/compact/auto.rs @@ -6,11 +6,10 @@ //! summary. A circuit breaker prevents runaway retries. use aion_config::compact::CompactConfig; -use aion_providers::{LlmProvider, ProviderError}; +use aion_providers::{LlmProvider, ProviderFailure, ProviderFailureKind, ProviderStreamReceiver}; use aion_types::compact::{CompactMetadata, CompactTrigger}; use aion_types::llm::{LlmEvent, LlmRequest, ThinkingConfig}; use aion_types::message::{ContentBlock, Message, Role, TokenUsage}; -use tokio::sync::mpsc; use super::prompt::{ COMPACT_MAX_OUTPUT_TOKENS, COMPACT_SYSTEM_PROMPT, build_compact_prompt, build_summary_content, @@ -42,13 +41,11 @@ pub struct CompactResult { #[derive(Debug, thiserror::Error)] pub enum CompactError { #[error("LLM provider error: {0}")] - Provider(#[from] ProviderError), + Provider(#[from] ProviderFailure), #[error("Prompt too long after {attempts} retries")] PromptTooLong { attempts: u32 }, #[error("Empty response from LLM")] EmptyResponse, - #[error("Stream error: {0}")] - StreamError(String), #[error("Circuit breaker tripped after {failures} consecutive failures")] CircuitBroken { failures: u32 }, } @@ -130,7 +127,7 @@ pub async fn autocompact( return Err(e); } }, - Err(ProviderError::PromptTooLong(_)) if ptl_attempts < MAX_PTL_RETRIES => { + Err(e) if is_context_too_large(&e) && ptl_attempts < MAX_PTL_RETRIES => { ptl_attempts += 1; // Remove the summary prompt (last msg), truncate, re-add prompt let conversation_part = &conv_messages[..conv_messages.len() - 1]; @@ -152,7 +149,7 @@ pub async fn autocompact( } } } - Err(ProviderError::PromptTooLong(_)) => { + Err(e) if is_context_too_large(&e) => { state.record_failure(); return Err(CompactError::PromptTooLong { attempts: ptl_attempts, @@ -212,15 +209,15 @@ pub async fn autocompact( /// Collect all text from a streaming LLM response. async fn collect_stream_text( - mut rx: mpsc::Receiver, + mut rx: ProviderStreamReceiver, ) -> Result<(String, TokenUsage), CompactError> { let mut text = String::new(); - while let Some(event) = rx.recv().await { + while let Some(item) = rx.recv().await { + let event = item?; match event { LlmEvent::TextDelta(delta) => text.push_str(&delta), LlmEvent::Done { usage, .. } => return Ok((text, usage)), - LlmEvent::Error(e) => return Err(CompactError::StreamError(e)), // Ignore thinking deltas and tool calls (shouldn't happen in compact) _ => {} } @@ -230,6 +227,10 @@ async fn collect_stream_text( Err(CompactError::EmptyResponse) } +fn is_context_too_large(error: &ProviderFailure) -> bool { + matches!(error.kind, ProviderFailureKind::ContextTooLarge) +} + /// Truncate the oldest ~20% of messages for PTL retry. /// /// Returns `None` if there are too few messages to truncate meaningfully. diff --git a/crates/aion-agent/src/engine.rs b/crates/aion-agent/src/engine.rs index 82a2bb61..54962eb8 100644 --- a/crates/aion-agent/src/engine.rs +++ b/crates/aion-agent/src/engine.rs @@ -6,7 +6,7 @@ use crate::cache_diagnostics::{CacheBreakDetector, CacheDiagnostic, CacheStats}; use crate::compact::state::CompactState; use crate::compact::{auto, emergency, estimate, micro}; use crate::confirm::ToolConfirmer; -use crate::error::AgentError; +use crate::error::{AgentError, CommandFailure}; use crate::orchestration::{ ExecutionControl, execute_tool_calls, execute_tool_calls_with_approval, }; @@ -466,7 +466,9 @@ impl AgentEngine { } Err(e) => { tracing::error!(command = cmd_name, error = %e, "Slash command failed"); - Err(AgentError::ApiError(e.to_string())) + Err(AgentError::Command(CommandFailure::Failed { + command: cmd_name.to_string(), + })) } }; } @@ -543,7 +545,8 @@ impl AgentEngine { let mut stop_reason = StopReason::EndTurn; let mut turn_usage = TokenUsage::default(); - while let Some(event) = rx.recv().await { + while let Some(item) = rx.recv().await { + let event = item?; match event { LlmEvent::TextDelta(text) => { self.output.emit_text_delta(&text, &self.current_msg_id); @@ -592,9 +595,6 @@ impl AgentEngine { stop_reason = sr; turn_usage = usage; } - LlmEvent::Error(e) => { - return Err(AgentError::ApiError(e)); - } } } @@ -1084,7 +1084,7 @@ impl AgentEngine { mod tests_set_config { use std::sync::{Arc, Mutex}; - use aion_providers::error::ProviderError; + use aion_providers::error::ProviderFailure; use aion_providers::provider::LlmProvider; use aion_tools::registry::ToolRegistry; use aion_types::llm::{LlmEvent, LlmRequest}; @@ -1110,7 +1110,8 @@ mod tests_set_config { async fn stream( &self, _: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result>, ProviderFailure> + { let (_tx, rx) = tokio::sync::mpsc::channel(1); Ok(rx) } @@ -1413,7 +1414,7 @@ mod tests_set_config { mod tests_phase6 { use std::sync::{Arc, Mutex}; - use aion_providers::error::ProviderError; + use aion_providers::error::ProviderFailure; use aion_providers::provider::LlmProvider; use aion_tools::registry::ToolRegistry; use aion_types::llm::{LlmEvent, LlmRequest}; @@ -1440,7 +1441,8 @@ mod tests_phase6 { async fn stream( &self, _: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result>, ProviderFailure> + { let (_tx, rx) = tokio::sync::mpsc::channel(1); Ok(rx) } @@ -1612,7 +1614,7 @@ mod tests_compact { use std::sync::{Arc, Mutex}; use aion_config::compact::CompactConfig; - use aion_providers::error::ProviderError; + use aion_providers::error::ProviderFailure; use aion_providers::provider::LlmProvider; use aion_tools::registry::ToolRegistry; use aion_types::llm::{LlmEvent, LlmRequest}; @@ -1664,7 +1666,8 @@ mod tests_compact { async fn stream( &self, _: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result>, ProviderFailure> + { let (_tx, rx) = tokio::sync::mpsc::channel(1); Ok(rx) } @@ -1995,7 +1998,7 @@ mod tests_plan_mode { use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; - use aion_providers::error::ProviderError; + use aion_providers::error::ProviderFailure; use aion_providers::provider::LlmProvider; use aion_tools::registry::ToolRegistry; use aion_types::llm::{LlmEvent, LlmRequest}; @@ -2024,7 +2027,8 @@ mod tests_plan_mode { async fn stream( &self, _: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result>, ProviderFailure> + { let (_tx, rx) = tokio::sync::mpsc::channel(1); Ok(rx) } @@ -2211,7 +2215,7 @@ mod tests_plan_mode { mod tests_handle_command { use std::sync::{Arc, Mutex}; - use aion_providers::error::ProviderError; + use aion_providers::error::ProviderFailure; use aion_providers::provider::LlmProvider; use aion_tools::registry::ToolRegistry; use aion_types::llm::{LlmEvent, LlmRequest}; @@ -2239,7 +2243,8 @@ mod tests_handle_command { async fn stream( &self, _: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result>, ProviderFailure> + { let (_tx, rx) = tokio::sync::mpsc::channel(1); Ok(rx) } diff --git a/crates/aion-agent/src/error.rs b/crates/aion-agent/src/error.rs index a3501f31..986e5b4c 100644 --- a/crates/aion-agent/src/error.rs +++ b/crates/aion-agent/src/error.rs @@ -1,17 +1,33 @@ -use aion_providers::error::ProviderError; +use aion_providers::error::ProviderFailure; #[derive(Debug, thiserror::Error)] -pub enum AgentError { - #[error("API error: {0}")] - ApiError(String), +pub enum AgentFailure { #[error( "provider repeatedly returned malformed tool calls ({count}/{limit}); stopped to avoid wasting tokens" )] RepeatedMalformedToolCall { count: usize, limit: usize }, #[error("Provider error: {0}")] - Provider(#[from] ProviderError), + Provider(#[from] ProviderFailure), #[error("User aborted the session")] UserAborted, #[error("Context window nearly full ({input_tokens} tokens used, limit {limit})")] ContextTooLong { input_tokens: u64, limit: usize }, + #[error("Command error: {0}")] + Command(CommandFailure), + #[error("Internal error: {0}")] + Internal(InternalFailure), +} + +pub type AgentError = AgentFailure; + +#[derive(Debug, thiserror::Error)] +pub enum CommandFailure { + #[error("slash command failed: {command}")] + Failed { command: String }, +} + +#[derive(Debug, thiserror::Error)] +pub enum InternalFailure { + #[error("{message}")] + Unexpected { message: String }, } diff --git a/crates/aion-agent/src/lib.rs b/crates/aion-agent/src/lib.rs index 75560816..9e70a3ec 100644 --- a/crates/aion-agent/src/lib.rs +++ b/crates/aion-agent/src/lib.rs @@ -12,6 +12,7 @@ pub mod error; pub mod orchestration; pub mod output; pub mod plan; +pub mod public_error; pub mod session; pub mod skill_tool; pub mod spawn_tool; diff --git a/crates/aion-agent/src/output/mod.rs b/crates/aion-agent/src/output/mod.rs index 563f9d80..90247e83 100644 --- a/crates/aion-agent/src/output/mod.rs +++ b/crates/aion-agent/src/output/mod.rs @@ -2,6 +2,7 @@ pub mod null_sink; pub mod protocol_sink; pub mod terminal; +use aion_protocol::events::PublicError; use crossterm::execute; use crossterm::style::{Attribute, Color, Print, ResetColor, SetAttribute, SetForegroundColor}; use std::io::{self, Write}; @@ -30,6 +31,10 @@ pub trait OutputSink: Send + Sync { ); /// Display error fn emit_error(&self, msg: &str); + /// Display a structured public error. + fn emit_public_error(&self, error: &PublicError) { + self.emit_error(&error.message); + } /// Display informational message fn emit_info(&self, msg: &str); } diff --git a/crates/aion-agent/src/output/protocol_sink.rs b/crates/aion-agent/src/output/protocol_sink.rs index b09bcac0..3b888d44 100644 --- a/crates/aion-agent/src/output/protocol_sink.rs +++ b/crates/aion-agent/src/output/protocol_sink.rs @@ -1,7 +1,9 @@ use std::sync::Arc; use aion_config::compat::ProviderCompat; -use aion_protocol::events::{Capabilities, ErrorInfo, ProtocolEvent, Usage}; +use aion_protocol::events::{ + Capabilities, ErrorInfo, ErrorOwnership, ProtocolEvent, PublicError, PublicErrorCode, Usage, +}; use aion_protocol::writer::{ProtocolEmitter, ProtocolWriter}; use super::OutputSink; @@ -132,13 +134,21 @@ impl OutputSink for ProtocolSink { let _ = self.writer.emit(&ProtocolEvent::Error { msg_id: None, error: ErrorInfo { - code: "engine_error".to_string(), + code: PublicErrorCode::InternalError, message: msg.to_string(), - retryable: false, + ownership: ErrorOwnership::Aionrs, + details: vec![], }, }); } + fn emit_public_error(&self, error: &PublicError) { + let _ = self.writer.emit(&ProtocolEvent::Error { + msg_id: None, + error: error.clone(), + }); + } + fn emit_info(&self, msg: &str) { let _ = self.writer.emit(&ProtocolEvent::Info { msg_id: String::new(), diff --git a/crates/aion-agent/src/public_error.rs b/crates/aion-agent/src/public_error.rs new file mode 100644 index 00000000..857fd028 --- /dev/null +++ b/crates/aion-agent/src/public_error.rs @@ -0,0 +1,503 @@ +use crate::error::{AgentFailure, CommandFailure, InternalFailure}; +use aion_protocol::events::{ + ErrorOwnership, PublicError, PublicErrorCode, PublicErrorDetail, PublicErrorDetailKey, +}; +use aion_providers::{ + InvalidProviderRequestSource, ProviderFailure, ProviderFailureKind, ProviderFailurePhase, +}; + +pub fn agent_failure_to_public_error(error: &AgentFailure) -> PublicError { + match error { + AgentFailure::Provider(error) => provider_failure_to_public_error(error), + AgentFailure::RepeatedMalformedToolCall { .. } => public_error( + PublicErrorCode::ProviderToolCallInvalid, + ErrorOwnership::Provider, + vec![], + ), + AgentFailure::UserAborted => { + public_error(PublicErrorCode::UserAborted, ErrorOwnership::User, vec![]) + } + AgentFailure::ContextTooLong { .. } => public_error( + PublicErrorCode::CompactionContextStillTooLarge, + ErrorOwnership::User, + vec![], + ), + AgentFailure::Command(CommandFailure::Failed { command }) => public_error( + PublicErrorCode::CommandFailed, + ErrorOwnership::User, + vec![PublicErrorDetail::new( + PublicErrorDetailKey::Command, + command, + )], + ), + AgentFailure::Internal(InternalFailure::Unexpected { .. }) => public_error( + PublicErrorCode::InternalError, + ErrorOwnership::Aionrs, + vec![], + ), + } +} + +fn provider_failure_to_public_error(error: &ProviderFailure) -> PublicError { + let code = provider_failure_kind_to_public_error_code(&error.kind, error.meta.phase); + public_error( + code, + provider_failure_ownership(&error.kind), + provider_details(error), + ) +} + +pub fn provider_failure_kind_to_public_error_code( + kind: &ProviderFailureKind, + phase: ProviderFailurePhase, +) -> PublicErrorCode { + match kind { + ProviderFailureKind::CredentialMissing => PublicErrorCode::ProviderCredentialMissing, + ProviderFailureKind::AuthFailed => PublicErrorCode::ProviderAuthFailed, + ProviderFailureKind::PermissionDenied => PublicErrorCode::ProviderPermissionDenied, + ProviderFailureKind::BillingRequired => PublicErrorCode::ProviderBillingRequired, + ProviderFailureKind::QuotaExceeded => PublicErrorCode::ProviderQuotaExceeded, + ProviderFailureKind::RateLimited => PublicErrorCode::ProviderRateLimited, + ProviderFailureKind::ModelNotFound => PublicErrorCode::ProviderModelNotFound, + ProviderFailureKind::ModelUnsupported => PublicErrorCode::ProviderModelUnsupported, + ProviderFailureKind::ModelUnavailable => PublicErrorCode::ProviderModelUnavailable, + ProviderFailureKind::EndpointNotFound => PublicErrorCode::ProviderEndpointNotFound, + ProviderFailureKind::ContextTooLarge => PublicErrorCode::ProviderContextTooLarge, + ProviderFailureKind::ToolSchema(_) => PublicErrorCode::ProviderToolSchemaInvalid, + ProviderFailureKind::ToolCall(_) => PublicErrorCode::ProviderToolCallInvalid, + ProviderFailureKind::ContentBlocked => PublicErrorCode::ProviderContentBlocked, + ProviderFailureKind::Timeout => PublicErrorCode::ProviderTimeout, + ProviderFailureKind::Transport + if matches!( + phase, + ProviderFailurePhase::AfterFirstDelta | ProviderFailurePhase::AfterToolCall + ) => + { + PublicErrorCode::ProviderStreamInterrupted + } + ProviderFailureKind::Transport => PublicErrorCode::ProviderTransportFailed, + ProviderFailureKind::Server => PublicErrorCode::ProviderServerError, + ProviderFailureKind::InvalidRequest(_) => PublicErrorCode::ProviderInvalidRequest, + ProviderFailureKind::ResponseParse => PublicErrorCode::ProviderResponseParseFailed, + ProviderFailureKind::EmptyResponse => PublicErrorCode::ProviderEmptyResponse, + ProviderFailureKind::Unknown => PublicErrorCode::ProviderUnknownError, + } +} + +fn provider_failure_ownership(kind: &ProviderFailureKind) -> ErrorOwnership { + match kind { + ProviderFailureKind::CredentialMissing + | ProviderFailureKind::AuthFailed + | ProviderFailureKind::PermissionDenied + | ProviderFailureKind::BillingRequired + | ProviderFailureKind::QuotaExceeded + | ProviderFailureKind::ModelNotFound + | ProviderFailureKind::ModelUnsupported + | ProviderFailureKind::EndpointNotFound + | ProviderFailureKind::ContextTooLarge => ErrorOwnership::User, + ProviderFailureKind::InvalidRequest(failure) => match failure.source { + InvalidProviderRequestSource::AionrsRequestBuilder => ErrorOwnership::Aionrs, + InvalidProviderRequestSource::Unknown => ErrorOwnership::User, + }, + ProviderFailureKind::ToolSchema(_) | ProviderFailureKind::ToolCall(_) => { + ErrorOwnership::Provider + } + ProviderFailureKind::RateLimited + | ProviderFailureKind::ModelUnavailable + | ProviderFailureKind::ContentBlocked + | ProviderFailureKind::Timeout + | ProviderFailureKind::Transport + | ProviderFailureKind::Server + | ProviderFailureKind::ResponseParse + | ProviderFailureKind::EmptyResponse => ErrorOwnership::Provider, + ProviderFailureKind::Unknown => ErrorOwnership::Unknown, + } +} + +fn provider_details(error: &ProviderFailure) -> Vec { + let mut details = vec![ + PublicErrorDetail::new(PublicErrorDetailKey::Provider, &error.meta.provider), + PublicErrorDetail::new( + PublicErrorDetailKey::Phase, + provider_phase_detail(error.meta.phase), + ), + ]; + + if let Some(model) = error.meta.model.as_deref() { + details.push(PublicErrorDetail::new(PublicErrorDetailKey::Model, model)); + } + if let Some(status) = error.meta.status { + details.push(PublicErrorDetail::new( + PublicErrorDetailKey::Status, + status.to_string(), + )); + } + if let Some(request_id) = error.meta.request_id.as_deref() { + details.push(PublicErrorDetail::new( + PublicErrorDetailKey::RequestId, + request_id, + )); + } + if let Some(raw_code) = error.meta.raw.raw_code.as_deref() { + details.push(PublicErrorDetail::new( + PublicErrorDetailKey::RawCode, + raw_code, + )); + } + if let Some(raw_type) = error.meta.raw.raw_type.as_deref() { + details.push(PublicErrorDetail::new( + PublicErrorDetailKey::RawType, + raw_type, + )); + } + + details +} + +fn public_error( + code: PublicErrorCode, + ownership: ErrorOwnership, + details: Vec, +) -> PublicError { + PublicError { + code, + message: public_message(code).to_string(), + ownership, + details, + } +} + +fn public_message(code: PublicErrorCode) -> &'static str { + match code { + PublicErrorCode::ProviderCredentialMissing => { + "The model provider credentials are not configured." + } + PublicErrorCode::ProviderAuthFailed => { + "The model provider rejected the configured credentials." + } + PublicErrorCode::ProviderPermissionDenied => { + "The model provider denied permission for this request." + } + PublicErrorCode::ProviderBillingRequired => { + "The model provider requires billing to complete this request." + } + PublicErrorCode::ProviderQuotaExceeded => "The model provider quota has been exceeded.", + PublicErrorCode::ProviderRateLimited => "The model provider rate-limited the request.", + PublicErrorCode::ProviderModelNotFound => { + "The requested model was not found by the provider." + } + PublicErrorCode::ProviderModelUnsupported => { + "The requested model does not support this request." + } + PublicErrorCode::ProviderModelUnavailable => { + "The requested model is currently unavailable." + } + PublicErrorCode::ProviderEndpointNotFound => "The model provider endpoint was not found.", + PublicErrorCode::ProviderContextTooLarge => { + "The request is too large for the model provider context window." + } + PublicErrorCode::ProviderToolSchemaInvalid => { + "The model provider rejected the tool schema." + } + PublicErrorCode::ProviderToolCallInvalid => { + "The model provider returned an invalid tool call." + } + PublicErrorCode::ProviderContentBlocked => { + "The model provider blocked the requested content." + } + PublicErrorCode::ProviderTimeout => "The model provider request timed out.", + PublicErrorCode::ProviderStreamInterrupted => "The model provider stream was interrupted.", + PublicErrorCode::ProviderTransportFailed => "The model provider transport request failed.", + PublicErrorCode::ProviderServerError => "The model provider returned a server error.", + PublicErrorCode::ProviderInvalidRequest => "The model provider rejected the request.", + PublicErrorCode::ProviderResponseParseFailed => { + "The model provider response could not be parsed." + } + PublicErrorCode::ProviderEmptyResponse => "The model provider returned an empty response.", + PublicErrorCode::ProviderUnknownError => "The model provider returned an unknown error.", + PublicErrorCode::CompactionContextStillTooLarge => { + "The conversation is still too large after compaction." + } + PublicErrorCode::CommandFailed => "The command failed.", + PublicErrorCode::UserAborted => "The user aborted the operation.", + PublicErrorCode::InternalInvariantViolation => { + "Aionrs encountered an internal invariant violation." + } + PublicErrorCode::InternalError => "Aionrs encountered an internal error.", + PublicErrorCode::ConfigInvalid + | PublicErrorCode::ConfigProfileNotFound + | PublicErrorCode::ConfigProviderAliasInvalid + | PublicErrorCode::ConfigEnvMissing + | PublicErrorCode::ConfigFileReadFailed + | PublicErrorCode::ConfigFileParseFailed + | PublicErrorCode::ConfigFileWriteFailed + | PublicErrorCode::BootstrapFailed + | PublicErrorCode::BootstrapProviderInitFailed + | PublicErrorCode::BootstrapToolInitFailed + | PublicErrorCode::BootstrapMcpInitFailed + | PublicErrorCode::BootstrapMemoryInitFailed + | PublicErrorCode::SessionCreateFailed + | PublicErrorCode::SessionLoadFailed + | PublicErrorCode::SessionSaveFailed + | PublicErrorCode::SessionListFailed + | PublicErrorCode::SessionIndexFailed + | PublicErrorCode::SessionNotFound + | PublicErrorCode::ToolNotFound + | PublicErrorCode::ToolInputInvalid + | PublicErrorCode::ToolExecutionFailed + | PublicErrorCode::ToolTimeout + | PublicErrorCode::ToolPermissionDenied + | PublicErrorCode::McpServerUnavailable + | PublicErrorCode::McpServerConfigInvalid + | PublicErrorCode::McpProtocolError + | PublicErrorCode::McpCapabilityDiscoveryFailed + | PublicErrorCode::McpInvocationFailed + | PublicErrorCode::ApprovalRejected + | PublicErrorCode::ApprovalTimeout + | PublicErrorCode::ApprovalChannelClosed + | PublicErrorCode::CompactionFailed + | PublicErrorCode::ProtocolInvalidCommand + | PublicErrorCode::ProtocolParseFailed + | PublicErrorCode::ProtocolClientDisconnected + | PublicErrorCode::ProtocolStateViolation => "Aionrs encountered an error.", + } +} + +fn provider_phase_detail(phase: ProviderFailurePhase) -> &'static str { + match phase { + ProviderFailurePhase::BeforeFirstDelta => "before_first_delta", + ProviderFailurePhase::AfterFirstDelta => "after_first_delta", + ProviderFailurePhase::AfterToolCall => "after_tool_call", + } +} + +#[cfg(test)] +mod tests { + use super::agent_failure_to_public_error; + use crate::error::{AgentFailure, CommandFailure}; + use aion_protocol::events::{ + ErrorOwnership, PublicErrorCode, PublicErrorDetail, PublicErrorDetailKey, + }; + use aion_providers::{ + ProviderFailurePhase, ProviderRawSignalSource, ProviderRawSignals, + classify_provider_failure, + }; + + fn detail_value(error: &aion_protocol::events::PublicError, key: PublicErrorDetailKey) -> &str { + error + .details + .iter() + .find(|detail| detail.key == key) + .map(|detail| detail.value.as_str()) + .expect("missing detail") + } + + #[test] + fn provider_auth_failed_maps_to_public_user_error_with_details() { + let failure = classify_provider_failure(ProviderRawSignals { + provider: "openai".to_string(), + model: Some("gpt-4.1".to_string()), + status: Some(401), + raw_code: Some("invalid_api_key".to_string()), + raw_type: Some("authentication_error".to_string()), + message: Some("raw provider auth text must not leak".to_string()), + request_id: Some("req_123".to_string()), + retry_after_ms: None, + source: ProviderRawSignalSource::Http, + phase: ProviderFailurePhase::BeforeFirstDelta, + }); + + let public = agent_failure_to_public_error(&AgentFailure::Provider(failure)); + + assert_eq!(public.code, PublicErrorCode::ProviderAuthFailed); + assert_eq!(public.ownership, ErrorOwnership::User); + assert_eq!( + detail_value(&public, PublicErrorDetailKey::Provider), + "openai" + ); + assert_eq!(detail_value(&public, PublicErrorDetailKey::Status), "401"); + assert_eq!( + detail_value(&public, PublicErrorDetailKey::Model), + "gpt-4.1" + ); + assert_eq!( + detail_value(&public, PublicErrorDetailKey::RequestId), + "req_123" + ); + assert_eq!( + detail_value(&public, PublicErrorDetailKey::RawCode), + "invalid_api_key" + ); + assert_eq!( + detail_value(&public, PublicErrorDetailKey::RawType), + "authentication_error" + ); + assert!(!public.message.contains("raw provider auth text")); + } + + #[test] + fn provider_transport_after_first_delta_maps_to_stream_interrupted() { + let failure = classify_provider_failure(ProviderRawSignals { + provider: "anthropic".to_string(), + model: None, + status: None, + raw_code: None, + raw_type: None, + message: Some("connection reset".to_string()), + request_id: None, + retry_after_ms: None, + source: ProviderRawSignalSource::Transport, + phase: ProviderFailurePhase::AfterFirstDelta, + }); + + let public = agent_failure_to_public_error(&AgentFailure::Provider(failure)); + + assert_eq!(public.code, PublicErrorCode::ProviderStreamInterrupted); + assert_eq!(public.ownership, ErrorOwnership::Provider); + } + + #[test] + fn provider_transport_before_first_delta_maps_to_transport_failed() { + let failure = classify_provider_failure(ProviderRawSignals { + provider: "anthropic".to_string(), + model: None, + status: None, + raw_code: None, + raw_type: None, + message: Some("dns failure".to_string()), + request_id: None, + retry_after_ms: None, + source: ProviderRawSignalSource::Transport, + phase: ProviderFailurePhase::BeforeFirstDelta, + }); + + let public = agent_failure_to_public_error(&AgentFailure::Provider(failure)); + + assert_eq!(public.code, PublicErrorCode::ProviderTransportFailed); + assert_eq!(public.ownership, ErrorOwnership::Provider); + } + + #[test] + fn provider_invalid_request_from_aionrs_request_builder_maps_to_aionrs_ownership() { + let failure = classify_provider_failure(ProviderRawSignals { + provider: "openai".to_string(), + model: None, + status: None, + raw_code: None, + raw_type: None, + message: Some("invalid request payload".to_string()), + request_id: None, + retry_after_ms: None, + source: ProviderRawSignalSource::RequestBuild, + phase: ProviderFailurePhase::BeforeFirstDelta, + }); + + let public = agent_failure_to_public_error(&AgentFailure::Provider(failure)); + + assert_eq!(public.code, PublicErrorCode::ProviderInvalidRequest); + assert_eq!(public.ownership, ErrorOwnership::Aionrs); + } + + #[test] + fn provider_rate_limited_maps_to_provider_rate_limited() { + let failure = classify_provider_failure(ProviderRawSignals { + provider: "openai".to_string(), + model: None, + status: Some(429), + raw_code: None, + raw_type: None, + message: None, + request_id: None, + retry_after_ms: Some(1000), + source: ProviderRawSignalSource::Http, + phase: ProviderFailurePhase::BeforeFirstDelta, + }); + + let public = agent_failure_to_public_error(&AgentFailure::Provider(failure)); + + assert_eq!(public.code, PublicErrorCode::ProviderRateLimited); + assert_eq!(public.ownership, ErrorOwnership::Provider); + } + + #[test] + fn provider_quota_exceeded_maps_to_provider_quota_exceeded() { + let failure = classify_provider_failure(ProviderRawSignals { + provider: "openai".to_string(), + model: None, + status: Some(429), + raw_code: Some("insufficient_quota".to_string()), + raw_type: None, + message: Some("quota exceeded".to_string()), + request_id: None, + retry_after_ms: None, + source: ProviderRawSignalSource::Http, + phase: ProviderFailurePhase::BeforeFirstDelta, + }); + + let public = agent_failure_to_public_error(&AgentFailure::Provider(failure)); + + assert_eq!(public.code, PublicErrorCode::ProviderQuotaExceeded); + assert_eq!(public.ownership, ErrorOwnership::User); + } + + #[test] + fn provider_response_parse_maps_to_response_parse_failed() { + let failure = classify_provider_failure(ProviderRawSignals { + provider: "openai".to_string(), + model: None, + status: None, + raw_code: None, + raw_type: None, + message: Some("invalid json".to_string()), + request_id: None, + retry_after_ms: None, + source: ProviderRawSignalSource::Json, + phase: ProviderFailurePhase::BeforeFirstDelta, + }); + + let public = agent_failure_to_public_error(&AgentFailure::Provider(failure)); + + assert_eq!(public.code, PublicErrorCode::ProviderResponseParseFailed); + assert_eq!(public.ownership, ErrorOwnership::Provider); + } + + #[test] + fn repeated_malformed_tool_call_maps_to_provider_tool_call_invalid() { + let public = agent_failure_to_public_error(&AgentFailure::RepeatedMalformedToolCall { + count: 3, + limit: 3, + }); + + assert_eq!(public.code, PublicErrorCode::ProviderToolCallInvalid); + assert_eq!(public.ownership, ErrorOwnership::Provider); + } + + #[test] + fn context_too_long_maps_to_compaction_context_still_too_large() { + let public = agent_failure_to_public_error(&AgentFailure::ContextTooLong { + input_tokens: 101, + limit: 100, + }); + + assert_eq!(public.code, PublicErrorCode::CompactionContextStillTooLarge); + assert_eq!(public.ownership, ErrorOwnership::User); + } + + #[test] + fn command_failed_maps_to_command_failed_with_command_detail() { + let public = + agent_failure_to_public_error(&AgentFailure::Command(CommandFailure::Failed { + command: "/compact".to_string(), + })); + + assert_eq!(public.code, PublicErrorCode::CommandFailed); + assert_eq!(public.ownership, ErrorOwnership::User); + assert_eq!( + public.details, + vec![PublicErrorDetail::new( + PublicErrorDetailKey::Command, + "/compact" + )] + ); + } +} diff --git a/crates/aion-agent/tests/autocompact_test.rs b/crates/aion-agent/tests/autocompact_test.rs index 425f36c7..7cda8d2a 100644 --- a/crates/aion-agent/tests/autocompact_test.rs +++ b/crates/aion-agent/tests/autocompact_test.rs @@ -19,7 +19,7 @@ use aion_agent::compact::prompt::{ }; use aion_agent::compact::state::CompactState; use aion_config::compact::CompactConfig; -use aion_providers::{LlmProvider, ProviderError}; +use aion_providers::{LlmProvider, ProviderError, ProviderFailure, ProviderStreamReceiver}; use aion_types::compact::CompactTrigger; use aion_types::llm::{LlmEvent, LlmRequest}; use aion_types::message::{ContentBlock, Message, Role, StopReason, TokenUsage}; @@ -28,11 +28,21 @@ use aion_types::message::{ContentBlock, Message, Role, StopReason, TokenUsage}; /// A mock LLM provider that returns pre-configured responses in order. struct MockProvider { - responses: Mutex, ProviderError>>>, + responses: Mutex>, ProviderFailure>>>, } impl MockProvider { - fn new(responses: Vec, ProviderError>>) -> Self { + fn new(responses: Vec, ProviderFailure>>) -> Self { + let responses = responses + .into_iter() + .map(|response| response.map(|events| events.into_iter().map(Ok).collect())) + .collect(); + Self::new_streams(responses) + } + + fn new_streams( + responses: Vec>, ProviderFailure>>, + ) -> Self { Self { responses: Mutex::new(VecDeque::from(responses)), } @@ -54,7 +64,7 @@ impl MockProvider { } /// Create a provider that returns an error. - fn with_error(error: ProviderError) -> Self { + fn with_error(error: ProviderFailure) -> Self { Self::new(vec![Err(error)]) } } @@ -64,7 +74,7 @@ impl LlmProvider for MockProvider { async fn stream( &self, _request: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result { let response = self .responses .lock() @@ -75,8 +85,8 @@ impl LlmProvider for MockProvider { match response { Ok(events) => { let (tx, rx) = mpsc::channel(events.len() + 1); - for event in events { - tx.send(event).await.ok(); + for item in events { + tx.send(item).await.ok(); } Ok(rx) } @@ -315,10 +325,13 @@ async fn tc_2_4_16_success_resets_failure_counter() { #[tokio::test] async fn tc_2_4_17_failure_increments_counter() { - let provider = MockProvider::with_error(ProviderError::Api { - status: 500, - message: "Internal error".to_string(), - }); + let provider = MockProvider::with_error( + ProviderError::Api { + status: 500, + message: "Internal error".to_string(), + } + .into(), + ); let messages = sample_conversation(10); let config = default_config(); let mut state = CompactState::new(); @@ -334,9 +347,7 @@ async fn tc_2_4_17_failure_increments_counter() { async fn tc_2_4_18_ptl_retry_succeeds() { let provider = MockProvider::new(vec![ // First attempt: prompt too long - Err(ProviderError::PromptTooLong( - "prompt exceeds limit".to_string(), - )), + Err(ProviderError::PromptTooLong("prompt exceeds limit".to_string()).into()), // Second attempt (after truncation): success Ok(vec![ LlmEvent::TextDelta("retried summary".to_string()), @@ -373,9 +384,9 @@ async fn tc_2_4_18_ptl_retry_succeeds() { #[tokio::test] async fn tc_2_4_19_ptl_retry_exhausted() { let provider = MockProvider::new(vec![ - Err(ProviderError::PromptTooLong("too long 1".to_string())), - Err(ProviderError::PromptTooLong("too long 2".to_string())), - Err(ProviderError::PromptTooLong("too long 3".to_string())), + Err(ProviderError::PromptTooLong("too long 1".to_string()).into()), + Err(ProviderError::PromptTooLong("too long 2".to_string()).into()), + Err(ProviderError::PromptTooLong("too long 3".to_string()).into()), ]); let messages = sample_conversation(20); @@ -407,7 +418,7 @@ async fn tc_2_4_20_ptl_retry_truncates_messages() { async fn stream( &self, request: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result { // Scope the lock so the MutexGuard is dropped before the await let current_attempt = { let mut attempt = self.attempt.lock().unwrap(); @@ -418,20 +429,20 @@ async fn tc_2_4_20_ptl_retry_truncates_messages() { }; if current_attempt == 0 { - return Err(ProviderError::PromptTooLong("too long".to_string())); + return Err(ProviderError::PromptTooLong("too long".to_string()).into()); } // Second attempt: succeed let (tx, rx) = mpsc::channel(2); - tx.send(LlmEvent::TextDelta( + tx.send(Ok(LlmEvent::TextDelta( "truncated summary".to_string(), - )) + ))) .await .ok(); - tx.send(LlmEvent::Done { + tx.send(Ok(LlmEvent::Done { stop_reason: StopReason::EndTurn, usage: TokenUsage::default(), - }) + })) .await .ok(); Ok(rx) @@ -484,9 +495,9 @@ async fn empty_response_fails() { #[tokio::test] async fn stream_error_fails() { - let provider = MockProvider::new(vec![Ok(vec![ - LlmEvent::TextDelta("partial".to_string()), - LlmEvent::Error("connection reset".to_string()), + let provider = MockProvider::new_streams(vec![Ok(vec![ + Ok(LlmEvent::TextDelta("partial".to_string())), + Err(ProviderError::Connection("connection reset".to_string()).into()), ])]); let messages = sample_conversation(10); @@ -494,7 +505,7 @@ async fn stream_error_fails() { let mut state = CompactState::new(); let result = autocompact(&provider, &messages, "test-model", &config, &mut state).await; - assert!(matches!(result, Err(CompactError::StreamError(_)))); + assert!(matches!(result, Err(CompactError::Provider(_)))); assert_eq!(state.consecutive_failures, 1); } diff --git a/crates/aion-agent/tests/common/mod.rs b/crates/aion-agent/tests/common/mod.rs index 3b17212c..34c728cc 100644 --- a/crates/aion-agent/tests/common/mod.rs +++ b/crates/aion-agent/tests/common/mod.rs @@ -12,7 +12,7 @@ use aion_config::config::{Config, ProviderType, SessionConfig, ToolsConfig}; use aion_config::hooks::HooksConfig; use aion_mcp::config::McpConfig; use aion_protocol::events::ToolCategory; -use aion_providers::{LlmProvider, ProviderError}; +use aion_providers::{LlmProvider, ProviderFailure, ProviderStreamReceiver}; use aion_tools::Tool; use aion_types::llm::{LlmEvent, LlmRequest}; use aion_types::message::{StopReason, TokenUsage}; @@ -26,7 +26,7 @@ use aion_types::tool::ToolResult; /// Each call to `stream` pops the first sequence from `responses`. /// When `responses` is empty it falls back to a single EndTurn with empty text. pub struct MockLlmProvider { - responses: Mutex>>, + responses: Mutex>>>, } impl MockLlmProvider { @@ -45,7 +45,7 @@ impl MockLlmProvider { }, ]; Self { - responses: Mutex::new(vec![events]), + responses: Mutex::new(vec![Self::ok_events(events)]), } } @@ -69,7 +69,7 @@ impl MockLlmProvider { }, ]; Self { - responses: Mutex::new(vec![events]), + responses: Mutex::new(vec![Self::ok_events(events)]), } } @@ -77,16 +77,26 @@ impl MockLlmProvider { /// Each call to `stream` consumes the next sequence. pub fn with_turns(turns: Vec>) -> Self { Self { - responses: Mutex::new(turns), + responses: Mutex::new(turns.into_iter().map(Self::ok_events).collect()), } } /// Create a provider that returns custom events. pub fn with_events(events: Vec) -> Self { Self { - responses: Mutex::new(vec![events]), + responses: Mutex::new(vec![Self::ok_events(events)]), } } + + pub fn with_stream_items(items: Vec>) -> Self { + Self { + responses: Mutex::new(vec![items]), + } + } + + fn ok_events(events: Vec) -> Vec> { + events.into_iter().map(Ok).collect() + } } #[async_trait] @@ -94,15 +104,15 @@ impl LlmProvider for MockLlmProvider { async fn stream( &self, _request: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result { let events = { let mut responses = self.responses.lock().unwrap(); if responses.is_empty() { // Fallback: end turn with empty text - vec![LlmEvent::Done { + vec![Ok(LlmEvent::Done { stop_reason: StopReason::EndTurn, usage: TokenUsage::default(), - }] + })] } else { responses.remove(0) } @@ -110,8 +120,8 @@ impl LlmProvider for MockLlmProvider { let (tx, rx) = mpsc::channel(64); tokio::spawn(async move { - for event in events { - let _ = tx.send(event).await; + for item in events { + let _ = tx.send(item).await; } }); Ok(rx) diff --git a/crates/aion-agent/tests/engine_compact_test.rs b/crates/aion-agent/tests/engine_compact_test.rs index 04fc82c2..554e96c1 100644 --- a/crates/aion-agent/tests/engine_compact_test.rs +++ b/crates/aion-agent/tests/engine_compact_test.rs @@ -18,7 +18,7 @@ use aion_agent::output::OutputSink; use aion_agent::output::terminal::TerminalSink; use aion_agent::session::SessionManager; use aion_config::compact::CompactConfig; -use aion_providers::{LlmProvider, ProviderError}; +use aion_providers::{LlmProvider, ProviderError, ProviderFailure, ProviderStreamReceiver}; use aion_tools::registry::ToolRegistry; use aion_types::llm::{LlmEvent, LlmRequest}; use aion_types::message::{StopReason, TokenUsage}; @@ -57,7 +57,7 @@ impl LlmProvider for CompactMockProvider { async fn stream( &self, _request: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result { *self.call_count.lock().unwrap() += 1; let events = self.turns.lock().unwrap().pop_front().unwrap_or_else(|| { vec![LlmEvent::Done { @@ -69,7 +69,7 @@ impl LlmProvider for CompactMockProvider { let (tx, rx) = mpsc::channel(64); tokio::spawn(async move { for event in events { - let _ = tx.send(event).await; + let _ = tx.send(Ok(event)).await; } }); Ok(rx) @@ -485,7 +485,7 @@ async fn tc_2_6_02_micro_before_auto_execution_order() { async fn stream( &self, request: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result { let is_compact = request.tools.is_empty(); if is_compact { @@ -506,7 +506,7 @@ async fn tc_2_6_02_micro_before_auto_execution_order() { let (tx, rx) = mpsc::channel(64); tokio::spawn(async move { for e in events { - let _ = tx.send(e).await; + let _ = tx.send(Ok(e)).await; } }); return Ok(rx); @@ -562,7 +562,7 @@ async fn tc_2_6_02_micro_before_auto_execution_order() { let (tx, rx) = mpsc::channel(64); tokio::spawn(async move { for e in events { - let _ = tx.send(e).await; + let _ = tx.send(Ok(e)).await; } }); Ok(rx) @@ -646,7 +646,7 @@ async fn tc_2_6_e2e_02_micro_and_auto_cooperative() { async fn stream( &self, request: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result { let is_compact = request.tools.is_empty(); if is_compact { @@ -666,7 +666,7 @@ async fn tc_2_6_e2e_02_micro_and_auto_cooperative() { let (tx, rx) = mpsc::channel(64); tokio::spawn(async move { for e in events { - let _ = tx.send(e).await; + let _ = tx.send(Ok(e)).await; } }); return Ok(rx); @@ -718,7 +718,7 @@ async fn tc_2_6_e2e_02_micro_and_auto_cooperative() { let (tx, rx) = mpsc::channel(64); tokio::spawn(async move { for e in events { - let _ = tx.send(e).await; + let _ = tx.send(Ok(e)).await; } }); Ok(rx) @@ -786,7 +786,7 @@ async fn tc_2_6_e2e_03_circuit_breaker_stops_retries() { async fn stream( &self, request: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result { let idx = { let mut i = self.call_index.lock().unwrap(); let v = *i; @@ -802,7 +802,8 @@ async fn tc_2_6_e2e_03_circuit_breaker_stops_retries() { return Err(ProviderError::Api { status: 500, message: "Internal error".to_string(), - }); + } + .into()); } // Regular conversation turns: tool use on odd calls, text on even @@ -842,7 +843,7 @@ async fn tc_2_6_e2e_03_circuit_breaker_stops_retries() { let (tx, rx) = mpsc::channel(64); tokio::spawn(async move { for event in events { - let _ = tx.send(event).await; + let _ = tx.send(Ok(event)).await; } }); Ok(rx) diff --git a/crates/aion-agent/tests/engine_test.rs b/crates/aion-agent/tests/engine_test.rs index ad3c860e..51d6b54f 100644 --- a/crates/aion-agent/tests/engine_test.rs +++ b/crates/aion-agent/tests/engine_test.rs @@ -7,7 +7,10 @@ use aion_agent::error::AgentError; use aion_agent::output::OutputSink; use aion_agent::output::terminal::TerminalSink; use aion_agent::session::SessionManager; -use aion_providers::{LlmProvider, ProviderError}; +use aion_providers::{ + LlmProvider, ProviderFailure, ProviderFailurePhase, ProviderRawSignalSource, + ProviderRawSignals, ProviderStreamReceiver, classify_provider_failure, +}; use aion_tools::registry::ToolRegistry; use aion_types::llm::{LlmEvent, LlmRequest}; use aion_types::message::{ContentBlock, Message, Role, StopReason, TokenUsage}; @@ -107,13 +110,13 @@ impl LlmProvider for RecordingRequestProvider { async fn stream( &self, request: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result { self.requests.lock().unwrap().push(request.messages.clone()); let events = self.responses.lock().unwrap().remove(0); let (tx, rx) = mpsc::channel(64); tokio::spawn(async move { for event in events { - let _ = tx.send(event).await; + let _ = tx.send(Ok(event)).await; } }); Ok(rx) @@ -793,17 +796,43 @@ async fn mixed_valid_and_malformed_tool_calls_do_not_trip_breaker() { ); } -// --------------------------------------------------------------------------- -// test_engine_api_error_handling -// -// Verifies that an LlmEvent::Error propagates as AgentError::ApiError with -// the original error message intact. -// --------------------------------------------------------------------------- #[tokio::test] -async fn test_engine_api_error_handling() { - let events = vec![LlmEvent::Error("test error".to_string())]; +async fn provider_stream_failure_maps_to_agent_provider_failure() { + struct StreamFailureProvider { + failure: ProviderFailure, + } - let provider = Arc::new(MockLlmProvider::with_events(events)); + #[async_trait] + impl LlmProvider for StreamFailureProvider { + async fn stream( + &self, + _request: &LlmRequest, + ) -> Result>, ProviderFailure> { + let (tx, rx) = mpsc::channel(1); + let failure = self.failure.clone(); + tokio::spawn(async move { + let _ = tx.send(Err(failure)).await; + }); + Ok(rx) + } + } + + let failure = classify_provider_failure(ProviderRawSignals { + provider: "test-provider".to_string(), + model: Some("test-model".to_string()), + status: None, + raw_code: None, + raw_type: None, + message: Some("connection reset".to_string()), + request_id: None, + retry_after_ms: None, + source: ProviderRawSignalSource::Transport, + phase: ProviderFailurePhase::AfterFirstDelta, + }); + + let provider = Arc::new(StreamFailureProvider { + failure: failure.clone(), + }); let config = test_config(); let registry = ToolRegistry::new(); let output = silent_output(); @@ -816,8 +845,5 @@ async fn test_engine_api_error_handling() { .map(|_| panic!("expected error, got Ok")) .unwrap_err(); - match err { - AgentError::ApiError(msg) => assert_eq!(msg, "test error"), - other => panic!("expected ApiError(\"test error\"), got: {:?}", other), - } + assert!(matches!(err, AgentError::Provider(actual) if actual == failure)); } diff --git a/crates/aion-agent/tests/output_compaction_test.rs b/crates/aion-agent/tests/output_compaction_test.rs index ad4cbf32..8180f8bd 100644 --- a/crates/aion-agent/tests/output_compaction_test.rs +++ b/crates/aion-agent/tests/output_compaction_test.rs @@ -3,7 +3,6 @@ mod common; use std::sync::{Arc, Mutex}; use async_trait::async_trait; -use tokio::sync::mpsc; use aion_agent::context::{SystemPromptCache, build_system_prompt}; use aion_agent::engine::AgentEngine; @@ -11,7 +10,7 @@ use aion_agent::orchestration::execute_tool_calls; use aion_agent::output::OutputSink; use aion_agent::output::null_sink::NullSink; use aion_compact::CompactionLevel; -use aion_providers::{LlmProvider, ProviderError}; +use aion_providers::{LlmProvider, ProviderFailure, ProviderStreamReceiver}; use aion_tools::registry::ToolRegistry; use aion_types::llm::{LlmEvent, LlmRequest}; use aion_types::message::{ContentBlock, StopReason, TokenUsage}; @@ -260,7 +259,7 @@ impl LlmProvider for CapturingProvider { async fn stream( &self, request: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result { self.captured.lock().unwrap().push(request.clone()); self.inner.stream(request).await } diff --git a/crates/aion-agent/tests/spawn_test.rs b/crates/aion-agent/tests/spawn_test.rs index c0b6fa9c..d5cf28e3 100644 --- a/crates/aion-agent/tests/spawn_test.rs +++ b/crates/aion-agent/tests/spawn_test.rs @@ -3,6 +3,7 @@ mod common; use std::sync::Arc; use aion_agent::spawner::{AgentSpawner, SubAgentConfig}; +use aion_providers::ProviderError; use aion_types::llm::LlmEvent; use aion_types::message::{StopReason, TokenUsage}; use common::{MockLlmProvider, test_config}; @@ -140,13 +141,11 @@ async fn test_spawn_shares_provider() { assert_eq!(result2.text, "second"); } -/// An LLM error event causes the sub-agent result to be marked as an error. +/// A provider stream failure causes the sub-agent result to be marked as an error. #[tokio::test] async fn test_spawn_agent_error_captured() { - // Emit an Error event — the engine converts this to AgentError::ApiError, - // which spawner catches and stores in SubAgentResult::is_error. - let provider = Arc::new(MockLlmProvider::with_events(vec![LlmEvent::Error( - "provider failed".to_string(), + let provider = Arc::new(MockLlmProvider::with_stream_items(vec![Err( + ProviderError::Connection("provider failed".to_string()).into(), )])); let spawner = AgentSpawner::new(provider, test_config(), std::env::temp_dir()); diff --git a/crates/aion-cli/src/main.rs b/crates/aion-cli/src/main.rs index 4596a6dc..74feb886 100644 --- a/crates/aion-cli/src/main.rs +++ b/crates/aion-cli/src/main.rs @@ -8,13 +8,17 @@ use aion_agent::bootstrap::AgentBootstrap; use aion_agent::output::OutputSink; use aion_agent::output::protocol_sink::ProtocolSink; use aion_agent::output::terminal::TerminalSink; +use aion_agent::public_error::agent_failure_to_public_error; use aion_agent::session; use aion_config::auth; use aion_config::config::{self, CliArgs, Config, McpServerConfig, TransportType}; use aion_mcp::manager::McpManager; use aion_mcp::tool_proxy::register_single_server_tools; use aion_protocol::commands::{ApprovalScope, ProtocolCommand}; -use aion_protocol::events::ProtocolEvent; +use aion_protocol::events::{ + ErrorOwnership, ProtocolEvent, PublicError, PublicErrorCode, PublicErrorDetail, + PublicErrorDetailKey, +}; use aion_protocol::reader::spawn_stdin_reader; use aion_protocol::writer::{ProtocolEmitter, ProtocolWriter}; use aion_protocol::{ToolApprovalManager, ToolApprovalResult}; @@ -285,18 +289,27 @@ async fn main() -> anyhow::Result<()> { } let prompt = cli.prompt.join(" "); + let mut prompt_error = None; if prompt.is_empty() { repl_loop(&mut engine, &terminal, &output).await?; } else { - let run_result = engine.run(&prompt, "").await?; - output.emit_stream_end( - "", - run_result.turns, - run_result.usage.input_tokens, - run_result.usage.output_tokens, - run_result.usage.cache_creation_tokens, - run_result.usage.cache_read_tokens, - ); + match engine.run(&prompt, "").await { + Ok(run_result) => { + output.emit_stream_end( + "", + run_result.turns, + run_result.usage.input_tokens, + run_result.usage.output_tokens, + run_result.usage.cache_creation_tokens, + run_result.usage.cache_read_tokens, + ); + } + Err(e) => { + let public = agent_failure_to_public_error(&e); + output.emit_public_error(&public); + prompt_error = Some(e); + } + } } engine.run_stop_hooks().await; @@ -305,7 +318,16 @@ async fn main() -> anyhow::Result<()> { mgr.shutdown().await; } - Ok(()) + non_interactive_prompt_run_result(prompt_error) +} + +fn non_interactive_prompt_run_result( + error: Option, +) -> anyhow::Result<()> { + match error { + Some(error) => Err(error.into()), + None => Ok(()), + } } async fn repl_loop( @@ -341,7 +363,8 @@ async fn repl_loop( } Err(aion_agent::error::AgentError::UserAborted) => break, Err(e) => { - output.emit_error(&e.to_string()); + let public = agent_failure_to_public_error(&e); + output.emit_public_error(&public); } } } @@ -418,6 +441,18 @@ fn to_mcp_server_config( }) } +fn add_mcp_after_message_public_error(name: &str) -> PublicError { + PublicError { + code: PublicErrorCode::ProtocolStateViolation, + message: format!("AddMcpServer '{name}': rejected - only allowed before first Message"), + ownership: ErrorOwnership::Host, + details: vec![PublicErrorDetail::new( + PublicErrorDetailKey::Command, + "add_mcp_server", + )], + } +} + /// Pending config fields: (model, thinking, thinking_budget, effort) type PendingConfig = ( Option, @@ -581,7 +616,8 @@ async fn run_json_stream_mode( ); } Err(e) => { - output.emit_error(&e.to_string()); + let public = agent_failure_to_public_error(&e); + output.emit_public_error(&public); output.emit_stream_end(&msg_id, 0, 0, 0, 0, 0); } } @@ -617,6 +653,9 @@ async fn run_json_stream_mode( ProtocolCommand::Ping => { let _ = writer.emit(&aion_protocol::events::ProtocolEvent::Pong); } + ProtocolCommand::AddMcpServer { name, .. } => { + output.emit_public_error(&add_mcp_after_message_public_error(&name)); + } _ => { tracing::debug!(target: "aion_protocol", "ignoring command during active message processing"); } @@ -717,9 +756,7 @@ async fn run_json_stream_mode( ); } ProtocolCommand::AddMcpServer { name, .. } => { - output.emit_error(&format!( - "AddMcpServer '{name}': rejected — only allowed before first Message" - )); + output.emit_public_error(&add_mcp_after_message_public_error(&name)); } ProtocolCommand::Ping => { let _ = writer.emit(&aion_protocol::events::ProtocolEvent::Pong); @@ -737,3 +774,36 @@ async fn run_json_stream_mode( Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn non_interactive_prompt_error_returns_error() { + let result = super::non_interactive_prompt_run_result(Some( + aion_agent::error::AgentError::UserAborted, + )); + + assert!(result.is_err()); + } + + #[test] + fn add_mcp_after_message_uses_protocol_state_violation_public_error() { + let error = add_mcp_after_message_public_error("tools"); + + assert_eq!(error.code, PublicErrorCode::ProtocolStateViolation); + assert_eq!(error.ownership, ErrorOwnership::Host); + assert_eq!( + error.details, + vec![PublicErrorDetail::new( + PublicErrorDetailKey::Command, + "add_mcp_server", + )] + ); + assert_eq!( + error.message, + "AddMcpServer 'tools': rejected - only allowed before first Message" + ); + } +} diff --git a/crates/aion-protocol/src/events.rs b/crates/aion-protocol/src/events.rs index b317ea76..c08b9b79 100644 --- a/crates/aion-protocol/src/events.rs +++ b/crates/aion-protocol/src/events.rs @@ -136,11 +136,128 @@ pub struct Usage { pub cache_write_tokens: Option, } -#[derive(Debug, Serialize)] -pub struct ErrorInfo { - pub code: String, +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct PublicError { + pub code: PublicErrorCode, pub message: String, - pub retryable: bool, + pub ownership: ErrorOwnership, + pub details: Vec, +} + +pub type ErrorInfo = PublicError; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum PublicErrorCode { + ProviderCredentialMissing, + ProviderAuthFailed, + ProviderPermissionDenied, + ProviderBillingRequired, + ProviderQuotaExceeded, + ProviderRateLimited, + ProviderModelNotFound, + ProviderModelUnsupported, + ProviderModelUnavailable, + ProviderEndpointNotFound, + ProviderContextTooLarge, + ProviderToolSchemaInvalid, + ProviderToolCallInvalid, + ProviderContentBlocked, + ProviderTimeout, + ProviderStreamInterrupted, + ProviderTransportFailed, + ProviderServerError, + ProviderInvalidRequest, + ProviderResponseParseFailed, + ProviderEmptyResponse, + ProviderUnknownError, + ConfigInvalid, + ConfigProfileNotFound, + ConfigProviderAliasInvalid, + ConfigEnvMissing, + ConfigFileReadFailed, + ConfigFileParseFailed, + ConfigFileWriteFailed, + BootstrapFailed, + BootstrapProviderInitFailed, + BootstrapToolInitFailed, + BootstrapMcpInitFailed, + BootstrapMemoryInitFailed, + SessionCreateFailed, + SessionLoadFailed, + SessionSaveFailed, + SessionListFailed, + SessionIndexFailed, + SessionNotFound, + ToolNotFound, + ToolInputInvalid, + ToolExecutionFailed, + ToolTimeout, + ToolPermissionDenied, + McpServerUnavailable, + McpServerConfigInvalid, + McpProtocolError, + McpCapabilityDiscoveryFailed, + McpInvocationFailed, + ApprovalRejected, + ApprovalTimeout, + ApprovalChannelClosed, + CompactionFailed, + CompactionContextStillTooLarge, + CommandFailed, + UserAborted, + ProtocolInvalidCommand, + ProtocolParseFailed, + ProtocolClientDisconnected, + ProtocolStateViolation, + InternalError, + InternalInvariantViolation, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum ErrorOwnership { + User, + Provider, + Aionrs, + Host, + Tool, + McpServer, + Unknown, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct PublicErrorDetail { + pub key: PublicErrorDetailKey, + pub value: String, +} + +impl PublicErrorDetail { + pub fn new(key: PublicErrorDetailKey, value: impl Into) -> Self { + Self { + key, + value: value.into(), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum PublicErrorDetailKey { + Provider, + Model, + Status, + RequestId, + Phase, + ConfigKey, + Profile, + SessionId, + ToolName, + McpServer, + McpCapability, + Command, + RawCode, + RawType, } #[cfg(test)] @@ -238,15 +355,97 @@ mod tests { let event = ProtocolEvent::Error { msg_id: None, error: ErrorInfo { - code: "rate_limit".to_string(), + code: PublicErrorCode::ProviderRateLimited, message: "Too many requests".to_string(), - retryable: true, + ownership: ErrorOwnership::Provider, + details: vec![], }, }; let json = serde_json::to_value(&event).unwrap(); assert_eq!(json["type"], "error"); assert!(json.get("msg_id").is_none()); - assert_eq!(json["error"]["retryable"], true); + assert_eq!(json["error"]["code"], "provider_rate_limited"); + assert_eq!(json["error"]["ownership"], "provider"); + assert!(json["error"].get("retryable").is_none()); + } + + #[test] + fn public_error_serializes_stable_snake_case_contract() { + let event = ProtocolEvent::Error { + msg_id: Some("msg-1".to_string()), + error: PublicError { + code: PublicErrorCode::ProviderAuthFailed, + message: "The model provider rejected the configured credentials.".to_string(), + ownership: ErrorOwnership::User, + details: vec![ + PublicErrorDetail::new(PublicErrorDetailKey::Provider, "openai"), + PublicErrorDetail::new(PublicErrorDetailKey::Status, "401"), + ], + }, + }; + + let json = serde_json::to_value(event).unwrap(); + assert_eq!(json["type"], "error"); + assert_eq!(json["msg_id"], "msg-1"); + assert_eq!(json["error"]["code"], "provider_auth_failed"); + assert_eq!( + json["error"]["message"], + "The model provider rejected the configured credentials." + ); + assert_eq!(json["error"]["ownership"], "user"); + assert_eq!(json["error"]["details"][0]["key"], "provider"); + assert_eq!(json["error"]["details"][0]["value"], "openai"); + assert!(json["error"].get("retryable").is_none()); + } + + #[test] + fn public_error_supports_mcp_capability_detail() { + let event = ProtocolEvent::Error { + msg_id: None, + error: PublicError { + code: PublicErrorCode::McpInvocationFailed, + message: "The MCP server failed while invoking a capability.".to_string(), + ownership: ErrorOwnership::McpServer, + details: vec![ + PublicErrorDetail::new(PublicErrorDetailKey::McpServer, "filesystem"), + PublicErrorDetail::new(PublicErrorDetailKey::McpCapability, "read_file"), + ], + }, + }; + + let json = serde_json::to_value(event).unwrap(); + assert_eq!(json["error"]["code"], "mcp_invocation_failed"); + assert_eq!(json["error"]["ownership"], "mcp_server"); + assert_eq!(json["error"]["details"][1]["key"], "mcp_capability"); + assert_eq!(json["error"]["details"][1]["value"], "read_file"); + } + + #[test] + fn public_error_values_support_equality() { + fn assert_eq_trait(_: &T) {} + + let left = PublicError { + code: PublicErrorCode::InternalError, + message: "Unexpected failure".to_string(), + ownership: ErrorOwnership::Aionrs, + details: vec![PublicErrorDetail::new( + PublicErrorDetailKey::Phase, + "streaming", + )], + }; + let right = PublicError { + code: PublicErrorCode::InternalError, + message: "Unexpected failure".to_string(), + ownership: ErrorOwnership::Aionrs, + details: vec![PublicErrorDetail::new( + PublicErrorDetailKey::Phase, + "streaming", + )], + }; + + assert_eq!(left, right); + assert_eq_trait(&left); + assert_eq_trait(&left.details[0]); } #[test] diff --git a/crates/aion-providers/src/anthropic.rs b/crates/aion-providers/src/anthropic.rs index db4383dc..96453727 100644 --- a/crates/aion-providers/src/anthropic.rs +++ b/crates/aion-providers/src/anthropic.rs @@ -3,10 +3,13 @@ use reqwest::header::{CONTENT_TYPE, HeaderMap, HeaderValue}; use serde_json::{Value, json}; use tokio::sync::mpsc; -use aion_types::llm::{LlmEvent, LlmRequest, ThinkingConfig}; +use aion_types::llm::{LlmRequest, ThinkingConfig}; use super::anthropic_shared; -use crate::{LlmProvider, ProviderError}; +use crate::{ + LlmProvider, ProviderError, ProviderFailure, ProviderFailurePhase, ProviderStreamReceiver, + provider_api_error, provider_failure_from_error, retry_after_ms_from_headers, +}; use aion_config::compat::ProviderCompat; pub struct AnthropicProvider { @@ -94,49 +97,78 @@ impl LlmProvider for AnthropicProvider { async fn stream( &self, request: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result { let url = format!("{}/v1/messages", self.base_url); let body = self.build_request_body(request); + let model = Some(request.model.clone()); tracing::debug!(target: "aion_providers", body = %serde_json::to_string_pretty(&body).unwrap_or_default(), "outgoing request"); let response = self .client .post(&url) - .headers(self.build_headers()?) + .headers(self.build_headers().map_err(|error| { + provider_failure_from_error( + "anthropic", + model.clone(), + error, + ProviderFailurePhase::BeforeFirstDelta, + ) + })?) .json(&body) .send() - .await?; + .await + .map_err(ProviderError::from) + .map_err(|error| { + provider_failure_from_error( + "anthropic", + model.clone(), + error, + ProviderFailurePhase::BeforeFirstDelta, + ) + })?; let status = response.status(); if !status.is_success() { + let retry_after_ms = retry_after_ms_from_headers(response.headers()); let body_text = response.text().await.unwrap_or_default(); - if status.as_u16() == 429 { - return Err(ProviderError::RateLimited { - retry_after_ms: 5000, - }); - } - return Err(ProviderError::Api { - status: status.as_u16(), - message: body_text, - }); + return Err(provider_failure_from_error( + "anthropic", + model, + provider_api_error(status.as_u16(), body_text, retry_after_ms), + ProviderFailurePhase::BeforeFirstDelta, + )); } let (tx, rx) = mpsc::channel(64); let client = self.client.clone(); - let headers = self.build_headers()?; + let headers = self.build_headers().map_err(|error| { + provider_failure_from_error( + "anthropic", + model.clone(), + error, + ProviderFailurePhase::BeforeFirstDelta, + ) + })?; let url_clone = url.clone(); + let stream_model = model.clone(); tokio::spawn(async move { match anthropic_shared::process_sse_stream(response, &tx).await { anthropic_shared::StreamOutcome::Ok => {} anthropic_shared::StreamOutcome::FailedPartial(e) => { - let _ = tx.send(LlmEvent::Error(e.to_string())).await; + let failure = provider_failure_from_error( + "anthropic", + stream_model.clone(), + e, + ProviderFailurePhase::AfterFirstDelta, + ); + let _ = tx.send(Err(failure)).await; } anthropic_shared::StreamOutcome::FailedEmpty(e) => { if e.is_retryable() { let mut backoff = std::time::Duration::from_secs(1); - let mut final_err = Some(e); + let mut final_err = Some((e, ProviderFailurePhase::BeforeFirstDelta)); for attempt in 1..=crate::retry::MAX_STREAM_RETRIES { backoff = crate::retry::backoff_sleep(attempt, backoff).await; match crate::retry::send_and_check(&client, &url_clone, &headers, &body) @@ -145,30 +177,51 @@ impl LlmProvider for AnthropicProvider { Ok(resp) => { let outcome = anthropic_shared::process_sse_stream(resp, &tx).await; + let phase = match &outcome { + anthropic_shared::StreamOutcome::FailedPartial(_) => { + ProviderFailurePhase::AfterFirstDelta + } + anthropic_shared::StreamOutcome::Ok + | anthropic_shared::StreamOutcome::FailedEmpty(_) => { + ProviderFailurePhase::BeforeFirstDelta + } + }; match crate::retry::evaluate_outcome(outcome, attempt) { Ok(None) => { final_err = None; break; } Ok(Some(e)) => { - final_err = Some(e); + final_err = Some((e, phase)); break; } Err(_) => continue, } } Err(e) if attempt == crate::retry::MAX_STREAM_RETRIES => { - final_err = Some(e); + final_err = Some((e, ProviderFailurePhase::BeforeFirstDelta)); break; } Err(_) => continue, } } - if let Some(err) = final_err { - let _ = tx.send(LlmEvent::Error(err.to_string())).await; + if let Some((err, phase)) = final_err { + let failure = provider_failure_from_error( + "anthropic", + stream_model.clone(), + err, + phase, + ); + let _ = tx.send(Err(failure)).await; } } else { - let _ = tx.send(LlmEvent::Error(e.to_string())).await; + let failure = provider_failure_from_error( + "anthropic", + stream_model.clone(), + e, + ProviderFailurePhase::BeforeFirstDelta, + ); + let _ = tx.send(Err(failure)).await; } } } diff --git a/crates/aion-providers/src/anthropic_shared.rs b/crates/aion-providers/src/anthropic_shared.rs index 45c579e7..7e1721c2 100644 --- a/crates/aion-providers/src/anthropic_shared.rs +++ b/crates/aion-providers/src/anthropic_shared.rs @@ -9,8 +9,8 @@ use aion_types::llm::LlmEvent; use aion_types::message::{ContentBlock, Message, Role, StopReason, TokenUsage}; use aion_types::tool::{ToolDef, truncate_deferred_description}; -use crate::error::ProviderError; use crate::tool_call_sanitize::{DroppedToolCallReason, format_dropped_tool_call}; +use crate::{ProviderError, ProviderStreamItem}; use aion_config::compat::ProviderCompat; /// Convert internal Message format to Anthropic API message format. @@ -322,7 +322,7 @@ pub enum StreamOutcome { /// Process the SSE stream from an Anthropic-compatible API pub async fn process_sse_stream( response: reqwest::Response, - tx: &mpsc::Sender, + tx: &mpsc::Sender, ) -> StreamOutcome { use futures::StreamExt; @@ -357,8 +357,18 @@ pub async fn process_sse_stream( current_event_type = event_type.to_string(); } else if let Some(data) = line.strip_prefix("data: ") { tracing::debug!(target: "aion_providers", chunk = %data, "sse chunk received"); - let events = parse_sse_data(¤t_event_type, data, &mut state); - for event in events { + let items = parse_sse_data(¤t_event_type, data, &mut state); + for item in items { + let event = match item { + Ok(event) => event, + Err(error) => { + return if emitted_content { + StreamOutcome::FailedPartial(error) + } else { + StreamOutcome::FailedEmpty(error) + }; + } + }; if matches!( event, LlmEvent::TextDelta(_) @@ -368,7 +378,7 @@ pub async fn process_sse_stream( ) { emitted_content = true; } - if tx.send(event).await.is_err() { + if tx.send(Ok(event)).await.is_err() { return StreamOutcome::Ok; // receiver dropped } } @@ -381,12 +391,20 @@ pub async fn process_sse_stream( } /// Parse a single SSE data payload into zero or more LlmEvents -pub fn parse_sse_data(event_type: &str, data: &str, state: &mut StreamState) -> Vec { +pub fn parse_sse_data( + event_type: &str, + data: &str, + state: &mut StreamState, +) -> Vec> { let mut events = Vec::new(); let json: Value = match serde_json::from_str(data) { Ok(v) => v, - Err(_) => return events, + Err(e) => { + return vec![Err(ProviderError::Parse(format!( + "Anthropic SSE JSON parse error: {e}" + )))]; + } }; match event_type { @@ -418,7 +436,7 @@ pub fn parse_sse_data(event_type: &str, data: &str, state: &mut StreamState) -> match delta_type { "text_delta" => { if let Some(text) = delta["text"].as_str() { - events.push(LlmEvent::TextDelta(text.to_string())); + events.push(Ok(LlmEvent::TextDelta(text.to_string()))); } } "input_json_delta" => { @@ -428,12 +446,12 @@ pub fn parse_sse_data(event_type: &str, data: &str, state: &mut StreamState) -> } "thinking_delta" => { if let Some(thinking) = delta["thinking"].as_str() { - events.push(LlmEvent::ThinkingDelta(thinking.to_string())); + events.push(Ok(LlmEvent::ThinkingDelta(thinking.to_string()))); } } "signature_delta" => { if let Some(signature) = delta["signature"].as_str() { - events.push(LlmEvent::ThinkingSignature(signature.to_string())); + events.push(Ok(LlmEvent::ThinkingSignature(signature.to_string()))); } } _ => {} @@ -451,12 +469,12 @@ pub fn parse_sse_data(event_type: &str, data: &str, state: &mut StreamState) -> "provider emitted tool_call with empty function name; recorded to history as-is" ); } - events.push(LlmEvent::ToolUse { + events.push(Ok(LlmEvent::ToolUse { id: state.tool_id.clone(), name: state.tool_name.clone(), input, extra: None, - }); + })); state.tool_input_json.clear(); } state.current_block_type = None; @@ -475,7 +493,7 @@ pub fn parse_sse_data(event_type: &str, data: &str, state: &mut StreamState) -> state.output_tokens = usage["output_tokens"].as_u64().unwrap_or(0); } - events.push(LlmEvent::Done { + events.push(Ok(LlmEvent::Done { stop_reason, usage: TokenUsage { input_tokens: state.input_tokens, @@ -483,7 +501,7 @@ pub fn parse_sse_data(event_type: &str, data: &str, state: &mut StreamState) -> cache_creation_tokens: state.cache_creation_tokens, cache_read_tokens: state.cache_read_tokens, }, - }); + })); } "message_stop" => { @@ -491,10 +509,15 @@ pub fn parse_sse_data(event_type: &str, data: &str, state: &mut StreamState) -> } "error" => { - let msg = json["error"]["message"] - .as_str() - .unwrap_or("Unknown API error"); - events.push(LlmEvent::Error(msg.to_string())); + let error = &json["error"]; + let msg = error["message"].as_str().unwrap_or("Unknown API error"); + events.push(Err(ProviderError::ApiSignal { + status: 400, + message: msg.to_string(), + raw_code: error["code"].as_str().map(str::to_string), + raw_type: error["type"].as_str().map(str::to_string), + retry_after_ms: None, + })); } _ => {} @@ -1270,7 +1293,7 @@ mod tests { let events = parse_sse_data("content_block_delta", data, &mut state); // assert assert_eq!(events.len(), 1); - match &events[0] { + match events[0].as_ref().unwrap() { LlmEvent::TextDelta(t) => assert_eq!(t, "Hello"), _ => panic!("expected TextDelta"), } @@ -1298,7 +1321,7 @@ mod tests { let events = parse_sse_data("content_block_stop", r#"{}"#, &mut state); // assert assert_eq!(events.len(), 1); - match &events[0] { + match events[0].as_ref().unwrap() { LlmEvent::ToolUse { id, name, input, .. } => { @@ -1325,7 +1348,7 @@ mod tests { let events = parse_sse_data("content_block_stop", r#"{}"#, &mut state); assert_eq!(events.len(), 1); - match &events[0] { + match events[0].as_ref().unwrap() { LlmEvent::ToolUse { id, name, .. } => { assert_eq!(id, "call_x"); assert_eq!(name, ""); @@ -1343,7 +1366,7 @@ mod tests { let events = parse_sse_data("message_delta", data, &mut state); // assert assert_eq!(events.len(), 1); - match &events[0] { + match events[0].as_ref().unwrap() { LlmEvent::Done { stop_reason, usage } => { assert_eq!(*stop_reason, StopReason::EndTurn); assert_eq!(usage.output_tokens, 42); @@ -1361,7 +1384,7 @@ mod tests { let events = parse_sse_data("content_block_delta", data, &mut state); // assert assert_eq!(events.len(), 1); - match &events[0] { + match events[0].as_ref().unwrap() { LlmEvent::ThinkingDelta(t) => assert_eq!(t, "reasoning step"), _ => panic!("expected ThinkingDelta"), } @@ -1375,7 +1398,7 @@ mod tests { let events = parse_sse_data("content_block_delta", data, &mut state); assert_eq!(events.len(), 1); - match &events[0] { + match events[0].as_ref().unwrap() { LlmEvent::ThinkingSignature(signature) => assert_eq!(signature, "sig-123"), _ => panic!("expected ThinkingSignature"), } @@ -1391,4 +1414,42 @@ mod tests { // assert assert!(events.is_empty()); } + + #[test] + fn test_parse_anthropic_malformed_json_returns_parse_error() { + let mut state = StreamState::new(); + + let events = parse_sse_data("content_block_delta", "{not-json}", &mut state); + + assert_eq!(events.len(), 1); + assert!(matches!(events[0], Err(ProviderError::Parse(_)))); + } + + #[test] + fn test_parse_anthropic_error_preserves_provider_error_type() { + let mut state = StreamState::new(); + let events = parse_sse_data( + "error", + r#"{"type":"error","error":{"type":"rate_limit_error","message":"rate limit exceeded"}}"#, + &mut state, + ); + + assert_eq!(events.len(), 1); + let error = events.into_iter().next().unwrap(); + let Err(error) = error else { + panic!("expected provider error") + }; + let failure = crate::provider_failure_from_error( + "anthropic", + Some("claude-3-5-sonnet".to_string()), + error, + crate::ProviderFailurePhase::BeforeFirstDelta, + ); + + assert_eq!(failure.kind, crate::ProviderFailureKind::RateLimited); + assert_eq!( + failure.meta.raw.raw_type.as_deref(), + Some("rate_limit_error") + ); + } } diff --git a/crates/aion-providers/src/bedrock.rs b/crates/aion-providers/src/bedrock.rs index 4cb02ea2..d0bef3df 100644 --- a/crates/aion-providers/src/bedrock.rs +++ b/crates/aion-providers/src/bedrock.rs @@ -18,7 +18,11 @@ use aion_types::llm::{LlmEvent, LlmRequest, ThinkingConfig}; use aion_types::message::{StopReason, TokenUsage}; use super::anthropic_shared; -use crate::{LlmProvider, ProviderError}; +use crate::{ + LlmProvider, ProviderError, ProviderFailure, ProviderFailurePhase, ProviderStreamItem, + ProviderStreamReceiver, provider_api_error, provider_failure_from_error, + retry_after_ms_from_headers, +}; use aion_config::compat::{self, ProviderCompat}; pub struct BedrockProvider { @@ -244,22 +248,46 @@ impl LlmProvider for BedrockProvider { async fn stream( &self, request: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result { let url = self.build_url(&request.model); let body = self.build_request_body(request); + let model = Some(request.model.clone()); tracing::debug!(target: "aion_providers", body = %serde_json::to_string_pretty(&body).unwrap_or_default(), "outgoing request"); let body_bytes = serde_json::to_vec(&body) - .map_err(|e| ProviderError::Connection(format!("JSON serialize error: {}", e)))?; + .map_err(|e| ProviderError::Connection(format!("JSON serialize error: {}", e))) + .map_err(|error| { + provider_failure_from_error( + "bedrock", + model.clone(), + error, + ProviderFailurePhase::BeforeFirstDelta, + ) + })?; - let credentials = self.resolve_credentials()?; + let credentials = self.resolve_credentials().map_err(|error| { + provider_failure_from_error( + "bedrock", + model.clone(), + error, + ProviderFailurePhase::BeforeFirstDelta, + ) + })?; let mut headers = HeaderMap::new(); headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); - let signed_headers = - self.sign_request("POST", &url, &headers, &body_bytes, &credentials)?; + let signed_headers = self + .sign_request("POST", &url, &headers, &body_bytes, &credentials) + .map_err(|error| { + provider_failure_from_error( + "bedrock", + model.clone(), + error, + ProviderFailurePhase::BeforeFirstDelta, + ) + })?; let response = self .client @@ -267,33 +295,55 @@ impl LlmProvider for BedrockProvider { .headers(signed_headers) .body(body_bytes) .send() - .await?; + .await + .map_err(ProviderError::from) + .map_err(|error| { + provider_failure_from_error( + "bedrock", + model.clone(), + error, + ProviderFailurePhase::BeforeFirstDelta, + ) + })?; let status = response.status(); if !status.is_success() { + let retry_after_ms = retry_after_ms_from_headers(response.headers()); let body_text = response.text().await.unwrap_or_default(); - if status.as_u16() == 429 { - return Err(ProviderError::RateLimited { - retry_after_ms: 5000, - }); - } let message = format_bedrock_error(status.as_u16(), &body_text); - return Err(ProviderError::Api { - status: status.as_u16(), - message, - }); + return Err(provider_failure_from_error( + "bedrock", + model, + provider_api_error(status.as_u16(), message, retry_after_ms), + ProviderFailurePhase::BeforeFirstDelta, + )); } let (tx, rx) = mpsc::channel(64); + let stream_model = model.clone(); // AWS event stream uses binary framing tokio::spawn(async move { match process_aws_event_stream(response, &tx).await { anthropic_shared::StreamOutcome::Ok => {} - anthropic_shared::StreamOutcome::FailedPartial(e) - | anthropic_shared::StreamOutcome::FailedEmpty(e) => { + anthropic_shared::StreamOutcome::FailedPartial(e) => { // Bedrock retry requires re-signing; not implemented yet. - let _ = tx.send(LlmEvent::Error(e.to_string())).await; + let failure = provider_failure_from_error( + "bedrock", + stream_model.clone(), + e, + ProviderFailurePhase::AfterFirstDelta, + ); + let _ = tx.send(Err(failure)).await; + } + anthropic_shared::StreamOutcome::FailedEmpty(e) => { + let failure = provider_failure_from_error( + "bedrock", + stream_model.clone(), + e, + ProviderFailurePhase::BeforeFirstDelta, + ); + let _ = tx.send(Err(failure)).await; } } }); @@ -305,7 +355,7 @@ impl LlmProvider for BedrockProvider { /// Process the AWS event stream (binary framed) from Bedrock async fn process_aws_event_stream( response: reqwest::Response, - tx: &mpsc::Sender, + tx: &mpsc::Sender, ) -> anthropic_shared::StreamOutcome { use futures::StreamExt; @@ -334,33 +384,74 @@ async fn process_aws_event_stream( if let Some(payload) = event_data { // The payload contains an SSE-like structure with "bytes" field - if let Ok(wrapper) = serde_json::from_slice::(&payload) { - // Bedrock wraps the payload in {"bytes": "base64-encoded-data"} - if let Some(b64) = wrapper["bytes"].as_str() - && let Ok(decoded) = base64::engine::general_purpose::STANDARD.decode(b64) - && let Ok(inner) = String::from_utf8(decoded) - { - tracing::debug!(target: "aion_providers", chunk = %inner, "bedrock event chunk"); - // Inner payload is JSON with event type hints - if let Ok(json_val) = serde_json::from_str::(&inner) { - let event_type = json_val["type"].as_str().unwrap_or(""); - let events = - anthropic_shared::parse_sse_data(event_type, &inner, &mut state); - for event in events { - if matches!( - event, - LlmEvent::TextDelta(_) - | LlmEvent::ThinkingDelta(_) - | LlmEvent::ThinkingSignature(_) - | LlmEvent::ToolUse { .. } - ) { - emitted_content = true; - } - if tx.send(event).await.is_err() { - return anthropic_shared::StreamOutcome::Ok; - } - } + let wrapper = match serde_json::from_slice::(&payload) { + Ok(wrapper) => wrapper, + Err(e) => { + return bedrock_parse_failure( + format!("Bedrock event wrapper JSON parse error: {e}"), + emitted_content, + ); + } + }; + + // Bedrock wraps the payload in {"bytes": "base64-encoded-data"}. + let Some(b64) = wrapper["bytes"].as_str() else { + continue; + }; + let decoded = match base64::engine::general_purpose::STANDARD.decode(b64) { + Ok(decoded) => decoded, + Err(e) => { + return bedrock_parse_failure( + format!("Bedrock event payload base64 decode error: {e}"), + emitted_content, + ); + } + }; + let inner = match String::from_utf8(decoded) { + Ok(inner) => inner, + Err(e) => { + return bedrock_parse_failure( + format!("Bedrock event payload UTF-8 decode error: {e}"), + emitted_content, + ); + } + }; + tracing::debug!(target: "aion_providers", chunk = %inner, "bedrock event chunk"); + + // Inner payload is JSON with event type hints. + let json_val = match serde_json::from_str::(&inner) { + Ok(json_val) => json_val, + Err(e) => { + return bedrock_parse_failure( + format!("Bedrock inner JSON parse error: {e}"), + emitted_content, + ); + } + }; + let event_type = json_val["type"].as_str().unwrap_or(""); + let items = anthropic_shared::parse_sse_data(event_type, &inner, &mut state); + for item in items { + let event = match item { + Ok(event) => event, + Err(error) => { + return if emitted_content { + anthropic_shared::StreamOutcome::FailedPartial(error) + } else { + anthropic_shared::StreamOutcome::FailedEmpty(error) + }; } + }; + if matches!( + event, + LlmEvent::TextDelta(_) + | LlmEvent::ThinkingDelta(_) + | LlmEvent::ThinkingSignature(_) + | LlmEvent::ToolUse { .. } + ) { + emitted_content = true; + } + if tx.send(Ok(event)).await.is_err() { + return anthropic_shared::StreamOutcome::Ok; } } } @@ -370,7 +461,7 @@ async fn process_aws_event_stream( // If we haven't sent a Done event, send one now if state.input_tokens > 0 || state.output_tokens > 0 { let _ = tx - .send(LlmEvent::Done { + .send(Ok(LlmEvent::Done { stop_reason: StopReason::EndTurn, usage: TokenUsage { input_tokens: state.input_tokens, @@ -378,13 +469,25 @@ async fn process_aws_event_stream( cache_creation_tokens: state.cache_creation_tokens, cache_read_tokens: state.cache_read_tokens, }, - }) + })) .await; } anthropic_shared::StreamOutcome::Ok } +fn bedrock_parse_failure( + message: String, + emitted_content: bool, +) -> anthropic_shared::StreamOutcome { + let error = ProviderError::Parse(message); + if emitted_content { + anthropic_shared::StreamOutcome::FailedPartial(error) + } else { + anthropic_shared::StreamOutcome::FailedEmpty(error) + } +} + /// Parse one AWS event stream message from the buffer. /// Returns (Some(payload), bytes_consumed) if a complete message is found, /// or None if more data is needed. @@ -477,3 +580,74 @@ pub fn credentials_from_config(bc: &aion_config::config::BedrockConfig) -> AwsCr AwsCredentials::Environment } } + +#[cfg(test)] +mod tests { + use super::*; + + use base64::Engine; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + + fn aws_event_frame(payload: &[u8]) -> Vec { + let total_len = 12 + payload.len() + 4; + let mut frame = Vec::with_capacity(total_len); + frame.extend_from_slice(&(total_len as u32).to_be_bytes()); + frame.extend_from_slice(&0_u32.to_be_bytes()); + frame.extend_from_slice(&0_u32.to_be_bytes()); + frame.extend_from_slice(payload); + frame.extend_from_slice(&0_u32.to_be_bytes()); + frame + } + + async fn response_for_body(body: Vec) -> reqwest::Response { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut buf = [0_u8; 4096]; + let _ = stream.read(&mut buf).await.unwrap(); + let headers = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/vnd.amazon.eventstream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n", + body.len() + ); + stream.write_all(headers.as_bytes()).await.unwrap(); + stream.write_all(&body).await.unwrap(); + }); + + reqwest::Client::new() + .get(format!("http://{addr}")) + .send() + .await + .unwrap() + } + + #[tokio::test] + async fn bedrock_invalid_event_wrapper_returns_parse_failure() { + let response = response_for_body(aws_event_frame(b"{not-json}")).await; + let (tx, _rx) = mpsc::channel(1); + + let outcome = process_aws_event_stream(response, &tx).await; + + assert!(matches!( + outcome, + anthropic_shared::StreamOutcome::FailedEmpty(ProviderError::Parse(_)) + )); + } + + #[tokio::test] + async fn bedrock_invalid_inner_json_returns_parse_failure() { + let inner = base64::engine::general_purpose::STANDARD.encode("{not-json}"); + let wrapper = serde_json::json!({ "bytes": inner }).to_string(); + let response = response_for_body(aws_event_frame(wrapper.as_bytes())).await; + let (tx, _rx) = mpsc::channel(1); + + let outcome = process_aws_event_stream(response, &tx).await; + + assert!(matches!( + outcome, + anthropic_shared::StreamOutcome::FailedEmpty(ProviderError::Parse(_)) + )); + } +} diff --git a/crates/aion-providers/src/error.rs b/crates/aion-providers/src/error.rs index 5e299f69..0c4b741e 100644 --- a/crates/aion-providers/src/error.rs +++ b/crates/aion-providers/src/error.rs @@ -1,9 +1,19 @@ +use reqwest::header::{HeaderMap, RETRY_AFTER}; + #[derive(Debug, thiserror::Error)] pub enum ProviderError { #[error("HTTP error: {0}")] Http(#[from] reqwest::Error), #[error("API error {status}: {message}")] Api { status: u16, message: String }, + #[error("API error {status}: {message}")] + ApiSignal { + status: u16, + message: String, + raw_code: Option, + raw_type: Option, + retry_after_ms: Option, + }, #[error("SSE parse error: {0}")] Parse(String), #[error("Rate limited, retry after {retry_after_ms}ms")] @@ -23,6 +33,563 @@ impl ProviderError { } } +pub fn retry_after_ms_from_headers(headers: &HeaderMap) -> Option { + headers + .get(RETRY_AFTER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + .map(|seconds| seconds.saturating_mul(1000)) +} + +pub fn provider_api_error( + status: u16, + message: impl Into, + retry_after_ms: Option, +) -> ProviderError { + ProviderError::ApiSignal { + status, + message: message.into(), + raw_code: None, + raw_type: None, + retry_after_ms: retry_after_ms.or_else(|| (status == 429).then_some(5000)), + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderRawSignals { + pub provider: String, + pub model: Option, + pub status: Option, + pub raw_code: Option, + pub raw_type: Option, + pub message: Option, + pub request_id: Option, + pub retry_after_ms: Option, + pub source: ProviderRawSignalSource, + pub phase: ProviderFailurePhase, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ProviderRawSignalSource { + Auth, + RequestBuild, + Http, + Sse, + Json, + Transport, + Unknown, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ProviderFailurePhase { + BeforeFirstDelta, + AfterFirstDelta, + AfterToolCall, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderFailure { + pub kind: ProviderFailureKind, + pub meta: ProviderFailureMeta, +} + +impl ProviderFailure { + pub fn tool_call_invalid( + provider: impl Into, + model: Option, + reason: impl Into, + phase: ProviderFailurePhase, + ) -> ProviderFailure { + let reason = summarize_raw_value(reason.into()); + ProviderFailure { + kind: ProviderFailureKind::ToolCall(ToolCallFailure { + reason: ToolCallFailureReason::InvalidProviderToolCall, + detail: Some(reason.clone()), + }), + meta: ProviderFailureMeta { + provider: provider.into(), + model, + status: None, + request_id: None, + retry_after_ms: None, + phase, + raw: RawProviderErrorSummary { + raw_code: None, + raw_type: None, + message: Some(summarize_raw_value(reason)), + }, + }, + } + } +} + +impl std::fmt::Display for ProviderFailure { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self.meta.raw.message.as_deref() { + Some(message) => write!( + f, + "{:?} from {} during {:?}: {}", + self.kind, self.meta.provider, self.meta.phase, message + ), + None => write!( + f, + "{:?} from {} during {:?}", + self.kind, self.meta.provider, self.meta.phase + ), + } + } +} + +impl std::error::Error for ProviderFailure {} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderFailureMeta { + pub provider: String, + pub model: Option, + pub status: Option, + pub request_id: Option, + pub retry_after_ms: Option, + pub phase: ProviderFailurePhase, + pub raw: RawProviderErrorSummary, +} + +impl ProviderFailureMeta { + fn from_raw(raw: ProviderRawSignals) -> ProviderFailureMeta { + ProviderFailureMeta { + provider: raw.provider, + model: raw.model, + status: raw.status, + request_id: raw.request_id, + retry_after_ms: raw.retry_after_ms, + phase: raw.phase, + raw: RawProviderErrorSummary { + raw_code: raw.raw_code.map(summarize_raw_value), + raw_type: raw.raw_type.map(summarize_raw_value), + message: raw.message.map(summarize_raw_value), + }, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RawProviderErrorSummary { + pub raw_code: Option, + pub raw_type: Option, + pub message: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ProviderFailureKind { + CredentialMissing, + AuthFailed, + PermissionDenied, + BillingRequired, + QuotaExceeded, + RateLimited, + ModelNotFound, + ModelUnsupported, + ModelUnavailable, + EndpointNotFound, + ContextTooLarge, + ToolSchema(ToolSchemaFailure), + ToolCall(ToolCallFailure), + ContentBlocked, + Timeout, + Transport, + Server, + InvalidRequest(InvalidProviderRequestFailure), + ResponseParse, + EmptyResponse, + Unknown, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ToolSchemaFailure { + pub source: ToolSchemaSource, + pub reason: ToolSchemaFailureReason, + pub detail: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ToolSchemaSource { + AionrsRequest, + ProviderResponse, + McpCapability, + Unknown, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ToolSchemaFailureReason { + InvalidSchema, + UnsupportedSchema, + BadRequest, + Unknown, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ToolCallFailure { + pub reason: ToolCallFailureReason, + pub detail: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ToolCallFailureReason { + InvalidProviderToolCall, + MissingName, + MissingId, + InvalidArguments, + Unknown, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct InvalidProviderRequestFailure { + pub source: InvalidProviderRequestSource, + pub reason: InvalidProviderRequestReason, + pub detail: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum InvalidProviderRequestSource { + AionrsRequestBuilder, + Unknown, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum InvalidProviderRequestReason { + BadRequest, + Unknown, +} + +pub fn classify_provider_failure(raw: ProviderRawSignals) -> ProviderFailure { + let kind = classify_provider_failure_kind(&raw); + let meta = ProviderFailureMeta::from_raw(raw); + ProviderFailure { kind, meta } +} + +pub fn provider_failure_from_error( + provider: impl Into, + model: Option, + error: ProviderError, + phase: ProviderFailurePhase, +) -> ProviderFailure { + let provider = provider.into(); + let raw = match error { + ProviderError::Http(error) => { + let status = error.status().map(|status| status.as_u16()); + ProviderRawSignals { + provider, + model, + status, + raw_code: None, + raw_type: None, + message: Some(error.to_string()), + request_id: None, + retry_after_ms: None, + source: if status.is_none() || error.is_connect() || error.is_timeout() { + ProviderRawSignalSource::Transport + } else { + ProviderRawSignalSource::Http + }, + phase, + } + } + ProviderError::Api { status, message } => ProviderRawSignals { + provider, + model, + status: Some(status), + raw_code: None, + raw_type: None, + message: Some(message), + request_id: None, + retry_after_ms: None, + source: ProviderRawSignalSource::Http, + phase, + }, + ProviderError::ApiSignal { + status, + message, + raw_code, + raw_type, + retry_after_ms, + } => ProviderRawSignals { + provider, + model, + status: Some(status), + raw_code, + raw_type, + message: Some(message), + request_id: None, + retry_after_ms, + source: ProviderRawSignalSource::Http, + phase, + }, + ProviderError::Parse(message) => ProviderRawSignals { + provider, + model, + status: None, + raw_code: None, + raw_type: None, + message: Some(message), + request_id: None, + retry_after_ms: None, + source: ProviderRawSignalSource::Sse, + phase, + }, + ProviderError::RateLimited { retry_after_ms } => ProviderRawSignals { + provider, + model, + status: Some(429), + raw_code: None, + raw_type: None, + message: None, + request_id: None, + retry_after_ms: Some(retry_after_ms), + source: ProviderRawSignalSource::Http, + phase, + }, + ProviderError::PromptTooLong(message) => ProviderRawSignals { + provider, + model, + status: Some(400), + raw_code: Some("prompt_too_long".to_string()), + raw_type: None, + message: Some(message), + request_id: None, + retry_after_ms: None, + source: ProviderRawSignalSource::Http, + phase, + }, + ProviderError::Connection(message) => ProviderRawSignals { + provider, + model, + status: None, + raw_code: None, + raw_type: None, + message: Some(message), + request_id: None, + retry_after_ms: None, + source: ProviderRawSignalSource::Transport, + phase, + }, + }; + + classify_provider_failure(raw) +} + +fn classify_provider_failure_kind(raw: &ProviderRawSignals) -> ProviderFailureKind { + let status = raw.status; + let text = signal_text(raw); + + if status == Some(401) || has_any(&text, auth_invalid_signals()) { + return ProviderFailureKind::AuthFailed; + } + + if raw.source == ProviderRawSignalSource::Auth { + return ProviderFailureKind::CredentialMissing; + } + + if has_any(&text, &["permission_error", "permission denied"]) { + return ProviderFailureKind::PermissionDenied; + } + + if has_any( + &text, + &[ + "rate_limit_error", + "rate_limit", + "rate limit", + "too many requests", + ], + ) { + return ProviderFailureKind::RateLimited; + } + + if has_any(&text, &["overloaded_error", "overloaded"]) { + return ProviderFailureKind::ModelUnavailable; + } + + if has_any(&text, &["api_error", "server_error"]) { + return ProviderFailureKind::Server; + } + + if has_any( + &text, + &[ + "context length", + "context_length", + "maximum context", + "maximum_context", + "too many tokens", + "too_many_tokens", + "prompt too long", + "prompt_too_long", + ], + ) { + return ProviderFailureKind::ContextTooLarge; + } + + if has_any( + &text, + &["insufficient_quota", "quota exceeded", "quota_exceeded"], + ) { + return ProviderFailureKind::QuotaExceeded; + } + + if has_any( + &text, + &["content_filter", "content blocked", "content_blocked"], + ) { + return ProviderFailureKind::ContentBlocked; + } + + if has_any(&text, &["timeout", "timed out"]) { + return ProviderFailureKind::Timeout; + } + + match status { + Some(403) => return ProviderFailureKind::PermissionDenied, + Some(402) => return ProviderFailureKind::BillingRequired, + Some(429) => return ProviderFailureKind::RateLimited, + Some(500..=599) => return ProviderFailureKind::Server, + Some(404) if has_model_not_found_signal(&text) => { + return ProviderFailureKind::ModelNotFound; + } + Some(404) => return ProviderFailureKind::EndpointNotFound, + _ => {} + } + + if matches!( + raw.source, + ProviderRawSignalSource::Sse | ProviderRawSignalSource::Json + ) { + return ProviderFailureKind::ResponseParse; + } + + if raw.source == ProviderRawSignalSource::RequestBuild { + return ProviderFailureKind::InvalidRequest(InvalidProviderRequestFailure { + source: InvalidProviderRequestSource::AionrsRequestBuilder, + reason: InvalidProviderRequestReason::BadRequest, + detail: raw.message.clone().map(summarize_raw_value), + }); + } + + if status == Some(400) { + return ProviderFailureKind::InvalidRequest(InvalidProviderRequestFailure { + source: InvalidProviderRequestSource::Unknown, + reason: InvalidProviderRequestReason::BadRequest, + detail: raw.message.clone().map(summarize_raw_value), + }); + } + + if raw.source == ProviderRawSignalSource::Transport { + return ProviderFailureKind::Transport; + } + + ProviderFailureKind::Unknown +} + +fn signal_text(raw: &ProviderRawSignals) -> String { + let mut parts = Vec::new(); + if let Some(raw_code) = raw.raw_code.as_deref() { + parts.push(raw_code); + } + if let Some(raw_type) = raw.raw_type.as_deref() { + parts.push(raw_type); + } + if let Some(message) = raw.message.as_deref() { + parts.push(message); + } + parts.join(" ").to_ascii_lowercase() +} + +fn has_any(text: &str, needles: &[&str]) -> bool { + needles.iter().any(|needle| text.contains(needle)) +} + +fn has_model_not_found_signal(text: &str) -> bool { + has_any( + text, + &[ + "model_not_found", + "model-not-found", + "model not found", + "model", + "模型不存在", + "模型未找到", + ], + ) +} + +fn auth_invalid_signals() -> &'static [&'static str] { + &[ + "invalid_api_key", + "invalid api key", + "incorrect api key", + "authentication failed", + "authentication_error", + "auth invalid", + "invalid auth", + "invalid token", + "unauthorized", + ] +} + +fn summarize_raw_value(value: String) -> String { + const MAX_RAW_SUMMARY_CHARS: usize = 512; + + let value = redact_known_secret_shapes(value); + if value.chars().count() <= MAX_RAW_SUMMARY_CHARS { + return value; + } + + value + .chars() + .take(MAX_RAW_SUMMARY_CHARS) + .collect::() +} + +fn redact_known_secret_shapes(value: String) -> String { + value + .split_whitespace() + .map(|part| { + if looks_like_provider_secret(part) { + "[redacted]" + } else { + part + } + }) + .collect::>() + .join(" ") +} + +fn looks_like_provider_secret(value: &str) -> bool { + let trimmed = value.trim_matches(|c: char| !c.is_ascii_alphanumeric() && c != '-' && c != '_'); + let lowered = trimmed.to_ascii_lowercase(); + + (lowered.starts_with("sk-") && trimmed.len() >= 20) + || (lowered.starts_with("ya29.") && trimmed.len() >= 20) + || looks_like_aws_access_key_id(trimmed) +} + +fn looks_like_aws_access_key_id(value: &str) -> bool { + value.len() == 20 + && (value.starts_with("AKIA") || value.starts_with("ASIA")) + && value + .chars() + .all(|c| c.is_ascii_uppercase() || c.is_ascii_digit()) +} + +impl From for ProviderFailure { + fn from(error: ProviderError) -> ProviderFailure { + provider_failure_from_error( + "unknown", + None, + error, + ProviderFailurePhase::BeforeFirstDelta, + ) + } +} + #[cfg(test)] mod retryable_tests { use super::*; diff --git a/crates/aion-providers/src/lib.rs b/crates/aion-providers/src/lib.rs index 8c4481c6..9937050f 100644 --- a/crates/aion-providers/src/lib.rs +++ b/crates/aion-providers/src/lib.rs @@ -8,5 +8,12 @@ pub mod retry; mod tool_call_sanitize; pub mod vertex; -pub use error::ProviderError; -pub use provider::{LlmProvider, create_provider}; +pub use error::{ + InvalidProviderRequestFailure, InvalidProviderRequestReason, InvalidProviderRequestSource, + ProviderError, ProviderFailure, ProviderFailureKind, ProviderFailureMeta, ProviderFailurePhase, + ProviderRawSignalSource, ProviderRawSignals, RawProviderErrorSummary, ToolCallFailure, + ToolCallFailureReason, ToolSchemaFailure, ToolSchemaFailureReason, ToolSchemaSource, + classify_provider_failure, provider_api_error, provider_failure_from_error, + retry_after_ms_from_headers, +}; +pub use provider::{LlmProvider, ProviderStreamItem, ProviderStreamReceiver, create_provider}; diff --git a/crates/aion-providers/src/openai.rs b/crates/aion-providers/src/openai.rs index 904f66fc..6c1676e8 100644 --- a/crates/aion-providers/src/openai.rs +++ b/crates/aion-providers/src/openai.rs @@ -10,7 +10,11 @@ use aion_types::tool::{ToolDef, truncate_deferred_description}; use crate::anthropic_shared::StreamOutcome; use crate::tool_call_sanitize::{DroppedToolCallReason, format_dropped_tool_call}; -use crate::{LlmProvider, ProviderError}; +use crate::{ + LlmProvider, ProviderError, ProviderFailure, ProviderFailurePhase, ProviderStreamItem, + ProviderStreamReceiver, provider_api_error, provider_failure_from_error, + retry_after_ms_from_headers, +}; use aion_config::compat::ProviderCompat; pub struct OpenAIProvider { @@ -604,10 +608,18 @@ impl LlmProvider for OpenAIProvider { async fn stream( &self, request: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result { let url = format!("{}{}", self.base_url, self.compat.api_path()); let body = self.build_request_body(request); - let headers = self.build_headers()?; + let model = Some(request.model.clone()); + let headers = self.build_headers().map_err(|error| { + provider_failure_from_error( + "openai", + model.clone(), + error, + ProviderFailurePhase::BeforeFirstDelta, + ) + })?; tracing::debug!(target: "aion_providers", body = %serde_json::to_string_pretty(&body).unwrap_or_default(), "outgoing request"); @@ -622,37 +634,49 @@ impl LlmProvider for OpenAIProvider { let status = response.status(); if !status.is_success() { + let retry_after_ms = retry_after_ms_from_headers(response.headers()); let body_text = response.text().await.unwrap_or_default(); - if status.as_u16() == 429 { - return Err(ProviderError::RateLimited { - retry_after_ms: 5000, - }); - } - return Err(ProviderError::Api { - status: status.as_u16(), - message: body_text, - }); + return Err(provider_api_error( + status.as_u16(), + body_text, + retry_after_ms, + )); } Ok(response) }) - .await?; + .await + .map_err(|error| { + provider_failure_from_error( + "openai", + model.clone(), + error, + ProviderFailurePhase::BeforeFirstDelta, + ) + })?; let (tx, rx) = mpsc::channel(64); let auto_tool_id = self.compat.auto_tool_id(); let client = self.client.clone(); let url_clone = url.clone(); + let stream_model = model.clone(); tokio::spawn(async move { match process_sse_stream(response, &tx, auto_tool_id).await { StreamOutcome::Ok => {} StreamOutcome::FailedPartial(e) => { - let _ = tx.send(LlmEvent::Error(e.to_string())).await; + let failure = provider_failure_from_error( + "openai", + stream_model.clone(), + e, + ProviderFailurePhase::AfterFirstDelta, + ); + let _ = tx.send(Err(failure)).await; } StreamOutcome::FailedEmpty(e) => { if e.is_retryable() { let mut backoff = std::time::Duration::from_secs(1); - let mut final_err = Some(e); + let mut final_err = Some((e, ProviderFailurePhase::BeforeFirstDelta)); for attempt in 1..=crate::retry::MAX_STREAM_RETRIES { backoff = crate::retry::backoff_sleep(attempt, backoff).await; match crate::retry::send_and_check(&client, &url_clone, &headers, &body) @@ -660,30 +684,43 @@ impl LlmProvider for OpenAIProvider { { Ok(resp) => { let outcome = process_sse_stream(resp, &tx, auto_tool_id).await; + let phase = stream_outcome_phase(&outcome); match crate::retry::evaluate_outcome(outcome, attempt) { Ok(None) => { final_err = None; break; } Ok(Some(e)) => { - final_err = Some(e); + final_err = Some((e, phase)); break; } Err(_) => continue, } } Err(e) if attempt == crate::retry::MAX_STREAM_RETRIES => { - final_err = Some(e); + final_err = Some((e, ProviderFailurePhase::BeforeFirstDelta)); break; } Err(_) => continue, } } - if let Some(err) = final_err { - let _ = tx.send(LlmEvent::Error(err.to_string())).await; + if let Some((err, phase)) = final_err { + let failure = provider_failure_from_error( + "openai", + stream_model.clone(), + err, + phase, + ); + let _ = tx.send(Err(failure)).await; } } else { - let _ = tx.send(LlmEvent::Error(e.to_string())).await; + let failure = provider_failure_from_error( + "openai", + stream_model.clone(), + e, + ProviderFailurePhase::BeforeFirstDelta, + ); + let _ = tx.send(Err(failure)).await; } } } @@ -695,7 +732,7 @@ impl LlmProvider for OpenAIProvider { async fn process_sse_stream( response: reqwest::Response, - tx: &mpsc::Sender, + tx: &mpsc::Sender, auto_tool_id: bool, ) -> StreamOutcome { use futures::StreamExt; @@ -735,12 +772,21 @@ async fn process_sse_stream( // Flush the deferred Done event now that the final // usage-only chunk (choices:[]) has updated token counts. if let Some(done) = state.flush_done() { - let _ = tx.send(done).await; + let _ = tx.send(Ok(done)).await; } return StreamOutcome::Ok; } - let events = parse_sse_chunk(data, &mut state, auto_tool_id); + let events = match parse_sse_chunk(data, &mut state, auto_tool_id) { + Ok(events) => events, + Err(err) => { + return if emitted_content { + StreamOutcome::FailedPartial(err) + } else { + StreamOutcome::FailedEmpty(err) + }; + } + }; for event in events { if matches!( event, @@ -750,7 +796,7 @@ async fn process_sse_stream( ) { emitted_content = true; } - if tx.send(event).await.is_err() { + if tx.send(Ok(event)).await.is_err() { return StreamOutcome::Ok; } } @@ -761,13 +807,22 @@ async fn process_sse_stream( StreamOutcome::Ok } -fn parse_sse_chunk(data: &str, state: &mut StreamState, auto_tool_id: bool) -> Vec { +fn stream_outcome_phase(outcome: &StreamOutcome) -> ProviderFailurePhase { + match outcome { + StreamOutcome::FailedPartial(_) => ProviderFailurePhase::AfterFirstDelta, + StreamOutcome::Ok | StreamOutcome::FailedEmpty(_) => ProviderFailurePhase::BeforeFirstDelta, + } +} + +fn parse_sse_chunk( + data: &str, + state: &mut StreamState, + auto_tool_id: bool, +) -> Result, ProviderError> { let mut events = Vec::new(); - let json: Value = match serde_json::from_str(data) { - Ok(v) => v, - Err(_) => return events, - }; + let json: Value = serde_json::from_str(data) + .map_err(|e| ProviderError::Parse(format!("OpenAI SSE JSON parse error: {e}")))?; // Extract usage if present if let Some(usage) = json.get("usage") { @@ -787,7 +842,7 @@ fn parse_sse_chunk(data: &str, state: &mut StreamState, auto_tool_id: bool) -> V } let Some(choice) = json["choices"].as_array().and_then(|c| c.first()) else { - return events; + return Ok(events); }; let delta = &choice["delta"]; @@ -887,7 +942,7 @@ fn parse_sse_chunk(data: &str, state: &mut StreamState, auto_tool_id: bool) -> V } } - events + Ok(events) } #[cfg(test)] @@ -1640,7 +1695,7 @@ mod tests { // chunk 1: finish_reason + text delta, no usage let chunk1 = r#"{"choices":[{"delta":{"content":"hi"},"finish_reason":"stop"}]}"#; - let events = parse_sse_chunk(chunk1, &mut state, false); + let events = parse_sse_chunk(chunk1, &mut state, false).unwrap(); // TextDelta is emitted immediately; Done is deferred. assert!( events.iter().all(|e| !matches!(e, LlmEvent::Done { .. })), @@ -1650,7 +1705,7 @@ mod tests { // chunk 2: trailing usage-only chunk (choices:[]) let chunk2 = r#"{"choices":[],"usage":{"prompt_tokens":10,"completion_tokens":5}}"#; - let events2 = parse_sse_chunk(chunk2, &mut state, false); + let events2 = parse_sse_chunk(chunk2, &mut state, false).unwrap(); assert!(events2.is_empty()); assert_eq!(state.input_tokens, 10); assert_eq!(state.output_tokens, 5); @@ -1675,7 +1730,7 @@ mod tests { // No text delta here, only finish_reason + usage in the same chunk. let chunk = r#"{"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":8,"completion_tokens":3}}"#; - let events = parse_sse_chunk(chunk, &mut state, false); + let events = parse_sse_chunk(chunk, &mut state, false).unwrap(); assert!( events.iter().all(|e| !matches!(e, LlmEvent::Done { .. })), "Done should be deferred even when usage is in the finish chunk" @@ -1727,7 +1782,7 @@ mod tests { let mut state = StreamState::new(); let chunk = r#"{"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":500,"completion_tokens":100,"prompt_cache_hit_tokens":999500}}"#; - let _ = parse_sse_chunk(chunk, &mut state, false); + let _ = parse_sse_chunk(chunk, &mut state, false).unwrap(); assert_eq!(state.input_tokens, 1_000_000); assert_eq!(state.output_tokens, 100); @@ -1740,7 +1795,7 @@ mod tests { let mut state = StreamState::new(); let chunk = r#"{"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":1000000,"completion_tokens":100,"prompt_tokens_details":{"cached_tokens":999000}}}"#; - let _ = parse_sse_chunk(chunk, &mut state, false); + let _ = parse_sse_chunk(chunk, &mut state, false).unwrap(); // prompt_tokens is already the full total for OpenAI assert_eq!(state.input_tokens, 1_000_000); @@ -1753,7 +1808,7 @@ mod tests { let mut state = StreamState::new(); let chunk = r#"{"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":50000,"completion_tokens":200}}"#; - let _ = parse_sse_chunk(chunk, &mut state, false); + let _ = parse_sse_chunk(chunk, &mut state, false).unwrap(); assert_eq!(state.input_tokens, 50_000); assert_eq!(state.output_tokens, 200); @@ -1767,14 +1822,14 @@ mod tests { // chunk 1: tool call delta (name + partial args) let chunk1 = r#"{"choices":[{"delta":{"role":"assistant","tool_calls":[{"extra_content":{},"function":{"arguments":"{\"skill\":\"test\",\"args\":\"hello\"}","name":"Skill"},"id":"call_abc123","type":"function"}]},"index":0}]}"#; - let events1 = parse_sse_chunk(chunk1, &mut state, false); + let events1 = parse_sse_chunk(chunk1, &mut state, false).unwrap(); assert!(events1.is_empty(), "no events until finish_reason"); assert_eq!(state.tool_calls.len(), 1); assert_eq!(state.tool_calls[0].name, "Skill"); // chunk 2: finish_reason:"stop" (not "tool_calls") let chunk2 = r#"{"choices":[{"delta":{"role":"assistant"},"finish_reason":"stop","index":0}],"usage":{"prompt_tokens":100,"completion_tokens":20,"total_tokens":120}}"#; - let events2 = parse_sse_chunk(chunk2, &mut state, false); + let events2 = parse_sse_chunk(chunk2, &mut state, false).unwrap(); // Tool call should be emitted let tool_events: Vec<_> = events2 @@ -1809,11 +1864,11 @@ mod tests { let mut state = StreamState::new(); let chunk1 = r#"{"choices":[{"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call_x","type":"function","function":{"name":"","arguments":"{}"}}]},"index":0}]}"#; - let events1 = parse_sse_chunk(chunk1, &mut state, false); + let events1 = parse_sse_chunk(chunk1, &mut state, false).unwrap(); assert!(events1.is_empty(), "no events until finish_reason"); let chunk2 = r#"{"choices":[{"delta":{},"finish_reason":"tool_calls","index":0}]}"#; - let events2 = parse_sse_chunk(chunk2, &mut state, false); + let events2 = parse_sse_chunk(chunk2, &mut state, false).unwrap(); let tool_use_name = events2.iter().find_map(|event| match event { LlmEvent::ToolUse { name, .. } => Some(name.clone()), @@ -1834,7 +1889,7 @@ mod tests { let chunk = r#"{"choices":[{"delta":{"content":"done"},"finish_reason":"stop","index":0}]}"#; - let events = parse_sse_chunk(chunk, &mut state, false); + let events = parse_sse_chunk(chunk, &mut state, false).unwrap(); let text_events: Vec<_> = events .iter() @@ -1857,7 +1912,7 @@ mod tests { // Simulate a provider that returns tool_calls without an id field let chunk = r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"name":"get_weather","arguments":"{\"city\":\"Beijing\"}"}}]},"finish_reason":"tool_calls","index":0}]}"#; - let events = parse_sse_chunk(chunk, &mut state, true); + let events = parse_sse_chunk(chunk, &mut state, true).unwrap(); let tool_use = events .iter() @@ -1876,7 +1931,7 @@ mod tests { let mut state = StreamState::new(); let chunk = r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_existing_123","function":{"name":"read_file","arguments":"{}"}}]},"finish_reason":"tool_calls","index":0}]}"#; - let events = parse_sse_chunk(chunk, &mut state, true); + let events = parse_sse_chunk(chunk, &mut state, true).unwrap(); let tool_use = events .iter() @@ -1893,7 +1948,7 @@ mod tests { let mut state = StreamState::new(); let chunk = r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"name":"get_weather","arguments":"{}"}}]},"finish_reason":"tool_calls","index":0}]}"#; - let events = parse_sse_chunk(chunk, &mut state, false); + let events = parse_sse_chunk(chunk, &mut state, false).unwrap(); let tool_use = events .iter() diff --git a/crates/aion-providers/src/provider.rs b/crates/aion-providers/src/provider.rs index ce3a32b3..53bbdc23 100644 --- a/crates/aion-providers/src/provider.rs +++ b/crates/aion-providers/src/provider.rs @@ -8,15 +8,18 @@ use aion_types::llm::{LlmEvent, LlmRequest}; use crate::anthropic; use crate::bedrock; -use crate::error::ProviderError; +use crate::error::ProviderFailure; use crate::openai; use crate::vertex; /// Unified interface for LLM API providers +pub type ProviderStreamItem = Result; +pub type ProviderStreamReceiver = mpsc::Receiver; + #[async_trait] pub trait LlmProvider: Send + Sync { async fn stream(&self, request: &LlmRequest) - -> Result, ProviderError>; + -> Result; } /// Create a provider from resolved config diff --git a/crates/aion-providers/src/retry.rs b/crates/aion-providers/src/retry.rs index a8776a75..8e004179 100644 --- a/crates/aion-providers/src/retry.rs +++ b/crates/aion-providers/src/retry.rs @@ -5,7 +5,7 @@ use reqwest::header::HeaderMap; use serde_json::Value; use super::anthropic_shared::StreamOutcome; -use crate::error::ProviderError; +use crate::error::{ProviderError, provider_api_error, retry_after_ms_from_headers}; /// Retry a fallible async operation with exponential backoff pub async fn with_retry(max_retries: u32, f: F) -> Result @@ -87,11 +87,13 @@ pub async fn send_and_check( let status = response.status(); if !status.is_success() { + let retry_after_ms = retry_after_ms_from_headers(response.headers()); let body_text = response.text().await.unwrap_or_default(); - return Err(ProviderError::Api { - status: status.as_u16(), - message: body_text, - }); + return Err(provider_api_error( + status.as_u16(), + body_text, + retry_after_ms, + )); } Ok(response) @@ -137,6 +139,8 @@ mod tests { use super::*; use crate::error::ProviderError; + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; #[tokio::test] async fn test_retry_succeeds_first_try() { @@ -245,6 +249,37 @@ mod tests { assert_eq!(counter.load(Ordering::SeqCst), 1); } + #[tokio::test] + async fn test_send_and_check_preserves_retry_after() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with( + ResponseTemplate::new(429) + .insert_header("retry-after", "7") + .set_body_string("Too Many Requests"), + ) + .mount(&server) + .await; + + let result = send_and_check( + &reqwest::Client::new(), + &format!("{}/v1/chat/completions", server.uri()), + &HeaderMap::new(), + &serde_json::json!({}), + ) + .await; + + assert!(matches!( + result.unwrap_err(), + ProviderError::ApiSignal { + status: 429, + retry_after_ms: Some(7000), + .. + } + )); + } + // --- evaluate_outcome tests --- #[test] diff --git a/crates/aion-providers/src/vertex.rs b/crates/aion-providers/src/vertex.rs index 3ec1167a..7c4c313e 100644 --- a/crates/aion-providers/src/vertex.rs +++ b/crates/aion-providers/src/vertex.rs @@ -10,10 +10,13 @@ use std::sync::Mutex; use std::time::{SystemTime, UNIX_EPOCH}; use tokio::sync::mpsc; -use aion_types::llm::{LlmEvent, LlmRequest, ThinkingConfig}; +use aion_types::llm::{LlmRequest, ThinkingConfig}; use super::anthropic_shared; -use crate::{LlmProvider, ProviderError}; +use crate::{ + LlmProvider, ProviderError, ProviderFailure, ProviderFailurePhase, ProviderStreamReceiver, + provider_api_error, provider_failure_from_error, retry_after_ms_from_headers, +}; use aion_config::compat::ProviderCompat; pub struct VertexProvider { @@ -255,20 +258,36 @@ impl LlmProvider for VertexProvider { async fn stream( &self, request: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result { let url = self.build_url(&request.model); let body = self.build_request_body(request); + let model = Some(request.model.clone()); tracing::debug!(target: "aion_providers", body = %serde_json::to_string_pretty(&body).unwrap_or_default(), "outgoing request"); - let access_token = self.get_access_token().await?; + let access_token = self.get_access_token().await.map_err(|error| { + provider_failure_from_error( + "vertex", + model.clone(), + error, + ProviderFailurePhase::BeforeFirstDelta, + ) + })?; let mut headers = HeaderMap::new(); headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); headers.insert( AUTHORIZATION, HeaderValue::from_str(&format!("Bearer {}", access_token)) - .map_err(|e| ProviderError::Connection(format!("Header error: {}", e)))?, + .map_err(|e| ProviderError::Connection(format!("Header error: {}", e))) + .map_err(|error| { + provider_failure_from_error( + "vertex", + model.clone(), + error, + ProviderFailurePhase::BeforeFirstDelta, + ) + })?, ); let response = self @@ -277,20 +296,27 @@ impl LlmProvider for VertexProvider { .headers(headers) .json(&body) .send() - .await?; + .await + .map_err(ProviderError::from) + .map_err(|error| { + provider_failure_from_error( + "vertex", + model.clone(), + error, + ProviderFailurePhase::BeforeFirstDelta, + ) + })?; let status = response.status(); if !status.is_success() { + let retry_after_ms = retry_after_ms_from_headers(response.headers()); let body_text = response.text().await.unwrap_or_default(); - if status.as_u16() == 429 { - return Err(ProviderError::RateLimited { - retry_after_ms: 5000, - }); - } - return Err(ProviderError::Api { - status: status.as_u16(), - message: body_text, - }); + return Err(provider_failure_from_error( + "vertex", + model, + provider_api_error(status.as_u16(), body_text, retry_after_ms), + ProviderFailurePhase::BeforeFirstDelta, + )); } let (tx, rx) = mpsc::channel(64); @@ -302,22 +328,37 @@ impl LlmProvider for VertexProvider { h.insert( AUTHORIZATION, HeaderValue::from_str(&format!("Bearer {}", access_token)) - .map_err(|e| ProviderError::Connection(format!("Header error: {}", e)))?, + .map_err(|e| ProviderError::Connection(format!("Header error: {}", e))) + .map_err(|error| { + provider_failure_from_error( + "vertex", + model.clone(), + error, + ProviderFailurePhase::BeforeFirstDelta, + ) + })?, ); h }; + let stream_model = model.clone(); // Vertex uses standard SSE (same as Anthropic) tokio::spawn(async move { match anthropic_shared::process_sse_stream(response, &tx).await { anthropic_shared::StreamOutcome::Ok => {} anthropic_shared::StreamOutcome::FailedPartial(e) => { - let _ = tx.send(LlmEvent::Error(e.to_string())).await; + let failure = provider_failure_from_error( + "vertex", + stream_model.clone(), + e, + ProviderFailurePhase::AfterFirstDelta, + ); + let _ = tx.send(Err(failure)).await; } anthropic_shared::StreamOutcome::FailedEmpty(e) => { if e.is_retryable() { let mut backoff = std::time::Duration::from_secs(1); - let mut final_err = Some(e); + let mut final_err = Some((e, ProviderFailurePhase::BeforeFirstDelta)); for attempt in 1..=crate::retry::MAX_STREAM_RETRIES { backoff = crate::retry::backoff_sleep(attempt, backoff).await; match crate::retry::send_and_check( @@ -331,30 +372,51 @@ impl LlmProvider for VertexProvider { Ok(resp) => { let outcome = anthropic_shared::process_sse_stream(resp, &tx).await; + let phase = match &outcome { + anthropic_shared::StreamOutcome::FailedPartial(_) => { + ProviderFailurePhase::AfterFirstDelta + } + anthropic_shared::StreamOutcome::Ok + | anthropic_shared::StreamOutcome::FailedEmpty(_) => { + ProviderFailurePhase::BeforeFirstDelta + } + }; match crate::retry::evaluate_outcome(outcome, attempt) { Ok(None) => { final_err = None; break; } Ok(Some(e)) => { - final_err = Some(e); + final_err = Some((e, phase)); break; } Err(_) => continue, } } Err(e) if attempt == crate::retry::MAX_STREAM_RETRIES => { - final_err = Some(e); + final_err = Some((e, ProviderFailurePhase::BeforeFirstDelta)); break; } Err(_) => continue, } } - if let Some(err) = final_err { - let _ = tx.send(LlmEvent::Error(err.to_string())).await; + if let Some((err, phase)) = final_err { + let failure = provider_failure_from_error( + "vertex", + stream_model.clone(), + err, + phase, + ); + let _ = tx.send(Err(failure)).await; } } else { - let _ = tx.send(LlmEvent::Error(e.to_string())).await; + let failure = provider_failure_from_error( + "vertex", + stream_model.clone(), + e, + ProviderFailurePhase::BeforeFirstDelta, + ); + let _ = tx.send(Err(failure)).await; } } } diff --git a/crates/aion-providers/tests/malformed_tool_call.rs b/crates/aion-providers/tests/malformed_tool_call.rs index 13561bf1..777b5f83 100644 --- a/crates/aion-providers/tests/malformed_tool_call.rs +++ b/crates/aion-providers/tests/malformed_tool_call.rs @@ -1,7 +1,7 @@ use aion_config::compat::ProviderCompat; -use aion_providers::LlmProvider; use aion_providers::anthropic_shared; use aion_providers::openai::OpenAIProvider; +use aion_providers::{LlmProvider, ProviderFailure}; use aion_types::llm::{LlmEvent, LlmRequest}; use aion_types::message::{ContentBlock, Message, Role}; use serde_json::{Value, json}; @@ -70,9 +70,12 @@ fn openai_request(messages: Vec) -> LlmRequest { } } -async fn collect_events(mut rx: tokio::sync::mpsc::Receiver) -> Vec { +async fn collect_events( + mut rx: tokio::sync::mpsc::Receiver>, +) -> Vec { let mut events = Vec::new(); - while let Some(event) = rx.recv().await { + while let Some(item) = rx.recv().await { + let event = item.expect("provider stream item should be an event"); events.push(event); } events diff --git a/crates/aion-providers/tests/provider_anthropic_test.rs b/crates/aion-providers/tests/provider_anthropic_test.rs index dcce169e..1bf8b9e1 100644 --- a/crates/aion-providers/tests/provider_anthropic_test.rs +++ b/crates/aion-providers/tests/provider_anthropic_test.rs @@ -5,7 +5,7 @@ use wiremock::{Mock, MockServer, ResponseTemplate}; use aion_config::compat::ProviderCompat; use aion_providers::anthropic::AnthropicProvider; -use aion_providers::{LlmProvider, ProviderError}; +use aion_providers::{LlmProvider, ProviderFailure, ProviderFailureKind}; use aion_types::llm::{LlmEvent, LlmRequest, ThinkingConfig}; use aion_types::message::{ContentBlock, Message, Role, StopReason}; @@ -49,9 +49,12 @@ fn text_sse_body(text: &str) -> String { } /// Collect all events from a receiver into a Vec, draining until closed. -async fn collect_events(mut rx: tokio::sync::mpsc::Receiver) -> Vec { +async fn collect_events( + mut rx: tokio::sync::mpsc::Receiver>, +) -> Vec { let mut events = Vec::new(); - while let Some(ev) = rx.recv().await { + while let Some(item) = rx.recv().await { + let ev = item.expect("provider stream item should be an event"); events.push(ev); } events @@ -283,7 +286,7 @@ data: {\"type\":\"message_stop\"}\n\n"; // test_anthropic_auth_error // --------------------------------------------------------------------------- -/// A 401 response from the API should produce a ProviderError::Api with status 401. +/// A 401 response from the API should produce an auth ProviderFailure. #[tokio::test] async fn test_anthropic_auth_error() { let server = MockServer::start().await; @@ -308,25 +311,21 @@ async fn test_anthropic_auth_error() { // Act let result = provider.stream(&request).await; - // Assert: returns an Api error with status 401 - match result { - Err(ProviderError::Api { status, message }) => { - assert_eq!(status, 401); - assert!( - message.contains("authentication_error") || message.contains("invalid x-api-key"), - "unexpected error message: {message}" - ); - } - Err(other) => panic!("expected Api error, got: {other:?}"), - Ok(_) => panic!("expected an error but stream succeeded"), - } + let failure = result.expect_err("expected an error but stream succeeded"); + assert_eq!(failure.kind, ProviderFailureKind::AuthFailed); + assert_eq!(failure.meta.status, Some(401)); + let message = failure.meta.raw.message.as_deref().unwrap_or_default(); + assert!( + message.contains("authentication_error") || message.contains("invalid x-api-key"), + "unexpected error message: {message}" + ); } // --------------------------------------------------------------------------- // test_anthropic_rate_limit_retryable // --------------------------------------------------------------------------- -/// A 429 response from the API should produce a ProviderError::RateLimited. +/// A 429 response from the API should produce a rate-limit ProviderFailure. #[tokio::test] async fn test_anthropic_rate_limit_retryable() { let server = MockServer::start().await; @@ -350,14 +349,12 @@ async fn test_anthropic_rate_limit_retryable() { // Act let result = provider.stream(&request).await; - // Assert: RateLimited error, which is retryable - match result { - Err(ProviderError::RateLimited { retry_after_ms }) => { - assert!(retry_after_ms > 0, "retry_after_ms should be positive"); - } - Err(other) => panic!("expected RateLimited error, got: {other:?}"), - Ok(_) => panic!("expected an error but stream succeeded"), - } + let failure = result.expect_err("expected an error but stream succeeded"); + assert_eq!(failure.kind, ProviderFailureKind::RateLimited); + assert!( + failure.meta.retry_after_ms.unwrap_or_default() > 0, + "retry_after_ms should be positive" + ); } // --------------------------------------------------------------------------- diff --git a/crates/aion-providers/tests/provider_failure_classification.rs b/crates/aion-providers/tests/provider_failure_classification.rs new file mode 100644 index 00000000..ad961c99 --- /dev/null +++ b/crates/aion-providers/tests/provider_failure_classification.rs @@ -0,0 +1,220 @@ +use aion_providers::error::{ + ProviderError, ProviderFailure, ProviderFailureKind, ProviderFailurePhase, + ProviderRawSignalSource, ProviderRawSignals, ToolCallFailureReason, ToolSchemaSource, + classify_provider_failure, provider_failure_from_error, +}; + +fn raw_signals(status: Option, message: impl Into) -> ProviderRawSignals { + ProviderRawSignals { + provider: "openai".to_string(), + model: Some("gpt-4o-mini".to_string()), + status, + raw_code: None, + raw_type: None, + message: Some(message.into()), + request_id: Some("req_123".to_string()), + retry_after_ms: None, + source: ProviderRawSignalSource::Http, + phase: ProviderFailurePhase::BeforeFirstDelta, + } +} + +#[test] +fn classifies_401_as_auth_failed_and_preserves_provider_status() { + let mut raw = raw_signals(Some(401), "Incorrect API key provided"); + raw.raw_code = Some("invalid_api_key".to_string()); + + let failure = classify_provider_failure(raw); + + assert_eq!(failure.kind, ProviderFailureKind::AuthFailed); + assert_eq!(failure.meta.provider, "openai"); + assert_eq!(failure.meta.model.as_deref(), Some("gpt-4o-mini")); + assert_eq!(failure.meta.status, Some(401)); + assert_eq!(failure.meta.request_id.as_deref(), Some("req_123")); + assert_eq!( + failure.meta.raw.raw_code.as_deref(), + Some("invalid_api_key") + ); + assert_eq!( + failure.meta.raw.message.as_deref(), + Some("Incorrect API key provided") + ); +} + +#[test] +fn classifies_429_as_rate_limited_and_preserves_retry_after() { + let mut raw = raw_signals(Some(429), "rate limit exceeded"); + raw.retry_after_ms = Some(2_500); + + let failure = classify_provider_failure(raw); + + assert_eq!(failure.kind, ProviderFailureKind::RateLimited); + assert_eq!(failure.meta.status, Some(429)); + assert_eq!(failure.meta.retry_after_ms, Some(2_500)); +} + +#[tokio::test] +async fn provider_error_http_without_status_is_transport_failure() { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + drop(listener); + + let error = reqwest::Client::new() + .post(format!("http://{addr}/v1/messages")) + .send() + .await + .unwrap_err(); + assert_eq!(error.status(), None); + + let failure = provider_failure_from_error( + "anthropic", + Some("claude-3-5-sonnet".to_string()), + ProviderError::Http(error), + ProviderFailurePhase::BeforeFirstDelta, + ); + + assert_eq!(failure.kind, ProviderFailureKind::Transport); + assert_eq!(failure.meta.phase, ProviderFailurePhase::BeforeFirstDelta); +} + +#[test] +fn classifies_context_window_messages_as_context_too_large() { + for message in [ + "maximum context length is 128000 tokens", + "too many tokens requested", + "prompt too long for the selected model", + ] { + let failure = classify_provider_failure(raw_signals(Some(400), message)); + + assert_eq!( + failure.kind, + ProviderFailureKind::ContextTooLarge, + "message should classify as ContextTooLarge: {message}" + ); + } +} + +#[test] +fn classifies_json_source_400_as_response_parse_when_no_specific_rule() { + let mut raw = raw_signals(Some(400), "malformed provider payload"); + raw.source = ProviderRawSignalSource::Json; + + let failure = classify_provider_failure(raw); + + assert_eq!(failure.kind, ProviderFailureKind::ResponseParse); +} + +#[test] +fn classifies_404_with_model_context_but_no_model_signal_as_endpoint_not_found() { + let failure = classify_provider_failure(raw_signals(Some(404), "not found")); + + assert_eq!(failure.kind, ProviderFailureKind::EndpointNotFound); +} + +#[test] +fn classifies_404_with_model_not_found_raw_code_as_model_not_found() { + let mut raw = raw_signals(Some(404), "not found"); + raw.raw_code = Some("model_not_found".to_string()); + + let failure = classify_provider_failure(raw); + + assert_eq!(failure.kind, ProviderFailureKind::ModelNotFound); +} + +#[test] +fn classifies_404_with_model_message_text_as_model_not_found() { + let failure = classify_provider_failure(raw_signals(Some(404), "model gpt-x not found")); + + assert_eq!(failure.kind, ProviderFailureKind::ModelNotFound); +} + +#[test] +fn tool_call_invalid_summarizes_and_redacts_detail_like_meta_raw_message() { + let secret = "sk-this-secret-value-should-not-survive"; + let reason = format!("malformed tool call {secret} {}", "x".repeat(900)); + + let failure = ProviderFailure::tool_call_invalid( + "anthropic".to_string(), + Some("claude-3-5-sonnet".to_string()), + reason, + ProviderFailurePhase::AfterToolCall, + ); + + let detail = match &failure.kind { + ProviderFailureKind::ToolCall(tool_call) => tool_call.detail.as_deref().unwrap(), + other => panic!("expected ToolCall failure, got {other:?}"), + }; + let raw_message = failure.meta.raw.message.as_deref().unwrap(); + + assert_eq!(detail, raw_message); + assert!(detail.chars().count() <= 512); + assert!(detail.contains("[redacted]")); + assert!(!detail.contains(secret)); +} + +#[test] +fn redacts_aws_access_key_shapes_in_raw_message_summary() { + let akia_key = "AKIAIOSFODNN7EXAMPLE"; + let asia_key = "ASIAIOSFODNN7EXAMPLE"; + let failure = classify_provider_failure(raw_signals( + None, + format!("credentials leaked: {akia_key} and {asia_key}"), + )); + + let summary = failure.meta.raw.message.as_deref().unwrap(); + + assert!(!summary.contains(akia_key)); + assert!(!summary.contains(asia_key)); + assert_eq!(summary.matches("[redacted]").count(), 2); +} + +#[test] +fn provider_error_conversion_defaults_phase_to_before_first_delta() { + for error in [ + ProviderError::Api { + status: 500, + message: "server error".to_string(), + }, + ProviderError::Parse("bad sse".to_string()), + ProviderError::RateLimited { + retry_after_ms: 1_000, + }, + ProviderError::PromptTooLong("prompt too long".to_string()), + ProviderError::Connection("disconnect".to_string()), + ] { + let failure = ProviderFailure::from(error); + + assert_eq!(failure.meta.phase, ProviderFailurePhase::BeforeFirstDelta); + } +} + +#[test] +fn classifies_provider_tool_call_invalid_without_mcp_tool_naming() { + let failure = ProviderFailure::tool_call_invalid( + "anthropic".to_string(), + Some("claude-3-5-sonnet".to_string()), + "missing tool call id".to_string(), + ProviderFailurePhase::AfterToolCall, + ); + + match &failure.kind { + ProviderFailureKind::ToolCall(tool_call) => { + assert_eq!( + tool_call.reason, + ToolCallFailureReason::InvalidProviderToolCall + ); + assert_eq!(tool_call.detail.as_deref(), Some("missing tool call id")); + } + other => panic!("expected ToolCall failure, got {other:?}"), + } + + assert_eq!(failure.meta.provider, "anthropic"); + assert_eq!(failure.meta.model.as_deref(), Some("claude-3-5-sonnet")); + assert_eq!(failure.meta.phase, ProviderFailurePhase::AfterToolCall); + assert_eq!( + format!("{:?}", ToolSchemaSource::McpCapability), + "McpCapability" + ); + let forbidden_mcp_tool_name = ["Mcp", "Tool"].join(""); + assert!(!format!("{failure:?}").contains(&forbidden_mcp_tool_name)); +} diff --git a/crates/aion-providers/tests/provider_openai_test.rs b/crates/aion-providers/tests/provider_openai_test.rs index 29160a21..1acd30f5 100644 --- a/crates/aion-providers/tests/provider_openai_test.rs +++ b/crates/aion-providers/tests/provider_openai_test.rs @@ -1,8 +1,8 @@ use std::time::Duration; use aion_config::compat::ProviderCompat; -use aion_providers::LlmProvider; use aion_providers::openai::OpenAIProvider; +use aion_providers::{LlmProvider, ProviderFailure, ProviderFailureKind, ProviderFailurePhase}; use aion_types::llm::{LlmEvent, LlmRequest}; use aion_types::message::{ContentBlock, Message, Role, StopReason}; use serde_json::json; @@ -34,9 +34,12 @@ fn make_request() -> LlmRequest { } /// Collect all events from the receiver until the channel closes. -async fn collect_events(mut rx: tokio::sync::mpsc::Receiver) -> Vec { +async fn collect_events( + mut rx: tokio::sync::mpsc::Receiver>, +) -> Vec { let mut events = Vec::new(); - while let Some(event) = rx.recv().await { + while let Some(item) = rx.recv().await { + let event = item.expect("provider stream item should be an event"); events.push(event); } events @@ -79,6 +82,22 @@ async fn start_server_after_initial_connect_refusal(sse_body: String) -> String format!("http://{addr}") } +async fn start_server_with_two_raw_responses(first: String, second: String) -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + tokio::spawn(async move { + for response in [first, second] { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut buf = [0_u8; 4096]; + let _ = stream.read(&mut buf).await.unwrap(); + stream.write_all(response.as_bytes()).await.unwrap(); + } + }); + + format!("http://{addr}") +} + // --------------------------------------------------------------------------- // test_openai_stream_text_response // --------------------------------------------------------------------------- @@ -526,7 +545,7 @@ async fn test_openai_stream_state_transitions() { // test_openai_api_error_non_success_status // --------------------------------------------------------------------------- -/// Verify that a non-2xx HTTP response is surfaced as a ProviderError::Api. +/// Verify that a non-2xx HTTP response is surfaced as a ProviderFailure. #[tokio::test] async fn test_openai_api_error_non_success_status() { let server = MockServer::start().await; @@ -543,19 +562,16 @@ async fn test_openai_api_error_non_success_status() { let result = provider.stream(&make_request()).await; assert!(result.is_err()); - match result.unwrap_err() { - aion_providers::ProviderError::Api { status, .. } => { - assert_eq!(status, 401); - } - e => panic!("expected Api error, got: {:?}", e), - } + let failure = result.unwrap_err(); + assert_eq!(failure.kind, ProviderFailureKind::AuthFailed); + assert_eq!(failure.meta.status, Some(401)); } // --------------------------------------------------------------------------- // test_openai_rate_limited // --------------------------------------------------------------------------- -/// Verify that a 429 response is surfaced as ProviderError::RateLimited. +/// Verify that a 429 response is surfaced as ProviderFailureKind::RateLimited. #[tokio::test] async fn test_openai_rate_limited() { let server = MockServer::start().await; @@ -571,12 +587,9 @@ async fn test_openai_rate_limited() { let result = provider.stream(&make_request()).await; assert!(result.is_err()); - match result.unwrap_err() { - aion_providers::ProviderError::RateLimited { retry_after_ms } => { - assert_eq!(retry_after_ms, 5000); - } - e => panic!("expected RateLimited error, got: {:?}", e), - } + let failure = result.unwrap_err(); + assert_eq!(failure.kind, ProviderFailureKind::RateLimited); + assert_eq!(failure.meta.retry_after_ms, Some(5000)); } // --------------------------------------------------------------------------- @@ -710,3 +723,67 @@ async fn test_openai_stream_empty_content_delta_skipped() { e => panic!("expected Done, got: {:?}", e), } } + +#[tokio::test] +async fn test_openai_malformed_sse_json_emits_response_parse_failure() { + let server = MockServer::start().await; + let sse_body = "data: {not-json}\n\n"; + + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_raw(sse_body, "text/event-stream")) + .mount(&server) + .await; + + let provider = + OpenAIProvider::new("test-key", &server.uri(), ProviderCompat::openai_defaults()); + let mut rx = provider.stream(&make_request()).await.unwrap(); + + let item = rx + .recv() + .await + .expect("malformed JSON should produce a stream item"); + let failure = item.expect_err("malformed JSON should be a provider failure"); + assert_eq!(failure.kind, ProviderFailureKind::ResponseParse); + assert_eq!(failure.meta.phase, ProviderFailurePhase::BeforeFirstDelta); +} + +#[tokio::test] +async fn test_openai_retry_partial_failure_preserves_after_first_delta_phase() { + let first_response = + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: 999\r\nconnection: close\r\n\r\n" + .to_string(); + + let chunk = json!({ + "id": "chatcmpl-retry-partial", + "object": "chat.completion.chunk", + "choices": [{ + "index": 0, + "delta": { "content": "partial" }, + "finish_reason": null + }] + }) + .to_string(); + let second_body = format!("data: {chunk}\n\ndata: {{not-json}}\n\n"); + let second_response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + second_body.len(), + second_body + ); + let base_url = start_server_with_two_raw_responses(first_response, second_response).await; + + let provider = OpenAIProvider::new("test-key", &base_url, ProviderCompat::openai_defaults()); + let mut rx = provider.stream(&make_request()).await.unwrap(); + + let mut failure = None; + while let Some(item) = rx.recv().await { + if let Err(err) = item { + failure = Some(err); + break; + } + } + + let failure = failure.expect("retry response partial parse failure should be emitted"); + assert_eq!(failure.kind, ProviderFailureKind::ResponseParse); + assert_eq!(failure.meta.phase, ProviderFailurePhase::AfterFirstDelta); +} diff --git a/crates/aion-types/src/llm.rs b/crates/aion-types/src/llm.rs index 25eab577..8e1b5d8d 100644 --- a/crates/aion-types/src/llm.rs +++ b/crates/aion-types/src/llm.rs @@ -45,8 +45,6 @@ pub enum LlmEvent { stop_reason: StopReason, usage: TokenUsage, }, - /// Error from the API - Error(String), } #[cfg(test)] diff --git a/docs/json-stream-protocol.md b/docs/json-stream-protocol.md index add9cc8a..a3acd82d 100644 --- a/docs/json-stream-protocol.md +++ b/docs/json-stream-protocol.md @@ -226,20 +226,36 @@ An error occurred. The agent may or may not continue depending on severity. "type": "error", "msg_id": "abc-123", "error": { - "code": "provider_error", - "message": "Rate limit exceeded", - "retryable": true + "code": "provider_rate_limited", + "message": "Rate limit exceeded. Retry after 30s.", + "ownership": "provider", + "details": [ + { + "key": "provider", + "value": "openai" + }, + { + "key": "status", + "value": "429" + } + ] } } ``` -| Error Code | Description | -|------------|-------------| -| `provider_error` | LLM API error (rate limit, auth, etc.) | -| `tool_error` | Built-in tool execution error | -| `config_error` | Configuration or initialization error | -| `protocol_error` | Invalid command from client | -| `internal_error` | Unexpected internal error | +| Field | Type | Description | +|-------|------|-------------| +| `code` | string | Stable snake_case public error code, for example `provider_rate_limited`, `tool_execution_failed`, `protocol_parse_failed`, or `internal_error` | +| `message` | string | Human-readable error message suitable for display | +| `ownership` | string | Responsible party: `user`, `provider`, `aionrs`, `host`, `tool`, `mcp_server`, or `unknown` | +| `details` | array | Structured key/value metadata entries; always present and empty when no metadata is available | + +Each `details` entry has: + +| Field | Type | Description | +|-------|------|-------------| +| `key` | string | Stable snake_case detail key, such as `provider`, `model`, `status`, `request_id`, `phase`, `config_key`, `profile`, `session_id`, `tool_name`, `mcp_server`, `mcp_capability`, `command`, `raw_code`, or `raw_type` | +| `value` | string | Detail value | ### 1.11 `info` @@ -490,9 +506,15 @@ After the first `message`, any further `add_mcp_server` commands are rejected: { "type": "error", "error": { - "code": "protocol_error", - "message": "AddMcpServer 'name': rejected — only allowed before first Message", - "retryable": false + "code": "protocol_state_violation", + "message": "AddMcpServer 'name': rejected - only allowed before first Message", + "ownership": "host", + "details": [ + { + "key": "command", + "value": "add_mcp_server" + } + ] } } ``` @@ -583,16 +605,23 @@ Client closes stdin (EOF) or sends SIGTERM. Agent cleans up and exits. ### 4.1 Invalid Command -If client sends malformed JSON or unknown command type: +`protocol_parse_failed` is reserved for host-visible command parse failures such as malformed JSON or an unknown command type. The current stdin reader logs malformed lines before command dispatch and continues; commands that parse successfully but violate the current state are emitted as `protocol_state_violation`. + +Reserved shape: ```json { "type": "error", - "msg_id": null, "error": { - "code": "protocol_error", + "code": "protocol_parse_failed", "message": "Unknown command type: foo", - "retryable": false + "ownership": "host", + "details": [ + { + "key": "raw_type", + "value": "foo" + } + ] } } ``` @@ -606,9 +635,19 @@ Agent should emit error and let the conversation continue if possible: "type": "error", "msg_id": "m3", "error": { - "code": "provider_error", + "code": "provider_rate_limited", "message": "Rate limit exceeded. Retry after 30s.", - "retryable": true + "ownership": "provider", + "details": [ + { + "key": "provider", + "value": "openai" + }, + { + "key": "status", + "value": "429" + } + ] } } ``` @@ -620,11 +659,16 @@ For unrecoverable errors, agent emits error and exits with non-zero status: ```json { "type": "error", - "msg_id": null, "error": { - "code": "config_error", + "code": "config_env_missing", "message": "ANTHROPIC_API_KEY not set", - "retryable": false + "ownership": "user", + "details": [ + { + "key": "config_key", + "value": "ANTHROPIC_API_KEY" + } + ] } } ``` diff --git a/tests/common/mod.rs b/tests/common/mod.rs index bacaadd8..cfff96cd 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -12,7 +12,7 @@ use aionrs::protocol::events::ToolCategory; use aionrs::config::{Config, ProviderType, ToolsConfig, SessionConfig}; use aionrs::hooks::HooksConfig; use aionrs::mcp::config::McpConfig; -use aionrs::provider::{LlmProvider, ProviderError}; +use aionrs::provider::{LlmProvider, ProviderFailure, ProviderStreamReceiver}; use aionrs::tools::Tool; use aionrs::types::llm::{LlmEvent, LlmRequest}; use aionrs::types::message::{StopReason, TokenUsage}; @@ -26,7 +26,7 @@ use aionrs::types::tool::ToolResult; /// Each call to `stream` pops the first sequence from `responses`. /// When `responses` is empty it falls back to a single EndTurn with empty text. pub struct MockLlmProvider { - responses: Mutex>>, + responses: Mutex>>>, } impl MockLlmProvider { @@ -45,7 +45,7 @@ impl MockLlmProvider { }, ]; Self { - responses: Mutex::new(vec![events]), + responses: Mutex::new(vec![Self::ok_events(events)]), } } @@ -68,7 +68,7 @@ impl MockLlmProvider { }, ]; Self { - responses: Mutex::new(vec![events]), + responses: Mutex::new(vec![Self::ok_events(events)]), } } @@ -76,16 +76,26 @@ impl MockLlmProvider { /// Each call to `stream` consumes the next sequence. pub fn with_turns(turns: Vec>) -> Self { Self { - responses: Mutex::new(turns), + responses: Mutex::new(turns.into_iter().map(Self::ok_events).collect()), } } /// Create a provider that returns custom events. pub fn with_events(events: Vec) -> Self { Self { - responses: Mutex::new(vec![events]), + responses: Mutex::new(vec![Self::ok_events(events)]), } } + + pub fn with_stream_items(items: Vec>) -> Self { + Self { + responses: Mutex::new(vec![items]), + } + } + + fn ok_events(events: Vec) -> Vec> { + events.into_iter().map(Ok).collect() + } } #[async_trait] @@ -93,15 +103,15 @@ impl LlmProvider for MockLlmProvider { async fn stream( &self, _request: &LlmRequest, - ) -> Result, ProviderError> { + ) -> Result { let events = { let mut responses = self.responses.lock().unwrap(); if responses.is_empty() { // Fallback: end turn with empty text - vec![LlmEvent::Done { + vec![Ok(LlmEvent::Done { stop_reason: StopReason::EndTurn, usage: TokenUsage::default(), - }] + })] } else { responses.remove(0) } @@ -109,8 +119,8 @@ impl LlmProvider for MockLlmProvider { let (tx, rx) = mpsc::channel(64); tokio::spawn(async move { - for event in events { - let _ = tx.send(event).await; + for item in events { + let _ = tx.send(item).await; } }); Ok(rx) diff --git a/tests/engine_test.rs b/tests/engine_test.rs index afee235c..abc12efb 100644 --- a/tests/engine_test.rs +++ b/tests/engine_test.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use aionrs::engine::{AgentEngine, AgentError}; use aionrs::output::terminal::TerminalSink; use aionrs::output::OutputSink; +use aionrs::provider::ProviderError; use aionrs::session::SessionManager; use aionrs::tools::registry::ToolRegistry; use aionrs::types::llm::LlmEvent; @@ -305,17 +306,12 @@ async fn test_engine_max_turns_returns_ok() { assert_eq!(result.turns, 1); } -// --------------------------------------------------------------------------- -// test_engine_api_error_handling -// -// Verifies that an LlmEvent::Error propagates as AgentError::ApiError with -// the original error message intact. // --------------------------------------------------------------------------- #[tokio::test] -async fn test_engine_api_error_handling() { - let events = vec![LlmEvent::Error("test error".to_string())]; - - let provider = Arc::new(MockLlmProvider::with_events(events)); +async fn provider_stream_failure_maps_to_agent_provider_failure() { + let provider = Arc::new(MockLlmProvider::with_stream_items(vec![Err( + ProviderError::Connection("test error".to_string()).into(), + )])); let config = test_config(); let registry = ToolRegistry::new(); let output = silent_output(); @@ -327,8 +323,5 @@ async fn test_engine_api_error_handling() { .map(|_| panic!("expected error, got Ok")) .unwrap_err(); - match err { - AgentError::ApiError(msg) => assert_eq!(msg, "test error"), - other => panic!("expected ApiError(\"test error\"), got: {:?}", other), - } + assert!(matches!(err, AgentError::Provider(_))); } diff --git a/tests/provider_anthropic_test.rs b/tests/provider_anthropic_test.rs index 84228c47..ff8a1c5d 100644 --- a/tests/provider_anthropic_test.rs +++ b/tests/provider_anthropic_test.rs @@ -6,7 +6,7 @@ use wiremock::{Mock, MockServer, ResponseTemplate}; use aionrs::provider::anthropic::AnthropicProvider; use aionrs::provider::compat::ProviderCompat; use aionrs::provider::debug::DebugConfig; -use aionrs::provider::{LlmProvider, ProviderError}; +use aionrs::provider::{LlmProvider, ProviderFailure, ProviderFailureKind}; use aionrs::types::llm::{LlmEvent, LlmRequest, ThinkingConfig}; use aionrs::types::message::{ContentBlock, Message, Role, StopReason}; @@ -50,9 +50,12 @@ fn text_sse_body(text: &str) -> String { } /// Collect all events from a receiver into a Vec, draining until closed. -async fn collect_events(mut rx: tokio::sync::mpsc::Receiver) -> Vec { +async fn collect_events( + mut rx: tokio::sync::mpsc::Receiver>, +) -> Vec { let mut events = Vec::new(); - while let Some(ev) = rx.recv().await { + while let Some(item) = rx.recv().await { + let ev = item.expect("provider stream item should be an event"); events.push(ev); } events @@ -259,7 +262,7 @@ data: {\"type\":\"message_stop\"}\n\n"; // test_anthropic_auth_error // --------------------------------------------------------------------------- -/// A 401 response from the API should produce a ProviderError::Api with status 401. +/// A 401 response from the API should produce an auth ProviderFailure. #[tokio::test] async fn test_anthropic_auth_error() { let server = MockServer::start().await; @@ -280,25 +283,21 @@ async fn test_anthropic_auth_error() { // Act let result = provider.stream(&request).await; - // Assert: returns an Api error with status 401 - match result { - Err(ProviderError::Api { status, message }) => { - assert_eq!(status, 401); - assert!( - message.contains("authentication_error") || message.contains("invalid x-api-key"), - "unexpected error message: {message}" - ); - } - Err(other) => panic!("expected Api error, got: {other:?}"), - Ok(_) => panic!("expected an error but stream succeeded"), - } + let failure = result.expect_err("expected an error but stream succeeded"); + assert_eq!(failure.kind, ProviderFailureKind::AuthFailed); + assert_eq!(failure.meta.status, Some(401)); + let message = failure.meta.raw.message.as_deref().unwrap_or_default(); + assert!( + message.contains("authentication_error") || message.contains("invalid x-api-key"), + "unexpected error message: {message}" + ); } // --------------------------------------------------------------------------- // test_anthropic_rate_limit_retryable // --------------------------------------------------------------------------- -/// A 429 response from the API should produce a ProviderError::RateLimited. +/// A 429 response from the API should produce a rate-limit ProviderFailure. #[tokio::test] async fn test_anthropic_rate_limit_retryable() { let server = MockServer::start().await; @@ -318,14 +317,12 @@ async fn test_anthropic_rate_limit_retryable() { // Act let result = provider.stream(&request).await; - // Assert: RateLimited error, which is retryable - match result { - Err(ProviderError::RateLimited { retry_after_ms }) => { - assert!(retry_after_ms > 0, "retry_after_ms should be positive"); - } - Err(other) => panic!("expected RateLimited error, got: {other:?}"), - Ok(_) => panic!("expected an error but stream succeeded"), - } + let failure = result.expect_err("expected an error but stream succeeded"); + assert_eq!(failure.kind, ProviderFailureKind::RateLimited); + assert!( + failure.meta.retry_after_ms.unwrap_or_default() > 0, + "retry_after_ms should be positive" + ); } // --------------------------------------------------------------------------- diff --git a/tests/provider_openai_test.rs b/tests/provider_openai_test.rs index 93f87c28..d2ab05f4 100644 --- a/tests/provider_openai_test.rs +++ b/tests/provider_openai_test.rs @@ -1,4 +1,4 @@ -use aionrs::provider::LlmProvider; +use aionrs::provider::{LlmProvider, ProviderFailure, ProviderFailureKind}; use aionrs::provider::compat::ProviderCompat; use aionrs::provider::debug::DebugConfig; use aionrs::provider::openai::OpenAIProvider; @@ -31,9 +31,12 @@ fn make_request() -> LlmRequest { } /// Collect all events from the receiver until the channel closes. -async fn collect_events(mut rx: tokio::sync::mpsc::Receiver) -> Vec { +async fn collect_events( + mut rx: tokio::sync::mpsc::Receiver>, +) -> Vec { let mut events = Vec::new(); - while let Some(event) = rx.recv().await { + while let Some(item) = rx.recv().await { + let event = item.expect("provider stream item should be an event"); events.push(event); } events @@ -449,7 +452,7 @@ async fn test_openai_stream_state_transitions() { // test_openai_api_error_non_success_status // --------------------------------------------------------------------------- -/// Verify that a non-2xx HTTP response is surfaced as a ProviderError::Api. +/// Verify that a non-2xx HTTP response is surfaced as a ProviderFailure. #[tokio::test] async fn test_openai_api_error_non_success_status() { let server = MockServer::start().await; @@ -468,19 +471,16 @@ async fn test_openai_api_error_non_success_status() { let result = provider.stream(&make_request()).await; assert!(result.is_err()); - match result.unwrap_err() { - aionrs::provider::ProviderError::Api { status, .. } => { - assert_eq!(status, 401); - } - e => panic!("expected Api error, got: {:?}", e), - } + let failure = result.unwrap_err(); + assert_eq!(failure.kind, ProviderFailureKind::AuthFailed); + assert_eq!(failure.meta.status, Some(401)); } // --------------------------------------------------------------------------- // test_openai_rate_limited // --------------------------------------------------------------------------- -/// Verify that a 429 response is surfaced as ProviderError::RateLimited. +/// Verify that a 429 response is surfaced as ProviderFailureKind::RateLimited. #[tokio::test] async fn test_openai_rate_limited() { let server = MockServer::start().await; @@ -495,12 +495,9 @@ async fn test_openai_rate_limited() { let result = provider.stream(&make_request()).await; assert!(result.is_err()); - match result.unwrap_err() { - aionrs::provider::ProviderError::RateLimited { retry_after_ms } => { - assert_eq!(retry_after_ms, 5000); - } - e => panic!("expected RateLimited error, got: {:?}", e), - } + let failure = result.unwrap_err(); + assert_eq!(failure.kind, ProviderFailureKind::RateLimited); + assert_eq!(failure.meta.retry_after_ms, Some(5000)); } // --------------------------------------------------------------------------- diff --git a/tests/spawn_test.rs b/tests/spawn_test.rs index 94886118..37215295 100644 --- a/tests/spawn_test.rs +++ b/tests/spawn_test.rs @@ -3,6 +3,7 @@ mod common; use std::sync::Arc; use aionrs::agent::spawner::{AgentSpawner, SubAgentConfig}; +use aionrs::provider::ProviderError; use aionrs::types::llm::LlmEvent; use aionrs::types::message::{StopReason, TokenUsage}; use common::{MockLlmProvider, test_config}; @@ -138,13 +139,11 @@ async fn test_spawn_shares_provider() { assert_eq!(result2.text, "second"); } -/// An LLM error event causes the sub-agent result to be marked as an error. +/// A provider stream failure causes the sub-agent result to be marked as an error. #[tokio::test] async fn test_spawn_agent_error_captured() { - // Emit an Error event — the engine converts this to AgentError::ApiError, - // which spawner catches and stores in SubAgentResult::is_error. - let provider = Arc::new(MockLlmProvider::with_events(vec![LlmEvent::Error( - "provider failed".to_string(), + let provider = Arc::new(MockLlmProvider::with_stream_items(vec![Err( + ProviderError::Connection("provider failed".to_string()).into(), )])); let spawner = AgentSpawner::new(provider, test_config());