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

feat: 全栈功能增强 - 扩展 provider/pool 管理、完善调度与数据层、重构前端 Pool 页面

Browse Source 
后端:
- 扩展 pool_admin payloads 和 provider query models，增强 endpoint key 管理
- 完善 scheduler-core 候选排序与请求候选逻辑
- 增强 usage-runtime 写入、provider-transport 网络层与 OAuth 刷新
- 改进 AI pipeline 响应转换与流式处理
- 扩展 global_models/provider_catalog 数据层查询能力
- 增强 video-tasks-core 多 provider 支持
- 新增大量集成测试覆盖 pool/keys/provider_query/frontdoor

前端:
- 重构 PoolManagement 页面，拆分状态管理/对话框逻辑到独立模块
- 新增 poolAdvancedDialog/poolSchedulingDialog/poolManagementState/poolMobilePresentation 工具函数及测试
- 改进 Dialog 组件与 provider tabs 显示

部署:
- 更新 Rust CI workflow 和 Dockerfile 构建配置

Closes #275
Co-authored-by: AAEE86 <ppk0227@hotmail.com>
This commit is contained in:
fawney19
2026-04-09 13:51:50 +08:00
parent 4fc95adfb9
commit b0b40c16ff
97 changed files with 5816 additions and 1881 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -8,5 +8,5 @@ pub use gemini_chat::convert_openai_chat_response_to_gemini_chat;
pub use openai_cli::convert_openai_chat_response_to_openai_cli;
pub use shared::{
build_openai_cli_response, build_openai_cli_response_with_content,
build_openai_cli_response_with_reasoning,
build_openai_cli_response_with_reasoning, OpenAiCliResponseUsage,
};

12
crates/aether-ai-pipeline/src/conversion/response/from_openai_chat/openai_cli.rs
View File

@@ -1,6 +1,8 @@
use serde_json::{json, Value};
use super::shared::{build_openai_cli_response_with_content, canonicalize_tool_arguments};
use super::shared::{
build_openai_cli_response_with_content, canonicalize_tool_arguments, OpenAiCliResponseUsage,
};
pub fn convert_openai_chat_response_to_openai_cli(
body_json: &Value,
@@ -143,8 +145,10 @@ pub fn convert_openai_chat_response_to_openai_cli(
message_content,
reasoning_summaries,
function_calls,
prompt_tokens,
output_tokens,
total_tokens,
OpenAiCliResponseUsage {
prompt_tokens,
output_tokens,
total_tokens,
},
))
}

33
crates/aether-ai-pipeline/src/conversion/response/from_openai_chat/shared.rs
View File

@@ -1,5 +1,12 @@
use serde_json::{json, Map, Value};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct OpenAiCliResponseUsage {
pub prompt_tokens: u64,
pub output_tokens: u64,
pub total_tokens: u64,
}
pub fn build_openai_cli_response(
response_id: &str,
model: &str,
@@ -24,9 +31,11 @@ pub fn build_openai_cli_response(
content,
Vec::new(),
function_calls,
prompt_tokens,
output_tokens,
total_tokens,
OpenAiCliResponseUsage {
prompt_tokens,
output_tokens,
total_tokens,
},
)
}
@@ -36,9 +45,7 @@ pub fn build_openai_cli_response_with_reasoning(
text: &str,
reasoning_summaries: Vec<String>,
function_calls: Vec<Value>,
prompt_tokens: u64,
output_tokens: u64,
total_tokens: u64,
usage: OpenAiCliResponseUsage,
) -> Value {
let content = if text.is_empty() {
Vec::new()
@@ -55,9 +62,7 @@ pub fn build_openai_cli_response_with_reasoning(
content,
reasoning_summaries,
function_calls,
prompt_tokens,
output_tokens,
total_tokens,
usage,
)
}
@@ -67,9 +72,7 @@ pub fn build_openai_cli_response_with_content(
content: Vec<Value>,
reasoning_summaries: Vec<String>,
function_calls: Vec<Value>,
prompt_tokens: u64,
output_tokens: u64,
total_tokens: u64,
usage: OpenAiCliResponseUsage,
) -> Value {
let mut output = Vec::new();
for (index, summary) in reasoning_summaries.into_iter().enumerate() {
@@ -104,9 +107,9 @@ pub fn build_openai_cli_response_with_content(
"model": model,
"output": output,
"usage": {
"input_tokens": prompt_tokens,
"output_tokens": output_tokens,
"total_tokens": total_tokens,
"input_tokens": usage.prompt_tokens,
"output_tokens": usage.output_tokens,
"total_tokens": usage.total_tokens,
}
})
}

12
crates/aether-ai-pipeline/src/conversion/response/to_openai_chat/claude_cli.rs
View File

@@ -1,6 +1,8 @@
use serde_json::{json, Value};
use super::super::from_openai_chat::build_openai_cli_response_with_reasoning;
use super::super::from_openai_chat::{
build_openai_cli_response_with_reasoning, OpenAiCliResponseUsage,
};
use super::shared::{build_generated_tool_call_id, canonicalize_tool_arguments};
pub fn convert_claude_cli_response_to_openai_cli(
@@ -76,8 +78,10 @@ pub fn convert_claude_cli_response_to_openai_cli(
&text,
reasoning_summaries,
function_calls,
prompt_tokens,
output_tokens,
total_tokens,
OpenAiCliResponseUsage {
prompt_tokens,
output_tokens,
total_tokens,
},
))
}

12
crates/aether-ai-pipeline/src/conversion/response/to_openai_chat/gemini_cli.rs
View File

@@ -1,6 +1,8 @@
use serde_json::{json, Value};
use super::super::from_openai_chat::build_openai_cli_response_with_content;
use super::super::from_openai_chat::{
build_openai_cli_response_with_content, OpenAiCliResponseUsage,
};
use super::shared::{
build_generated_tool_call_id, canonicalize_tool_arguments, extract_gemini_image_url,
};
@@ -99,8 +101,10 @@ pub fn convert_gemini_cli_response_to_openai_cli(
message_content,
reasoning_summaries,
function_calls,
prompt_tokens,
output_tokens,
total_tokens,
OpenAiCliResponseUsage {
prompt_tokens,
output_tokens,
total_tokens,
},
))
}

4
crates/aether-ai-pipeline/src/finalize/standard/claude/stream.rs
View File

