feat: provider api_formats 可空继承、OpenAI 图片 edit/variation 与用量配额多项补强

- 鉴权: provider_api_keys.api_formats 改为可空,OAuth 托管 key 自动继承 provider endpoints 激活格式,相关 handler/测试同步更新
- 图片 planner: OpenAI 图片路由新增 edit/variation 操作并完善参数校验、响应合并与流式处理
- 用量: user me usage 返回区分 client_requested_stream/upstream_is_stream,前端 usage 列表筛选与展示增强
- 统计: stats_daily_model 新增 cache_creation_ephemeral_5m/1h tokens 字段与回填链路
- 配额/observability: quota repository 新增内存与 SQL 扩展,admin observability usage 字段扩充
- 其它: OAuth 导入/轮询收敛、provider 汇总与 pool admin 读写链路小修、新增 system_config 缓存与 provider template handler

Closes #318

Co-authored-by: Entropy.Xu <53283266+Entropy-Xu@users.noreply.github.com>
This commit is contained in:
fawney19
2026-04-23 14:42:51 +08:00
parent f55f22d2e8
commit fa328e18a1
128 changed files with 5583 additions and 690 deletions

View File

@@ -17,9 +17,61 @@ const GEMINI_VIDEO_TASK_NOT_FOUND_DETAIL: &str = "Video task not found";
const AI_PUBLIC_METHOD_NOT_ALLOWED_DETAIL: &str = "Method not allowed";
const AI_PUBLIC_UNAUTHORIZED_DETAIL: &str = "Unauthorized";
const OPENAI_CHAT_IMAGE_MODEL_DETAIL: &str =
"gpt-image-2 仅支持通过 /v1/images/generations 或 /v1/images/edits 调用";
const OPENAI_IMAGE_MODEL_DETAIL: &str = "图片接口当前仅支持模型 gpt-image-2";
const OPENAI_IMAGE_N_DETAIL: &str = "图片接口当前仅支持 n=1";
"图片模型仅支持通过 /v1/images/generations、/v1/images/edits 或 /v1/images/variations 调用";
const OPENAI_IMAGE_PROMPT_DETAIL: &str = "图片生成/编辑请求缺少 prompt";
const OPENAI_IMAGE_EDIT_INPUT_DETAIL: &str = "图片编辑请求至少需要 1 张输入图片";
const OPENAI_IMAGE_VARIATION_INPUT_DETAIL: &str = "图片变体请求需要 image 文件";
const OPENAI_IMAGE_N_DETAIL: &str = "当前 Codex 图片反代仅支持 n=1";
const OPENAI_IMAGE_STREAM_VARIATION_DETAIL: &str = "图片变体接口当前仅支持同步响应";
const OPENAI_IMAGE_STREAM_MODEL_DETAIL: &str = "stream/partial_images 仅支持 GPT Image 系列模型";
const OPENAI_IMAGE_PARTIAL_IMAGES_DETAIL: &str =
"partial_images 仅支持 0-3且必须配合 stream=true";
const OPENAI_IMAGE_STYLE_DETAIL: &str = "当前 Codex 图片反代暂不支持 style 参数";
const OPENAI_IMAGE_RESPONSE_FORMAT_DETAIL: &str = "response_format 仅支持 url 或 b64_json";
const OPENAI_IMAGE_OUTPUT_FORMAT_DETAIL: &str = "output_format 仅支持 png、jpeg 或 webp";
const OPENAI_IMAGE_QUALITY_DETAIL: &str = "quality 仅支持 low、medium、high、standard 或 hd";
const OPENAI_IMAGE_BACKGROUND_DETAIL: &str = "background 仅支持 auto、opaque 或 transparent";
const OPENAI_IMAGE_MODERATION_DETAIL: &str = "moderation 仅支持 auto 或 low";
const OPENAI_IMAGE_INPUT_FIDELITY_DETAIL: &str = "input_fidelity 仅支持 low 或 high";
const OPENAI_IMAGE_OUTPUT_COMPRESSION_DETAIL: &str = "output_compression 必须是 0-100 的整数";
const OPENAI_IMAGE_INVALID_JSON_DETAIL: &str = "图片接口 JSON 请求体无效";
const OPENAI_IMAGE_INVALID_MULTIPART_DETAIL: &str = "图片接口 multipart/form-data 请求体无效";
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum OpenAiImageOperation {
Generate,
Edit,
Variation,
}
impl OpenAiImageOperation {
fn from_path(path: &str) -> Option<Self> {
match path {
"/v1/images/generations" => Some(Self::Generate),
"/v1/images/edits" => Some(Self::Edit),
"/v1/images/variations" => Some(Self::Variation),
_ => None,
}
}
}
#[derive(Debug, Default)]
struct OpenAiImageValidationInput {
model: Option<String>,
prompt: Option<String>,
image_count: usize,
n: Option<u64>,
stream: bool,
partial_images: Option<u64>,
response_format: Option<String>,
output_format: Option<String>,
quality: Option<String>,
background: Option<String>,
moderation: Option<String>,
input_fidelity: Option<String>,
output_compression: Option<u64>,
style_present: bool,
}
pub(crate) fn ai_public_local_requires_buffered_body(
request_context: &GatewayPublicRequestContext,
@@ -76,17 +128,13 @@ fn maybe_build_local_openai_request_validation_response(
}
let request_body = request_body?;
let payload = serde_json::from_slice::<Value>(request_body).ok()?;
if decision.route_kind.as_deref() == Some("chat")
&& request_context.request_path == "/v1/chat/completions"
{
let model = payload
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
if model.eq_ignore_ascii_case("gpt-image-2") {
let payload = serde_json::from_slice::<Value>(request_body).ok()?;
let model = payload.get("model").and_then(Value::as_str)?;
if is_openai_image_model(model) {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_CHAT_IMAGE_MODEL_DETAIL,
@@ -98,33 +146,179 @@ fn maybe_build_local_openai_request_validation_response(
if decision.route_kind.as_deref() != Some("image")
|| !matches!(
request_context.request_path.as_str(),
"/v1/images/generations" | "/v1/images/edits"
"/v1/images/generations" | "/v1/images/edits" | "/v1/images/variations"
)
{
return None;
}
if let Some(model) = payload
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
if !model.eq_ignore_ascii_case("gpt-image-2") {
let Some(operation) = OpenAiImageOperation::from_path(&request_context.request_path) else {
return None;
};
let validation = match parse_openai_image_validation_input(
operation,
request_context.request_content_type.as_deref(),
request_body,
) {
Ok(validation) => validation,
Err(detail) => {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_MODEL_DETAIL,
detail,
));
}
};
if validation
.model
.as_deref()
.is_some_and(|model| !image_model_supported_for_operation(operation, model))
{
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
format!(
"该接口不支持模型 {}",
validation.model.as_deref().unwrap_or_default()
),
));
}
match operation {
OpenAiImageOperation::Generate | OpenAiImageOperation::Edit
if validation.prompt.is_none() =>
{
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_PROMPT_DETAIL,
));
}
OpenAiImageOperation::Edit if validation.image_count == 0 => {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_EDIT_INPUT_DETAIL,
));
}
OpenAiImageOperation::Variation if validation.image_count == 0 => {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_VARIATION_INPUT_DETAIL,
));
}
_ => {}
}
if validation.n.is_some_and(|value| value != 1) {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_N_DETAIL,
));
}
if validation.partial_images.is_some_and(|value| value > 3)
|| (validation.partial_images.is_some() && !validation.stream)
{
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_PARTIAL_IMAGES_DETAIL,
));
}
if validation.style_present {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_STYLE_DETAIL,
));
}
if validation.stream {
if operation == OpenAiImageOperation::Variation {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_STREAM_VARIATION_DETAIL,
));
}
if !image_model_supports_streaming(validation.model.as_deref()) {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_STREAM_MODEL_DETAIL,
));
}
}
if let Some(n) = payload.get("n").and_then(image_request_count) {
if n != 1 {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_N_DETAIL,
));
}
if validation
.response_format
.as_deref()
.is_some_and(|value| !matches!(value, "url" | "b64_json"))
{
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_RESPONSE_FORMAT_DETAIL,
));
}
if validation
.output_format
.as_deref()
.is_some_and(|value| !matches!(value, "png" | "jpeg" | "jpg" | "webp"))
{
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_OUTPUT_FORMAT_DETAIL,
));
}
if validation
.quality
.as_deref()
.is_some_and(|value| !matches!(value, "low" | "medium" | "high" | "standard" | "hd"))
{
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_QUALITY_DETAIL,
));
}
if validation
.background
.as_deref()
.is_some_and(|value| !matches!(value, "auto" | "opaque" | "transparent"))
{
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_BACKGROUND_DETAIL,
));
}
if validation
.moderation
.as_deref()
.is_some_and(|value| !matches!(value, "auto" | "low"))
{
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_MODERATION_DETAIL,
));
}
if validation
.input_fidelity
.as_deref()
.is_some_and(|value| !matches!(value, "low" | "high"))
{
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_INPUT_FIDELITY_DETAIL,
));
}
if validation
.output_compression
.is_some_and(|value| value > 100)
{
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_OUTPUT_COMPRESSION_DETAIL,
));
}
None
@@ -141,6 +335,306 @@ fn image_request_count(value: &Value) -> Option<u64> {
})
}
fn is_openai_image_model(model: &str) -> bool {
canonicalize_openai_image_model(model).is_some()
}
fn canonicalize_openai_image_model(model: &str) -> Option<&'static str> {
match model.trim().to_ascii_lowercase().as_str() {
"gpt-image-1" => Some("gpt-image-1"),
"gpt-image-1.5" => Some("gpt-image-1.5"),
"gpt-image-1-mini" => Some("gpt-image-1-mini"),
"gpt-image-2" => Some("gpt-image-2"),
"chatgpt-image-latest" => Some("chatgpt-image-latest"),
"dall-e-2" => Some("dall-e-2"),
"dall-e-3" => Some("dall-e-3"),
_ => None,
}
}
fn image_model_supported_for_operation(operation: OpenAiImageOperation, model: &str) -> bool {
match operation {
OpenAiImageOperation::Generate => true,
OpenAiImageOperation::Edit => !matches!(model, "dall-e-3"),
OpenAiImageOperation::Variation => model == "dall-e-2",
}
}
fn image_model_supports_streaming(model: Option<&str>) -> bool {
!matches!(model, Some("dall-e-2" | "dall-e-3"))
}
fn parse_openai_image_validation_input(
operation: OpenAiImageOperation,
content_type: Option<&str>,
request_body: &Bytes,
) -> Result<OpenAiImageValidationInput, &'static str> {
if request_body.is_empty() {
return Err(match operation {
OpenAiImageOperation::Generate | OpenAiImageOperation::Edit => {
OPENAI_IMAGE_PROMPT_DETAIL
}
OpenAiImageOperation::Variation => OPENAI_IMAGE_VARIATION_INPUT_DETAIL,
});
}
let content_type = content_type.unwrap_or_default().to_ascii_lowercase();
if content_type.contains("multipart/form-data") {
parse_openai_image_validation_input_from_multipart(request_body, &content_type)
} else {
parse_openai_image_validation_input_from_json(request_body)
}
}
fn parse_openai_image_validation_input_from_json(
request_body: &Bytes,
) -> Result<OpenAiImageValidationInput, &'static str> {
let payload = serde_json::from_slice::<Value>(request_body)
.map_err(|_| OPENAI_IMAGE_INVALID_JSON_DETAIL)?;
let object = payload
.as_object()
.ok_or(OPENAI_IMAGE_INVALID_JSON_DETAIL)?;
Ok(OpenAiImageValidationInput {
model: normalize_openai_image_model_for_operation(
object.get("model").and_then(Value::as_str),
)
.ok_or(OPENAI_IMAGE_INVALID_JSON_DETAIL)?,
prompt: object
.get("prompt")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
image_count: count_json_images(object),
n: object.get("n").and_then(image_request_count),
stream: object
.get("stream")
.and_then(value_as_bool)
.unwrap_or(false),
partial_images: object.get("partial_images").and_then(image_request_count),
response_format: object
.get("response_format")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.to_ascii_lowercase()),
output_format: object
.get("output_format")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.to_ascii_lowercase()),
quality: object
.get("quality")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.to_ascii_lowercase()),
background: object
.get("background")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.to_ascii_lowercase()),
moderation: object
.get("moderation")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.to_ascii_lowercase()),
input_fidelity: object
.get("input_fidelity")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.to_ascii_lowercase()),
output_compression: object
.get("output_compression")
.and_then(image_request_count),
style_present: object
.get("style")
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty()),
})
}
fn parse_openai_image_validation_input_from_multipart(
request_body: &Bytes,
content_type: &str,
) -> Result<OpenAiImageValidationInput, &'static str> {
let boundary = content_type
.split(';')
.find_map(|segment| segment.trim().strip_prefix("boundary="))
.map(|value| value.trim_matches('"').to_string())
.ok_or(OPENAI_IMAGE_INVALID_MULTIPART_DETAIL)?;
let fields = parse_multipart_fields(request_body, &boundary);
if fields.is_empty() {
return Err(OPENAI_IMAGE_INVALID_MULTIPART_DETAIL);
}
let model = fields
.iter()
.find(|field| field.name.trim() == "model")
.map(|field| String::from_utf8_lossy(&field.data).trim().to_string());
Ok(OpenAiImageValidationInput {
model: normalize_openai_image_model_for_operation(model.as_deref())
.ok_or(OPENAI_IMAGE_INVALID_MULTIPART_DETAIL)?,
prompt: multipart_text_field(&fields, "prompt"),
image_count: fields
.iter()
.filter(|field| {
matches!(
field.name.trim(),
"image" | "image[]" | "images" | "images[]"
)
})
.count(),
n: multipart_text_field(&fields, "n").and_then(|value| value.trim().parse::<u64>().ok()),
stream: multipart_text_field(&fields, "stream")
.and_then(|value| parse_bool_string(&value))
.unwrap_or(false),
partial_images: multipart_text_field(&fields, "partial_images")
.and_then(|value| value.trim().parse::<u64>().ok()),
response_format: multipart_text_field(&fields, "response_format")
.map(|value| value.to_ascii_lowercase()),
output_format: multipart_text_field(&fields, "output_format")
.map(|value| value.to_ascii_lowercase()),
quality: multipart_text_field(&fields, "quality").map(|value| value.to_ascii_lowercase()),
background: multipart_text_field(&fields, "background")
.map(|value| value.to_ascii_lowercase()),
moderation: multipart_text_field(&fields, "moderation")
.map(|value| value.to_ascii_lowercase()),
input_fidelity: multipart_text_field(&fields, "input_fidelity")
.map(|value| value.to_ascii_lowercase()),
output_compression: multipart_text_field(&fields, "output_compression")
.and_then(|value| value.trim().parse::<u64>().ok()),
style_present: multipart_text_field(&fields, "style").is_some(),
})
}
fn normalize_openai_image_model_for_operation(model: Option<&str>) -> Option<Option<String>> {
let Some(model) = model.map(str::trim).filter(|value| !value.is_empty()) else {
return Some(None);
};
canonicalize_openai_image_model(model).map(|canonical| Some(canonical.to_string()))
}
fn count_json_images(object: &serde_json::Map<String, Value>) -> usize {
let mut count = 0usize;
if let Some(value) = object.get("image") {
count += json_image_count(value);
}
if let Some(values) = object.get("images").and_then(Value::as_array) {
count += values.iter().map(json_image_count).sum::<usize>();
}
count
}
fn json_image_count(value: &Value) -> usize {
match value {
Value::Array(values) => values.iter().map(json_image_count).sum(),
Value::String(text) => (!text.trim().is_empty()) as usize,
Value::Object(_) => 1,
_ => 0,
}
}
fn value_as_bool(value: &Value) -> Option<bool> {
value
.as_bool()
.or_else(|| value.as_str().and_then(parse_bool_string))
}
fn parse_bool_string(value: &str) -> Option<bool> {
match value.trim().to_ascii_lowercase().as_str() {
"true" | "1" | "yes" => Some(true),
"false" | "0" | "no" => Some(false),
_ => None,
}
}
#[derive(Debug)]
struct MultipartField {
name: String,
data: Vec<u8>,
}
fn multipart_text_field(fields: &[MultipartField], name: &str) -> Option<String> {
fields
.iter()
.find(|field| field.name.trim() == name)
.map(|field| String::from_utf8_lossy(&field.data).trim().to_string())
.filter(|value| !value.is_empty())
}
fn parse_multipart_fields(body: &[u8], boundary: &str) -> Vec<MultipartField> {
let delimiter = format!("--{boundary}").into_bytes();
let mut parts = Vec::new();
let mut cursor = 0usize;
while let Some(index) = find_subslice(&body[cursor..], &delimiter) {
let start = cursor + index + delimiter.len();
if body.get(start..start + 2) == Some(b"--") {
break;
}
let mut part = &body[start..];
if part.starts_with(b"\r\n") {
part = &part[2..];
}
let Some(next) = find_subslice(part, &delimiter) else {
break;
};
let raw = &part[..next];
let raw = raw.strip_suffix(b"\r\n").unwrap_or(raw);
if let Some(field) = parse_multipart_field(raw) {
parts.push(field);
}
cursor = start + next;
}
parts
}
fn parse_multipart_field(raw: &[u8]) -> Option<MultipartField> {
let header_end = find_subslice(raw, b"\r\n\r\n")?;
let headers = &raw[..header_end];
let data = raw.get(header_end + 4..)?.to_vec();
let header_text = String::from_utf8_lossy(headers);
let mut name = None;
for line in header_text.lines() {
let trimmed = line.trim();
if trimmed
.to_ascii_lowercase()
.starts_with("content-disposition:")
{
name = extract_quoted_header_value(trimmed, "name");
}
}
Some(MultipartField { name: name?, data })
}
fn extract_quoted_header_value(header: &str, key: &str) -> Option<String> {
let pattern = format!("{key}=\"");
let start = header.find(&pattern)? + pattern.len();
let rest = &header[start..];
let end = rest.find('"')?;
Some(rest[..end].to_string())
}
fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option<usize> {
if needle.is_empty() || haystack.len() < needle.len() {
return None;
}
haystack
.windows(needle.len())
.position(|window| window == needle)
}
fn maybe_build_local_ai_public_route_guard_response(
request_context: &GatewayPublicRequestContext,
) -> Option<Response<Body>> {

View File

@@ -2,6 +2,9 @@ use crate::api::ai::public_api_format_local_path;
use crate::handlers::shared::{
query_param_optional_bool, query_param_value, unix_ms_to_rfc3339, unix_secs_to_rfc3339,
};
use crate::provider_key_auth::{
provider_key_configured_api_formats, provider_key_effective_api_formats,
};
use crate::AppState;
use aether_data_contracts::repository::candidates::{
PublicHealthTimelineBucket, RequestCandidateStatus, StoredRequestCandidate,
@@ -67,19 +70,7 @@ pub(crate) struct ApiFormatHealthMonitorOptions {
}
pub(crate) fn provider_key_api_formats(key: &StoredProviderCatalogKey) -> Vec<String> {
key.api_formats
.as_ref()
.and_then(|value| value.as_array())
.map(|items| {
items
.iter()
.filter_map(|item| item.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.collect::<Vec<_>>()
})
.unwrap_or_default()
provider_key_configured_api_formats(key)
}
pub(crate) async fn build_public_providers_payload(
@@ -336,16 +327,25 @@ pub(crate) async fn build_api_format_health_monitor_payload(
let mut endpoint_ids_by_format = BTreeMap::<String, Vec<String>>::new();
let mut endpoint_to_format = BTreeMap::<String, String>::new();
let mut provider_ids_by_format = BTreeMap::<String, BTreeSet<String>>::new();
let mut active_endpoints_by_provider = BTreeMap::<String, Vec<_>>::new();
let provider_type_by_id = providers
.iter()
.map(|provider| (provider.id.clone(), provider.provider_type.clone()))
.collect::<BTreeMap<_, _>>();
for endpoint in active_endpoints {
endpoint_to_format.insert(endpoint.id.clone(), endpoint.api_format.clone());
endpoint_ids_by_format
.entry(endpoint.api_format.clone())
.or_default()
.push(endpoint.id);
.push(endpoint.id.clone());
provider_ids_by_format
.entry(endpoint.api_format.clone())
.or_default()
.insert(endpoint.provider_id.clone());
active_endpoints_by_provider
.entry(endpoint.provider_id.clone())
.or_default()
.push(endpoint);
}
let all_endpoint_ids = endpoint_to_format.keys().cloned().collect::<Vec<_>>();
@@ -357,7 +357,15 @@ pub(crate) async fn build_api_format_health_monitor_payload(
.ok()
.unwrap_or_default();
for key in keys.into_iter().filter(|key| key.is_active) {
for api_format in provider_key_api_formats(&key) {
let provider_type = provider_type_by_id
.get(&key.provider_id)
.map(String::as_str)
.unwrap_or("");
let endpoints = active_endpoints_by_provider
.get(&key.provider_id)
.map(Vec::as_slice)
.unwrap_or(&[]);
for api_format in provider_key_effective_api_formats(&key, provider_type, endpoints) {
if provider_ids_by_format
.get(&api_format)
.is_some_and(|provider_ids| provider_ids.contains(key.provider_id.as_str()))

View File

@@ -124,7 +124,13 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
}
let Some(key) = active_keys
.iter()
.find(|key| provider_catalog_key_supports_format(key, &format_value))
.find(|key| {
provider_catalog_key_supports_format(
key,
provider.provider_type.as_str(),
&format_value,
)
})
.cloned()
.or_else(|| active_keys.into_iter().next())
else {

View File

@@ -194,6 +194,55 @@ fn build_users_me_usage_api_key_payload(
}
}
fn users_me_usage_request_body_stream_flag(item: &StoredRequestUsageAudit) -> Option<bool> {
item.request_body
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|body| body.get("stream"))
.and_then(serde_json::Value::as_bool)
}
fn users_me_usage_api_format_defaults_to_non_stream(item: &StoredRequestUsageAudit) -> bool {
let api_format = item
.api_format
.as_deref()
.or(item.endpoint_api_format.as_deref())
.map(str::trim)
.filter(|value| !value.is_empty());
matches!(
api_format,
Some(value)
if value.eq_ignore_ascii_case("openai:chat")
|| value.eq_ignore_ascii_case("openai:cli")
|| value.eq_ignore_ascii_case("openai:compact")
|| value.eq_ignore_ascii_case("openai:image")
|| value.eq_ignore_ascii_case("claude:chat")
|| value.eq_ignore_ascii_case("claude:cli")
)
}
fn users_me_usage_request_body_implies_default_non_stream(item: &StoredRequestUsageAudit) -> bool {
let Some(body) = item
.request_body
.as_ref()
.and_then(serde_json::Value::as_object)
else {
return false;
};
!body.contains_key("stream") && users_me_usage_api_format_defaults_to_non_stream(item)
}
fn users_me_usage_client_is_stream(item: &StoredRequestUsageAudit) -> bool {
item.request_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get("client_requested_stream"))
.and_then(serde_json::Value::as_bool)
.or_else(|| users_me_usage_request_body_stream_flag(item))
.or_else(|| users_me_usage_request_body_implies_default_non_stream(item).then_some(false))
.unwrap_or(item.is_stream)
}
fn build_users_me_usage_record_payload(
item: &StoredRequestUsageAudit,
include_actual_cost: bool,
@@ -206,6 +255,7 @@ fn build_users_me_usage_record_payload(
let cache_read_price_per_1m = item.settlement_cache_read_price_per_1m();
let cache_creation_input_tokens = users_me_usage_cache_creation_tokens(item);
let rate_multiplier = item.settlement_rate_multiplier();
let client_is_stream = users_me_usage_client_is_stream(item);
let mut payload = json!({
"id": item.id,
"model": item.model,
@@ -221,6 +271,9 @@ fn build_users_me_usage_record_payload(
"response_time_ms": item.response_time_ms,
"first_byte_time_ms": item.first_byte_time_ms,
"is_stream": item.is_stream,
"upstream_is_stream": item.is_stream,
"client_requested_stream": client_is_stream,
"client_is_stream": client_is_stream,
"status": item.status,
"has_fallback": item.has_fallback(),
"created_at": unix_secs_to_rfc3339(item.created_at_unix_ms),
@@ -253,6 +306,7 @@ fn build_users_me_usage_record_payload(
fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_json::Value {
let cache_creation_input_tokens = users_me_usage_cache_creation_tokens(item);
let client_is_stream = users_me_usage_client_is_stream(item);
let mut payload = json!({
"id": item.id,
"status": item.status,
@@ -270,6 +324,10 @@ fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_
"first_byte_time_ms": item.first_byte_time_ms,
"api_format": item.api_format,
"endpoint_api_format": item.endpoint_api_format,
"is_stream": item.is_stream,
"upstream_is_stream": item.is_stream,
"client_requested_stream": client_is_stream,
"client_is_stream": client_is_stream,
"has_format_conversion": item.has_format_conversion,
"target_model": item.target_model,
"has_fallback": item.has_fallback(),
@@ -543,7 +601,7 @@ pub(super) async fn handle_users_me_usage_get(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("user usage summary lookup failed: {err:?}"),
false,
)
);
}
};
summary_by_model = match state
@@ -561,7 +619,7 @@ pub(super) async fn handle_users_me_usage_get(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("user usage model breakdown lookup failed: {err:?}"),
false,
)
);
}
};
summary_by_provider = match state
@@ -579,7 +637,7 @@ pub(super) async fn handle_users_me_usage_get(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("user usage provider breakdown lookup failed: {err:?}"),
false,
)
);
}
};
summary_by_api_format = match state
@@ -597,7 +655,7 @@ pub(super) async fn handle_users_me_usage_get(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("user usage api_format breakdown lookup failed: {err:?}"),
false,
)
);
}
};
@@ -614,7 +672,7 @@ pub(super) async fn handle_users_me_usage_get(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("user api key search context lookup failed: {err:?}"),
false,
)
);
}
};
let keyword_query = UsageAuditKeywordSearchQuery {
@@ -648,7 +706,7 @@ pub(super) async fn handle_users_me_usage_get(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("user usage search count lookup failed: {err:?}"),
false,
)
);
}
};
record_items = match state
@@ -665,7 +723,7 @@ pub(super) async fn handle_users_me_usage_get(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("user usage search lookup failed: {err:?}"),
false,
)
);
}
};
} else {
@@ -692,7 +750,7 @@ pub(super) async fn handle_users_me_usage_get(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("user usage count lookup failed: {err:?}"),
false,
)
);
}
};
record_items = match state
@@ -718,7 +776,7 @@ pub(super) async fn handle_users_me_usage_get(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("user usage records lookup failed: {err:?}"),
false,
)
);
}
};
api_key_names = match resolve_users_me_api_key_names(state, &record_items).await {
@@ -728,7 +786,7 @@ pub(super) async fn handle_users_me_usage_get(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("user api key name lookup failed: {err:?}"),
false,
)
);
}
};
}
@@ -826,7 +884,7 @@ pub(super) async fn handle_users_me_usage_active_get(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("user active usage lookup failed: {err:?}"),
false,
)
);
}
},
None => match state
@@ -852,7 +910,7 @@ pub(super) async fn handle_users_me_usage_active_get(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("user active usage lookup failed: {err:?}"),
false,
)
);
}
},
};
@@ -907,7 +965,7 @@ pub(super) async fn handle_users_me_usage_interval_timeline_get(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("user interval timeline lookup failed: {err:?}"),
false,
)
);
}
};
@@ -976,7 +1034,7 @@ pub(super) async fn handle_users_me_usage_heatmap_get(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("user heatmap lookup failed: {err:?}"),
false,
)
);
}
};
@@ -1089,8 +1147,12 @@ mod tests {
use std::collections::BTreeMap;
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
use serde_json::json;
use super::{build_users_me_usage_active_payload, build_users_me_usage_record_payload};
use super::{
build_users_me_usage_active_payload, build_users_me_usage_record_payload,
users_me_usage_client_is_stream,
};
fn sample_usage(status: &str) -> StoredRequestUsageAudit {
StoredRequestUsageAudit::new(
@@ -1165,4 +1227,80 @@ mod tests {
assert_eq!(payload["cache_creation_ephemeral_5m_input_tokens"], 4);
assert_eq!(payload["cache_creation_ephemeral_1h_input_tokens"], 6);
}
#[test]
fn user_usage_payloads_include_symmetric_stream_fields() {
let item = StoredRequestUsageAudit {
is_stream: true,
request_metadata: Some(json!({
"client_requested_stream": false
})),
..sample_usage("completed")
};
assert!(!users_me_usage_client_is_stream(&item));
let record_payload =
build_users_me_usage_record_payload(&item, false, &BTreeMap::new(), false);
assert_eq!(record_payload["is_stream"], true);
assert_eq!(record_payload["upstream_is_stream"], true);
assert_eq!(record_payload["client_requested_stream"], false);
assert_eq!(record_payload["client_is_stream"], false);
let active_payload = build_users_me_usage_active_payload(&item);
assert_eq!(active_payload["is_stream"], true);
assert_eq!(active_payload["upstream_is_stream"], true);
assert_eq!(active_payload["client_requested_stream"], false);
assert_eq!(active_payload["client_is_stream"], false);
}
#[test]
fn user_usage_stream_inference_falls_back_to_request_body_stream_flag() {
let item = StoredRequestUsageAudit {
is_stream: true,
request_body: Some(json!({
"model": "gpt-5.4",
"stream": false
})),
..sample_usage("completed")
};
assert!(!users_me_usage_client_is_stream(&item));
let record_payload =
build_users_me_usage_record_payload(&item, false, &BTreeMap::new(), false);
assert_eq!(record_payload["is_stream"], true);
assert_eq!(record_payload["upstream_is_stream"], true);
assert_eq!(record_payload["client_requested_stream"], false);
assert_eq!(record_payload["client_is_stream"], false);
}
#[test]
fn user_usage_stream_defaults_to_non_stream_for_openai_cli_request_body_without_flag() {
let item = StoredRequestUsageAudit {
is_stream: true,
api_format: Some("openai:cli".to_string()),
request_body: Some(json!({
"model": "gpt-5.4",
"input": [{"role": "user", "content": "hi"}],
"store": false
})),
..sample_usage("completed")
};
assert!(!users_me_usage_client_is_stream(&item));
let record_payload =
build_users_me_usage_record_payload(&item, false, &BTreeMap::new(), false);
assert_eq!(record_payload["is_stream"], true);
assert_eq!(record_payload["upstream_is_stream"], true);
assert_eq!(record_payload["client_requested_stream"], false);
assert_eq!(record_payload["client_is_stream"], false);
let active_payload = build_users_me_usage_active_payload(&item);
assert_eq!(active_payload["is_stream"], true);
assert_eq!(active_payload["upstream_is_stream"], true);
assert_eq!(active_payload["client_requested_stream"], false);
assert_eq!(active_payload["client_is_stream"], false);
}
}

View File

@@ -1,5 +1,6 @@
use super::enabled_key_capability_short_names;
use crate::handlers::shared::{json_string_list, unix_secs_to_rfc3339};
use crate::handlers::shared::unix_secs_to_rfc3339;
use crate::provider_key_auth::provider_key_effective_api_formats;
use crate::AppState;
use serde_json::json;
use std::collections::{BTreeMap, HashMap};
@@ -34,7 +35,11 @@ pub(crate) async fn build_admin_keys_grouped_by_format_payload(
.map(|provider| {
(
provider.id.clone(),
(provider.name.clone(), provider.is_active),
(
provider.name.clone(),
provider.is_active,
provider.provider_type.clone(),
),
)
})
.collect::<HashMap<_, _>>();
@@ -44,18 +49,28 @@ pub(crate) async fn build_admin_keys_grouped_by_format_payload(
state.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids),
);
let endpoint_base_url_by_provider_and_format = endpoints_result
let active_endpoints = endpoints_result
.ok()
.unwrap_or_default()
.into_iter()
.filter(|endpoint| endpoint.is_active)
.collect::<Vec<_>>();
let endpoint_base_url_by_provider_and_format = active_endpoints
.iter()
.map(|endpoint| {
(
(endpoint.provider_id, endpoint.api_format),
endpoint.base_url,
(endpoint.provider_id.clone(), endpoint.api_format.clone()),
endpoint.base_url.clone(),
)
})
.collect::<HashMap<_, _>>();
let mut endpoints_by_provider = HashMap::<String, Vec<_>>::new();
for endpoint in active_endpoints {
endpoints_by_provider
.entry(endpoint.provider_id.clone())
.or_default()
.push(endpoint);
}
let mut keys = keys_result.ok().unwrap_or_default();
keys.sort_by(|left, right| {
@@ -72,7 +87,7 @@ pub(crate) async fn build_admin_keys_grouped_by_format_payload(
let mut grouped = BTreeMap::<String, Vec<serde_json::Value>>::new();
for key in keys {
let Some((provider_name, provider_is_active)) =
let Some((provider_name, provider_is_active, provider_type)) =
provider_metadata_by_id.get(&key.provider_id)
else {
continue;
@@ -108,7 +123,14 @@ pub(crate) async fn build_admin_keys_grouped_by_format_payload(
.cloned()
.unwrap_or_default();
let capability_names = enabled_key_capability_short_names(key.capabilities.as_ref());
let api_formats = json_string_list(key.api_formats.as_ref());
let api_formats = provider_key_effective_api_formats(
&key,
provider_type,
endpoints_by_provider
.get(&key.provider_id)
.map(Vec::as_slice)
.unwrap_or(&[]),
);
if api_formats.is_empty() {
continue;
}