mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 08:50:21 +08:00
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:
14
Cargo.toml
14
Cargo.toml
@@ -79,7 +79,17 @@ uuid = { version = "1", features = ["serde", "v4", "v5"] }
|
||||
webpki-roots = "0.26"
|
||||
url = "2"
|
||||
|
||||
[profile.dev]
|
||||
# Keep file/line information for backtraces while avoiding full debug info
|
||||
# generation on very large crates during local development builds.
|
||||
debug = "line-tables-only"
|
||||
|
||||
[profile.test]
|
||||
# The gateway test target pulls in a very large in-crate test tree, so use the
|
||||
# lighter debug format here as well to reduce rustc peak memory.
|
||||
debug = "line-tables-only"
|
||||
|
||||
[profile.release]
|
||||
lto = true
|
||||
lto = "thin"
|
||||
strip = true
|
||||
codegen-units = 1
|
||||
codegen-units = 8
|
||||
|
||||
@@ -29,7 +29,8 @@ pub(crate) use crate::ai_pipeline::{
|
||||
OPENAI_CLI_SYNC_FINALIZE_REPORT_KIND, OPENAI_CLI_SYNC_PLAN_KIND,
|
||||
OPENAI_CLI_SYNC_SUCCESS_REPORT_KIND, OPENAI_COMPACT_STREAM_PLAN_KIND,
|
||||
OPENAI_COMPACT_SYNC_ERROR_REPORT_KIND, OPENAI_COMPACT_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_COMPACT_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_COMPACT_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
|
||||
OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
|
||||
|
||||
@@ -2,12 +2,14 @@ use serde_json::Value;
|
||||
|
||||
use crate::ai_pipeline::adaptation::private_envelope::transform_provider_private_stream_line as transform_envelope_line;
|
||||
use crate::ai_pipeline::adaptation::KiroToClaudeCliStreamState;
|
||||
use crate::ai_pipeline::finalize::sse::encode_json_sse;
|
||||
use crate::ai_pipeline::finalize::standard::StreamingStandardConversionState;
|
||||
use crate::ai_pipeline::{resolve_finalize_stream_rewrite_mode, FinalizeStreamRewriteMode};
|
||||
use crate::GatewayError;
|
||||
|
||||
enum RewriteMode {
|
||||
EnvelopeUnwrap,
|
||||
OpenAiImage(OpenAiImageStreamState),
|
||||
Standard(StreamingStandardConversionState),
|
||||
KiroToClaudeCli(KiroToClaudeCliStreamState),
|
||||
}
|
||||
@@ -24,6 +26,9 @@ pub(crate) fn maybe_build_local_stream_rewriter<'a>(
|
||||
let report_context = report_context?;
|
||||
let mode = match resolve_finalize_stream_rewrite_mode(report_context)? {
|
||||
FinalizeStreamRewriteMode::EnvelopeUnwrap => RewriteMode::EnvelopeUnwrap,
|
||||
FinalizeStreamRewriteMode::OpenAiImage => {
|
||||
RewriteMode::OpenAiImage(OpenAiImageStreamState::default())
|
||||
}
|
||||
FinalizeStreamRewriteMode::Standard => {
|
||||
RewriteMode::Standard(StreamingStandardConversionState::default())
|
||||
}
|
||||
@@ -41,6 +46,9 @@ pub(crate) fn maybe_build_local_stream_rewriter<'a>(
|
||||
|
||||
impl LocalStreamRewriter<'_> {
|
||||
pub(crate) fn push_chunk(&mut self, chunk: &[u8]) -> Result<Vec<u8>, GatewayError> {
|
||||
if let RewriteMode::OpenAiImage(state) = &mut self.mode {
|
||||
return state.push_chunk(self.report_context, chunk);
|
||||
}
|
||||
if let RewriteMode::KiroToClaudeCli(state) = &mut self.mode {
|
||||
return state.push_chunk(self.report_context, chunk);
|
||||
}
|
||||
@@ -54,12 +62,16 @@ impl LocalStreamRewriter<'_> {
|
||||
}
|
||||
|
||||
pub(crate) fn finish(&mut self) -> Result<Vec<u8>, GatewayError> {
|
||||
if let RewriteMode::OpenAiImage(state) = &mut self.mode {
|
||||
return state.finish(self.report_context);
|
||||
}
|
||||
if let RewriteMode::KiroToClaudeCli(state) = &mut self.mode {
|
||||
return state.finish(self.report_context);
|
||||
}
|
||||
if self.buffered.is_empty() {
|
||||
match &mut self.mode {
|
||||
RewriteMode::Standard(state) => return state.finish(self.report_context),
|
||||
RewriteMode::OpenAiImage(_) => {}
|
||||
RewriteMode::KiroToClaudeCli(_) => {}
|
||||
RewriteMode::EnvelopeUnwrap => {}
|
||||
}
|
||||
@@ -71,6 +83,7 @@ impl LocalStreamRewriter<'_> {
|
||||
RewriteMode::Standard(state) => {
|
||||
output.extend(state.finish(self.report_context)?);
|
||||
}
|
||||
RewriteMode::OpenAiImage(_) => {}
|
||||
RewriteMode::KiroToClaudeCli(_) => {}
|
||||
RewriteMode::EnvelopeUnwrap => {}
|
||||
}
|
||||
@@ -81,12 +94,208 @@ impl LocalStreamRewriter<'_> {
|
||||
match &mut self.mode {
|
||||
RewriteMode::EnvelopeUnwrap => transform_envelope_line(self.report_context, line)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string())),
|
||||
RewriteMode::OpenAiImage(_) => Ok(Vec::new()),
|
||||
RewriteMode::Standard(state) => state.transform_line(self.report_context, line),
|
||||
RewriteMode::KiroToClaudeCli(_) => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct OpenAiImageStreamState {
|
||||
buffered: Vec<u8>,
|
||||
latest_image: Option<OpenAiImageFrame>,
|
||||
emitted_partial_count: u64,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct OpenAiImageFrame {
|
||||
b64_json: String,
|
||||
}
|
||||
|
||||
impl OpenAiImageStreamState {
|
||||
fn push_chunk(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
chunk: &[u8],
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
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)
|
||||
}
|
||||
|
||||
fn finish(&mut self, report_context: &Value) -> Result<Vec<u8>, GatewayError> {
|
||||
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>, GatewayError> {
|
||||
let text =
|
||||
std::str::from_utf8(block).map_err(|err| GatewayError::Internal(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).map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let event_type = event
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.or(event_name.as_deref())
|
||||
.unwrap_or_default();
|
||||
match event_type {
|
||||
"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_output_item_done(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
event: &Value,
|
||||
) -> Result<Vec<u8>, GatewayError> {
|
||||
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 {
|
||||
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>, GatewayError> {
|
||||
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 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_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);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "../tests_stream.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_pipeline::CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT;
|
||||
use crate::ai_pipeline::{build_generated_tool_call_id, canonicalize_tool_arguments};
|
||||
use crate::{usage::GatewaySyncReportRequest, GatewayError};
|
||||
use base64::Engine as _;
|
||||
@@ -105,6 +106,20 @@ fn maybe_build_local_openai_image_sync_finalize_response(
|
||||
let Some(body_base64) = payload.body_base64.as_deref() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let response_format = report_context
|
||||
.get("image_request")
|
||||
.and_then(|value| value.get("response_format"))
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("url");
|
||||
let default_output_format = report_context
|
||||
.get("image_request")
|
||||
.and_then(|value| value.get("output_format"))
|
||||
.and_then(serde_json::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)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
@@ -157,6 +172,7 @@ fn maybe_build_local_openai_image_sync_finalize_response(
|
||||
};
|
||||
images.push(serde_json::json!({
|
||||
"b64_json": result,
|
||||
"output_format": item.get("output_format").cloned().unwrap_or(serde_json::Value::String(default_output_format.to_string())),
|
||||
"revised_prompt": item.get("revised_prompt").cloned().unwrap_or(serde_json::Value::Null),
|
||||
}));
|
||||
}
|
||||
@@ -191,13 +207,46 @@ fn maybe_build_local_openai_image_sync_finalize_response(
|
||||
.iter()
|
||||
.map(|image| serde_json::json!({
|
||||
"type": "image_generation_call",
|
||||
"output_format": image.get("output_format").cloned().unwrap_or(serde_json::Value::Null),
|
||||
"revised_prompt": image.get("revised_prompt").cloned().unwrap_or(serde_json::Value::Null),
|
||||
}))
|
||||
.collect::<Vec<_>>(),
|
||||
});
|
||||
let client_images = images
|
||||
.iter()
|
||||
.map(|image| {
|
||||
let revised_prompt = image
|
||||
.get("revised_prompt")
|
||||
.cloned()
|
||||
.unwrap_or(serde_json::Value::Null);
|
||||
let b64_json = image
|
||||
.get("b64_json")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let output_format = image
|
||||
.get("output_format")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.unwrap_or(default_output_format);
|
||||
if response_format.eq_ignore_ascii_case("b64_json") {
|
||||
serde_json::json!({
|
||||
"b64_json": b64_json,
|
||||
"revised_prompt": revised_prompt,
|
||||
})
|
||||
} else {
|
||||
serde_json::json!({
|
||||
"url": format!(
|
||||
"data:{};base64,{}",
|
||||
image_output_mime_type(output_format),
|
||||
b64_json
|
||||
),
|
||||
"revised_prompt": revised_prompt,
|
||||
})
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let client_body_json = serde_json::json!({
|
||||
"created": created.unwrap_or_default(),
|
||||
"data": images,
|
||||
"data": client_images,
|
||||
"usage": provider_body_json.get("usage").cloned().unwrap_or(serde_json::Value::Null),
|
||||
});
|
||||
|
||||
@@ -210,6 +259,14 @@ fn maybe_build_local_openai_image_sync_finalize_response(
|
||||
)?))
|
||||
}
|
||||
|
||||
fn image_output_mime_type(output_format: &str) -> &'static str {
|
||||
match output_format.trim().to_ascii_lowercase().as_str() {
|
||||
"jpeg" | "jpg" => "image/jpeg",
|
||||
"webp" => "image/webp",
|
||||
_ => "image/png",
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "../tests_sync.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -29,6 +29,94 @@ fn antigravity_stream_rewriter_unwraps_and_injects_tool_ids() {
|
||||
assert!(output_text.contains("\"modelVersion\":\"claude-sonnet-4-5\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_image_stream_rewriter_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 =
|
||||
maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
|
||||
|
||||
let first = rewriter
|
||||
.push_chunk(
|
||||
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(
|
||||
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().expect("finish should succeed").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_image_stream_rewriter_emits_partial_and_completed_events_for_edit() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:image",
|
||||
"needs_conversion": false,
|
||||
"image_request": {
|
||||
"operation": "edit",
|
||||
"partial_images": 2
|
||||
}
|
||||
});
|
||||
let mut rewriter =
|
||||
maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
|
||||
|
||||
let partial = rewriter
|
||||
.push_chunk(
|
||||
concat!(
|
||||
"event: response.output_item.done\n",
|
||||
"data: {\"type\":\"response.output_item.done\",\"output_index\":1,\"item\":{\"id\":\"ig_edit_123\",\"type\":\"image_generation_call\",\"result\":\"d29ybGQ=\"}}\n\n"
|
||||
)
|
||||
.as_bytes(),
|
||||
)
|
||||
.expect("rewrite should succeed");
|
||||
let partial_text = utf8(partial);
|
||||
assert!(partial_text.contains("event: image_edit.partial_image"));
|
||||
assert!(partial_text.contains("\"type\":\"image_edit.partial_image\""));
|
||||
assert!(partial_text.contains("\"b64_json\":\"d29ybGQ=\""));
|
||||
assert!(partial_text.contains("\"partial_image_index\":1"));
|
||||
|
||||
let completed = rewriter
|
||||
.push_chunk(
|
||||
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_edit.completed"));
|
||||
assert!(completed_text.contains("\"type\":\"image_edit.completed\""));
|
||||
assert!(completed_text.contains("\"b64_json\":\"d29ybGQ=\""));
|
||||
assert!(completed_text.contains("\"total_tokens\":9"));
|
||||
assert!(rewriter.finish().expect("finish should succeed").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_cli_v1internal_stream_rewriter_unwraps_response_object() {
|
||||
let report_context = json!({
|
||||
|
||||
@@ -1773,7 +1773,12 @@ async fn local_finalize_handles_openai_image_stream_response_from_output_item_do
|
||||
"client_api_format": "openai:image",
|
||||
"provider_api_format": "openai:image",
|
||||
"model": "gpt-image-2",
|
||||
"mapped_model": "gpt-5.4"
|
||||
"mapped_model": "gpt-5.4",
|
||||
"image_request": {
|
||||
"operation": "generate",
|
||||
"response_format": "b64_json",
|
||||
"output_format": "png"
|
||||
}
|
||||
})),
|
||||
status_code: 200,
|
||||
headers: BTreeMap::from([(
|
||||
@@ -1830,3 +1835,63 @@ async fn local_finalize_handles_openai_image_stream_response_from_output_item_do
|
||||
"aGVsbG8="
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_finalize_handles_openai_image_stream_response_with_url_response_format() {
|
||||
let payload = GatewaySyncReportRequest {
|
||||
trace_id: "trace-openai-image-finalize-url-123".to_string(),
|
||||
report_kind: "openai_image_sync_finalize".to_string(),
|
||||
report_context: Some(json!({
|
||||
"client_api_format": "openai:image",
|
||||
"provider_api_format": "openai:image",
|
||||
"model": "gpt-image-1",
|
||||
"mapped_model": "gpt-5.4",
|
||||
"image_request": {
|
||||
"operation": "generate",
|
||||
"response_format": "url",
|
||||
"output_format": "webp"
|
||||
}
|
||||
})),
|
||||
status_code: 200,
|
||||
headers: BTreeMap::from([(
|
||||
"content-type".to_string(),
|
||||
"text/event-stream".to_string(),
|
||||
)]),
|
||||
body_json: None,
|
||||
client_body_json: None,
|
||||
body_base64: Some(base64::engine::general_purpose::STANDARD.encode(
|
||||
concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_img_url_123\",\"object\":\"response\",\"created_at\":1776839946,\"status\":\"in_progress\",\"model\":\"gpt-5.4\"}}\n\n",
|
||||
"event: response.output_item.done\n",
|
||||
"data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ig_url_123\",\"type\":\"image_generation_call\",\"status\":\"completed\",\"output_format\":\"webp\",\"revised_prompt\":\"revised webp prompt\",\"result\":\"aGVsbG8=\"}}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_img_url_123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"tool_usage\":{\"image_gen\":{\"input_tokens\":11,\"output_tokens\":22,\"total_tokens\":33}}}}\n\n"
|
||||
)
|
||||
.as_bytes(),
|
||||
)),
|
||||
telemetry: None,
|
||||
};
|
||||
|
||||
let outcome = maybe_build_local_core_sync_finalize_response(
|
||||
"trace-openai-image-finalize-url-123",
|
||||
&test_decision(),
|
||||
&payload,
|
||||
)
|
||||
.expect("image finalize should succeed")
|
||||
.expect("image finalize should match");
|
||||
|
||||
let response_body = to_bytes(outcome.response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("response body should read");
|
||||
let response_json: serde_json::Value =
|
||||
serde_json::from_slice(&response_body).expect("response should be json");
|
||||
assert_eq!(
|
||||
response_json["data"][0]["url"],
|
||||
"data:image/webp;base64,aGVsbG8="
|
||||
);
|
||||
assert_eq!(
|
||||
response_json["data"][0]["revised_prompt"],
|
||||
"revised webp prompt"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -26,6 +26,7 @@ pub(crate) use self::planner::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
build_local_gemini_files_stream_plan_and_reports_for_kind,
|
||||
build_local_gemini_files_sync_plan_and_reports_for_kind,
|
||||
build_local_image_stream_plan_and_reports_for_kind,
|
||||
build_local_image_sync_plan_and_reports_for_kind,
|
||||
build_local_openai_chat_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_chat_sync_plan_and_reports_for_kind,
|
||||
@@ -38,9 +39,9 @@ pub(crate) use self::planner::{
|
||||
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
|
||||
extract_pool_sticky_session_token, maybe_build_stream_decision_payload,
|
||||
maybe_build_stream_plan_payload, maybe_build_sync_decision_payload,
|
||||
maybe_build_sync_plan_payload, set_local_openai_chat_execution_exhausted_diagnostic,
|
||||
GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
|
||||
PlannerAppState,
|
||||
maybe_build_sync_plan_payload, planner_is_matching_stream_request,
|
||||
set_local_openai_chat_execution_exhausted_diagnostic, GatewayAuthApiKeySnapshot,
|
||||
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
|
||||
};
|
||||
pub(crate) use self::pure::*;
|
||||
pub(crate) use crate::control::GatewayControlDecision;
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_provider_transport::provider_types::provider_type_is_fixed;
|
||||
use tracing::warn;
|
||||
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
@@ -295,6 +296,22 @@ fn transport_key_supports_api_format(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
endpoint_api_format: &str,
|
||||
) -> bool {
|
||||
let provider_type = transport.provider.provider_type.trim();
|
||||
let auth_type = transport.key.auth_type.trim();
|
||||
let inherits_provider_api_formats = provider_type_is_fixed(provider_type)
|
||||
&& (auth_type.eq_ignore_ascii_case("oauth")
|
||||
|| (provider_type.eq_ignore_ascii_case("kiro")
|
||||
&& auth_type.eq_ignore_ascii_case("bearer")
|
||||
&& transport
|
||||
.key
|
||||
.decrypted_auth_config
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty())));
|
||||
if inherits_provider_api_formats {
|
||||
return true;
|
||||
}
|
||||
|
||||
match transport.key.api_formats.as_deref() {
|
||||
None => true,
|
||||
Some(formats) => formats
|
||||
|
||||
@@ -10,9 +10,10 @@ pub(crate) use crate::ai_pipeline::contracts::{
|
||||
GEMINI_FILES_LIST_PLAN_KIND, GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
OPENAI_CLI_STREAM_PLAN_KIND, OPENAI_CLI_SYNC_PLAN_KIND, OPENAI_COMPACT_STREAM_PLAN_KIND,
|
||||
OPENAI_COMPACT_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
OPENAI_COMPACT_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_pipeline::{
|
||||
|
||||
@@ -7,9 +7,10 @@ use crate::ai_pipeline::planner::common::{
|
||||
GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
OPENAI_CLI_STREAM_PLAN_KIND, OPENAI_CLI_SYNC_PLAN_KIND, OPENAI_COMPACT_STREAM_PLAN_KIND,
|
||||
OPENAI_COMPACT_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
OPENAI_COMPACT_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_pipeline::planner::plan_builders::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
@@ -63,13 +64,20 @@ pub(crate) async fn maybe_build_stream_plan_payload_impl(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> Result<Option<GatewayControlPlanResponse>, GatewayError> {
|
||||
let Some(plan_kind) = resolve_stream_plan_kind(parts, decision) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(payload) =
|
||||
super::maybe_build_stream_decision_payload(state, parts, trace_id, decision, body_json)
|
||||
.await?
|
||||
let Some(payload) = super::maybe_build_stream_decision_payload(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
body_base64,
|
||||
)
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -132,6 +140,9 @@ fn build_stream_plan_payload_from_decision(
|
||||
OPENAI_CLI_STREAM_PLAN_KIND => {
|
||||
build_openai_cli_stream_plan_from_decision(parts, body_json, payload, false)?
|
||||
}
|
||||
OPENAI_IMAGE_STREAM_PLAN_KIND => {
|
||||
build_standard_stream_plan_from_decision(parts, body_json, payload, false)?
|
||||
}
|
||||
OPENAI_COMPACT_STREAM_PLAN_KIND => {
|
||||
build_openai_cli_stream_plan_from_decision(parts, body_json, payload, true)?
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@ pub(crate) use super::passthrough::{
|
||||
};
|
||||
pub(crate) use super::specialized::{
|
||||
maybe_build_stream_local_gemini_files_decision_payload,
|
||||
maybe_build_stream_local_image_decision_payload,
|
||||
maybe_build_sync_local_gemini_files_decision_payload,
|
||||
maybe_build_sync_local_image_decision_payload, maybe_build_sync_local_video_decision_payload,
|
||||
};
|
||||
|
||||
@@ -18,12 +18,13 @@ pub(crate) async fn maybe_build_stream_decision_payload(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> Result<Option<GatewayControlSyncDecisionResponse>, GatewayError> {
|
||||
let Some(plan_kind) = resolve_execution_runtime_stream_plan_kind(parts, decision) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
if !is_matching_stream_request(plan_kind, parts, body_json) {
|
||||
if !is_matching_stream_request(plan_kind, parts, body_json, body_base64) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
@@ -35,6 +36,20 @@ pub(crate) async fn maybe_build_stream_decision_payload(
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
if let Some(payload) = super::maybe_build_stream_local_image_decision_payload(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
if let Some(payload) = super::maybe_build_stream_local_decision_payload(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
|
||||
@@ -36,9 +36,11 @@ pub(crate) use self::plan_builders::{
|
||||
build_passthrough_sync_plan_from_decision, build_standard_stream_plan_from_decision,
|
||||
build_standard_sync_plan_from_decision, LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
};
|
||||
pub(crate) use self::route::is_matching_stream_request as planner_is_matching_stream_request;
|
||||
pub(crate) use self::specialized::{
|
||||
build_local_gemini_files_stream_plan_and_reports_for_kind,
|
||||
build_local_gemini_files_sync_plan_and_reports_for_kind,
|
||||
build_local_image_stream_plan_and_reports_for_kind,
|
||||
build_local_image_sync_plan_and_reports_for_kind,
|
||||
build_local_video_sync_plan_and_reports_for_kind,
|
||||
};
|
||||
@@ -83,8 +85,17 @@ pub(crate) async fn maybe_build_stream_decision_payload(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> Result<Option<GatewayControlSyncDecisionResponse>, GatewayError> {
|
||||
decision::maybe_build_stream_decision_payload(state, parts, trace_id, decision, body_json).await
|
||||
decision::maybe_build_stream_decision_payload(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
body_base64,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_sync_plan_payload(
|
||||
@@ -114,7 +125,15 @@ pub(crate) async fn maybe_build_stream_plan_payload(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> Result<Option<GatewayControlPlanResponse>, GatewayError> {
|
||||
decision::maybe_build_stream_plan_payload_impl(state, parts, trace_id, decision, body_json)
|
||||
.await
|
||||
decision::maybe_build_stream_plan_payload_impl(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
body_base64,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
@@ -101,6 +101,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
original_headers: &parts.headers,
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: spec_metadata.require_streaming,
|
||||
has_envelope: resolved.is_kiro || resolved.is_antigravity,
|
||||
needs_conversion: false,
|
||||
extra_fields,
|
||||
|
||||
@@ -26,6 +26,7 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
|
||||
pub(crate) original_headers: &'a http::HeaderMap,
|
||||
pub(crate) original_request_body_json: Option<&'a Value>,
|
||||
pub(crate) original_request_body_base64: Option<&'a str>,
|
||||
pub(crate) client_requested_stream: bool,
|
||||
pub(crate) has_envelope: bool,
|
||||
pub(crate) needs_conversion: bool,
|
||||
pub(crate) extra_fields: Map<String, Value>,
|
||||
@@ -117,6 +118,10 @@ pub(crate) fn build_local_execution_report_context(
|
||||
)
|
||||
.unwrap_or(Value::Null),
|
||||
);
|
||||
object.insert(
|
||||
"client_requested_stream".to_string(),
|
||||
Value::Bool(parts.client_requested_stream),
|
||||
);
|
||||
object.insert("has_envelope".to_string(), Value::Bool(parts.has_envelope));
|
||||
object.insert(
|
||||
"needs_conversion".to_string(),
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
use super::specialized::is_openai_image_stream_request;
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_pipeline::{
|
||||
is_matching_stream_request as is_matching_stream_request_impl,
|
||||
@@ -5,6 +6,7 @@ use crate::ai_pipeline::{
|
||||
resolve_execution_runtime_sync_plan_kind as resolve_execution_runtime_sync_plan_kind_impl,
|
||||
supports_stream_scheduler_decision_kind as supports_stream_scheduler_decision_kind_impl,
|
||||
supports_sync_scheduler_decision_kind as supports_sync_scheduler_decision_kind_impl,
|
||||
OPENAI_IMAGE_STREAM_PLAN_KIND,
|
||||
};
|
||||
|
||||
pub(crate) fn resolve_execution_runtime_stream_plan_kind(
|
||||
@@ -37,7 +39,11 @@ pub(crate) fn is_matching_stream_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_impl(plan_kind, parts.uri.path(), body_json)
|
||||
}
|
||||
|
||||
@@ -52,6 +58,7 @@ pub(crate) fn supports_stream_scheduler_decision_kind(plan_kind: &str) -> bool {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use axum::http::{Method, Request};
|
||||
use base64::Engine as _;
|
||||
|
||||
use super::{
|
||||
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
|
||||
@@ -108,15 +115,45 @@ mod tests {
|
||||
"openai_chat_stream",
|
||||
&parts,
|
||||
&serde_json::json!({"stream": false}),
|
||||
None,
|
||||
));
|
||||
assert!(is_matching_stream_request(
|
||||
"openai_chat_stream",
|
||||
&parts,
|
||||
&serde_json::json!({"stream": true}),
|
||||
None,
|
||||
));
|
||||
assert!(supports_sync_scheduler_decision_kind("openai_chat_sync"));
|
||||
assert!(supports_stream_scheduler_decision_kind(
|
||||
"openai_chat_stream"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_stream_matching_parses_multipart_stream_flag() {
|
||||
let request = 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_request(
|
||||
"openai_image_stream",
|
||||
&parts,
|
||||
&serde_json::json!({}),
|
||||
Some(body_base64.as_str()),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -84,7 +84,7 @@ pub(crate) fn local_openai_image_spec_metadata(
|
||||
api_format: spec.api_format,
|
||||
decision_kind: spec.decision_kind,
|
||||
report_kind: Some(spec.report_kind),
|
||||
require_streaming: false,
|
||||
require_streaming: spec.require_streaming,
|
||||
requested_model_family: Some(RequestedModelFamily::Standard),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -86,6 +86,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
||||
original_headers: &parts.headers,
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: resolved.provider_request_body_base64.as_deref(),
|
||||
client_requested_stream: spec_metadata.require_streaming,
|
||||
has_envelope: false,
|
||||
needs_conversion: false,
|
||||
extra_fields,
|
||||
|
||||
@@ -5,11 +5,15 @@ mod support;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_pipeline::planner::plan_builders::{
|
||||
build_passthrough_sync_plan_from_decision, LocalSyncPlanAndReport,
|
||||
build_passthrough_sync_plan_from_decision, build_standard_stream_plan_from_decision,
|
||||
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
};
|
||||
use crate::ai_pipeline::planner::spec_metadata::local_openai_image_spec_metadata;
|
||||
use crate::ai_pipeline::resolve_local_image_sync_spec as resolve_sync_spec;
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::ai_pipeline::{
|
||||
resolve_local_image_stream_spec as resolve_stream_spec,
|
||||
resolve_local_image_sync_spec as resolve_sync_spec,
|
||||
};
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
|
||||
|
||||
use self::decision::maybe_build_local_openai_image_decision_payload_for_candidate;
|
||||
@@ -17,6 +21,7 @@ use self::support::{
|
||||
list_local_openai_image_candidate_attempts, resolve_local_openai_image_decision_input,
|
||||
};
|
||||
|
||||
pub(crate) use self::request::is_openai_image_stream_request;
|
||||
pub(super) use crate::ai_pipeline::LocalOpenAiImageSpec;
|
||||
|
||||
pub(crate) async fn build_local_image_sync_plan_and_reports_for_kind(
|
||||
@@ -44,6 +49,31 @@ pub(crate) async fn build_local_image_sync_plan_and_reports_for_kind(
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_image_stream_plan_and_reports_for_kind(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Vec<LocalStreamPlanAndReport>, GatewayError> {
|
||||
let Some(spec) = resolve_stream_spec(plan_kind) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
build_local_stream_plan_and_reports(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_sync_local_image_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
@@ -58,8 +88,75 @@ pub(crate) async fn maybe_build_sync_local_image_decision_payload(
|
||||
};
|
||||
let spec_metadata = local_openai_image_spec_metadata(spec);
|
||||
|
||||
let Some(input) =
|
||||
resolve_local_openai_image_decision_input(state, trace_id, decision, body_json).await
|
||||
let Some(input) = resolve_local_openai_image_decision_input(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let Some(attempts) = list_local_openai_image_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.decision_kind,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
&input,
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_stream_local_image_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<GatewayControlSyncDecisionResponse>, GatewayError> {
|
||||
let Some(spec) = resolve_stream_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let spec_metadata = local_openai_image_spec_metadata(spec);
|
||||
|
||||
let Some(input) = resolve_local_openai_image_decision_input(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -107,8 +204,15 @@ async fn build_local_sync_plan_and_reports(
|
||||
spec: LocalOpenAiImageSpec,
|
||||
) -> Result<Vec<LocalSyncPlanAndReport>, GatewayError> {
|
||||
let spec_metadata = local_openai_image_spec_metadata(spec);
|
||||
let Some(input) =
|
||||
resolve_local_openai_image_decision_input(state, trace_id, decision, body_json).await
|
||||
let Some(input) = resolve_local_openai_image_decision_input(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
@@ -159,3 +263,73 @@ async fn build_local_sync_plan_and_reports(
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
async fn build_local_stream_plan_and_reports(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
spec: LocalOpenAiImageSpec,
|
||||
) -> Result<Vec<LocalStreamPlanAndReport>, GatewayError> {
|
||||
let spec_metadata = local_openai_image_spec_metadata(spec);
|
||||
let Some(input) = resolve_local_openai_image_decision_input(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let Some(attempts) = list_local_openai_image_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.decision_kind,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
&input,
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_standard_stream_plan_from_decision(parts, body_json, payload, false) {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind = spec_metadata.decision_kind,
|
||||
error = ?err,
|
||||
"gateway local openai image stream decision plan build failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
@@ -78,6 +78,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
|
||||
original_headers: &parts.headers,
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: body_base64,
|
||||
client_requested_stream: spec_metadata.require_streaming,
|
||||
has_envelope: false,
|
||||
needs_conversion: false,
|
||||
extra_fields,
|
||||
@@ -85,7 +86,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
|
||||
|
||||
Some(build_local_execution_decision_response(
|
||||
LocalExecutionDecisionResponseParts {
|
||||
decision_is_stream: false,
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy: ExecutionStrategy::LocalSameFormat,
|
||||
conversion_mode: ConversionMode::None,
|
||||
@@ -112,7 +113,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream: false,
|
||||
upstream_is_stream: spec_metadata.require_streaming,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -30,28 +30,24 @@ use crate::clock::current_unix_secs;
|
||||
use crate::AppState;
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
pub(super) const OPENAI_IMAGE_DEFAULT_MODEL: &str = "gpt-image-2";
|
||||
|
||||
pub(super) use crate::ai_pipeline::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalOpenAiImageCandidateAttempt;
|
||||
pub(super) use crate::ai_pipeline::planner::decision_input::LocalRequestedModelDecisionInput as LocalOpenAiImageDecisionInput;
|
||||
|
||||
use super::request::resolve_requested_image_model_for_request;
|
||||
|
||||
pub(super) async fn resolve_local_openai_image_decision_input(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
) -> Option<LocalOpenAiImageDecisionInput> {
|
||||
let Some(auth_context) = resolve_local_openai_image_auth_context(decision) else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let requested_model = body_json
|
||||
.get("model")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(OPENAI_IMAGE_DEFAULT_MODEL)
|
||||
.to_string();
|
||||
let requested_model = resolve_requested_image_model_for_request(parts, body_json, body_base64)?;
|
||||
|
||||
let resolved_input = match resolve_local_authenticated_decision_input(
|
||||
state,
|
||||
|
||||
@@ -11,7 +11,9 @@ pub(crate) use self::files::{
|
||||
maybe_build_sync_local_gemini_files_decision_payload,
|
||||
};
|
||||
pub(crate) use self::image::{
|
||||
build_local_image_sync_plan_and_reports_for_kind, maybe_build_sync_local_image_decision_payload,
|
||||
build_local_image_stream_plan_and_reports_for_kind,
|
||||
build_local_image_sync_plan_and_reports_for_kind, is_openai_image_stream_request,
|
||||
maybe_build_stream_local_image_decision_payload, maybe_build_sync_local_image_decision_payload,
|
||||
};
|
||||
pub(crate) use self::video::{
|
||||
build_local_video_sync_plan_and_reports_for_kind, maybe_build_sync_local_video_decision_payload,
|
||||
|
||||
@@ -69,6 +69,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
|
||||
original_headers: &parts.headers,
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: false,
|
||||
has_envelope: false,
|
||||
needs_conversion: false,
|
||||
extra_fields,
|
||||
|
||||
@@ -22,7 +22,7 @@ fn applies_codex_defaults_when_body_rules_do_not_handle_fields() {
|
||||
assert!(body.get("top_p").is_none());
|
||||
assert!(body.get("metadata").is_none());
|
||||
assert_eq!(body["store"], false);
|
||||
assert_eq!(body["instructions"], "You are GPT-5.");
|
||||
assert_eq!(body["instructions"], "You are ChatGPT.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -125,10 +125,12 @@ fn injects_chatgpt_account_id_and_session_headers_for_codex_requests() {
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("user-agent"),
|
||||
Some(&"codex-tui/0.122.0 (Aether; x86_64) vscode/3.0.12 (codex-tui; 0.122.0)".to_string())
|
||||
Some(
|
||||
&"codex-tui/0.122.0 (Mac OS 15.2.0; arm64) vscode/2.6.11 (codex-tui; 0.122.0)"
|
||||
.to_string()
|
||||
)
|
||||
);
|
||||
assert_eq!(headers.get("version"), Some(&"0.122.0".to_string()));
|
||||
assert_eq!(headers.get("originator"), Some(&"codex_cli_rs".to_string()));
|
||||
assert_eq!(headers.get("originator"), Some(&"codex-tui".to_string()));
|
||||
assert_eq!(
|
||||
headers.get("session_id"),
|
||||
Some(&"ab5ecce4f0d110fe".to_string())
|
||||
@@ -168,10 +170,6 @@ fn respects_existing_codex_request_and_session_headers() {
|
||||
"user-agent",
|
||||
HeaderValue::from_static("user-specified-agent"),
|
||||
);
|
||||
original_headers.insert(
|
||||
"version",
|
||||
HeaderValue::from_static("user-specified-version"),
|
||||
);
|
||||
original_headers.insert(
|
||||
"originator",
|
||||
HeaderValue::from_static("user-specified-originator"),
|
||||
@@ -192,7 +190,6 @@ fn respects_existing_codex_request_and_session_headers() {
|
||||
Some(&"kept-by-rule-request".to_string())
|
||||
);
|
||||
assert!(!headers.contains_key("user-agent"));
|
||||
assert!(!headers.contains_key("version"));
|
||||
assert!(!headers.contains_key("originator"));
|
||||
assert_eq!(headers.get("session_id"), Some(&"kept-by-rule".to_string()));
|
||||
assert!(!headers.contains_key("conversation_id"));
|
||||
@@ -226,10 +223,12 @@ fn skips_conversation_id_for_compact_codex_requests() {
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("user-agent"),
|
||||
Some(&"codex-tui/0.122.0 (Aether; x86_64) vscode/3.0.12 (codex-tui; 0.122.0)".to_string())
|
||||
Some(
|
||||
&"codex-tui/0.122.0 (Mac OS 15.2.0; arm64) vscode/2.6.11 (codex-tui; 0.122.0)"
|
||||
.to_string()
|
||||
)
|
||||
);
|
||||
assert_eq!(headers.get("version"), Some(&"0.122.0".to_string()));
|
||||
assert_eq!(headers.get("originator"), Some(&"codex_cli_rs".to_string()));
|
||||
assert_eq!(headers.get("originator"), Some(&"codex-tui".to_string()));
|
||||
assert_eq!(
|
||||
headers.get("session_id"),
|
||||
Some(&"ab5ecce4f0d110fe".to_string())
|
||||
|
||||
@@ -75,6 +75,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
original_headers: &parts.headers,
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: spec_metadata.require_streaming,
|
||||
has_envelope: false,
|
||||
needs_conversion: true,
|
||||
extra_fields,
|
||||
|
||||
@@ -265,7 +265,7 @@ mod tests {
|
||||
|
||||
assert!(converted.get("metadata").is_none());
|
||||
assert_eq!(converted["store"], false);
|
||||
assert_eq!(converted["instructions"], "You are GPT-5.");
|
||||
assert_eq!(converted["instructions"], "You are ChatGPT.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -92,6 +92,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
|
||||
original_headers: &parts.headers,
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: upstream_is_stream,
|
||||
has_envelope: false,
|
||||
needs_conversion: matches!(
|
||||
resolved.conversion_mode,
|
||||
|
||||
@@ -96,6 +96,7 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
|
||||
original_headers: &parts.headers,
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: spec_metadata.require_streaming,
|
||||
has_envelope: resolved.is_antigravity,
|
||||
needs_conversion: matches!(
|
||||
resolved.conversion_mode,
|
||||
|
||||
@@ -53,13 +53,14 @@ pub(crate) use aether_ai_pipeline::api::{
|
||||
resolve_execution_runtime_stream_plan_kind, resolve_execution_runtime_sync_plan_kind,
|
||||
resolve_finalize_stream_rewrite_mode, resolve_gemini_files_stream_spec,
|
||||
resolve_gemini_files_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec,
|
||||
resolve_local_image_sync_spec, resolve_local_same_format_stream_spec,
|
||||
resolve_local_same_format_sync_spec, resolve_local_video_sync_spec,
|
||||
resolve_openai_chat_max_tokens, resolve_openai_cli_stream_spec, resolve_openai_cli_sync_spec,
|
||||
stream_body_contains_error_event, supports_stream_scheduler_decision_kind,
|
||||
supports_sync_scheduler_decision_kind, sync_chat_response_conversion_kind,
|
||||
sync_cli_response_conversion_kind, transform_provider_private_stream_line, value_as_u64,
|
||||
CanonicalStreamFrame, ClaudeClientEmitter, ClaudeProviderState, ExecutionRuntimeAuthContext,
|
||||
resolve_local_image_stream_spec, resolve_local_image_sync_spec,
|
||||
resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
|
||||
resolve_local_video_sync_spec, resolve_openai_chat_max_tokens, resolve_openai_cli_stream_spec,
|
||||
resolve_openai_cli_sync_spec, stream_body_contains_error_event,
|
||||
supports_stream_scheduler_decision_kind, supports_sync_scheduler_decision_kind,
|
||||
sync_chat_response_conversion_kind, sync_cli_response_conversion_kind,
|
||||
transform_provider_private_stream_line, value_as_u64, CanonicalStreamFrame,
|
||||
ClaudeClientEmitter, ClaudeProviderState, ExecutionRuntimeAuthContext,
|
||||
FinalizeStreamRewriteMode, GatewayControlPlanRequest, GatewayControlPlanResponse,
|
||||
GatewayControlSyncDecisionResponse, GeminiClientEmitter, GeminiProviderState,
|
||||
LocalCoreSyncErrorKind, LocalGeminiFilesSpec, LocalOpenAiCliSpec, LocalOpenAiImageSpec,
|
||||
@@ -76,12 +77,14 @@ pub(crate) use aether_ai_pipeline::api::{
|
||||
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,
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_ACTION,
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_CHAT_STREAM_PLAN_KIND,
|
||||
GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND, GEMINI_CHAT_SYNC_ERROR_REPORT_KIND,
|
||||
GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
|
||||
GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_STREAM_PLAN_KIND,
|
||||
CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND, CODEX_OPENAI_IMAGE_DEFAULT_MODEL,
|
||||
CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT, CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL,
|
||||
CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT, CODEX_OPENAI_IMAGE_INTERNAL_MODEL,
|
||||
EXECUTION_RUNTIME_STREAM_ACTION, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
|
||||
EXECUTION_RUNTIME_SYNC_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
|
||||
GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND,
|
||||
GEMINI_CHAT_SYNC_ERROR_REPORT_KIND, GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND,
|
||||
GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_STREAM_PLAN_KIND,
|
||||
GEMINI_CLI_STREAM_SUCCESS_REPORT_KIND, GEMINI_CLI_SYNC_ERROR_REPORT_KIND,
|
||||
GEMINI_CLI_SYNC_FINALIZE_REPORT_KIND, GEMINI_CLI_SYNC_PLAN_KIND,
|
||||
GEMINI_CLI_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME,
|
||||
@@ -96,7 +99,8 @@ pub(crate) use aether_ai_pipeline::api::{
|
||||
OPENAI_CLI_SYNC_FINALIZE_REPORT_KIND, OPENAI_CLI_SYNC_PLAN_KIND,
|
||||
OPENAI_CLI_SYNC_SUCCESS_REPORT_KIND, OPENAI_COMPACT_STREAM_PLAN_KIND,
|
||||
OPENAI_COMPACT_SYNC_ERROR_REPORT_KIND, OPENAI_COMPACT_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_COMPACT_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_COMPACT_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
|
||||
OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND,
|
||||
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
|
||||
|
||||
@@ -4,6 +4,7 @@ pub(crate) use crate::ai_pipeline::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
build_local_gemini_files_stream_plan_and_reports_for_kind,
|
||||
build_local_gemini_files_sync_plan_and_reports_for_kind,
|
||||
build_local_image_stream_plan_and_reports_for_kind,
|
||||
build_local_image_sync_plan_and_reports_for_kind,
|
||||
build_local_openai_chat_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_chat_sync_plan_and_reports_for_kind,
|
||||
@@ -30,14 +31,15 @@ pub(crate) use aether_ai_pipeline::api::{
|
||||
normalize_provider_private_report_context, normalize_provider_private_response_value,
|
||||
provider_private_response_allows_sync_finalize, resolve_claude_stream_spec,
|
||||
resolve_claude_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec,
|
||||
resolve_local_image_sync_spec, resolve_local_same_format_stream_spec,
|
||||
resolve_local_same_format_sync_spec, ExecutionRuntimeAuthContext, GatewayControlPlanRequest,
|
||||
GatewayControlPlanResponse, GatewayControlSyncDecisionResponse, LocalCoreSyncErrorKind,
|
||||
LocalOpenAiImageSpec, LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
|
||||
LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec,
|
||||
LocalStreamPlanAndReport, LocalSyncPlanAndReport, StreamingStandardTerminalObserver,
|
||||
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
|
||||
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
resolve_local_image_stream_spec, resolve_local_image_sync_spec,
|
||||
resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
|
||||
ExecutionRuntimeAuthContext, GatewayControlPlanRequest, GatewayControlPlanResponse,
|
||||
GatewayControlSyncDecisionResponse, LocalCoreSyncErrorKind, LocalOpenAiImageSpec,
|
||||
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec, LocalStandardSourceFamily,
|
||||
LocalStandardSourceMode, LocalStandardSpec, LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
StreamingStandardTerminalObserver, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_FILES_DOWNLOAD_PLAN_KIND,
|
||||
GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
|
||||
OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
@@ -83,8 +85,9 @@ pub(crate) fn is_matching_stream_request(
|
||||
plan_kind: &str,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
) -> bool {
|
||||
aether_ai_pipeline::api::is_matching_stream_request(plan_kind, parts.uri.path(), body_json)
|
||||
crate::ai_pipeline::planner_is_matching_stream_request(plan_kind, parts, body_json, body_base64)
|
||||
}
|
||||
|
||||
pub(crate) fn supports_sync_scheduler_decision_kind(plan_kind: &str) -> bool {
|
||||
|
||||
@@ -15,6 +15,7 @@ const AI_POST_ROUTE_PATTERNS: &[&str] = &[
|
||||
"/v1/responses/compact",
|
||||
"/v1/images/generations",
|
||||
"/v1/images/edits",
|
||||
"/v1/images/variations",
|
||||
];
|
||||
|
||||
const AI_ANY_ROUTE_PATTERNS: &[&str] = &[
|
||||
|
||||
2
apps/aether-gateway/src/cache/mod.rs
vendored
2
apps/aether-gateway/src/cache/mod.rs
vendored
@@ -3,6 +3,7 @@ mod auth_context;
|
||||
mod dashboard_response;
|
||||
mod direct_plan_bypass;
|
||||
mod scheduler_affinity;
|
||||
mod system_config;
|
||||
|
||||
pub(crate) use auth_api_key_last_used::AuthApiKeyLastUsedCache;
|
||||
pub(crate) use auth_context::AuthContextCache;
|
||||
@@ -11,3 +12,4 @@ pub(crate) use direct_plan_bypass::DirectPlanBypassCache;
|
||||
pub(crate) use scheduler_affinity::{
|
||||
SchedulerAffinityCache, SchedulerAffinitySnapshotEntry, SchedulerAffinityTarget,
|
||||
};
|
||||
pub(crate) use system_config::SystemConfigCache;
|
||||
|
||||
32
apps/aether-gateway/src/cache/system_config.rs
vendored
Normal file
32
apps/aether-gateway/src/cache/system_config.rs
vendored
Normal file
@@ -0,0 +1,32 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use aether_cache::ExpiringMap;
|
||||
|
||||
const MAX_ENTRIES: usize = 512;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct SystemConfigCache {
|
||||
entries: ExpiringMap<String, Option<serde_json::Value>>,
|
||||
}
|
||||
|
||||
impl Default for SystemConfigCache {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
entries: ExpiringMap::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SystemConfigCache {
|
||||
pub(crate) fn get(&self, key: &str, ttl: Duration) -> Option<Option<serde_json::Value>> {
|
||||
self.entries.get_fresh(&key.to_string(), ttl)
|
||||
}
|
||||
|
||||
pub(crate) fn insert(&self, key: String, value: Option<serde_json::Value>, ttl: Duration) {
|
||||
self.entries.insert(key, value, ttl, MAX_ENTRIES);
|
||||
}
|
||||
|
||||
pub(crate) fn clear(&self) {
|
||||
self.entries.clear();
|
||||
}
|
||||
}
|
||||
@@ -111,6 +111,7 @@ pub(crate) const RUST_FRONTDOOR_OWNED_ROUTE_PATTERNS: &[&str] = &[
|
||||
"/v1/chat/completions",
|
||||
"/v1/images/generations",
|
||||
"/v1/images/edits",
|
||||
"/v1/images/variations",
|
||||
"/v1/messages",
|
||||
"/v1/messages/count_tokens",
|
||||
"/v1/responses",
|
||||
|
||||
@@ -33,7 +33,7 @@ pub(super) fn classify_ai_public_route(
|
||||
} else if method == http::Method::POST
|
||||
&& matches!(
|
||||
normalized_path,
|
||||
"/v1/images/generations" | "/v1/images/edits"
|
||||
"/v1/images/generations" | "/v1/images/edits" | "/v1/images/variations"
|
||||
)
|
||||
{
|
||||
Some(classified(
|
||||
|
||||
@@ -710,6 +710,16 @@ impl GatewayDataState {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn find_provider_quotas_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderQuotaSnapshot>, DataLayerError> {
|
||||
match &self.provider_quota_reader {
|
||||
Some(repository) => repository.find_by_provider_ids(provider_ids).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
|
||||
pub(crate) async fn upsert_usage(
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use crate::ai_pipeline_api::{
|
||||
build_local_gemini_files_stream_plan_and_reports_for_kind,
|
||||
build_local_gemini_files_sync_plan_and_reports_for_kind,
|
||||
build_local_image_stream_plan_and_reports_for_kind,
|
||||
build_local_image_sync_plan_and_reports_for_kind,
|
||||
build_local_openai_chat_stream_plan_and_reports_for_kind,
|
||||
build_local_openai_chat_sync_plan_and_reports_for_kind,
|
||||
@@ -464,6 +465,33 @@ pub(crate) async fn maybe_execute_stream_via_local_gemini_files_decision(
|
||||
execute_stream_plan_and_reports(state, trace_id, decision, plan_kind, plan_and_reports).await
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_execute_stream_via_local_image_decision(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
body_base64: Option<&str>,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let plan_and_reports: Vec<LocalStreamPlanAndReport> =
|
||||
build_local_image_stream_plan_and_reports_for_kind(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
)
|
||||
.await?;
|
||||
if plan_and_reports.is_empty() {
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
execute_stream_plan_and_reports(state, trace_id, decision, plan_kind, plan_and_reports).await
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_execute_sync_via_local_video_decision(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
|
||||
@@ -74,8 +74,15 @@ pub(crate) async fn maybe_execute_stream_via_plan_fallback(
|
||||
_bypass_cache_key: String,
|
||||
_fallback_reason: GatewayFallbackReason,
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let Some(payload) =
|
||||
maybe_build_stream_plan_payload(state, parts, trace_id, decision, body_json).await?
|
||||
let Some(payload) = maybe_build_stream_plan_payload(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
body_json,
|
||||
body_base64.as_deref(),
|
||||
)
|
||||
.await?
|
||||
else {
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
@@ -14,6 +14,7 @@ use crate::{AppState, GatewayError, GatewayFallbackReason};
|
||||
use super::{
|
||||
build_direct_plan_bypass_cache_key, execute_stream_plan_and_reports,
|
||||
maybe_execute_stream_via_local_decision, maybe_execute_stream_via_local_gemini_files_decision,
|
||||
maybe_execute_stream_via_local_image_decision,
|
||||
maybe_execute_stream_via_local_openai_cli_decision,
|
||||
maybe_execute_stream_via_local_same_format_provider_decision,
|
||||
maybe_execute_stream_via_local_standard_decision, maybe_execute_stream_via_plan_fallback,
|
||||
@@ -36,7 +37,7 @@ pub(crate) async fn maybe_execute_via_stream_decision_path(
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
if !is_matching_stream_request(plan_kind, parts, &body_json) {
|
||||
if !is_matching_stream_request(plan_kind, parts, &body_json, body_base64.as_deref()) {
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
@@ -59,6 +60,24 @@ pub(crate) async fn maybe_execute_via_stream_decision_path(
|
||||
}
|
||||
|
||||
if supports_stream_scheduler_decision_kind(plan_kind) {
|
||||
match maybe_execute_stream_via_local_image_decision(
|
||||
state,
|
||||
parts,
|
||||
&body_json,
|
||||
body_base64.as_deref(),
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
|
||||
match maybe_execute_stream_via_local_decision(
|
||||
state, parts, trace_id, decision, &body_json, plan_kind,
|
||||
)
|
||||
|
||||
@@ -45,7 +45,7 @@ pub(crate) async fn maybe_execute_via_sync_decision_path(
|
||||
};
|
||||
|
||||
if let Some(stream_plan_kind) = resolve_execution_runtime_stream_plan_kind(parts, decision) {
|
||||
if is_matching_stream_request(stream_plan_kind, parts, &body_json) {
|
||||
if is_matching_stream_request(stream_plan_kind, parts, &body_json, body_base64.as_deref()) {
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::handlers::public::provider_key_api_formats;
|
||||
use crate::provider_key_auth::provider_key_effective_api_formats;
|
||||
use aether_scheduler_core::count_recent_rpm_requests_for_provider_key_since;
|
||||
use serde_json::json;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
@@ -18,6 +18,16 @@ pub(crate) async fn build_admin_key_health_payload(
|
||||
.await
|
||||
.ok()
|
||||
.and_then(|mut keys| keys.drain(..).next())?;
|
||||
let provider = state
|
||||
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&key.provider_id))
|
||||
.await
|
||||
.ok()
|
||||
.and_then(|mut providers| providers.drain(..).next())?;
|
||||
let endpoints = state
|
||||
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&key.provider_id))
|
||||
.await
|
||||
.ok()
|
||||
.unwrap_or_default();
|
||||
|
||||
let request_count = key.request_count.unwrap_or(0);
|
||||
let success_count = key.success_count.unwrap_or(0);
|
||||
@@ -97,7 +107,9 @@ pub(crate) async fn build_admin_key_health_payload(
|
||||
.unwrap_or(0));
|
||||
} else {
|
||||
let mut formats_payload = serde_json::Map::new();
|
||||
for format_name in provider_key_api_formats(&key) {
|
||||
for format_name in
|
||||
provider_key_effective_api_formats(&key, &provider.provider_type, &endpoints)
|
||||
{
|
||||
let health_data = health_by_format.and_then(|formats| formats.get(&format_name));
|
||||
let circuit_data = circuit_by_format.and_then(|formats| formats.get(&format_name));
|
||||
let is_open = circuit_data
|
||||
@@ -363,11 +375,27 @@ pub(crate) async fn recover_all_admin_key_health(
|
||||
if !updated {
|
||||
continue;
|
||||
}
|
||||
let provider = state
|
||||
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&key.provider_id))
|
||||
.await
|
||||
.ok()
|
||||
.and_then(|mut providers| providers.drain(..).next());
|
||||
let endpoints = state
|
||||
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&key.provider_id))
|
||||
.await
|
||||
.ok()
|
||||
.unwrap_or_default();
|
||||
let api_formats = provider
|
||||
.as_ref()
|
||||
.map(|provider| {
|
||||
provider_key_effective_api_formats(&key, &provider.provider_type, &endpoints)
|
||||
})
|
||||
.unwrap_or_default();
|
||||
payload_items.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"provider_id": key.provider_id,
|
||||
"api_formats": key.api_formats.unwrap_or_else(|| json!([])),
|
||||
"api_formats": api_formats,
|
||||
}));
|
||||
}
|
||||
|
||||
|
||||
@@ -1,9 +1,8 @@
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::handlers::admin::shared::unix_secs_to_rfc3339;
|
||||
use crate::handlers::public::{
|
||||
api_format_display_name, build_public_health_timeline, provider_key_api_formats,
|
||||
};
|
||||
use crate::handlers::public::{api_format_display_name, build_public_health_timeline};
|
||||
use crate::handlers::shared::unix_ms_to_rfc3339;
|
||||
use crate::provider_key_auth::provider_key_effective_api_formats;
|
||||
use aether_data_contracts::repository::candidates::PublicHealthTimelineBucket;
|
||||
use aether_scheduler_core::{is_provider_key_circuit_open, provider_key_health_score};
|
||||
use serde_json::json;
|
||||
@@ -50,17 +49,26 @@ pub(crate) async fn build_admin_endpoint_health_status_payload(
|
||||
let mut endpoint_to_format = BTreeMap::<String, String>::new();
|
||||
let mut provider_ids_by_format = BTreeMap::<String, BTreeSet<String>>::new();
|
||||
let mut active_provider_formats = BTreeSet::<(String, String)>::new();
|
||||
let provider_type_by_id = providers
|
||||
.iter()
|
||||
.map(|provider| (provider.id.clone(), provider.provider_type.clone()))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let mut active_endpoints_by_provider = BTreeMap::<String, Vec<_>>::new();
|
||||
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_provider_formats.insert((endpoint.provider_id, endpoint.api_format));
|
||||
active_provider_formats.insert((endpoint.provider_id.clone(), endpoint.api_format.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<_>>();
|
||||
|
||||
@@ -74,7 +82,15 @@ pub(crate) async fn build_admin_endpoint_health_status_payload(
|
||||
.ok()
|
||||
.unwrap_or_default();
|
||||
for key in keys {
|
||||
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 !active_provider_formats.contains(&(key.provider_id.clone(), api_format.clone()))
|
||||
{
|
||||
continue;
|
||||
|
||||
@@ -122,7 +122,13 @@ pub(crate) async fn build_admin_global_model_routing_payload(
|
||||
.cloned()
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.filter(|key| provider_catalog_key_supports_format(key, &endpoint.api_format))
|
||||
.filter(|key| {
|
||||
provider_catalog_key_supports_format(
|
||||
key,
|
||||
provider.provider_type.as_str(),
|
||||
&endpoint.api_format,
|
||||
)
|
||||
})
|
||||
.filter(|key| {
|
||||
key_allowed_models_match_global_model_for_routing(
|
||||
key.allowed_models.as_ref(),
|
||||
|
||||
@@ -5,7 +5,9 @@ use crate::handlers::admin::provider::shared::paths::{
|
||||
use crate::handlers::admin::provider::shared::payloads::{
|
||||
AdminProviderCreateRequest, AdminProviderUpdatePatch,
|
||||
};
|
||||
use crate::handlers::admin::provider::write::provider::build_admin_fixed_provider_endpoint_record;
|
||||
use crate::handlers::admin::provider::write::provider::{
|
||||
reconcile_admin_fixed_provider_template_endpoints, reconcile_admin_fixed_provider_template_keys,
|
||||
};
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::handlers::admin::shared::attach_admin_audit_response;
|
||||
use crate::GatewayError;
|
||||
@@ -73,24 +75,12 @@ pub(crate) async fn maybe_build_local_admin_provider_writes_response(
|
||||
return Ok(Some(build_admin_providers_data_unavailable_response()));
|
||||
};
|
||||
|
||||
if let Some((base_url, endpoint_signatures)) =
|
||||
state.fixed_provider_template(&created_provider.provider_type)
|
||||
if state
|
||||
.fixed_provider_template(&created_provider.provider_type)
|
||||
.is_some()
|
||||
{
|
||||
for endpoint_signature in endpoint_signatures {
|
||||
let endpoint = match build_admin_fixed_provider_endpoint_record(
|
||||
&created_provider,
|
||||
endpoint_signature,
|
||||
base_url,
|
||||
) {
|
||||
Ok(endpoint) => endpoint,
|
||||
Err(message) => {
|
||||
return Ok(Some(build_admin_provider_bad_request_response(message)));
|
||||
}
|
||||
};
|
||||
let Some(_) = state.create_provider_catalog_endpoint(&endpoint).await? else {
|
||||
return Ok(Some(build_admin_providers_data_unavailable_response()));
|
||||
};
|
||||
}
|
||||
reconcile_admin_fixed_provider_template_endpoints(state, &created_provider).await?;
|
||||
reconcile_admin_fixed_provider_template_keys(state, &created_provider).await?;
|
||||
}
|
||||
return Ok(Some(attach_admin_audit_response(
|
||||
Json(json!({
|
||||
@@ -167,6 +157,13 @@ pub(crate) async fn maybe_build_local_admin_provider_writes_response(
|
||||
else {
|
||||
return Ok(Some(build_admin_providers_data_unavailable_response()));
|
||||
};
|
||||
if state
|
||||
.fixed_provider_template(&updated_record.provider_type)
|
||||
.is_some()
|
||||
{
|
||||
reconcile_admin_fixed_provider_template_endpoints(state, &updated_record).await?;
|
||||
reconcile_admin_fixed_provider_template_keys(state, &updated_record).await?;
|
||||
}
|
||||
return Ok(Some(
|
||||
match state
|
||||
.build_admin_provider_summary_payload(&provider_id)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use crate::handlers::admin::provider::shared::paths::admin_provider_id_for_keys;
|
||||
use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyCreateRequest;
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::provider_key_auth::provider_key_effective_api_formats;
|
||||
use crate::{model_fetch::perform_model_fetch_for_key, GatewayError};
|
||||
use axum::{
|
||||
body::{Body, Bytes},
|
||||
@@ -94,11 +95,17 @@ pub(super) async fn maybe_handle(
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or(0);
|
||||
let endpoints = state
|
||||
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
|
||||
.await?;
|
||||
let api_formats =
|
||||
provider_key_effective_api_formats(&created, &provider.provider_type, &endpoints);
|
||||
|
||||
Ok(Some(
|
||||
Json(state.build_admin_provider_key_response(
|
||||
&created,
|
||||
&provider.provider_type,
|
||||
&api_formats,
|
||||
now_unix_secs,
|
||||
))
|
||||
.into_response(),
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use crate::handlers::admin::provider::shared::paths::admin_update_key_id;
|
||||
use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyUpdatePatch;
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::provider_key_auth::provider_key_effective_api_formats;
|
||||
use crate::{model_fetch::perform_model_fetch_for_key, GatewayError};
|
||||
use axum::{
|
||||
body::{Body, Bytes},
|
||||
@@ -117,11 +118,17 @@ pub(super) async fn maybe_handle(
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or(0);
|
||||
let endpoints = state
|
||||
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
|
||||
.await?;
|
||||
let api_formats =
|
||||
provider_key_effective_api_formats(&updated, &provider.provider_type, &endpoints);
|
||||
|
||||
Ok(Some(
|
||||
Json(state.build_admin_provider_key_response(
|
||||
&updated,
|
||||
&provider.provider_type,
|
||||
&api_formats,
|
||||
now_unix_secs,
|
||||
))
|
||||
.into_response(),
|
||||
|
||||
@@ -201,8 +201,10 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
|
||||
match update_existing_provider_oauth_catalog_key(
|
||||
state,
|
||||
&existing_key,
|
||||
provider_type,
|
||||
&access_token,
|
||||
&auth_config,
|
||||
&api_formats,
|
||||
None,
|
||||
expires_at,
|
||||
)
|
||||
@@ -242,6 +244,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
|
||||
match create_provider_oauth_catalog_key(
|
||||
state,
|
||||
provider_id,
|
||||
provider_type,
|
||||
key_name.as_str(),
|
||||
&access_token,
|
||||
&auth_config,
|
||||
|
||||
@@ -57,6 +57,7 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
|
||||
let endpoints = state
|
||||
.list_provider_catalog_endpoints_by_provider_ids(&[provider_id.to_string()])
|
||||
.await?;
|
||||
let api_formats = provider_oauth_active_api_formats(&endpoints);
|
||||
let runtime_endpoint = provider_oauth_runtime_endpoint_for_provider("kiro", &endpoints);
|
||||
let request_proxy = state
|
||||
.resolve_admin_provider_oauth_operation_proxy_snapshot(
|
||||
@@ -191,8 +192,10 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
|
||||
match update_existing_provider_oauth_catalog_key(
|
||||
state,
|
||||
&existing_key,
|
||||
provider.provider_type.as_str(),
|
||||
&access_token,
|
||||
&auth_config,
|
||||
&api_formats,
|
||||
None,
|
||||
refreshed_auth_config.expires_at,
|
||||
)
|
||||
@@ -223,10 +226,11 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
|
||||
match create_provider_oauth_catalog_key(
|
||||
state,
|
||||
provider_id,
|
||||
provider.provider_type.as_str(),
|
||||
&key_name,
|
||||
&access_token,
|
||||
&auth_config,
|
||||
&provider_oauth_active_api_formats(&endpoints),
|
||||
&api_formats,
|
||||
None,
|
||||
refreshed_auth_config.expires_at,
|
||||
)
|
||||
|
||||
@@ -174,8 +174,10 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
|
||||
match state
|
||||
.update_existing_provider_oauth_catalog_key(
|
||||
&existing_key,
|
||||
&provider_type,
|
||||
&access_token,
|
||||
&auth_config,
|
||||
&api_formats,
|
||||
None,
|
||||
expires_at,
|
||||
)
|
||||
@@ -213,6 +215,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
|
||||
match state
|
||||
.create_provider_oauth_catalog_key(
|
||||
&provider_id,
|
||||
&provider_type,
|
||||
&name,
|
||||
&access_token,
|
||||
&auth_config,
|
||||
|
||||
@@ -351,8 +351,10 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
|
||||
match state
|
||||
.update_existing_provider_oauth_catalog_key(
|
||||
&existing_key,
|
||||
&provider.provider_type,
|
||||
&access_token,
|
||||
&auth_config,
|
||||
&api_formats,
|
||||
key_proxy.clone(),
|
||||
Some(expires_at),
|
||||
)
|
||||
@@ -374,6 +376,7 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
|
||||
match state
|
||||
.create_provider_oauth_catalog_key(
|
||||
&provider_id,
|
||||
&provider.provider_type,
|
||||
&key_name,
|
||||
&access_token,
|
||||
&auth_config,
|
||||
|
||||
@@ -166,8 +166,10 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
||||
match state
|
||||
.update_existing_provider_oauth_catalog_key(
|
||||
&existing_key,
|
||||
&provider_type,
|
||||
&access_token,
|
||||
&auth_config,
|
||||
&api_formats,
|
||||
None,
|
||||
expires_at,
|
||||
)
|
||||
@@ -204,6 +206,7 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
||||
match state
|
||||
.create_provider_oauth_catalog_key(
|
||||
&provider_id,
|
||||
&provider_type,
|
||||
&name,
|
||||
&access_token,
|
||||
&auth_config,
|
||||
|
||||
@@ -2,12 +2,13 @@ use super::state::{
|
||||
enrich_admin_provider_oauth_auth_config, json_non_empty_string, json_u64_value,
|
||||
};
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::provider_key_auth::provider_active_api_formats;
|
||||
use crate::GatewayError;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
};
|
||||
use aether_provider_transport::provider_types::provider_type_is_fixed;
|
||||
use serde_json::json;
|
||||
use std::collections::BTreeSet;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use uuid::Uuid;
|
||||
|
||||
@@ -23,16 +24,7 @@ pub(crate) fn provider_oauth_key_proxy_value(
|
||||
pub(crate) fn provider_oauth_active_api_formats(
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
) -> Vec<String> {
|
||||
let mut formats = Vec::new();
|
||||
let mut seen = BTreeSet::new();
|
||||
for endpoint in endpoints.iter().filter(|endpoint| endpoint.is_active) {
|
||||
let api_format = endpoint.api_format.trim();
|
||||
if api_format.is_empty() || !seen.insert(api_format.to_string()) {
|
||||
continue;
|
||||
}
|
||||
formats.push(api_format.to_string());
|
||||
}
|
||||
formats
|
||||
provider_active_api_formats(endpoints)
|
||||
}
|
||||
|
||||
pub(crate) fn build_provider_oauth_auth_config_from_token_payload(
|
||||
@@ -76,6 +68,7 @@ pub(crate) fn build_provider_oauth_auth_config_from_token_payload(
|
||||
pub(crate) async fn create_provider_oauth_catalog_key(
|
||||
state: &AdminAppState<'_>,
|
||||
provider_id: &str,
|
||||
provider_type: &str,
|
||||
name: &str,
|
||||
access_token: &str,
|
||||
auth_config: &serde_json::Map<String, serde_json::Value>,
|
||||
@@ -108,7 +101,7 @@ pub(crate) async fn create_provider_oauth_catalog_key(
|
||||
)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?
|
||||
.with_transport_fields(
|
||||
Some(json!(api_formats)),
|
||||
provider_oauth_catalog_key_api_formats(provider_type, api_formats),
|
||||
encrypted_api_key,
|
||||
Some(encrypted_auth_config),
|
||||
None,
|
||||
@@ -136,8 +129,10 @@ pub(crate) async fn create_provider_oauth_catalog_key(
|
||||
pub(crate) async fn update_existing_provider_oauth_catalog_key(
|
||||
state: &AdminAppState<'_>,
|
||||
existing_key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
access_token: &str,
|
||||
auth_config: &serde_json::Map<String, serde_json::Value>,
|
||||
api_formats: &[String],
|
||||
proxy: Option<serde_json::Value>,
|
||||
expires_at_unix_secs: Option<u64>,
|
||||
) -> Result<Option<StoredProviderCatalogKey>, GatewayError> {
|
||||
@@ -159,6 +154,7 @@ pub(crate) async fn update_existing_provider_oauth_catalog_key(
|
||||
let mut updated = existing_key.clone();
|
||||
updated.encrypted_api_key = encrypted_api_key;
|
||||
updated.encrypted_auth_config = Some(encrypted_auth_config);
|
||||
updated.api_formats = provider_oauth_catalog_key_api_formats(provider_type, api_formats);
|
||||
updated.is_active = true;
|
||||
updated.expires_at_unix_secs = expires_at_unix_secs;
|
||||
updated.oauth_invalid_at_unix_secs = None;
|
||||
@@ -172,3 +168,14 @@ pub(crate) async fn update_existing_provider_oauth_catalog_key(
|
||||
updated.updated_at_unix_secs = Some(now_unix_secs);
|
||||
state.update_provider_catalog_key(&updated).await
|
||||
}
|
||||
|
||||
fn provider_oauth_catalog_key_api_formats(
|
||||
provider_type: &str,
|
||||
api_formats: &[String],
|
||||
) -> Option<serde_json::Value> {
|
||||
if provider_type_is_fixed(provider_type) {
|
||||
None
|
||||
} else {
|
||||
Some(json!(api_formats))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,26 +3,14 @@ use crate::handlers::admin::provider::shared::support::{
|
||||
};
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::handlers::admin::shared::{provider_key_status_snapshot_payload, unix_secs_to_rfc3339};
|
||||
use crate::provider_key_auth::provider_key_auth_semantics;
|
||||
use crate::provider_key_auth::{provider_key_auth_semantics, provider_key_effective_api_formats};
|
||||
use aether_admin::provider::pool as admin_provider_pool_pure;
|
||||
use aether_admin::provider::quota as admin_provider_quota_pure;
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
pub(super) fn admin_pool_api_formats(key: &StoredProviderCatalogKey) -> Vec<String> {
|
||||
key.api_formats
|
||||
.as_ref()
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.map(|values| {
|
||||
values
|
||||
.iter()
|
||||
.filter_map(serde_json::Value::as_str)
|
||||
.map(ToOwned::to_owned)
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn admin_pool_string_list(value: Option<&serde_json::Value>) -> Option<Vec<String>> {
|
||||
let values = value
|
||||
.and_then(serde_json::Value::as_array)
|
||||
@@ -749,6 +737,7 @@ fn admin_pool_scheduling_payload(
|
||||
pub(super) fn build_admin_pool_key_payload(
|
||||
state: &AdminAppState<'_>,
|
||||
provider_type: &str,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
key: &StoredProviderCatalogKey,
|
||||
runtime: &AdminProviderPoolRuntimeState,
|
||||
pool_config: Option<AdminProviderPoolConfig>,
|
||||
@@ -915,7 +904,11 @@ pub(super) fn build_admin_pool_key_payload(
|
||||
);
|
||||
payload.insert(
|
||||
"api_formats".to_string(),
|
||||
json!(admin_pool_api_formats(key)),
|
||||
json!(provider_key_effective_api_formats(
|
||||
key,
|
||||
provider_type,
|
||||
endpoints,
|
||||
)),
|
||||
);
|
||||
payload.insert(
|
||||
"rate_multipliers".to_string(),
|
||||
|
||||
@@ -124,6 +124,9 @@ pub(super) async fn build_admin_pool_list_keys_response(
|
||||
};
|
||||
|
||||
let key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
|
||||
let endpoints = state
|
||||
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
|
||||
.await?;
|
||||
let runtime = match (state.redis_kv_runner(), pool_config.as_ref()) {
|
||||
(Some(runner), Some(pool_config)) if !key_ids.is_empty() => {
|
||||
read_admin_provider_pool_runtime_state(
|
||||
@@ -143,6 +146,7 @@ pub(super) async fn build_admin_pool_list_keys_response(
|
||||
pool_payloads::build_admin_pool_key_payload(
|
||||
state,
|
||||
&provider.provider_type,
|
||||
&endpoints,
|
||||
&key,
|
||||
&runtime,
|
||||
pool_config.clone(),
|
||||
|
||||
@@ -14,6 +14,9 @@ use crate::ai_pipeline::{maybe_build_sync_finalize_outcome, GatewayControlDecisi
|
||||
use crate::execution_runtime;
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
|
||||
use crate::model_fetch::ModelFetchRuntimeState;
|
||||
use crate::provider_key_auth::{
|
||||
provider_key_configured_api_formats, provider_key_inherits_provider_api_formats,
|
||||
};
|
||||
use crate::provider_transport::kiro::{
|
||||
build_kiro_generate_assistant_response_url, build_kiro_provider_headers,
|
||||
build_kiro_provider_request_body, supports_local_kiro_request_transport_with_network,
|
||||
@@ -272,9 +275,13 @@ fn provider_query_select_kiro_endpoint<'a>(
|
||||
|
||||
fn provider_query_key_supports_endpoint(
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
endpoint_api_format: &str,
|
||||
) -> bool {
|
||||
let formats = json_string_list(key.api_formats.as_ref());
|
||||
if provider_key_inherits_provider_api_formats(key, provider_type) {
|
||||
return true;
|
||||
}
|
||||
let formats = provider_key_configured_api_formats(key);
|
||||
formats.is_empty()
|
||||
|| formats
|
||||
.iter()
|
||||
@@ -321,7 +328,11 @@ async fn provider_query_select_preferred_non_kiro_endpoint(
|
||||
for key in keys {
|
||||
if !key.is_active
|
||||
|| selected_key_id.is_some_and(|value| value != key.id.as_str())
|
||||
|| !provider_query_key_supports_endpoint(key, &endpoint.api_format)
|
||||
|| !provider_query_key_supports_endpoint(
|
||||
key,
|
||||
&provider.provider_type,
|
||||
&endpoint.api_format,
|
||||
)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
@@ -348,7 +359,11 @@ async fn provider_query_select_preferred_non_kiro_endpoint(
|
||||
for key in keys {
|
||||
if !key.is_active
|
||||
|| selected_key_id.is_some_and(|value| value != key.id.as_str())
|
||||
|| !provider_query_key_supports_endpoint(key, &endpoint.api_format)
|
||||
|| !provider_query_key_supports_endpoint(
|
||||
key,
|
||||
&provider.provider_type,
|
||||
&endpoint.api_format,
|
||||
)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
@@ -375,7 +390,11 @@ async fn provider_query_select_preferred_non_kiro_endpoint(
|
||||
&& keys.iter().any(|key| {
|
||||
key.is_active
|
||||
&& selected_key_id.is_none_or(|value| value == key.id.as_str())
|
||||
&& provider_query_key_supports_endpoint(key, &endpoint.api_format)
|
||||
&& provider_query_key_supports_endpoint(
|
||||
key,
|
||||
&provider.provider_type,
|
||||
&endpoint.api_format,
|
||||
)
|
||||
})
|
||||
})
|
||||
.or_else(|| endpoints.iter().find(|endpoint| endpoint.is_active))
|
||||
@@ -531,7 +550,13 @@ async fn provider_query_build_kiro_test_candidates(
|
||||
ADMIN_PROVIDER_QUERY_API_KEY_NOT_FOUND_DETAIL,
|
||||
));
|
||||
};
|
||||
if !key.is_active || !provider_query_key_supports_endpoint(key, &endpoint.api_format) {
|
||||
if !key.is_active
|
||||
|| !provider_query_key_supports_endpoint(
|
||||
key,
|
||||
&provider.provider_type,
|
||||
&endpoint.api_format,
|
||||
)
|
||||
{
|
||||
return Err(build_admin_provider_query_not_found_response(
|
||||
ADMIN_PROVIDER_QUERY_NO_ACTIVE_TEST_CANDIDATE_DETAIL,
|
||||
));
|
||||
@@ -567,7 +592,9 @@ async fn provider_query_build_kiro_test_candidates(
|
||||
.as_deref()
|
||||
.is_none_or(|value| value == key.id.as_str())
|
||||
})
|
||||
.filter(|key| provider_query_key_supports_endpoint(key, &endpoint.api_format))
|
||||
.filter(|key| {
|
||||
provider_query_key_supports_endpoint(key, &provider.provider_type, &endpoint.api_format)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
keys.sort_by_key(|key| {
|
||||
provider_query_test_key_sort_key(provider.provider_type.as_str(), key, &endpoint.api_format)
|
||||
|
||||
@@ -3,8 +3,6 @@ use crate::handlers::admin::request::AdminAppState;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
};
|
||||
use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot;
|
||||
use futures_util::future::join_all;
|
||||
use serde_json::json;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
@@ -235,17 +233,6 @@ pub(crate) async fn build_admin_providers_summary_payload(
|
||||
.or_default()
|
||||
.insert(row.global_model_id);
|
||||
}
|
||||
let quota_snapshots_by_provider = join_all(provider_ids.iter().map(|provider_id| async {
|
||||
let quota_snapshot = state
|
||||
.read_provider_quota_snapshot(provider_id)
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
(provider_id.clone(), quota_snapshot)
|
||||
}))
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<BTreeMap<String, Option<StoredProviderQuotaSnapshot>>>();
|
||||
let now_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
@@ -268,9 +255,7 @@ pub(crate) async fn build_admin_providers_summary_payload(
|
||||
.get(&provider.id)
|
||||
.map(Vec::as_slice)
|
||||
.unwrap_or(&[]),
|
||||
quota_snapshots_by_provider
|
||||
.get(&provider.id)
|
||||
.and_then(Option::as_ref),
|
||||
None,
|
||||
model_stats_by_provider.get(&provider.id),
|
||||
active_global_model_ids,
|
||||
now_unix_secs,
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
use crate::handlers::admin::shared::unix_secs_to_rfc3339;
|
||||
use crate::handlers::public::{
|
||||
provider_key_api_formats, request_candidate_event_unix_ms, request_candidate_status_label,
|
||||
};
|
||||
use crate::handlers::public::{request_candidate_event_unix_ms, request_candidate_status_label};
|
||||
use crate::provider_key_auth::provider_key_effective_api_formats;
|
||||
use aether_data_contracts::repository::candidates::{
|
||||
RequestCandidateStatus, StoredRequestCandidate,
|
||||
};
|
||||
@@ -66,7 +65,9 @@ pub(crate) fn build_admin_provider_summary_value(
|
||||
keys_by_endpoint.entry(endpoint.id.clone()).or_default();
|
||||
}
|
||||
for key in keys {
|
||||
for api_format in provider_key_api_formats(key) {
|
||||
for api_format in
|
||||
provider_key_effective_api_formats(key, &provider.provider_type, endpoints)
|
||||
{
|
||||
if let Some(endpoint_id) = format_to_endpoint_id.get(&api_format) {
|
||||
keys_by_endpoint
|
||||
.entry(endpoint_id.clone())
|
||||
@@ -131,6 +132,26 @@ pub(crate) fn build_admin_provider_summary_value(
|
||||
.and_then(|cfg| cfg.get("architecture_id"))
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(ToOwned::to_owned);
|
||||
let billing_type = quota_snapshot
|
||||
.map(|quota| quota.billing_type.clone())
|
||||
.or_else(|| provider.billing_type.clone());
|
||||
let monthly_quota_usd = quota_snapshot
|
||||
.and_then(|quota| quota.monthly_quota_usd)
|
||||
.or(provider.monthly_quota_usd);
|
||||
let monthly_used_usd = quota_snapshot
|
||||
.map(|quota| quota.monthly_used_usd)
|
||||
.or(provider.monthly_used_usd);
|
||||
let quota_reset_day = quota_snapshot
|
||||
.and_then(|quota| quota.quota_reset_day)
|
||||
.or(provider.quota_reset_day);
|
||||
let quota_last_reset_at = quota_snapshot
|
||||
.and_then(|quota| quota.quota_last_reset_at_unix_secs)
|
||||
.or(provider.quota_last_reset_at_unix_secs)
|
||||
.and_then(unix_secs_to_rfc3339);
|
||||
let quota_expires_at = quota_snapshot
|
||||
.and_then(|quota| quota.quota_expires_at_unix_secs)
|
||||
.or(provider.quota_expires_at_unix_secs)
|
||||
.and_then(unix_secs_to_rfc3339);
|
||||
|
||||
json!({
|
||||
"id": provider.id.clone(),
|
||||
@@ -142,16 +163,12 @@ pub(crate) fn build_admin_provider_summary_value(
|
||||
"keep_priority_on_conversion": provider.keep_priority_on_conversion,
|
||||
"enable_format_conversion": provider.enable_format_conversion,
|
||||
"is_active": provider.is_active,
|
||||
"billing_type": quota_snapshot.map(|quota| quota.billing_type.clone()),
|
||||
"monthly_quota_usd": quota_snapshot.and_then(|quota| quota.monthly_quota_usd),
|
||||
"monthly_used_usd": quota_snapshot.map(|quota| quota.monthly_used_usd),
|
||||
"quota_reset_day": quota_snapshot.and_then(|quota| quota.quota_reset_day),
|
||||
"quota_last_reset_at": quota_snapshot
|
||||
.and_then(|quota| quota.quota_last_reset_at_unix_secs)
|
||||
.and_then(unix_secs_to_rfc3339),
|
||||
"quota_expires_at": quota_snapshot
|
||||
.and_then(|quota| quota.quota_expires_at_unix_secs)
|
||||
.and_then(unix_secs_to_rfc3339),
|
||||
"billing_type": billing_type,
|
||||
"monthly_quota_usd": monthly_quota_usd,
|
||||
"monthly_used_usd": monthly_used_usd,
|
||||
"quota_reset_day": quota_reset_day,
|
||||
"quota_last_reset_at": quota_last_reset_at,
|
||||
"quota_expires_at": quota_expires_at,
|
||||
"max_retries": provider.max_retries,
|
||||
"proxy": provider.proxy.clone(),
|
||||
"stream_first_byte_timeout": provider.stream_first_byte_timeout_secs,
|
||||
|
||||
@@ -10,6 +10,7 @@ use crate::handlers::admin::shared::{
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_provider_transport::provider_types::provider_type_is_fixed;
|
||||
use serde_json::json;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use uuid::Uuid;
|
||||
@@ -129,6 +130,8 @@ pub(crate) async fn build_admin_create_provider_key_record(
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or(0);
|
||||
let inherits_provider_api_formats =
|
||||
auth_type == "oauth" && provider_type_is_fixed(&provider.provider_type);
|
||||
let mut key = StoredProviderCatalogKey::new(
|
||||
Uuid::new_v4().to_string(),
|
||||
provider.id.clone(),
|
||||
@@ -139,7 +142,11 @@ pub(crate) async fn build_admin_create_provider_key_record(
|
||||
)
|
||||
.map_err(|err| err.to_string())?
|
||||
.with_transport_fields(
|
||||
Some(json!(api_formats)),
|
||||
if inherits_provider_api_formats {
|
||||
None
|
||||
} else {
|
||||
Some(json!(api_formats))
|
||||
},
|
||||
encrypted_api_key,
|
||||
encrypted_auth_config,
|
||||
normalize_json_object(payload.rate_multipliers, "rate_multipliers")?,
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::provider_key_auth::provider_key_effective_api_formats;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery,
|
||||
};
|
||||
@@ -29,6 +30,11 @@ pub(crate) async fn build_admin_provider_keys_payload(
|
||||
})
|
||||
.await
|
||||
.ok()?;
|
||||
let endpoints = state
|
||||
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
|
||||
.await
|
||||
.ok()
|
||||
.unwrap_or_default();
|
||||
let now_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
@@ -39,9 +45,12 @@ pub(crate) async fn build_admin_provider_keys_payload(
|
||||
.items
|
||||
.into_iter()
|
||||
.map(|key| {
|
||||
let api_formats =
|
||||
provider_key_effective_api_formats(&key, &provider.provider_type, &endpoints);
|
||||
state.build_admin_provider_key_response(
|
||||
&key,
|
||||
&provider.provider_type,
|
||||
&api_formats,
|
||||
now_unix_secs,
|
||||
)
|
||||
})
|
||||
|
||||
@@ -7,9 +7,11 @@ use crate::handlers::admin::shared::{
|
||||
decrypt_catalog_secret_with_fallbacks, encrypt_catalog_secret_with_fallbacks, json_string_list,
|
||||
normalize_json_object, normalize_string_list, parse_catalog_auth_config_json,
|
||||
};
|
||||
use crate::provider_key_auth::provider_key_is_oauth_managed;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_provider_transport::provider_types::provider_type_is_fixed;
|
||||
use serde_json::json;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
@@ -35,6 +37,9 @@ pub(crate) async fn build_admin_update_provider_key_record(
|
||||
.auth_type
|
||||
.as_deref()
|
||||
.is_some_and(|_| target_auth_type != current_auth_type);
|
||||
let managed_fixed_oauth_key = provider_type_is_fixed(&provider.provider_type)
|
||||
&& (provider_key_is_oauth_managed(existing, &provider.provider_type)
|
||||
|| target_auth_type.eq_ignore_ascii_case("oauth"));
|
||||
|
||||
let api_key_present = fields.contains("api_key");
|
||||
let api_key_value = payload
|
||||
@@ -187,11 +192,19 @@ pub(crate) async fn build_admin_update_provider_key_record(
|
||||
if fields.contains("api_formats") {
|
||||
let api_formats = normalize_string_list(payload.api_formats)
|
||||
.ok_or_else(|| "api_formats 为必填字段".to_string())?;
|
||||
validate_vertex_api_formats(&provider.provider_type, &target_auth_type, &api_formats)?;
|
||||
updated.api_formats = Some(json!(api_formats));
|
||||
if managed_fixed_oauth_key {
|
||||
updated.api_formats = None;
|
||||
} else {
|
||||
validate_vertex_api_formats(&provider.provider_type, &target_auth_type, &api_formats)?;
|
||||
updated.api_formats = Some(json!(api_formats));
|
||||
}
|
||||
} else if payload.auth_type.is_some() {
|
||||
let api_formats = json_string_list(existing.api_formats.as_ref());
|
||||
validate_vertex_api_formats(&provider.provider_type, &target_auth_type, &api_formats)?;
|
||||
if managed_fixed_oauth_key {
|
||||
updated.api_formats = None;
|
||||
} else {
|
||||
let api_formats = json_string_list(existing.api_formats.as_ref());
|
||||
validate_vertex_api_formats(&provider.provider_type, &target_auth_type, &api_formats)?;
|
||||
}
|
||||
}
|
||||
|
||||
updated.auth_type = target_auth_type;
|
||||
|
||||
@@ -1,7 +1,13 @@
|
||||
mod create;
|
||||
mod endpoint;
|
||||
mod template;
|
||||
mod update;
|
||||
|
||||
pub(crate) use self::create::build_admin_create_provider_record;
|
||||
pub(crate) use self::endpoint::build_admin_fixed_provider_endpoint_record;
|
||||
pub(crate) use self::template::{
|
||||
apply_admin_fixed_provider_endpoint_template_overrides,
|
||||
reconcile_admin_fixed_provider_template_endpoints,
|
||||
reconcile_admin_fixed_provider_template_keys,
|
||||
};
|
||||
pub(crate) use self::update::build_admin_update_provider_record;
|
||||
|
||||
@@ -3,30 +3,65 @@ use crate::handlers::public::normalize_admin_base_url;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
|
||||
};
|
||||
use serde_json::json;
|
||||
use aether_provider_transport::provider_types::{
|
||||
FixedProviderEndpointTemplate, FixedProviderTemplate,
|
||||
};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use uuid::Uuid;
|
||||
|
||||
pub(crate) fn build_admin_fixed_provider_endpoint_record(
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub(crate) struct AdminFixedProviderEndpointDefaults {
|
||||
pub(crate) api_format: String,
|
||||
pub(crate) api_family: String,
|
||||
pub(crate) endpoint_kind: String,
|
||||
pub(crate) is_active: bool,
|
||||
pub(crate) base_url: String,
|
||||
pub(crate) header_rules: Option<serde_json::Value>,
|
||||
pub(crate) body_rules: Option<serde_json::Value>,
|
||||
pub(crate) max_retries: Option<i32>,
|
||||
pub(crate) custom_path: Option<String>,
|
||||
pub(crate) config: Option<serde_json::Value>,
|
||||
pub(crate) format_acceptance_config: Option<serde_json::Value>,
|
||||
pub(crate) proxy: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
pub(crate) fn build_admin_fixed_provider_endpoint_defaults(
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
api_format: &str,
|
||||
base_url: &str,
|
||||
) -> Result<StoredProviderCatalogEndpoint, String> {
|
||||
template: &FixedProviderTemplate,
|
||||
endpoint_template: &FixedProviderEndpointTemplate,
|
||||
) -> Result<AdminFixedProviderEndpointDefaults, String> {
|
||||
let (normalized_api_format, api_family, endpoint_kind) =
|
||||
admin_endpoint_signature_parts(api_format)
|
||||
.ok_or_else(|| format!("无效的 api_format: {api_format}"))?;
|
||||
admin_endpoint_signature_parts(endpoint_template.api_format)
|
||||
.ok_or_else(|| format!("无效的 api_format: {}", endpoint_template.api_format))?;
|
||||
let body_rules = admin_default_body_rules_for_signature(
|
||||
normalized_api_format,
|
||||
Some(provider.provider_type.as_str()),
|
||||
)
|
||||
.and_then(|(_, rules)| (!rules.is_empty()).then_some(serde_json::Value::Array(rules)));
|
||||
let endpoint_config = if provider.provider_type == "codex"
|
||||
&& matches!(normalized_api_format, "openai:cli" | "openai:image")
|
||||
{
|
||||
Some(json!({ "upstream_stream_policy": "force_stream" }))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
Ok(AdminFixedProviderEndpointDefaults {
|
||||
api_format: normalized_api_format.to_string(),
|
||||
api_family: api_family.to_string(),
|
||||
endpoint_kind: endpoint_kind.to_string(),
|
||||
is_active: true,
|
||||
base_url: normalize_admin_base_url(template.base_url)?,
|
||||
header_rules: None,
|
||||
body_rules,
|
||||
max_retries: Some(provider.max_retries.unwrap_or(2)),
|
||||
custom_path: endpoint_template.custom_path.map(ToOwned::to_owned),
|
||||
config: fixed_provider_endpoint_default_config(endpoint_template),
|
||||
format_acceptance_config: None,
|
||||
proxy: None,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn build_admin_fixed_provider_endpoint_record(
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
template: &FixedProviderTemplate,
|
||||
endpoint_template: &FixedProviderEndpointTemplate,
|
||||
) -> Result<StoredProviderCatalogEndpoint, String> {
|
||||
let defaults =
|
||||
build_admin_fixed_provider_endpoint_defaults(provider, template, endpoint_template)?;
|
||||
let now_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
@@ -36,22 +71,32 @@ pub(crate) fn build_admin_fixed_provider_endpoint_record(
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
Uuid::new_v4().to_string(),
|
||||
provider.id.clone(),
|
||||
normalized_api_format.to_string(),
|
||||
Some(api_family.to_string()),
|
||||
Some(endpoint_kind.to_string()),
|
||||
true,
|
||||
defaults.api_format,
|
||||
Some(defaults.api_family),
|
||||
Some(defaults.endpoint_kind),
|
||||
defaults.is_active,
|
||||
)
|
||||
.map_err(|err| err.to_string())?
|
||||
.with_timestamps(Some(now_unix_secs), Some(now_unix_secs))
|
||||
.with_transport_fields(
|
||||
normalize_admin_base_url(base_url)?,
|
||||
None,
|
||||
body_rules,
|
||||
Some(provider.max_retries.unwrap_or(2)),
|
||||
None,
|
||||
endpoint_config,
|
||||
None,
|
||||
None,
|
||||
defaults.base_url,
|
||||
defaults.header_rules,
|
||||
defaults.body_rules,
|
||||
defaults.max_retries,
|
||||
defaults.custom_path,
|
||||
defaults.config,
|
||||
defaults.format_acceptance_config,
|
||||
defaults.proxy,
|
||||
)
|
||||
.map_err(|err| err.to_string())
|
||||
}
|
||||
|
||||
fn fixed_provider_endpoint_default_config(
|
||||
endpoint_template: &FixedProviderEndpointTemplate,
|
||||
) -> Option<serde_json::Value> {
|
||||
let mut config = serde_json::Map::new();
|
||||
for default in endpoint_template.config_defaults {
|
||||
config.insert(default.key.to_string(), default.value.to_json_value());
|
||||
}
|
||||
(!config.is_empty()).then_some(serde_json::Value::Object(config))
|
||||
}
|
||||
|
||||
@@ -0,0 +1,530 @@
|
||||
use super::endpoint::{
|
||||
build_admin_fixed_provider_endpoint_defaults, build_admin_fixed_provider_endpoint_record,
|
||||
};
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::provider_key_auth::provider_key_is_oauth_managed;
|
||||
use crate::GatewayError;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_provider_transport::provider_types::{
|
||||
fixed_provider_template, FixedProviderEndpointTemplate, FixedProviderTemplate,
|
||||
};
|
||||
use serde_json::{json, Map, Value};
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
const FIXED_PROVIDER_TEMPLATE_METADATA_KEY: &str = "_aether_fixed_provider_template";
|
||||
const OVERRIDE_BODY_RULES: &str = "body_rules";
|
||||
const OVERRIDE_FORMAT_ACCEPTANCE_CONFIG: &str = "format_acceptance_config";
|
||||
const OVERRIDE_HEADER_RULES: &str = "header_rules";
|
||||
const OVERRIDE_IS_ACTIVE: &str = "is_active";
|
||||
const OVERRIDE_MAX_RETRIES: &str = "max_retries";
|
||||
const OVERRIDE_PROXY: &str = "proxy";
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
struct FixedProviderEndpointMetadata {
|
||||
provider_type: String,
|
||||
item_key: String,
|
||||
version: u32,
|
||||
retired: bool,
|
||||
overrides: BTreeSet<String>,
|
||||
config_keys: BTreeSet<String>,
|
||||
}
|
||||
|
||||
pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
) -> Result<(), GatewayError> {
|
||||
let Some(template) = state.fixed_provider_template(&provider.provider_type) else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let existing_endpoints = state
|
||||
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
|
||||
.await?;
|
||||
let mut matched_endpoint_ids = BTreeSet::new();
|
||||
|
||||
for endpoint_template in template.endpoints {
|
||||
let existing_endpoint = existing_endpoints
|
||||
.iter()
|
||||
.find(|endpoint| endpoint_matches_fixed_provider_template(endpoint, endpoint_template));
|
||||
match existing_endpoint {
|
||||
Some(existing_endpoint) => {
|
||||
matched_endpoint_ids.insert(existing_endpoint.id.clone());
|
||||
let updated = reconcile_fixed_provider_endpoint(
|
||||
provider,
|
||||
existing_endpoint,
|
||||
template,
|
||||
endpoint_template,
|
||||
)
|
||||
.map_err(GatewayError::Internal)?;
|
||||
if updated != *existing_endpoint {
|
||||
let Some(_) = state.update_provider_catalog_endpoint(&updated).await? else {
|
||||
return Err(GatewayError::Internal(
|
||||
"provider catalog endpoint writer unavailable".to_string(),
|
||||
));
|
||||
};
|
||||
}
|
||||
}
|
||||
None => {
|
||||
let mut created = build_admin_fixed_provider_endpoint_record(
|
||||
provider,
|
||||
template,
|
||||
endpoint_template,
|
||||
)
|
||||
.map_err(GatewayError::Internal)?;
|
||||
let metadata =
|
||||
managed_fixed_provider_endpoint_metadata(template, endpoint_template);
|
||||
upsert_fixed_provider_endpoint_metadata(&mut created, &metadata);
|
||||
let Some(_) = state.create_provider_catalog_endpoint(&created).await? else {
|
||||
return Err(GatewayError::Internal(
|
||||
"provider catalog endpoint writer unavailable".to_string(),
|
||||
));
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for existing_endpoint in &existing_endpoints {
|
||||
if matched_endpoint_ids.contains(&existing_endpoint.id) {
|
||||
continue;
|
||||
}
|
||||
let Some(metadata) = fixed_provider_endpoint_metadata(existing_endpoint) else {
|
||||
continue;
|
||||
};
|
||||
if metadata.retired && !existing_endpoint.is_active {
|
||||
continue;
|
||||
}
|
||||
let mut retired = existing_endpoint.clone();
|
||||
let mut retired_metadata = metadata;
|
||||
retired.is_active = false;
|
||||
retired_metadata.retired = true;
|
||||
upsert_fixed_provider_endpoint_metadata(&mut retired, &retired_metadata);
|
||||
if retired != *existing_endpoint {
|
||||
retired.updated_at_unix_secs = Some(current_unix_secs());
|
||||
let Some(_) = state.update_provider_catalog_endpoint(&retired).await? else {
|
||||
return Err(GatewayError::Internal(
|
||||
"provider catalog endpoint writer unavailable".to_string(),
|
||||
));
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn reconcile_admin_fixed_provider_template_keys(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
) -> Result<(), GatewayError> {
|
||||
let Some(_) = state.fixed_provider_template(&provider.provider_type) else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let existing_keys = state
|
||||
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
|
||||
.await?;
|
||||
for existing_key in existing_keys {
|
||||
let Some(updated_key) = reconcile_fixed_provider_key(provider, &existing_key) else {
|
||||
continue;
|
||||
};
|
||||
let Some(_) = state.update_provider_catalog_key(&updated_key).await? else {
|
||||
return Err(GatewayError::Internal(
|
||||
"provider catalog key writer unavailable".to_string(),
|
||||
));
|
||||
};
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn apply_admin_fixed_provider_endpoint_template_overrides(
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
existing_endpoint: &StoredProviderCatalogEndpoint,
|
||||
updated_endpoint: &mut StoredProviderCatalogEndpoint,
|
||||
) -> Result<(), String> {
|
||||
let Some(template) = fixed_provider_template(&provider.provider_type) else {
|
||||
return Ok(());
|
||||
};
|
||||
let Some(endpoint_template) =
|
||||
resolve_fixed_provider_endpoint_template(template, existing_endpoint, updated_endpoint)
|
||||
else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let defaults =
|
||||
build_admin_fixed_provider_endpoint_defaults(provider, template, endpoint_template)?;
|
||||
let mut metadata = fixed_provider_endpoint_metadata(existing_endpoint)
|
||||
.unwrap_or_else(|| managed_fixed_provider_endpoint_metadata(template, endpoint_template));
|
||||
let mut overrides = metadata.overrides.clone();
|
||||
|
||||
sync_override_if_changed(
|
||||
&mut overrides,
|
||||
OVERRIDE_HEADER_RULES,
|
||||
&existing_endpoint.header_rules,
|
||||
&updated_endpoint.header_rules,
|
||||
&defaults.header_rules,
|
||||
);
|
||||
sync_override_if_changed(
|
||||
&mut overrides,
|
||||
OVERRIDE_BODY_RULES,
|
||||
&existing_endpoint.body_rules,
|
||||
&updated_endpoint.body_rules,
|
||||
&defaults.body_rules,
|
||||
);
|
||||
sync_override_if_changed(
|
||||
&mut overrides,
|
||||
OVERRIDE_MAX_RETRIES,
|
||||
&existing_endpoint.max_retries,
|
||||
&updated_endpoint.max_retries,
|
||||
&defaults.max_retries,
|
||||
);
|
||||
sync_override_if_changed(
|
||||
&mut overrides,
|
||||
OVERRIDE_IS_ACTIVE,
|
||||
&existing_endpoint.is_active,
|
||||
&updated_endpoint.is_active,
|
||||
&defaults.is_active,
|
||||
);
|
||||
sync_override_if_changed(
|
||||
&mut overrides,
|
||||
OVERRIDE_PROXY,
|
||||
&existing_endpoint.proxy,
|
||||
&updated_endpoint.proxy,
|
||||
&defaults.proxy,
|
||||
);
|
||||
sync_override_if_changed(
|
||||
&mut overrides,
|
||||
OVERRIDE_FORMAT_ACCEPTANCE_CONFIG,
|
||||
&existing_endpoint.format_acceptance_config,
|
||||
&updated_endpoint.format_acceptance_config,
|
||||
&defaults.format_acceptance_config,
|
||||
);
|
||||
|
||||
let current_config_defaults = fixed_provider_endpoint_config_defaults(endpoint_template);
|
||||
let config = endpoint_config_without_metadata(updated_endpoint.config.as_ref());
|
||||
let existing_config = endpoint_config_without_metadata(existing_endpoint.config.as_ref());
|
||||
let current_config_keys = current_config_defaults
|
||||
.keys()
|
||||
.cloned()
|
||||
.collect::<BTreeSet<_>>();
|
||||
let mut tracked_config_keys = metadata.config_keys.clone();
|
||||
tracked_config_keys.extend(current_config_keys.iter().cloned());
|
||||
|
||||
for key in tracked_config_keys {
|
||||
let before = existing_config.get(&key);
|
||||
let actual = config.get(&key);
|
||||
let desired = current_config_defaults.get(&key);
|
||||
sync_override_if_changed(
|
||||
&mut overrides,
|
||||
&config_override_key(&key),
|
||||
&before.cloned(),
|
||||
&actual.cloned(),
|
||||
&desired.cloned(),
|
||||
);
|
||||
}
|
||||
|
||||
metadata.provider_type = template.provider_type.to_string();
|
||||
metadata.item_key = endpoint_template.item_key.to_string();
|
||||
metadata.version = template.version;
|
||||
metadata.retired = false;
|
||||
metadata.overrides = overrides;
|
||||
metadata.config_keys = current_config_keys;
|
||||
updated_endpoint.config = materialize_endpoint_config(config, &metadata);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn reconcile_fixed_provider_endpoint(
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
existing_endpoint: &StoredProviderCatalogEndpoint,
|
||||
template: &FixedProviderTemplate,
|
||||
endpoint_template: &FixedProviderEndpointTemplate,
|
||||
) -> Result<StoredProviderCatalogEndpoint, String> {
|
||||
let defaults =
|
||||
build_admin_fixed_provider_endpoint_defaults(provider, template, endpoint_template)?;
|
||||
let mut updated = existing_endpoint.clone();
|
||||
let metadata = fixed_provider_endpoint_metadata(existing_endpoint)
|
||||
.unwrap_or_else(|| managed_fixed_provider_endpoint_metadata(template, endpoint_template));
|
||||
|
||||
updated.api_format = defaults.api_format.clone();
|
||||
updated.api_family = Some(defaults.api_family.clone());
|
||||
updated.endpoint_kind = Some(defaults.endpoint_kind.clone());
|
||||
updated.base_url = defaults.base_url;
|
||||
updated.custom_path = defaults.custom_path;
|
||||
|
||||
if !metadata.overrides.contains(OVERRIDE_HEADER_RULES) {
|
||||
updated.header_rules = defaults.header_rules;
|
||||
}
|
||||
if !metadata.overrides.contains(OVERRIDE_BODY_RULES) {
|
||||
updated.body_rules = defaults.body_rules;
|
||||
}
|
||||
if !metadata.overrides.contains(OVERRIDE_MAX_RETRIES) {
|
||||
updated.max_retries = defaults.max_retries;
|
||||
}
|
||||
if !metadata.overrides.contains(OVERRIDE_IS_ACTIVE) {
|
||||
updated.is_active = defaults.is_active;
|
||||
}
|
||||
if !metadata.overrides.contains(OVERRIDE_PROXY) {
|
||||
updated.proxy = defaults.proxy;
|
||||
}
|
||||
if !metadata
|
||||
.overrides
|
||||
.contains(OVERRIDE_FORMAT_ACCEPTANCE_CONFIG)
|
||||
{
|
||||
updated.format_acceptance_config = defaults.format_acceptance_config;
|
||||
}
|
||||
|
||||
let mut config = endpoint_config_without_metadata(updated.config.as_ref());
|
||||
let current_config_defaults = fixed_provider_endpoint_config_defaults(endpoint_template);
|
||||
let current_config_keys = current_config_defaults
|
||||
.keys()
|
||||
.cloned()
|
||||
.collect::<BTreeSet<_>>();
|
||||
|
||||
for old_key in metadata.config_keys.difference(¤t_config_keys) {
|
||||
if !metadata
|
||||
.overrides
|
||||
.contains(config_override_key(old_key.as_str()).as_str())
|
||||
{
|
||||
config.remove(old_key);
|
||||
}
|
||||
}
|
||||
for (key, value) in ¤t_config_defaults {
|
||||
if !metadata
|
||||
.overrides
|
||||
.contains(config_override_key(key.as_str()).as_str())
|
||||
{
|
||||
config.insert(key.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
|
||||
let mut next_metadata = metadata;
|
||||
next_metadata.provider_type = template.provider_type.to_string();
|
||||
next_metadata.item_key = endpoint_template.item_key.to_string();
|
||||
next_metadata.version = template.version;
|
||||
next_metadata.retired = false;
|
||||
next_metadata.config_keys = current_config_keys;
|
||||
updated.config = materialize_endpoint_config(config, &next_metadata);
|
||||
|
||||
if updated != *existing_endpoint {
|
||||
updated.updated_at_unix_secs = Some(current_unix_secs());
|
||||
}
|
||||
Ok(updated)
|
||||
}
|
||||
|
||||
fn resolve_fixed_provider_endpoint_template<'a>(
|
||||
template: &'a FixedProviderTemplate,
|
||||
existing_endpoint: &StoredProviderCatalogEndpoint,
|
||||
updated_endpoint: &StoredProviderCatalogEndpoint,
|
||||
) -> Option<&'a FixedProviderEndpointTemplate> {
|
||||
if let Some(metadata) = fixed_provider_endpoint_metadata(existing_endpoint) {
|
||||
if let Some(item) = template
|
||||
.endpoints
|
||||
.iter()
|
||||
.find(|item| item.item_key == metadata.item_key)
|
||||
{
|
||||
return Some(item);
|
||||
}
|
||||
}
|
||||
|
||||
template.endpoints.iter().find(|item| {
|
||||
item.api_format
|
||||
.eq_ignore_ascii_case(updated_endpoint.api_format.trim())
|
||||
|| item
|
||||
.api_format
|
||||
.eq_ignore_ascii_case(existing_endpoint.api_format.trim())
|
||||
})
|
||||
}
|
||||
|
||||
fn endpoint_matches_fixed_provider_template(
|
||||
endpoint: &StoredProviderCatalogEndpoint,
|
||||
endpoint_template: &FixedProviderEndpointTemplate,
|
||||
) -> bool {
|
||||
if let Some(metadata) = fixed_provider_endpoint_metadata(endpoint) {
|
||||
if metadata.item_key == endpoint_template.item_key {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
endpoint
|
||||
.api_format
|
||||
.trim()
|
||||
.eq_ignore_ascii_case(endpoint_template.api_format)
|
||||
}
|
||||
|
||||
fn fixed_provider_endpoint_metadata(
|
||||
endpoint: &StoredProviderCatalogEndpoint,
|
||||
) -> Option<FixedProviderEndpointMetadata> {
|
||||
let config = endpoint.config.as_ref()?.as_object()?;
|
||||
let metadata = config
|
||||
.get(FIXED_PROVIDER_TEMPLATE_METADATA_KEY)?
|
||||
.as_object()?;
|
||||
let provider_type = metadata
|
||||
.get("provider_type")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())?;
|
||||
let item_key = metadata
|
||||
.get("item_key")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())?;
|
||||
if !metadata
|
||||
.get("managed")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(FixedProviderEndpointMetadata {
|
||||
provider_type: provider_type.to_string(),
|
||||
item_key: item_key.to_string(),
|
||||
version: metadata
|
||||
.get("version")
|
||||
.and_then(Value::as_u64)
|
||||
.and_then(|value| u32::try_from(value).ok())
|
||||
.unwrap_or(0),
|
||||
retired: metadata
|
||||
.get("retired")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false),
|
||||
overrides: metadata
|
||||
.get("overrides")
|
||||
.and_then(Value::as_array)
|
||||
.map(|items| {
|
||||
items
|
||||
.iter()
|
||||
.filter_map(Value::as_str)
|
||||
.map(ToOwned::to_owned)
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default(),
|
||||
config_keys: metadata
|
||||
.get("config_keys")
|
||||
.and_then(Value::as_array)
|
||||
.map(|items| {
|
||||
items
|
||||
.iter()
|
||||
.filter_map(Value::as_str)
|
||||
.map(ToOwned::to_owned)
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default(),
|
||||
})
|
||||
}
|
||||
|
||||
fn managed_fixed_provider_endpoint_metadata(
|
||||
template: &FixedProviderTemplate,
|
||||
endpoint_template: &FixedProviderEndpointTemplate,
|
||||
) -> FixedProviderEndpointMetadata {
|
||||
FixedProviderEndpointMetadata {
|
||||
provider_type: template.provider_type.to_string(),
|
||||
item_key: endpoint_template.item_key.to_string(),
|
||||
version: template.version,
|
||||
retired: false,
|
||||
overrides: BTreeSet::new(),
|
||||
config_keys: fixed_provider_endpoint_config_defaults(endpoint_template)
|
||||
.into_keys()
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
fn upsert_fixed_provider_endpoint_metadata(
|
||||
endpoint: &mut StoredProviderCatalogEndpoint,
|
||||
metadata: &FixedProviderEndpointMetadata,
|
||||
) {
|
||||
let config = endpoint_config_without_metadata(endpoint.config.as_ref());
|
||||
endpoint.config = materialize_endpoint_config(config, metadata);
|
||||
}
|
||||
|
||||
fn endpoint_config_without_metadata(config: Option<&Value>) -> Map<String, Value> {
|
||||
let mut config = config
|
||||
.and_then(Value::as_object)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
config.remove(FIXED_PROVIDER_TEMPLATE_METADATA_KEY);
|
||||
config
|
||||
}
|
||||
|
||||
fn materialize_endpoint_config(
|
||||
mut config: Map<String, Value>,
|
||||
metadata: &FixedProviderEndpointMetadata,
|
||||
) -> Option<Value> {
|
||||
config.insert(
|
||||
FIXED_PROVIDER_TEMPLATE_METADATA_KEY.to_string(),
|
||||
json!({
|
||||
"managed": true,
|
||||
"provider_type": metadata.provider_type,
|
||||
"item_key": metadata.item_key,
|
||||
"version": metadata.version,
|
||||
"retired": metadata.retired,
|
||||
"overrides": metadata.overrides.iter().cloned().collect::<Vec<_>>(),
|
||||
"config_keys": metadata.config_keys.iter().cloned().collect::<Vec<_>>(),
|
||||
}),
|
||||
);
|
||||
Some(Value::Object(config))
|
||||
}
|
||||
|
||||
fn fixed_provider_endpoint_config_defaults(
|
||||
endpoint_template: &FixedProviderEndpointTemplate,
|
||||
) -> BTreeMap<String, Value> {
|
||||
endpoint_template
|
||||
.config_defaults
|
||||
.iter()
|
||||
.map(|item| (item.key.to_string(), item.value.to_json_value()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn config_override_key(key: &str) -> String {
|
||||
format!("config.{key}")
|
||||
}
|
||||
|
||||
fn current_unix_secs() -> u64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn reconcile_fixed_provider_key(
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
existing_key: &StoredProviderCatalogKey,
|
||||
) -> Option<StoredProviderCatalogKey> {
|
||||
if !provider_key_is_oauth_managed(existing_key, &provider.provider_type)
|
||||
|| existing_key.api_formats.is_none()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut updated = existing_key.clone();
|
||||
updated.api_formats = None;
|
||||
updated.updated_at_unix_secs = Some(current_unix_secs());
|
||||
Some(updated)
|
||||
}
|
||||
|
||||
fn sync_override<T>(overrides: &mut BTreeSet<String>, key: &str, actual: &T, desired: &T)
|
||||
where
|
||||
T: PartialEq,
|
||||
{
|
||||
if actual == desired {
|
||||
overrides.remove(key);
|
||||
} else {
|
||||
overrides.insert(key.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
fn sync_override_if_changed<T>(
|
||||
overrides: &mut BTreeSet<String>,
|
||||
key: &str,
|
||||
before: &T,
|
||||
actual: &T,
|
||||
desired: &T,
|
||||
) where
|
||||
T: PartialEq,
|
||||
{
|
||||
if before == actual {
|
||||
return;
|
||||
}
|
||||
sync_override(overrides, key, actual, desired);
|
||||
}
|
||||
@@ -36,12 +36,14 @@ impl<'a> AdminAppState<'a> {
|
||||
&self,
|
||||
key: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
api_formats: &[String],
|
||||
now_unix_secs: u64,
|
||||
) -> serde_json::Value {
|
||||
crate::handlers::admin::shared::build_admin_provider_key_response(
|
||||
self.app,
|
||||
key,
|
||||
provider_type,
|
||||
api_formats,
|
||||
now_unix_secs,
|
||||
)
|
||||
}
|
||||
@@ -274,6 +276,7 @@ impl<'a> AdminAppState<'a> {
|
||||
String,
|
||||
> {
|
||||
use crate::api::ai::admin_endpoint_signature_parts;
|
||||
use crate::handlers::admin::provider::write::provider::apply_admin_fixed_provider_endpoint_template_overrides;
|
||||
use crate::handlers::public::{admin_requested_force_stream, normalize_admin_base_url};
|
||||
use aether_admin::provider::endpoints as admin_provider_endpoints_pure;
|
||||
let (fields, payload) = patch.into_parts();
|
||||
@@ -346,6 +349,11 @@ impl<'a> AdminAppState<'a> {
|
||||
.ok_or_else(|| format!("无效的 api_format: {}", updated.api_format))?;
|
||||
updated.api_family = Some(api_family.to_string());
|
||||
updated.endpoint_kind = Some(endpoint_kind.to_string());
|
||||
apply_admin_fixed_provider_endpoint_template_overrides(
|
||||
provider,
|
||||
existing_endpoint,
|
||||
&mut updated,
|
||||
)?;
|
||||
updated.updated_at_unix_secs = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.ok()
|
||||
|
||||
@@ -113,6 +113,16 @@ impl<'a> AdminAppState<'a> {
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn read_provider_quota_snapshots(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<
|
||||
Vec<aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot>,
|
||||
GatewayError,
|
||||
> {
|
||||
self.app.read_provider_quota_snapshots(provider_ids).await
|
||||
}
|
||||
|
||||
pub(crate) async fn update_provider_catalog_key_health_state(
|
||||
&self,
|
||||
key_id: &str,
|
||||
|
||||
@@ -465,6 +465,7 @@ impl<'a> AdminAppState<'a> {
|
||||
pub(crate) async fn create_provider_oauth_catalog_key(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
provider_type: &str,
|
||||
name: &str,
|
||||
access_token: &str,
|
||||
auth_config: &serde_json::Map<String, serde_json::Value>,
|
||||
@@ -478,6 +479,7 @@ impl<'a> AdminAppState<'a> {
|
||||
crate::handlers::admin::provider::oauth::provisioning::create_provider_oauth_catalog_key(
|
||||
self,
|
||||
provider_id,
|
||||
provider_type,
|
||||
name,
|
||||
access_token,
|
||||
auth_config,
|
||||
@@ -491,8 +493,10 @@ impl<'a> AdminAppState<'a> {
|
||||
pub(crate) async fn update_existing_provider_oauth_catalog_key(
|
||||
&self,
|
||||
existing_key: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
access_token: &str,
|
||||
auth_config: &serde_json::Map<String, serde_json::Value>,
|
||||
api_formats: &[String],
|
||||
proxy: Option<serde_json::Value>,
|
||||
expires_at_unix_secs: Option<u64>,
|
||||
) -> Result<
|
||||
@@ -502,8 +506,10 @@ impl<'a> AdminAppState<'a> {
|
||||
crate::handlers::admin::provider::oauth::provisioning::update_existing_provider_oauth_catalog_key(
|
||||
self,
|
||||
existing_key,
|
||||
provider_type,
|
||||
access_token,
|
||||
auth_config,
|
||||
api_formats,
|
||||
proxy,
|
||||
expires_at_unix_secs,
|
||||
)
|
||||
|
||||
@@ -94,7 +94,7 @@ impl<'a> AdminAppState<'a> {
|
||||
pub(crate) fn fixed_provider_template(
|
||||
&self,
|
||||
provider_type: &str,
|
||||
) -> Option<(&'static str, &'static [&'static str])> {
|
||||
) -> Option<&'static crate::provider_transport::provider_types::FixedProviderTemplate> {
|
||||
crate::provider_transport::provider_types::fixed_provider_template(provider_type)
|
||||
}
|
||||
|
||||
|
||||
@@ -336,6 +336,7 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
||||
trace_id.as_str(),
|
||||
&resolved,
|
||||
&payload.body_json,
|
||||
payload.body_base64.as_deref(),
|
||||
)
|
||||
.await?
|
||||
else {
|
||||
@@ -484,6 +485,7 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
||||
trace_id.as_str(),
|
||||
&resolved,
|
||||
&payload.body_json,
|
||||
payload.body_base64.as_deref(),
|
||||
)
|
||||
.await?
|
||||
{
|
||||
@@ -638,6 +640,7 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
||||
trace_id.as_str(),
|
||||
&resolved,
|
||||
&payload.body_json,
|
||||
payload.body_base64.as_deref(),
|
||||
)
|
||||
.await?
|
||||
{
|
||||
|
||||
@@ -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>> {
|
||||
|
||||
@@ -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()))
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -1,5 +1,8 @@
|
||||
use crate::handlers::shared::{json_string_list, unix_secs_to_rfc3339};
|
||||
use crate::provider_key_auth::provider_key_auth_semantics;
|
||||
use crate::provider_key_auth::{
|
||||
provider_key_auth_semantics, provider_key_configured_api_formats,
|
||||
provider_key_inherits_provider_api_formats,
|
||||
};
|
||||
use crate::AppState;
|
||||
use aether_admin::provider::quota as admin_provider_quota_pure;
|
||||
use aether_admin::provider::status as admin_provider_status_pure;
|
||||
@@ -18,17 +21,18 @@ const OAUTH_REQUEST_FAILED_PREFIX: &str = "[REQUEST_FAILED] ";
|
||||
|
||||
pub(crate) fn provider_catalog_key_supports_format(
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
let Some(value) = key.api_formats.as_ref() else {
|
||||
if provider_key_inherits_provider_api_formats(key, provider_type) {
|
||||
return true;
|
||||
};
|
||||
let Some(values) = value.as_array() else {
|
||||
}
|
||||
let formats = provider_key_configured_api_formats(key);
|
||||
if formats.is_empty() {
|
||||
return true;
|
||||
};
|
||||
values
|
||||
}
|
||||
formats
|
||||
.iter()
|
||||
.filter_map(serde_json::Value::as_str)
|
||||
.any(|candidate| candidate.trim().eq_ignore_ascii_case(api_format))
|
||||
}
|
||||
|
||||
@@ -1188,6 +1192,7 @@ pub(crate) fn build_admin_provider_key_response(
|
||||
state: &AppState,
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
api_formats: &[String],
|
||||
now_unix_secs: u64,
|
||||
) -> serde_json::Value {
|
||||
let request_count = u64::from(key.request_count.unwrap_or(0));
|
||||
@@ -1245,8 +1250,9 @@ pub(crate) fn build_admin_provider_key_response(
|
||||
payload.insert(
|
||||
"api_formats".to_string(),
|
||||
serde_json::Value::Array(
|
||||
json_string_list(key.api_formats.as_ref())
|
||||
.into_iter()
|
||||
api_formats
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(serde_json::Value::String)
|
||||
.collect(),
|
||||
),
|
||||
|
||||
@@ -170,14 +170,14 @@ struct GatewayDataArgs {
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DATA_POSTGRES_MAX_CONNECTIONS",
|
||||
default_value_t = 30
|
||||
default_value_t = 100
|
||||
)]
|
||||
postgres_max_connections: u32,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DATA_POSTGRES_ACQUIRE_TIMEOUT_MS",
|
||||
default_value_t = 3_000
|
||||
default_value_t = 10_000
|
||||
)]
|
||||
postgres_acquire_timeout_ms: u64,
|
||||
|
||||
@@ -1243,8 +1243,8 @@ mod tests {
|
||||
redis_url: None,
|
||||
redis_key_prefix: None,
|
||||
postgres_min_connections: 1,
|
||||
postgres_max_connections: 30,
|
||||
postgres_acquire_timeout_ms: 3_000,
|
||||
postgres_max_connections: 100,
|
||||
postgres_acquire_timeout_ms: 10_000,
|
||||
postgres_idle_timeout_ms: 60_000,
|
||||
postgres_max_lifetime_ms: 1_800_000,
|
||||
postgres_statement_cache_capacity: 100,
|
||||
|
||||
@@ -1330,6 +1330,8 @@ SELECT
|
||||
aggregated.input_tokens,
|
||||
aggregated.output_tokens,
|
||||
aggregated.cache_creation_tokens,
|
||||
aggregated.cache_creation_ephemeral_5m_tokens,
|
||||
aggregated.cache_creation_ephemeral_1h_tokens,
|
||||
aggregated.cache_read_tokens,
|
||||
aggregated.total_cost,
|
||||
aggregated.response_time_sum_ms,
|
||||
@@ -1344,6 +1346,8 @@ DO UPDATE SET
|
||||
input_tokens = EXCLUDED.input_tokens,
|
||||
output_tokens = EXCLUDED.output_tokens,
|
||||
cache_creation_tokens = EXCLUDED.cache_creation_tokens,
|
||||
cache_creation_ephemeral_5m_tokens = EXCLUDED.cache_creation_ephemeral_5m_tokens,
|
||||
cache_creation_ephemeral_1h_tokens = EXCLUDED.cache_creation_ephemeral_1h_tokens,
|
||||
cache_read_tokens = EXCLUDED.cache_read_tokens,
|
||||
total_cost = EXCLUDED.total_cost,
|
||||
response_time_sum_ms = EXCLUDED.response_time_sum_ms,
|
||||
|
||||
@@ -1,4 +1,8 @@
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
};
|
||||
use aether_provider_transport::provider_types::provider_type_is_fixed;
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum ProviderKeyCredentialKind {
|
||||
@@ -146,12 +150,68 @@ pub(crate) fn provider_key_is_oauth_managed(
|
||||
provider_key_auth_semantics(key, provider_type).oauth_managed()
|
||||
}
|
||||
|
||||
pub(crate) fn provider_key_configured_api_formats(key: &StoredProviderCatalogKey) -> Vec<String> {
|
||||
key.api_formats
|
||||
.as_ref()
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.map(|items| {
|
||||
items
|
||||
.iter()
|
||||
.filter_map(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
pub(crate) fn provider_active_api_formats(
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
) -> Vec<String> {
|
||||
let mut formats = Vec::new();
|
||||
let mut seen = BTreeSet::new();
|
||||
for endpoint in endpoints.iter().filter(|endpoint| endpoint.is_active) {
|
||||
let api_format = endpoint.api_format.trim();
|
||||
if api_format.is_empty() || !seen.insert(api_format.to_string()) {
|
||||
continue;
|
||||
}
|
||||
formats.push(api_format.to_string());
|
||||
}
|
||||
formats
|
||||
}
|
||||
|
||||
pub(crate) fn provider_key_inherits_provider_api_formats(
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
) -> bool {
|
||||
provider_type_is_fixed(provider_type) && provider_key_is_oauth_managed(key, provider_type)
|
||||
}
|
||||
|
||||
pub(crate) fn provider_key_effective_api_formats(
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
) -> Vec<String> {
|
||||
if provider_key_inherits_provider_api_formats(key, provider_type) {
|
||||
provider_active_api_formats(endpoints)
|
||||
} else {
|
||||
provider_key_configured_api_formats(key)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
provider_key_auth_semantics, ProviderKeyCredentialKind, ProviderKeyRuntimeAuthKind,
|
||||
provider_active_api_formats, provider_key_auth_semantics,
|
||||
provider_key_configured_api_formats, provider_key_effective_api_formats,
|
||||
provider_key_inherits_provider_api_formats, ProviderKeyCredentialKind,
|
||||
ProviderKeyRuntimeAuthKind,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
fn sample_key(auth_type: &str) -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
@@ -165,6 +225,21 @@ mod tests {
|
||||
.expect("key should build")
|
||||
}
|
||||
|
||||
fn sample_endpoint(api_format: &str, is_active: bool) -> StoredProviderCatalogEndpoint {
|
||||
let mut endpoint = StoredProviderCatalogEndpoint::new(
|
||||
format!("endpoint-{api_format}"),
|
||||
"provider-1".to_string(),
|
||||
api_format.to_string(),
|
||||
None,
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build");
|
||||
endpoint.is_active = is_active;
|
||||
endpoint.base_url = "https://example.invalid".to_string();
|
||||
endpoint
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recognizes_oauth_managed_key() {
|
||||
let semantics = provider_key_auth_semantics(&sample_key("oauth"), "codex");
|
||||
@@ -227,4 +302,70 @@ mod tests {
|
||||
ProviderKeyRuntimeAuthKind::ServiceAccount
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deduplicates_active_provider_api_formats() {
|
||||
let endpoints = vec![
|
||||
sample_endpoint("openai:cli", true),
|
||||
sample_endpoint("openai:image", true),
|
||||
sample_endpoint("openai:cli", true),
|
||||
sample_endpoint("openai:compact", false),
|
||||
];
|
||||
|
||||
assert_eq!(
|
||||
provider_active_api_formats(&endpoints),
|
||||
vec!["openai:cli".to_string(), "openai:image".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fixed_oauth_key_with_null_formats_inherits_provider_formats() {
|
||||
let key = sample_key("oauth");
|
||||
let endpoints = vec![
|
||||
sample_endpoint("openai:cli", true),
|
||||
sample_endpoint("openai:image", true),
|
||||
];
|
||||
|
||||
assert!(provider_key_inherits_provider_api_formats(&key, "codex"));
|
||||
assert_eq!(
|
||||
provider_key_effective_api_formats(&key, "codex", &endpoints),
|
||||
vec!["openai:cli".to_string(), "openai:image".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fixed_oauth_key_with_legacy_explicit_formats_still_inherits_provider_formats() {
|
||||
let mut key = sample_key("oauth");
|
||||
key.api_formats = Some(json!(["openai:compact"]));
|
||||
let endpoints = vec![
|
||||
sample_endpoint("openai:cli", true),
|
||||
sample_endpoint("openai:image", true),
|
||||
];
|
||||
|
||||
assert!(provider_key_inherits_provider_api_formats(&key, "codex"));
|
||||
assert_eq!(
|
||||
provider_key_effective_api_formats(&key, "codex", &endpoints),
|
||||
vec!["openai:cli".to_string(), "openai:image".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_formats_do_not_inherit_for_non_fixed_key() {
|
||||
let mut key = sample_key("oauth");
|
||||
key.api_formats = Some(json!(["openai:compact"]));
|
||||
let endpoints = vec![
|
||||
sample_endpoint("openai:cli", true),
|
||||
sample_endpoint("openai:image", true),
|
||||
];
|
||||
|
||||
assert!(!provider_key_inherits_provider_api_formats(&key, "openai"));
|
||||
assert_eq!(
|
||||
provider_key_configured_api_formats(&key),
|
||||
vec!["openai:compact".to_string()]
|
||||
);
|
||||
assert_eq!(
|
||||
provider_key_effective_api_formats(&key, "openai", &endpoints),
|
||||
vec!["openai:compact".to_string()]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ use aether_runtime::{ConcurrencyGate, DistributedConcurrencyGate};
|
||||
use super::super::async_task::{VideoTaskPollerConfig, VideoTaskService};
|
||||
use super::super::cache::{
|
||||
AuthApiKeyLastUsedCache, AuthContextCache, DashboardResponseCache, DirectPlanBypassCache,
|
||||
SchedulerAffinityCache,
|
||||
SchedulerAffinityCache, SystemConfigCache,
|
||||
};
|
||||
use super::super::data::GatewayDataState;
|
||||
use super::super::fallback_metrics;
|
||||
@@ -59,6 +59,7 @@ pub struct AppState {
|
||||
pub(crate) direct_plan_bypass_cache: Arc<DirectPlanBypassCache>,
|
||||
pub(crate) scheduler_affinity_cache: Arc<SchedulerAffinityCache>,
|
||||
pub(crate) dashboard_response_cache: Arc<DashboardResponseCache>,
|
||||
pub(crate) system_config_cache: Arc<SystemConfigCache>,
|
||||
pub(crate) fallback_metrics: Arc<fallback_metrics::GatewayFallbackMetrics>,
|
||||
pub(crate) frontdoor_cors: Option<Arc<FrontdoorCorsConfig>>,
|
||||
pub(crate) frontdoor_user_rpm: Arc<FrontdoorUserRpmLimiter>,
|
||||
|
||||
@@ -24,6 +24,7 @@ use super::super::async_task::{
|
||||
use super::super::cache::{
|
||||
AuthApiKeyLastUsedCache, AuthContextCache, DashboardResponseCache, DirectPlanBypassCache,
|
||||
SchedulerAffinityCache, SchedulerAffinitySnapshotEntry, SchedulerAffinityTarget,
|
||||
SystemConfigCache,
|
||||
};
|
||||
use super::super::data::{GatewayDataConfig, GatewayDataState};
|
||||
use super::super::fallback_metrics;
|
||||
@@ -48,6 +49,8 @@ use crate::maintenance::spawn_stats_hourly_aggregation_worker;
|
||||
use crate::maintenance::spawn_usage_cleanup_worker;
|
||||
use crate::maintenance::spawn_wallet_daily_usage_aggregation_worker;
|
||||
|
||||
const SYSTEM_CONFIG_CACHE_TTL: Duration = Duration::from_secs(3);
|
||||
|
||||
impl AppState {
|
||||
fn spawn_scheduler_affinity_redis_write(
|
||||
&self,
|
||||
@@ -90,6 +93,7 @@ impl AppState {
|
||||
|
||||
pub(crate) fn replace_data_state(&mut self, data: Arc<GatewayDataState>) {
|
||||
self.clear_provider_transport_snapshot_cache();
|
||||
self.system_config_cache.clear();
|
||||
self.tunnel = crate::tunnel::EmbeddedTunnelState::with_data(Arc::clone(&data));
|
||||
self.data = data;
|
||||
}
|
||||
@@ -147,6 +151,7 @@ impl AppState {
|
||||
direct_plan_bypass_cache: Arc::new(DirectPlanBypassCache::default()),
|
||||
scheduler_affinity_cache: Arc::new(SchedulerAffinityCache::default()),
|
||||
dashboard_response_cache: Arc::new(DashboardResponseCache::default()),
|
||||
system_config_cache: Arc::new(SystemConfigCache::default()),
|
||||
fallback_metrics: Arc::new(fallback_metrics::GatewayFallbackMetrics::default()),
|
||||
frontdoor_cors: None,
|
||||
frontdoor_user_rpm: Arc::new(FrontdoorUserRpmLimiter::new(
|
||||
@@ -396,10 +401,18 @@ impl AppState {
|
||||
&self,
|
||||
key: &str,
|
||||
) -> Result<Option<serde_json::Value>, GatewayError> {
|
||||
self.data
|
||||
if let Some(value) = self.system_config_cache.get(key, SYSTEM_CONFIG_CACHE_TTL) {
|
||||
return Ok(value);
|
||||
}
|
||||
|
||||
let value = self
|
||||
.data
|
||||
.find_system_config_value(key)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
self.system_config_cache
|
||||
.insert(key.to_string(), value.clone(), SYSTEM_CONFIG_CACHE_TTL);
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
pub(crate) async fn upsert_system_config_json_value(
|
||||
@@ -408,10 +421,17 @@ impl AppState {
|
||||
value: &serde_json::Value,
|
||||
description: Option<&str>,
|
||||
) -> Result<serde_json::Value, GatewayError> {
|
||||
self.data
|
||||
let value = self
|
||||
.data
|
||||
.upsert_system_config_value(key, value, description)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
self.system_config_cache.insert(
|
||||
key.to_string(),
|
||||
Some(value.clone()),
|
||||
SYSTEM_CONFIG_CACHE_TTL,
|
||||
);
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
pub(crate) async fn list_system_config_entries(
|
||||
@@ -436,10 +456,14 @@ impl AppState {
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_system_config_value(&self, key: &str) -> Result<bool, GatewayError> {
|
||||
self.data
|
||||
let deleted = self
|
||||
.data
|
||||
.delete_system_config_value(key)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
self.system_config_cache
|
||||
.insert(key.to_string(), None, SYSTEM_CONFIG_CACHE_TTL);
|
||||
Ok(deleted)
|
||||
}
|
||||
|
||||
pub(crate) async fn read_admin_system_stats(
|
||||
@@ -842,3 +866,89 @@ fn runtime_miss_diagnostic_has_candidate_signal(
|
||||
|| diagnostic.skipped_candidate_count.unwrap_or(0) > 0
|
||||
|| !diagnostic.skip_reasons.is_empty()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
use super::AppState;
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
#[tokio::test]
|
||||
async fn system_config_reads_use_short_lived_cache_until_app_invalidation() {
|
||||
let state = AppState::new()
|
||||
.expect("app state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::disabled()
|
||||
.with_system_config_values_for_tests([("site_name".to_string(), json!("old"))]),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
state
|
||||
.read_system_config_json_value("site_name")
|
||||
.await
|
||||
.expect("system config read should succeed"),
|
||||
Some(json!("old"))
|
||||
);
|
||||
|
||||
state
|
||||
.data
|
||||
.upsert_system_config_value("site_name", &json!("bypassed"), None)
|
||||
.await
|
||||
.expect("direct data write should succeed");
|
||||
|
||||
assert_eq!(
|
||||
state
|
||||
.read_system_config_json_value("site_name")
|
||||
.await
|
||||
.expect("cached system config read should succeed"),
|
||||
Some(json!("old"))
|
||||
);
|
||||
|
||||
state
|
||||
.upsert_system_config_json_value("site_name", &json!("fresh"), None)
|
||||
.await
|
||||
.expect("app system config write should succeed");
|
||||
|
||||
assert_eq!(
|
||||
state
|
||||
.read_system_config_json_value("site_name")
|
||||
.await
|
||||
.expect("refreshed system config read should succeed"),
|
||||
Some(json!("fresh"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn replacing_data_state_clears_system_config_cache() {
|
||||
let mut state = AppState::new()
|
||||
.expect("app state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::disabled()
|
||||
.with_system_config_values_for_tests([("site_name".to_string(), json!("old"))]),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
state
|
||||
.read_system_config_json_value("site_name")
|
||||
.await
|
||||
.expect("system config read should succeed"),
|
||||
Some(json!("old"))
|
||||
);
|
||||
|
||||
state.replace_data_state(Arc::new(
|
||||
GatewayDataState::disabled()
|
||||
.with_system_config_values_for_tests([("site_name".to_string(), json!("new"))]),
|
||||
));
|
||||
|
||||
assert_eq!(
|
||||
state
|
||||
.read_system_config_json_value("site_name")
|
||||
.await
|
||||
.expect("system config read should reflect replaced data"),
|
||||
Some(json!("new"))
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -33,6 +33,16 @@ impl AppState {
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn read_provider_quota_snapshots(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<quota::StoredProviderQuotaSnapshot>, GatewayError> {
|
||||
self.data
|
||||
.find_provider_quotas_by_provider_ids(provider_ids)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn read_recent_request_candidates(
|
||||
&self,
|
||||
limit: usize,
|
||||
|
||||
427
apps/aether-gateway/src/tests/ai_execute/stream/image.rs
Normal file
427
apps/aether-gateway/src/tests/ai_execute/stream/image.rs
Normal file
@@ -0,0 +1,427 @@
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
|
||||
to_bytes, Arc, Body, Json, Mutex, Request, Router, StatusCode, TRACE_ID_HEADER,
|
||||
};
|
||||
use crate::ai_pipeline::CODEX_OPENAI_IMAGE_INTERNAL_MODEL;
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
|
||||
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data_contracts::repository::candidate_selection::{
|
||||
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth_refresh() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeStreamRequest {
|
||||
trace_id: String,
|
||||
url: String,
|
||||
model: String,
|
||||
authorization: String,
|
||||
x_client_request_id: String,
|
||||
tool_type: String,
|
||||
tool_partial_images: Option<u64>,
|
||||
request_stream: bool,
|
||||
plan_stream: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenRefreshRequest {
|
||||
content_type: String,
|
||||
body: String,
|
||||
}
|
||||
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["openai", "codex"])),
|
||||
Some(serde_json::json!(["openai:image"])),
|
||||
Some(serde_json::json!(["gpt-image-2"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800_i64),
|
||||
Some(serde_json::json!(["openai", "codex"])),
|
||||
Some(serde_json::json!(["openai:image"])),
|
||||
Some(serde_json::json!(["gpt-image-2"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-codex-image-stream-local-1".to_string(),
|
||||
provider_name: "codex".to_string(),
|
||||
provider_type: "codex".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-codex-image-stream-local-1".to_string(),
|
||||
endpoint_api_format: "openai:image".to_string(),
|
||||
endpoint_api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("image".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-codex-image-stream-local-1".to_string(),
|
||||
key_name: "oauth".to_string(),
|
||||
key_auth_type: "oauth".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["openai:image".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"openai:image": 1})),
|
||||
model_id: "model-codex-image-stream-local-1".to_string(),
|
||||
global_model_id: "global-model-codex-image-stream-local-1".to_string(),
|
||||
global_model_name: "gpt-image-2".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "gpt-image-2".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-image-2".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["openai:image".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-codex-image-stream-local-1".to_string(),
|
||||
"codex".to_string(),
|
||||
Some("https://chatgpt.com".to_string()),
|
||||
"codex".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-codex-image-stream-local-1".to_string(),
|
||||
"provider-codex-image-stream-local-1".to_string(),
|
||||
"openai:image".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("image".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://chatgpt.com/backend-api/codex".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
Some(serde_json::json!({"upstream_stream_policy":"force_stream"})),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
let encrypted_auth_config = encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"provider_type":"codex","refresh_token":"rt-codex-image-stream-local-123"}"#,
|
||||
)
|
||||
.expect("auth config should encrypt");
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-codex-image-stream-local-1".to_string(),
|
||||
"provider-codex-image-stream-local-1".to_string(),
|
||||
"oauth".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["openai:image"])),
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__")
|
||||
.expect("placeholder api key should encrypt"),
|
||||
Some(encrypted_auth_config),
|
||||
None,
|
||||
Some(serde_json::json!({"openai:image": 1})),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeStreamRequest>));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let seen_refresh = Arc::new(Mutex::new(None::<SeenRefreshRequest>));
|
||||
let seen_refresh_clone = Arc::clone(&seen_refresh);
|
||||
|
||||
let refresh = Router::new().route(
|
||||
"/oauth/token",
|
||||
any(move |request: Request| {
|
||||
let seen_refresh_inner = Arc::clone(&seen_refresh_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
*seen_refresh_inner.lock().expect("mutex should lock") = Some(SeenRefreshRequest {
|
||||
content_type: parts
|
||||
.headers
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
body: String::from_utf8(raw_body.to_vec())
|
||||
.expect("refresh body should be utf8"),
|
||||
});
|
||||
Json(json!({
|
||||
"access_token": "refreshed-codex-image-stream-access-token",
|
||||
"refresh_token": "rt-codex-image-stream-local-456",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/stream",
|
||||
any(move |request: Request| {
|
||||
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value = serde_json::from_slice(&raw_body)
|
||||
.expect("execution runtime payload should parse");
|
||||
*seen_execution_runtime_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(SeenExecutionRuntimeStreamRequest {
|
||||
trace_id: parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
url: payload
|
||||
.get("url")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
model: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("model"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
authorization: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("authorization"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
x_client_request_id: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-client-request-id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
tool_type: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("tools"))
|
||||
.and_then(|value| value.get(0))
|
||||
.and_then(|value| value.get("type"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
tool_partial_images: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("tools"))
|
||||
.and_then(|value| value.get(0))
|
||||
.and_then(|value| value.get("partial_images"))
|
||||
.and_then(|value| value.as_u64()),
|
||||
request_stream: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("stream"))
|
||||
.and_then(|value| value.as_bool())
|
||||
.unwrap_or(false),
|
||||
plan_stream: payload
|
||||
.get("stream")
|
||||
.and_then(|value| value.as_bool())
|
||||
.unwrap_or(false),
|
||||
});
|
||||
let frames = concat!(
|
||||
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.output_item.done\\ndata: {\\\"type\\\":\\\"response.output_item.done\\\",\\\"output_index\\\":0,\\\"item\\\":{\\\"id\\\":\\\"ig_123\\\",\\\"type\\\":\\\"image_generation_call\\\",\\\"result\\\":\\\"aGVsbG8=\\\"}}\\n\\n\"}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"tool_usage\\\":{\\\"image_gen\\\":{\\\"input_tokens\\\":11,\\\"output_tokens\\\":22,\\\"total_tokens\\\":33}}}}\\n\\n\"}}\n",
|
||||
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41}}}\n",
|
||||
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
|
||||
);
|
||||
let mut response = http::Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from(frames))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
http::HeaderValue::from_static("application/x-ndjson"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let client_api_key = "sk-client-codex-image-stream-local";
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key(client_api_key)),
|
||||
sample_auth_snapshot(
|
||||
"key-codex-image-stream-client-123",
|
||||
"user-codex-image-stream-client-123",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (refresh_url, refresh_handle) = start_server(refresh).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let oauth_refresh =
|
||||
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
|
||||
Arc::new(
|
||||
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
|
||||
.with_token_url_for_tests("codex", format!("{refresh_url}/oauth/token")),
|
||||
),
|
||||
]);
|
||||
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::new(InMemoryRequestCandidateRepository::default()),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
)
|
||||
.with_oauth_refresh_coordinator_for_tests(oauth_refresh);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/images/generations"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
format!("Bearer {client_api_key}"),
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-codex-image-stream-local-123")
|
||||
.body(
|
||||
"{\"model\":\"gpt-image-2\",\"prompt\":\"生成一张中国历史视觉海报\",\"stream\":true,\"partial_images\":1}",
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("text/event-stream")
|
||||
);
|
||||
let response_text = response.text().await.expect("body should read");
|
||||
assert!(response_text.contains("event: image_generation.partial_image"));
|
||||
assert!(response_text.contains("\"type\":\"image_generation.partial_image\""));
|
||||
assert!(response_text.contains("\"b64_json\":\"aGVsbG8=\""));
|
||||
assert!(response_text.contains("event: image_generation.completed"));
|
||||
assert!(response_text.contains("\"type\":\"image_generation.completed\""));
|
||||
assert!(response_text.contains("\"total_tokens\":33"));
|
||||
assert!(!response_text.contains("response.completed"));
|
||||
|
||||
let seen_refresh_request = seen_refresh
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("refresh request should be captured");
|
||||
assert_eq!(
|
||||
seen_refresh_request.content_type,
|
||||
"application/x-www-form-urlencoded"
|
||||
);
|
||||
assert!(seen_refresh_request
|
||||
.body
|
||||
.contains("refresh_token=rt-codex-image-stream-local-123"));
|
||||
|
||||
let seen_execution_runtime_request = seen_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution runtime stream should be captured");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.trace_id,
|
||||
"trace-codex-image-stream-local-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.url,
|
||||
"https://chatgpt.com/backend-api/codex/responses"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.model,
|
||||
CODEX_OPENAI_IMAGE_INTERNAL_MODEL
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.authorization,
|
||||
"Bearer refreshed-codex-image-stream-access-token"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.x_client_request_id,
|
||||
"trace-codex-image-stream-local-123"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.tool_type, "image_generation");
|
||||
assert_eq!(seen_execution_runtime_request.tool_partial_images, Some(1));
|
||||
assert!(seen_execution_runtime_request.request_stream);
|
||||
assert!(seen_execution_runtime_request.plan_stream);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
refresh_handle.abort();
|
||||
}
|
||||
@@ -24,3 +24,4 @@ use super::{
|
||||
};
|
||||
|
||||
mod decision;
|
||||
mod image;
|
||||
|
||||
@@ -2,6 +2,7 @@ use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
|
||||
to_bytes, Arc, Body, Json, Mutex, Request, Router, StatusCode, TRACE_ID_HEADER,
|
||||
};
|
||||
use crate::ai_pipeline::CODEX_OPENAI_IMAGE_INTERNAL_MODEL;
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
@@ -27,13 +28,13 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r
|
||||
model: String,
|
||||
authorization: String,
|
||||
x_client_request_id: String,
|
||||
user_agent: String,
|
||||
version: String,
|
||||
originator: String,
|
||||
prompt: String,
|
||||
content_is_string: bool,
|
||||
tool_type: String,
|
||||
tool_size: String,
|
||||
tool_quality: String,
|
||||
tool_background: String,
|
||||
tool_choice_type: String,
|
||||
tool_has_n: bool,
|
||||
request_stream: bool,
|
||||
plan_stream: bool,
|
||||
@@ -268,24 +269,6 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
user_agent: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("user-agent"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
version: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("version"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
originator: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("originator"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
prompt: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
@@ -320,6 +303,32 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
tool_quality: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("tools"))
|
||||
.and_then(|value| value.get(0))
|
||||
.and_then(|value| value.get("quality"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
tool_background: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("tools"))
|
||||
.and_then(|value| value.get(0))
|
||||
.and_then(|value| value.get("background"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
tool_choice_type: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("tool_choice"))
|
||||
.and_then(|value| value.get("type"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
tool_has_n: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
@@ -349,9 +358,13 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r
|
||||
concat!(
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_img_123\",\"created_at\":1776839946}}\n\n",
|
||||
"data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ig_123\",\"type\":\"image_generation_call\",\"status\":\"generating\",\"output_format\":\"png\",\"quality\":\"medium\",\"size\":\"1024x1024\",\"revised_prompt\":\"中国历史视觉海报\",\"result\":\"aGVsbG8=\"}}\n\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_img_123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":2440,\"output_tokens\":184,\"total_tokens\":2624},\"tool_usage\":{\"image_gen\":{\"input_tokens\":171,\"input_tokens_details\":{\"image_tokens\":0,\"text_tokens\":171},\"output_tokens\":1372,\"output_tokens_details\":{\"image_tokens\":1372,\"text_tokens\":0},\"total_tokens\":1543}}}}\n\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_img_123\",\"object\":\"response\",\"model\":\"__CODEX_IMAGE_MODEL__\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":2440,\"output_tokens\":184,\"total_tokens\":2624},\"tool_usage\":{\"image_gen\":{\"input_tokens\":171,\"input_tokens_details\":{\"image_tokens\":0,\"text_tokens\":171},\"output_tokens\":1372,\"output_tokens_details\":{\"image_tokens\":1372,\"text_tokens\":0},\"total_tokens\":1543}}}}\n\n",
|
||||
"data: [DONE]\n\n"
|
||||
)
|
||||
.replace(
|
||||
"__CODEX_IMAGE_MODEL__",
|
||||
CODEX_OPENAI_IMAGE_INTERNAL_MODEL,
|
||||
)
|
||||
)
|
||||
},
|
||||
"telemetry": {
|
||||
@@ -457,7 +470,10 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r
|
||||
seen_execution_runtime_request.url,
|
||||
"https://chatgpt.com/backend-api/codex/responses"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.model, "gpt-5.4");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.model,
|
||||
CODEX_OPENAI_IMAGE_INTERNAL_MODEL
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.authorization,
|
||||
"Bearer refreshed-codex-image-access-token"
|
||||
@@ -466,11 +482,6 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r
|
||||
seen_execution_runtime_request.x_client_request_id,
|
||||
"trace-codex-image-local-123"
|
||||
);
|
||||
assert!(seen_execution_runtime_request
|
||||
.user_agent
|
||||
.starts_with("codex-tui/0.122.0"));
|
||||
assert_eq!(seen_execution_runtime_request.version, "0.122.0");
|
||||
assert_eq!(seen_execution_runtime_request.originator, "codex_cli_rs");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.prompt,
|
||||
"生成一张中国历史视觉海报"
|
||||
@@ -478,6 +489,12 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r
|
||||
assert!(seen_execution_runtime_request.content_is_string);
|
||||
assert_eq!(seen_execution_runtime_request.tool_type, "image_generation");
|
||||
assert_eq!(seen_execution_runtime_request.tool_size, "1024x1024");
|
||||
assert_eq!(seen_execution_runtime_request.tool_quality, "high");
|
||||
assert_eq!(seen_execution_runtime_request.tool_background, "auto");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.tool_choice_type,
|
||||
"image_generation"
|
||||
);
|
||||
assert!(!seen_execution_runtime_request.tool_has_n);
|
||||
assert!(seen_execution_runtime_request.request_stream);
|
||||
assert!(!seen_execution_runtime_request.plan_stream);
|
||||
|
||||
@@ -985,9 +985,11 @@ fn admin_provider_write_uses_specific_local_owners() {
|
||||
for pattern in [
|
||||
"mod create;",
|
||||
"mod endpoint;",
|
||||
"mod template;",
|
||||
"mod update;",
|
||||
"pub(crate) use self::create::build_admin_create_provider_record;",
|
||||
"pub(crate) use self::endpoint::build_admin_fixed_provider_endpoint_record;",
|
||||
"pub(crate) use self::template::{",
|
||||
"pub(crate) use self::update::build_admin_update_provider_record;",
|
||||
] {
|
||||
assert!(
|
||||
@@ -1044,7 +1046,7 @@ fn admin_provider_write_uses_specific_local_owners() {
|
||||
for pattern in [
|
||||
"pub(crate) fn build_admin_fixed_provider_endpoint_record(",
|
||||
"admin_endpoint_signature_parts(",
|
||||
"normalize_admin_base_url(base_url)?",
|
||||
"normalize_admin_base_url(template.base_url)?",
|
||||
] {
|
||||
assert!(
|
||||
write_provider_endpoint.contains(pattern),
|
||||
@@ -1052,6 +1054,20 @@ fn admin_provider_write_uses_specific_local_owners() {
|
||||
);
|
||||
}
|
||||
|
||||
let write_provider_template = read_workspace_file(
|
||||
"apps/aether-gateway/src/handlers/admin/provider/write/provider/template.rs",
|
||||
);
|
||||
for pattern in [
|
||||
"pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints(",
|
||||
"pub(crate) fn apply_admin_fixed_provider_endpoint_template_overrides(",
|
||||
"const FIXED_PROVIDER_TEMPLATE_METADATA_KEY: &str = \"_aether_fixed_provider_template\";",
|
||||
] {
|
||||
assert!(
|
||||
write_provider_template.contains(pattern),
|
||||
"handlers/admin/provider/write/provider/template.rs should own {pattern}"
|
||||
);
|
||||
}
|
||||
|
||||
let write_key_create =
|
||||
read_workspace_file("apps/aether-gateway/src/handlers/admin/provider/write/keys/create.rs");
|
||||
for pattern in [
|
||||
|
||||
@@ -1481,6 +1481,83 @@ async fn gateway_marks_exhausted_codex_pool_key_as_blocked_when_flag_enabled() {
|
||||
assert_eq!(keys[0]["account_quota"], json!("5H剩余 0.0%"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_lists_inherited_fixed_provider_api_formats_for_pool_keys() {
|
||||
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
true,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(json!({
|
||||
"pool_advanced": {
|
||||
"enabled": true
|
||||
}
|
||||
})),
|
||||
);
|
||||
provider.provider_type = "codex".to_string();
|
||||
|
||||
let mut key = sample_key(
|
||||
"key-codex-inherited",
|
||||
"provider-codex",
|
||||
"openai:cli",
|
||||
"oauth-placeholder",
|
||||
);
|
||||
key.name = "codex inherited".to_string();
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.api_formats = Some(json!(["openai:cli", "openai:compact"]));
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
vec![
|
||||
sample_endpoint(
|
||||
"endpoint-codex-cli",
|
||||
"provider-codex",
|
||||
"openai:cli",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
),
|
||||
sample_endpoint(
|
||||
"endpoint-codex-image",
|
||||
"provider-codex",
|
||||
"openai:image",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
),
|
||||
],
|
||||
vec![key],
|
||||
));
|
||||
let state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
));
|
||||
|
||||
let response = local_admin_pool_response(
|
||||
&state,
|
||||
http::Method::GET,
|
||||
"/api/admin/pool/provider-codex/keys?page=1&page_size=50&status=all",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = serde_json::from_slice(
|
||||
&to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should read"),
|
||||
)
|
||||
.expect("json body should parse");
|
||||
let keys = payload["keys"].as_array().expect("keys should be array");
|
||||
|
||||
assert_eq!(keys.len(), 1);
|
||||
assert_eq!(
|
||||
keys[0]["api_formats"],
|
||||
json!(["openai:cli", "openai:image"])
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_prefers_status_snapshot_codex_quota_over_stale_metadata() {
|
||||
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(
|
||||
|
||||
@@ -1564,7 +1564,7 @@ async fn gateway_handles_openai_cli_test_model_locally() {
|
||||
.json_body
|
||||
.as_ref()
|
||||
.and_then(|body| body.get("instructions")),
|
||||
Some(&json!("You are GPT-5."))
|
||||
Some(&json!("You are ChatGPT."))
|
||||
);
|
||||
assert_eq!(
|
||||
plan.body
|
||||
|
||||
@@ -942,7 +942,7 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
|
||||
.list_endpoints_by_provider_ids(std::slice::from_ref(&created.id))
|
||||
.await
|
||||
.expect("endpoints should list");
|
||||
assert_eq!(endpoints.len(), 2);
|
||||
assert_eq!(endpoints.len(), 3);
|
||||
let cli_endpoint = endpoints
|
||||
.iter()
|
||||
.find(|endpoint| endpoint.api_format == "openai:cli")
|
||||
@@ -951,6 +951,10 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
|
||||
.iter()
|
||||
.find(|endpoint| endpoint.api_format == "openai:compact")
|
||||
.expect("compact endpoint should exist");
|
||||
let image_endpoint = endpoints
|
||||
.iter()
|
||||
.find(|endpoint| endpoint.api_format == "openai:image")
|
||||
.expect("image endpoint should exist");
|
||||
assert_eq!(
|
||||
cli_endpoint.base_url,
|
||||
"https://chatgpt.com/backend-api/codex"
|
||||
@@ -959,8 +963,13 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
|
||||
compact_endpoint.base_url,
|
||||
"https://chatgpt.com/backend-api/codex"
|
||||
);
|
||||
assert_eq!(
|
||||
image_endpoint.base_url,
|
||||
"https://chatgpt.com/backend-api/codex"
|
||||
);
|
||||
assert_eq!(cli_endpoint.max_retries, Some(7));
|
||||
assert_eq!(compact_endpoint.max_retries, Some(7));
|
||||
assert_eq!(image_endpoint.max_retries, Some(7));
|
||||
assert_eq!(
|
||||
cli_endpoint
|
||||
.config
|
||||
@@ -969,14 +978,220 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
|
||||
.and_then(serde_json::Value::as_str),
|
||||
Some("force_stream")
|
||||
);
|
||||
assert_eq!(
|
||||
image_endpoint
|
||||
.config
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("upstream_stream_policy"))
|
||||
.and_then(serde_json::Value::as_str),
|
||||
Some("force_stream")
|
||||
);
|
||||
assert!(cli_endpoint.body_rules.is_none());
|
||||
assert!(compact_endpoint.body_rules.is_none());
|
||||
assert!(image_endpoint.body_rules.is_none());
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoints() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/providers/provider-codex",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
true,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
);
|
||||
provider.provider_type = "codex".to_string();
|
||||
let mut cli_endpoint = sample_endpoint(
|
||||
"endpoint-codex-cli",
|
||||
"provider-codex",
|
||||
"openai:cli",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
);
|
||||
cli_endpoint.max_retries = Some(2);
|
||||
cli_endpoint.config = Some(json!({"upstream_stream_policy": "force_stream"}));
|
||||
let mut key = sample_key(
|
||||
"key-codex-oauth",
|
||||
"provider-codex",
|
||||
"openai:cli",
|
||||
"oauth-placeholder",
|
||||
);
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.api_formats = Some(json!(["openai:cli"]));
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
vec![cli_endpoint],
|
||||
vec![key],
|
||||
));
|
||||
|
||||
let (_upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.patch(format!("{gateway_url}/api/admin/providers/provider-codex"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"max_retries": 9
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let status = response.status();
|
||||
let body = response.text().await.expect("body should read");
|
||||
assert_eq!(status, StatusCode::OK, "body={body}");
|
||||
|
||||
let endpoints = provider_catalog_repository
|
||||
.list_endpoints_by_provider_ids(&["provider-codex".to_string()])
|
||||
.await
|
||||
.expect("endpoints should list");
|
||||
assert_eq!(endpoints.len(), 3);
|
||||
let cli_endpoint = endpoints
|
||||
.iter()
|
||||
.find(|endpoint| endpoint.api_format == "openai:cli")
|
||||
.expect("cli endpoint should exist");
|
||||
let compact_endpoint = endpoints
|
||||
.iter()
|
||||
.find(|endpoint| endpoint.api_format == "openai:compact")
|
||||
.expect("compact endpoint should exist");
|
||||
let image_endpoint = endpoints
|
||||
.iter()
|
||||
.find(|endpoint| endpoint.api_format == "openai:image")
|
||||
.expect("image endpoint should exist");
|
||||
|
||||
assert_eq!(cli_endpoint.max_retries, Some(9));
|
||||
assert_eq!(compact_endpoint.max_retries, Some(9));
|
||||
assert_eq!(image_endpoint.max_retries, Some(9));
|
||||
assert_eq!(
|
||||
cli_endpoint
|
||||
.config
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("_aether_fixed_provider_template"))
|
||||
.and_then(|value| value.get("managed"))
|
||||
.and_then(serde_json::Value::as_bool),
|
||||
Some(true)
|
||||
);
|
||||
assert_eq!(
|
||||
image_endpoint
|
||||
.config
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("upstream_stream_policy"))
|
||||
.and_then(serde_json::Value::as_str),
|
||||
Some("force_stream")
|
||||
);
|
||||
let keys = provider_catalog_repository
|
||||
.list_keys_by_provider_ids(&["provider-codex".to_string()])
|
||||
.await
|
||||
.expect("keys should list");
|
||||
assert_eq!(keys.len(), 1);
|
||||
assert!(keys[0].api_formats.is_none());
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_lists_effective_api_formats_for_fixed_oauth_provider_keys() {
|
||||
let mut provider = sample_provider("provider-codex", "codex", 10)
|
||||
.with_transport_fields(true, false, true, None, None, None, None, None, None);
|
||||
provider.provider_type = "codex".to_string();
|
||||
|
||||
let mut key = sample_key(
|
||||
"key-codex-legacy",
|
||||
"provider-codex",
|
||||
"openai:cli",
|
||||
"oauth-placeholder",
|
||||
);
|
||||
key.name = "codex legacy".to_string();
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.api_formats = Some(json!(["openai:cli", "openai:compact"]));
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
vec![
|
||||
sample_endpoint(
|
||||
"endpoint-codex-cli",
|
||||
"provider-codex",
|
||||
"openai:cli",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
),
|
||||
sample_endpoint(
|
||||
"endpoint-codex-image",
|
||||
"provider-codex",
|
||||
"openai:image",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
),
|
||||
],
|
||||
vec![key],
|
||||
));
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-codex/keys?skip=0&limit=100"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
let keys = payload.as_array().expect("keys payload should be array");
|
||||
|
||||
assert_eq!(keys.len(), 1);
|
||||
assert_eq!(
|
||||
keys[0]["api_formats"],
|
||||
json!(["openai:cli", "openai:image"])
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_health_monitor_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
|
||||
@@ -578,7 +578,7 @@ async fn gateway_rejects_gpt_image_2_on_chat_completions_without_hitting_fallbac
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(
|
||||
payload["detail"],
|
||||
"gpt-image-2 仅支持通过 /v1/images/generations 或 /v1/images/edits 调用"
|
||||
"图片模型仅支持通过 /v1/images/generations、/v1/images/edits 或 /v1/images/variations 调用"
|
||||
);
|
||||
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
@@ -640,7 +640,66 @@ async fn gateway_rejects_image_request_with_n_greater_than_one_without_hitting_f
|
||||
Some(EXECUTION_PATH_LOCAL_AI_PUBLIC)
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["detail"], "图片接口当前仅支持 n=1");
|
||||
assert_eq!(payload["detail"], "当前 Codex 图片反代仅支持 n=1");
|
||||
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
fallback_probe_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_variation_request_without_image_without_hitting_fallback_probe() {
|
||||
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
|
||||
let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
|
||||
let fallback_probe = Router::new().route(
|
||||
"/{*path}",
|
||||
any(move |_request: Request| {
|
||||
let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone);
|
||||
async move {
|
||||
*fallback_probe_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Json(json!({"proxied": true}))).into_response()
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-openai-image-variation")),
|
||||
unrestricted_models_snapshot("key-openai-image-variation", "user-openai-image-variation"),
|
||||
)]));
|
||||
|
||||
let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_auth_api_key_data_reader_for_tests(auth_repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/images/variations"))
|
||||
.header("authorization", "Bearer sk-openai-image-variation")
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.body(
|
||||
serde_json::to_vec(&json!({
|
||||
"model": "dall-e-2",
|
||||
"response_format": "url"
|
||||
}))
|
||||
.expect("request body should encode"),
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_AI_PUBLIC)
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["detail"], "图片变体请求需要 image 文件");
|
||||
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
|
||||
@@ -444,6 +444,94 @@ pub fn admin_usage_provider_key_name(
|
||||
.or_else(|| item.routing_key_name().map(ToOwned::to_owned))
|
||||
}
|
||||
|
||||
fn admin_usage_request_body_stream_flag(item: &StoredRequestUsageAudit) -> Option<bool> {
|
||||
item.request_body
|
||||
.as_ref()
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|body| body.get("stream"))
|
||||
.and_then(Value::as_bool)
|
||||
}
|
||||
|
||||
fn admin_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 admin_usage_request_body_implies_default_non_stream(item: &StoredRequestUsageAudit) -> bool {
|
||||
let Some(body) = item.request_body.as_ref().and_then(Value::as_object) else {
|
||||
return false;
|
||||
};
|
||||
!body.contains_key("stream") && admin_usage_api_format_defaults_to_non_stream(item)
|
||||
}
|
||||
|
||||
pub fn admin_usage_client_is_stream(item: &StoredRequestUsageAudit) -> bool {
|
||||
item.request_metadata
|
||||
.as_ref()
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|metadata| metadata.get("client_requested_stream"))
|
||||
.and_then(Value::as_bool)
|
||||
.or_else(|| admin_usage_request_body_stream_flag(item))
|
||||
.or_else(|| admin_usage_request_body_implies_default_non_stream(item).then_some(false))
|
||||
.unwrap_or(item.is_stream)
|
||||
}
|
||||
|
||||
fn admin_usage_active_request_json(
|
||||
item: &StoredRequestUsageAudit,
|
||||
api_key_name: Option<String>,
|
||||
provider_key_name: Option<String>,
|
||||
) -> Value {
|
||||
let cache_creation_input_tokens = admin_usage_cache_creation_tokens(item);
|
||||
let client_is_stream = admin_usage_client_is_stream(item);
|
||||
let mut value = json!({
|
||||
"id": item.id,
|
||||
"status": item.status,
|
||||
"input_tokens": item.input_tokens,
|
||||
"effective_input_tokens": admin_usage_effective_input_tokens(item),
|
||||
"output_tokens": item.output_tokens,
|
||||
"cache_creation_input_tokens": cache_creation_input_tokens,
|
||||
"cache_creation_ephemeral_5m_input_tokens": item.cache_creation_ephemeral_5m_input_tokens,
|
||||
"cache_creation_ephemeral_1h_input_tokens": item.cache_creation_ephemeral_1h_input_tokens,
|
||||
"cache_read_input_tokens": item.cache_read_input_tokens,
|
||||
"cost": round_to(item.total_cost_usd, 6),
|
||||
"actual_cost": round_to(item.actual_total_cost_usd, 6),
|
||||
"response_time_ms": item.response_time_ms,
|
||||
"first_byte_time_ms": item.first_byte_time_ms,
|
||||
"provider": item.provider_name,
|
||||
"api_key_name": api_key_name,
|
||||
"provider_key_name": provider_key_name,
|
||||
"is_stream": item.is_stream,
|
||||
"upstream_is_stream": item.is_stream,
|
||||
"client_requested_stream": client_is_stream,
|
||||
"client_is_stream": client_is_stream,
|
||||
"has_fallback": admin_usage_has_fallback(item),
|
||||
});
|
||||
if let Some(api_format) = item.api_format.as_ref() {
|
||||
value["api_format"] = json!(api_format);
|
||||
}
|
||||
if let Some(endpoint_api_format) = item.endpoint_api_format.as_ref() {
|
||||
value["endpoint_api_format"] = json!(endpoint_api_format);
|
||||
}
|
||||
value["has_format_conversion"] = json!(item.has_format_conversion);
|
||||
if let Some(target_model) = item.target_model.as_ref() {
|
||||
value["target_model"] = json!(target_model);
|
||||
}
|
||||
value
|
||||
}
|
||||
|
||||
pub fn admin_usage_record_json(
|
||||
item: &StoredRequestUsageAudit,
|
||||
users_by_id: &BTreeMap<String, StoredUserSummary>,
|
||||
@@ -469,8 +557,9 @@ pub fn admin_usage_record_json(
|
||||
let user_email = user
|
||||
.and_then(|value| value.email.clone())
|
||||
.unwrap_or_else(|| "已删除用户".to_string());
|
||||
let client_is_stream = admin_usage_client_is_stream(item);
|
||||
|
||||
json!({
|
||||
let mut payload = json!({
|
||||
"id": item.id,
|
||||
"user_id": item.user_id,
|
||||
"user_email": user_email,
|
||||
@@ -497,7 +586,6 @@ pub fn admin_usage_record_json(
|
||||
"response_time_ms": item.response_time_ms,
|
||||
"first_byte_time_ms": item.first_byte_time_ms,
|
||||
"created_at": unix_secs_to_rfc3339(item.created_at_unix_ms),
|
||||
"is_stream": item.is_stream,
|
||||
"input_price_per_1m": input_price_per_1m,
|
||||
"output_price_per_1m": output_price_per_1m,
|
||||
"cache_creation_price_per_1m": cache_creation_price_per_1m,
|
||||
@@ -515,7 +603,18 @@ pub fn admin_usage_record_json(
|
||||
"api_key_name": api_key_name,
|
||||
"provider_key_name": provider_key_name,
|
||||
"model_version": Value::Null,
|
||||
})
|
||||
});
|
||||
let object = payload
|
||||
.as_object_mut()
|
||||
.expect("admin usage record payload should be an object");
|
||||
object.insert("is_stream".to_string(), json!(item.is_stream));
|
||||
object.insert("upstream_is_stream".to_string(), json!(item.is_stream));
|
||||
object.insert(
|
||||
"client_requested_stream".to_string(),
|
||||
json!(client_is_stream),
|
||||
);
|
||||
object.insert("client_is_stream".to_string(), json!(client_is_stream));
|
||||
payload
|
||||
}
|
||||
|
||||
pub fn admin_usage_total_tokens(item: &StoredRequestUsageAudit) -> u64 {
|
||||
@@ -1250,12 +1349,13 @@ pub fn admin_usage_resolve_request_capture_body(
|
||||
body_override: Option<Value>,
|
||||
) -> Option<Value> {
|
||||
let resolved_model = item.model.clone();
|
||||
let client_is_stream = admin_usage_client_is_stream(item);
|
||||
let mut request_body = body_override.or_else(|| item.request_body.clone())?;
|
||||
if let Some(body) = request_body.as_object_mut() {
|
||||
body.entry("model".to_string())
|
||||
.or_insert_with(|| json!(resolved_model));
|
||||
if !body.contains_key("stream") {
|
||||
body.insert("stream".to_string(), json!(item.is_stream));
|
||||
body.insert("stream".to_string(), json!(client_is_stream));
|
||||
}
|
||||
if let Some(target_model) = item
|
||||
.target_model
|
||||
@@ -1419,37 +1519,7 @@ pub fn build_admin_usage_active_requests_response(
|
||||
let provider_key_name = admin_usage_provider_key_name(item, provider_key_names);
|
||||
let api_key_name =
|
||||
admin_usage_api_key_name(item, api_key_names, auth_api_key_reader_available);
|
||||
let cache_creation_input_tokens = admin_usage_cache_creation_tokens(item);
|
||||
let mut value = json!({
|
||||
"id": item.id,
|
||||
"status": item.status,
|
||||
"input_tokens": item.input_tokens,
|
||||
"effective_input_tokens": admin_usage_effective_input_tokens(item),
|
||||
"output_tokens": item.output_tokens,
|
||||
"cache_creation_input_tokens": cache_creation_input_tokens,
|
||||
"cache_creation_ephemeral_5m_input_tokens": item.cache_creation_ephemeral_5m_input_tokens,
|
||||
"cache_creation_ephemeral_1h_input_tokens": item.cache_creation_ephemeral_1h_input_tokens,
|
||||
"cache_read_input_tokens": item.cache_read_input_tokens,
|
||||
"cost": round_to(item.total_cost_usd, 6),
|
||||
"actual_cost": round_to(item.actual_total_cost_usd, 6),
|
||||
"response_time_ms": item.response_time_ms,
|
||||
"first_byte_time_ms": item.first_byte_time_ms,
|
||||
"provider": item.provider_name,
|
||||
"api_key_name": api_key_name,
|
||||
"provider_key_name": provider_key_name,
|
||||
"has_fallback": admin_usage_has_fallback(item),
|
||||
});
|
||||
if let Some(api_format) = item.api_format.as_ref() {
|
||||
value["api_format"] = json!(api_format);
|
||||
}
|
||||
if let Some(endpoint_api_format) = item.endpoint_api_format.as_ref() {
|
||||
value["endpoint_api_format"] = json!(endpoint_api_format);
|
||||
}
|
||||
value["has_format_conversion"] = json!(item.has_format_conversion);
|
||||
if let Some(target_model) = item.target_model.as_ref() {
|
||||
value["target_model"] = json!(target_model);
|
||||
}
|
||||
value
|
||||
admin_usage_active_request_json(item, api_key_name, provider_key_name)
|
||||
})
|
||||
.collect();
|
||||
|
||||
@@ -1674,9 +1744,10 @@ mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
admin_usage_has_body_value, admin_usage_has_fallback, admin_usage_is_failed,
|
||||
admin_usage_matches_search, admin_usage_matches_status, admin_usage_matches_username,
|
||||
admin_usage_record_json, build_admin_usage_detail_payload,
|
||||
admin_usage_active_request_json, admin_usage_client_is_stream, admin_usage_has_body_value,
|
||||
admin_usage_has_fallback, admin_usage_is_failed, admin_usage_matches_search,
|
||||
admin_usage_matches_status, admin_usage_matches_username, admin_usage_record_json,
|
||||
admin_usage_resolve_request_capture_body, build_admin_usage_detail_payload,
|
||||
};
|
||||
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UsageBodyField};
|
||||
|
||||
@@ -1738,6 +1809,100 @@ mod tests {
|
||||
assert!(admin_usage_matches_status(&item, Some("completed")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_requested_stream_prefers_request_metadata_flag() {
|
||||
let item = StoredRequestUsageAudit {
|
||||
is_stream: true,
|
||||
request_metadata: Some(json!({
|
||||
"client_requested_stream": false
|
||||
})),
|
||||
..sample_usage("completed", Some(200), None)
|
||||
};
|
||||
|
||||
assert!(!admin_usage_client_is_stream(&item));
|
||||
|
||||
let record = admin_usage_record_json(
|
||||
&item,
|
||||
&BTreeMap::new(),
|
||||
&BTreeMap::new(),
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
);
|
||||
assert_eq!(record["is_stream"], true);
|
||||
assert_eq!(record["upstream_is_stream"], true);
|
||||
assert_eq!(record["client_requested_stream"], false);
|
||||
assert_eq!(record["client_is_stream"], false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_requested_stream_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", Some(200), None)
|
||||
};
|
||||
|
||||
assert!(!admin_usage_client_is_stream(&item));
|
||||
|
||||
let active = admin_usage_active_request_json(&item, None, None);
|
||||
assert_eq!(active["is_stream"], true);
|
||||
assert_eq!(active["upstream_is_stream"], true);
|
||||
assert_eq!(active["client_requested_stream"], false);
|
||||
assert_eq!(active["client_is_stream"], false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_requested_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", Some(200), None)
|
||||
};
|
||||
|
||||
assert!(!admin_usage_client_is_stream(&item));
|
||||
|
||||
let record = admin_usage_record_json(
|
||||
&item,
|
||||
&BTreeMap::new(),
|
||||
&BTreeMap::new(),
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
);
|
||||
assert_eq!(record["is_stream"], true);
|
||||
assert_eq!(record["upstream_is_stream"], true);
|
||||
assert_eq!(record["client_requested_stream"], false);
|
||||
assert_eq!(record["client_is_stream"], false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replay_body_defaults_stream_to_client_requested_mode() {
|
||||
let item = StoredRequestUsageAudit {
|
||||
is_stream: true,
|
||||
request_body: Some(json!({
|
||||
"model": "gpt-5.4",
|
||||
"input": "hello"
|
||||
})),
|
||||
request_metadata: Some(json!({
|
||||
"client_requested_stream": false
|
||||
})),
|
||||
..sample_usage("completed", Some(200), None)
|
||||
};
|
||||
|
||||
let body = admin_usage_resolve_request_capture_body(&item, None)
|
||||
.expect("replay body should resolve");
|
||||
assert_eq!(body["stream"], false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn legacy_failure_signals_still_work_when_status_is_missing() {
|
||||
let item = StoredRequestUsageAudit {
|
||||
|
||||
@@ -48,7 +48,8 @@ pub use crate::contracts::{
|
||||
OPENAI_CLI_SYNC_FINALIZE_REPORT_KIND, OPENAI_CLI_SYNC_PLAN_KIND,
|
||||
OPENAI_CLI_SYNC_SUCCESS_REPORT_KIND, OPENAI_COMPACT_STREAM_PLAN_KIND,
|
||||
OPENAI_COMPACT_SYNC_ERROR_REPORT_KIND, OPENAI_COMPACT_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_COMPACT_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_COMPACT_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
|
||||
OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND,
|
||||
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
|
||||
@@ -135,7 +136,10 @@ pub use crate::planner::specialized::{
|
||||
resolve_stream_spec as resolve_gemini_files_stream_spec,
|
||||
resolve_sync_spec as resolve_gemini_files_sync_spec, LocalGeminiFilesSpec,
|
||||
},
|
||||
image::{resolve_sync_spec as resolve_local_image_sync_spec, LocalOpenAiImageSpec},
|
||||
image::{
|
||||
resolve_stream_spec as resolve_local_image_stream_spec,
|
||||
resolve_sync_spec as resolve_local_image_sync_spec, LocalOpenAiImageSpec,
|
||||
},
|
||||
video::{
|
||||
resolve_sync_spec as resolve_local_video_sync_spec, LocalVideoCreateFamily,
|
||||
LocalVideoCreateSpec,
|
||||
@@ -160,4 +164,7 @@ pub use crate::planner::standard::{
|
||||
resolve_sync_spec as resolve_openai_cli_sync_spec, LocalOpenAiCliSpec,
|
||||
},
|
||||
LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec,
|
||||
CODEX_OPENAI_IMAGE_DEFAULT_MODEL, CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT,
|
||||
CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL, CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT,
|
||||
CODEX_OPENAI_IMAGE_INTERNAL_MODEL,
|
||||
};
|
||||
|
||||
@@ -23,26 +23,28 @@ pub use plan_kinds::{
|
||||
GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
OPENAI_CLI_STREAM_PLAN_KIND, OPENAI_CLI_SYNC_PLAN_KIND, OPENAI_COMPACT_STREAM_PLAN_KIND,
|
||||
OPENAI_COMPACT_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
OPENAI_COMPACT_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
pub use report_kinds::{
|
||||
core_error_background_report_kind, core_error_default_client_api_format,
|
||||
core_success_background_report_kind, implicit_sync_finalize_report_kind,
|
||||
CLAUDE_CHAT_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CHAT_SYNC_ERROR_REPORT_KIND,
|
||||
CLAUDE_CHAT_SYNC_FINALIZE_REPORT_KIND, CLAUDE_CHAT_SYNC_SUCCESS_REPORT_KIND,
|
||||
CLAUDE_CLI_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CLI_SYNC_ERROR_REPORT_KIND,
|
||||
CLAUDE_CLI_SYNC_FINALIZE_REPORT_KIND, CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND,
|
||||
GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND, GEMINI_CHAT_SYNC_ERROR_REPORT_KIND,
|
||||
GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND, GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND,
|
||||
GEMINI_CLI_STREAM_SUCCESS_REPORT_KIND, GEMINI_CLI_SYNC_ERROR_REPORT_KIND,
|
||||
GEMINI_CLI_SYNC_FINALIZE_REPORT_KIND, GEMINI_CLI_SYNC_SUCCESS_REPORT_KIND,
|
||||
GEMINI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND, OPENAI_CHAT_STREAM_SUCCESS_REPORT_KIND,
|
||||
OPENAI_CHAT_SYNC_ERROR_REPORT_KIND, OPENAI_CHAT_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND, OPENAI_CLI_STREAM_SUCCESS_REPORT_KIND,
|
||||
OPENAI_CLI_SYNC_ERROR_REPORT_KIND, OPENAI_CLI_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_CLI_SYNC_SUCCESS_REPORT_KIND, OPENAI_COMPACT_SYNC_ERROR_REPORT_KIND,
|
||||
OPENAI_COMPACT_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
|
||||
core_success_background_report_kind, implicit_stream_success_report_kind,
|
||||
implicit_sync_finalize_report_kind, CLAUDE_CHAT_STREAM_SUCCESS_REPORT_KIND,
|
||||
CLAUDE_CHAT_SYNC_ERROR_REPORT_KIND, CLAUDE_CHAT_SYNC_FINALIZE_REPORT_KIND,
|
||||
CLAUDE_CHAT_SYNC_SUCCESS_REPORT_KIND, CLAUDE_CLI_STREAM_SUCCESS_REPORT_KIND,
|
||||
CLAUDE_CLI_SYNC_ERROR_REPORT_KIND, CLAUDE_CLI_SYNC_FINALIZE_REPORT_KIND,
|
||||
CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND, GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND,
|
||||
GEMINI_CHAT_SYNC_ERROR_REPORT_KIND, GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND,
|
||||
GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_STREAM_SUCCESS_REPORT_KIND,
|
||||
GEMINI_CLI_SYNC_ERROR_REPORT_KIND, GEMINI_CLI_SYNC_FINALIZE_REPORT_KIND,
|
||||
GEMINI_CLI_SYNC_SUCCESS_REPORT_KIND, GEMINI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_CHAT_STREAM_SUCCESS_REPORT_KIND, OPENAI_CHAT_SYNC_ERROR_REPORT_KIND,
|
||||
OPENAI_CHAT_SYNC_FINALIZE_REPORT_KIND, OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND,
|
||||
OPENAI_CLI_STREAM_SUCCESS_REPORT_KIND, OPENAI_CLI_SYNC_ERROR_REPORT_KIND,
|
||||
OPENAI_CLI_SYNC_FINALIZE_REPORT_KIND, OPENAI_CLI_SYNC_SUCCESS_REPORT_KIND,
|
||||
OPENAI_COMPACT_SYNC_ERROR_REPORT_KIND, OPENAI_COMPACT_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
|
||||
OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND, OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND,
|
||||
};
|
||||
|
||||
@@ -3,6 +3,7 @@ pub const GEMINI_FILES_UPLOAD_PLAN_KIND: &str = "gemini_files_upload";
|
||||
pub const GEMINI_FILES_LIST_PLAN_KIND: &str = "gemini_files_list";
|
||||
pub const GEMINI_FILES_DELETE_PLAN_KIND: &str = "gemini_files_delete";
|
||||
pub const GEMINI_FILES_DOWNLOAD_PLAN_KIND: &str = "gemini_files_download";
|
||||
pub const OPENAI_IMAGE_STREAM_PLAN_KIND: &str = "openai_image_stream";
|
||||
pub const OPENAI_IMAGE_SYNC_PLAN_KIND: &str = "openai_image_sync";
|
||||
pub const OPENAI_VIDEO_CONTENT_PLAN_KIND: &str = "openai_video_content";
|
||||
pub const OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND: &str = "openai_video_cancel_sync";
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use crate::contracts::{
|
||||
CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
|
||||
GEMINI_CLI_SYNC_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_CLI_SYNC_PLAN_KIND,
|
||||
OPENAI_COMPACT_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
|
||||
OPENAI_COMPACT_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
|
||||
};
|
||||
|
||||
pub const OPENAI_CHAT_SYNC_FINALIZE_REPORT_KIND: &str = "openai_chat_sync_finalize";
|
||||
@@ -27,6 +27,7 @@ pub const OPENAI_CHAT_STREAM_SUCCESS_REPORT_KIND: &str = "openai_chat_stream_suc
|
||||
pub const CLAUDE_CHAT_STREAM_SUCCESS_REPORT_KIND: &str = "claude_chat_stream_success";
|
||||
pub const GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND: &str = "gemini_chat_stream_success";
|
||||
pub const OPENAI_CLI_STREAM_SUCCESS_REPORT_KIND: &str = "openai_cli_stream_success";
|
||||
pub const OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND: &str = "openai_image_stream_success";
|
||||
pub const CLAUDE_CLI_STREAM_SUCCESS_REPORT_KIND: &str = "claude_cli_stream_success";
|
||||
pub const GEMINI_CLI_STREAM_SUCCESS_REPORT_KIND: &str = "gemini_cli_stream_success";
|
||||
|
||||
@@ -93,3 +94,10 @@ pub fn core_success_background_report_kind(report_kind: &str) -> Option<&'static
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn implicit_stream_success_report_kind(plan_kind: &str) -> Option<&'static str> {
|
||||
match plan_kind {
|
||||
OPENAI_IMAGE_STREAM_PLAN_KIND => Some(OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,6 +7,7 @@ use crate::adaptation::surfaces::{
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum FinalizeStreamRewriteMode {
|
||||
EnvelopeUnwrap,
|
||||
OpenAiImage,
|
||||
Standard,
|
||||
KiroToClaudeCli,
|
||||
}
|
||||
@@ -45,6 +46,10 @@ pub fn resolve_finalize_stream_rewrite_mode(
|
||||
.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:cli" && client_api_format == "claude:cli")
|
||||
.then_some(FinalizeStreamRewriteMode::KiroToClaudeCli);
|
||||
@@ -144,4 +149,17 @@ mod tests {
|
||||
});
|
||||
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)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,9 +8,10 @@ use crate::contracts::{
|
||||
GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
OPENAI_CLI_STREAM_PLAN_KIND, OPENAI_CLI_SYNC_PLAN_KIND, OPENAI_COMPACT_STREAM_PLAN_KIND,
|
||||
OPENAI_COMPACT_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
OPENAI_COMPACT_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
|
||||
pub fn resolve_execution_runtime_stream_plan_kind(
|
||||
@@ -88,6 +89,14 @@ pub fn resolve_execution_runtime_stream_plan_kind(
|
||||
return Some(OPENAI_COMPACT_STREAM_PLAN_KIND);
|
||||
}
|
||||
|
||||
if route_family == Some("openai")
|
||||
&& route_kind == Some("image")
|
||||
&& *method == Method::POST
|
||||
&& matches!(path, "/v1/images/generations" | "/v1/images/edits")
|
||||
{
|
||||
return Some(OPENAI_IMAGE_STREAM_PLAN_KIND);
|
||||
}
|
||||
|
||||
if route_family == Some("openai")
|
||||
&& route_kind == Some("video")
|
||||
&& *method == Method::GET
|
||||
@@ -171,7 +180,10 @@ pub fn resolve_execution_runtime_sync_plan_kind(
|
||||
if route_family == Some("openai")
|
||||
&& route_kind == Some("image")
|
||||
&& *method == Method::POST
|
||||
&& matches!(path, "/v1/images/generations" | "/v1/images/edits")
|
||||
&& matches!(
|
||||
path,
|
||||
"/v1/images/generations" | "/v1/images/edits" | "/v1/images/variations"
|
||||
)
|
||||
{
|
||||
return Some(OPENAI_IMAGE_SYNC_PLAN_KIND);
|
||||
}
|
||||
@@ -258,7 +270,8 @@ pub fn is_matching_stream_request(
|
||||
| CLAUDE_CHAT_STREAM_PLAN_KIND
|
||||
| OPENAI_CLI_STREAM_PLAN_KIND
|
||||
| OPENAI_COMPACT_STREAM_PLAN_KIND
|
||||
| CLAUDE_CLI_STREAM_PLAN_KIND => body_json
|
||||
| CLAUDE_CLI_STREAM_PLAN_KIND
|
||||
| OPENAI_IMAGE_STREAM_PLAN_KIND => body_json
|
||||
.get("stream")
|
||||
.and_then(|value| value.as_bool())
|
||||
.unwrap_or(false),
|
||||
@@ -300,6 +313,7 @@ pub fn supports_stream_scheduler_decision_kind(plan_kind: &str) -> bool {
|
||||
| CLAUDE_CHAT_STREAM_PLAN_KIND
|
||||
| GEMINI_CHAT_STREAM_PLAN_KIND
|
||||
| OPENAI_CLI_STREAM_PLAN_KIND
|
||||
| OPENAI_IMAGE_STREAM_PLAN_KIND
|
||||
| OPENAI_COMPACT_STREAM_PLAN_KIND
|
||||
| CLAUDE_CLI_STREAM_PLAN_KIND
|
||||
| GEMINI_CLI_STREAM_PLAN_KIND
|
||||
@@ -318,7 +332,8 @@ mod tests {
|
||||
supports_sync_scheduler_decision_kind,
|
||||
};
|
||||
use crate::contracts::{
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
|
||||
OPENAI_IMAGE_SYNC_PLAN_KIND,
|
||||
};
|
||||
|
||||
#[test]
|
||||
@@ -387,8 +402,59 @@ mod tests {
|
||||
),
|
||||
Some(OPENAI_IMAGE_SYNC_PLAN_KIND)
|
||||
);
|
||||
assert_eq!(
|
||||
resolve_execution_runtime_sync_plan_kind(
|
||||
Some("ai_public"),
|
||||
Some("openai"),
|
||||
Some("image"),
|
||||
&Method::POST,
|
||||
"/v1/images/variations",
|
||||
),
|
||||
Some(OPENAI_IMAGE_SYNC_PLAN_KIND)
|
||||
);
|
||||
assert!(supports_sync_scheduler_decision_kind(
|
||||
OPENAI_IMAGE_SYNC_PLAN_KIND
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_openai_image_stream_plan_kind() {
|
||||
assert_eq!(
|
||||
resolve_execution_runtime_stream_plan_kind(
|
||||
Some("ai_public"),
|
||||
Some("openai"),
|
||||
Some("image"),
|
||||
&Method::POST,
|
||||
"/v1/images/generations",
|
||||
),
|
||||
Some(OPENAI_IMAGE_STREAM_PLAN_KIND)
|
||||
);
|
||||
assert_eq!(
|
||||
resolve_execution_runtime_stream_plan_kind(
|
||||
Some("ai_public"),
|
||||
Some("openai"),
|
||||
Some("image"),
|
||||
&Method::POST,
|
||||
"/v1/images/edits",
|
||||
),
|
||||
Some(OPENAI_IMAGE_STREAM_PLAN_KIND)
|
||||
);
|
||||
assert!(supports_stream_scheduler_decision_kind(
|
||||
OPENAI_IMAGE_STREAM_PLAN_KIND
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stream_matching_requires_openai_image_stream_flag() {
|
||||
assert!(!is_matching_stream_request(
|
||||
OPENAI_IMAGE_STREAM_PLAN_KIND,
|
||||
"/v1/images/generations",
|
||||
&serde_json::json!({"stream": false}),
|
||||
));
|
||||
assert!(is_matching_stream_request(
|
||||
OPENAI_IMAGE_STREAM_PLAN_KIND,
|
||||
"/v1/images/generations",
|
||||
&serde_json::json!({"stream": true}),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
use crate::contracts::OPENAI_IMAGE_SYNC_PLAN_KIND;
|
||||
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> {
|
||||
@@ -13,6 +14,19 @@ pub fn resolve_sync_spec(plan_kind: &str) -> Option<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,
|
||||
}
|
||||
@@ -20,12 +34,21 @@ pub fn resolve_sync_spec(plan_kind: &str) -> Option<LocalOpenAiImageSpec> {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::resolve_sync_spec;
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,10 +8,19 @@ use sha2::Sha256;
|
||||
use uuid::Uuid;
|
||||
|
||||
const CODEX_PROMPT_CACHE_NAMESPACE_VERSION: &str = "v3";
|
||||
const CODEX_DEFAULT_INSTRUCTIONS: &str = "You are ChatGPT.";
|
||||
const CODEX_DEFAULT_USER_AGENT: &str =
|
||||
"codex-tui/0.122.0 (Aether; x86_64) vscode/3.0.12 (codex-tui; 0.122.0)";
|
||||
const CODEX_DEFAULT_VERSION: &str = "0.122.0";
|
||||
const CODEX_DEFAULT_ORIGINATOR: &str = "codex_cli_rs";
|
||||
"codex-tui/0.122.0 (Mac OS 15.2.0; arm64) vscode/2.6.11 (codex-tui; 0.122.0)";
|
||||
const CODEX_DEFAULT_ORIGINATOR: &str = "codex-tui";
|
||||
pub const CODEX_OPENAI_IMAGE_INTERNAL_MODEL: &str = "gpt-5.4-mini";
|
||||
pub const CODEX_OPENAI_IMAGE_DEFAULT_MODEL: &str = "gpt-image-2";
|
||||
pub const CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL: &str = "dall-e-2";
|
||||
pub const CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT: &str = "png";
|
||||
pub const CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT: &str =
|
||||
"Create a faithful variation of the provided image.";
|
||||
const CODEX_IMAGE_TOOL_DEFAULT_SIZE: &str = "1024x1024";
|
||||
const CODEX_IMAGE_TOOL_DEFAULT_QUALITY: &str = "high";
|
||||
const CODEX_IMAGE_TOOL_DEFAULT_BACKGROUND: &str = "auto";
|
||||
const UUID_NAMESPACE_OID_BYTES: [u8; 16] = [
|
||||
0x6b, 0xa7, 0xb8, 0x12, 0x9d, 0xad, 0x11, 0xd1, 0x80, 0xb4, 0x00, 0xc0, 0x4f, 0xd4, 0x30, 0xc8,
|
||||
];
|
||||
@@ -30,6 +39,100 @@ fn is_openai_compact_request(provider_api_format: &str) -> bool {
|
||||
.eq_ignore_ascii_case("openai:compact")
|
||||
}
|
||||
|
||||
fn is_openai_image_request(provider_api_format: &str) -> bool {
|
||||
provider_api_format
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("openai:image")
|
||||
}
|
||||
|
||||
fn apply_codex_openai_image_tool_overrides(body_object: &mut serde_json::Map<String, Value>) {
|
||||
let mut tool = body_object
|
||||
.get("tools")
|
||||
.and_then(Value::as_array)
|
||||
.and_then(|tools| tools.first())
|
||||
.and_then(Value::as_object)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
|
||||
tool.insert("type".to_string(), json!("image_generation"));
|
||||
tool.entry("output_format".to_string())
|
||||
.or_insert_with(|| json!(CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT));
|
||||
if !tool.contains_key("action") {
|
||||
tool.entry("size".to_string())
|
||||
.or_insert_with(|| json!(CODEX_IMAGE_TOOL_DEFAULT_SIZE));
|
||||
tool.entry("quality".to_string())
|
||||
.or_insert_with(|| json!(CODEX_IMAGE_TOOL_DEFAULT_QUALITY));
|
||||
tool.entry("background".to_string())
|
||||
.or_insert_with(|| json!(CODEX_IMAGE_TOOL_DEFAULT_BACKGROUND));
|
||||
}
|
||||
|
||||
body_object.insert("tools".to_string(), json!([tool]));
|
||||
body_object.insert(
|
||||
"tool_choice".to_string(),
|
||||
json!({
|
||||
"type": "image_generation"
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
fn codex_openai_image_has_prompt(body_object: &serde_json::Map<String, Value>) -> bool {
|
||||
body_object
|
||||
.get("input")
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(Value::as_object)
|
||||
.filter_map(|item| item.get("content"))
|
||||
.any(|content| match content {
|
||||
Value::String(text) => !text.trim().is_empty(),
|
||||
Value::Array(items) => items.iter().any(|item| {
|
||||
item.as_object()
|
||||
.filter(|item| item.get("type").and_then(Value::as_str) == Some("input_text"))
|
||||
.and_then(|item| item.get("text").and_then(Value::as_str))
|
||||
.map(str::trim)
|
||||
.is_some_and(|text| !text.is_empty())
|
||||
}),
|
||||
_ => false,
|
||||
})
|
||||
}
|
||||
|
||||
fn inject_codex_default_variation_prompt(body_object: &mut serde_json::Map<String, Value>) {
|
||||
let Some(action) = body_object
|
||||
.get("tools")
|
||||
.and_then(Value::as_array)
|
||||
.and_then(|tools| tools.first())
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|tool| tool.get("action"))
|
||||
.and_then(Value::as_str)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
if action != "edit" || codex_openai_image_has_prompt(body_object) {
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(input) = body_object.get_mut("input").and_then(Value::as_array_mut) else {
|
||||
return;
|
||||
};
|
||||
let Some(first_message) = input.first_mut().and_then(Value::as_object_mut) else {
|
||||
return;
|
||||
};
|
||||
let Some(content) = first_message
|
||||
.get_mut("content")
|
||||
.and_then(Value::as_array_mut)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
|
||||
content.insert(
|
||||
0,
|
||||
json!({
|
||||
"type": "input_text",
|
||||
"text": CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
fn build_stable_codex_prompt_cache_key(user_api_key_id: &str) -> Option<String> {
|
||||
let normalized = user_api_key_id.trim();
|
||||
if normalized.is_empty() {
|
||||
@@ -111,6 +214,21 @@ fn extract_codex_account_id(decrypted_auth_config_raw: Option<&str>) -> Option<S
|
||||
})
|
||||
}
|
||||
|
||||
fn maybe_insert_default_codex_header(
|
||||
provider_request_headers: &mut BTreeMap<String, String>,
|
||||
original_headers: &http::HeaderMap,
|
||||
header_name: &str,
|
||||
header_value: &str,
|
||||
) {
|
||||
if header_map_has_non_empty_value(original_headers, header_name)
|
||||
|| btree_map_has_non_empty_value(provider_request_headers, header_name)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
provider_request_headers.insert(header_name.to_string(), header_value.to_string());
|
||||
}
|
||||
|
||||
fn maybe_inject_codex_prompt_cache_key(
|
||||
provider_request_body: &mut Value,
|
||||
provider_type: &str,
|
||||
@@ -196,7 +314,19 @@ pub fn apply_codex_openai_cli_special_body_edits(
|
||||
if !body_rules_handle_path(body_rules, "instructions")
|
||||
&& !body_object.contains_key("instructions")
|
||||
{
|
||||
body_object.insert("instructions".to_string(), json!("You are GPT-5."));
|
||||
body_object.insert(
|
||||
"instructions".to_string(),
|
||||
json!(CODEX_DEFAULT_INSTRUCTIONS),
|
||||
);
|
||||
}
|
||||
if is_openai_image_request(provider_api_format) {
|
||||
body_object.insert(
|
||||
"model".to_string(),
|
||||
json!(CODEX_OPENAI_IMAGE_INTERNAL_MODEL),
|
||||
);
|
||||
body_object.insert("stream".to_string(), json!(true));
|
||||
apply_codex_openai_image_tool_overrides(body_object);
|
||||
inject_codex_default_variation_prompt(body_object);
|
||||
}
|
||||
|
||||
maybe_inject_codex_prompt_cache_key(
|
||||
@@ -243,27 +373,18 @@ pub fn apply_codex_openai_cli_special_headers(
|
||||
}
|
||||
}
|
||||
|
||||
if !header_map_has_non_empty_value(original_headers, "user-agent")
|
||||
&& !btree_map_has_non_empty_value(provider_request_headers, "user-agent")
|
||||
{
|
||||
provider_request_headers.insert(
|
||||
"user-agent".to_string(),
|
||||
CODEX_DEFAULT_USER_AGENT.to_string(),
|
||||
if !is_openai_image_request(provider_api_format) {
|
||||
maybe_insert_default_codex_header(
|
||||
provider_request_headers,
|
||||
original_headers,
|
||||
"user-agent",
|
||||
CODEX_DEFAULT_USER_AGENT,
|
||||
);
|
||||
}
|
||||
|
||||
if !header_map_has_non_empty_value(original_headers, "version")
|
||||
&& !btree_map_has_non_empty_value(provider_request_headers, "version")
|
||||
{
|
||||
provider_request_headers.insert("version".to_string(), CODEX_DEFAULT_VERSION.to_string());
|
||||
}
|
||||
|
||||
if !header_map_has_non_empty_value(original_headers, "originator")
|
||||
&& !btree_map_has_non_empty_value(provider_request_headers, "originator")
|
||||
{
|
||||
provider_request_headers.insert(
|
||||
"originator".to_string(),
|
||||
CODEX_DEFAULT_ORIGINATOR.to_string(),
|
||||
maybe_insert_default_codex_header(
|
||||
provider_request_headers,
|
||||
original_headers,
|
||||
"originator",
|
||||
CODEX_DEFAULT_ORIGINATOR,
|
||||
);
|
||||
}
|
||||
|
||||
@@ -289,3 +410,100 @@ pub fn apply_codex_openai_cli_special_headers(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{apply_codex_openai_cli_special_body_edits, CODEX_OPENAI_IMAGE_INTERNAL_MODEL};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn codex_image_body_edits_force_tool_choice_and_default_generate_tool_fields() {
|
||||
let mut provider_request_body = json!({
|
||||
"input": [{
|
||||
"role": "user",
|
||||
"content": "generate image"
|
||||
}],
|
||||
"tools": [{
|
||||
"type": "image_generation"
|
||||
}],
|
||||
"tool_choice": "auto"
|
||||
});
|
||||
|
||||
apply_codex_openai_cli_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
"codex",
|
||||
"openai:image",
|
||||
None,
|
||||
None,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
provider_request_body["tools"][0]["size"],
|
||||
json!("1024x1024")
|
||||
);
|
||||
assert_eq!(provider_request_body["tools"][0]["quality"], json!("high"));
|
||||
assert_eq!(
|
||||
provider_request_body["tools"][0]["background"],
|
||||
json!("auto")
|
||||
);
|
||||
assert_eq!(
|
||||
provider_request_body["tools"][0]["output_format"],
|
||||
json!("png")
|
||||
);
|
||||
assert_eq!(
|
||||
provider_request_body["model"],
|
||||
json!(CODEX_OPENAI_IMAGE_INTERNAL_MODEL)
|
||||
);
|
||||
assert_eq!(provider_request_body["stream"], json!(true));
|
||||
assert_eq!(
|
||||
provider_request_body["tool_choice"]["type"],
|
||||
json!("image_generation")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_image_body_edits_preserve_edit_action_without_generate_defaults() {
|
||||
let mut provider_request_body = json!({
|
||||
"tools": [{
|
||||
"type": "image_generation",
|
||||
"action": "edit",
|
||||
"input_image_mask": { "image_url": "data:image/png;base64,mask" }
|
||||
}],
|
||||
"input": [{
|
||||
"role": "user",
|
||||
"content": [{
|
||||
"type": "input_image",
|
||||
"image_url": "data:image/png;base64,image"
|
||||
}]
|
||||
}],
|
||||
"tool_choice": "auto"
|
||||
});
|
||||
|
||||
apply_codex_openai_cli_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
"codex",
|
||||
"openai:image",
|
||||
None,
|
||||
None,
|
||||
);
|
||||
|
||||
assert_eq!(provider_request_body["tools"][0]["action"], json!("edit"));
|
||||
assert!(provider_request_body["tools"][0].get("size").is_none());
|
||||
assert!(provider_request_body["tools"][0].get("quality").is_none());
|
||||
assert!(provider_request_body["tools"][0]
|
||||
.get("background")
|
||||
.is_none());
|
||||
assert_eq!(
|
||||
provider_request_body["tools"][0]["output_format"],
|
||||
json!("png")
|
||||
);
|
||||
assert_eq!(
|
||||
provider_request_body["input"][0]["content"][0]["text"],
|
||||
json!("Create a faithful variation of the provided image.")
|
||||
);
|
||||
assert_eq!(
|
||||
provider_request_body["tool_choice"]["type"],
|
||||
json!("image_generation")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,7 +8,9 @@ pub mod openai_cli;
|
||||
|
||||
pub use codex::{
|
||||
apply_codex_openai_cli_special_body_edits, apply_codex_openai_cli_special_headers,
|
||||
apply_openai_compact_special_body_edits,
|
||||
apply_openai_compact_special_body_edits, CODEX_OPENAI_IMAGE_DEFAULT_MODEL,
|
||||
CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT, CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL,
|
||||
CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT, CODEX_OPENAI_IMAGE_INTERNAL_MODEL,
|
||||
};
|
||||
pub use family::{LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec};
|
||||
pub use matrix::{
|
||||
|
||||
@@ -80,6 +80,13 @@ where
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len() == 0
|
||||
}
|
||||
|
||||
pub fn clear(&self) {
|
||||
let Ok(mut entries) = self.entries.lock() else {
|
||||
return;
|
||||
};
|
||||
entries.clear();
|
||||
}
|
||||
}
|
||||
|
||||
impl<K, V> ExpiringMap<K, V>
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user