feat(rust): 支持 image 同步转流式 SSE、free_team_first 调度预设及账户错误自动重试

- openai:image 同步响应桥接为流式 SSE(image_generation.completed / image_edit.completed 事件)
- image 请求解析与前置校验移除模型白名单,支持任意自定义模型名
- 调度器新增 free_team_first 预设(mode: both / free_only / team_only)
- pool config 解析重构:新增 POOL_ALLOWED_SCHEDULING_PRESETS 白名单,规范化 mode 字段
- 错误分类器新增账户/账单错误模式集,升级为 RetryUpstreamFailure 而非 StopSemanticClientError
- 修复 stream execution 中上游 headers 与输出 headers 混用导致 content-type 判断错误的问题
- codex image 工具始终写入 action 字段(generate/edit),仅 generate 操作填充默认 size/quality
- 前端:pool 节点组只展示实际执行过的候选节点,全部 skipped 时折叠为最后一个节点
- Redis 测试就绪检测从 TCP 连接改为 PING/PONG 协议验证
This commit is contained in:
fawney19
2026-04-24 10:05:55 +08:00
parent a9d10163af
commit 5d710d9d10
16 changed files with 1388 additions and 204 deletions

View File

@@ -309,6 +309,7 @@ pub(crate) fn resolve_core_stream_error_finalize_report_kind(
pub(crate) fn resolve_core_stream_direct_finalize_report_kind(plan_kind: &str) -> Option<String> {
let report_kind = match plan_kind {
"openai_chat_stream" => "openai_chat_sync_finalize",
"openai_image_stream" => "openai_image_sync_finalize",
"claude_chat_stream" => "claude_chat_sync_finalize",
"gemini_chat_stream" => "gemini_chat_sync_finalize",
"openai_cli_stream" => "openai_cli_sync_finalize",

View File

@@ -895,6 +895,7 @@ async fn execute_stream_from_frame_stream(
let direct_stream_finalize_kind = resolve_core_stream_direct_finalize_report_kind(plan_kind);
let normalized_stream_report_context =
normalize_provider_private_report_context(report_context.as_ref());
let upstream_headers = headers.clone();
let mut private_stream_normalizer =
maybe_build_provider_private_stream_normalizer(report_context.as_ref());
let mut local_stream_rewriter =
@@ -904,10 +905,10 @@ async fn execute_stream_from_frame_stream(
headers.remove("content-length");
headers.insert("content-type".to_string(), "text/event-stream".to_string());
}
let content_type = headers.get("content-type").map(String::as_str);
let upstream_content_type = upstream_headers.get("content-type").map(String::as_str);
let skip_direct_finalize_prefetch = should_skip_direct_finalize_prefetch(
direct_stream_finalize_kind.as_deref(),
content_type,
upstream_content_type,
plan.provider_api_format.as_str(),
plan.client_api_format.as_str(),
private_stream_normalizer.is_some(),
@@ -933,7 +934,7 @@ async fn execute_stream_from_frame_stream(
key_id = %plan.key_id,
model_name,
candidate_index = candidate_index.as_str(),
content_type = content_type.unwrap_or("-"),
content_type = upstream_content_type.unwrap_or("-"),
provider_api_format = plan.provider_api_format.as_str(),
client_api_format = plan.client_api_format.as_str(),
"gateway skipped direct finalize prefetch for same-format passthrough stream"
@@ -1012,8 +1013,10 @@ async fn execute_stream_from_frame_stream(
provider_prefetched_body.extend_from_slice(&chunk);
prefetched_inspection_body.extend_from_slice(&chunk);
let inspection =
inspect_prefetched_stream_body(&headers, &prefetched_inspection_body);
let inspection = inspect_prefetched_stream_body(
&upstream_headers,
&prefetched_inspection_body,
);
match inspection {
StreamPrefetchInspection::EmbeddedError(body_json) => {
debug!(
@@ -1062,7 +1065,8 @@ async fn execute_stream_from_frame_stream(
StreamPrefetchInspection::NonError => {}
}
if !response_headers_indicate_sse(&headers) && (200..300).contains(&status_code)
if !response_headers_indicate_sse(&upstream_headers)
&& (200..300).contains(&status_code)
{
if let Some(body_json) =
parse_prefetched_sync_json_body(&prefetched_inspection_body)
@@ -1285,7 +1289,7 @@ async fn execute_stream_from_frame_stream(
let request_id_for_report_log = short_request_id(&request_id);
let candidate_id_for_report = candidate_id.clone();
let emit_passthrough_sse_terminal_error =
skip_direct_finalize_prefetch && response_headers_indicate_sse(&headers);
skip_direct_finalize_prefetch && response_headers_indicate_sse(&upstream_headers);
let body_capture_policy = match UsageRuntimeAccess::body_capture_policy(state.data.as_ref())
.await
{
@@ -2138,6 +2142,119 @@ mod tests {
server.abort();
}
#[tokio::test]
async fn execute_execution_runtime_stream_bridges_openai_image_sync_json_from_remote_runtime_to_image_sse(
) {
let listener = crate::test_support::bind_loopback_listener()
.await
.expect("listener should bind");
let addr = listener.local_addr().expect("local addr should resolve");
let server = tokio::spawn(async move {
let app = Router::new().route(
"/v1/execute/stream",
any(|_request: Request| async move {
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"application/json\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"{\\\"created\\\":1776972364,\\\"data\\\":[{\\\"b64_json\\\":\\\"aGVsbG8=\\\"}],\\\"usage\\\":{\\\"total_tokens\\\":100,\\\"input_tokens\\\":50,\\\"output_tokens\\\":50,\\\"input_tokens_details\\\":{\\\"text_tokens\\\":10,\\\"image_tokens\\\":40}}}\"}}\n",
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41}}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
);
let mut response = axum::http::Response::new(Body::from(frames));
response.headers_mut().insert(
header::CONTENT_TYPE,
HeaderValue::from_static("application/x-ndjson"),
);
response
}),
);
axum::serve(listener, app)
.await
.expect("server should start");
});
let state = AppState::new()
.expect("app state should build")
.with_execution_runtime_override_base_url(format!("http://{addr}"));
let plan = ExecutionPlan {
request_id: "req-remote-runtime-image-sync-json-stream".into(),
candidate_id: Some("cand-remote-runtime-image-sync-json-stream".into()),
provider_name: Some("openai".into()),
provider_id: "prov-1".into(),
endpoint_id: "ep-1".into(),
key_id: "key-1".into(),
method: "POST".into(),
url: "https://chatgpt.com/backend-api/codex/responses".into(),
headers: BTreeMap::from([
("content-type".into(), "application/json".into()),
("accept".into(), "text/event-stream".into()),
]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gpt-image-1",
"prompt": "hello",
"stream": true
})),
stream: true,
client_api_format: "openai:image".into(),
provider_api_format: "openai:image".into(),
model_name: Some("gpt-image-1".into()),
proxy: None,
tls_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(5_000),
total_ms: Some(5_000),
..ExecutionTimeouts::default()
}),
};
let decision = GatewayControlDecision::synthetic(
"/v1/images/generations",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("image".to_string()),
Some("openai:image".to_string()),
)
.with_execution_runtime_candidate(true);
let response = execute_execution_runtime_stream(
&state,
plan,
"trace-remote-runtime-image-sync-json-stream",
&decision,
"openai_image_stream",
None,
Some(json!({
"provider_api_format": "openai:image",
"client_api_format": "openai:image",
"mapped_model": "gpt-image-1",
"image_request": {
"operation": "generate"
}
})),
)
.await
.expect("execution should succeed")
.expect("execution should return a client response");
assert_eq!(
response
.headers()
.get(header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok()),
Some("text/event-stream")
);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body should read");
let text = String::from_utf8(body.to_vec()).expect("response body should be utf8");
assert!(text.contains("event: image_generation.completed"));
assert!(text.contains("\"type\":\"image_generation.completed\""));
assert!(text.contains("\"b64_json\":\"aGVsbG8=\""));
assert!(text.contains("\"total_tokens\":100"));
server.abort();
}
#[tokio::test]
async fn execute_execution_runtime_stream_returns_client_error_with_local_tunnel_message_before_first_data(
) {

View File

@@ -1014,6 +1014,146 @@ mod tests {
);
}
#[tokio::test]
async fn direct_execution_frame_stream_bridges_openai_image_sync_json_to_image_sse() {
let listener = crate::test_support::bind_loopback_listener()
.await
.expect("listener should bind");
let addr = listener.local_addr().expect("local addr should resolve");
let server = tokio::spawn(async move {
let app = Router::new().route(
"/responses",
post(|| async {
let body = serde_json::json!({
"created": 1776971267_u64,
"data": [{
"b64_json": "aGVsbG8="
}],
"usage": {
"total_tokens": 100,
"input_tokens": 50,
"output_tokens": 50,
"input_tokens_details": {
"text_tokens": 10,
"image_tokens": 40
}
}
});
let mut response = axum::http::Response::new(Body::from(
serde_json::to_vec(&body).expect("json should encode"),
));
response.headers_mut().insert(
header::CONTENT_TYPE,
HeaderValue::from_static("application/json"),
);
response
}),
);
axum::serve(listener, app)
.await
.expect("server should start");
});
let runtime = DirectSyncExecutionRuntime::new();
let execution = runtime
.execute_stream(&ExecutionPlan {
request_id: "req-image-sync-bridge".to_string(),
candidate_id: Some("cand-image-sync-bridge".to_string()),
provider_name: Some("OpenAI".to_string()),
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "key-1".to_string(),
method: "POST".to_string(),
url: format!("http://{addr}/responses"),
headers: BTreeMap::new(),
content_type: None,
content_encoding: None,
body: RequestBody::from_json(serde_json::json!({
"model": "gpt-image-1",
"prompt": "poster",
"stream": true
})),
stream: true,
client_api_format: "openai:image".to_string(),
provider_api_format: "openai:image".to_string(),
model_name: Some("gpt-image-1".into()),
proxy: None,
tls_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(5_000),
total_ms: Some(5_000),
..ExecutionTimeouts::default()
}),
})
.await
.expect("stream execution should succeed");
let frames = build_direct_execution_frame_stream(execution)
.map(|item| item.expect("frame should encode"))
.collect::<Vec<_>>()
.await
.into_iter()
.map(|bytes| String::from_utf8(bytes.to_vec()).expect("frame should be utf8"))
.collect::<Vec<_>>();
server.abort();
let header_frame: Value =
serde_json::from_str(&frames[0]).expect("headers frame should parse");
assert_eq!(
header_frame
.get("payload")
.and_then(|payload| payload.get("headers"))
.and_then(|headers| headers.get("content-type"))
.and_then(Value::as_str),
Some("text/event-stream")
);
let data_frame = frames
.iter()
.map(|line| serde_json::from_str::<Value>(line).expect("frame should parse"))
.find(|frame| frame.get("type").and_then(Value::as_str) == Some("data"))
.expect("data frame should exist");
let bridged_body = base64::engine::general_purpose::STANDARD
.decode(
data_frame
.get("payload")
.and_then(|payload| payload.get("chunk_b64"))
.and_then(Value::as_str)
.expect("chunk_b64 should exist"),
)
.expect("data frame should decode");
let bridged_text = String::from_utf8(bridged_body).expect("bridged body should be utf8");
assert!(bridged_text.contains("event: image_generation.completed"));
assert!(bridged_text.contains("\"type\":\"image_generation.completed\""));
assert!(bridged_text.contains("\"b64_json\":\"aGVsbG8=\""));
assert!(bridged_text.contains("\"total_tokens\":100"));
let eof_frame = frames
.iter()
.map(|line| serde_json::from_str::<Value>(line).expect("frame should parse"))
.find(|frame| frame.get("type").and_then(Value::as_str) == Some("eof"))
.expect("eof frame should exist");
assert_eq!(
eof_frame
.get("payload")
.and_then(|payload| payload.get("summary"))
.and_then(|summary| summary.get("model"))
.and_then(Value::as_str),
Some("gpt-image-1")
);
assert_eq!(
eof_frame
.get("payload")
.and_then(|payload| payload.get("summary"))
.and_then(|summary| summary.get("standardized_usage"))
.and_then(|usage| usage.get("dimensions"))
.and_then(|dimensions| dimensions.get("total_tokens"))
.and_then(Value::as_i64),
Some(100)
);
}
#[tokio::test]
async fn direct_execution_frame_stream_preserves_local_tunnel_stream_error_message() {
let state = AppState::new().expect("app state should build");