diff --git a/.github/workflows/rust-ci.yml b/.github/workflows/rust-ci.yml index bce689902..71b2d8832 100644 --- a/.github/workflows/rust-ci.yml +++ b/.github/workflows/rust-ci.yml @@ -191,6 +191,7 @@ jobs: env: RUSTC_WRAPPER: sccache SCCACHE_GHA_ENABLED: "true" + RUST_MIN_STACK: "16777216" run: cargo nextest run -p aether-gateway - name: Show sccache stats diff --git a/apps/aether-gateway/build.rs b/apps/aether-gateway/build.rs index b4af98622..fba2fe1f9 100644 --- a/apps/aether-gateway/build.rs +++ b/apps/aether-gateway/build.rs @@ -11,19 +11,18 @@ fn main() { let package_version = env::var("CARGO_PKG_VERSION").unwrap_or_else(|_| "unknown".to_string()); let version = env::var("AETHER_BUILD_VERSION") .ok() - .filter(|value| !value.trim().is_empty()) + .and_then(|value| normalize_gateway_version_source(&value)) .or_else(|| { env::var("AETHER_VERSION") .ok() - .filter(|value| !value.trim().is_empty()) + .and_then(|value| normalize_gateway_version_source(&value)) }) .or_else(|| { env::var("GITHUB_REF_NAME") .ok() - .filter(|value| value.trim().starts_with('v')) + .and_then(|value| normalize_gateway_version_source(&value)) }) .or_else(git_describe_version) - .map(|value| normalize_version(&value)) .filter(|value| !value.is_empty()) .unwrap_or(package_version); @@ -38,7 +37,9 @@ fn main() { fn git_describe_version() -> Option { let output = Command::new("git") - .args(["describe", "--tags", "--always", "--dirty"]) + .args([ + "describe", "--tags", "--match", "v[0-9]*", "--always", "--dirty", + ]) .output() .ok()?; if !output.status.success() { @@ -46,17 +47,13 @@ fn git_describe_version() -> Option { } let version = String::from_utf8(output.stdout).ok()?; let version = version.trim(); - if version.is_empty() { - None - } else { - Some(version.to_string()) - } + normalize_gateway_version_source(version) } -fn normalize_version(value: &str) -> String { - value - .trim() - .strip_prefix('v') - .unwrap_or(value.trim()) - .to_string() +fn normalize_gateway_version_source(value: &str) -> Option { + let trimmed = value.trim(); + if trimmed.is_empty() || trimmed.starts_with("tunnel-v") { + return None; + } + Some(trimmed.strip_prefix('v').unwrap_or(trimmed).to_string()) } diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs index 5a04068f5..cd42341cb 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs @@ -56,6 +56,7 @@ pub(crate) fn resolve_same_format_provider_transport_unsupported_reason_for_trac "jina:embedding" => "jina:embedding", "jina:rerank" => "jina:rerank", "doubao:embedding" => "doubao:embedding", + "aliyun:multimodal_embedding" => "aliyun:multimodal_embedding", _ => return Some("transport_api_format_unsupported"), }; let behavior = policy::classify_same_format_provider_request_behavior( diff --git a/apps/aether-gateway/src/api/ai/aliyun.rs b/apps/aether-gateway/src/api/ai/aliyun.rs new file mode 100644 index 000000000..c1db71102 --- /dev/null +++ b/apps/aether-gateway/src/api/ai/aliyun.rs @@ -0,0 +1,15 @@ +pub(crate) fn normalized_signature(api_format: &str) -> Option<&'static str> { + match crate::ai_serving::normalize_api_format_alias(api_format).as_str() { + "aliyun:multimodal_embedding" => Some("aliyun:multimodal_embedding"), + _ => None, + } +} + +pub(crate) fn local_path(api_format: &str) -> Option<&'static str> { + match crate::ai_serving::normalize_api_format_alias(api_format).as_str() { + "aliyun:multimodal_embedding" => { + Some("/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding") + } + _ => None, + } +} diff --git a/apps/aether-gateway/src/api/ai/mod.rs b/apps/aether-gateway/src/api/ai/mod.rs index 49b602fcf..45159b721 100644 --- a/apps/aether-gateway/src/api/ai/mod.rs +++ b/apps/aether-gateway/src/api/ai/mod.rs @@ -1,3 +1,4 @@ +mod aliyun; mod claude; mod doubao; mod gemini; diff --git a/apps/aether-gateway/src/api/ai/registry.rs b/apps/aether-gateway/src/api/ai/registry.rs index b5d212bba..119eabc00 100644 --- a/apps/aether-gateway/src/api/ai/registry.rs +++ b/apps/aether-gateway/src/api/ai/registry.rs @@ -1,7 +1,7 @@ use axum::routing::{any, post}; use axum::Router; -use super::{claude, doubao, gemini, jina, openai}; +use super::{aliyun, claude, doubao, gemini, jina, openai}; use crate::{handlers::proxy::proxy_request, state::AppState}; // Router registration patterns live here so AI public ingress has a single mount registry. @@ -56,6 +56,7 @@ pub(crate) fn public_api_format_local_path(api_format: &str) -> &'static str { .or_else(|| gemini::local_path(&normalized)) .or_else(|| jina::local_path(&normalized)) .or_else(|| doubao::local_path(&normalized)) + .or_else(|| aliyun::local_path(&normalized)) .unwrap_or("/") } @@ -66,6 +67,7 @@ pub(crate) fn normalize_admin_endpoint_signature(api_format: &str) -> Option<&'s .or_else(|| gemini::normalized_signature(&normalized)) .or_else(|| jina::normalized_signature(&normalized)) .or_else(|| doubao::normalized_signature(&normalized)) + .or_else(|| aliyun::normalized_signature(&normalized)) } pub(crate) fn admin_endpoint_signature_parts( @@ -101,6 +103,12 @@ mod tests { ), ("jina:embedding", "jina", "embedding", "/v1/embeddings"), ("doubao:embedding", "doubao", "embedding", "/v1/embeddings"), + ( + "aliyun:multimodal_embedding", + "aliyun", + "multimodal_embedding", + "/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding", + ), ("openai:rerank", "openai", "rerank", "/v1/rerank"), ("jina:rerank", "jina", "rerank", "/v1/rerank"), ] { diff --git a/apps/aether-gateway/src/data/state/runtime.rs b/apps/aether-gateway/src/data/state/runtime.rs index 23aedc87f..cbbce5584 100644 --- a/apps/aether-gateway/src/data/state/runtime.rs +++ b/apps/aether-gateway/src/data/state/runtime.rs @@ -715,6 +715,21 @@ impl GatewayDataState { } } + pub(crate) async fn find_pending_plan_purchase_order_by_user_id( + &self, + user_id: &str, + product_id: &str, + ) -> Result, DataLayerError> { + match &self.wallet_reader { + Some(repository) => { + repository + .find_pending_plan_purchase_order_by_user_id(user_id, product_id) + .await + } + None => Ok(None), + } + } + pub(crate) async fn find_wallet_refund( &self, wallet_id: &str, diff --git a/apps/aether-gateway/src/data/state/testing/mod.rs b/apps/aether-gateway/src/data/state/testing/mod.rs index 27537e3d8..aaeec82ac 100644 --- a/apps/aether-gateway/src/data/state/testing/mod.rs +++ b/apps/aether-gateway/src/data/state/testing/mod.rs @@ -958,6 +958,64 @@ impl GatewayDataState { } } + #[cfg(test)] + pub(crate) fn with_user_billing_and_wallet_for_tests( + user_repository: Arc, + billing_repository: Arc, + wallet_repository: Arc, + ) -> Self + where + T: aether_data::repository::wallet::WalletRepository + 'static, + { + let wallet_reader: Arc = wallet_repository.clone(); + let wallet_writer: Arc = wallet_repository; + Self { + config: GatewayDataConfig::disabled(), + backends: None, + auth_api_key_reader: None, + auth_api_key_writer: None, + auth_module_reader: None, + auth_module_writer: None, + announcement_reader: None, + announcement_writer: None, + management_token_reader: None, + management_token_writer: None, + oauth_provider_reader: None, + oauth_provider_writer: None, + proxy_node_reader: None, + proxy_node_writer: None, + billing_reader: Some(billing_repository), + gemini_file_mapping_reader: None, + gemini_file_mapping_writer: None, + global_model_reader: None, + global_model_writer: None, + minimal_candidate_selection_reader: None, + request_candidate_reader: None, + request_candidate_writer: None, + provider_catalog_reader: None, + provider_catalog_writer: None, + pool_score_reader: None, + pool_score_writer: None, + provider_quota_reader: None, + provider_quota_writer: None, + routing_group_reader: None, + routing_group_writer: None, + usage_reader: None, + usage_writer: None, + user_reader: Some(user_repository), + user_preferences: None, + usage_worker_queue: None, + video_task_reader: None, + video_task_writer: None, + background_task_reader: None, + background_task_writer: None, + wallet_reader: Some(wallet_reader), + wallet_writer: Some(wallet_writer), + settlement_writer: None, + system_config_values: None, + } + } + #[cfg(test)] pub(crate) fn with_user_wallet_and_usage_for_tests( user_repository: Arc, diff --git a/apps/aether-gateway/src/execution_runtime/kiro_cache.rs b/apps/aether-gateway/src/execution_runtime/kiro_cache.rs index c677e7074..839c08fa5 100644 --- a/apps/aether-gateway/src/execution_runtime/kiro_cache.rs +++ b/apps/aether-gateway/src/execution_runtime/kiro_cache.rs @@ -1,14 +1,21 @@ use std::collections::{BTreeMap, HashMap}; +use std::fmt::Write as _; use std::sync::{Mutex, OnceLock}; use std::time::{Duration, Instant}; +use aether_runtime_state::{DataLayerError, RuntimeState}; +use serde::{Deserialize, Serialize}; use serde_json::Value; use sha2::{Digest, Sha256}; +use tracing::warn; + +use crate::clock::current_unix_ms; const DEFAULT_CACHE_TTL: Duration = Duration::from_secs(300); const ONE_HOUR_CACHE_TTL: Duration = Duration::from_secs(3600); const MAX_ENTRIES: usize = 2048; -const PREFIX_LOOKBACK_LIMIT: usize = 10; +const KIRO_PROMPT_CACHE_INDEX_KEY: &str = "kiro:prompt-cache:index"; +const PREFIX_LOOKBACK_WINDOW: usize = 20; const TOKENS_PER_TOOL: u64 = 150; const TOKENS_PER_MESSAGE: u64 = 4; const INLINE_IMAGE_DATA_TOKEN_PLACEHOLDER: &str = "[inline-image-data]"; @@ -21,6 +28,7 @@ pub(crate) struct KiroPromptCacheProfile { total_input_tokens: u64, min_cacheable_tokens: u64, breakpoints: Vec, + match_candidates: Vec, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -37,6 +45,12 @@ struct KiroPromptCacheEntry { expires_at: Instant, } +#[derive(Debug, Clone, Copy, Deserialize, Serialize)] +struct KiroPromptCacheRuntimeEntry { + token_count: u64, + ttl_secs: u64, +} + #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub(crate) struct KiroPromptCacheUsage { pub(crate) cache_creation_input_tokens: u64, @@ -53,11 +67,10 @@ struct PendingBlock { value: Value, tokens: u64, breakpoint_ttl: Option, - is_message_end: bool, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct PrefixCandidate { +struct KiroPromptCacheCandidate { fingerprint: [u8; 32], cumulative_tokens: u64, } @@ -66,6 +79,230 @@ pub(crate) fn kiro_prompt_cache_tracker() -> &'static KiroPromptCacheTracker { KIRO_PROMPT_CACHE_TRACKER.get_or_init(KiroPromptCacheTracker::default) } +pub(crate) async fn compute_kiro_prompt_cache_usage( + runtime_state: &RuntimeState, + credential_id: String, + profile: &KiroPromptCacheProfile, +) -> KiroPromptCacheUsage { + match compute_kiro_prompt_cache_usage_with_runtime_state( + runtime_state, + credential_id.as_str(), + profile, + ) + .await + { + Ok(usage) => usage, + Err(err) => { + warn!( + event_name = "kiro_simulated_cache_runtime_state_failed", + log_type = "event", + error = ?err, + "failed to update Kiro simulated cache runtime state; falling back to process-local tracker" + ); + kiro_prompt_cache_tracker().compute_and_update(credential_id, profile) + } + } +} + +async fn compute_kiro_prompt_cache_usage_with_runtime_state( + runtime_state: &RuntimeState, + credential_id: &str, + profile: &KiroPromptCacheProfile, +) -> Result { + let last_breakpoint = match profile.breakpoints.last().copied() { + Some(last_breakpoint) => last_breakpoint, + None => return Ok(KiroPromptCacheUsage::default()), + }; + + let reversed_candidates = profile + .match_candidates + .iter() + .rev() + .copied() + .collect::>(); + let candidate_keys = reversed_candidates + .iter() + .map(|candidate| kiro_prompt_cache_runtime_key(credential_id, &candidate.fingerprint)) + .collect::>(); + let candidate_values = runtime_state.kv_get_many(&candidate_keys).await?; + let mut existing_entries = HashMap::::new(); + let mut matched_tokens = 0u64; + let mut matched_refresh: Option<(String, KiroPromptCacheRuntimeEntry)> = None; + + for ((candidate, key), value) in reversed_candidates + .iter() + .zip(candidate_keys.iter()) + .zip(candidate_values) + { + let Some(entry) = value + .as_deref() + .and_then(parse_kiro_prompt_cache_runtime_entry) + else { + continue; + }; + existing_entries.insert(key.clone(), entry); + if matched_tokens == 0 { + matched_tokens = entry + .token_count + .min(candidate.cumulative_tokens) + .min(profile.total_input_tokens); + matched_refresh = Some((key.clone(), entry)); + } + } + + if let Some((key, entry)) = matched_refresh { + store_kiro_prompt_cache_runtime_entry(runtime_state, key.as_str(), entry).await?; + } + + let creation_tokens = last_breakpoint + .cumulative_tokens + .min(profile.total_input_tokens) + .saturating_sub(matched_tokens); + + for breakpoint in &profile.breakpoints { + let key = kiro_prompt_cache_runtime_key(credential_id, &breakpoint.fingerprint); + let ttl_secs = breakpoint.ttl.as_secs().max(1); + let entry = existing_entries + .get(&key) + .copied() + .map(|existing| KiroPromptCacheRuntimeEntry { + token_count: existing.token_count.max(breakpoint.cumulative_tokens), + ttl_secs: existing.ttl_secs.max(ttl_secs), + }) + .unwrap_or(KiroPromptCacheRuntimeEntry { + token_count: breakpoint.cumulative_tokens, + ttl_secs, + }); + store_kiro_prompt_cache_runtime_entry(runtime_state, key.as_str(), entry).await?; + } + + trim_kiro_prompt_cache_runtime_state(runtime_state, MAX_ENTRIES).await; + + Ok(KiroPromptCacheUsage { + cache_creation_input_tokens: creation_tokens, + cache_read_input_tokens: matched_tokens, + }) +} + +async fn store_kiro_prompt_cache_runtime_entry( + runtime_state: &RuntimeState, + key: &str, + entry: KiroPromptCacheRuntimeEntry, +) -> Result<(), DataLayerError> { + let ttl = Duration::from_secs(entry.ttl_secs.max(1)); + runtime_state + .kv_set( + key, + encode_kiro_prompt_cache_runtime_entry(entry), + Some(ttl), + ) + .await?; + + let expires_at_ms = current_unix_ms().saturating_add(entry.ttl_secs.saturating_mul(1000)); + if let Err(err) = runtime_state + .score_set(KIRO_PROMPT_CACHE_INDEX_KEY, key, expires_at_ms as f64) + .await + { + warn!( + event_name = "kiro_simulated_cache_index_update_failed", + log_type = "event", + cache_key = %key, + error = ?err, + "failed to update Kiro simulated cache index; cache entry was persisted but cleanup may lag" + ); + } + + Ok(()) +} + +async fn trim_kiro_prompt_cache_runtime_state(runtime_state: &RuntimeState, max_entries: usize) { + if let Err(err) = runtime_state + .score_remove_by_score(KIRO_PROMPT_CACHE_INDEX_KEY, current_unix_ms() as f64) + .await + { + warn!( + event_name = "kiro_simulated_cache_index_expiry_trim_failed", + log_type = "event", + error = ?err, + "failed to trim expired Kiro simulated cache index entries" + ); + return; + } + + let Ok(index_len) = runtime_state.score_len(KIRO_PROMPT_CACHE_INDEX_KEY).await else { + return; + }; + if index_len <= max_entries { + return; + } + + let Ok(all_members) = runtime_state + .score_range_by_min(KIRO_PROMPT_CACHE_INDEX_KEY, 0.0) + .await + else { + return; + }; + let trim_count = index_len.saturating_sub(max_entries); + if trim_count == 0 { + return; + } + + let trimmed_members = all_members.into_iter().take(trim_count).collect::>(); + if let Err(err) = runtime_state.kv_delete_many(&trimmed_members).await { + warn!( + event_name = "kiro_simulated_cache_kv_trim_failed", + log_type = "event", + error = ?err, + trim_count, + "failed to delete trimmed Kiro simulated cache KV entries" + ); + } + if let Err(err) = runtime_state + .score_remove_by_rank(KIRO_PROMPT_CACHE_INDEX_KEY, 0, trim_count as i64 - 1) + .await + { + warn!( + event_name = "kiro_simulated_cache_index_trim_failed", + log_type = "event", + error = ?err, + trim_count, + "failed to delete trimmed Kiro simulated cache index entries" + ); + } +} + +fn parse_kiro_prompt_cache_runtime_entry(value: &str) -> Option { + serde_json::from_str::(value) + .ok() + .filter(|entry| entry.token_count > 0 && entry.ttl_secs > 0) +} + +fn encode_kiro_prompt_cache_runtime_entry(entry: KiroPromptCacheRuntimeEntry) -> String { + serde_json::to_string(&entry).unwrap_or_else(|_| { + format!( + r#"{{"token_count":{},"ttl_secs":{}}}"#, + entry.token_count, entry.ttl_secs + ) + }) +} + +fn kiro_prompt_cache_runtime_key(credential_id: &str, fingerprint: &[u8; 32]) -> String { + let credential_hash: [u8; 32] = Sha256::digest(credential_id.as_bytes()).into(); + format!( + "kiro:prompt-cache:{}:{}", + hex_digest(&credential_hash), + hex_digest(fingerprint) + ) +} + +fn hex_digest(bytes: &[u8]) -> String { + let mut output = String::with_capacity(bytes.len() * 2); + for byte in bytes { + let _ = write!(&mut output, "{byte:02x}"); + } + output +} + pub(crate) fn build_kiro_prompt_cache_profile( request_body: &Value, total_input_tokens: u64, @@ -75,7 +312,8 @@ pub(crate) fn build_kiro_prompt_cache_profile( .and_then(Value::as_str) .unwrap_or_default(); let flattened = flatten_cacheable_blocks(request_body); - if flattened.iter().all(|block| block.breakpoint_ttl.is_none()) { + let automatic_ttl = extract_cache_ttl(request_body); + if automatic_ttl.is_none() && flattened.iter().all(|block| block.breakpoint_ttl.is_none()) { return None; } @@ -89,12 +327,15 @@ pub(crate) fn build_kiro_prompt_cache_profile( prefix_hasher.update(prelude_bytes); let mut cumulative_tokens = 0u64; - let mut active_ttl: Option = None; let mut breakpoints = Vec::new(); let mut seen_fingerprints = std::collections::BTreeSet::<[u8; 32]>::new(); - let mut lookback_candidates = Vec::new(); + let mut match_candidates = Vec::new(); + let last_block_index = flattened.len().saturating_sub(1); - for block in flattened { + for (block_index, mut block) in flattened.into_iter().enumerate() { + if block.breakpoint_ttl.is_none() && block_index == last_block_index { + block.breakpoint_ttl = automatic_ttl; + } cumulative_tokens = cumulative_tokens.saturating_add(block.tokens); let block_bytes = serde_json::to_vec(&block.value).unwrap_or_default(); let block_hash: [u8; 32] = Sha256::digest(block_bytes).into(); @@ -105,13 +346,6 @@ pub(crate) fn build_kiro_prompt_cache_profile( prefix_hasher.update(fingerprint); if let Some(ttl) = block.breakpoint_ttl { - push_lookback_breakpoints( - &mut breakpoints, - &mut seen_fingerprints, - &lookback_candidates, - ttl, - ); - active_ttl = Some(ttl); push_breakpoint( &mut breakpoints, &mut seen_fingerprints, @@ -120,18 +354,7 @@ pub(crate) fn build_kiro_prompt_cache_profile( ttl, ); } - if block.is_message_end { - if let Some(ttl) = active_ttl { - push_breakpoint( - &mut breakpoints, - &mut seen_fingerprints, - fingerprint, - cumulative_tokens, - ttl, - ); - } - } - push_prefix_candidate(&mut lookback_candidates, fingerprint, cumulative_tokens); + push_match_candidate(&mut match_candidates, fingerprint, cumulative_tokens); } let min_cacheable_tokens = minimum_cacheable_tokens_for_model(model); @@ -139,13 +362,54 @@ pub(crate) fn build_kiro_prompt_cache_profile( .into_iter() .filter(|breakpoint| breakpoint.cumulative_tokens >= min_cacheable_tokens) .collect::>(); - (!cacheable_breakpoints.is_empty()).then_some(KiroPromptCacheProfile { + if cacheable_breakpoints.is_empty() { + return None; + } + let match_candidates = build_lookback_match_candidates( + &match_candidates, + &cacheable_breakpoints, + min_cacheable_tokens, + ); + Some(KiroPromptCacheProfile { total_input_tokens, min_cacheable_tokens, breakpoints: cacheable_breakpoints, + match_candidates, }) } +fn build_lookback_match_candidates( + candidates: &[KiroPromptCacheCandidate], + breakpoints: &[KiroPromptCacheBreakpoint], + min_cacheable_tokens: u64, +) -> Vec { + let mut out = Vec::new(); + let mut seen_fingerprints = std::collections::BTreeSet::<[u8; 32]>::new(); + + for breakpoint in breakpoints { + let Some(index) = candidates + .iter() + .position(|candidate| candidate.fingerprint == breakpoint.fingerprint) + else { + continue; + }; + let start = index + .saturating_add(1) + .saturating_sub(PREFIX_LOOKBACK_WINDOW); + for candidate in &candidates[start..=index] { + if candidate.cumulative_tokens < min_cacheable_tokens + || candidate.cumulative_tokens > breakpoint.cumulative_tokens + || !seen_fingerprints.insert(candidate.fingerprint) + { + continue; + } + out.push(*candidate); + } + } + + out +} + pub(crate) fn kiro_simulated_cache_enabled_from_provider_config(config: Option<&Value>) -> bool { config .and_then(Value::as_object) @@ -269,35 +533,15 @@ fn push_breakpoint( } } -fn push_lookback_breakpoints( - breakpoints: &mut Vec, - seen_fingerprints: &mut std::collections::BTreeSet<[u8; 32]>, - candidates: &[PrefixCandidate], - ttl: Duration, -) { - for candidate in candidates { - push_breakpoint( - breakpoints, - seen_fingerprints, - candidate.fingerprint, - candidate.cumulative_tokens, - ttl, - ); - } -} - -fn push_prefix_candidate( - candidates: &mut Vec, +fn push_match_candidate( + candidates: &mut Vec, fingerprint: [u8; 32], cumulative_tokens: u64, ) { - candidates.push(PrefixCandidate { + candidates.push(KiroPromptCacheCandidate { fingerprint, cumulative_tokens, }); - if candidates.len() > PREFIX_LOOKBACK_LIMIT { - candidates.remove(0); - } } fn flatten_cacheable_blocks(request_body: &Value) -> Vec { @@ -316,7 +560,6 @@ fn flatten_cacheable_blocks(request_body: &Value) -> Vec { tokens: TOKENS_PER_TOOL, value, breakpoint_ttl, - is_message_end: false, }); } } @@ -337,7 +580,6 @@ fn flatten_cacheable_blocks(request_body: &Value) -> Vec { tokens: count_system_block_tokens(item), value, breakpoint_ttl, - is_message_end: false, }); } } @@ -351,7 +593,6 @@ fn flatten_cacheable_blocks(request_body: &Value) -> Vec { tokens: count_text_tokens(text), value, breakpoint_ttl: None, - is_message_end: false, }); } other => { @@ -364,7 +605,6 @@ fn flatten_cacheable_blocks(request_body: &Value) -> Vec { tokens: count_system_block_tokens(other), value, breakpoint_ttl: None, - is_message_end: false, }); } } @@ -376,11 +616,17 @@ fn flatten_cacheable_blocks(request_body: &Value) -> Vec { .get("role") .and_then(Value::as_str) .unwrap_or_default(); + let message_breakpoint_ttl = extract_cache_ttl(message); match message.get("content") { Some(Value::Array(items)) => { let last_block_index = items.len().saturating_sub(1); for (block_index, item) in items.iter().enumerate() { - let breakpoint_ttl = extract_cache_ttl(item); + let breakpoint_ttl = + extract_cache_ttl(item).or(if block_index == last_block_index { + message_breakpoint_ttl + } else { + None + }); let mut normalized = item.clone(); strip_cache_control(&mut normalized); let value = canonicalize_json(serde_json::json!({ @@ -394,7 +640,6 @@ fn flatten_cacheable_blocks(request_body: &Value) -> Vec { tokens: count_message_content_tokens(item), value, breakpoint_ttl, - is_message_end: block_index == last_block_index, }); } } @@ -409,8 +654,7 @@ fn flatten_cacheable_blocks(request_body: &Value) -> Vec { blocks.push(PendingBlock { tokens: count_text_tokens(text), value, - breakpoint_ttl: None, - is_message_end: true, + breakpoint_ttl: message_breakpoint_ttl, }); } Some(other) => { @@ -424,8 +668,7 @@ fn flatten_cacheable_blocks(request_body: &Value) -> Vec { blocks.push(PendingBlock { tokens: count_message_content_tokens(other), value, - breakpoint_ttl: None, - is_message_end: true, + breakpoint_ttl: message_breakpoint_ttl, }); } None => {} @@ -634,20 +877,16 @@ impl KiroPromptCacheTracker { }; let mut matched_tokens = 0; - for breakpoint in profile - .breakpoints - .iter() - .rev() - .take(PREFIX_LOOKBACK_LIMIT.saturating_add(1)) - { - let key = (credential_id.clone(), breakpoint.fingerprint); - let Some(entry) = entries.get(&key) else { + for candidate in profile.match_candidates.iter().rev() { + let key = (credential_id.clone(), candidate.fingerprint); + let Some(entry) = entries.get_mut(&key) else { continue; }; if entry.expires_at > now { + entry.expires_at = entry.expires_at.max(now + entry.ttl); matched_tokens = entry .token_count - .min(breakpoint.cumulative_tokens) + .min(candidate.cumulative_tokens) .min(profile.total_input_tokens); break; } @@ -664,6 +903,7 @@ impl KiroPromptCacheTracker { Some(existing) => { existing.token_count = existing.token_count.max(breakpoint.cumulative_tokens); existing.ttl = existing.ttl.max(breakpoint.ttl); + existing.expires_at = existing.expires_at.max(now + existing.ttl); } None => { self.evict_to_capacity(&mut entries); @@ -702,6 +942,7 @@ impl KiroPromptCacheTracker { #[cfg(test)] mod tests { use super::*; + use aether_runtime_state::MemoryRuntimeStateConfig; fn long_text(label: &str) -> String { format!("{} {}", label, "cacheable prompt chunk ".repeat(300)) @@ -756,7 +997,138 @@ mod tests { } #[test] - fn tracker_supports_prefix_hits_without_extending_expiry() { + fn profile_reads_top_level_automatic_cache_control() { + let request = serde_json::json!({ + "model": "claude-sonnet-4.6", + "cache_control": {"type": "ephemeral"}, + "messages": [{ + "role": "user", + "content": long_text("automatic cached turn") + }] + }); + + let profile = + build_kiro_prompt_cache_profile(&request, estimate_kiro_prompt_input_tokens(&request)) + .expect("top-level cache_control should create an automatic cache profile"); + let tracker = KiroPromptCacheTracker::default(); + let usage = tracker.compute_and_update("cred".to_string(), &profile); + + assert_eq!(profile.breakpoints.len(), 1); + assert!(usage.cache_creation_input_tokens > 0); + } + + #[test] + fn profile_does_not_create_message_end_breakpoints_from_explicit_cache_control() { + let request = serde_json::json!({ + "model": "claude-sonnet-4.6", + "system": [{ + "type": "text", + "text": long_text("explicit cached system"), + "cache_control": {"type": "ephemeral"} + }], + "messages": [{ + "role": "user", + "content": long_text("uncached later turn") + }] + }); + + let profile = + build_kiro_prompt_cache_profile(&request, estimate_kiro_prompt_input_tokens(&request)) + .expect("explicit cache_control should create a cache profile"); + + assert_eq!(profile.breakpoints.len(), 1); + assert!(profile.breakpoints[0].cumulative_tokens < profile.total_input_tokens); + } + + #[tokio::test] + async fn runtime_state_tracker_reads_cached_prefix_across_calls() { + let request = serde_json::json!({ + "model": "claude-sonnet-4.6", + "system": [{ + "type": "text", + "text": long_text("runtime shared system"), + "cache_control": {"type": "ephemeral"} + }], + "messages": [{"role": "user", "content": "reuse runtime cache"}] + }); + let profile = + build_kiro_prompt_cache_profile(&request, estimate_kiro_prompt_input_tokens(&request)) + .expect("cacheable request should create a cache profile"); + let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default()); + + let first = + compute_kiro_prompt_cache_usage(&runtime, "runtime-cred".to_string(), &profile).await; + let second = + compute_kiro_prompt_cache_usage(&runtime, "runtime-cred".to_string(), &profile).await; + + assert!(first.cache_creation_input_tokens > 0); + assert_eq!(first.cache_read_input_tokens, 0); + assert_eq!(second.cache_creation_input_tokens, 0); + assert!(second.cache_read_input_tokens > 0); + } + + #[tokio::test] + async fn runtime_state_tracker_trims_oldest_entries_to_capacity() { + let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default()); + let now_ms = current_unix_ms(); + let keys = [ + "kiro:prompt-cache:test-oldest".to_string(), + "kiro:prompt-cache:test-middle".to_string(), + "kiro:prompt-cache:test-newest".to_string(), + ]; + + for (index, key) in keys.iter().enumerate() { + runtime + .kv_set( + key, + encode_kiro_prompt_cache_runtime_entry(KiroPromptCacheRuntimeEntry { + token_count: 100 + index as u64, + ttl_secs: 120, + }), + Some(Duration::from_secs(120)), + ) + .await + .expect("cache entry should store"); + runtime + .score_set( + KIRO_PROMPT_CACHE_INDEX_KEY, + key, + now_ms.saturating_add(60_000 + index as u64 * 1_000) as f64, + ) + .await + .expect("cache index should store"); + } + + trim_kiro_prompt_cache_runtime_state(&runtime, 2).await; + + assert_eq!( + runtime + .kv_get(&keys[0]) + .await + .expect("oldest entry should read"), + None + ); + assert!(runtime + .kv_get(&keys[1]) + .await + .expect("middle entry should read") + .is_some()); + assert!(runtime + .kv_get(&keys[2]) + .await + .expect("newest entry should read") + .is_some()); + assert_eq!( + runtime + .score_range_by_min(KIRO_PROMPT_CACHE_INDEX_KEY, 0.0) + .await + .expect("cache index should read"), + vec![keys[1].clone(), keys[2].clone()] + ); + } + + #[test] + fn tracker_refreshes_cached_prefix_ttl_on_read() { let base = serde_json::json!({ "model": "claude-sonnet-4.6", "system": [{ @@ -799,13 +1171,13 @@ mod tests { ); assert!(hit.cache_read_input_tokens > 0); - let expired = tracker.compute_and_update_at( + let refreshed = tracker.compute_and_update_at( "cred".to_string(), &base_profile, start + Duration::from_secs(301), ); - assert!(expired.cache_creation_input_tokens > 0); - assert_eq!(expired.cache_read_input_tokens, 0); + assert_eq!(refreshed.cache_creation_input_tokens, 0); + assert!(refreshed.cache_read_input_tokens > 0); } #[test] @@ -867,6 +1239,146 @@ mod tests { assert!(hit.cache_creation_input_tokens > 0); } + #[test] + fn tracker_reads_cached_prefix_within_prompt_cache_lookback_window() { + let first = serde_json::json!({ + "model": "claude-sonnet-4.6", + "messages": [{ + "role": "user", + "content": [{ + "type": "text", + "text": long_text("shared first turn"), + "cache_control": {"type": "ephemeral"} + }] + }] + }); + let mut second_messages = vec![serde_json::json!({ + "role": "user", + "content": [{ + "type": "text", + "text": long_text("shared first turn") + }] + })]; + for index in 0..12 { + second_messages.push(serde_json::json!({ + "role": if index % 2 == 0 { "assistant" } else { "user" }, + "content": format!("intermediate turn {index}") + })); + } + second_messages.push(serde_json::json!({ + "role": "user", + "content": [{ + "type": "text", + "text": long_text("new tail turn"), + "cache_control": {"type": "ephemeral"} + }] + })); + let second = serde_json::json!({ + "model": "claude-sonnet-4.6", + "messages": second_messages + }); + let first_profile = + build_kiro_prompt_cache_profile(&first, estimate_kiro_prompt_input_tokens(&first)) + .expect("first request should be cacheable"); + let second_profile = + build_kiro_prompt_cache_profile(&second, estimate_kiro_prompt_input_tokens(&second)) + .expect("second request should be cacheable"); + let tracker = KiroPromptCacheTracker::default(); + let start = Instant::now(); + + let created = tracker.compute_and_update_at("cred".to_string(), &first_profile, start); + assert!(created.cache_creation_input_tokens > 0); + assert_eq!(created.cache_read_input_tokens, 0); + + let hit = tracker.compute_and_update_at( + "cred".to_string(), + &second_profile, + start + Duration::from_secs(60), + ); + assert!(hit.cache_read_input_tokens > 0); + assert!(hit.cache_creation_input_tokens > 0); + } + + #[test] + fn tracker_does_not_read_cached_prefix_outside_prompt_cache_lookback_window() { + let first = serde_json::json!({ + "model": "claude-sonnet-4.6", + "messages": [{ + "role": "user", + "content": [{ + "type": "text", + "text": long_text("shared first turn"), + "cache_control": {"type": "ephemeral"} + }] + }] + }); + let mut second_messages = vec![serde_json::json!({ + "role": "user", + "content": [{ + "type": "text", + "text": long_text("shared first turn") + }] + })]; + for index in 0..20 { + second_messages.push(serde_json::json!({ + "role": if index % 2 == 0 { "assistant" } else { "user" }, + "content": format!("intermediate turn {index}") + })); + } + second_messages.push(serde_json::json!({ + "role": "user", + "content": [{ + "type": "text", + "text": long_text("new tail turn"), + "cache_control": {"type": "ephemeral"} + }] + })); + let second = serde_json::json!({ + "model": "claude-sonnet-4.6", + "messages": second_messages + }); + let first_profile = + build_kiro_prompt_cache_profile(&first, estimate_kiro_prompt_input_tokens(&first)) + .expect("first request should be cacheable"); + let second_profile = + build_kiro_prompt_cache_profile(&second, estimate_kiro_prompt_input_tokens(&second)) + .expect("second request should be cacheable"); + let tracker = KiroPromptCacheTracker::default(); + let start = Instant::now(); + + let created = tracker.compute_and_update_at("cred".to_string(), &first_profile, start); + assert!(created.cache_creation_input_tokens > 0); + assert_eq!(created.cache_read_input_tokens, 0); + + let miss = tracker.compute_and_update_at( + "cred".to_string(), + &second_profile, + start + Duration::from_secs(60), + ); + assert!(miss.cache_creation_input_tokens > 0); + assert_eq!(miss.cache_read_input_tokens, 0); + } + + #[test] + fn profile_reads_message_level_cache_control() { + let request = serde_json::json!({ + "model": "claude-sonnet-4.6", + "messages": [{ + "role": "system", + "content": long_text("message level cached system"), + "cache_control": {"type": "ephemeral"} + }] + }); + + let profile = + build_kiro_prompt_cache_profile(&request, estimate_kiro_prompt_input_tokens(&request)) + .expect("message-level cache_control should create a cache profile"); + let tracker = KiroPromptCacheTracker::default(); + let usage = tracker.compute_and_update("cred".to_string(), &profile); + + assert!(usage.cache_creation_input_tokens > 0); + } + #[test] fn billed_input_tokens_subtracts_cache_usage() { assert_eq!( diff --git a/apps/aether-gateway/src/execution_runtime/kiro_web_search.rs b/apps/aether-gateway/src/execution_runtime/kiro_web_search.rs index a64e978c9..1ffa21955 100644 --- a/apps/aether-gateway/src/execution_runtime/kiro_web_search.rs +++ b/apps/aether-gateway/src/execution_runtime/kiro_web_search.rs @@ -15,8 +15,8 @@ use tracing::{debug, warn}; use uuid::Uuid; use crate::execution_runtime::kiro_cache::{ - billed_input_tokens, build_kiro_prompt_cache_profile, estimate_kiro_prompt_input_tokens, - kiro_prompt_cache_tracker, kiro_simulated_cache_enabled_from_provider_config, + billed_input_tokens, build_kiro_prompt_cache_profile, compute_kiro_prompt_cache_usage, + estimate_kiro_prompt_input_tokens, kiro_simulated_cache_enabled_from_provider_config, KiroPromptCacheProfile, KiroPromptCacheUsage, }; use crate::execution_runtime::ndjson::encode_stream_frame_ndjson; @@ -160,14 +160,17 @@ pub(crate) async fn maybe_execute_kiro_web_search_stream( let search_results = parse_mcp_search_results(&mcp_execution.result); let cache_usage = if kiro_simulated_cache_enabled(state, plan).await { - request - .cache_profile - .as_ref() - .map(|profile| { - kiro_prompt_cache_tracker() - .compute_and_update(kiro_cache_credential_id(plan), profile) - }) - .unwrap_or_default() + match request.cache_profile.as_ref() { + Some(profile) => { + compute_kiro_prompt_cache_usage( + state.runtime_state(), + kiro_cache_credential_id(plan), + profile, + ) + .await + } + None => KiroPromptCacheUsage::default(), + } } else { KiroPromptCacheUsage::default() }; diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index 0cc6ab491..e6bfc9dfa 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -65,7 +65,7 @@ use crate::execution_runtime::chatgpt_web_image::maybe_execute_chatgpt_web_image use crate::execution_runtime::grok::maybe_execute_grok_stream; use crate::execution_runtime::kiro_cache::{ billed_input_tokens as kiro_billed_input_tokens, build_kiro_prompt_cache_profile, - estimate_kiro_prompt_input_tokens, kiro_prompt_cache_tracker, + compute_kiro_prompt_cache_usage, estimate_kiro_prompt_input_tokens, kiro_simulated_cache_enabled_from_provider_config, kiro_simulated_cache_enabled_from_report_context, KiroPromptCacheUsage, KIRO_SIMULATED_CACHE_ENABLED_CONTEXT_FIELD, @@ -402,7 +402,8 @@ async fn seed_kiro_simulated_cache_enabled( } } -fn seed_kiro_report_context_prompt_cache_usage( +async fn seed_kiro_report_context_prompt_cache_usage( + state: &AppState, plan: &ExecutionPlan, report_context: &mut Option, ) { @@ -450,8 +451,12 @@ fn seed_kiro_report_context_prompt_cache_usage( return; }; - let cache_usage = kiro_prompt_cache_tracker() - .compute_and_update(kiro_stream_cache_credential_id(plan), &profile); + let cache_usage = compute_kiro_prompt_cache_usage( + state.runtime_state(), + kiro_stream_cache_credential_id(plan), + &profile, + ) + .await; if cache_usage.cache_creation_input_tokens == 0 && cache_usage.cache_read_input_tokens == 0 { return; } @@ -494,7 +499,8 @@ fn kiro_cache_usage_from_report_context(report_context: &Value) -> Option, summary: &mut Option, @@ -572,8 +578,12 @@ fn maybe_apply_kiro_prompt_cache_usage_to_stream_summary( return; }; - let cache_usage = kiro_prompt_cache_tracker() - .compute_and_update(kiro_stream_cache_credential_id(plan), &profile); + let cache_usage = compute_kiro_prompt_cache_usage( + state.runtime_state(), + kiro_stream_cache_credential_id(plan), + &profile, + ) + .await; if cache_usage.cache_creation_input_tokens == 0 && cache_usage.cache_read_input_tokens == 0 { return; } @@ -831,7 +841,6 @@ pub(crate) async fn execute_execution_runtime_stream( ) -> Result>, GatewayError> { let stream_started_at = Instant::now(); ensure_execution_request_candidate_slot(state, &mut plan, &mut report_context).await; - seed_kiro_report_context_input_tokens(&plan, &mut report_context); let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref()); let request_candidate_status_snapshot = snapshot_local_request_candidate_status(&plan, report_context.as_ref()); @@ -1981,10 +1990,12 @@ async fn execute_stream_from_frame_stream( }; let mut report_context = attach_provider_response_headers_to_report_context(report_context, &headers); - seed_kiro_report_context_input_tokens(&plan, &mut report_context); if status_code == 200 { seed_kiro_simulated_cache_enabled(state, &plan, &mut report_context).await; - seed_kiro_report_context_prompt_cache_usage(&plan, &mut report_context); + if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) { + seed_kiro_report_context_input_tokens(&plan, &mut report_context); + } + seed_kiro_report_context_prompt_cache_usage(state, &plan, &mut report_context).await; } let mut buffered_frames = VecDeque::new(); let mut stream_terminal_summary: Option = None; @@ -3357,7 +3368,13 @@ async fn execute_stream_from_frame_stream( "gateway skipped client stream flush after downstream disconnect" ); } - if let Some(normalizer) = private_stream_normalizer.as_mut() { + // Buffered stream state is partial after a terminal failure; normal + // finish paths may synthesize successful terminal events. + let should_finish_stream_rewriters = terminal_failure.is_none(); + if let Some(normalizer) = private_stream_normalizer + .as_mut() + .filter(|_| should_finish_stream_rewriters) + { match normalizer.finish() { Ok(normalized_chunk) if !normalized_chunk.is_empty() => { let provider_private_error_body_json = @@ -3391,13 +3408,12 @@ async fn execute_stream_from_frame_stream( error = ?err, "gateway failed to rewrite normalized private stream chunk during flush" ); - terminal_failure.get_or_insert_with(|| { - build_stream_failure_report( - "execution_runtime_stream_rewrite_flush_error", - format!("failed to rewrite normalized private stream chunk during flush: {err:?}"), - 502, - ) - }); + let failure = build_stream_failure_report( + "execution_runtime_stream_rewrite_flush_error", + format!("failed to rewrite normalized private stream chunk during flush: {err:?}"), + 502, + ); + terminal_failure.get_or_insert(failure); Vec::new() } } @@ -3472,7 +3488,7 @@ async fn execute_stream_from_frame_stream( } } } - if !downstream_dropped { + if !downstream_dropped && terminal_failure.is_none() { if let Some(rewriter) = local_stream_rewriter.as_mut() { match rewriter.finish() { Ok(flushed_chunk) if !flushed_chunk.is_empty() => { @@ -3700,10 +3716,12 @@ async fn execute_stream_from_frame_stream( } maybe_apply_kiro_prompt_cache_usage_to_stream_summary( + &state_for_report, &plan_for_report, report_context_owned.as_ref(), &mut stream_terminal_summary, - ); + ) + .await; let requires_observed_terminal_event = stream_requires_observed_terminal_event( plan_for_report.provider_api_format.as_str(), stream_usage_report_context.as_ref(), @@ -3898,8 +3916,9 @@ mod tests { use std::time::{Duration, Instant}; use aether_contracts::{ - ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTimeouts, RequestBody, - StandardizedUsage, + ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, + ExecutionStreamTerminalSummary, ExecutionTimeouts, RequestBody, StandardizedUsage, + StreamFrame, StreamFramePayload, StreamFrameType, }; use aether_data::repository::candidates::InMemoryRequestCandidateRepository; use aether_data::repository::usage::InMemoryUsageReadRepository; @@ -3946,6 +3965,10 @@ mod tests { .with_execution_runtime_candidate(true) } + fn test_state() -> AppState { + AppState::new().expect("gateway state should build") + } + #[test] fn detects_client_visible_sse_terminal_events() { assert!(stream_chunk_contains_sse_done(b"data: [DONE]\n\n")); @@ -3996,6 +4019,12 @@ mod tests { out } + fn ndjson_frame(frame: StreamFrame) -> Bytes { + let mut bytes = serde_json::to_vec(&frame).expect("stream frame should serialize"); + bytes.push(b'\n'); + Bytes::from(bytes) + } + #[test] fn merge_stream_terminal_summary_prefers_more_complete_observed_usage() { let mut runtime_usage = StandardizedUsage::new(); @@ -4124,8 +4153,8 @@ mod tests { )); } - #[test] - fn kiro_stream_summary_applies_prompt_cache_usage_from_original_request() { + #[tokio::test] + async fn kiro_stream_summary_applies_prompt_cache_usage_from_original_request() { let request_body = json!({ "model": "claude-opus-4-7", "system": [ @@ -4173,6 +4202,7 @@ mod tests { transport_profile: None, timeouts: None, }; + let state = test_state(); let mut first_summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { @@ -4183,10 +4213,12 @@ mod tests { ..ExecutionStreamTerminalSummary::default() }); maybe_apply_kiro_prompt_cache_usage_to_stream_summary( + &state, &plan, Some(&report_context), &mut first_summary, - ); + ) + .await; let first_usage = first_summary .as_ref() .and_then(|summary| summary.standardized_usage.as_ref()) @@ -4203,10 +4235,12 @@ mod tests { ..ExecutionStreamTerminalSummary::default() }); maybe_apply_kiro_prompt_cache_usage_to_stream_summary( + &state, &plan, Some(&report_context), &mut second_summary, - ); + ) + .await; let second_usage = second_summary .as_ref() .and_then(|summary| summary.standardized_usage.as_ref()) @@ -4217,8 +4251,126 @@ mod tests { assert_eq!(second_usage.output_tokens, 19); } - #[test] - fn kiro_stream_summary_seeds_input_tokens_without_cache_control() { + #[tokio::test] + async fn kiro_stream_summary_reads_cached_prefix_within_prompt_cache_lookback_window() { + let first_request_body = json!({ + "model": "claude-sonnet-4.6", + "messages": [{ + "role": "user", + "content": [{ + "type": "text", + "text": "shared first turn ".repeat(600), + "cache_control": {"type": "ephemeral"} + }] + }] + }); + let mut second_messages = vec![json!({ + "role": "user", + "content": [{ + "type": "text", + "text": "shared first turn ".repeat(600) + }] + })]; + for index in 0..12 { + second_messages.push(json!({ + "role": if index % 2 == 0 { "assistant" } else { "user" }, + "content": format!("intermediate stream turn {index}") + })); + } + second_messages.push(json!({ + "role": "user", + "content": [{ + "type": "text", + "text": "new tail turn ".repeat(600), + "cache_control": {"type": "ephemeral"} + }] + })); + let second_request_body = json!({ + "model": "claude-sonnet-4.6", + "messages": second_messages + }); + let plan = ExecutionPlan { + request_id: "req-kiro-cache-stream-long-tail".into(), + candidate_id: Some("cand-kiro-cache-stream-long-tail".into()), + provider_name: Some("Kiro".into()), + provider_id: "provider-kiro-cache-stream-long-tail".into(), + endpoint_id: "endpoint-kiro-cache-stream-long-tail".into(), + key_id: "key-kiro-cache-stream-long-tail".into(), + method: "POST".into(), + url: "https://q.us-east-1.amazonaws.com/generateAssistantResponse?beta=true".into(), + headers: BTreeMap::new(), + content_type: Some("application/json".into()), + content_encoding: None, + body: RequestBody::from_json(json!({"conversationState": {}})), + stream: true, + client_api_format: "claude:messages".into(), + provider_api_format: "claude:messages".into(), + model_name: Some("claude-sonnet-4.6".into()), + proxy: None, + transport_profile: None, + timeouts: None, + }; + let first_report_context = json!({ + "original_request_body": first_request_body, + "kiro_simulated_cache_enabled": true, + }); + let second_report_context = json!({ + "original_request_body": second_request_body, + "kiro_simulated_cache_enabled": true, + }); + let state = test_state(); + + let mut first_summary = Some(ExecutionStreamTerminalSummary { + standardized_usage: Some(StandardizedUsage { + input_tokens: 4_000, + output_tokens: 17, + ..StandardizedUsage::new() + }), + ..ExecutionStreamTerminalSummary::default() + }); + maybe_apply_kiro_prompt_cache_usage_to_stream_summary( + &state, + &plan, + Some(&first_report_context), + &mut first_summary, + ) + .await; + let first_usage = first_summary + .as_ref() + .and_then(|summary| summary.standardized_usage.as_ref()) + .expect("first usage should exist"); + assert!(first_usage.cache_creation_tokens > 0); + assert_eq!(first_usage.cache_read_tokens, 0); + + let mut second_summary = Some(ExecutionStreamTerminalSummary { + standardized_usage: Some(StandardizedUsage { + input_tokens: 8_000, + output_tokens: 19, + ..StandardizedUsage::new() + }), + ..ExecutionStreamTerminalSummary::default() + }); + maybe_apply_kiro_prompt_cache_usage_to_stream_summary( + &state, + &plan, + Some(&second_report_context), + &mut second_summary, + ) + .await; + let second_usage = second_summary + .as_ref() + .and_then(|summary| summary.standardized_usage.as_ref()) + .expect("second usage should exist"); + assert!( + second_usage.cache_read_tokens > 0, + "stream summary should reuse the far earlier cached prefix" + ); + assert!(second_usage.cache_creation_tokens > 0); + assert_eq!(second_usage.output_tokens, 19); + } + + #[tokio::test] + async fn kiro_stream_summary_seeds_input_tokens_without_cache_control() { let request_body = json!({ "model": "claude-opus-4-7", "system": [ @@ -4264,6 +4416,7 @@ mod tests { transport_profile: None, timeouts: None, }; + let state = test_state(); let mut summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { @@ -4275,10 +4428,12 @@ mod tests { }); maybe_apply_kiro_prompt_cache_usage_to_stream_summary( + &state, &plan, Some(&report_context), &mut summary, - ); + ) + .await; let usage = summary .as_ref() @@ -4291,8 +4446,8 @@ mod tests { assert_eq!(usage.output_tokens, 13); } - #[test] - fn kiro_stream_summary_bills_existing_cache_usage_when_input_is_zero() { + #[tokio::test] + async fn kiro_stream_summary_bills_existing_cache_usage_when_input_is_zero() { let request_body = json!({ "model": "claude-opus-4-7", "system": [ @@ -4338,6 +4493,7 @@ mod tests { transport_profile: None, timeouts: None, }; + let state = test_state(); let mut summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { @@ -4350,10 +4506,12 @@ mod tests { }); maybe_apply_kiro_prompt_cache_usage_to_stream_summary( + &state, &plan, Some(&report_context), &mut summary, - ); + ) + .await; let usage = summary .as_ref() @@ -4366,8 +4524,8 @@ mod tests { assert_eq!(usage.output_tokens, 23); } - #[test] - fn kiro_stream_summary_clears_cache_usage_when_simulated_cache_disabled() { + #[tokio::test] + async fn kiro_stream_summary_clears_cache_usage_when_simulated_cache_disabled() { let request_body = json!({ "model": "claude-opus-4-7", "system": [ @@ -4414,6 +4572,7 @@ mod tests { transport_profile: None, timeouts: None, }; + let state = test_state(); let mut summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { @@ -4427,10 +4586,12 @@ mod tests { }); maybe_apply_kiro_prompt_cache_usage_to_stream_summary( + &state, &plan, Some(&report_context), &mut summary, - ); + ) + .await; let usage = summary .as_ref() @@ -4443,8 +4604,8 @@ mod tests { assert_eq!(usage.output_tokens, 23); } - #[test] - fn kiro_stream_summary_does_not_subtract_cache_from_already_billed_input() { + #[tokio::test] + async fn kiro_stream_summary_does_not_subtract_cache_from_already_billed_input() { let request_body = json!({ "model": "claude-opus-4-7", "messages": [ @@ -4492,6 +4653,7 @@ mod tests { transport_profile: None, timeouts: None, }; + let state = test_state(); let mut summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { @@ -4505,10 +4667,12 @@ mod tests { }); maybe_apply_kiro_prompt_cache_usage_to_stream_summary( + &state, &plan, Some(&report_context), &mut summary, - ); + ) + .await; let usage = summary .as_ref() @@ -4628,9 +4792,11 @@ mod tests { "original_request_body": request_body, "kiro_simulated_cache_enabled": true, })); + let state = AppState::new().expect("gateway state should build"); super::seed_kiro_report_context_input_tokens(&plan, &mut report_context); - super::seed_kiro_report_context_prompt_cache_usage(&plan, &mut report_context); + super::seed_kiro_report_context_prompt_cache_usage(&state, &plan, &mut report_context) + .await; let context = report_context.as_ref().expect("context should exist"); assert!(context @@ -4697,9 +4863,11 @@ mod tests { let mut report_context = Some(json!({ "original_request_body": request_body, })); + let state = AppState::new().expect("gateway state should build"); super::seed_kiro_report_context_input_tokens(&plan, &mut report_context); - super::seed_kiro_report_context_prompt_cache_usage(&plan, &mut report_context); + super::seed_kiro_report_context_prompt_cache_usage(&state, &plan, &mut report_context) + .await; let context = report_context.as_ref().expect("context should exist"); assert!(context @@ -5017,6 +5185,154 @@ mod tests { assert_eq!(first.as_ref(), b": aether-keepalive\n\n"); } + #[tokio::test] + async fn execute_stream_from_frame_stream_does_not_finalize_rewritten_tool_call_after_midstream_error( + ) { + let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); + let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); + let state = AppState::new() + .expect("app state should build") + .with_data_state_for_tests( + crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests( + Arc::clone(&request_candidate_repository), + Arc::clone(&usage_repository), + ), + ) + .with_usage_runtime_for_tests(UsageRuntimeConfig { + enabled: true, + ..UsageRuntimeConfig::default() + }); + let plan = ExecutionPlan { + request_id: "req-responses-tool-midstream-error".into(), + candidate_id: Some("cand-responses-tool-midstream-error".into()), + provider_name: Some("openai".into()), + provider_id: "provider-openai-responses".into(), + endpoint_id: "endpoint-openai-responses".into(), + key_id: "key-openai-responses".into(), + method: "POST".into(), + url: "https://api.openai.com/v1/responses".into(), + headers: BTreeMap::from([ + ("content-type".into(), "application/json".into()), + ("accept".into(), "text/event-stream".into()), + ]), + content_type: Some("application/json".into()), + content_encoding: None, + body: RequestBody::from_json(json!({ + "model": "gpt-5.5", + "input": [], + "stream": true + })), + stream: true, + client_api_format: "claude:messages".into(), + provider_api_format: "openai:responses".into(), + model_name: Some("gpt-5.5".into()), + proxy: None, + transport_profile: None, + timeouts: None, + }; + let upstream_chunk = concat!( + "event: response.created\n", + "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_midstream_error\",\"model\":\"gpt-5.5\",\"status\":\"in_progress\"}}\n\n", + "event: response.output_item.added\n", + "data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"id\":\"fc_1\",\"call_id\":\"call_1\",\"name\":\"lookup\",\"arguments\":\"\",\"status\":\"in_progress\"}}\n\n", + "event: response.function_call_arguments.delta\n", + "data: {\"type\":\"response.function_call_arguments.delta\",\"output_index\":0,\"item_id\":\"fc_1\",\"call_id\":\"call_1\",\"delta\":\"{\\\"query\\\":\\\"abc\"}\n\n" + ); + let frame_stream = stream! { + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Headers, + payload: StreamFramePayload::Headers { + status_code: 200, + headers: BTreeMap::from([( + "content-type".to_string(), + "text/event-stream".to_string(), + )]), + }, + })); + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Data, + payload: StreamFramePayload::Data { + chunk_b64: None, + text: Some(upstream_chunk.to_string()), + }, + })); + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Error, + payload: StreamFramePayload::Error { + error: ExecutionError { + kind: ExecutionErrorKind::Internal, + phase: ExecutionPhase::StreamRead, + message: "error reading a body from connection: stream error received: unexpected internal error encountered".to_string(), + upstream_status: Some(200), + retryable: false, + failover_recommended: false, + }, + }, + })); + } + .boxed(); + + let response = execute_stream_from_frame_stream( + &state, + plan, + "trace-responses-tool-midstream-error", + &test_decision(), + "openai_responses_stream", + Some("openai_responses_stream_success".to_string()), + Some(json!({ + "request_id": "req-responses-tool-midstream-error", + "candidate_id": "cand-responses-tool-midstream-error", + "candidate_index": 0, + "retry_index": 0, + "provider_api_format": "openai:responses", + "client_api_format": "claude:messages", + "needs_conversion": true, + })), + crate::clock::current_unix_ms(), + Instant::now(), + frame_stream, + None, + ) + .await + .expect("execution should succeed") + .expect("execution should return a client response"); + + let body = to_bytes(response.into_body(), usize::MAX) + .await + .expect("response body should read"); + let body_text = String::from_utf8(body.to_vec()).expect("body should be utf8"); + assert!(body_text.contains("event: content_block_start")); + assert!(body_text.contains("event: content_block_delta")); + assert!(body_text.contains("\"type\":\"tool_use\"")); + assert!(!body_text.contains("event: content_block_stop")); + assert!(!body_text.contains("event: message_delta")); + assert!(!body_text.contains("event: message_stop")); + assert!(!body_text.contains("\"stop_reason\":\"tool_use\"")); + assert!(body_text.contains("\"error\"")); + assert!(body_text.contains("unexpected internal error encountered")); + assert!(body_text.contains("data: [DONE]")); + + let candidates = tokio::time::timeout(Duration::from_secs(1), async { + loop { + let candidates = request_candidate_repository + .list_by_request_id("req-responses-tool-midstream-error") + .await + .expect("request candidates should read"); + if candidates + .first() + .is_some_and(|candidate| candidate.status == RequestCandidateStatus::Failed) + { + break candidates; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("candidate should be marked failed"); + assert_eq!(candidates[0].status_code, Some(200)); + assert_eq!(candidates[0].error_type.as_deref(), Some("internal")); + } + #[tokio::test] async fn openai_image_stream_ignores_plan_total_timeout() { let state = AppState::new().expect("app state should build"); diff --git a/apps/aether-gateway/src/execution_runtime/sync/execution.rs b/apps/aether-gateway/src/execution_runtime/sync/execution.rs index 26e14d9c1..8666e9aea 100644 --- a/apps/aether-gateway/src/execution_runtime/sync/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/sync/execution.rs @@ -41,6 +41,12 @@ use crate::clock::current_unix_ms as current_request_candidate_unix_ms; use crate::control::GatewayControlDecision; use crate::execution_runtime::chatgpt_web_image::maybe_execute_chatgpt_web_image_sync; use crate::execution_runtime::grok::maybe_execute_grok_sync; +use crate::execution_runtime::kiro_cache::{ + build_kiro_prompt_cache_profile, compute_kiro_prompt_cache_usage, + estimate_kiro_prompt_input_tokens, kiro_simulated_cache_enabled_from_provider_config, + kiro_simulated_cache_enabled_from_report_context, KiroPromptCacheUsage, + KIRO_SIMULATED_CACHE_ENABLED_CONTEXT_FIELD, +}; use crate::execution_runtime::oauth_retry::refresh_oauth_plan_auth_for_retry; #[cfg(test)] use crate::execution_runtime::remote_compat::post_sync_plan_to_remote_execution_runtime; @@ -355,6 +361,180 @@ fn build_sync_report_payload( } } +fn seed_kiro_sync_report_context_input_tokens( + plan: &ExecutionPlan, + report_context: &mut Option, +) { + if !plan + .provider_name + .as_deref() + .is_some_and(|provider_name| provider_name.eq_ignore_ascii_case("Kiro")) + { + return; + } + + let Some(context) = report_context.as_mut().and_then(Value::as_object_mut) else { + return; + }; + if context + .get("input_tokens") + .and_then(Value::as_u64) + .is_some_and(|input_tokens| input_tokens > 0) + { + return; + } + + let Some(original_request_body) = context.get("original_request_body").cloned() else { + return; + }; + let estimated_input_tokens = estimate_kiro_prompt_input_tokens(&original_request_body); + context.insert( + "input_tokens".to_string(), + Value::from(estimated_input_tokens), + ); +} + +async fn seed_kiro_sync_simulated_cache_enabled( + state: &AppState, + plan: &ExecutionPlan, + report_context: &mut Option, +) { + if !plan + .provider_name + .as_deref() + .is_some_and(|provider_name| provider_name.eq_ignore_ascii_case("Kiro")) + { + return; + } + + let enabled = match state + .read_provider_catalog_providers_by_ids(std::slice::from_ref(&plan.provider_id)) + .await + { + Ok(providers) => providers + .iter() + .find(|provider| provider.id == plan.provider_id) + .filter(|provider| provider.provider_type.eq_ignore_ascii_case("kiro")) + .is_some_and(|provider| { + kiro_simulated_cache_enabled_from_provider_config(provider.config.as_ref()) + }), + Err(err) => { + warn!( + event_name = "kiro_simulated_cache_config_read_failed", + log_type = "event", + request_id = %plan.request_id, + provider_id = %plan.provider_id, + error = ?err, + "failed to read Kiro simulated cache provider config; defaulting disabled" + ); + false + } + }; + + let Some(context) = report_context.as_mut().and_then(Value::as_object_mut) else { + return; + }; + if enabled { + context.insert( + KIRO_SIMULATED_CACHE_ENABLED_CONTEXT_FIELD.to_string(), + Value::Bool(true), + ); + } else { + context.remove(KIRO_SIMULATED_CACHE_ENABLED_CONTEXT_FIELD); + } +} + +async fn seed_kiro_sync_report_context_prompt_cache_usage( + state: &AppState, + plan: &ExecutionPlan, + report_context: &mut Option, +) { + if !plan + .provider_name + .as_deref() + .is_some_and(|provider_name| provider_name.eq_ignore_ascii_case("Kiro")) + { + return; + } + + let simulated_cache_enabled = + kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()); + let Some(context) = report_context.as_mut().and_then(Value::as_object_mut) else { + return; + }; + if context + .get("kiro_web_search_mcp") + .and_then(Value::as_bool) + .unwrap_or(false) + { + return; + } + if !simulated_cache_enabled { + return; + } + if kiro_cache_usage_from_context_object(context).is_some() { + return; + } + + let Some(original_request_body) = context.get("original_request_body").cloned() else { + return; + }; + let input_tokens = context + .get("input_tokens") + .and_then(Value::as_u64) + .filter(|value| *value > 0) + .unwrap_or_else(|| { + let estimated = estimate_kiro_prompt_input_tokens(&original_request_body); + context.insert("input_tokens".to_string(), Value::from(estimated)); + estimated + }); + let Some(profile) = build_kiro_prompt_cache_profile(&original_request_body, input_tokens) + else { + return; + }; + + let cache_usage = compute_kiro_prompt_cache_usage( + state.runtime_state(), + kiro_sync_cache_credential_id(plan), + &profile, + ) + .await; + if cache_usage.cache_creation_input_tokens == 0 && cache_usage.cache_read_input_tokens == 0 { + return; + } + context.insert( + "cache_creation_input_tokens".to_string(), + Value::from(cache_usage.cache_creation_input_tokens), + ); + context.insert( + "cache_read_input_tokens".to_string(), + Value::from(cache_usage.cache_read_input_tokens), + ); +} + +fn kiro_sync_cache_credential_id(plan: &ExecutionPlan) -> String { + format!("{}:{}:{}", plan.provider_id, plan.endpoint_id, plan.key_id) +} + +fn kiro_cache_usage_from_context_object( + context: &serde_json::Map, +) -> Option { + let cache_creation_input_tokens = context + .get("cache_creation_input_tokens") + .and_then(Value::as_u64) + .unwrap_or(0); + let cache_read_input_tokens = context + .get("cache_read_input_tokens") + .and_then(Value::as_u64) + .unwrap_or(0); + (cache_creation_input_tokens > 0 || cache_read_input_tokens > 0).then_some( + KiroPromptCacheUsage { + cache_creation_input_tokens, + cache_read_input_tokens, + }, + ) +} + fn invalid_gemini_provider_success_message( plan: &ExecutionPlan, report_context: Option<&Value>, @@ -1182,12 +1362,18 @@ fn build_json_whitespace_heartbeat_stream( } } -pub(crate) fn build_openai_image_sync_json_whitespace_heartbeat_stream( +pub(crate) fn build_sync_json_whitespace_heartbeat_stream( rx: mpsc::Receiver>, ) -> impl futures_util::Stream> + Send + 'static { build_json_whitespace_heartbeat_stream(rx, OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_INTERVAL, None) } +pub(crate) fn build_openai_image_sync_json_whitespace_heartbeat_stream( + rx: mpsc::Receiver>, +) -> impl futures_util::Stream> + Send + 'static { + build_sync_json_whitespace_heartbeat_stream(rx) +} + async fn openai_image_sync_json_heartbeat_final_bytes( result: Result>, GatewayError>, ) -> Vec { @@ -1934,8 +2120,15 @@ async fn execute_execution_runtime_sync_impl( } let status_code = result.status_code; let has_body_bytes = body_base64.is_some(); - let report_context = + let mut report_context = attach_provider_response_headers_to_report_context(report_context, &headers); + if (200..300).contains(&status_code) { + seed_kiro_sync_simulated_cache_enabled(state, &plan, &mut report_context).await; + if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) { + seed_kiro_sync_report_context_input_tokens(&plan, &mut report_context); + } + seed_kiro_sync_report_context_prompt_cache_usage(state, &plan, &mut report_context).await; + } let mut client_headers = headers.clone(); apply_endpoint_response_header_rules(state, &plan, &mut client_headers, body_json.as_ref()) .await?; @@ -2504,6 +2697,45 @@ mod tests { plan } + fn test_kiro_sync_plan() -> ExecutionPlan { + ExecutionPlan { + request_id: "req-kiro-sync-cache-1".to_string(), + candidate_id: Some("candidate-kiro-sync-cache-1".to_string()), + provider_name: Some("Kiro".to_string()), + provider_id: "provider-kiro-sync-1".to_string(), + endpoint_id: "endpoint-kiro-sync-1".to_string(), + key_id: "key-kiro-sync-1".to_string(), + method: "POST".to_string(), + url: "https://kiro.example/generateAssistantResponse".to_string(), + headers: BTreeMap::new(), + content_type: Some("application/json".to_string()), + content_encoding: None, + body: aether_contracts::RequestBody::from_json(json!({ + "model": "claude-sonnet-4", + "messages": [{"role": "user", "content": "hello kiro"}], + })), + stream: false, + client_api_format: "claude:messages".to_string(), + provider_api_format: "claude:messages".to_string(), + model_name: Some("claude-sonnet-4".to_string()), + proxy: None, + transport_profile: None, + timeouts: None, + } + } + + fn test_kiro_sync_cacheable_request_body() -> serde_json::Value { + json!({ + "model": "claude-sonnet-4", + "system": [{ + "type": "text", + "text": format!("sync cacheable prompt {}", "cacheable prompt chunk ".repeat(300)), + "cache_control": {"type": "ephemeral"} + }], + "messages": [{"role": "user", "content": "reuse this Kiro prompt"}] + }) + } + #[test] fn invalid_gemini_provider_success_uses_plan_format_when_context_is_missing() { let plan = test_gemini_chat_plan(); @@ -2659,6 +2891,68 @@ mod tests { ); } + #[test] + fn kiro_sync_report_context_seeds_input_tokens_from_original_request_body() { + let plan = test_kiro_sync_plan(); + let mut report_context = Some(json!({ + "original_request_body": test_kiro_sync_cacheable_request_body(), + })); + + seed_kiro_sync_report_context_input_tokens(&plan, &mut report_context); + + assert!(report_context + .as_ref() + .and_then(|value| value.get("input_tokens")) + .and_then(Value::as_u64) + .is_some_and(|tokens| tokens > 0)); + } + + #[tokio::test] + async fn kiro_sync_report_context_applies_prompt_cache_usage_from_tracker() { + let state = AppState::new().expect("gateway state should build"); + let plan = test_kiro_sync_plan(); + + let mut first_report_context = Some(json!({ + "original_request_body": test_kiro_sync_cacheable_request_body(), + "kiro_simulated_cache_enabled": true, + })); + seed_kiro_sync_report_context_input_tokens(&plan, &mut first_report_context); + seed_kiro_sync_report_context_prompt_cache_usage(&state, &plan, &mut first_report_context) + .await; + let first_creation = first_report_context + .as_ref() + .and_then(|value| value.get("cache_creation_input_tokens")) + .and_then(Value::as_u64) + .unwrap_or_default(); + let first_read = first_report_context + .as_ref() + .and_then(|value| value.get("cache_read_input_tokens")) + .and_then(Value::as_u64) + .unwrap_or_default(); + assert!(first_creation > 0); + assert_eq!(first_read, 0); + + let mut second_report_context = Some(json!({ + "original_request_body": test_kiro_sync_cacheable_request_body(), + "kiro_simulated_cache_enabled": true, + })); + seed_kiro_sync_report_context_input_tokens(&plan, &mut second_report_context); + seed_kiro_sync_report_context_prompt_cache_usage(&state, &plan, &mut second_report_context) + .await; + let second_creation = second_report_context + .as_ref() + .and_then(|value| value.get("cache_creation_input_tokens")) + .and_then(Value::as_u64) + .unwrap_or_default(); + let second_read = second_report_context + .as_ref() + .and_then(|value| value.get("cache_read_input_tokens")) + .and_then(Value::as_u64) + .unwrap_or_default(); + assert_eq!(second_creation, 0); + assert!(second_read > 0); + } + #[tokio::test] async fn json_whitespace_heartbeat_stream_prefixes_final_json() { let (tx, rx) = mpsc::channel::>(1); diff --git a/apps/aether-gateway/src/execution_runtime/sync/mod.rs b/apps/aether-gateway/src/execution_runtime/sync/mod.rs index 954650e92..72a4ca5e4 100644 --- a/apps/aether-gateway/src/execution_runtime/sync/mod.rs +++ b/apps/aether-gateway/src/execution_runtime/sync/mod.rs @@ -1,7 +1,8 @@ mod execution; pub(crate) use execution::{ - build_openai_image_sync_json_whitespace_heartbeat_stream, execute_execution_runtime_sync, + build_openai_image_sync_json_whitespace_heartbeat_stream, + build_sync_json_whitespace_heartbeat_stream, execute_execution_runtime_sync, }; #[allow(unused_imports)] diff --git a/apps/aether-gateway/src/executor/orchestration.rs b/apps/aether-gateway/src/executor/orchestration.rs index 6fede309a..f4a894b20 100644 --- a/apps/aether-gateway/src/executor/orchestration.rs +++ b/apps/aether-gateway/src/executor/orchestration.rs @@ -9,6 +9,7 @@ use serde_json::{json, Value}; use tokio::sync::mpsc; use crate::ai_serving::api::{ + build_core_error_body_for_client_format, build_local_gemini_files_stream_attempt_source_for_kind, build_local_gemini_files_sync_attempt_source_for_kind, build_local_image_stream_attempt_source_for_kind, @@ -29,17 +30,18 @@ use crate::ai_serving::api::{ resolve_gemini_sync_spec, resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec, set_local_openai_chat_execution_exhausted_diagnostic, set_local_openai_image_execution_exhausted_diagnostic, AiStreamAttempt, AiSyncAttempt, - LocalStandardSpec, EXECUTION_RUNTIME_STREAM_DECISION_ACTION, + LocalCoreSyncErrorKind, LocalStandardSpec, EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION, }; use crate::ai_serving::LocalExecutionAttemptSource; use crate::api::response::{ attach_control_metadata_headers, build_client_response_from_parts_with_mutator, }; -use crate::constants::EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS; +use crate::constants::{CONTROL_CANDIDATE_ID_HEADER, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS}; use crate::control::GatewayControlDecision; use crate::execution_runtime::sync::{ - build_openai_image_sync_json_whitespace_heartbeat_stream, execute_execution_runtime_sync, + build_openai_image_sync_json_whitespace_heartbeat_stream, + build_sync_json_whitespace_heartbeat_stream, execute_execution_runtime_sync, }; use crate::executor::candidate_loop::{ execute_stream_attempt_source, execute_sync_attempt_source, execute_sync_plan_and_reports, @@ -47,15 +49,19 @@ use crate::executor::candidate_loop::{ }; use crate::executor::{ build_local_execution_exhaustion, record_failed_usage_for_exhausted_request, - LocalExecutionRequestOutcome, + LocalExecutionExhaustion, LocalExecutionRequestOutcome, }; use crate::handlers::shared::system_config_bool; use crate::{AiExecutionDecision, AppState, GatewayError}; const ENABLE_OPENAI_IMAGE_SYNC_HEARTBEAT_CONFIG_KEY: &str = "enable_openai_image_sync_heartbeat"; +const ENABLE_STANDARD_TEXT_SYNC_HEARTBEAT_CONFIG_KEY: &str = "enable_standard_text_sync_heartbeat"; const OPENAI_IMAGE_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS: u16 = 502; const OPENAI_IMAGE_SYNC_HEARTBEAT_EXHAUSTED_STATUS: u16 = 503; const OPENAI_IMAGE_SYNC_HEARTBEAT_ERROR_MESSAGE_LIMIT: usize = 4096; +const STANDARD_TEXT_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS: u16 = 502; +const STANDARD_TEXT_SYNC_HEARTBEAT_EXHAUSTED_STATUS: u16 = 503; +const STANDARD_TEXT_SYNC_HEARTBEAT_ERROR_MESSAGE_LIMIT: usize = 4096; pub(crate) async fn maybe_execute_sync_local_path( state: &AppState, @@ -95,6 +101,65 @@ pub(crate) async fn maybe_execute_sync_via_local_decision( return Ok(LocalExecutionRequestOutcome::NoPath); }; + if standard_text_sync_heartbeat_should_wrap(state, plan_kind).await { + let parts_for_task = parts.clone(); + let body_json_for_task = body_json.clone(); + return Ok(LocalExecutionRequestOutcome::responded( + build_standard_text_sync_heartbeat_shell_response( + state.clone(), + parts_for_task, + trace_id.to_string(), + decision.clone(), + plan_kind.to_string(), + move |state, parts, trace_id, decision, plan_kind, started_at| async move { + let Some((attempt_source, candidate_count)) = + build_local_openai_chat_sync_attempt_source_for_kind( + &state, + &parts, + trace_id.as_str(), + &decision, + &body_json_for_task, + plan_kind.as_str(), + ) + .await? + else { + return Ok(LocalExecutionRequestOutcome::NoPath); + }; + + let outcome = execute_sync_attempt_source::( + &state, + &parts, + trace_id.as_str(), + &decision, + plan_kind.as_str(), + attempt_source, + ) + .await?; + match outcome { + LocalExecutionRequestOutcome::Exhausted(exhaustion) => { + set_local_openai_chat_execution_exhausted_diagnostic( + &state, + trace_id.as_str(), + &decision, + plan_kind.as_str(), + &body_json_for_task, + candidate_count, + ); + record_standard_text_sync_heartbeat_exhaustion( + &state, + exhaustion, + &started_at, + ) + .await; + Ok(LocalExecutionRequestOutcome::NoPath) + } + outcome => Ok(outcome), + } + }, + )?, + )); + } + let outcome = execute_sync_attempt_source::( state, parts, @@ -176,6 +241,57 @@ pub(crate) async fn maybe_execute_sync_via_local_openai_responses_decision( return Ok(LocalExecutionRequestOutcome::NoPath); }; + if standard_text_sync_heartbeat_should_wrap(state, plan_kind).await { + let parts_for_task = parts.clone(); + let body_json_for_task = body_json.clone(); + return Ok(LocalExecutionRequestOutcome::responded( + build_standard_text_sync_heartbeat_shell_response( + state.clone(), + parts_for_task, + trace_id.to_string(), + decision.clone(), + plan_kind.to_string(), + move |state, parts, trace_id, decision, plan_kind, started_at| async move { + let Some((attempt_source, _candidate_count)) = + build_local_openai_responses_sync_attempt_source_for_kind( + &state, + &parts, + trace_id.as_str(), + &decision, + &body_json_for_task, + plan_kind.as_str(), + ) + .await? + else { + return Ok(LocalExecutionRequestOutcome::NoPath); + }; + + let outcome = execute_sync_attempt_source::( + &state, + &parts, + trace_id.as_str(), + &decision, + plan_kind.as_str(), + attempt_source, + ) + .await?; + match outcome { + LocalExecutionRequestOutcome::Exhausted(exhaustion) => { + record_standard_text_sync_heartbeat_exhaustion( + &state, + exhaustion, + &started_at, + ) + .await; + Ok(LocalExecutionRequestOutcome::NoPath) + } + outcome => Ok(outcome), + } + }, + )?, + )); + } + execute_sync_attempt_source::( state, parts, @@ -235,6 +351,57 @@ pub(crate) async fn maybe_execute_sync_via_standard_family_decision( return Ok(LocalExecutionRequestOutcome::NoPath); }; + if standard_text_sync_heartbeat_should_wrap(state, plan_kind).await { + let parts_for_task = parts.clone(); + let body_json_for_task = body_json.clone(); + return Ok(LocalExecutionRequestOutcome::responded( + build_standard_text_sync_heartbeat_shell_response( + state.clone(), + parts_for_task, + trace_id.to_string(), + decision.clone(), + plan_kind.to_string(), + move |state, parts, trace_id, decision, plan_kind, started_at| async move { + let Some((attempt_source, _candidate_count)) = + build_standard_family_sync_attempt_source( + &state, + &parts, + trace_id.as_str(), + &decision, + &body_json_for_task, + spec, + ) + .await? + else { + return Ok(LocalExecutionRequestOutcome::NoPath); + }; + + let outcome = execute_sync_attempt_source::( + &state, + &parts, + trace_id.as_str(), + &decision, + plan_kind.as_str(), + attempt_source, + ) + .await?; + match outcome { + LocalExecutionRequestOutcome::Exhausted(exhaustion) => { + record_standard_text_sync_heartbeat_exhaustion( + &state, + exhaustion, + &started_at, + ) + .await; + Ok(LocalExecutionRequestOutcome::NoPath) + } + outcome => Ok(outcome), + } + }, + )?, + )); + } + execute_sync_attempt_source::( state, parts, @@ -399,6 +566,57 @@ pub(crate) async fn maybe_execute_sync_via_local_same_format_provider_decision( return Ok(LocalExecutionRequestOutcome::NoPath); }; + if standard_text_sync_heartbeat_should_wrap(state, plan_kind).await { + let parts_for_task = parts.clone(); + let body_json_for_task = body_json.clone(); + return Ok(LocalExecutionRequestOutcome::responded( + build_standard_text_sync_heartbeat_shell_response( + state.clone(), + parts_for_task, + trace_id.to_string(), + decision.clone(), + plan_kind.to_string(), + move |state, parts, trace_id, decision, plan_kind, started_at| async move { + let Some((attempt_source, _candidate_count)) = + build_local_same_format_sync_attempt_source( + &state, + &parts, + trace_id.as_str(), + &decision, + &body_json_for_task, + spec, + ) + .await? + else { + return Ok(LocalExecutionRequestOutcome::NoPath); + }; + + let outcome = execute_sync_attempt_source::( + &state, + &parts, + trace_id.as_str(), + &decision, + plan_kind.as_str(), + attempt_source, + ) + .await?; + match outcome { + LocalExecutionRequestOutcome::Exhausted(exhaustion) => { + record_standard_text_sync_heartbeat_exhaustion( + &state, + exhaustion, + &started_at, + ) + .await; + Ok(LocalExecutionRequestOutcome::NoPath) + } + outcome => Ok(outcome), + } + }, + )?, + )); + } + execute_sync_attempt_source::( state, parts, @@ -495,6 +713,341 @@ async fn openai_image_sync_heartbeat_enabled(state: &AppState) -> bool { } } +async fn standard_text_sync_heartbeat_enabled(state: &AppState) -> bool { + match state + .read_system_config_json_value(ENABLE_STANDARD_TEXT_SYNC_HEARTBEAT_CONFIG_KEY) + .await + { + Ok(value) => system_config_bool(value.as_ref(), false), + Err(err) => { + tracing::warn!( + event_name = "standard_text_sync_heartbeat_config_read_failed", + log_type = "ops", + error = ?err, + "gateway failed to read standard text sync heartbeat config; defaulting disabled" + ); + false + } + } +} + +fn standard_text_sync_heartbeat_applies_to_plan_kind(plan_kind: &str) -> bool { + matches!( + plan_kind, + "openai_chat_sync" + | "openai_responses_sync" + | "openai_responses_compact_sync" + | "claude_chat_sync" + | "claude_cli_sync" + | "gemini_chat_sync" + | "gemini_cli_sync" + ) +} + +async fn standard_text_sync_heartbeat_should_wrap(state: &AppState, plan_kind: &str) -> bool { + standard_text_sync_heartbeat_applies_to_plan_kind(plan_kind) + && standard_text_sync_heartbeat_enabled(state).await +} + +fn standard_text_sync_heartbeat_client_api_format_for_plan_kind(plan_kind: &str) -> &'static str { + match plan_kind { + "openai_responses_sync" => "openai:responses", + "openai_responses_compact_sync" => "openai:responses:compact", + "claude_chat_sync" | "claude_cli_sync" => "claude:messages", + "gemini_chat_sync" | "gemini_cli_sync" => "gemini:generate_content", + _ => "openai:chat", + } +} + +fn build_standard_text_sync_heartbeat_shell_response( + state: AppState, + parts: http::request::Parts, + trace_id: String, + decision: GatewayControlDecision, + plan_kind: String, + execute: F, +) -> Result, GatewayError> +where + F: FnOnce( + AppState, + http::request::Parts, + String, + GatewayControlDecision, + String, + Instant, + ) -> Fut + + Send + + 'static, + Fut: std::future::Future> + + Send + + 'static, +{ + let request_id = (!trace_id.trim().is_empty()).then(|| trace_id.clone()); + let client_api_format = + standard_text_sync_heartbeat_client_api_format_for_plan_kind(plan_kind.as_str()) + .to_string(); + let redaction_slot = parts + .extensions + .get::() + .cloned(); + let trace_id_for_response = trace_id.clone(); + let decision_for_response = decision.clone(); + let started_at = Instant::now(); + let (tx, rx) = mpsc::channel::>(1); + + tokio::spawn(async move { + let bytes = standard_text_sync_heartbeat_final_bytes( + client_api_format.as_str(), + redaction_slot.as_ref(), + execute(state, parts, trace_id, decision, plan_kind, started_at).await, + ) + .await; + let _ = tx.send(Ok(Bytes::from(bytes))).await; + }); + + let headers = BTreeMap::from([( + CONTENT_TYPE.as_str().to_string(), + "application/json".to_string(), + )]); + let response = build_client_response_from_parts_with_mutator( + StatusCode::OK.as_u16(), + &headers, + Body::from_stream(build_sync_json_whitespace_heartbeat_stream(rx)), + trace_id_for_response.as_str(), + Some(&decision_for_response), + |headers| { + headers.remove(CONTENT_LENGTH); + headers.remove(CONTENT_ENCODING); + headers.insert( + CACHE_CONTROL, + HeaderValue::from_static("no-cache, no-transform"), + ); + headers.insert( + HeaderName::from_static("x-accel-buffering"), + HeaderValue::from_static("no"), + ); + Ok(()) + }, + )?; + attach_control_metadata_headers(response, request_id.as_deref(), None) +} + +async fn record_standard_text_sync_heartbeat_exhaustion( + state: &AppState, + exhaustion: LocalExecutionExhaustion, + started_at: &Instant, +) { + record_failed_usage_for_exhausted_request( + state, + exhaustion, + started_at, + "Standard text sync heartbeat exhausted all local candidates", + EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS, + None, + ) + .await; +} + +async fn standard_text_sync_heartbeat_final_bytes( + client_api_format: &str, + redaction_slot: Option<&crate::privacy::RedactionSessionSlot>, + result: Result, +) -> Vec { + match result { + Ok(LocalExecutionRequestOutcome::Responded(response)) => { + standard_text_sync_heartbeat_response_body_bytes( + client_api_format, + redaction_slot, + response, + ) + .await + } + Ok(LocalExecutionRequestOutcome::Exhausted(_)) + | Ok(LocalExecutionRequestOutcome::NoPath) => standard_text_sync_heartbeat_error_body( + client_api_format, + STANDARD_TEXT_SYNC_HEARTBEAT_EXHAUSTED_STATUS, + "standard text sync exhausted all local candidates", + ), + Err(err) => standard_text_sync_heartbeat_error_body( + client_api_format, + STANDARD_TEXT_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS, + &format!("{err:?}"), + ), + } +} + +async fn standard_text_sync_heartbeat_response_body_bytes( + client_api_format: &str, + redaction_slot: Option<&crate::privacy::RedactionSessionSlot>, + response: Response, +) -> Vec { + let status_code = response.status().as_u16(); + let (parts, body) = response.into_parts(); + match to_bytes(body, usize::MAX).await { + Ok(bytes) => { + let body = match standard_text_sync_heartbeat_restore_response_body( + redaction_slot, + &parts.headers, + bytes.as_ref(), + ) { + Ok(body) => body, + Err(err) => { + return standard_text_sync_heartbeat_error_body( + client_api_format, + STANDARD_TEXT_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS, + &format!("{err:?}"), + ); + } + }; + if (200..300).contains(&status_code) && !body.is_empty() { + return body; + } + if !(200..300).contains(&status_code) { + return standard_text_sync_heartbeat_error_body_from_response( + client_api_format, + status_code, + body.as_ref(), + ); + } + standard_text_sync_heartbeat_error_body( + client_api_format, + STANDARD_TEXT_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS, + "empty standard text sync response", + ) + } + Err(err) => standard_text_sync_heartbeat_error_body( + client_api_format, + STANDARD_TEXT_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS, + &err.to_string(), + ), + } +} + +fn standard_text_sync_heartbeat_restore_response_body( + redaction_slot: Option<&crate::privacy::RedactionSessionSlot>, + headers: &http::HeaderMap, + body: &[u8], +) -> Result, GatewayError> { + let Some(redaction_slot) = redaction_slot else { + return Ok(body.to_vec()); + }; + let candidate_id = headers + .get(CONTROL_CANDIDATE_ID_HEADER) + .and_then(|value| value.to_str().ok()) + .map(str::trim) + .filter(|value| !value.is_empty()); + let Some(session) = redaction_slot.take_for_candidate(candidate_id) else { + return Ok(body.to_vec()); + }; + let mut header_values = headers + .iter() + .map(|(name, value)| { + ( + name.as_str().to_string(), + value.to_str().unwrap_or_default().to_string(), + ) + }) + .collect::>(); + crate::privacy::restore_sync_response_body(&mut header_values, body, &session) + .map(|restored| restored.body) +} + +fn standard_text_sync_heartbeat_error_body_from_response( + client_api_format: &str, + status_code: u16, + body: &[u8], +) -> Vec { + if let Ok(mut value) = serde_json::from_slice::(body) { + if standard_text_sync_heartbeat_insert_upstream_status(&mut value, status_code) { + return serde_json::to_vec(&value).unwrap_or_else(|_| { + standard_text_sync_heartbeat_error_body( + client_api_format, + status_code, + &format!("upstream returned status {status_code}"), + ) + }); + } + } + + let message = standard_text_sync_heartbeat_error_message_from_body(status_code, body); + standard_text_sync_heartbeat_error_body(client_api_format, status_code, message.as_str()) +} + +fn standard_text_sync_heartbeat_insert_upstream_status( + value: &mut Value, + status_code: u16, +) -> bool { + let Some(error) = value.get_mut("error").and_then(Value::as_object_mut) else { + return false; + }; + error.insert("upstream_status".to_string(), Value::from(status_code)); + error + .entry("message".to_string()) + .or_insert_with(|| Value::String(format!("upstream returned status {status_code}"))); + true +} + +fn standard_text_sync_heartbeat_error_message_from_body(status_code: u16, body: &[u8]) -> String { + let text = String::from_utf8_lossy(body).trim().to_string(); + if text.is_empty() { + return format!("upstream returned status {status_code}"); + } + text.chars() + .take(STANDARD_TEXT_SYNC_HEARTBEAT_ERROR_MESSAGE_LIMIT) + .collect() +} + +fn standard_text_sync_heartbeat_error_body( + client_api_format: &str, + status_code: u16, + message: &str, +) -> Vec { + let mut body = build_core_error_body_for_client_format( + client_api_format, + message, + Some("upstream_error"), + standard_text_sync_heartbeat_error_kind(status_code), + ) + .unwrap_or_else(|| { + json!({ + "error": { + "type": "upstream_error", + "message": message, + "code": status_code, + } + }) + }); + if !standard_text_sync_heartbeat_insert_upstream_status(&mut body, status_code) { + body = json!({ + "error": { + "type": "upstream_error", + "message": message, + "code": status_code, + "upstream_status": status_code, + } + }); + } + serde_json::to_vec(&body).unwrap_or_else(|_| { + format!( + "{{\"error\":{{\"type\":\"upstream_error\",\"code\":{status_code},\"upstream_status\":{status_code}}}}}" + ) + .into_bytes() + }) +} + +fn standard_text_sync_heartbeat_error_kind(status_code: u16) -> LocalCoreSyncErrorKind { + match status_code { + 400 => LocalCoreSyncErrorKind::InvalidRequest, + 401 => LocalCoreSyncErrorKind::Authentication, + 403 => LocalCoreSyncErrorKind::PermissionDenied, + 404 => LocalCoreSyncErrorKind::NotFound, + 413 => LocalCoreSyncErrorKind::ContextLengthExceeded, + 429 => LocalCoreSyncErrorKind::RateLimit, + 503 => LocalCoreSyncErrorKind::Overloaded, + _ => LocalCoreSyncErrorKind::ServerError, + } +} + fn build_openai_image_sync_heartbeat_shell_response( state: AppState, request_path: String, @@ -944,10 +1497,35 @@ pub(crate) fn decision_payload_is_direct_execution(payload: &AiExecutionDecision #[cfg(test)] mod tests { use super::*; + use futures_util::StreamExt; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; const TEST_OPENAI_IMAGE_SYNC_PLAN_KIND: &str = "openai_image_sync"; + const TEST_STANDARD_TEXT_SYNC_PLAN_KIND: &str = "openai_responses_compact_sync"; + + struct TestSyncAttemptSource { + attempts: VecDeque, + } + + impl TestSyncAttemptSource { + fn new(attempts: Vec) -> Self { + Self { + attempts: VecDeque::from(attempts), + } + } + } + + #[async_trait::async_trait] + impl LocalExecutionAttemptSource for TestSyncAttemptSource { + async fn next_execution_attempt(&mut self) -> Result, GatewayError> { + Ok(self.attempts.pop_front()) + } + + async fn drain_execution_attempts(&mut self) -> Result, GatewayError> { + Ok(self.attempts.drain(..).collect()) + } + } fn test_openai_image_heartbeat_decision() -> GatewayControlDecision { GatewayControlDecision::synthetic( @@ -1028,6 +1606,63 @@ mod tests { } } + fn test_standard_text_heartbeat_decision() -> GatewayControlDecision { + GatewayControlDecision::synthetic( + "/v1/responses", + Some("ai_public".to_string()), + Some("openai".to_string()), + Some("responses".to_string()), + Some("openai:responses:compact".to_string()), + ) + .with_execution_runtime_candidate(true) + } + + fn test_standard_text_heartbeat_plan( + endpoint_id: &str, + candidate_id: &str, + client_api_format: &str, + ) -> aether_contracts::ExecutionPlan { + aether_contracts::ExecutionPlan { + request_id: "trace-standard-text-heartbeat-retry".to_string(), + candidate_id: Some(candidate_id.to_string()), + provider_name: Some("OpenAI".to_string()), + provider_id: "provider-openai".to_string(), + endpoint_id: endpoint_id.to_string(), + key_id: "key-openai".to_string(), + method: "POST".to_string(), + url: "https://api.openai.com/v1/responses".to_string(), + headers: BTreeMap::new(), + content_type: Some("application/json".to_string()), + content_encoding: None, + body: aether_contracts::RequestBody::from_json(json!({"model": "gpt-5"})), + stream: false, + client_api_format: client_api_format.to_string(), + provider_api_format: client_api_format.to_string(), + model_name: Some("gpt-5".to_string()), + proxy: None, + transport_profile: None, + timeouts: None, + } + } + + fn test_standard_text_heartbeat_attempt( + candidate_index: u32, + endpoint_id: &str, + candidate_id: &str, + client_api_format: &str, + ) -> AiSyncAttempt { + AiSyncAttempt { + plan: test_standard_text_heartbeat_plan(endpoint_id, candidate_id, client_api_format), + report_kind: None, + report_context: Some(json!({ + "candidate_index": candidate_index, + "retry_index": 0, + "client_api_format": client_api_format, + "provider_api_format": client_api_format, + })), + } + } + #[tokio::test] async fn openai_image_sync_heartbeat_success_body_is_unchanged() { let response = Response::builder() @@ -1133,4 +1768,239 @@ mod tests { assert_eq!(call_count.load(Ordering::SeqCst), 2); assert_eq!(body, json!({"data": [{"b64_json": "second-candidate"}]})); } + + #[tokio::test] + async fn standard_text_sync_heartbeat_missing_config_defaults_disabled() { + let state = AppState::new().expect("state should build"); + + assert!(!standard_text_sync_heartbeat_enabled(&state).await); + } + + #[tokio::test] + async fn standard_text_sync_heartbeat_no_local_candidates_preserves_no_path() { + let state = AppState::new() + .expect("state should build") + .with_data_state_for_tests( + crate::data::GatewayDataState::disabled().with_system_config_values_for_tests([( + ENABLE_STANDARD_TEXT_SYNC_HEARTBEAT_CONFIG_KEY.to_string(), + json!(true), + )]), + ); + let (parts, _) = http::Request::builder() + .method(http::Method::POST) + .uri("/v1/responses") + .body(()) + .expect("request should build") + .into_parts(); + + let outcome = maybe_execute_sync_via_local_openai_responses_decision( + &state, + &parts, + "trace-standard-text-heartbeat-no-path", + &test_standard_text_heartbeat_decision(), + &json!({"model": "missing-local-candidate"}), + TEST_STANDARD_TEXT_SYNC_PLAN_KIND, + ) + .await + .expect("heartbeat no-path check should execute"); + + assert!(matches!(outcome, LocalExecutionRequestOutcome::NoPath)); + } + + #[tokio::test] + async fn standard_text_sync_heartbeat_success_body_is_unchanged() { + let response = Response::builder() + .status(StatusCode::OK) + .body(Body::from(r#"{"id":"resp_123","output":[]}"#)) + .expect("response should build"); + + let bytes = standard_text_sync_heartbeat_response_body_bytes( + "openai:responses:compact", + None, + response, + ) + .await; + let body: Value = serde_json::from_slice(&bytes).expect("body should decode"); + + assert_eq!(body, json!({"id": "resp_123", "output": []})); + } + + #[tokio::test] + async fn standard_text_sync_heartbeat_claude_error_body_includes_upstream_status() { + let response = Response::builder() + .status(StatusCode::TOO_MANY_REQUESTS) + .body(Body::from( + r#"{"type":"error","error":{"type":"rate_limit_error","message":"slow down"}}"#, + )) + .expect("response should build"); + + let bytes = + standard_text_sync_heartbeat_response_body_bytes("claude:messages", None, response) + .await; + let body: Value = serde_json::from_slice(&bytes).expect("body should decode"); + + assert_eq!(body["type"], json!("error")); + assert_eq!(body["error"]["type"], json!("rate_limit_error")); + assert_eq!(body["error"]["message"], json!("slow down")); + assert_eq!(body["error"]["upstream_status"], json!(429)); + } + + #[test] + fn standard_text_sync_heartbeat_applies_to_chat_and_cli_plan_kinds() { + assert!(standard_text_sync_heartbeat_applies_to_plan_kind( + "claude_chat_sync" + )); + assert!(standard_text_sync_heartbeat_applies_to_plan_kind( + "claude_cli_sync" + )); + assert!(standard_text_sync_heartbeat_applies_to_plan_kind( + "gemini_chat_sync" + )); + assert!(standard_text_sync_heartbeat_applies_to_plan_kind( + "gemini_cli_sync" + )); + assert!(!standard_text_sync_heartbeat_applies_to_plan_kind( + "openai_embedding_sync" + )); + } + + #[tokio::test] + async fn standard_text_sync_heartbeat_redirect_status_is_wrapped_as_error() { + let response = Response::builder() + .status(StatusCode::TEMPORARY_REDIRECT) + .body(Body::from(r#"{"location":"https://upstream.example"}"#)) + .expect("response should build"); + + let bytes = standard_text_sync_heartbeat_response_body_bytes( + "openai:responses:compact", + None, + response, + ) + .await; + let body: Value = serde_json::from_slice(&bytes).expect("body should decode"); + + assert_eq!(body["error"]["type"], json!("server_error")); + assert_eq!(body["error"]["upstream_status"], json!(307)); + } + + #[tokio::test] + async fn standard_text_sync_heartbeat_shell_sends_whitespace_before_background_finishes() { + let state = AppState::new().expect("state should build"); + let (parts, _) = http::Request::builder() + .method(http::Method::POST) + .uri("/v1/responses") + .body(()) + .expect("request should build") + .into_parts(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel::<()>(); + + let response = build_standard_text_sync_heartbeat_shell_response( + state, + parts, + "trace-standard-text-heartbeat-shell".to_string(), + test_standard_text_heartbeat_decision(), + TEST_STANDARD_TEXT_SYNC_PLAN_KIND.to_string(), + move |_state, _parts, _trace_id, _decision, _plan_kind, _started_at| async move { + let _ = release_rx.await; + Ok(LocalExecutionRequestOutcome::responded( + Response::builder() + .status(StatusCode::OK) + .body(Body::from(r#"{"id":"resp_done","output":[]}"#)) + .expect("response should build"), + )) + }, + ) + .expect("heartbeat shell should build"); + let mut body_stream = response.into_body().into_data_stream(); + + let first = body_stream + .next() + .await + .expect("heartbeat stream should yield") + .expect("heartbeat chunk should be ok"); + assert_eq!(first.as_ref(), b"\n"); + let _ = release_tx.send(()); + } + + #[test] + fn standard_text_sync_heartbeat_compact_non_json_error_body_is_wrapped_in_client_format() { + let bytes = standard_text_sync_heartbeat_error_body_from_response( + "openai:responses:compact", + 502, + b"bad gateway from upstream", + ); + let body: Value = serde_json::from_slice(&bytes).expect("body should decode"); + + assert_eq!(body["error"]["type"], json!("server_error")); + assert_eq!(body["error"]["message"], json!("bad gateway from upstream")); + assert_eq!(body["error"]["upstream_status"], json!(502)); + } + + #[tokio::test] + async fn standard_text_sync_heartbeat_attempts_retry_first_candidate_then_return_second() { + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_for_override = Arc::clone(&call_count); + let state = AppState::new() + .expect("state should build") + .with_execution_runtime_sync_override_for_tests(move |plan| { + call_count_for_override.fetch_add(1, Ordering::SeqCst); + if plan.endpoint_id == "endpoint-retry" { + Ok(test_openai_image_execution_result( + plan, + StatusCode::TOO_MANY_REQUESTS.as_u16(), + json!({"error": {"message": "retry this candidate"}}), + )) + } else { + Ok(test_openai_image_execution_result( + plan, + StatusCode::OK.as_u16(), + json!({"id": "resp_second_candidate", "output": []}), + )) + } + }); + let attempts = vec![ + test_standard_text_heartbeat_attempt( + 0, + "endpoint-retry", + "candidate-retry", + "openai:responses:compact", + ), + test_standard_text_heartbeat_attempt( + 1, + "endpoint-success", + "candidate-success", + "openai:responses:compact", + ), + ]; + + let (parts, _) = http::Request::builder() + .method(http::Method::POST) + .uri("/v1/responses") + .body(()) + .expect("request should build") + .into_parts(); + let outcome = execute_sync_attempt_source::( + &state, + &parts, + "trace-standard-text-heartbeat-retry", + &test_standard_text_heartbeat_decision(), + TEST_STANDARD_TEXT_SYNC_PLAN_KIND, + TestSyncAttemptSource::new(attempts), + ) + .await + .expect("heartbeat attempts should execute"); + let LocalExecutionRequestOutcome::Responded(response) = outcome else { + panic!("second candidate should return a response"); + }; + let bytes = standard_text_sync_heartbeat_response_body_bytes( + "openai:responses:compact", + None, + response, + ) + .await; + let body: Value = serde_json::from_slice(&bytes).expect("body should decode"); + + assert_eq!(call_count.load(Ordering::SeqCst), 2); + assert_eq!(body, json!({"id": "resp_second_candidate", "output": []})); + } } diff --git a/apps/aether-gateway/src/handlers/admin/model/global/providers.rs b/apps/aether-gateway/src/handlers/admin/model/global/providers.rs index 73e24b703..c41d3fba9 100644 --- a/apps/aether-gateway/src/handlers/admin/model/global/providers.rs +++ b/apps/aether-gateway/src/handlers/admin/model/global/providers.rs @@ -13,6 +13,7 @@ const EMBEDDING_API_FORMATS: &[&str] = &[ "jina:embedding", "gemini:embedding", "doubao:embedding", + "aliyun:multimodal_embedding", ]; fn json_value_contains_string(value: &serde_json::Value, expected: &str) -> bool { diff --git a/apps/aether-gateway/src/handlers/admin/model/payloads.rs b/apps/aether-gateway/src/handlers/admin/model/payloads.rs index fc2258fab..7c11f7bf5 100644 --- a/apps/aether-gateway/src/handlers/admin/model/payloads.rs +++ b/apps/aether-gateway/src/handlers/admin/model/payloads.rs @@ -6,6 +6,7 @@ const EMBEDDING_API_FORMATS: &[&str] = &[ "jina:embedding", "gemini:embedding", "doubao:embedding", + "aliyun:multimodal_embedding", ]; pub(crate) fn model_tiered_pricing_first_tier_value( diff --git a/apps/aether-gateway/src/handlers/admin/observability/monitoring/cache_affinity.rs b/apps/aether-gateway/src/handlers/admin/observability/monitoring/cache_affinity.rs index 35ea59cd8..a1a569ce8 100644 --- a/apps/aether-gateway/src/handlers/admin/observability/monitoring/cache_affinity.rs +++ b/apps/aether-gateway/src/handlers/admin/observability/monitoring/cache_affinity.rs @@ -94,7 +94,7 @@ fn split_admin_monitoring_api_format_and_model( fn is_known_admin_monitoring_api_format_family(value: &str) -> bool { matches!( value.trim().to_ascii_lowercase().as_str(), - "openai" | "claude" | "gemini" | "jina" | "doubao" + "openai" | "claude" | "gemini" | "jina" | "doubao" | "aliyun" ) } @@ -116,6 +116,7 @@ fn is_known_admin_monitoring_api_format(value: &str) -> bool { | "jina:embedding" | "jina:rerank" | "doubao:embedding" + | "aliyun:multimodal_embedding" ) } diff --git a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs index 25e5686f1..5c0b621db 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs @@ -597,6 +597,22 @@ fn provider_query_build_test_request_body_for_api_format( let message = provider_query_extract_message(payload) .unwrap_or_else(|| DEFAULT_PROVIDER_QUERY_TEST_MESSAGE.to_string()); match client_api_format.as_str() { + "openai:embedding" => json!({ + "model": model, + "input": message, + }), + "openai:rerank" => json!({ + "model": model, + "query": message, + "documents": [ + "apple", + "banana", + "fruit", + "vegetable" + ], + "return_documents": true, + "top_n": 4, + }), "openai:responses" | "openai:responses:compact" => json!({ "model": model, "input": message, @@ -649,6 +665,21 @@ fn provider_query_insert_default_test_conversation( let message = provider_query_extract_message(payload) .unwrap_or_else(|| DEFAULT_PROVIDER_QUERY_TEST_MESSAGE.to_string()); match client_api_format { + "openai:embedding" => { + object.insert("input".to_string(), Value::String(message)); + } + "openai:rerank" => { + object.insert("query".to_string(), Value::String(message)); + object + .entry("documents".to_string()) + .or_insert_with(|| json!(["apple", "banana", "fruit", "vegetable"])); + object + .entry("return_documents".to_string()) + .or_insert(Value::Bool(true)); + object + .entry("top_n".to_string()) + .or_insert_with(|| Value::from(4_u64)); + } "openai:responses" | "openai:responses:compact" => { object.insert("input".to_string(), Value::String(message)); } @@ -2920,8 +2951,13 @@ async fn provider_query_execute_standard_test_candidate( ); provider_request_body } - "openai:embedding" | "gemini:embedding" | "jina:embedding" | "doubao:embedding" - | "openai:rerank" | "jina:rerank" => { + "openai:embedding" + | "gemini:embedding" + | "jina:embedding" + | "doubao:embedding" + | "aliyun:multimodal_embedding" + | "openai:rerank" + | "jina:rerank" => { let Some(mut provider_request_body) = crate::ai_serving::build_standard_request_body_with_model_directives_and_request_headers( &request_body, @@ -3051,6 +3087,7 @@ async fn provider_query_execute_standard_test_candidate( | "gemini:embedding" | "jina:embedding" | "doubao:embedding" + | "aliyun:multimodal_embedding" | "openai:rerank" | "jina:rerank" => state.resolve_local_oauth_header_auth(&transport).await?, _ => None, @@ -3062,6 +3099,7 @@ async fn provider_query_execute_standard_test_candidate( | "openai:embedding" | "jina:embedding" | "doubao:embedding" + | "aliyun:multimodal_embedding" | "openai:rerank" | "jina:rerank" => { crate::provider_transport::auth::resolve_local_openai_bearer_auth(&transport) diff --git a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/adapter.rs b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/adapter.rs index 8c062832d..876ebd1d9 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/adapter.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/adapter.rs @@ -74,6 +74,7 @@ pub(super) fn provider_query_standard_test_unsupported_reason( | "openai:embedding" | "jina:embedding" | "doubao:embedding" + | "aliyun:multimodal_embedding" | "openai:rerank" | "jina:rerank" => { crate::provider_transport::policy::local_standard_transport_unsupported_reason_with_network( @@ -253,6 +254,7 @@ pub(super) fn provider_query_test_adapter_for_provider_api_format( | "gemini:embedding" | "jina:embedding" | "doubao:embedding" + | "aliyun:multimodal_embedding" | "openai:rerank" | "jina:rerank" ) { @@ -352,6 +354,7 @@ pub(super) fn provider_query_transport_supports_model_test_execution( | "openai:embedding" | "jina:embedding" | "doubao:embedding" + | "aliyun:multimodal_embedding" | "openai:rerank" | "jina:rerank" => { crate::provider_transport::policy::supports_local_standard_transport_with_network( diff --git a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/summary.rs b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/summary.rs index 5ef3f4c63..881ec8879 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/summary.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/summary.rs @@ -138,6 +138,12 @@ fn provider_query_endpoint_route_payload( "embeddings", "openai_batch", ), + "aliyun:multimodal_embedding" => ( + "Aliyun DashScope", + "dashscope_native", + "multimodal-embedding", + "dashscope_contents", + ), "openai:chat" if is_vertex && is_openai_compat => ( "Vertex AI OpenAI-compatible", "openai_compatible", diff --git a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs index 24c8215bf..867fb69cc 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs @@ -470,6 +470,26 @@ fn provider_query_compact_test_request_body_defaults_to_responses_input() { assert!(body.get("messages").is_none()); } +#[test] +fn provider_query_embedding_test_request_body_defaults_to_embedding_input() { + let payload = json!({"message": "hello from embedding"}); + + let client_api_format = + provider_query_standard_test_client_api_format("aliyun:multimodal_embedding"); + let body = provider_query_build_test_request_body_for_api_format( + &payload, + "qwen3-vl-embedding", + "/api/admin/provider-query/test-model", + client_api_format, + ); + + assert_eq!(client_api_format, "openai:embedding"); + assert_eq!(body["model"], json!("qwen3-vl-embedding")); + assert_eq!(body["input"], json!("hello from embedding")); + assert!(body.get("messages").is_none()); + assert!(body.get("stream").is_none()); +} + #[test] fn provider_query_compact_test_request_body_promotes_prompt_to_input() { let payload = json!({ @@ -658,6 +678,13 @@ fn provider_query_test_adapter_routes_fixed_provider_endpoint_types() { provider_query_test_adapter_for_provider_api_format("custom", "gemini:embedding"), Some(ProviderQueryTestAdapter::Standard) ); + assert_eq!( + provider_query_test_adapter_for_provider_api_format( + "aliyun", + "aliyun:multimodal_embedding" + ), + Some(ProviderQueryTestAdapter::Standard) + ); assert_eq!( provider_query_test_adapter_for_provider_api_format("jina", "jina:rerank"), Some(ProviderQueryTestAdapter::Standard) diff --git a/apps/aether-gateway/src/handlers/admin/users/shared.rs b/apps/aether-gateway/src/handlers/admin/users/shared.rs index 615f2daf2..dbae70da8 100644 --- a/apps/aether-gateway/src/handlers/admin/users/shared.rs +++ b/apps/aether-gateway/src/handlers/admin/users/shared.rs @@ -278,6 +278,9 @@ pub(crate) fn normalize_admin_user_api_formats( if item.is_empty() { return Err("allowed_api_formats 不能为空".to_string()); } + if !looks_like_admin_api_format_signature(item) { + return Err(format!("allowed_api_formats 格式无效: {item}")); + } let Some(normalized_item) = crate::api::ai::normalize_admin_endpoint_signature(item) else { return Err(format!("allowed_api_formats 格式无效: {item}")); }; @@ -289,6 +292,12 @@ pub(crate) fn normalize_admin_user_api_formats( Ok(Some(normalized)) } +fn looks_like_admin_api_format_signature(value: &str) -> bool { + value + .split_once(':') + .is_some_and(|(family, kind)| !family.trim().is_empty() && !kind.trim().is_empty()) +} + pub(crate) fn normalize_admin_user_ip_rules( value: Option>, ) -> Result>, String> { diff --git a/apps/aether-gateway/src/handlers/public/ai_public.rs b/apps/aether-gateway/src/handlers/public/ai_public.rs index 589556dcd..dcc84a5c0 100644 --- a/apps/aether-gateway/src/handlers/public/ai_public.rs +++ b/apps/aether-gateway/src/handlers/public/ai_public.rs @@ -522,12 +522,45 @@ fn embedding_array_input_is_non_empty(items: &[Value]) -> bool { item.as_array() .is_some_and(|items| embedding_token_array_is_non_empty(items)) }) + || items.iter().all(embedding_multimodal_content_is_non_empty) } fn embedding_token_array_is_non_empty(items: &[Value]) -> bool { !items.is_empty() && items.iter().all(|item| item.as_u64().is_some()) } +fn embedding_multimodal_content_is_non_empty(value: &Value) -> bool { + let Some(object) = value.as_object() else { + return false; + }; + let valid_text = object + .get("text") + .map(|value| value.as_str().is_some_and(|text| !text.trim().is_empty())); + let valid_image = object + .get("image") + .map(|value| value.as_str().is_some_and(|image| !image.trim().is_empty())); + let valid_video = object + .get("video") + .map(|value| value.as_str().is_some_and(|video| !video.trim().is_empty())); + let valid_multi_images = object.get("multi_images").map(|value| { + value.as_array().is_some_and(|items| { + !items.is_empty() + && items + .iter() + .all(|item| item.as_str().is_some_and(|image| !image.trim().is_empty())) + }) + }); + + [valid_text, valid_image, valid_video, valid_multi_images] + .into_iter() + .flatten() + .all(|valid| valid) + && [valid_text, valid_image, valid_video, valid_multi_images] + .into_iter() + .flatten() + .any(|valid| valid) +} + fn image_request_count(value: &Value) -> Option { value .as_u64() diff --git a/apps/aether-gateway/src/handlers/public/support/billing.rs b/apps/aether-gateway/src/handlers/public/support/billing.rs index c84f5366d..531a5aa71 100644 --- a/apps/aether-gateway/src/handlers/public/support/billing.rs +++ b/apps/aether-gateway/src/handlers/public/support/billing.rs @@ -327,6 +327,32 @@ pub(super) async fn handle_billing_plan_checkout( false, ); } + match state + .find_pending_plan_purchase_order_by_user_id(&auth.user.id, &plan.id) + .await + { + Ok(Some(order)) => { + return build_auth_json_response( + http::StatusCode::OK, + json!({ + "order": payment_order_payload(&order, &plan), + "payment_instructions": sanitize_wallet_gateway_response( + order.gateway_response.clone() + ), + "reused_pending_order": true, + }), + None, + ) + } + Ok(None) => {} + Err(err) => { + return build_auth_error_response( + http::StatusCode::INTERNAL_SERVER_ERROR, + format!("pending billing checkout lookup failed: {err:?}"), + false, + ) + } + } let now = Utc::now(); let order_no = billing_order_no(now); let expires_at = now + chrono::Duration::minutes(30); diff --git a/apps/aether-gateway/src/handlers/public/support/models/shared.rs b/apps/aether-gateway/src/handlers/public/support/models/shared.rs index 856b1e135..d57cf6500 100644 --- a/apps/aether-gateway/src/handlers/public/support/models/shared.rs +++ b/apps/aether-gateway/src/handlers/public/support/models/shared.rs @@ -25,6 +25,7 @@ pub(crate) fn models_api_format(request_context: &GatewayPublicRequestContext) - "jina:embedding" => Some("jina:embedding"), "jina:rerank" => Some("jina:rerank"), "doubao:embedding" => Some("doubao:embedding"), + "aliyun:multimodal_embedding" => Some("aliyun:multimodal_embedding"), _ => None, } } @@ -43,6 +44,7 @@ const MODELS_EMBEDDING_QUERY_API_FORMATS: &[&str] = &[ "jina:embedding", "gemini:embedding", "doubao:embedding", + "aliyun:multimodal_embedding", ]; const MODELS_RERANK_QUERY_API_FORMATS: &[&str] = &["openai:rerank", "jina:rerank"]; @@ -54,9 +56,11 @@ pub(super) fn models_query_api_formats(api_format: &str) -> &'static [&'static s | "claude:messages" | "gemini:generate_content" => MODELS_CROSS_FORMAT_QUERY_API_FORMATS, "openai:image" => &["openai:image"], - "openai:embedding" | "jina:embedding" | "gemini:embedding" | "doubao:embedding" => { - MODELS_EMBEDDING_QUERY_API_FORMATS - } + "openai:embedding" + | "jina:embedding" + | "gemini:embedding" + | "doubao:embedding" + | "aliyun:multimodal_embedding" => MODELS_EMBEDDING_QUERY_API_FORMATS, "openai:rerank" | "jina:rerank" => MODELS_RERANK_QUERY_API_FORMATS, _ => &[], } diff --git a/apps/aether-gateway/src/state/runtime/wallet/reads.rs b/apps/aether-gateway/src/state/runtime/wallet/reads.rs index d582e97bb..d3e8a356f 100644 --- a/apps/aether-gateway/src/state/runtime/wallet/reads.rs +++ b/apps/aether-gateway/src/state/runtime/wallet/reads.rs @@ -138,6 +138,18 @@ impl AppState { .map_err(|err| GatewayError::Internal(err.to_string())) } + pub(crate) async fn find_pending_plan_purchase_order_by_user_id( + &self, + user_id: &str, + product_id: &str, + ) -> Result, GatewayError> + { + self.data + .find_pending_plan_purchase_order_by_user_id(user_id, product_id) + .await + .map_err(|err| GatewayError::Internal(err.to_string())) + } + pub(crate) async fn find_wallet_refund( &self, wallet_id: &str, diff --git a/apps/aether-gateway/src/tests/ai_execute/finalize_local.rs b/apps/aether-gateway/src/tests/ai_execute/finalize_local.rs index 9805fef10..19bb338c5 100644 --- a/apps/aether-gateway/src/tests/ai_execute/finalize_local.rs +++ b/apps/aether-gateway/src/tests/ai_execute/finalize_local.rs @@ -24,8 +24,39 @@ use aether_data_contracts::repository::provider_catalog::{ }; use sha2::{Digest, Sha256}; -#[tokio::test] -async fn gateway_executes_openai_chat_sync_upstream_stream_via_local_finalize_response() { +const OPENAI_CHAT_FINALIZE_TEST_STACK_BYTES: usize = 16 * 1024 * 1024; + +fn run_openai_chat_finalize_test(test_name: &'static str, make_future: F) +where + F: FnOnce() -> Fut + Send + 'static, + Fut: std::future::Future + 'static, +{ + let handle = std::thread::Builder::new() + .name(test_name.to_string()) + .stack_size(OPENAI_CHAT_FINALIZE_TEST_STACK_BYTES) + .spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime should build"); + runtime.block_on(make_future()); + }) + .expect("openai chat finalize test thread should spawn"); + + if let Err(payload) = handle.join() { + std::panic::resume_unwind(payload); + } +} + +#[test] +fn gateway_executes_openai_chat_sync_upstream_stream_via_local_finalize_response() { + run_openai_chat_finalize_test( + "gateway_executes_openai_chat_sync_upstream_stream_via_local_finalize_response", + gateway_executes_openai_chat_sync_upstream_stream_via_local_finalize_response_impl, + ); +} + +async fn gateway_executes_openai_chat_sync_upstream_stream_via_local_finalize_response_impl() { use base64::Engine as _; #[derive(Debug, Clone)] @@ -521,8 +552,16 @@ async fn gateway_executes_openai_chat_sync_upstream_stream_via_local_finalize_re upstream_handle.abort(); } -#[tokio::test] -async fn gateway_executes_openai_chat_cross_format_upstream_stream_via_local_finalize_response() { +#[test] +fn gateway_executes_openai_chat_cross_format_upstream_stream_via_local_finalize_response() { + run_openai_chat_finalize_test( + "gateway_executes_openai_chat_cross_format_upstream_stream_via_local_finalize_response", + gateway_executes_openai_chat_cross_format_upstream_stream_via_local_finalize_response_impl, + ); +} + +async fn gateway_executes_openai_chat_cross_format_upstream_stream_via_local_finalize_response_impl( +) { use base64::Engine as _; #[derive(Debug, Clone)] struct SeenRemoteExecutionRuntimeRequest { @@ -955,8 +994,16 @@ async fn gateway_executes_openai_chat_cross_format_upstream_stream_via_local_fin upstream_handle.abort(); } -#[tokio::test] -async fn gateway_executes_openai_chat_cross_format_tool_use_upstream_stream_via_local_finalize_response( +#[test] +fn gateway_executes_openai_chat_cross_format_tool_use_upstream_stream_via_local_finalize_response() +{ + run_openai_chat_finalize_test( + "gateway_executes_openai_chat_cross_format_tool_use_upstream_stream_via_local_finalize_response", + gateway_executes_openai_chat_cross_format_tool_use_upstream_stream_via_local_finalize_response_impl, + ); +} + +async fn gateway_executes_openai_chat_cross_format_tool_use_upstream_stream_via_local_finalize_response_impl( ) { use base64::Engine as _; @@ -1782,8 +1829,15 @@ async fn gateway_skips_openai_chat_antigravity_cross_format_sync_candidate_as_tr upstream_handle.abort(); } -#[tokio::test] -async fn gateway_executes_openai_chat_cross_format_claude_upstream_sync_via_local_finalize_response( +#[test] +fn gateway_executes_openai_chat_cross_format_claude_upstream_sync_via_local_finalize_response() { + run_openai_chat_finalize_test( + "gateway_executes_openai_chat_cross_format_claude_upstream_sync_via_local_finalize_response", + gateway_executes_openai_chat_cross_format_claude_upstream_sync_via_local_finalize_response_impl, + ); +} + +async fn gateway_executes_openai_chat_cross_format_claude_upstream_sync_via_local_finalize_response_impl( ) { fn hash_api_key(value: &str) -> String { let mut hasher = Sha256::new(); @@ -2134,8 +2188,15 @@ async fn gateway_executes_openai_chat_cross_format_claude_upstream_sync_via_loca upstream_handle.abort(); } -#[tokio::test] -async fn gateway_executes_openai_chat_cross_format_gemini_upstream_sync_via_local_finalize_response( +#[test] +fn gateway_executes_openai_chat_cross_format_gemini_upstream_sync_via_local_finalize_response() { + run_openai_chat_finalize_test( + "gateway_executes_openai_chat_cross_format_gemini_upstream_sync_via_local_finalize_response", + gateway_executes_openai_chat_cross_format_gemini_upstream_sync_via_local_finalize_response_impl, + ); +} + +async fn gateway_executes_openai_chat_cross_format_gemini_upstream_sync_via_local_finalize_response_impl( ) { fn hash_api_key(value: &str) -> String { let mut hasher = Sha256::new(); diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/claude/kiro.rs b/apps/aether-gateway/src/tests/ai_execute/sync/claude/kiro.rs index e97bed3b2..6239a7a1c 100644 --- a/apps/aether-gateway/src/tests/ai_execute/sync/claude/kiro.rs +++ b/apps/aether-gateway/src/tests/ai_execute/sync/claude/kiro.rs @@ -8,6 +8,9 @@ use super::{ StoredProviderCatalogKey, StoredProviderCatalogProvider, StoredProviderModelMapping, DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER, }; +use aether_data::repository::usage::InMemoryUsageReadRepository; +use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UsageReadRepository}; +use aether_usage_runtime::UsageRuntimeConfig; const KIRO_CLAUDE_CLI_SYNC_TEST_STACK_BYTES: usize = 16 * 1024 * 1024; @@ -33,6 +36,30 @@ where } } +async fn wait_for_completed_usage(repository: &T, request_id: &str) -> StoredRequestUsageAudit +where + T: UsageReadRepository + ?Sized, +{ + let timeout = std::time::Duration::from_secs(60); + let deadline = tokio::time::Instant::now() + timeout; + loop { + if let Some(usage) = repository + .find_by_request_id(request_id) + .await + .expect("usage should read") + { + if usage.status == "completed" { + return usage; + } + } + assert!( + tokio::time::Instant::now() < deadline, + "usage {request_id} should complete within {timeout:?}" + ); + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } +} + #[test] fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candidate() { run_kiro_claude_cli_sync_test( @@ -194,7 +221,7 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid Some(serde_json::json!({"url":"http://provider-proxy.internal:8080"})), Some(20.0), None, - None, + Some(serde_json::json!({"kiro": {"simulated_cache_enabled": true}})), ) } @@ -354,14 +381,15 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid let raw_body = to_bytes(body, usize::MAX).await.expect("body should read"); let payload: serde_json::Value = serde_json::from_slice(&raw_body).expect("execution runtime payload should parse"); + let trace_id = parts + .headers + .get(TRACE_ID_HEADER) + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); *seen_execution_runtime_inner.lock().expect("mutex should lock") = Some(SeenExecutionRuntimeSyncRequest { - trace_id: parts - .headers - .get(TRACE_ID_HEADER) - .and_then(|value| value.to_str().ok()) - .unwrap_or_default() - .to_string(), + trace_id: trace_id.clone(), url: payload .get("url") .and_then(|value| value.as_str()) @@ -453,7 +481,7 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid .concat(); Json(json!({ - "request_id": "trace-kiro-cli-local-sync-123", + "request_id": trace_id, "status_code": 200, "headers": { "content-type": "application/vnd.amazon.eventstream" @@ -481,6 +509,7 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid sample_candidate_row(), ])); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); + let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider_catalog_provider()], vec![sample_provider_catalog_endpoint()], @@ -491,34 +520,51 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone()) .with_data_state_for_tests( - crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests( + crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests( auth_repository, candidate_selection_repository, provider_catalog_repository, Arc::clone(&request_candidate_repository), + Arc::clone(&usage_repository), DEVELOPMENT_ENCRYPTION_KEY, ), - ); + ) + .with_usage_runtime_for_tests(UsageRuntimeConfig { + enabled: true, + ..UsageRuntimeConfig::default() + }); let gateway = build_router_with_state(gateway_state); let (gateway_url, gateway_handle) = start_server(gateway).await; - let response = reqwest::Client::new() - .post(format!("{gateway_url}/v1/messages")) - .header(http::header::CONTENT_TYPE, "application/json") - .header( - http::header::AUTHORIZATION, - "Bearer sk-client-kiro-cli-local-sync", - ) - .header(TRACE_ID_HEADER, "trace-kiro-cli-local-sync-123") - .body( - "{\"model\":\"claude-sonnet-4\",\"messages\":[{\"role\":\"user\",\"content\":\"hello kiro\"}],\"thinking\":{\"type\":\"enabled\",\"budget_tokens\":64}}", - ) - .send() - .await - .expect("request should succeed"); + async fn send_kiro_request( + gateway_url: &str, + trace_id: &str, + body: String, + ) -> (StatusCode, String) { + let response = reqwest::Client::new() + .post(format!("{gateway_url}/v1/messages")) + .header(http::header::CONTENT_TYPE, "application/json") + .header( + http::header::AUTHORIZATION, + "Bearer sk-client-kiro-cli-local-sync", + ) + .header(TRACE_ID_HEADER, trace_id) + .body(body) + .send() + .await + .expect("request should succeed"); - let status = response.status(); - let response_body = response.text().await.expect("body should read"); + let status = response.status(); + let response_body = response.text().await.expect("body should read"); + (status, response_body) + } + + let (status, response_body) = send_kiro_request( + &gateway_url, + "trace-kiro-cli-local-sync-123", + "{\"model\":\"claude-sonnet-4\",\"messages\":[{\"role\":\"user\",\"content\":\"hello kiro\"}],\"thinking\":{\"type\":\"enabled\",\"budget_tokens\":64}}".to_string(), + ) + .await; assert!( status == StatusCode::OK, "unexpected status={status} body={response_body} decision_hits={} plan_hits={} public_hits={}", @@ -598,6 +644,58 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid "report-sync should stay local when request candidate persistence is available" ); + let cacheable_request_body = serde_json::json!({ + "model": "claude-sonnet-4", + "system": [{ + "type": "text", + "text": format!("sync cacheable prompt {}", "cacheable prompt chunk ".repeat(300)), + "cache_control": {"type": "ephemeral"} + }], + "messages": [{"role": "user", "content": "reuse this Kiro prompt"}] + }) + .to_string(); + let (first_cache_status, first_cache_body) = send_kiro_request( + &gateway_url, + "trace-kiro-cli-local-sync-cache-1", + cacheable_request_body.clone(), + ) + .await; + assert!( + first_cache_status == StatusCode::OK, + "unexpected first cache status={first_cache_status} body={first_cache_body}" + ); + let first_usage = wait_for_completed_usage( + usage_repository.as_ref(), + "trace-kiro-cli-local-sync-cache-1", + ) + .await; + assert!( + first_usage.cache_creation_input_tokens > 0, + "first Kiro sync cacheable request should create simulated cache" + ); + assert_eq!(first_usage.cache_read_input_tokens, 0); + + let (second_cache_status, second_cache_body) = send_kiro_request( + &gateway_url, + "trace-kiro-cli-local-sync-cache-2", + cacheable_request_body, + ) + .await; + assert!( + second_cache_status == StatusCode::OK, + "unexpected second cache status={second_cache_status} body={second_cache_body}" + ); + let second_usage = wait_for_completed_usage( + usage_repository.as_ref(), + "trace-kiro-cli-local-sync-cache-2", + ) + .await; + assert!( + second_usage.cache_read_input_tokens > 0, + "second Kiro sync cacheable request should read simulated cache" + ); + assert_eq!(second_usage.cache_creation_input_tokens, 0); + assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0); assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0); assert_eq!(*public_hits.lock().expect("mutex should lock"), 0); diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/cli.rs b/apps/aether-gateway/src/tests/ai_execute/sync/cli.rs index 625445c9b..b5ebb4942 100644 --- a/apps/aether-gateway/src/tests/ai_execute/sync/cli.rs +++ b/apps/aether-gateway/src/tests/ai_execute/sync/cli.rs @@ -26,8 +26,40 @@ use aether_data_contracts::repository::provider_catalog::{ use base64::Engine as _; use sha2::{Digest, Sha256}; -#[tokio::test] -async fn gateway_executes_openai_responses_sync_via_local_decision_gate_with_local_sync_decision() { +const CLI_SYNC_TEST_STACK_BYTES: usize = 16 * 1024 * 1024; + +fn run_cli_sync_test(test_name: &'static str, make_future: F) +where + F: FnOnce() -> Fut + Send + 'static, + Fut: std::future::Future + 'static, +{ + let handle = std::thread::Builder::new() + .name(test_name.to_string()) + .stack_size(CLI_SYNC_TEST_STACK_BYTES) + .spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime should build"); + runtime.block_on(make_future()); + }) + .expect("cli sync test thread should spawn"); + + if let Err(payload) = handle.join() { + std::panic::resume_unwind(payload); + } +} + +#[test] +fn gateway_executes_openai_responses_sync_via_local_decision_gate_with_local_sync_decision() { + run_cli_sync_test( + "gateway_executes_openai_responses_sync_via_local_decision_gate_with_local_sync_decision", + gateway_executes_openai_responses_sync_via_local_decision_gate_with_local_sync_decision_impl, + ); +} + +async fn gateway_executes_openai_responses_sync_via_local_decision_gate_with_local_sync_decision_impl( +) { #[derive(Debug, Clone)] struct SeenExecutionRuntimeSyncRequest { trace_id: String, @@ -546,8 +578,15 @@ async fn gateway_executes_openai_responses_sync_via_local_decision_gate_with_loc upstream_handle.abort(); } -#[tokio::test] -async fn gateway_waits_for_api_key_concurrency_slot_then_executes_openai_responses_sync() { +#[test] +fn gateway_waits_for_api_key_concurrency_slot_then_executes_openai_responses_sync() { + run_cli_sync_test( + "gateway_waits_for_api_key_concurrency_slot_then_executes_openai_responses_sync", + gateway_waits_for_api_key_concurrency_slot_then_executes_openai_responses_sync_impl, + ); +} + +async fn gateway_waits_for_api_key_concurrency_slot_then_executes_openai_responses_sync_impl() { fn hash_api_key(value: &str) -> String { let mut hasher = Sha256::new(); hasher.update(value.as_bytes()); @@ -929,8 +968,16 @@ async fn gateway_waits_for_api_key_concurrency_slot_then_executes_openai_respons upstream_handle.abort(); } -#[tokio::test] -async fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_openai_responses_sync() { +#[test] +fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_openai_responses_sync() { + run_cli_sync_test( + "gateway_returns_concurrency_limited_after_wait_budget_expires_for_openai_responses_sync", + gateway_returns_concurrency_limited_after_wait_budget_expires_for_openai_responses_sync_impl, + ); +} + +async fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_openai_responses_sync_impl( +) { fn hash_api_key(value: &str) -> String { let mut hasher = Sha256::new(); hasher.update(value.as_bytes()); @@ -1278,8 +1325,15 @@ async fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_opena upstream_handle.abort(); } -#[tokio::test] -async fn gateway_returns_openai_responses_error_for_local_sync_failure() { +#[test] +fn gateway_returns_openai_responses_error_for_local_sync_failure() { + run_cli_sync_test( + "gateway_returns_openai_responses_error_for_local_sync_failure", + gateway_returns_openai_responses_error_for_local_sync_failure_impl, + ); +} + +async fn gateway_returns_openai_responses_error_for_local_sync_failure_impl() { fn hash_api_key(value: &str) -> String { let mut hasher = Sha256::new(); hasher.update(value.as_bytes()); @@ -1571,8 +1625,16 @@ async fn gateway_returns_openai_responses_error_for_local_sync_failure() { upstream_handle.abort(); } -#[tokio::test] -async fn gateway_returns_openai_responses_error_for_local_cross_format_gemini_cli_sync_failure() { +#[test] +fn gateway_returns_openai_responses_error_for_local_cross_format_gemini_cli_sync_failure() { + run_cli_sync_test( + "gateway_returns_openai_responses_error_for_local_cross_format_gemini_cli_sync_failure", + gateway_returns_openai_responses_error_for_local_cross_format_gemini_cli_sync_failure_impl, + ); +} + +async fn gateway_returns_openai_responses_error_for_local_cross_format_gemini_cli_sync_failure_impl( +) { #[derive(Debug, Clone)] struct SeenExecutionRuntimeSyncRequest { trace_id: String, @@ -1972,8 +2034,15 @@ async fn gateway_returns_openai_responses_error_for_local_cross_format_gemini_cl upstream_handle.abort(); } -#[tokio::test] -async fn gateway_returns_openai_responses_error_for_local_cross_format_claude_sync_failure() { +#[test] +fn gateway_returns_openai_responses_error_for_local_cross_format_claude_sync_failure() { + run_cli_sync_test( + "gateway_returns_openai_responses_error_for_local_cross_format_claude_sync_failure", + gateway_returns_openai_responses_error_for_local_cross_format_claude_sync_failure_impl, + ); +} + +async fn gateway_returns_openai_responses_error_for_local_cross_format_claude_sync_failure_impl() { #[derive(Debug, Clone)] struct SeenExecutionRuntimeSyncRequest { trace_id: String, @@ -2349,8 +2418,16 @@ async fn gateway_returns_openai_responses_error_for_local_cross_format_claude_sy upstream_handle.abort(); } -#[tokio::test] -async fn gateway_returns_openai_responses_error_for_local_cross_format_claude_chat_sync_failure() { +#[test] +fn gateway_returns_openai_responses_error_for_local_cross_format_claude_chat_sync_failure() { + run_cli_sync_test( + "gateway_returns_openai_responses_error_for_local_cross_format_claude_chat_sync_failure", + gateway_returns_openai_responses_error_for_local_cross_format_claude_chat_sync_failure_impl, + ); +} + +async fn gateway_returns_openai_responses_error_for_local_cross_format_claude_chat_sync_failure_impl( +) { #[derive(Debug, Clone)] struct SeenExecutionRuntimeSyncRequest { trace_id: String, @@ -2729,8 +2806,16 @@ async fn gateway_returns_openai_responses_error_for_local_cross_format_claude_ch upstream_handle.abort(); } -#[tokio::test] -async fn gateway_returns_openai_responses_error_for_local_cross_format_gemini_chat_sync_failure() { +#[test] +fn gateway_returns_openai_responses_error_for_local_cross_format_gemini_chat_sync_failure() { + run_cli_sync_test( + "gateway_returns_openai_responses_error_for_local_cross_format_gemini_chat_sync_failure", + gateway_returns_openai_responses_error_for_local_cross_format_gemini_chat_sync_failure_impl, + ); +} + +async fn gateway_returns_openai_responses_error_for_local_cross_format_gemini_chat_sync_failure_impl( +) { #[derive(Debug, Clone)] struct SeenExecutionRuntimeSyncRequest { trace_id: String, @@ -3109,8 +3194,15 @@ async fn gateway_returns_openai_responses_error_for_local_cross_format_gemini_ch upstream_handle.abort(); } -#[tokio::test] -async fn gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_refresh() { +#[test] +fn gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_refresh() { + run_cli_sync_test( + "gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_refresh", + gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_refresh_impl, + ); +} + +async fn gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_refresh_impl() { #[derive(Debug, Clone)] struct SeenExecutionRuntimeSyncRequest { trace_id: String, diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/pii_redaction_formats.rs b/apps/aether-gateway/src/tests/ai_execute/sync/pii_redaction_formats.rs index f313fff78..e8f555847 100644 --- a/apps/aether-gateway/src/tests/ai_execute/sync/pii_redaction_formats.rs +++ b/apps/aether-gateway/src/tests/ai_execute/sync/pii_redaction_formats.rs @@ -226,19 +226,24 @@ impl StandardFormat { } } -#[tokio::test] -async fn ai_execute_openai_responses_pii_redaction_round_trip_same_format() { - let (response_json, seen) = run_redaction_format_case(RedactionFormatCase { - test_id: "openai-responses-pii-redaction-same-format", - trace_id: "trace-openai-responses-pii-redaction-same-format", - client_format: StandardFormat::OpenAiResponses, - provider_format: StandardFormat::OpenAiResponses, - }) - .await; +#[test] +fn ai_execute_openai_responses_pii_redaction_round_trip_same_format() { + run_async_test_on_large_stack( + "ai_execute_openai_responses_pii_redaction_round_trip_same_format", + async { + let (response_json, seen) = run_redaction_format_case(RedactionFormatCase { + test_id: "openai-responses-pii-redaction-same-format", + trace_id: "trace-openai-responses-pii-redaction-same-format", + client_format: StandardFormat::OpenAiResponses, + provider_format: StandardFormat::OpenAiResponses, + }) + .await; - assert_provider_request_redacted(&seen, StandardFormat::OpenAiResponses); - assert!(seen.body.get("input").is_some()); - assert_restored_response(&response_json, StandardFormat::OpenAiResponses); + assert_provider_request_redacted(&seen, StandardFormat::OpenAiResponses); + assert!(seen.body.get("input").is_some()); + assert_restored_response(&response_json, StandardFormat::OpenAiResponses); + }, + ); } #[test] @@ -261,42 +266,52 @@ fn ai_execute_claude_messages_pii_redaction_round_trip_same_format() { ); } -#[tokio::test] -async fn ai_execute_openai_chat_pii_redaction_before_claude_conversion() { - let (response_json, seen) = run_redaction_format_case(RedactionFormatCase { - test_id: "openai-chat-pii-redaction-before-claude-conversion", - trace_id: "trace-openai-chat-pii-redaction-before-claude-conversion", - client_format: StandardFormat::OpenAiChat, - provider_format: StandardFormat::ClaudeMessages, - }) - .await; +#[test] +fn ai_execute_openai_chat_pii_redaction_before_claude_conversion() { + run_async_test_on_large_stack( + "ai_execute_openai_chat_pii_redaction_before_claude_conversion", + async { + let (response_json, seen) = run_redaction_format_case(RedactionFormatCase { + test_id: "openai-chat-pii-redaction-before-claude-conversion", + trace_id: "trace-openai-chat-pii-redaction-before-claude-conversion", + client_format: StandardFormat::OpenAiChat, + provider_format: StandardFormat::ClaudeMessages, + }) + .await; - assert_provider_request_redacted(&seen, StandardFormat::ClaudeMessages); - assert!(seen.body.get("messages").is_some()); - assert_eq!( - seen.body["model"], - StandardFormat::ClaudeMessages.provider_model() + assert_provider_request_redacted(&seen, StandardFormat::ClaudeMessages); + assert!(seen.body.get("messages").is_some()); + assert_eq!( + seen.body["model"], + StandardFormat::ClaudeMessages.provider_model() + ); + assert_restored_response(&response_json, StandardFormat::OpenAiChat); + }, ); - assert_restored_response(&response_json, StandardFormat::OpenAiChat); } -#[tokio::test] -async fn ai_execute_openai_responses_pii_redaction_before_claude_conversion() { - let (response_json, seen) = run_redaction_format_case(RedactionFormatCase { - test_id: "openai-responses-pii-redaction-before-claude-conversion", - trace_id: "trace-openai-responses-pii-redaction-before-claude-conversion", - client_format: StandardFormat::OpenAiResponses, - provider_format: StandardFormat::ClaudeMessages, - }) - .await; +#[test] +fn ai_execute_openai_responses_pii_redaction_before_claude_conversion() { + run_async_test_on_large_stack( + "ai_execute_openai_responses_pii_redaction_before_claude_conversion", + async { + let (response_json, seen) = run_redaction_format_case(RedactionFormatCase { + test_id: "openai-responses-pii-redaction-before-claude-conversion", + trace_id: "trace-openai-responses-pii-redaction-before-claude-conversion", + client_format: StandardFormat::OpenAiResponses, + provider_format: StandardFormat::ClaudeMessages, + }) + .await; - assert_provider_request_redacted(&seen, StandardFormat::ClaudeMessages); - assert!(seen.body.get("messages").is_some()); - assert_eq!( - seen.body["model"], - StandardFormat::ClaudeMessages.provider_model() + assert_provider_request_redacted(&seen, StandardFormat::ClaudeMessages); + assert!(seen.body.get("messages").is_some()); + assert_eq!( + seen.body["model"], + StandardFormat::ClaudeMessages.provider_model() + ); + assert_restored_response(&response_json, StandardFormat::OpenAiResponses); + }, ); - assert_restored_response(&response_json, StandardFormat::OpenAiResponses); } #[test] diff --git a/apps/aether-gateway/src/tests/architecture/admin_system.rs b/apps/aether-gateway/src/tests/architecture/admin_system.rs index c02e09f67..76f4207ea 100644 --- a/apps/aether-gateway/src/tests/architecture/admin_system.rs +++ b/apps/aether-gateway/src/tests/architecture/admin_system.rs @@ -28,7 +28,7 @@ fn admin_system_build_version_contract_uses_explicit_local_build_arg() { let deploy = read_workspace_file("deploy.sh"); for pattern in [ "detect_build_version()", - "git describe --tags --always --dirty", + "git describe --tags --match 'v[0-9]*' --always --dirty", "AETHER_BUILD_VERSION=\"${AETHER_BUILD_VERSION:-$(detect_build_version)}\"", "--build-arg \"AETHER_BUILD_VERSION=$AETHER_BUILD_VERSION\"", ">>> AETHER_BUILD_VERSION", @@ -43,6 +43,8 @@ fn admin_system_build_version_contract_uses_explicit_local_build_arg() { for pattern in [ "process.env.AETHER_BUILD_VERSION", "process.env.AETHER_VERSION", + "git describe --tags --match \"v[0-9]*\" --always --dirty", + "trimmed.startsWith('tunnel-v')", ] { assert!( vite_config.contains(pattern), @@ -60,6 +62,18 @@ fn admin_system_build_version_contract_uses_explicit_local_build_arg() { "api/core.rs should expose build version pattern {pattern}" ); } + + let build_rs = read_workspace_file("apps/aether-gateway/build.rs"); + for pattern in [ + "\"--match\"", + "\"v[0-9]*\"", + "trimmed.starts_with(\"tunnel-v\")", + ] { + assert!( + build_rs.contains(pattern), + "apps/aether-gateway/build.rs should ignore tunnel release tags for gateway version pattern {pattern}" + ); + } } #[test] diff --git a/apps/aether-gateway/src/tests/control/admin/system.rs b/apps/aether-gateway/src/tests/control/admin/system.rs index 577f7c627..557dcee59 100644 --- a/apps/aether-gateway/src/tests/control/admin/system.rs +++ b/apps/aether-gateway/src/tests/control/admin/system.rs @@ -1452,6 +1452,14 @@ async fn gateway_handles_admin_system_api_formats_locally_with_trusted_admin_pri assert!(formats.iter().any(|item| item["value"] == "jina:embedding")); assert!(formats.iter().any(|item| item["value"] == "jina:rerank")); assert!(formats.iter().any(|item| item["value"] == "gemini:video")); + let aliyun_embedding = formats + .iter() + .find(|item| item["value"] == "aliyun:multimodal_embedding") + .expect("aliyun multimodal embedding format should exist"); + assert_eq!( + aliyun_embedding["default_path"], + "/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding" + ); assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0); gateway_handle.abort(); diff --git a/apps/aether-gateway/src/tests/control/proxy/embeddings.rs b/apps/aether-gateway/src/tests/control/proxy/embeddings.rs index c51b3dedf..be9cf31b9 100644 --- a/apps/aether-gateway/src/tests/control/proxy/embeddings.rs +++ b/apps/aether-gateway/src/tests/control/proxy/embeddings.rs @@ -179,6 +179,129 @@ fn vertex_gemini_embedding_success_state(execution_runtime_url: String) -> AppSt .with_data_state_for_tests(data_state) } +fn aliyun_embedding_success_state(execution_runtime_url: String) -> AppState { + let mut snapshot = sample_currently_usable_auth_snapshot( + "key-aliyun-embedding-success", + "user-aliyun-embedding-success", + ); + snapshot.user_allowed_providers = None; + snapshot.api_key_allowed_providers = None; + snapshot.user_allowed_api_formats = Some(vec!["openai:embedding".to_string()]); + snapshot.api_key_allowed_api_formats = Some(vec!["openai:embedding".to_string()]); + snapshot.user_allowed_models = Some(vec!["qwen3-vl-embedding".to_string()]); + snapshot.api_key_allowed_models = Some(vec!["qwen3-vl-embedding".to_string()]); + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some(hash_api_key("sk-aliyun-embedding-success")), + snapshot, + )])); + let candidate_repository = + Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ + aliyun_embedding_candidate_row(), + ])); + let mut provider = sample_provider("provider-aliyun-embedding", "Aliyun DashScope", 1); + provider.provider_type = "aliyun".to_string(); + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![provider], + vec![sample_endpoint( + "endpoint-aliyun-embedding", + "provider-aliyun-embedding", + "aliyun:multimodal_embedding", + "https://dashscope.aliyuncs.com", + )], + vec![sample_key( + "key-upstream-aliyun-embedding", + "provider-aliyun-embedding", + "aliyun:multimodal_embedding", + "sk-upstream-aliyun-embedding", + )], + )); + let data_state = + GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests( + provider_catalog_repository, + candidate_repository, + ) + .with_auth_api_key_reader(auth_repository) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY); + + build_state_with_execution_runtime_override(execution_runtime_url) + .with_data_state_for_tests(data_state) +} + +fn mixed_embedding_success_state(execution_runtime_url: String) -> AppState { + let mut snapshot = sample_currently_usable_auth_snapshot( + "key-mixed-embedding-success", + "user-mixed-embedding-success", + ); + snapshot.user_allowed_providers = None; + snapshot.api_key_allowed_providers = None; + snapshot.user_allowed_api_formats = Some(vec!["openai:embedding".to_string()]); + snapshot.api_key_allowed_api_formats = Some(vec!["openai:embedding".to_string()]); + snapshot.user_allowed_models = Some(vec!["qwen3-vl-embedding".to_string()]); + snapshot.api_key_allowed_models = Some(vec!["qwen3-vl-embedding".to_string()]); + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some(hash_api_key("sk-mixed-embedding-success")), + snapshot, + )])); + + let mut openai_candidate = embedding_candidate_row(); + openai_candidate.model_id = "model-openai-qwen-vl-embedding".to_string(); + openai_candidate.global_model_id = "global-qwen3-vl-embedding".to_string(); + openai_candidate.global_model_name = "qwen3-vl-embedding".to_string(); + openai_candidate.model_provider_model_name = "openai-qwen-fallback".to_string(); + let candidate_repository = + Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ + openai_candidate, + aliyun_embedding_candidate_row(), + ])); + + let mut aliyun_provider = sample_provider("provider-aliyun-embedding", "Aliyun DashScope", 1); + aliyun_provider.provider_type = "aliyun".to_string(); + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![ + sample_provider("provider-embedding", "OpenAI Embeddings", 1), + aliyun_provider, + ], + vec![ + sample_endpoint( + "endpoint-embedding", + "provider-embedding", + "openai:embedding", + "https://api.openai.example", + ), + sample_endpoint( + "endpoint-aliyun-embedding", + "provider-aliyun-embedding", + "aliyun:multimodal_embedding", + "https://dashscope.aliyuncs.com", + ), + ], + vec![ + sample_key( + "key-upstream-embedding", + "provider-embedding", + "openai:embedding", + "sk-upstream-embedding", + ), + sample_key( + "key-upstream-aliyun-embedding", + "provider-aliyun-embedding", + "aliyun:multimodal_embedding", + "sk-upstream-aliyun-embedding", + ), + ], + )); + let data_state = + GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests( + provider_catalog_repository, + candidate_repository, + ) + .with_auth_api_key_reader(auth_repository) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY); + + build_state_with_execution_runtime_override(execution_runtime_url) + .with_data_state_for_tests(data_state) +} + fn gemini_embedding_conversion_execution_runtime() -> Router { Router::new().route( "/v1/execute/sync", @@ -189,6 +312,29 @@ fn gemini_embedding_conversion_execution_runtime() -> Router { ) } +fn aliyun_embedding_conversion_execution_runtime( + expected_contents: serde_json::Value, + expected_parameters: Option, +) -> Router { + let expected_contents = Arc::new(expected_contents); + let expected_parameters = Arc::new(expected_parameters); + Router::new().route( + "/v1/execute/sync", + any(move |Json(plan): Json| { + let expected_contents = Arc::clone(&expected_contents); + let expected_parameters = Arc::clone(&expected_parameters); + async move { + assert_openai_to_aliyun_embedding_execution_plan( + &plan, + &expected_contents, + expected_parameters.as_ref().as_ref(), + ); + Json(aliyun_embedding_execution_result(&plan)) + } + }), + ) +} + fn vertex_gemini_embedding_conversion_execution_runtime() -> Router { Router::new().route( "/v1/execute/sync", @@ -300,6 +446,40 @@ fn vertex_gemini_embedding_candidate_row() -> StoredMinimalCandidateSelectionRow row } +fn aliyun_embedding_candidate_row() -> StoredMinimalCandidateSelectionRow { + StoredMinimalCandidateSelectionRow { + provider_id: "provider-aliyun-embedding".to_string(), + provider_name: "Aliyun DashScope".to_string(), + provider_type: "aliyun".to_string(), + provider_priority: 1, + provider_is_active: true, + endpoint_id: "endpoint-aliyun-embedding".to_string(), + endpoint_api_format: "aliyun:multimodal_embedding".to_string(), + endpoint_api_family: Some("aliyun".to_string()), + endpoint_kind: Some("multimodal_embedding".to_string()), + endpoint_is_active: true, + key_id: "key-upstream-aliyun-embedding".to_string(), + key_name: "default".to_string(), + key_auth_type: "api_key".to_string(), + key_is_active: true, + key_api_formats: Some(vec!["aliyun:multimodal_embedding".to_string()]), + key_allowed_models: None, + key_capabilities: None, + key_internal_priority: 50, + key_global_priority_by_format: None, + model_id: "model-qwen3-vl-embedding".to_string(), + global_model_id: "global-qwen3-vl-embedding".to_string(), + global_model_name: "qwen3-vl-embedding".to_string(), + global_model_mappings: None, + global_model_supports_streaming: Some(false), + model_provider_model_name: "qwen3-vl-embedding".to_string(), + model_provider_model_mappings: None, + model_supports_streaming: Some(false), + model_is_active: true, + model_is_available: true, + } +} + fn assert_embedding_execution_plan(plan: &ExecutionPlan) { assert_eq!(plan.client_api_format, "openai:embedding"); assert_eq!(plan.provider_api_format, "openai:embedding"); @@ -311,6 +491,34 @@ fn assert_embedding_execution_plan(plan: &ExecutionPlan) { assert!(body.get("input").is_some()); } +fn assert_openai_to_aliyun_embedding_execution_plan( + plan: &ExecutionPlan, + expected_contents: &serde_json::Value, + expected_parameters: Option<&serde_json::Value>, +) { + assert_eq!(plan.client_api_format, "openai:embedding"); + assert_eq!(plan.provider_api_format, "aliyun:multimodal_embedding"); + assert_eq!(plan.method, "POST"); + assert_eq!( + plan.url, + "https://dashscope.aliyuncs.com/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding" + ); + assert_eq!( + plan.headers.get("authorization").map(String::as_str), + Some("Bearer sk-upstream-aliyun-embedding") + ); + assert_eq!(plan.model_name.as_deref(), Some("qwen3-vl-embedding")); + assert!(!plan.stream); + let body = plan.body.json_body.as_ref().expect("json request body"); + assert_eq!(body["model"], "qwen3-vl-embedding"); + assert_eq!(&body["input"]["contents"], expected_contents); + match expected_parameters { + Some(expected) => assert_eq!(&body["parameters"], expected), + None => assert!(body.get("parameters").is_none()), + } + assert!(body.get("messages").is_none()); +} + fn assert_openai_to_gemini_embedding_execution_plan(plan: &ExecutionPlan) { assert_eq!(plan.client_api_format, "openai:embedding"); assert_eq!(plan.provider_api_format, "gemini:embedding"); @@ -503,6 +711,41 @@ fn gemini_batch_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionRes } } +fn aliyun_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult { + ExecutionResult { + request_id: plan.request_id.clone(), + candidate_id: plan.candidate_id.clone(), + status_code: 200, + headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]), + body: Some(ResponseBody { + json_body: Some(json!({ + "output": { + "embeddings": [ + { + "index": 0, + "embedding": [0.1, 0.2, 0.3], + "type": "fusion" + } + ] + }, + "usage": { + "input_tokens": 432, + "input_tokens_details": { + "image_tokens": 402, + "text_tokens": 30 + }, + "output_tokens": 1, + "total_tokens": 433 + }, + "request_id": "aliyun-request-1" + })), + body_bytes_b64: None, + }), + telemetry: None, + error: None, + } +} + #[tokio::test] async fn embeddings_route_accepts_openai_payload() { let (execution_runtime_url, execution_runtime_handle) = @@ -714,6 +957,183 @@ async fn embeddings_route_converts_openai_batch_payload_to_gemini_batch_endpoint execution_runtime_handle.abort(); } +#[tokio::test] +async fn embeddings_route_converts_text_payload_to_aliyun_embedding_provider() { + let (execution_runtime_url, execution_runtime_handle) = + start_server(aliyun_embedding_conversion_execution_runtime( + json!([{ "text": "hello" }]), + Some(json!({ "dimension": 1024 })), + )) + .await; + let gateway = build_router_with_state(aliyun_embedding_success_state(execution_runtime_url)); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let response = reqwest::Client::new() + .post(format!("{gateway_url}/v1/embeddings")) + .header( + http::header::AUTHORIZATION, + "Bearer sk-aliyun-embedding-success", + ) + .json(&json!({ + "model": "qwen3-vl-embedding", + "input": "hello", + "dimensions": 1024 + })) + .send() + .await + .expect("request should succeed"); + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response + .headers() + .get(CONTROL_ENDPOINT_SIGNATURE_HEADER) + .and_then(|value| value.to_str().ok()), + Some("openai:embedding") + ); + let payload: serde_json::Value = response.json().await.expect("body should parse"); + assert_eq!(payload["object"], "list"); + assert_eq!(payload["request_id"], "aliyun-request-1"); + assert_eq!(payload["model"], "qwen3-vl-embedding"); + assert_eq!(payload["data"][0]["object"], "embedding"); + assert_eq!(payload["data"][0]["embedding"], json!([0.1, 0.2, 0.3])); + assert_eq!(payload["data"][0]["type"], "fusion"); + assert_eq!(payload["usage"]["prompt_tokens"], json!(432)); + assert_eq!(payload["usage"]["completion_tokens"], json!(1)); + assert_eq!(payload["usage"]["total_tokens"], json!(433)); + + gateway_handle.abort(); + execution_runtime_handle.abort(); +} + +#[tokio::test] +async fn embeddings_route_converts_multimodal_payload_to_aliyun_embedding_provider() { + let expected_contents = json!([ + { "text": "white running shoes" }, + { "image": "https://example.com/shoe.png" }, + { "multi_images": ["https://example.com/a.png", "https://example.com/b.png"] } + ]); + let (execution_runtime_url, execution_runtime_handle) = + start_server(aliyun_embedding_conversion_execution_runtime( + expected_contents.clone(), + Some(json!({ "res_level": 2, "max_video_frames": 64 })), + )) + .await; + let gateway = build_router_with_state(aliyun_embedding_success_state(execution_runtime_url)); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let response = reqwest::Client::new() + .post(format!("{gateway_url}/v1/embeddings")) + .header( + http::header::AUTHORIZATION, + "Bearer sk-aliyun-embedding-success", + ) + .json(&json!({ + "model": "qwen3-vl-embedding", + "input": expected_contents, + "parameters": { + "res_level": 2, + "max_video_frames": 64 + } + })) + .send() + .await + .expect("request should succeed"); + + assert_eq!(response.status(), StatusCode::OK); + let payload: serde_json::Value = response.json().await.expect("body should parse"); + assert_eq!(payload["data"][0]["embedding"], json!([0.1, 0.2, 0.3])); + assert_eq!(payload["data"][0]["type"], "fusion"); + + gateway_handle.abort(); + execution_runtime_handle.abort(); +} + +#[tokio::test] +async fn embeddings_route_skips_openai_candidate_for_multimodal_payload() { + let expected_contents = json!([ + { "text": "white running shoes" }, + { "image": "https://example.com/shoe.png" } + ]); + let (execution_runtime_url, execution_runtime_handle) = + start_server(aliyun_embedding_conversion_execution_runtime( + expected_contents.clone(), + Some(json!({ "enable_fusion": true })), + )) + .await; + let gateway = build_router_with_state(mixed_embedding_success_state(execution_runtime_url)); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let response = reqwest::Client::new() + .post(format!("{gateway_url}/v1/embeddings")) + .header( + http::header::AUTHORIZATION, + "Bearer sk-mixed-embedding-success", + ) + .json(&json!({ + "model": "qwen3-vl-embedding", + "input": expected_contents, + "parameters": { + "enable_fusion": true + } + })) + .send() + .await + .expect("request should succeed"); + + assert_eq!(response.status(), StatusCode::OK); + let payload: serde_json::Value = response.json().await.expect("body should parse"); + assert_eq!(payload["data"][0]["embedding"], json!([0.1, 0.2, 0.3])); + assert_eq!(payload["data"][0]["type"], "fusion"); + + gateway_handle.abort(); + execution_runtime_handle.abort(); +} + +#[tokio::test] +async fn embeddings_route_converts_fusion_payload_to_aliyun_embedding_provider() { + let expected_contents = json!([ + { + "text": "white running shoes", + "image": "https://example.com/shoe.png" + }, + { "video": "https://example.com/demo.mp4" } + ]); + let (execution_runtime_url, execution_runtime_handle) = + start_server(aliyun_embedding_conversion_execution_runtime( + expected_contents.clone(), + Some(json!({ "enable_fusion": true })), + )) + .await; + let gateway = build_router_with_state(aliyun_embedding_success_state(execution_runtime_url)); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let response = reqwest::Client::new() + .post(format!("{gateway_url}/v1/embeddings")) + .header( + http::header::AUTHORIZATION, + "Bearer sk-aliyun-embedding-success", + ) + .json(&json!({ + "model": "qwen3-vl-embedding", + "input": expected_contents, + "parameters": { + "enable_fusion": true + } + })) + .send() + .await + .expect("request should succeed"); + + assert_eq!(response.status(), StatusCode::OK); + let payload: serde_json::Value = response.json().await.expect("body should parse"); + assert_eq!(payload["data"][0]["embedding"], json!([0.1, 0.2, 0.3])); + assert_eq!(payload["data"][0]["type"], "fusion"); + + gateway_handle.abort(); + execution_runtime_handle.abort(); +} + #[tokio::test] async fn gemini_embed_content_route_uses_native_gemini_embedding_provider() { let (execution_runtime_url, execution_runtime_handle) = @@ -836,6 +1256,14 @@ async fn embeddings_route_rejects_invalid_local_payloads() { r#"{"model":"text-embedding-3-small","input":[[1],[]]}"#, "Embedding request input is required", ), + ( + r#"{"model":"text-embedding-3-small","input":[{}]}"#, + "Embedding request input is required", + ), + ( + r#"{"model":"text-embedding-3-small","input":[{"image":" "} ]}"#, + "Embedding request input is required", + ), ( r#"{"model":"text-embedding-3-small","input":"hello","stream":true}"#, "Embedding requests do not support streaming", diff --git a/apps/aether-gateway/src/tests/frontdoor/public_support.rs b/apps/aether-gateway/src/tests/frontdoor/public_support.rs index 61b50cdad..d1891e23b 100644 --- a/apps/aether-gateway/src/tests/frontdoor/public_support.rs +++ b/apps/aether-gateway/src/tests/frontdoor/public_support.rs @@ -24,6 +24,7 @@ use aether_data::repository::auth::{ use aether_data::repository::auth_modules::{ InMemoryAuthModuleReadRepository, StoredLdapModuleConfig, StoredOAuthProviderModuleConfig, }; +use aether_data::repository::billing::InMemoryBillingReadRepository; use aether_data::repository::management_tokens::{ InMemoryManagementTokenRepository, StoredManagementToken, StoredManagementTokenUserSummary, StoredManagementTokenWithUser, @@ -36,6 +37,10 @@ use aether_data::repository::users::{ use aether_data::repository::wallet::{ InMemoryWalletRepository, StoredWalletSnapshot, WalletWriteRepository, }; +use aether_data_contracts::repository::billing::{ + AdminBillingMutationOutcome, BillingPlanWriteInput, BillingReadRepository, + PaymentGatewayConfigWriteInput, +}; use aether_data_contracts::repository::global_models::StoredProviderActiveGlobalModel; use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository; use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UsageRepository}; @@ -3901,6 +3906,209 @@ async fn gateway_creates_wallet_recharge_orders_locally_without_proxying_upstrea upstream_handle.abort(); } +#[tokio::test] +async fn gateway_reuses_pending_billing_plan_checkout_order_without_proxying_upstream() { + let now = Utc::now(); + let user = StoredUserAuthRecord::new( + "user-billing-checkout-reuse".to_string(), + Some("billing-checkout-reuse@example.com".to_string()), + true, + "billing_checkout_reuse_user".to_string(), + Some("$2y$10$.OBQfixAECpsb8V/VS3csOMf00x2E/jD/gnud20t6RG0yiQosyOZ2".to_string()), + "user".to_string(), + "local".to_string(), + Some(json!(["openai"])), + Some(json!(["openai:chat"])), + Some(json!(["gpt-5"])), + true, + false, + Some(now), + Some(now), + ) + .expect("auth user should build"); + let wallet = StoredWalletSnapshot::new( + "wallet-billing-checkout-reuse".to_string(), + Some(user.id.clone()), + None, + 12.5, + 3.0, + "finite".to_string(), + "USD".to_string(), + "active".to_string(), + 20.0, + 4.5, + 0.0, + 0.0, + now.timestamp(), + ) + .expect("wallet should build"); + let access_token = build_test_auth_token( + "access", + serde_json::Map::from_iter([ + ("user_id".to_string(), json!(user.id.clone())), + ("role".to_string(), json!(user.role.clone())), + ( + "created_at".to_string(), + json!(user.created_at.map(|value| value.to_rfc3339())), + ), + ( + "session_id".to_string(), + json!("session-billing-checkout-reuse"), + ), + ]), + now + chrono::Duration::hours(1), + ); + let billing_repository = Arc::new(InMemoryBillingReadRepository::seed(Vec::new())); + let encrypted_key = encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "epay-secret") + .expect("merchant key should encrypt"); + let AdminBillingMutationOutcome::Applied(_) = billing_repository + .upsert_payment_gateway_config(&PaymentGatewayConfigWriteInput { + provider: "epay".to_string(), + enabled: true, + endpoint_url: "https://pay.example.com/".to_string(), + callback_base_url: Some("https://app.example.com".to_string()), + merchant_id: "merchant-1".to_string(), + merchant_key_encrypted: Some(encrypted_key), + preserve_existing_secret: false, + pay_currency: "CNY".to_string(), + usd_exchange_rate: 7.25, + min_recharge_usd: 1.0, + channels_json: json!([ + { + "channel": "alipay", + "display_name": "支付宝" + } + ]), + }) + .await + .expect("gateway config should create") + else { + panic!("gateway config should apply"); + }; + let plan = match billing_repository + .create_billing_plan(&BillingPlanWriteInput { + title: "每日额度月卡".to_string(), + description: Some("测试套餐".to_string()), + price_amount: 100.0, + price_currency: "CNY".to_string(), + duration_unit: "month".to_string(), + duration_value: 1, + enabled: true, + sort_order: 1, + max_active_per_user: 1, + purchase_limit_scope: "active_period".to_string(), + entitlements_json: json!([ + { + "type": "daily_quota", + "daily_quota_usd": 50.0, + "reset_timezone": "Asia/Shanghai", + "allow_wallet_overage": false + } + ]), + }) + .await + .expect("billing plan should create") + { + AdminBillingMutationOutcome::Applied(plan) => plan, + other => panic!("billing plan should apply, got {other:?}"), + }; + + let upstream_hits = Arc::new(Mutex::new(0usize)); + let upstream_hits_clone = Arc::clone(&upstream_hits); + let upstream = Router::new().route( + "/{*path}", + any(move |_request: Request| { + let upstream_hits_inner = Arc::clone(&upstream_hits_clone); + async move { + *upstream_hits_inner.lock().expect("mutex should lock") += 1; + (StatusCode::OK, Body::from("proxied")) + } + }), + ); + + let (_upstream_url, upstream_handle) = start_server(upstream).await; + let user_repository = Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![user])); + let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![wallet])); + let state = AppState::new() + .expect("gateway should build") + .with_data_state_for_tests(GatewayDataState::with_user_billing_and_wallet_for_tests( + user_repository, + billing_repository, + wallet_repository, + )) + .with_auth_sessions_for_tests([sample_auth_session( + "user-billing-checkout-reuse", + "session-billing-checkout-reuse", + "device-billing-checkout-reuse", + "refresh-token-placeholder", + now, + )]); + let gateway = build_router_with_state(state); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let client = reqwest::Client::new(); + let checkout_body = json!({ + "payment_provider": "epay", + "payment_method": "epay", + "payment_channel": "alipay", + }); + let first_response = client + .post(format!( + "{gateway_url}/api/billing/plans/{}/checkout", + plan.id + )) + .header("authorization", format!("Bearer {access_token}")) + .header("x-client-device-id", "device-billing-checkout-reuse") + .header("user-agent", "AetherTest/1.0") + .json(&checkout_body) + .send() + .await + .expect("first checkout request should succeed"); + assert_eq!(first_response.status(), StatusCode::OK); + let first_payload: serde_json::Value = first_response + .json() + .await + .expect("first checkout json should parse"); + let first_order_id = first_payload["order"]["id"] + .as_str() + .expect("first order id should exist") + .to_string(); + assert_eq!(first_payload["order"]["status"], "pending"); + assert_eq!(first_payload["order"]["product_id"], plan.id); + assert_eq!( + first_payload["reused_pending_order"], + serde_json::Value::Null + ); + + let second_response = client + .post(format!( + "{gateway_url}/api/billing/plans/{}/checkout", + plan.id + )) + .header("authorization", format!("Bearer {access_token}")) + .header("x-client-device-id", "device-billing-checkout-reuse") + .header("user-agent", "AetherTest/1.0") + .json(&checkout_body) + .send() + .await + .expect("second checkout request should succeed"); + assert_eq!(second_response.status(), StatusCode::OK); + let second_payload: serde_json::Value = second_response + .json() + .await + .expect("second checkout json should parse"); + assert_eq!(second_payload["order"]["id"], first_order_id); + assert_eq!(second_payload["reused_pending_order"], true); + assert_eq!( + second_payload["payment_instructions"], + first_payload["payment_instructions"] + ); + assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0); + + gateway_handle.abort(); + upstream_handle.abort(); +} + #[tokio::test] async fn gateway_creates_wallet_refunds_locally_without_proxying_upstream() { let now = Utc::now(); diff --git a/apps/aether-gateway/src/tests/usage/direct.rs b/apps/aether-gateway/src/tests/usage/direct.rs index e0b6e7d51..0dae5fa64 100644 --- a/apps/aether-gateway/src/tests/usage/direct.rs +++ b/apps/aether-gateway/src/tests/usage/direct.rs @@ -154,74 +154,79 @@ async fn gateway_records_usage_for_execution_runtime_sync_when_runtime_enabled() #[test] fn gateway_records_pending_usage_before_execution_runtime_sync_result_arrives() { - run_async_test_on_large_stack("pending-usage-sync-before-runtime-result", async move { - let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); - let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); - let execution_request_started = Arc::new(tokio::sync::Notify::new()); - let allow_execution_response = Arc::new(tokio::sync::Notify::new()); + run_async_test_on_large_stack( + "gateway_records_pending_usage_before_execution_runtime_sync_result_arrives", + gateway_records_pending_usage_before_execution_runtime_sync_result_arrives_impl(), + ); +} - let upstream = Router::new().route( - "/api/internal/gateway/report-sync", - any(|_request: Request| async move { Json(json!({"ok": true})) }), - ); +async fn gateway_records_pending_usage_before_execution_runtime_sync_result_arrives_impl() { + let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); + let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); + let execution_request_started = Arc::new(tokio::sync::Notify::new()); + let allow_execution_response = Arc::new(tokio::sync::Notify::new()); - let execution_runtime = Router::new().route( - "/v1/execute/sync", - any({ + let upstream = Router::new().route( + "/api/internal/gateway/report-sync", + any(|_request: Request| async move { Json(json!({"ok": true})) }), + ); + + let execution_runtime = Router::new().route( + "/v1/execute/sync", + any({ + let execution_request_started = Arc::clone(&execution_request_started); + let allow_execution_response = Arc::clone(&allow_execution_response); + move |_request: Request| { let execution_request_started = Arc::clone(&execution_request_started); let allow_execution_response = Arc::clone(&allow_execution_response); - move |_request: Request| { - let execution_request_started = Arc::clone(&execution_request_started); - let allow_execution_response = Arc::clone(&allow_execution_response); - async move { - execution_request_started.notify_one(); - allow_execution_response.notified().await; - Json(json!({ - "request_id": "req-usage-sync-pending-123", - "status_code": 200, - "headers": { - "content-type": "application/json" - }, - "body": { - "json_body": { - "id": "chatcmpl-usage-sync-pending-123", - "usage": { - "input_tokens": 3, - "output_tokens": 5, - "total_tokens": 8 - } + async move { + execution_request_started.notify_one(); + allow_execution_response.notified().await; + Json(json!({ + "request_id": "req-usage-sync-pending-123", + "status_code": 200, + "headers": { + "content-type": "application/json" + }, + "body": { + "json_body": { + "id": "chatcmpl-usage-sync-pending-123", + "usage": { + "input_tokens": 3, + "output_tokens": 5, + "total_tokens": 8 } - }, - "telemetry": { - "elapsed_ms": 45 } - })) - } + }, + "telemetry": { + "elapsed_ms": 45 + } + })) } - }), - ); + } + }), + ); - let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( - Some(hash_api_key("sk-client-openai-usage-sync-pending")), - sample_local_openai_auth_snapshot( - "api-key-usage-sync-pending-123", - "user-usage-sync-pending-123", - ), - )])); - let candidate_selection_repository = - Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ - sample_local_openai_candidate_row(), - ])); - let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( - vec![sample_local_openai_provider()], - vec![sample_local_openai_endpoint()], - vec![sample_local_openai_key()], - )); + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some(hash_api_key("sk-client-openai-usage-sync-pending")), + sample_local_openai_auth_snapshot( + "api-key-usage-sync-pending-123", + "user-usage-sync-pending-123", + ), + )])); + let candidate_selection_repository = + Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ + sample_local_openai_candidate_row(), + ])); + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![sample_local_openai_provider()], + vec![sample_local_openai_endpoint()], + vec![sample_local_openai_key()], + )); - let (upstream_url, upstream_handle) = start_server(upstream).await; - let (execution_runtime_url, execution_runtime_handle) = - start_server(execution_runtime).await; - let gateway_state = + let (upstream_url, upstream_handle) = start_server(upstream).await; + let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; + let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url) .with_data_state_for_tests( GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests( @@ -237,76 +242,74 @@ fn gateway_records_pending_usage_before_execution_runtime_sync_result_arrives() enabled: true, ..UsageRuntimeConfig::default() }); - let gateway = build_router_with_state(gateway_state); - let (gateway_url, gateway_handle) = start_server(gateway).await; + let gateway = build_router_with_state(gateway_state); + let (gateway_url, gateway_handle) = start_server(gateway).await; - let request_task = tokio::spawn({ - let gateway_url = gateway_url.clone(); - async move { - let response = reqwest::Client::new() - .post(format!("{gateway_url}/v1/chat/completions")) - .header(http::header::CONTENT_TYPE, "application/json") - .header( - http::header::AUTHORIZATION, - "Bearer sk-client-openai-usage-sync-pending", - ) - .header(TRACE_ID_HEADER, "req-usage-sync-pending-123") - .body("{\"model\":\"gpt-5\",\"messages\":[]}") - .send() - .await - .expect("request should succeed"); - let status = response.status(); - let body = response.text().await.expect("body should read"); - (status, body) - } - }); - - execution_request_started.notified().await; - - let mut pending = None; - for _ in 0..50 { - pending = usage_repository - .find_by_request_id("req-usage-sync-pending-123") + let request_task = tokio::spawn({ + let gateway_url = gateway_url.clone(); + async move { + let response = reqwest::Client::new() + .post(format!("{gateway_url}/v1/chat/completions")) + .header(http::header::CONTENT_TYPE, "application/json") + .header( + http::header::AUTHORIZATION, + "Bearer sk-client-openai-usage-sync-pending", + ) + .header(TRACE_ID_HEADER, "req-usage-sync-pending-123") + .body("{\"model\":\"gpt-5\",\"messages\":[]}") + .send() .await - .expect("usage lookup should succeed"); - if pending - .as_ref() - .is_some_and(|stored| stored.status == "pending") - { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(10)).await; + .expect("request should succeed"); + let status = response.status(); + let body = response.text().await.expect("body should read"); + (status, body) } - let pending = - pending.expect("pending usage should be recorded before sync result resolves"); - assert_eq!(pending.status, "pending"); - assert_eq!(pending.billing_status, "pending"); - assert_eq!(pending.response_time_ms, None); - - allow_execution_response.notify_one(); - - let (status, _body) = request_task.await.expect("request task should join"); - assert_eq!(status, StatusCode::OK); - - let mut stored = None; - for _ in 0..50 { - stored = usage_repository - .find_by_request_id("req-usage-sync-pending-123") - .await - .expect("usage lookup should succeed"); - if stored.as_ref().is_some_and(|row| row.status == "completed") { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(10)).await; - } - let stored = stored.expect("usage should be finalized"); - assert_eq!(stored.status, "completed"); - assert_eq!(stored.response_time_ms, Some(45)); - - gateway_handle.abort(); - execution_runtime_handle.abort(); - upstream_handle.abort(); }); + + execution_request_started.notified().await; + + let mut pending = None; + for _ in 0..50 { + pending = usage_repository + .find_by_request_id("req-usage-sync-pending-123") + .await + .expect("usage lookup should succeed"); + if pending + .as_ref() + .is_some_and(|stored| stored.status == "pending") + { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + let pending = pending.expect("pending usage should be recorded before sync result resolves"); + assert_eq!(pending.status, "pending"); + assert_eq!(pending.billing_status, "pending"); + assert_eq!(pending.response_time_ms, None); + + allow_execution_response.notify_one(); + + let (status, _body) = request_task.await.expect("request task should join"); + assert_eq!(status, StatusCode::OK); + + let mut stored = None; + for _ in 0..50 { + stored = usage_repository + .find_by_request_id("req-usage-sync-pending-123") + .await + .expect("usage lookup should succeed"); + if stored.as_ref().is_some_and(|row| row.status == "completed") { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + let stored = stored.expect("usage should be finalized"); + assert_eq!(stored.status, "completed"); + assert_eq!(stored.response_time_ms, Some(45)); + + gateway_handle.abort(); + execution_runtime_handle.abort(); + upstream_handle.abort(); } #[tokio::test] diff --git a/crates/aether-admin/src/provider/models.rs b/crates/aether-admin/src/provider/models.rs index 72e039fee..0f51c811b 100644 --- a/crates/aether-admin/src/provider/models.rs +++ b/crates/aether-admin/src/provider/models.rs @@ -8,6 +8,7 @@ const EMBEDDING_API_FORMATS: &[&str] = &[ "jina:embedding", "gemini:embedding", "doubao:embedding", + "aliyun:multimodal_embedding", ]; fn unix_secs_to_rfc3339(unix_secs: u64) -> Option { diff --git a/crates/aether-admin/src/system.rs b/crates/aether-admin/src/system.rs index 2c0a09843..8d0b65a23 100644 --- a/crates/aether-admin/src/system.rs +++ b/crates/aether-admin/src/system.rs @@ -835,6 +835,18 @@ const ADMIN_API_FORMAT_DEFINITIONS: &[AdminApiFormatDefinition] = &[ default_path: "/v1/embeddings", aliases: &["doubao_embedding"], }, + AdminApiFormatDefinition { + value: "aliyun:multimodal_embedding", + label: "Aliyun Multimodal Embedding", + default_path: "/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding", + aliases: &[ + "aliyun_embedding", + "aliyun_multimodal_embedding", + "dashscope_embedding", + "dashscope_multimodal_embedding", + "dashscope:multimodal_embedding", + ], + }, ]; pub fn build_admin_system_check_update_payload(current_version: String) -> serde_json::Value { diff --git a/crates/aether-ai-formats/src/formats/aliyun/embedding/mod.rs b/crates/aether-ai-formats/src/formats/aliyun/embedding/mod.rs new file mode 100644 index 000000000..e0062185c --- /dev/null +++ b/crates/aether-ai-formats/src/formats/aliyun/embedding/mod.rs @@ -0,0 +1,2 @@ +pub mod request; +pub mod response; diff --git a/crates/aether-ai-formats/src/formats/aliyun/embedding/request.rs b/crates/aether-ai-formats/src/formats/aliyun/embedding/request.rs new file mode 100644 index 000000000..f9d09c768 --- /dev/null +++ b/crates/aether-ai-formats/src/formats/aliyun/embedding/request.rs @@ -0,0 +1,256 @@ +use serde_json::{Map, Value}; + +use crate::formats::context::FormatContext; +use crate::formats::openai::embedding::request::mapped_embedding_model; +use crate::protocol::canonical::{ + CanonicalEmbeddingContent, CanonicalEmbeddingInput, CanonicalRequest, +}; + +pub fn to(request: &CanonicalRequest, ctx: &FormatContext) -> Option { + let embedding = request.embedding.as_ref()?; + let contents = embedding_input_to_contents(&embedding.input)?; + if contents.is_empty() { + return None; + } + + let mut output = Map::new(); + output.insert( + "model".to_string(), + Value::String(mapped_embedding_model( + request, + ctx.mapped_model_or(request.model.as_str()), + )), + ); + output.insert( + "input".to_string(), + Value::Object(Map::from_iter([( + "contents".to_string(), + Value::Array(contents), + )])), + ); + + let mut parameters = embedding.parameters.clone().unwrap_or_default(); + if let Some(dimensions) = embedding.dimensions { + parameters + .entry("dimension".to_string()) + .or_insert_with(|| Value::from(dimensions)); + } + if !parameters.is_empty() { + output.insert("parameters".to_string(), Value::Object(parameters)); + } + + Some(Value::Object(output)) +} + +fn embedding_input_to_contents(input: &CanonicalEmbeddingInput) -> Option> { + match input { + CanonicalEmbeddingInput::String(text) => { + non_empty_text_content(text).map(|content| vec![content]) + } + CanonicalEmbeddingInput::StringArray(items) => items + .iter() + .map(|text| non_empty_text_content(text)) + .collect(), + CanonicalEmbeddingInput::Multimodal(items) => { + items.iter().map(multimodal_content_to_value).collect() + } + CanonicalEmbeddingInput::TokenArray(_) | CanonicalEmbeddingInput::TokenArrayArray(_) => { + None + } + } +} + +fn non_empty_text_content(text: &str) -> Option { + let text = text.trim(); + if text.is_empty() { + None + } else { + Some(Value::Object(Map::from_iter([( + "text".to_string(), + Value::String(text.to_string()), + )]))) + } +} + +fn multimodal_content_to_value(content: &CanonicalEmbeddingContent) -> Option { + if content.is_empty() { + return None; + } + let mut object = Map::new(); + if let Some(text) = content + .text + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + object.insert("text".to_string(), Value::String(text.to_string())); + } + if let Some(image) = content + .image + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + object.insert("image".to_string(), Value::String(image.to_string())); + } + if let Some(video) = content + .video + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + object.insert("video".to_string(), Value::String(video.to_string())); + } + if let Some(multi_images) = content + .multi_images + .as_ref() + .filter(|values| !values.is_empty() && values.iter().all(|value| !value.trim().is_empty())) + { + object.insert( + "multi_images".to_string(), + Value::Array( + multi_images + .iter() + .map(|value| Value::String(value.trim().to_string())) + .collect(), + ), + ); + } + if object.is_empty() { + None + } else { + Some(Value::Object(object)) + } +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use serde_json::{json, Map, Value}; + + use super::to; + use crate::formats::context::FormatContext; + use crate::protocol::canonical::{ + CanonicalEmbeddingContent, CanonicalEmbeddingInput, CanonicalEmbeddingRequest, + CanonicalRequest, + }; + + fn canonical_embedding(input: CanonicalEmbeddingInput) -> CanonicalRequest { + CanonicalRequest { + model: "text-embedding-3-small".to_string(), + embedding: Some(CanonicalEmbeddingRequest { + input, + encoding_format: None, + dimensions: None, + task: None, + user: None, + parameters: None, + extensions: BTreeMap::new(), + }), + ..CanonicalRequest::default() + } + } + + #[test] + fn text_input_uses_dashscope_contents() { + let request = canonical_embedding(CanonicalEmbeddingInput::StringArray(vec![ + "alpha".to_string(), + "beta".to_string(), + ])); + + let body = to( + &request, + &FormatContext::default().with_mapped_model("qwen3-vl-embedding"), + ) + .expect("aliyun request"); + + assert_eq!(body["model"], "qwen3-vl-embedding"); + assert_eq!( + body["input"]["contents"], + json!([{ "text": "alpha" }, { "text": "beta" }]) + ); + } + + #[test] + fn multimodal_input_and_parameters_use_dashscope_contract() { + let mut request = canonical_embedding(CanonicalEmbeddingInput::Multimodal(vec![ + CanonicalEmbeddingContent { + text: Some("white running shoes".to_string()), + image: None, + video: None, + multi_images: None, + }, + CanonicalEmbeddingContent { + text: None, + image: Some("https://example.com/shoe.png".to_string()), + video: None, + multi_images: None, + }, + CanonicalEmbeddingContent { + text: None, + image: None, + video: None, + multi_images: Some(vec![ + "https://example.com/a.png".to_string(), + "https://example.com/b.png".to_string(), + ]), + }, + ])); + let embedding = request.embedding.as_mut().expect("embedding request"); + embedding.dimensions = Some(1024); + embedding.parameters = Some(Map::from_iter([ + ("enable_fusion".to_string(), Value::Bool(true)), + ("res_level".to_string(), Value::from(2_u64)), + ("max_video_frames".to_string(), Value::from(64_u64)), + ])); + + let body = to( + &request, + &FormatContext::default().with_mapped_model("qwen3-vl-embedding"), + ) + .expect("aliyun request"); + + assert_eq!( + body["input"]["contents"], + json!([ + { "text": "white running shoes" }, + { "image": "https://example.com/shoe.png" }, + { "multi_images": ["https://example.com/a.png", "https://example.com/b.png"] } + ]) + ); + assert_eq!(body["parameters"]["dimension"], 1024); + assert_eq!(body["parameters"]["enable_fusion"], true); + assert_eq!(body["parameters"]["res_level"], 2); + assert_eq!(body["parameters"]["max_video_frames"], 64); + } + + #[test] + fn parameter_dimension_wins_over_openai_dimensions() { + let mut request = canonical_embedding(CanonicalEmbeddingInput::String("alpha".to_string())); + let embedding = request.embedding.as_mut().expect("embedding request"); + embedding.dimensions = Some(1024); + embedding.parameters = Some(Map::from_iter([( + "dimension".to_string(), + Value::from(512_u64), + )])); + + let body = to( + &request, + &FormatContext::default().with_mapped_model("qwen3-vl-embedding"), + ) + .expect("aliyun request"); + + assert_eq!(body["parameters"]["dimension"], 512); + } + + #[test] + fn token_arrays_are_not_convertible() { + let request = canonical_embedding(CanonicalEmbeddingInput::TokenArray(vec![1, 2, 3])); + assert!(to( + &request, + &FormatContext::default().with_mapped_model("qwen3-vl-embedding"), + ) + .is_none()); + } +} diff --git a/crates/aether-ai-formats/src/formats/aliyun/embedding/response.rs b/crates/aether-ai-formats/src/formats/aliyun/embedding/response.rs new file mode 100644 index 000000000..4a8b5b2fc --- /dev/null +++ b/crates/aether-ai-formats/src/formats/aliyun/embedding/response.rs @@ -0,0 +1,134 @@ +use std::collections::BTreeMap; + +use serde_json::{Map, Value}; + +use crate::formats::openai::embedding::request::namespace_extensions; +use crate::protocol::canonical::{CanonicalEmbedding, CanonicalEmbeddingResponse, CanonicalUsage}; + +pub fn from(body_json: &Value) -> Option { + let body = body_json.as_object()?; + if body.contains_key("error") || body.contains_key("code") && body.contains_key("message") { + return None; + } + let data = body + .get("output")? + .as_object()? + .get("embeddings")? + .as_array()?; + let mut embeddings = Vec::new(); + for (fallback_index, item) in data.iter().enumerate() { + let item_object = item.as_object()?; + let values = item_object.get("embedding")?.as_array()?; + let embedding = values + .iter() + .map(Value::as_f64) + .collect::>>()?; + let mut extensions = + namespace_extensions("aliyun", item_object, &["index", "embedding", "type"]); + if let Some(value) = item_object.get("type").cloned() { + extensions.insert( + "openai".to_string(), + Value::Object(Map::from_iter([("type".to_string(), value)])), + ); + } + embeddings.push(CanonicalEmbedding { + index: item_object + .get("index") + .and_then(Value::as_u64) + .and_then(|value| usize::try_from(value).ok()) + .unwrap_or(fallback_index), + embedding, + extensions, + }); + } + + let request_id = body.get("request_id").and_then(Value::as_str); + let mut extensions = + namespace_extensions("aliyun", body, &["output", "usage", "request_id", "model"]); + if let Some(request_id) = request_id { + extensions.insert( + "openai".to_string(), + Value::Object(Map::from_iter([( + "request_id".to_string(), + Value::String(request_id.to_string()), + )])), + ); + } + + Some(CanonicalEmbeddingResponse { + id: request_id.unwrap_or("aliyun-request-unknown").to_string(), + model: body + .get("model") + .and_then(Value::as_str) + .unwrap_or("unknown") + .to_string(), + embeddings, + usage: aliyun_usage_to_canonical(body.get("usage")), + extensions, + }) +} + +fn aliyun_usage_to_canonical(value: Option<&Value>) -> Option { + let usage = value?.as_object()?; + let input_tokens = usage + .get("input_tokens") + .and_then(Value::as_u64) + .unwrap_or(0); + let output_tokens = usage + .get("output_tokens") + .and_then(Value::as_u64) + .unwrap_or(0); + Some(CanonicalUsage { + input_tokens, + output_tokens, + total_tokens: usage + .get("total_tokens") + .and_then(Value::as_u64) + .unwrap_or(input_tokens.saturating_add(output_tokens)), + extensions: BTreeMap::from([("aliyun".to_string(), Value::Object(usage.clone()))]), + ..CanonicalUsage::default() + }) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::from; + use crate::formats::openai::embedding::response::to as to_openai; + + #[test] + fn parses_dashscope_embeddings_to_openai_compatible_shape() { + let body = json!({ + "output": { + "embeddings": [ + { + "index": 0, + "embedding": [0.1, 0.2, 0.3], + "type": "fused" + } + ] + }, + "usage": { + "input_tokens": 432, + "input_tokens_details": { + "image_tokens": 402, + "text_tokens": 30 + }, + "output_tokens": 1, + "total_tokens": 433 + }, + "request_id": "aliyun-request-1" + }); + + let canonical = from(&body).expect("aliyun response"); + let emitted = to_openai(&canonical).expect("openai response"); + + assert_eq!(emitted["request_id"], "aliyun-request-1"); + assert_eq!(emitted["data"][0]["embedding"], json!([0.1, 0.2, 0.3])); + assert_eq!(emitted["data"][0]["type"], "fused"); + assert_eq!(emitted["usage"]["prompt_tokens"], 432); + assert_eq!(emitted["usage"]["completion_tokens"], 1); + assert_eq!(emitted["usage"]["total_tokens"], 433); + } +} diff --git a/crates/aether-ai-formats/src/formats/aliyun/mod.rs b/crates/aether-ai-formats/src/formats/aliyun/mod.rs new file mode 100644 index 000000000..1768e8be6 --- /dev/null +++ b/crates/aether-ai-formats/src/formats/aliyun/mod.rs @@ -0,0 +1 @@ +pub mod embedding; diff --git a/crates/aether-ai-formats/src/formats/claude/messages/request.rs b/crates/aether-ai-formats/src/formats/claude/messages/request.rs index 798b89159..e8a3e2408 100644 --- a/crates/aether-ai-formats/src/formats/claude/messages/request.rs +++ b/crates/aether-ai-formats/src/formats/claude/messages/request.rs @@ -16,8 +16,8 @@ use crate::{ claude_generation_config, claude_messages_to_canonical, claude_parallel_tool_calls, claude_system_to_canonical_instructions, claude_thinking_to_canonical, claude_tool_choice_to_canonical, claude_tools_to_canonical, - compact_canonical_claude_messages, insert_f64, namespace_extension_object, - CanonicalRequest, + compact_canonical_claude_messages, insert_f64, mark_claude_messages_request_source, + namespace_extension_object, CanonicalRequest, }, }; @@ -85,6 +85,7 @@ pub fn from_raw(body_json: &Value) -> Option { "output_config", ], ); + mark_claude_messages_request_source(&mut canonical.extensions); if !builtin_tools.is_empty() { canonical_extension_object_mut(&mut canonical.extensions, "claude") .insert("builtin_tools".to_string(), Value::Array(builtin_tools)); @@ -115,7 +116,7 @@ pub fn to_raw( ); output.insert( "max_tokens".to_string(), - Value::from(canonical.generation.max_tokens.unwrap_or(1024)), + Value::from(canonical.generation.max_tokens.unwrap_or(8192)), ); if let Some(system) = canonical_instructions_to_claude_system(&canonical.instructions) { output.insert("system".to_string(), system); diff --git a/crates/aether-ai-formats/src/formats/conversion/request.rs b/crates/aether-ai-formats/src/formats/conversion/request.rs index 16b7b44cd..c556c963f 100644 --- a/crates/aether-ai-formats/src/formats/conversion/request.rs +++ b/crates/aether-ai-formats/src/formats/conversion/request.rs @@ -159,10 +159,14 @@ fn request_context(mapped_model: &str, upstream_is_stream: bool) -> FormatContex mod tests { use serde_json::{json, Value}; + use crate::formats::{context::FormatContext, registry}; + use super::{ convert_openai_chat_request_to_claude_request, convert_openai_chat_request_to_openai_responses_request, normalize_claude_request_to_openai_chat_request, + normalize_gemini_request_to_openai_chat_request, + normalize_openai_responses_request_to_openai_chat_request, }; #[test] @@ -216,6 +220,185 @@ mod tests { assert_eq!(converted["messages"][0]["content"], "hello"); } + #[test] + fn claude_request_to_chat_clamps_max_reasoning_effort_to_high() { + let body = json!({ + "model": "claude-sonnet", + "messages": [{"role": "user", "content": "hello"}], + "thinking": {"type": "enabled", "budget_tokens": 1024}, + "output_config": {"effort": "max"}, + "max_tokens": 128, + }); + + let converted = + normalize_claude_request_to_openai_chat_request(&body).expect("openai chat request"); + + assert_eq!(converted["reasoning_effort"], "high"); + } + + #[test] + fn gemini_request_to_chat_clamps_xhigh_reasoning_effort_to_high() { + let body = json!({ + "contents": [{ + "role": "user", + "parts": [{"text": "hello"}] + }], + "generationConfig": { + "thinkingConfig": {"thinkingBudget": 8192} + } + }); + + let converted = normalize_gemini_request_to_openai_chat_request( + &body, + "/v1beta/models/gemini-2.5-pro:generateContent", + ) + .expect("openai chat request"); + + assert_eq!(converted["reasoning_effort"], "high"); + } + + #[test] + fn responses_request_normalizer_keeps_tool_history_chat_safe() { + let call_id_one = "call_weather_123"; + let call_id_two = "call_lookup_456"; + let tool_output_one = json!({ + "toolCallId": call_id_one, + "input": {"city": "Hangzhou"}, + "output": { + "content": [{"type": "text", "text": "sunny"}], + "isError": false, + }, + }); + let body = json!({ + "model": "glm-5.1", + "input": [ + "weather now", + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "thinking first"}] + }, + { + "type": "message", + "role": "assistant", + "content": "planning" + }, + { + "type": "function_call", + "call_id": call_id_one, + "id": call_id_one, + "name": "mcp__mapsWeather", + "arguments": "{\"city\":\"Hangzhou\"}" + }, + { + "type": "web_search_call", + "id": "ignored_web_search", + "action": {"query": "should be skipped"} + }, + { + "type": "function_call", + "call_id": call_id_two, + "id": call_id_two, + "name": "mcp__lookupData", + "arguments": "{\"query\":\"museum\"}" + }, + { + "type": "function_call_output", + "call_id": call_id_one, + "output": tool_output_one.to_string() + }, + { + "type": "function_call_output", + "call_id": call_id_two, + "output": "done-2" + } + ] + }); + + let converted = normalize_openai_responses_request_to_openai_chat_request(&body) + .expect("openai chat request"); + let messages = converted["messages"].as_array().expect("messages"); + + assert_eq!(messages.len(), 4); + assert_eq!(messages[0]["role"], "user"); + assert_eq!(messages[0]["content"], "weather now"); + assert_eq!(messages[1]["role"], "assistant"); + assert_eq!(messages[1]["reasoning_content"], "thinking first"); + assert_eq!(messages[1]["content"], "planning"); + assert_eq!(messages[1]["tool_calls"].as_array().unwrap().len(), 2); + assert_eq!(messages[1]["tool_calls"][0]["id"], call_id_one); + assert_eq!( + messages[1]["tool_calls"][0]["function"]["name"], + "mcp__mapsWeather" + ); + assert_eq!(messages[1]["tool_calls"][1]["id"], call_id_two); + assert_eq!( + messages[1]["tool_calls"][1]["function"]["name"], + "mcp__lookupData" + ); + assert_eq!(messages[2]["role"], "tool"); + assert_eq!(messages[2]["tool_call_id"], call_id_one); + let content = messages[2]["content"] + .as_str() + .expect("tool result content should stay a string"); + assert_eq!( + serde_json::from_str::(content).expect("tool output json"), + tool_output_one + ); + assert_eq!(messages[3]["role"], "tool"); + assert_eq!(messages[3]["tool_call_id"], call_id_two); + assert_eq!(messages[3]["content"], "done-2"); + } + + #[test] + fn responses_request_normalizer_emits_empty_message_content_as_empty_string() { + let body = json!({ + "model": "glm-5.1", + "input": [ + { + "type": "message", + "role": "assistant", + "content": null + } + ] + }); + + let converted = normalize_openai_responses_request_to_openai_chat_request(&body) + .expect("openai chat request"); + let messages = converted["messages"].as_array().expect("messages"); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0]["role"], "assistant"); + assert_eq!(messages[0]["content"], ""); + } + + #[test] + fn responses_request_normalizer_clamps_chat_reasoning_effort_and_filters_extensions() { + let body = json!({ + "model": "gpt-5.1", + "input": "hello", + "reasoning": {"effort": "xhigh"}, + "text": {"verbosity": "high"}, + "include": ["reasoning.encrypted_content"], + "store": false, + "service_tier": "priority", + "prompt_cache_key": "cache_123", + "safety_identifier": "user_123" + }); + + let converted = normalize_openai_responses_request_to_openai_chat_request(&body) + .expect("openai chat request"); + + assert_eq!(converted["reasoning_effort"], "high"); + assert_eq!(converted["verbosity"], "high"); + assert_eq!(converted["service_tier"], "priority"); + assert_eq!(converted["prompt_cache_key"], "cache_123"); + assert_eq!(converted["safety_identifier"], "user_123"); + assert!(converted.get("include").is_none()); + assert!(converted.get("store").is_none()); + assert!(converted.get("text").is_none()); + assert!(converted.get("reasoning").is_none()); + } + #[test] fn request_normalizer_preserves_multiple_claude_tool_results() { let body = json!({ @@ -550,4 +733,203 @@ mod tests { assert!(!block_content_json.contains("\"source\"")); assert!(!block_content_json.contains("document body")); } + + #[test] + fn claude_request_to_responses_uses_developer_system_and_sub2api_defaults() { + let body = json!({ + "model": "claude-sonnet", + "system": [{ + "type": "text", + "text": "Be exact.", + "cache_control": {"type": "ephemeral"} + }], + "messages": [ + {"role": "user", "content": "hello"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "private plan", "signature": "sig_hidden"}, + {"type": "text", "text": "visible answer"}, + { + "type": "tool_use", + "id": "toolu_calc", + "name": "calc", + "input": {"x": 1} + } + ] + } + ], + "tools": [ + {"name": "implicit_empty", "description": "empty"}, + {"name": "object_empty", "input_schema": {"type": "object"}} + ], + "thinking": {"type": "enabled", "budget_tokens": 4096}, + "temperature": 0.2, + "top_p": 0.9, + "max_tokens": 10, + }); + + let converted = registry::convert_request( + "claude:messages", + "openai:responses", + &body, + &FormatContext::default().with_mapped_model("gpt-5.1"), + ) + .expect("responses request"); + + assert_eq!(converted["model"], "gpt-5.1"); + assert!(converted.get("temperature").is_none()); + assert!(converted.get("top_p").is_none()); + assert!(converted.get("instructions").is_none()); + assert_eq!(converted["text"]["verbosity"], "medium"); + assert_eq!(converted["reasoning"]["effort"], "medium"); + assert_eq!(converted["reasoning"]["summary"], "auto"); + assert_eq!(converted["max_output_tokens"], 128); + assert_eq!(converted["store"], false); + assert_eq!(converted["parallel_tool_calls"], true); + assert!(converted["include"] + .as_array() + .expect("include") + .iter() + .any(|value| value.as_str() == Some("reasoning.encrypted_content"))); + + let input = converted["input"].as_array().expect("responses input"); + assert_eq!(input[0]["role"], "developer"); + assert_eq!(input[0]["content"][0]["type"], "input_text"); + assert_eq!(input[0]["content"][0]["text"], "Be exact."); + assert_eq!( + input[0]["content"][0]["cache_control"], + json!({"type": "ephemeral"}) + ); + let input_json = Value::Array(input.clone()).to_string(); + assert!(input_json.contains("visible answer")); + assert!(!input_json.contains("private plan")); + assert!(!input_json.contains("sig_hidden")); + + let tools = converted["tools"].as_array().expect("tools"); + assert_eq!(tools.len(), 2); + for tool in tools { + assert_eq!(tool["parameters"]["type"], "object"); + assert!(tool["parameters"]["properties"].is_object()); + } + } + + #[test] + fn claude_output_config_effort_controls_responses_reasoning() { + let body = json!({ + "model": "claude-sonnet", + "messages": [{"role": "user", "content": "hello"}], + "thinking": {"type": "enabled", "budget_tokens": 1024}, + "output_config": {"effort": "max"}, + "max_tokens": 128, + }); + + let converted = registry::convert_request( + "claude:messages", + "openai:responses", + &body, + &FormatContext::default(), + ) + .expect("responses request"); + + assert_eq!(converted["reasoning"]["effort"], "xhigh"); + assert_eq!(converted["reasoning"]["summary"], "auto"); + } + + #[test] + fn responses_to_claude_defaults_max_tokens_and_omits_false_is_error() { + let body = json!({ + "model": "gpt-5", + "input": [ + { + "type": "function_call_output", + "call_id": "toolu_ok", + "output": "ok", + "is_error": false + }, + { + "type": "function_call_output", + "call_id": "toolu_bad", + "output": "bad", + "is_error": true + } + ] + }); + + let converted = registry::convert_request( + "openai:responses", + "claude:messages", + &body, + &FormatContext::default(), + ) + .expect("claude request"); + + assert_eq!(converted["max_tokens"], 8192); + let messages_json = converted["messages"].to_string(); + assert!(!messages_json.contains("\"is_error\":false")); + assert!(messages_json.contains("\"is_error\":true")); + } + + #[test] + fn claude_request_to_responses_splits_tool_result_media_from_output() { + let body = json!({ + "model": "claude-sonnet", + "messages": [ + { + "role": "user", + "content": "Describe the file" + }, + { + "role": "assistant", + "content": [{ + "type": "tool_use", + "id": "toolu_read", + "name": "Read", + "input": {"file_path": "/tmp/photo.png"} + }] + }, + { + "role": "user", + "content": [{ + "type": "tool_result", + "tool_use_id": "toolu_read", + "content": [ + {"type": "text", "text": "File metadata: 800x600 PNG"}, + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": "AAAA" + } + } + ] + }] + } + ], + "max_tokens": 128, + }); + + let converted = registry::convert_request( + "claude:messages", + "openai:responses", + &body, + &FormatContext::default(), + ) + .expect("responses request"); + let input = converted["input"].as_array().expect("responses input"); + + assert_eq!(input.len(), 4); + assert_eq!(input[1]["type"], "function_call"); + assert_eq!(input[1]["call_id"], "toolu_read"); + assert_eq!(input[2]["type"], "function_call_output"); + assert_eq!(input[2]["call_id"], "toolu_read"); + assert_eq!(input[2]["output"], "File metadata: 800x600 PNG"); + assert_eq!(input[3]["role"], "user"); + assert_eq!(input[3]["content"][0]["type"], "input_image"); + assert_eq!( + input[3]["content"][0]["image_url"], + "data:image/png;base64,AAAA" + ); + } } diff --git a/crates/aether-ai-formats/src/formats/gemini/embedding/request.rs b/crates/aether-ai-formats/src/formats/gemini/embedding/request.rs index 83416a830..7d69e5236 100644 --- a/crates/aether-ai-formats/src/formats/gemini/embedding/request.rs +++ b/crates/aether-ai-formats/src/formats/gemini/embedding/request.rs @@ -113,6 +113,7 @@ mod tests { dimensions: None, task: None, user: None, + parameters: None, extensions: BTreeMap::new(), }), ..CanonicalRequest::default() diff --git a/crates/aether-ai-formats/src/formats/id.rs b/crates/aether-ai-formats/src/formats/id.rs index f181032bc..f5f96e3ac 100644 --- a/crates/aether-ai-formats/src/formats/id.rs +++ b/crates/aether-ai-formats/src/formats/id.rs @@ -9,6 +9,7 @@ pub enum FormatFamily { Gemini, Jina, Doubao, + Aliyun, } #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] @@ -30,6 +31,7 @@ pub enum FormatId { JinaEmbedding, JinaRerank, DoubaoEmbedding, + AliyunMultimodalEmbedding, } impl FormatId { @@ -52,6 +54,7 @@ impl FormatId { Self::GeminiGenerateContent | Self::GeminiEmbedding => FormatFamily::Gemini, Self::JinaEmbedding | Self::JinaRerank => FormatFamily::Jina, Self::DoubaoEmbedding => FormatFamily::Doubao, + Self::AliyunMultimodalEmbedding => FormatFamily::Aliyun, } } @@ -75,6 +78,7 @@ impl FormatId { Self::JinaEmbedding => "jina:embedding", Self::JinaRerank => "jina:rerank", Self::DoubaoEmbedding => "doubao:embedding", + Self::AliyunMultimodalEmbedding => "aliyun:multimodal_embedding", } } } @@ -103,13 +107,22 @@ impl FromStr for FormatId { "jina:embedding" | "/jina/v1/embeddings" => Ok(Self::JinaEmbedding), "jina:rerank" | "/jina/v1/rerank" => Ok(Self::JinaRerank), "doubao:embedding" => Ok(Self::DoubaoEmbedding), + "aliyun:multimodal_embedding" + | "aliyun_embedding" + | "aliyun_multimodal_embedding" + | "dashscope:multimodal_embedding" + | "dashscope_embedding" + | "dashscope_multimodal_embedding" => Ok(Self::AliyunMultimodalEmbedding), _ => Err(()), } } } pub fn normalize_api_format_alias(value: &str) -> String { - value.trim().to_ascii_lowercase() + let normalized = value.trim().to_ascii_lowercase(); + FormatId::parse(&normalized) + .map(|format| format.as_str().to_string()) + .unwrap_or(normalized) } pub fn api_format_alias_matches(left: &str, right: &str) -> bool { @@ -117,7 +130,13 @@ pub fn api_format_alias_matches(left: &str, right: &str) -> bool { } pub fn api_format_storage_aliases(value: &str) -> Vec { - vec![normalize_api_format_alias(value)] + match FormatId::parse(value).map(FormatId::canonical) { + Some(FormatId::AliyunMultimodalEmbedding) => vec![ + "aliyun:multimodal_embedding".to_string(), + "dashscope:multimodal_embedding".to_string(), + ], + _ => vec![normalize_api_format_alias(value)], + } } pub fn is_openai_responses_format(value: &str) -> bool { @@ -185,6 +204,18 @@ mod tests { FormatId::parse("doubao:embedding"), Some(FormatId::DoubaoEmbedding) ); + assert_eq!( + FormatId::parse("aliyun:multimodal_embedding").map(|format| format.to_string()), + Some("aliyun:multimodal_embedding".to_string()) + ); + assert_eq!( + FormatId::parse("dashscope:multimodal_embedding").map(|format| format.to_string()), + Some("aliyun:multimodal_embedding".to_string()) + ); + assert_eq!( + FormatId::parse("dashscope_embedding").map(|format| format.to_string()), + Some("aliyun:multimodal_embedding".to_string()) + ); assert_eq!(FormatId::OpenAiEmbedding.to_string(), "openai:embedding"); } @@ -197,6 +228,7 @@ mod tests { (FormatId::GeminiEmbedding, FormatFamily::Gemini), (FormatId::JinaEmbedding, FormatFamily::Jina), (FormatId::DoubaoEmbedding, FormatFamily::Doubao), + (FormatId::AliyunMultimodalEmbedding, FormatFamily::Aliyun), ] { assert_eq!(format.family(), family); assert_eq!(format.profile(), FormatProfile::Default); @@ -315,6 +347,13 @@ mod tests { api_format_storage_aliases("doubao:embedding"), vec!["doubao:embedding".to_string()] ); + assert_eq!( + api_format_storage_aliases("dashscope:multimodal_embedding"), + vec![ + "aliyun:multimodal_embedding".to_string(), + "dashscope:multimodal_embedding".to_string(), + ] + ); } #[test] diff --git a/crates/aether-ai-formats/src/formats/matrix.rs b/crates/aether-ai-formats/src/formats/matrix.rs index 501ab0eba..f563eb4ab 100644 --- a/crates/aether-ai-formats/src/formats/matrix.rs +++ b/crates/aether-ai-formats/src/formats/matrix.rs @@ -42,6 +42,7 @@ const EMBEDDING_CANDIDATE_API_FORMATS: &[&str] = &[ "gemini:embedding", "jina:embedding", "doubao:embedding", + "aliyun:multimodal_embedding", ]; const RERANK_CANDIDATE_API_FORMATS: &[&str] = &["openai:rerank", "jina:rerank"]; @@ -238,7 +239,11 @@ pub fn is_standard_api_format(api_format: &str) -> bool { pub fn is_embedding_api_format(api_format: &str) -> bool { matches!( normalize_api_format_alias(api_format).as_str(), - "openai:embedding" | "gemini:embedding" | "jina:embedding" | "doubao:embedding" + "openai:embedding" + | "gemini:embedding" + | "jina:embedding" + | "doubao:embedding" + | "aliyun:multimodal_embedding" ) } @@ -267,9 +272,11 @@ pub fn api_data_format_id(api_format: &str) -> Option<&'static str> { "gemini:generate_content" => Some("gemini"), "openai:chat" => Some("openai_chat"), "openai:responses" | "openai:responses:compact" => Some("openai_responses"), - "openai:embedding" | "gemini:embedding" | "jina:embedding" | "doubao:embedding" => { - Some("embedding") - } + "openai:embedding" + | "gemini:embedding" + | "jina:embedding" + | "doubao:embedding" + | "aliyun:multimodal_embedding" => Some("embedding"), "openai:rerank" | "jina:rerank" => Some("rerank"), _ => None, } @@ -442,6 +449,7 @@ mod tests { "gemini:embedding", "jina:embedding", "doubao:embedding", + "aliyun:multimodal_embedding", ] ); assert_eq!( @@ -451,6 +459,7 @@ mod tests { "openai:embedding", "gemini:embedding", "doubao:embedding", + "aliyun:multimodal_embedding", ] ); assert!(!request_candidate_api_formats("openai:embedding", false).contains(&"openai:chat")); @@ -479,6 +488,7 @@ mod tests { "openai:embedding", "jina:embedding", "doubao:embedding", + "aliyun:multimodal_embedding", ] ); assert_eq!( @@ -488,6 +498,17 @@ mod tests { "openai:embedding", "gemini:embedding", "jina:embedding", + "aliyun:multimodal_embedding", + ] + ); + assert_eq!( + request_candidate_api_formats("aliyun:multimodal_embedding", false), + vec![ + "aliyun:multimodal_embedding", + "openai:embedding", + "gemini:embedding", + "jina:embedding", + "doubao:embedding", ] ); @@ -496,6 +517,7 @@ mod tests { "gemini:embedding", "jina:embedding", "doubao:embedding", + "aliyun:multimodal_embedding", ]; for client_api_format in embedding_formats { for provider_api_format in embedding_formats { @@ -520,6 +542,7 @@ mod tests { "gemini:embedding", "jina:embedding", "doubao:embedding", + "aliyun:multimodal_embedding", ]; let standard_formats = [ "openai:chat", diff --git a/crates/aether-ai-formats/src/formats/mod.rs b/crates/aether-ai-formats/src/formats/mod.rs index acd9973ab..6dd9a9bf1 100644 --- a/crates/aether-ai-formats/src/formats/mod.rs +++ b/crates/aether-ai-formats/src/formats/mod.rs @@ -1,3 +1,4 @@ +pub mod aliyun; pub mod claude; pub mod context; pub mod conversion; diff --git a/crates/aether-ai-formats/src/formats/openai/chat/request.rs b/crates/aether-ai-formats/src/formats/openai/chat/request.rs index dffc7e166..e1369f0fe 100644 --- a/crates/aether-ai-formats/src/formats/openai/chat/request.rs +++ b/crates/aether-ai-formats/src/formats/openai/chat/request.rs @@ -1,4 +1,4 @@ -use serde_json::{json, Value}; +use serde_json::{json, Map, Value}; use crate::{ formats::context::FormatContext, @@ -194,6 +194,7 @@ pub fn to_raw(canonical: &CanonicalRequest) -> Value { .and_then(|value| value.get("effort")) .and_then(Value::as_str) }) + .and_then(openai_chat_reasoning_effort) { output.insert( "reasoning_effort".to_string(), @@ -206,12 +207,12 @@ pub fn to_raw(canonical: &CanonicalRequest) -> Value { "openai", &output, )); - output.extend(namespace_extension_object( + output.extend(chat_compatible_openai_responses_extension_object( &canonical.extensions, OPENAI_RESPONSES_EXTENSION_NAMESPACE, &output, )); - output.extend(namespace_extension_object( + output.extend(chat_compatible_openai_responses_extension_object( &canonical.extensions, OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE, &output, @@ -219,6 +220,31 @@ pub fn to_raw(canonical: &CanonicalRequest) -> Value { Value::Object(output) } +fn openai_chat_reasoning_effort(value: &str) -> Option<&'static str> { + match value.trim().to_ascii_lowercase().as_str() { + "low" => Some("low"), + "medium" => Some("medium"), + "high" | "xhigh" | "max" => Some("high"), + _ => None, + } +} + +fn chat_compatible_openai_responses_extension_object( + extensions: &std::collections::BTreeMap, + namespace: &str, + existing: &Map, +) -> Map { + namespace_extension_object(extensions, namespace, existing) + .into_iter() + .filter(|(key, _)| { + matches!( + key.as_str(), + "verbosity" | "service_tier" | "prompt_cache_key" | "safety_identifier" | "user" + ) + }) + .collect() +} + fn force_stream_options(body: &mut Value, upstream_is_stream: bool) { if !upstream_is_stream { return; diff --git a/crates/aether-ai-formats/src/formats/openai/chat/stream.rs b/crates/aether-ai-formats/src/formats/openai/chat/stream.rs index 5f2da5369..231ca34fe 100644 --- a/crates/aether-ai-formats/src/formats/openai/chat/stream.rs +++ b/crates/aether-ai-formats/src/formats/openai/chat/stream.rs @@ -45,7 +45,7 @@ pub struct OpenAIResponsesProviderState { model: Option, started: bool, finished: bool, - text: String, + text_parts: BTreeMap, reasoning: String, reasoning_parts: BTreeMap, tool_calls: BTreeMap, @@ -420,24 +420,87 @@ impl OpenAIResponsesProviderState { index } + fn text_part_key_from_event(value: &Value) -> String { + let item_key = value + .get("output_index") + .and_then(Value::as_u64) + .map(|value| format!("output:{value}")) + .or_else(|| { + value + .get("item_id") + .or_else(|| value.get("id")) + .and_then(Value::as_str) + .map(|value| format!("item:{value}")) + }) + .unwrap_or_else(|| "output:default".to_string()); + let content_index = value + .get("content_index") + .and_then(Value::as_u64) + .unwrap_or(0); + format!("{item_key}:content:{content_index}") + } + + fn text_part_key_from_message_item( + output_index: Option, + item: &Map, + content_index: usize, + ) -> String { + let item_key = output_index + .map(|value| format!("output:{value}")) + .or_else(|| { + item.get("id") + .and_then(Value::as_str) + .map(|value| format!("item:{value}")) + }) + .unwrap_or_else(|| "output:default".to_string()); + format!("{item_key}:content:{content_index}") + } + + fn emit_text_delta( + &mut self, + report_context: &Value, + out: &mut Vec, + key: String, + text: &str, + ) { + if text.is_empty() { + return; + } + self.text_parts.entry(key).or_default().push_str(text); + self.ensure_started(report_context, out); + let (id, model) = self.identity(report_context); + out.push(CanonicalStreamFrame { + id, + model, + event: CanonicalStreamEvent::TextDelta(text.to_string()), + }); + } + fn emit_missing_text( &mut self, report_context: &Value, out: &mut Vec, + key: String, text: &str, ) { - let missing = if text.starts_with(&self.text) { - text[self.text.len()..].to_string() - } else if self.text == text || self.text.starts_with(text) { - String::new() - } else { - text.to_string() + let missing = { + let current = self.text_parts.entry(key).or_default(); + let missing = if text.starts_with(current.as_str()) { + text[current.len()..].to_string() + } else if current.as_str() == text || current.starts_with(text) { + String::new() + } else { + text.to_string() + }; + if !missing.is_empty() { + current.push_str(&missing); + } + missing }; if missing.is_empty() { return; } self.ensure_started(report_context, out); - self.text.push_str(&missing); let (id, model) = self.identity(report_context); out.push(CanonicalStreamFrame { id, @@ -695,29 +758,34 @@ impl OpenAIResponsesProviderState { report_context: &Value, out: &mut Vec, item: &Map, + output_index: Option, ) { if item.get("type").and_then(Value::as_str) != Some("message") { return; } - let mut completed_text = String::new(); - for raw_content in item + for (content_index, raw_content) in item .get("content") .and_then(Value::as_array) .into_iter() .flatten() + .enumerate() { let Some(content) = raw_content.as_object() else { continue; }; if content.get("type").and_then(Value::as_str) == Some("output_text") { if let Some(text) = content.get("text").and_then(Value::as_str) { - completed_text.push_str(text); + if !text.is_empty() { + let key = Self::text_part_key_from_message_item( + output_index, + item, + content_index, + ); + self.emit_missing_text(report_context, out, key, text); + } } } } - if !completed_text.is_empty() { - self.emit_missing_text(report_context, out, &completed_text); - } } fn emit_reasoning_item( @@ -831,18 +899,13 @@ impl OpenAIResponsesProviderState { } "response.output_text.delta" | "response.outtext.delta" => match value.get("delta") { Some(Value::String(piece)) if !piece.is_empty() => { - self.ensure_started(report_context, &mut out); - self.text.push_str(piece); - let (id, model) = self.identity(report_context); - out.push(CanonicalStreamFrame { - id, - model, - event: CanonicalStreamEvent::TextDelta(piece.clone()), - }); + let key = Self::text_part_key_from_event(&value); + self.emit_text_delta(report_context, &mut out, key, piece); } Some(Value::Object(delta)) => { if let Some(text) = delta.get("text").and_then(Value::as_str) { - self.emit_missing_text(report_context, &mut out, text); + let key = Self::text_part_key_from_event(&value); + self.emit_missing_text(report_context, &mut out, key, text); } } _ => {} @@ -852,7 +915,8 @@ impl OpenAIResponsesProviderState { if part.get("type").and_then(Value::as_str) == Some("output_text") { if let Some(text) = part.get("text").and_then(Value::as_str) { if !text.is_empty() { - self.emit_missing_text(report_context, &mut out, text); + let key = Self::text_part_key_from_event(&value); + self.emit_missing_text(report_context, &mut out, key, text); } } } @@ -890,7 +954,8 @@ impl OpenAIResponsesProviderState { }) .unwrap_or_default(); if !text.is_empty() { - self.emit_missing_text(report_context, &mut out, text); + let key = Self::text_part_key_from_event(&value); + self.emit_missing_text(report_context, &mut out, key, text); } } "response.reasoning_summary_text.delta" => { @@ -967,7 +1032,7 @@ impl OpenAIResponsesProviderState { self.emit_tool_result_item(report_context, &mut out, item, output_index); } "message" => { - self.emit_message_item(report_context, &mut out, item); + self.emit_message_item(report_context, &mut out, item, output_index); } "reasoning" => { self.ensure_started(report_context, &mut out); @@ -1139,7 +1204,7 @@ impl OpenAIResponsesProviderState { self.emit_tool_result_item(report_context, &mut out, item, output_index); } "message" => { - self.emit_message_item(report_context, &mut out, item); + self.emit_message_item(report_context, &mut out, item, output_index); } "reasoning" => { self.emit_reasoning_item(report_context, &mut out, item); @@ -1194,7 +1259,12 @@ impl OpenAIResponsesProviderState { }; match item.get("type").and_then(Value::as_str).unwrap_or_default() { "message" => { - self.emit_message_item(report_context, &mut out, item); + self.emit_message_item( + report_context, + &mut out, + item, + Some(output_index), + ); } "function_call" => { self.emit_tool_call_item( @@ -3235,6 +3305,102 @@ mod tests { assert_eq!(text, "Hello world"); } + #[test] + fn openai_responses_provider_state_dedupes_text_snapshots_per_output_item() { + let mut state = OpenAIResponsesProviderState::default(); + let report_context = json!({}); + let mut frames = Vec::new(); + + for event in [ + json!({ + "type": "response.output_text.delta", + "response_id": "resp_multi_message", + "output_index": 0, + "item_id": "msg_1", + "content_index": 0, + "delta": "First message.", + }), + json!({ + "type": "response.output_text.done", + "response_id": "resp_multi_message", + "output_index": 0, + "item_id": "msg_1", + "content_index": 0, + "text": "First message.", + }), + json!({ + "type": "response.output_item.done", + "response_id": "resp_multi_message", + "output_index": 0, + "item": { + "type": "message", + "id": "msg_1", + "status": "completed", + "content": [{ + "type": "output_text", + "text": "First message.", + }], + }, + }), + json!({ + "type": "response.output_text.delta", + "response_id": "resp_multi_message", + "output_index": 1, + "item_id": "msg_2", + "content_index": 0, + "delta": "Second message.", + }), + json!({ + "type": "response.output_text.done", + "response_id": "resp_multi_message", + "output_index": 1, + "item_id": "msg_2", + "content_index": 0, + "text": "Second message.", + }), + json!({ + "type": "response.content_part.done", + "response_id": "resp_multi_message", + "output_index": 1, + "item_id": "msg_2", + "content_index": 0, + "part": { + "type": "output_text", + "text": "Second message.", + }, + }), + json!({ + "type": "response.output_item.done", + "response_id": "resp_multi_message", + "output_index": 1, + "item": { + "type": "message", + "id": "msg_2", + "status": "completed", + "content": [{ + "type": "output_text", + "text": "Second message.", + }], + }, + }), + ] { + frames.extend( + state + .push_line(&report_context, data_line(event)) + .expect("responses text event should parse"), + ); + } + + let text = frames + .iter() + .filter_map(|frame| match &frame.event { + CanonicalStreamEvent::TextDelta(text) => Some(text.as_str()), + _ => None, + }) + .collect::(); + assert_eq!(text, "First message.Second message."); + } + #[test] fn openai_responses_provider_state_delays_arguments_until_tool_name_is_known() { let mut state = OpenAIResponsesProviderState::default(); diff --git a/crates/aether-ai-formats/src/formats/openai/embedding/request.rs b/crates/aether-ai-formats/src/formats/openai/embedding/request.rs index 3223f438e..891bdd8d9 100644 --- a/crates/aether-ai-formats/src/formats/openai/embedding/request.rs +++ b/crates/aether-ai-formats/src/formats/openai/embedding/request.rs @@ -50,6 +50,10 @@ pub(crate) fn from_namespace(body_json: &Value, namespace: &str) -> Option Option { + object.remove("cache_control"); + for child in object.values_mut() { + strip_codex_cache_control_fields(child); + } + } + Value::Array(items) => { + for child in items { + strip_codex_cache_control_fields(child); + } + } + _ => {} + } +} + fn extract_codex_prompt_cache_control_seed(provider_request_body: &Value) -> Option { let mut anchors = Vec::new(); collect_codex_prompt_cache_control_anchors(provider_request_body, &mut anchors); @@ -780,6 +797,7 @@ pub fn apply_codex_openai_responses_special_body_edits( inject_codex_default_variation_prompt(body_object); } + strip_codex_cache_control_fields(provider_request_body); insert_codex_prompt_cache_key(provider_request_body, prompt_cache_key); } @@ -1206,6 +1224,49 @@ mod tests { assert_eq!(body_a["prompt_cache_key"], body_b["prompt_cache_key"]); assert_ne!(body_a["prompt_cache_key"], body_c["prompt_cache_key"]); + assert!(!body_a.to_string().contains("\"cache_control\"")); + assert!(!body_b.to_string().contains("\"cache_control\"")); + assert!(!body_c.to_string().contains("\"cache_control\"")); + } + + #[test] + fn codex_responses_body_edits_strip_developer_cache_control_before_upstream() { + let mut provider_request_body = json!({ + "input": [{ + "type": "message", + "role": "developer", + "content": [{ + "type": "input_text", + "text": "stable system brief", + "cache_control": {"type": "ephemeral"} + }] + }, { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "new turn"}] + }], + "model": "gpt-5.4" + }); + + apply_codex_openai_responses_special_body_edits( + &mut provider_request_body, + "codex", + "openai:responses", + None, + Some("key-a"), + ); + + assert!(provider_request_body + .get("prompt_cache_key") + .and_then(|value| value.as_str()) + .is_some_and(|value| !value.trim().is_empty())); + assert!(!provider_request_body + .to_string() + .contains("\"cache_control\"")); + assert_eq!( + provider_request_body["input"][0]["content"][0]["text"], + json!("stable system brief") + ); } #[test] diff --git a/crates/aether-ai-formats/src/formats/openai/responses/request.rs b/crates/aether-ai-formats/src/formats/openai/responses/request.rs index 12a05f171..9fed5a930 100644 --- a/crates/aether-ai-formats/src/formats/openai/responses/request.rs +++ b/crates/aether-ai-formats/src/formats/openai/responses/request.rs @@ -1,12 +1,15 @@ +use std::collections::BTreeMap; + use serde_json::{json, Map, Value}; use crate::{ formats::context::FormatContext, formats::openai::shared::map_thinking_budget_to_openai_reasoning_effort, protocol::canonical::{ - canonical_response_format_to_openai, canonicalize_tool_arguments, media_data_or_url, - namespace_extension_object, openai_content_text, openai_extensions, - openai_response_format_to_canonical, openai_responses_extension, + canonical_response_format_to_openai, canonicalize_tool_arguments, + is_claude_messages_request, is_claude_system_instruction, is_claude_thinking_block, + is_claude_tool_result, media_data_or_url, namespace_extension_object, openai_content_text, + openai_extensions, openai_response_format_to_canonical, openai_responses_extension, openai_responses_generation_config, openai_responses_input_to_canonical_messages, openai_responses_tool_choice_to_canonical, openai_responses_tools_to_canonical, CanonicalContentBlock, CanonicalInstruction, CanonicalRequest, CanonicalRole, @@ -135,13 +138,18 @@ pub fn to_raw( output.insert("instructions".to_string(), instructions); } let mut input = canonical_messages_to_responses_input(canonical)?; + if let Some(developer_message) = + claude_system_instructions_to_responses_developer_message(canonical) + { + input.insert(0, developer_message); + } ensure_json_object_response_input_mentions_json(canonical, instructions.as_ref(), &mut input); output.insert("input".to_string(), Value::Array(input)); if upstream_is_stream && !compact { output.insert("stream".to_string(), Value::Bool(true)); } - if let Some(max_tokens) = canonical.generation.max_tokens { + if let Some(max_tokens) = responses_max_output_tokens(canonical) { output.insert("max_output_tokens".to_string(), Value::from(max_tokens)); } insert_number(&mut output, "temperature", canonical.generation.temperature); @@ -170,11 +178,7 @@ pub fn to_raw( canonical_tool_choice_to_responses(tool_choice), ); } - if let Some(reasoning) = canonical - .thinking - .as_ref() - .and_then(reasoning_config_to_responses) - { + if let Some(reasoning) = canonical_reasoning_config_to_responses(canonical) { output.insert("reasoning".to_string(), reasoning); } @@ -188,6 +192,7 @@ pub fn to_raw( OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE, &output, )); + apply_claude_responses_request_defaults(canonical, mapped_model, &mut output); if compact { output.remove("stream"); } @@ -199,6 +204,7 @@ fn canonical_instructions_to_responses(canonical: &CanonicalRequest) -> Option>() @@ -206,6 +212,13 @@ fn canonical_instructions_to_responses(canonical: &CanonicalRequest) -> Option Option Option { + let content = canonical + .instructions + .iter() + .filter(|instruction| is_claude_system_instruction(instruction)) + .filter_map(claude_system_instruction_to_responses_part) + .collect::>(); + (!content.is_empty()).then(|| { + json!({ + "type": "message", + "role": "developer", + "content": content, + }) + }) +} + +fn claude_system_instruction_to_responses_part( + instruction: &CanonicalInstruction, +) -> Option { + if instruction.text.trim().is_empty() { + return None; + } + let mut part = Map::new(); + part.insert("type".to_string(), Value::String("input_text".to_string())); + part.insert("text".to_string(), Value::String(instruction.text.clone())); + part.extend(namespace_extension_object( + &instruction.extensions, + "claude", + &part, + )); + Some(Value::Object(part)) +} + fn canonical_messages_to_responses_input(canonical: &CanonicalRequest) -> Option> { let mut input = Vec::new(); for message in &canonical.messages { @@ -223,6 +271,7 @@ fn canonical_messages_to_responses_input(canonical: &CanonicalRequest) -> Option CanonicalRole::System | CanonicalRole::Developer => continue, }; let mut content = Vec::new(); + let mut saw_tool_item = false; for block in &message.content { match block { CanonicalContentBlock::ToolUse { @@ -232,6 +281,7 @@ fn canonical_messages_to_responses_input(canonical: &CanonicalRequest) -> Option .. } => { flush_responses_message(&mut input, role, &mut content); + saw_tool_item = true; input.push(json!({ "type": "function_call", "call_id": id, @@ -243,16 +293,42 @@ fn canonical_messages_to_responses_input(canonical: &CanonicalRequest) -> Option tool_use_id, output, content_text, + extensions, .. } => { flush_responses_message(&mut input, role, &mut content); + saw_tool_item = true; + let (tool_output, extra_user_content) = responses_tool_result_payload( + output.as_ref(), + content_text.as_deref(), + extensions, + ); input.push(json!({ "type": "function_call_output", "call_id": tool_use_id, - "output": responses_tool_result_output(output.as_ref(), content_text.as_deref()), + "output": tool_output, })); + if !extra_user_content.is_empty() { + input.push(json!({ + "type": "message", + "role": "user", + "content": extra_user_content, + })); + } + } + CanonicalContentBlock::Thinking { + text, extensions, .. + } => { + if is_claude_thinking_block(extensions) { + continue; + } + if role == "assistant" && !text.trim().is_empty() { + content.push(json!({ + "type": "output_text", + "text": format!("{text}"), + })); + } } - CanonicalContentBlock::Thinking { .. } => {} other => { if let Some(part) = canonical_block_to_responses_input_part(other, role) { content.push(part); @@ -260,11 +336,73 @@ fn canonical_messages_to_responses_input(canonical: &CanonicalRequest) -> Option } } } + if content.is_empty() && !saw_tool_item { + if role == "assistant" { + input.push(json!({ + "type": "message", + "role": role, + "content": [{ + "type": "output_text", + "text": "", + }], + })); + } else { + input.push(json!({ + "type": "message", + "role": role, + "content": "", + })); + } + continue; + } flush_responses_message(&mut input, role, &mut content); } Some(input) } +fn responses_max_output_tokens(canonical: &CanonicalRequest) -> Option { + canonical.generation.max_tokens.map(|max_tokens| { + if is_claude_messages_request(&canonical.extensions) && max_tokens < 128 { + 128 + } else { + max_tokens + } + }) +} + +fn apply_claude_responses_request_defaults( + canonical: &CanonicalRequest, + mapped_model: &str, + output: &mut Map, +) { + if !is_claude_messages_request(&canonical.extensions) { + return; + } + if mapped_model + .trim() + .to_ascii_lowercase() + .starts_with("gpt-5") + { + output.remove("temperature"); + output.remove("top_p"); + } + output + .entry("store".to_string()) + .or_insert_with(|| Value::Bool(false)); + output + .entry("parallel_tool_calls".to_string()) + .or_insert_with(|| Value::Bool(true)); + let include = output + .entry("include".to_string()) + .or_insert_with(|| Value::Array(Vec::new())); + if let Some(include) = include.as_array_mut() { + let encrypted_content = Value::String("reasoning.encrypted_content".to_string()); + if !include.iter().any(|value| value == &encrypted_content) { + include.push(encrypted_content); + } + } +} + fn ensure_json_object_response_input_mentions_json( canonical: &CanonicalRequest, instructions: Option<&Value>, @@ -419,6 +557,42 @@ fn canonical_tools_to_responses(canonical: &CanonicalRequest) -> Vec { tools } +fn canonical_reasoning_config_to_responses(canonical: &CanonicalRequest) -> Option { + let is_claude_request = is_claude_messages_request(&canonical.extensions); + if !is_claude_request { + return canonical + .thinking + .as_ref() + .and_then(reasoning_config_to_responses); + } + + let mut object = canonical + .thinking + .as_ref() + .and_then(|thinking| openai_responses_extension(&thinking.extensions).cloned()) + .and_then(|value| match value { + Value::Object(object) => Some(object), + _ => None, + }) + .unwrap_or_default(); + let effort = canonical + .thinking + .as_ref() + .and_then(|thinking| thinking.extensions.get("claude")) + .and_then(|value| value.get("output_config")) + .and_then(|value| value.get("effort")) + .and_then(Value::as_str) + .map(openai_responses_reasoning_effort) + .unwrap_or("medium"); + object + .entry("effort".to_string()) + .or_insert_with(|| Value::String(effort.to_string())); + object + .entry("summary".to_string()) + .or_insert_with(|| Value::String("auto".to_string())); + Some(Value::Object(object)) +} + fn reasoning_config_to_responses(thinking: &CanonicalThinkingConfig) -> Option { openai_responses_extension(&thinking.extensions) .cloned() @@ -481,6 +655,10 @@ fn canonical_text_config_to_responses(canonical: &CanonicalRequest) -> Option Value { Value::String(description.clone()), ); } - if let Some(parameters) = &tool.parameters { - out.insert("parameters".to_string(), parameters.clone()); - } + out.insert( + "parameters".to_string(), + responses_tool_parameters_schema(tool.parameters.as_ref()), + ); out.extend(namespace_extension_object( &tool.extensions, OPENAI_RESPONSES_EXTENSION_NAMESPACE, @@ -523,6 +702,25 @@ fn canonical_tool_to_responses(tool: &CanonicalToolDefinition) -> Value { Value::Object(out) } +fn responses_tool_parameters_schema(parameters: Option<&Value>) -> Value { + match parameters { + Some(Value::Object(schema)) => { + let mut schema = schema.clone(); + if schema + .get("type") + .and_then(Value::as_str) + .is_some_and(|value| value == "object") + && !schema.contains_key("properties") + { + schema.insert("properties".to_string(), json!({})); + } + Value::Object(schema) + } + Some(Value::Null) | None => json!({"type": "object", "properties": {}}), + Some(value) => value.clone(), + } +} + fn canonical_tool_choice_to_responses(choice: &CanonicalToolChoice) -> Value { match choice { CanonicalToolChoice::Auto => Value::String("auto".to_string()), @@ -535,13 +733,169 @@ fn canonical_tool_choice_to_responses(choice: &CanonicalToolChoice) -> Value { } } +fn responses_tool_result_payload( + output: Option<&Value>, + content_text: Option<&str>, + extensions: &BTreeMap, +) -> (Value, Vec) { + if is_claude_tool_result(extensions) { + if let Some(Value::Array(parts)) = output { + return claude_tool_result_parts_to_responses_payload(parts); + } + } + ( + responses_tool_result_output(output, content_text), + Vec::new(), + ) +} + fn responses_tool_result_output(output: Option<&Value>, content_text: Option<&str>) -> Value { - match output { - Some(Value::String(text)) => Value::String(text.clone()), - Some(value) => serde_json::to_string(value) - .map(Value::String) - .unwrap_or_else(|_| Value::String(String::new())), - None => Value::String(content_text.unwrap_or_default().to_string()), + let text = match output { + Some(Value::String(text)) => text.clone(), + Some(Value::Null) => String::new(), + Some(value) => serde_json::to_string(value).unwrap_or_default(), + None => content_text.unwrap_or_default().to_string(), + }; + Value::String(non_empty_responses_tool_output(&text)) +} + +fn claude_tool_result_parts_to_responses_payload(parts: &[Value]) -> (Value, Vec) { + let mut output_texts = Vec::new(); + let mut extra_user_content = Vec::new(); + + for part in parts { + let Some(part_object) = part.as_object() else { + output_texts.push("[Claude tool_result non-text content omitted]".to_string()); + continue; + }; + match part_object + .get("type") + .and_then(Value::as_str) + .unwrap_or_default() + { + "text" => { + if let Some(text) = part_object.get("text").and_then(Value::as_str) { + if !text.is_empty() { + output_texts.push(text.to_string()); + } + } + } + "image" => { + if let Some(part) = claude_image_block_to_responses_input_part(part_object) { + extra_user_content.push(part); + } else { + output_texts.push(claude_tool_result_media_summary("image", part_object)); + } + } + "document" | "file" => { + if let Some(part) = claude_document_block_to_responses_input_part(part_object) { + extra_user_content.push(part); + } else { + output_texts.push(claude_tool_result_media_summary("document", part_object)); + } + } + "" => output_texts.push("[Claude tool_result object content omitted]".to_string()), + raw_type => { + output_texts.push(format!("[Claude tool_result {raw_type} content omitted]")) + } + } + } + + ( + Value::String(non_empty_responses_tool_output(&output_texts.join("\n\n"))), + extra_user_content, + ) +} + +fn claude_image_block_to_responses_input_part(block: &Map) -> Option { + let source = block.get("source")?.as_object()?; + match source + .get("type") + .and_then(Value::as_str) + .unwrap_or_default() + { + "base64" => { + let media_type = claude_source_media_type(source).unwrap_or("image/png"); + let data = claude_source_str(source, "data")?; + Some(json!({ + "type": "input_image", + "image_url": format!("data:{media_type};base64,{data}"), + })) + } + "url" => { + let url = claude_source_str(source, "url")?; + Some(json!({ + "type": "input_image", + "image_url": url, + })) + } + _ => None, + } +} + +fn claude_document_block_to_responses_input_part(block: &Map) -> Option { + let source = block.get("source")?.as_object()?; + let file_data = match source + .get("type") + .and_then(Value::as_str) + .unwrap_or_default() + { + "base64" => { + let media_type = claude_source_media_type(source).unwrap_or("application/octet-stream"); + let data = claude_source_str(source, "data")?; + format!("data:{media_type};base64,{data}") + } + "url" => claude_source_str(source, "url")?.to_string(), + _ => return None, + }; + + let mut part = Map::new(); + part.insert("type".to_string(), Value::String("input_file".to_string())); + part.insert("file_data".to_string(), Value::String(file_data)); + if let Some(filename) = block + .get("title") + .or_else(|| block.get("name")) + .and_then(Value::as_str) + .filter(|value| !value.trim().is_empty()) + { + part.insert("filename".to_string(), Value::String(filename.to_string())); + } + Some(Value::Object(part)) +} + +fn claude_tool_result_media_summary(kind: &str, block: &Map) -> String { + let media_type = block + .get("source") + .and_then(Value::as_object) + .and_then(claude_source_media_type); + match media_type { + Some(media_type) if !media_type.trim().is_empty() => { + format!("[Claude tool_result {kind} content omitted: {media_type}]") + } + _ => format!("[Claude tool_result {kind} content omitted]"), + } +} + +fn claude_source_media_type(source: &Map) -> Option<&str> { + source + .get("media_type") + .or_else(|| source.get("mime_type")) + .and_then(Value::as_str) + .filter(|value| !value.trim().is_empty()) +} + +fn claude_source_str<'a>(source: &'a Map, key: &str) -> Option<&'a str> { + source + .get(key) + .and_then(Value::as_str) + .filter(|value| !value.trim().is_empty()) +} + +fn non_empty_responses_tool_output(text: &str) -> String { + if text.is_empty() { + "(empty)".to_string() + } else { + text.to_string() } } @@ -591,4 +945,62 @@ mod tests { .to_ascii_lowercase() .contains("json")); } + + #[test] + fn responses_request_preserves_empty_chat_messages() { + let request = CanonicalRequest { + model: "gpt-5.5".to_string(), + messages: vec![ + CanonicalMessage { + role: CanonicalRole::User, + content: vec![CanonicalContentBlock::Text { + text: String::new(), + extensions: Default::default(), + }], + extensions: Default::default(), + }, + CanonicalMessage { + role: CanonicalRole::Assistant, + content: Vec::new(), + extensions: Default::default(), + }, + ], + ..CanonicalRequest::default() + }; + + let body = to_raw(&request, "gpt-5.5", false, false).expect("responses body"); + + assert_eq!(body["input"][0]["role"], "user"); + assert_eq!(body["input"][0]["content"], ""); + assert_eq!(body["input"][1]["role"], "assistant"); + assert_eq!(body["input"][1]["content"][0]["type"], "output_text"); + assert_eq!(body["input"][1]["content"][0]["text"], ""); + } + + #[test] + fn responses_request_uses_empty_marker_for_empty_tool_output() { + let request = CanonicalRequest { + model: "gpt-5.5".to_string(), + messages: vec![CanonicalMessage { + role: CanonicalRole::Tool, + content: vec![CanonicalContentBlock::ToolResult { + tool_use_id: "call_empty".to_string(), + name: None, + output: Some(json!("")), + content_text: None, + is_error: false, + extensions: Default::default(), + }], + extensions: Default::default(), + }], + ..CanonicalRequest::default() + }; + + let body = to_raw(&request, "gpt-5.5", false, false).expect("responses body"); + + assert_eq!(body["input"].as_array().expect("input").len(), 1); + assert_eq!(body["input"][0]["type"], "function_call_output"); + assert_eq!(body["input"][0]["call_id"], "call_empty"); + assert_eq!(body["input"][0]["output"], "(empty)"); + } } diff --git a/crates/aether-ai-formats/src/formats/registry.rs b/crates/aether-ai-formats/src/formats/registry.rs index ff06e8c29..1f810565d 100644 --- a/crates/aether-ai-formats/src/formats/registry.rs +++ b/crates/aether-ai-formats/src/formats/registry.rs @@ -1,6 +1,7 @@ use serde_json::Value; use crate::formats::{ + aliyun, claude::messages as claude_messages, doubao, gemini::{self, generate_content as gemini_generate_content}, @@ -29,7 +30,9 @@ pub fn parse_request( FormatId::JinaEmbedding => jina::embedding::request::from(body, ctx), FormatId::OpenAiRerank => openai::rerank::request::from(body, ctx), FormatId::JinaRerank => jina::rerank::request::from(body, ctx), - FormatId::GeminiEmbedding | FormatId::DoubaoEmbedding => None, + FormatId::GeminiEmbedding + | FormatId::DoubaoEmbedding + | FormatId::AliyunMultimodalEmbedding => None, } .ok_or_else(|| FormatError::RequestParseFailed { format: source.as_str().to_string(), @@ -62,6 +65,7 @@ pub fn emit_request( FormatId::JinaRerank => jina::rerank::request::to(&request, ctx), FormatId::GeminiEmbedding => gemini::embedding::request::to(&request, ctx), FormatId::DoubaoEmbedding => doubao::embedding::request::to(&request, ctx), + FormatId::AliyunMultimodalEmbedding => aliyun::embedding::request::to(&request, ctx), } .ok_or_else(|| FormatError::RequestEmitFailed { format: target.as_str().to_string(), @@ -96,7 +100,8 @@ pub fn parse_response( | FormatId::OpenAiRerank | FormatId::JinaRerank | FormatId::GeminiEmbedding - | FormatId::DoubaoEmbedding => None, + | FormatId::DoubaoEmbedding + | FormatId::AliyunMultimodalEmbedding => None, } .ok_or_else(|| FormatError::ResponseParseFailed { format: source.as_str().to_string(), @@ -120,7 +125,8 @@ pub fn emit_response( | FormatId::OpenAiRerank | FormatId::JinaRerank | FormatId::GeminiEmbedding - | FormatId::DoubaoEmbedding => None, + | FormatId::DoubaoEmbedding + | FormatId::AliyunMultimodalEmbedding => None, } .ok_or_else(|| FormatError::ResponseEmitFailed { format: target.as_str().to_string(), @@ -252,6 +258,119 @@ mod tests { assert!(doubao.get("messages").is_none()); } + #[test] + fn converts_openai_embedding_to_aliyun_multimodal_payload_shape() { + let body = json!({ + "model": "text-embedding-3-small", + "input": [ + {"text": "white running shoes"}, + {"image": "https://example.com/shoe.png"}, + {"multi_images": ["https://example.com/a.png", "https://example.com/b.png"]} + ], + "dimensions": 1024, + "parameters": { + "enable_fusion": true, + "res_level": 2, + "max_video_frames": 64 + } + }); + + let converted = convert_request( + "openai:embedding", + "aliyun:multimodal_embedding", + &body, + &FormatContext::default().with_mapped_model("qwen3-vl-embedding"), + ) + .expect("aliyun multimodal embedding conversion should succeed"); + + assert_eq!(converted["model"], "qwen3-vl-embedding"); + assert_eq!(converted["input"]["contents"], body["input"]); + assert_eq!(converted["parameters"]["dimension"], 1024); + assert_eq!(converted["parameters"]["enable_fusion"], true); + assert_eq!(converted["parameters"]["res_level"], 2); + assert_eq!(converted["parameters"]["max_video_frames"], 64); + assert!(converted.get("messages").is_none()); + } + + #[test] + fn aliyun_embedding_conversion_rejects_token_arrays() { + let body = json!({ + "model": "text-embedding-3-small", + "input": [1, 2, 3] + }); + + assert!(convert_request( + "openai:embedding", + "aliyun:multimodal_embedding", + &body, + &FormatContext::default().with_mapped_model("qwen3-vl-embedding"), + ) + .is_err()); + } + + #[test] + fn multimodal_embedding_conversion_is_aliyun_only() { + let body = json!({ + "model": "qwen3-vl-embedding", + "input": [ + {"text": "white running shoes"}, + {"image": "https://example.com/shoe.png"} + ] + }); + let ctx = FormatContext::default().with_mapped_model("qwen3-vl-embedding"); + + assert!(convert_request("openai:embedding", "openai:embedding", &body, &ctx).is_err()); + assert!(convert_request("openai:embedding", "jina:embedding", &body, &ctx).is_err()); + assert!(convert_request("openai:embedding", "gemini:embedding", &body, &ctx).is_err()); + assert!(convert_request("openai:embedding", "doubao:embedding", &body, &ctx).is_err()); + assert!(convert_request( + "openai:embedding", + "aliyun:multimodal_embedding", + &body, + &ctx + ) + .is_ok()); + } + + #[test] + fn parses_aliyun_embedding_response_to_openai_shape() { + let body = json!({ + "output": { + "embeddings": [ + { + "index": 0, + "embedding": [0.1, 0.2, 0.3], + "type": "fused" + } + ] + }, + "usage": { + "input_tokens": 432, + "input_tokens_details": { + "image_tokens": 402, + "text_tokens": 30 + }, + "output_tokens": 1, + "total_tokens": 433 + }, + "request_id": "aliyun-request-1" + }); + + let canonical = + crate::protocol::canonical::from_embedding_to_canonical_response(&body, "aliyun") + .expect("aliyun embedding response should parse"); + let emitted = + crate::protocol::canonical::canonical_to_embedding_response(&canonical, "openai") + .expect("openai embedding response should emit"); + + assert_eq!(emitted["request_id"], "aliyun-request-1"); + assert_eq!(emitted["data"][0]["embedding"], json!([0.1, 0.2, 0.3])); + assert_eq!(emitted["data"][0]["type"], "fused"); + assert_eq!(emitted["usage"]["prompt_tokens"], 432); + assert_eq!(emitted["usage"]["completion_tokens"], 1); + assert_eq!(emitted["usage"]["total_tokens"], 433); + } + #[test] fn embedding_registry_keeps_gemini_and_doubao_emit_only() { let body = json!({ diff --git a/crates/aether-ai-formats/src/formats/shared/model_directives.rs b/crates/aether-ai-formats/src/formats/shared/model_directives.rs index 345cbd354..5eb46d1c5 100644 --- a/crates/aether-ai-formats/src/formats/shared/model_directives.rs +++ b/crates/aether-ai-formats/src/formats/shared/model_directives.rs @@ -38,8 +38,7 @@ impl ReasoningEffort { Self::Low => "low", Self::Medium => "medium", Self::High => "high", - Self::XHigh => "xhigh", - Self::Max => "xhigh", + Self::XHigh | Self::Max => "high", } } @@ -524,7 +523,7 @@ mod tests { "gpt-5.4-xhigh", ) .expect("directive should apply"); - assert_eq!(openai_chat["reasoning_effort"], "xhigh"); + assert_eq!(openai_chat["reasoning_effort"], "high"); let mut responses = json!({ "model": "gpt-5-upstream", @@ -597,7 +596,7 @@ mod tests { "gpt-5.4-fast-xhigh", ) .expect("directive should apply"); - assert_eq!(openai_chat["reasoning_effort"], "xhigh"); + assert_eq!(openai_chat["reasoning_effort"], "high"); assert_eq!(openai_chat["service_tier"], "priority"); let mut reversed = json!({"model": "gpt-5-upstream", "reasoning_effort": "low"}); diff --git a/crates/aether-ai-formats/src/formats/shared/standard_matrix.rs b/crates/aether-ai-formats/src/formats/shared/standard_matrix.rs index 5de3ade2d..441445943 100644 --- a/crates/aether-ai-formats/src/formats/shared/standard_matrix.rs +++ b/crates/aether-ai-formats/src/formats/shared/standard_matrix.rs @@ -1079,6 +1079,33 @@ mod tests { None, ) .expect("typed canonical claude route should build"); + if matches!( + provider_api_format, + "openai:responses" | "openai:responses:compact" + ) { + assert!(converted.get("instructions").is_none()); + assert_eq!(converted["input"][0]["role"], "developer"); + assert_eq!(converted["input"][0]["content"][0]["text"], "Be exact."); + assert_eq!(converted["max_output_tokens"], 128); + assert_eq!(converted["text"]["verbosity"], "medium"); + assert_eq!(converted["reasoning"]["effort"], "medium"); + assert_eq!(converted["reasoning"]["summary"], "auto"); + if provider_api_format == "openai:responses" { + assert_eq!(converted["store"], false); + assert!(converted["include"] + .as_array() + .expect("include") + .iter() + .any(|value| value.as_str() == Some("reasoning.encrypted_content"))); + } else { + assert!(converted.get("store").is_none()); + assert!(converted.get("include").is_none()); + } + let input_json = converted["input"].to_string(); + assert!(!input_json.contains("plan")); + assert!(!input_json.contains("sig_123")); + continue; + } let legacy = legacy_claude_request_body(&request, provider_api_format, upstream_is_stream); assert_eq!( diff --git a/crates/aether-ai-formats/src/formats/shared/standard_normalize.rs b/crates/aether-ai-formats/src/formats/shared/standard_normalize.rs index 14d8b0cd5..0187685d8 100644 --- a/crates/aether-ai-formats/src/formats/shared/standard_normalize.rs +++ b/crates/aether-ai-formats/src/formats/shared/standard_normalize.rs @@ -738,7 +738,7 @@ mod tests { .expect("openai chat body should build"); assert_eq!(provider_request_body["model"], "gpt-5-upstream"); - assert_eq!(provider_request_body["reasoning_effort"], "xhigh"); + assert_eq!(provider_request_body["reasoning_effort"], "high"); } #[test] diff --git a/crates/aether-ai-formats/src/formats/shared/stream_core/format_matrix.rs b/crates/aether-ai-formats/src/formats/shared/stream_core/format_matrix.rs index 6f095bde0..54615b91a 100644 --- a/crates/aether-ai-formats/src/formats/shared/stream_core/format_matrix.rs +++ b/crates/aether-ai-formats/src/formats/shared/stream_core/format_matrix.rs @@ -260,7 +260,8 @@ impl ProviderStreamParser { | FormatId::GeminiEmbedding | FormatId::JinaEmbedding | FormatId::JinaRerank - | FormatId::DoubaoEmbedding => return None, + | FormatId::DoubaoEmbedding + | FormatId::AliyunMultimodalEmbedding => return None, }) } @@ -350,7 +351,8 @@ impl ClientStreamEmitter { | FormatId::GeminiEmbedding | FormatId::JinaEmbedding | FormatId::JinaRerank - | FormatId::DoubaoEmbedding => return None, + | FormatId::DoubaoEmbedding + | FormatId::AliyunMultimodalEmbedding => return None, }) } @@ -414,7 +416,8 @@ fn parse_provider_error( | FormatId::GeminiEmbedding | FormatId::JinaEmbedding | FormatId::JinaRerank - | FormatId::DoubaoEmbedding => None, + | FormatId::DoubaoEmbedding + | FormatId::AliyunMultimodalEmbedding => None, } } diff --git a/crates/aether-ai-formats/src/formats/shared/sync_products.rs b/crates/aether-ai-formats/src/formats/shared/sync_products.rs index 5700d4835..41f502eb1 100644 --- a/crates/aether-ai-formats/src/formats/shared/sync_products.rs +++ b/crates/aether-ai-formats/src/formats/shared/sync_products.rs @@ -439,6 +439,7 @@ fn embedding_response_namespace_for_api_format(api_format: &str) -> Option<&'sta "openai:embedding" => Some("openai"), "jina:embedding" => Some("jina"), "gemini:embedding" => Some("gemini"), + "aliyun:multimodal_embedding" => Some("aliyun"), _ => None, } } diff --git a/crates/aether-ai-formats/src/lib.rs b/crates/aether-ai-formats/src/lib.rs index 67160b7a5..24034bb68 100644 --- a/crates/aether-ai-formats/src/lib.rs +++ b/crates/aether-ai-formats/src/lib.rs @@ -43,9 +43,9 @@ pub use protocol::canonical::{ from_gemini_to_canonical_response, from_openai_chat_to_canonical_request, from_openai_chat_to_canonical_response, from_openai_responses_to_canonical_request, from_openai_responses_to_canonical_response, CanonicalContentBlock, CanonicalEmbedding, - CanonicalEmbeddingInput, CanonicalEmbeddingRequest, CanonicalEmbeddingResponse, - CanonicalGenerationConfig, CanonicalInstruction, CanonicalMessage, CanonicalRequest, - CanonicalResponse, CanonicalResponseFormat, CanonicalResponseOutput, CanonicalRole, - CanonicalStopReason, CanonicalStreamEvent, CanonicalStreamFrame, CanonicalThinkingConfig, - CanonicalToolChoice, CanonicalToolDefinition, CanonicalUsage, + CanonicalEmbeddingContent, CanonicalEmbeddingInput, CanonicalEmbeddingRequest, + CanonicalEmbeddingResponse, CanonicalGenerationConfig, CanonicalInstruction, CanonicalMessage, + CanonicalRequest, CanonicalResponse, CanonicalResponseFormat, CanonicalResponseOutput, + CanonicalRole, CanonicalStopReason, CanonicalStreamEvent, CanonicalStreamFrame, + CanonicalThinkingConfig, CanonicalToolChoice, CanonicalToolDefinition, CanonicalUsage, }; diff --git a/crates/aether-ai-formats/src/protocol/canonical.rs b/crates/aether-ai-formats/src/protocol/canonical.rs index f6a46d557..ed5c81a8a 100644 --- a/crates/aether-ai-formats/src/protocol/canonical.rs +++ b/crates/aether-ai-formats/src/protocol/canonical.rs @@ -11,8 +11,12 @@ pub use crate::protocol::stream::{CanonicalStreamEvent, CanonicalStreamFrame}; pub(crate) const OPENAI_RESPONSES_EXTENSION_NAMESPACE: &str = "openai_responses"; pub(crate) const OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE: &str = "openai_cli"; const AETHER_EXTENSION_NAMESPACE: &str = "aether"; +const CLAUDE_MESSAGES_REQUEST_SOURCE_MARKER: &str = "claude_messages_request"; +const CLAUDE_SYSTEM_SOURCE_MARKER: &str = "claude_system"; +const CLAUDE_THINKING_SOURCE_MARKER: &str = "claude_thinking"; const CLAUDE_TOOL_RESULT_SOURCE_MARKER: &str = "claude_tool_result"; const OPENAI_CHAT_TOOL_RESULT_SOURCE_MARKER: &str = "openai_chat_tool_result"; +const OPENAI_RESPONSES_TOOL_RESULT_SOURCE_MARKER: &str = "openai_responses_tool_result"; const OPENAI_CHAT_TOOL_ERROR_PREFIX: &str = "[tool error]"; #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] @@ -242,6 +246,19 @@ pub enum CanonicalEmbeddingInput { StringArray(Vec), TokenArray(Vec), TokenArrayArray(Vec>), + Multimodal(Vec), +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct CanonicalEmbeddingContent { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub text: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub image: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub video: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub multi_images: Option>, } impl CanonicalEmbeddingInput { @@ -253,6 +270,9 @@ impl CanonicalEmbeddingInput { } Self::TokenArray(values) => values.is_empty(), Self::TokenArrayArray(values) => values.is_empty() || values.iter().any(Vec::is_empty), + Self::Multimodal(values) => { + values.is_empty() || values.iter().any(CanonicalEmbeddingContent::is_empty) + } } } @@ -260,11 +280,47 @@ impl CanonicalEmbeddingInput { match self { Self::String(value) => Some(vec![value.as_str()]), Self::StringArray(values) => Some(values.iter().map(String::as_str).collect()), - Self::TokenArray(_) | Self::TokenArrayArray(_) => None, + Self::TokenArray(_) | Self::TokenArrayArray(_) | Self::Multimodal(_) => None, } } } +impl CanonicalEmbeddingContent { + pub(crate) fn is_empty(&self) -> bool { + let text_empty = self + .text + .as_ref() + .is_some_and(|value| value.trim().is_empty()); + let image_empty = self + .image + .as_ref() + .is_some_and(|value| value.trim().is_empty()); + let video_empty = self + .video + .as_ref() + .is_some_and(|value| value.trim().is_empty()); + let multi_images_empty = self.multi_images.as_ref().is_some_and(|values| { + values.is_empty() || values.iter().any(|value| value.trim().is_empty()) + }); + let has_any = self + .text + .as_ref() + .is_some_and(|value| !value.trim().is_empty()) + || self + .image + .as_ref() + .is_some_and(|value| !value.trim().is_empty()) + || self + .video + .as_ref() + .is_some_and(|value| !value.trim().is_empty()) + || self.multi_images.as_ref().is_some_and(|values| { + !values.is_empty() && values.iter().all(|value| !value.trim().is_empty()) + }); + !has_any || text_empty || image_empty || video_empty || multi_images_empty + } +} + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct CanonicalEmbeddingRequest { pub input: CanonicalEmbeddingInput, @@ -276,6 +332,8 @@ pub struct CanonicalEmbeddingRequest { pub task: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub user: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub parameters: Option>, #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] pub extensions: BTreeMap, } @@ -498,6 +556,7 @@ pub(crate) fn canonical_to_embedding_request( "jina" => crate::formats::jina::embedding::request::to(canonical, &ctx), "gemini" => crate::formats::gemini::embedding::request::to(canonical, &ctx), "doubao" => crate::formats::doubao::embedding::request::to(canonical, &ctx), + "aliyun" => crate::formats::aliyun::embedding::request::to(canonical, &ctx), _ => None, } } @@ -695,6 +754,7 @@ pub fn from_embedding_to_canonical_response( } "jina" => crate::formats::openai::embedding::response::from_namespace(body_json, "jina"), "gemini" => crate::formats::gemini::embedding::response::from(body_json), + "aliyun" => crate::formats::aliyun::embedding::response::from(body_json), _ => None, } } @@ -1051,7 +1111,7 @@ pub(crate) fn claude_system_to_canonical_instructions( Some(vec![CanonicalInstruction { role: CanonicalRole::System, text, - extensions: BTreeMap::new(), + extensions: claude_system_instruction_extensions(BTreeMap::new()), }]) } } @@ -1070,7 +1130,10 @@ pub(crate) fn claude_system_to_canonical_instructions( instructions.push(CanonicalInstruction { role: CanonicalRole::System, text: strip_claude_billing_header(text), - extensions: claude_extensions(block, &["type", "text"]), + extensions: claude_system_instruction_extensions(claude_extensions( + block, + &["type", "text"], + )), }); } } @@ -1080,6 +1143,40 @@ pub(crate) fn claude_system_to_canonical_instructions( } } +fn claude_system_instruction_extensions( + mut extensions: BTreeMap, +) -> BTreeMap { + canonical_extension_object_mut(&mut extensions, AETHER_EXTENSION_NAMESPACE).insert( + "source".to_string(), + Value::String(CLAUDE_SYSTEM_SOURCE_MARKER.to_string()), + ); + extensions +} + +pub(crate) fn mark_claude_messages_request_source(extensions: &mut BTreeMap) { + canonical_extension_object_mut(extensions, AETHER_EXTENSION_NAMESPACE).insert( + "source".to_string(), + Value::String(CLAUDE_MESSAGES_REQUEST_SOURCE_MARKER.to_string()), + ); +} + +pub(crate) fn is_claude_messages_request(extensions: &BTreeMap) -> bool { + extensions + .get(AETHER_EXTENSION_NAMESPACE) + .and_then(|value| value.get("source")) + .and_then(Value::as_str) + == Some(CLAUDE_MESSAGES_REQUEST_SOURCE_MARKER) +} + +pub(crate) fn is_claude_system_instruction(instruction: &CanonicalInstruction) -> bool { + instruction + .extensions + .get(AETHER_EXTENSION_NAMESPACE) + .and_then(|value| value.get("source")) + .and_then(Value::as_str) + == Some(CLAUDE_SYSTEM_SOURCE_MARKER) +} + pub(crate) fn claude_messages_to_canonical( messages: Option<&Value>, ) -> Option> { @@ -1173,7 +1270,10 @@ pub(crate) fn claude_block_to_canonical_block(block: &Value) -> Option Some(CanonicalContentBlock::Thinking { text: String::new(), @@ -1182,7 +1282,10 @@ pub(crate) fn claude_block_to_canonical_block(block: &Value) -> Option claude_media_block_to_canonical(block_object, true), "document" => claude_media_block_to_canonical(block_object, false), @@ -1337,6 +1440,23 @@ pub(crate) fn openai_message_content_blocks( let reasoning_blocks = openai_reasoning_blocks(message); if !reasoning_blocks.is_empty() { blocks.splice(0..0, reasoning_blocks); + } else if let Some(reasoning_content) = message + .get("reasoning_content") + .and_then(Value::as_str) + .filter(|value| !value.trim().is_empty()) + { + let mut extensions = BTreeMap::new(); + canonical_extension_object_mut(&mut extensions, "openai") + .insert("omit_reasoning_parts".to_string(), Value::Bool(true)); + blocks.insert( + 0, + CanonicalContentBlock::Thinking { + text: reasoning_content.to_string(), + signature: None, + encrypted_content: None, + extensions, + }, + ); } } let mut saw_tool_calls = false; @@ -1504,6 +1624,7 @@ pub(crate) fn openai_responses_input_to_canonical_messages( Value::Array(items) => { let mut messages = Vec::new(); let mut next_generated_tool_call_index = 0usize; + let mut pending_reasoning: Option = None; for item in items { if let Some(text) = item.as_str() { if !text.trim().is_empty() { @@ -1516,9 +1637,13 @@ pub(crate) fn openai_responses_input_to_canonical_messages( extensions: BTreeMap::new(), }); } + pending_reasoning = None; continue; } - let item_object = item.as_object()?; + let Some(item_object) = item.as_object() else { + pending_reasoning = None; + continue; + }; let item_type = item_object .get("type") .and_then(Value::as_str) @@ -1526,6 +1651,12 @@ pub(crate) fn openai_responses_input_to_canonical_messages( .trim() .to_ascii_lowercase(); match item_type.as_str() { + "reasoning" => { + let reasoning = openai_responses_reasoning_text(item_object); + if !reasoning.is_empty() { + pending_reasoning = Some(reasoning); + } + } "message" => { let role = openai_role_to_canonical( item_object @@ -1548,11 +1679,13 @@ pub(crate) fn openai_responses_input_to_canonical_messages( ), }); } + pending_reasoning = None; continue; } + let is_assistant = role == CanonicalRole::Assistant; messages.push(CanonicalMessage { role, - content: openai_responses_content_to_blocks( + content: openai_responses_chat_safe_content_to_blocks( item_object.get("content"), )?, extensions: openai_responses_extensions( @@ -1560,13 +1693,17 @@ pub(crate) fn openai_responses_input_to_canonical_messages( &["type", "role", "content"], ), }); + if !is_assistant { + pending_reasoning = None; + } } "function_call" => { let name = item_object .get("name") .and_then(Value::as_str) .map(str::trim) - .filter(|value| !value.is_empty())?; + .unwrap_or_default() + .to_string(); let id = item_object .get("call_id") .or_else(|| item_object.get("id")) @@ -1580,52 +1717,20 @@ pub(crate) fn openai_responses_input_to_canonical_messages( next_generated_tool_call_index += 1; generated }); - messages.push(CanonicalMessage { - role: CanonicalRole::Assistant, - content: vec![CanonicalContentBlock::ToolUse { - id, - name: name.to_string(), - input: parse_jsonish_value(item_object.get("arguments")), - extensions: openai_responses_extensions( - item_object, - &["type", "call_id", "id", "name", "arguments"], - ), - }], - extensions: BTreeMap::new(), - }); - } - "web_search_call" => { - let id = item_object - .get("id") - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToOwned::to_owned) - .unwrap_or_else(|| { - let generated = - format!("call_auto_{next_generated_tool_call_index}"); - next_generated_tool_call_index += 1; - generated - }); - let query = item_object - .get("action") - .and_then(Value::as_object) - .and_then(|action| action.get("query")) - .and_then(Value::as_str) - .unwrap_or_default(); - messages.push(CanonicalMessage { - role: CanonicalRole::Assistant, - content: vec![CanonicalContentBlock::ToolUse { - id, - name: "web_search".to_string(), - input: json!({ "query": query }), - extensions: openai_responses_extensions( - item_object, - &["type", "id", "status", "action"], - ), - }], - extensions: BTreeMap::new(), - }); + let tool_use = CanonicalContentBlock::ToolUse { + id, + name, + input: parse_jsonish_value(item_object.get("arguments")), + extensions: openai_responses_extensions( + item_object, + &["type", "call_id", "id", "name", "arguments"], + ), + }; + append_openai_responses_tool_use( + &mut messages, + tool_use, + &mut pending_reasoning, + ); } "function_call_output" => { let id = item_object @@ -1644,6 +1749,21 @@ pub(crate) fn openai_responses_input_to_canonical_messages( }); let raw_output = item_object.get("output"); let output = Some(parse_jsonish_value(raw_output)); + let mut extensions = openai_responses_extensions( + item_object, + &[ + "type", + "call_id", + "tool_call_id", + "id", + "output", + "is_error", + ], + ); + extensions.insert( + AETHER_EXTENSION_NAMESPACE.to_string(), + json!({ "source": OPENAI_RESPONSES_TOOL_RESULT_SOURCE_MARKER }), + ); messages.push(CanonicalMessage { role: CanonicalRole::Tool, content: vec![CanonicalContentBlock::ToolResult { @@ -1655,30 +1775,15 @@ pub(crate) fn openai_responses_input_to_canonical_messages( .get("is_error") .and_then(Value::as_bool) .unwrap_or(false), - extensions: openai_responses_extensions( - item_object, - &[ - "type", - "call_id", - "tool_call_id", - "id", - "output", - "is_error", - ], - ), + extensions, }], extensions: BTreeMap::new(), }); + pending_reasoning = None; + } + _ => { + pending_reasoning = None; } - _ => messages.push(CanonicalMessage { - role: CanonicalRole::Unknown, - content: vec![CanonicalContentBlock::Unknown { - raw_type: item_type, - payload: item.clone(), - extensions: BTreeMap::new(), - }], - extensions: BTreeMap::new(), - }), } } Some(messages) @@ -1687,6 +1792,102 @@ pub(crate) fn openai_responses_input_to_canonical_messages( } } +fn append_openai_responses_tool_use( + messages: &mut Vec, + tool_use: CanonicalContentBlock, + pending_reasoning: &mut Option, +) { + let reasoning = pending_reasoning.take().filter(|value| !value.is_empty()); + if let Some(last_message) = messages.last_mut() { + if last_message.role == CanonicalRole::Assistant { + if let Some(reasoning) = reasoning { + prepend_openai_responses_reasoning_block(last_message, reasoning); + } + last_message.content.push(tool_use); + return; + } + } + + let mut content = Vec::new(); + if let Some(reasoning) = reasoning { + content.push(openai_responses_reasoning_block(reasoning)); + } + content.push(tool_use); + messages.push(CanonicalMessage { + role: CanonicalRole::Assistant, + content, + extensions: BTreeMap::new(), + }); +} + +fn prepend_openai_responses_reasoning_block(message: &mut CanonicalMessage, reasoning: String) { + if message + .content + .iter() + .any(|block| matches!(block, CanonicalContentBlock::Thinking { .. })) + { + return; + } + message + .content + .insert(0, openai_responses_reasoning_block(reasoning)); +} + +fn openai_responses_reasoning_block(text: String) -> CanonicalContentBlock { + let mut extensions = BTreeMap::new(); + canonical_extension_object_mut(&mut extensions, "openai") + .insert("omit_reasoning_parts".to_string(), Value::Bool(true)); + CanonicalContentBlock::Thinking { + text, + signature: None, + encrypted_content: None, + extensions, + } +} + +fn openai_responses_reasoning_text(item_object: &Map) -> String { + let mut parts = openai_responses_reasoning_text_parts(item_object.get("summary")); + if parts.is_empty() { + parts = openai_responses_reasoning_text_parts(item_object.get("content")); + } + parts.join("\n") +} + +fn openai_responses_reasoning_text_parts(raw: Option<&Value>) -> Vec { + let Some(raw) = raw else { + return Vec::new(); + }; + match raw { + Value::Array(items) => items + .iter() + .filter_map(openai_responses_reasoning_text_part) + .collect(), + other => openai_responses_reasoning_text_part(other) + .into_iter() + .collect(), + } +} + +fn openai_responses_reasoning_text_part(raw: &Value) -> Option { + if let Some(text) = raw.as_str() { + return (!text.is_empty()).then(|| text.to_string()); + } + let raw_object = raw.as_object()?; + let text = raw_object.get("text").and_then(Value::as_str)?; + (!text.is_empty()).then(|| text.to_string()) +} + +fn openai_responses_chat_safe_content_to_blocks( + content: Option<&Value>, +) -> Option> { + Some( + openai_responses_content_to_blocks(content)? + .into_iter() + .filter(|block| !matches!(block, CanonicalContentBlock::Unknown { .. })) + .collect(), + ) +} + pub(crate) fn openai_responses_content_to_blocks( content: Option<&Value>, ) -> Option> { @@ -1836,6 +2037,21 @@ pub(crate) fn openai_responses_output_to_canonical_blocks( .unwrap_or_else(|| format!("call_auto_{index}")); let raw_output = item_object.get("output"); let output = Some(parse_jsonish_value(raw_output)); + let mut extensions = openai_responses_extensions( + item_object, + &[ + "type", + "id", + "call_id", + "tool_call_id", + "output", + "is_error", + ], + ); + extensions.insert( + AETHER_EXTENSION_NAMESPACE.to_string(), + json!({ "source": OPENAI_RESPONSES_TOOL_RESULT_SOURCE_MARKER }), + ); blocks.push(CanonicalContentBlock::ToolResult { tool_use_id: id, name: None, @@ -1845,17 +2061,7 @@ pub(crate) fn openai_responses_output_to_canonical_blocks( .get("is_error") .and_then(Value::as_bool) .unwrap_or(false), - extensions: openai_responses_extensions( - item_object, - &[ - "type", - "id", - "call_id", - "tool_call_id", - "output", - "is_error", - ], - ), + extensions, }); } "image_generation_call" => { @@ -2381,11 +2587,7 @@ fn canonical_message_blocks_to_openai_chat( output.insert( "content".to_string(), if !tool_calls.is_empty() && content_parts.is_empty() { - if reasoning_parts.is_empty() { - Value::Array(Vec::new()) - } else { - Value::Null - } + Value::Null } else { openai_content_value_from_parts(content_parts, false) }, @@ -2432,6 +2634,11 @@ fn canonical_tool_result_to_openai_chat(block: &CanonicalContentBlock) -> Value } else { content } + } else if is_openai_responses_tool_result(extensions) { + openai_responses_tool_result_content_for_chat( + result_output.as_ref(), + content_text.as_deref(), + ) } else { result_output .clone() @@ -2441,7 +2648,23 @@ fn canonical_tool_result_to_openai_chat(block: &CanonicalContentBlock) -> Value Value::Object(output) } -fn is_claude_tool_result(extensions: &BTreeMap) -> bool { +fn claude_thinking_extensions(mut extensions: BTreeMap) -> BTreeMap { + canonical_extension_object_mut(&mut extensions, AETHER_EXTENSION_NAMESPACE).insert( + "source".to_string(), + Value::String(CLAUDE_THINKING_SOURCE_MARKER.to_string()), + ); + extensions +} + +pub(crate) fn is_claude_thinking_block(extensions: &BTreeMap) -> bool { + extensions + .get(AETHER_EXTENSION_NAMESPACE) + .and_then(|value| value.get("source")) + .and_then(Value::as_str) + == Some(CLAUDE_THINKING_SOURCE_MARKER) +} + +pub(crate) fn is_claude_tool_result(extensions: &BTreeMap) -> bool { extensions .get(AETHER_EXTENSION_NAMESPACE) .and_then(|value| value.get("source")) @@ -2449,6 +2672,28 @@ fn is_claude_tool_result(extensions: &BTreeMap) -> bool { == Some(CLAUDE_TOOL_RESULT_SOURCE_MARKER) } +fn is_openai_responses_tool_result(extensions: &BTreeMap) -> bool { + extensions + .get(AETHER_EXTENSION_NAMESPACE) + .and_then(|value| value.get("source")) + .and_then(Value::as_str) + == Some(OPENAI_RESPONSES_TOOL_RESULT_SOURCE_MARKER) +} + +fn openai_responses_tool_result_content_for_chat( + output: Option<&Value>, + content_text: Option<&str>, +) -> Value { + if let Some(text) = content_text { + return Value::String(text.to_string()); + } + match output { + Some(Value::String(text)) => Value::String(text.clone()), + Some(value) => Value::String(value.to_string()), + None => Value::String(String::new()), + } +} + fn openai_chat_tool_result_content(output: Option<&Value>, content_text: Option<&str>) -> Value { match output { Some(Value::String(text)) => Value::String(text.clone()), @@ -2828,6 +3073,9 @@ pub(crate) fn openai_content_value_from_parts(parts: Vec, tool_only: bool if parts.is_empty() && tool_only { return Value::Null; } + if parts.is_empty() { + return Value::String(String::new()); + } if parts.len() == 1 { if let Some(text) = parts[0] .as_object() @@ -3978,7 +4226,9 @@ pub(crate) fn canonical_block_to_claude( extensions, ), ); - out.insert("is_error".to_string(), Value::Bool(*is_error)); + if *is_error { + out.insert("is_error".to_string(), Value::Bool(true)); + } out.extend(namespace_extension_object(extensions, "claude", &out)); Some(Some(Value::Object(out))) } @@ -5050,8 +5300,8 @@ mod tests { from_gemini_to_canonical_request, from_gemini_to_canonical_response, from_openai_chat_to_canonical_request, from_openai_chat_to_canonical_response, from_openai_responses_to_canonical_request, from_openai_responses_to_canonical_response, - CanonicalContentBlock, CanonicalEmbedding, CanonicalEmbeddingInput, - CanonicalEmbeddingRequest, CanonicalRole, CanonicalUsage, + CanonicalContentBlock, CanonicalEmbedding, CanonicalEmbeddingContent, + CanonicalEmbeddingInput, CanonicalEmbeddingRequest, CanonicalRole, CanonicalUsage, }; use serde_json::{json, Value}; @@ -5107,6 +5357,44 @@ mod tests { "nested token array", CanonicalEmbeddingInput::TokenArrayArray(vec![vec![1, 2], vec![3, 4]]), ), + ( + json!([ + {"text": "white running shoes"}, + {"image": "https://example.com/shoe.png"}, + {"video": "https://example.com/demo.mp4"}, + {"multi_images": ["https://example.com/a.png", "https://example.com/b.png"]} + ]), + "multimodal array", + CanonicalEmbeddingInput::Multimodal(vec![ + CanonicalEmbeddingContent { + text: Some("white running shoes".to_string()), + image: None, + video: None, + multi_images: None, + }, + CanonicalEmbeddingContent { + text: None, + image: Some("https://example.com/shoe.png".to_string()), + video: None, + multi_images: None, + }, + CanonicalEmbeddingContent { + text: None, + image: None, + video: Some("https://example.com/demo.mp4".to_string()), + multi_images: None, + }, + CanonicalEmbeddingContent { + text: None, + image: None, + video: None, + multi_images: Some(vec![ + "https://example.com/a.png".to_string(), + "https://example.com/b.png".to_string(), + ]), + }, + ]), + ), ]; for (input, label, expected_input) in cases { @@ -5133,6 +5421,9 @@ mod tests { json!({"model": "text-embedding-3-small", "input": []}), json!({"model": "text-embedding-3-small", "input": [1, "two"]}), json!({"model": "text-embedding-3-small", "input": [[1], []]}), + json!({"model": "text-embedding-3-small", "input": [{"image": " "}]}), + json!({"model": "text-embedding-3-small", "input": [{"multi_images": []}]}), + json!({"model": "text-embedding-3-small", "input": ["hello", {"image": "https://example.com/a.png"}]}), json!({"model": "", "input": "hello"}), json!({"input": "hello"}), json!({"model": "text-embedding-3-small", "messages": []}), @@ -5202,6 +5493,7 @@ mod tests { dimensions: Some(2), task: None, user: None, + parameters: None, extensions: Default::default(), }), ..Default::default() @@ -5247,6 +5539,7 @@ mod tests { dimensions: Some(1536), task: Some("retrieval.passage".to_string()), user: Some("user-1".to_string()), + parameters: None, extensions: Default::default(), }), ..Default::default() @@ -5299,6 +5592,7 @@ mod tests { dimensions: None, task: None, user: None, + parameters: None, extensions: Default::default(), }), ..Default::default() @@ -5534,6 +5828,41 @@ mod tests { assert_eq!(rebuilt["n"], 2); } + #[test] + fn openai_chat_request_adapter_preserves_reasoning_content_for_responses() { + let request = json!({ + "model": "gpt-5", + "messages": [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "reasoning_content": "internal plan", + "content": "final answer" + } + ] + }); + + let canonical = from_openai_chat_to_canonical_request(&request).expect("canonical request"); + assert!(matches!( + canonical.messages[1].content.first(), + Some(CanonicalContentBlock::Thinking { text, .. }) if text == "internal plan" + )); + + let rebuilt = canonical_to_openai_responses_request(&canonical, "gpt-5-upstream", false) + .expect("openai responses request"); + let parts = rebuilt["input"][1]["content"] + .as_array() + .expect("content parts"); + + assert_eq!(parts[0]["type"], "output_text"); + assert!(parts[0]["text"] + .as_str() + .expect("reasoning text") + .contains("internal plan")); + assert_eq!(parts[1]["type"], "output_text"); + assert_eq!(parts[1]["text"], "final answer"); + } + #[test] fn openai_responses_request_adapter_preserves_audio_reasoning_tools_and_text_config() { let request = json!({ diff --git a/crates/aether-data-contracts/src/repository/global_models/types.rs b/crates/aether-data-contracts/src/repository/global_models/types.rs index 5e84d5913..97ee84425 100644 --- a/crates/aether-data-contracts/src/repository/global_models/types.rs +++ b/crates/aether-data-contracts/src/repository/global_models/types.rs @@ -7,6 +7,7 @@ const EMBEDDING_API_FORMATS: &[&str] = &[ "gemini:embedding", "jina:embedding", "doubao:embedding", + "aliyun:multimodal_embedding", "/v1/embeddings", "/jina/v1/embeddings", ]; diff --git a/crates/aether-data/src/repository/global_models/mod.rs b/crates/aether-data/src/repository/global_models/mod.rs index c9391f441..70532df20 100644 --- a/crates/aether-data/src/repository/global_models/mod.rs +++ b/crates/aether-data/src/repository/global_models/mod.rs @@ -25,6 +25,7 @@ const EMBEDDING_API_FORMATS: &[&str] = &[ "gemini:embedding", "jina:embedding", "doubao:embedding", + "aliyun:multimodal_embedding", "/v1/embeddings", "/jina/v1/embeddings", ]; diff --git a/crates/aether-data/src/repository/global_models/postgres.rs b/crates/aether-data/src/repository/global_models/postgres.rs index 5c9e03abd..54c158571 100644 --- a/crates/aether-data/src/repository/global_models/postgres.rs +++ b/crates/aether-data/src/repository/global_models/postgres.rs @@ -71,6 +71,7 @@ SELECT OR COALESCE(gm.config->'api_formats' @> '["jina:embedding"]'::jsonb, FALSE) OR COALESCE(gm.config->'api_formats' @> '["gemini:embedding"]'::jsonb, FALSE) OR COALESCE(gm.config->'api_formats' @> '["doubao:embedding"]'::jsonb, FALSE) + OR COALESCE(gm.config->'api_formats' @> '["aliyun:multimodal_embedding"]'::jsonb, FALSE) OR LOWER(COALESCE(m.config->>'embedding', 'false')) = 'true' OR LOWER(COALESCE(m.config->>'model_type', '')) = 'embedding' OR LOWER(COALESCE(m.config->>'type', '')) = 'embedding' @@ -80,6 +81,7 @@ SELECT OR COALESCE(m.config::jsonb->'api_formats' @> '["jina:embedding"]'::jsonb, FALSE) OR COALESCE(m.config::jsonb->'api_formats' @> '["gemini:embedding"]'::jsonb, FALSE) OR COALESCE(m.config::jsonb->'api_formats' @> '["doubao:embedding"]'::jsonb, FALSE) + OR COALESCE(m.config::jsonb->'api_formats' @> '["aliyun:multimodal_embedding"]'::jsonb, FALSE) ) AS supports_embedding, m.is_active FROM models m diff --git a/crates/aether-data/src/repository/wallet/memory.rs b/crates/aether-data/src/repository/wallet/memory.rs index ce511997d..93754eb5c 100644 --- a/crates/aether-data/src/repository/wallet/memory.rs +++ b/crates/aether-data/src/repository/wallet/memory.rs @@ -769,6 +769,41 @@ impl WalletReadRepository for InMemoryWalletRepository { .cloned()) } + async fn find_pending_plan_purchase_order_by_user_id( + &self, + user_id: &str, + product_id: &str, + ) -> Result, DataLayerError> { + let now = current_unix_secs(); + Ok(self + .payment_orders_by_id + .read() + .expect("wallet repo lock") + .values() + .filter(|order| { + order.user_id.as_deref() == Some(user_id) + && order.status == "pending" + && order + .expires_at_unix_secs + .is_some_and(|expires_at| expires_at > now) + && order + .gateway_response + .as_ref() + .is_some_and(|gateway_response| { + gateway_response + .get("order_kind") + .and_then(serde_json::Value::as_str) + == Some("plan_purchase") + && gateway_response + .get("product_id") + .and_then(serde_json::Value::as_str) + == Some(product_id) + }) + }) + .max_by_key(|order| order.created_at_unix_ms) + .cloned()) + } + async fn find_wallet_refund( &self, wallet_id: &str, diff --git a/crates/aether-data/src/repository/wallet/mysql.rs b/crates/aether-data/src/repository/wallet/mysql.rs index 4433fe41d..779d2aca7 100644 --- a/crates/aether-data/src/repository/wallet/mysql.rs +++ b/crates/aether-data/src/repository/wallet/mysql.rs @@ -536,6 +536,32 @@ WHERE wallet_id = ? .await } + async fn find_pending_plan_purchase_order_by_user_id( + &self, + user_id: &str, + product_id: &str, + ) -> Result, DataLayerError> { + let sql = payment_order_select_sql( + r#" +WHERE user_id = ? + AND product_id = ? + AND order_kind = 'plan_purchase' + AND status = 'pending' + AND expires_at > ? +ORDER BY created_at DESC +LIMIT 1 +"#, + ); + let row = sqlx::query(&sql) + .bind(user_id) + .bind(product_id) + .bind(current_unix_secs_i64()) + .fetch_optional(&self.pool) + .await + .map_sql_err()?; + row.as_ref().map(map_payment_order_row).transpose() + } + async fn find_wallet_refund( &self, wallet_id: &str, diff --git a/crates/aether-data/src/repository/wallet/postgres.rs b/crates/aether-data/src/repository/wallet/postgres.rs index b8fb5cc98..6fc0e3b66 100644 --- a/crates/aether-data/src/repository/wallet/postgres.rs +++ b/crates/aether-data/src/repository/wallet/postgres.rs @@ -577,6 +577,41 @@ WHERE user_id = $1 LIMIT 1 "#; +const FIND_PENDING_PLAN_PURCHASE_ORDER_BY_USER_SQL: &str = r#" +SELECT + id, + order_no, + wallet_id, + user_id, + CAST(amount_usd AS DOUBLE PRECISION) AS amount_usd, + CAST(pay_amount AS DOUBLE PRECISION) AS pay_amount, + pay_currency, + CAST(exchange_rate AS DOUBLE PRECISION) AS exchange_rate, + CAST(refunded_amount_usd AS DOUBLE PRECISION) AS refunded_amount_usd, + CAST(refundable_amount_usd AS DOUBLE PRECISION) AS refundable_amount_usd, + payment_method, + payment_provider, + payment_channel, + order_kind, + product_id, + product_snapshot, + gateway_order_id, + gateway_response, + status, + CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_ms, + CAST(EXTRACT(EPOCH FROM paid_at) AS BIGINT) AS paid_at_unix_secs, + CAST(EXTRACT(EPOCH FROM credited_at) AS BIGINT) AS credited_at_unix_secs, + CAST(EXTRACT(EPOCH FROM expires_at) AS BIGINT) AS expires_at_unix_secs +FROM payment_orders +WHERE user_id = $1 + AND product_id = $2 + AND order_kind = 'plan_purchase' + AND status = 'pending' + AND expires_at > NOW() +ORDER BY created_at DESC +LIMIT 1 +"#; + const FIND_WALLET_REFUND_SQL: &str = r#" SELECT id, @@ -1112,6 +1147,20 @@ impl WalletReadRepository for SqlxWalletRepository { row.as_ref().map(map_admin_payment_order_row).transpose() } + async fn find_pending_plan_purchase_order_by_user_id( + &self, + user_id: &str, + product_id: &str, + ) -> Result, DataLayerError> { + let row = sqlx::query(FIND_PENDING_PLAN_PURCHASE_ORDER_BY_USER_SQL) + .bind(user_id) + .bind(product_id) + .fetch_optional(&self.pool) + .await + .map_postgres_err()?; + row.as_ref().map(map_admin_payment_order_row).transpose() + } + async fn find_wallet_refund( &self, wallet_id: &str, diff --git a/crates/aether-data/src/repository/wallet/sqlite.rs b/crates/aether-data/src/repository/wallet/sqlite.rs index 77eedce56..7469af7b2 100644 --- a/crates/aether-data/src/repository/wallet/sqlite.rs +++ b/crates/aether-data/src/repository/wallet/sqlite.rs @@ -783,6 +783,32 @@ LIMIT 1 row.as_ref().map(map_payment_order_row).transpose() } + async fn find_pending_plan_purchase_order_by_user_id( + &self, + user_id: &str, + product_id: &str, + ) -> Result, DataLayerError> { + let sql = payment_order_select_sql( + r#" +WHERE user_id = ? + AND product_id = ? + AND order_kind = 'plan_purchase' + AND status = 'pending' + AND expires_at > ? +ORDER BY created_at DESC +LIMIT 1 +"#, + ); + let row = sqlx::query(&sql) + .bind(user_id) + .bind(product_id) + .bind(current_unix_secs_i64()) + .fetch_optional(&self.pool) + .await + .map_sql_err()?; + row.as_ref().map(map_payment_order_row).transpose() + } + async fn find_wallet_refund( &self, wallet_id: &str, @@ -5460,6 +5486,175 @@ INSERT INTO billing_plans ( assert_eq!(wallet_balance, 0.0); } + #[tokio::test] + async fn sqlite_finds_reusable_pending_plan_purchase_order() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_sqlite_migrations(&pool) + .await + .expect("sqlite migrations should run"); + + let repository = SqliteWalletReadRepository::new(pool); + sqlx::query( + "INSERT INTO users (id, username, email, auth_source, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)", + ) + .bind("user-pending-plan-1") + .bind("Pending Buyer") + .bind("pending-plan@example.com") + .bind("local") + .bind(1_i64) + .bind(1_i64) + .execute(repository.pool()) + .await + .expect("user should seed"); + + let _wallet_order = match repository + .create_wallet_recharge_order(CreateWalletRechargeOrderInput { + preferred_wallet_id: Some("wallet-pending-plan-1".to_string()), + user_id: "user-pending-plan-1".to_string(), + amount_usd: 1.0, + pay_amount: Some(1.0), + pay_currency: Some("USD".to_string()), + exchange_rate: Some(1.0), + payment_method: "bootstrap".to_string(), + payment_provider: None, + payment_channel: None, + gateway_order_id: "gateway-bootstrap-pending-plan-1".to_string(), + gateway_response: json!({ "bootstrap": true }), + order_no: "order-bootstrap-pending-plan-1".to_string(), + expires_at_unix_secs: 4_102_444_800, + }) + .await + .expect("wallet should be created") + { + CreateWalletRechargeOrderOutcome::Created(order) => order, + CreateWalletRechargeOrderOutcome::WalletInactive => { + panic!("new wallet should be active") + } + }; + + let plan_snapshot = json!({ + "id": "pending-plan", + "title": "每日额度月卡", + "duration_unit": "month", + "duration_value": 1, + "max_active_per_user": 1, + "purchase_limit_scope": "active_period", + "entitlements": [ + { + "type": "daily_quota", + "daily_quota_usd": 50.0, + "reset_timezone": "Asia/Shanghai", + "allow_wallet_overage": false + } + ] + }); + let pending_order = match repository + .create_plan_purchase_order(CreatePlanPurchaseOrderInput { + preferred_wallet_id: None, + user_id: "user-pending-plan-1".to_string(), + amount_usd: 13.8, + pay_amount: 100.0, + pay_currency: "CNY".to_string(), + exchange_rate: 7.24637681, + payment_method: "alipay".to_string(), + payment_provider: Some("epay".to_string()), + payment_channel: Some("alipay".to_string()), + gateway_order_id: "gateway-pending-plan-1".to_string(), + gateway_response: json!({ "checkout": true }), + order_no: "order-pending-plan-1".to_string(), + product_id: "pending-plan".to_string(), + product_snapshot: plan_snapshot.clone(), + expires_at_unix_secs: 4_102_444_800, + }) + .await + .expect("pending plan order should create") + { + CreatePlanPurchaseOrderOutcome::Created(order) => order, + other => panic!("pending plan order should be created, got {other:?}"), + }; + let now = chrono::Utc::now().timestamp().max(0); + for (id, order_no, status, product_id, user_id, expires_at, created_at) in [ + ( + "expired-pending-plan-order", + "order-expired-pending-plan", + "pending", + "pending-plan", + "user-pending-plan-1", + now - 10, + now + 10, + ), + ( + "credited-pending-plan-order", + "order-credited-pending-plan", + "credited", + "pending-plan", + "user-pending-plan-1", + now + 3_600, + now + 20, + ), + ( + "other-user-pending-plan-order", + "order-other-user-pending-plan", + "pending", + "pending-plan", + "other-user", + now + 3_600, + now + 30, + ), + ] { + sqlx::query( + r#" +INSERT INTO payment_orders ( + id, order_no, wallet_id, user_id, amount_usd, pay_amount, pay_currency, + exchange_rate, refunded_amount_usd, refundable_amount_usd, payment_method, + payment_provider, payment_channel, order_kind, product_id, product_snapshot, + fulfillment_status, gateway_order_id, gateway_response, status, created_at, expires_at +) VALUES (?, ?, ?, ?, 13.8, 100.0, 'CNY', 7.24637681, 0, 0, 'alipay', + 'epay', 'alipay', 'plan_purchase', ?, ?, 'pending', ?, ?, ?, ?, ?) + "#, + ) + .bind(id) + .bind(order_no) + .bind("wallet-pending-plan-1") + .bind(user_id) + .bind(product_id) + .bind(plan_snapshot.to_string()) + .bind(format!("gateway-{id}")) + .bind(json!({ "checkout": id }).to_string()) + .bind(status) + .bind(created_at) + .bind(expires_at) + .execute(repository.pool()) + .await + .expect("extra payment order should seed"); + } + + let found = repository + .find_pending_plan_purchase_order_by_user_id("user-pending-plan-1", "pending-plan") + .await + .expect("pending plan lookup should run") + .expect("pending plan order should be found"); + assert_eq!(found.id, pending_order.id); + assert_eq!( + repository + .find_pending_plan_purchase_order_by_user_id("user-pending-plan-1", "missing-plan") + .await + .expect("missing plan lookup should run"), + None + ); + assert_eq!( + repository + .find_pending_plan_purchase_order_by_user_id("missing-user", "pending-plan") + .await + .expect("missing user lookup should run"), + None + ); + } + #[tokio::test] async fn sqlite_plan_purchase_replaces_same_class_entitlements_on_manual_credit() { let pool = sqlx::sqlite::SqlitePoolOptions::new() diff --git a/crates/aether-data/src/repository/wallet/types.rs b/crates/aether-data/src/repository/wallet/types.rs index d7f1368c0..50313cad3 100644 --- a/crates/aether-data/src/repository/wallet/types.rs +++ b/crates/aether-data/src/repository/wallet/types.rs @@ -927,6 +927,12 @@ pub trait WalletReadRepository: Send + Sync { order_id: &str, ) -> Result, crate::DataLayerError>; + async fn find_pending_plan_purchase_order_by_user_id( + &self, + user_id: &str, + product_id: &str, + ) -> Result, crate::DataLayerError>; + async fn find_wallet_refund( &self, wallet_id: &str, diff --git a/crates/aether-provider-transport/src/conversion.rs b/crates/aether-provider-transport/src/conversion.rs index 21961e4ff..193514fc9 100644 --- a/crates/aether-provider-transport/src/conversion.rs +++ b/crates/aether-provider-transport/src/conversion.rs @@ -184,7 +184,11 @@ pub fn request_pair_transport_unsupported_reason( ) } } - "openai:embedding" | "jina:embedding" | "doubao:embedding" | "openai:rerank" + "openai:embedding" + | "jina:embedding" + | "doubao:embedding" + | "aliyun:multimodal_embedding" + | "openai:rerank" | "jina:rerank" => local_standard_transport_unsupported_reason_with_network( transport, provider_api_format.as_str(), @@ -221,6 +225,7 @@ fn request_direct_auth_for_provider_format( | "openai:embedding" | "jina:embedding" | "doubao:embedding" + | "aliyun:multimodal_embedding" | "openai:rerank" | "jina:rerank" => resolve_local_openai_bearer_auth(transport), "gemini:generate_content" | "gemini:embedding" => { diff --git a/crates/aether-provider-transport/src/policy.rs b/crates/aether-provider-transport/src/policy.rs index c08366182..75d9da03e 100644 --- a/crates/aether-provider-transport/src/policy.rs +++ b/crates/aether-provider-transport/src/policy.rs @@ -225,7 +225,7 @@ fn endpoint_kind_allows_embedding(endpoint_kind: Option<&str>) -> bool { .map(|value| { matches!( value.to_ascii_lowercase().as_str(), - "embedding" | "embeddings" + "embedding" | "embeddings" | "multimodal_embedding" | "multimodal_embeddings" ) }) .unwrap_or(true) @@ -349,10 +349,13 @@ mod tests { ("jina", "jina:embedding"), ("doubao", "doubao:embedding"), ("volcengine", "doubao:embedding"), + ("aliyun", "aliyun:multimodal_embedding"), + ("dashscope", "aliyun:multimodal_embedding"), ("custom", "openai:embedding"), ("custom", "gemini:embedding"), ("custom", "jina:embedding"), ("custom", "doubao:embedding"), + ("custom", "aliyun:multimodal_embedding"), ] { let transport = sample_transport(provider_type, api_format, Some("embedding")); assert_eq!( diff --git a/crates/aether-provider-transport/src/provider_types.rs b/crates/aether-provider-transport/src/provider_types.rs index 74c5bbb70..42916ada7 100644 --- a/crates/aether-provider-transport/src/provider_types.rs +++ b/crates/aether-provider-transport/src/provider_types.rs @@ -84,6 +84,7 @@ pub enum ProviderLocalEmbeddingSupport { Gemini, Jina, Doubao, + Aliyun, } impl ProviderLocalEmbeddingSupport { @@ -99,11 +100,13 @@ impl ProviderLocalEmbeddingSupport { | "jina:embedding" | "jina:rerank" | "doubao:embedding" + | "aliyun:multimodal_embedding" ), Self::OpenAi => matches!(api_format.as_str(), "openai:embedding" | "openai:rerank"), Self::Gemini => api_format == "gemini:embedding", Self::Jina => matches!(api_format.as_str(), "jina:embedding" | "jina:rerank"), Self::Doubao => api_format == "doubao:embedding", + Self::Aliyun => api_format == "aliyun:multimodal_embedding", } } } @@ -192,6 +195,10 @@ const DOUBAO_RUNTIME_POLICY: ProviderRuntimePolicy = ProviderRuntimePolicy { local_embedding_support: ProviderLocalEmbeddingSupport::Doubao, ..STANDARD_RUNTIME_POLICY }; +const ALIYUN_RUNTIME_POLICY: ProviderRuntimePolicy = ProviderRuntimePolicy { + local_embedding_support: ProviderLocalEmbeddingSupport::Aliyun, + ..STANDARD_RUNTIME_POLICY +}; const CLAUDE_CODE_RUNTIME_POLICY: ProviderRuntimePolicy = ProviderRuntimePolicy { fixed_provider: true, @@ -470,6 +477,7 @@ pub fn provider_runtime_policy(provider_type: &str) -> ProviderRuntimePolicy { "gemini" | "google" => GEMINI_RUNTIME_POLICY, "jina" => JINA_RUNTIME_POLICY, "doubao" | "volcengine" => DOUBAO_RUNTIME_POLICY, + "aliyun" | "dashscope" => ALIYUN_RUNTIME_POLICY, _ => STANDARD_RUNTIME_POLICY, } } @@ -875,6 +883,8 @@ mod tests { ("jina", "jina:embedding"), ("doubao", "doubao:embedding"), ("volcengine", "doubao:embedding"), + ("aliyun", "aliyun:multimodal_embedding"), + ("dashscope", "aliyun:multimodal_embedding"), ] { assert!( provider_type_supports_local_embedding_transport(provider_type, api_format), @@ -888,6 +898,8 @@ mod tests { ("vertex_ai", "openai:embedding"), ("jina", "doubao:embedding"), ("doubao", "jina:embedding"), + ("aliyun", "openai:embedding"), + ("openai", "aliyun:multimodal_embedding"), ("claude_code", "openai:embedding"), ("openai", "openai:chat"), ] { diff --git a/crates/aether-provider-transport/src/request_url/mod.rs b/crates/aether-provider-transport/src/request_url/mod.rs index 65e88407c..e9d5b4ac4 100644 --- a/crates/aether-provider-transport/src/request_url/mod.rs +++ b/crates/aether-provider-transport/src/request_url/mod.rs @@ -127,6 +127,10 @@ fn build_transport_request_url_inner( "openai:embedding" | "jina:embedding" => { build_provider_embedding_v1_url(&transport.endpoint.base_url, params.request_query) } + "aliyun:multimodal_embedding" => build_aliyun_multimodal_embedding_url( + &transport.endpoint.base_url, + params.request_query, + ), "openai:rerank" | "jina:rerank" => { build_provider_rerank_v1_url(&transport.endpoint.base_url, params.request_query) } @@ -425,6 +429,18 @@ fn build_provider_embedding_v1_url(upstream_base_url: &str, query: Option<&str>) build_provider_api_root_url(upstream_base_url, "/embeddings", query) } +fn build_aliyun_multimodal_embedding_url( + upstream_base_url: &str, + query: Option<&str>, +) -> Option { + build_passthrough_path_url( + upstream_base_url, + "/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding", + query, + &[], + ) +} + fn build_provider_rerank_v1_url(upstream_base_url: &str, query: Option<&str>) -> Option { build_provider_api_root_url(upstream_base_url, "/rerank", query) } @@ -1019,6 +1035,12 @@ mod tests { "https://ark.volces.example/api/v3", None, ); + let aliyun = sample_transport( + "aliyun", + "aliyun:multimodal_embedding", + "https://dashscope.aliyuncs.com", + None, + ); assert_eq!( build_transport_request_url( @@ -1078,6 +1100,20 @@ mod tests { .as_deref(), Some("https://ark.volces.example/api/v3/embeddings") ); + assert_eq!( + build_transport_request_url( + &aliyun, + TransportRequestUrlParams { + provider_api_format: "aliyun:multimodal_embedding", + mapped_model: Some("qwen3-vl-embedding"), + upstream_is_stream: false, + request_query: None, + kiro_api_region: None, + }, + ) + .as_deref(), + Some("https://dashscope.aliyuncs.com/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding") + ); } #[test] diff --git a/crates/aether-provider-transport/src/same_format_provider/mod.rs b/crates/aether-provider-transport/src/same_format_provider/mod.rs index 4a9c7796d..fcffc2dd0 100644 --- a/crates/aether-provider-transport/src/same_format_provider/mod.rs +++ b/crates/aether-provider-transport/src/same_format_provider/mod.rs @@ -169,6 +169,14 @@ pub fn build_same_format_provider_request_body( ); } + if embedding_multimodal_input_requires_aliyun_provider( + input.client_api_format, + input.provider_api_format, + input.body_json, + ) { + return None; + } + let mut provider_request_body = if aether_ai_formats::api_format_alias_matches( input.client_api_format, input.provider_api_format, @@ -245,6 +253,31 @@ pub fn build_same_format_provider_request_body( Some(provider_request_body) } +fn embedding_multimodal_input_requires_aliyun_provider( + client_api_format: &str, + provider_api_format: &str, + body_json: &Value, +) -> bool { + aether_ai_formats::is_embedding_api_format(client_api_format) + && embedding_input_is_multimodal(body_json.get("input")) + && aether_ai_formats::normalize_api_format_alias(provider_api_format) + != "aliyun:multimodal_embedding" +} + +fn embedding_input_is_multimodal(value: Option<&Value>) -> bool { + value + .and_then(Value::as_array) + .is_some_and(|items| !items.is_empty() && items.iter().all(embedding_content_is_multimodal)) +} + +fn embedding_content_is_multimodal(value: &Value) -> bool { + value.as_object().is_some_and(|object| { + ["text", "image", "video", "multi_images"] + .iter() + .any(|key| object.contains_key(*key)) + }) +} + fn strip_gemini_function_response_ids(value: &mut Value) { match value { Value::Object(object) => { @@ -849,6 +882,33 @@ mod tests { assert_eq!(body.get("stream"), Some(&json!(true))); } + #[test] + fn same_format_embedding_body_rejects_multimodal_for_openai_like_provider() { + let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput { + body_json: &json!({ + "model": "qwen3-vl-embedding", + "input": [ + {"text": "white running shoes"}, + {"image": "https://example.com/shoe.png"} + ] + }), + mapped_model: "openai-qwen-fallback", + client_api_format: "openai:embedding", + provider_api_format: "openai:embedding", + source_model: Some("qwen3-vl-embedding"), + family: SameFormatProviderFamily::Standard, + body_rules: None, + request_headers: None, + upstream_is_stream: false, + force_body_stream_field: false, + kiro_auth_config: None, + is_claude_code: false, + enable_model_directives: false, + }); + + assert!(body.is_none()); + } + #[test] fn same_format_standard_body_overrides_client_stream_for_non_stream_upstream() { let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput { diff --git a/crates/aether-usage-runtime/src/write.rs b/crates/aether-usage-runtime/src/write.rs index 183229098..3fbc2a37b 100644 --- a/crates/aether-usage-runtime/src/write.rs +++ b/crates/aether-usage-runtime/src/write.rs @@ -146,6 +146,7 @@ pub struct SyncTerminalUsagePayloadSeed { pub provider_response_body_state: Option, pub client_response: Option, pub client_response_body_state: Option, + pub standardized_usage: Option, pub capture_metadata: Option, } @@ -735,7 +736,10 @@ pub fn build_sync_terminal_usage_payload_seed( .and_then(|context| context.get(UPSTREAM_IS_STREAM_KEY)) .and_then(Value::as_bool) .unwrap_or(false); - let provider_response_full = if upstream_is_stream && payload.body_base64.is_some() { + let provider_response_full = if upstream_is_stream + && payload.body_base64.is_some() + && !body_json_has_terminal_error(payload.body_json.as_ref()) + { decode_body_for_storage(payload.body_base64.as_deref()) .or_else(|| payload.body_json.as_ref().cloned()) } else { @@ -763,6 +767,7 @@ pub fn build_sync_terminal_usage_payload_seed( .or_else(|| headers_to_json(&payload.headers)); let client_response_headers = context_usage_value(context, "client_response_headers") .or_else(|| headers_to_json(&payload.headers)); + let standardized_usage = kiro_simulated_cache_standardized_usage_from_context(context); SyncTerminalUsagePayloadSeed { report_kind: payload.report_kind.clone(), status_code: payload.status_code, @@ -777,6 +782,7 @@ pub fn build_sync_terminal_usage_payload_seed( provider_response_body_state, client_response, client_response_body_state, + standardized_usage, capture_metadata: build_payload_body_capture_metadata( payload.body_base64.as_deref(), None, @@ -786,6 +792,12 @@ pub fn build_sync_terminal_usage_payload_seed( } } +fn body_json_has_terminal_error(body_json: Option<&Value>) -> bool { + body_json + .and_then(|value| value.get("error")) + .is_some_and(|error| !error.is_null()) +} + pub fn build_stream_terminal_usage_payload_seed( payload: &GatewayStreamReportRequest, ) -> StreamTerminalUsagePayloadSeed { @@ -847,11 +859,14 @@ pub fn build_sync_terminal_usage_seed( provider_response_body_state, client_response, client_response_body_state, + standardized_usage, capture_metadata, } = payload_seed; - let standardized_usage = provider_response_full + let derived_standardized_usage = provider_response_full .as_ref() .map(|response| map_usage_from_response(response, context_seed.provider_contract.as_str())); + let standardized_usage = + merge_standardized_usage_with_context_cache(standardized_usage, derived_standardized_usage); let terminal_state = infer_sync_terminal_state( report_kind.as_str(), status_code, @@ -905,6 +920,25 @@ pub fn build_sync_terminal_usage_seed( } } +fn merge_standardized_usage_with_context_cache( + context_usage: Option, + derived_usage: Option, +) -> Option { + let Some(context_usage) = context_usage else { + return derived_usage; + }; + + let mut usage = derived_usage.unwrap_or_default(); + usage.input_tokens = context_usage.input_tokens; + if context_usage.cache_creation_tokens > 0 { + usage.cache_creation_tokens = context_usage.cache_creation_tokens; + } + if context_usage.cache_read_tokens > 0 { + usage.cache_read_tokens = context_usage.cache_read_tokens; + } + Some(usage) +} + pub fn build_stream_terminal_usage_seed( context_seed: TerminalUsageContextSeed, payload_seed: StreamTerminalUsagePayloadSeed, @@ -1751,6 +1785,31 @@ fn context_body_value(context: Option<&Map>, key: &str) -> Option } } +fn kiro_simulated_cache_standardized_usage_from_context( + context: Option<&Map>, +) -> Option { + let enabled = context_bool(context, "kiro_simulated_cache_enabled").unwrap_or(false); + if !enabled { + return None; + } + + let input_tokens = context_u64(context, "input_tokens")?; + let cache_creation_tokens = context_u64(context, "cache_creation_input_tokens").unwrap_or(0); + let cache_read_tokens = context_u64(context, "cache_read_input_tokens").unwrap_or(0); + if cache_creation_tokens == 0 && cache_read_tokens == 0 { + return None; + } + + let billed_input_tokens = input_tokens + .saturating_sub(cache_creation_tokens) + .saturating_sub(cache_read_tokens); + let mut usage = StandardizedUsage::new(); + usage.input_tokens = billed_input_tokens as i64; + usage.cache_creation_tokens = cache_creation_tokens as i64; + usage.cache_read_tokens = cache_read_tokens as i64; + Some(usage) +} + fn context_has_inline_body(context: Option<&Map>, key: &str) -> bool { matches!(context_value_ref(context, key), Some(value) if !value.is_null()) } @@ -5108,6 +5167,87 @@ mod tests { ); } + #[test] + fn sync_terminal_usage_prefers_error_body_over_partial_upstream_stream_body() { + let partial_sse_body = concat!( + "event: response.created\n", + "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_partial_123\",\"object\":\"response\",\"model\":\"gpt-5.5\",\"status\":\"in_progress\",\"output\":[]}}\n\n", + "event: response.output_item.added\n", + "data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"name\":\"exec_command\"}}\n\n", + "event: response.function_call_arguments.delta\n", + "data: {\"type\":\"response.function_call_arguments.delta\",\"delta\":\"{\\\"cmd\\\":\"}\n\n", + ); + let plan = ExecutionPlan { + request_id: "req-sync-upstream-stream-error-1".to_string(), + candidate_id: Some("cand-sync-upstream-stream-error-1".to_string()), + provider_name: Some("OpenAI".to_string()), + provider_id: "provider-1".to_string(), + endpoint_id: "endpoint-1".to_string(), + key_id: "key-1".to_string(), + method: "POST".to_string(), + url: "https://example.com/v1/responses".to_string(), + headers: BTreeMap::new(), + content_type: None, + content_encoding: None, + body: RequestBody { + json_body: None, + body_bytes_b64: None, + body_ref: None, + }, + stream: false, + client_api_format: "claude:messages".to_string(), + provider_api_format: "openai:responses".to_string(), + model_name: Some("gpt-5.5".to_string()), + proxy: None, + transport_profile: None, + timeouts: None, + }; + let payload = GatewaySyncReportRequest { + trace_id: "trace-sync-upstream-stream-error-1".to_string(), + report_kind: "openai_responses_sync_error".to_string(), + report_context: Some(json!({ + "client_api_format": "claude:messages", + "provider_api_format": "openai:responses", + "upstream_is_stream": true, + "needs_conversion": true + })), + status_code: 200, + headers: BTreeMap::from([( + "content-type".to_string(), + "text/event-stream".to_string(), + )]), + body_json: Some(json!({ + "error": { + "type": "internal", + "message": "error decoding response body: stream error received" + } + })), + client_body_json: None, + body_base64: Some(base64::engine::general_purpose::STANDARD.encode(partial_sse_body)), + telemetry: None, + }; + + let event = + build_sync_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload) + .expect("usage event should build"); + + assert_eq!(event.event_type, UsageEventType::Failed); + assert_eq!(event.data.status_code, Some(200)); + assert_eq!( + event.data.error_message.as_deref(), + Some("error decoding response body: stream error received") + ); + assert_eq!( + event + .data + .response_body + .as_ref() + .and_then(|value| value.pointer("/error/type")) + .and_then(Value::as_str), + Some("internal") + ); + } + #[test] fn sync_terminal_seed_path_matches_legacy_wrapper_event() { let plan = ExecutionPlan { @@ -5453,6 +5593,64 @@ mod tests { ); } + #[test] + fn sync_terminal_usage_applies_kiro_simulated_cache_context() { + let plan = ExecutionPlan { + request_id: "req-sync-kiro-cache-context-1".to_string(), + candidate_id: Some("cand-sync-kiro-cache-context-1".to_string()), + provider_name: Some("Kiro".to_string()), + provider_id: "provider-kiro-1".to_string(), + endpoint_id: "endpoint-kiro-1".to_string(), + key_id: "key-kiro-1".to_string(), + method: "POST".to_string(), + url: "https://kiro.example/generateAssistantResponse".to_string(), + headers: BTreeMap::new(), + content_type: Some("application/json".to_string()), + content_encoding: None, + body: RequestBody::from_json(json!({ + "model": "claude-sonnet-4", + "messages": [{"role": "user", "content": "hello kiro"}], + })), + stream: false, + client_api_format: "claude:messages".to_string(), + provider_api_format: "claude:messages".to_string(), + model_name: Some("claude-sonnet-4".to_string()), + proxy: None, + transport_profile: None, + timeouts: None, + }; + let payload = GatewaySyncReportRequest { + trace_id: "trace-sync-kiro-cache-context-1".to_string(), + report_kind: "claude_cli_sync_success".to_string(), + report_context: Some(json!({ + "client_api_format": "claude:messages", + "provider_api_format": "claude:messages", + "provider_name": "Kiro", + "model": "claude-sonnet-4", + "input_tokens": 1800, + "kiro_simulated_cache_enabled": true, + "cache_creation_input_tokens": 1200, + "cache_read_input_tokens": 300, + })), + status_code: 200, + headers: BTreeMap::new(), + body_json: Some(json!({"id": "kiro-sync-response-1"})), + client_body_json: None, + body_base64: None, + telemetry: None, + }; + + let event = + build_sync_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload) + .expect("usage event should build"); + + assert_eq!(event.event_type, UsageEventType::Completed); + assert_eq!(event.data.input_tokens, Some(300)); + assert_eq!(event.data.cache_creation_input_tokens, Some(1200)); + assert_eq!(event.data.cache_read_input_tokens, Some(300)); + assert_eq!(event.data.total_tokens, Some(300)); + } + #[test] fn sync_terminal_usage_treats_null_error_field_as_success() { let plan = ExecutionPlan { diff --git a/deploy.sh b/deploy.sh index bbb411178..509523d34 100755 --- a/deploy.sh +++ b/deploy.sh @@ -14,7 +14,7 @@ export LOCAL_APP_IMAGE detect_build_version() { if command -v git >/dev/null 2>&1; then local version - if version=$(git describe --tags --always --dirty 2>/dev/null); then + if version=$(git describe --tags --match 'v[0-9]*' --always --dirty 2>/dev/null); then if [ -n "$version" ]; then printf '%s\n' "$version" return 0 @@ -58,7 +58,7 @@ Options: Environment: LOCAL_APP_IMAGE 本地构建镜像名，默认 aether-app:latest - AETHER_BUILD_VERSION 应用显示版本，默认 git describe --tags --always --dirty + AETHER_BUILD_VERSION 应用显示版本，默认 git describe --tags --match 'v[0-9]*' --always --dirty EOF } diff --git a/docs/api/embeddings.md b/docs/api/embeddings.md index ccedcdf68..77cdb25ec 100644 --- a/docs/api/embeddings.md +++ b/docs/api/embeddings.md @@ -32,7 +32,7 @@ Required fields: | Field | Type | Notes | | --- | --- | --- | | `model` | string | Must name a model allowed for the API key and user. Blank strings are rejected. | -| `input` | string, string array, integer token array, or nested integer token arrays | Must be non empty. Empty strings, empty arrays, and nested arrays with empty token arrays are rejected. | +| `input` | string, string array, integer token array, nested integer token arrays, or multimodal object array | Must be non empty. Empty strings, empty arrays, empty token arrays, and empty multimodal objects are rejected. | Optional fields that pass through the embedding conversion path when supported by the provider: @@ -40,6 +40,7 @@ Optional fields that pass through the embedding conversion path when supported b | --- | --- | | `encoding_format` | Passed to OpenAI compatible providers. | | `dimensions` | Passed to providers whose embedding request shape supports it. | +| `parameters` | Provider-specific embedding parameters. For Aliyun DashScope this maps to DashScope `parameters`; `dimensions` is emitted as `parameters.dimension` unless `parameters.dimension` is already set. | | `user` | Passed to OpenAI compatible providers. | | `task` | Passed to Jina and OpenAI compatible embedding requests. Jina defaults to `text-matching` when no task is supplied. | @@ -61,7 +62,18 @@ Accepted `input` shapes: { "model": "text-embedding-3-small", "input": [[1, 2], [3, 4]] } ``` -Use string or string array input when routing to Gemini or Doubao embedding providers. Token arrays are accepted by the OpenAI compatible public endpoint, but Gemini and Doubao provider request emitters require text input. +```json +{ + "model": "qwen3-vl-embedding", + "input": [ + { "text": "white running shoes" }, + { "image": "https://dashscope.oss-cn-beijing.aliyuncs.com/images/256_1.png" } + ], + "parameters": { "enable_fusion": true } +} +``` + +Use string or string array input when routing to Gemini or Doubao embedding providers. Token arrays are accepted by the OpenAI compatible public endpoint, but Gemini, Doubao, and Aliyun provider request emitters require text or multimodal content input. ## Provider Format Mapping @@ -73,6 +85,7 @@ Embedding routes can select only embedding provider API formats. Chat, responses | `jina:embedding` | `/v1/embeddings` | OpenAI compatible payload with a Jina `task`. Defaults to `text-matching` if omitted. | | `gemini:embedding` | `models/{model}:embedContent` | Single text input uses `content.parts[].text`. Multiple text inputs use `requests[].content.parts[].text`. | | `doubao:embedding` | `/embeddings/multimodal` | Text input is emitted as `input` items like `{ "type": "text", "text": "..." }`. | +| `aliyun:multimodal_embedding` | `/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding` | Text and multimodal inputs are emitted as DashScope `input.contents`. Supports `text`, `image`, `video`, `multi_images`, `parameters.enable_fusion`, `parameters.res_level`, and `parameters.max_video_frames`. Alias: `dashscope:multimodal_embedding`. | Custom provider endpoint paths are available when the endpoint is configured for an embedding API format. Gemini custom paths can use `{model}` and `{action}`. For `gemini:embedding`, `{action}` expands to `embedContent`. @@ -81,7 +94,7 @@ Custom provider endpoint paths are available when the endpoint is configured for To use embeddings through the gateway: 1. The global model should include embedding metadata, for example `supported_capabilities: ["embedding"]`, `config.model_type: "embedding"`, or `config.api_formats` with one of the embedding formats. -2. The provider model or mapping must expose an embedding API format, one of `openai:embedding`, `gemini:embedding`, `jina:embedding`, or `doubao:embedding`. +2. The provider model or mapping must expose an embedding API format, one of `openai:embedding`, `gemini:embedding`, `jina:embedding`, `doubao:embedding`, or `aliyun:multimodal_embedding`. 3. The user and API key must be allowed to access the model and the `openai:embedding` client API format. 4. Public and admin catalog responses expose `supports_embedding` so clients can display embedding capability separately from chat. @@ -89,6 +102,71 @@ Billing fails closed for embedding global models. A model marked as embedding ca No schema migration is needed for embedding metadata. Existing model capability, config, provider mapping, API format, and pricing fields carry the data. +## Aliyun Qwen3-VL Examples + +Text request through Aether: + +```bash +curl -sS "http://localhost:8084/v1/embeddings" \ + -H "Authorization: Bearer sk-your-aether-key" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "qwen3-vl-embedding", + "input": "white running shoes", + "dimensions": 1024 + }' +``` + +Image and text fusion request: + +```bash +curl -sS "http://localhost:8084/v1/embeddings" \ + -H "Authorization: Bearer sk-your-aether-key" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "qwen3-vl-embedding", + "input": [ + { "text": "white running shoes, lightweight and breathable" }, + { "image": "https://dashscope.oss-cn-beijing.aliyuncs.com/images/256_1.png" } + ], + "parameters": { "enable_fusion": true } + }' +``` + +Video request: + +```bash +curl -sS "http://localhost:8084/v1/embeddings" \ + -H "Authorization: Bearer sk-your-aether-key" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "qwen3-vl-embedding", + "input": [ + { "video": "https://help-static-aliyun-doc.aliyuncs.com/file-manage-files/zh-CN/20250107/lbcemt/new+video.mp4" } + ], + "parameters": { "max_video_frames": 64 } + }' +``` + +Multi-image fusion request: + +```bash +curl -sS "http://localhost:8084/v1/embeddings" \ + -H "Authorization: Bearer sk-your-aether-key" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "qwen3-vl-embedding", + "input": [ + { "text": "product photos from multiple angles" }, + { "multi_images": [ + "https://example.com/front.png", + "https://example.com/side.png" + ] } + ], + "parameters": { "enable_fusion": true } + }' +``` + ## Failure Behavior The gateway validates deterministic request errors before local execution or provider transport. diff --git a/frontend/src/api/billing.ts b/frontend/src/api/billing.ts index 002c890c4..c276c9c28 100644 --- a/frontend/src/api/billing.ts +++ b/frontend/src/api/billing.ts @@ -126,6 +126,7 @@ export interface BillingCheckoutResponse { product?: BillingPlan | null } payment_instructions: Record + reused_pending_order?: boolean } export interface UserPlanEntitlement { diff --git a/frontend/src/api/endpoints/types/__tests__/api-format.spec.ts b/frontend/src/api/endpoints/types/__tests__/api-format.spec.ts index e0d6215de..8dbe52f81 100644 --- a/frontend/src/api/endpoints/types/__tests__/api-format.spec.ts +++ b/frontend/src/api/endpoints/types/__tests__/api-format.spec.ts @@ -21,6 +21,8 @@ describe('api format display helpers', () => { expect(normalizeApiFormatAlias('JINA_EMBEDDING')).toBe(API_FORMATS.JINA_EMBEDDING) expect(normalizeApiFormatAlias('JINA_RERANK')).toBe(API_FORMATS.JINA_RERANK) expect(normalizeApiFormatAlias('DOUBAO_EMBEDDING')).toBe(API_FORMATS.DOUBAO_EMBEDDING) + expect(normalizeApiFormatAlias('ALIYUN_MULTIMODAL_EMBEDDING')).toBe(API_FORMATS.ALIYUN_MULTIMODAL_EMBEDDING) + expect(normalizeApiFormatAlias('dashscope:multimodal_embedding')).toBe(API_FORMATS.ALIYUN_MULTIMODAL_EMBEDDING) }) it('formats rerank api format ids distinctly from chat formats', () => { @@ -35,10 +37,12 @@ describe('api format display helpers', () => { expect(formatApiFormat(API_FORMATS.GEMINI_EMBEDDING)).toBe('Gemini Embedding') expect(formatApiFormat(API_FORMATS.JINA_EMBEDDING)).toBe('Jina Embedding') expect(formatApiFormat(API_FORMATS.DOUBAO_EMBEDDING)).toBe('Doubao Embedding') + expect(formatApiFormat(API_FORMATS.ALIYUN_MULTIMODAL_EMBEDDING)).toBe('Aliyun Multimodal Embedding') expect(formatApiFormatShort(API_FORMATS.OPENAI_EMBEDDING)).toBe('OE') expect(formatApiFormatShort(API_FORMATS.GEMINI_EMBEDDING)).toBe('GE') expect(formatApiFormatShort(API_FORMATS.JINA_EMBEDDING)).toBe('JE') expect(formatApiFormatShort(API_FORMATS.DOUBAO_EMBEDDING)).toBe('DE') + expect(formatApiFormatShort(API_FORMATS.ALIYUN_MULTIMODAL_EMBEDDING)).toBe('AE') }) it('does not remap retired api format ids', () => { @@ -72,6 +76,7 @@ describe('api format display helpers', () => { API_FORMATS.GEMINI_EMBEDDING, API_FORMATS.JINA_EMBEDDING, API_FORMATS.JINA_RERANK, + API_FORMATS.ALIYUN_MULTIMODAL_EMBEDDING, ])).toEqual([ API_FORMATS.OPENAI, API_FORMATS.OPENAI_RESPONSES, @@ -81,6 +86,7 @@ describe('api format display helpers', () => { API_FORMATS.JINA_EMBEDDING, API_FORMATS.JINA_RERANK, API_FORMATS.DOUBAO_EMBEDDING, + API_FORMATS.ALIYUN_MULTIMODAL_EMBEDDING, 'openai:compact', ]) }) @@ -104,6 +110,7 @@ describe('api format display helpers', () => { it('groups embedding api formats by provider family', () => { expect(groupApiFormats([ API_FORMATS.DOUBAO_EMBEDDING, + API_FORMATS.ALIYUN_MULTIMODAL_EMBEDDING, API_FORMATS.JINA_RERANK, API_FORMATS.JINA_EMBEDDING, API_FORMATS.GEMINI_EMBEDDING, @@ -114,6 +121,7 @@ describe('api format display helpers', () => { { family: 'gemini', label: 'Gemini', formats: [API_FORMATS.GEMINI_EMBEDDING] }, { family: 'jina', label: 'Jina', formats: [API_FORMATS.JINA_EMBEDDING, API_FORMATS.JINA_RERANK] }, { family: 'doubao', label: 'Doubao', formats: [API_FORMATS.DOUBAO_EMBEDDING] }, + { family: 'aliyun', label: 'Aliyun', formats: [API_FORMATS.ALIYUN_MULTIMODAL_EMBEDDING] }, ]) }) diff --git a/frontend/src/api/endpoints/types/api-format.ts b/frontend/src/api/endpoints/types/api-format.ts index 3d6e55b38..e947e982b 100644 --- a/frontend/src/api/endpoints/types/api-format.ts +++ b/frontend/src/api/endpoints/types/api-format.ts @@ -18,6 +18,7 @@ export const API_FORMATS = { JINA_EMBEDDING: 'jina:embedding', JINA_RERANK: 'jina:rerank', DOUBAO_EMBEDDING: 'doubao:embedding', + ALIYUN_MULTIMODAL_EMBEDDING: 'aliyun:multimodal_embedding', } as const export type APIFormat = typeof API_FORMATS[keyof typeof API_FORMATS] @@ -39,6 +40,7 @@ export const API_FORMAT_LABELS: Record = { [API_FORMATS.JINA_EMBEDDING]: 'Jina Embedding', [API_FORMATS.JINA_RERANK]: 'Jina Rerank', [API_FORMATS.DOUBAO_EMBEDDING]: 'Doubao Embedding', + [API_FORMATS.ALIYUN_MULTIMODAL_EMBEDDING]: 'Aliyun Multimodal Embedding', CLAUDE: 'Claude Messages', CLAUDE_MESSAGES: 'Claude Messages', OPENAI: 'OpenAI Chat', @@ -56,6 +58,7 @@ export const API_FORMAT_LABELS: Record = { JINA_EMBEDDING: 'Jina Embedding', JINA_RERANK: 'Jina Rerank', DOUBAO_EMBEDDING: 'Doubao Embedding', + ALIYUN_MULTIMODAL_EMBEDDING: 'Aliyun Multimodal Embedding', } // API 格式缩写映射（用于空间紧凑的显示场景） @@ -75,6 +78,7 @@ export const API_FORMAT_SHORT: Record = { [API_FORMATS.JINA_EMBEDDING]: 'JE', [API_FORMATS.JINA_RERANK]: 'JR', [API_FORMATS.DOUBAO_EMBEDDING]: 'DE', + [API_FORMATS.ALIYUN_MULTIMODAL_EMBEDDING]: 'AE', OPENAI: 'O', OPENAI_RESPONSES: 'OR', OPENAI_RESPONSES_COMPACT: 'ORC', @@ -92,6 +96,7 @@ export const API_FORMAT_SHORT: Record = { JINA_EMBEDDING: 'JE', JINA_RERANK: 'JR', DOUBAO_EMBEDDING: 'DE', + ALIYUN_MULTIMODAL_EMBEDDING: 'AE', } // API 格式排序顺序（统一的显示顺序） @@ -111,6 +116,7 @@ export const API_FORMAT_ORDER: string[] = [ API_FORMATS.JINA_EMBEDDING, API_FORMATS.JINA_RERANK, API_FORMATS.DOUBAO_EMBEDDING, + API_FORMATS.ALIYUN_MULTIMODAL_EMBEDDING, ] // Family 显示名称映射 @@ -120,6 +126,7 @@ export const API_FORMAT_FAMILY_LABELS: Record = { gemini: 'Gemini', jina: 'Jina', doubao: 'Doubao', + aliyun: 'Aliyun', } // Kind 显示名称映射 @@ -137,7 +144,7 @@ export const API_FORMAT_KIND_LABELS: Record = { } // Family 排序顺序 -const FAMILY_ORDER = ['openai', 'claude', 'gemini', 'jina', 'doubao'] +const FAMILY_ORDER = ['openai', 'claude', 'gemini', 'jina', 'doubao', 'aliyun'] // 工具函数：从 API 格式中提取 family 和 kind export function parseApiFormat(format: string): { family: string; kind: string } { @@ -183,8 +190,20 @@ export function normalizeApiFormatAlias(format: string | null | undefined): stri return API_FORMATS.JINA_RERANK case 'DOUBAO_EMBEDDING': return API_FORMATS.DOUBAO_EMBEDDING + case 'ALIYUN_MULTIMODAL_EMBEDDING': + case 'ALIYUN_EMBEDDING': + case 'DASHSCOPE_MULTIMODAL_EMBEDDING': + case 'DASHSCOPE_EMBEDDING': + return API_FORMATS.ALIYUN_MULTIMODAL_EMBEDDING default: - return raw.toLowerCase() + switch (raw.toLowerCase()) { + case 'dashscope:multimodal_embedding': + case 'aliyun_multimodal_embedding': + case 'dashscope_multimodal_embedding': + return API_FORMATS.ALIYUN_MULTIMODAL_EMBEDDING + default: + return raw.toLowerCase() + } } } diff --git a/frontend/src/features/models/components/__tests__/global-model-form-helpers.spec.ts b/frontend/src/features/models/components/__tests__/global-model-form-helpers.spec.ts index 1e26602c7..aab3ac4ca 100644 --- a/frontend/src/features/models/components/__tests__/global-model-form-helpers.spec.ts +++ b/frontend/src/features/models/components/__tests__/global-model-form-helpers.spec.ts @@ -32,7 +32,13 @@ describe('global model form embedding payload helpers', () => { streaming: false, embedding: true, model_type: 'embedding', - api_formats: ['openai:embedding', 'gemini:embedding', 'jina:embedding', 'doubao:embedding'], + api_formats: [ + 'openai:embedding', + 'gemini:embedding', + 'jina:embedding', + 'doubao:embedding', + 'aliyun:multimodal_embedding', + ], }, }) }) diff --git a/frontend/src/features/models/components/global-model-form-helpers.ts b/frontend/src/features/models/components/global-model-form-helpers.ts index 5ea20bf8f..6f9da07aa 100644 --- a/frontend/src/features/models/components/global-model-form-helpers.ts +++ b/frontend/src/features/models/components/global-model-form-helpers.ts @@ -6,6 +6,7 @@ export const EMBEDDING_API_FORMATS = [ 'gemini:embedding', 'jina:embedding', 'doubao:embedding', + 'aliyun:multimodal_embedding', ] as const export const RERANK_API_FORMATS = [ diff --git a/frontend/src/features/providers/components/__tests__/provider-model-form-helpers.spec.ts b/frontend/src/features/providers/components/__tests__/provider-model-form-helpers.spec.ts index 6579bc561..10753db41 100644 --- a/frontend/src/features/providers/components/__tests__/provider-model-form-helpers.spec.ts +++ b/frontend/src/features/providers/components/__tests__/provider-model-form-helpers.spec.ts @@ -16,6 +16,7 @@ describe('provider model form embedding helpers', () => { { supported_capabilities: null, config: { embedding: true } }, { supported_capabilities: null, config: { model_type: 'embedding' } }, { supported_capabilities: null, config: { api_formats: ['doubao:embedding'] } }, + { supported_capabilities: null, config: { api_formats: ['aliyun:multimodal_embedding'] } }, { supports_embedding: true, effective_supports_embedding: null, config: {} }, { supports_embedding: null, effective_supports_embedding: true, config: {} }, ])('detects embedding metadata from %o', (model) => { diff --git a/frontend/src/features/providers/components/provider-model-form-helpers.ts b/frontend/src/features/providers/components/provider-model-form-helpers.ts index 3798a9319..141443976 100644 --- a/frontend/src/features/providers/components/provider-model-form-helpers.ts +++ b/frontend/src/features/providers/components/provider-model-form-helpers.ts @@ -7,6 +7,11 @@ interface EmbeddingMetadataCarrier { config?: Record | null } +function isEmbeddingApiFormat(format: unknown): boolean { + const value = String(format).trim().toLowerCase() + return value.endsWith(':embedding') || value === 'aliyun:multimodal_embedding' +} + export interface ProviderModelCreatePayloadInput { globalModelId: string providerModelName: string @@ -45,7 +50,7 @@ export function modelSupportsEmbedding(model: EmbeddingMetadataCarrier | null | return supportedCapabilities?.includes('embedding') === true || config.embedding === true || config.model_type === 'embedding' - || (Array.isArray(config.api_formats) && config.api_formats.some((format) => String(format).endsWith(':embedding'))) + || (Array.isArray(config.api_formats) && config.api_formats.some(isEmbeddingApiFormat)) } export function buildProviderModelCreatePayload(input: ProviderModelCreatePayloadInput): ModelCreate { diff --git a/frontend/src/features/providers/components/provider-tabs/__tests__/model-test-request.spec.ts b/frontend/src/features/providers/components/provider-tabs/__tests__/model-test-request.spec.ts index 19d8ac042..9622695c9 100644 --- a/frontend/src/features/providers/components/provider-tabs/__tests__/model-test-request.spec.ts +++ b/frontend/src/features/providers/components/provider-tabs/__tests__/model-test-request.spec.ts @@ -34,6 +34,21 @@ describe('buildDefaultModelTestRequestBody', () => { expect(body.stream).toBeUndefined() }) + it('uses multimodal input and qwen3-vl defaults for Aliyun embedding tests', () => { + const body = JSON.parse(buildDefaultModelTestRequestBody('qwen3-vl-embedding', 'aliyun:multimodal_embedding')) + + expect(body).toEqual({ + model: 'qwen3-vl-embedding', + input: [ + { text: 'White running shoes, lightweight and breathable.' }, + { image: 'https://dashscope.oss-cn-beijing.aliyuncs.com/images/256_1.png' }, + ], + parameters: { + enable_fusion: true, + }, + }) + }) + it.each([ 'openai:rerank', 'jina:rerank', diff --git a/frontend/src/features/providers/components/provider-tabs/model-test-request.ts b/frontend/src/features/providers/components/provider-tabs/model-test-request.ts index 002a67701..0e0b576fe 100644 --- a/frontend/src/features/providers/components/provider-tabs/model-test-request.ts +++ b/frontend/src/features/providers/components/provider-tabs/model-test-request.ts @@ -151,6 +151,19 @@ export function buildDefaultModelTestRequestBody( ): string { const normalizedApiFormat = normalizeApiFormatAlias(apiFormat ?? '') + if (normalizedApiFormat === 'aliyun:multimodal_embedding') { + return JSON.stringify({ + model: modelName, + input: [ + { text: 'White running shoes, lightweight and breathable.' }, + { image: 'https://dashscope.oss-cn-beijing.aliyuncs.com/images/256_1.png' }, + ], + parameters: { + enable_fusion: true, + }, + }, null, 2) + } + if (normalizedApiFormat.endsWith(':embedding')) { return JSON.stringify({ model: modelName, @@ -273,4 +286,4 @@ export function parseModelTestRequestHeadersDraft( emptyError: null, invalidTypeError: '测试请求头必须是 JSON 对象', }) -} \ No newline at end of file +} diff --git a/frontend/src/features/providers/composables/useEndpointStatus.ts b/frontend/src/features/providers/composables/useEndpointStatus.ts index a0cff1e50..db79a9886 100644 --- a/frontend/src/features/providers/composables/useEndpointStatus.ts +++ b/frontend/src/features/providers/composables/useEndpointStatus.ts @@ -18,6 +18,7 @@ const ENDPOINT_SORT_ORDER = [ 'jina:embedding', 'jina:rerank', 'doubao:embedding', + 'aliyun:multimodal_embedding', ] /** diff --git a/frontend/src/features/providers/composables/useProviderFilters.ts b/frontend/src/features/providers/composables/useProviderFilters.ts index 909e3b1d8..28e88e857 100644 --- a/frontend/src/features/providers/composables/useProviderFilters.ts +++ b/frontend/src/features/providers/composables/useProviderFilters.ts @@ -34,6 +34,7 @@ export function useProviderFilters( { value: 'jina:embedding', label: 'Jina Embedding' }, { value: 'jina:rerank', label: 'Jina Rerank' }, { value: 'doubao:embedding', label: 'Doubao Embedding' }, + { value: 'aliyun:multimodal_embedding', label: 'Aliyun Multimodal Embedding' }, ] const modelFilters = computed(() => { diff --git a/frontend/src/features/usage/conversation/__tests__/stream.spec.ts b/frontend/src/features/usage/conversation/__tests__/stream.spec.ts index c449c10bc..e1b579413 100644 --- a/frontend/src/features/usage/conversation/__tests__/stream.spec.ts +++ b/frontend/src/features/usage/conversation/__tests__/stream.spec.ts @@ -110,6 +110,178 @@ describe('Conversation stream compatibility', () => { }) }) + it('renders OpenAI Responses custom tool calls without text output', () => { + const requestBody = { + model: 'gpt-5.5', + stream: true, + input: 'Patch a file', + } + const toolInput = '*** Begin Patch\n*** Update File: demo.rs\n*** End Patch\n' + const rawSse = [ + 'event: response.created', + 'data: {"type":"response.created","response":{"id":"resp_custom_123","object":"response","model":"gpt-5.5","status":"in_progress"}}', + '', + 'event: response.output_item.added', + 'data: {"type":"response.output_item.added","output_index":0,"item":{"id":"ctc_123","type":"custom_tool_call","status":"in_progress","call_id":"call_123","input":"","name":"apply_patch"}}', + '', + 'event: response.custom_tool_call_input.delta', + 'data: {"type":"response.custom_tool_call_input.delta","output_index":0,"item_id":"ctc_123","delta":"*** Begin Patch\\n"}', + '', + 'event: response.custom_tool_call_input.delta', + 'data: {"type":"response.custom_tool_call_input.delta","output_index":0,"item_id":"ctc_123","delta":"*** Update File: demo.rs\\n*** End Patch\\n"}', + '', + 'event: response.custom_tool_call_input.done', + `data: ${JSON.stringify({ type: 'response.custom_tool_call_input.done', output_index: 0, item_id: 'ctc_123', input: toolInput })}`, + '', + 'event: response.output_item.done', + `data: ${JSON.stringify({ type: 'response.output_item.done', output_index: 0, item: { id: 'ctc_123', type: 'custom_tool_call', status: 'completed', call_id: 'call_123', input: toolInput, name: 'apply_patch' } })}`, + '', + 'event: response.completed', + 'data: {"type":"response.completed","response":{"id":"resp_custom_123","object":"response","model":"gpt-5.5","status":"completed","output":[]}}', + '', + 'data: [DONE]', + '', + ].join('\n') + + const parsed = parseResponse(rawSse, requestBody, 'openai:responses') + expect(parsed.messages[0]?.content[0]).toMatchObject({ + type: 'tool_use', + toolName: 'apply_patch', + toolId: 'call_123', + input: toolInput, + }) + + const rendered = renderResponse(rawSse, requestBody, 'openai:responses') + expect(rendered.error).toBeUndefined() + expect(rendered.isStream).toBe(true) + expect(rendered.blocks).toHaveLength(1) + + const firstBlock = rendered.blocks[0] + if (!firstBlock || firstBlock.type !== 'message') { + throw new Error('expected first render block to be message') + } + + expect(firstBlock.content[0]).toMatchObject({ + type: 'tool_use', + toolName: 'apply_patch', + toolId: 'call_123', + input: toolInput, + }) + }) + + it('keeps OpenAI Responses custom tool calls when text output is present', () => { + const requestBody = { + model: 'gpt-5.5', + stream: true, + input: 'Explain and patch', + } + const rawSse = [ + 'event: response.output_text.delta', + 'data: {"type":"response.output_text.delta","delta":"I will patch it."}', + '', + 'event: response.output_item.added', + 'data: {"type":"response.output_item.added","output_index":1,"item":{"id":"ctc_456","type":"custom_tool_call","status":"in_progress","call_id":"call_456","input":"","name":"apply_patch"}}', + '', + 'event: response.custom_tool_call_input.delta', + 'data: {"type":"response.custom_tool_call_input.delta","output_index":1,"item_id":"ctc_456","delta":"patch text"}', + '', + 'event: response.output_item.done', + 'data: {"type":"response.output_item.done","output_index":1,"item":{"id":"ctc_456","type":"custom_tool_call","status":"completed","call_id":"call_456","input":"patch text","name":"apply_patch"}}', + '', + ].join('\n') + + const rendered = renderResponse(rawSse, requestBody, 'openai:responses') + const firstBlock = rendered.blocks[0] + if (!firstBlock || firstBlock.type !== 'message') { + throw new Error('expected first render block to be message') + } + + expect(firstBlock.content.map(block => block.type)).toEqual(['text', 'tool_use']) + expect(firstBlock.content[1]).toMatchObject({ + type: 'tool_use', + toolName: 'apply_patch', + input: 'patch text', + }) + }) + + it('renders future OpenAI Responses call items through the generic call fallback', () => { + const requestBody = { + model: 'gpt-5.5', + stream: true, + input: 'Run a command', + } + const action = { command: 'npm test', timeout_ms: 1000 } + const expectedInput = JSON.stringify(action, null, 2) + const rawSse = [ + 'event: response.output_item.added', + 'data: {"type":"response.output_item.added","output_index":0,"item":{"id":"shell_123","type":"shell_call","status":"in_progress"}}', + '', + 'event: response.output_item.done', + `data: ${JSON.stringify({ type: 'response.output_item.done', output_index: 0, item: { id: 'shell_123', type: 'shell_call', status: 'completed', action } })}`, + '', + ].join('\n') + + const parsed = parseResponse(rawSse, requestBody, 'openai:responses') + expect(parsed.messages[0]?.content[0]).toMatchObject({ + type: 'tool_use', + toolName: 'shell_call', + toolId: 'shell_123', + input: expectedInput, + }) + + const rendered = renderResponse(rawSse, requestBody, 'openai:responses') + const firstBlock = rendered.blocks[0] + if (!firstBlock || firstBlock.type !== 'message') { + throw new Error('expected first render block to be message') + } + + expect(firstBlock.content[0]).toMatchObject({ + type: 'tool_use', + toolName: 'shell_call', + toolId: 'shell_123', + input: expectedInput, + }) + }) + + it('keeps streamed function_call arguments when response.completed omits them', () => { + const requestBody = { + model: 'gpt-5.5', + stream: true, + input: 'What is the weather?', + } + const rawSse = [ + 'event: response.created', + `data: ${JSON.stringify({ type: 'response.created', response: { id: 'resp_fc_1', object: 'response', model: 'gpt-5.5', status: 'in_progress' } })}`, + '', + 'event: response.output_item.added', + `data: ${JSON.stringify({ type: 'response.output_item.added', output_index: 0, item: { id: 'fc_1', type: 'function_call', status: 'in_progress', call_id: 'call_1', name: 'get_weather', arguments: '' } })}`, + '', + 'event: response.function_call_arguments.delta', + `data: ${JSON.stringify({ type: 'response.function_call_arguments.delta', output_index: 0, item_id: 'fc_1', delta: '{"city":' })}`, + '', + 'event: response.function_call_arguments.delta', + `data: ${JSON.stringify({ type: 'response.function_call_arguments.delta', output_index: 0, item_id: 'fc_1', delta: '"SF"}' })}`, + '', + // 最终项故意不带 arguments：解析器不应用 '{}' 冲掉已收集的增量参数 + 'event: response.completed', + `data: ${JSON.stringify({ type: 'response.completed', response: { id: 'resp_fc_1', object: 'response', model: 'gpt-5.5', status: 'completed', output: [{ id: 'fc_1', type: 'function_call', status: 'completed', call_id: 'call_1', name: 'get_weather' }] } })}`, + '', + 'data: [DONE]', + '', + ].join('\n') + + const parsed = parseResponse(rawSse, requestBody, 'openai:responses') + // 命中同一 key，不重复渲染 + expect(parsed.messages).toHaveLength(1) + expect(parsed.messages[0]?.content).toHaveLength(1) + expect(parsed.messages[0]?.content[0]).toMatchObject({ + type: 'tool_use', + toolName: 'get_weather', + toolId: 'call_1', + input: '{"city":"SF"}', + }) + }) + it('renders HTML-entity encoded OpenAI tool arguments as formatted JSON', () => { const requestBody = { model: 'gpt-5.4', diff --git a/frontend/src/features/usage/conversation/openai.ts b/frontend/src/features/usage/conversation/openai.ts index a13139579..c575f489e 100644 --- a/frontend/src/features/usage/conversation/openai.ts +++ b/frontend/src/features/usage/conversation/openai.ts @@ -313,11 +313,11 @@ export class OpenAIParser implements ApiFormatParser { return createMessage(role, contentBlocks) } - // function_call -> 工具调用 - if (itemType === 'function_call') { - const toolId = String(item.call_id || item.id || '') - const toolName = String(item.name || '') - const args = String(item.arguments || '{}') + // Responses API call item -> 工具调用 + if (this.isResponsesCallItemType(itemType)) { + const toolId = this.responsesCallId(item) + const toolName = this.responsesCallName(item) + const args = this.responsesCallInput(item) return createMessage('assistant', [createToolUseBlock(toolId, toolName, args)]) } @@ -439,6 +439,14 @@ export class OpenAIParser implements ApiFormatParser { if (contentBlocks.length > 0) { result.messages.push(createMessage('assistant', contentBlocks)) } + } else if (item && this.isResponsesCallItemType(item.type)) { + result.messages.push(createMessage('assistant', [ + createToolUseBlock( + this.responsesCallId(item), + this.responsesCallName(item), + this.responsesCallInput(item) + ), + ])) } } @@ -566,8 +574,39 @@ export class OpenAIParser implements ApiFormatParser { const textParts: string[] = [] const toolCalls = new Map() - let currentToolId = '' - let currentToolName = '' + const outputIndexToToolKey = new Map() + let currentToolKey = '' + + const ensureToolCall = ( + key: string, + id: string, + name: string, + initialInput?: string + ) => { + if (!key) return + const existing = toolCalls.get(key) + if (existing) { + if (id) existing.id = id + if (name) existing.name = name + if (initialInput) existing.args = [initialInput] + return + } + toolCalls.set(key, { + name, + id, + args: initialInput ? [initialInput] : [], + }) + } + + const resolveToolKey = (chunk: RawObject): string => { + const itemId = typeof chunk.item_id === 'string' ? chunk.item_id : '' + if (itemId) return itemId + const outputIndex = typeof chunk.output_index === 'number' ? chunk.output_index : null + if (outputIndex != null) { + return outputIndexToToolKey.get(outputIndex) || currentToolKey + } + return currentToolKey + } for (const rawChunk of chunks) { const chunk = rawChunk as RawObject @@ -596,28 +635,55 @@ export class OpenAIParser implements ApiFormatParser { continue } - // 处理函数调用输出项添加: response.output_item.added - if (eventType === 'response.output_item.added') { + // 处理 Responses call 输出项添加/完成: response.output_item.added / done + if (eventType === 'response.output_item.added' || eventType === 'response.output_item.done') { const item = chunk.item as RawObject | undefined - if (item?.type === 'function_call') { - currentToolId = String(item.call_id || item.id || '') - currentToolName = String(item.name || '') - if (currentToolId && !toolCalls.has(currentToolId)) { - toolCalls.set(currentToolId, { - name: currentToolName, - id: currentToolId, - args: [], - }) + if (item && this.isResponsesCallItemType(item.type)) { + const itemId = typeof item.id === 'string' ? item.id : '' + const toolId = this.responsesCallId(item) + const key = itemId || toolId || String(chunk.output_index ?? '') + const input = eventType === 'response.output_item.done' && this.responsesCallHasInput(item) + ? this.responsesCallInput(item) + : '' + ensureToolCall(key, toolId, this.responsesCallName(item), input) + currentToolKey = key + if (typeof chunk.output_index === 'number') { + outputIndexToToolKey.set(chunk.output_index, key) } } continue } - // 处理函数调用参数增量: response.function_call_arguments.delta - if (eventType === 'response.function_call_arguments.delta') { + // 处理已知 call 输入增量 + if ( + eventType === 'response.function_call_arguments.delta' || + eventType === 'response.custom_tool_call_input.delta' + ) { const delta = chunk.delta - if (typeof delta === 'string' && currentToolId && toolCalls.has(currentToolId)) { - toolCalls.get(currentToolId)?.args.push(delta) + const key = resolveToolKey(chunk) + if (typeof delta === 'string' && key && toolCalls.has(key)) { + toolCalls.get(key)?.args.push(delta) + } + continue + } + + if (eventType === 'response.function_call_arguments.done') { + const key = resolveToolKey(chunk) + const args = typeof chunk.arguments === 'string' + ? chunk.arguments + : typeof chunk.delta === 'string' + ? chunk.delta + : null + if (key && toolCalls.has(key) && args != null) { + toolCalls.get(key)!.args = [args] + } + continue + } + + if (eventType === 'response.custom_tool_call_input.done') { + const key = resolveToolKey(chunk) + if (key && toolCalls.has(key) && typeof chunk.input === 'string') { + toolCalls.get(key)!.args = [chunk.input] } continue } @@ -630,17 +696,29 @@ export class OpenAIParser implements ApiFormatParser { result.model = response.model } - // 从 output 中提取文本（备用方案） - if (textParts.length === 0 && Array.isArray(response?.output)) { - for (const rawItem of response.output as unknown[]) { - const item = rawItem as RawObject - if (item?.type === 'message' && Array.isArray(item?.content)) { + // 从 output 中提取文本和工具调用（备用方案） + if (Array.isArray(response?.output)) { + const output = response.output as unknown[] + for (let index = 0; index < output.length; index++) { + const item = output[index] as RawObject + if (textParts.length === 0 && item?.type === 'message' && Array.isArray(item?.content)) { for (const rawContent of item.content as unknown[]) { const content = rawContent as RawObject if (content?.type === 'output_text' && typeof content?.text === 'string') { textParts.push(content.text) } } + } else if (this.isResponsesCallItemType(item.type)) { + const itemId = typeof item.id === 'string' ? item.id : '' + const toolId = this.responsesCallId(item) + // 与流式阶段使用同一套 key 命中同一条工具调用，避免重复渲染 + const key = itemId || toolId || outputIndexToToolKey.get(index) || String(index) + // 仅在最终项确实带有输入时才覆盖，避免用 '{}' 等默认值 + // 冲掉已通过增量事件收集到的参数 + const input = this.responsesCallHasInput(item) + ? this.responsesCallInput(item) + : '' + ensureToolCall(key, toolId, this.responsesCallName(item), input) } } } @@ -731,6 +809,49 @@ export class OpenAIParser implements ApiFormatParser { return createMessage(role, contentBlocks) } + private isResponsesCallItemType(itemType: unknown): boolean { + return typeof itemType === 'string' && itemType.endsWith('_call') + } + + private responsesCallId(item: RawObject): string { + return String(item.call_id || item.id || '') + } + + private responsesCallName(item: RawObject): string { + const name = typeof item.name === 'string' ? item.name.trim() : '' + if (name) return name + return typeof item.type === 'string' ? item.type : 'tool_call' + } + + private responsesCallInputCandidate(item: RawObject): unknown { + if (item.type === 'function_call') return item.arguments + if (item.type === 'custom_tool_call') return item.input + for (const key of ['input', 'arguments', 'action', 'query', 'code', 'prompt']) { + if (item[key] != null) return item[key] + } + return undefined + } + + private responsesCallInput(item: RawObject): string { + const input = this.responsesCallInputCandidate(item) + if (typeof input === 'string') return input + if (input == null) { + if (item.type === 'function_call') return '{}' + if (item.type === 'custom_tool_call') return '' + return JSON.stringify(item, null, 2) + } + return JSON.stringify(input, null, 2) + } + + private responsesCallHasInput(item: RawObject): boolean { + const input = this.responsesCallInputCandidate(item) + if (input == null) { + return item.type !== 'function_call' && item.type !== 'custom_tool_call' + } + if (typeof input === 'string') return input.length > 0 + return true + } + /** * 映射角色 */ @@ -887,12 +1008,12 @@ export class OpenAIParser implements ApiFormatParser { return createMessageBlock(role, contentBlocks, { roleLabel: this.getRoleLabel(role) }) } - // function_call -> 工具调用 - if (itemType === 'function_call') { - const toolName = String(item.name || '工具调用') - const args = this.formatJson(item.arguments) + // Responses API call item -> 工具调用 + if (this.isResponsesCallItemType(itemType)) { + const toolName = this.responsesCallName(item) + const args = this.formatJson(this.responsesCallInput(item)) return createMessageBlock('assistant', [ - createToolUseRenderBlock(toolName, args, String(item.call_id || item.id || '')), + createToolUseRenderBlock(toolName, args, this.responsesCallId(item)), ], { roleLabel: 'Assistant', badges: [createBadgeBlock('工具调用', 'outline')] }) } @@ -1015,6 +1136,17 @@ export class OpenAIParser implements ApiFormatParser { roleLabel: 'Assistant', })) } + } else if (this.isResponsesCallItemType(item.type)) { + blocks.push(createMessageBlock('assistant', [ + createToolUseRenderBlock( + this.responsesCallName(item), + this.formatJson(this.responsesCallInput(item)), + this.responsesCallId(item) + ), + ], { + roleLabel: 'Assistant', + badges: [createBadgeBlock('工具调用', 'outline')], + })) } } diff --git a/frontend/src/mocks/__tests__/embedding-metadata.spec.ts b/frontend/src/mocks/__tests__/embedding-metadata.spec.ts index a4a941c22..803b57a7d 100644 --- a/frontend/src/mocks/__tests__/embedding-metadata.spec.ts +++ b/frontend/src/mocks/__tests__/embedding-metadata.spec.ts @@ -20,7 +20,7 @@ describe('embedding mock metadata', () => { it('includes all embedding API formats as distinct catalog formats', () => { const embeddingFormats = MOCK_API_FORMATS.formats - .filter(format => format.value.endsWith(':embedding')) + .filter(format => format.value.endsWith(':embedding') || format.value.endsWith('_embedding')) .map(format => [format.value, format.label]) expect(embeddingFormats).toEqual([ @@ -28,6 +28,7 @@ describe('embedding mock metadata', () => { ['gemini:embedding', 'Gemini Embedding'], ['jina:embedding', 'Jina Embedding'], ['doubao:embedding', 'Doubao Embedding'], + ['aliyun:multimodal_embedding', 'Aliyun Multimodal Embedding'], ]) }) diff --git a/frontend/src/mocks/data.ts b/frontend/src/mocks/data.ts index a50e44adb..c9581ae14 100644 --- a/frontend/src/mocks/data.ts +++ b/frontend/src/mocks/data.ts @@ -1196,6 +1196,12 @@ export const MOCK_API_FORMATS = { { value: 'gemini:video', label: 'Gemini Video', default_path: '/v1beta/models/{model}:predictLongRunning', aliases: [] }, { value: 'jina:embedding', label: 'Jina Embedding', default_path: '/v1/embeddings', aliases: [] }, { value: 'jina:rerank', label: 'Jina Rerank', default_path: '/v1/rerank', aliases: [] }, - { value: 'doubao:embedding', label: 'Doubao Embedding', default_path: '/embeddings/multimodal', aliases: [] } + { value: 'doubao:embedding', label: 'Doubao Embedding', default_path: '/embeddings/multimodal', aliases: [] }, + { + value: 'aliyun:multimodal_embedding', + label: 'Aliyun Multimodal Embedding', + default_path: '/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding', + aliases: ['dashscope:multimodal_embedding'], + }, ] } diff --git a/frontend/src/mocks/handler.ts b/frontend/src/mocks/handler.ts index 3fdc53249..7437b295d 100644 --- a/frontend/src/mocks/handler.ts +++ b/frontend/src/mocks/handler.ts @@ -2084,7 +2084,7 @@ function generateMockModelsForProvider(providerId: string) { const hasClaude = provider.api_formats.some(f => f.includes('claude')) const hasOpenAI = provider.api_formats.some(f => f.includes('openai')) const hasGemini = provider.api_formats.some(f => f.includes('gemini')) - const hasEmbedding = provider.api_formats.some(f => f.endsWith(':embedding')) + const hasEmbedding = provider.api_formats.some(f => f.endsWith(':embedding') || f === 'aliyun:multimodal_embedding') const hasRerank = provider.api_formats.some(f => f.endsWith(':rerank')) const models: Record[] = [] diff --git a/frontend/src/views/admin/SystemSettings.vue b/frontend/src/views/admin/SystemSettings.vue index 035668779..8bf172b2e 100644 --- a/frontend/src/views/admin/SystemSettings.vue +++ b/frontend/src/views/admin/SystemSettings.vue @@ -82,6 +82,7 @@ :auto-delete-expired-keys="systemConfig.auto_delete_expired_keys" :enable-format-conversion="systemConfig.enable_format_conversion" :enable-openai-image-sync-heartbeat="systemConfig.enable_openai_image_sync_heartbeat" + :enable-standard-text-sync-heartbeat="systemConfig.enable_standard_text_sync_heartbeat" :loading="systemConfigLoading || basicConfigLoading" :has-changes="hasBasicConfigChanges" @save="saveBasicConfig" @@ -106,6 +107,7 @@ @update:auto-delete-expired-keys="systemConfig.auto_delete_expired_keys = $event" @update:enable-format-conversion="systemConfig.enable_format_conversion = $event" @update:enable-openai-image-sync-heartbeat="systemConfig.enable_openai_image_sync_heartbeat = $event" + @update:enable-standard-text-sync-heartbeat="systemConfig.enable_standard_text_sync_heartbeat = $event" /> diff --git a/frontend/src/views/admin/system-settings/BasicConfigSection.vue b/frontend/src/views/admin/system-settings/BasicConfigSection.vue index d19330b04..90d3d5859 100644 --- a/frontend/src/views/admin/system-settings/BasicConfigSection.vue +++ b/frontend/src/views/admin/system-settings/BasicConfigSection.vue @@ -172,6 +172,27 @@ +
+
+ +
+ +

+ 开启后标准文本非流式接口外层 HTTP 状态固定为 200，上游失败需读取响应体 error.upstream_status +

+
+
+
+
@@ -498,6 +519,7 @@ defineProps<{ autoDeleteExpiredKeys: boolean enableFormatConversion: boolean enableOpenaiImageSyncHeartbeat: boolean + enableStandardTextSyncHeartbeat: boolean loading: boolean hasChanges: boolean }>() @@ -525,5 +547,6 @@ defineEmits<{ 'update:autoDeleteExpiredKeys': [value: boolean] 'update:enableFormatConversion': [value: boolean] 'update:enableOpenaiImageSyncHeartbeat': [value: boolean] + 'update:enableStandardTextSyncHeartbeat': [value: boolean] }>() diff --git a/frontend/src/views/admin/system-settings/__tests__/useSystemConfig.spec.ts b/frontend/src/views/admin/system-settings/__tests__/useSystemConfig.spec.ts index 798232bc5..f4cc53b38 100644 --- a/frontend/src/views/admin/system-settings/__tests__/useSystemConfig.spec.ts +++ b/frontend/src/views/admin/system-settings/__tests__/useSystemConfig.spec.ts @@ -1,13 +1,14 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' -const { getSystemConfigMock } = vi.hoisted(() => ({ +const { getSystemConfigMock, updateSystemConfigMock } = vi.hoisted(() => ({ getSystemConfigMock: vi.fn(), + updateSystemConfigMock: vi.fn(), })) vi.mock('@/api/admin', () => ({ adminApi: { getSystemConfig: getSystemConfigMock, - updateSystemConfig: vi.fn(), + updateSystemConfig: updateSystemConfigMock, getSystemVersion: vi.fn(), }, })) @@ -40,6 +41,7 @@ interface DeferredConfigResponse { describe('useSystemConfig', () => { beforeEach(() => { getSystemConfigMock.mockReset() + updateSystemConfigMock.mockReset() }) it('loads config keys in parallel and keeps change detection disabled until the baseline is ready', async () => { @@ -53,6 +55,7 @@ describe('useSystemConfig', () => { expect(getSystemConfigMock.mock.calls.map(([key]) => key)).toContain('request_record_level') expect(getSystemConfigMock.mock.calls.map(([key]) => key)).toContain('proxy_node_metrics_cleanup_batch_size') + expect(getSystemConfigMock.mock.calls.map(([key]) => key)).toContain('enable_standard_text_sync_heartbeat') state.systemConfig.value.request_record_level = 'headers' expect(state.systemConfigLoading.value).toBe(true) @@ -74,4 +77,29 @@ describe('useSystemConfig', () => { state.systemConfig.value.request_record_level = 'full' expect(state.hasLogConfigChanges.value).toBe(true) }) + + it('loads and saves the standard text sync heartbeat flag as a basic config item', async () => { + getSystemConfigMock.mockImplementation(async (key: string) => ({ + key, + value: key === 'enable_standard_text_sync_heartbeat' ? false : undefined, + is_set: key === 'enable_standard_text_sync_heartbeat', + })) + updateSystemConfigMock.mockResolvedValue({}) + + const state = useSystemConfig() + await state.loadSystemConfig() + + expect(state.systemConfig.value.enable_standard_text_sync_heartbeat).toBe(false) + state.systemConfig.value.enable_standard_text_sync_heartbeat = true + expect(state.hasBasicConfigChanges.value).toBe(true) + + await state.saveBasicConfig() + + expect(updateSystemConfigMock).toHaveBeenCalledWith( + 'enable_standard_text_sync_heartbeat', + true, + '标准文本非流式心跳开关：开启后外层 HTTP 状态固定为 200，上游失败写入响应体' + ) + expect(state.hasBasicConfigChanges.value).toBe(false) + }) }) diff --git a/frontend/src/views/admin/system-settings/composables/useSystemConfig.ts b/frontend/src/views/admin/system-settings/composables/useSystemConfig.ts index 2d5c89c49..c0b50e3c8 100644 --- a/frontend/src/views/admin/system-settings/composables/useSystemConfig.ts +++ b/frontend/src/views/admin/system-settings/composables/useSystemConfig.ts @@ -35,6 +35,8 @@ export interface SystemConfig { enable_format_conversion: boolean // 同步生图心跳 enable_openai_image_sync_heartbeat: boolean + // 标准文本非流式心跳 + enable_standard_text_sync_heartbeat: boolean // 请求记录 request_record_level: string max_request_body_size: number @@ -89,6 +91,8 @@ const CONFIG_KEYS = [ 'enable_format_conversion', // 同步生图心跳 'enable_openai_image_sync_heartbeat', + // 标准文本非流式心跳 + 'enable_standard_text_sync_heartbeat', // 请求记录 'request_record_level', 'max_request_body_size', @@ -145,6 +149,8 @@ function createDefaultConfig(): SystemConfig { enable_format_conversion: false, // 同步生图心跳 enable_openai_image_sync_heartbeat: false, + // 标准文本非流式心跳 + enable_standard_text_sync_heartbeat: false, // 请求记录 request_record_level: 'basic', max_request_body_size: 1048576, @@ -230,7 +236,10 @@ export function useSystemConfig() { originalConfig.value.registration_privacy_policy_version || systemConfig.value.auto_delete_expired_keys !== originalConfig.value.auto_delete_expired_keys || systemConfig.value.enable_format_conversion !== originalConfig.value.enable_format_conversion || - systemConfig.value.enable_openai_image_sync_heartbeat !== originalConfig.value.enable_openai_image_sync_heartbeat + systemConfig.value.enable_openai_image_sync_heartbeat !== + originalConfig.value.enable_openai_image_sync_heartbeat || + systemConfig.value.enable_standard_text_sync_heartbeat !== + originalConfig.value.enable_standard_text_sync_heartbeat ) }) @@ -513,6 +522,11 @@ export function useSystemConfig() { value: systemConfig.value.enable_openai_image_sync_heartbeat, description: '同步生图心跳开关：开启后外层 HTTP 状态固定为 200，上游失败写入响应体', }, + { + key: 'enable_standard_text_sync_heartbeat', + value: systemConfig.value.enable_standard_text_sync_heartbeat, + description: '标准文本非流式心跳开关：开启后外层 HTTP 状态固定为 200，上游失败写入响应体', + }, ] const turnstileSecret = systemConfig.value.turnstile_secret_key.trim() if (turnstileSecret) { @@ -565,6 +579,8 @@ export function useSystemConfig() { systemConfig.value.enable_format_conversion originalConfig.value.enable_openai_image_sync_heartbeat = systemConfig.value.enable_openai_image_sync_heartbeat + originalConfig.value.enable_standard_text_sync_heartbeat = + systemConfig.value.enable_standard_text_sync_heartbeat } success('基础配置已保存') } catch (err) { diff --git a/frontend/src/views/user/BillingPlans.vue b/frontend/src/views/user/BillingPlans.vue index c778bf314..72410e076 100644 --- a/frontend/src/views/user/BillingPlans.vue +++ b/frontend/src/views/user/BillingPlans.vue @@ -351,7 +351,11 @@ async function checkoutPlan(plan: BillingPlan) { payment_channel: option.payment_channel, }) latestCheckout.value = response - success('套餐订单已创建') + success( + response.reused_pending_order + ? '已有待支付订单，已打开原支付链接' + : '套餐订单已创建' + ) submitPaymentInstructions(response.payment_instructions) } catch (err) { log.error('创建套餐订单失败:', err) diff --git a/frontend/src/views/user/__tests__/model-catalog-helpers.spec.ts b/frontend/src/views/user/__tests__/model-catalog-helpers.spec.ts index e65c2cb41..ecef4b031 100644 --- a/frontend/src/views/user/__tests__/model-catalog-helpers.spec.ts +++ b/frontend/src/views/user/__tests__/model-catalog-helpers.spec.ts @@ -35,6 +35,7 @@ describe('model catalog embedding helpers', () => { expect(supportsEmbedding(model({ config: { embedding: true } }))).toBe(true) expect(supportsEmbedding(model({ config: { model_type: 'embedding' } }))).toBe(true) expect(supportsEmbedding(model({ config: { api_formats: ['jina:embedding'] } }))).toBe(true) + expect(supportsEmbedding(model({ config: { api_formats: ['aliyun:multimodal_embedding'] } }))).toBe(true) expect(supportsEmbedding(model({ config: { api_formats: ['openai:chat'] } }))).toBe(false) }) diff --git a/frontend/src/views/user/components/UserModelDetailDrawer.vue b/frontend/src/views/user/components/UserModelDetailDrawer.vue index f046ac98d..afbc91e3a 100644 --- a/frontend/src/views/user/components/UserModelDetailDrawer.vue +++ b/frontend/src/views/user/components/UserModelDetailDrawer.vue @@ -395,7 +395,10 @@ function supportsEmbedding(model: PublicGlobalModel): boolean { || model.supported_capabilities?.includes('embedding') === true || model.config?.embedding === true || model.config?.model_type === 'embedding' - || (Array.isArray(model.config?.api_formats) && model.config.api_formats.some((format) => String(format).endsWith(':embedding'))) + || (Array.isArray(model.config?.api_formats) && model.config.api_formats.some((format) => { + const value = String(format).trim().toLowerCase() + return value.endsWith(':embedding') || value === 'aliyun:multimodal_embedding' + })) } // 添加 ESC 键监听 diff --git a/frontend/src/views/user/model-catalog-helpers.ts b/frontend/src/views/user/model-catalog-helpers.ts index fd0b57cc0..d766e7db1 100644 --- a/frontend/src/views/user/model-catalog-helpers.ts +++ b/frontend/src/views/user/model-catalog-helpers.ts @@ -1,11 +1,16 @@ import type { PublicGlobalModel } from '@/api/public-models' +function isEmbeddingApiFormat(format: unknown): boolean { + const value = String(format).trim().toLowerCase() + return value.endsWith(':embedding') || value === 'aliyun:multimodal_embedding' +} + export function supportsEmbedding(model: PublicGlobalModel): boolean { return model.supports_embedding === true || model.supported_capabilities?.includes('embedding') === true || model.config?.embedding === true || model.config?.model_type === 'embedding' - || (Array.isArray(model.config?.api_formats) && model.config.api_formats.some((format) => String(format).endsWith(':embedding'))) + || (Array.isArray(model.config?.api_formats) && model.config.api_formats.some(isEmbeddingApiFormat)) } export function supportsRerank(model: PublicGlobalModel): boolean { diff --git a/frontend/vite.config.ts b/frontend/vite.config.ts index 8e4155b22..70919f04a 100644 --- a/frontend/vite.config.ts +++ b/frontend/vite.config.ts @@ -5,17 +5,24 @@ import { execSync } from 'child_process' function normalizeVersion(version: string): string { const trimmed = version.trim() + if (!trimmed || trimmed.startsWith('tunnel-v')) { + return '' + } return trimmed.startsWith('v') || trimmed.startsWith('V') ? trimmed.slice(1) : trimmed } function getGitVersion(): string { const envVersion = process.env.AETHER_BUILD_VERSION || process.env.AETHER_VERSION if (envVersion?.trim()) { - return normalizeVersion(envVersion) + const version = normalizeVersion(envVersion) + if (version) { + return version + } } try { - return normalizeVersion(execSync('git describe --tags --always --dirty').toString()) + const version = normalizeVersion(execSync('git describe --tags --match "v[0-9]*" --always --dirty').toString()) + return version || '0.0.0.dev0' } catch { return '0.0.0.dev0' }