From 0f6d4b9146d25cdf96b0c811bdaab07b624c98a2 Mon Sep 17 00:00:00 2001 From: "Entropy.Xu" Date: Thu, 28 May 2026 16:05:36 +0800 Subject: [PATCH 01/29] =?UTF-8?q?feat(embedding):=20=E6=8E=A5=E5=85=A5?= =?UTF-8?q?=E9=98=BF=E9=87=8C=E4=BA=91=E5=A4=9A=E6=A8=A1=E6=80=81=E5=90=91?= =?UTF-8?q?=E9=87=8F=E7=AB=AF=E7=82=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../passthrough/provider/family/request.rs | 1 + apps/aether-gateway/src/api/ai/aliyun.rs | 15 + apps/aether-gateway/src/api/ai/mod.rs | 1 + apps/aether-gateway/src/api/ai/registry.rs | 10 +- .../handlers/admin/model/global/providers.rs | 1 + .../src/handlers/admin/model/payloads.rs | 1 + .../monitoring/cache_affinity.rs | 3 +- .../admin/provider/query/models/model_test.rs | 42 +- .../query/models/model_test/adapter.rs | 3 + .../query/models/model_test/summary.rs | 6 + .../provider/query/models/model_test/tests.rs | 27 ++ .../src/handlers/public/ai_public.rs | 33 ++ .../handlers/public/support/models/shared.rs | 10 +- .../src/tests/control/admin/system.rs | 8 + .../src/tests/control/proxy/embeddings.rs | 428 ++++++++++++++++++ crates/aether-admin/src/provider/models.rs | 1 + crates/aether-admin/src/system.rs | 12 + .../src/formats/aliyun/embedding/mod.rs | 2 + .../src/formats/aliyun/embedding/request.rs | 256 +++++++++++ .../src/formats/aliyun/embedding/response.rs | 134 ++++++ .../src/formats/aliyun/mod.rs | 1 + .../src/formats/gemini/embedding/request.rs | 1 + crates/aether-ai-formats/src/formats/id.rs | 43 +- .../aether-ai-formats/src/formats/matrix.rs | 31 +- crates/aether-ai-formats/src/formats/mod.rs | 1 + .../src/formats/openai/embedding/request.rs | 11 + .../aether-ai-formats/src/formats/registry.rs | 125 ++++- .../shared/stream_core/format_matrix.rs | 9 +- .../src/formats/shared/sync_products.rs | 1 + crates/aether-ai-formats/src/lib.rs | 10 +- .../src/protocol/canonical.rs | 106 ++++- .../src/repository/global_models/types.rs | 1 + .../src/repository/global_models/mod.rs | 1 + .../src/repository/global_models/postgres.rs | 2 + .../src/conversion.rs | 7 +- .../aether-provider-transport/src/policy.rs | 5 +- .../src/provider_types.rs | 12 + .../src/request_url/mod.rs | 36 ++ .../src/same_format_provider/mod.rs | 60 +++ docs/api/embeddings.md | 84 +++- .../types/__tests__/api-format.spec.ts | 8 + .../src/api/endpoints/types/api-format.ts | 23 +- .../global-model-form-helpers.spec.ts | 8 +- .../components/global-model-form-helpers.ts | 1 + .../provider-model-form-helpers.spec.ts | 1 + .../components/provider-model-form-helpers.ts | 7 +- .../__tests__/model-test-request.spec.ts | 15 + .../provider-tabs/model-test-request.ts | 15 +- .../composables/useEndpointStatus.ts | 1 + .../composables/useProviderFilters.ts | 1 + .../__tests__/embedding-metadata.spec.ts | 3 +- frontend/src/mocks/data.ts | 8 +- frontend/src/mocks/handler.ts | 2 +- .../__tests__/model-catalog-helpers.spec.ts | 1 + .../user/components/UserModelDetailDrawer.vue | 5 +- .../src/views/user/model-catalog-helpers.ts | 7 +- 56 files changed, 1605 insertions(+), 42 deletions(-) create mode 100644 apps/aether-gateway/src/api/ai/aliyun.rs create mode 100644 crates/aether-ai-formats/src/formats/aliyun/embedding/mod.rs create mode 100644 crates/aether-ai-formats/src/formats/aliyun/embedding/request.rs create mode 100644 crates/aether-ai-formats/src/formats/aliyun/embedding/response.rs create mode 100644 crates/aether-ai-formats/src/formats/aliyun/mod.rs 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 6748727d9..4d870a39e 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 @@ -54,6 +54,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/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 c186dc7f9..62491f72c 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)); } @@ -2914,8 +2945,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, @@ -2982,6 +3018,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, @@ -2993,6 +3030,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 9224b2c98..eb82f84a9 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 @@ -464,6 +464,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!({ @@ -650,6 +670,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/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/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/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 349ff0061..106ba1ef5 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/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/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/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 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/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..83780fdf9 100644 --- a/crates/aether-ai-formats/src/protocol/canonical.rs +++ b/crates/aether-ai-formats/src/protocol/canonical.rs @@ -242,6 +242,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 +266,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 +276,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 +328,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 +552,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 +750,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, } } @@ -5050,8 +5106,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 +5163,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 +5227,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 +5299,7 @@ mod tests { dimensions: Some(2), task: None, user: None, + parameters: None, extensions: Default::default(), }), ..Default::default() @@ -5247,6 +5345,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 +5398,7 @@ mod tests { dimensions: None, task: None, user: None, + parameters: None, extensions: Default::default(), }), ..Default::default() 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-provider-transport/src/conversion.rs b/crates/aether-provider-transport/src/conversion.rs index d8beb9ed7..6f1f79a6d 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 a6b586e76..c7d676f45 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) @@ -348,10 +348,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 82cd7ed73..4fa14b187 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", } } } @@ -187,6 +190,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, @@ -465,6 +472,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, } } @@ -846,6 +854,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), @@ -859,6 +869,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 0d3c33b97..6bd306785 100644 --- a/crates/aether-provider-transport/src/request_url/mod.rs +++ b/crates/aether-provider-transport/src/request_url/mod.rs @@ -125,6 +125,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) } @@ -386,6 +390,18 @@ fn build_provider_embedding_v1_url(upstream_base_url: &str, query: Option<&str>) build_provider_v1_url(upstream_base_url, "/embeddings", "/v1/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_v1_url(upstream_base_url, "/rerank", "/v1/rerank", query) } @@ -936,6 +952,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( @@ -995,6 +1017,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 74bde355a..73aeac4e9 100644 --- a/crates/aether-provider-transport/src/same_format_provider/mod.rs +++ b/crates/aether-provider-transport/src/same_format_provider/mod.rs @@ -159,6 +159,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, @@ -235,6 +243,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) => { @@ -728,6 +761,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/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/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/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 6d1c12fed..7e064d46a 100644 --- a/frontend/src/mocks/handler.ts +++ b/frontend/src/mocks/handler.ts @@ -1622,7 +1622,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/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 { From 906742e3c4fa341f752a8f926b4263b26d73c3d4 Mon Sep 17 00:00:00 2001 From: "Entropy.Xu" Date: Thu, 28 May 2026 23:09:36 +0800 Subject: [PATCH 02/29] =?UTF-8?q?fix(billing):=20=E5=A4=8D=E7=94=A8?= =?UTF-8?q?=E5=BE=85=E6=94=AF=E4=BB=98=E5=A5=97=E9=A4=90=E8=AE=A2=E5=8D=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/aether-gateway/src/data/state/runtime.rs | 15 ++ .../src/data/state/testing/mod.rs | 58 +++++ .../src/handlers/public/support/billing.rs | 26 +++ .../src/state/runtime/wallet/reads.rs | 12 + .../src/tests/frontdoor/public_support.rs | 208 ++++++++++++++++++ .../src/repository/wallet/memory.rs | 35 +++ .../src/repository/wallet/mysql.rs | 26 +++ .../src/repository/wallet/postgres.rs | 49 +++++ .../src/repository/wallet/sqlite.rs | 195 ++++++++++++++++ .../src/repository/wallet/types.rs | 6 + frontend/src/api/billing.ts | 1 + frontend/src/views/user/BillingPlans.vue | 6 +- 12 files changed, 636 insertions(+), 1 deletion(-) 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/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/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/frontdoor/public_support.rs b/apps/aether-gateway/src/tests/frontdoor/public_support.rs index 2540c2642..1079e8b4f 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}; @@ -3879,6 +3884,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/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/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/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) From bb5a4454a5469f85037a21e6ce25b428213b875a Mon Sep 17 00:00:00 2001 From: "Entropy.Xu" Date: Fri, 29 May 2026 14:35:16 +0800 Subject: [PATCH 03/29] =?UTF-8?q?feat(gateway):=20=E6=B7=BB=E5=8A=A0?= =?UTF-8?q?=E6=A0=87=E5=87=86=E6=96=87=E6=9C=AC=E9=9D=9E=E6=B5=81=E5=BC=8F?= =?UTF-8?q?=E5=BF=83=E8=B7=B3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../src/execution_runtime/sync/execution.rs | 8 +- .../src/execution_runtime/sync/mod.rs | 3 +- .../src/executor/orchestration.rs | 847 +++++++++++++++++- frontend/src/views/admin/SystemSettings.vue | 2 + .../system-settings/BasicConfigSection.vue | 23 + .../__tests__/useSystemConfig.spec.ts | 32 +- .../composables/useSystemConfig.ts | 18 +- 7 files changed, 924 insertions(+), 9 deletions(-) diff --git a/apps/aether-gateway/src/execution_runtime/sync/execution.rs b/apps/aether-gateway/src/execution_runtime/sync/execution.rs index 26e14d9c1..8f38e67c8 100644 --- a/apps/aether-gateway/src/execution_runtime/sync/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/sync/execution.rs @@ -1182,12 +1182,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 { 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..d29f5d587 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, @@ -86,6 +92,65 @@ pub(crate) async fn maybe_execute_sync_via_local_decision( body_json: &serde_json::Value, plan_kind: &str, ) -> Result { + 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 Some((attempt_source, candidate_count)) = build_local_openai_chat_sync_attempt_source_for_kind( state, parts, trace_id, decision, body_json, plan_kind, @@ -167,6 +232,57 @@ pub(crate) async fn maybe_execute_sync_via_local_openai_responses_decision( body_json: &serde_json::Value, plan_kind: &str, ) -> Result { + 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), + } + }, + )?, + )); + } + let Some((attempt_source, _candidate_count)) = build_local_openai_responses_sync_attempt_source_for_kind( state, parts, trace_id, decision, body_json, plan_kind, @@ -227,6 +343,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), + } + }, + )?, + )); + } + let Some((attempt_source, _candidate_count)) = build_standard_family_sync_attempt_source( state, parts, trace_id, decision, body_json, spec, ) @@ -391,6 +558,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), + } + }, + )?, + )); + } + let Some((attempt_source, _candidate_count)) = build_local_same_format_sync_attempt_source( state, parts, trace_id, decision, body_json, spec, ) @@ -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,208 @@ 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_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/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) { From 6bfd56b54fbdaabe5737969275c5eeddf3a03cab Mon Sep 17 00:00:00 2001 From: stabey Date: Sun, 31 May 2026 01:57:44 +0800 Subject: [PATCH 04/29] =?UTF-8?q?fix(usage):=20=E9=81=BF=E5=85=8D=E4=B8=8A?= =?UTF-8?q?=E6=B8=B8=E6=B5=81=E5=BC=8F=E9=94=99=E8=AF=AF=E8=AF=AF=E8=AE=B0?= =?UTF-8?q?=E4=B8=BA=E6=88=90=E5=8A=9F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 当上游流式响应中途失败时，sync error payload 可能同时包含合成错误体和部分上游流 body。 优先使用合成错误体生成 usage 终态，避免只因为上游先返回过 200 和部分 SSE 内容就把失败请求记录为 completed/settled。 --- crates/aether-usage-runtime/src/write.rs | 92 +++++++++++++++++++++++- 1 file changed, 91 insertions(+), 1 deletion(-) diff --git a/crates/aether-usage-runtime/src/write.rs b/crates/aether-usage-runtime/src/write.rs index 183229098..df61093cb 100644 --- a/crates/aether-usage-runtime/src/write.rs +++ b/crates/aether-usage-runtime/src/write.rs @@ -735,7 +735,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 { @@ -786,6 +789,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 { @@ -5108,6 +5117,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 { From 8d1e54eba62d7757336cd92b46b4df6748650f5c Mon Sep 17 00:00:00 2001 From: stabey Date: Sun, 31 May 2026 01:57:50 +0800 Subject: [PATCH 05/29] =?UTF-8?q?fix(stream):=20=E4=B8=AD=E9=80=94?= =?UTF-8?q?=E5=A4=B1=E8=B4=A5=E6=97=B6=E4=B8=8D=E5=90=88=E6=88=90=E6=AD=A3?= =?UTF-8?q?=E5=B8=B8=E6=94=B6=E5=B0=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 上游流式读取失败后，已经缓冲的局部转换状态可能是不完整的工具调用。 在 terminal failure 存在时跳过 normalizer 和 rewriter 的 finish 路径，避免把半截 tool_use 补成正常的 Claude message_stop。 --- .../src/execution_runtime/stream/execution.rs | 182 ++++++++++++++++-- 1 file changed, 171 insertions(+), 11 deletions(-) diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index 0cc6ab491..4653fe2b4 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -3357,7 +3357,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 +3397,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 +3477,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() => { @@ -3898,8 +3903,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; @@ -3996,6 +4002,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(); @@ -5017,6 +5029,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"); From 3dfafbc3797796ecc01c69b15a329a6271cb736c Mon Sep 17 00:00:00 2001 From: stabey Date: Sun, 31 May 2026 01:57:55 +0800 Subject: [PATCH 06/29] =?UTF-8?q?fix(ai):=20=E6=8C=89=20Responses=20?= =?UTF-8?q?=E6=96=87=E6=9C=AC=E5=88=86=E7=89=87=E5=8E=BB=E9=87=8D=E5=BF=AB?= =?UTF-8?q?=E7=85=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit upstream 已有 8abedecb 处理单个 OpenAI Responses 文本流中 delta 与 done/completed 快照重复输出的问题。 本提交保留该方向，并把去重状态从全局文本扩展为按 output_index/item_id 与 content_index 分片记录，避免多个 message item 或多个 text content part 共用同一段快照状态。 --- .../src/formats/openai/chat/stream.rs | 222 +++++++++++++++--- 1 file changed, 194 insertions(+), 28 deletions(-) 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(); From eda94913cf72f2ee7327fc6b5131767efd2cbb35 Mon Sep 17 00:00:00 2001 From: stabey Date: Sun, 31 May 2026 02:49:43 +0800 Subject: [PATCH 07/29] =?UTF-8?q?test(gateway):=20=E4=BF=AE=E5=A4=8D=20usa?= =?UTF-8?q?ge=20=E5=90=8C=E6=AD=A5=E6=B5=8B=E8=AF=95=E6=A0=88=E6=BA=A2?= =?UTF-8?q?=E5=87=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI 中 gateway_records_pending_usage_before_execution_runtime_sync_result_arrives 仍会在默认测试栈上溢出。 复用 large-stack tokio runtime 包装该测试，避免 gateway 全量测试在无业务失败时被 SIGABRT 中断。 --- apps/aether-gateway/src/tests/usage/direct.rs | 32 +++++++++++++++++-- 1 file changed, 30 insertions(+), 2 deletions(-) diff --git a/apps/aether-gateway/src/tests/usage/direct.rs b/apps/aether-gateway/src/tests/usage/direct.rs index dd4d1596d..0dae5fa64 100644 --- a/apps/aether-gateway/src/tests/usage/direct.rs +++ b/apps/aether-gateway/src/tests/usage/direct.rs @@ -19,6 +19,27 @@ fn large_request_body(stream: bool) -> String { .expect("request body should encode") } +fn run_async_test_on_large_stack(name: &'static str, future: F) +where + F: std::future::Future + Send + 'static, +{ + let handle = std::thread::Builder::new() + .name(name.to_string()) + .stack_size(16 * 1024 * 1024) + .spawn(move || { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("tokio runtime should build") + .block_on(future); + }) + .expect("large-stack usage direct test thread should spawn"); + + if let Err(payload) = handle.join() { + std::panic::resume_unwind(payload); + } +} + #[tokio::test] async fn gateway_records_usage_for_execution_runtime_sync_when_runtime_enabled() { let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); @@ -131,8 +152,15 @@ async fn gateway_records_usage_for_execution_runtime_sync_when_runtime_enabled() upstream_handle.abort(); } -#[tokio::test] -async fn gateway_records_pending_usage_before_execution_runtime_sync_result_arrives() { +#[test] +fn gateway_records_pending_usage_before_execution_runtime_sync_result_arrives() { + 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(), + ); +} + +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()); From 1c05a722c19e16ce85bf8a3a91956f47e8ae3654 Mon Sep 17 00:00:00 2001 From: stabey Date: Sun, 31 May 2026 03:09:11 +0800 Subject: [PATCH 08/29] =?UTF-8?q?test(gateway):=20=E4=BF=AE=E5=A4=8D=20usa?= =?UTF-8?q?ge=20local=20=E5=90=8C=E6=AD=A5=E6=B5=8B=E8=AF=95=E6=A0=88?= =?UTF-8?q?=E6=BA=A2=E5=87=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/aether-gateway/src/tests/usage/local.rs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/apps/aether-gateway/src/tests/usage/local.rs b/apps/aether-gateway/src/tests/usage/local.rs index b988f1ef6..118ec2507 100644 --- a/apps/aether-gateway/src/tests/usage/local.rs +++ b/apps/aether-gateway/src/tests/usage/local.rs @@ -85,8 +85,15 @@ where stored.expect("usage should be present once the expected status is observed") } -#[tokio::test] -async fn gateway_handles_local_openai_chat_sync_report_with_local_reporting_when_usage_runtime_enabled( +#[test] +fn gateway_handles_local_openai_chat_sync_report_with_local_reporting_when_usage_runtime_enabled() { + run_async_test_on_large_stack( + "gateway_handles_local_openai_chat_sync_report_with_local_reporting_when_usage_runtime_enabled", + gateway_handles_local_openai_chat_sync_report_with_local_reporting_when_usage_runtime_enabled_impl(), + ); +} + +async fn gateway_handles_local_openai_chat_sync_report_with_local_reporting_when_usage_runtime_enabled_impl( ) { let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); From ba4e194cb55fa0d1428b9dcf13529bd2a380d785 Mon Sep 17 00:00:00 2001 From: stabey Date: Sun, 31 May 2026 03:21:29 +0800 Subject: [PATCH 09/29] =?UTF-8?q?test(gateway):=20=E4=BF=AE=E5=A4=8D=20usa?= =?UTF-8?q?ge=20base=20=E8=AE=B0=E5=BD=95=E6=B5=8B=E8=AF=95=E6=A0=88?= =?UTF-8?q?=E6=BA=A2=E5=87=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/aether-gateway/src/tests/usage/local.rs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/apps/aether-gateway/src/tests/usage/local.rs b/apps/aether-gateway/src/tests/usage/local.rs index 118ec2507..3e19d8308 100644 --- a/apps/aether-gateway/src/tests/usage/local.rs +++ b/apps/aether-gateway/src/tests/usage/local.rs @@ -559,8 +559,15 @@ async fn gateway_applies_system_max_request_body_size_to_local_openai_chat_sync_ upstream_handle.abort(); } -#[tokio::test] -async fn gateway_strips_request_and_response_bodies_when_request_record_level_is_base() { +#[test] +fn gateway_strips_request_and_response_bodies_when_request_record_level_is_base() { + run_async_test_on_large_stack( + "gateway_strips_request_and_response_bodies_when_request_record_level_is_base", + gateway_strips_request_and_response_bodies_when_request_record_level_is_base_impl(), + ); +} + +async fn gateway_strips_request_and_response_bodies_when_request_record_level_is_base_impl() { let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); From de0a880ca655c1e7671afa9fae351cf31f4f4fad Mon Sep 17 00:00:00 2001 From: stabey Date: Sun, 31 May 2026 03:35:01 +0800 Subject: [PATCH 10/29] =?UTF-8?q?test(gateway):=20=E4=BF=AE=E5=A4=8D=20usa?= =?UTF-8?q?ge=20wallet=20=E6=B5=8B=E8=AF=95=E6=A0=88=E6=BA=A2=E5=87=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/aether-gateway/src/tests/usage/wallet.rs | 32 +++++++++++++++++-- 1 file changed, 30 insertions(+), 2 deletions(-) diff --git a/apps/aether-gateway/src/tests/usage/wallet.rs b/apps/aether-gateway/src/tests/usage/wallet.rs index 39efe4382..744e10294 100644 --- a/apps/aether-gateway/src/tests/usage/wallet.rs +++ b/apps/aether-gateway/src/tests/usage/wallet.rs @@ -11,8 +11,36 @@ use super::{ }; use aether_data::repository::settlement::InMemorySettlementRepository; -#[tokio::test] -async fn gateway_settles_wallet_for_completed_execution_runtime_sync_usage() { +fn run_async_test_on_large_stack(name: &'static str, future: F) +where + F: std::future::Future + Send + 'static, +{ + let handle = std::thread::Builder::new() + .name(name.to_string()) + .stack_size(8 * 1024 * 1024) + .spawn(move || { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("tokio runtime should build") + .block_on(future); + }) + .expect("large-stack test thread should spawn"); + + if let Err(payload) = handle.join() { + std::panic::resume_unwind(payload); + } +} + +#[test] +fn gateway_settles_wallet_for_completed_execution_runtime_sync_usage() { + run_async_test_on_large_stack( + "gateway_settles_wallet_for_completed_execution_runtime_sync_usage", + gateway_settles_wallet_for_completed_execution_runtime_sync_usage_impl(), + ); +} + +async fn gateway_settles_wallet_for_completed_execution_runtime_sync_usage_impl() { let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let billing_repository = Arc::new(InMemoryBillingReadRepository::seed(vec![ From 2e24e5f358d1c6c18074a3e7be519ce8d1352ad3 Mon Sep 17 00:00:00 2001 From: "Entropy.Xu" Date: Sun, 31 May 2026 19:52:18 +0800 Subject: [PATCH 11/29] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=EF=BC=9A=E6=89=A9?= =?UTF-8?q?=E5=A4=A7=20Kiro=20=E6=A8=A1=E6=8B=9F=E7=BC=93=E5=AD=98?= =?UTF-8?q?=E5=89=8D=E7=BC=80=E8=AF=BB=E5=8F=96=E8=8C=83=E5=9B=B4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../src/execution_runtime/kiro_cache.rs | 122 ++++++++++++------ .../src/execution_runtime/stream/execution.rs | 113 ++++++++++++++++ 2 files changed, 193 insertions(+), 42 deletions(-) diff --git a/apps/aether-gateway/src/execution_runtime/kiro_cache.rs b/apps/aether-gateway/src/execution_runtime/kiro_cache.rs index c677e7074..afe25eb84 100644 --- a/apps/aether-gateway/src/execution_runtime/kiro_cache.rs +++ b/apps/aether-gateway/src/execution_runtime/kiro_cache.rs @@ -8,7 +8,6 @@ use sha2::{Digest, Sha256}; 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 TOKENS_PER_TOOL: u64 = 150; const TOKENS_PER_MESSAGE: u64 = 4; const INLINE_IMAGE_DATA_TOKEN_PLACEHOLDER: &str = "[inline-image-data]"; @@ -21,6 +20,7 @@ pub(crate) struct KiroPromptCacheProfile { total_input_tokens: u64, min_cacheable_tokens: u64, breakpoints: Vec, + match_candidates: Vec, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -57,7 +57,7 @@ struct PendingBlock { } #[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct PrefixCandidate { +struct KiroPromptCacheCandidate { fingerprint: [u8; 32], cumulative_tokens: u64, } @@ -92,7 +92,7 @@ pub(crate) fn build_kiro_prompt_cache_profile( 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(); for block in flattened { cumulative_tokens = cumulative_tokens.saturating_add(block.tokens); @@ -105,12 +105,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, @@ -131,7 +125,7 @@ pub(crate) fn build_kiro_prompt_cache_profile( ); } } - 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,10 +133,19 @@ 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 { + let last_cacheable_tokens = cacheable_breakpoints.last()?.cumulative_tokens; + let match_candidates = match_candidates + .into_iter() + .filter(|candidate| { + candidate.cumulative_tokens >= min_cacheable_tokens + && candidate.cumulative_tokens <= last_cacheable_tokens + }) + .collect(); + Some(KiroPromptCacheProfile { total_input_tokens, min_cacheable_tokens, breakpoints: cacheable_breakpoints, + match_candidates, }) } @@ -269,35 +272,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 { @@ -634,20 +617,15 @@ 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); + for candidate in profile.match_candidates.iter().rev() { + let key = (credential_id.clone(), candidate.fingerprint); let Some(entry) = entries.get(&key) else { continue; }; if entry.expires_at > now { matched_tokens = entry .token_count - .min(breakpoint.cumulative_tokens) + .min(candidate.cumulative_tokens) .min(profile.total_input_tokens); break; } @@ -867,6 +845,66 @@ mod tests { assert!(hit.cache_creation_input_tokens > 0); } + #[test] + fn tracker_reads_cached_prefix_beyond_tail_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 billed_input_tokens_subtracts_cache_usage() { assert_eq!( diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index 0cc6ab491..a0312a435 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -4217,6 +4217,119 @@ mod tests { assert_eq!(second_usage.output_tokens, 19); } + #[test] + fn kiro_stream_summary_reads_cached_prefix_beyond_tail_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 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( + &plan, + Some(&first_report_context), + &mut first_summary, + ); + 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( + &plan, + Some(&second_report_context), + &mut second_summary, + ); + 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); + } + #[test] fn kiro_stream_summary_seeds_input_tokens_without_cache_control() { let request_body = json!({ From a3b8a997097e9fa025966774686d5ea29bfd8e10 Mon Sep 17 00:00:00 2001 From: "Entropy.Xu" Date: Sun, 31 May 2026 20:22:18 +0800 Subject: [PATCH 12/29] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=EF=BC=9A=E8=A1=A5?= =?UTF-8?q?=E9=BD=90=20Kiro=20=E6=A8=A1=E6=8B=9F=E7=BC=93=E5=AD=98=20TTL?= =?UTF-8?q?=20=E5=92=8C=E6=B6=88=E6=81=AF=E7=BA=A7=E6=96=AD=E7=82=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../src/execution_runtime/kiro_cache.rs | 44 +++++++++++++++---- 1 file changed, 36 insertions(+), 8 deletions(-) diff --git a/apps/aether-gateway/src/execution_runtime/kiro_cache.rs b/apps/aether-gateway/src/execution_runtime/kiro_cache.rs index afe25eb84..624700e6c 100644 --- a/apps/aether-gateway/src/execution_runtime/kiro_cache.rs +++ b/apps/aether-gateway/src/execution_runtime/kiro_cache.rs @@ -359,11 +359,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!({ @@ -392,7 +398,7 @@ fn flatten_cacheable_blocks(request_body: &Value) -> Vec { blocks.push(PendingBlock { tokens: count_text_tokens(text), value, - breakpoint_ttl: None, + breakpoint_ttl: message_breakpoint_ttl, is_message_end: true, }); } @@ -407,7 +413,7 @@ fn flatten_cacheable_blocks(request_body: &Value) -> Vec { blocks.push(PendingBlock { tokens: count_message_content_tokens(other), value, - breakpoint_ttl: None, + breakpoint_ttl: message_breakpoint_ttl, is_message_end: true, }); } @@ -619,10 +625,11 @@ impl KiroPromptCacheTracker { let mut matched_tokens = 0; for candidate in profile.match_candidates.iter().rev() { let key = (credential_id.clone(), candidate.fingerprint); - let Some(entry) = entries.get(&key) else { + 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(candidate.cumulative_tokens) @@ -642,6 +649,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); @@ -734,7 +742,7 @@ mod tests { } #[test] - fn tracker_supports_prefix_hits_without_extending_expiry() { + fn tracker_refreshes_cached_prefix_ttl_on_read() { let base = serde_json::json!({ "model": "claude-sonnet-4.6", "system": [{ @@ -777,13 +785,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] @@ -905,6 +913,26 @@ mod tests { assert!(hit.cache_creation_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!( From 9734be31cfebd723b2174981b6dd434421f97b5c Mon Sep 17 00:00:00 2001 From: "Entropy.Xu" Date: Sun, 31 May 2026 20:45:16 +0800 Subject: [PATCH 13/29] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=EF=BC=9A=E6=94=B6?= =?UTF-8?q?=E6=95=9B=20Kiro=20=E6=A8=A1=E6=8B=9F=E7=BC=93=E5=AD=98?= =?UTF-8?q?=E6=96=AD=E7=82=B9=E8=AF=AD=E4=B9=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../src/execution_runtime/kiro_cache.rs | 185 +++++++++++++++--- .../src/execution_runtime/stream/execution.rs | 2 +- 2 files changed, 154 insertions(+), 33 deletions(-) diff --git a/apps/aether-gateway/src/execution_runtime/kiro_cache.rs b/apps/aether-gateway/src/execution_runtime/kiro_cache.rs index 624700e6c..a62bc6776 100644 --- a/apps/aether-gateway/src/execution_runtime/kiro_cache.rs +++ b/apps/aether-gateway/src/execution_runtime/kiro_cache.rs @@ -8,6 +8,7 @@ use sha2::{Digest, Sha256}; 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_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]"; @@ -53,7 +54,6 @@ struct PendingBlock { value: Value, tokens: u64, breakpoint_ttl: Option, - is_message_end: bool, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -75,7 +75,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 +90,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 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,7 +109,6 @@ pub(crate) fn build_kiro_prompt_cache_profile( prefix_hasher.update(fingerprint); if let Some(ttl) = block.breakpoint_ttl { - active_ttl = Some(ttl); push_breakpoint( &mut breakpoints, &mut seen_fingerprints, @@ -114,17 +117,6 @@ 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_match_candidate(&mut match_candidates, fingerprint, cumulative_tokens); } @@ -133,14 +125,14 @@ pub(crate) fn build_kiro_prompt_cache_profile( .into_iter() .filter(|breakpoint| breakpoint.cumulative_tokens >= min_cacheable_tokens) .collect::>(); - let last_cacheable_tokens = cacheable_breakpoints.last()?.cumulative_tokens; - let match_candidates = match_candidates - .into_iter() - .filter(|candidate| { - candidate.cumulative_tokens >= min_cacheable_tokens - && candidate.cumulative_tokens <= last_cacheable_tokens - }) - .collect(); + 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, @@ -149,6 +141,38 @@ pub(crate) fn build_kiro_prompt_cache_profile( }) } +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) @@ -299,7 +323,6 @@ fn flatten_cacheable_blocks(request_body: &Value) -> Vec { tokens: TOKENS_PER_TOOL, value, breakpoint_ttl, - is_message_end: false, }); } } @@ -320,7 +343,6 @@ fn flatten_cacheable_blocks(request_body: &Value) -> Vec { tokens: count_system_block_tokens(item), value, breakpoint_ttl, - is_message_end: false, }); } } @@ -334,7 +356,6 @@ fn flatten_cacheable_blocks(request_body: &Value) -> Vec { tokens: count_text_tokens(text), value, breakpoint_ttl: None, - is_message_end: false, }); } other => { @@ -347,7 +368,6 @@ fn flatten_cacheable_blocks(request_body: &Value) -> Vec { tokens: count_system_block_tokens(other), value, breakpoint_ttl: None, - is_message_end: false, }); } } @@ -383,7 +403,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, }); } } @@ -399,7 +418,6 @@ fn flatten_cacheable_blocks(request_body: &Value) -> Vec { tokens: count_text_tokens(text), value, breakpoint_ttl: message_breakpoint_ttl, - is_message_end: true, }); } Some(other) => { @@ -414,7 +432,6 @@ fn flatten_cacheable_blocks(request_body: &Value) -> Vec { tokens: count_message_content_tokens(other), value, breakpoint_ttl: message_breakpoint_ttl, - is_message_end: true, }); } None => {} @@ -741,6 +758,50 @@ mod tests { ); } + #[test] + 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); + } + #[test] fn tracker_refreshes_cached_prefix_ttl_on_read() { let base = serde_json::json!({ @@ -854,7 +915,7 @@ mod tests { } #[test] - fn tracker_reads_cached_prefix_beyond_tail_lookback_window() { + fn tracker_reads_cached_prefix_within_prompt_cache_lookback_window() { let first = serde_json::json!({ "model": "claude-sonnet-4.6", "messages": [{ @@ -913,6 +974,66 @@ mod tests { 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!({ diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index a0312a435..f85f1fe7f 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -4218,7 +4218,7 @@ mod tests { } #[test] - fn kiro_stream_summary_reads_cached_prefix_beyond_tail_lookback_window() { + fn kiro_stream_summary_reads_cached_prefix_within_prompt_cache_lookback_window() { let first_request_body = json!({ "model": "claude-sonnet-4.6", "messages": [{ From c2bcfab7d436dcefb22cdc71297028b71ae1f828 Mon Sep 17 00:00:00 2001 From: "Entropy.Xu" Date: Sun, 31 May 2026 22:13:45 +0800 Subject: [PATCH 14/29] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=EF=BC=9AKiro=20?= =?UTF-8?q?=E6=A8=A1=E6=8B=9F=E7=BC=93=E5=AD=98=E6=8E=A5=E5=85=A5=E5=85=B1?= =?UTF-8?q?=E4=BA=AB=E8=BF=90=E8=A1=8C=E6=97=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../src/execution_runtime/kiro_cache.rs | 185 +++++++++++ .../src/execution_runtime/kiro_web_search.rs | 23 +- .../src/execution_runtime/stream/execution.rs | 23 +- .../src/execution_runtime/sync/execution.rs | 288 +++++++++++++++++- .../src/tests/ai_execute/sync/claude/kiro.rs | 150 +++++++-- crates/aether-usage-runtime/src/write.rs | 110 ++++++- 6 files changed, 734 insertions(+), 45 deletions(-) diff --git a/apps/aether-gateway/src/execution_runtime/kiro_cache.rs b/apps/aether-gateway/src/execution_runtime/kiro_cache.rs index a62bc6776..3e019462d 100644 --- a/apps/aether-gateway/src/execution_runtime/kiro_cache.rs +++ b/apps/aether-gateway/src/execution_runtime/kiro_cache.rs @@ -1,9 +1,13 @@ 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; const DEFAULT_CACHE_TTL: Duration = Duration::from_secs(300); const ONE_HOUR_CACHE_TTL: Duration = Duration::from_secs(3600); @@ -38,6 +42,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, @@ -66,6 +76,153 @@ 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 { + runtime_state + .kv_set( + key.as_str(), + encode_kiro_prompt_cache_runtime_entry(entry), + Some(Duration::from_secs(entry.ttl_secs.max(1))), + ) + .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, + }); + runtime_state + .kv_set( + key.as_str(), + encode_kiro_prompt_cache_runtime_entry(entry), + Some(Duration::from_secs(entry.ttl_secs.max(1))), + ) + .await?; + } + + Ok(KiroPromptCacheUsage { + cache_creation_input_tokens: creation_tokens, + cache_read_input_tokens: matched_tokens, + }) +} + +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, @@ -705,6 +862,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)) @@ -802,6 +960,33 @@ mod tests { 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); + } + #[test] fn tracker_refreshes_cached_prefix_ttl_on_read() { let base = serde_json::json!({ 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 f85f1fe7f..774553074 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_prompt_cache_tracker, 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; } @@ -1984,7 +1989,7 @@ async fn execute_stream_from_frame_stream( 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); + 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; @@ -4741,9 +4746,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 @@ -4810,9 +4817,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 diff --git a/apps/aether-gateway/src/execution_runtime/sync/execution.rs b/apps/aether-gateway/src/execution_runtime/sync/execution.rs index 26e14d9c1..feae9586c 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>, @@ -1934,8 +2114,13 @@ 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_report_context_input_tokens(&plan, &mut report_context); + seed_kiro_sync_simulated_cache_enabled(state, &plan, &mut report_context).await; + 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 +2689,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 +2883,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/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/crates/aether-usage-runtime/src/write.rs b/crates/aether-usage-runtime/src/write.rs index 183229098..096930731 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, } @@ -763,6 +764,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 +779,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, @@ -847,11 +850,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 +911,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 +1776,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()) } @@ -5453,6 +5503,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 { From d1b64b67487b1200021b7b4ef2dbc186860f6e4f Mon Sep 17 00:00:00 2001 From: "Entropy.Xu" Date: Sun, 31 May 2026 22:43:27 +0800 Subject: [PATCH 15/29] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=EF=BC=9A=E5=AE=8C?= =?UTF-8?q?=E5=96=84=20Kiro=20=E6=A8=A1=E6=8B=9F=E7=BC=93=E5=AD=98?= =?UTF-8?q?=E5=85=B1=E4=BA=AB=E5=9B=9E=E6=94=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../src/execution_runtime/kiro_cache.rs | 168 ++++++++++++++++-- .../src/execution_runtime/stream/execution.rs | 83 ++++++--- 2 files changed, 212 insertions(+), 39 deletions(-) diff --git a/apps/aether-gateway/src/execution_runtime/kiro_cache.rs b/apps/aether-gateway/src/execution_runtime/kiro_cache.rs index 3e019462d..839c08fa5 100644 --- a/apps/aether-gateway/src/execution_runtime/kiro_cache.rs +++ b/apps/aether-gateway/src/execution_runtime/kiro_cache.rs @@ -9,9 +9,12 @@ 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 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; @@ -148,13 +151,7 @@ async fn compute_kiro_prompt_cache_usage_with_runtime_state( } if let Some((key, entry)) = matched_refresh { - runtime_state - .kv_set( - key.as_str(), - encode_kiro_prompt_cache_runtime_entry(entry), - Some(Duration::from_secs(entry.ttl_secs.max(1))), - ) - .await?; + store_kiro_prompt_cache_runtime_entry(runtime_state, key.as_str(), entry).await?; } let creation_tokens = last_breakpoint @@ -176,21 +173,104 @@ async fn compute_kiro_prompt_cache_usage_with_runtime_state( token_count: breakpoint.cumulative_tokens, ttl_secs, }); - runtime_state - .kv_set( - key.as_str(), - encode_kiro_prompt_cache_runtime_entry(entry), - Some(Duration::from_secs(entry.ttl_secs.max(1))), - ) - .await?; + 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() @@ -987,6 +1067,66 @@ mod tests { 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!({ diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index 774553074..80ac484e5 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, - compute_kiro_prompt_cache_usage, 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, @@ -499,7 +499,8 @@ fn kiro_cache_usage_from_report_context(report_context: &Value) -> Option, summary: &mut Option, @@ -577,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; } @@ -3705,10 +3710,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(), @@ -3951,6 +3958,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")); @@ -4129,8 +4140,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": [ @@ -4178,6 +4189,7 @@ mod tests { transport_profile: None, timeouts: None, }; + let state = test_state(); let mut first_summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { @@ -4188,10 +4200,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()) @@ -4208,10 +4222,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()) @@ -4222,8 +4238,8 @@ mod tests { assert_eq!(second_usage.output_tokens, 19); } - #[test] - fn kiro_stream_summary_reads_cached_prefix_within_prompt_cache_lookback_window() { + #[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": [{ @@ -4289,6 +4305,7 @@ mod tests { "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 { @@ -4299,10 +4316,12 @@ mod tests { ..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()) @@ -4319,10 +4338,12 @@ mod tests { ..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()) @@ -4335,8 +4356,8 @@ 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_seeds_input_tokens_without_cache_control() { let request_body = json!({ "model": "claude-opus-4-7", "system": [ @@ -4382,6 +4403,7 @@ mod tests { transport_profile: None, timeouts: None, }; + let state = test_state(); let mut summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { @@ -4393,10 +4415,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() @@ -4409,8 +4433,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": [ @@ -4456,6 +4480,7 @@ mod tests { transport_profile: None, timeouts: None, }; + let state = test_state(); let mut summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { @@ -4468,10 +4493,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() @@ -4484,8 +4511,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": [ @@ -4532,6 +4559,7 @@ mod tests { transport_profile: None, timeouts: None, }; + let state = test_state(); let mut summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { @@ -4545,10 +4573,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() @@ -4561,8 +4591,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": [ @@ -4610,6 +4640,7 @@ mod tests { transport_profile: None, timeouts: None, }; + let state = test_state(); let mut summary = Some(ExecutionStreamTerminalSummary { standardized_usage: Some(StandardizedUsage { @@ -4623,10 +4654,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() From 98dc5925a510662cf6a88255249fd4eaa2230ae6 Mon Sep 17 00:00:00 2001 From: zhefox Date: Tue, 2 Jun 2026 00:35:19 +0800 Subject: [PATCH 16/29] fix(provider): preserve openai responses tool history in chat conversion --- .../src/formats/conversion/request.rs | 59 +++++++++++++ .../src/protocol/canonical.rs | 88 +++++++++++++------ 2 files changed, 120 insertions(+), 27 deletions(-) diff --git a/crates/aether-ai-formats/src/formats/conversion/request.rs b/crates/aether-ai-formats/src/formats/conversion/request.rs index 16b7b44cd..6340de445 100644 --- a/crates/aether-ai-formats/src/formats/conversion/request.rs +++ b/crates/aether-ai-formats/src/formats/conversion/request.rs @@ -163,6 +163,7 @@ mod tests { convert_openai_chat_request_to_claude_request, convert_openai_chat_request_to_openai_responses_request, normalize_claude_request_to_openai_chat_request, + normalize_openai_responses_request_to_openai_chat_request, }; #[test] @@ -216,6 +217,64 @@ mod tests { assert_eq!(converted["messages"][0]["content"], "hello"); } + #[test] + fn responses_request_normalizer_keeps_tool_history_chat_safe() { + let call_id = "call_weather_123"; + let tool_output = json!({ + "toolCallId": call_id, + "input": {"city": "Hangzhou"}, + "output": { + "content": [{"type": "text", "text": "sunny"}], + "isError": false, + }, + }); + let body = json!({ + "model": "glm-5.1", + "input": [ + { + "role": "user", + "content": [{"type": "input_text", "text": "weather now"}] + }, + { + "type": "function_call", + "call_id": call_id, + "id": call_id, + "name": "mcp__mapsWeather", + "arguments": "{\"city\":\"Hangzhou\"}" + }, + { + "type": "function_call_output", + "call_id": call_id, + "output": tool_output.to_string() + } + ] + }); + + 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(), 3); + assert_eq!(messages[0]["role"], "user"); + assert_eq!(messages[0]["content"], "weather now"); + assert_eq!(messages[1]["role"], "assistant"); + assert!(messages[1]["content"].is_null()); + assert_eq!(messages[1]["tool_calls"][0]["id"], call_id); + assert_eq!( + messages[1]["tool_calls"][0]["function"]["name"], + "mcp__mapsWeather" + ); + assert_eq!(messages[2]["role"], "tool"); + assert_eq!(messages[2]["tool_call_id"], call_id); + 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 + ); + } + #[test] fn request_normalizer_preserves_multiple_claude_tool_results() { let body = json!({ diff --git a/crates/aether-ai-formats/src/protocol/canonical.rs b/crates/aether-ai-formats/src/protocol/canonical.rs index f6a46d557..f9f441a2b 100644 --- a/crates/aether-ai-formats/src/protocol/canonical.rs +++ b/crates/aether-ai-formats/src/protocol/canonical.rs @@ -13,6 +13,7 @@ pub(crate) const OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE: &str = "openai_cli const AETHER_EXTENSION_NAMESPACE: &str = "aether"; 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)] @@ -1644,6 +1645,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,17 +1671,7 @@ 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(), }); @@ -1836,6 +1842,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 +1866,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 +2392,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 +2439,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() @@ -2449,6 +2461,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()), From 0daa8c196bf8b0e08969846c91e6b832a3a8f3f5 Mon Sep 17 00:00:00 2001 From: zhefox Date: Tue, 2 Jun 2026 09:04:55 +0800 Subject: [PATCH 17/29] fix(provider): preserve reasoning and Claude tool results in responses conversion --- .../src/formats/conversion/request.rs | 149 +++++++++- .../src/formats/openai/responses/request.rs | 278 +++++++++++++++++- .../src/protocol/canonical.rs | 249 ++++++++++++---- 3 files changed, 593 insertions(+), 83 deletions(-) diff --git a/crates/aether-ai-formats/src/formats/conversion/request.rs b/crates/aether-ai-formats/src/formats/conversion/request.rs index 6340de445..c27435bef 100644 --- a/crates/aether-ai-formats/src/formats/conversion/request.rs +++ b/crates/aether-ai-formats/src/formats/conversion/request.rs @@ -159,6 +159,8 @@ 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, @@ -219,9 +221,10 @@ mod tests { #[test] fn responses_request_normalizer_keeps_tool_history_chat_safe() { - let call_id = "call_weather_123"; - let tool_output = json!({ - "toolCallId": call_id, + 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"}], @@ -231,21 +234,44 @@ mod tests { let body = json!({ "model": "glm-5.1", "input": [ + "weather now", { - "role": "user", - "content": [{"type": "input_text", "text": "weather now"}] + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "thinking first"}] + }, + { + "type": "message", + "role": "assistant", + "content": "planning" }, { "type": "function_call", - "call_id": call_id, - "id": call_id, + "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, - "output": tool_output.to_string() + "call_id": call_id_one, + "output": tool_output_one.to_string() + }, + { + "type": "function_call_output", + "call_id": call_id_two, + "output": "done-2" } ] }); @@ -254,25 +280,57 @@ mod tests { .expect("openai chat request"); let messages = converted["messages"].as_array().expect("messages"); - assert_eq!(messages.len(), 3); + 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!(messages[1]["content"].is_null()); - assert_eq!(messages[1]["tool_calls"][0]["id"], call_id); + 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); + 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 + 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] @@ -609,4 +667,67 @@ mod tests { assert!(!block_content_json.contains("\"source\"")); assert!(!block_content_json.contains("document body")); } + + #[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/openai/responses/request.rs b/crates/aether-ai-formats/src/formats/openai/responses/request.rs index 12a05f171..42d87b7e6 100644 --- a/crates/aether-ai-formats/src/formats/openai/responses/request.rs +++ b/crates/aether-ai-formats/src/formats/openai/responses/request.rs @@ -1,11 +1,13 @@ +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, + canonical_response_format_to_openai, canonicalize_tool_arguments, 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, @@ -223,6 +225,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 +235,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 +247,37 @@ 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, .. } => { + 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,6 +285,25 @@ 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) @@ -535,13 +579,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 +791,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/protocol/canonical.rs b/crates/aether-ai-formats/src/protocol/canonical.rs index f9f441a2b..e606e13b6 100644 --- a/crates/aether-ai-formats/src/protocol/canonical.rs +++ b/crates/aether-ai-formats/src/protocol/canonical.rs @@ -1338,6 +1338,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; @@ -1505,6 +1522,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() { @@ -1517,9 +1535,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) @@ -1527,6 +1549,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 @@ -1549,11 +1577,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( @@ -1561,13 +1591,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")) @@ -1581,52 +1615,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 @@ -1675,16 +1677,11 @@ pub(crate) fn openai_responses_input_to_canonical_messages( }], 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) @@ -1693,6 +1690,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> { @@ -2453,7 +2546,7 @@ fn canonical_tool_result_to_openai_chat(block: &CanonicalContentBlock) -> Value Value::Object(output) } -fn is_claude_tool_result(extensions: &BTreeMap) -> bool { +pub(crate) fn is_claude_tool_result(extensions: &BTreeMap) -> bool { extensions .get(AETHER_EXTENSION_NAMESPACE) .and_then(|value| value.get("source")) @@ -2862,6 +2955,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() @@ -5568,6 +5664,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!({ From 0a2c674ad8f0029cd6ab169f7b84af8ae3d02820 Mon Sep 17 00:00:00 2001 From: elky Date: Tue, 2 Jun 2026 09:43:12 +0800 Subject: [PATCH 18/29] Ignore tunnel release tags for app build version --- apps/aether-gateway/build.rs | 29 +++++++++---------- .../src/tests/architecture/admin_system.rs | 16 +++++++++- deploy.sh | 4 +-- frontend/vite.config.ts | 11 +++++-- 4 files changed, 39 insertions(+), 21 deletions(-) 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/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/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/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' } From 03e436707d28cb2a039adb5eaa11dc8aca7490a3 Mon Sep 17 00:00:00 2001 From: elky Date: Tue, 2 Jun 2026 16:59:13 +0800 Subject: [PATCH 19/29] Fix gateway usage nextest stack overflow --- apps/aether-gateway/src/tests/usage/local.rs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/apps/aether-gateway/src/tests/usage/local.rs b/apps/aether-gateway/src/tests/usage/local.rs index b988f1ef6..118ec2507 100644 --- a/apps/aether-gateway/src/tests/usage/local.rs +++ b/apps/aether-gateway/src/tests/usage/local.rs @@ -85,8 +85,15 @@ where stored.expect("usage should be present once the expected status is observed") } -#[tokio::test] -async fn gateway_handles_local_openai_chat_sync_report_with_local_reporting_when_usage_runtime_enabled( +#[test] +fn gateway_handles_local_openai_chat_sync_report_with_local_reporting_when_usage_runtime_enabled() { + run_async_test_on_large_stack( + "gateway_handles_local_openai_chat_sync_report_with_local_reporting_when_usage_runtime_enabled", + gateway_handles_local_openai_chat_sync_report_with_local_reporting_when_usage_runtime_enabled_impl(), + ); +} + +async fn gateway_handles_local_openai_chat_sync_report_with_local_reporting_when_usage_runtime_enabled_impl( ) { let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); From 6cabcd85aac9223f2899bf60efc0dca630197934 Mon Sep 17 00:00:00 2001 From: zhefox Date: Tue, 2 Jun 2026 17:14:26 +0800 Subject: [PATCH 20/29] fix(provider): preserve Claude messages defaults in responses conversion --- .../src/formats/claude/messages/request.rs | 7 +- .../src/formats/conversion/request.rs | 136 +++++++++++++ .../src/formats/openai/responses/request.rs | 180 ++++++++++++++++-- .../src/formats/shared/standard_matrix.rs | 27 +++ .../src/protocol/canonical.rs | 74 ++++++- 5 files changed, 403 insertions(+), 21 deletions(-) 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 c27435bef..49cf9132f 100644 --- a/crates/aether-ai-formats/src/formats/conversion/request.rs +++ b/crates/aether-ai-formats/src/formats/conversion/request.rs @@ -668,6 +668,142 @@ mod tests { 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!({ 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 42d87b7e6..9fed5a930 100644 --- a/crates/aether-ai-formats/src/formats/openai/responses/request.rs +++ b/crates/aether-ai-formats/src/formats/openai/responses/request.rs @@ -6,9 +6,10 @@ 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, is_claude_tool_result, - 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, @@ -137,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); @@ -172,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); } @@ -190,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"); } @@ -201,6 +204,7 @@ fn canonical_instructions_to_responses(canonical: &CanonicalRequest) -> Option>() @@ -208,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 { @@ -270,7 +316,12 @@ fn canonical_messages_to_responses_input(canonical: &CanonicalRequest) -> Option })); } } - CanonicalContentBlock::Thinking { text, .. } => { + CanonicalContentBlock::Thinking { + text, extensions, .. + } => { + if is_claude_thinking_block(extensions) { + continue; + } if role == "assistant" && !text.trim().is_empty() { content.push(json!({ "type": "output_text", @@ -309,6 +360,49 @@ fn canonical_messages_to_responses_input(canonical: &CanonicalRequest) -> Option 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>, @@ -463,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() @@ -525,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, @@ -567,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()), 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/protocol/canonical.rs b/crates/aether-ai-formats/src/protocol/canonical.rs index e606e13b6..afd59d043 100644 --- a/crates/aether-ai-formats/src/protocol/canonical.rs +++ b/crates/aether-ai-formats/src/protocol/canonical.rs @@ -11,6 +11,9 @@ 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"; @@ -1052,7 +1055,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()), }]) } } @@ -1071,7 +1074,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"], + )), }); } } @@ -1081,6 +1087,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> { @@ -1174,7 +1214,10 @@ pub(crate) fn claude_block_to_canonical_block(block: &Value) -> Option Some(CanonicalContentBlock::Thinking { text: String::new(), @@ -1183,7 +1226,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), @@ -2546,6 +2592,22 @@ fn canonical_tool_result_to_openai_chat(block: &CanonicalContentBlock) -> Value Value::Object(output) } +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) @@ -4108,7 +4170,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))) } From 746af0d93ecde6c22c2182307e827e61c4cd11e7 Mon Sep 17 00:00:00 2001 From: elky Date: Tue, 2 Jun 2026 23:06:29 +0800 Subject: [PATCH 21/29] Fix Kiro cache usage reporting --- .../src/execution_runtime/stream/execution.rs | 5 +++-- .../src/execution_runtime/sync/execution.rs | 4 +++- apps/aether-gateway/src/handlers/admin/users/shared.rs | 9 +++++++++ 3 files changed, 15 insertions(+), 3 deletions(-) diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index f40e078b2..e6bfc9dfa 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -841,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()); @@ -1991,9 +1990,11 @@ 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; + 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(); diff --git a/apps/aether-gateway/src/execution_runtime/sync/execution.rs b/apps/aether-gateway/src/execution_runtime/sync/execution.rs index feae9586c..2e0c0b929 100644 --- a/apps/aether-gateway/src/execution_runtime/sync/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/sync/execution.rs @@ -2117,8 +2117,10 @@ async fn execute_execution_runtime_sync_impl( let mut report_context = attach_provider_response_headers_to_report_context(report_context, &headers); if (200..300).contains(&status_code) { - seed_kiro_sync_report_context_input_tokens(&plan, &mut report_context); 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(); 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> { From ef6a11c1463cd00a09a7415f61505a53715a3ec3 Mon Sep 17 00:00:00 2001 From: elky Date: Wed, 3 Jun 2026 00:25:01 +0800 Subject: [PATCH 22/29] fix(gateway): preserve heartbeat no-path fallback --- .../src/executor/orchestration.rs | 99 ++++++++++++------- 1 file changed, 65 insertions(+), 34 deletions(-) diff --git a/apps/aether-gateway/src/executor/orchestration.rs b/apps/aether-gateway/src/executor/orchestration.rs index d29f5d587..f4a894b20 100644 --- a/apps/aether-gateway/src/executor/orchestration.rs +++ b/apps/aether-gateway/src/executor/orchestration.rs @@ -92,6 +92,15 @@ pub(crate) async fn maybe_execute_sync_via_local_decision( body_json: &serde_json::Value, plan_kind: &str, ) -> Result { + let Some((attempt_source, candidate_count)) = + build_local_openai_chat_sync_attempt_source_for_kind( + state, parts, trace_id, decision, body_json, plan_kind, + ) + .await? + else { + 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(); @@ -151,15 +160,6 @@ pub(crate) async fn maybe_execute_sync_via_local_decision( )); } - let Some((attempt_source, candidate_count)) = - build_local_openai_chat_sync_attempt_source_for_kind( - state, parts, trace_id, decision, body_json, plan_kind, - ) - .await? - else { - return Ok(LocalExecutionRequestOutcome::NoPath); - }; - let outcome = execute_sync_attempt_source::( state, parts, @@ -232,6 +232,15 @@ pub(crate) async fn maybe_execute_sync_via_local_openai_responses_decision( body_json: &serde_json::Value, plan_kind: &str, ) -> Result { + let Some((attempt_source, _candidate_count)) = + build_local_openai_responses_sync_attempt_source_for_kind( + state, parts, trace_id, decision, body_json, plan_kind, + ) + .await? + else { + 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(); @@ -283,15 +292,6 @@ pub(crate) async fn maybe_execute_sync_via_local_openai_responses_decision( )); } - let Some((attempt_source, _candidate_count)) = - build_local_openai_responses_sync_attempt_source_for_kind( - state, parts, trace_id, decision, body_json, plan_kind, - ) - .await? - else { - return Ok(LocalExecutionRequestOutcome::NoPath); - }; - execute_sync_attempt_source::( state, parts, @@ -343,6 +343,14 @@ pub(crate) async fn maybe_execute_sync_via_standard_family_decision( return Ok(LocalExecutionRequestOutcome::NoPath); }; + let Some((attempt_source, _candidate_count)) = build_standard_family_sync_attempt_source( + state, parts, trace_id, decision, body_json, spec, + ) + .await? + else { + 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(); @@ -394,14 +402,6 @@ pub(crate) async fn maybe_execute_sync_via_standard_family_decision( )); } - let Some((attempt_source, _candidate_count)) = build_standard_family_sync_attempt_source( - state, parts, trace_id, decision, body_json, spec, - ) - .await? - else { - return Ok(LocalExecutionRequestOutcome::NoPath); - }; - execute_sync_attempt_source::( state, parts, @@ -558,6 +558,14 @@ pub(crate) async fn maybe_execute_sync_via_local_same_format_provider_decision( return Ok(LocalExecutionRequestOutcome::NoPath); }; + let Some((attempt_source, _candidate_count)) = build_local_same_format_sync_attempt_source( + state, parts, trace_id, decision, body_json, spec, + ) + .await? + else { + 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(); @@ -609,14 +617,6 @@ pub(crate) async fn maybe_execute_sync_via_local_same_format_provider_decision( )); } - let Some((attempt_source, _candidate_count)) = build_local_same_format_sync_attempt_source( - state, parts, trace_id, decision, body_json, spec, - ) - .await? - else { - return Ok(LocalExecutionRequestOutcome::NoPath); - }; - execute_sync_attempt_source::( state, parts, @@ -1776,6 +1776,37 @@ mod tests { 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() From 869c3d3037fb5cb6675a119638b50161393b2dd4 Mon Sep 17 00:00:00 2001 From: elky Date: Wed, 3 Jun 2026 00:48:36 +0800 Subject: [PATCH 23/29] Fix finalize local test stack overflow --- .../src/tests/ai_execute/finalize_local.rs | 81 ++++++++++++++++--- 1 file changed, 71 insertions(+), 10 deletions(-) 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(); From c7a5155ce4cdbf137f9bb8aa87da4ae0906b755b Mon Sep 17 00:00:00 2001 From: elky Date: Wed, 3 Jun 2026 01:12:40 +0800 Subject: [PATCH 24/29] Fix sync CLI test stack overflow --- .../src/tests/ai_execute/sync/cli.rs | 128 +++++++++++++++--- 1 file changed, 110 insertions(+), 18 deletions(-) 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, From 5ac153a0bb50286601d56bc64332bd64c9796aaa Mon Sep 17 00:00:00 2001 From: elky Date: Wed, 3 Jun 2026 01:32:25 +0800 Subject: [PATCH 25/29] Fix gateway nextest stack limit --- .github/workflows/rust-ci.yml | 1 + .../ai_execute/sync/pii_redaction_formats.rs | 99 +++++++++++-------- 2 files changed, 58 insertions(+), 42 deletions(-) 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/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] From 5e39801bba119c4e9fe3be08d8bf14821f8d6cae Mon Sep 17 00:00:00 2001 From: zhefox Date: Wed, 3 Jun 2026 10:52:26 +0800 Subject: [PATCH 26/29] fix(provider): clamp reasoning effort and filter chat extensions --- .../src/formats/conversion/request.rs | 66 +++++++++++++++++++ .../src/formats/openai/chat/request.rs | 32 ++++++++- .../src/formats/shared/model_directives.rs | 7 +- .../src/formats/shared/standard_normalize.rs | 2 +- 4 files changed, 99 insertions(+), 8 deletions(-) diff --git a/crates/aether-ai-formats/src/formats/conversion/request.rs b/crates/aether-ai-formats/src/formats/conversion/request.rs index 49cf9132f..c556c963f 100644 --- a/crates/aether-ai-formats/src/formats/conversion/request.rs +++ b/crates/aether-ai-formats/src/formats/conversion/request.rs @@ -165,6 +165,7 @@ mod tests { 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, }; @@ -219,6 +220,43 @@ 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"; @@ -333,6 +371,34 @@ mod tests { 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!({ 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/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_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] From ada690624bbc383397ff994ba8d7d08bd773c6b0 Mon Sep 17 00:00:00 2001 From: stabey Date: Thu, 4 Jun 2026 11:06:36 +0800 Subject: [PATCH 27/29] fix(usage): render responses call items in conversation view --- .../conversation/__tests__/stream.spec.ts | 133 ++++++++++++ .../src/features/usage/conversation/openai.ts | 189 +++++++++++++++--- 2 files changed, 292 insertions(+), 30 deletions(-) diff --git a/frontend/src/features/usage/conversation/__tests__/stream.spec.ts b/frontend/src/features/usage/conversation/__tests__/stream.spec.ts index c449c10bc..7c4df9966 100644 --- a/frontend/src/features/usage/conversation/__tests__/stream.spec.ts +++ b/frontend/src/features/usage/conversation/__tests__/stream.spec.ts @@ -110,6 +110,139 @@ 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('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..0c6937a1f 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,26 @@ export class OpenAIParser implements ApiFormatParser { result.model = response.model } - // 从 output 中提取文本（备用方案） - if (textParts.length === 0 && Array.isArray(response?.output)) { + // 从 output 中提取文本和工具调用（备用方案） + if (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)) { + 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) + ensureToolCall( + itemId || toolId, + toolId, + this.responsesCallName(item), + this.responsesCallInput(item) + ) } } } @@ -731,6 +806,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 +1005,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 +1133,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')], + })) } } From 2de2a792f6a2c63f9c86177db6cb4bb1d677446f Mon Sep 17 00:00:00 2001 From: stabey Date: Thu, 4 Jun 2026 12:01:36 +0800 Subject: [PATCH 28/29] fix(codex): strip cache_control before responses upstream --- .../src/formats/openai/responses/codex.rs | 61 +++++++++++++++++++ 1 file changed, 61 insertions(+) diff --git a/crates/aether-ai-formats/src/formats/openai/responses/codex.rs b/crates/aether-ai-formats/src/formats/openai/responses/codex.rs index 0a5766aea..1b4a9fad5 100644 --- a/crates/aether-ai-formats/src/formats/openai/responses/codex.rs +++ b/crates/aether-ai-formats/src/formats/openai/responses/codex.rs @@ -398,6 +398,23 @@ fn collect_codex_prompt_cache_control_anchors(value: &Value, anchors: &mut Vec { + 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] From 9f19ca575438f070ee9ff900981e76292d3eb35e Mon Sep 17 00:00:00 2001 From: stabey <36232531+stabey@users.noreply.github.com> Date: Thu, 4 Jun 2026 13:20:47 +0800 Subject: [PATCH 29/29] fix(usage): keep streamed call args in responses completion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The response.completed fallback rebuilt every call item with responsesCallInput(), which returns '{}' for a function_call lacking arguments. Since '{}' is truthy, ensureToolCall overwrote arguments already collected from streamed delta events. Guard the completed branch with responsesCallHasInput (matching the output_item.done branch) so empty/default inputs no longer clobber streamed args, and align its dedupe key with the streaming phase to avoid duplicate tool-call rendering when an item has no id. Drop the now-dead '工具调用' fallbacks since responsesCallName never returns empty. --- .../conversation/__tests__/stream.spec.ts | 39 +++++++++++++++++++ .../src/features/usage/conversation/openai.ts | 23 ++++++----- 2 files changed, 52 insertions(+), 10 deletions(-) diff --git a/frontend/src/features/usage/conversation/__tests__/stream.spec.ts b/frontend/src/features/usage/conversation/__tests__/stream.spec.ts index 7c4df9966..e1b579413 100644 --- a/frontend/src/features/usage/conversation/__tests__/stream.spec.ts +++ b/frontend/src/features/usage/conversation/__tests__/stream.spec.ts @@ -243,6 +243,45 @@ describe('Conversation stream compatibility', () => { }) }) + 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 0c6937a1f..c575f489e 100644 --- a/frontend/src/features/usage/conversation/openai.ts +++ b/frontend/src/features/usage/conversation/openai.ts @@ -698,8 +698,9 @@ export class OpenAIParser implements ApiFormatParser { // 从 output 中提取文本和工具调用（备用方案） if (Array.isArray(response?.output)) { - for (const rawItem of response.output as unknown[]) { - const item = rawItem as RawObject + 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 @@ -710,12 +711,14 @@ export class OpenAIParser implements ApiFormatParser { } else if (this.isResponsesCallItemType(item.type)) { const itemId = typeof item.id === 'string' ? item.id : '' const toolId = this.responsesCallId(item) - ensureToolCall( - itemId || toolId, - toolId, - this.responsesCallName(item), - this.responsesCallInput(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) } } } @@ -1007,7 +1010,7 @@ export class OpenAIParser implements ApiFormatParser { // Responses API call item -> 工具调用 if (this.isResponsesCallItemType(itemType)) { - const toolName = this.responsesCallName(item) || '工具调用' + const toolName = this.responsesCallName(item) const args = this.formatJson(this.responsesCallInput(item)) return createMessageBlock('assistant', [ createToolUseRenderBlock(toolName, args, this.responsesCallId(item)), @@ -1136,7 +1139,7 @@ export class OpenAIParser implements ApiFormatParser { } else if (this.isResponsesCallItemType(item.type)) { blocks.push(createMessageBlock('assistant', [ createToolUseRenderBlock( - this.responsesCallName(item) || '工具调用', + this.responsesCallName(item), this.formatJson(this.responsesCallInput(item)), this.responsesCallId(item) ),