@@ -187,12 +187,12 @@ impl ClaudeProviderState {
tool_state.call_id = block
.get("id")
.and_then(Value::as_str)
.unwrap_or_else(|| tool_state.call_id.as_str())
.unwrap_or(tool_state.call_id.as_str())
.to_string();
tool_state.name = block
.get("name")
.and_then(Value::as_str)
.unwrap_or_else(|| tool_state.name.as_str())
.unwrap_or(tool_state.name.as_str())
.to_string();
if !tool_state.started_emitted {
out.push(CanonicalStreamFrame {

4
crates/aether-ai-pipeline/src/finalize/standard/gemini/stream.rs
View File

@@ -125,12 +125,12 @@ impl GeminiProviderState {
tool_state.call_id = function_call
.get("id")
.and_then(Value::as_str)
.unwrap_or_else(|| tool_state.call_id.as_str())
.unwrap_or(tool_state.call_id.as_str())
.to_string();
tool_state.name = function_call
.get("name")
.and_then(Value::as_str)
.unwrap_or_else(|| tool_state.name.as_str())
.unwrap_or(tool_state.name.as_str())
.to_string();
if !tool_state.started_emitted {
out.push(CanonicalStreamFrame {

12
crates/aether-ai-pipeline/src/finalize/standard/openai/stream.rs
View File

@@ -367,12 +367,12 @@ impl OpenAICliProviderState {
.get("call_id")
.or_else(|| item.get("id"))
.and_then(Value::as_str)
.unwrap_or_else(|| state.call_id.as_str())
.unwrap_or(state.call_id.as_str())
.to_string();
state.name = item
.get("name")
.and_then(Value::as_str)
.unwrap_or_else(|| state.name.as_str())
.unwrap_or(state.name.as_str())
.to_string();
if !state.started_emitted {
out.push(CanonicalStreamFrame {
@@ -525,12 +525,12 @@ impl OpenAICliProviderState {
.get("call_id")
.or_else(|| item.get("id"))
.and_then(Value::as_str)
.unwrap_or_else(|| state.call_id.as_str())
.unwrap_or(state.call_id.as_str())
.to_string();
state.name = item
.get("name")
.and_then(Value::as_str)
.unwrap_or_else(|| state.name.as_str())
.unwrap_or(state.name.as_str())
.to_string();
if !state.started_emitted {
out.push(CanonicalStreamFrame {
@@ -662,7 +662,7 @@ impl OpenAIChatClientEmitter {
return Ok(Vec::new());
}
self.started = true;
Ok(encode_json_sse(
encode_json_sse(
None,
&build_openai_chat_role_chunk(
self.response_id
@@ -670,7 +670,7 @@ impl OpenAIChatClientEmitter {
.unwrap_or("chatcmpl-local-stream"),
self.model.as_deref().unwrap_or("unknown"),
),
)?)
)
}
pub fn emit(&mut self, frame: CanonicalStreamFrame) -> Result<Vec<u8>, PipelineFinalizeError> {

9
crates/aether-data-contracts/src/repository/global_models/types.rs
View File

@@ -184,6 +184,9 @@ pub struct StoredAdminGlobalModel {
pub default_tiered_pricing: Option<Value>,
pub supported_capabilities: Option<Value>,
pub config: Option<Value>,
pub provider_count: u64,
pub active_provider_count: u64,
pub usage_count: u64,
pub created_at_unix_secs: Option<u64>,
pub updated_at_unix_secs: Option<u64>,
}
@@ -199,6 +202,9 @@ impl StoredAdminGlobalModel {
default_tiered_pricing: Option<Value>,
supported_capabilities: Option<Value>,
config: Option<Value>,
provider_count: u64,
active_provider_count: u64,
usage_count: u64,
created_at_unix_secs: Option<u64>,
updated_at_unix_secs: Option<u64>,
) -> Result<Self, crate::DataLayerError> {
@@ -227,6 +233,9 @@ impl StoredAdminGlobalModel {
default_tiered_pricing,
supported_capabilities,
config,
provider_count,
active_provider_count,
usage_count,
created_at_unix_secs,
updated_at_unix_secs,
})

14
crates/aether-data-contracts/src/repository/provider_catalog/types.rs
View File

@@ -270,6 +270,8 @@ pub struct StoredProviderCatalogKey {
pub utilization_samples: Option<serde_json::Value>,
pub last_probe_increase_at_unix_secs: Option<u64>,
pub request_count: Option<u32>,
pub total_tokens: u64,
pub total_cost_usd: f64,
pub success_count: Option<u32>,
pub error_count: Option<u32>,
pub total_response_time_ms: Option<u32>,
@@ -340,6 +342,8 @@ impl StoredProviderCatalogKey {
utilization_samples: None,
last_probe_increase_at_unix_secs: None,
request_count: None,
total_tokens: 0,
total_cost_usd: 0.0,
success_count: None,
error_count: None,
total_response_time_ms: None,
@@ -425,6 +429,16 @@ impl StoredProviderCatalogKey {
self
}
pub fn with_usage_totals(mut self, total_tokens: u64, total_cost_usd: f64) -> Self {
self.total_tokens = total_tokens;
self.total_cost_usd = if total_cost_usd.is_finite() {
total_cost_usd
} else {
0.0
};
self
}
pub fn with_health_fields(
mut self,
health_by_format: Option<serde_json::Value>,

8
crates/aether-data/src/redis/kv.rs
View File

@@ -85,13 +85,13 @@ impl RedisKvRunner {
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
Ok(redis::cmd("SETEX")
redis::cmd("SETEX")
.arg(&namespaced_key)
.arg(resolved_ttl)
.arg(value)
.query_async(&mut connection)
.await
.map_redis_err()?)
.map_redis_err()
})
.await
}
@@ -104,11 +104,11 @@ impl RedisKvRunner {
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
Ok(redis::cmd("DEL")
redis::cmd("DEL")
.arg(&namespaced_key)
.query_async(&mut connection)
.await
.map_redis_err()?)
.map_redis_err()
})
.await
}

12
crates/aether-data/src/redis/stream.rs
View File

@@ -214,10 +214,10 @@ impl RedisStreamRunner {
for (key, value) in fields {
command.arg(key).arg(value);
}
Ok(command
command
.query_async::<String>(&mut connection)
.await
.map_redis_err()?)
.map_redis_err()
})
.await
}
@@ -323,10 +323,10 @@ impl RedisStreamRunner {
for id in ids {
command.arg(id);
}
Ok(command
command
.query_async::<usize>(&mut connection)
.await
.map_redis_err()?)
.map_redis_err()
})
.await
}
@@ -352,10 +352,10 @@ impl RedisStreamRunner {
for id in ids {
command.arg(id);
}
Ok(command
command
.query_async::<usize>(&mut connection)
.await
.map_redis_err()?)
.map_redis_err()
})
.await
}

2
crates/aether-data/src/repository/candidates/memory.rs
View File

@@ -388,7 +388,7 @@ impl RequestCandidateWriteRepository for InMemoryRequestCandidateRepository {
.filter(|row| row.created_at_unix_secs < created_before_unix_secs)
.map(|row| (row.created_at_unix_secs, row.id.clone()))
.collect::<Vec<_>>();
ids.sort_by(|left, right| left.cmp(right));
ids.sort();
let mut deleted = 0usize;
for (_, id) in ids.into_iter().take(limit) {

77
crates/aether-data/src/repository/global_models/memory.rs
View File

@@ -1,3 +1,4 @@
use std::collections::BTreeSet;
use std::sync::RwLock;
use async_trait::async_trait;
@@ -91,6 +92,35 @@ impl InMemoryGlobalModelReadRepository {
.expect("admin global model repository lock") = items.into_iter().collect();
self
}
fn admin_global_model_provider_counts(&self, global_model_id: &str) -> (u64, u64) {
let items = self
.admin_provider_model_items
.read()
.expect("admin provider model repository lock");
let provider_count = items
.iter()
.filter(|item| item.global_model_id == global_model_id)
.map(|item| item.provider_id.clone())
.collect::<BTreeSet<_>>()
.len() as u64;
let active_provider_count = items
.iter()
.filter(|item| item.global_model_id == global_model_id && item.is_active)
.map(|item| item.provider_id.clone())
.collect::<BTreeSet<_>>()
.len() as u64;
(provider_count, active_provider_count)
}
fn enrich_admin_global_model(&self, item: &StoredAdminGlobalModel) -> StoredAdminGlobalModel {
let mut enriched = item.clone();
let (provider_count, active_provider_count) =
self.admin_global_model_provider_counts(&item.id);
enriched.provider_count = provider_count;
enriched.active_provider_count = active_provider_count;
enriched
}
}
#[async_trait]
@@ -260,6 +290,7 @@ impl GlobalModelReadRepository for InMemoryGlobalModelReadRepository {
.into_iter()
.skip(query.offset)
.take(query.limit)
.map(|item| self.enrich_admin_global_model(&item))
.collect();
Ok(StoredAdminGlobalModelPage { items, total })
}
@@ -354,7 +385,7 @@ impl GlobalModelReadRepository for InMemoryGlobalModelReadRepository {
Ok(items
.iter()
.find(|item| item.id == global_model_id)
.cloned())
.map(|item| self.enrich_admin_global_model(item)))
}
async fn get_admin_global_model_by_name(
@@ -365,7 +396,10 @@ impl GlobalModelReadRepository for InMemoryGlobalModelReadRepository {
.admin_global_model_items
.read()
.expect("admin global model repository lock");
Ok(items.iter().find(|item| item.name == model_name).cloned())
Ok(items
.iter()
.find(|item| item.name == model_name)
.map(|item| self.enrich_admin_global_model(item)))
}
async fn list_admin_provider_models_by_global_model_id(
@@ -537,35 +571,40 @@ impl GlobalModelWriteRepository for InMemoryGlobalModelReadRepository {
record.default_tiered_pricing.clone(),
record.supported_capabilities.clone(),
record.config.clone(),
0,
0,
0,
Some(1_711_000_000),
Some(1_711_000_000),
)?;
self.admin_global_model_items
.write()
.expect("admin global model repository lock")
.push(stored.clone());
Ok(Some(stored))
.push(stored);
self.get_admin_global_model_by_id(&record.id).await
}
async fn update_admin_global_model(
&self,
record: &UpdateAdminGlobalModelRecord,
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
let mut items = self
.admin_global_model_items
.write()
.expect("admin global model repository lock");
let Some(existing) = items.iter_mut().find(|item| item.id == record.id) else {
return Ok(None);
};
existing.display_name = record.display_name.clone();
existing.is_active = record.is_active;
existing.default_price_per_request = record.default_price_per_request;
existing.default_tiered_pricing = record.default_tiered_pricing.clone();
existing.supported_capabilities = record.supported_capabilities.clone();
existing.config = record.config.clone();
existing.updated_at_unix_secs = Some(1_711_000_100);
Ok(Some(existing.clone()))
{
let mut items = self
.admin_global_model_items
.write()
.expect("admin global model repository lock");
let Some(existing) = items.iter_mut().find(|item| item.id == record.id) else {
return Ok(None);
};
existing.display_name = record.display_name.clone();
existing.is_active = record.is_active;
existing.default_price_per_request = record.default_price_per_request;
existing.default_tiered_pricing = record.default_tiered_pricing.clone();
existing.supported_capabilities = record.supported_capabilities.clone();
existing.config = record.config.clone();
existing.updated_at_unix_secs = Some(1_711_000_100);
}
self.get_admin_global_model_by_id(&record.id).await
}
async fn delete_admin_global_model(

146
crates/aether-data/src/repository/global_models/sql.rs
View File

@@ -104,22 +104,39 @@ LEFT JOIN global_models gm ON gm.id = m.global_model_id
const LIST_ADMIN_GLOBAL_MODELS_PREFIX: &str = r#"
SELECT
id,
name,
COALESCE(NULLIF(display_name, ''), name) AS display_name,
is_active,
CAST(default_price_per_request AS DOUBLE PRECISION) AS default_price_per_request,
default_tiered_pricing,
supported_capabilities,
config,
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
FROM global_models
gm.id,
gm.name,
COALESCE(NULLIF(gm.display_name, ''), gm.name) AS display_name,
gm.is_active,
CAST(gm.default_price_per_request AS DOUBLE PRECISION) AS default_price_per_request,
gm.default_tiered_pricing,
gm.supported_capabilities,
gm.config,
COALESCE(gm_stats.provider_count, 0) AS provider_count,
COALESCE(gm_stats.active_provider_count, 0) AS active_provider_count,
COALESCE(gm.usage_count, 0)::bigint AS usage_count,
EXTRACT(EPOCH FROM gm.created_at)::bigint AS created_at_unix_secs,
EXTRACT(EPOCH FROM gm.updated_at)::bigint AS updated_at_unix_secs
FROM global_models gm
LEFT JOIN (
SELECT
m.global_model_id,
COUNT(DISTINCT m.provider_id)::bigint AS provider_count,
COUNT(
DISTINCT CASE
WHEN m.is_active = TRUE AND p.is_active = TRUE THEN m.provider_id
ELSE NULL
END
)::bigint AS active_provider_count
FROM models m
JOIN providers p ON p.id = m.provider_id
GROUP BY m.global_model_id
) gm_stats ON gm_stats.global_model_id = gm.id
"#;
const COUNT_ADMIN_GLOBAL_MODELS_PREFIX: &str = r#"
SELECT COUNT(id) AS total
FROM global_models
FROM global_models gm
"#;
const LIST_ACTIVE_GLOBAL_MODEL_IDS_BY_PROVIDER_IDS_PREFIX: &str = r#"
@@ -374,19 +391,35 @@ ORDER BY gm.name ASC, m.created_at DESC, m.id ASC
let row = sqlx::query(
r#"
SELECT
id,
name,
COALESCE(NULLIF(display_name, ''), name) AS display_name,
is_active,
CAST(default_price_per_request AS DOUBLE PRECISION) AS default_price_per_request,
default_tiered_pricing,
supported_capabilities,
config
,
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
FROM global_models
WHERE id = $1
gm.id,
gm.name,
COALESCE(NULLIF(gm.display_name, ''), gm.name) AS display_name,
gm.is_active,
CAST(gm.default_price_per_request AS DOUBLE PRECISION) AS default_price_per_request,
gm.default_tiered_pricing,
gm.supported_capabilities,
gm.config,
COALESCE(gm_stats.provider_count, 0) AS provider_count,
COALESCE(gm_stats.active_provider_count, 0) AS active_provider_count,
COALESCE(gm.usage_count, 0)::bigint AS usage_count,
EXTRACT(EPOCH FROM gm.created_at)::bigint AS created_at_unix_secs,
EXTRACT(EPOCH FROM gm.updated_at)::bigint AS updated_at_unix_secs
FROM global_models gm
LEFT JOIN (
SELECT
m.global_model_id,
COUNT(DISTINCT m.provider_id)::bigint AS provider_count,
COUNT(
DISTINCT CASE
WHEN m.is_active = TRUE AND p.is_active = TRUE THEN m.provider_id
ELSE NULL
END
)::bigint AS active_provider_count
FROM models m
JOIN providers p ON p.id = m.provider_id
GROUP BY m.global_model_id
) gm_stats ON gm_stats.global_model_id = gm.id
WHERE gm.id = $1
LIMIT 1
"#,
)
@@ -405,19 +438,35 @@ LIMIT 1
let row = sqlx::query(
r#"
SELECT
id,
name,
COALESCE(NULLIF(display_name, ''), name) AS display_name,
is_active,
CAST(default_price_per_request AS DOUBLE PRECISION) AS default_price_per_request,
default_tiered_pricing,
supported_capabilities,
config
,
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_secs,
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
FROM global_models
WHERE name = $1
gm.id,
gm.name,
COALESCE(NULLIF(gm.display_name, ''), gm.name) AS display_name,
gm.is_active,
CAST(gm.default_price_per_request AS DOUBLE PRECISION) AS default_price_per_request,
gm.default_tiered_pricing,
gm.supported_capabilities,
gm.config,
COALESCE(gm_stats.provider_count, 0) AS provider_count,
COALESCE(gm_stats.active_provider_count, 0) AS active_provider_count,
COALESCE(gm.usage_count, 0)::bigint AS usage_count,
EXTRACT(EPOCH FROM gm.created_at)::bigint AS created_at_unix_secs,
EXTRACT(EPOCH FROM gm.updated_at)::bigint AS updated_at_unix_secs
FROM global_models gm
LEFT JOIN (
SELECT
m.global_model_id,
COUNT(DISTINCT m.provider_id)::bigint AS provider_count,
COUNT(
DISTINCT CASE
WHEN m.is_active = TRUE AND p.is_active = TRUE THEN m.provider_id
ELSE NULL
END
)::bigint AS active_provider_count
FROM models m
JOIN providers p ON p.id = m.provider_id
GROUP BY m.global_model_id
) gm_stats ON gm_stats.global_model_id = gm.id
WHERE gm.name = $1
LIMIT 1
"#,
)
@@ -928,7 +977,7 @@ fn apply_admin_global_model_filters(
) {
builder.push(" WHERE 1=1");
if let Some(is_active) = query.is_active {
builder.push(" AND is_active = ").push_bind(is_active);
builder.push(" AND gm.is_active = ").push_bind(is_active);
}
if let Some(search) = query
.search
@@ -938,9 +987,9 @@ fn apply_admin_global_model_filters(
{
let pattern = format!("%{search}%");
builder
.push(" AND (name ILIKE ")
.push(" AND (gm.name ILIKE ")
.push_bind(pattern.clone())
.push(" OR display_name ILIKE ")
.push(" OR gm.display_name ILIKE ")
.push_bind(pattern)
.push(")");
}
@@ -1062,6 +1111,18 @@ fn map_admin_global_model_row(row: &PgRow) -> Result<StoredAdminGlobalModel, Dat
.try_get::<Option<i64>, _>("updated_at_unix_secs")
.map_postgres_err()?
.map(|value| value.max(0) as u64);
let provider_count = row
.try_get::<i64, _>("provider_count")
.map_postgres_err()?
.max(0) as u64;
let active_provider_count = row
.try_get::<i64, _>("active_provider_count")
.map_postgres_err()?
.max(0) as u64;
let usage_count = row
.try_get::<i64, _>("usage_count")
.map_postgres_err()?
.max(0) as u64;
StoredAdminGlobalModel::new(
row.try_get("id").map_postgres_err()?,
row.try_get("name").map_postgres_err()?,
@@ -1072,6 +1133,9 @@ fn map_admin_global_model_row(row: &PgRow) -> Result<StoredAdminGlobalModel, Dat
row.try_get("default_tiered_pricing").map_postgres_err()?,
row.try_get("supported_capabilities").map_postgres_err()?,
row.try_get("config").map_postgres_err()?,
provider_count,
active_provider_count,
usage_count,
created_at_unix_secs,
updated_at_unix_secs,
)

61
crates/aether-data/src/repository/provider_catalog/sql.rs
View File

@@ -161,6 +161,8 @@ SELECT
utilization_samples,
EXTRACT(EPOCH FROM last_probe_increase_at)::bigint AS last_probe_increase_at_unix_secs,
request_count,
total_tokens,
CAST(total_cost_usd AS DOUBLE PRECISION) AS total_cost_usd,
success_count,
error_count,
total_response_time_ms,
@@ -214,6 +216,8 @@ SELECT
utilization_samples,
EXTRACT(EPOCH FROM last_probe_increase_at)::bigint AS last_probe_increase_at_unix_secs,
request_count,
total_tokens,
CAST(total_cost_usd AS DOUBLE PRECISION) AS total_cost_usd,
success_count,
error_count,
total_response_time_ms,
@@ -503,6 +507,8 @@ SELECT
utilization_samples,
EXTRACT(EPOCH FROM last_probe_increase_at)::bigint AS last_probe_increase_at_unix_secs,
request_count,
total_tokens,
CAST(total_cost_usd AS DOUBLE PRECISION) AS total_cost_usd,
success_count,
error_count,
total_response_time_ms,
@@ -1164,30 +1170,30 @@ INSERT INTO provider_api_keys (
ELSE TO_TIMESTAMP($35::double precision)
END,
COALESCE($36, 0),
0,
0,
COALESCE($37, 0),
COALESCE($38, 0),
COALESCE($39, 0),
COALESCE($40, 0),
COALESCE($41, 0),
CASE
WHEN $40::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($40::double precision)
WHEN $42::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($42::double precision)
END,
CASE
WHEN $41::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($41::double precision)
WHEN $43::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($43::double precision)
END,
$42,
$43,
$44,
$45,
$46,
$47,
CASE
WHEN $46::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($46::double precision)
WHEN $48::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($48::double precision)
END,
CASE
WHEN $47::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($47::double precision)
WHEN $49::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($49::double precision)
END
)
"#,
@@ -1231,6 +1237,13 @@ INSERT INTO provider_api_keys (
.map(|value| value as f64),
)
.bind(key.request_count.map(|value| value as i32))
.bind(Some(i64::try_from(key.total_tokens).map_err(|_| {
DataLayerError::InvalidInput(format!(
"provider catalog key.total_tokens exceeds i64: {}",
key.total_tokens
))
})?))
.bind(key.total_cost_usd)
.bind(key.success_count.map(|value| value as i32))
.bind(key.error_count.map(|value| value as i32))
.bind(key.total_response_time_ms.map(|value| value as i32))
@@ -2092,6 +2105,18 @@ fn map_key_row(row: &PgRow) -> Result<StoredProviderCatalogKey, DataLayerError>
})
})
.transpose()?;
let total_tokens = row_get::<Option<i64>>(row, "total_tokens")?
.unwrap_or(0)
.try_into()
.map_err(|_| {
DataLayerError::UnexpectedValue("invalid provider_api_keys.total_tokens".to_string())
})?;
let total_cost_usd = row_get::<Option<f64>>(row, "total_cost_usd")?.unwrap_or(0.0);
if !total_cost_usd.is_finite() {
return Err(DataLayerError::UnexpectedValue(
"invalid provider_api_keys.total_cost_usd".to_string(),
));
}
let success_count = row_get::<Option<i32>>(row, "success_count")?
.map(|value| {
u32::try_from(value).map_err(|_| {
@@ -2212,6 +2237,7 @@ fn map_key_row(row: &PgRow) -> Result<StoredProviderCatalogKey, DataLayerError>
success_count,
)
.with_usage_fields(error_count, total_response_time_ms)
.with_usage_totals(total_tokens, total_cost_usd)
.with_health_fields(
row.try_get("health_by_format").ok(),
row.try_get("circuit_breaker_by_format").ok(),
@@ -2263,4 +2289,15 @@ mod tests {
let repository = SqlxProviderCatalogReadRepository::new(pool);
let _ = repository.pool();
}
#[test]
fn key_queries_include_usage_totals() {
for sql in [
super::LIST_KEYS_BY_IDS_PREFIX,
super::LIST_KEYS_BY_PROVIDER_IDS_PREFIX,
] {
assert!(sql.contains("total_tokens"));
assert!(sql.contains("total_cost_usd"));
}
}
}

2
crates/aether-data/src/repository/wallet/types.rs
View File

@@ -459,6 +459,7 @@ pub struct CreateWalletRechargeOrderInput {
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
#[allow(clippy::large_enum_variant)]
pub enum CreateWalletRechargeOrderOutcome {
Created(StoredAdminPaymentOrder),
WalletInactive,
@@ -506,6 +507,7 @@ pub struct ProcessPaymentCallbackInput {
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
#[allow(clippy::large_enum_variant)]
pub enum ProcessPaymentCallbackOutcome {
DuplicateProcessed {
order_id: Option<String>,

25
crates/aether-model-fetch/src/logic.rs
View File

@@ -246,6 +246,12 @@ pub fn aggregate_models_for_cache(models: &[Value]) -> Vec<Value> {
.collect::<BTreeSet<_>>()
})
.unwrap_or_default();
let legacy_api_format = object
.get("api_format")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let existing_formats = entry
.get("api_formats")
.and_then(Value::as_array)
@@ -259,9 +265,15 @@ pub fn aggregate_models_for_cache(models: &[Value]) -> Vec<Value> {
.collect::<BTreeSet<_>>()
})
.unwrap_or_default();
let merged_formats = existing_formats
let mut merged_formats = existing_formats
.union(&api_formats)
.cloned()
.collect::<BTreeSet<_>>();
if let Some(api_format) = legacy_api_format {
merged_formats.insert(api_format);
}
let merged_formats = merged_formats
.into_iter()
.map(Value::String)
.collect::<Vec<_>>();
entry.insert("api_formats".to_string(), Value::Array(merged_formats));
@@ -453,6 +465,17 @@ mod tests {
);
}
#[test]
fn aggregate_models_for_cache_preserves_legacy_api_format_field() {
let aggregated = aggregate_models_for_cache(&[json!({
"id":"gpt-5",
"api_format":"openai:chat"
})]);
assert_eq!(aggregated.len(), 1);
assert_eq!(aggregated[0]["api_formats"], json!(["openai:chat"]));
assert!(aggregated[0].get("api_format").is_none());
}
#[test]
fn build_gemini_models_url_preserves_base_query() {
let url =

1
crates/aether-provider-transport/src/kiro.rs
View File

@@ -24,6 +24,7 @@ pub use request::{
apply_local_body_rules, apply_local_header_rules, body_rules_are_locally_supported,
build_kiro_provider_headers, build_kiro_provider_request_body,
header_rules_are_locally_supported, supports_local_kiro_request_shape,
KiroProviderHeadersInput,
};
pub use url::{
build_kiro_generate_assistant_response_url, resolve_kiro_base_url,

54
crates/aether-provider-transport/src/kiro/request.rs
View File

@@ -77,16 +77,32 @@ pub fn build_kiro_provider_request_body(
Some(provider_request_body)
}
#[derive(Clone, Copy, Debug)]
pub struct KiroProviderHeadersInput<'a> {
pub headers: &'a http::HeaderMap,
pub provider_request_body: &'a Value,
pub original_request_body: &'a Value,
pub header_rules: Option<&'a Value>,
pub auth_header: &'a str,
pub auth_value: &'a str,
pub auth_config: &'a KiroAuthConfig,
pub machine_id: &'a str,
}
pub fn build_kiro_provider_headers(
headers: &http::HeaderMap,
provider_request_body: &Value,
original_request_body: &Value,
header_rules: Option<&Value>,
auth_header: &str,
auth_value: &str,
auth_config: &KiroAuthConfig,
machine_id: &str,
input: KiroProviderHeadersInput<'_>,
) -> Option<BTreeMap<String, String>> {
let KiroProviderHeadersInput {
headers,
provider_request_body,
original_request_body,
header_rules,
auth_header,
auth_value,
auth_config,
machine_id,
} = input;
let mut out = BTreeMap::new();
for (name, value) in headers {
let Ok(value) = value.to_str() else {
@@ -133,7 +149,7 @@ mod tests {
use super::super::credentials::KiroAuthConfig;
use super::{
build_kiro_provider_headers, build_kiro_provider_request_body,
supports_local_kiro_request_shape,
supports_local_kiro_request_shape, KiroProviderHeadersInput,
};
#[test]
@@ -221,19 +237,19 @@ mod tests {
node_version: None,
access_token: Some("cached-token".to_string()),
};
let headers = build_kiro_provider_headers(
&http::HeaderMap::new(),
&json!({"conversationState": {}}),
&json!({"messages": []}),
Some(&json!([
let headers = build_kiro_provider_headers(KiroProviderHeadersInput {
headers: &http::HeaderMap::new(),
provider_request_body: &json!({"conversationState": {}}),
original_request_body: &json!({"messages": []}),
header_rules: Some(&json!([
{"action":"set","key":"accept","value":"text/plain"},
{"action":"set","key":"x-endpoint-tag","value":"kiro-local"}
])),
"authorization",
"Bearer cached-token",
&auth_config,
"machine-123",
)
auth_header: "authorization",
auth_value: "Bearer cached-token",
auth_config: &auth_config,
machine_id: "machine-123",
})
.expect("headers should build");
assert_eq!(

9
crates/aether-provider-transport/src/network.rs
View File

@@ -153,19 +153,14 @@ pub fn resolve_transport_tls_profile(
}
fn effective_proxy_config(transport: &GatewayProviderTransportSnapshot) -> Option<&Value> {
for candidate in [
[
transport.key.proxy.as_ref(),
transport.endpoint.proxy.as_ref(),
transport.provider.proxy.as_ref(),
]
.into_iter()
.flatten()
{
if proxy_enabled(candidate) {
return Some(candidate);
}
}
None
.find(|candidate| proxy_enabled(candidate))
}
fn proxy_enabled(value: &Value) -> bool {

1
crates/aether-provider-transport/src/oauth_refresh.rs
View File

@@ -16,6 +16,7 @@ use super::kiro::{
use super::snapshot::GatewayProviderTransportSnapshot;
#[derive(Debug, Clone, PartialEq, Eq)]
#[allow(clippy::large_enum_variant)]
pub enum LocalResolvedOAuthRequestAuth {
#[allow(dead_code)]
Header {

88
crates/aether-scheduler-core/src/candidate.rs
View File

@@ -258,16 +258,32 @@ pub fn reorder_candidates_by_scheduler_health(
});
}
#[derive(Clone, Copy, Debug)]
pub struct CandidateRuntimeSelectabilityInput<'a> {
pub candidate: &'a SchedulerMinimalCandidateSelectionCandidate,
pub recent_candidates: &'a [StoredRequestCandidate],
pub provider_concurrent_limits: &'a BTreeMap<String, usize>,
pub provider_key_rpm_states: &'a BTreeMap<String, StoredProviderCatalogKey>,
pub now_unix_secs: u64,
pub cached_affinity_target: Option<&'a crate::SchedulerAffinityTarget>,
pub provider_quota_blocks_requests: bool,
pub rpm_reset_at: Option<u64>,
}
pub fn candidate_is_selectable_with_runtime_state(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
recent_candidates: &[StoredRequestCandidate],
provider_concurrent_limits: &BTreeMap<String, usize>,
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
now_unix_secs: u64,
cached_affinity_target: Option<&crate::SchedulerAffinityTarget>,
provider_quota_blocks_requests: bool,
rpm_reset_at: Option<u64>,
input: CandidateRuntimeSelectabilityInput<'_>,
) -> bool {
let CandidateRuntimeSelectabilityInput {
candidate,
recent_candidates,
provider_concurrent_limits,
provider_key_rpm_states,
now_unix_secs,
cached_affinity_target,
provider_quota_blocks_requests,
rpm_reset_at,
} = input;
if provider_quota_blocks_requests {
return false;
}
@@ -375,7 +391,7 @@ mod tests {
candidate_is_selectable_with_runtime_state, candidate_supports_required_capability,
collect_global_model_names_for_required_capability,
collect_selectable_candidates_from_keys, reorder_candidates_by_scheduler_health,
SchedulerMinimalCandidateSelectionCandidate,
CandidateRuntimeSelectabilityInput, SchedulerMinimalCandidateSelectionCandidate,
};
use crate::SchedulerAuthConstraints;
@@ -627,14 +643,16 @@ mod tests {
let provider_concurrent_limits = BTreeMap::from([("provider-1".to_string(), 1)]);
assert!(!candidate_is_selectable_with_runtime_state(
&sample_candidate("1", None),
&recent_candidates,
&provider_concurrent_limits,
&BTreeMap::new(),
100,
None,
false,
None,
CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &recent_candidates,
provider_concurrent_limits: &provider_concurrent_limits,
provider_key_rpm_states: &BTreeMap::new(),
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
rpm_reset_at: None,
},
));
}
@@ -643,24 +661,28 @@ mod tests {
let provider_key_rpm_states = BTreeMap::from([("key-1".to_string(), sample_key("1", 0.0))]);
assert!(!candidate_is_selectable_with_runtime_state(
&sample_candidate("1", None),
&[],
&BTreeMap::new(),
&provider_key_rpm_states,
100,
None,
false,
None,
CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &[],
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &provider_key_rpm_states,
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
rpm_reset_at: None,
},
));
assert!(!candidate_is_selectable_with_runtime_state(
&sample_candidate("1", None),
&[],
&BTreeMap::new(),
&BTreeMap::new(),
100,
None,
true,
None,
CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &[],
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &BTreeMap::new(),
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: true,
rpm_reset_at: None,
},
));
}

6
crates/aether-scheduler-core/src/lib.rs
View File

@@ -18,7 +18,8 @@ pub use candidate::{
auth_api_key_concurrency_limit_reached, build_minimal_candidate_selection,
candidate_is_selectable_with_runtime_state, candidate_supports_required_capability,
collect_global_model_names_for_required_capability, collect_selectable_candidates_from_keys,
reorder_candidates_by_scheduler_health, SchedulerMinimalCandidateSelectionCandidate,
reorder_candidates_by_scheduler_health, CandidateRuntimeSelectabilityInput,
SchedulerMinimalCandidateSelectionCandidate,
};
pub use health::{
aggregate_provider_key_health_score, count_recent_active_requests_for_api_key,
@@ -40,6 +41,7 @@ pub use request_candidate::{
build_report_request_candidate_status_record, execution_error_details,
finalize_execution_request_candidate_report_context, is_terminal_candidate_status,
parse_request_candidate_report_context, resolve_report_request_candidate_slot,
LocalRequestCandidateStatusRecordInput, ReportRequestCandidateStatusRecordInput,
SchedulerExecutionRequestCandidateSeed, SchedulerRequestCandidateReportContext,
SchedulerResolvedReportRequestCandidateSlot,
SchedulerRequestCandidateStatusUpdate, SchedulerResolvedReportRequestCandidateSlot,
};

4
crates/aether-scheduler-core/src/model.rs
View File

@@ -90,9 +90,7 @@ pub fn resolve_provider_model_name(
}
}
let Some(global_model_mappings) = row.global_model_mappings.as_ref() else {
return None;
};
let global_model_mappings = row.global_model_mappings.as_ref()?;
for allowed_model in sorted_allowed_models {
for pattern in global_model_mappings {
if matches_model_mapping(pattern, &allowed_model) {

168
crates/aether-scheduler-core/src/request_candidate.rs
View File

@@ -41,6 +41,31 @@ pub struct SchedulerExecutionRequestCandidateSeed {
pub report_context: Value,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SchedulerRequestCandidateStatusUpdate {
pub status: RequestCandidateStatus,
pub status_code: Option<u16>,
pub error_type: Option<String>,
pub error_message: Option<String>,
pub latency_ms: Option<u64>,
pub started_at_unix_secs: Option<u64>,
pub finished_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone)]
pub struct LocalRequestCandidateStatusRecordInput<'a> {
pub plan: &'a ExecutionPlan,
pub report_context: Option<&'a Value>,
pub status_update: SchedulerRequestCandidateStatusUpdate,
}
#[derive(Debug, Clone)]
pub struct ReportRequestCandidateStatusRecordInput {
pub slot: SchedulerResolvedReportRequestCandidateSlot,
pub status_update: SchedulerRequestCandidateStatusUpdate,
pub now_unix_secs: u64,
}
pub fn execution_error_details(
error: Option<&ExecutionError>,
body_json: Option<&Value>,
@@ -231,16 +256,23 @@ pub fn build_execution_request_candidate_seed(
}
pub fn build_local_request_candidate_status_record(
plan: &ExecutionPlan,
report_context: Option<&Value>,
status: RequestCandidateStatus,
status_code: Option<u16>,
error_type: Option<String>,
error_message: Option<String>,
latency_ms: Option<u64>,
started_at_unix_secs: Option<u64>,
finished_at_unix_secs: Option<u64>,
input: LocalRequestCandidateStatusRecordInput<'_>,
) -> Option<UpsertRequestCandidateRecord> {
let LocalRequestCandidateStatusRecordInput {
plan,
report_context,
status_update,
} = input;
let SchedulerRequestCandidateStatusUpdate {
status,
status_code,
error_type,
error_message,
latency_ms,
started_at_unix_secs,
finished_at_unix_secs,
} = status_update;
let candidate_id = plan
.candidate_id
.as_deref()
@@ -278,16 +310,23 @@ pub fn build_local_request_candidate_status_record(
}
pub fn build_report_request_candidate_status_record(
slot: SchedulerResolvedReportRequestCandidateSlot,
status: RequestCandidateStatus,
status_code: Option<u16>,
error_type: Option<String>,
error_message: Option<String>,
latency_ms: Option<u64>,
started_at_unix_secs: Option<u64>,
finished_at_unix_secs: Option<u64>,
now_unix_secs: u64,
input: ReportRequestCandidateStatusRecordInput,
) -> UpsertRequestCandidateRecord {
let ReportRequestCandidateStatusRecordInput {
slot,
status_update,
now_unix_secs,
} = input;
let SchedulerRequestCandidateStatusUpdate {
status,
status_code,
error_type,
error_message,
latency_ms,
started_at_unix_secs,
finished_at_unix_secs,
} = status_update;
let terminal_unix_secs = finished_at_unix_secs.unwrap_or(now_unix_secs);
let started_at_unix_secs = started_at_unix_secs
.or(slot.started_at_unix_secs)
@@ -454,7 +493,8 @@ mod tests {
build_report_request_candidate_status_record, execution_error_details,
finalize_execution_request_candidate_report_context,
parse_request_candidate_report_context, resolve_report_request_candidate_slot,
SchedulerResolvedReportRequestCandidateSlot,
LocalRequestCandidateStatusRecordInput, ReportRequestCandidateStatusRecordInput,
SchedulerRequestCandidateStatusUpdate, SchedulerResolvedReportRequestCandidateSlot,
};
fn sample_candidate(
@@ -605,23 +645,26 @@ mod tests {
let mut plan = sample_plan();
plan.candidate_id = Some("cand-1".to_string());
let record = build_local_request_candidate_status_record(
&plan,
Some(&json!({
"candidate_index": 1,
"retry_index": 2,
"user_id": "user-1",
"api_key_id": "api-key-1"
})),
RequestCandidateStatus::Failed,
Some(500),
Some("Upstream5xx".to_string()),
Some("boom".to_string()),
Some(42),
Some(100),
Some(101),
)
.expect("record should build");
let record =
build_local_request_candidate_status_record(LocalRequestCandidateStatusRecordInput {
plan: &plan,
report_context: Some(&json!({
"candidate_index": 1,
"retry_index": 2,
"user_id": "user-1",
"api_key_id": "api-key-1"
})),
status_update: SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(500),
error_type: Some("Upstream5xx".to_string()),
error_message: Some("boom".to_string()),
latency_ms: Some(42),
started_at_unix_secs: Some(100),
finished_at_unix_secs: Some(101),
},
})
.expect("record should build");
assert_eq!(record.id, "cand-1");
assert_eq!(record.candidate_index, 1);
@@ -632,31 +675,34 @@ mod tests {
#[test]
fn builds_report_request_candidate_status_record_with_terminal_timestamps() {
let record = build_report_request_candidate_status_record(
SchedulerResolvedReportRequestCandidateSlot {
id: "cand-1".to_string(),
request_id: "req-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("api-key-1".to_string()),
candidate_index: 1,
retry_index: 0,
provider_id: Some("provider-1".to_string()),
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("key-1".to_string()),
extra_data: None,
created_at_unix_secs: 10,
started_at_unix_secs: None,
finished_at_unix_secs: None,
},
RequestCandidateStatus::Success,
Some(200),
None,
None,
Some(12),
None,
None,
123,
);
let record =
build_report_request_candidate_status_record(ReportRequestCandidateStatusRecordInput {
slot: SchedulerResolvedReportRequestCandidateSlot {
id: "cand-1".to_string(),
request_id: "req-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("api-key-1".to_string()),
candidate_index: 1,
retry_index: 0,
provider_id: Some("provider-1".to_string()),
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("key-1".to_string()),
extra_data: None,
created_at_unix_secs: 10,
started_at_unix_secs: None,
finished_at_unix_secs: None,
},
status_update: SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Success,
status_code: Some(200),
error_type: None,
error_message: None,
latency_ms: Some(12),
started_at_unix_secs: None,
finished_at_unix_secs: None,
},
now_unix_secs: 123,
});
assert_eq!(record.started_at_unix_secs, Some(123));
assert_eq!(record.finished_at_unix_secs, Some(123));

61
crates/aether-usage-runtime/src/write.rs
View File

@@ -58,6 +58,18 @@ pub struct TerminalUsageOutcome {
pub audit_payload: Option<Value>,
}
struct TerminalUsageOutcomeBaseInput<'a> {
plan: &'a ExecutionPlan,
report_context: Option<&'a Value>,
terminal_state: UsageTerminalState,
status_code: u16,
telemetry: Option<&'a ExecutionTelemetry>,
provider_response: Option<Value>,
client_response: Option<Value>,
provider_response_headers: Option<Value>,
client_response_headers: Option<Value>,
}
pub fn build_pending_usage_record(
plan: &ExecutionPlan,
report_context: Option<&Value>,
@@ -132,17 +144,17 @@ pub fn build_sync_terminal_usage_outcome(
.clone()
.or_else(|| decode_body_for_storage(payload.body_base64.as_deref()));
let client_response = payload.client_body_json.clone();
build_terminal_usage_outcome_base(
build_terminal_usage_outcome_base(TerminalUsageOutcomeBaseInput {
plan,
report_context,
infer_sync_terminal_state(payload, provider_response.as_ref()),
payload.status_code,
payload.telemetry.as_ref(),
terminal_state: infer_sync_terminal_state(payload, provider_response.as_ref()),
status_code: payload.status_code,
telemetry: payload.telemetry.as_ref(),
provider_response,
client_response,
Some(headers_to_json(&payload.headers)),
Some(headers_to_json(&payload.headers)),
)
provider_response_headers: Some(headers_to_json(&payload.headers)),
client_response_headers: Some(headers_to_json(&payload.headers)),
})
}
pub fn build_stream_terminal_usage_outcome(
@@ -152,17 +164,17 @@ pub fn build_stream_terminal_usage_outcome(
) -> TerminalUsageOutcome {
let provider_response = decode_body_for_storage(payload.provider_body_base64.as_deref());
let client_response = decode_body_for_storage(payload.client_body_base64.as_deref());
build_terminal_usage_outcome_base(
build_terminal_usage_outcome_base(TerminalUsageOutcomeBaseInput {
plan,
report_context,
infer_stream_terminal_state(payload),
payload.status_code,
payload.telemetry.as_ref(),
terminal_state: infer_stream_terminal_state(payload),
status_code: payload.status_code,
telemetry: payload.telemetry.as_ref(),
provider_response,
client_response,
Some(headers_to_json(&payload.headers)),
Some(headers_to_json(&payload.headers)),
)
provider_response_headers: Some(headers_to_json(&payload.headers)),
client_response_headers: Some(headers_to_json(&payload.headers)),
})
}
pub fn build_terminal_usage_event_from_outcome(
@@ -239,16 +251,19 @@ pub fn build_terminal_usage_event_from_outcome(
}
fn build_terminal_usage_outcome_base(
plan: &ExecutionPlan,
report_context: Option<&Value>,
terminal_state: UsageTerminalState,
status_code: u16,
telemetry: Option<&ExecutionTelemetry>,
provider_response: Option<Value>,
client_response: Option<Value>,
provider_response_headers: Option<Value>,
client_response_headers: Option<Value>,
input: TerminalUsageOutcomeBaseInput<'_>,
) -> TerminalUsageOutcome {
let TerminalUsageOutcomeBaseInput {
plan,
report_context,
terminal_state,
status_code,
telemetry,
provider_response,
client_response,
provider_response_headers,
client_response_headers,
} = input;
let context = report_context.and_then(Value::as_object);
let client_contract = context_string(context, "client_contract")
.or_else(|| context_string(context, "client_api_format"))

74
crates/aether-video-tasks-core/src/follow_up.rs
View File

@@ -1,18 +1,35 @@
use serde_json::{Map, Value};
pub fn build_video_follow_up_report_context(
request_id: &str,
user_id: &str,
api_key_id: &str,
task_id: &str,
provider_id: &str,
endpoint_id: &str,
key_id: &str,
provider_name: Option<&str>,
model_name: Option<&str>,
client_api_format: &str,
provider_api_format: &str,
) -> Value {
#[derive(Clone, Copy, Debug)]
pub struct VideoFollowUpReportContextInput<'a> {
pub request_id: &'a str,
pub user_id: &'a str,
pub api_key_id: &'a str,
pub task_id: &'a str,
pub provider_id: &'a str,
pub endpoint_id: &'a str,
pub key_id: &'a str,
pub provider_name: Option<&'a str>,
pub model_name: Option<&'a str>,
pub client_api_format: &'a str,
pub provider_api_format: &'a str,
}
pub fn build_video_follow_up_report_context(input: VideoFollowUpReportContextInput<'_>) -> Value {
let VideoFollowUpReportContextInput {
request_id,
user_id,
api_key_id,
task_id,
provider_id,
endpoint_id,
key_id,
provider_name,
model_name,
client_api_format,
provider_api_format,
} = input;
let mut context = Map::new();
context.insert(
"request_id".to_string(),
@@ -84,23 +101,26 @@ pub fn resolve_follow_up_auth(
#[cfg(test)]
mod tests {
use super::{build_video_follow_up_report_context, resolve_follow_up_auth};
use super::{
build_video_follow_up_report_context, resolve_follow_up_auth,
VideoFollowUpReportContextInput,
};
#[test]
fn builds_follow_up_report_context_with_transport_metadata() {
let context = build_video_follow_up_report_context(
"req_123",
"user_123",
"key_123",
"task_123",
"provider_123",
"endpoint_123",
"transport_key_123",
Some("provider-name"),
Some("model-name"),
"openai:video",
"openai:video",
);
let context = build_video_follow_up_report_context(VideoFollowUpReportContextInput {
request_id: "req_123",
user_id: "user_123",
api_key_id: "key_123",
task_id: "task_123",
provider_id: "provider_123",
endpoint_id: "endpoint_123",
key_id: "transport_key_123",
provider_name: Some("provider-name"),
model_name: Some("model-name"),
client_api_format: "openai:video",
provider_api_format: "openai:video",
});
assert_eq!(context["request_id"].as_str(), Some("req_123"));
assert_eq!(context["provider_id"].as_str(), Some("provider_123"));

26
crates/aether-video-tasks-core/src/gemini.rs
View File

@@ -8,7 +8,7 @@ use crate::{
build_video_follow_up_report_context, current_unix_timestamp_secs, gemini_metadata_video_url,
request_body_string, request_body_u32, resolve_follow_up_auth, GeminiVideoTaskSeed,
LocalVideoTaskFollowUpPlan, LocalVideoTaskReadResponse, LocalVideoTaskSnapshot,
LocalVideoTaskStatus, DEFAULT_VIDEO_TASK_MAX_POLL_COUNT,
LocalVideoTaskStatus, VideoFollowUpReportContextInput, DEFAULT_VIDEO_TASK_MAX_POLL_COUNT,
DEFAULT_VIDEO_TASK_POLL_INTERVAL_SECONDS,
};
@@ -233,17 +233,19 @@ impl GeminiVideoTaskSeed {
},
report_kind: Some("gemini_video_cancel_sync_finalize".to_string()),
report_context: Some(build_video_follow_up_report_context(
&self.persistence.request_id,
&user_id,
&api_key_id,
&self.local_short_id,
&self.transport.provider_id,
&self.transport.endpoint_id,
&self.transport.key_id,
self.transport.provider_name.as_deref(),
Some(self.model.as_str()),
"gemini:video",
"gemini:video",
VideoFollowUpReportContextInput {
request_id: &self.persistence.request_id,
user_id: &user_id,
api_key_id: &api_key_id,
task_id: &self.local_short_id,
provider_id: &self.transport.provider_id,
endpoint_id: &self.transport.endpoint_id,
key_id: &self.transport.key_id,
provider_name: self.transport.provider_name.as_deref(),
model_name: Some(self.model.as_str()),
client_api_format: "gemini:video",
provider_api_format: "gemini:video",
},
)),
})
}

4
crates/aether-video-tasks-core/src/lib.rs
View File

@@ -18,7 +18,9 @@ mod util;
pub use body::{
context_text, context_u64, request_body_string, request_body_text, request_body_u32,
};
pub use follow_up::{build_video_follow_up_report_context, resolve_follow_up_auth};
pub use follow_up::{
build_video_follow_up_report_context, resolve_follow_up_auth, VideoFollowUpReportContextInput,
};
pub use gemini::map_gemini_stored_task_to_read_response;
pub use openai::map_openai_stored_task_to_read_response;
pub use path::{

134
crates/aether-video-tasks-core/src/openai.rs
View File

@@ -11,7 +11,8 @@ use crate::{
parse_video_content_variant, request_body_string, request_body_u32, resolve_follow_up_auth,
LocalVideoTaskContentAction, LocalVideoTaskFollowUpPlan, LocalVideoTaskReadResponse,
LocalVideoTaskSnapshot, LocalVideoTaskStatus, OpenAiVideoTaskSeed,
DEFAULT_VIDEO_TASK_MAX_POLL_COUNT, DEFAULT_VIDEO_TASK_POLL_INTERVAL_SECONDS,
VideoFollowUpReportContextInput, DEFAULT_VIDEO_TASK_MAX_POLL_COUNT,
DEFAULT_VIDEO_TASK_POLL_INTERVAL_SECONDS,
};
pub fn map_openai_stored_task_to_read_response(
@@ -208,34 +209,36 @@ impl OpenAiVideoTaskSeed {
)
};
Some(LocalVideoTaskContentAction::StreamPlan(ExecutionPlan {
request_id: trace_id.to_string(),
candidate_id: None,
provider_name: self.transport.provider_name.clone(),
provider_id: self.transport.provider_id.clone(),
endpoint_id: self.transport.endpoint_id.clone(),
key_id: self.transport.key_id.clone(),
method: "GET".to_string(),
url,
headers,
content_type: None,
content_encoding: None,
body: RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
Some(LocalVideoTaskContentAction::StreamPlan(Box::new(
ExecutionPlan {
request_id: trace_id.to_string(),
candidate_id: None,
provider_name: self.transport.provider_name.clone(),
provider_id: self.transport.provider_id.clone(),
endpoint_id: self.transport.endpoint_id.clone(),
key_id: self.transport.key_id.clone(),
method: "GET".to_string(),
url,
headers,
content_type: None,
content_encoding: None,
body: RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
},
stream: true,
client_api_format: "openai:video".to_string(),
provider_api_format: "openai:video".to_string(),
model_name: self
.model
.clone()
.or_else(|| self.transport.model_name.clone()),
proxy: self.transport.proxy.clone(),
tls_profile: self.transport.tls_profile.clone(),
timeouts: self.transport.timeouts.clone(),
},
stream: true,
client_api_format: "openai:video".to_string(),
provider_api_format: "openai:video".to_string(),
model_name: self
.model
.clone()
.or_else(|| self.transport.model_name.clone()),
proxy: self.transport.proxy.clone(),
tls_profile: self.transport.tls_profile.clone(),
timeouts: self.transport.timeouts.clone(),
}))
)))
}
pub fn client_body_json(&self) -> Value {
@@ -342,17 +345,19 @@ impl OpenAiVideoTaskSeed {
},
report_kind: Some("openai_video_delete_sync_finalize".to_string()),
report_context: Some(build_video_follow_up_report_context(
&self.persistence.request_id,
&user_id,
&api_key_id,
&self.local_task_id,
&self.transport.provider_id,
&self.transport.endpoint_id,
&self.transport.key_id,
self.transport.provider_name.as_deref(),
model_name.as_deref(),
"openai:video",
"openai:video",
VideoFollowUpReportContextInput {
request_id: &self.persistence.request_id,
user_id: &user_id,
api_key_id: &api_key_id,
task_id: &self.local_task_id,
provider_id: &self.transport.provider_id,
endpoint_id: &self.transport.endpoint_id,
key_id: &self.transport.key_id,
provider_name: self.transport.provider_name.as_deref(),
model_name: model_name.as_deref(),
client_api_format: "openai:video",
provider_api_format: "openai:video",
},
)),
})
}
@@ -466,17 +471,19 @@ impl OpenAiVideoTaskSeed {
},
report_kind: Some("openai_video_cancel_sync_finalize".to_string()),
report_context: Some(build_video_follow_up_report_context(
&self.persistence.request_id,
&user_id,
&api_key_id,
&self.local_task_id,
&self.transport.provider_id,
&self.transport.endpoint_id,
&self.transport.key_id,
self.transport.provider_name.as_deref(),
model_name.as_deref(),
"openai:video",
"openai:video",
VideoFollowUpReportContextInput {
request_id: &self.persistence.request_id,
user_id: &user_id,
api_key_id: &api_key_id,
task_id: &self.local_task_id,
provider_id: &self.transport.provider_id,
endpoint_id: &self.transport.endpoint_id,
key_id: &self.transport.key_id,
provider_name: self.transport.provider_name.as_deref(),
model_name: model_name.as_deref(),
client_api_format: "openai:video",
provider_api_format: "openai:video",
},
)),
})
}
@@ -513,19 +520,20 @@ impl OpenAiVideoTaskSeed {
.entry("content-type".to_string())
.or_insert_with(|| content_type.clone());
let mut report_context = build_video_follow_up_report_context(
&self.persistence.request_id,
&user_id,
&api_key_id,
&self.local_task_id,
&self.transport.provider_id,
&self.transport.endpoint_id,
&self.transport.key_id,
self.transport.provider_name.as_deref(),
model_name.as_deref(),
"openai:video",
"openai:video",
);
let mut report_context =
build_video_follow_up_report_context(VideoFollowUpReportContextInput {
request_id: &self.persistence.request_id,
user_id: &user_id,
api_key_id: &api_key_id,
task_id: &self.local_task_id,
provider_id: &self.transport.provider_id,
endpoint_id: &self.transport.endpoint_id,
key_id: &self.transport.key_id,
provider_name: self.transport.provider_name.as_deref(),
model_name: model_name.as_deref(),
client_api_format: "openai:video",
provider_api_format: "openai:video",
});
if let Some(report_context_object) = report_context.as_object_mut() {
report_context_object.insert("original_request_body".to_string(), body_json.clone());
}

2
crates/aether-video-tasks-core/src/types.rs
View File

@@ -42,7 +42,7 @@ pub struct LocalVideoTaskReadRefreshPlan {
#[derive(Debug, Clone, PartialEq)]
pub enum LocalVideoTaskContentAction {
Immediate { status_code: u16, body_json: Value },
StreamPlan(ExecutionPlan),
StreamPlan(Box<ExecutionPlan>),
}
#[derive(Debug, Clone, PartialEq, Eq)]