i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor ai serving modules and crates

Browse Source
This commit is contained in:
fawney19
2026-05-02 13:23:54 +08:00
parent 4fc7cecf30
commit c130d0e2c9
309 changed files with 21549 additions and 14265 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
crates/aether-ai-formats/src/conversion/mod.rs
View File

@@ -1,9 +1,8 @@
//! Legacy pairwise conversion exports.
//! Pairwise format conversion entry points.
//!
//! Primary routing lives under `crate::formats::<wire_format>` and must pass
//! through canonical IR. This module remains so older pipeline/gateway call
//! sites and focused golden tests can keep their existing function names while
//! the cleanup proceeds.
//! The public helpers in this module route through the registry, so request and
//! response conversion still pass through the typed canonical IR before a target
//! wire format is emitted.
pub mod request;
pub mod response;

19
crates/aether-ai-formats/src/conversion/request/mod.rs
View File

@@ -1,13 +1,8 @@
//! Pairwise request adapters kept for compatibility and focused tests.
//! Pairwise request conversion helpers.
//!
//! New request routing should use the registry so every conversion passes
//! through the typed canonical IR.
//! Legacy request conversion function names.
//!
//! This module is intentionally a compatibility facade. Real wire-format
//! parsing and emitting lives under `formats::<format>::request`, and all
//! conversion goes through the registry's canonical IR path.
//! These helpers keep the call sites readable while delegating wire-format
//! parsing and emitting to `formats::<format>::request` through the registry's
//! canonical IR path.
use serde_json::Value;
@@ -171,7 +166,7 @@ mod tests {
};
#[test]
fn legacy_request_facade_routes_through_registry() {
fn pairwise_request_helper_routes_through_registry() {
let body = json!({
"model": "gpt-source",
"messages": [{"role": "user", "content": "hello"}],
@@ -191,7 +186,7 @@ mod tests {
}
#[test]
fn legacy_request_facade_keeps_claude_alias_shape() {
fn pairwise_request_helper_keeps_claude_shape() {
let body = json!({
"model": "gpt-source",
"messages": [{"role": "user", "content": "hello"}],
@@ -206,7 +201,7 @@ mod tests {
}
#[test]
fn legacy_normalizer_uses_format_adapter() {
fn request_normalizer_uses_format_adapter() {
let body = json!({
"model": "claude-sonnet",
"messages": [{"role": "user", "content": [{"type": "text", "text": "hello"}]}],

17
crates/aether-ai-formats/src/conversion/response/mod.rs
View File

@@ -1,13 +1,8 @@
//! Pairwise response adapters kept for compatibility and focused tests.
//! Pairwise response conversion helpers.
//!
//! New response routing should use the registry so every conversion passes
//! through the typed canonical IR.
//! Legacy response conversion function names.
//!
//! This module is intentionally a compatibility facade. Real wire-format
//! parsing and emitting lives under `formats::<format>::response`, and all
//! conversion goes through the registry's canonical IR path.
//! These helpers keep the call sites readable while delegating wire-format
//! parsing and emitting to `formats::<format>::response` through the registry's
//! canonical IR path.
use serde_json::{json, Value};
@@ -260,7 +255,7 @@ mod tests {
};
#[test]
fn legacy_response_facade_routes_through_registry() {
fn pairwise_response_helper_routes_through_registry() {
let body = json!({
"id": "chatcmpl-test",
"object": "chat.completion",
@@ -281,7 +276,7 @@ mod tests {
}
#[test]
fn legacy_response_facade_uses_report_context_model_fallback() {
fn pairwise_response_helper_uses_report_context_model_fallback() {
let body = json!({
"id": "msg-test",
"type": "message",

3
crates/aether-ai-pipeline/src/conversion/canonical.rs
View File

@@ -1,3 +0,0 @@
#![allow(deprecated)]
pub use aether_ai_formats::canonical::*;

40
crates/aether-ai-pipeline/src/conversion/mod.rs
View File

@@ -1,40 +0,0 @@
pub mod canonical;
mod error;
mod registry;
pub mod request;
pub mod response;
pub use aether_ai_formats::{
convert_request, convert_response, FormatContext, FormatError, FormatFamily, FormatId,
FormatProfile,
};
pub use canonical::{
canonical_request_unknown_block_count, canonical_response_unknown_block_count,
canonical_to_claude_request, canonical_to_claude_response, canonical_to_gemini_request,
canonical_to_gemini_response, canonical_to_openai_chat_request,
canonical_to_openai_chat_response, canonical_to_openai_responses_compact_request,
canonical_to_openai_responses_compact_response, canonical_to_openai_responses_request,
canonical_to_openai_responses_response, canonical_unknown_block_count,
from_claude_to_canonical_request, from_claude_to_canonical_response,
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, CanonicalGenerationConfig, CanonicalInstruction, CanonicalMessage,
CanonicalRequest, CanonicalResponse, CanonicalResponseFormat, CanonicalResponseOutput,
CanonicalRole, CanonicalStopReason, CanonicalStreamEvent, CanonicalStreamFrame,
CanonicalThinkingConfig, CanonicalToolChoice, CanonicalToolDefinition, CanonicalUsage,
};
pub use error::{
build_core_error_body_for_client_format, core_error_background_report_kind,
core_error_default_client_api_format, core_success_background_report_kind,
is_core_error_finalize_kind, LocalCoreSyncErrorKind,
};
pub use registry::{
request_candidate_api_format_preference, request_candidate_api_formats,
request_conversion_direct_auth, request_conversion_enabled_for_transport,
request_conversion_kind, request_conversion_requires_enable_flag,
request_conversion_transport_supported, request_conversion_transport_unsupported_reason,
request_pair_allowed_for_transport, sync_chat_response_conversion_kind,
sync_cli_response_conversion_kind, RequestConversionKind, SyncChatResponseConversionKind,
SyncCliResponseConversionKind,
};

861
crates/aether-ai-pipeline/src/conversion/registry.rs
View File

@@ -1,861 +0,0 @@
#![allow(dead_code)]
use aether_ai_formats::normalize_api_format_alias;
use aether_provider_transport::auth::{
resolve_local_gemini_auth, resolve_local_openai_bearer_auth, resolve_local_standard_auth,
};
use aether_provider_transport::kiro::local_kiro_request_transport_unsupported_reason_with_network;
use aether_provider_transport::policy::{
local_gemini_transport_unsupported_reason_with_network,
local_openai_chat_transport_unsupported_reason,
local_standard_transport_unsupported_reason_with_network,
};
use aether_provider_transport::vertex::{
is_vertex_api_key_transport_context,
local_vertex_api_key_gemini_transport_unsupported_reason_with_network,
resolve_local_vertex_api_key_query_auth, VERTEX_API_KEY_QUERY_PARAM,
};
use aether_provider_transport::GatewayProviderTransportSnapshot;
pub use aether_ai_formats::matrix::{
RequestConversionKind, SyncChatResponseConversionKind, SyncCliResponseConversionKind,
};
pub fn request_candidate_api_format_preference(
client_api_format: &str,
provider_api_format: &str,
) -> Option<(u8, u8)> {
aether_ai_formats::matrix::request_candidate_api_format_preference(
client_api_format,
provider_api_format,
)
}
pub fn request_candidate_api_formats(
client_api_format: &str,
require_streaming: bool,
) -> Vec<&'static str> {
aether_ai_formats::matrix::request_candidate_api_formats(client_api_format, require_streaming)
}
pub fn request_conversion_kind(
client_api_format: &str,
provider_api_format: &str,
) -> Option<RequestConversionKind> {
aether_ai_formats::matrix::request_conversion_kind(client_api_format, provider_api_format)
}
pub fn sync_chat_response_conversion_kind(
provider_api_format: &str,
client_api_format: &str,
) -> Option<SyncChatResponseConversionKind> {
aether_ai_formats::matrix::sync_chat_response_conversion_kind(
provider_api_format,
client_api_format,
)
}
pub fn sync_cli_response_conversion_kind(
provider_api_format: &str,
client_api_format: &str,
) -> Option<SyncCliResponseConversionKind> {
aether_ai_formats::matrix::sync_cli_response_conversion_kind(
provider_api_format,
client_api_format,
)
}
pub fn request_conversion_requires_enable_flag(
client_api_format: &str,
provider_api_format: &str,
) -> bool {
aether_ai_formats::matrix::request_conversion_requires_enable_flag(
client_api_format,
provider_api_format,
)
}
pub fn request_conversion_enabled_for_transport(
transport: &GatewayProviderTransportSnapshot,
client_api_format: &str,
provider_api_format: &str,
) -> bool {
let client_api_format = normalize_api_format_alias(client_api_format);
let provider_api_format = normalize_api_format_alias(provider_api_format);
if client_api_format == provider_api_format {
return true;
}
if request_conversion_kind(client_api_format.as_str(), provider_api_format.as_str()).is_none() {
return false;
}
if !request_conversion_requires_enable_flag(
client_api_format.as_str(),
provider_api_format.as_str(),
) {
return true;
}
transport.provider.enable_format_conversion
|| endpoint_accepts_client_api_format(transport, client_api_format.as_str())
}
pub fn request_pair_allowed_for_transport(
transport: &GatewayProviderTransportSnapshot,
client_api_format: &str,
provider_api_format: &str,
) -> bool {
let client_api_format = normalize_api_format_alias(client_api_format);
let provider_api_format = normalize_api_format_alias(provider_api_format);
if client_api_format == provider_api_format {
return true;
}
if request_conversion_kind(client_api_format.as_str(), provider_api_format.as_str()).is_none() {
return false;
}
if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("kiro")
&& aether_ai_formats::api_format_alias_matches(&provider_api_format, "claude:messages")
{
return request_conversion_enabled_for_transport(
transport,
client_api_format.as_str(),
provider_api_format.as_str(),
) && local_kiro_request_transport_unsupported_reason_with_network(transport)
.is_none();
}
request_conversion_enabled_for_transport(
transport,
client_api_format.as_str(),
provider_api_format.as_str(),
)
}
pub fn request_conversion_transport_supported(
transport: &GatewayProviderTransportSnapshot,
kind: RequestConversionKind,
) -> bool {
request_conversion_transport_unsupported_reason(transport, kind).is_none()
}
pub fn request_conversion_transport_unsupported_reason(
transport: &GatewayProviderTransportSnapshot,
_kind: RequestConversionKind,
) -> Option<&'static str> {
if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("kiro")
&& aether_ai_formats::api_format_alias_matches(
&transport.endpoint.api_format,
"claude:messages",
)
{
return local_kiro_request_transport_unsupported_reason_with_network(transport);
}
match normalize_api_format_alias(&transport.endpoint.api_format).as_str() {
"openai:chat" => local_openai_chat_transport_unsupported_reason(transport),
"openai:responses" | "openai:responses:compact" => {
local_standard_transport_unsupported_reason_with_network(
transport,
transport.endpoint.api_format.trim(),
)
}
"claude:messages" => {
local_standard_transport_unsupported_reason_with_network(transport, "claude:messages")
}
"gemini:generate_content" if is_vertex_api_key_transport_context(transport) => {
local_vertex_api_key_gemini_transport_unsupported_reason_with_network(transport)
}
"gemini:generate_content" => local_gemini_transport_unsupported_reason_with_network(
transport,
"gemini:generate_content",
),
_ => Some("transport_api_format_unsupported"),
}
}
pub fn request_conversion_direct_auth(
transport: &GatewayProviderTransportSnapshot,
_kind: RequestConversionKind,
) -> Option<(String, String)> {
match normalize_api_format_alias(&transport.endpoint.api_format).as_str() {
"openai:chat" | "openai:responses" | "openai:responses:compact" => {
resolve_local_openai_bearer_auth(transport)
}
"gemini:generate_content" => {
if is_vertex_api_key_transport_context(transport) {
resolve_local_vertex_api_key_query_auth(transport)
.map(|auth| (VERTEX_API_KEY_QUERY_PARAM.to_string(), auth.value))
} else {
resolve_local_gemini_auth(transport)
}
}
"claude:messages" => resolve_local_standard_auth(transport),
_ => None,
}
}
fn endpoint_accepts_client_api_format(
transport: &GatewayProviderTransportSnapshot,
client_api_format: &str,
) -> bool {
let Some(config) = transport
.endpoint
.format_acceptance_config
.as_ref()
.and_then(serde_json::Value::as_object)
else {
return false;
};
if !config
.get("enabled")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
{
return false;
}
if config
.get("reject_formats")
.is_some_and(|value| json_format_list_contains(value, client_api_format))
{
return false;
}
match config.get("accept_formats") {
Some(value) => json_format_list_contains(value, client_api_format),
None => true,
}
}
fn json_format_list_contains(value: &serde_json::Value, api_format: &str) -> bool {
let Some(items) = value.as_array() else {
return false;
};
items.iter().any(|item| {
item.as_str().is_some_and(|candidate| {
aether_ai_formats::api_format_alias_matches(candidate, api_format)
})
})
}
#[cfg(test)]
mod tests {
use super::{
request_candidate_api_format_preference, request_candidate_api_formats,
request_conversion_direct_auth, request_conversion_enabled_for_transport,
request_conversion_kind, request_conversion_requires_enable_flag,
request_conversion_transport_supported, request_pair_allowed_for_transport,
sync_chat_response_conversion_kind, sync_cli_response_conversion_kind,
RequestConversionKind, SyncChatResponseConversionKind, SyncCliResponseConversionKind,
};
use aether_provider_transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
const STANDARD_SURFACES: &[&str] = &[
"openai:chat",
"openai:responses",
"claude:messages",
"gemini:generate_content",
];
fn expected_request_conversion_kind(provider_api_format: &str) -> RequestConversionKind {
match provider_api_format {
"openai:chat" => RequestConversionKind::ToOpenAIChat,
"openai:responses" => RequestConversionKind::ToOpenAiResponses,
"claude:messages" => RequestConversionKind::ToClaudeStandard,
"gemini:generate_content" => RequestConversionKind::ToGeminiStandard,
other => panic!("unexpected provider api format: {other}"),
}
}
#[test]
fn request_conversion_registry_supports_bidirectional_standard_matrix() {
assert_eq!(
request_conversion_kind("openai:chat", "openai:responses"),
Some(RequestConversionKind::ToOpenAiResponses)
);
assert_eq!(
request_conversion_kind("openai:chat", "claude:messages"),
Some(RequestConversionKind::ToClaudeStandard)
);
assert_eq!(
request_conversion_kind("openai:responses", "openai:chat"),
Some(RequestConversionKind::ToOpenAIChat)
);
assert_eq!(
request_conversion_kind("openai:responses:compact", "gemini:generate_content"),
None
);
assert_eq!(
request_conversion_kind("gemini:generate_content", "openai:responses:compact"),
None
);
assert_eq!(
request_conversion_kind("openai:chat", "openai:responses:compact"),
None
);
assert_eq!(
request_conversion_kind("openai:responses", "openai:cli"),
None
);
assert_eq!(
request_conversion_kind("openai:compact", "openai:responses:compact"),
None
);
assert_eq!(
request_conversion_kind("gemini:generate_content", "claude:messages"),
Some(RequestConversionKind::ToClaudeStandard)
);
assert_eq!(request_conversion_kind("claude:chat", "claude:cli"), None);
assert_eq!(
request_conversion_kind("claude:messages", "claude:messages"),
None
);
}
#[test]
fn request_conversion_registry_covers_all_standard_surface_pairs() {
for client_api_format in STANDARD_SURFACES {
for provider_api_format in STANDARD_SURFACES {
let actual = request_conversion_kind(client_api_format, provider_api_format);
if client_api_format == provider_api_format {
assert_eq!(
actual, None,
"{client_api_format} -> {provider_api_format} should be same-format"
);
} else {
assert_eq!(
actual,
Some(expected_request_conversion_kind(provider_api_format)),
"{client_api_format} -> {provider_api_format} should be routable"
);
}
}
}
}
#[test]
fn sync_response_conversion_registry_supports_bidirectional_standard_matrix() {
assert_eq!(
sync_chat_response_conversion_kind("openai:chat", "claude:messages"),
Some(SyncChatResponseConversionKind::ToClaudeChat)
);
assert_eq!(
sync_chat_response_conversion_kind("claude:messages", "gemini:generate_content"),
Some(SyncChatResponseConversionKind::ToGeminiChat)
);
assert_eq!(
sync_chat_response_conversion_kind("gemini:generate_content", "openai:chat"),
Some(SyncChatResponseConversionKind::ToOpenAIChat)
);
assert_eq!(
sync_cli_response_conversion_kind("openai:responses", "gemini:generate_content"),
Some(SyncCliResponseConversionKind::ToGeminiCli)
);
assert_eq!(
sync_cli_response_conversion_kind("claude:messages", "openai:responses"),
Some(SyncCliResponseConversionKind::ToOpenAiResponses)
);
assert_eq!(
sync_cli_response_conversion_kind("claude:messages", "openai:responses:compact"),
Some(SyncCliResponseConversionKind::ToOpenAiResponses)
);
assert_eq!(
sync_cli_response_conversion_kind("openai:responses:compact", "claude:messages"),
None
);
assert_eq!(
sync_cli_response_conversion_kind("gemini:generate_content", "claude:messages"),
Some(SyncCliResponseConversionKind::ToClaudeCli)
);
assert_eq!(
sync_cli_response_conversion_kind("openai:responses", "openai:cli"),
None
);
assert_eq!(
sync_cli_response_conversion_kind("openai:compact", "openai:responses:compact"),
None
);
}
#[test]
fn sync_response_conversion_registry_covers_all_standard_surface_pairs() {
for provider_api_format in STANDARD_SURFACES {
for client_api_format in ["openai:chat", "claude:messages", "gemini:generate_content"] {
let actual =
sync_chat_response_conversion_kind(provider_api_format, client_api_format);
if *provider_api_format == client_api_format {
assert_eq!(
actual, None,
"{provider_api_format} -> {client_api_format} should be same-format"
);
} else {
let expected = match client_api_format {
"openai:chat" => SyncChatResponseConversionKind::ToOpenAIChat,
"claude:messages" => SyncChatResponseConversionKind::ToClaudeChat,
"gemini:generate_content" => SyncChatResponseConversionKind::ToGeminiChat,
other => panic!("unexpected chat client api format: {other}"),
};
assert_eq!(
actual,
Some(expected),
"{provider_api_format} -> {client_api_format} should finalize to chat"
);
}
}
for client_api_format in [
"openai:responses",
"claude:messages",
"gemini:generate_content",
] {
let actual =
sync_cli_response_conversion_kind(provider_api_format, client_api_format);
if *provider_api_format == client_api_format {
assert_eq!(
actual, None,
"{provider_api_format} -> {client_api_format} should be same-format"
);
} else {
let expected = match client_api_format {
"openai:responses" => SyncCliResponseConversionKind::ToOpenAiResponses,
"claude:messages" => SyncCliResponseConversionKind::ToClaudeCli,
"gemini:generate_content" => SyncCliResponseConversionKind::ToGeminiCli,
other => panic!("unexpected cli client api format: {other}"),
};
assert_eq!(
actual,
Some(expected),
"{provider_api_format} -> {client_api_format} should finalize to cli"
);
}
}
}
}
#[test]
fn request_candidate_registry_excludes_compact_as_cross_format_target() {
assert_eq!(
request_candidate_api_formats("openai:chat", false),
vec![
"openai:chat",
"openai:responses",
"claude:messages",
"gemini:generate_content",
]
);
assert_eq!(
request_candidate_api_formats("openai:responses", false),
vec![
"openai:responses",
"openai:chat",
"claude:messages",
"gemini:generate_content",
]
);
assert_eq!(
request_candidate_api_formats("claude:messages", false),
vec![
"claude:messages",
"openai:chat",
"openai:responses",
"gemini:generate_content",
]
);
assert_eq!(
request_candidate_api_formats("openai:cli", false),
Vec::<&'static str>::new()
);
assert_eq!(
request_candidate_api_formats("claude:cli", false),
Vec::<&'static str>::new()
);
assert_eq!(
request_candidate_api_formats("openai:compact", false),
Vec::<&'static str>::new()
);
}
#[test]
fn request_candidate_registry_prefers_same_kind_before_same_family_fallbacks() {
assert_eq!(
request_candidate_api_format_preference("claude:cli", "openai:responses"),
None
);
assert_eq!(
request_candidate_api_format_preference("claude:cli", "claude:chat"),
None
);
assert_eq!(
request_candidate_api_format_preference("claude:cli", "openai:chat"),
None
);
}
#[test]
fn request_conversion_enable_flag_only_applies_to_real_data_format_conversions() {
assert!(!request_conversion_requires_enable_flag(
"claude:messages",
"claude:messages"
));
assert!(request_conversion_requires_enable_flag(
"claude:chat",
"claude:cli"
));
assert!(request_conversion_requires_enable_flag(
"openai:chat",
"openai:responses"
));
assert!(request_conversion_requires_enable_flag(
"openai:responses",
"openai:chat"
));
assert!(request_conversion_requires_enable_flag(
"openai:chat",
"gemini:generate_content"
));
}
#[test]
fn request_conversion_helpers_follow_transport_api_format() {
let transport = GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "provider".to_string(),
provider_type: "openai".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: true,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "openai:chat".to_string(),
api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
is_active: true,
base_url: "https://api.openai.com".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "bearer".to_string(),
is_active: true,
api_formats: None,
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},
};
assert!(request_conversion_transport_supported(
&transport,
RequestConversionKind::ToOpenAIChat
));
assert_eq!(
request_conversion_direct_auth(&transport, RequestConversionKind::ToOpenAIChat),
Some(("authorization".to_string(), "Bearer secret".to_string()))
);
}
#[test]
fn endpoint_level_format_acceptance_enables_cross_format_pair_without_provider_flag() {
let transport = GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "provider".to_string(),
provider_type: "custom".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "openai:responses".to_string(),
api_family: Some("openai".to_string()),
endpoint_kind: Some("responses".to_string()),
is_active: true,
base_url: "https://right.codes/codex".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: Some("/v1/messages".to_string()),
config: None,
format_acceptance_config: Some(serde_json::json!({
"enabled": true,
"accept_formats": ["claude:messages"],
})),
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "bearer".to_string(),
is_active: true,
api_formats: Some(vec!["openai:responses".to_string()]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},
};
assert!(request_conversion_enabled_for_transport(
&transport,
"claude:messages",
"openai:responses"
));
assert!(request_pair_allowed_for_transport(
&transport,
"claude:messages",
"openai:responses"
));
assert!(!request_pair_allowed_for_transport(
&transport,
"gemini:generate_content",
"openai:responses"
));
}
#[test]
fn endpoint_reject_formats_override_endpoint_cross_format_enablement() {
let transport = GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "provider".to_string(),
provider_type: "custom".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "openai:responses".to_string(),
api_family: Some("openai".to_string()),
endpoint_kind: Some("responses".to_string()),
is_active: true,
base_url: "https://right.codes/codex".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: Some("/v1/messages".to_string()),
config: None,
format_acceptance_config: Some(serde_json::json!({
"enabled": true,
"reject_formats": ["claude:messages"],
})),
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "bearer".to_string(),
is_active: true,
api_formats: Some(vec!["openai:responses".to_string()]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},
};
assert!(!request_conversion_enabled_for_transport(
&transport,
"claude:messages",
"openai:responses"
));
}
#[test]
fn vertex_gemini_transport_supports_cross_format_conversion_with_query_auth() {
let transport = GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-vertex".to_string(),
name: "vertex".to_string(),
provider_type: "vertex_ai".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: true,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-vertex".to_string(),
provider_id: "provider-vertex".to_string(),
api_format: "gemini:generate_content".to_string(),
api_family: Some("gemini".to_string()),
endpoint_kind: Some("generate_content".to_string()),
is_active: true,
base_url: "https://aiplatform.googleapis.com".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-vertex".to_string(),
provider_id: "provider-vertex".to_string(),
name: "key".to_string(),
auth_type: "api_key".to_string(),
is_active: true,
api_formats: Some(vec!["gemini:generate_content".to_string()]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "vertex-secret".to_string(),
decrypted_auth_config: None,
},
};
assert!(request_conversion_transport_supported(
&transport,
RequestConversionKind::ToGeminiStandard
));
assert_eq!(
request_conversion_direct_auth(&transport, RequestConversionKind::ToGeminiStandard),
Some(("key".to_string(), "vertex-secret".to_string()))
);
}
#[test]
fn kiro_claude_messages_transport_supports_cross_format_conversion_via_envelope() {
let transport = GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-kiro".to_string(),
name: "kiro".to_string(),
provider_type: "kiro".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: true,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-kiro".to_string(),
provider_id: "provider-kiro".to_string(),
api_format: "claude:messages".to_string(),
api_family: Some("claude".to_string()),
endpoint_kind: Some("messages".to_string()),
is_active: true,
base_url: "https://q.{region}.amazonaws.com".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-kiro".to_string(),
provider_id: "provider-kiro".to_string(),
name: "key".to_string(),
auth_type: "bearer".to_string(),
is_active: true,
api_formats: Some(vec!["claude:messages".to_string()]),
auth_type_by_format: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "kiro-secret".to_string(),
decrypted_auth_config: None,
},
};
assert!(request_pair_allowed_for_transport(
&transport,
"openai:chat",
"claude:messages"
));
assert!(request_conversion_transport_supported(
&transport,
RequestConversionKind::ToClaudeStandard
));
}
}

1
crates/aether-ai-pipeline/src/conversion/request/mod.rs
View File

@@ -1 +0,0 @@
pub use aether_ai_formats::conversion::request::*;

1
crates/aether-ai-pipeline/src/conversion/response/mod.rs
View File

@@ -1 +0,0 @@
pub use aether_ai_formats::conversion::response::*;

182
crates/aether-ai-pipeline/src/finalize/stream_rewrite.rs
View File

@@ -1,182 +0,0 @@
use serde_json::Value;
use crate::adaptation::surfaces::{
provider_adaptation_should_unwrap_stream_envelope, KIRO_ENVELOPE_NAME,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FinalizeStreamRewriteMode {
EnvelopeUnwrap,
OpenAiImage,
Standard,
KiroToClaudeCli,
KiroToClaudeCliThenStandard,
}
pub fn resolve_finalize_stream_rewrite_mode(
report_context: &Value,
) -> Option<FinalizeStreamRewriteMode> {
let needs_conversion = report_context
.get("needs_conversion")
.and_then(Value::as_bool)
.unwrap_or(false);
let envelope_name = report_context
.get("envelope_name")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
let provider_api_format = report_context
.get("provider_api_format")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
let client_api_format = report_context
.get("client_api_format")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
if needs_conversion
&& envelope_name.eq_ignore_ascii_case(KIRO_ENVELOPE_NAME)
&& provider_api_format == "claude:messages"
{
return supports_standard_stream_rewrite(
provider_api_format.as_str(),
client_api_format.as_str(),
)
.then_some(FinalizeStreamRewriteMode::KiroToClaudeCliThenStandard);
}
if needs_conversion {
return supports_standard_stream_rewrite(
provider_api_format.as_str(),
client_api_format.as_str(),
)
.then_some(FinalizeStreamRewriteMode::Standard);
}
if provider_api_format == "openai:image" && client_api_format == "openai:image" {
return Some(FinalizeStreamRewriteMode::OpenAiImage);
}
if envelope_name.eq_ignore_ascii_case(KIRO_ENVELOPE_NAME) {
return (provider_api_format == "claude:messages"
&& client_api_format == "claude:messages")
.then_some(FinalizeStreamRewriteMode::KiroToClaudeCli);
}
(provider_api_format == client_api_format
&& provider_adaptation_should_unwrap_stream_envelope(
envelope_name.as_str(),
provider_api_format.as_str(),
))
.then_some(FinalizeStreamRewriteMode::EnvelopeUnwrap)
}
fn supports_standard_stream_rewrite(provider_api_format: &str, client_api_format: &str) -> bool {
is_standard_provider_api_format(provider_api_format)
&& (is_standard_chat_client_api_format(client_api_format)
|| is_standard_cli_client_api_format(client_api_format))
}
fn is_standard_provider_api_format(api_format: &str) -> bool {
matches!(
aether_ai_formats::normalize_api_format_alias(api_format).as_str(),
"openai:chat"
| "openai:responses"
| "openai:responses:compact"
| "claude:messages"
| "gemini:generate_content"
)
}
fn is_standard_chat_client_api_format(api_format: &str) -> bool {
matches!(
api_format,
"openai:chat" | "claude:messages" | "gemini:generate_content"
)
}
fn is_standard_cli_client_api_format(api_format: &str) -> bool {
matches!(
aether_ai_formats::normalize_api_format_alias(api_format).as_str(),
"openai:responses"
| "openai:responses:compact"
| "claude:messages"
| "gemini:generate_content"
)
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{resolve_finalize_stream_rewrite_mode, FinalizeStreamRewriteMode};
#[test]
fn resolves_standard_mode_for_cross_format_standard_streams() {
let report_context = json!({
"provider_api_format": "claude:messages",
"client_api_format": "openai:chat",
"needs_conversion": true,
});
assert_eq!(
resolve_finalize_stream_rewrite_mode(&report_context),
Some(FinalizeStreamRewriteMode::Standard)
);
}
#[test]
fn resolves_envelope_unwrap_for_same_format_private_envelopes() {
let report_context = json!({
"provider_api_format": "gemini:generate_content",
"client_api_format": "gemini:generate_content",
"envelope_name": "antigravity:v1internal",
"needs_conversion": false,
});
assert_eq!(
resolve_finalize_stream_rewrite_mode(&report_context),
Some(FinalizeStreamRewriteMode::EnvelopeUnwrap)
);
}
#[test]
fn resolves_kiro_same_format_streams_to_kiro_mode() {
let report_context = json!({
"provider_api_format": "claude:messages",
"client_api_format": "claude:messages",
"envelope_name": "kiro:generateAssistantResponse",
"needs_conversion": false,
});
assert_eq!(
resolve_finalize_stream_rewrite_mode(&report_context),
Some(FinalizeStreamRewriteMode::KiroToClaudeCli)
);
}
#[test]
fn rejects_unsupported_non_conversion_streams() {
let report_context = json!({
"provider_api_format": "openai:chat",
"client_api_format": "openai:chat",
"needs_conversion": false,
});
assert_eq!(resolve_finalize_stream_rewrite_mode(&report_context), None);
}
#[test]
fn resolves_openai_image_mode_for_same_format_image_streams() {
let report_context = json!({
"provider_api_format": "openai:image",
"client_api_format": "openai:image",
"needs_conversion": false,
});
assert_eq!(
resolve_finalize_stream_rewrite_mode(&report_context),
Some(FinalizeStreamRewriteMode::OpenAiImage)
);
}
}

54
crates/aether-ai-pipeline/src/planner/specialized/image.rs
View File

@@ -1,54 +0,0 @@
use crate::contracts::{OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND};
#[derive(Debug, Clone, Copy)]
pub struct LocalOpenAiImageSpec {
pub api_format: &'static str,
pub decision_kind: &'static str,
pub report_kind: &'static str,
pub require_streaming: bool,
}
pub fn resolve_sync_spec(plan_kind: &str) -> Option<LocalOpenAiImageSpec> {
match plan_kind {
OPENAI_IMAGE_SYNC_PLAN_KIND => Some(LocalOpenAiImageSpec {
api_format: "openai:image",
decision_kind: OPENAI_IMAGE_SYNC_PLAN_KIND,
report_kind: "openai_image_sync_finalize",
require_streaming: false,
}),
_ => None,
}
}
pub fn resolve_stream_spec(plan_kind: &str) -> Option<LocalOpenAiImageSpec> {
match plan_kind {
OPENAI_IMAGE_STREAM_PLAN_KIND => Some(LocalOpenAiImageSpec {
api_format: "openai:image",
decision_kind: OPENAI_IMAGE_STREAM_PLAN_KIND,
report_kind: "openai_image_stream_success",
require_streaming: true,
}),
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::{resolve_stream_spec, resolve_sync_spec};
#[test]
fn resolves_openai_image_sync_spec() {
let spec = resolve_sync_spec("openai_image_sync").expect("spec");
assert_eq!(spec.api_format, "openai:image");
assert_eq!(spec.report_kind, "openai_image_sync_finalize");
assert!(!spec.require_streaming);
}
#[test]
fn resolves_openai_image_stream_spec() {
let spec = resolve_stream_spec("openai_image_stream").expect("spec");
assert_eq!(spec.api_format, "openai:image");
assert_eq!(spec.report_kind, "openai_image_stream_success");
assert!(spec.require_streaming);
}
}

57
crates/aether-ai-pipeline/src/transport.rs
View File

@@ -1,57 +0,0 @@
pub mod antigravity {
pub use aether_provider_transport::antigravity::*;
}
pub mod auth {
pub use aether_provider_transport::auth::*;
}
pub mod claude_code {
pub use aether_provider_transport::claude_code::*;
}
pub mod kiro {
pub use aether_provider_transport::kiro::*;
}
pub mod oauth_refresh {
pub use aether_provider_transport::oauth_refresh::*;
}
pub mod policy {
pub use aether_provider_transport::policy::*;
}
pub mod provider_types {
pub use aether_provider_transport::provider_types::*;
}
pub mod rules {
pub use aether_provider_transport::rules::*;
}
pub mod snapshot {
pub use aether_provider_transport::snapshot::*;
}
pub mod url {
pub use aether_provider_transport::url::*;
}
pub mod vertex {
pub use aether_provider_transport::vertex::*;
}
pub use aether_provider_transport::{
apply_local_body_rules, apply_local_header_rules, body_rules_are_locally_supported,
body_rules_handle_path, build_passthrough_headers, ensure_upstream_auth_header,
header_rules_are_locally_supported, local_gemini_transport_unsupported_reason_with_network,
local_openai_chat_transport_unsupported_reason,
local_standard_transport_unsupported_reason_with_network, resolve_transport_execution_timeouts,
resolve_transport_proxy_snapshot, resolve_transport_proxy_snapshot_with_tunnel_affinity,
resolve_transport_tls_profile, should_skip_upstream_passthrough_header,
supports_local_gemini_transport_with_network,
supports_local_generic_oauth_request_auth_resolution,
supports_local_oauth_request_auth_resolution, transport_proxy_is_locally_supported,
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
};

21
crates/aether-ai-serving/Cargo.toml Normal file
View File

@@ -0,0 +1,21 @@
[package]
name = "aether-ai-serving"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "AI serving application contracts and ports for Aether"
[dependencies]
aether-ai-formats.workspace = true
aether-ai-surfaces.workspace = true
aether-contracts.workspace = true
aether-scheduler-core.workspace = true
async-trait.workspace = true
http.workspace = true
serde.workspace = true
serde_json.workspace = true
url.workspace = true
[dev-dependencies]
tokio.workspace = true

96
crates/aether-ai-serving/src/attempt_loop.rs Normal file
View File

@@ -0,0 +1,96 @@
use async_trait::async_trait;
pub trait AiExecutionAttempt {
fn execution_plan(&self) -> &aether_contracts::ExecutionPlan;
fn report_kind(&self) -> Option<String>;
fn report_context(&self) -> Option<serde_json::Value>;
}
#[derive(Debug)]
pub enum AiAttemptLoopOutcome<Response, Exhaustion> {
Responded(Response),
Exhausted(Exhaustion),
NoPath,
}
#[async_trait]
pub trait AiAttemptLoopPort<Attempt>: Send + Sync
where
Attempt: AiExecutionAttempt + Send + Sync + 'static,
{
type Response: Send;
type Exhaustion: Send;
type Error: Send;
async fn execute_attempt(
&self,
attempt: &Attempt,
) -> Result<Option<Self::Response>, Self::Error>;
async fn mark_unused_attempts(&self, attempts: Vec<Attempt>) -> Result<(), Self::Error>;
async fn build_exhaustion(
&self,
last_plan: aether_contracts::ExecutionPlan,
last_report_context: Option<serde_json::Value>,
) -> Result<Self::Exhaustion, Self::Error>;
}
pub async fn run_ai_attempt_loop<Port, Attempt>(
port: &Port,
attempts: Vec<Attempt>,
) -> Result<AiAttemptLoopOutcome<Port::Response, Port::Exhaustion>, Port::Error>
where
Port: AiAttemptLoopPort<Attempt>,
Attempt: AiExecutionAttempt + Send + Sync + 'static,
{
let mut remaining = attempts.into_iter();
let mut last_attempted = None;
while let Some(attempt) = remaining.next() {
last_attempted = Some((attempt.execution_plan().clone(), attempt.report_context()));
if let Some(response) = port.execute_attempt(&attempt).await? {
port.mark_unused_attempts(remaining.collect()).await?;
return Ok(AiAttemptLoopOutcome::Responded(response));
}
}
let Some((last_plan, last_report_context)) = last_attempted else {
return Ok(AiAttemptLoopOutcome::NoPath);
};
Ok(AiAttemptLoopOutcome::Exhausted(
port.build_exhaustion(last_plan, last_report_context)
.await?,
))
}
impl AiExecutionAttempt for crate::dto::AiSyncAttempt {
fn execution_plan(&self) -> &aether_contracts::ExecutionPlan {
&self.plan
}
fn report_kind(&self) -> Option<String> {
self.report_kind.clone()
}
fn report_context(&self) -> Option<serde_json::Value> {
self.report_context.clone()
}
}
impl AiExecutionAttempt for crate::dto::AiStreamAttempt {
fn execution_plan(&self) -> &aether_contracts::ExecutionPlan {
&self.plan
}
fn report_kind(&self) -> Option<String> {
self.report_kind.clone()
}
fn report_context(&self) -> Option<serde_json::Value> {
self.report_context.clone()
}
}

523
crates/aether-ai-serving/src/attempt_plan.rs Normal file
View File

@@ -0,0 +1,523 @@
use std::collections::BTreeMap;
use aether_ai_surfaces::api::ExecutionRuntimeAuthContext;
use aether_contracts::{ExecutionPlan, RequestBody};
use url::Url;
use crate::dto::AiExecutionDecision;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AiDecisionPlanCore {
pub request_id: String,
pub provider_id: String,
pub endpoint_id: String,
pub key_id: String,
pub provider_api_format: String,
pub client_api_format: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AiUpstreamAuthPair {
pub header: String,
pub value: String,
}
#[derive(Debug)]
pub struct AiExecutionPlanFromDecisionParts {
pub core: AiDecisionPlanCore,
pub method: String,
pub url: String,
pub headers: BTreeMap<String, String>,
pub content_type: Option<String>,
pub body: RequestBody,
pub stream: bool,
}
#[derive(Debug)]
pub struct AiExecutionDecisionFromPlanParts {
pub action: String,
pub decision_kind: Option<String>,
pub request_id: Option<String>,
pub upstream_base_url: Option<String>,
pub include_auth_pair: bool,
pub plan: ExecutionPlan,
pub report_kind: Option<String>,
pub report_context: Option<serde_json::Value>,
pub auth_context: Option<ExecutionRuntimeAuthContext>,
}
pub fn take_ai_non_empty_string(value: &mut Option<String>) -> Option<String> {
value.take().filter(|value| !value.trim().is_empty())
}
pub fn trim_ai_owned_non_empty_string(value: String) -> Option<String> {
let trimmed = value.trim();
if trimmed.is_empty() {
return None;
}
if trimmed.len() == value.len() {
return Some(value);
}
Some(trimmed.to_owned())
}
pub fn take_ai_decision_plan_core(payload: &mut AiExecutionDecision) -> Option<AiDecisionPlanCore> {
Some(AiDecisionPlanCore {
request_id: take_ai_non_empty_string(&mut payload.request_id)?,
provider_id: take_ai_non_empty_string(&mut payload.provider_id)?,
endpoint_id: take_ai_non_empty_string(&mut payload.endpoint_id)?,
key_id: take_ai_non_empty_string(&mut payload.key_id)?,
provider_api_format: take_ai_non_empty_string(&mut payload.provider_api_format)?,
client_api_format: take_ai_non_empty_string(&mut payload.client_api_format)?,
})
}
pub fn take_ai_upstream_auth_pair(
payload: &mut AiExecutionDecision,
) -> Option<Option<AiUpstreamAuthPair>> {
let header = take_ai_non_empty_string(&mut payload.auth_header);
let value = take_ai_non_empty_string(&mut payload.auth_value);
match (header, value) {
(Some(header), Some(value)) => Some(Some(AiUpstreamAuthPair { header, value })),
(None, None) => Some(None),
_ => None,
}
}
pub fn resolve_ai_passthrough_sync_request_body(
provider_request_body: Option<serde_json::Value>,
provider_request_body_base64: Option<String>,
) -> RequestBody {
if let Some(body_bytes_b64) =
provider_request_body_base64.and_then(trim_ai_owned_non_empty_string)
{
return RequestBody {
json_body: None,
body_bytes_b64: Some(body_bytes_b64),
body_ref: None,
};
}
match provider_request_body.unwrap_or(serde_json::Value::Null) {
serde_json::Value::Null => RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
},
other => RequestBody::from_json(other),
}
}
pub fn build_ai_execution_plan_from_decision(
payload: &mut AiExecutionDecision,
parts: AiExecutionPlanFromDecisionParts,
) -> ExecutionPlan {
ExecutionPlan {
request_id: parts.core.request_id,
candidate_id: payload.candidate_id.take(),
provider_name: payload.provider_name.take(),
provider_id: parts.core.provider_id,
endpoint_id: parts.core.endpoint_id,
key_id: parts.core.key_id,
method: parts.method,
url: parts.url,
headers: parts.headers,
content_type: parts.content_type,
content_encoding: None,
body: parts.body,
stream: parts.stream,
client_api_format: parts.core.client_api_format,
provider_api_format: parts.core.provider_api_format,
model_name: payload.model_name.take(),
proxy: payload.proxy.take(),
tls_profile: payload.tls_profile.take(),
timeouts: payload.timeouts.take(),
}
}
pub fn build_ai_execution_decision_from_plan(
parts: AiExecutionDecisionFromPlanParts,
) -> AiExecutionDecision {
let ExecutionPlan {
request_id,
candidate_id,
provider_name,
provider_id,
endpoint_id,
key_id,
method,
url,
headers,
content_type,
content_encoding: _content_encoding,
body,
stream,
client_api_format,
provider_api_format,
model_name,
proxy,
tls_profile,
timeouts,
} = parts.plan;
let auth_pair = parts
.include_auth_pair
.then(|| extract_ai_auth_header_pair(&headers))
.flatten();
let provider_contract = provider_api_format.clone();
let client_contract = client_api_format.clone();
let request_id = parts.request_id.unwrap_or(request_id);
let auth_header = auth_pair.map(|(name, _)| name.to_string());
let auth_value = auth_pair.map(|(_, value)| value.to_string());
let RequestBody {
json_body,
body_bytes_b64,
body_ref: _body_ref,
} = body;
AiExecutionDecision {
action: parts.action,
decision_kind: parts.decision_kind,
execution_strategy: Some(ai_execution_strategy_for_formats(
provider_api_format.as_str(),
client_api_format.as_str(),
)),
conversion_mode: Some(ai_conversion_mode_for_formats(
provider_api_format.as_str(),
client_api_format.as_str(),
)),
request_id: Some(request_id),
candidate_id,
provider_name,
provider_id: Some(provider_id),
endpoint_id: Some(endpoint_id),
key_id: Some(key_id),
upstream_base_url: parts.upstream_base_url,
upstream_url: Some(url),
provider_request_method: Some(method),
auth_header,
auth_value,
provider_api_format: Some(provider_api_format),
client_api_format: Some(client_api_format),
provider_contract: Some(provider_contract),
client_contract: Some(client_contract),
model_name,
mapped_model: None,
prompt_cache_key: None,
extra_headers: BTreeMap::new(),
provider_request_headers: headers,
provider_request_body: json_body,
provider_request_body_base64: body_bytes_b64,
content_type,
proxy,
tls_profile,
timeouts,
upstream_is_stream: stream,
report_kind: parts.report_kind,
report_context: parts.report_context,
auth_context: parts.auth_context,
}
}
pub fn extract_ai_auth_header_pair(headers: &BTreeMap<String, String>) -> Option<(&str, &str)> {
[
"authorization",
"x-api-key",
"api-key",
"x-goog-api-key",
"proxy-authorization",
]
.into_iter()
.find_map(|name| {
headers
.iter()
.find(|(header_name, _)| header_name.eq_ignore_ascii_case(name))
.map(|(header_name, value)| (header_name.as_str(), value.as_str()))
})
}
pub fn infer_ai_upstream_base_url(upstream_url: &str) -> Option<String> {
let parsed = Url::parse(upstream_url).ok()?;
let host = parsed.host_str()?;
let mut base = format!("{}://{}", parsed.scheme(), host);
if let Some(port) = parsed.port() {
base.push(':');
base.push_str(port.to_string().as_str());
}
let base_path = infer_ai_upstream_base_path(parsed.path());
if !base_path.is_empty() {
base.push_str(base_path);
}
Some(base)
}
fn infer_ai_upstream_base_path(path: &str) -> &str {
let trimmed = path.trim_end_matches('/');
if trimmed.is_empty() || trimmed == "/" {
return "";
}
for suffix in [
"/responses/compact",
"/responses",
"/chat/completions",
"/messages",
] {
if let Some(prefix) = trimmed.strip_suffix(suffix) {
return normalize_inferred_ai_base_path(prefix);
}
}
for marker in ["/v1/videos", "/v1beta/"] {
if let Some((prefix, _)) = trimmed.split_once(marker) {
return normalize_inferred_ai_base_path(prefix);
}
}
normalize_inferred_ai_base_path(trimmed)
}
fn normalize_inferred_ai_base_path(path: &str) -> &str {
let trimmed = path.trim_end_matches('/');
if trimmed.is_empty() || trimmed == "/" {
""
} else {
trimmed
}
}
fn ai_execution_strategy_for_formats(provider_api_format: &str, client_api_format: &str) -> String {
if provider_api_format == client_api_format {
"local_same_format"
} else {
"local_cross_format"
}
.to_string()
}
fn ai_conversion_mode_for_formats(provider_api_format: &str, client_api_format: &str) -> String {
if provider_api_format == client_api_format {
"none"
} else {
"bidirectional"
}
.to_string()
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use serde_json::json;
use super::*;
#[test]
fn take_ai_decision_plan_core_consumes_required_non_empty_fields() {
let mut payload = test_decision();
let core =
take_ai_decision_plan_core(&mut payload).expect("core fields should be available");
assert_eq!(core.request_id, "req_1");
assert_eq!(core.provider_id, "provider_1");
assert_eq!(core.endpoint_id, "endpoint_1");
assert_eq!(core.key_id, "key_1");
assert_eq!(core.provider_api_format, "openai:chat");
assert_eq!(core.client_api_format, "openai:chat");
assert!(payload.request_id.is_none());
assert!(payload.provider_api_format.is_none());
}
#[test]
fn take_ai_decision_plan_core_rejects_blank_required_fields() {
let mut payload = test_decision();
payload.endpoint_id = Some(" ".to_string());
assert!(take_ai_decision_plan_core(&mut payload).is_none());
}
#[test]
fn take_ai_upstream_auth_pair_rejects_incomplete_auth() {
let mut payload = test_decision();
payload.auth_header = Some("authorization".to_string());
payload.auth_value = Some(" ".to_string());
assert!(take_ai_upstream_auth_pair(&mut payload).is_none());
}
#[test]
fn resolve_ai_passthrough_sync_request_body_prefers_trimmed_base64() {
let body = resolve_ai_passthrough_sync_request_body(
Some(json!({"ignored": true})),
Some(" YWJj ".to_string()),
);
assert_eq!(body.body_bytes_b64.as_deref(), Some("YWJj"));
assert!(body.json_body.is_none());
}
#[test]
fn resolve_ai_passthrough_sync_request_body_uses_json_when_no_base64() {
let body = resolve_ai_passthrough_sync_request_body(Some(json!({"ok": true})), None);
assert_eq!(body.json_body, Some(json!({"ok": true})));
assert!(body.body_bytes_b64.is_none());
}
#[test]
fn build_ai_execution_plan_from_decision_merges_core_and_remaining_payload_fields() {
let mut payload = test_decision();
let core =
take_ai_decision_plan_core(&mut payload).expect("core fields should be available");
let plan = build_ai_execution_plan_from_decision(
&mut payload,
AiExecutionPlanFromDecisionParts {
core,
method: "POST".to_string(),
url: "https://example.com/v1/chat/completions".to_string(),
headers: BTreeMap::from([(
"content-type".to_string(),
"application/json".to_string(),
)]),
content_type: Some("application/json".to_string()),
body: RequestBody::from_json(json!({"model": "gpt-test"})),
stream: true,
},
);
assert_eq!(plan.request_id, "req_1");
assert_eq!(plan.candidate_id.as_deref(), Some("candidate_1"));
assert_eq!(plan.provider_id, "provider_1");
assert_eq!(plan.endpoint_id, "endpoint_1");
assert_eq!(plan.key_id, "key_1");
assert!(plan.stream);
assert_eq!(plan.provider_api_format, "openai:chat");
assert_eq!(plan.client_api_format, "openai:chat");
assert_eq!(plan.model_name.as_deref(), Some("gpt-test"));
assert!(payload.candidate_id.is_none());
assert!(payload.model_name.is_none());
}
#[test]
fn infer_ai_upstream_base_url_preserves_codex_base_path() {
assert_eq!(
infer_ai_upstream_base_url("https://tiger.bookapi.cc/codex/responses").as_deref(),
Some("https://tiger.bookapi.cc/codex")
);
assert_eq!(
infer_ai_upstream_base_url("https://chatgpt.com/backend-api/codex/responses")
.as_deref(),
Some("https://chatgpt.com/backend-api/codex")
);
}
#[test]
fn infer_ai_upstream_base_url_preserves_nested_v1_prefix() {
assert_eq!(
infer_ai_upstream_base_url(
"https://api.openai.example/custom/v1/chat/completions?mode=1"
)
.as_deref(),
Some("https://api.openai.example/custom/v1")
);
}
#[test]
fn infer_ai_upstream_base_url_strips_video_operation_path() {
assert_eq!(
infer_ai_upstream_base_url("https://video.example/nested/v1/videos/task-123/content")
.as_deref(),
Some("https://video.example/nested")
);
}
#[test]
fn build_ai_execution_decision_from_plan_maps_plan_fields() {
let plan = ExecutionPlan {
request_id: "plan-request".to_string(),
candidate_id: Some("candidate-1".to_string()),
provider_name: Some("provider".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://api.example.com/v1/chat/completions".to_string(),
headers: BTreeMap::from([("Authorization".to_string(), "Bearer secret".to_string())]),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(json!({"model": "mapped"})),
stream: false,
client_api_format: "openai:chat".to_string(),
provider_api_format: "claude:messages".to_string(),
model_name: Some("mapped".to_string()),
proxy: None,
tls_profile: None,
timeouts: None,
};
let decision = build_ai_execution_decision_from_plan(AiExecutionDecisionFromPlanParts {
action: "execution_runtime.sync_decision".to_string(),
decision_kind: Some("openai_chat_sync".to_string()),
request_id: Some("trace-1".to_string()),
upstream_base_url: Some("https://api.example.com".to_string()),
include_auth_pair: true,
plan,
report_kind: Some("report".to_string()),
report_context: Some(json!({"candidate_index": 0})),
auth_context: None,
});
assert_eq!(decision.request_id.as_deref(), Some("trace-1"));
assert_eq!(
decision.execution_strategy.as_deref(),
Some("local_cross_format")
);
assert_eq!(decision.conversion_mode.as_deref(), Some("bidirectional"));
assert_eq!(decision.auth_header.as_deref(), Some("Authorization"));
assert_eq!(decision.auth_value.as_deref(), Some("Bearer secret"));
assert_eq!(
decision.provider_request_body,
Some(json!({"model": "mapped"}))
);
assert_eq!(decision.report_kind.as_deref(), Some("report"));
}
fn test_decision() -> AiExecutionDecision {
AiExecutionDecision {
action: "sync".to_string(),
decision_kind: Some("test".to_string()),
execution_strategy: None,
conversion_mode: None,
request_id: Some("req_1".to_string()),
candidate_id: Some("candidate_1".to_string()),
provider_name: Some("provider".to_string()),
provider_id: Some("provider_1".to_string()),
endpoint_id: Some("endpoint_1".to_string()),
key_id: Some("key_1".to_string()),
upstream_base_url: Some("https://example.com".to_string()),
upstream_url: Some("https://example.com/v1/chat/completions".to_string()),
provider_request_method: None,
auth_header: Some("authorization".to_string()),
auth_value: Some("Bearer token".to_string()),
provider_api_format: Some("openai:chat".to_string()),
client_api_format: Some("openai:chat".to_string()),
provider_contract: Some("openai:chat".to_string()),
client_contract: Some("openai:chat".to_string()),
model_name: Some("gpt-test".to_string()),
mapped_model: Some("gpt-test".to_string()),
prompt_cache_key: None,
extra_headers: BTreeMap::new(),
provider_request_headers: BTreeMap::new(),
provider_request_body: None,
provider_request_body_base64: None,
content_type: None,
proxy: None,
tls_profile: None,
timeouts: None,
upstream_is_stream: false,
report_kind: None,
report_context: None,
auth_context: None,
}
}
}

159
crates/aether-ai-serving/src/candidate_materialization.rs Normal file
View File

@@ -0,0 +1,159 @@
use async_trait::async_trait;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AiCandidateMaterializationOutcome<Attempt> {
pub attempts: Vec<Attempt>,
pub candidate_count: usize,
}
#[async_trait]
pub trait AiCandidateMaterializationPort: Send + Sync {
type Candidate: Send;
type Eligible: Send + Sync;
type Skipped: Send;
type Attempt: Send;
type Error: Send;
async fn resolve_and_rank_candidates(
&self,
candidates: Vec<Self::Candidate>,
) -> Result<(Vec<Self::Eligible>, Vec<Self::Skipped>), Self::Error>;
fn decorate_skipped_candidate(&self, skipped: Self::Skipped) -> Self::Skipped {
skipped
}
fn remember_first_candidate_affinity(&self, candidates: &[Self::Eligible]);
async fn persist_available_candidates(
&self,
candidates: Vec<Self::Eligible>,
) -> Result<Vec<Self::Attempt>, Self::Error>;
async fn persist_skipped_candidates(
&self,
starting_candidate_index: u32,
skipped_candidates: Vec<Self::Skipped>,
) -> Result<(), Self::Error>;
}
pub async fn run_ai_candidate_materialization<Port>(
port: &Port,
candidates: Vec<Port::Candidate>,
preselection_skipped_candidates: Vec<Port::Skipped>,
) -> Result<AiCandidateMaterializationOutcome<Port::Attempt>, Port::Error>
where
Port: AiCandidateMaterializationPort,
{
let (candidates, skipped_candidates) = port.resolve_and_rank_candidates(candidates).await?;
let skipped_candidates = preselection_skipped_candidates
.into_iter()
.chain(skipped_candidates)
.map(|candidate| port.decorate_skipped_candidate(candidate))
.collect::<Vec<_>>();
let candidate_count = candidates.len() + skipped_candidates.len();
port.remember_first_candidate_affinity(&candidates);
let available_candidate_count = candidates.len() as u32;
let attempts = port.persist_available_candidates(candidates).await?;
port.persist_skipped_candidates(available_candidate_count, skipped_candidates)
.await?;
Ok(AiCandidateMaterializationOutcome {
attempts,
candidate_count,
})
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Mutex;
#[derive(Default)]
struct TestPort {
calls: Mutex<Vec<String>>,
}
#[async_trait]
impl AiCandidateMaterializationPort for TestPort {
type Candidate = &'static str;
type Eligible = &'static str;
type Skipped = &'static str;
type Attempt = &'static str;
type Error = std::convert::Infallible;
async fn resolve_and_rank_candidates(
&self,
candidates: Vec<Self::Candidate>,
) -> Result<(Vec<Self::Eligible>, Vec<Self::Skipped>), Self::Error> {
self.calls
.lock()
.unwrap()
.push(format!("resolve:{}", candidates.join(",")));
Ok((vec!["eligible-a", "eligible-b"], vec!["resolved-skip"]))
}
fn decorate_skipped_candidate(&self, skipped: Self::Skipped) -> Self::Skipped {
self.calls
.lock()
.unwrap()
.push(format!("decorate:{skipped}"));
skipped
}
fn remember_first_candidate_affinity(&self, candidates: &[Self::Eligible]) {
self.calls
.lock()
.unwrap()
.push(format!("affinity:{}", candidates.join(",")));
}
async fn persist_available_candidates(
&self,
candidates: Vec<Self::Eligible>,
) -> Result<Vec<Self::Attempt>, Self::Error> {
self.calls
.lock()
.unwrap()
.push(format!("available:{}", candidates.join(",")));
Ok(vec!["attempt-a", "attempt-b"])
}
async fn persist_skipped_candidates(
&self,
starting_candidate_index: u32,
skipped_candidates: Vec<Self::Skipped>,
) -> Result<(), Self::Error> {
self.calls.lock().unwrap().push(format!(
"skipped:{starting_candidate_index}:{}",
skipped_candidates.join(",")
));
Ok(())
}
}
#[tokio::test]
async fn materialization_runs_in_serving_order_and_counts_all_candidates() {
let port = TestPort::default();
let outcome =
run_ai_candidate_materialization(&port, vec!["candidate-a"], vec!["pre-skip"])
.await
.unwrap();
assert_eq!(outcome.attempts, ["attempt-a", "attempt-b"]);
assert_eq!(outcome.candidate_count, 4);
assert_eq!(
port.calls.lock().unwrap().as_slice(),
[
"resolve:candidate-a",
"decorate:pre-skip",
"decorate:resolved-skip",
"affinity:eligible-a,eligible-b",
"available:eligible-a,eligible-b",
"skipped:2:pre-skip,resolved-skip",
]
);
}
}

210
crates/aether-ai-serving/src/candidate_metadata.rs Normal file
View File

@@ -0,0 +1,210 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::{Map, Value};
use crate::{ConversionMode, ExecutionStrategy};
pub struct AiCandidateMetadataParts<'a> {
pub provider_api_format: &'a str,
pub client_api_format: &'a str,
pub global_model_id: &'a str,
pub global_model_name: &'a str,
pub model_id: &'a str,
pub selected_provider_model_name: &'a str,
pub mapping_matched_model: Option<&'a str>,
pub provider_name: &'a str,
pub key_name: &'a str,
pub extra_fields: Map<String, Value>,
}
pub fn build_ai_candidate_metadata(parts: AiCandidateMetadataParts<'_>) -> Value {
let mut object = Map::new();
object.insert(
"provider_api_format".to_string(),
Value::String(parts.provider_api_format.to_string()),
);
object.insert(
"client_api_format".to_string(),
Value::String(parts.client_api_format.to_string()),
);
object.insert(
"global_model_id".to_string(),
Value::String(parts.global_model_id.to_string()),
);
object.insert(
"global_model_name".to_string(),
Value::String(parts.global_model_name.to_string()),
);
object.insert(
"model_id".to_string(),
Value::String(parts.model_id.to_string()),
);
object.insert(
"selected_provider_model_name".to_string(),
Value::String(parts.selected_provider_model_name.to_string()),
);
object.insert(
"mapping_matched_model".to_string(),
parts
.mapping_matched_model
.map(|value| Value::String(value.to_string()))
.unwrap_or(Value::Null),
);
object.insert(
"provider_name".to_string(),
Value::String(parts.provider_name.to_string()),
);
object.insert(
"key_name".to_string(),
Value::String(parts.key_name.to_string()),
);
object.extend(parts.extra_fields);
Value::Object(object)
}
pub fn build_ai_candidate_metadata_from_candidate(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
provider_api_format: &str,
client_api_format: &str,
extra_fields: Map<String, Value>,
) -> Value {
build_ai_candidate_metadata(AiCandidateMetadataParts {
provider_api_format,
client_api_format,
global_model_id: candidate.global_model_id.as_str(),
global_model_name: candidate.global_model_name.as_str(),
model_id: candidate.model_id.as_str(),
selected_provider_model_name: candidate.selected_provider_model_name.as_str(),
mapping_matched_model: candidate.mapping_matched_model.as_deref(),
provider_name: candidate.provider_name.as_str(),
key_name: candidate.key_name.as_str(),
extra_fields,
})
}
pub fn append_ai_execution_contract_fields_to_value(
value: Value,
execution_strategy: &str,
conversion_mode: &str,
client_contract: &str,
provider_contract: &str,
) -> Value {
match value {
Value::Object(mut object) => {
object.insert(
"execution_strategy".to_string(),
Value::String(execution_strategy.to_string()),
);
object.insert(
"conversion_mode".to_string(),
Value::String(conversion_mode.to_string()),
);
object.insert(
"client_contract".to_string(),
Value::String(client_contract.to_string()),
);
object.insert(
"provider_contract".to_string(),
Value::String(provider_contract.to_string()),
);
Value::Object(object)
}
other => other,
}
}
pub fn ai_local_execution_contract_for_formats(
client_api_format: &str,
provider_api_format: &str,
) -> (ExecutionStrategy, ConversionMode) {
if aether_ai_formats::api_format_alias_matches(client_api_format, provider_api_format) {
return (ExecutionStrategy::LocalSameFormat, ConversionMode::None);
}
let conversion_mode =
if aether_ai_formats::request_conversion_kind(client_api_format, provider_api_format)
.is_some()
{
ConversionMode::Bidirectional
} else {
ConversionMode::None
};
(ExecutionStrategy::LocalCrossFormat, conversion_mode)
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn candidate_metadata_builds_base_candidate_fields_and_extra_data() {
let mut extra_fields = Map::new();
extra_fields.insert("source".to_string(), json!("test"));
let metadata = build_ai_candidate_metadata(AiCandidateMetadataParts {
provider_api_format: "openai:responses",
client_api_format: "claude:messages",
global_model_id: "global-1",
global_model_name: "gpt-5.4",
model_id: "model-1",
selected_provider_model_name: "gpt-5.4",
mapping_matched_model: Some("gpt-5"),
provider_name: "RightCode",
key_name: "key-a",
extra_fields,
});
assert_eq!(metadata["provider_api_format"], "openai:responses");
assert_eq!(metadata["client_api_format"], "claude:messages");
assert_eq!(metadata["global_model_id"], "global-1");
assert_eq!(metadata["mapping_matched_model"], "gpt-5");
assert_eq!(metadata["source"], "test");
}
#[test]
fn execution_contract_fields_append_to_object_values_only() {
let value = append_ai_execution_contract_fields_to_value(
json!({"existing": true}),
"local_cross_format",
"bidirectional",
"openai:chat",
"claude:messages",
);
assert_eq!(value["existing"], true);
assert_eq!(value["execution_strategy"], "local_cross_format");
assert_eq!(value["conversion_mode"], "bidirectional");
assert_eq!(value["client_contract"], "openai:chat");
assert_eq!(value["provider_contract"], "claude:messages");
assert_eq!(
append_ai_execution_contract_fields_to_value(
Value::Null,
"local_same_format",
"none",
"openai:chat",
"openai:chat",
),
Value::Null
);
}
#[test]
fn local_execution_contract_is_derived_from_client_and_provider_formats() {
assert_eq!(
ai_local_execution_contract_for_formats(" OPENAI:CHAT ", "openai:chat"),
(ExecutionStrategy::LocalSameFormat, ConversionMode::None)
);
assert_eq!(
ai_local_execution_contract_for_formats("openai:chat", "claude:messages"),
(
ExecutionStrategy::LocalCrossFormat,
ConversionMode::Bidirectional,
)
);
assert_eq!(
ai_local_execution_contract_for_formats("openai:chat", "unknown:format"),
(ExecutionStrategy::LocalCrossFormat, ConversionMode::None)
);
}
}

433
crates/aether-ai-serving/src/candidate_persistence.rs Normal file
View File

@@ -0,0 +1,433 @@
use aether_scheduler_core::SchedulerRankingOutcome;
use async_trait::async_trait;
use serde_json::Value;
#[async_trait]
pub trait AiAvailableCandidatePersistencePort: Send + Sync {
type Candidate: Clone + Send + Sync;
type Attempt: Send;
type ExtraData: Clone + Send + Sync;
type Error: Send;
fn attempt_slot_count(&self, candidate: &Self::Candidate) -> u32;
fn build_extra_data(&self, candidate: &Self::Candidate) -> Option<Self::ExtraData>;
fn generate_candidate_id(&self) -> String;
fn should_persist_available_candidate(&self, candidate: &Self::Candidate) -> bool;
async fn persist_available_candidate(
&self,
candidate: &Self::Candidate,
candidate_index: u32,
retry_index: u32,
generated_candidate_id: &str,
extra_data: Option<Self::ExtraData>,
) -> Result<String, Self::Error>;
fn build_attempt(
&self,
candidate: Self::Candidate,
candidate_index: u32,
retry_index: u32,
candidate_id: String,
) -> Self::Attempt;
}
pub async fn run_ai_available_candidate_persistence<Port>(
port: &Port,
candidates: Vec<Port::Candidate>,
) -> Result<Vec<Port::Attempt>, Port::Error>
where
Port: AiAvailableCandidatePersistencePort,
{
let total_attempts = candidates
.iter()
.map(|candidate| port.attempt_slot_count(candidate) as usize)
.sum();
let mut materialized = Vec::with_capacity(total_attempts);
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let candidate_index = candidate_index as u32;
let attempt_slots = port.attempt_slot_count(&candidate).max(1);
let extra_data = port.build_extra_data(&candidate);
let mut owned_candidate = Some(candidate);
for retry_index in 0..attempt_slots {
let candidate = owned_candidate
.as_ref()
.expect("candidate should remain available until final retry");
let generated_candidate_id = port.generate_candidate_id();
let candidate_id = if port.should_persist_available_candidate(candidate) {
port.persist_available_candidate(
candidate,
candidate_index,
retry_index,
generated_candidate_id.as_str(),
extra_data.clone(),
)
.await?
} else {
generated_candidate_id
};
let candidate = if retry_index + 1 == attempt_slots {
owned_candidate
.take()
.expect("final retry should consume owned candidate")
} else {
candidate.clone()
};
materialized.push(port.build_attempt(
candidate,
candidate_index,
retry_index,
candidate_id,
));
}
}
Ok(materialized)
}
#[async_trait]
pub trait AiSkippedCandidatePersistencePort: Send + Sync {
type Skipped: Send + Sync;
type ExtraData: Send + Sync;
type Error: Send;
fn should_persist_skipped_candidate(&self, candidate: &Self::Skipped) -> bool;
fn build_extra_data(&self, candidate: &Self::Skipped) -> Option<Self::ExtraData>;
fn generate_candidate_id(&self) -> String;
async fn persist_skipped_candidate(
&self,
candidate: &Self::Skipped,
candidate_index: u32,
generated_candidate_id: &str,
extra_data: Option<Self::ExtraData>,
) -> Result<(), Self::Error>;
}
pub async fn run_ai_skipped_candidate_persistence<Port>(
port: &Port,
starting_candidate_index: u32,
skipped_candidates: Vec<Port::Skipped>,
) -> Result<(), Port::Error>
where
Port: AiSkippedCandidatePersistencePort,
{
let mut next_candidate_index = starting_candidate_index;
for skipped_candidate in skipped_candidates {
if !port.should_persist_skipped_candidate(&skipped_candidate) {
continue;
}
let generated_candidate_id = port.generate_candidate_id();
let extra_data = port.build_extra_data(&skipped_candidate);
port.persist_skipped_candidate(
&skipped_candidate,
next_candidate_index,
generated_candidate_id.as_str(),
extra_data,
)
.await?;
next_candidate_index = next_candidate_index.saturating_add(1);
}
Ok(())
}
pub fn ai_should_persist_available_candidate_for_pool_key(pool_key_index: Option<u32>) -> bool {
pool_key_index.is_none_or(|index| index == 0)
}
pub fn ai_should_persist_skipped_candidate_for_pool_membership(is_pool_candidate: bool) -> bool {
!is_pool_candidate
}
pub fn ai_candidate_extra_data_with_ranking(
extra_data: Option<Value>,
ranking: Option<&SchedulerRankingOutcome>,
) -> Option<Value> {
let Some(ranking) = ranking else {
return extra_data;
};
let mut object = match extra_data {
Some(Value::Object(object)) => object,
Some(value) => {
let mut object = serde_json::Map::new();
object.insert("extra".to_string(), value);
object
}
None => serde_json::Map::new(),
};
crate::append_ai_ranking_metadata_to_object(&mut object, ranking);
Some(Value::Object(object))
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Mutex;
#[derive(Debug, Clone, PartialEq, Eq)]
struct TestCandidate {
id: &'static str,
attempt_slots: u32,
persist: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct TestAttempt {
id: &'static str,
candidate_index: u32,
retry_index: u32,
candidate_id: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct TestSkipped {
id: &'static str,
persist: bool,
}
#[derive(Default)]
struct TestPort {
next_id: Mutex<u32>,
calls: Mutex<Vec<String>>,
}
impl TestPort {
fn next_candidate_id(&self) -> String {
let mut next_id = self.next_id.lock().unwrap();
*next_id += 1;
format!("candidate-{next_id}")
}
}
#[async_trait]
impl AiAvailableCandidatePersistencePort for TestPort {
type Candidate = TestCandidate;
type Attempt = TestAttempt;
type ExtraData = String;
type Error = std::convert::Infallible;
fn attempt_slot_count(&self, candidate: &Self::Candidate) -> u32 {
candidate.attempt_slots
}
fn build_extra_data(&self, candidate: &Self::Candidate) -> Option<Self::ExtraData> {
Some(format!("extra:{}", candidate.id))
}
fn generate_candidate_id(&self) -> String {
self.next_candidate_id()
}
fn should_persist_available_candidate(&self, candidate: &Self::Candidate) -> bool {
candidate.persist
}
async fn persist_available_candidate(
&self,
candidate: &Self::Candidate,
candidate_index: u32,
retry_index: u32,
generated_candidate_id: &str,
extra_data: Option<Self::ExtraData>,
) -> Result<String, Self::Error> {
self.calls.lock().unwrap().push(format!(
"available:{}:{candidate_index}:{retry_index}:{generated_candidate_id}:{}",
candidate.id,
extra_data.unwrap_or_default()
));
Ok(format!("stored-{generated_candidate_id}"))
}
fn build_attempt(
&self,
candidate: Self::Candidate,
candidate_index: u32,
retry_index: u32,
candidate_id: String,
) -> Self::Attempt {
TestAttempt {
id: candidate.id,
candidate_index,
retry_index,
candidate_id,
}
}
}
#[async_trait]
impl AiSkippedCandidatePersistencePort for TestPort {
type Skipped = TestSkipped;
type ExtraData = String;
type Error = std::convert::Infallible;
fn should_persist_skipped_candidate(&self, candidate: &Self::Skipped) -> bool {
candidate.persist
}
fn build_extra_data(&self, candidate: &Self::Skipped) -> Option<Self::ExtraData> {
Some(format!("extra:{}", candidate.id))
}
fn generate_candidate_id(&self) -> String {
self.next_candidate_id()
}
async fn persist_skipped_candidate(
&self,
candidate: &Self::Skipped,
candidate_index: u32,
generated_candidate_id: &str,
extra_data: Option<Self::ExtraData>,
) -> Result<(), Self::Error> {
self.calls.lock().unwrap().push(format!(
"skipped:{}:{candidate_index}:{generated_candidate_id}:{}",
candidate.id,
extra_data.unwrap_or_default()
));
Ok(())
}
}
#[tokio::test]
async fn available_persistence_expands_candidates_into_retry_attempts() {
let port = TestPort::default();
let attempts = run_ai_available_candidate_persistence(
&port,
vec![
TestCandidate {
id: "a",
attempt_slots: 2,
persist: true,
},
TestCandidate {
id: "b",
attempt_slots: 1,
persist: false,
},
],
)
.await
.unwrap();
assert_eq!(
attempts,
[
TestAttempt {
id: "a",
candidate_index: 0,
retry_index: 0,
candidate_id: "stored-candidate-1".to_string(),
},
TestAttempt {
id: "a",
candidate_index: 0,
retry_index: 1,
candidate_id: "stored-candidate-2".to_string(),
},
TestAttempt {
id: "b",
candidate_index: 1,
retry_index: 0,
candidate_id: "candidate-3".to_string(),
},
]
);
assert_eq!(
port.calls.lock().unwrap().as_slice(),
[
"available:a:0:0:candidate-1:extra:a",
"available:a:0:1:candidate-2:extra:a",
]
);
}
#[tokio::test]
async fn skipped_persistence_keeps_indices_for_persisted_candidates_only() {
let port = TestPort::default();
run_ai_skipped_candidate_persistence(
&port,
3,
vec![
TestSkipped {
id: "ignored",
persist: false,
},
TestSkipped {
id: "a",
persist: true,
},
TestSkipped {
id: "b",
persist: true,
},
],
)
.await
.unwrap();
assert_eq!(
port.calls.lock().unwrap().as_slice(),
[
"skipped:a:3:candidate-1:extra:a",
"skipped:b:4:candidate-2:extra:b",
]
);
}
#[test]
fn pool_candidate_persistence_policy_persists_representatives_only() {
assert!(ai_should_persist_available_candidate_for_pool_key(None));
assert!(ai_should_persist_available_candidate_for_pool_key(Some(0)));
assert!(!ai_should_persist_available_candidate_for_pool_key(Some(1)));
assert!(ai_should_persist_skipped_candidate_for_pool_membership(
false
));
assert!(!ai_should_persist_skipped_candidate_for_pool_membership(
true
));
}
#[test]
fn candidate_extra_data_with_ranking_preserves_existing_shapes() {
let ranking = SchedulerRankingOutcome {
original_index: 1,
ranking_index: 0,
priority_mode: aether_scheduler_core::SchedulerPriorityMode::Provider,
ranking_mode: aether_scheduler_core::SchedulerRankingMode::CacheAffinity,
priority_slot: 3,
promoted_by: Some("cached_affinity"),
demoted_by: None,
};
let object = ai_candidate_extra_data_with_ranking(
Some(serde_json::json!({"source": "test"})),
Some(&ranking),
)
.expect("extra data should exist");
assert_eq!(object.get("source"), Some(&serde_json::json!("test")));
assert_eq!(object.get("ranking_index"), Some(&serde_json::json!(0)));
assert_eq!(
object.get("promoted_by"),
Some(&serde_json::json!("cached_affinity"))
);
let scalar =
ai_candidate_extra_data_with_ranking(Some(serde_json::json!("raw")), Some(&ranking))
.expect("scalar extra data should be wrapped");
assert_eq!(scalar.get("extra"), Some(&serde_json::json!("raw")));
assert_eq!(scalar.get("priority_slot"), Some(&serde_json::json!(3)));
}
}

117
crates/aether-ai-serving/src/candidate_persistence_policy.rs Normal file
View File

@@ -0,0 +1,117 @@
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[allow(clippy::enum_variant_names)]
pub enum AiCandidatePersistencePolicyKind {
StandardDecision,
SameFormatProviderDecision,
OpenAiChatDecision,
OpenAiResponsesDecision,
ImageDecision,
GeminiFilesDecision,
VideoDecision,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AiCandidatePersistencePolicySpec {
pub available_error_context: &'static str,
pub skipped_error_context: &'static str,
pub record_runtime_miss_diagnostic: bool,
}
pub fn ai_candidate_persistence_policy_spec(
kind: AiCandidatePersistencePolicyKind,
) -> AiCandidatePersistencePolicySpec {
match kind {
AiCandidatePersistencePolicyKind::StandardDecision => AiCandidatePersistencePolicySpec {
available_error_context:
"gateway local standard decision request candidate upsert failed",
skipped_error_context:
"gateway local standard decision failed to persist skipped candidate",
record_runtime_miss_diagnostic: true,
},
AiCandidatePersistencePolicyKind::SameFormatProviderDecision => {
AiCandidatePersistencePolicySpec {
available_error_context:
"gateway local same-format decision request candidate upsert failed",
skipped_error_context:
"gateway local same-format decision failed to persist skipped candidate",
record_runtime_miss_diagnostic: true,
}
}
AiCandidatePersistencePolicyKind::OpenAiChatDecision => AiCandidatePersistencePolicySpec {
available_error_context:
"gateway local openai chat decision request candidate upsert failed",
skipped_error_context:
"gateway local openai chat decision failed to persist skipped candidate",
record_runtime_miss_diagnostic: true,
},
AiCandidatePersistencePolicyKind::OpenAiResponsesDecision => {
AiCandidatePersistencePolicySpec {
available_error_context:
"gateway local openai responses decision request candidate upsert failed",
skipped_error_context:
"gateway local openai responses decision failed to persist skipped candidate",
record_runtime_miss_diagnostic: true,
}
}
AiCandidatePersistencePolicyKind::ImageDecision => AiCandidatePersistencePolicySpec {
available_error_context:
"gateway local openai image decision request candidate upsert failed",
skipped_error_context:
"gateway local openai image decision failed to persist skipped candidate",
record_runtime_miss_diagnostic: false,
},
AiCandidatePersistencePolicyKind::GeminiFilesDecision => AiCandidatePersistencePolicySpec {
available_error_context: "gateway local gemini files request candidate upsert failed",
skipped_error_context: "gateway local gemini files failed to persist skipped candidate",
record_runtime_miss_diagnostic: false,
},
AiCandidatePersistencePolicyKind::VideoDecision => AiCandidatePersistencePolicySpec {
available_error_context: "gateway local video decision request candidate upsert failed",
skipped_error_context:
"gateway local video decision failed to persist skipped candidate",
record_runtime_miss_diagnostic: false,
},
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn decision_policies_define_candidate_persistence_side_effects() {
let standard = ai_candidate_persistence_policy_spec(
AiCandidatePersistencePolicyKind::StandardDecision,
);
assert_eq!(
standard.available_error_context,
"gateway local standard decision request candidate upsert failed"
);
assert_eq!(
standard.skipped_error_context,
"gateway local standard decision failed to persist skipped candidate"
);
assert!(standard.record_runtime_miss_diagnostic);
let same_format = ai_candidate_persistence_policy_spec(
AiCandidatePersistencePolicyKind::SameFormatProviderDecision,
);
assert_eq!(
same_format.available_error_context,
"gateway local same-format decision request candidate upsert failed"
);
assert_eq!(
same_format.skipped_error_context,
"gateway local same-format decision failed to persist skipped candidate"
);
assert!(same_format.record_runtime_miss_diagnostic);
let image =
ai_candidate_persistence_policy_spec(AiCandidatePersistencePolicyKind::ImageDecision);
assert_eq!(
image.available_error_context,
"gateway local openai image decision request candidate upsert failed"
);
assert!(!image.record_runtime_miss_diagnostic);
}
}

91
crates/aether-ai-serving/src/candidate_preparation.rs Normal file
View File

@@ -0,0 +1,91 @@
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AiPreparedHeaderAuthenticatedCandidate {
pub auth_header: String,
pub auth_value: String,
pub mapped_model: String,
}
pub fn prepare_ai_header_authenticated_candidate(
direct_auth: Option<(String, String)>,
oauth_header_auth: Option<(String, String)>,
selected_provider_model_name: &str,
) -> Result<AiPreparedHeaderAuthenticatedCandidate, &'static str> {
let Some((auth_header, auth_value)) = direct_auth.or(oauth_header_auth) else {
return Err("transport_auth_unavailable");
};
let mapped_model = resolve_ai_candidate_mapped_model(selected_provider_model_name)?;
Ok(AiPreparedHeaderAuthenticatedCandidate {
auth_header,
auth_value,
mapped_model,
})
}
pub fn resolve_ai_candidate_mapped_model(
selected_provider_model_name: &str,
) -> Result<String, &'static str> {
let mapped_model = selected_provider_model_name.trim().to_string();
if mapped_model.is_empty() {
return Err("mapped_model_missing");
}
Ok(mapped_model)
}
#[cfg(test)]
mod tests {
use super::{prepare_ai_header_authenticated_candidate, resolve_ai_candidate_mapped_model};
#[test]
fn mapped_model_trims_selected_provider_model_name() {
assert_eq!(
resolve_ai_candidate_mapped_model(" gpt-test-upstream "),
Ok("gpt-test-upstream".to_string())
);
}
#[test]
fn mapped_model_rejects_empty_selected_provider_model_name() {
assert_eq!(
resolve_ai_candidate_mapped_model(" "),
Err("mapped_model_missing")
);
}
#[test]
fn header_auth_preparation_prefers_direct_auth_and_allows_empty_value() {
let prepared = prepare_ai_header_authenticated_candidate(
Some(("authorization".to_string(), String::new())),
Some(("x-oauth".to_string(), "oauth".to_string())),
"gpt-test-upstream",
)
.expect("direct auth should prepare candidate");
assert_eq!(prepared.auth_header, "authorization");
assert_eq!(prepared.auth_value, "");
assert_eq!(prepared.mapped_model, "gpt-test-upstream");
}
#[test]
fn header_auth_preparation_falls_back_to_oauth_header_auth() {
let prepared = prepare_ai_header_authenticated_candidate(
None,
Some(("authorization".to_string(), "Bearer token".to_string())),
" gpt-test-upstream ",
)
.expect("oauth header auth should prepare candidate");
assert_eq!(prepared.auth_header, "authorization");
assert_eq!(prepared.auth_value, "Bearer token");
assert_eq!(prepared.mapped_model, "gpt-test-upstream");
}
#[test]
fn header_auth_preparation_requires_some_auth() {
assert_eq!(
prepare_ai_header_authenticated_candidate(None, None, "gpt-test-upstream"),
Err("transport_auth_unavailable")
);
}
}

185
crates/aether-ai-serving/src/candidate_preselection.rs Normal file
View File

@@ -0,0 +1,185 @@
use std::collections::BTreeSet;
use async_trait::async_trait;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AiCandidatePreselectionOutcome<Candidate, Skipped> {
pub candidates: Vec<Candidate>,
pub skipped_candidates: Vec<Skipped>,
}
#[async_trait]
pub trait AiCandidatePreselectionPort: Send + Sync {
type Candidate: Send;
type Skipped: Send;
type Error: Send;
fn candidate_api_formats(&self) -> Vec<String>;
fn candidate_api_format_matches_client(&self, candidate_api_format: &str) -> bool;
async fn list_candidates_for_api_format(
&self,
candidate_api_format: &str,
matches_client_format: bool,
) -> Result<(Vec<Self::Candidate>, Vec<Self::Skipped>), Self::Error>;
fn candidate_allowed(
&self,
_candidate: &Self::Candidate,
_candidate_api_format: &str,
_matches_client_format: bool,
) -> bool {
true
}
fn skipped_candidate_allowed(
&self,
_skipped_candidate: &Self::Skipped,
_candidate_api_format: &str,
_matches_client_format: bool,
) -> bool {
true
}
fn candidate_key(&self, candidate: &Self::Candidate) -> String;
fn skipped_candidate_key(&self, skipped_candidate: &Self::Skipped) -> String;
}
pub async fn run_ai_candidate_preselection<Port>(
port: &Port,
) -> Result<AiCandidatePreselectionOutcome<Port::Candidate, Port::Skipped>, Port::Error>
where
Port: AiCandidatePreselectionPort,
{
let mut candidates = Vec::new();
let mut skipped_candidates = Vec::new();
let mut seen_candidates = BTreeSet::new();
let mut seen_skipped_candidates = BTreeSet::new();
for candidate_api_format in port.candidate_api_formats() {
let matches_client_format =
port.candidate_api_format_matches_client(candidate_api_format.as_str());
let (selected, skipped) = port
.list_candidates_for_api_format(candidate_api_format.as_str(), matches_client_format)
.await?;
for skipped_candidate in skipped {
if !port.skipped_candidate_allowed(
&skipped_candidate,
candidate_api_format.as_str(),
matches_client_format,
) {
continue;
}
let candidate_key = port.skipped_candidate_key(&skipped_candidate);
if seen_skipped_candidates.insert(candidate_key) {
skipped_candidates.push(skipped_candidate);
}
}
for candidate in selected {
if !port.candidate_allowed(
&candidate,
candidate_api_format.as_str(),
matches_client_format,
) {
continue;
}
let candidate_key = port.candidate_key(&candidate);
if seen_candidates.insert(candidate_key) {
candidates.push(candidate);
}
}
}
Ok(AiCandidatePreselectionOutcome {
candidates,
skipped_candidates,
})
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Mutex;
#[derive(Default)]
struct TestPort {
calls: Mutex<Vec<String>>,
}
#[async_trait]
impl AiCandidatePreselectionPort for TestPort {
type Candidate = &'static str;
type Skipped = &'static str;
type Error = std::convert::Infallible;
fn candidate_api_formats(&self) -> Vec<String> {
vec!["same".to_string(), "cross".to_string()]
}
fn candidate_api_format_matches_client(&self, candidate_api_format: &str) -> bool {
candidate_api_format == "same"
}
async fn list_candidates_for_api_format(
&self,
candidate_api_format: &str,
matches_client_format: bool,
) -> Result<(Vec<Self::Candidate>, Vec<Self::Skipped>), Self::Error> {
self.calls.lock().unwrap().push(format!(
"list:{candidate_api_format}:{matches_client_format}"
));
Ok(match candidate_api_format {
"same" => (vec!["candidate-a"], vec!["skip-a"]),
"cross" => (
vec!["candidate-a", "candidate-b", "blocked-candidate"],
vec!["skip-a", "skip-b", "blocked-skip"],
),
_ => (Vec::new(), Vec::new()),
})
}
fn candidate_allowed(
&self,
candidate: &Self::Candidate,
_candidate_api_format: &str,
matches_client_format: bool,
) -> bool {
matches_client_format || *candidate != "blocked-candidate"
}
fn skipped_candidate_allowed(
&self,
skipped_candidate: &Self::Skipped,
_candidate_api_format: &str,
matches_client_format: bool,
) -> bool {
matches_client_format || *skipped_candidate != "blocked-skip"
}
fn candidate_key(&self, candidate: &Self::Candidate) -> String {
(*candidate).to_string()
}
fn skipped_candidate_key(&self, skipped_candidate: &Self::Skipped) -> String {
(*skipped_candidate).to_string()
}
}
#[tokio::test]
async fn preselection_runs_formats_in_order_filters_cross_format_and_dedupes() {
let port = TestPort::default();
let outcome = run_ai_candidate_preselection(&port).await.unwrap();
assert_eq!(outcome.candidates, ["candidate-a", "candidate-b"]);
assert_eq!(outcome.skipped_candidates, ["skip-a", "skip-b"]);
assert_eq!(
port.calls.lock().unwrap().as_slice(),
["list:same:true", "list:cross:false"]
);
}
}

301
crates/aether-ai-serving/src/candidate_ranking.rs Normal file
View File

@@ -0,0 +1,301 @@
use aether_scheduler_core::{
apply_scheduler_candidate_ranking, requested_capability_priority_for_candidate,
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode, SchedulerRankableCandidate,
SchedulerRankingContext, SchedulerRankingMode, SchedulerRankingOutcome,
SchedulerTunnelAffinityBucket,
};
use async_trait::async_trait;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AiRankingSchedulingMode {
FixedOrder,
CacheAffinity,
LoadBalance,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AiRankingContextConfig {
pub priority_mode: SchedulerPriorityMode,
pub scheduling_mode: AiRankingSchedulingMode,
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct AiRankableCandidateParts<'a> {
pub candidate: &'a SchedulerMinimalCandidateSelectionCandidate,
pub original_index: usize,
pub normalized_client_api_format: &'a str,
pub provider_api_format: &'a str,
pub required_capabilities: Option<&'a serde_json::Value>,
pub cached_affinity_match: bool,
pub tunnel_bucket: SchedulerTunnelAffinityBucket,
pub keep_priority_on_conversion: bool,
}
#[async_trait]
pub trait AiCandidateRankingPort: Send + Sync {
type Candidate: Send + Sync;
type AffinityTarget: Send + Sync;
type Error: Send;
fn affinity_requested_model(&self, candidates: &[Self::Candidate]) -> Option<String>;
async fn read_cached_affinity_target(
&self,
normalized_client_api_format: &str,
affinity_requested_model: Option<&str>,
) -> Result<Option<Self::AffinityTarget>, Self::Error>;
fn cached_affinity_matches(
&self,
candidate: &Self::Candidate,
target: &Self::AffinityTarget,
) -> bool;
async fn build_rankable_candidate(
&self,
candidate: &Self::Candidate,
original_index: usize,
normalized_client_api_format: &str,
cached_affinity_match: bool,
) -> Result<SchedulerRankableCandidate, Self::Error>;
fn ranking_context(&self) -> SchedulerRankingContext;
fn apply_ranking_outcome(
&self,
candidate: &mut Self::Candidate,
outcome: SchedulerRankingOutcome,
);
}
pub fn build_ai_rankable_candidate(
parts: AiRankableCandidateParts<'_>,
) -> SchedulerRankableCandidate {
let is_same_format = aether_ai_formats::api_format_alias_matches(
parts.provider_api_format,
parts.normalized_client_api_format,
);
let format_preference = aether_ai_formats::request_candidate_api_format_preference(
parts.normalized_client_api_format,
parts.provider_api_format,
)
.unwrap_or((u8::MAX, u8::MAX));
let mut rankable =
SchedulerRankableCandidate::from_candidate(parts.candidate, parts.original_index);
// The scheduler order is the upstream tie-breaker; AI serving only adds transport facts.
rankable.provider_id.clear();
rankable.endpoint_id.clear();
rankable.key_id.clear();
rankable.selected_provider_model_name.clear();
rankable
.with_capability_priority(requested_capability_priority_for_candidate(
parts.required_capabilities,
parts.candidate,
))
.with_cached_affinity_match(parts.cached_affinity_match)
.with_tunnel_bucket(parts.tunnel_bucket)
.with_format_state(
!is_same_format && !parts.keep_priority_on_conversion,
format_preference,
)
}
pub fn ai_ranking_context(config: AiRankingContextConfig) -> SchedulerRankingContext {
SchedulerRankingContext {
priority_mode: config.priority_mode,
ranking_mode: ai_ranking_mode(config.scheduling_mode),
include_health: false,
load_balance_seed: 0,
}
}
fn ai_ranking_mode(mode: AiRankingSchedulingMode) -> SchedulerRankingMode {
match mode {
AiRankingSchedulingMode::FixedOrder => SchedulerRankingMode::FixedOrder,
AiRankingSchedulingMode::CacheAffinity => SchedulerRankingMode::CacheAffinity,
AiRankingSchedulingMode::LoadBalance => SchedulerRankingMode::LoadBalance,
}
}
pub async fn run_ai_candidate_ranking<Port>(
port: &Port,
mut candidates: Vec<Port::Candidate>,
normalized_client_api_format: &str,
) -> Result<Vec<Port::Candidate>, Port::Error>
where
Port: AiCandidateRankingPort,
{
let affinity_requested_model = port.affinity_requested_model(&candidates);
let cached_affinity_target = port
.read_cached_affinity_target(
normalized_client_api_format,
affinity_requested_model.as_deref(),
)
.await?;
let mut rankables = Vec::with_capacity(candidates.len());
for (original_index, candidate) in candidates.iter().enumerate() {
let cached_affinity_match = cached_affinity_target
.as_ref()
.is_some_and(|target| port.cached_affinity_matches(candidate, target));
rankables.push(
port.build_rankable_candidate(
candidate,
original_index,
normalized_client_api_format,
cached_affinity_match,
)
.await?,
);
}
let outcomes =
apply_scheduler_candidate_ranking(&mut candidates, &rankables, port.ranking_context());
for outcome in outcomes {
let ranking_index = outcome.ranking_index;
if let Some(candidate) = candidates.get_mut(ranking_index) {
port.apply_ranking_outcome(candidate, outcome);
}
}
Ok(candidates)
}
#[cfg(test)]
mod tests {
use super::*;
use aether_scheduler_core::{SchedulerPriorityMode, SchedulerRankingMode};
use std::sync::Mutex;
#[derive(Debug, Clone, PartialEq, Eq)]
struct TestCandidate {
id: &'static str,
priority: i32,
ranking_index: Option<usize>,
cached_affinity: bool,
}
#[derive(Default)]
struct TestPort {
calls: Mutex<Vec<String>>,
}
#[async_trait]
impl AiCandidateRankingPort for TestPort {
type Candidate = TestCandidate;
type AffinityTarget = &'static str;
type Error = std::convert::Infallible;
fn affinity_requested_model(&self, candidates: &[Self::Candidate]) -> Option<String> {
candidates.first().map(|_| "model-a".to_string())
}
async fn read_cached_affinity_target(
&self,
normalized_client_api_format: &str,
affinity_requested_model: Option<&str>,
) -> Result<Option<Self::AffinityTarget>, Self::Error> {
self.calls.lock().unwrap().push(format!(
"affinity:{normalized_client_api_format}:{}",
affinity_requested_model.unwrap_or_default()
));
Ok(Some("candidate-b"))
}
fn cached_affinity_matches(
&self,
candidate: &Self::Candidate,
target: &Self::AffinityTarget,
) -> bool {
candidate.id == *target
}
async fn build_rankable_candidate(
&self,
candidate: &Self::Candidate,
original_index: usize,
_normalized_client_api_format: &str,
cached_affinity_match: bool,
) -> Result<SchedulerRankableCandidate, Self::Error> {
self.calls
.lock()
.unwrap()
.push(format!("rankable:{}:{cached_affinity_match}", candidate.id));
Ok(SchedulerRankableCandidate {
provider_id: candidate.id.to_string(),
endpoint_id: String::new(),
key_id: String::new(),
selected_provider_model_name: String::new(),
provider_priority: candidate.priority,
key_internal_priority: 0,
key_global_priority_for_format: None,
capability_priority: (0, 0),
cached_affinity_match,
affinity_hash: None,
tunnel_bucket: Default::default(),
demote_cross_format: false,
format_preference: (0, 0),
health_bucket: None,
health_score: 1.0,
original_index,
})
}
fn ranking_context(&self) -> SchedulerRankingContext {
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::CacheAffinity,
include_health: false,
load_balance_seed: 0,
}
}
fn apply_ranking_outcome(
&self,
candidate: &mut Self::Candidate,
outcome: SchedulerRankingOutcome,
) {
candidate.ranking_index = Some(outcome.ranking_index);
candidate.cached_affinity = outcome.promoted_by.is_some();
}
}
#[tokio::test]
async fn ranking_builds_rankables_applies_scheduler_order_and_writes_outcomes() {
let port = TestPort::default();
let candidates = vec![
TestCandidate {
id: "candidate-a",
priority: 10,
ranking_index: None,
cached_affinity: false,
},
TestCandidate {
id: "candidate-b",
priority: 20,
ranking_index: None,
cached_affinity: false,
},
];
let ranked = run_ai_candidate_ranking(&port, candidates, "openai:chat")
.await
.unwrap();
assert_eq!(ranked[0].id, "candidate-b");
assert_eq!(ranked[0].ranking_index, Some(0));
assert!(ranked[0].cached_affinity);
assert_eq!(ranked[1].id, "candidate-a");
assert_eq!(ranked[1].ranking_index, Some(1));
assert_eq!(
port.calls.lock().unwrap().as_slice(),
[
"affinity:openai:chat:model-a",
"rankable:candidate-a:false",
"rankable:candidate-b:true",
]
);
}
}

415
crates/aether-ai-serving/src/candidate_resolution.rs Normal file
View File

@@ -0,0 +1,415 @@
use async_trait::async_trait;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AiCandidateResolutionMode {
Standard,
WithoutTransportPairGate,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AiCandidateResolutionRequest<'a> {
pub client_api_format: &'a str,
pub requested_model: Option<&'a str>,
pub mode: AiCandidateResolutionMode,
}
impl<'a> AiCandidateResolutionRequest<'a> {
pub fn standard(client_api_format: &'a str, requested_model: Option<&'a str>) -> Self {
Self {
client_api_format,
requested_model,
mode: AiCandidateResolutionMode::Standard,
}
}
pub fn without_transport_pair_gate(
client_api_format: &'a str,
requested_model: Option<&'a str>,
) -> Self {
Self {
client_api_format,
requested_model,
mode: AiCandidateResolutionMode::WithoutTransportPairGate,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AiCandidateResolutionOutcome<Eligible, Skipped> {
pub eligible_candidates: Vec<Eligible>,
pub skipped_candidates: Vec<Skipped>,
}
#[async_trait]
pub trait AiCandidateResolutionPort: Send + Sync {
type Candidate: Send;
type Transport: Send + Sync;
type Eligible: Send + Sync;
type Skipped: Send;
type Error: Send;
async fn read_candidate_transport(
&self,
candidate: &Self::Candidate,
) -> Result<Option<Self::Transport>, Self::Error>;
fn build_missing_transport_skipped_candidate(
&self,
candidate: Self::Candidate,
) -> Self::Skipped;
fn candidate_common_skip_reason(
&self,
candidate: &Self::Candidate,
transport: &Self::Transport,
requested_model: Option<&str>,
) -> Option<&'static str>;
fn candidate_transport_pair_skip_reason(
&self,
candidate: &Self::Candidate,
transport: &Self::Transport,
normalized_client_api_format: &str,
requested_model: &str,
) -> Option<&'static str>;
fn build_skipped_candidate(
&self,
candidate: Self::Candidate,
transport: Self::Transport,
skip_reason: &'static str,
) -> Self::Skipped;
fn build_eligible_candidate(
&self,
candidate: Self::Candidate,
transport: Self::Transport,
) -> Self::Eligible;
async fn rank_eligible_candidates(
&self,
candidates: Vec<Self::Eligible>,
normalized_client_api_format: &str,
) -> Result<Vec<Self::Eligible>, Self::Error>;
async fn apply_pool_scheduler(
&self,
candidates: Vec<Self::Eligible>,
) -> Result<(Vec<Self::Eligible>, Vec<Self::Skipped>), Self::Error>;
}
pub async fn run_ai_candidate_resolution<Port>(
port: &Port,
candidates: Vec<Port::Candidate>,
request: AiCandidateResolutionRequest<'_>,
) -> Result<AiCandidateResolutionOutcome<Port::Eligible, Port::Skipped>, Port::Error>
where
Port: AiCandidateResolutionPort,
{
let normalized_client_api_format = request.client_api_format.trim().to_ascii_lowercase();
let requested_model = request
.requested_model
.map(str::trim)
.filter(|value| !value.is_empty());
let mut eligible = Vec::with_capacity(candidates.len());
let mut skipped = Vec::with_capacity(candidates.len());
for candidate in candidates {
let Some(transport) = port.read_candidate_transport(&candidate).await? else {
skipped.push(port.build_missing_transport_skipped_candidate(candidate));
continue;
};
match candidate_skip_reason_for_mode(
port,
request.mode,
&candidate,
&transport,
normalized_client_api_format.as_str(),
requested_model,
) {
Some(skip_reason) => {
skipped.push(port.build_skipped_candidate(candidate, transport, skip_reason));
}
None => {
eligible.push(port.build_eligible_candidate(candidate, transport));
}
}
}
let ranked = port
.rank_eligible_candidates(eligible, normalized_client_api_format.as_str())
.await?;
let (ranked, pool_skipped) = port.apply_pool_scheduler(ranked).await?;
skipped.extend(pool_skipped);
Ok(AiCandidateResolutionOutcome {
eligible_candidates: ranked,
skipped_candidates: skipped,
})
}
fn candidate_skip_reason_for_mode<Port>(
port: &Port,
mode: AiCandidateResolutionMode,
candidate: &Port::Candidate,
transport: &Port::Transport,
normalized_client_api_format: &str,
requested_model: Option<&str>,
) -> Option<&'static str>
where
Port: AiCandidateResolutionPort,
{
port.candidate_common_skip_reason(candidate, transport, requested_model)
.or_else(|| match mode {
AiCandidateResolutionMode::Standard => port.candidate_transport_pair_skip_reason(
candidate,
transport,
normalized_client_api_format,
requested_model.unwrap_or_default(),
),
AiCandidateResolutionMode::WithoutTransportPairGate => None,
})
}
pub fn extract_ai_pool_sticky_session_token(body_json: &serde_json::Value) -> Option<String> {
fn non_empty_str(value: Option<&serde_json::Value>) -> Option<&str> {
value
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
let object = body_json.as_object()?;
non_empty_str(object.get("prompt_cache_key"))
.or_else(|| non_empty_str(object.get("conversation_id")))
.or_else(|| non_empty_str(object.get("conversationId")))
.or_else(|| non_empty_str(object.get("session_id")))
.or_else(|| non_empty_str(object.get("sessionId")))
.or_else(|| {
object
.get("metadata")
.and_then(serde_json::Value::as_object)
.and_then(|metadata| {
non_empty_str(metadata.get("session_id"))
.or_else(|| non_empty_str(metadata.get("conversation_id")))
})
})
.or_else(|| {
object
.get("conversationState")
.and_then(serde_json::Value::as_object)
.and_then(|state| {
non_empty_str(state.get("conversationId"))
.or_else(|| non_empty_str(state.get("sessionId")))
})
})
.map(ToOwned::to_owned)
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use std::sync::Mutex;
#[derive(Default)]
struct TestPort {
calls: Mutex<Vec<String>>,
}
#[async_trait]
impl AiCandidateResolutionPort for TestPort {
type Candidate = &'static str;
type Transport = &'static str;
type Eligible = String;
type Skipped = String;
type Error = std::convert::Infallible;
async fn read_candidate_transport(
&self,
candidate: &Self::Candidate,
) -> Result<Option<Self::Transport>, Self::Error> {
self.calls
.lock()
.unwrap()
.push(format!("transport:{candidate}"));
Ok(match *candidate {
"missing" => None,
"inactive" => Some("inactive-transport"),
_ => Some("active-transport"),
})
}
fn build_missing_transport_skipped_candidate(
&self,
candidate: Self::Candidate,
) -> Self::Skipped {
format!("{candidate}:transport_snapshot_missing")
}
fn candidate_common_skip_reason(
&self,
candidate: &Self::Candidate,
_transport: &Self::Transport,
requested_model: Option<&str>,
) -> Option<&'static str> {
self.calls.lock().unwrap().push(format!(
"common:{candidate}:{}",
requested_model.unwrap_or_default()
));
(*candidate == "inactive").then_some("provider_inactive")
}
fn candidate_transport_pair_skip_reason(
&self,
candidate: &Self::Candidate,
_transport: &Self::Transport,
normalized_client_api_format: &str,
requested_model: &str,
) -> Option<&'static str> {
self.calls.lock().unwrap().push(format!(
"pair:{candidate}:{normalized_client_api_format}:{requested_model}"
));
(*candidate == "unsupported").then_some("transport_unsupported")
}
fn build_skipped_candidate(
&self,
candidate: Self::Candidate,
_transport: Self::Transport,
skip_reason: &'static str,
) -> Self::Skipped {
format!("{candidate}:{skip_reason}")
}
fn build_eligible_candidate(
&self,
candidate: Self::Candidate,
_transport: Self::Transport,
) -> Self::Eligible {
format!("eligible:{candidate}")
}
async fn rank_eligible_candidates(
&self,
mut candidates: Vec<Self::Eligible>,
normalized_client_api_format: &str,
) -> Result<Vec<Self::Eligible>, Self::Error> {
self.calls
.lock()
.unwrap()
.push(format!("rank:{normalized_client_api_format}"));
candidates.reverse();
Ok(candidates)
}
async fn apply_pool_scheduler(
&self,
candidates: Vec<Self::Eligible>,
) -> Result<(Vec<Self::Eligible>, Vec<Self::Skipped>), Self::Error> {
self.calls.lock().unwrap().push("pool".to_string());
Ok((candidates, vec!["pool:cooldown".to_string()]))
}
}
#[tokio::test]
async fn resolution_reads_transport_gates_candidates_then_ranks_and_applies_pool() {
let port = TestPort::default();
let outcome = run_ai_candidate_resolution(
&port,
vec!["first", "missing", "inactive", "unsupported", "second"],
AiCandidateResolutionRequest::standard(" OpenAI:Chat ", Some(" gpt-4.1 ")),
)
.await
.unwrap();
assert_eq!(
outcome.eligible_candidates,
["eligible:second", "eligible:first"]
);
assert_eq!(
outcome.skipped_candidates,
[
"missing:transport_snapshot_missing",
"inactive:provider_inactive",
"unsupported:transport_unsupported",
"pool:cooldown",
]
);
assert_eq!(
port.calls.lock().unwrap().as_slice(),
[
"transport:first",
"common:first:gpt-4.1",
"pair:first:openai:chat:gpt-4.1",
"transport:missing",
"transport:inactive",
"common:inactive:gpt-4.1",
"transport:unsupported",
"common:unsupported:gpt-4.1",
"pair:unsupported:openai:chat:gpt-4.1",
"transport:second",
"common:second:gpt-4.1",
"pair:second:openai:chat:gpt-4.1",
"rank:openai:chat",
"pool",
]
);
}
#[tokio::test]
async fn resolution_mode_can_skip_transport_pair_gate() {
let port = TestPort::default();
let outcome = run_ai_candidate_resolution(
&port,
vec!["unsupported"],
AiCandidateResolutionRequest::without_transport_pair_gate("openai:chat", None),
)
.await
.unwrap();
assert_eq!(outcome.eligible_candidates, ["eligible:unsupported"]);
assert_eq!(outcome.skipped_candidates, ["pool:cooldown"]);
assert_eq!(
port.calls.lock().unwrap().as_slice(),
[
"transport:unsupported",
"common:unsupported:",
"rank:openai:chat",
"pool",
]
);
}
#[test]
fn sticky_session_token_is_extracted_from_known_request_fields() {
assert_eq!(
extract_ai_pool_sticky_session_token(&json!({
"prompt_cache_key": " cache-a ",
"conversation_id": "conversation-b"
}))
.as_deref(),
Some("cache-a")
);
assert_eq!(
extract_ai_pool_sticky_session_token(&json!({
"metadata": {"conversation_id": " conversation-c "}
}))
.as_deref(),
Some("conversation-c")
);
assert_eq!(
extract_ai_pool_sticky_session_token(&json!({
"conversationState": {"sessionId": " session-d "}
}))
.as_deref(),
Some("session-d")
);
}
}

196
crates/aether-ai-serving/src/decision_input.rs Normal file
View File

@@ -0,0 +1,196 @@
use async_trait::async_trait;
#[async_trait]
pub trait AiAuthenticatedDecisionInputPort: Send + Sync {
type AuthContext: Send + Sync;
type AuthSnapshot: Send;
type RequiredCapabilities: Send + Sync;
type ResolvedInput: Send;
type Error: Send;
async fn read_auth_snapshot(
&self,
auth_context: &Self::AuthContext,
) -> Result<Option<Self::AuthSnapshot>, Self::Error>;
async fn resolve_required_capabilities(
&self,
auth_context: &Self::AuthContext,
requested_model: Option<&str>,
explicit_required_capabilities: Option<&Self::RequiredCapabilities>,
) -> Result<Option<Self::RequiredCapabilities>, Self::Error>;
fn build_resolved_input(
&self,
auth_context: Self::AuthContext,
auth_snapshot: Self::AuthSnapshot,
required_capabilities: Option<Self::RequiredCapabilities>,
) -> Self::ResolvedInput;
}
pub async fn run_ai_authenticated_decision_input<Port>(
port: &Port,
auth_context: Port::AuthContext,
requested_model: Option<&str>,
explicit_required_capabilities: Option<&Port::RequiredCapabilities>,
) -> Result<Option<Port::ResolvedInput>, Port::Error>
where
Port: AiAuthenticatedDecisionInputPort,
{
let auth_snapshot = match port.read_auth_snapshot(&auth_context).await? {
Some(snapshot) => snapshot,
None => return Ok(None),
};
let required_capabilities = port
.resolve_required_capabilities(
&auth_context,
requested_model,
explicit_required_capabilities,
)
.await?;
Ok(Some(port.build_resolved_input(
auth_context,
auth_snapshot,
required_capabilities,
)))
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Mutex;
#[derive(Debug, Clone, PartialEq, Eq)]
struct TestAuthContext {
user_id: &'static str,
api_key_id: &'static str,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct TestResolvedInput {
auth_context: TestAuthContext,
auth_snapshot: &'static str,
required_capabilities: Option<String>,
}
struct TestPort {
auth_snapshot: Option<&'static str>,
calls: Mutex<Vec<String>>,
}
#[async_trait]
impl AiAuthenticatedDecisionInputPort for TestPort {
type AuthContext = TestAuthContext;
type AuthSnapshot = &'static str;
type RequiredCapabilities = String;
type ResolvedInput = TestResolvedInput;
type Error = std::convert::Infallible;
async fn read_auth_snapshot(
&self,
auth_context: &Self::AuthContext,
) -> Result<Option<Self::AuthSnapshot>, Self::Error> {
self.calls.lock().unwrap().push(format!(
"snapshot:{}:{}",
auth_context.user_id, auth_context.api_key_id
));
Ok(self.auth_snapshot)
}
async fn resolve_required_capabilities(
&self,
auth_context: &Self::AuthContext,
requested_model: Option<&str>,
explicit_required_capabilities: Option<&Self::RequiredCapabilities>,
) -> Result<Option<Self::RequiredCapabilities>, Self::Error> {
self.calls.lock().unwrap().push(format!(
"capabilities:{}:{}:{}",
auth_context.user_id,
requested_model.unwrap_or_default(),
explicit_required_capabilities
.map(String::as_str)
.unwrap_or_default()
));
Ok(Some("merged-capabilities".to_string()))
}
fn build_resolved_input(
&self,
auth_context: Self::AuthContext,
auth_snapshot: Self::AuthSnapshot,
required_capabilities: Option<Self::RequiredCapabilities>,
) -> Self::ResolvedInput {
TestResolvedInput {
auth_context,
auth_snapshot,
required_capabilities,
}
}
}
#[tokio::test]
async fn authenticated_decision_input_resolves_snapshot_and_capabilities() {
let port = TestPort {
auth_snapshot: Some("snapshot-a"),
calls: Mutex::new(Vec::new()),
};
let auth_context = TestAuthContext {
user_id: "user-a",
api_key_id: "key-a",
};
let explicit = "explicit-capability".to_string();
let resolved = run_ai_authenticated_decision_input(
&port,
auth_context.clone(),
Some("model-a"),
Some(&explicit),
)
.await
.unwrap();
assert_eq!(
resolved,
Some(TestResolvedInput {
auth_context,
auth_snapshot: "snapshot-a",
required_capabilities: Some("merged-capabilities".to_string()),
})
);
assert_eq!(
port.calls.lock().unwrap().as_slice(),
[
"snapshot:user-a:key-a",
"capabilities:user-a:model-a:explicit-capability",
]
);
}
#[tokio::test]
async fn authenticated_decision_input_stops_when_snapshot_is_missing() {
let port = TestPort {
auth_snapshot: None,
calls: Mutex::new(Vec::new()),
};
let resolved = run_ai_authenticated_decision_input(
&port,
TestAuthContext {
user_id: "user-a",
api_key_id: "key-a",
},
Some("model-a"),
None,
)
.await
.unwrap();
assert_eq!(resolved, None);
assert_eq!(
port.calls.lock().unwrap().as_slice(),
["snapshot:user-a:key-a"]
);
}
}

215
crates/aether-ai-serving/src/decision_path.rs Normal file
View File

@@ -0,0 +1,215 @@
use async_trait::async_trait;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum AiSyncDecisionStep {
VideoTaskFollowUp,
LocalVideo,
LocalImage,
LocalOpenAiChat,
LocalOpenAiResponses,
LocalStandardFamily,
LocalSameFormatProvider,
LocalGeminiFiles,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AiStreamDecisionStep {
LocalVideoContent,
LocalImage,
LocalOpenAiChat,
LocalOpenAiResponses,
LocalStandardFamily,
LocalSameFormatProvider,
LocalGeminiFiles,
}
#[async_trait]
pub trait AiSyncDecisionPathPort: Send + Sync {
type Decision: Send;
type Error: Send;
fn sync_decision_step_enabled(&self, step: AiSyncDecisionStep) -> bool;
async fn build_sync_decision_step(
&self,
step: AiSyncDecisionStep,
) -> Result<Option<Self::Decision>, Self::Error>;
}
#[async_trait]
pub trait AiStreamDecisionPathPort: Send + Sync {
type Decision: Send;
type Error: Send;
async fn build_stream_decision_step(
&self,
step: AiStreamDecisionStep,
) -> Result<Option<Self::Decision>, Self::Error>;
}
pub async fn run_ai_sync_decision_path<Port>(
port: &Port,
) -> Result<Option<Port::Decision>, Port::Error>
where
Port: AiSyncDecisionPathPort,
{
for step in [
AiSyncDecisionStep::VideoTaskFollowUp,
AiSyncDecisionStep::LocalVideo,
AiSyncDecisionStep::LocalImage,
AiSyncDecisionStep::LocalOpenAiChat,
AiSyncDecisionStep::LocalOpenAiResponses,
AiSyncDecisionStep::LocalStandardFamily,
AiSyncDecisionStep::LocalSameFormatProvider,
AiSyncDecisionStep::LocalGeminiFiles,
] {
if !port.sync_decision_step_enabled(step) {
continue;
}
if let Some(decision) = port.build_sync_decision_step(step).await? {
return Ok(Some(decision));
}
}
Ok(None)
}
pub async fn run_ai_stream_decision_path<Port>(
port: &Port,
) -> Result<Option<Port::Decision>, Port::Error>
where
Port: AiStreamDecisionPathPort,
{
for step in [
AiStreamDecisionStep::LocalVideoContent,
AiStreamDecisionStep::LocalImage,
AiStreamDecisionStep::LocalOpenAiChat,
AiStreamDecisionStep::LocalOpenAiResponses,
AiStreamDecisionStep::LocalStandardFamily,
AiStreamDecisionStep::LocalSameFormatProvider,
AiStreamDecisionStep::LocalGeminiFiles,
] {
if let Some(decision) = port.build_stream_decision_step(step).await? {
return Ok(Some(decision));
}
}
Ok(None)
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::{BTreeSet, VecDeque};
use std::sync::Mutex;
#[derive(Default)]
struct TestSyncDecisionPort {
disabled: BTreeSet<AiSyncDecisionStep>,
outcomes: Mutex<VecDeque<Option<&'static str>>>,
calls: Mutex<Vec<String>>,
}
#[async_trait]
impl AiSyncDecisionPathPort for TestSyncDecisionPort {
type Decision = &'static str;
type Error = std::convert::Infallible;
fn sync_decision_step_enabled(&self, step: AiSyncDecisionStep) -> bool {
!self.disabled.contains(&step)
}
async fn build_sync_decision_step(
&self,
step: AiSyncDecisionStep,
) -> Result<Option<Self::Decision>, Self::Error> {
self.calls.lock().unwrap().push(format!("{step:?}"));
Ok(self.outcomes.lock().unwrap().pop_front().unwrap_or(None))
}
}
#[derive(Default)]
struct TestStreamDecisionPort {
outcomes: Mutex<VecDeque<Option<&'static str>>>,
calls: Mutex<Vec<String>>,
}
#[async_trait]
impl AiStreamDecisionPathPort for TestStreamDecisionPort {
type Decision = &'static str;
type Error = std::convert::Infallible;
async fn build_stream_decision_step(
&self,
step: AiStreamDecisionStep,
) -> Result<Option<Self::Decision>, Self::Error> {
self.calls.lock().unwrap().push(format!("{step:?}"));
Ok(self.outcomes.lock().unwrap().pop_front().unwrap_or(None))
}
}
#[tokio::test]
async fn sync_decision_path_runs_steps_in_serving_order() {
let port = TestSyncDecisionPort::default();
let decision = run_ai_sync_decision_path(&port).await.unwrap();
assert_eq!(decision, None);
assert_eq!(
port.calls.lock().unwrap().as_slice(),
[
"VideoTaskFollowUp",
"LocalVideo",
"LocalImage",
"LocalOpenAiChat",
"LocalOpenAiResponses",
"LocalStandardFamily",
"LocalSameFormatProvider",
"LocalGeminiFiles",
]
);
}
#[tokio::test]
async fn sync_decision_path_skips_disabled_steps_and_stops_at_first_decision() {
let port = TestSyncDecisionPort {
disabled: BTreeSet::from([AiSyncDecisionStep::LocalGeminiFiles]),
outcomes: Mutex::new(VecDeque::from([
None,
None,
Some("image_decision"),
Some("should_not_run"),
])),
calls: Mutex::default(),
};
let decision = run_ai_sync_decision_path(&port).await.unwrap();
assert_eq!(decision, Some("image_decision"));
assert_eq!(
port.calls.lock().unwrap().as_slice(),
["VideoTaskFollowUp", "LocalVideo", "LocalImage"]
);
}
#[tokio::test]
async fn stream_decision_path_stops_at_first_decision() {
let port = TestStreamDecisionPort {
outcomes: Mutex::new(VecDeque::from([
None,
None,
Some("chat_decision"),
Some("should_not_run"),
])),
calls: Mutex::default(),
};
let decision = run_ai_stream_decision_path(&port).await.unwrap();
assert_eq!(decision, Some("chat_decision"));
assert_eq!(
port.calls.lock().unwrap().as_slice(),
["LocalVideoContent", "LocalImage", "LocalOpenAiChat"]
);
}
}

92
crates/aether-ai-serving/src/decision_payload.rs Normal file
View File

@@ -0,0 +1,92 @@
use std::collections::BTreeMap;
use aether_ai_surfaces::api::{
ExecutionRuntimeAuthContext, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
};
use aether_contracts::{ExecutionTimeouts, ProxySnapshot};
use crate::{AiExecutionDecision, ConversionMode, ExecutionStrategy};
pub struct AiExecutionDecisionResponseParts {
pub decision_is_stream: bool,
pub decision_kind: String,
pub execution_strategy: ExecutionStrategy,
pub conversion_mode: ConversionMode,
pub request_id: String,
pub candidate_id: String,
pub provider_name: String,
pub provider_id: String,
pub endpoint_id: String,
pub key_id: String,
pub upstream_base_url: String,
pub upstream_url: String,
pub provider_request_method: Option<String>,
pub auth_header: Option<String>,
pub auth_value: Option<String>,
pub provider_api_format: String,
pub client_api_format: String,
pub model_name: String,
pub mapped_model: String,
pub prompt_cache_key: Option<String>,
pub provider_request_headers: BTreeMap<String, String>,
pub provider_request_body: Option<serde_json::Value>,
pub provider_request_body_base64: Option<String>,
pub content_type: Option<String>,
pub proxy: Option<ProxySnapshot>,
pub tls_profile: Option<String>,
pub timeouts: Option<ExecutionTimeouts>,
pub upstream_is_stream: bool,
pub report_kind: Option<String>,
pub report_context: Option<serde_json::Value>,
pub auth_context: ExecutionRuntimeAuthContext,
}
pub fn build_ai_execution_decision_response(
parts: AiExecutionDecisionResponseParts,
) -> AiExecutionDecision {
AiExecutionDecision {
action: ai_execution_decision_action(parts.decision_is_stream).to_string(),
decision_kind: Some(parts.decision_kind),
execution_strategy: Some(parts.execution_strategy.as_str().to_string()),
conversion_mode: Some(parts.conversion_mode.as_str().to_string()),
request_id: Some(parts.request_id),
candidate_id: Some(parts.candidate_id),
provider_name: Some(parts.provider_name),
provider_id: Some(parts.provider_id),
endpoint_id: Some(parts.endpoint_id),
key_id: Some(parts.key_id),
upstream_base_url: Some(parts.upstream_base_url),
upstream_url: Some(parts.upstream_url),
provider_request_method: parts.provider_request_method,
auth_header: parts.auth_header,
auth_value: parts.auth_value,
provider_api_format: Some(parts.provider_api_format.clone()),
client_api_format: Some(parts.client_api_format.clone()),
provider_contract: Some(parts.provider_api_format),
client_contract: Some(parts.client_api_format),
model_name: Some(parts.model_name),
mapped_model: Some(parts.mapped_model),
prompt_cache_key: parts.prompt_cache_key,
extra_headers: BTreeMap::new(),
provider_request_headers: parts.provider_request_headers,
provider_request_body: parts.provider_request_body,
provider_request_body_base64: parts.provider_request_body_base64,
content_type: parts.content_type,
proxy: parts.proxy,
tls_profile: parts.tls_profile,
timeouts: parts.timeouts,
upstream_is_stream: parts.upstream_is_stream,
report_kind: parts.report_kind,
report_context: parts.report_context,
auth_context: Some(parts.auth_context),
}
}
pub const fn ai_execution_decision_action(decision_is_stream: bool) -> &'static str {
if decision_is_stream {
EXECUTION_RUNTIME_STREAM_DECISION_ACTION
} else {
EXECUTION_RUNTIME_SYNC_DECISION_ACTION
}
}

147
crates/aether-ai-pipeline/src/contracts/control_payloads.rs → crates/aether-ai-serving/src/dto.rs
View File

@@ -1,25 +1,51 @@
use std::collections::BTreeMap;
use aether_ai_surfaces::api::ExecutionRuntimeAuthContext;
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, ProxySnapshot};
use serde::{Deserialize, Serialize};
use crate::contracts::ExecutionRuntimeAuthContext;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ExecutionStrategy {
GatewayAffinityForward,
RawPublicProxy,
LocalSameFormat,
LocalCrossFormat,
}
#[derive(Debug, Serialize)]
pub struct GatewayControlPlanRequest {
pub trace_id: String,
pub method: String,
pub path: String,
pub query_string: Option<String>,
pub headers: BTreeMap<String, String>,
pub body_json: serde_json::Value,
#[serde(skip_serializing_if = "Option::is_none")]
pub body_base64: Option<String>,
pub auth_context: Option<ExecutionRuntimeAuthContext>,
impl ExecutionStrategy {
pub const fn as_str(self) -> &'static str {
match self {
Self::GatewayAffinityForward => "gateway_affinity_forward",
Self::RawPublicProxy => "raw_public_proxy",
Self::LocalSameFormat => "local_same_format",
Self::LocalCrossFormat => "local_cross_format",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ConversionMode {
None,
RequestOnly,
ResponseOnly,
Bidirectional,
}
impl ConversionMode {
pub const fn as_str(self) -> &'static str {
match self {
Self::None => "none",
Self::RequestOnly => "request_only",
Self::ResponseOnly => "response_only",
Self::Bidirectional => "bidirectional",
}
}
}
#[derive(Debug, Deserialize, Serialize)]
pub struct GatewayControlPlanResponse {
pub struct AiExecutionPlanPayload {
pub action: String,
#[serde(default)]
pub plan_kind: Option<String>,
@@ -34,7 +60,7 @@ pub struct GatewayControlPlanResponse {
}
#[derive(Debug, Deserialize, Serialize)]
pub struct GatewayControlSyncDecisionResponse {
pub struct AiExecutionDecision {
pub action: String,
#[serde(default)]
pub decision_kind: Option<String>,
@@ -105,42 +131,19 @@ pub struct GatewayControlSyncDecisionResponse {
}
#[derive(Debug)]
pub struct LocalSyncPlanAndReport {
pub struct AiSyncAttempt {
pub plan: ExecutionPlan,
pub report_kind: Option<String>,
pub report_context: Option<serde_json::Value>,
}
#[derive(Debug)]
pub struct LocalStreamPlanAndReport {
pub struct AiStreamAttempt {
pub plan: ExecutionPlan,
pub report_kind: Option<String>,
pub report_context: Option<serde_json::Value>,
}
#[allow(clippy::too_many_arguments)]
pub fn build_gateway_control_plan_request(
trace_id: &str,
method: &str,
path: &str,
query_string: Option<&str>,
headers: BTreeMap<String, String>,
body_json: serde_json::Value,
body_base64: Option<String>,
auth_context: Option<ExecutionRuntimeAuthContext>,
) -> GatewayControlPlanRequest {
GatewayControlPlanRequest {
trace_id: trace_id.to_string(),
method: method.to_string(),
path: path.to_string(),
query_string: query_string.map(ToOwned::to_owned),
headers,
body_json,
body_base64,
auth_context,
}
}
pub fn augment_sync_report_context(
report_context: Option<serde_json::Value>,
provider_request_headers: &BTreeMap<String, String>,
@@ -160,7 +163,7 @@ pub fn augment_sync_report_context(
Ok(Some(serde_json::Value::Object(report_context)))
}
fn decision_has_exact_provider_request(payload: &GatewayControlSyncDecisionResponse) -> bool {
fn decision_has_exact_provider_request(payload: &AiExecutionDecision) -> bool {
!payload.provider_request_headers.is_empty()
&& (payload.provider_request_body.is_some()
|| payload
@@ -170,9 +173,7 @@ fn decision_has_exact_provider_request(payload: &GatewayControlSyncDecisionRespo
.unwrap_or(false))
}
pub fn generic_decision_missing_exact_provider_request(
payload: &GatewayControlSyncDecisionResponse,
) -> bool {
pub fn generic_decision_missing_exact_provider_request(payload: &AiExecutionDecision) -> bool {
!decision_has_exact_provider_request(payload)
}
@@ -181,14 +182,13 @@ mod tests {
use std::collections::BTreeMap;
use super::{
augment_sync_report_context, build_gateway_control_plan_request,
generic_decision_missing_exact_provider_request, ExecutionRuntimeAuthContext,
GatewayControlSyncDecisionResponse,
augment_sync_report_context, generic_decision_missing_exact_provider_request,
AiExecutionDecision,
};
#[test]
fn generic_decision_detects_missing_exact_provider_request() {
let payload = GatewayControlSyncDecisionResponse {
let payload = AiExecutionDecision {
action: "local".to_string(),
decision_kind: Some("sync".to_string()),
execution_strategy: None,
@@ -239,55 +239,12 @@ mod tests {
.expect("context should exist");
assert_eq!(
report_context.get("trace_id"),
Some(&serde_json::json!("abc"))
report_context["provider_request_headers"]["content-type"],
"application/json"
);
assert_eq!(
report_context
.get("provider_request_headers")
.and_then(|value| value.get("content-type")),
Some(&serde_json::json!("application/json"))
);
assert!(report_context.get("provider_request_body").is_none());
}
#[test]
fn build_gateway_control_plan_request_preserves_request_shape() {
let payload = build_gateway_control_plan_request(
"trace-123",
"POST",
"/v1/chat/completions",
Some("stream=true"),
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
serde_json::json!({"model": "gpt-5"}),
Some("eyJmb28iOiJiYXIifQ==".to_string()),
Some(ExecutionRuntimeAuthContext {
user_id: "user-1".to_string(),
api_key_id: "key-1".to_string(),
username: None,
api_key_name: None,
balance_remaining: Some(12.5),
access_allowed: true,
api_key_is_standalone: false,
}),
);
assert_eq!(payload.trace_id, "trace-123");
assert_eq!(payload.method, "POST");
assert_eq!(payload.path, "/v1/chat/completions");
assert_eq!(payload.query_string.as_deref(), Some("stream=true"));
assert_eq!(
payload.headers.get("content-type").map(String::as_str),
Some("application/json")
);
assert_eq!(payload.body_json, serde_json::json!({"model": "gpt-5"}));
assert_eq!(payload.body_base64.as_deref(), Some("eyJmb28iOiJiYXIifQ=="));
assert_eq!(
payload
.auth_context
.as_ref()
.map(|ctx| ctx.user_id.as_str()),
Some("user-1")
assert!(
report_context.get("provider_request_body").is_none(),
"provider request body should not be copied into report context"
);
}
}

408
crates/aether-ai-serving/src/execution_path.rs Normal file
View File

@@ -0,0 +1,408 @@
use async_trait::async_trait;
#[derive(Debug)]
pub enum AiServingExecutionOutcome<Response, Exhaustion> {
Responded(Response),
Exhausted(Exhaustion),
NoPath,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AiPlanFallbackReason {
RemoteDecisionMiss,
SchedulerDecisionUnsupported,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AiSyncExecutionStep {
VideoTaskFollowUp,
LocalVideo,
LocalImage,
LocalOpenAiChat,
LocalOpenAiResponses,
LocalStandardFamily,
LocalSameFormatProvider,
LocalGeminiFiles,
RemoteDecision,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AiStreamExecutionStep {
LocalVideoContent,
LocalImage,
LocalOpenAiChat,
LocalOpenAiResponses,
LocalStandardFamily,
LocalSameFormatProvider,
LocalGeminiFiles,
RemoteDecision,
}
#[async_trait]
pub trait AiSyncExecutionPathPort: Send + Sync {
type Response: Send;
type Exhaustion: Send;
type Error: Send;
fn scheduler_decision_supported(&self) -> bool;
async fn execute_sync_step(
&self,
step: AiSyncExecutionStep,
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error>;
async fn execute_sync_plan_fallback(
&self,
reason: AiPlanFallbackReason,
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error>;
}
#[async_trait]
pub trait AiStreamExecutionPathPort: Send + Sync {
type Response: Send;
type Exhaustion: Send;
type Error: Send;
fn scheduler_decision_supported(&self) -> bool;
async fn execute_stream_step(
&self,
step: AiStreamExecutionStep,
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error>;
async fn execute_stream_plan_fallback(
&self,
reason: AiPlanFallbackReason,
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error>;
}
pub async fn run_ai_sync_execution_path<Port>(
port: &Port,
) -> Result<AiServingExecutionOutcome<Port::Response, Port::Exhaustion>, Port::Error>
where
Port: AiSyncExecutionPathPort,
{
let mut exhausted = None;
if let Some(response) =
absorb_sync_step(port, AiSyncExecutionStep::VideoTaskFollowUp, &mut exhausted).await?
{
return Ok(response);
}
if port.scheduler_decision_supported() {
for step in [
AiSyncExecutionStep::LocalVideo,
AiSyncExecutionStep::LocalImage,
AiSyncExecutionStep::LocalOpenAiChat,
AiSyncExecutionStep::LocalOpenAiResponses,
AiSyncExecutionStep::LocalStandardFamily,
AiSyncExecutionStep::LocalSameFormatProvider,
AiSyncExecutionStep::LocalGeminiFiles,
AiSyncExecutionStep::RemoteDecision,
] {
if let Some(response) = absorb_sync_step(port, step, &mut exhausted).await? {
return Ok(response);
}
}
}
let fallback_reason = if port.scheduler_decision_supported() {
AiPlanFallbackReason::RemoteDecisionMiss
} else {
AiPlanFallbackReason::SchedulerDecisionUnsupported
};
match port.execute_sync_plan_fallback(fallback_reason).await? {
AiServingExecutionOutcome::Responded(response) => {
Ok(AiServingExecutionOutcome::Responded(response))
}
AiServingExecutionOutcome::Exhausted(outcome) => {
Ok(AiServingExecutionOutcome::Exhausted(outcome))
}
AiServingExecutionOutcome::NoPath => Ok(exhausted
.map(AiServingExecutionOutcome::Exhausted)
.unwrap_or(AiServingExecutionOutcome::NoPath)),
}
}
pub async fn run_ai_stream_execution_path<Port>(
port: &Port,
) -> Result<AiServingExecutionOutcome<Port::Response, Port::Exhaustion>, Port::Error>
where
Port: AiStreamExecutionPathPort,
{
let mut exhausted = None;
if let Some(response) = absorb_stream_step(
port,
AiStreamExecutionStep::LocalVideoContent,
&mut exhausted,
)
.await?
{
return Ok(response);
}
if port.scheduler_decision_supported() {
for step in [
AiStreamExecutionStep::LocalImage,
AiStreamExecutionStep::LocalOpenAiChat,
AiStreamExecutionStep::LocalOpenAiResponses,
AiStreamExecutionStep::LocalStandardFamily,
AiStreamExecutionStep::LocalSameFormatProvider,
AiStreamExecutionStep::LocalGeminiFiles,
AiStreamExecutionStep::RemoteDecision,
] {
if let Some(response) = absorb_stream_step(port, step, &mut exhausted).await? {
return Ok(response);
}
}
}
let fallback_reason = if port.scheduler_decision_supported() {
AiPlanFallbackReason::RemoteDecisionMiss
} else {
AiPlanFallbackReason::SchedulerDecisionUnsupported
};
match port.execute_stream_plan_fallback(fallback_reason).await? {
AiServingExecutionOutcome::Responded(response) => {
Ok(AiServingExecutionOutcome::Responded(response))
}
AiServingExecutionOutcome::Exhausted(outcome) => {
Ok(AiServingExecutionOutcome::Exhausted(outcome))
}
AiServingExecutionOutcome::NoPath => Ok(exhausted
.map(AiServingExecutionOutcome::Exhausted)
.unwrap_or(AiServingExecutionOutcome::NoPath)),
}
}
async fn absorb_sync_step<Port>(
port: &Port,
step: AiSyncExecutionStep,
exhausted: &mut Option<Port::Exhaustion>,
) -> Result<Option<AiServingExecutionOutcome<Port::Response, Port::Exhaustion>>, Port::Error>
where
Port: AiSyncExecutionPathPort,
{
match port.execute_sync_step(step).await? {
AiServingExecutionOutcome::Responded(response) => {
Ok(Some(AiServingExecutionOutcome::Responded(response)))
}
AiServingExecutionOutcome::Exhausted(outcome) => {
*exhausted = Some(outcome);
Ok(None)
}
AiServingExecutionOutcome::NoPath => Ok(None),
}
}
async fn absorb_stream_step<Port>(
port: &Port,
step: AiStreamExecutionStep,
exhausted: &mut Option<Port::Exhaustion>,
) -> Result<Option<AiServingExecutionOutcome<Port::Response, Port::Exhaustion>>, Port::Error>
where
Port: AiStreamExecutionPathPort,
{
match port.execute_stream_step(step).await? {
AiServingExecutionOutcome::Responded(response) => {
Ok(Some(AiServingExecutionOutcome::Responded(response)))
}
AiServingExecutionOutcome::Exhausted(outcome) => {
*exhausted = Some(outcome);
Ok(None)
}
AiServingExecutionOutcome::NoPath => Ok(None),
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::VecDeque;
use std::sync::Mutex;
#[derive(Default)]
struct TestSyncPort {
scheduler_supported: bool,
outcomes: Mutex<VecDeque<AiServingExecutionOutcome<&'static str, &'static str>>>,
calls: Mutex<Vec<String>>,
}
#[async_trait]
impl AiSyncExecutionPathPort for TestSyncPort {
type Response = &'static str;
type Exhaustion = &'static str;
type Error = std::convert::Infallible;
fn scheduler_decision_supported(&self) -> bool {
self.scheduler_supported
}
async fn execute_sync_step(
&self,
step: AiSyncExecutionStep,
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error>
{
self.calls.lock().unwrap().push(format!("{step:?}"));
Ok(self
.outcomes
.lock()
.unwrap()
.pop_front()
.unwrap_or(AiServingExecutionOutcome::NoPath))
}
async fn execute_sync_plan_fallback(
&self,
reason: AiPlanFallbackReason,
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error>
{
self.calls
.lock()
.unwrap()
.push(format!("Fallback:{reason:?}"));
Ok(self
.outcomes
.lock()
.unwrap()
.pop_front()
.unwrap_or(AiServingExecutionOutcome::NoPath))
}
}
#[derive(Default)]
struct TestStreamPort {
scheduler_supported: bool,
outcomes: Mutex<VecDeque<AiServingExecutionOutcome<&'static str, &'static str>>>,
calls: Mutex<Vec<String>>,
}
#[async_trait]
impl AiStreamExecutionPathPort for TestStreamPort {
type Response = &'static str;
type Exhaustion = &'static str;
type Error = std::convert::Infallible;
fn scheduler_decision_supported(&self) -> bool {
self.scheduler_supported
}
async fn execute_stream_step(
&self,
step: AiStreamExecutionStep,
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error>
{
self.calls.lock().unwrap().push(format!("{step:?}"));
Ok(self
.outcomes
.lock()
.unwrap()
.pop_front()
.unwrap_or(AiServingExecutionOutcome::NoPath))
}
async fn execute_stream_plan_fallback(
&self,
reason: AiPlanFallbackReason,
) -> Result<AiServingExecutionOutcome<Self::Response, Self::Exhaustion>, Self::Error>
{
self.calls
.lock()
.unwrap()
.push(format!("Fallback:{reason:?}"));
Ok(self
.outcomes
.lock()
.unwrap()
.pop_front()
.unwrap_or(AiServingExecutionOutcome::NoPath))
}
}
#[tokio::test]
async fn sync_path_runs_scheduler_steps_before_remote_and_fallback() {
let port = TestSyncPort {
scheduler_supported: true,
..TestSyncPort::default()
};
let outcome = run_ai_sync_execution_path(&port).await.unwrap();
assert!(matches!(outcome, AiServingExecutionOutcome::NoPath));
assert_eq!(
port.calls.lock().unwrap().as_slice(),
[
"VideoTaskFollowUp",
"LocalVideo",
"LocalImage",
"LocalOpenAiChat",
"LocalOpenAiResponses",
"LocalStandardFamily",
"LocalSameFormatProvider",
"LocalGeminiFiles",
"RemoteDecision",
"Fallback:RemoteDecisionMiss",
]
);
}
#[tokio::test]
async fn sync_path_returns_last_exhaustion_when_fallback_has_no_path() {
let port = TestSyncPort {
scheduler_supported: true,
outcomes: Mutex::new(VecDeque::from([
AiServingExecutionOutcome::NoPath,
AiServingExecutionOutcome::Exhausted("local_video_exhausted"),
])),
calls: Mutex::default(),
};
let outcome = run_ai_sync_execution_path(&port).await.unwrap();
assert!(matches!(
outcome,
AiServingExecutionOutcome::Exhausted("local_video_exhausted")
));
}
#[tokio::test]
async fn stream_path_skips_scheduler_steps_when_unsupported() {
let port = TestStreamPort {
scheduler_supported: false,
..TestStreamPort::default()
};
let outcome = run_ai_stream_execution_path(&port).await.unwrap();
assert!(matches!(outcome, AiServingExecutionOutcome::NoPath));
assert_eq!(
port.calls.lock().unwrap().as_slice(),
["LocalVideoContent", "Fallback:SchedulerDecisionUnsupported",]
);
}
#[tokio::test]
async fn stream_path_stops_at_first_response() {
let port = TestStreamPort {
scheduler_supported: true,
outcomes: Mutex::new(VecDeque::from([
AiServingExecutionOutcome::NoPath,
AiServingExecutionOutcome::Responded("image_response"),
])),
calls: Mutex::default(),
};
let outcome = run_ai_stream_execution_path(&port).await.unwrap();
assert!(matches!(
outcome,
AiServingExecutionOutcome::Responded("image_response")
));
assert_eq!(
port.calls.lock().unwrap().as_slice(),
["LocalVideoContent", "LocalImage"]
);
}
}

190
crates/aether-ai-serving/src/failure_diagnostic.rs Normal file
View File

@@ -0,0 +1,190 @@
use serde_json::{json, Value};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CandidateFailureDiagnosticKind {
RequestBodyBuild,
RequestConversion,
BodyRules,
HeaderRules,
UrlBuild,
TransportAuth,
EnvelopeBuild,
}
impl CandidateFailureDiagnosticKind {
pub const fn as_str(self) -> &'static str {
match self {
Self::RequestBodyBuild => "request_body_build",
Self::RequestConversion => "request_conversion",
Self::BodyRules => "body_rules",
Self::HeaderRules => "header_rules",
Self::UrlBuild => "url_build",
Self::TransportAuth => "transport_auth",
Self::EnvelopeBuild => "envelope_build",
}
}
}
#[derive(Debug, Clone)]
pub struct CandidateFailureDiagnostic {
kind: CandidateFailureDiagnosticKind,
path: String,
message: String,
source: Option<String>,
client_api_format: Option<String>,
provider_api_format: Option<String>,
safe_to_show: bool,
}
impl CandidateFailureDiagnostic {
pub fn new(
kind: CandidateFailureDiagnosticKind,
path: impl Into<String>,
message: impl Into<String>,
) -> Self {
Self {
kind,
path: path.into(),
message: message.into(),
source: None,
client_api_format: None,
provider_api_format: None,
safe_to_show: true,
}
}
pub fn source(mut self, source: impl Into<String>) -> Self {
self.source = Some(source.into());
self
}
pub fn formats(
mut self,
client_api_format: impl Into<String>,
provider_api_format: impl Into<String>,
) -> Self {
self.client_api_format = Some(client_api_format.into());
self.provider_api_format = Some(provider_api_format.into());
self
}
pub fn to_extra_data(&self) -> Value {
let diagnostic = self.to_value();
let mut extra_data = json!({
"failure_diagnostic": diagnostic,
});
// Compatibility for current usage UI and already persisted trace readers.
if self.kind == CandidateFailureDiagnosticKind::RequestBodyBuild {
if let Some(object) = extra_data.as_object_mut() {
object.insert(
"request_body_build_error".to_string(),
json!({
"path": self.path,
"message": self.message,
"client_api_format": self.client_api_format,
"provider_api_format": self.provider_api_format,
}),
);
}
}
extra_data
}
pub fn upstream_url_missing(
client_api_format: impl Into<String>,
provider_api_format: impl Into<String>,
source: impl Into<String>,
) -> Self {
Self::new(
CandidateFailureDiagnosticKind::UrlBuild,
"$.endpoint",
"无法构建上游请求地址；请检查 base_url、custom_path、API 格式和模型映射",
)
.formats(client_api_format, provider_api_format)
.source(source)
}
pub fn header_rules_apply_failed(
client_api_format: impl Into<String>,
provider_api_format: impl Into<String>,
source: impl Into<String>,
) -> Self {
Self::new(
CandidateFailureDiagnosticKind::HeaderRules,
"$.endpoint.header_rules",
"Header 规则应用失败；请检查规则格式、条件配置，或是否试图覆盖受保护认证头",
)
.formats(client_api_format, provider_api_format)
.source(source)
}
pub fn body_rules_apply_failed(
client_api_format: impl Into<String>,
provider_api_format: impl Into<String>,
source: impl Into<String>,
) -> Self {
Self::new(
CandidateFailureDiagnosticKind::BodyRules,
"$.endpoint.body_rules",
"Body 规则应用失败；请检查规则格式、条件配置，或规则输出是否仍是当前上游支持的请求体",
)
.formats(client_api_format, provider_api_format)
.source(source)
}
pub fn body_rules_unsupported_for_binary_upload(
client_api_format: impl Into<String>,
provider_api_format: impl Into<String>,
source: impl Into<String>,
) -> Self {
Self::new(
CandidateFailureDiagnosticKind::BodyRules,
"$.endpoint.body_rules",
"二进制上传暂不支持本地应用 Body 规则；请移除该 Endpoint 的 Body 规则或改用 JSON 请求体",
)
.formats(client_api_format, provider_api_format)
.source(source)
}
pub fn provider_request_body_missing(
client_api_format: impl Into<String>,
provider_api_format: impl Into<String>,
source: impl Into<String>,
) -> Self {
Self::new(
CandidateFailureDiagnosticKind::RequestBodyBuild,
"$",
"无法构建上游请求体；请检查请求体是否为支持的 JSON object，以及该任务类型必需字段是否存在且取值受支持",
)
.formats(client_api_format, provider_api_format)
.source(source)
}
pub fn envelope_build_failed(
client_api_format: impl Into<String>,
provider_api_format: impl Into<String>,
source: impl Into<String>,
) -> Self {
Self::new(
CandidateFailureDiagnosticKind::EnvelopeBuild,
"$",
"无法构建上游请求封装；请检查该 Provider 的认证配置、模型映射、Endpoint Body 规则和当前请求体是否兼容",
)
.formats(client_api_format, provider_api_format)
.source(source)
}
fn to_value(&self) -> Value {
json!({
"kind": self.kind.as_str(),
"path": self.path,
"message": self.message,
"source": self.source,
"client_api_format": self.client_api_format,
"provider_api_format": self.provider_api_format,
"safe_to_show": self.safe_to_show,
})
}
}

131
crates/aether-ai-serving/src/lib.rs Normal file
View File

@@ -0,0 +1,131 @@
pub mod attempt_loop;
pub mod attempt_plan;
pub mod candidate_materialization;
pub mod candidate_metadata;
pub mod candidate_persistence;
pub mod candidate_persistence_policy;
pub mod candidate_preparation;
pub mod candidate_preselection;
pub mod candidate_ranking;
pub mod candidate_resolution;
pub mod decision_input;
pub mod decision_path;
pub mod decision_payload;
pub mod dto;
pub mod execution_path;
pub mod failure_diagnostic;
pub mod plan_payload;
pub mod pool_scheduler;
pub mod ports;
pub mod ranking_metadata;
pub mod report_context;
pub mod request_body_diagnostics;
pub mod runtime_miss;
pub mod surface_spec;
pub use attempt_loop::{
run_ai_attempt_loop, AiAttemptLoopOutcome, AiAttemptLoopPort, AiExecutionAttempt,
};
pub use attempt_plan::{
build_ai_execution_decision_from_plan, build_ai_execution_plan_from_decision,
extract_ai_auth_header_pair, infer_ai_upstream_base_url,
resolve_ai_passthrough_sync_request_body, take_ai_decision_plan_core, take_ai_non_empty_string,
take_ai_upstream_auth_pair, trim_ai_owned_non_empty_string, AiDecisionPlanCore,
AiExecutionDecisionFromPlanParts, AiExecutionPlanFromDecisionParts, AiUpstreamAuthPair,
};
pub use candidate_materialization::{
run_ai_candidate_materialization, AiCandidateMaterializationOutcome,
AiCandidateMaterializationPort,
};
pub use candidate_metadata::{
ai_local_execution_contract_for_formats, append_ai_execution_contract_fields_to_value,
build_ai_candidate_metadata, build_ai_candidate_metadata_from_candidate,
AiCandidateMetadataParts,
};
pub use candidate_persistence::{
ai_candidate_extra_data_with_ranking, ai_should_persist_available_candidate_for_pool_key,
ai_should_persist_skipped_candidate_for_pool_membership,
run_ai_available_candidate_persistence, run_ai_skipped_candidate_persistence,
AiAvailableCandidatePersistencePort, AiSkippedCandidatePersistencePort,
};
pub use candidate_persistence_policy::{
ai_candidate_persistence_policy_spec, AiCandidatePersistencePolicyKind,
AiCandidatePersistencePolicySpec,
};
pub use candidate_preparation::{
prepare_ai_header_authenticated_candidate, resolve_ai_candidate_mapped_model,
AiPreparedHeaderAuthenticatedCandidate,
};
pub use candidate_preselection::{
run_ai_candidate_preselection, AiCandidatePreselectionOutcome, AiCandidatePreselectionPort,
};
pub use candidate_ranking::{
ai_ranking_context, build_ai_rankable_candidate, run_ai_candidate_ranking,
AiCandidateRankingPort, AiRankableCandidateParts, AiRankingContextConfig,
AiRankingSchedulingMode,
};
pub use candidate_resolution::{
extract_ai_pool_sticky_session_token, run_ai_candidate_resolution, AiCandidateResolutionMode,
AiCandidateResolutionOutcome, AiCandidateResolutionPort, AiCandidateResolutionRequest,
};
pub use decision_input::{run_ai_authenticated_decision_input, AiAuthenticatedDecisionInputPort};
pub use decision_path::{
run_ai_stream_decision_path, run_ai_sync_decision_path, AiStreamDecisionPathPort,
AiStreamDecisionStep, AiSyncDecisionPathPort, AiSyncDecisionStep,
};
pub use decision_payload::{
ai_execution_decision_action, build_ai_execution_decision_response,
AiExecutionDecisionResponseParts,
};
pub use dto::{
augment_sync_report_context, generic_decision_missing_exact_provider_request,
AiExecutionDecision, AiExecutionPlanPayload, AiStreamAttempt, AiSyncAttempt, ConversionMode,
ExecutionStrategy,
};
pub use execution_path::{
run_ai_stream_execution_path, run_ai_sync_execution_path, AiPlanFallbackReason,
AiServingExecutionOutcome, AiStreamExecutionPathPort, AiStreamExecutionStep,
AiSyncExecutionPathPort, AiSyncExecutionStep,
};
pub use failure_diagnostic::{CandidateFailureDiagnostic, CandidateFailureDiagnosticKind};
pub use plan_payload::{
build_ai_stream_execution_plan_payload, build_ai_sync_execution_plan_payload,
};
pub use pool_scheduler::{
normalize_enabled_ai_pool_presets, run_ai_pool_scheduler, AiPoolCandidateFacts,
AiPoolCandidateInput, AiPoolCandidateOrchestration, AiPoolCatalogKeyContext,
AiPoolRuntimeState, AiPoolScheduledCandidate, AiPoolSchedulerOutcome, AiPoolSchedulingConfig,
AiPoolSchedulingPreset, AiPoolSkippedCandidate, AI_POOL_ACCOUNT_BLOCKED_SKIP_REASON,
AI_POOL_ACCOUNT_EXHAUSTED_SKIP_REASON, AI_POOL_COOLDOWN_SKIP_REASON,
AI_POOL_COST_LIMIT_REACHED_SKIP_REASON,
};
pub use ranking_metadata::append_ai_ranking_metadata_to_object;
pub use report_context::{
build_ai_execution_report_context, build_ai_report_context_original_request_echo,
insert_provider_stream_event_api_format, provider_stream_event_api_format_for_provider_type,
AiExecutionReportContextParts, AiRequestOrigin,
};
pub use request_body_diagnostics::{
request_body_build_failure_extra_data, same_format_provider_request_body_failure_extra_data,
};
pub use runtime_miss::{
apply_ai_runtime_candidate_evaluation_progress,
apply_ai_runtime_candidate_evaluation_progress_preserving_candidate_signal,
apply_ai_runtime_candidate_evaluation_progress_to_diagnostic,
apply_ai_runtime_candidate_terminal_plan_reason_to_diagnostic,
apply_ai_runtime_candidate_terminal_reason, build_ai_runtime_candidate_evaluation_diagnostic,
build_ai_runtime_execution_exhausted_diagnostic, record_ai_runtime_candidate_skip_reason,
record_ai_runtime_candidate_skip_reason_on_diagnostic,
set_ai_runtime_candidate_evaluation_diagnostic, set_ai_runtime_execution_exhausted_diagnostic,
set_ai_runtime_miss_diagnostic_reason, AiRuntimeMissDiagnosticFields,
AiRuntimeMissDiagnosticPort,
};
pub use surface_spec::{
ai_gemini_files_spec_metadata, ai_openai_image_spec_metadata,
ai_openai_responses_spec_metadata, ai_requested_model_family_for_same_format_provider,
ai_requested_model_family_for_standard_source, ai_requested_model_family_for_video_create,
ai_same_format_provider_spec_metadata, ai_standard_spec_metadata,
ai_video_create_spec_metadata, extract_ai_gemini_model_from_path,
extract_ai_requested_model_from_request_path, extract_ai_standard_requested_model,
AiExecutionSurfaceSpecMetadata, AiRequestedModelFamily,
};

108
crates/aether-ai-serving/src/plan_payload.rs Normal file
View File

@@ -0,0 +1,108 @@
use aether_ai_surfaces::api::{
ExecutionRuntimeAuthContext, EXECUTION_RUNTIME_STREAM_ACTION, EXECUTION_RUNTIME_SYNC_ACTION,
};
use crate::dto::{AiExecutionPlanPayload, AiStreamAttempt, AiSyncAttempt};
pub fn build_ai_sync_execution_plan_payload(
plan_kind: &str,
attempt: AiSyncAttempt,
auth_context: Option<ExecutionRuntimeAuthContext>,
) -> AiExecutionPlanPayload {
AiExecutionPlanPayload {
action: EXECUTION_RUNTIME_SYNC_ACTION.to_string(),
plan_kind: Some(plan_kind.to_string()),
plan: Some(attempt.plan),
report_kind: attempt.report_kind,
report_context: attempt.report_context,
auth_context,
}
}
pub fn build_ai_stream_execution_plan_payload(
plan_kind: &str,
attempt: AiStreamAttempt,
auth_context: Option<ExecutionRuntimeAuthContext>,
) -> AiExecutionPlanPayload {
AiExecutionPlanPayload {
action: EXECUTION_RUNTIME_STREAM_ACTION.to_string(),
plan_kind: Some(plan_kind.to_string()),
plan: Some(attempt.plan),
report_kind: attempt.report_kind,
report_context: attempt.report_context,
auth_context,
}
}
#[cfg(test)]
mod tests {
use super::*;
use aether_contracts::{ExecutionPlan, RequestBody};
#[test]
fn sync_plan_payload_uses_sync_action_and_attempt_report_fields() {
let payload = build_ai_sync_execution_plan_payload(
"openai_chat_sync",
AiSyncAttempt {
plan: test_plan(),
report_kind: Some("sync_success".to_string()),
report_context: Some(serde_json::json!({"candidate_index": 0})),
},
None,
);
assert_eq!(payload.action, EXECUTION_RUNTIME_SYNC_ACTION);
assert_eq!(payload.plan_kind.as_deref(), Some("openai_chat_sync"));
assert_eq!(payload.report_kind.as_deref(), Some("sync_success"));
assert_eq!(
payload
.report_context
.as_ref()
.and_then(|value| value.get("candidate_index"))
.and_then(serde_json::Value::as_u64),
Some(0)
);
assert!(payload.plan.is_some());
}
#[test]
fn stream_plan_payload_uses_stream_action() {
let payload = build_ai_stream_execution_plan_payload(
"openai_chat_stream",
AiStreamAttempt {
plan: test_plan(),
report_kind: None,
report_context: None,
},
None,
);
assert_eq!(payload.action, EXECUTION_RUNTIME_STREAM_ACTION);
assert_eq!(payload.plan_kind.as_deref(), Some("openai_chat_stream"));
assert!(payload.plan.is_some());
}
fn test_plan() -> ExecutionPlan {
ExecutionPlan {
request_id: "req_1".to_string(),
candidate_id: None,
provider_name: Some("provider".to_string()),
provider_id: "provider_id".to_string(),
endpoint_id: "endpoint_id".to_string(),
key_id: "key_id".to_string(),
method: "POST".to_string(),
url: "https://example.com/v1/chat/completions".to_string(),
headers: Default::default(),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(serde_json::json!({"model": "model"})),
stream: false,
client_api_format: "openai:chat".to_string(),
provider_api_format: "openai:chat".to_string(),
model_name: Some("model".to_string()),
proxy: None,
tls_profile: None,
timeouts: None,
}
}
}

1048
crates/aether-ai-serving/src/pool_scheduler.rs Normal file
View File

File diff suppressed because it is too large Load Diff

12
crates/aether-ai-serving/src/ports.rs Normal file
View File

@@ -0,0 +1,12 @@
pub use crate::attempt_loop::AiAttemptLoopPort;
pub use crate::candidate_materialization::AiCandidateMaterializationPort;
pub use crate::candidate_persistence::{
AiAvailableCandidatePersistencePort, AiSkippedCandidatePersistencePort,
};
pub use crate::candidate_preselection::AiCandidatePreselectionPort;
pub use crate::candidate_ranking::AiCandidateRankingPort;
pub use crate::candidate_resolution::AiCandidateResolutionPort;
pub use crate::decision_input::AiAuthenticatedDecisionInputPort;
pub use crate::decision_path::{AiStreamDecisionPathPort, AiSyncDecisionPathPort};
pub use crate::execution_path::{AiStreamExecutionPathPort, AiSyncExecutionPathPort};
pub use crate::runtime_miss::AiRuntimeMissDiagnosticPort;

66
crates/aether-ai-serving/src/ranking_metadata.rs Normal file
View File

@@ -0,0 +1,66 @@
use aether_scheduler_core::SchedulerRankingOutcome;
use serde_json::{Map, Value};
pub fn append_ai_ranking_metadata_to_object(
object: &mut Map<String, Value>,
ranking: &SchedulerRankingOutcome,
) {
object.insert(
"ranking_mode".to_string(),
Value::String(format!("{:?}", ranking.ranking_mode)),
);
object.insert(
"priority_mode".to_string(),
Value::String(format!("{:?}", ranking.priority_mode)),
);
object.insert(
"ranking_index".to_string(),
Value::Number(serde_json::Number::from(ranking.ranking_index as u64)),
);
object.insert(
"priority_slot".to_string(),
Value::Number(serde_json::Number::from(i64::from(ranking.priority_slot))),
);
if let Some(promoted_by) = ranking.promoted_by {
object.insert(
"promoted_by".to_string(),
Value::String(promoted_by.to_string()),
);
}
if let Some(demoted_by) = ranking.demoted_by {
object.insert(
"demoted_by".to_string(),
Value::String(demoted_by.to_string()),
);
}
}
#[cfg(test)]
mod tests {
use super::*;
use aether_scheduler_core::{SchedulerPriorityMode, SchedulerRankingMode};
use serde_json::json;
#[test]
fn ranking_metadata_appends_scheduler_outcome_fields() {
let ranking = SchedulerRankingOutcome {
original_index: 2,
ranking_index: 1,
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::CacheAffinity,
priority_slot: 7,
promoted_by: Some("cached_affinity"),
demoted_by: Some("cross_format"),
};
let mut object = Map::new();
append_ai_ranking_metadata_to_object(&mut object, &ranking);
assert_eq!(object.get("ranking_mode"), Some(&json!("CacheAffinity")));
assert_eq!(object.get("priority_mode"), Some(&json!("Provider")));
assert_eq!(object.get("ranking_index"), Some(&json!(1)));
assert_eq!(object.get("priority_slot"), Some(&json!(7)));
assert_eq!(object.get("promoted_by"), Some(&json!("cached_affinity")));
assert_eq!(object.get("demoted_by"), Some(&json!("cross_format")));
}
}

388
crates/aether-ai-serving/src/report_context.rs Normal file
View File

@@ -0,0 +1,388 @@
use std::collections::BTreeMap;
use aether_ai_surfaces::api::ExecutionRuntimeAuthContext;
use aether_scheduler_core::SchedulerRankingOutcome;
use serde_json::{Map, Value};
use crate::append_ai_ranking_metadata_to_object;
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct AiRequestOrigin {
pub client_ip: Option<String>,
pub user_agent: Option<String>,
}
pub struct AiExecutionReportContextParts<'a> {
pub auth_context: &'a ExecutionRuntimeAuthContext,
pub request_id: &'a str,
pub candidate_id: &'a str,
pub candidate_index: u32,
pub retry_index: u32,
pub pool_key_index: Option<u32>,
pub model: &'a str,
pub provider_name: &'a str,
pub provider_id: &'a str,
pub endpoint_id: &'a str,
pub key_id: &'a str,
pub key_name: Option<&'a str>,
pub model_id: Option<&'a str>,
pub global_model_id: Option<&'a str>,
pub global_model_name: Option<&'a str>,
pub provider_api_format: &'a str,
pub client_api_format: &'a str,
pub mapped_model: Option<&'a str>,
pub candidate_group_id: Option<&'a str>,
pub ranking: Option<&'a SchedulerRankingOutcome>,
pub upstream_url: Option<&'a str>,
pub header_rules: Option<&'a Value>,
pub body_rules: Option<&'a Value>,
pub provider_request_method: Option<Value>,
pub provider_request_headers: Option<&'a BTreeMap<String, String>>,
pub original_headers: &'a BTreeMap<String, String>,
pub original_request_body: Option<Value>,
pub request_origin: AiRequestOrigin,
pub client_requested_stream: bool,
pub upstream_is_stream: bool,
pub has_envelope: bool,
pub needs_conversion: bool,
pub extra_fields: Map<String, Value>,
}
pub fn build_ai_execution_report_context(parts: AiExecutionReportContextParts<'_>) -> Value {
let mut object = Map::new();
object.insert(
"user_id".to_string(),
Value::String(parts.auth_context.user_id.clone()),
);
object.insert(
"api_key_id".to_string(),
Value::String(parts.auth_context.api_key_id.clone()),
);
object.insert(
"api_key_is_standalone".to_string(),
Value::Bool(parts.auth_context.api_key_is_standalone),
);
object.insert(
"username".to_string(),
parts
.auth_context
.username
.clone()
.map(Value::String)
.unwrap_or(Value::Null),
);
object.insert(
"api_key_name".to_string(),
parts
.auth_context
.api_key_name
.clone()
.map(Value::String)
.unwrap_or(Value::Null),
);
object.insert(
"request_id".to_string(),
Value::String(parts.request_id.to_string()),
);
object.insert(
"candidate_id".to_string(),
Value::String(parts.candidate_id.to_string()),
);
object.insert(
"candidate_index".to_string(),
Value::Number(parts.candidate_index.into()),
);
object.insert(
"retry_index".to_string(),
Value::Number(parts.retry_index.into()),
);
object.insert("model".to_string(), Value::String(parts.model.to_string()));
object.insert(
"provider_name".to_string(),
Value::String(parts.provider_name.to_string()),
);
object.insert(
"provider_id".to_string(),
Value::String(parts.provider_id.to_string()),
);
object.insert(
"endpoint_id".to_string(),
Value::String(parts.endpoint_id.to_string()),
);
object.insert(
"key_id".to_string(),
Value::String(parts.key_id.to_string()),
);
object.insert(
"provider_api_format".to_string(),
Value::String(parts.provider_api_format.to_string()),
);
object.insert(
"client_api_format".to_string(),
Value::String(parts.client_api_format.to_string()),
);
object.insert(
"original_headers".to_string(),
serde_json::to_value(parts.original_headers).expect("control headers should serialize"),
);
object.insert(
"original_request_body".to_string(),
parts.original_request_body.unwrap_or(Value::Null),
);
if let Some(client_ip) = parts.request_origin.client_ip {
object.insert("client_ip".to_string(), Value::String(client_ip));
}
if let Some(user_agent) = parts.request_origin.user_agent {
object.insert("user_agent".to_string(), Value::String(user_agent));
}
object.insert(
"client_requested_stream".to_string(),
Value::Bool(parts.client_requested_stream),
);
object.insert(
"upstream_is_stream".to_string(),
Value::Bool(parts.upstream_is_stream),
);
object.insert("has_envelope".to_string(), Value::Bool(parts.has_envelope));
object.insert(
"needs_conversion".to_string(),
Value::Bool(parts.needs_conversion),
);
if let Some(key_name) = parts.key_name {
object.insert("key_name".to_string(), Value::String(key_name.to_string()));
}
if let Some(model_id) = parts.model_id {
object.insert("model_id".to_string(), Value::String(model_id.to_string()));
}
if let Some(global_model_id) = parts.global_model_id {
object.insert(
"global_model_id".to_string(),
Value::String(global_model_id.to_string()),
);
}
if let Some(global_model_name) = parts.global_model_name {
object.insert(
"global_model_name".to_string(),
Value::String(global_model_name.to_string()),
);
}
if let Some(mapped_model) = parts.mapped_model {
object.insert(
"mapped_model".to_string(),
Value::String(mapped_model.to_string()),
);
}
if let Some(candidate_group_id) = parts.candidate_group_id {
object.insert(
"candidate_group_id".to_string(),
Value::String(candidate_group_id.to_string()),
);
}
if let Some(ranking) = parts.ranking {
append_ai_ranking_metadata_to_object(&mut object, ranking);
}
if let Some(upstream_url) = parts.upstream_url {
object.insert(
"upstream_url".to_string(),
Value::String(upstream_url.to_string()),
);
}
if let Some(header_rules) = parts.header_rules {
object.insert("header_rules".to_string(), header_rules.clone());
}
if let Some(body_rules) = parts.body_rules {
object.insert("body_rules".to_string(), body_rules.clone());
}
if let Some(provider_request_method) = parts.provider_request_method {
object.insert(
"provider_request_method".to_string(),
provider_request_method,
);
}
if let Some(provider_request_headers) = parts.provider_request_headers {
object.insert(
"provider_request_headers".to_string(),
serde_json::to_value(provider_request_headers)
.expect("provider request headers should serialize"),
);
}
if let Some(pool_key_index) = parts.pool_key_index {
object.insert(
"pool_key_index".to_string(),
Value::Number(pool_key_index.into()),
);
}
object.extend(parts.extra_fields);
Value::Object(object)
}
pub fn provider_stream_event_api_format_for_provider_type(
provider_type: &str,
) -> Option<&'static str> {
match provider_type.trim().to_ascii_lowercase().as_str() {
"codex" => Some("openai:responses"),
_ => None,
}
}
pub fn insert_provider_stream_event_api_format(
extra_fields: &mut Map<String, Value>,
provider_type: &str,
) {
if let Some(api_format) = provider_stream_event_api_format_for_provider_type(provider_type) {
extra_fields.insert(
"provider_stream_event_api_format".to_string(),
Value::String(api_format.to_string()),
);
}
}
pub fn build_ai_report_context_original_request_echo(
body_json: Option<&Value>,
body_bytes_b64: Option<&str>,
) -> Option<Value> {
if let Some(body_bytes_b64) = body_bytes_b64
.map(str::trim)
.filter(|value| !value.is_empty())
{
return Some(serde_json::json!({ "body_bytes_b64": body_bytes_b64 }));
}
body_json.filter(|body| !body.is_null()).cloned()
}
#[cfg(test)]
mod tests {
use super::*;
use aether_scheduler_core::{SchedulerPriorityMode, SchedulerRankingMode};
use serde_json::json;
fn sample_auth_context() -> ExecutionRuntimeAuthContext {
ExecutionRuntimeAuthContext {
user_id: "user-1".to_string(),
api_key_id: "key-1".to_string(),
username: Some("alice".to_string()),
api_key_name: Some("primary".to_string()),
balance_remaining: Some(42.0),
access_allowed: true,
api_key_is_standalone: false,
}
}
#[test]
fn report_context_builds_core_execution_fields() {
let auth_context = sample_auth_context();
let original_headers = BTreeMap::from([("x-trace-id".to_string(), "trace-a".to_string())]);
let provider_headers =
BTreeMap::from([("authorization".to_string(), "Bearer token".to_string())]);
let mut extra_fields = Map::new();
extra_fields.insert("extra".to_string(), json!("value"));
let ranking = SchedulerRankingOutcome {
original_index: 2,
ranking_index: 1,
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::CacheAffinity,
priority_slot: 7,
promoted_by: Some("cached_affinity"),
demoted_by: None,
};
let report = build_ai_execution_report_context(AiExecutionReportContextParts {
auth_context: &auth_context,
request_id: "trace-a",
candidate_id: "candidate-a",
candidate_index: 3,
retry_index: 1,
pool_key_index: Some(0),
model: "gpt-5",
provider_name: "RightCode",
provider_id: "provider-1",
endpoint_id: "endpoint-1",
key_id: "key-1",
key_name: Some("primary"),
model_id: Some("model-1"),
global_model_id: Some("global-1"),
global_model_name: Some("GPT-5"),
provider_api_format: "openai:responses",
client_api_format: "openai:chat",
mapped_model: Some("gpt-5"),
candidate_group_id: Some("group-1"),
ranking: Some(&ranking),
upstream_url: Some("https://example.com/v1/responses"),
header_rules: Some(&json!({"set": []})),
body_rules: None,
provider_request_method: Some(json!("POST")),
provider_request_headers: Some(&provider_headers),
original_headers: &original_headers,
original_request_body: Some(json!({"model": "gpt-5"})),
request_origin: AiRequestOrigin {
client_ip: Some("127.0.0.1".to_string()),
user_agent: Some("test-agent".to_string()),
},
client_requested_stream: false,
upstream_is_stream: true,
has_envelope: false,
needs_conversion: true,
extra_fields,
});
assert_eq!(report["user_id"], "user-1");
assert_eq!(report["candidate_index"], 3);
assert_eq!(report["retry_index"], 1);
assert_eq!(report["pool_key_index"], 0);
assert_eq!(report["original_headers"]["x-trace-id"], "trace-a");
assert_eq!(report["original_request_body"]["model"], "gpt-5");
assert_eq!(report["ranking_index"], 1);
assert_eq!(
report["provider_request_headers"]["authorization"],
"Bearer token"
);
assert_eq!(report["extra"], "value");
}
#[test]
fn provider_stream_event_api_format_is_codex_only() {
assert_eq!(
provider_stream_event_api_format_for_provider_type("codex"),
Some("openai:responses")
);
assert_eq!(
provider_stream_event_api_format_for_provider_type(" CODEX "),
Some("openai:responses")
);
assert_eq!(
provider_stream_event_api_format_for_provider_type("openai"),
None
);
}
#[test]
fn original_request_echo_preserves_full_request_body() {
let body = json!({
"messages": [{"role": "user", "content": "large payload should be omitted"}],
"service_tier": "default",
"instructions": "Be concise.",
"thinking": {"type": "enabled", "budget_tokens": 512},
"metadata": {"trace": "keep"},
"body_bytes_b64": "aGVsbG8=",
});
let echo = build_ai_report_context_original_request_echo(Some(&body), None)
.expect("echo should be produced");
assert_eq!(echo, body);
}
#[test]
fn original_request_echo_prefers_binary_body_bytes() {
let echo = build_ai_report_context_original_request_echo(
Some(&json!({"ignored": true})),
Some("aGVsbG8="),
)
.expect("echo should be produced");
assert_eq!(echo, json!({"body_bytes_b64": "aGVsbG8="}));
}
}

741
crates/aether-ai-serving/src/request_body_diagnostics.rs Normal file
View File

@@ -0,0 +1,741 @@
use serde_json::Value;
use aether_ai_surfaces::api::is_openai_responses_family_format;
use crate::{CandidateFailureDiagnostic, CandidateFailureDiagnosticKind};
pub fn request_body_build_failure_extra_data(
body_json: &Value,
client_api_format: &str,
provider_api_format: &str,
) -> Option<Value> {
let diagnostic =
diagnose_request_body_build_failure(body_json, client_api_format, provider_api_format)?;
Some(
diagnostic
.formats(client_api_format, provider_api_format)
.source(request_body_build_source(
client_api_format,
provider_api_format,
))
.to_extra_data(),
)
}
pub fn same_format_provider_request_body_failure_extra_data(
body_json: &Value,
provider_api_format: &str,
body_rules: Option<&Value>,
context: &str,
) -> Option<Value> {
let diagnostic =
diagnose_same_format_provider_request_body_failure(body_json, body_rules, context)?;
Some(
diagnostic
.formats(provider_api_format, provider_api_format)
.source(context)
.to_extra_data(),
)
}
type RequestBodyBuildDiagnostic = CandidateFailureDiagnostic;
fn diagnose_request_body_build_failure(
body_json: &Value,
client_api_format: &str,
provider_api_format: &str,
) -> Option<RequestBodyBuildDiagnostic> {
if !body_json.is_object() {
return Some(diagnostic("$", "请求体必须是 JSON object"));
}
if is_openai_responses_client_format(client_api_format) {
if let Some(diagnostic) = diagnose_openai_responses_request(body_json) {
return Some(diagnostic);
}
return Some(diagnostic(
"$",
"OpenAI Responses 请求体初步结构检查通过；失败可能发生在后续跨格式转换或 Body 规则应用",
));
}
if client_api_format == "openai:chat"
&& (provider_api_format.starts_with("claude:")
|| provider_api_format.starts_with("gemini:"))
{
return diagnose_openai_chat_cross_format_request(body_json, provider_api_format);
}
Some(diagnostic(
"$",
"请求体转换失败；当前转换器未返回更细的字段路径",
))
}
fn is_openai_responses_client_format(client_api_format: &str) -> bool {
is_openai_responses_family_format(client_api_format)
}
fn diagnose_same_format_provider_request_body_failure(
body_json: &Value,
body_rules: Option<&Value>,
context: &str,
) -> Option<RequestBodyBuildDiagnostic> {
if !body_json.is_object() {
return Some(diagnostic("$", "反代请求体必须是 JSON object"));
}
if body_rules.is_some_and(|rules| !rules.is_array()) {
return Some(diagnostic(
"$.endpoint.body_rules",
"Endpoint Body 规则必须是数组，本地反代无法应用该配置",
));
}
match context {
"kiro_envelope" => Some(diagnostic(
"$",
"Kiro 反代请求体包装失败；请检查 Kiro auth_config 与 Endpoint Body 规则",
)),
"antigravity_envelope" => Some(diagnostic(
"$",
"Antigravity 反代请求体包装失败；请检查请求体是否满足该传输封装要求",
)),
_ => Some(diagnostic(
"$",
"反代请求体构建失败；当前路径未返回更细的字段信息",
)),
}
}
fn diagnose_openai_chat_cross_format_request(
body_json: &Value,
provider_api_format: &str,
) -> Option<RequestBodyBuildDiagnostic> {
let request = body_json.as_object()?;
if let Some(messages) = request.get("messages") {
let Some(messages) = messages.as_array() else {
return Some(diagnostic(
"$.messages",
"OpenAI Chat 的 messages 必须是数组",
));
};
for (message_index, message) in messages.iter().enumerate() {
let Some(message_object) = message.as_object() else {
return Some(diagnostic(
format!("$.messages[{message_index}]"),
"message 必须是 object",
));
};
let role = message_object
.get("role")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
match role.as_str() {
"system" | "developer" => {
if let Some(diagnostic) = diagnose_openai_text_content(
message_object.get("content"),
format!("$.messages[{message_index}].content"),
) {
return Some(diagnostic);
}
}
"user" | "assistant" => {
if let Some(diagnostic) = diagnose_openai_content_blocks(
message_object.get("content"),
format!("$.messages[{message_index}].content"),
role.as_str(),
) {
return Some(diagnostic);
}
if role == "assistant" {
if let Some(diagnostic) = diagnose_openai_assistant_tool_calls(
message_object.get("tool_calls"),
format!("$.messages[{message_index}].tool_calls"),
) {
return Some(diagnostic);
}
}
}
"tool" => {
let valid_tool_call_id = message_object
.get("tool_call_id")
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty());
if !valid_tool_call_id {
return Some(diagnostic(
format!("$.messages[{message_index}].tool_call_id"),
"tool 消息必须包含非空 tool_call_id",
));
}
}
_ => {}
}
}
}
if let Some(diagnostic) = diagnose_openai_tools(request.get("tools"), provider_api_format) {
return Some(diagnostic);
}
diagnose_openai_tool_choice(request.get("tool_choice"))
}
fn diagnose_openai_responses_request(body_json: &Value) -> Option<RequestBodyBuildDiagnostic> {
let request = body_json.as_object()?;
if let Some(diagnostic) = diagnose_openai_responses_text_content(
request.get("instructions"),
"$.instructions".to_string(),
) {
return Some(diagnostic);
}
if let Some(diagnostic) = diagnose_openai_responses_input(request.get("input")) {
return Some(diagnostic);
}
if let Some(diagnostic) = diagnose_openai_responses_tools(request.get("tools")) {
return Some(diagnostic);
}
diagnose_openai_responses_tool_choice(request.get("tool_choice"))
}
fn diagnose_openai_responses_input(input: Option<&Value>) -> Option<RequestBodyBuildDiagnostic> {
let input = input?;
match input {
Value::Null | Value::String(_) => None,
Value::Array(items) => {
for (item_index, item) in items.iter().enumerate() {
if item.is_string() {
continue;
}
let item_path = format!("$.input[{item_index}]");
let Some(item_object) = item.as_object() else {
return Some(diagnostic(
item_path,
"OpenAI Responses input 数组项必须是 string 或 object",
));
};
let item_type = item_object
.get("type")
.and_then(Value::as_str)
.unwrap_or("message")
.trim()
.to_ascii_lowercase();
match item_type.as_str() {
"message" => {
let role = item_object
.get("role")
.and_then(Value::as_str)
.unwrap_or("user")
.trim()
.to_ascii_lowercase();
if role == "system" || role == "developer" {
if let Some(diagnostic) = diagnose_openai_responses_text_content(
item_object.get("content"),
format!("{item_path}.content"),
) {
return Some(diagnostic);
}
} else if let Some(diagnostic) = diagnose_openai_responses_message_content(
item_object.get("content"),
format!("{item_path}.content"),
) {
return Some(diagnostic);
}
}
"function_call" => {
let valid_name = item_object
.get("name")
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty());
if !valid_name {
return Some(diagnostic(
format!("{item_path}.name"),
"function_call 必须包含非空 name",
));
}
}
_ => {}
}
}
None
}
_ => Some(diagnostic(
"$.input",
"OpenAI Responses input 必须是 string、array 或 null",
)),
}
}
fn diagnose_openai_responses_text_content(
content: Option<&Value>,
path: String,
) -> Option<RequestBodyBuildDiagnostic> {
match content {
None | Some(Value::Null) | Some(Value::String(_)) => None,
Some(Value::Array(parts)) => {
for (part_index, part) in parts.iter().enumerate() {
if !part.is_object() {
return Some(diagnostic(
format!("{path}[{part_index}]"),
"文本 content 数组项必须是 object",
));
}
}
None
}
Some(_) => Some(diagnostic(
path,
"文本 content 必须是 string、array 或 null",
)),
}
}
fn diagnose_openai_responses_message_content(
content: Option<&Value>,
path: String,
) -> Option<RequestBodyBuildDiagnostic> {
match content {
None | Some(Value::Null) | Some(Value::String(_)) => None,
Some(Value::Array(parts)) => {
for (part_index, part) in parts.iter().enumerate() {
let part_path = format!("{path}[{part_index}]");
let Some(part_object) = part.as_object() else {
return Some(diagnostic(part_path, "message content 数组项必须是 object"));
};
let part_type = part_object
.get("type")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
if matches!(
part_type.as_str(),
"input_image" | "output_image" | "image_url"
) && image_part_url(part_object).is_none()
{
return Some(diagnostic(
part_path,
"图片 content 缺少 image_url/url，无法规范化为 OpenAI Chat 图片内容",
));
}
}
None
}
Some(_) => None,
}
}
fn diagnose_openai_text_content(
content: Option<&Value>,
path: String,
) -> Option<RequestBodyBuildDiagnostic> {
match content {
None | Some(Value::Null) | Some(Value::String(_)) => None,
Some(Value::Array(parts)) => {
for (part_index, part) in parts.iter().enumerate() {
if !part.is_object() {
return Some(diagnostic(
format!("{path}[{part_index}]"),
"content 数组项必须是 object",
));
}
}
None
}
Some(_) => Some(diagnostic(path, "content 必须是 string、array 或 null")),
}
}
fn diagnose_openai_content_blocks(
content: Option<&Value>,
path: String,
role: &str,
) -> Option<RequestBodyBuildDiagnostic> {
match content {
None | Some(Value::Null) | Some(Value::String(_)) => None,
Some(Value::Array(parts)) => {
for (part_index, part) in parts.iter().enumerate() {
let part_path = format!("{path}[{part_index}]");
let Some(part_object) = part.as_object() else {
return Some(diagnostic(part_path, "content 数组项必须是 object"));
};
let part_type = part_object
.get("type")
.and_then(Value::as_str)
.unwrap_or_default();
if matches!(part_type, "image_url" | "input_image" | "output_image")
&& role == "user"
&& image_part_url(part_object).is_none()
{
return Some(diagnostic(
part_path,
"图片 content 缺少 image_url/url，无法转换为 Claude image block",
));
}
}
None
}
Some(_) => Some(diagnostic(path, "content 必须是 string、array 或 null")),
}
}
fn diagnose_openai_assistant_tool_calls(
tool_calls: Option<&Value>,
path: String,
) -> Option<RequestBodyBuildDiagnostic> {
let tool_calls = tool_calls?;
let Some(tool_calls) = tool_calls.as_array() else {
return Some(diagnostic(path, "assistant.tool_calls 必须是数组"));
};
for (tool_call_index, tool_call) in tool_calls.iter().enumerate() {
let tool_call_path = format!("{path}[{tool_call_index}]");
let Some(tool_call_object) = tool_call.as_object() else {
return Some(diagnostic(tool_call_path, "tool_call 必须是 object"));
};
let Some(function) = tool_call_object.get("function").and_then(Value::as_object) else {
return Some(diagnostic(
format!("{tool_call_path}.function"),
"tool_call 必须包含 function object",
));
};
let valid_name = function
.get("name")
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty());
if !valid_name {
return Some(diagnostic(
format!("{tool_call_path}.function.name"),
"tool_call.function.name 必须是非空字符串",
));
}
}
None
}
fn diagnose_openai_tools(
tools: Option<&Value>,
provider_api_format: &str,
) -> Option<RequestBodyBuildDiagnostic> {
let tools = tools?;
let Some(tools) = tools.as_array() else {
return Some(diagnostic("$.tools", "OpenAI Chat 的 tools 必须是数组"));
};
for (tool_index, tool) in tools.iter().enumerate() {
let tool_path = format!("$.tools[{tool_index}]");
let Some(tool_object) = tool.as_object() else {
return Some(diagnostic(tool_path, "tool 必须是 object"));
};
if tool_object
.get("type")
.and_then(Value::as_str)
.is_some_and(|value| value != "function")
{
continue;
}
let Some(function) = tool_object.get("function").and_then(Value::as_object) else {
let native_tool_hint = if provider_api_format.starts_with("claude:") {
"；如果这是 Claude 原生 tool，请改为 OpenAI function tool 格式"
} else if provider_api_format.starts_with("gemini:") {
"；如果这是 Gemini 原生 tool，请改为 OpenAI function tool 格式"
} else {
""
};
return Some(diagnostic(
format!("{tool_path}.function"),
format!("OpenAI tool 必须包含 function object{native_tool_hint}"),
));
};
let valid_name = function
.get("name")
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty());
if !valid_name {
return Some(diagnostic(
format!("{tool_path}.function.name"),
"OpenAI tool 的 function.name 必须是非空字符串",
));
}
}
None
}
fn diagnose_openai_responses_tools(tools: Option<&Value>) -> Option<RequestBodyBuildDiagnostic> {
let tools = tools?;
let tool_values = tools.as_array()?;
for (tool_index, tool) in tool_values.iter().enumerate() {
let tool_path = format!("$.tools[{tool_index}]");
let Some(tool_object) = tool.as_object() else {
return Some(diagnostic(tool_path, "OpenAI Responses tool 必须是 object"));
};
let tool_type = tool_object
.get("type")
.and_then(Value::as_str)
.unwrap_or("function")
.trim()
.to_ascii_lowercase();
if tool_type.starts_with("web_search")
|| tool_object.get("function").is_some()
|| tool_type != "function"
{
continue;
}
let valid_name = tool_object
.get("name")
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty());
if !valid_name {
return Some(diagnostic(
format!("{tool_path}.name"),
"OpenAI Responses function tool 必须包含非空 name",
));
}
}
None
}
fn diagnose_openai_responses_tool_choice(
tool_choice: Option<&Value>,
) -> Option<RequestBodyBuildDiagnostic> {
let Some(Value::Object(object)) = tool_choice else {
return None;
};
let is_cli_function_choice = object.get("function").is_none()
&& object
.get("type")
.and_then(Value::as_str)
.is_some_and(|value| value.eq_ignore_ascii_case("function"));
if !is_cli_function_choice {
return None;
}
let valid_name = object
.get("name")
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty());
if valid_name {
None
} else {
Some(diagnostic(
"$.tool_choice.name",
"OpenAI Responses tool_choice 指定 function 时必须包含非空 name",
))
}
}
fn diagnose_openai_tool_choice(tool_choice: Option<&Value>) -> Option<RequestBodyBuildDiagnostic> {
let Some(Value::Object(object)) = tool_choice else {
return None;
};
let valid_name = object
.get("function")
.and_then(Value::as_object)
.and_then(|function| function.get("name"))
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty());
if valid_name {
None
} else {
Some(diagnostic(
"$.tool_choice.function.name",
"tool_choice 指定具体工具时必须包含非空 function.name",
))
}
}
fn image_part_url(part_object: &serde_json::Map<String, Value>) -> Option<&str> {
part_object
.get("image_url")
.and_then(|value| {
value.as_str().or_else(|| {
value
.as_object()
.and_then(|object| object.get("url"))
.and_then(Value::as_str)
})
})
.or_else(|| part_object.get("url").and_then(Value::as_str))
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn diagnostic(path: impl Into<String>, message: impl Into<String>) -> RequestBodyBuildDiagnostic {
CandidateFailureDiagnostic::new(
CandidateFailureDiagnosticKind::RequestBodyBuild,
path,
message,
)
}
fn request_body_build_source(client_api_format: &str, provider_api_format: &str) -> String {
format!("{client_api_format}_to_{provider_api_format}")
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::request_body_build_failure_extra_data;
#[test]
fn openai_chat_to_claude_reports_claude_native_tool_shape() {
let body = json!({
"model": "gpt-5.4",
"messages": [{ "role": "user", "content": "hello" }],
"tools": [{
"name": "read_file",
"description": "Read a file",
"input_schema": { "type": "object" }
}]
});
let diagnostic =
request_body_build_failure_extra_data(&body, "openai:chat", "claude:messages")
.expect("diagnostic");
assert_eq!(
diagnostic["request_body_build_error"]["path"],
"$.tools[0].function"
);
assert_eq!(
diagnostic["failure_diagnostic"]["kind"],
"request_body_build"
);
assert_eq!(
diagnostic["failure_diagnostic"]["source"],
"openai:chat_to_claude:messages"
);
assert!(diagnostic["request_body_build_error"]["message"]
.as_str()
.expect("message")
.contains("Claude 原生 tool"));
}
#[test]
fn openai_chat_to_claude_reports_invalid_message_content_part() {
let body = json!({
"model": "gpt-5.4",
"messages": [{
"role": "user",
"content": ["not-an-object"]
}]
});
let diagnostic =
request_body_build_failure_extra_data(&body, "openai:chat", "claude:messages")
.expect("diagnostic");
assert_eq!(
diagnostic["request_body_build_error"]["path"],
"$.messages[0].content[0]"
);
}
#[test]
fn openai_chat_to_gemini_reports_gemini_native_tool_shape() {
let body = json!({
"model": "gpt-5.4",
"messages": [{ "role": "user", "content": "hello" }],
"tools": [{
"functionDeclarations": [{
"name": "search",
"parameters": { "type": "object" }
}]
}]
});
let diagnostic =
request_body_build_failure_extra_data(&body, "openai:chat", "gemini:generate_content")
.expect("diagnostic");
assert_eq!(
diagnostic["request_body_build_error"]["path"],
"$.tools[0].function"
);
assert!(diagnostic["request_body_build_error"]["message"]
.as_str()
.expect("message")
.contains("Gemini 原生 tool"));
}
#[test]
fn openai_responses_reports_invalid_function_call_name() {
let body = json!({
"model": "gpt-5.4",
"input": [{
"type": "function_call",
"arguments": "{}"
}]
});
let diagnostic =
request_body_build_failure_extra_data(&body, "openai:responses", "claude:messages")
.expect("diagnostic");
assert_eq!(
diagnostic["request_body_build_error"]["path"],
"$.input[0].name"
);
}
#[test]
fn openai_responses_reports_invalid_tool_choice_name() {
let body = json!({
"model": "gpt-5.4",
"input": "hello",
"tool_choice": { "type": "function" }
});
let diagnostic = request_body_build_failure_extra_data(
&body,
"openai:responses",
"gemini:generate_content",
)
.expect("diagnostic");
assert_eq!(
diagnostic["request_body_build_error"]["path"],
"$.tool_choice.name"
);
}
#[test]
fn same_format_provider_reports_non_object_body() {
let diagnostic = super::same_format_provider_request_body_failure_extra_data(
&json!("raw"),
"openai:chat",
None,
"same_format",
)
.expect("diagnostic");
assert_eq!(diagnostic["request_body_build_error"]["path"], "$");
assert!(diagnostic["request_body_build_error"]["message"]
.as_str()
.expect("message")
.contains("反代请求体必须是 JSON object"));
}
#[test]
fn same_format_provider_reports_invalid_body_rules_shape() {
let diagnostic = super::same_format_provider_request_body_failure_extra_data(
&json!({ "model": "gpt-5.4" }),
"openai:chat",
Some(&json!({ "action": "set" })),
"same_format",
)
.expect("diagnostic");
assert_eq!(
diagnostic["request_body_build_error"]["path"],
"$.endpoint.body_rules"
);
}
}

468
crates/aether-ai-serving/src/runtime_miss.rs Normal file
View File

@@ -0,0 +1,468 @@
pub trait AiRuntimeMissDiagnosticPort: Send + Sync {
type Decision: Send + Sync;
type Diagnostic: Send;
fn build_runtime_miss_diagnostic(
&self,
decision: &Self::Decision,
plan_kind: &str,
requested_model: Option<&str>,
reason: &str,
) -> Self::Diagnostic;
fn set_candidate_count(&self, diagnostic: &mut Self::Diagnostic, candidate_count: usize);
fn apply_candidate_evaluation_progress(
&self,
diagnostic: &mut Self::Diagnostic,
candidate_count: usize,
);
fn apply_candidate_terminal_plan_reason(
&self,
diagnostic: &mut Self::Diagnostic,
no_plan_reason: &'static str,
);
fn record_candidate_skip_reason(
&self,
diagnostic: &mut Self::Diagnostic,
skip_reason: &'static str,
);
fn set_runtime_miss_diagnostic(&self, trace_id: &str, diagnostic: Self::Diagnostic);
fn mutate_runtime_miss_diagnostic<F>(&self, trace_id: &str, apply: F)
where
F: FnOnce(&mut Self::Diagnostic) + Send;
fn runtime_miss_diagnostic_has_candidate_signal(&self, trace_id: &str) -> bool;
}
pub trait AiRuntimeMissDiagnosticFields {
fn set_reason(&mut self, reason: String);
fn set_candidate_count(&mut self, candidate_count: usize);
fn candidate_count(&self) -> Option<usize>;
fn skipped_candidate_count(&self) -> Option<usize>;
fn skip_reason_count(&self, skip_reason: &str) -> usize;
fn skip_reason_len(&self) -> usize;
fn record_skip_reason(&mut self, skip_reason: &'static str);
}
pub fn apply_ai_runtime_candidate_evaluation_progress_to_diagnostic<Diagnostic>(
diagnostic: &mut Diagnostic,
candidate_count: usize,
) where
Diagnostic: AiRuntimeMissDiagnosticFields,
{
diagnostic.set_candidate_count(candidate_count);
diagnostic.set_reason(if candidate_count == 0 {
"candidate_list_empty".to_string()
} else {
"candidate_evaluation_incomplete".to_string()
});
}
pub fn apply_ai_runtime_candidate_terminal_plan_reason_to_diagnostic<Diagnostic>(
diagnostic: &mut Diagnostic,
no_plan_reason: &'static str,
) where
Diagnostic: AiRuntimeMissDiagnosticFields,
{
let candidate_count = diagnostic.candidate_count().unwrap_or(0);
let skipped_candidate_count = diagnostic.skipped_candidate_count().unwrap_or(0);
diagnostic.set_reason(if candidate_count == 0 {
"candidate_list_empty".to_string()
} else if skipped_candidate_count >= candidate_count
&& diagnostic.skip_reason_len() == 1
&& diagnostic.skip_reason_count("api_key_concurrency_limit_reached") > 0
{
"api_key_concurrency_limit_reached".to_string()
} else if skipped_candidate_count >= candidate_count {
"all_candidates_skipped".to_string()
} else {
no_plan_reason.to_string()
});
}
pub fn record_ai_runtime_candidate_skip_reason_on_diagnostic<Diagnostic>(
diagnostic: &mut Diagnostic,
skip_reason: &'static str,
) where
Diagnostic: AiRuntimeMissDiagnosticFields,
{
diagnostic.record_skip_reason(skip_reason);
}
pub fn set_ai_runtime_miss_diagnostic_reason<Port>(
port: &Port,
trace_id: &str,
decision: &Port::Decision,
plan_kind: &str,
requested_model: Option<&str>,
reason: &str,
) where
Port: AiRuntimeMissDiagnosticPort,
{
port.set_runtime_miss_diagnostic(
trace_id,
port.build_runtime_miss_diagnostic(decision, plan_kind, requested_model, reason),
);
}
pub fn build_ai_runtime_execution_exhausted_diagnostic<Port>(
port: &Port,
decision: &Port::Decision,
plan_kind: &str,
requested_model: Option<&str>,
candidate_count: usize,
) -> Port::Diagnostic
where
Port: AiRuntimeMissDiagnosticPort,
{
let mut diagnostic = port.build_runtime_miss_diagnostic(
decision,
plan_kind,
requested_model,
"execution_runtime_candidates_exhausted",
);
port.set_candidate_count(&mut diagnostic, candidate_count);
diagnostic
}
pub fn set_ai_runtime_execution_exhausted_diagnostic<Port>(
port: &Port,
trace_id: &str,
decision: &Port::Decision,
plan_kind: &str,
requested_model: Option<&str>,
candidate_count: usize,
) where
Port: AiRuntimeMissDiagnosticPort,
{
port.set_runtime_miss_diagnostic(
trace_id,
build_ai_runtime_execution_exhausted_diagnostic(
port,
decision,
plan_kind,
requested_model,
candidate_count,
),
);
}
pub fn build_ai_runtime_candidate_evaluation_diagnostic<Port>(
port: &Port,
decision: &Port::Decision,
plan_kind: &str,
requested_model: Option<&str>,
candidate_count: usize,
) -> Port::Diagnostic
where
Port: AiRuntimeMissDiagnosticPort,
{
let mut diagnostic = port.build_runtime_miss_diagnostic(
decision,
plan_kind,
requested_model,
"candidate_evaluation_incomplete",
);
port.apply_candidate_evaluation_progress(&mut diagnostic, candidate_count);
diagnostic
}
pub fn set_ai_runtime_candidate_evaluation_diagnostic<Port>(
port: &Port,
trace_id: &str,
decision: &Port::Decision,
plan_kind: &str,
requested_model: Option<&str>,
candidate_count: usize,
) where
Port: AiRuntimeMissDiagnosticPort,
{
port.set_runtime_miss_diagnostic(
trace_id,
build_ai_runtime_candidate_evaluation_diagnostic(
port,
decision,
plan_kind,
requested_model,
candidate_count,
),
);
}
pub fn apply_ai_runtime_candidate_evaluation_progress<Port>(
port: &Port,
trace_id: &str,
candidate_count: usize,
) where
Port: AiRuntimeMissDiagnosticPort,
{
port.mutate_runtime_miss_diagnostic(trace_id, |diagnostic| {
port.apply_candidate_evaluation_progress(diagnostic, candidate_count);
});
}
pub fn apply_ai_runtime_candidate_evaluation_progress_preserving_candidate_signal<Port>(
port: &Port,
trace_id: &str,
candidate_count: usize,
) where
Port: AiRuntimeMissDiagnosticPort,
{
let preserve_existing_candidate_signal =
candidate_count == 0 && port.runtime_miss_diagnostic_has_candidate_signal(trace_id);
if preserve_existing_candidate_signal {
return;
}
apply_ai_runtime_candidate_evaluation_progress(port, trace_id, candidate_count);
}
pub fn apply_ai_runtime_candidate_terminal_reason<Port>(
port: &Port,
trace_id: &str,
no_plan_reason: &'static str,
) where
Port: AiRuntimeMissDiagnosticPort,
{
port.mutate_runtime_miss_diagnostic(trace_id, |diagnostic| {
port.apply_candidate_terminal_plan_reason(diagnostic, no_plan_reason);
});
}
pub fn record_ai_runtime_candidate_skip_reason<Port>(
port: &Port,
trace_id: &str,
skip_reason: &'static str,
) where
Port: AiRuntimeMissDiagnosticPort,
{
port.mutate_runtime_miss_diagnostic(trace_id, |diagnostic| {
port.record_candidate_skip_reason(diagnostic, skip_reason);
});
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::BTreeMap;
use std::sync::Mutex;
#[derive(Debug, Clone, PartialEq, Eq)]
struct TestDecision {
id: &'static str,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct TestDiagnostic {
decision_id: String,
plan_kind: String,
requested_model: Option<String>,
reason: String,
candidate_count: Option<usize>,
terminal_reason: Option<&'static str>,
skip_reasons: BTreeMap<&'static str, usize>,
}
#[derive(Default)]
struct TestPort {
diagnostics: Mutex<BTreeMap<String, TestDiagnostic>>,
}
impl AiRuntimeMissDiagnosticPort for TestPort {
type Decision = TestDecision;
type Diagnostic = TestDiagnostic;
fn build_runtime_miss_diagnostic(
&self,
decision: &Self::Decision,
plan_kind: &str,
requested_model: Option<&str>,
reason: &str,
) -> Self::Diagnostic {
TestDiagnostic {
decision_id: decision.id.to_string(),
plan_kind: plan_kind.to_string(),
requested_model: requested_model.map(str::to_string),
reason: reason.to_string(),
..Default::default()
}
}
fn set_candidate_count(&self, diagnostic: &mut Self::Diagnostic, candidate_count: usize) {
diagnostic.candidate_count = Some(candidate_count);
}
fn apply_candidate_evaluation_progress(
&self,
diagnostic: &mut Self::Diagnostic,
candidate_count: usize,
) {
diagnostic.candidate_count = Some(candidate_count);
}
fn apply_candidate_terminal_plan_reason(
&self,
diagnostic: &mut Self::Diagnostic,
no_plan_reason: &'static str,
) {
diagnostic.terminal_reason = Some(no_plan_reason);
}
fn record_candidate_skip_reason(
&self,
diagnostic: &mut Self::Diagnostic,
skip_reason: &'static str,
) {
*diagnostic.skip_reasons.entry(skip_reason).or_insert(0) += 1;
}
fn set_runtime_miss_diagnostic(&self, trace_id: &str, diagnostic: Self::Diagnostic) {
self.diagnostics
.lock()
.unwrap()
.insert(trace_id.to_string(), diagnostic);
}
fn mutate_runtime_miss_diagnostic<F>(&self, trace_id: &str, apply: F)
where
F: FnOnce(&mut Self::Diagnostic) + Send,
{
let mut diagnostics = self.diagnostics.lock().unwrap();
let diagnostic = diagnostics.entry(trace_id.to_string()).or_default();
apply(diagnostic);
}
fn runtime_miss_diagnostic_has_candidate_signal(&self, trace_id: &str) -> bool {
self.diagnostics
.lock()
.unwrap()
.get(trace_id)
.is_some_and(|diagnostic| diagnostic.candidate_count.unwrap_or_default() > 0)
}
}
impl AiRuntimeMissDiagnosticFields for TestDiagnostic {
fn set_reason(&mut self, reason: String) {
self.reason = reason;
}
fn set_candidate_count(&mut self, candidate_count: usize) {
self.candidate_count = Some(candidate_count);
}
fn candidate_count(&self) -> Option<usize> {
self.candidate_count
}
fn skipped_candidate_count(&self) -> Option<usize> {
self.skip_reasons.values().copied().sum::<usize>().into()
}
fn skip_reason_count(&self, skip_reason: &str) -> usize {
self.skip_reasons.get(skip_reason).copied().unwrap_or(0)
}
fn skip_reason_len(&self) -> usize {
self.skip_reasons.len()
}
fn record_skip_reason(&mut self, skip_reason: &'static str) {
*self.skip_reasons.entry(skip_reason).or_insert(0) += 1;
}
}
#[test]
fn runtime_miss_builds_and_sets_execution_exhausted_diagnostic() {
let port = TestPort::default();
set_ai_runtime_execution_exhausted_diagnostic(
&port,
"trace-a",
&TestDecision { id: "decision-a" },
"openai_chat",
Some("gpt-5"),
3,
);
let diagnostic = port.diagnostics.lock().unwrap().get("trace-a").cloned();
assert_eq!(
diagnostic,
Some(TestDiagnostic {
decision_id: "decision-a".to_string(),
plan_kind: "openai_chat".to_string(),
requested_model: Some("gpt-5".to_string()),
reason: "execution_runtime_candidates_exhausted".to_string(),
candidate_count: Some(3),
..Default::default()
})
);
}
#[test]
fn runtime_miss_preserves_candidate_signal_and_records_terminal_updates() {
let port = TestPort::default();
apply_ai_runtime_candidate_evaluation_progress(&port, "trace-a", 2);
apply_ai_runtime_candidate_evaluation_progress_preserving_candidate_signal(
&port, "trace-a", 0,
);
apply_ai_runtime_candidate_terminal_reason(&port, "trace-a", "no_local_sync_plans");
record_ai_runtime_candidate_skip_reason(&port, "trace-a", "transport_missing");
let diagnostic = port.diagnostics.lock().unwrap().get("trace-a").cloned();
assert_eq!(
diagnostic,
Some(TestDiagnostic {
candidate_count: Some(2),
terminal_reason: Some("no_local_sync_plans"),
skip_reasons: BTreeMap::from([("transport_missing", 1)]),
..Default::default()
})
);
}
#[test]
fn runtime_miss_diagnostic_field_helpers_apply_candidate_reason_state_machine() {
let mut diagnostic = TestDiagnostic::default();
apply_ai_runtime_candidate_evaluation_progress_to_diagnostic(&mut diagnostic, 0);
assert_eq!(diagnostic.candidate_count, Some(0));
assert_eq!(diagnostic.reason, "candidate_list_empty");
apply_ai_runtime_candidate_evaluation_progress_to_diagnostic(&mut diagnostic, 3);
assert_eq!(diagnostic.candidate_count, Some(3));
assert_eq!(diagnostic.reason, "candidate_evaluation_incomplete");
record_ai_runtime_candidate_skip_reason_on_diagnostic(
&mut diagnostic,
"api_key_concurrency_limit_reached",
);
record_ai_runtime_candidate_skip_reason_on_diagnostic(
&mut diagnostic,
"api_key_concurrency_limit_reached",
);
apply_ai_runtime_candidate_terminal_plan_reason_to_diagnostic(
&mut diagnostic,
"no_local_sync_plans",
);
assert_eq!(diagnostic.reason, "no_local_sync_plans");
diagnostic.candidate_count = Some(2);
apply_ai_runtime_candidate_terminal_plan_reason_to_diagnostic(
&mut diagnostic,
"no_local_sync_plans",
);
assert_eq!(diagnostic.reason, "api_key_concurrency_limit_reached");
record_ai_runtime_candidate_skip_reason_on_diagnostic(&mut diagnostic, "transport_missing");
apply_ai_runtime_candidate_terminal_plan_reason_to_diagnostic(
&mut diagnostic,
"no_local_sync_plans",
);
assert_eq!(diagnostic.reason, "all_candidates_skipped");
}
}

236
crates/aether-ai-serving/src/surface_spec.rs Normal file
View File

@@ -0,0 +1,236 @@
use aether_ai_surfaces::api::{
LocalGeminiFilesSpec, LocalOpenAiImageSpec, LocalOpenAiResponsesSpec,
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec, LocalStandardSourceFamily,
LocalStandardSpec, LocalVideoCreateFamily, LocalVideoCreateSpec,
};
use serde_json::Value;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AiRequestedModelFamily {
Standard,
Gemini,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AiExecutionSurfaceSpecMetadata {
pub api_format: &'static str,
pub decision_kind: &'static str,
pub report_kind: Option<&'static str>,
pub require_streaming: bool,
pub requested_model_family: Option<AiRequestedModelFamily>,
}
pub const fn ai_requested_model_family_for_standard_source(
family: LocalStandardSourceFamily,
) -> AiRequestedModelFamily {
match family {
LocalStandardSourceFamily::Standard => AiRequestedModelFamily::Standard,
LocalStandardSourceFamily::Gemini => AiRequestedModelFamily::Gemini,
}
}
pub const fn ai_standard_spec_metadata(spec: LocalStandardSpec) -> AiExecutionSurfaceSpecMetadata {
AiExecutionSurfaceSpecMetadata {
api_format: spec.api_format,
decision_kind: spec.decision_kind,
report_kind: Some(spec.report_kind),
require_streaming: spec.require_streaming,
requested_model_family: Some(ai_requested_model_family_for_standard_source(spec.family)),
}
}
pub const fn ai_same_format_provider_spec_metadata(
spec: LocalSameFormatProviderSpec,
) -> AiExecutionSurfaceSpecMetadata {
AiExecutionSurfaceSpecMetadata {
api_format: spec.api_format,
decision_kind: spec.decision_kind,
report_kind: Some(spec.report_kind),
require_streaming: spec.require_streaming,
requested_model_family: Some(ai_requested_model_family_for_same_format_provider(
spec.family,
)),
}
}
pub const fn ai_openai_responses_spec_metadata(
spec: LocalOpenAiResponsesSpec,
) -> AiExecutionSurfaceSpecMetadata {
AiExecutionSurfaceSpecMetadata {
api_format: spec.api_format,
decision_kind: spec.decision_kind,
report_kind: Some(spec.report_kind),
require_streaming: spec.require_streaming,
requested_model_family: None,
}
}
pub const fn ai_gemini_files_spec_metadata(
spec: LocalGeminiFilesSpec,
) -> AiExecutionSurfaceSpecMetadata {
AiExecutionSurfaceSpecMetadata {
api_format: "gemini:files",
decision_kind: spec.decision_kind,
report_kind: spec.report_kind,
require_streaming: spec.require_streaming,
requested_model_family: None,
}
}
pub const fn ai_openai_image_spec_metadata(
spec: LocalOpenAiImageSpec,
) -> AiExecutionSurfaceSpecMetadata {
AiExecutionSurfaceSpecMetadata {
api_format: spec.api_format,
decision_kind: spec.decision_kind,
report_kind: Some(spec.report_kind),
require_streaming: spec.require_streaming,
requested_model_family: Some(AiRequestedModelFamily::Standard),
}
}
pub const fn ai_video_create_spec_metadata(
spec: LocalVideoCreateSpec,
) -> AiExecutionSurfaceSpecMetadata {
AiExecutionSurfaceSpecMetadata {
api_format: spec.api_format,
decision_kind: spec.decision_kind,
report_kind: Some(spec.report_kind),
require_streaming: false,
requested_model_family: Some(ai_requested_model_family_for_video_create(spec.family)),
}
}
pub const fn ai_requested_model_family_for_same_format_provider(
family: LocalSameFormatProviderFamily,
) -> AiRequestedModelFamily {
match family {
LocalSameFormatProviderFamily::Standard => AiRequestedModelFamily::Standard,
LocalSameFormatProviderFamily::Gemini => AiRequestedModelFamily::Gemini,
}
}
pub const fn ai_requested_model_family_for_video_create(
family: LocalVideoCreateFamily,
) -> AiRequestedModelFamily {
match family {
LocalVideoCreateFamily::OpenAi => AiRequestedModelFamily::Standard,
LocalVideoCreateFamily::Gemini => AiRequestedModelFamily::Gemini,
}
}
pub fn extract_ai_gemini_model_from_path(path: &str) -> Option<String> {
let (_, suffix) = path.split_once("/models/")?;
let model = suffix
.split_once(':')
.map(|(value, _)| value)
.unwrap_or(suffix);
let model = model.trim();
if model.is_empty() {
None
} else {
Some(model.to_string())
}
}
pub fn extract_ai_standard_requested_model(body_json: &Value) -> Option<String> {
body_json
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
pub fn extract_ai_requested_model_from_request_path(
request_path: &str,
body_json: &Value,
family: AiRequestedModelFamily,
) -> Option<String> {
match family {
AiRequestedModelFamily::Standard => extract_ai_standard_requested_model(body_json),
AiRequestedModelFamily::Gemini => extract_ai_gemini_model_from_path(request_path),
}
}
#[cfg(test)]
mod tests {
use super::*;
use aether_ai_surfaces::api::{LocalStandardSourceMode, OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND};
#[test]
fn standard_spec_metadata_maps_model_family_and_report_kind() {
let metadata = ai_standard_spec_metadata(LocalStandardSpec {
mode: LocalStandardSourceMode::Chat,
family: LocalStandardSourceFamily::Standard,
api_format: "openai:chat",
decision_kind: "openai_chat_sync",
report_kind: OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND,
require_streaming: false,
});
assert_eq!(metadata.api_format, "openai:chat");
assert_eq!(
metadata.requested_model_family,
Some(AiRequestedModelFamily::Standard)
);
}
#[test]
fn video_spec_metadata_maps_gemini_family_without_stream_requirement() {
let metadata = ai_video_create_spec_metadata(LocalVideoCreateSpec {
family: LocalVideoCreateFamily::Gemini,
api_format: "gemini:video",
decision_kind: "gemini_video_create",
report_kind: "gemini_video_create_success",
});
assert!(!metadata.require_streaming);
assert_eq!(
metadata.requested_model_family,
Some(AiRequestedModelFamily::Gemini)
);
}
#[test]
fn gemini_model_path_parser_trims_method_suffix() {
let model = extract_ai_gemini_model_from_path(
"/v1beta/models/gemini-2.5-pro:streamGenerateContent",
);
assert_eq!(model.as_deref(), Some("gemini-2.5-pro"));
}
#[test]
fn standard_requested_model_parser_reads_request_body_model() {
let requested_model = extract_ai_standard_requested_model(
&serde_json::json!({ "model": " claude-sonnet-4 " }),
);
assert_eq!(requested_model.as_deref(), Some("claude-sonnet-4"));
}
#[test]
fn requested_model_parser_delegates_by_family() {
let body = serde_json::json!({ "model": " claude-sonnet-4 " });
assert_eq!(
extract_ai_requested_model_from_request_path(
"/v1/chat/completions",
&body,
AiRequestedModelFamily::Standard,
)
.as_deref(),
Some("claude-sonnet-4")
);
assert_eq!(
extract_ai_requested_model_from_request_path(
"/v1beta/models/gemini-2.5-pro:generateContent",
&serde_json::json!({}),
AiRequestedModelFamily::Gemini,
)
.as_deref(),
Some("gemini-2.5-pro")
);
}
}

6
crates/aether-ai-pipeline/Cargo.toml → crates/aether-ai-surfaces/Cargo.toml
View File

@@ -1,20 +1,18 @@
[package]
name = "aether-ai-pipeline"
name = "aether-ai-surfaces"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Shared AI pipeline contracts and planner logic for Aether"
description = "Pure AI API surface contracts, routing, adaptation, and finalize logic for Aether"
[dependencies]
aether-ai-formats.workspace = true
aether-contracts.workspace = true
aether-usage-runtime.workspace = true
base64.workspace = true
http.workspace = true
serde.workspace = true
serde_json.workspace = true
aether-provider-transport.workspace = true
sha1 = "0.10"
sha2.workspace = true
uuid.workspace = true

4
crates/aether-ai-pipeline/src/adaptation/kiro_stream.rs → crates/aether-ai-surfaces/src/adaptation/kiro_stream.rs
View File

@@ -1,5 +1,9 @@
use serde_json::{json, Value};
pub use self::state::KiroToClaudeCliStreamState;
mod state;
pub const KIRO_CONTEXT_WINDOW_TOKENS: f64 = 200_000.0;
pub const KIRO_MAX_THINKING_BUFFER: usize = 1024 * 1024;

71
crates/aether-ai-surfaces/src/adaptation/kiro_stream/state.rs Normal file
View File

@@ -0,0 +1,71 @@
use std::collections::BTreeMap;
const MAX_MESSAGE_SIZE: usize = 16 * 1024 * 1024;
const MAX_BUFFER_SIZE: usize = MAX_MESSAGE_SIZE;
const MAX_ERRORS: usize = 5;
#[derive(Default)]
pub struct KiroToClaudeCliStreamState {
decoder: EventStreamDecoder,
state: KiroClaudeStreamState,
started: bool,
}
#[derive(Default)]
struct KiroClaudeStreamState {
model: String,
thinking_enabled: bool,
estimated_input_tokens: usize,
message_id: String,
output_tokens: usize,
context_input_tokens: Option<usize>,
next_block_index: usize,
open_blocks: BTreeMap<usize, String>,
text_block_index: Option<usize>,
thinking_block_index: Option<usize>,
tool_block_indices: BTreeMap<String, usize>,
thinking_buffer: String,
in_thinking_block: bool,
thinking_extracted: bool,
strip_thinking_leading_newline: bool,
has_tool_use: bool,
stop_reason_override: Option<String>,
had_error: bool,
last_content: String,
}
#[derive(Default)]
struct EventStreamDecoder {
buffer: Vec<u8>,
error_count: usize,
stopped: bool,
}
#[derive(Default)]
struct AwsHeaders {
values: BTreeMap<String, AwsHeaderValue>,
}
enum AwsHeaderValue {
Ignored,
String(String),
}
struct AwsEventFrame {
headers: AwsHeaders,
payload: Vec<u8>,
}
enum FrameParseError {
Incomplete,
Invalid(String),
}
#[path = "stream/decoder.rs"]
mod decoder;
#[path = "stream/state.rs"]
mod stream_state;
#[cfg(test)]
#[path = "stream/tests.rs"]
mod tests;

222
crates/aether-ai-surfaces/src/adaptation/kiro_stream/stream/decoder.rs Normal file
View File

@@ -0,0 +1,222 @@
use std::collections::BTreeMap;
use super::{
AwsEventFrame, AwsHeaderValue, AwsHeaders, EventStreamDecoder, FrameParseError,
MAX_BUFFER_SIZE, MAX_ERRORS, MAX_MESSAGE_SIZE,
};
use crate::adaptation::kiro_stream::kiro_crc32 as crc32;
impl EventStreamDecoder {
pub(super) fn feed(&mut self, data: &[u8]) -> Result<(), String> {
if self.stopped || data.is_empty() {
return Ok(());
}
let new_size = self.buffer.len() + data.len();
if new_size > MAX_BUFFER_SIZE {
self.stopped = true;
return Err(format!(
"buffer overflow: size={new_size} max={MAX_BUFFER_SIZE}"
));
}
self.buffer.extend_from_slice(data);
Ok(())
}
pub(super) fn decode_available(&mut self) -> Result<Vec<AwsEventFrame>, String> {
let mut out = Vec::new();
if self.stopped {
return Ok(out);
}
loop {
match parse_frame(&self.buffer) {
Ok(Some((frame, consumed))) => {
if consumed == 0 {
break;
}
out.push(frame);
self.buffer.drain(..consumed);
self.error_count = 0;
}
Ok(None) => break,
Err(FrameParseError::Incomplete) => break,
Err(FrameParseError::Invalid(message)) => {
self.error_count += 1;
if self.error_count >= MAX_ERRORS {
self.stopped = true;
return Err(message);
}
if self.buffer.is_empty() {
break;
}
self.buffer.drain(..1);
}
}
}
Ok(out)
}
}
impl AwsHeaders {
fn get_string(&self, name: &str) -> Option<&str> {
match self.values.get(name) {
Some(AwsHeaderValue::String(value)) => Some(value.as_str()),
_ => None,
}
}
pub(super) fn message_type(&self) -> Option<&str> {
self.get_string(":message-type")
}
pub(super) fn event_type(&self) -> Option<&str> {
self.get_string(":event-type")
}
pub(super) fn exception_type(&self) -> Option<&str> {
self.get_string(":exception-type")
}
pub(super) fn error_code(&self) -> Option<&str> {
self.get_string(":error-code")
}
}
fn parse_frame(buffer: &[u8]) -> Result<Option<(AwsEventFrame, usize)>, FrameParseError> {
if buffer.len() < 12 {
return Ok(None);
}
let total_length = u32::from_be_bytes(buffer[0..4].try_into().expect("slice size")) as usize;
let header_length = u32::from_be_bytes(buffer[4..8].try_into().expect("slice size")) as usize;
let prelude_crc = u32::from_be_bytes(buffer[8..12].try_into().expect("slice size"));
if total_length < 16 {
return Err(FrameParseError::Invalid(format!(
"message too small: length={total_length}"
)));
}
if total_length > MAX_MESSAGE_SIZE {
return Err(FrameParseError::Invalid(format!(
"message too large: length={total_length}"
)));
}
if buffer.len() < total_length {
return Ok(None);
}
if crc32(&buffer[0..8]) != prelude_crc {
return Err(FrameParseError::Invalid("prelude crc mismatch".to_string()));
}
let message_crc = u32::from_be_bytes(
buffer[total_length - 4..total_length]
.try_into()
.expect("slice size"),
);
if crc32(&buffer[..total_length - 4]) != message_crc {
return Err(FrameParseError::Invalid("message crc mismatch".to_string()));
}
let headers_start = 12;
let headers_end = headers_start + header_length;
if headers_end > total_length - 4 {
return Err(FrameParseError::Invalid(
"header length exceeds frame boundary".to_string(),
));
}
let headers = parse_headers(&buffer[headers_start..headers_end], header_length)?;
let payload = buffer[headers_end..total_length - 4].to_vec();
Ok(Some((AwsEventFrame { headers, payload }, total_length)))
}
fn parse_headers(data: &[u8], header_length: usize) -> Result<AwsHeaders, FrameParseError> {
if data.len() < header_length {
return Err(FrameParseError::Incomplete);
}
let mut values = BTreeMap::new();
let mut offset = 0usize;
while offset < header_length {
ensure_header_bytes(data, offset, 1)?;
let name_len = data[offset] as usize;
offset += 1;
if name_len == 0 {
return Err(FrameParseError::Invalid(
"header name length cannot be 0".to_string(),
));
}
ensure_header_bytes(data, offset, name_len)?;
let name = String::from_utf8_lossy(&data[offset..offset + name_len]).to_string();
offset += name_len;
ensure_header_bytes(data, offset, 1)?;
let value_type = data[offset];
offset += 1;
let value = match value_type {
0 => AwsHeaderValue::Ignored,
1 => AwsHeaderValue::Ignored,
2 => {
ensure_header_bytes(data, offset, 1)?;
let _ = i8::from_be_bytes([data[offset]]);
offset += 1;
AwsHeaderValue::Ignored
}
3 => {
ensure_header_bytes(data, offset, 2)?;
let _ = i16::from_be_bytes(data[offset..offset + 2].try_into().expect("slice"));
offset += 2;
AwsHeaderValue::Ignored
}
4 => {
ensure_header_bytes(data, offset, 4)?;
let _ = i32::from_be_bytes(data[offset..offset + 4].try_into().expect("slice"));
offset += 4;
AwsHeaderValue::Ignored
}
5 | 8 => {
ensure_header_bytes(data, offset, 8)?;
let _ = i64::from_be_bytes(data[offset..offset + 8].try_into().expect("slice"));
offset += 8;
AwsHeaderValue::Ignored
}
6 => {
ensure_header_bytes(data, offset, 2)?;
let length = u16::from_be_bytes(data[offset..offset + 2].try_into().expect("slice"))
as usize;
offset += 2;
ensure_header_bytes(data, offset, length)?;
offset += length;
AwsHeaderValue::Ignored
}
7 => {
ensure_header_bytes(data, offset, 2)?;
let length = u16::from_be_bytes(data[offset..offset + 2].try_into().expect("slice"))
as usize;
offset += 2;
ensure_header_bytes(data, offset, length)?;
let out = String::from_utf8_lossy(&data[offset..offset + length]).to_string();
offset += length;
AwsHeaderValue::String(out)
}
9 => {
ensure_header_bytes(data, offset, 16)?;
offset += 16;
AwsHeaderValue::Ignored
}
other => {
return Err(FrameParseError::Invalid(format!(
"invalid header type: {other}"
)));
}
};
values.insert(name, value);
}
Ok(AwsHeaders { values })
}
fn ensure_header_bytes(data: &[u8], offset: usize, needed: usize) -> Result<(), FrameParseError> {
let available = data.len().saturating_sub(offset);
if available < needed {
return Err(FrameParseError::Incomplete);
}
Ok(())
}

8
crates/aether-ai-surfaces/src/adaptation/kiro_stream/stream/state.rs Normal file
View File

@@ -0,0 +1,8 @@
#[path = "state/blocks.rs"]
mod blocks;
#[path = "state/events.rs"]
mod events;
#[path = "state/finalize.rs"]
mod finalize;
#[path = "state/lifecycle.rs"]
mod lifecycle;

97
crates/aether-ai-surfaces/src/adaptation/kiro_stream/stream/state/blocks.rs Normal file
View File

@@ -0,0 +1,97 @@
use serde_json::{json, Value};
use super::super::KiroClaudeStreamState;
impl KiroClaudeStreamState {
pub(super) fn ensure_text_block_open(&mut self) -> Vec<Value> {
if let Some(idx) = self.text_block_index {
if self
.open_blocks
.get(&idx)
.map(|value| value == "text")
.unwrap_or(false)
{
return Vec::new();
}
}
let idx = self.next_block_index;
self.next_block_index += 1;
self.text_block_index = Some(idx);
self.open_blocks.insert(idx, "text".to_string());
vec![json!({
"type": "content_block_start",
"index": idx,
"content_block": {"type": "text", "text": ""}
})]
}
pub(super) fn ensure_thinking_block_open(&mut self) -> Vec<Value> {
if let Some(idx) = self.thinking_block_index {
if self
.open_blocks
.get(&idx)
.map(|value| value == "thinking")
.unwrap_or(false)
{
return Vec::new();
}
}
let idx = self.next_block_index;
self.next_block_index += 1;
self.thinking_block_index = Some(idx);
self.open_blocks.insert(idx, "thinking".to_string());
vec![json!({
"type": "content_block_start",
"index": idx,
"content_block": {"type": "thinking", "thinking": ""}
})]
}
pub(super) fn close_block(&mut self, idx: usize) -> Vec<Value> {
if self.open_blocks.remove(&idx).is_none() {
return Vec::new();
}
vec![json!({"type": "content_block_stop", "index": idx})]
}
pub(super) fn emit_text_delta(&mut self, text: &str) -> Vec<Value> {
if text.is_empty() {
return Vec::new();
}
let mut events = self.ensure_text_block_open();
let idx = self.text_block_index.unwrap_or_default();
events.push(json!({
"type": "content_block_delta",
"index": idx,
"delta": {"type": "text_delta", "text": text}
}));
events
}
pub(super) fn emit_thinking_delta(&mut self, thinking: &str) -> Vec<Value> {
if thinking.is_empty() {
return Vec::new();
}
let mut events = self.ensure_thinking_block_open();
let idx = self.thinking_block_index.unwrap_or_default();
events.push(json!({
"type": "content_block_delta",
"index": idx,
"delta": {"type": "thinking_delta", "thinking": thinking}
}));
events
}
pub(super) fn close_thinking_block(&mut self) -> Vec<Value> {
let Some(idx) = self.thinking_block_index else {
return Vec::new();
};
let mut events = vec![json!({
"type": "content_block_delta",
"index": idx,
"delta": {"type": "thinking_delta", "thinking": ""}
})];
events.extend(self.close_block(idx));
events
}
}

331
crates/aether-ai-surfaces/src/adaptation/kiro_stream/stream/state/events.rs Normal file
View File

@@ -0,0 +1,331 @@
use serde_json::{json, Value};
use crate::adaptation::kiro_stream::{
calculate_kiro_context_input_tokens, encode_kiro_sse_events, estimate_kiro_tokens,
find_kiro_real_thinking_end_tag, find_kiro_real_thinking_end_tag_at_buffer_end,
find_kiro_real_thinking_start_tag, KIRO_MAX_THINKING_BUFFER,
};
use crate::finalize::AiSurfaceFinalizeError;
use super::super::AwsEventFrame;
use super::super::KiroClaudeStreamState;
fn floor_char_boundary(text: &str, index: usize) -> usize {
let mut boundary = index.min(text.len());
while boundary > 0 && !text.is_char_boundary(boundary) {
boundary -= 1;
}
boundary
}
fn split_preserving_trailing_bytes(
buffer: &str,
trailing_bytes: usize,
) -> Option<(String, String)> {
if buffer.len() <= trailing_bytes {
return None;
}
let split = floor_char_boundary(buffer, buffer.len() - trailing_bytes);
if split == 0 {
return None;
}
Some((buffer[..split].to_string(), buffer[split..].to_string()))
}
impl KiroClaudeStreamState {
pub(super) fn process_frame(
&mut self,
frame: AwsEventFrame,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let message_type = frame.headers.message_type().unwrap_or("event");
match message_type {
"event" => self.process_event_frame(frame),
"exception" => self.process_exception_frame(frame),
"error" => self.process_error_frame(frame),
_ => Ok(Vec::new()),
}
}
pub(super) fn process_event_frame(
&mut self,
frame: AwsEventFrame,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let event_type = frame.headers.event_type().unwrap_or_default();
let payload: Value = if frame.payload.is_empty() {
json!({})
} else {
serde_json::from_slice(&frame.payload).unwrap_or_else(|_| json!({}))
};
let payload_object = payload.as_object();
let mut events = Vec::new();
match event_type {
"assistantResponseEvent" => {
if let Some(content) = payload_object
.and_then(|value| value.get("content"))
.and_then(Value::as_str)
{
events.extend(self.process_assistant_response(content));
}
}
"toolUseEvent" => {
if let Some(payload_object) = payload_object {
let name = payload_object
.get("name")
.and_then(Value::as_str)
.unwrap_or_default();
let tool_use_id = payload_object
.get("toolUseId")
.or_else(|| payload_object.get("tool_use_id"))
.and_then(Value::as_str)
.unwrap_or_default();
let input_json = match payload_object.get("input") {
None | Some(Value::Null) => String::new(),
Some(Value::String(text)) => text.clone(),
Some(other) => {
serde_json::to_string(other).map_err(AiSurfaceFinalizeError::from)?
}
};
let stop = payload_object
.get("stop")
.and_then(Value::as_bool)
.unwrap_or(false);
events.extend(self.process_tool_use(name, tool_use_id, &input_json, stop));
}
}
"contextUsageEvent" => {
if let Some(percentage) = payload_object
.and_then(|value| value.get("contextUsagePercentage"))
.and_then(Value::as_f64)
{
self.context_input_tokens =
Some(calculate_kiro_context_input_tokens(percentage));
}
}
_ => {}
}
encode_kiro_sse_events(events).map_err(AiSurfaceFinalizeError::from)
}
pub(super) fn process_exception_frame(
&mut self,
frame: AwsEventFrame,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let exception_type = frame
.headers
.exception_type()
.unwrap_or("UnknownException")
.to_string();
if exception_type == "ContentLengthExceededException" {
self.stop_reason_override = Some("max_tokens".to_string());
return Ok(Vec::new());
}
self.emit_stream_error("upstream_exception", &exception_type)
}
pub(super) fn process_error_frame(
&mut self,
frame: AwsEventFrame,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let error_code = frame
.headers
.error_code()
.unwrap_or("UnknownError")
.to_string();
self.emit_stream_error("upstream_error", &error_code)
}
pub(super) fn process_assistant_response(&mut self, content: &str) -> Vec<Value> {
if content.is_empty() || content == self.last_content {
return Vec::new();
}
self.last_content = content.to_string();
self.output_tokens += estimate_kiro_tokens(content);
if !self.thinking_enabled {
return self.emit_text_delta(content);
}
self.thinking_buffer.push_str(content);
if self.thinking_buffer.len() > KIRO_MAX_THINKING_BUFFER {
let overflow = std::mem::take(&mut self.thinking_buffer);
if self.in_thinking_block {
let mut events = self.emit_thinking_delta(&overflow);
events.extend(self.close_thinking_block());
self.in_thinking_block = false;
self.thinking_extracted = true;
return events;
}
return self.emit_text_delta(&overflow);
}
let mut events = Vec::new();
loop {
if !self.in_thinking_block && !self.thinking_extracted {
if let Some(start_pos) = find_kiro_real_thinking_start_tag(&self.thinking_buffer) {
let before = self.thinking_buffer[..start_pos].to_string();
if !before.trim().is_empty() {
events.extend(self.emit_text_delta(&before));
}
self.in_thinking_block = true;
self.strip_thinking_leading_newline = true;
self.thinking_buffer =
self.thinking_buffer[start_pos + "<thinking>".len()..].to_string();
events.extend(self.ensure_thinking_block_open());
continue;
}
let keep = "<thinking>".len();
if let Some((safe, remaining)) =
split_preserving_trailing_bytes(&self.thinking_buffer, keep)
{
if !safe.trim().is_empty() {
events.extend(self.emit_text_delta(&safe));
self.thinking_buffer = remaining;
}
}
break;
}
if self.in_thinking_block {
if self.strip_thinking_leading_newline {
if self.thinking_buffer.starts_with('\n') {
self.thinking_buffer.remove(0);
self.strip_thinking_leading_newline = false;
} else if !self.thinking_buffer.is_empty() {
self.strip_thinking_leading_newline = false;
}
}
if let Some(end_pos) = find_kiro_real_thinking_end_tag(&self.thinking_buffer) {
let thinking_text = self.thinking_buffer[..end_pos].to_string();
if !thinking_text.is_empty() {
events.extend(self.emit_thinking_delta(&thinking_text));
}
events.extend(self.close_thinking_block());
self.in_thinking_block = false;
self.thinking_extracted = true;
self.thinking_buffer =
self.thinking_buffer[end_pos + "</thinking>".len()..].to_string();
continue;
}
let keep = "</thinking>".len();
if let Some((safe, remaining)) =
split_preserving_trailing_bytes(&self.thinking_buffer, keep)
{
if !safe.is_empty() {
events.extend(self.emit_thinking_delta(&safe));
self.thinking_buffer = remaining;
}
}
break;
}
if !self.thinking_buffer.is_empty() {
let remaining = std::mem::take(&mut self.thinking_buffer);
events.extend(self.emit_text_delta(&remaining));
}
break;
}
events
}
pub(super) fn process_tool_use(
&mut self,
name: &str,
tool_use_id: &str,
input_json: &str,
stop: bool,
) -> Vec<Value> {
if tool_use_id.is_empty() {
return Vec::new();
}
self.has_tool_use = true;
let mut events = Vec::new();
if self.thinking_enabled && self.in_thinking_block && !self.thinking_buffer.is_empty() {
if let Some(end_pos) =
find_kiro_real_thinking_end_tag_at_buffer_end(&self.thinking_buffer)
{
let thinking_text = self.thinking_buffer[..end_pos].to_string();
if !thinking_text.is_empty() {
events.extend(self.emit_thinking_delta(&thinking_text));
}
events.extend(self.close_thinking_block());
let remaining = self.thinking_buffer[end_pos + "</thinking>".len()..].to_string();
self.thinking_buffer.clear();
self.in_thinking_block = false;
self.thinking_extracted = true;
if !remaining.is_empty() {
events.extend(self.emit_text_delta(&remaining));
}
} else {
let thinking = std::mem::take(&mut self.thinking_buffer);
events.extend(self.emit_thinking_delta(&thinking));
events.extend(self.close_thinking_block());
self.in_thinking_block = false;
self.thinking_extracted = true;
}
}
if self.thinking_enabled
&& !self.in_thinking_block
&& !self.thinking_extracted
&& !self.thinking_buffer.is_empty()
{
let buffered = std::mem::take(&mut self.thinking_buffer);
events.extend(self.emit_text_delta(&buffered));
}
if let Some(idx) = self.text_block_index.take() {
events.extend(self.close_block(idx));
}
let block_index = if let Some(block_index) = self.tool_block_indices.get(tool_use_id) {
*block_index
} else {
let block_index = self.next_block_index;
self.next_block_index += 1;
self.tool_block_indices
.insert(tool_use_id.to_string(), block_index);
block_index
};
if let std::collections::btree_map::Entry::Vacant(e) = self.open_blocks.entry(block_index) {
e.insert("tool_use".to_string());
events.push(json!({
"type": "content_block_start",
"index": block_index,
"content_block": {
"type": "tool_use",
"id": tool_use_id,
"name": name,
"input": {},
}
}));
}
if !input_json.is_empty() {
self.output_tokens += estimate_kiro_tokens(input_json);
events.push(json!({
"type": "content_block_delta",
"index": block_index,
"delta": {
"type": "input_json_delta",
"partial_json": input_json,
}
}));
}
if stop {
events.extend(self.close_block(block_index));
}
events
}
}

96
crates/aether-ai-surfaces/src/adaptation/kiro_stream/stream/state/finalize.rs Normal file
View File

@@ -0,0 +1,96 @@
use crate::adaptation::kiro_stream::{
build_kiro_final_message_sse_events, encode_kiro_sse_events,
find_kiro_real_thinking_end_tag_at_buffer_end,
};
use crate::finalize::AiSurfaceFinalizeError;
use super::super::KiroClaudeStreamState;
impl KiroClaudeStreamState {
pub(super) fn finalize(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.thinking_enabled && !self.thinking_buffer.is_empty() {
let flush_events = if self.in_thinking_block {
if let Some(end_pos) =
find_kiro_real_thinking_end_tag_at_buffer_end(&self.thinking_buffer)
{
let thinking_text = self.thinking_buffer[..end_pos].to_string();
let mut events = Vec::new();
if !thinking_text.is_empty() {
events.extend(self.emit_thinking_delta(&thinking_text));
}
events.extend(self.close_thinking_block());
let remaining =
self.thinking_buffer[end_pos + "</thinking>".len()..].to_string();
if !remaining.is_empty() {
events.extend(self.emit_text_delta(&remaining));
}
events
} else {
let mut events = self.emit_thinking_delta(&self.thinking_buffer.clone());
events.extend(self.close_thinking_block());
events
}
} else {
self.emit_text_delta(&self.thinking_buffer.clone())
};
self.thinking_buffer.clear();
self.in_thinking_block = false;
self.thinking_extracted = true;
let mut output =
encode_kiro_sse_events(flush_events).map_err(AiSurfaceFinalizeError::from)?;
for idx in self
.open_blocks
.keys()
.cloned()
.collect::<Vec<_>>()
.into_iter()
.rev()
{
output.extend(
encode_kiro_sse_events(self.close_block(idx))
.map_err(AiSurfaceFinalizeError::from)?,
);
}
output.extend(self.final_message_bytes()?);
return Ok(output);
}
let mut output = Vec::new();
for idx in self
.open_blocks
.keys()
.cloned()
.collect::<Vec<_>>()
.into_iter()
.rev()
{
output.extend(
encode_kiro_sse_events(self.close_block(idx))
.map_err(AiSurfaceFinalizeError::from)?,
);
}
output.extend(self.final_message_bytes()?);
Ok(output)
}
pub(super) fn final_message_bytes(&self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let stop_reason = self.stop_reason_override.clone().unwrap_or_else(|| {
if self.has_tool_use {
"tool_use"
} else {
"end_turn"
}
.to_string()
});
let input_tokens = self
.context_input_tokens
.unwrap_or(self.estimated_input_tokens) as u64;
encode_kiro_sse_events(build_kiro_final_message_sse_events(
&stop_reason,
input_tokens as usize,
self.output_tokens,
))
.map_err(AiSurfaceFinalizeError::from)
}
}

129
crates/aether-ai-surfaces/src/adaptation/kiro_stream/stream/state/lifecycle.rs Normal file
View File

@@ -0,0 +1,129 @@
use serde_json::Value;
use uuid::Uuid;
use crate::adaptation::kiro_stream::{
build_kiro_initial_sse_events, build_kiro_stream_error_sse_events, encode_kiro_sse_events,
};
use crate::finalize::AiSurfaceFinalizeError;
use super::super::{EventStreamDecoder, KiroClaudeStreamState, KiroToClaudeCliStreamState};
impl KiroToClaudeCliStreamState {
pub fn new(report_context: &Value) -> Self {
Self {
decoder: EventStreamDecoder::default(),
state: KiroClaudeStreamState::new(report_context),
started: false,
}
}
pub fn push_chunk(
&mut self,
_report_context: &Value,
chunk: &[u8],
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut output = Vec::new();
if !self.started {
self.started = true;
output.extend(self.state.generate_initial_bytes()?);
}
if let Err(err) = self.decoder.feed(chunk) {
output.extend(
self.state
.emit_stream_error("upstream_stream_error", &err)?,
);
return Ok(output);
}
match self.decoder.decode_available() {
Ok(frames) => {
for frame in frames {
output.extend(self.state.process_frame(frame)?);
}
}
Err(err) => {
output.extend(
self.state
.emit_stream_error("upstream_stream_error", &err)?,
);
}
}
Ok(output)
}
pub fn finish(&mut self, _report_context: &Value) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if !self.started || self.state.had_error {
return Ok(Vec::new());
}
self.state.finalize()
}
}
impl KiroClaudeStreamState {
pub(super) fn new(report_context: &Value) -> Self {
let model = report_context
.get("mapped_model")
.and_then(Value::as_str)
.filter(|value| !value.is_empty())
.or_else(|| {
report_context
.get("model")
.and_then(Value::as_str)
.filter(|value| !value.is_empty())
})
.unwrap_or("unknown")
.to_string();
let thinking_enabled = report_context
.get("original_request_body")
.and_then(Value::as_object)
.and_then(|body| body.get("thinking"))
.and_then(Value::as_object)
.and_then(|thinking| thinking.get("type"))
.and_then(Value::as_str)
.map(|value| {
value.trim().eq_ignore_ascii_case("enabled")
|| value.trim().eq_ignore_ascii_case("adaptive")
})
.unwrap_or(false);
let estimated_input_tokens = report_context
.get("input_tokens")
.and_then(Value::as_u64)
.map(|value| value as usize)
.unwrap_or(0);
Self {
model,
thinking_enabled,
estimated_input_tokens,
message_id: format!("msg_{}", Uuid::new_v4().simple()),
..Self::default()
}
}
pub(super) fn generate_initial_bytes(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let events = build_kiro_initial_sse_events(
&self.message_id,
&self.model,
self.estimated_input_tokens,
);
let mut events = events;
if !self.thinking_enabled {
events.extend(self.ensure_text_block_open());
}
encode_kiro_sse_events(events).map_err(AiSurfaceFinalizeError::from)
}
pub(super) fn emit_stream_error(
&mut self,
error_type: &str,
message: &str,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.had_error {
return Ok(Vec::new());
}
self.had_error = true;
encode_kiro_sse_events(build_kiro_stream_error_sse_events(error_type, message))
.map_err(AiSurfaceFinalizeError::from)
}
}

166
crates/aether-ai-surfaces/src/adaptation/kiro_stream/stream/tests.rs Normal file
View File

@@ -0,0 +1,166 @@
use crate::adaptation::kiro_stream::kiro_crc32 as crc32;
use serde_json::{json, Value};
use super::KiroToClaudeCliStreamState;
fn encode_string_header(name: &str, value: &str) -> Vec<u8> {
let mut out = Vec::new();
out.push(name.len() as u8);
out.extend_from_slice(name.as_bytes());
out.push(7);
out.extend_from_slice(&(value.len() as u16).to_be_bytes());
out.extend_from_slice(value.as_bytes());
out
}
fn encode_event_frame(message_type: &str, event_type: Option<&str>, payload: &Value) -> Vec<u8> {
let mut headers = encode_string_header(":message-type", message_type);
if let Some(event_type) = event_type {
headers.extend_from_slice(&encode_string_header(":event-type", event_type));
}
let payload_bytes = serde_json::to_vec(payload).expect("payload should encode");
encode_frame(headers, payload_bytes)
}
fn encode_frame(headers: Vec<u8>, payload: Vec<u8>) -> Vec<u8> {
let total_len = 12 + headers.len() + payload.len() + 4;
let header_len = headers.len();
let mut out = Vec::with_capacity(total_len);
out.extend_from_slice(&(total_len as u32).to_be_bytes());
out.extend_from_slice(&(header_len as u32).to_be_bytes());
let prelude_crc = crc32(&out[..8]);
out.extend_from_slice(&prelude_crc.to_be_bytes());
out.extend_from_slice(&headers);
out.extend_from_slice(&payload);
let message_crc = crc32(&out);
out.extend_from_slice(&message_crc.to_be_bytes());
out
}
fn kiro_report_context(thinking_enabled: bool) -> Value {
let mut context = json!({
"provider_api_format": "claude:messages",
"client_api_format": "claude:messages",
"envelope_name": "kiro:generateAssistantResponse",
"mapped_model": "claude-sonnet-4.5"
});
if thinking_enabled {
context["original_request_body"] = json!({
"thinking": {
"type": "enabled"
}
});
}
context
}
#[test]
fn kiro_stream_rewriter_converts_text_events_to_claude_sse() {
let report_context = kiro_report_context(false);
let mut rewriter = KiroToClaudeCliStreamState::new(&report_context);
let chunk = [
encode_event_frame(
"event",
Some("assistantResponseEvent"),
&json!({"content": "Hello from Kiro"}),
),
encode_event_frame(
"event",
Some("contextUsageEvent"),
&json!({"contextUsagePercentage": 1.0}),
),
]
.concat();
let first = rewriter
.push_chunk(&report_context, &chunk)
.expect("rewrite should succeed");
let rest = rewriter
.finish(&report_context)
.expect("finish should succeed");
let text = String::from_utf8([first, rest].concat()).expect("utf8 should decode");
assert!(text.contains("event: message_start"));
assert!(text.contains("\"type\":\"content_block_delta\""));
assert!(text.contains("Hello from Kiro"));
assert!(text.contains("\"stop_reason\":\"end_turn\""));
assert!(text.contains("\"input_tokens\":2000"));
}
#[test]
fn kiro_stream_rewriter_converts_tool_use_to_claude_events() {
let report_context = kiro_report_context(false);
let mut rewriter = KiroToClaudeCliStreamState::new(&report_context);
let chunk = [
encode_event_frame(
"event",
Some("assistantResponseEvent"),
&json!({"content": "Need a tool."}),
),
encode_event_frame(
"event",
Some("toolUseEvent"),
&json!({
"name": "get_weather",
"toolUseId": "tool_123",
"input": {"city": "SF"},
"stop": true
}),
),
]
.concat();
let first = rewriter
.push_chunk(&report_context, &chunk)
.expect("rewrite should succeed");
let rest = rewriter
.finish(&report_context)
.expect("finish should succeed");
let text = String::from_utf8([first, rest].concat()).expect("utf8 should decode");
assert!(text.contains("\"type\":\"tool_use\""));
assert!(text.contains("\"id\":\"tool_123\""));
assert!(text.contains("\"name\":\"get_weather\""));
assert!(text.contains("\"partial_json\":\"{\\\"city\\\":\\\"SF\\\"}\""));
assert!(text.contains("\"stop_reason\":\"tool_use\""));
}
#[test]
fn kiro_stream_rewriter_handles_multibyte_text_without_thinking_tag() {
let report_context = kiro_report_context(true);
let mut rewriter = KiroToClaudeCliStreamState::new(&report_context);
let chunk = encode_event_frame(
"event",
Some("assistantResponseEvent"),
&json!({"content": "\n\n你好！有"}),
);
let first = rewriter
.push_chunk(&report_context, &chunk)
.expect("rewrite should succeed");
let rest = rewriter
.finish(&report_context)
.expect("finish should succeed");
let text = String::from_utf8([first, rest].concat()).expect("utf8 should decode");
assert!(text.contains("\"type\":\"text_delta\""));
assert!(text.contains("你好！有"));
}
#[test]
fn kiro_stream_rewriter_handles_multibyte_text_inside_thinking_block() {
let report_context = kiro_report_context(true);
let mut rewriter = KiroToClaudeCliStreamState::new(&report_context);
let chunk = encode_event_frame(
"event",
Some("assistantResponseEvent"),
&json!({"content": "<thinking>\n\n你好！有"}),
);
let first = rewriter
.push_chunk(&report_context, &chunk)
.expect("rewrite should succeed");
let rest = rewriter
.finish(&report_context)
.expect("finish should succeed");
let text = String::from_utf8([first, rest].concat()).expect("utf8 should decode");
assert!(text.contains("\"type\":\"thinking_delta\""));
assert!(text.contains("你好！有"));
}

0
crates/aether-ai-pipeline/src/adaptation/mod.rs → crates/aether-ai-surfaces/src/adaptation/mod.rs
View File

118
crates/aether-ai-pipeline/src/adaptation/private_envelope.rs → crates/aether-ai-surfaces/src/adaptation/private_envelope.rs
View File

@@ -2,6 +2,9 @@ use std::collections::BTreeMap;
use serde_json::Value;
use crate::adaptation::kiro_stream::KiroToClaudeCliStreamState;
use crate::finalize::AiSurfaceFinalizeError;
use super::surfaces::{
provider_adaptation_allows_sync_finalize_envelope, provider_adaptation_descriptor_for_envelope,
provider_adaptation_should_unwrap_stream_envelope, ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME,
@@ -153,6 +156,95 @@ pub fn transform_provider_private_stream_line(
Ok(out)
}
enum ProviderPrivateStreamNormalizeMode {
EnvelopeUnwrap,
KiroToClaudeCli(Box<KiroToClaudeCliStreamState>),
}
pub struct ProviderPrivateStreamNormalizer<'a> {
report_context: &'a Value,
buffered: Vec<u8>,
mode: ProviderPrivateStreamNormalizeMode,
}
pub fn maybe_build_provider_private_stream_normalizer<'a>(
report_context: Option<&'a Value>,
) -> Option<ProviderPrivateStreamNormalizer<'a>> {
let report_context = report_context?;
if !report_context
.get("has_envelope")
.and_then(Value::as_bool)
.unwrap_or(false)
{
return None;
}
let envelope_name = report_context
.get("envelope_name")
.and_then(Value::as_str)
.unwrap_or_default();
let provider_api_format = report_context
.get("provider_api_format")
.and_then(Value::as_str)
.unwrap_or_default();
let descriptor =
provider_adaptation_descriptor_for_envelope(envelope_name, provider_api_format)?;
let mode = if descriptor
.envelope_name
.eq_ignore_ascii_case(KIRO_ENVELOPE_NAME)
{
ProviderPrivateStreamNormalizeMode::KiroToClaudeCli(Box::new(
KiroToClaudeCliStreamState::new(report_context),
))
} else if descriptor.unwraps_response_envelope {
ProviderPrivateStreamNormalizeMode::EnvelopeUnwrap
} else {
return None;
};
Some(ProviderPrivateStreamNormalizer {
report_context,
buffered: Vec::new(),
mode,
})
}
impl ProviderPrivateStreamNormalizer<'_> {
pub fn push_chunk(&mut self, chunk: &[u8]) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
match &mut self.mode {
ProviderPrivateStreamNormalizeMode::KiroToClaudeCli(state) => {
state.push_chunk(self.report_context, chunk)
}
ProviderPrivateStreamNormalizeMode::EnvelopeUnwrap => {
self.buffered.extend_from_slice(chunk);
let mut output = Vec::new();
while let Some(line_end) = self.buffered.iter().position(|byte| *byte == b'\n') {
let line = self.buffered.drain(..=line_end).collect::<Vec<_>>();
output.extend(
transform_provider_private_stream_line(self.report_context, line)
.map_err(AiSurfaceFinalizeError::from)?,
);
}
Ok(output)
}
}
}
pub fn finish(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
match &mut self.mode {
ProviderPrivateStreamNormalizeMode::KiroToClaudeCli(state) => {
state.finish(self.report_context)
}
ProviderPrivateStreamNormalizeMode::EnvelopeUnwrap => {
if self.buffered.is_empty() {
return Ok(Vec::new());
}
let line = std::mem::take(&mut self.buffered);
transform_provider_private_stream_line(self.report_context, line)
.map_err(AiSurfaceFinalizeError::from)
}
}
}
}
pub fn stream_body_contains_error_event(body: &[u8]) -> bool {
let Ok(text) = std::str::from_utf8(body) else {
return false;
@@ -335,8 +427,9 @@ mod tests {
use serde_json::json;
use super::{
normalize_provider_private_report_context, normalize_provider_private_response_value,
stream_body_contains_error_event, transform_provider_private_stream_line,
maybe_build_provider_private_stream_normalizer, normalize_provider_private_report_context,
normalize_provider_private_response_value, stream_body_contains_error_event,
transform_provider_private_stream_line,
};
#[test]
@@ -404,6 +497,27 @@ mod tests {
assert!(output_text.contains("\"id\":\"call_get_weather_0\""));
}
#[test]
fn private_stream_normalizer_unwraps_antigravity_stream() {
let report_context = json!({
"has_envelope": true,
"provider_api_format": "gemini:generate_content",
"client_api_format": "gemini:generate_content",
"envelope_name": "antigravity:v1internal",
"mapped_model": "claude-sonnet-4-5",
});
let mut normalizer = maybe_build_provider_private_stream_normalizer(Some(&report_context))
.expect("normalizer should exist");
let output = normalizer
.push_chunk(
b"data: {\"response\":{\"candidates\":[{\"content\":{\"parts\":[{\"functionCall\":{\"name\":\"get_weather\",\"args\":{\"city\":\"SF\"}}}],\"role\":\"model\"},\"index\":0}],\"modelVersion\":\"claude-sonnet-4-5\"},\"responseId\":\"resp_123\"}\n\n",
)
.expect("unwrap should succeed");
let output_text = String::from_utf8(output).expect("text should decode");
assert!(output_text.contains("\"_v1internal_response_id\":\"resp_123\""));
assert!(output_text.contains("\"id\":\"call_get_weather_0\""));
}
#[test]
fn detects_sse_error_events_without_explicit_type_field() {
let body = br#"event: error

0
crates/aether-ai-pipeline/src/adaptation/surfaces.rs → crates/aether-ai-surfaces/src/adaptation/surfaces.rs
View File

123
crates/aether-ai-pipeline/src/api.rs → crates/aether-ai-surfaces/src/api.rs
View File

@@ -3,12 +3,13 @@ pub use crate::adaptation::kiro_stream::{
build_kiro_stream_error_sse_events, calculate_kiro_context_input_tokens,
encode_kiro_sse_events, estimate_kiro_tokens, find_kiro_real_thinking_end_tag,
find_kiro_real_thinking_end_tag_at_buffer_end, find_kiro_real_thinking_start_tag, kiro_crc32,
KIRO_MAX_THINKING_BUFFER,
KiroToClaudeCliStreamState, KIRO_MAX_THINKING_BUFFER,
};
pub use crate::adaptation::private_envelope::{
normalize_provider_private_report_context, normalize_provider_private_response_value,
provider_private_response_allows_sync_finalize, stream_body_contains_error_event,
transform_provider_private_stream_line,
maybe_build_provider_private_stream_normalizer, normalize_provider_private_report_context,
normalize_provider_private_response_value, provider_private_response_allows_sync_finalize,
stream_body_contains_error_event, transform_provider_private_stream_line,
ProviderPrivateStreamNormalizer,
};
pub use crate::adaptation::surfaces::{
provider_adaptation_allows_sync_finalize_envelope, provider_adaptation_anchor_api_format,
@@ -18,16 +19,14 @@ pub use crate::adaptation::surfaces::{
ProviderAdaptationSurface, ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME,
GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME, KIRO_ENVELOPE_NAME,
};
pub use crate::contracts::augment_sync_report_context;
pub use crate::contracts::{
core_error_background_report_kind, core_error_default_client_api_format,
core_success_background_report_kind, generic_decision_missing_exact_provider_request,
implicit_sync_finalize_report_kind, is_openai_responses_stream_plan_kind,
is_openai_responses_sync_plan_kind, ExecutionRuntimeAuthContext, GatewayControlPlanRequest,
GatewayControlPlanResponse, GatewayControlSyncDecisionResponse, LocalStreamPlanAndReport,
LocalSyncPlanAndReport, CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_STREAM_SUCCESS_REPORT_KIND,
CLAUDE_CHAT_SYNC_ERROR_REPORT_KIND, CLAUDE_CHAT_SYNC_FINALIZE_REPORT_KIND,
CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CHAT_SYNC_SUCCESS_REPORT_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
core_success_background_report_kind, implicit_sync_finalize_report_kind,
is_openai_responses_stream_plan_kind, is_openai_responses_sync_plan_kind,
ExecutionRuntimeAuthContext, GatewayControlPlanRequest, CLAUDE_CHAT_STREAM_PLAN_KIND,
CLAUDE_CHAT_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CHAT_SYNC_ERROR_REPORT_KIND,
CLAUDE_CHAT_SYNC_FINALIZE_REPORT_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND,
CLAUDE_CHAT_SYNC_SUCCESS_REPORT_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
CLAUDE_CLI_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CLI_SYNC_ERROR_REPORT_KIND,
CLAUDE_CLI_SYNC_FINALIZE_REPORT_KIND, CLAUDE_CLI_SYNC_PLAN_KIND,
CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND, EXECUTION_RUNTIME_STREAM_ACTION,
@@ -59,50 +58,18 @@ pub use crate::contracts::{
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
};
pub use crate::conversion::request::{
convert_openai_chat_request_to_claude_request, convert_openai_chat_request_to_gemini_request,
convert_openai_chat_request_to_openai_responses_request, extract_openai_text_content,
normalize_claude_request_to_openai_chat_request,
normalize_gemini_request_to_openai_chat_request,
normalize_openai_responses_request_to_openai_chat_request, parse_openai_tool_result_content,
};
pub use crate::conversion::response::{
build_openai_responses_response, build_openai_responses_response_with_content,
build_openai_responses_response_with_reasoning, convert_claude_chat_response_to_openai_chat,
convert_claude_response_to_openai_responses, convert_gemini_chat_response_to_openai_chat,
convert_gemini_response_to_openai_responses, convert_openai_chat_response_to_claude_chat,
convert_openai_chat_response_to_gemini_chat, convert_openai_chat_response_to_openai_responses,
convert_openai_responses_response_to_openai_chat, OpenAiResponsesResponseUsage,
};
pub use crate::conversion::{
build_core_error_body_for_client_format, canonical_request_unknown_block_count,
canonical_response_unknown_block_count, canonical_to_claude_request,
canonical_to_claude_response, canonical_to_gemini_request, canonical_to_gemini_response,
canonical_to_openai_chat_request, canonical_to_openai_chat_response,
canonical_to_openai_responses_compact_request, canonical_to_openai_responses_compact_response,
canonical_to_openai_responses_request, canonical_to_openai_responses_response,
canonical_unknown_block_count, convert_request, convert_response,
from_claude_to_canonical_request, from_claude_to_canonical_response,
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,
is_core_error_finalize_kind, request_candidate_api_format_preference,
request_candidate_api_formats, request_conversion_direct_auth,
request_conversion_enabled_for_transport, request_conversion_kind,
request_conversion_requires_enable_flag, request_conversion_transport_supported,
request_conversion_transport_unsupported_reason, request_pair_allowed_for_transport,
sync_chat_response_conversion_kind, sync_cli_response_conversion_kind, CanonicalContentBlock,
CanonicalGenerationConfig, CanonicalInstruction, CanonicalMessage, CanonicalRequest,
CanonicalResponse, CanonicalResponseFormat, CanonicalResponseOutput, CanonicalRole,
CanonicalStopReason, CanonicalThinkingConfig, CanonicalToolChoice, CanonicalToolDefinition,
CanonicalUsage, FormatContext, FormatError, FormatFamily, FormatId, FormatProfile,
LocalCoreSyncErrorKind, RequestConversionKind, SyncChatResponseConversionKind,
SyncCliResponseConversionKind,
};
pub use crate::finalize::common::{
build_generated_tool_call_id, build_local_success_background_report,
build_local_success_conversion_background_report, canonicalize_tool_arguments,
prepare_local_success_response_parts, prepare_local_success_response_parts_owned,
LocalSyncReportParts,
};
pub use crate::finalize::error_body::{
build_core_error_body_for_client_format, is_core_error_finalize_kind, LocalCoreSyncErrorKind,
};
pub use crate::finalize::openai_image_stream::{
maybe_build_openai_image_sync_finalize_product, OpenAiImageStreamState,
OpenAiImageSyncFinalizeProduct,
};
pub use crate::finalize::sse::{encode_done_sse, encode_json_sse, map_claude_stop_reason};
pub use crate::finalize::standard::claude::stream::{ClaudeClientEmitter, ClaudeProviderState};
@@ -129,8 +96,12 @@ pub use crate::finalize::sync_products::{
maybe_build_standard_sync_finalize_product_from_normalized_payload,
StandardCrossFormatSyncProduct, StandardSyncFinalizeNormalizedProduct,
};
pub use crate::finalize::sync_to_stream::{
maybe_bridge_standard_sync_json_to_stream, SyncToStreamBridgeOutcome,
};
pub use crate::finalize::{
resolve_finalize_stream_rewrite_mode, FinalizeStreamRewriteMode, PipelineFinalizeError,
maybe_build_ai_surface_stream_rewriter, resolve_finalize_stream_rewrite_mode,
AiSurfaceFinalizeError, AiSurfaceStreamRewriter, FinalizeStreamRewriteMode,
};
pub use crate::planner::common::{
force_upstream_streaming_for_provider, parse_direct_request_body,
@@ -147,9 +118,9 @@ pub use crate::planner::passthrough::provider::{
LocalSameFormatProviderSpec,
};
pub use crate::planner::route::{
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
resolve_execution_runtime_sync_plan_kind, supports_stream_scheduler_decision_kind,
supports_sync_scheduler_decision_kind,
is_matching_stream_http_request, is_matching_stream_request,
resolve_execution_runtime_stream_plan_kind, resolve_execution_runtime_sync_plan_kind,
supports_stream_scheduler_decision_kind, supports_sync_scheduler_decision_kind,
};
pub use crate::planner::specialized::{
files::{
@@ -157,8 +128,12 @@ pub use crate::planner::specialized::{
resolve_sync_spec as resolve_gemini_files_sync_spec, LocalGeminiFilesSpec,
},
image::{
build_openai_image_provider_request_body, default_model_for_openai_image_operation,
is_openai_image_stream_request, normalize_openai_image_request,
openai_image_operation_from_path, resolve_requested_openai_image_model_for_request,
resolve_stream_spec as resolve_local_image_stream_spec,
resolve_sync_spec as resolve_local_image_sync_spec, LocalOpenAiImageSpec,
NormalizedOpenAiImageRequest, OpenAiImageOperation, OpenAiImageResponseFormat,
},
video::{
resolve_sync_spec as resolve_local_video_sync_spec, LocalVideoCreateFamily,
@@ -170,7 +145,6 @@ pub use crate::planner::standard::{
apply_openai_responses_compact_special_body_edits, build_cross_format_openai_chat_request_body,
build_cross_format_openai_responses_request_body, build_local_openai_chat_request_body,
build_local_openai_responses_request_body, build_standard_request_body,
build_standard_upstream_url,
claude::{
resolve_stream_spec as resolve_claude_stream_spec,
resolve_sync_spec as resolve_claude_sync_spec,
@@ -189,7 +163,42 @@ pub use crate::planner::standard::{
CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL, CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT,
CODEX_OPENAI_IMAGE_INTERNAL_MODEL,
};
pub use aether_ai_formats::conversion::request::{
convert_openai_chat_request_to_claude_request, convert_openai_chat_request_to_gemini_request,
convert_openai_chat_request_to_openai_responses_request, extract_openai_text_content,
normalize_claude_request_to_openai_chat_request,
normalize_gemini_request_to_openai_chat_request,
normalize_openai_responses_request_to_openai_chat_request, parse_openai_tool_result_content,
};
pub use aether_ai_formats::conversion::response::{
build_openai_responses_response, build_openai_responses_response_with_content,
build_openai_responses_response_with_reasoning, convert_claude_chat_response_to_openai_chat,
convert_claude_response_to_openai_responses, convert_gemini_chat_response_to_openai_chat,
convert_gemini_response_to_openai_responses, convert_openai_chat_response_to_claude_chat,
convert_openai_chat_response_to_gemini_chat, convert_openai_chat_response_to_openai_responses,
convert_openai_responses_response_to_openai_chat, OpenAiResponsesResponseUsage,
};
pub use aether_ai_formats::{
api_format_alias_matches, api_format_storage_aliases, is_openai_responses_compact_format,
is_openai_responses_family_format, is_openai_responses_format, normalize_api_format_alias,
};
pub use aether_ai_formats::{
canonical_request_unknown_block_count, canonical_response_unknown_block_count,
canonical_to_claude_request, canonical_to_claude_response, canonical_to_gemini_request,
canonical_to_gemini_response, canonical_to_openai_chat_request,
canonical_to_openai_chat_response, canonical_to_openai_responses_compact_request,
canonical_to_openai_responses_compact_response, canonical_to_openai_responses_request,
canonical_to_openai_responses_response, canonical_unknown_block_count, convert_request,
convert_response, from_claude_to_canonical_request, from_claude_to_canonical_response,
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,
request_candidate_api_format_preference, request_candidate_api_formats,
request_conversion_kind, request_conversion_requires_enable_flag,
sync_chat_response_conversion_kind, sync_cli_response_conversion_kind, CanonicalContentBlock,
CanonicalGenerationConfig, CanonicalInstruction, CanonicalMessage, CanonicalRequest,
CanonicalResponse, CanonicalResponseFormat, CanonicalResponseOutput, CanonicalRole,
CanonicalStopReason, CanonicalThinkingConfig, CanonicalToolChoice, CanonicalToolDefinition,
CanonicalUsage, FormatContext, FormatError, FormatFamily, FormatId, FormatProfile,
RequestConversionKind, SyncChatResponseConversionKind, SyncCliResponseConversionKind,
};

0
crates/aether-ai-pipeline/src/contracts/actions.rs → crates/aether-ai-surfaces/src/contracts/actions.rs
View File

0
crates/aether-ai-pipeline/src/contracts/auth_context.rs → crates/aether-ai-surfaces/src/contracts/auth_context.rs
View File

88
crates/aether-ai-surfaces/src/contracts/control_payloads.rs Normal file
View File

@@ -0,0 +1,88 @@
use std::collections::BTreeMap;
use serde::Serialize;
use crate::contracts::ExecutionRuntimeAuthContext;
#[derive(Debug, Serialize)]
pub struct GatewayControlPlanRequest {
pub trace_id: String,
pub method: String,
pub path: String,
pub query_string: Option<String>,
pub headers: BTreeMap<String, String>,
pub body_json: serde_json::Value,
#[serde(skip_serializing_if = "Option::is_none")]
pub body_base64: Option<String>,
pub auth_context: Option<ExecutionRuntimeAuthContext>,
}
#[allow(clippy::too_many_arguments)]
pub fn build_gateway_control_plan_request(
trace_id: &str,
method: &str,
path: &str,
query_string: Option<&str>,
headers: BTreeMap<String, String>,
body_json: serde_json::Value,
body_base64: Option<String>,
auth_context: Option<ExecutionRuntimeAuthContext>,
) -> GatewayControlPlanRequest {
GatewayControlPlanRequest {
trace_id: trace_id.to_string(),
method: method.to_string(),
path: path.to_string(),
query_string: query_string.map(ToOwned::to_owned),
headers,
body_json,
body_base64,
auth_context,
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use super::{build_gateway_control_plan_request, ExecutionRuntimeAuthContext};
#[test]
fn build_gateway_control_plan_request_preserves_request_shape() {
let payload = build_gateway_control_plan_request(
"trace-123",
"POST",
"/v1/chat/completions",
Some("stream=true"),
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
serde_json::json!({"model": "gpt-5"}),
Some("eyJmb28iOiJiYXIifQ==".to_string()),
Some(ExecutionRuntimeAuthContext {
user_id: "user-1".to_string(),
api_key_id: "key-1".to_string(),
username: None,
api_key_name: None,
balance_remaining: Some(12.5),
access_allowed: true,
api_key_is_standalone: false,
}),
);
assert_eq!(payload.trace_id, "trace-123");
assert_eq!(payload.method, "POST");
assert_eq!(payload.path, "/v1/chat/completions");
assert_eq!(payload.query_string.as_deref(), Some("stream=true"));
assert_eq!(
payload.headers.get("content-type").map(String::as_str),
Some("application/json")
);
assert_eq!(payload.body_json, serde_json::json!({"model": "gpt-5"}));
assert_eq!(payload.body_base64.as_deref(), Some("eyJmb28iOiJiYXIifQ=="));
assert_eq!(
payload
.auth_context
.as_ref()
.map(|ctx| ctx.user_id.as_str()),
Some("user-1")
);
}
}

7
crates/aether-ai-pipeline/src/contracts/mod.rs → crates/aether-ai-surfaces/src/contracts/mod.rs
View File

@@ -9,12 +9,7 @@ pub use actions::{
EXECUTION_RUNTIME_SYNC_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
};
pub use auth_context::ExecutionRuntimeAuthContext;
pub use control_payloads::{
augment_sync_report_context, build_gateway_control_plan_request,
generic_decision_missing_exact_provider_request, GatewayControlPlanRequest,
GatewayControlPlanResponse, GatewayControlSyncDecisionResponse, LocalStreamPlanAndReport,
LocalSyncPlanAndReport,
};
pub use control_payloads::{build_gateway_control_plan_request, GatewayControlPlanRequest};
pub use plan_kinds::{
is_openai_responses_stream_plan_kind, is_openai_responses_sync_plan_kind,
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,

0
crates/aether-ai-pipeline/src/contracts/plan_kinds.rs → crates/aether-ai-surfaces/src/contracts/plan_kinds.rs
View File

0
crates/aether-ai-pipeline/src/contracts/report_kinds.rs → crates/aether-ai-surfaces/src/contracts/report_kinds.rs
View File

40
crates/aether-ai-pipeline/src/finalize/common.rs → crates/aether-ai-surfaces/src/finalize/common.rs
View File

@@ -1,10 +1,21 @@
use std::collections::BTreeMap;
use aether_usage_runtime::GatewaySyncReportRequest;
use serde_json::Value;
use crate::contracts::core_success_background_report_kind;
#[derive(Debug, Clone, PartialEq)]
pub struct LocalSyncReportParts {
pub trace_id: String,
pub report_kind: String,
pub report_context: Option<Value>,
pub status_code: u16,
pub headers: BTreeMap<String, String>,
pub body_json: Option<Value>,
pub client_body_json: Option<Value>,
pub body_base64: Option<String>,
}
pub fn build_generated_tool_call_id(index: usize) -> String {
format!("call_auto_{index}")
}
@@ -36,7 +47,7 @@ pub fn prepare_local_success_response_parts_owned(
Ok((body_bytes, headers))
}
fn should_capture_client_sync_success_body(payload: &GatewaySyncReportRequest) -> bool {
fn should_capture_client_sync_success_body(payload: &LocalSyncReportParts) -> bool {
payload
.report_context
.as_ref()
@@ -47,10 +58,10 @@ fn should_capture_client_sync_success_body(payload: &GatewaySyncReportRequest) -
}
pub fn build_local_success_background_report(
payload: &GatewaySyncReportRequest,
payload: &LocalSyncReportParts,
body_json: Value,
headers: BTreeMap<String, String>,
) -> Option<GatewaySyncReportRequest> {
) -> Option<LocalSyncReportParts> {
let report_kind = core_success_background_report_kind(payload.report_kind.as_str())?;
let upstream_is_stream = should_capture_client_sync_success_body(payload);
let client_body_json = upstream_is_stream.then(|| body_json.clone());
@@ -65,7 +76,7 @@ pub fn build_local_success_background_report(
None
};
Some(GatewaySyncReportRequest {
Some(LocalSyncReportParts {
trace_id: payload.trace_id.clone(),
report_kind: report_kind.to_string(),
report_context: payload.report_context.clone(),
@@ -74,18 +85,17 @@ pub fn build_local_success_background_report(
body_json: provider_body_json,
client_body_json,
body_base64: provider_body_base64,
telemetry: payload.telemetry.clone(),
})
}
pub fn build_local_success_conversion_background_report(
payload: &GatewaySyncReportRequest,
payload: &LocalSyncReportParts,
client_body_json: Value,
provider_body_json: Value,
) -> Option<GatewaySyncReportRequest> {
) -> Option<LocalSyncReportParts> {
let report_kind = core_success_background_report_kind(payload.report_kind.as_str())?;
Some(GatewaySyncReportRequest {
Some(LocalSyncReportParts {
trace_id: payload.trace_id.clone(),
report_kind: report_kind.to_string(),
report_context: payload.report_context.clone(),
@@ -94,7 +104,6 @@ pub fn build_local_success_conversion_background_report(
body_json: Some(provider_body_json),
client_body_json: Some(client_body_json),
body_base64: None,
telemetry: payload.telemetry.clone(),
})
}
@@ -107,8 +116,8 @@ mod tests {
build_generated_tool_call_id, build_local_success_background_report,
build_local_success_conversion_background_report, canonicalize_tool_arguments,
prepare_local_success_response_parts, prepare_local_success_response_parts_owned,
LocalSyncReportParts,
};
use aether_usage_runtime::GatewaySyncReportRequest;
use std::collections::BTreeMap;
#[test]
@@ -189,7 +198,7 @@ mod tests {
#[test]
fn build_local_success_background_report_maps_finalize_kind() {
let payload = GatewaySyncReportRequest {
let payload = LocalSyncReportParts {
trace_id: "trace-1".to_string(),
report_kind: "openai_chat_sync_finalize".to_string(),
report_context: Some(serde_json::json!({"request_id": "req-1"})),
@@ -198,7 +207,6 @@ mod tests {
body_json: None,
client_body_json: None,
body_base64: None,
telemetry: None,
};
let report = build_local_success_background_report(
@@ -215,7 +223,7 @@ mod tests {
#[test]
fn build_local_success_background_report_preserves_provider_stream_for_upstream_stream_sync() {
let payload = GatewaySyncReportRequest {
let payload = LocalSyncReportParts {
trace_id: "trace-1b".to_string(),
report_kind: "openai_chat_sync_finalize".to_string(),
report_context: Some(serde_json::json!({
@@ -236,7 +244,6 @@ mod tests {
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1b\",\"object\":\"response\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n",
)
)),
telemetry: None,
};
let report = build_local_success_background_report(
@@ -256,7 +263,7 @@ mod tests {
#[test]
fn build_local_success_conversion_background_report_maps_provider_body() {
let payload = GatewaySyncReportRequest {
let payload = LocalSyncReportParts {
trace_id: "trace-2".to_string(),
report_kind: "openai_chat_sync_finalize".to_string(),
report_context: Some(serde_json::json!({"request_id": "req-2"})),
@@ -265,7 +272,6 @@ mod tests {
body_json: None,
client_body_json: None,
body_base64: None,
telemetry: None,
};
let report = build_local_success_conversion_background_report(

0
crates/aether-ai-pipeline/src/conversion/error.rs → crates/aether-ai-surfaces/src/finalize/error_body.rs
View File

22
crates/aether-ai-pipeline/src/finalize/mod.rs → crates/aether-ai-surfaces/src/finalize/mod.rs
View File

@@ -3,38 +3,44 @@ use std::fmt;
pub use self::sse::{encode_done_sse, encode_json_sse, map_claude_stop_reason};
pub use self::standard::stream_core::CanonicalStreamEvent;
pub use self::standard::stream_core::CanonicalStreamFrame;
pub use self::stream_rewrite::{resolve_finalize_stream_rewrite_mode, FinalizeStreamRewriteMode};
pub use self::stream_rewrite::{
maybe_build_ai_surface_stream_rewriter, resolve_finalize_stream_rewrite_mode,
AiSurfaceStreamRewriter, FinalizeStreamRewriteMode,
};
pub mod common;
pub mod error_body;
pub mod openai_image_stream;
pub mod sse;
pub mod standard;
pub mod stream_rewrite;
pub mod sync_products;
pub mod sync_to_stream;
#[derive(Debug)]
pub struct PipelineFinalizeError(pub String);
pub struct AiSurfaceFinalizeError(pub String);
impl PipelineFinalizeError {
impl AiSurfaceFinalizeError {
pub fn new(message: impl Into<String>) -> Self {
Self(message.into())
}
}
impl fmt::Display for PipelineFinalizeError {
impl fmt::Display for AiSurfaceFinalizeError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "Pipeline finalize error: {}", self.0)
write!(f, "AI surface finalize error: {}", self.0)
}
}
impl std::error::Error for PipelineFinalizeError {}
impl std::error::Error for AiSurfaceFinalizeError {}
impl From<serde_json::Error> for PipelineFinalizeError {
impl From<serde_json::Error> for AiSurfaceFinalizeError {
fn from(source: serde_json::Error) -> Self {
Self(source.to_string())
}
}
impl From<base64::DecodeError> for PipelineFinalizeError {
impl From<base64::DecodeError> for AiSurfaceFinalizeError {
fn from(source: base64::DecodeError) -> Self {
Self(source.to_string())
}

714
crates/aether-ai-surfaces/src/finalize/openai_image_stream.rs Normal file
View File

@@ -0,0 +1,714 @@
use base64::Engine as _;
use serde_json::Value;
use crate::contracts::OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND;
use crate::finalize::sse::encode_json_sse;
use crate::finalize::AiSurfaceFinalizeError;
use crate::planner::standard::CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT;
#[derive(Default)]
pub struct OpenAiImageStreamState {
buffered: Vec<u8>,
latest_image: Option<OpenAiImageFrame>,
emitted_partial_count: u64,
saw_upstream_partial: bool,
emitted_failure: bool,
}
#[derive(Clone)]
struct OpenAiImageFrame {
b64_json: String,
}
impl OpenAiImageStreamState {
pub fn push_chunk(
&mut self,
report_context: &Value,
chunk: &[u8],
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
self.buffered.extend_from_slice(chunk);
let mut output = Vec::new();
while let Some(block_end) = find_sse_block_end(&self.buffered) {
let block = self.buffered.drain(..block_end).collect::<Vec<_>>();
output.extend(self.transform_block(report_context, &block)?);
drain_sse_separator(&mut self.buffered);
}
Ok(output)
}
pub fn finish(&mut self, report_context: &Value) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.buffered.is_empty() {
return Ok(Vec::new());
}
let block = std::mem::take(&mut self.buffered);
self.transform_block(report_context, &block)
}
fn transform_block(
&mut self,
report_context: &Value,
block: &[u8],
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let text = std::str::from_utf8(block)
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?;
let mut event_name = None::<String>;
let mut data_lines = Vec::new();
for raw_line in text.lines() {
let line = raw_line.trim_end_matches('\r');
if let Some(value) = line.strip_prefix("event:") {
event_name = Some(value.trim().to_string());
} else if let Some(value) = line.strip_prefix("data:") {
data_lines.push(value.trim().to_string());
}
}
let data = data_lines.join("\n");
if data.is_empty() || data == "[DONE]" {
return Ok(Vec::new());
}
let event: Value = serde_json::from_str(&data)?;
let event_type = event
.get("type")
.and_then(Value::as_str)
.or(event_name.as_deref())
.unwrap_or_default();
match event_type {
"error" | "response.failed" => self.handle_failed(report_context, &event),
"response.image_generation_call.partial_image" => {
self.handle_image_generation_partial(report_context, &event)
}
"response.output_item.done" => self.handle_output_item_done(report_context, &event),
"response.completed" => self.handle_completed(report_context, &event),
_ => Ok(Vec::new()),
}
}
fn handle_image_generation_partial(
&mut self,
report_context: &Value,
event: &Value,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.emitted_failure {
return Ok(Vec::new());
}
if requested_partial_images(report_context) == 0 {
return Ok(Vec::new());
}
let Some(result) = event
.get("partial_image_b64")
.or_else(|| event.get("b64_json"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Ok(Vec::new());
};
let partial_image_index = event
.get("partial_image_index")
.or_else(|| event.get("output_index"))
.and_then(Value::as_u64)
.unwrap_or(self.emitted_partial_count);
self.emitted_partial_count = self
.emitted_partial_count
.max(partial_image_index.saturating_add(1));
self.saw_upstream_partial = true;
self.latest_image = Some(OpenAiImageFrame {
b64_json: result.to_string(),
});
encode_json_sse(
Some(image_partial_event_name(report_context)),
&serde_json::json!({
"type": image_partial_event_name(report_context),
"b64_json": result,
"partial_image_index": partial_image_index,
}),
)
}
fn handle_output_item_done(
&mut self,
report_context: &Value,
event: &Value,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.emitted_failure {
return Ok(Vec::new());
}
let Some(item) = event.get("item").and_then(Value::as_object) else {
return Ok(Vec::new());
};
if item.get("type").and_then(Value::as_str) != Some("image_generation_call") {
return Ok(Vec::new());
}
let Some(result) = item.get("result").and_then(Value::as_str).map(str::trim) else {
return Ok(Vec::new());
};
if result.is_empty() {
return Ok(Vec::new());
}
self.latest_image = Some(OpenAiImageFrame {
b64_json: result.to_string(),
});
if requested_partial_images(report_context) == 0 || self.saw_upstream_partial {
return Ok(Vec::new());
}
let partial_image_index = event
.get("output_index")
.and_then(Value::as_u64)
.unwrap_or(self.emitted_partial_count);
self.emitted_partial_count = partial_image_index.saturating_add(1);
encode_json_sse(
Some(image_partial_event_name(report_context)),
&serde_json::json!({
"type": image_partial_event_name(report_context),
"b64_json": result,
"partial_image_index": partial_image_index,
}),
)
}
fn handle_completed(
&mut self,
report_context: &Value,
event: &Value,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.emitted_failure {
return Ok(Vec::new());
}
if self.latest_image.is_none() {
if let Some(result) = completed_response_image_result(event) {
self.latest_image = Some(OpenAiImageFrame {
b64_json: result.to_string(),
});
}
}
let Some(latest_image) = self.latest_image.clone() else {
return Ok(Vec::new());
};
let usage = event
.get("response")
.and_then(Value::as_object)
.and_then(|response| {
response
.get("tool_usage")
.and_then(|value| value.get("image_gen"))
.cloned()
.or_else(|| response.get("usage").cloned())
})
.unwrap_or(Value::Null);
encode_json_sse(
Some(image_completed_event_name(report_context)),
&serde_json::json!({
"type": image_completed_event_name(report_context),
"b64_json": latest_image.b64_json,
"usage": usage,
}),
)
}
fn handle_failed(
&mut self,
report_context: &Value,
event: &Value,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.emitted_failure {
return Ok(Vec::new());
}
self.emitted_failure = true;
let error = image_failure_error(event);
encode_json_sse(
Some(image_failed_event_name(report_context)),
&serde_json::json!({
"type": image_failed_event_name(report_context),
"error": error,
}),
)
}
}
fn image_failure_error(event: &Value) -> Value {
let mut error = event
.get("error")
.or_else(|| event.get("response").and_then(|value| value.get("error")))
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
if !error.contains_key("message") {
if let Some(message) = event
.get("message")
.and_then(Value::as_str)
.or_else(|| {
event
.get("response")
.and_then(|value| value.get("error"))
.and_then(|value| value.get("message"))
.and_then(Value::as_str)
})
.map(str::trim)
.filter(|value| !value.is_empty())
{
error.insert("message".to_string(), Value::String(message.to_string()));
}
}
if !error.contains_key("code") {
if let Some(code) = event
.get("code")
.or_else(|| {
event
.get("response")
.and_then(|value| value.get("error"))
.and_then(|value| value.get("code"))
})
.cloned()
{
error.insert("code".to_string(), code);
}
}
if !error.contains_key("type") {
let inferred_type = error
.get("code")
.and_then(Value::as_str)
.filter(|value| !value.is_empty())
.unwrap_or("upstream_error");
error.insert("type".to_string(), Value::String(inferred_type.to_string()));
}
if !error.contains_key("message") {
error.insert(
"message".to_string(),
Value::String("Image generation failed".to_string()),
);
}
Value::Object(error)
}
fn completed_response_image_result(event: &Value) -> Option<&str> {
event
.get("response")
.and_then(|value| value.get("output"))
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter(|item| item.get("type").and_then(Value::as_str) == Some("image_generation_call"))
.filter_map(|item| item.get("result").and_then(Value::as_str))
.map(str::trim)
.find(|value| !value.is_empty())
}
fn requested_partial_images(report_context: &Value) -> u64 {
report_context
.get("image_request")
.and_then(|value| value.get("partial_images"))
.and_then(Value::as_u64)
.unwrap_or(0)
}
fn image_partial_event_name(report_context: &Value) -> &'static str {
if image_request_operation(report_context) == Some("edit") {
"image_edit.partial_image"
} else {
"image_generation.partial_image"
}
}
fn image_completed_event_name(report_context: &Value) -> &'static str {
if image_request_operation(report_context) == Some("edit") {
"image_edit.completed"
} else {
"image_generation.completed"
}
}
fn image_failed_event_name(report_context: &Value) -> &'static str {
if image_request_operation(report_context) == Some("edit") {
"image_edit.failed"
} else {
"image_generation.failed"
}
}
fn image_request_operation(report_context: &Value) -> Option<&str> {
report_context
.get("image_request")
.and_then(|value| value.get("operation"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn find_sse_block_end(buffer: &[u8]) -> Option<usize> {
buffer
.windows(2)
.position(|window| window == b"\n\n")
.map(|index| index + 2)
.or_else(|| {
buffer
.windows(4)
.position(|window| window == b"\r\n\r\n")
.map(|index| index + 4)
})
}
fn drain_sse_separator(buffer: &mut Vec<u8>) {
while matches!(buffer.first(), Some(b'\n' | b'\r')) {
buffer.remove(0);
}
}
pub struct OpenAiImageSyncFinalizeProduct {
pub client_body_json: Value,
pub provider_body_json: Value,
}
pub fn maybe_build_openai_image_sync_finalize_product(
report_kind: &str,
status_code: u16,
report_context: Option<&Value>,
body_base64: Option<&str>,
) -> Result<Option<OpenAiImageSyncFinalizeProduct>, AiSurfaceFinalizeError> {
if report_kind != OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND || status_code >= 400 {
return Ok(None);
}
let Some(report_context) = report_context else {
return Ok(None);
};
if report_context
.get("client_api_format")
.and_then(Value::as_str)
.map(str::trim)
!= Some("openai:image")
{
return Ok(None);
}
let Some(body_base64) = body_base64 else {
return Ok(None);
};
let default_output_format = report_context
.get("image_request")
.and_then(|value| value.get("output_format"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT);
let body_bytes = base64::engine::general_purpose::STANDARD.decode(body_base64)?;
let text = std::str::from_utf8(&body_bytes)
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?;
let mut created = None;
let mut completed_response = None;
let mut images = Vec::new();
for raw_block in text.split("\n\n") {
let block = raw_block.trim();
if block.is_empty() {
continue;
}
let data_line = block
.lines()
.find_map(|line| line.trim().strip_prefix("data:").map(str::trim));
let Some(data_line) = data_line else {
continue;
};
if data_line.is_empty() || data_line == "[DONE]" {
continue;
}
let event: Value = serde_json::from_str(data_line)?;
match event
.get("type")
.and_then(Value::as_str)
.unwrap_or_default()
{
"response.created" => {
created = event
.get("response")
.and_then(|value| value.get("created_at"))
.and_then(Value::as_i64)
.or(created);
}
"response.output_item.done" => {
let Some(item) = event.get("item").and_then(Value::as_object) else {
continue;
};
if item.get("type").and_then(Value::as_str) != Some("image_generation_call") {
continue;
}
let Some(result) = item.get("result").and_then(Value::as_str) else {
continue;
};
images.push(serde_json::json!({
"b64_json": result,
"output_format": item.get("output_format").cloned().unwrap_or(Value::String(default_output_format.to_string())),
"revised_prompt": item.get("revised_prompt").cloned().unwrap_or(Value::Null),
}));
}
"response.completed" => {
completed_response = event.get("response").and_then(Value::as_object).cloned();
}
_ => {}
}
}
if images.is_empty() {
return Ok(None);
}
let completed_response = completed_response.unwrap_or_default();
let provider_usage = completed_response
.get("tool_usage")
.and_then(|value| value.get("image_gen"))
.cloned()
.or_else(|| completed_response.get("usage").cloned());
let provider_body_json = serde_json::json!({
"id": completed_response.get("id").cloned().unwrap_or(Value::Null),
"object": "response",
"model": completed_response.get("model").cloned().unwrap_or(Value::Null),
"status": completed_response.get("status").cloned().unwrap_or(Value::String("completed".to_string())),
"usage": provider_usage,
"tool_usage": completed_response.get("tool_usage").cloned().unwrap_or(Value::Null),
"output": images
.iter()
.map(|image| serde_json::json!({
"type": "image_generation_call",
"output_format": image.get("output_format").cloned().unwrap_or(Value::Null),
"revised_prompt": image.get("revised_prompt").cloned().unwrap_or(Value::Null),
}))
.collect::<Vec<_>>(),
});
let client_images = images
.iter()
.map(|image| {
let revised_prompt = image.get("revised_prompt").cloned().unwrap_or(Value::Null);
let b64_json = image
.get("b64_json")
.and_then(Value::as_str)
.unwrap_or_default();
serde_json::json!({
"b64_json": b64_json,
"revised_prompt": revised_prompt,
})
})
.collect::<Vec<_>>();
let client_body_json = serde_json::json!({
"created": created.unwrap_or_default(),
"data": client_images,
"usage": provider_body_json.get("usage").cloned().unwrap_or(Value::Null),
});
Ok(Some(OpenAiImageSyncFinalizeProduct {
client_body_json,
provider_body_json,
}))
}
#[cfg(test)]
mod tests {
use base64::Engine as _;
use serde_json::json;
use super::{maybe_build_openai_image_sync_finalize_product, OpenAiImageStreamState};
fn utf8(bytes: Vec<u8>) -> String {
String::from_utf8(bytes).expect("utf8 should decode")
}
#[test]
fn emits_completed_event_for_generate() {
let report_context = json!({
"provider_api_format": "openai:image",
"client_api_format": "openai:image",
"needs_conversion": false,
"image_request": {
"operation": "generate"
}
});
let mut rewriter = OpenAiImageStreamState::default();
let first = rewriter
.push_chunk(
&report_context,
concat!(
"event: response.output_item.done\n",
"data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ig_123\",\"type\":\"image_generation_call\",\"result\":\"aGVsbG8=\"}}\n\n"
)
.as_bytes(),
)
.expect("rewrite should succeed");
assert!(first.is_empty());
let second = rewriter
.push_chunk(
&report_context,
concat!(
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"tool_usage\":{\"image_gen\":{\"input_tokens\":1,\"output_tokens\":2,\"total_tokens\":3}}}}\n\n"
)
.as_bytes(),
)
.expect("rewrite should succeed");
let output_text = utf8(second);
assert!(output_text.contains("event: image_generation.completed"));
assert!(output_text.contains("\"type\":\"image_generation.completed\""));
assert!(output_text.contains("\"b64_json\":\"aGVsbG8=\""));
assert!(output_text.contains("\"input_tokens\":1"));
assert!(!output_text.contains("data: [DONE]"));
assert!(rewriter
.finish(&report_context)
.expect("finish should succeed")
.is_empty());
}
#[test]
fn maps_responses_partial_image_events() {
let report_context = json!({
"provider_api_format": "openai:image",
"client_api_format": "openai:image",
"needs_conversion": false,
"image_request": {
"operation": "generate",
"partial_images": 1
}
});
let mut rewriter = OpenAiImageStreamState::default();
let partial = rewriter
.push_chunk(
&report_context,
concat!(
"event: response.image_generation_call.partial_image\n",
"data: {\"type\":\"response.image_generation_call.partial_image\",\"partial_image_index\":0,\"partial_image_b64\":\"cGFydGlhbA==\"}\n\n"
)
.as_bytes(),
)
.expect("rewrite should succeed");
let partial_text = utf8(partial);
assert!(partial_text.contains("event: image_generation.partial_image"));
assert!(partial_text.contains("\"type\":\"image_generation.partial_image\""));
assert!(partial_text.contains("\"b64_json\":\"cGFydGlhbA==\""));
assert!(partial_text.contains("\"partial_image_index\":0"));
assert!(!partial_text.contains("response.image_generation_call.partial_image"));
let done = rewriter
.push_chunk(
&report_context,
concat!(
"event: response.output_item.done\n",
"data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ig_123\",\"type\":\"image_generation_call\",\"result\":\"ZmluYWw=\"}}\n\n"
)
.as_bytes(),
)
.expect("rewrite should succeed");
assert!(done.is_empty());
let completed = rewriter
.push_chunk(
&report_context,
concat!(
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":4,\"output_tokens\":5,\"total_tokens\":9}}}\n\n"
)
.as_bytes(),
)
.expect("rewrite should succeed");
let completed_text = utf8(completed);
assert!(completed_text.contains("event: image_generation.completed"));
assert!(completed_text.contains("\"type\":\"image_generation.completed\""));
assert!(completed_text.contains("\"b64_json\":\"ZmluYWw=\""));
assert!(completed_text.contains("\"total_tokens\":9"));
}
#[test]
fn maps_upstream_error_to_generation_failed_once() {
let report_context = json!({
"provider_api_format": "openai:image",
"client_api_format": "openai:image",
"needs_conversion": false,
"image_request": {
"operation": "generate"
}
});
let mut rewriter = OpenAiImageStreamState::default();
let output = rewriter
.push_chunk(
&report_context,
concat!(
"event: error\n",
"data: {\"type\":\"error\",\"error\":{\"type\":\"input-images\",\"code\":\"rate_limit_exceeded\",\"message\":\"Rate limit reached for gpt-image-2\",\"param\":null}}\n\n",
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"Rate limit reached for gpt-image-2\"}}}\n\n"
)
.as_bytes(),
)
.expect("rewrite should succeed");
let output_text = utf8(output);
assert!(output_text.contains("event: image_generation.failed"));
assert_eq!(
output_text
.matches("event: image_generation.failed")
.count(),
1
);
assert!(output_text.contains("\"type\":\"image_generation.failed\""));
assert!(output_text.contains("\"type\":\"input-images\""));
assert!(output_text.contains("\"code\":\"rate_limit_exceeded\""));
assert!(output_text.contains("\"message\":\"Rate limit reached for gpt-image-2\""));
assert!(!output_text.contains("response.failed"));
assert!(rewriter
.finish(&report_context)
.expect("finish should succeed")
.is_empty());
}
#[test]
fn sync_finalize_product_maps_stream_response_to_client_and_provider_bodies() {
let report_context = json!({
"client_api_format": "openai:image",
"provider_api_format": "openai:image",
"image_request": {
"operation": "generate",
"output_format": "png"
}
});
let body_base64 = base64::engine::general_purpose::STANDARD.encode(
concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"created_at\":1776839946}}\n\n",
"event: response.output_item.done\n",
"data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"image_generation_call\",\"output_format\":\"png\",\"revised_prompt\":\"revised history prompt\",\"result\":\"aGVsbG8=\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_img_123\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"tool_usage\":{\"image_gen\":{\"input_tokens\":171,\"output_tokens\":1372,\"total_tokens\":1543}}}}\n\n"
)
.as_bytes(),
);
let product = maybe_build_openai_image_sync_finalize_product(
"openai_image_sync_finalize",
200,
Some(&report_context),
Some(&body_base64),
)
.expect("finalize should succeed")
.expect("finalize should match");
assert_eq!(product.client_body_json["created"], 1776839946);
assert_eq!(product.client_body_json["data"][0]["b64_json"], "aGVsbG8=");
assert_eq!(
product.client_body_json["data"][0]["revised_prompt"],
"revised history prompt"
);
assert_eq!(product.client_body_json["usage"]["input_tokens"], 171);
assert_eq!(product.provider_body_json["id"], "resp_img_123");
assert_eq!(
product.provider_body_json["output"][0]["output_format"],
"png"
);
assert_eq!(
product.provider_body_json["output"][0]["revised_prompt"],
"revised history prompt"
);
}
}

6
crates/aether-ai-pipeline/src/finalize/sse.rs → crates/aether-ai-surfaces/src/finalize/sse.rs
View File

@@ -1,6 +1,6 @@
use serde_json::Value;
use crate::finalize::PipelineFinalizeError;
use crate::finalize::AiSurfaceFinalizeError;
pub fn map_claude_stop_reason(
stop_reason: Option<&str>,
@@ -27,7 +27,7 @@ pub fn encode_done_sse() -> Vec<u8> {
pub fn encode_json_sse(
event: Option<&str>,
value: &Value,
) -> Result<Vec<u8>, PipelineFinalizeError> {
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut out = Vec::new();
if let Some(event) = event.filter(|value| !value.trim().is_empty()) {
out.extend_from_slice(b"event: ");
@@ -35,7 +35,7 @@ pub fn encode_json_sse(
out.push(b'\n');
}
out.extend_from_slice(b"data: ");
out.extend(serde_json::to_vec(value).map_err(PipelineFinalizeError::from)?);
out.extend(serde_json::to_vec(value).map_err(AiSurfaceFinalizeError::from)?);
out.extend_from_slice(b"\n\n");
Ok(out)
}

0
crates/aether-ai-pipeline/src/finalize/standard/claude/mod.rs → crates/aether-ai-surfaces/src/finalize/standard/claude/mod.rs
View File

24
crates/aether-ai-pipeline/src/finalize/standard/claude/stream.rs → crates/aether-ai-surfaces/src/finalize/standard/claude/stream.rs
View File

@@ -5,7 +5,7 @@ use serde_json::{json, Map, Value};
use crate::finalize::common::{build_generated_tool_call_id, canonicalize_tool_arguments};
use crate::finalize::sse::{encode_json_sse, map_claude_stop_reason};
use crate::finalize::standard::stream_core::common::*;
use crate::finalize::PipelineFinalizeError;
use crate::finalize::AiSurfaceFinalizeError;
#[derive(Default)]
struct ClaudeProviderToolState {
@@ -60,7 +60,7 @@ impl ClaudeProviderState {
&mut self,
report_context: &Value,
line: Vec<u8>,
) -> Result<Vec<CanonicalStreamFrame>, PipelineFinalizeError> {
) -> Result<Vec<CanonicalStreamFrame>, AiSurfaceFinalizeError> {
let Some(value) = decode_json_data_line(&line) else {
return Ok(Vec::new());
};
@@ -339,7 +339,7 @@ impl ClaudeProviderState {
pub fn finish(
&mut self,
report_context: &Value,
) -> Result<Vec<CanonicalStreamFrame>, PipelineFinalizeError> {
) -> Result<Vec<CanonicalStreamFrame>, AiSurfaceFinalizeError> {
if !self.started || self.finished {
return Ok(Vec::new());
}
@@ -406,7 +406,7 @@ impl ClaudeClientEmitter {
self.model = Some(frame.model.clone());
}
fn ensure_started(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
fn ensure_started(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.started {
return Ok(Vec::new());
}
@@ -432,7 +432,7 @@ impl ClaudeClientEmitter {
)
}
fn close_open_block(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
fn close_open_block(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let Some(open_block) = self.open_block.take() else {
return Ok(Vec::new());
};
@@ -450,7 +450,7 @@ impl ClaudeClientEmitter {
)
}
fn ensure_text_block(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
fn ensure_text_block(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut out = Vec::new();
if let Some(ClaudeOpenBlock::Text { .. }) = self.open_block {
return Ok(out);
@@ -473,7 +473,7 @@ impl ClaudeClientEmitter {
Ok(out)
}
fn ensure_thinking_block(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
fn ensure_thinking_block(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut out = Vec::new();
if let Some(ClaudeOpenBlock::Thinking { .. }) = self.open_block {
return Ok(out);
@@ -501,7 +501,7 @@ impl ClaudeClientEmitter {
tool_index: usize,
call_id: &str,
name: &str,
) -> Result<Vec<u8>, PipelineFinalizeError> {
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut out = Vec::new();
if let Some(ClaudeOpenBlock::Tool {
tool_index: current_tool_index,
@@ -549,7 +549,7 @@ impl ClaudeClientEmitter {
tool_use_id: String,
name: Option<String>,
content: String,
) -> Result<Vec<u8>, PipelineFinalizeError> {
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut out = self.ensure_started()?;
out.extend(self.close_open_block()?);
let block_index = self.next_block_index;
@@ -580,7 +580,7 @@ impl ClaudeClientEmitter {
Ok(out)
}
pub fn emit(&mut self, frame: CanonicalStreamFrame) -> Result<Vec<u8>, PipelineFinalizeError> {
pub fn emit(&mut self, frame: CanonicalStreamFrame) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
self.update_identity(&frame);
match frame.event {
CanonicalStreamEvent::Start => self.ensure_started(),
@@ -742,7 +742,7 @@ impl ClaudeClientEmitter {
}
}
pub fn finish(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
pub fn finish(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if !self.started || self.finished {
return Ok(Vec::new());
}
@@ -762,7 +762,7 @@ impl ClaudeClientEmitter {
fn emit_content_part(
&mut self,
part: CanonicalContentPart,
) -> Result<Vec<u8>, PipelineFinalizeError> {
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut out = self.ensure_started()?;
out.extend(self.close_open_block()?);
let block_index = self.next_block_index;

0
crates/aether-ai-pipeline/src/finalize/standard/gemini/mod.rs → crates/aether-ai-surfaces/src/finalize/standard/gemini/mod.rs
View File

14
crates/aether-ai-pipeline/src/finalize/standard/gemini/stream.rs → crates/aether-ai-surfaces/src/finalize/standard/gemini/stream.rs
View File

@@ -5,7 +5,7 @@ use serde_json::{json, Map, Value};
use crate::finalize::common::{build_generated_tool_call_id, canonicalize_tool_arguments};
use crate::finalize::sse::encode_json_sse;
use crate::finalize::standard::stream_core::common::*;
use crate::finalize::PipelineFinalizeError;
use crate::finalize::AiSurfaceFinalizeError;
#[derive(Default)]
struct GeminiProviderToolState {
@@ -71,7 +71,7 @@ impl GeminiProviderState {
&mut self,
report_context: &Value,
line: Vec<u8>,
) -> Result<Vec<CanonicalStreamFrame>, PipelineFinalizeError> {
) -> Result<Vec<CanonicalStreamFrame>, AiSurfaceFinalizeError> {
let Some(value) = decode_json_data_line(&line) else {
return Ok(Vec::new());
};
@@ -326,7 +326,7 @@ impl GeminiProviderState {
pub fn finish(
&mut self,
report_context: &Value,
) -> Result<Vec<CanonicalStreamFrame>, PipelineFinalizeError> {
) -> Result<Vec<CanonicalStreamFrame>, AiSurfaceFinalizeError> {
if !self.started || self.finished {
return Ok(Vec::new());
}
@@ -370,7 +370,7 @@ impl GeminiClientEmitter {
parts: Vec<Value>,
finish_reason: Option<&str>,
usage: Option<CanonicalUsage>,
) -> Result<Vec<u8>, PipelineFinalizeError> {
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut candidate = Map::new();
candidate.insert(
"content".to_string(),
@@ -429,7 +429,7 @@ impl GeminiClientEmitter {
encode_json_sse(None, &Value::Object(response))
}
fn flush_pending_tool_calls(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
fn flush_pending_tool_calls(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut out = Vec::new();
let mut pending = Vec::new();
for (index, tool_call) in &mut self.tool_calls {
@@ -461,7 +461,7 @@ impl GeminiClientEmitter {
Ok(out)
}
pub fn emit(&mut self, frame: CanonicalStreamFrame) -> Result<Vec<u8>, PipelineFinalizeError> {
pub fn emit(&mut self, frame: CanonicalStreamFrame) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
self.update_identity(&frame);
match frame.event {
CanonicalStreamEvent::Start => Ok(Vec::new()),
@@ -554,7 +554,7 @@ impl GeminiClientEmitter {
}
}
pub fn finish(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
pub fn finish(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.finished {
return Ok(Vec::new());
}

0
crates/aether-ai-pipeline/src/finalize/standard/mod.rs → crates/aether-ai-surfaces/src/finalize/standard/mod.rs
View File

0
crates/aether-ai-pipeline/src/finalize/standard/openai/mod.rs → crates/aether-ai-surfaces/src/finalize/standard/openai/mod.rs
View File

38
crates/aether-ai-pipeline/src/finalize/standard/openai/stream.rs → crates/aether-ai-surfaces/src/finalize/standard/openai/stream.rs
View File

@@ -5,7 +5,7 @@ use serde_json::{json, Map, Value};
use crate::finalize::common::build_generated_tool_call_id;
use crate::finalize::sse::{encode_done_sse, encode_json_sse};
use crate::finalize::standard::stream_core::common::*;
use crate::finalize::PipelineFinalizeError;
use crate::finalize::AiSurfaceFinalizeError;
#[derive(Default)]
struct OpenAIChatProviderToolState {
@@ -105,7 +105,7 @@ impl OpenAIChatProviderState {
&mut self,
report_context: &Value,
line: Vec<u8>,
) -> Result<Vec<CanonicalStreamFrame>, PipelineFinalizeError> {
) -> Result<Vec<CanonicalStreamFrame>, AiSurfaceFinalizeError> {
let Some(value) = decode_json_data_line(&line) else {
return Ok(Vec::new());
};
@@ -344,7 +344,7 @@ impl OpenAIChatProviderState {
pub fn finish(
&mut self,
report_context: &Value,
) -> Result<Vec<CanonicalStreamFrame>, PipelineFinalizeError> {
) -> Result<Vec<CanonicalStreamFrame>, AiSurfaceFinalizeError> {
if !self.started || self.finished {
return Ok(Vec::new());
}
@@ -682,7 +682,7 @@ impl OpenAIResponsesProviderState {
&mut self,
report_context: &Value,
line: Vec<u8>,
) -> Result<Vec<CanonicalStreamFrame>, PipelineFinalizeError> {
) -> Result<Vec<CanonicalStreamFrame>, AiSurfaceFinalizeError> {
let Some(value) = decode_json_data_line(&line) else {
return Ok(Vec::new());
};
@@ -1112,7 +1112,7 @@ impl OpenAIResponsesProviderState {
pub fn finish(
&mut self,
report_context: &Value,
) -> Result<Vec<CanonicalStreamFrame>, PipelineFinalizeError> {
) -> Result<Vec<CanonicalStreamFrame>, AiSurfaceFinalizeError> {
if !self.started || self.finished {
return Ok(Vec::new());
}
@@ -1187,7 +1187,7 @@ impl OpenAIChatClientEmitter {
self.model = Some(frame.model.clone());
}
fn ensure_started(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
fn ensure_started(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.started {
return Ok(Vec::new());
}
@@ -1203,7 +1203,7 @@ impl OpenAIChatClientEmitter {
)
}
pub fn emit(&mut self, frame: CanonicalStreamFrame) -> Result<Vec<u8>, PipelineFinalizeError> {
pub fn emit(&mut self, frame: CanonicalStreamFrame) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
self.update_identity(&frame);
match frame.event {
CanonicalStreamEvent::Start => self.ensure_started(),
@@ -1388,7 +1388,7 @@ impl OpenAIChatClientEmitter {
}
}
pub fn finish(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
pub fn finish(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if !self.started || self.finished {
return Ok(Vec::new());
}
@@ -1469,7 +1469,7 @@ impl OpenAIResponsesClientEmitter {
&mut self,
event: &str,
mut payload: Value,
) -> Result<Vec<u8>, PipelineFinalizeError> {
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if let Some(object) = payload.as_object_mut() {
object.insert(
"sequence_number".to_string(),
@@ -1484,7 +1484,7 @@ impl OpenAIResponsesClientEmitter {
self.model = Some(frame.model.clone());
}
fn ensure_started(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
fn ensure_started(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.started {
return Ok(Vec::new());
}
@@ -1550,7 +1550,7 @@ impl OpenAIResponsesClientEmitter {
output_index
}
fn ensure_reasoning_item_started(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
fn ensure_reasoning_item_started(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut out = self.ensure_started()?;
let output_index = self.ensure_reasoning_output_index();
let item_id = self.ensure_reasoning_item_id();
@@ -1590,7 +1590,7 @@ impl OpenAIResponsesClientEmitter {
Ok(out)
}
fn ensure_text_item_started(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
fn ensure_text_item_started(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut out = self.ensure_started()?;
let output_index = self.ensure_message_output_index();
let item_id = self.ensure_message_item_id();
@@ -1633,7 +1633,7 @@ impl OpenAIResponsesClientEmitter {
Ok(out)
}
fn finish_text_item(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
fn finish_text_item(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if !self.text_item_started {
return Ok(Vec::new());
}
@@ -1690,7 +1690,7 @@ impl OpenAIResponsesClientEmitter {
Ok(out)
}
fn finish_reasoning_item(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
fn finish_reasoning_item(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if !self.reasoning_item_started {
return Ok(Vec::new());
}
@@ -1743,7 +1743,7 @@ impl OpenAIResponsesClientEmitter {
Ok(out)
}
fn finish_tool_items(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
fn finish_tool_items(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut out = Vec::new();
let indices = self.tool_calls.keys().copied().collect::<Vec<_>>();
for index in indices {
@@ -1790,7 +1790,7 @@ impl OpenAIResponsesClientEmitter {
Ok(out)
}
fn finish_tool_result_items(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
fn finish_tool_result_items(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut out = Vec::new();
let indices = self.tool_results.keys().copied().collect::<Vec<_>>();
for index in indices {
@@ -1939,7 +1939,7 @@ impl OpenAIResponsesClientEmitter {
})
}
pub fn emit(&mut self, frame: CanonicalStreamFrame) -> Result<Vec<u8>, PipelineFinalizeError> {
pub fn emit(&mut self, frame: CanonicalStreamFrame) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
self.update_identity(&frame);
match frame.event {
CanonicalStreamEvent::Start => self.ensure_started(),
@@ -2147,7 +2147,7 @@ impl OpenAIResponsesClientEmitter {
}
}
pub fn emit_error(&mut self, error_body: Value) -> Result<Vec<u8>, PipelineFinalizeError> {
pub fn emit_error(&mut self, error_body: Value) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let Some(error) = error_body.get("error").cloned() else {
return Ok(Vec::new());
};
@@ -2161,7 +2161,7 @@ impl OpenAIResponsesClientEmitter {
)
}
pub fn finish(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
pub fn finish(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if !self.started || self.finished {
return Ok(Vec::new());
}

0
crates/aether-ai-pipeline/src/finalize/standard/stream_core/common.rs → crates/aether-ai-surfaces/src/finalize/standard/stream_core/common.rs
View File

28
crates/aether-ai-pipeline/src/finalize/standard/stream_core/format_matrix.rs → crates/aether-ai-surfaces/src/finalize/standard/stream_core/format_matrix.rs
View File

@@ -2,7 +2,9 @@ use aether_ai_formats::FormatId;
use aether_contracts::{ExecutionStreamTerminalSummary, StandardizedUsage};
use serde_json::Value;
use crate::conversion::{build_core_error_body_for_client_format, LocalCoreSyncErrorKind};
use crate::finalize::error_body::{
build_core_error_body_for_client_format, LocalCoreSyncErrorKind,
};
use crate::finalize::sse::encode_json_sse;
use crate::finalize::standard::claude::stream::{ClaudeClientEmitter, ClaudeProviderState};
use crate::finalize::standard::gemini::stream::{GeminiClientEmitter, GeminiProviderState};
@@ -13,7 +15,7 @@ use crate::finalize::standard::openai::stream::{
use crate::finalize::standard::stream_core::common::{
decode_json_data_line, CanonicalStreamEvent, CanonicalStreamFrame, CanonicalUsage,
};
use crate::finalize::PipelineFinalizeError;
use crate::finalize::AiSurfaceFinalizeError;
#[derive(Default)]
pub struct StreamingStandardFormatMatrix {
@@ -27,7 +29,7 @@ impl StreamingStandardFormatMatrix {
&mut self,
report_context: &Value,
line: Vec<u8>,
) -> Result<Vec<u8>, PipelineFinalizeError> {
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.terminated {
return Ok(Vec::new());
}
@@ -43,7 +45,7 @@ impl StreamingStandardFormatMatrix {
self.emit_frames(frames)
}
pub fn finish(&mut self, report_context: &Value) -> Result<Vec<u8>, PipelineFinalizeError> {
pub fn finish(&mut self, report_context: &Value) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.terminated {
return Ok(Vec::new());
}
@@ -74,7 +76,7 @@ impl StreamingStandardFormatMatrix {
fn emit_frames(
&mut self,
frames: Vec<CanonicalStreamFrame>,
) -> Result<Vec<u8>, PipelineFinalizeError> {
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let Some(client) = self.client.as_mut() else {
return Ok(Vec::new());
};
@@ -85,7 +87,7 @@ impl StreamingStandardFormatMatrix {
Ok(out)
}
fn emit_error(&mut self, error_body: Value) -> Result<Vec<u8>, PipelineFinalizeError> {
fn emit_error(&mut self, error_body: Value) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let Some(client) = self.client.as_mut() else {
return Ok(Vec::new());
};
@@ -104,7 +106,7 @@ impl StreamingStandardTerminalObserver {
&mut self,
report_context: &Value,
line: Vec<u8>,
) -> Result<(), PipelineFinalizeError> {
) -> Result<(), AiSurfaceFinalizeError> {
self.ensure_initialized(report_context);
let Some(provider) = self.provider.as_mut() else {
return Ok(());
@@ -117,7 +119,7 @@ impl StreamingStandardTerminalObserver {
pub fn finish(
&mut self,
report_context: &Value,
) -> Result<Option<ExecutionStreamTerminalSummary>, PipelineFinalizeError> {
) -> Result<Option<ExecutionStreamTerminalSummary>, AiSurfaceFinalizeError> {
self.ensure_initialized(report_context);
let Some(provider) = self.provider.as_mut() else {
return Ok(self.latest_summary.clone());
@@ -215,7 +217,7 @@ impl ProviderStreamParser {
&mut self,
report_context: &Value,
line: Vec<u8>,
) -> Result<Vec<CanonicalStreamFrame>, PipelineFinalizeError> {
) -> Result<Vec<CanonicalStreamFrame>, AiSurfaceFinalizeError> {
match self {
ProviderStreamParser::OpenAIChat(state) => state.push_line(report_context, line),
ProviderStreamParser::OpenAIResponses(state) => state.push_line(report_context, line),
@@ -227,7 +229,7 @@ impl ProviderStreamParser {
fn finish(
&mut self,
report_context: &Value,
) -> Result<Vec<CanonicalStreamFrame>, PipelineFinalizeError> {
) -> Result<Vec<CanonicalStreamFrame>, AiSurfaceFinalizeError> {
match self {
ProviderStreamParser::OpenAIChat(state) => state.finish(report_context),
ProviderStreamParser::OpenAIResponses(state) => state.finish(report_context),
@@ -295,7 +297,7 @@ impl ClientStreamEmitter {
})
}
fn emit(&mut self, frame: CanonicalStreamFrame) -> Result<Vec<u8>, PipelineFinalizeError> {
fn emit(&mut self, frame: CanonicalStreamFrame) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
match self {
ClientStreamEmitter::OpenAIChat(state) => state.emit(frame),
ClientStreamEmitter::OpenAIResponses(state) => state.emit(frame),
@@ -304,7 +306,7 @@ impl ClientStreamEmitter {
}
}
fn finish(&mut self) -> Result<Vec<u8>, PipelineFinalizeError> {
fn finish(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
match self {
ClientStreamEmitter::OpenAIChat(state) => state.finish(),
ClientStreamEmitter::OpenAIResponses(state) => state.finish(),
@@ -313,7 +315,7 @@ impl ClientStreamEmitter {
}
}
fn emit_error(&mut self, error_body: Value) -> Result<Vec<u8>, PipelineFinalizeError> {
fn emit_error(&mut self, error_body: Value) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
match self {
ClientStreamEmitter::OpenAIResponses(state) => state.emit_error(error_body),
ClientStreamEmitter::Claude(_) => {

0
crates/aether-ai-pipeline/src/finalize/standard/stream_core/mod.rs → crates/aether-ai-surfaces/src/finalize/standard/stream_core/mod.rs
View File

356
crates/aether-ai-surfaces/src/finalize/stream_rewrite.rs Normal file
View File

@@ -0,0 +1,356 @@
use serde_json::Value;
use crate::adaptation::kiro_stream::KiroToClaudeCliStreamState;
use crate::adaptation::private_envelope::transform_provider_private_stream_line;
use crate::adaptation::surfaces::{
provider_adaptation_should_unwrap_stream_envelope, KIRO_ENVELOPE_NAME,
};
use crate::finalize::openai_image_stream::OpenAiImageStreamState;
use crate::finalize::standard::stream_core::StreamingStandardFormatMatrix;
use crate::finalize::AiSurfaceFinalizeError;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FinalizeStreamRewriteMode {
EnvelopeUnwrap,
OpenAiImage,
Standard,
KiroToClaudeCli,
KiroToClaudeCliThenStandard,
}
pub fn resolve_finalize_stream_rewrite_mode(
report_context: &Value,
) -> Option<FinalizeStreamRewriteMode> {
let needs_conversion = report_context
.get("needs_conversion")
.and_then(Value::as_bool)
.unwrap_or(false);
let envelope_name = report_context
.get("envelope_name")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
let provider_api_format = report_context
.get("provider_api_format")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
let client_api_format = report_context
.get("client_api_format")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
if needs_conversion
&& envelope_name.eq_ignore_ascii_case(KIRO_ENVELOPE_NAME)
&& provider_api_format == "claude:messages"
{
return supports_standard_stream_rewrite(
provider_api_format.as_str(),
client_api_format.as_str(),
)
.then_some(FinalizeStreamRewriteMode::KiroToClaudeCliThenStandard);
}
if needs_conversion {
return supports_standard_stream_rewrite(
provider_api_format.as_str(),
client_api_format.as_str(),
)
.then_some(FinalizeStreamRewriteMode::Standard);
}
if provider_api_format == "openai:image" && client_api_format == "openai:image" {
return Some(FinalizeStreamRewriteMode::OpenAiImage);
}
if envelope_name.eq_ignore_ascii_case(KIRO_ENVELOPE_NAME) {
return (provider_api_format == "claude:messages"
&& client_api_format == "claude:messages")
.then_some(FinalizeStreamRewriteMode::KiroToClaudeCli);
}
(provider_api_format == client_api_format
&& provider_adaptation_should_unwrap_stream_envelope(
envelope_name.as_str(),
provider_api_format.as_str(),
))
.then_some(FinalizeStreamRewriteMode::EnvelopeUnwrap)
}
enum AiSurfaceStreamRewriteState {
EnvelopeUnwrap,
OpenAiImage(Box<OpenAiImageStreamState>),
Standard(Box<StreamingStandardFormatMatrix>),
KiroToClaudeCli(Box<KiroToClaudeCliStreamState>),
KiroToClaudeCliThenStandard {
kiro: Box<KiroToClaudeCliStreamState>,
standard: Box<StreamingStandardFormatMatrix>,
},
}
pub struct AiSurfaceStreamRewriter<'a> {
report_context: &'a Value,
buffered: Vec<u8>,
state: AiSurfaceStreamRewriteState,
}
pub fn maybe_build_ai_surface_stream_rewriter<'a>(
report_context: Option<&'a Value>,
) -> Option<AiSurfaceStreamRewriter<'a>> {
let report_context = report_context?;
let state = match resolve_finalize_stream_rewrite_mode(report_context)? {
FinalizeStreamRewriteMode::EnvelopeUnwrap => AiSurfaceStreamRewriteState::EnvelopeUnwrap,
FinalizeStreamRewriteMode::OpenAiImage => {
AiSurfaceStreamRewriteState::OpenAiImage(Box::<OpenAiImageStreamState>::default())
}
FinalizeStreamRewriteMode::Standard => {
AiSurfaceStreamRewriteState::Standard(Box::<StreamingStandardFormatMatrix>::default())
}
FinalizeStreamRewriteMode::KiroToClaudeCli => AiSurfaceStreamRewriteState::KiroToClaudeCli(
Box::new(KiroToClaudeCliStreamState::new(report_context)),
),
FinalizeStreamRewriteMode::KiroToClaudeCliThenStandard => {
AiSurfaceStreamRewriteState::KiroToClaudeCliThenStandard {
kiro: Box::new(KiroToClaudeCliStreamState::new(report_context)),
standard: Box::<StreamingStandardFormatMatrix>::default(),
}
}
};
Some(AiSurfaceStreamRewriter {
report_context,
buffered: Vec::new(),
state,
})
}
impl AiSurfaceStreamRewriter<'_> {
pub fn push_chunk(&mut self, chunk: &[u8]) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
match &mut self.state {
AiSurfaceStreamRewriteState::OpenAiImage(state) => {
state.push_chunk(self.report_context, chunk)
}
AiSurfaceStreamRewriteState::KiroToClaudeCli(state) => {
state.push_chunk(self.report_context, chunk)
}
AiSurfaceStreamRewriteState::KiroToClaudeCliThenStandard { kiro, standard } => {
let claude_bytes = kiro.push_chunk(self.report_context, chunk)?;
transform_standard_bytes(standard, self.report_context, claude_bytes)
}
AiSurfaceStreamRewriteState::EnvelopeUnwrap
| AiSurfaceStreamRewriteState::Standard(_) => {
self.buffered.extend_from_slice(chunk);
let mut output = Vec::new();
while let Some(line_end) = self.buffered.iter().position(|byte| *byte == b'\n') {
let line = self.buffered.drain(..=line_end).collect::<Vec<_>>();
output.extend(self.transform_line(line)?);
}
Ok(output)
}
}
}
pub fn finish(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
match &mut self.state {
AiSurfaceStreamRewriteState::OpenAiImage(state) => state.finish(self.report_context),
AiSurfaceStreamRewriteState::KiroToClaudeCli(state) => {
state.finish(self.report_context)
}
AiSurfaceStreamRewriteState::KiroToClaudeCliThenStandard { kiro, standard } => {
let mut output = transform_standard_bytes(
standard,
self.report_context,
kiro.finish(self.report_context)?,
)?;
output.extend(standard.finish(self.report_context)?);
Ok(output)
}
AiSurfaceStreamRewriteState::EnvelopeUnwrap
| AiSurfaceStreamRewriteState::Standard(_) => {
if self.buffered.is_empty() {
if let AiSurfaceStreamRewriteState::Standard(state) = &mut self.state {
return state.finish(self.report_context);
}
return Ok(Vec::new());
}
let line = std::mem::take(&mut self.buffered);
let mut output = self.transform_line(line)?;
if let AiSurfaceStreamRewriteState::Standard(state) = &mut self.state {
output.extend(state.finish(self.report_context)?);
}
Ok(output)
}
}
}
fn transform_line(&mut self, line: Vec<u8>) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
match &mut self.state {
AiSurfaceStreamRewriteState::EnvelopeUnwrap => {
transform_provider_private_stream_line(self.report_context, line)
.map_err(AiSurfaceFinalizeError::from)
}
AiSurfaceStreamRewriteState::Standard(state) => {
transform_standard_line(state, self.report_context, line)
}
AiSurfaceStreamRewriteState::OpenAiImage(_)
| AiSurfaceStreamRewriteState::KiroToClaudeCli(_)
| AiSurfaceStreamRewriteState::KiroToClaudeCliThenStandard { .. } => Ok(Vec::new()),
}
}
}
fn transform_standard_bytes(
standard: &mut StreamingStandardFormatMatrix,
report_context: &Value,
bytes: Vec<u8>,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if bytes.is_empty() {
return Ok(Vec::new());
}
let mut output = Vec::new();
for line in bytes.split_inclusive(|byte| *byte == b'\n') {
output.extend(transform_standard_line(
standard,
report_context,
line.to_vec(),
)?);
}
Ok(output)
}
fn transform_standard_line(
standard: &mut StreamingStandardFormatMatrix,
report_context: &Value,
line: Vec<u8>,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let line = if should_unwrap_envelope(report_context) {
transform_provider_private_stream_line(report_context, line)?
} else {
line
};
if line.is_empty() {
return Ok(Vec::new());
}
standard.transform_line(report_context, line)
}
fn should_unwrap_envelope(report_context: &Value) -> bool {
let envelope_name = report_context
.get("envelope_name")
.and_then(Value::as_str)
.unwrap_or_default();
let provider_api_format = report_context
.get("provider_api_format")
.and_then(Value::as_str)
.unwrap_or_default();
provider_adaptation_should_unwrap_stream_envelope(envelope_name, provider_api_format)
}
fn supports_standard_stream_rewrite(provider_api_format: &str, client_api_format: &str) -> bool {
is_standard_provider_api_format(provider_api_format)
&& (is_standard_chat_client_api_format(client_api_format)
|| is_standard_cli_client_api_format(client_api_format))
}
fn is_standard_provider_api_format(api_format: &str) -> bool {
matches!(
aether_ai_formats::normalize_api_format_alias(api_format).as_str(),
"openai:chat"
| "openai:responses"
| "openai:responses:compact"
| "claude:messages"
| "gemini:generate_content"
)
}
fn is_standard_chat_client_api_format(api_format: &str) -> bool {
matches!(
api_format,
"openai:chat" | "claude:messages" | "gemini:generate_content"
)
}
fn is_standard_cli_client_api_format(api_format: &str) -> bool {
matches!(
aether_ai_formats::normalize_api_format_alias(api_format).as_str(),
"openai:responses"
| "openai:responses:compact"
| "claude:messages"
| "gemini:generate_content"
)
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{resolve_finalize_stream_rewrite_mode, FinalizeStreamRewriteMode};
#[test]
fn resolves_standard_mode_for_cross_format_standard_streams() {
let report_context = json!({
"provider_api_format": "claude:messages",
"client_api_format": "openai:chat",
"needs_conversion": true,
});
assert_eq!(
resolve_finalize_stream_rewrite_mode(&report_context),
Some(FinalizeStreamRewriteMode::Standard)
);
}
#[test]
fn resolves_envelope_unwrap_for_same_format_private_envelopes() {
let report_context = json!({
"provider_api_format": "gemini:generate_content",
"client_api_format": "gemini:generate_content",
"envelope_name": "antigravity:v1internal",
"needs_conversion": false,
});
assert_eq!(
resolve_finalize_stream_rewrite_mode(&report_context),
Some(FinalizeStreamRewriteMode::EnvelopeUnwrap)
);
}
#[test]
fn resolves_kiro_same_format_streams_to_kiro_mode() {
let report_context = json!({
"provider_api_format": "claude:messages",
"client_api_format": "claude:messages",
"envelope_name": "kiro:generateAssistantResponse",
"needs_conversion": false,
});
assert_eq!(
resolve_finalize_stream_rewrite_mode(&report_context),
Some(FinalizeStreamRewriteMode::KiroToClaudeCli)
);
}
#[test]
fn rejects_unsupported_non_conversion_streams() {
let report_context = json!({
"provider_api_format": "openai:chat",
"client_api_format": "openai:chat",
"needs_conversion": false,
});
assert_eq!(resolve_finalize_stream_rewrite_mode(&report_context), None);
}
#[test]
fn resolves_openai_image_mode_for_same_format_image_streams() {
let report_context = json!({
"provider_api_format": "openai:image",
"client_api_format": "openai:image",
"needs_conversion": false,
});
assert_eq!(
resolve_finalize_stream_rewrite_mode(&report_context),
Some(FinalizeStreamRewriteMode::OpenAiImage)
);
}
}

34
crates/aether-ai-pipeline/src/finalize/sync_products.rs → crates/aether-ai-surfaces/src/finalize/sync_products.rs
View File

@@ -1,23 +1,23 @@
use base64::Engine as _;
use std::collections::BTreeMap;
use aether_ai_formats::registry::{convert_response, FormatContext};
use serde_json::{json, Map, Value};
use super::PipelineFinalizeError;
use crate::conversion::response::{
use aether_ai_formats::conversion::response::{
convert_claude_response_to_openai_responses, convert_gemini_chat_response_to_openai_chat,
convert_openai_chat_response_to_claude_chat, convert_openai_chat_response_to_gemini_chat,
convert_openai_chat_response_to_openai_responses,
convert_openai_responses_response_to_openai_chat,
};
use crate::conversion::{
use aether_ai_formats::registry::{convert_response, FormatContext};
use aether_ai_formats::{
canonical_to_claude_response, canonical_to_gemini_response, canonical_to_openai_chat_response,
canonical_to_openai_responses_compact_response, canonical_to_openai_responses_response,
from_claude_to_canonical_response, from_gemini_to_canonical_response,
from_openai_chat_to_canonical_response, from_openai_responses_to_canonical_response,
sync_chat_response_conversion_kind, sync_cli_response_conversion_kind,
};
use serde_json::{json, Map, Value};
use super::AiSurfaceFinalizeError;
use crate::finalize::standard::gemini::stream::GeminiProviderState;
use crate::finalize::standard::stream_core::common::{
map_openai_finish_reason_to_gemini, parse_json_arguments_value, CanonicalContentPart,
@@ -42,7 +42,7 @@ pub fn maybe_build_standard_cross_format_sync_product_from_normalized_payload(
report_context: Option<&Value>,
body_json: Option<&Value>,
body_base64: Option<&str>,
) -> Result<Option<StandardCrossFormatSyncProduct>, PipelineFinalizeError> {
) -> Result<Option<StandardCrossFormatSyncProduct>, AiSurfaceFinalizeError> {
if status_code >= 400 {
return Ok(None);
}
@@ -92,7 +92,7 @@ pub fn maybe_build_standard_same_format_sync_body_from_normalized_payload(
report_context: Option<&Value>,
body_json: Option<&Value>,
body_base64: Option<&str>,
) -> Result<Option<Value>, PipelineFinalizeError> {
) -> Result<Option<Value>, AiSurfaceFinalizeError> {
let stream_body = maybe_build_standard_same_format_stream_sync_body(
report_kind,
status_code,
@@ -115,7 +115,7 @@ pub fn maybe_build_openai_responses_same_family_sync_body_from_normalized_payloa
report_context: Option<&Value>,
body_json: Option<&Value>,
body_base64: Option<&str>,
) -> Result<Option<Value>, PipelineFinalizeError> {
) -> Result<Option<Value>, AiSurfaceFinalizeError> {
let stream_body = maybe_build_openai_responses_same_family_stream_sync_body(
report_kind,
status_code,
@@ -138,7 +138,7 @@ pub fn maybe_build_openai_chat_cross_format_sync_product_from_normalized_payload
report_context: Option<&Value>,
body_json: Option<&Value>,
body_base64: Option<&str>,
) -> Result<Option<StandardCrossFormatSyncProduct>, PipelineFinalizeError> {
) -> Result<Option<StandardCrossFormatSyncProduct>, AiSurfaceFinalizeError> {
if report_kind != "openai_chat_sync_finalize" || status_code >= 400 {
return Ok(None);
}
@@ -202,7 +202,7 @@ pub fn maybe_build_openai_responses_cross_format_sync_product_from_normalized_pa
report_context: Option<&Value>,
body_json: Option<&Value>,
body_base64: Option<&str>,
) -> Result<Option<StandardCrossFormatSyncProduct>, PipelineFinalizeError> {
) -> Result<Option<StandardCrossFormatSyncProduct>, AiSurfaceFinalizeError> {
if !is_openai_responses_finalize_kind(report_kind) || status_code >= 400 {
return Ok(None);
}
@@ -273,7 +273,7 @@ pub fn maybe_build_standard_sync_finalize_product_from_normalized_payload(
report_context: Option<&Value>,
body_json: Option<&Value>,
body_base64: Option<&str>,
) -> Result<Option<StandardSyncFinalizeNormalizedProduct>, PipelineFinalizeError> {
) -> Result<Option<StandardSyncFinalizeNormalizedProduct>, AiSurfaceFinalizeError> {
if let Some(body_json) = maybe_build_standard_same_format_sync_body_from_normalized_payload(
report_kind,
status_code,
@@ -389,7 +389,7 @@ fn maybe_build_standard_same_format_stream_sync_body(
status_code: u16,
report_context: Option<&Value>,
body_base64: Option<&str>,
) -> Result<Option<Value>, PipelineFinalizeError> {
) -> Result<Option<Value>, AiSurfaceFinalizeError> {
if status_code >= 400 {
return Ok(None);
}
@@ -487,7 +487,7 @@ fn maybe_build_openai_responses_same_family_stream_sync_body(
status_code: u16,
report_context: Option<&Value>,
body_base64: Option<&str>,
) -> Result<Option<Value>, PipelineFinalizeError> {
) -> Result<Option<Value>, AiSurfaceFinalizeError> {
if status_code >= 400 || !is_openai_responses_finalize_kind(report_kind) {
return Ok(None);
}
@@ -534,7 +534,7 @@ fn maybe_build_openai_cross_format_provider_body_from_normalized_payload(
body_json: Option<&Value>,
body_base64: Option<&str>,
provider_api_format: &str,
) -> Result<Option<Value>, PipelineFinalizeError> {
) -> Result<Option<Value>, AiSurfaceFinalizeError> {
let aggregated_stream_body = match body_base64 {
Some(body_base64) => {
let body_bytes = base64::engine::general_purpose::STANDARD.decode(body_base64)?;
@@ -2873,13 +2873,13 @@ mod tests {
maybe_build_standard_sync_finalize_product_from_normalized_payload,
StandardSyncFinalizeNormalizedProduct,
};
use crate::conversion::response::{
use aether_ai_formats::conversion::response::{
convert_claude_chat_response_to_openai_chat, convert_gemini_chat_response_to_openai_chat,
convert_openai_chat_response_to_claude_chat, convert_openai_chat_response_to_gemini_chat,
convert_openai_chat_response_to_openai_responses,
convert_openai_responses_response_to_openai_chat,
};
use crate::conversion::{sync_cli_response_conversion_kind, SyncCliResponseConversionKind};
use aether_ai_formats::{sync_cli_response_conversion_kind, SyncCliResponseConversionKind};
use base64::Engine as _;
use serde_json::json;

569
crates/aether-ai-surfaces/src/finalize/sync_to_stream.rs Normal file
View File

@@ -0,0 +1,569 @@
use aether_ai_formats::conversion::response::{
convert_claude_response_to_openai_responses, convert_gemini_response_to_openai_responses,
convert_openai_chat_response_to_openai_responses,
};
use aether_contracts::{ExecutionStreamTerminalSummary, StandardizedUsage};
use serde_json::{json, Value};
use crate::finalize::sse::encode_json_sse;
use crate::finalize::standard::claude::stream::ClaudeClientEmitter;
use crate::finalize::standard::gemini::stream::GeminiClientEmitter;
use crate::finalize::standard::openai::stream::{
OpenAIChatClientEmitter, OpenAIResponsesClientEmitter, OpenAIResponsesProviderState,
};
use crate::finalize::standard::stream_core::CanonicalStreamFrame;
use crate::finalize::AiSurfaceFinalizeError;
pub struct SyncToStreamBridgeOutcome {
pub sse_body: Vec<u8>,
pub terminal_summary: Option<ExecutionStreamTerminalSummary>,
}
pub fn maybe_bridge_standard_sync_json_to_stream(
provider_body_json: &Value,
provider_api_format: &str,
client_api_format: &str,
report_context: Option<&Value>,
) -> Result<Option<SyncToStreamBridgeOutcome>, AiSurfaceFinalizeError> {
let provider_api_format = normalize_api_format(provider_api_format);
let client_api_format = normalize_api_format(client_api_format);
if provider_api_format == "openai:image" && client_api_format == "openai:image" {
return maybe_bridge_openai_image_sync_json_to_stream(provider_body_json, report_context);
}
if !is_standard_api_format(provider_api_format.as_str())
|| !is_standard_api_format(client_api_format.as_str())
{
return Ok(None);
}
let bridge_context = build_bridge_report_context(
report_context,
provider_api_format.as_str(),
client_api_format.as_str(),
);
let Some(openai_responses_response) = convert_provider_sync_response_to_openai_responses(
provider_body_json,
provider_api_format.as_str(),
&bridge_context,
) else {
return Ok(None);
};
let terminal_summary =
build_terminal_summary_from_openai_responses_response(&openai_responses_response);
let canonical_frames = build_canonical_frames_from_openai_responses_response(
&openai_responses_response,
&bridge_context,
)?;
let sse_body =
emit_client_stream_from_canonical_frames(canonical_frames, client_api_format.as_str())?;
Ok(Some(SyncToStreamBridgeOutcome {
sse_body,
terminal_summary,
}))
}
fn maybe_bridge_openai_image_sync_json_to_stream(
provider_body_json: &Value,
report_context: Option<&Value>,
) -> Result<Option<SyncToStreamBridgeOutcome>, AiSurfaceFinalizeError> {
let Some(response) = provider_body_json.as_object() else {
return Ok(None);
};
let Some(image) = response
.get("data")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
.find_map(extract_openai_image_sync_b64_json)
else {
return Ok(None);
};
let usage = response.get("usage").cloned().unwrap_or(Value::Null);
let event_name = openai_image_completed_event_name(report_context);
let sse_body = encode_json_sse(
Some(event_name),
&json!({
"type": event_name,
"b64_json": image,
"usage": usage,
}),
)?;
Ok(Some(SyncToStreamBridgeOutcome {
sse_body,
terminal_summary: Some(ExecutionStreamTerminalSummary {
standardized_usage: response
.get("usage")
.and_then(standardized_usage_from_openai_usage),
finish_reason: Some("stop".to_string()),
response_id: response
.get("id")
.and_then(Value::as_str)
.map(ToOwned::to_owned),
model: response
.get("model")
.and_then(Value::as_str)
.map(ToOwned::to_owned)
.or_else(|| image_bridge_model(report_context)),
observed_finish: true,
unknown_event_count: 0,
parser_error: None,
}),
}))
}
fn normalize_api_format(value: &str) -> String {
aether_ai_formats::normalize_api_format_alias(value)
}
fn is_standard_api_format(value: &str) -> bool {
matches!(
value,
"openai:chat"
| "openai:responses"
| "openai:responses:compact"
| "claude:messages"
| "gemini:generate_content"
)
}
fn extract_openai_image_sync_b64_json(item: &serde_json::Map<String, Value>) -> Option<String> {
item.get("b64_json")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.or_else(|| {
item.get("url")
.and_then(Value::as_str)
.and_then(extract_base64_from_data_url)
})
}
fn extract_base64_from_data_url(value: &str) -> Option<String> {
let trimmed = value.trim();
let (metadata, payload) = trimmed.split_once(',')?;
if !metadata.starts_with("data:") || !metadata.ends_with(";base64") {
return None;
}
(!payload.trim().is_empty()).then(|| payload.trim().to_string())
}
fn openai_image_completed_event_name(report_context: Option<&Value>) -> &'static str {
if openai_image_request_operation(report_context) == Some("edit") {
"image_edit.completed"
} else {
"image_generation.completed"
}
}
fn openai_image_request_operation(report_context: Option<&Value>) -> Option<&str> {
report_context
.and_then(|value| value.get("image_request"))
.and_then(|value| value.get("operation"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn image_bridge_model(report_context: Option<&Value>) -> Option<String> {
report_context.and_then(|context| {
context
.get("mapped_model")
.or_else(|| context.get("model"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
})
}
fn build_bridge_report_context(
report_context: Option<&Value>,
provider_api_format: &str,
client_api_format: &str,
) -> Value {
let mut context = report_context
.cloned()
.filter(Value::is_object)
.unwrap_or_else(|| json!({}));
let object = context
.as_object_mut()
.expect("bridge report context should stay object");
object
.entry("provider_api_format".to_string())
.or_insert_with(|| Value::String(provider_api_format.to_string()));
object
.entry("client_api_format".to_string())
.or_insert_with(|| Value::String(client_api_format.to_string()));
context
}
fn convert_provider_sync_response_to_openai_responses(
provider_body_json: &Value,
provider_api_format: &str,
report_context: &Value,
) -> Option<Value> {
match provider_api_format {
"openai:responses" | "openai:responses:compact" => Some(provider_body_json.clone()),
"openai:chat" => convert_openai_chat_response_to_openai_responses(
provider_body_json,
report_context,
false,
),
"claude:messages" => {
convert_claude_response_to_openai_responses(provider_body_json, report_context)
}
"gemini:generate_content" => {
convert_gemini_response_to_openai_responses(provider_body_json, report_context)
}
_ => None,
}
}
fn build_canonical_frames_from_openai_responses_response(
openai_responses_response: &Value,
report_context: &Value,
) -> Result<Vec<CanonicalStreamFrame>, AiSurfaceFinalizeError> {
let mut state = OpenAIResponsesProviderState::default();
let line = format!(
"data: {}\n",
serde_json::to_string(&json!({
"type": "response.completed",
"response": openai_responses_response,
}))
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?
);
let mut frames = state
.push_line(report_context, line.into_bytes())
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?;
frames.extend(
state
.finish(report_context)
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
);
Ok(frames)
}
fn emit_client_stream_from_canonical_frames(
canonical_frames: Vec<CanonicalStreamFrame>,
client_api_format: &str,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
match client_api_format {
"openai:chat" => {
let mut emitter = OpenAIChatClientEmitter::default();
emit_with_openai_chat_emitter(&mut emitter, canonical_frames)
}
"openai:responses" | "openai:responses:compact" => {
let mut emitter = OpenAIResponsesClientEmitter::default();
emit_with_openai_responses_emitter(&mut emitter, canonical_frames)
}
"claude:messages" => {
let mut emitter = ClaudeClientEmitter::default();
emit_with_claude_emitter(&mut emitter, canonical_frames)
}
"gemini:generate_content" => {
let mut emitter = GeminiClientEmitter::default();
emit_with_gemini_emitter(&mut emitter, canonical_frames)
}
_ => Ok(Vec::new()),
}
}
fn emit_with_openai_chat_emitter(
emitter: &mut OpenAIChatClientEmitter,
canonical_frames: Vec<CanonicalStreamFrame>,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut output = Vec::new();
for frame in canonical_frames {
output.extend(
emitter
.emit(frame)
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
);
}
output.extend(
emitter
.finish()
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
);
Ok(output)
}
fn emit_with_openai_responses_emitter(
emitter: &mut OpenAIResponsesClientEmitter,
canonical_frames: Vec<CanonicalStreamFrame>,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut output = Vec::new();
for frame in canonical_frames {
output.extend(
emitter
.emit(frame)
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
);
}
output.extend(
emitter
.finish()
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
);
Ok(output)
}
fn emit_with_claude_emitter(
emitter: &mut ClaudeClientEmitter,
canonical_frames: Vec<CanonicalStreamFrame>,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut output = Vec::new();
for frame in canonical_frames {
output.extend(
emitter
.emit(frame)
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
);
}
output.extend(
emitter
.finish()
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
);
Ok(output)
}
fn emit_with_gemini_emitter(
emitter: &mut GeminiClientEmitter,
canonical_frames: Vec<CanonicalStreamFrame>,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut output = Vec::new();
for frame in canonical_frames {
output.extend(
emitter
.emit(frame)
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
);
}
output.extend(
emitter
.finish()
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
);
Ok(output)
}
fn build_terminal_summary_from_openai_responses_response(
openai_responses_response: &Value,
) -> Option<ExecutionStreamTerminalSummary> {
let response = openai_responses_response.as_object()?;
let response_id = response
.get("id")
.and_then(Value::as_str)
.map(ToOwned::to_owned);
let model = response
.get("model")
.and_then(Value::as_str)
.map(ToOwned::to_owned);
let finish_reason = response
.get("output")
.and_then(Value::as_array)
.map(|output| resolve_openai_responses_finish_reason(output))
.filter(|value| !value.trim().is_empty());
let standardized_usage = response
.get("usage")
.and_then(standardized_usage_from_openai_usage);
Some(ExecutionStreamTerminalSummary {
standardized_usage,
finish_reason,
response_id,
model,
observed_finish: true,
unknown_event_count: 0,
parser_error: None,
})
}
fn resolve_openai_responses_finish_reason(output: &[Value]) -> String {
let has_tool_calls = output.iter().filter_map(Value::as_object).any(|item| {
item.get("type")
.and_then(Value::as_str)
.is_some_and(|value| value == "function_call")
});
if has_tool_calls {
"tool_calls".to_string()
} else {
"stop".to_string()
}
}
fn standardized_usage_from_openai_usage(value: &Value) -> Option<StandardizedUsage> {
let usage = value.as_object()?;
let mut input_tokens = usage
.get("input_tokens")
.or_else(|| usage.get("prompt_tokens"))
.and_then(Value::as_i64)
.unwrap_or(0);
let output_tokens = usage
.get("output_tokens")
.or_else(|| usage.get("completion_tokens"))
.and_then(Value::as_i64)
.unwrap_or(0);
let cache_creation_tokens = usage
.get("cache_creation_input_tokens")
.and_then(Value::as_i64)
.or_else(|| {
usage
.get("input_tokens_details")
.or_else(|| usage.get("prompt_tokens_details"))
.and_then(Value::as_object)
.and_then(|details| details.get("cached_creation_tokens"))
.and_then(Value::as_i64)
})
.unwrap_or(0);
let cache_read_tokens = usage
.get("cache_read_input_tokens")
.and_then(Value::as_i64)
.or_else(|| {
usage
.get("input_tokens_details")
.or_else(|| usage.get("prompt_tokens_details"))
.and_then(Value::as_object)
.and_then(|details| details.get("cached_tokens"))
.and_then(Value::as_i64)
})
.unwrap_or(0);
let total_tokens = usage.get("total_tokens").and_then(Value::as_i64).unwrap_or(
input_tokens
.saturating_add(output_tokens)
.saturating_add(cache_creation_tokens)
.saturating_add(cache_read_tokens),
);
if input_tokens == 0 && total_tokens > output_tokens {
input_tokens = total_tokens.saturating_sub(output_tokens);
}
let mut standardized_usage = StandardizedUsage::new();
standardized_usage.input_tokens = input_tokens;
standardized_usage.output_tokens = output_tokens;
standardized_usage.cache_creation_tokens = cache_creation_tokens;
standardized_usage.cache_read_tokens = cache_read_tokens;
standardized_usage
.dimensions
.insert("total_tokens".to_string(), json!(total_tokens));
Some(standardized_usage.normalize_cache_creation_breakdown())
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{maybe_bridge_standard_sync_json_to_stream, standardized_usage_from_openai_usage};
fn utf8(bytes: Vec<u8>) -> String {
String::from_utf8(bytes).expect("utf8 should decode")
}
#[test]
fn openai_sync_usage_derives_missing_input_tokens_from_total() {
let usage = standardized_usage_from_openai_usage(&json!({
"output_tokens": 177,
"total_tokens": 20_612,
"input_tokens_details": {
"cached_tokens": 19_840,
},
}))
.expect("usage should parse");
assert_eq!(usage.input_tokens, 20_435);
assert_eq!(usage.output_tokens, 177);
assert_eq!(usage.cache_read_tokens, 19_840);
}
#[test]
fn bridges_openai_image_sync_json_to_generation_completed_sse() {
let report_context = json!({
"provider_api_format": "openai:image",
"client_api_format": "openai:image",
"mapped_model": "gpt-image-1",
"image_request": {
"operation": "generate"
}
});
let outcome = maybe_bridge_standard_sync_json_to_stream(
&json!({
"created": 1776971267,
"data": [{
"b64_json": "aGVsbG8="
}],
"usage": {
"total_tokens": 100,
"input_tokens": 50,
"output_tokens": 50,
"input_tokens_details": {
"text_tokens": 10,
"image_tokens": 40
}
}
}),
"openai:image",
"openai:image",
Some(&report_context),
)
.expect("bridge should succeed")
.expect("bridge should produce sse");
let output = utf8(outcome.sse_body);
assert!(output.contains("event: image_generation.completed"));
assert!(output.contains("\"type\":\"image_generation.completed\""));
assert!(output.contains("\"b64_json\":\"aGVsbG8=\""));
assert!(output.contains("\"total_tokens\":100"));
let summary = outcome
.terminal_summary
.expect("terminal summary should exist");
assert_eq!(summary.model.as_deref(), Some("gpt-image-1"));
assert_eq!(summary.finish_reason.as_deref(), Some("stop"));
assert_eq!(
summary
.standardized_usage
.as_ref()
.and_then(|usage| usage.dimensions.get("total_tokens"))
.cloned(),
Some(json!(100))
);
}
#[test]
fn bridges_openai_image_sync_data_url_to_edit_completed_sse() {
let report_context = json!({
"provider_api_format": "openai:image",
"client_api_format": "openai:image",
"image_request": {
"operation": "edit"
}
});
let outcome = maybe_bridge_standard_sync_json_to_stream(
&json!({
"created": 1776971267,
"data": [{
"url": "data:image/webp;base64,d29ybGQ="
}],
"usage": {
"total_tokens": 9,
"input_tokens": 4,
"output_tokens": 5
}
}),
"openai:image",
"openai:image",
Some(&report_context),
)
.expect("bridge should succeed")
.expect("bridge should produce sse");
let output = utf8(outcome.sse_body);
assert!(output.contains("event: image_edit.completed"));
assert!(output.contains("\"type\":\"image_edit.completed\""));
assert!(output.contains("\"b64_json\":\"d29ybGQ=\""));
assert!(output.contains("\"total_tokens\":9"));
}
}

2
crates/aether-ai-pipeline/src/lib.rs → crates/aether-ai-surfaces/src/lib.rs
View File

@@ -1,7 +1,5 @@
pub mod adaptation;
pub mod api;
pub mod contracts;
pub mod conversion;
pub mod finalize;
pub mod planner;
pub mod transport;

0
crates/aether-ai-pipeline/src/planner/common.rs → crates/aether-ai-surfaces/src/planner/common.rs
View File

0
crates/aether-ai-pipeline/src/planner/matrix.rs → crates/aether-ai-surfaces/src/planner/matrix.rs
View File

0
crates/aether-ai-pipeline/src/planner/mod.rs → crates/aether-ai-surfaces/src/planner/mod.rs
View File

0
crates/aether-ai-pipeline/src/planner/openai.rs → crates/aether-ai-surfaces/src/planner/openai.rs
View File

0
crates/aether-ai-pipeline/src/planner/passthrough/mod.rs → crates/aether-ai-surfaces/src/planner/passthrough/mod.rs
View File

0
crates/aether-ai-pipeline/src/planner/passthrough/provider.rs → crates/aether-ai-surfaces/src/planner/passthrough/provider.rs
View File

49
crates/aether-ai-pipeline/src/planner/route.rs → crates/aether-ai-surfaces/src/planner/route.rs
View File

@@ -14,6 +14,7 @@ use crate::contracts::{
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
};
use crate::planner::specialized::image::is_openai_image_stream_request;
pub fn resolve_execution_runtime_stream_plan_kind(
route_class: Option<&str>,
@@ -299,6 +300,19 @@ pub fn is_matching_stream_request(
}
}
pub fn is_matching_stream_http_request(
plan_kind: &str,
parts: &http::request::Parts,
body_json: &serde_json::Value,
body_base64: Option<&str>,
) -> bool {
if plan_kind == OPENAI_IMAGE_STREAM_PLAN_KIND {
return is_openai_image_stream_request(parts, body_json, body_base64);
}
is_matching_stream_request(plan_kind, parts.uri.path(), body_json)
}
pub fn supports_sync_scheduler_decision_kind(plan_kind: &str) -> bool {
matches!(
plan_kind,
@@ -341,12 +355,13 @@ pub fn supports_stream_scheduler_decision_kind(plan_kind: &str) -> bool {
#[cfg(test)]
mod tests {
use base64::Engine as _;
use http::Method;
use super::{
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
resolve_execution_runtime_sync_plan_kind, supports_stream_scheduler_decision_kind,
supports_sync_scheduler_decision_kind,
is_matching_stream_http_request, is_matching_stream_request,
resolve_execution_runtime_stream_plan_kind, resolve_execution_runtime_sync_plan_kind,
supports_stream_scheduler_decision_kind, supports_sync_scheduler_decision_kind,
};
use crate::contracts::{
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
@@ -556,4 +571,32 @@ mod tests {
&serde_json::json!({"stream": true}),
));
}
#[test]
fn http_stream_matching_detects_openai_image_multipart_stream_flag() {
let request = http::Request::builder()
.method(Method::POST)
.uri("/v1/images/edits")
.header(
http::header::CONTENT_TYPE,
"multipart/form-data; boundary=image-stream-boundary",
)
.body(())
.expect("request should build");
let (parts, _) = request.into_parts();
let body = concat!(
"--image-stream-boundary\r\n",
"Content-Disposition: form-data; name=\"stream\"\r\n\r\n",
"true\r\n",
"--image-stream-boundary--\r\n"
);
let body_base64 = base64::engine::general_purpose::STANDARD.encode(body.as_bytes());
assert!(is_matching_stream_http_request(
OPENAI_IMAGE_STREAM_PLAN_KIND,
&parts,
&serde_json::json!({}),
Some(body_base64.as_str()),
));
}
}

0
crates/aether-ai-pipeline/src/planner/specialized/files.rs → crates/aether-ai-surfaces/src/planner/specialized/files.rs
View File

1127
crates/aether-ai-surfaces/src/planner/specialized/image.rs Normal file
View File

File diff suppressed because it is too large Load Diff

0
crates/aether-ai-pipeline/src/planner/specialized/mod.rs → crates/aether-ai-surfaces/src/planner/specialized/mod.rs
View File

0
crates/aether-ai-pipeline/src/planner/specialized/video.rs → crates/aether-ai-surfaces/src/planner/specialized/video.rs
View File

0
crates/aether-ai-pipeline/src/planner/standard/claude/chat.rs → crates/aether-ai-surfaces/src/planner/standard/claude/chat.rs
View File

0
crates/aether-ai-pipeline/src/planner/standard/claude/cli.rs → crates/aether-ai-surfaces/src/planner/standard/claude/cli.rs
View File

0
crates/aether-ai-pipeline/src/planner/standard/claude/mod.rs → crates/aether-ai-surfaces/src/planner/standard/claude/mod.rs
View File

2
crates/aether-ai-pipeline/src/planner/standard/codex.rs → crates/aether-ai-surfaces/src/planner/standard/codex.rs
View File

@@ -1,7 +1,7 @@
use std::collections::BTreeMap;
use std::fmt::Write;
use aether_provider_transport::body_rules_handle_path;
use aether_ai_formats::proxy::rules::body_rules_handle_path;
use serde_json::{json, Value};
use sha1::{Digest as Sha1Digest, Sha1};
use sha2::Sha256;

0
crates/aether-ai-pipeline/src/planner/standard/family.rs → crates/aether-ai-surfaces/src/planner/standard/family.rs
View File

0
crates/aether-ai-pipeline/src/planner/standard/gemini/chat.rs → crates/aether-ai-surfaces/src/planner/standard/gemini/chat.rs
View File

0
crates/aether-ai-pipeline/src/planner/standard/gemini/cli.rs → crates/aether-ai-surfaces/src/planner/standard/gemini/cli.rs
View File

0
crates/aether-ai-pipeline/src/planner/standard/gemini/mod.rs → crates/aether-ai-surfaces/src/planner/standard/gemini/mod.rs
View File

69
crates/aether-ai-pipeline/src/planner/standard/matrix.rs → crates/aether-ai-surfaces/src/planner/standard/matrix.rs
View File

@@ -1,10 +1,14 @@
use std::borrow::Cow;
use aether_ai_formats::registry::{convert_request, FormatContext};
use aether_provider_transport::{
apply_local_body_rules, build_transport_request_url, GatewayProviderTransportSnapshot,
TransportRequestUrlParams,
use aether_ai_formats::conversion::request::{
convert_openai_chat_request_to_claude_request, convert_openai_chat_request_to_gemini_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,
};
use aether_ai_formats::proxy::rules::apply_local_body_rules;
use aether_ai_formats::registry::{convert_request, FormatContext};
use serde_json::Value;
use super::{
@@ -12,13 +16,6 @@ use super::{
codex::apply_codex_openai_responses_special_body_edits,
normalize::build_local_openai_chat_request_body,
};
use crate::conversion::request::{
convert_openai_chat_request_to_claude_request, convert_openai_chat_request_to_gemini_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,
};
#[allow(clippy::too_many_arguments)]
pub fn build_standard_request_body(
@@ -132,25 +129,6 @@ fn normalize_standard_request_to_openai_chat_request_cow<'a>(
}
}
pub fn build_standard_upstream_url(
parts: &http::request::Parts,
transport: &GatewayProviderTransportSnapshot,
mapped_model: &str,
provider_api_format: &str,
upstream_is_stream: bool,
) -> Option<String> {
build_transport_request_url(
transport,
TransportRequestUrlParams {
provider_api_format,
mapped_model: Some(mapped_model),
upstream_is_stream,
request_query: parts.uri.query(),
kiro_api_region: None,
},
)
}
#[cfg(test)]
mod tests {
use super::{
@@ -1108,35 +1086,10 @@ mod tests {
assert_eq!(claude["tools"][0]["name"], "get_weather");
assert_eq!(claude["tool_choice"]["name"], "get_weather");
let auth_config = aether_provider_transport::kiro::KiroAuthConfig {
auth_method: None,
refresh_token: None,
expires_at: None,
profile_arn: Some("arn:aws:bedrock:demo".to_string()),
region: None,
auth_region: None,
api_region: Some("us-east-1".to_string()),
client_id: None,
client_secret: None,
machine_id: None,
kiro_version: None,
system_version: None,
node_version: None,
access_token: Some("token".to_string()),
};
let kiro = aether_provider_transport::kiro::build_kiro_provider_request_body(
&claude,
"claude-sonnet-4.6",
&auth_config,
None,
)
.expect("kiro envelope should build");
let tool_spec = &kiro["conversationState"]["currentMessage"]["userInputMessage"]
["userInputMessageContext"]["tools"][0]["toolSpecification"];
assert_eq!(tool_spec["name"], "get_weather");
assert_eq!(claude["tools"][0]["name"], "get_weather");
assert!(
tool_spec["inputSchema"]["json"].get("required").is_none(),
"Kiro envelope should strip empty required arrays from tool schema"
claude["tools"][0]["input_schema"].get("required").is_some(),
"surface conversion should preserve the Claude tool schema before transport envelopes"
);
}
}

5
crates/aether-ai-pipeline/src/planner/standard/mod.rs → crates/aether-ai-surfaces/src/planner/standard/mod.rs
View File

@@ -13,10 +13,7 @@ pub use codex::{
CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT, CODEX_OPENAI_IMAGE_INTERNAL_MODEL,
};
pub use family::{LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec};
pub use matrix::{
build_standard_request_body, build_standard_upstream_url,
normalize_standard_request_to_openai_chat_request,
};
pub use matrix::{build_standard_request_body, normalize_standard_request_to_openai_chat_request};
pub use normalize::{
build_cross_format_openai_chat_request_body, build_cross_format_openai_responses_request_body,
build_local_openai_chat_request_body, build_local_openai_responses_request_body,

7
crates/aether-ai-pipeline/src/planner/standard/normalize.rs → crates/aether-ai-surfaces/src/planner/standard/normalize.rs
View File

@@ -1,11 +1,10 @@
use serde_json::{json, Value};
use crate::conversion::request::{
use aether_ai_formats::conversion::request::{
convert_openai_chat_request_to_claude_request, convert_openai_chat_request_to_gemini_request,
convert_openai_chat_request_to_openai_responses_request,
normalize_openai_responses_request_to_openai_chat_request,
};
use crate::conversion::{request_conversion_kind, RequestConversionKind};
use aether_ai_formats::{request_conversion_kind, RequestConversionKind};
use serde_json::{json, Value};
pub fn build_local_openai_chat_request_body(
body_json: &Value,

0
crates/aether-ai-pipeline/src/planner/standard/openai_responses.rs → crates/aether-ai-surfaces/src/planner/standard/openai_responses.rs
View File

5
crates/aether-provider-transport/src/antigravity.rs
View File

@@ -9,8 +9,9 @@ pub use auth::{
ANTIGRAVITY_PROVIDER_TYPE, ANTIGRAVITY_REQUEST_USER_AGENT,
};
pub use policy::{
classify_local_antigravity_request_support, AntigravityRequestSideSpec,
AntigravityRequestSideSupport, AntigravityRequestSideUnsupportedReason,
classify_local_antigravity_request_support, is_antigravity_provider_transport,
AntigravityRequestSideSpec, AntigravityRequestSideSupport,
AntigravityRequestSideUnsupportedReason,
};
pub use request::{
build_antigravity_safe_v1internal_request, classify_antigravity_safe_request_body,

Some files were not shown because too many files have changed in this diff Show More