mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-08 10:27:46 +08:00
Merge PR #669: align GPT-5.6 and Codex request protocols
This commit is contained in:
elky
@@ -463,8 +463,8 @@ async fn gateway_executes_openai_chat_stream_via_local_openai_responses_cross_fo
|
||||
accept: String,
|
||||
authorization: String,
|
||||
x_client_request_id: String,
|
||||
session_id: String,
|
||||
conversation_id: String,
|
||||
codex_session_id: String,
|
||||
codex_thread_id: String,
|
||||
instructions: String,
|
||||
user_text: String,
|
||||
prompt_cache_key: String,
|
||||
@@ -742,15 +742,15 @@ async fn gateway_executes_openai_chat_stream_via_local_openai_responses_cross_fo
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
session_id: payload
|
||||
codex_session_id: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("session_id"))
|
||||
.and_then(|value| value.get("session-id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
conversation_id: payload
|
||||
codex_thread_id: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("conversation_id"))
|
||||
.and_then(|value| value.get("thread-id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
@@ -889,19 +889,16 @@ async fn gateway_executes_openai_chat_stream_via_local_openai_responses_cross_fo
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.x_client_request_id,
|
||||
"trace-openai-chat-cli-local-123"
|
||||
seen_execution_runtime_request.codex_thread_id
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.prompt_cache_key,
|
||||
"bc749eb7-a9e2-5793-8d14-abd659c700b0"
|
||||
seen_execution_runtime_request.codex_session_id,
|
||||
seen_execution_runtime_request.codex_thread_id
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.session_id,
|
||||
"d1e9b802644e1f52"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.conversation_id,
|
||||
"d1e9b802644e1f52"
|
||||
assert!(seen_execution_runtime_request.prompt_cache_key.is_empty());
|
||||
assert_ne!(
|
||||
seen_execution_runtime_request.codex_thread_id,
|
||||
seen_execution_runtime_request.trace_id
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.instructions,
|
||||
|
||||
@@ -2,7 +2,6 @@ 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_serving::CODEX_OPENAI_IMAGE_INTERNAL_MODEL;
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
@@ -55,13 +54,9 @@ async fn gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth
|
||||
struct SeenExecutionRuntimeStreamRequest {
|
||||
trace_id: String,
|
||||
url: String,
|
||||
model: String,
|
||||
authorization: String,
|
||||
x_client_request_id: String,
|
||||
tool_type: String,
|
||||
tool_action: String,
|
||||
tool_partial_images: Option<u64>,
|
||||
request_stream: bool,
|
||||
headers: serde_json::Value,
|
||||
body: serde_json::Value,
|
||||
plan_stream: bool,
|
||||
}
|
||||
|
||||
@@ -180,7 +175,7 @@ async fn gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
Some(serde_json::json!({"upstream_stream_policy":"force_stream"})),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
@@ -272,65 +267,26 @@ async fn gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth
|
||||
.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
|
||||
headers: payload.get("headers").cloned().unwrap_or_default(),
|
||||
body: 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_action: 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("action"))
|
||||
.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),
|
||||
.cloned()
|
||||
.unwrap_or_default(),
|
||||
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\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"application/json\"}}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"{\\\"created\\\":1776991097,\\\"data\\\":[{\\\"b64_json\\\":\\\"aGVsbG8=\\\",\\\"revised_prompt\\\":\\\"水墨视觉海报\\\"}],\\\"usage\\\":{\\\"input_tokens\\\":11,\\\"output_tokens\\\":22,\\\"total_tokens\\\":33}}\"}}\n",
|
||||
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41}}}\n",
|
||||
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
|
||||
);
|
||||
@@ -397,26 +353,28 @@ async fn gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-codex-image-stream-local-123")
|
||||
.body(
|
||||
"{\"model\":\"gpt-image-2\",\"prompt\":\"生成一张中国历史视觉海报\",\"stream\":true,\"partial_images\":1}",
|
||||
"{\"model\":\"gpt-image-2\",\"prompt\":\"生成一张水墨视觉海报\",\"background\":\"auto\",\"quality\":\"auto\",\"size\":\"auto\",\"stream\":true,\"response_format\":\"b64_json\"}",
|
||||
)
|
||||
.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_status = response.status();
|
||||
let response_content_type = response
|
||||
.headers()
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_string);
|
||||
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_eq!(response_status, StatusCode::OK, "{response_text}");
|
||||
assert_eq!(
|
||||
response_content_type.as_deref(),
|
||||
Some("text/event-stream"),
|
||||
"{response_text}"
|
||||
);
|
||||
assert!(response_text.contains("event: image_generation.completed"));
|
||||
assert!(response_text.contains("\"type\":\"image_generation.completed\""));
|
||||
assert!(response_text.contains("\"b64_json\":\"aGVsbG8=\""));
|
||||
assert!(response_text.contains("\"total_tokens\":33"));
|
||||
assert!(!response_text.contains("response.completed"));
|
||||
|
||||
@@ -444,25 +402,34 @@ async fn gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth
|
||||
);
|
||||
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
|
||||
"https://chatgpt.com/backend-api/codex/images/generations"
|
||||
);
|
||||
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"
|
||||
seen_execution_runtime_request.body,
|
||||
json!({
|
||||
"prompt": "生成一张水墨视觉海报",
|
||||
"background": "auto",
|
||||
"model": "gpt-image-2",
|
||||
"quality": "auto",
|
||||
"size": "auto"
|
||||
})
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.tool_type, "image_generation");
|
||||
assert_eq!(seen_execution_runtime_request.tool_action, "generate");
|
||||
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);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.headers["user-agent"],
|
||||
"codex_cli_rs/0.144.1"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.headers["originator"],
|
||||
"codex_cli_rs"
|
||||
);
|
||||
for header in ["x-client-request-id", "session-id", "thread-id"] {
|
||||
assert!(seen_execution_runtime_request.headers.get(header).is_none());
|
||||
}
|
||||
assert!(!seen_execution_runtime_request.plan_stream);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
@@ -594,7 +561,7 @@ async fn gateway_bridges_codex_image_sync_json_to_streaming_image_sse_impl() {
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
Some(serde_json::json!({"upstream_stream_policy":"force_stream"})),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
@@ -772,8 +739,8 @@ async fn gateway_bridges_codex_image_sync_json_to_streaming_image_sse_impl() {
|
||||
seen_execution_runtime_request.trace_id,
|
||||
"trace-codex-image-stream-json-123"
|
||||
);
|
||||
assert!(seen_execution_runtime_request.request_stream);
|
||||
assert!(seen_execution_runtime_request.plan_stream);
|
||||
assert!(!seen_execution_runtime_request.request_stream);
|
||||
assert!(!seen_execution_runtime_request.plan_stream);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
@@ -1152,12 +1119,14 @@ async fn gateway_routes_openai_responses_stream_image_intent_to_openai_image_pla
|
||||
seen_plan.url,
|
||||
"https://images.example.com/v1/images/generations"
|
||||
);
|
||||
assert!(seen_plan.plan_stream);
|
||||
assert!(!seen_plan.plan_stream);
|
||||
assert_eq!(seen_plan.auth_header, "Bearer sk-upstream-image-bridge");
|
||||
assert_eq!(seen_plan.body_json["stream"], true);
|
||||
assert_eq!(seen_plan.body_json["input"], "Draw a mountain observatory");
|
||||
assert_eq!(seen_plan.body_json["tools"][0]["type"], "image_generation");
|
||||
assert_eq!(seen_plan.body_json["tools"][0]["size"], "1024x1024");
|
||||
assert_eq!(seen_plan.body_json["model"], "gpt-image-2");
|
||||
assert_eq!(seen_plan.body_json["prompt"], "Draw a mountain observatory");
|
||||
assert_eq!(seen_plan.body_json["size"], "1024x1024");
|
||||
assert!(seen_plan.body_json.get("stream").is_none());
|
||||
assert!(seen_plan.body_json.get("input").is_none());
|
||||
assert!(seen_plan.body_json.get("tools").is_none());
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
|
||||
@@ -2,7 +2,7 @@ use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override, json,
|
||||
run_stream_cli_test, start_server, to_bytes, Arc, Body, Bytes, HeaderName, HeaderValue,
|
||||
Infallible, Json, Mutex, Request, Response, Router, StatusCode,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_HEADER, TRACE_ID_HEADER,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_SYNC, EXECUTION_PATH_HEADER, TRACE_ID_HEADER,
|
||||
};
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
@@ -23,33 +23,37 @@ use aether_data_contracts::repository::provider_catalog::{
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[test]
|
||||
fn gateway_executes_openai_responses_compact_stream_via_local_decision_gate_with_local_stream_decision(
|
||||
) {
|
||||
fn gateway_executes_openai_responses_compact_as_unary_request() {
|
||||
run_stream_cli_test(
|
||||
"gateway_executes_openai_responses_compact_stream_via_local_decision_gate_with_local_stream_decision",
|
||||
gateway_executes_openai_responses_compact_stream_via_local_decision_gate_with_local_stream_decision_impl,
|
||||
"gateway_executes_openai_responses_compact_as_unary_request",
|
||||
gateway_executes_openai_responses_compact_as_unary_request_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gate_with_local_stream_decision_impl(
|
||||
) {
|
||||
async fn gateway_executes_openai_responses_compact_as_unary_request_impl() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeStreamRequest {
|
||||
trace_id: String,
|
||||
url: String,
|
||||
model: String,
|
||||
content_encoding: String,
|
||||
stream: bool,
|
||||
accept: String,
|
||||
turn_state: String,
|
||||
authorization: String,
|
||||
chatgpt_account_id: String,
|
||||
fedramp: String,
|
||||
responses_lite: String,
|
||||
session_id: String,
|
||||
thread_id: String,
|
||||
x_client_request_id_present: bool,
|
||||
endpoint_tag: String,
|
||||
conditional_header: String,
|
||||
renamed_header: String,
|
||||
dropped_header_present: bool,
|
||||
metadata_mode: String,
|
||||
metadata_source: String,
|
||||
metadata_origin: String,
|
||||
instructions: String,
|
||||
store_present: bool,
|
||||
body: serde_json::Value,
|
||||
proxy_node_id: String,
|
||||
transport_profile_id: String,
|
||||
}
|
||||
@@ -71,7 +75,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat
|
||||
false,
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:responses:compact"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
Some(serde_json::json!(["gpt-5.6-sol"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
@@ -82,7 +86,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat
|
||||
Some(4_102_444_800_i64),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:responses:compact"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
Some(serde_json::json!(["gpt-5.6-sol"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
@@ -91,7 +95,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-openai-compact-local-1".to_string(),
|
||||
provider_name: "openai".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_type: "codex".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-openai-compact-local-1".to_string(),
|
||||
@@ -101,7 +105,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-openai-compact-local-1".to_string(),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: "bearer".to_string(),
|
||||
key_auth_type: "oauth".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["openai:responses:compact".to_string()]),
|
||||
key_allowed_models: None,
|
||||
@@ -110,12 +114,12 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat
|
||||
key_global_priority_by_format: Some(serde_json::json!({"openai:responses:compact": 1})),
|
||||
model_id: "model-openai-compact-local-1".to_string(),
|
||||
global_model_id: "global-model-openai-compact-local-1".to_string(),
|
||||
global_model_name: "gpt-5".to_string(),
|
||||
global_model_name: "gpt-5.6-sol".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "gpt-5-upstream".to_string(),
|
||||
model_provider_model_name: "deployment-production".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-5-upstream".to_string(),
|
||||
name: "deployment-production".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["openai:responses:compact".to_string()]),
|
||||
endpoint_ids: None,
|
||||
@@ -131,7 +135,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat
|
||||
"provider-openai-compact-local-1".to_string(),
|
||||
"openai".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"custom".to_string(),
|
||||
"codex".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
@@ -161,15 +165,12 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat
|
||||
"https://api.openai.example".to_string(),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","key":"x-endpoint-tag","value":"openai-compact-local"},
|
||||
{"action":"set","key":"x-conditional-tag","value":"header-condition-hit","condition":{"path":"instructions","op":"exists","source":"current"}},
|
||||
{"action":"set","key":"x-conditional-tag","value":"header-condition-hit","condition":{"path":"reasoning","op":"exists","source":"current"}},
|
||||
{"action":"rename","from":"x-client-rename","to":"x-upstream-rename"},
|
||||
{"action":"drop","key":"x-drop-me"}
|
||||
])),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","path":"instructions","value":"You are GPT-5.","condition":{"path":"instructions","op":"not_exists","source":"current"}},
|
||||
{"action":"set","path":"metadata.mode","value":"safe","condition":{"path":"metadata.mode","op":"not_exists","source":"current"}},
|
||||
{"action":"rename","from":"metadata.client","to":"metadata.source"},
|
||||
{"action":"set","path":"metadata.origin","value":"from-original","condition":{"path":"metadata.client","op":"exists","source":"original"}},
|
||||
{"action":"set","path":"instructions","value":"Use the configured tools.","condition":{"path":"instructions","op":"not_exists","source":"current"}},
|
||||
{"action":"drop","path":"store"}
|
||||
])),
|
||||
Some(2),
|
||||
@@ -186,7 +187,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat
|
||||
"key-openai-compact-local-1".to_string(),
|
||||
"provider-openai-compact-local-1".to_string(),
|
||||
"prod".to_string(),
|
||||
"bearer".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
@@ -198,7 +199,13 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat
|
||||
"sk-upstream-openai-compact",
|
||||
)
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
Some(
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"account_id":"acc-compact-local-123","is_fedramp":true}"#,
|
||||
)
|
||||
.expect("auth config should encrypt"),
|
||||
),
|
||||
None,
|
||||
Some(serde_json::json!({"openai:responses:compact": 1})),
|
||||
None,
|
||||
@@ -242,7 +249,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/decision-stream",
|
||||
"/api/internal/gateway/decision-sync",
|
||||
any(move |_request: Request| {
|
||||
let decision_hits_inner = Arc::clone(&decision_hits_clone);
|
||||
async move {
|
||||
@@ -252,7 +259,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/plan-stream",
|
||||
"/api/internal/gateway/plan-sync",
|
||||
any(move |_request: Request| {
|
||||
let plan_hits_inner = Arc::clone(&plan_hits_clone);
|
||||
async move {
|
||||
@@ -262,7 +269,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-stream",
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(move |request: Request| {
|
||||
let seen_report_inner = Arc::clone(&seen_report_clone);
|
||||
async move {
|
||||
@@ -299,135 +306,166 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/stream",
|
||||
"/v1/execute/sync",
|
||||
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(),
|
||||
stream: payload
|
||||
.get("stream")
|
||||
.and_then(|value| value.as_bool())
|
||||
.unwrap_or(false),
|
||||
accept: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("accept"))
|
||||
.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(),
|
||||
endpoint_tag: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-endpoint-tag"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
conditional_header: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-conditional-tag"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
renamed_header: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-upstream-rename"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
dropped_header_present: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-drop-me"))
|
||||
.is_some(),
|
||||
metadata_mode: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("metadata"))
|
||||
.and_then(|value| value.get("mode"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
metadata_source: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("metadata"))
|
||||
.and_then(|value| value.get("source"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
metadata_origin: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("metadata"))
|
||||
.and_then(|value| value.get("origin"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
instructions: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("instructions"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
store_present: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("store"))
|
||||
.is_some(),
|
||||
proxy_node_id: payload
|
||||
.get("proxy")
|
||||
.and_then(|value| value.get("node_id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
transport_profile_id: payload
|
||||
.get("transport_profile").and_then(|value| value.get("profile_id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
let stream = concat!(
|
||||
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp-compact-local-123\\\",\\\"object\\\":\\\"response\\\",\\\"model\\\":\\\"gpt-5-upstream\\\",\\\"output\\\":[]}}\\n\\n\"}}\n",
|
||||
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41,\"ttfb_ms\":11}}}\n",
|
||||
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
|
||||
);
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from(stream))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/x-ndjson"),
|
||||
);
|
||||
response
|
||||
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(),
|
||||
content_encoding: payload
|
||||
.get("content_encoding")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
stream: payload
|
||||
.get("stream")
|
||||
.and_then(|value| value.as_bool())
|
||||
.unwrap_or(false),
|
||||
accept: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("accept"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
turn_state: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-codex-turn-state"))
|
||||
.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(),
|
||||
chatgpt_account_id: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("chatgpt-account-id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
fedramp: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-openai-fedramp"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
responses_lite: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-openai-internal-codex-responses-lite"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
session_id: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("session-id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
thread_id: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("thread-id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
x_client_request_id_present: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-client-request-id"))
|
||||
.is_some(),
|
||||
endpoint_tag: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-endpoint-tag"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
conditional_header: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-conditional-tag"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
renamed_header: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-upstream-rename"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
dropped_header_present: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-drop-me"))
|
||||
.is_some(),
|
||||
instructions: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("instructions"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
store_present: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("store"))
|
||||
.is_some(),
|
||||
body: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.cloned()
|
||||
.unwrap_or(serde_json::Value::Null),
|
||||
proxy_node_id: payload
|
||||
.get("proxy")
|
||||
.and_then(|value| value.get("node_id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
transport_profile_id: payload
|
||||
.get("transport_profile")
|
||||
.and_then(|value| value.get("profile_id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
Json(json!({
|
||||
"request_id": "trace-openai-compact-local-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json",
|
||||
"x-codex-turn-state": "turn-state-compact-123"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"output": [{
|
||||
"type": "compaction",
|
||||
"id": "cmp-compact-local-123",
|
||||
"encrypted_content": "encrypted-compact-history"
|
||||
}]
|
||||
}
|
||||
},
|
||||
"telemetry": {"elapsed_ms": 41, "ttfb_ms": 11}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
@@ -474,23 +512,40 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat
|
||||
)
|
||||
.header("x-client-rename", "rename-openai-compact")
|
||||
.header("x-drop-me", "drop-openai-compact")
|
||||
.header("x-codex-turn-state", "turn-state-inbound-123")
|
||||
.header("session-id", "session-compact-local-123")
|
||||
.header("thread-id", "thread-compact-local-123")
|
||||
.header(TRACE_ID_HEADER, "trace-openai-compact-local-123")
|
||||
.body("{\"model\":\"gpt-5\",\"input\":\"hello\",\"stream\":true,\"metadata\":{\"client\":\"desktop-openai-compact\"},\"store\":false}")
|
||||
.body(r#"{"model":"gpt-5.6-sol","input":"hello","client_metadata":{"origin":"codex"},"include":["reasoning.encrypted_content"],"store":false,"stream":true,"stream_options":{"reasoning_summary_delivery":"sequential_cutoff"},"tool_choice":"auto","parallel_tool_calls":true,"reasoning":{"effort":"high"},"text":{"verbosity":"medium"},"tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}],"prompt_cache_key":"session:compact-e2e"}"#)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
if response.status() != StatusCode::OK {
|
||||
let status = response.status();
|
||||
let headers = response.headers().clone();
|
||||
let body = response.text().await.expect("error body should read");
|
||||
panic!("Compact request failed: status={status}, headers={headers:?}, body={body}");
|
||||
}
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get("x-codex-turn-state")
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("turn-state-compact-123")
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM)
|
||||
Some(EXECUTION_PATH_EXECUTION_RUNTIME_SYNC)
|
||||
);
|
||||
let body: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(
|
||||
body["output"][0]["encrypted_content"],
|
||||
"encrypted-compact-history"
|
||||
);
|
||||
let body = response.text().await.expect("body should read");
|
||||
assert!(body.contains("event: response.completed"));
|
||||
assert!(body.contains("\"model\":\"gpt-5-upstream\""));
|
||||
|
||||
let seen_execution_runtime_request = seen_execution_runtime
|
||||
.lock()
|
||||
@@ -505,13 +560,36 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat
|
||||
seen_execution_runtime_request.url,
|
||||
"https://api.openai.example/custom/v1/responses/compact"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.model, "gpt-5-upstream");
|
||||
assert!(seen_execution_runtime_request.stream);
|
||||
assert_eq!(seen_execution_runtime_request.accept, "text/event-stream");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.model,
|
||||
"deployment-production"
|
||||
);
|
||||
assert!(seen_execution_runtime_request.content_encoding.is_empty());
|
||||
assert!(!seen_execution_runtime_request.stream);
|
||||
assert_ne!(seen_execution_runtime_request.accept, "text/event-stream");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.turn_state,
|
||||
"turn-state-inbound-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.authorization,
|
||||
"Bearer sk-upstream-openai-compact"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.chatgpt_account_id,
|
||||
"acc-compact-local-123"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.fedramp, "true");
|
||||
assert_eq!(seen_execution_runtime_request.responses_lite, "true");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.session_id,
|
||||
"session-compact-local-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.thread_id,
|
||||
"thread-compact-local-123"
|
||||
);
|
||||
assert!(!seen_execution_runtime_request.x_client_request_id_present);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.endpoint_tag,
|
||||
"openai-compact-local"
|
||||
@@ -525,20 +603,61 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat
|
||||
"rename-openai-compact"
|
||||
);
|
||||
assert!(!seen_execution_runtime_request.dropped_header_present);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.instructions,
|
||||
"You are GPT-5."
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.metadata_mode, "safe");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.metadata_source,
|
||||
"desktop-openai-compact"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.metadata_origin,
|
||||
"from-original"
|
||||
);
|
||||
assert!(seen_execution_runtime_request.instructions.is_empty());
|
||||
assert!(!seen_execution_runtime_request.store_present);
|
||||
for field in [
|
||||
"client_metadata",
|
||||
"include",
|
||||
"store",
|
||||
"stream",
|
||||
"stream_options",
|
||||
"tool_choice",
|
||||
] {
|
||||
assert!(
|
||||
seen_execution_runtime_request.body.get(field).is_none(),
|
||||
"Compact request must omit {field}"
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.body["parallel_tool_calls"],
|
||||
json!(false)
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.body["reasoning"]["effort"],
|
||||
json!("high")
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.body["text"]["verbosity"],
|
||||
json!("medium")
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.body["reasoning"]["context"],
|
||||
json!("all_turns")
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.body["input"][0]["type"],
|
||||
json!("additional_tools")
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.body["input"][0]["tools"][0]["name"],
|
||||
json!("lookup")
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.body["input"][1]["role"],
|
||||
json!("developer")
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.body["input"][1]["content"][0]["text"],
|
||||
json!("Use the configured tools.")
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.body["input"][2]["content"][0]["text"],
|
||||
json!("hello")
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.body["prompt_cache_key"],
|
||||
json!("session:compact-e2e")
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.proxy_node_id,
|
||||
"proxy-node-openai-compact-local"
|
||||
@@ -558,7 +677,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert!(
|
||||
!*seen_report.lock().expect("mutex should lock"),
|
||||
"report-stream should stay local when request candidate persistence is available"
|
||||
"report-sync should stay local when request candidate persistence is available"
|
||||
);
|
||||
|
||||
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
@@ -38,10 +38,24 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
|
||||
trace_id: String,
|
||||
url: String,
|
||||
model: String,
|
||||
content_encoding: String,
|
||||
stream: bool,
|
||||
accept: String,
|
||||
authorization: String,
|
||||
chatgpt_account_id: String,
|
||||
fedramp: String,
|
||||
x_client_request_id: String,
|
||||
session_id: String,
|
||||
thread_id: String,
|
||||
prompt_cache_key: String,
|
||||
responses_lite: String,
|
||||
has_top_level_tools: bool,
|
||||
has_top_level_instructions: bool,
|
||||
has_additional_tools: bool,
|
||||
parallel_tool_calls: bool,
|
||||
reasoning_effort: String,
|
||||
reasoning_context: String,
|
||||
has_compaction_trigger: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -74,7 +88,7 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
|
||||
false,
|
||||
Some(serde_json::json!(["openai", "codex"])),
|
||||
Some(serde_json::json!(["openai:responses"])),
|
||||
Some(serde_json::json!(["gpt-5.4"])),
|
||||
Some(serde_json::json!(["gpt-5.6-sol"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
@@ -85,7 +99,7 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
|
||||
Some(4_102_444_800_i64),
|
||||
Some(serde_json::json!(["openai", "codex"])),
|
||||
Some(serde_json::json!(["openai:responses"])),
|
||||
Some(serde_json::json!(["gpt-5.4"])),
|
||||
Some(serde_json::json!(["gpt-5.6-sol"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
@@ -113,12 +127,12 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
|
||||
key_global_priority_by_format: Some(serde_json::json!({"openai:responses": 1})),
|
||||
model_id: "model-codex-cli-stream-local-1".to_string(),
|
||||
global_model_id: "global-model-codex-cli-stream-local-1".to_string(),
|
||||
global_model_name: "gpt-5.4".to_string(),
|
||||
global_model_name: "gpt-5.6-sol".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "gpt-5.4".to_string(),
|
||||
model_provider_model_name: "gpt-5.6-sol".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-5.4".to_string(),
|
||||
name: "gpt-5.6-sol".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["openai:responses".to_string()]),
|
||||
endpoint_ids: None,
|
||||
@@ -176,7 +190,7 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
|
||||
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-stream-local-123"}"#,
|
||||
r#"{"provider_type":"codex","refresh_token":"rt-codex-stream-local-123","account_id":"acc-codex-stream-local-123","is_fedramp":true}"#,
|
||||
)
|
||||
.expect("auth config should encrypt");
|
||||
StoredProviderCatalogKey::new(
|
||||
@@ -367,6 +381,11 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
content_encoding: payload
|
||||
.get("content_encoding")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
stream: payload
|
||||
.get("stream")
|
||||
.and_then(|value| value.as_bool())
|
||||
@@ -383,16 +402,106 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
chatgpt_account_id: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("chatgpt-account-id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
fedramp: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-openai-fedramp"))
|
||||
.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(),
|
||||
session_id: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("session-id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
thread_id: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("thread-id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
prompt_cache_key: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("prompt_cache_key"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
responses_lite: payload
|
||||
.get("headers")
|
||||
.and_then(|value| {
|
||||
value.get("x-openai-internal-codex-responses-lite")
|
||||
})
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
has_top_level_tools: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.is_some_and(|body| body.get("tools").is_some()),
|
||||
has_top_level_instructions: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.is_some_and(|body| body.get("instructions").is_some()),
|
||||
has_additional_tools: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("input"))
|
||||
.and_then(|value| value.as_array())
|
||||
.and_then(|input| input.first())
|
||||
.and_then(|item| item.get("type"))
|
||||
.and_then(|value| value.as_str())
|
||||
== Some("additional_tools"),
|
||||
parallel_tool_calls: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("parallel_tool_calls"))
|
||||
.and_then(|value| value.as_bool())
|
||||
.unwrap_or(true),
|
||||
reasoning_effort: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("reasoning"))
|
||||
.and_then(|value| value.get("effort"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
reasoning_context: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("reasoning"))
|
||||
.and_then(|value| value.get("context"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
has_compaction_trigger: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("input"))
|
||||
.and_then(|value| value.as_array())
|
||||
.is_some_and(|input| {
|
||||
input.iter().any(|item| {
|
||||
item.get("type").and_then(|value| value.as_str())
|
||||
== Some("compaction_trigger")
|
||||
})
|
||||
}),
|
||||
});
|
||||
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.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp_codex_cli_stream_local_123\\\",\\\"object\\\":\\\"response\\\",\\\"model\\\":\\\"gpt-5.4\\\",\\\"status\\\":\\\"completed\\\",\\\"usage\\\":{\\\"input_tokens\\\":1,\\\"output_tokens\\\":2,\\\"total_tokens\\\":3}}}\\n\\n\"}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.output_item.done\\ndata: {\\\"type\\\":\\\"response.output_item.done\\\",\\\"item\\\":{\\\"type\\\":\\\"compaction\\\",\\\"encrypted_content\\\":\\\"ENCRYPTED_CONTEXT_COMPACTION_SUMMARY\\\"}}\\n\\n\"}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp_codex_cli_stream_local_123\\\",\\\"object\\\":\\\"response\\\",\\\"model\\\":\\\"gpt-5.6-sol\\\",\\\"status\\\":\\\"completed\\\",\\\"usage\\\":{\\\"input_tokens\\\":1,\\\"output_tokens\\\":2,\\\"total_tokens\\\":3}}}\\n\\n\"}}\n",
|
||||
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41}}}\n",
|
||||
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
|
||||
);
|
||||
@@ -469,8 +578,16 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
|
||||
http::header::AUTHORIZATION,
|
||||
format!("Bearer {client_api_key}"),
|
||||
)
|
||||
.header("session-id", "session-codex-stream-local-123")
|
||||
.header("thread-id", "thread-codex-stream-local-123")
|
||||
.header(
|
||||
"x-client-request-id",
|
||||
"thread-codex-stream-local-123",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-codex-cli-stream-local-123")
|
||||
.body("{\"model\":\"gpt-5.4\",\"input\":\"hello\",\"stream\":true}")
|
||||
.body(
|
||||
r#"{"model":"gpt-5.6-sol","instructions":"Use the configured tools.","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"compact"}]},{"type":"compaction_trigger"}],"tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}],"parallel_tool_calls":true,"prompt_cache_key":"thread-codex-stream-local-123","client_metadata":{"session_id":"session-codex-stream-local-123","thread_id":"thread-codex-stream-local-123"},"stream":true}"#,
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
@@ -478,9 +595,20 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let response_body =
|
||||
strip_sse_keepalive_comments(&response.text().await.expect("body should read"));
|
||||
assert!(response_body.contains("event: response.output_item.done\n"));
|
||||
assert!(response_body.contains("\"type\":\"compaction\""));
|
||||
assert!(response_body.contains("ENCRYPTED_CONTEXT_COMPACTION_SUMMARY"));
|
||||
let data_line = response_body
|
||||
.lines()
|
||||
.find_map(|line| line.strip_prefix("data: "))
|
||||
.filter_map(|line| line.strip_prefix("data: "))
|
||||
.find(|line| {
|
||||
serde_json::from_str::<serde_json::Value>(line)
|
||||
.ok()
|
||||
.and_then(|event| event.get("type").cloned())
|
||||
.and_then(|value| value.as_str().map(ToOwned::to_owned))
|
||||
.as_deref()
|
||||
== Some("response.completed")
|
||||
})
|
||||
.expect("completed event data should exist");
|
||||
let completed_event: serde_json::Value =
|
||||
serde_json::from_str(data_line).expect("completed event should parse");
|
||||
@@ -495,7 +623,7 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
|
||||
"response": {
|
||||
"id": "resp_codex_cli_stream_local_123",
|
||||
"object": "response",
|
||||
"model": "gpt-5.4",
|
||||
"model": "gpt-5.6-sol",
|
||||
"status": "completed",
|
||||
"usage": {
|
||||
"input_tokens": 1,
|
||||
@@ -542,7 +670,8 @@ async fn gateway_executes_codex_cli_stream_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, "gpt-5.6-sol");
|
||||
assert_eq!(seen_execution_runtime_request.content_encoding, "zstd");
|
||||
assert!(seen_execution_runtime_request.stream);
|
||||
assert_eq!(seen_execution_runtime_request.accept, "text/event-stream");
|
||||
assert_eq!(
|
||||
@@ -550,9 +679,37 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
|
||||
"Bearer refreshed-codex-stream-access-token"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.x_client_request_id,
|
||||
"trace-codex-cli-stream-local-123"
|
||||
seen_execution_runtime_request.chatgpt_account_id,
|
||||
"acc-codex-stream-local-123"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.fedramp, "true");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.x_client_request_id,
|
||||
"thread-codex-stream-local-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.session_id,
|
||||
"session-codex-stream-local-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.thread_id,
|
||||
"thread-codex-stream-local-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.thread_id,
|
||||
seen_execution_runtime_request.prompt_cache_key
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.responses_lite, "true");
|
||||
assert!(!seen_execution_runtime_request.has_top_level_tools);
|
||||
assert!(!seen_execution_runtime_request.has_top_level_instructions);
|
||||
assert!(seen_execution_runtime_request.has_additional_tools);
|
||||
assert!(!seen_execution_runtime_request.parallel_tool_calls);
|
||||
assert_eq!(seen_execution_runtime_request.reasoning_effort, "low");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.reasoning_context,
|
||||
"all_turns"
|
||||
);
|
||||
assert!(seen_execution_runtime_request.has_compaction_trigger);
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-codex-cli-stream-local-123")
|
||||
|
||||
@@ -3223,6 +3223,9 @@ async fn gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_ref
|
||||
model: String,
|
||||
authorization: String,
|
||||
x_client_request_id: String,
|
||||
session_id: String,
|
||||
thread_id: String,
|
||||
prompt_cache_key: String,
|
||||
stream_present: bool,
|
||||
plan_stream: bool,
|
||||
}
|
||||
@@ -3509,6 +3512,25 @@ async fn gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_ref
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
session_id: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("session-id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
thread_id: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("thread-id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
prompt_cache_key: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("prompt_cache_key"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
stream_present: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
@@ -3646,7 +3668,16 @@ async fn gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_ref
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.x_client_request_id,
|
||||
"trace-codex-cli-local-123"
|
||||
seen_execution_runtime_request.thread_id
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.session_id,
|
||||
seen_execution_runtime_request.thread_id
|
||||
);
|
||||
assert!(seen_execution_runtime_request.prompt_cache_key.is_empty());
|
||||
assert_ne!(
|
||||
seen_execution_runtime_request.thread_id,
|
||||
seen_execution_runtime_request.trace_id
|
||||
);
|
||||
assert!(seen_execution_runtime_request.stream_present);
|
||||
assert!(seen_execution_runtime_request.plan_stream);
|
||||
|
||||
@@ -2,7 +2,6 @@ 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_serving::CODEX_OPENAI_IMAGE_INTERNAL_MODEL;
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
@@ -411,10 +410,10 @@ async fn gateway_converts_gemini_image_sync_to_openai_image_provider_impl() {
|
||||
url: String,
|
||||
authorization: String,
|
||||
model: String,
|
||||
action: String,
|
||||
prompt: String,
|
||||
image_url: String,
|
||||
request_stream: bool,
|
||||
body_stream: Option<bool>,
|
||||
}
|
||||
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
@@ -574,12 +573,6 @@ async fn gateway_converts_gemini_image_sync_to_openai_image_provider_impl() {
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.cloned()
|
||||
.unwrap_or_else(|| json!({}));
|
||||
let content = body_json
|
||||
.get("input")
|
||||
.and_then(|value| value.get(0))
|
||||
.and_then(|value| value.get("content"))
|
||||
.cloned()
|
||||
.unwrap_or_else(|| json!([]));
|
||||
*seen_execution_runtime_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(SeenExecutionRuntimeSyncRequest {
|
||||
@@ -605,39 +598,22 @@ async fn gateway_converts_gemini_image_sync_to_openai_image_provider_impl() {
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
action: body_json
|
||||
.get("tools")
|
||||
.and_then(|value| value.get(0))
|
||||
.and_then(|value| value.get("action"))
|
||||
prompt: body_json
|
||||
.get("prompt")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
prompt: content
|
||||
.as_array()
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.find(|item| {
|
||||
item.get("type").and_then(|value| value.as_str()) == Some("input_text")
|
||||
})
|
||||
.and_then(|item| item.get("text"))
|
||||
image_url: body_json
|
||||
.get("image")
|
||||
.and_then(|value| value.get("image_url"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
image_url: content
|
||||
.as_array()
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.find(|item| {
|
||||
item.get("type").and_then(|value| value.as_str()) == Some("input_image")
|
||||
})
|
||||
.and_then(|item| item.get("image_url"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
request_stream: body_json
|
||||
request_stream: payload
|
||||
.get("stream")
|
||||
.and_then(|value| value.as_bool())
|
||||
.unwrap_or(true),
|
||||
body_stream: body_json.get("stream").and_then(serde_json::Value::as_bool),
|
||||
});
|
||||
Json(json!({
|
||||
"request_id": "trace-gemini-image-to-openai-123",
|
||||
@@ -647,21 +623,16 @@ async fn gateway_converts_gemini_image_sync_to_openai_image_provider_impl() {
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"id": "resp_img_bridge_123",
|
||||
"object": "response",
|
||||
"created": 1776839946,
|
||||
"model": "gpt-image-2-upstream",
|
||||
"status": "completed",
|
||||
"usage": {
|
||||
"input_tokens": 3,
|
||||
"output_tokens": 4,
|
||||
"total_tokens": 7
|
||||
},
|
||||
"output": [{
|
||||
"type": "image_generation_call",
|
||||
"status": "completed",
|
||||
"output_format": "png",
|
||||
"data": [{
|
||||
"revised_prompt": "converted gemini prompt",
|
||||
"result": "aGVsbG8="
|
||||
"b64_json": "aGVsbG8="
|
||||
}]
|
||||
}
|
||||
},
|
||||
@@ -750,14 +721,13 @@ async fn gateway_converts_gemini_image_sync_to_openai_image_provider_impl() {
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.url,
|
||||
"https://api.openai.com/v1/images/generations"
|
||||
"https://api.openai.com/v1/images/edits"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.authorization,
|
||||
"Bearer sk-upstream-openai-image"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.model, "gpt-image-2-upstream");
|
||||
assert_eq!(seen_execution_runtime_request.action, "edit");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.prompt,
|
||||
"Change the background"
|
||||
@@ -767,6 +737,7 @@ async fn gateway_converts_gemini_image_sync_to_openai_image_provider_impl() {
|
||||
"data:image/png;base64,aGVsbG8="
|
||||
);
|
||||
assert!(!seen_execution_runtime_request.request_stream);
|
||||
assert_eq!(seen_execution_runtime_request.body_stream, None);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
@@ -785,18 +756,9 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: String,
|
||||
url: String,
|
||||
model: String,
|
||||
authorization: String,
|
||||
x_client_request_id: 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,
|
||||
headers: serde_json::Value,
|
||||
body: serde_json::Value,
|
||||
plan_stream: bool,
|
||||
}
|
||||
|
||||
@@ -1011,98 +973,18 @@ 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(),
|
||||
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(),
|
||||
prompt: payload
|
||||
headers: payload.get("headers").cloned().unwrap_or_default(),
|
||||
body: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("input"))
|
||||
.and_then(|value| value.get(0))
|
||||
.and_then(|value| value.get("content"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
content_is_string: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("input"))
|
||||
.and_then(|value| value.get(0))
|
||||
.and_then(|value| value.get("content"))
|
||||
.is_some_and(|value| value.is_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_size: 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("size"))
|
||||
.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"))
|
||||
.and_then(|value| value.get("tools"))
|
||||
.and_then(|value| value.get(0))
|
||||
.and_then(|value| value.as_object())
|
||||
.is_some_and(|object| object.contains_key("n")),
|
||||
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),
|
||||
.cloned()
|
||||
.unwrap_or_default(),
|
||||
plan_stream: payload
|
||||
.get("stream")
|
||||
.and_then(|value| value.as_bool())
|
||||
@@ -1112,20 +994,11 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r
|
||||
"request_id": "trace-codex-image-local-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "text/event-stream"
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"body_bytes_b64": base64::engine::general_purpose::STANDARD.encode(
|
||||
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\":\"__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,
|
||||
)
|
||||
r#"{"created":1776839946,"data":[{"b64_json":"aGVsbG8=","revised_prompt":"水墨视觉海报"}],"usage":{"input_tokens":171,"output_tokens":1372,"total_tokens":1543}}"#
|
||||
)
|
||||
},
|
||||
"telemetry": {
|
||||
@@ -1182,7 +1055,7 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r
|
||||
format!("Bearer {client_api_key}"),
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-codex-image-local-123")
|
||||
.body("{\"model\":\"gpt-image-2\",\"prompt\":\"生成一张中国历史视觉海报\",\"size\":\"1024x1024\",\"n\":1,\"response_format\":\"b64_json\"}")
|
||||
.body("{\"model\":\"gpt-image-2\",\"prompt\":\"生成一张水墨视觉海报\",\"background\":\"auto\",\"quality\":\"auto\",\"size\":\"auto\",\"n\":1,\"response_format\":\"b64_json\"}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
@@ -1191,10 +1064,7 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r
|
||||
let response_json: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(response_json["created"], 1776839946);
|
||||
assert_eq!(response_json["data"][0]["b64_json"], "aGVsbG8=");
|
||||
assert_eq!(
|
||||
response_json["data"][0]["revised_prompt"],
|
||||
"中国历史视觉海报"
|
||||
);
|
||||
assert_eq!(response_json["data"][0]["revised_prompt"], "水墨视觉海报");
|
||||
assert_eq!(response_json["usage"]["input_tokens"], 171);
|
||||
assert_eq!(response_json["usage"]["output_tokens"], 1372);
|
||||
|
||||
@@ -1229,35 +1099,34 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r
|
||||
);
|
||||
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
|
||||
"https://chatgpt.com/backend-api/codex/images/generations"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.authorization,
|
||||
"Bearer refreshed-codex-image-access-token"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.x_client_request_id,
|
||||
"trace-codex-image-local-123"
|
||||
seen_execution_runtime_request.body,
|
||||
json!({
|
||||
"prompt": "生成一张水墨视觉海报",
|
||||
"background": "auto",
|
||||
"model": "gpt-image-2",
|
||||
"n": 1,
|
||||
"quality": "auto",
|
||||
"size": "auto"
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.prompt,
|
||||
"生成一张中国历史视觉海报"
|
||||
seen_execution_runtime_request.headers["user-agent"],
|
||||
"codex_cli_rs/0.144.1"
|
||||
);
|
||||
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"
|
||||
seen_execution_runtime_request.headers["originator"],
|
||||
"codex_cli_rs"
|
||||
);
|
||||
assert!(!seen_execution_runtime_request.tool_has_n);
|
||||
assert!(seen_execution_runtime_request.request_stream);
|
||||
for header in ["x-client-request-id", "session-id", "thread-id"] {
|
||||
assert!(seen_execution_runtime_request.headers.get(header).is_none());
|
||||
}
|
||||
assert!(!seen_execution_runtime_request.plan_stream);
|
||||
|
||||
let persisted_transport_state =
|
||||
|
||||
@@ -29,3 +29,4 @@ mod cli;
|
||||
mod gemini;
|
||||
mod image;
|
||||
mod pii_redaction_formats;
|
||||
mod search;
|
||||
|
||||
@@ -0,0 +1,597 @@
|
||||
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,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_SYNC, EXECUTION_PATH_HEADER, TRACE_ID_HEADER,
|
||||
};
|
||||
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::candidates::{
|
||||
RequestCandidateReadRepository, RequestCandidateStatus,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
const SEARCH_SYNC_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
|
||||
|
||||
fn run_search_sync_test<F, Fut>(test_name: &'static str, make_future: F)
|
||||
where
|
||||
F: FnOnce() -> Fut + Send + 'static,
|
||||
Fut: std::future::Future<Output = ()> + 'static,
|
||||
{
|
||||
let handle = std::thread::Builder::new()
|
||||
.name(test_name.to_string())
|
||||
.stack_size(SEARCH_SYNC_TEST_STACK_BYTES)
|
||||
.spawn(move || {
|
||||
let runtime = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("test runtime should build");
|
||||
runtime.block_on(make_future());
|
||||
})
|
||||
.expect("search sync test thread should spawn");
|
||||
|
||||
if let Err(payload) = handle.join() {
|
||||
std::panic::resume_unwind(payload);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_executes_codex_search_with_responses_permission_and_search_contract() {
|
||||
run_search_sync_test(
|
||||
"gateway_executes_codex_search_with_responses_permission_and_search_contract",
|
||||
gateway_executes_codex_search_with_responses_permission_and_search_contract_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_executes_codex_search_with_responses_permission_and_search_contract_impl() {
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn auth_snapshot() -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
"user-search-1".to_string(),
|
||||
"alice".to_string(),
|
||||
Some("[email protected]".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(json!(["openai", "codex"])),
|
||||
Some(json!(["openai:responses"])),
|
||||
None,
|
||||
"api-key-search-1".to_string(),
|
||||
Some("search-client".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800_i64),
|
||||
Some(json!(["openai", "codex"])),
|
||||
Some(json!(["openai:responses"])),
|
||||
None,
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-codex-search-1".to_string(),
|
||||
provider_name: "codex".to_string(),
|
||||
provider_type: "codex".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-codex-search-1".to_string(),
|
||||
endpoint_api_format: "openai:search".to_string(),
|
||||
endpoint_api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("search".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-codex-search-1".to_string(),
|
||||
key_name: "oauth".to_string(),
|
||||
key_auth_type: "oauth".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["openai:responses".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(json!({"openai:search": 1})),
|
||||
model_id: "model-codex-search-1".to_string(),
|
||||
global_model_id: "global-model-codex-search-1".to_string(),
|
||||
global_model_name: "gpt-5.6-sol".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(false),
|
||||
model_provider_model_name: "gpt-5.6-sol".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-5.6-sol".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["openai:responses".to_string()]),
|
||||
endpoint_ids: None,
|
||||
}]),
|
||||
model_supports_streaming: Some(false),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-codex-search-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(900.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-codex-search-1".to_string(),
|
||||
"provider-codex-search-1".to_string(),
|
||||
"openai:search".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("search".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://chatgpt.com/backend-api/codex".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn key() -> StoredProviderCatalogKey {
|
||||
let auth_config = encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"provider_type":"codex","account_id":"account-search-1","is_fedramp":true}"#,
|
||||
)
|
||||
.expect("auth config should encrypt");
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-codex-search-1".to_string(),
|
||||
"provider-codex-search-1".to_string(),
|
||||
"oauth".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(json!(["openai:responses"])),
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
"codex-search-access-token",
|
||||
)
|
||||
.expect("access token should encrypt"),
|
||||
Some(auth_config),
|
||||
None,
|
||||
Some(json!({"openai:search": 1})),
|
||||
None,
|
||||
Some(4_102_444_800),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let seen_plans = Arc::new(Mutex::new(Vec::<serde_json::Value>::new()));
|
||||
let seen_plans_clone = Arc::clone(&seen_plans);
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |request: Request| {
|
||||
let seen_plans_inner = Arc::clone(&seen_plans_clone);
|
||||
async move {
|
||||
let (_, body) = request.into_parts();
|
||||
let bytes = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value =
|
||||
serde_json::from_slice(&bytes).expect("execution payload should parse");
|
||||
let request_id = payload["request_id"]
|
||||
.as_str()
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
let provider_id = payload["provider_id"]
|
||||
.as_str()
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
seen_plans_inner
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.push(payload);
|
||||
let execution_result = if request_id == "trace-search-error-1" {
|
||||
json!({
|
||||
"request_id": request_id,
|
||||
"status_code": 400,
|
||||
"headers": {
|
||||
"content-type": "application/json",
|
||||
"x-search-upstream": "rate-limited"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"error": {
|
||||
"type": "rate_limit_error",
|
||||
"message": "Search capacity reached",
|
||||
"param": null,
|
||||
"code": "rate_limit_exceeded"
|
||||
},
|
||||
"future_error_field": {"retryable": true}
|
||||
}
|
||||
},
|
||||
"telemetry": {"elapsed_ms": 17}
|
||||
})
|
||||
} else if request_id == "trace-search-failover-1"
|
||||
&& provider_id == "provider-codex-search-1"
|
||||
{
|
||||
json!({
|
||||
"request_id": request_id,
|
||||
"status_code": 500,
|
||||
"headers": {
|
||||
"content-type": "application/json",
|
||||
"x-search-upstream": "primary"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"error": {
|
||||
"type": "server_error",
|
||||
"message": "Search backend unavailable"
|
||||
}
|
||||
}
|
||||
},
|
||||
"telemetry": {"elapsed_ms": 11}
|
||||
})
|
||||
} else if request_id == "trace-search-failover-1" {
|
||||
json!({
|
||||
"request_id": request_id,
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json",
|
||||
"x-search-upstream": "backup"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"output": "search fallback result"
|
||||
}
|
||||
},
|
||||
"telemetry": {"elapsed_ms": 23}
|
||||
})
|
||||
} else {
|
||||
json!({
|
||||
"request_id": request_id,
|
||||
"status_code": 201,
|
||||
"headers": {
|
||||
"content-type": "application/json",
|
||||
"x-search-upstream": "alpha"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"output": "search result",
|
||||
"encrypted_output": "encrypted-search-result",
|
||||
"future_response_field": {"enabled": true}
|
||||
}
|
||||
},
|
||||
"telemetry": {"elapsed_ms": 42}
|
||||
})
|
||||
};
|
||||
(
|
||||
StatusCode::OK,
|
||||
[("x-search-source", "codex-alpha")],
|
||||
Json(execution_result),
|
||||
)
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let client_api_key = "sk-client-search";
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key(client_api_key)),
|
||||
auth_snapshot(),
|
||||
)]));
|
||||
let candidate_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed({
|
||||
let primary = candidate_row();
|
||||
let mut backup = primary.clone();
|
||||
backup.provider_id = "provider-codex-search-2".to_string();
|
||||
backup.provider_name = "codex-backup".to_string();
|
||||
backup.provider_priority = 20;
|
||||
backup.endpoint_id = "endpoint-codex-search-2".to_string();
|
||||
backup.key_id = "key-codex-search-2".to_string();
|
||||
backup.key_name = "oauth-backup".to_string();
|
||||
backup.key_internal_priority = 6;
|
||||
backup.key_global_priority_by_format = Some(json!({"openai:search": 2}));
|
||||
backup.model_id = "model-codex-search-2".to_string();
|
||||
vec![primary, backup]
|
||||
}));
|
||||
let catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
{
|
||||
let primary = provider();
|
||||
let mut backup = primary.clone();
|
||||
backup.id = "provider-codex-search-2".to_string();
|
||||
backup.name = "codex-backup".to_string();
|
||||
vec![primary, backup]
|
||||
},
|
||||
{
|
||||
let primary = endpoint();
|
||||
let mut backup = primary.clone();
|
||||
backup.id = "endpoint-codex-search-2".to_string();
|
||||
backup.provider_id = "provider-codex-search-2".to_string();
|
||||
vec![primary, backup]
|
||||
},
|
||||
{
|
||||
let primary = key();
|
||||
let mut backup = primary.clone();
|
||||
backup.id = "key-codex-search-2".to_string();
|
||||
backup.provider_id = "provider-codex-search-2".to_string();
|
||||
backup.name = "oauth-backup".to_string();
|
||||
backup.global_priority_by_format = Some(json!({"openai:search": 2}));
|
||||
vec![primary, backup]
|
||||
},
|
||||
));
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let data_state =
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_repository,
|
||||
catalog_repository,
|
||||
Arc::clone(&request_candidates),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
)
|
||||
.with_system_config_values_for_tests([(
|
||||
crate::system_features::ENABLE_MODEL_DIRECTIVES_CONFIG_KEY.to_string(),
|
||||
json!(true),
|
||||
)]);
|
||||
let state = build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(data_state);
|
||||
let gateway = build_router_with_state(state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/alpha/search"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
format!("Bearer {client_api_key}"),
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-search-1")
|
||||
.json(&json!({
|
||||
"id": "session-search-1",
|
||||
"model": "gpt-5.6-sol-ultra-fast",
|
||||
"reasoning": {"effort": "low", "summary": "auto"},
|
||||
"input": "find current OpenAI documentation",
|
||||
"commands": {
|
||||
"search_query": [{"q": "OpenAI Codex search"}],
|
||||
"open": [{"ref_id": "turn0search0"}]
|
||||
},
|
||||
"settings": {
|
||||
"search_context_size": "high",
|
||||
"allowed_callers": ["direct"]
|
||||
},
|
||||
"max_output_tokens": 4096,
|
||||
"store": false,
|
||||
"stream": true,
|
||||
"future_request_field": {"enabled": true}
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("search request should succeed");
|
||||
|
||||
if response.status() != StatusCode::CREATED {
|
||||
let status = response.status();
|
||||
let body = response.text().await.expect("error response should read");
|
||||
panic!("Search request returned {status}: {body}");
|
||||
}
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get("x-search-upstream")
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("alpha")
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_EXECUTION_RUNTIME_SYNC)
|
||||
);
|
||||
let response_json: serde_json::Value = response.json().await.expect("response should parse");
|
||||
assert_eq!(response_json["output"], "search result");
|
||||
assert_eq!(response_json["encrypted_output"], "encrypted-search-result");
|
||||
assert_eq!(response_json["future_response_field"]["enabled"], true);
|
||||
|
||||
let plan = seen_plans
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.first()
|
||||
.cloned()
|
||||
.expect("execution plan should be captured");
|
||||
assert_eq!(
|
||||
plan["url"],
|
||||
"https://chatgpt.com/backend-api/codex/alpha/search"
|
||||
);
|
||||
assert_eq!(plan["client_api_format"], "openai:search");
|
||||
assert_eq!(plan["provider_api_format"], "openai:search");
|
||||
assert_eq!(plan["stream"], false);
|
||||
assert_eq!(plan["timeouts"]["total_ms"], 900_000);
|
||||
assert_eq!(
|
||||
plan["headers"]["authorization"],
|
||||
"Bearer codex-search-access-token"
|
||||
);
|
||||
assert_eq!(plan["headers"]["chatgpt-account-id"], "account-search-1");
|
||||
assert_eq!(plan["headers"]["x-openai-fedramp"], "true");
|
||||
assert_eq!(plan["headers"]["originator"], "codex_cli_rs");
|
||||
assert!(plan["headers"]["user-agent"]
|
||||
.as_str()
|
||||
.is_some_and(|value| value.starts_with("codex_cli_rs/")));
|
||||
assert!(plan["headers"].get("openai-beta").is_none());
|
||||
assert!(plan["headers"]
|
||||
.get("x-openai-internal-codex-responses-lite")
|
||||
.is_none());
|
||||
assert_ne!(plan["headers"]["accept"], "text/event-stream");
|
||||
|
||||
let body = &plan["body"]["json_body"];
|
||||
assert_eq!(body["id"], "session-search-1");
|
||||
assert_eq!(body["model"], "gpt-5.6-sol");
|
||||
assert_eq!(body["reasoning"]["effort"], "max");
|
||||
assert_eq!(body["reasoning"]["summary"], "auto");
|
||||
assert_eq!(
|
||||
body["commands"]["search_query"][0]["q"],
|
||||
"OpenAI Codex search"
|
||||
);
|
||||
assert_eq!(body["commands"]["open"][0]["ref_id"], "turn0search0");
|
||||
assert_eq!(body["settings"]["search_context_size"], "high");
|
||||
assert_eq!(body["max_output_tokens"], 4096);
|
||||
assert!(body.get("store").is_none());
|
||||
assert!(body.get("future_request_field").is_none());
|
||||
assert!(body.get("stream").is_none());
|
||||
assert!(body.get("service_tier").is_none());
|
||||
|
||||
let candidates = request_candidates
|
||||
.list_by_request_id("trace-search-1")
|
||||
.await
|
||||
.expect("request candidates should read");
|
||||
assert_eq!(candidates.len(), 1);
|
||||
assert_eq!(candidates[0].status, RequestCandidateStatus::Success);
|
||||
|
||||
let expected_error_body = json!({
|
||||
"error": {
|
||||
"type": "rate_limit_error",
|
||||
"message": "Search capacity reached",
|
||||
"param": null,
|
||||
"code": "rate_limit_exceeded"
|
||||
},
|
||||
"future_error_field": {"retryable": true}
|
||||
});
|
||||
let error_response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/alpha/search"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
format!("Bearer {client_api_key}"),
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-search-error-1")
|
||||
.json(&json!({
|
||||
"id": "session-search-error-1",
|
||||
"model": "gpt-5.6-sol",
|
||||
"input": "find current OpenAI documentation",
|
||||
"commands": {"search_query": [{"q": "OpenAI documentation"}]}
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("search error response should return");
|
||||
|
||||
assert_eq!(error_response.status(), StatusCode::BAD_REQUEST);
|
||||
assert_eq!(
|
||||
error_response
|
||||
.headers()
|
||||
.get("x-search-upstream")
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("rate-limited")
|
||||
);
|
||||
assert_eq!(
|
||||
error_response
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("error response should parse"),
|
||||
expected_error_body
|
||||
);
|
||||
let error_candidates = request_candidates
|
||||
.list_by_request_id("trace-search-error-1")
|
||||
.await
|
||||
.expect("error request candidates should read");
|
||||
assert_eq!(error_candidates.len(), 1);
|
||||
assert_eq!(error_candidates[0].status, RequestCandidateStatus::Failed);
|
||||
|
||||
let failover_response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/alpha/search"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
format!("Bearer {client_api_key}"),
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-search-failover-1")
|
||||
.json(&json!({
|
||||
"id": "session-search-failover-1",
|
||||
"model": "gpt-5.6-sol",
|
||||
"input": "find current OpenAI documentation",
|
||||
"commands": {"search_query": [{"q": "OpenAI documentation"}]}
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("search failover response should return");
|
||||
|
||||
assert_eq!(failover_response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
failover_response
|
||||
.headers()
|
||||
.get("x-search-upstream")
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("backup")
|
||||
);
|
||||
assert_eq!(
|
||||
failover_response
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("failover response should parse")["output"],
|
||||
"search fallback result"
|
||||
);
|
||||
let failover_plans = seen_plans
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.iter()
|
||||
.filter(|plan| plan["request_id"] == "trace-search-failover-1")
|
||||
.map(|plan| plan["provider_id"].clone())
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
failover_plans,
|
||||
vec![
|
||||
json!("provider-codex-search-1"),
|
||||
json!("provider-codex-search-2")
|
||||
]
|
||||
);
|
||||
let failover_candidates = request_candidates
|
||||
.list_by_request_id("trace-search-failover-1")
|
||||
.await
|
||||
.expect("failover request candidates should read");
|
||||
assert_eq!(failover_candidates.len(), 2);
|
||||
assert_eq!(
|
||||
failover_candidates[0].status,
|
||||
RequestCandidateStatus::Failed
|
||||
);
|
||||
assert_eq!(failover_candidates[0].status_code, Some(500));
|
||||
assert_eq!(
|
||||
failover_candidates[1].status,
|
||||
RequestCandidateStatus::Success
|
||||
);
|
||||
assert_eq!(failover_candidates[1].status_code, Some(200));
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
}
|
||||
@@ -66,6 +66,7 @@ fn sample_decision() -> crate::control::GatewayControlDecision {
|
||||
auth_context: None,
|
||||
admin_principal: None,
|
||||
local_auth_rejection: None,
|
||||
model_directive_policy: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -559,6 +559,103 @@ async fn gateway_creates_admin_provider_endpoint_locally_with_trusted_admin_prin
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_streaming_policy_for_search_endpoint_before_catalog_write() {
|
||||
let mut create_provider = sample_provider("provider-search-create", "search-create", 10);
|
||||
create_provider.provider_type = "custom".to_string();
|
||||
let mut update_provider = sample_provider("provider-search-update", "search-update", 20);
|
||||
update_provider.provider_type = "custom".to_string();
|
||||
let mut existing_endpoint = sample_endpoint(
|
||||
"endpoint-search-update",
|
||||
"provider-search-update",
|
||||
"openai:search",
|
||||
"https://search.example/v1",
|
||||
);
|
||||
existing_endpoint.config = Some(json!({"marker": "kept"}));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![create_provider, update_provider],
|
||||
vec![existing_endpoint],
|
||||
vec![],
|
||||
));
|
||||
|
||||
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 client = reqwest::Client::new();
|
||||
|
||||
let create_response = client
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-search-create/endpoints"
|
||||
))
|
||||
.header(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!({
|
||||
"provider_id": "provider-search-create",
|
||||
"api_format": "openai:search",
|
||||
"base_url": "https://search.example/v1",
|
||||
"config": {"upstream_stream_policy": "force_stream"}
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(create_response.status(), StatusCode::BAD_REQUEST);
|
||||
let create_payload: serde_json::Value = create_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
assert_eq!(
|
||||
create_payload["detail"],
|
||||
"OpenAI Search 端点仅支持非流式上游请求"
|
||||
);
|
||||
|
||||
let update_response = client
|
||||
.put(format!(
|
||||
"{gateway_url}/api/admin/endpoints/endpoint-search-update"
|
||||
))
|
||||
.header(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!({
|
||||
"config": {"upstreamStreamPolicy": true}
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(update_response.status(), StatusCode::BAD_REQUEST);
|
||||
let update_payload: serde_json::Value = update_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
assert_eq!(
|
||||
update_payload["detail"],
|
||||
"OpenAI Search 端点仅支持非流式上游请求"
|
||||
);
|
||||
|
||||
let created = provider_catalog_repository
|
||||
.list_endpoints_by_provider_ids(&["provider-search-create".to_string()])
|
||||
.await
|
||||
.expect("endpoints should read");
|
||||
assert!(created.is_empty());
|
||||
let unchanged = provider_catalog_repository
|
||||
.list_endpoints_by_ids(&["endpoint-search-update".to_string()])
|
||||
.await
|
||||
.expect("endpoint should read");
|
||||
assert_eq!(unchanged.len(), 1);
|
||||
assert_eq!(unchanged[0].config, Some(json!({"marker": "kept"})));
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_updates_admin_provider_endpoint_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
|
||||
@@ -6086,6 +6086,9 @@ async fn gateway_manual_codex_oauth_refresh_reconciles_missing_fixed_endpoint_im
|
||||
.iter()
|
||||
.find(|endpoint| endpoint.api_format == "openai:responses")
|
||||
.expect("openai responses endpoint should be reconciled");
|
||||
assert!(endpoints
|
||||
.iter()
|
||||
.any(|endpoint| endpoint.api_format == "openai:search"));
|
||||
assert_eq!(
|
||||
responses_endpoint.base_url,
|
||||
"https://chatgpt.com/backend-api/codex"
|
||||
|
||||
@@ -559,7 +559,7 @@ async fn gateway_handles_admin_provider_query_models_falls_back_to_codex_preset_
|
||||
.expect("mutex should lock") += 1;
|
||||
assert_eq!(
|
||||
plan.url,
|
||||
"https://chatgpt.com/backend-api/codex/models?client_version=0.128.0-alpha.1"
|
||||
"https://chatgpt.com/backend-api/codex/models?client_version=0.144.1"
|
||||
);
|
||||
Json(json!({
|
||||
"request_id": "req-provider-query-codex-invalidated",
|
||||
@@ -625,6 +625,11 @@ async fn gateway_handles_admin_provider_query_models_falls_back_to_codex_preset_
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["success"], json!(true));
|
||||
assert_eq!(payload["data"]["error"], serde_json::Value::Null);
|
||||
let warning = payload["data"]["warning"]
|
||||
.as_str()
|
||||
.expect("Codex fallback warning should be present");
|
||||
assert!(warning.contains("Codex 动态模型目录不可用"));
|
||||
assert!(warning.contains("invalidated"));
|
||||
let model_ids = payload["data"]["models"]
|
||||
.as_array()
|
||||
.expect("models should be an array")
|
||||
@@ -634,11 +639,14 @@ async fn gateway_handles_admin_provider_query_models_falls_back_to_codex_preset_
|
||||
assert_eq!(
|
||||
model_ids,
|
||||
vec![
|
||||
"gpt-5.3-codex",
|
||||
"gpt-5.3-codex-spark",
|
||||
"codex-auto-review",
|
||||
"gpt-5.2",
|
||||
"gpt-5.4",
|
||||
"gpt-5.4-mini",
|
||||
"gpt-5.5",
|
||||
"gpt-5.6-luna",
|
||||
"gpt-5.6-sol",
|
||||
"gpt-5.6-terra",
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
@@ -2391,6 +2399,188 @@ async fn gateway_streams_codex_openai_responses_upstream_for_admin_pool_model_te
|
||||
execution_runtime_handle.abort();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_executes_codex_search_admin_pool_model_test_with_search_contract() {
|
||||
run_provider_query_test(
|
||||
"gateway_executes_codex_search_admin_pool_model_test_with_search_contract",
|
||||
gateway_executes_codex_search_admin_pool_model_test_with_search_contract_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_executes_codex_search_admin_pool_model_test_with_search_contract_impl() {
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |Json(plan): Json<ExecutionPlan>| async move {
|
||||
assert_eq!(plan.provider_id, "provider-codex-search");
|
||||
assert_eq!(plan.endpoint_id, "endpoint-codex-search");
|
||||
assert_eq!(plan.key_id, "key-codex-search");
|
||||
assert_eq!(plan.client_api_format, "openai:search");
|
||||
assert_eq!(plan.provider_api_format, "openai:search");
|
||||
assert_eq!(
|
||||
plan.url,
|
||||
"https://chatgpt.com/backend-api/codex/alpha/search"
|
||||
);
|
||||
assert_eq!(plan.model_name.as_deref(), Some("gpt-5.6-sol"));
|
||||
assert!(!plan.stream, "Codex Search is a synchronous JSON protocol");
|
||||
assert_eq!(
|
||||
plan.timeouts
|
||||
.as_ref()
|
||||
.and_then(|timeouts| timeouts.total_ms),
|
||||
Some(900_000)
|
||||
);
|
||||
assert_eq!(
|
||||
plan.headers.get("authorization").map(String::as_str),
|
||||
Some("Bearer codex-search-access-token")
|
||||
);
|
||||
assert_eq!(
|
||||
plan.headers.get("chatgpt-account-id").map(String::as_str),
|
||||
Some("account-search-admin")
|
||||
);
|
||||
assert_eq!(
|
||||
plan.headers.get("x-openai-fedramp").map(String::as_str),
|
||||
Some("true")
|
||||
);
|
||||
assert_eq!(
|
||||
plan.headers.get("originator").map(String::as_str),
|
||||
Some("codex_cli_rs")
|
||||
);
|
||||
assert!(plan
|
||||
.headers
|
||||
.get("user-agent")
|
||||
.is_some_and(|value| value.starts_with("codex_cli_rs/")));
|
||||
assert!(!plan.headers.contains_key("openai-beta"));
|
||||
assert!(!plan
|
||||
.headers
|
||||
.contains_key("x-openai-internal-codex-responses-lite"));
|
||||
assert_ne!(
|
||||
plan.headers.get("accept").map(String::as_str),
|
||||
Some("text/event-stream")
|
||||
);
|
||||
|
||||
let body = plan.body.json_body.as_ref().expect("search json body");
|
||||
assert_eq!(
|
||||
body["id"],
|
||||
json!("aether-model-test-provider-query-search-trace")
|
||||
);
|
||||
assert_eq!(body["model"], json!("gpt-5.6-sol"));
|
||||
assert_eq!(body["input"], json!("find current OpenAI documentation"));
|
||||
assert_eq!(
|
||||
body["commands"]["search_query"][0]["q"],
|
||||
json!("OpenAI Codex Search")
|
||||
);
|
||||
assert!(body.get("stream").is_none());
|
||||
assert!(body.get("store").is_none());
|
||||
assert!(body.get("service_tier").is_none());
|
||||
assert!(body.get("unknown_field").is_none());
|
||||
|
||||
Json(json!({
|
||||
"request_id": plan.request_id,
|
||||
"candidate_id": plan.candidate_id,
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"output": "search result"
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 21
|
||||
}
|
||||
}))
|
||||
}),
|
||||
);
|
||||
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let mut provider = sample_provider("provider-codex-search", "Codex Search", 10);
|
||||
provider.provider_type = "codex".to_string();
|
||||
provider.request_timeout_secs = Some(900.0);
|
||||
let mut endpoint = sample_endpoint(
|
||||
"endpoint-codex-search",
|
||||
"provider-codex-search",
|
||||
"openai:search",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
);
|
||||
endpoint.config = Some(json!({"upstream_stream_policy": "force_stream"}));
|
||||
let mut key = sample_key(
|
||||
"key-codex-search",
|
||||
"provider-codex-search",
|
||||
"openai:search",
|
||||
"codex-search-access-token",
|
||||
);
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.encrypted_auth_config = Some(
|
||||
aether_crypto::encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"provider_type":"codex","account_id":"account-search-admin","is_fedramp":true}"#,
|
||||
)
|
||||
.expect("auth config should encrypt"),
|
||||
);
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
vec![endpoint],
|
||||
vec![key],
|
||||
));
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/provider-query/test-model-failover"
|
||||
))
|
||||
.header(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!({
|
||||
"provider_id": "provider-codex-search",
|
||||
"mode": "pool",
|
||||
"model": "gpt-5.6-sol",
|
||||
"failover_models": ["gpt-5.6-sol"],
|
||||
"api_format": "openai:search",
|
||||
"endpoint_id": "endpoint-codex-search",
|
||||
"request_id": "provider-query-search-trace",
|
||||
"request_body": {
|
||||
"model": "gpt-5.6-sol",
|
||||
"input": "find current OpenAI documentation",
|
||||
"commands": {
|
||||
"search_query": [{"q": "OpenAI Codex Search"}]
|
||||
},
|
||||
"max_output_tokens": 256,
|
||||
"stream": true,
|
||||
"store": false,
|
||||
"service_tier": "priority",
|
||||
"unknown_field": true
|
||||
}
|
||||
}))
|
||||
.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");
|
||||
assert_eq!(payload["success"], json!(true), "payload={payload}");
|
||||
assert_eq!(
|
||||
payload["attempts"][0]["request_body"]["id"],
|
||||
json!("aether-model-test-provider-query-search-trace")
|
||||
);
|
||||
assert_eq!(
|
||||
payload["attempts"][0]["response_body"]["output"],
|
||||
json!("search result")
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_routes_grok_responses_admin_pool_model_test_through_grok_runtime() {
|
||||
run_provider_query_test(
|
||||
@@ -3837,13 +4027,12 @@ async fn gateway_handles_openai_responses_test_model_locally_impl() {
|
||||
.and_then(|value| value.as_str()),
|
||||
Some(prompt)
|
||||
);
|
||||
assert_eq!(
|
||||
plan.body
|
||||
.json_body
|
||||
.as_ref()
|
||||
.and_then(|body| body.get("instructions")),
|
||||
Some(&json!(""))
|
||||
);
|
||||
assert!(plan
|
||||
.body
|
||||
.json_body
|
||||
.as_ref()
|
||||
.and_then(|body| body.get("instructions"))
|
||||
.is_none());
|
||||
assert_eq!(
|
||||
plan.body
|
||||
.json_body
|
||||
@@ -3856,7 +4045,7 @@ async fn gateway_handles_openai_responses_test_model_locally_impl() {
|
||||
.json_body
|
||||
.as_ref()
|
||||
.and_then(|body| body.get("prompt_cache_key"))
|
||||
.is_some());
|
||||
.is_none());
|
||||
Json(json!({
|
||||
"request_id": plan.request_id,
|
||||
"candidate_id": plan.candidate_id,
|
||||
@@ -3960,8 +4149,8 @@ async fn gateway_handles_openai_image_test_model_locally_impl() {
|
||||
assert_eq!(plan.client_api_format, "openai:image");
|
||||
assert_eq!(plan.provider_api_format, "openai:image");
|
||||
assert_eq!(plan.model_name.as_deref(), Some("gpt-image-1"));
|
||||
assert_eq!(plan.url, "https://api.openai.example/v1/responses");
|
||||
assert!(plan.stream);
|
||||
assert_eq!(plan.url, "https://api.openai.example/v1/images/generations");
|
||||
assert!(!plan.stream);
|
||||
assert_eq!(
|
||||
plan.headers.get("authorization").map(String::as_str),
|
||||
Some("Bearer sk-test-image")
|
||||
@@ -3971,38 +4160,42 @@ async fn gateway_handles_openai_image_test_model_locally_impl() {
|
||||
.json_body
|
||||
.as_ref()
|
||||
.and_then(|body| body.get("model")),
|
||||
Some(&json!(crate::ai_serving::CODEX_OPENAI_IMAGE_INTERNAL_MODEL))
|
||||
Some(&json!("gpt-image-1"))
|
||||
);
|
||||
assert_eq!(
|
||||
plan.body
|
||||
.json_body
|
||||
.as_ref()
|
||||
.and_then(|body| body.get("input"))
|
||||
.and_then(|input| input.as_array())
|
||||
.and_then(|items| items.first())
|
||||
.and_then(|item| item.get("content"))
|
||||
.and_then(|body| body.get("prompt"))
|
||||
.and_then(|value| value.as_str()),
|
||||
Some("Draw a small blue square")
|
||||
);
|
||||
assert!(plan
|
||||
.body
|
||||
.json_body
|
||||
.as_ref()
|
||||
.is_some_and(|body| body.get("stream").is_none()));
|
||||
Json(json!({
|
||||
"request_id": plan.request_id,
|
||||
"candidate_id": plan.candidate_id,
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "text/event-stream"
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"body_bytes_b64": base64::engine::general_purpose::STANDARD.encode(
|
||||
concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"created_at\":1776839946}}\n\n",
|
||||
"event: response.output_item.done\n",
|
||||
"data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"image_generation_call\",\"output_format\":\"png\",\"revised_prompt\":\"revised prompt\",\"result\":\"aGVsbG8=\"}}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_img_123\",\"model\":\"gpt-image-1\",\"status\":\"completed\",\"tool_usage\":{\"image_gen\":{\"input_tokens\":171,\"output_tokens\":1372,\"total_tokens\":1543}}}}\n\n"
|
||||
)
|
||||
.as_bytes()
|
||||
)
|
||||
"json_body": {
|
||||
"created": 1776839946,
|
||||
"model": "gpt-image-1",
|
||||
"data": [{
|
||||
"b64_json": "aGVsbG8=",
|
||||
"revised_prompt": "revised prompt"
|
||||
}],
|
||||
"usage": {
|
||||
"input_tokens": 171,
|
||||
"output_tokens": 1372,
|
||||
"total_tokens": 1543
|
||||
}
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 19
|
||||
|
||||
@@ -832,7 +832,7 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() {
|
||||
"is_active": false,
|
||||
"concurrent_limit": 8,
|
||||
"max_retries": 6,
|
||||
"request_timeout": 55.0,
|
||||
"request_timeout": aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS,
|
||||
"stream_first_byte_timeout": 11.0,
|
||||
"enable_format_conversion": false,
|
||||
"config": {
|
||||
@@ -860,7 +860,10 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() {
|
||||
assert_eq!(payload["enable_format_conversion"], false);
|
||||
assert_eq!(payload["is_active"], false);
|
||||
assert_eq!(payload["max_retries"], 6);
|
||||
assert_eq!(payload["request_timeout"], 55.0);
|
||||
assert_eq!(
|
||||
payload["request_timeout"].as_f64(),
|
||||
Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64)
|
||||
);
|
||||
assert_eq!(payload["stream_first_byte_timeout"], 11.0);
|
||||
assert_eq!(payload["proxy"], json!({"url": "https://proxy.example"}));
|
||||
assert_eq!(payload["claude_code_advanced"], json!({"pool_size": 2}));
|
||||
@@ -870,6 +873,21 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() {
|
||||
assert_eq!(payload["ops_configured"], true);
|
||||
assert_eq!(payload["ops_architecture_id"], "cubence");
|
||||
|
||||
let invalid_timeout_response = reqwest::Client::new()
|
||||
.patch(format!("{gateway_url}/api/admin/providers/provider-openai"))
|
||||
.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!({
|
||||
"request_timeout":
|
||||
aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS + 1
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(invalid_timeout_response.status(), StatusCode::BAD_REQUEST);
|
||||
|
||||
let disable_response = reqwest::Client::new()
|
||||
.patch(format!("{gateway_url}/api/admin/providers/provider-openai"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
@@ -924,6 +942,10 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() {
|
||||
.iter()
|
||||
.find(|provider| provider.id == "provider-openai")
|
||||
.expect("provider should exist");
|
||||
assert_eq!(
|
||||
updated_provider.request_timeout_secs,
|
||||
Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64)
|
||||
);
|
||||
assert_eq!(
|
||||
updated_provider
|
||||
.config
|
||||
@@ -1000,6 +1022,7 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
|
||||
"website": "codex.example",
|
||||
"keep_priority_on_conversion": true,
|
||||
"max_retries": 7,
|
||||
"request_timeout": aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS,
|
||||
"config": {"chat_pii_redaction": {"enabled": true}},
|
||||
"pool_advanced": {},
|
||||
"failover_rules": {"strategy": "ordered"},
|
||||
@@ -1034,6 +1057,10 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
|
||||
assert_eq!(created.website.as_deref(), Some("https://codex.example"));
|
||||
assert!(created.enable_format_conversion);
|
||||
assert_eq!(created.max_retries, Some(7));
|
||||
assert_eq!(
|
||||
created.request_timeout_secs,
|
||||
Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64)
|
||||
);
|
||||
assert_eq!(created.keep_priority_on_conversion, true);
|
||||
assert_eq!(
|
||||
created
|
||||
@@ -1080,7 +1107,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(), 3);
|
||||
assert_eq!(endpoints.len(), 4);
|
||||
let responses_endpoint = endpoints
|
||||
.iter()
|
||||
.find(|endpoint| endpoint.api_format == "openai:responses")
|
||||
@@ -1089,6 +1116,10 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
|
||||
.iter()
|
||||
.find(|endpoint| endpoint.api_format == "openai:responses:compact")
|
||||
.expect("compact endpoint should exist");
|
||||
let search_endpoint = endpoints
|
||||
.iter()
|
||||
.find(|endpoint| endpoint.api_format == "openai:search")
|
||||
.expect("search endpoint should exist");
|
||||
let image_endpoint = endpoints
|
||||
.iter()
|
||||
.find(|endpoint| endpoint.api_format == "openai:image")
|
||||
@@ -1101,12 +1132,17 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
|
||||
compact_endpoint.base_url,
|
||||
"https://chatgpt.com/backend-api/codex"
|
||||
);
|
||||
assert_eq!(
|
||||
search_endpoint.base_url,
|
||||
"https://chatgpt.com/backend-api/codex"
|
||||
);
|
||||
assert_eq!(
|
||||
image_endpoint.base_url,
|
||||
"https://chatgpt.com/backend-api/codex"
|
||||
);
|
||||
assert_eq!(responses_endpoint.max_retries, Some(7));
|
||||
assert_eq!(compact_endpoint.max_retries, Some(7));
|
||||
assert_eq!(search_endpoint.max_retries, Some(7));
|
||||
assert_eq!(image_endpoint.max_retries, Some(7));
|
||||
assert_eq!(
|
||||
responses_endpoint
|
||||
@@ -1116,16 +1152,25 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
|
||||
.and_then(serde_json::Value::as_str),
|
||||
Some("force_stream")
|
||||
);
|
||||
assert_eq!(
|
||||
search_endpoint
|
||||
.config
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("upstream_stream_policy"))
|
||||
.and_then(serde_json::Value::as_str),
|
||||
None
|
||||
);
|
||||
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")
|
||||
None
|
||||
);
|
||||
assert!(responses_endpoint.body_rules.is_none());
|
||||
assert!(compact_endpoint.body_rules.is_none());
|
||||
assert!(search_endpoint.body_rules.is_none());
|
||||
assert!(image_endpoint.body_rules.is_none());
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
@@ -1216,7 +1261,7 @@ async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoint
|
||||
.list_endpoints_by_provider_ids(&["provider-codex".to_string()])
|
||||
.await
|
||||
.expect("endpoints should list");
|
||||
assert_eq!(endpoints.len(), 3);
|
||||
assert_eq!(endpoints.len(), 4);
|
||||
let responses_endpoint = endpoints
|
||||
.iter()
|
||||
.find(|endpoint| endpoint.api_format == "openai:responses")
|
||||
@@ -1225,6 +1270,10 @@ async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoint
|
||||
.iter()
|
||||
.find(|endpoint| endpoint.api_format == "openai:responses:compact")
|
||||
.expect("compact endpoint should exist");
|
||||
let search_endpoint = endpoints
|
||||
.iter()
|
||||
.find(|endpoint| endpoint.api_format == "openai:search")
|
||||
.expect("search endpoint should exist");
|
||||
let image_endpoint = endpoints
|
||||
.iter()
|
||||
.find(|endpoint| endpoint.api_format == "openai:image")
|
||||
@@ -1232,6 +1281,7 @@ async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoint
|
||||
|
||||
assert_eq!(responses_endpoint.max_retries, Some(9));
|
||||
assert_eq!(compact_endpoint.max_retries, Some(9));
|
||||
assert_eq!(search_endpoint.max_retries, Some(9));
|
||||
assert_eq!(image_endpoint.max_retries, Some(9));
|
||||
assert_eq!(
|
||||
responses_endpoint
|
||||
@@ -1242,20 +1292,37 @@ async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoint
|
||||
.and_then(serde_json::Value::as_bool),
|
||||
Some(true)
|
||||
);
|
||||
assert_eq!(
|
||||
search_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!(
|
||||
search_endpoint
|
||||
.config
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("upstream_stream_policy"))
|
||||
.and_then(serde_json::Value::as_str),
|
||||
None
|
||||
);
|
||||
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")
|
||||
None
|
||||
);
|
||||
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!(keys[0].api_formats, Some(json!(["openai:responses"])));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
|
||||
@@ -1379,7 +1379,7 @@ async fn gateway_handles_admin_stats_leaderboard_models_locally_with_trusted_adm
|
||||
assert_eq!(payload["metric"], "tokens");
|
||||
assert_eq!(payload["items"][0]["rank"], 1);
|
||||
assert_eq!(payload["items"][0]["id"], "gpt-5");
|
||||
assert_eq!(payload["items"][0]["value"], 160);
|
||||
assert_eq!(payload["items"][0]["value"], 150);
|
||||
assert_eq!(payload["items"][1]["id"], "claude-3-5-sonnet");
|
||||
assert_eq!(payload["items"][1]["value"], 100);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
@@ -641,7 +641,7 @@ async fn gateway_handles_admin_usage_aggregation_stats_locally_with_trusted_admi
|
||||
assert_eq!(items[0]["model"], "gpt-5");
|
||||
assert_eq!(items[0]["request_count"], 2);
|
||||
assert_eq!(items[0]["output_tokens"], 40);
|
||||
assert_eq!(items[0]["effective_input_tokens"], 150);
|
||||
assert_eq!(items[0]["effective_input_tokens"], 120);
|
||||
assert_eq!(items[0]["total_input_context"], 160);
|
||||
assert_eq!(items[0]["cache_creation_tokens"], 30);
|
||||
assert_eq!(items[0]["cache_creation_ephemeral_5m_tokens"], 12);
|
||||
@@ -1026,7 +1026,7 @@ async fn gateway_handles_admin_usage_active_locally_with_trusted_admin_principal
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["requests"].as_array().expect("array").len(), 1);
|
||||
assert_eq!(payload["requests"][0]["id"], "usage-pending");
|
||||
assert_eq!(payload["requests"][0]["effective_input_tokens"], 5);
|
||||
assert_eq!(payload["requests"][0]["effective_input_tokens"], 0);
|
||||
assert_eq!(payload["requests"][0]["provider"], "OpenAI");
|
||||
assert_eq!(payload["requests"][0]["api_key_name"], "fresh-primary");
|
||||
assert_eq!(payload["requests"][0]["has_fallback"], true);
|
||||
@@ -1326,7 +1326,7 @@ async fn gateway_handles_admin_usage_records_locally_with_trusted_admin_principa
|
||||
payload["records"][0]["provider_key_name"],
|
||||
"upstream-primary"
|
||||
);
|
||||
assert_eq!(payload["records"][0]["effective_input_tokens"], 35);
|
||||
assert_eq!(payload["records"][0]["effective_input_tokens"], 20);
|
||||
assert_eq!(payload["records"][0]["first_byte_time_ms"], 120);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
@@ -2053,8 +2053,8 @@ async fn gateway_handles_admin_usage_detail_locally_with_trusted_admin_principal
|
||||
assert_eq!(payload["api_key"]["name"], "primary");
|
||||
assert_eq!(payload["provider"], "OpenAI");
|
||||
assert_eq!(payload["model"], "gpt-5");
|
||||
assert_eq!(payload["effective_input_tokens"], 115);
|
||||
assert_eq!(payload["total_tokens"], 165);
|
||||
assert_eq!(payload["effective_input_tokens"], 100);
|
||||
assert_eq!(payload["total_tokens"], 150);
|
||||
assert_eq!(payload["cache_creation_cost"], 0.0);
|
||||
assert_eq!(payload["cache_read_cost"], 0.0);
|
||||
assert_eq!(
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use std::io;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_contracts::tunnel::RequestMeta;
|
||||
use aether_data::repository::proxy_nodes::ProxyNodeReadRepository;
|
||||
use axum::body::Body;
|
||||
use axum::routing::{any, post};
|
||||
@@ -10,12 +11,45 @@ use futures_util::stream;
|
||||
use http::header::HeaderValue;
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
use std::collections::HashMap;
|
||||
use std::time::Duration;
|
||||
|
||||
use super::{
|
||||
build_router_with_state, sample_proxy_node, start_server, AppState, GatewayDataState,
|
||||
InMemoryProxyNodeRepository, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
fn relay_request_meta(
|
||||
stream: bool,
|
||||
request_timeout_ms: Option<u64>,
|
||||
stream_first_byte_timeout_ms: Option<u64>,
|
||||
) -> RequestMeta {
|
||||
RequestMeta {
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
endpoint_id: Some("endpoint-1".to_string()),
|
||||
key_id: Some("key-1".to_string()),
|
||||
method: "POST".to_string(),
|
||||
url: "https://example.com/responses".to_string(),
|
||||
headers: HashMap::new(),
|
||||
stream,
|
||||
request_timeout_ms,
|
||||
stream_first_byte_timeout_ms,
|
||||
timeout: 60,
|
||||
follow_redirects: None,
|
||||
http1_only: false,
|
||||
transport_profile: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn relay_envelope(meta: &RequestMeta, body: &[u8]) -> Vec<u8> {
|
||||
let encoded_meta = serde_json::to_vec(meta).expect("metadata should encode");
|
||||
let mut envelope = Vec::with_capacity(4 + encoded_meta.len() + body.len());
|
||||
envelope.extend_from_slice(&(encoded_meta.len() as u32).to_be_bytes());
|
||||
envelope.extend_from_slice(&encoded_meta);
|
||||
envelope.extend_from_slice(body);
|
||||
envelope
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_internal_tunnel_heartbeat_locally_with_loopback() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
@@ -292,10 +326,13 @@ async fn gateway_forwards_tunnel_relay_to_attachment_owner() {
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let envelope = relay_envelope(&relay_request_meta(false, None, None), b"relay-envelope");
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/api/internal/tunnel/relay/node-123"))
|
||||
.header(TRACE_ID_HEADER, "trace-owner-forward")
|
||||
.body("relay-envelope")
|
||||
.header(http::header::CONTENT_TYPE, "application/octet-stream")
|
||||
.body(envelope.clone())
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
@@ -309,8 +346,8 @@ async fn gateway_forwards_tunnel_relay_to_attachment_owner() {
|
||||
Some("trace-owner-forward")
|
||||
);
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
"relay-envelope"
|
||||
response.bytes().await.expect("body should read"),
|
||||
Bytes::from(envelope)
|
||||
);
|
||||
assert_eq!(*owner_hits.lock().expect("mutex should lock"), 1);
|
||||
|
||||
@@ -318,6 +355,71 @@ async fn gateway_forwards_tunnel_relay_to_attachment_owner() {
|
||||
owner_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_owner_relay_uses_non_stream_timeout_from_envelope() {
|
||||
let owner = Router::new().route(
|
||||
"/api/internal/tunnel/relay/node-123",
|
||||
post(|body: Body| async move {
|
||||
let body = axum::body::to_bytes(body, usize::MAX)
|
||||
.await
|
||||
.expect("body should read");
|
||||
tokio::time::sleep(Duration::from_millis(40)).await;
|
||||
(StatusCode::OK, Body::from(body))
|
||||
}),
|
||||
);
|
||||
|
||||
let (owner_url, owner_handle) = start_server(owner).await;
|
||||
let data_state = GatewayDataState::disabled().with_system_config_values_for_tests(vec![(
|
||||
"tunnel.attachments.node-123".to_string(),
|
||||
json!({
|
||||
"gateway_instance_id": "gateway-b",
|
||||
"relay_base_url": owner_url,
|
||||
"conn_count": 1,
|
||||
"observed_at_unix_secs": 4_102_444_800u64,
|
||||
}),
|
||||
)]);
|
||||
let mut state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(data_state)
|
||||
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a.internal"));
|
||||
let short_timeout_client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_millis(10))
|
||||
.build()
|
||||
.expect("test client should build");
|
||||
state.client = short_timeout_client.clone();
|
||||
state.owner_forward_client = short_timeout_client;
|
||||
let gateway = build_router_with_state(state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let meta = relay_request_meta(false, Some(100), None);
|
||||
let envelope = relay_envelope(&meta, b"relay-body");
|
||||
let encoded_meta = serde_json::to_vec(&meta).expect("metadata should encode");
|
||||
let split_at = 4 + encoded_meta.len() / 2;
|
||||
let request_body = reqwest::Body::wrap_stream(stream::iter(vec![
|
||||
Ok::<Bytes, io::Error>(Bytes::copy_from_slice(&envelope[..split_at])),
|
||||
Ok::<Bytes, io::Error>(Bytes::copy_from_slice(&envelope[split_at..])),
|
||||
]));
|
||||
|
||||
let response = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(1))
|
||||
.build()
|
||||
.expect("request client should build")
|
||||
.post(format!("{gateway_url}/api/internal/tunnel/relay/node-123"))
|
||||
.body(request_body)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.bytes().await.expect("response body should read"),
|
||||
Bytes::from(envelope)
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
owner_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_streams_tunnel_relay_body_to_attachment_owner() {
|
||||
let owner_hits = Arc::new(Mutex::new(0usize));
|
||||
@@ -331,8 +433,12 @@ async fn gateway_streams_tunnel_relay_body_to_attachment_owner() {
|
||||
let body = axum::body::to_bytes(body, usize::MAX)
|
||||
.await
|
||||
.expect("body should read");
|
||||
assert_eq!(body, Bytes::from_static(b"relay-stream-envelope"));
|
||||
(StatusCode::OK, Body::from("stream-ok"))
|
||||
let response_body = Body::from_stream(async_stream::stream! {
|
||||
yield Ok::<_, io::Error>(Bytes::from_static(b"stream-"));
|
||||
tokio::time::sleep(Duration::from_millis(40)).await;
|
||||
yield Ok::<_, io::Error>(Bytes::from_static(b"ok"));
|
||||
});
|
||||
(StatusCode::OK, response_body)
|
||||
}
|
||||
}),
|
||||
);
|
||||
@@ -347,19 +453,23 @@ async fn gateway_streams_tunnel_relay_body_to_attachment_owner() {
|
||||
"observed_at_unix_secs": 4_102_444_800u64,
|
||||
}),
|
||||
)]);
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(data_state)
|
||||
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a.internal")),
|
||||
);
|
||||
let mut state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(data_state)
|
||||
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a.internal"));
|
||||
state.client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_millis(10))
|
||||
.build()
|
||||
.expect("short shared client should build");
|
||||
let gateway = build_router_with_state(state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let request_body = reqwest::Body::wrap_stream(stream::iter(vec![
|
||||
Ok::<Bytes, io::Error>(Bytes::from_static(b"relay-")),
|
||||
Ok::<Bytes, io::Error>(Bytes::from_static(b"stream-")),
|
||||
Ok::<Bytes, io::Error>(Bytes::from_static(b"envelope")),
|
||||
]));
|
||||
let meta = relay_request_meta(true, Some(900_000), Some(100));
|
||||
let envelope = relay_envelope(&meta, b"relay-stream-envelope");
|
||||
let expected_envelope = Bytes::copy_from_slice(&envelope);
|
||||
let request_body = reqwest::Body::wrap_stream(stream::iter(vec![Ok::<Bytes, io::Error>(
|
||||
expected_envelope.clone(),
|
||||
)]));
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/api/internal/tunnel/relay/node-123"))
|
||||
.body(request_body)
|
||||
@@ -369,8 +479,8 @@ async fn gateway_streams_tunnel_relay_body_to_attachment_owner() {
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
"stream-ok"
|
||||
response.bytes().await.expect("body should read"),
|
||||
Bytes::from_static(b"stream-ok")
|
||||
);
|
||||
assert_eq!(*owner_hits.lock().expect("mutex should lock"), 1);
|
||||
|
||||
|
||||
@@ -35,6 +35,27 @@ use aether_data_contracts::repository::video_tasks::{
|
||||
use base64::Engine as _;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
fn run_frontdoor_async_test<F>(name: &'static str, future: F)
|
||||
where
|
||||
F: std::future::Future<Output = ()> + Send + 'static,
|
||||
{
|
||||
let handle = std::thread::Builder::new()
|
||||
.name(name.to_string())
|
||||
.stack_size(16 * 1024 * 1024)
|
||||
.spawn(move || {
|
||||
tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("frontdoor test runtime should build")
|
||||
.block_on(future);
|
||||
})
|
||||
.expect("large-stack frontdoor test thread should spawn");
|
||||
|
||||
if let Err(payload) = handle.join() {
|
||||
std::panic::resume_unwind(payload);
|
||||
}
|
||||
}
|
||||
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
use super::{
|
||||
hash_api_key, sample_models_candidate_row, unrestricted_models_snapshot,
|
||||
InMemoryAuthApiKeySnapshotRepository, InMemoryMinimalCandidateSelectionReadRepository,
|
||||
InMemoryVideoTaskRepository, UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository,
|
||||
VideoTaskStatus, VideoTaskWriteRepository, DEVELOPMENT_ENCRYPTION_KEY,
|
||||
InMemoryVideoTaskRepository, StoredAuthApiKeySnapshot, UpsertVideoTask, VideoTaskLookupKey,
|
||||
VideoTaskReadRepository, VideoTaskStatus, VideoTaskWriteRepository, DEVELOPMENT_ENCRYPTION_KEY,
|
||||
};
|
||||
use crate::image_capabilities::openai_image_gateway_max_generation_count;
|
||||
use crate::tests::{
|
||||
@@ -26,6 +26,94 @@ use std::collections::HashMap;
|
||||
use std::future::pending;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
fn codex_models_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("[email protected]".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(json!(["codex"])),
|
||||
Some(json!(["openai:responses"])),
|
||||
Some(json!(["frontier-sol", "broken-luna"])),
|
||||
api_key_id.to_string(),
|
||||
Some("codex-models".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(10),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(json!(["codex"])),
|
||||
Some(json!(["openai:responses"])),
|
||||
Some(json!(["frontier-sol", "broken-luna"])),
|
||||
)
|
||||
.expect("Codex models auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_codex_models_candidate_row(
|
||||
provider_id: &str,
|
||||
global_model_name: &str,
|
||||
source_model_name: &str,
|
||||
) -> StoredMinimalCandidateSelectionRow {
|
||||
let mut row = sample_models_candidate_row(
|
||||
provider_id,
|
||||
"codex",
|
||||
"openai:responses",
|
||||
global_model_name,
|
||||
10,
|
||||
);
|
||||
row.provider_type = "codex".to_string();
|
||||
row.key_auth_type = "oauth".to_string();
|
||||
row.model_provider_model_name = source_model_name.to_string();
|
||||
row.model_provider_model_mappings = Some(vec![
|
||||
aether_data_contracts::repository::candidate_selection::StoredProviderModelMapping {
|
||||
name: source_model_name.to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["openai:responses".to_string()]),
|
||||
endpoint_ids: None,
|
||||
},
|
||||
]);
|
||||
row
|
||||
}
|
||||
|
||||
fn complete_codex_model_card(source_model_name: &str) -> serde_json::Value {
|
||||
json!({
|
||||
"id": source_model_name,
|
||||
"api_formats": ["openai:responses"],
|
||||
"slug": source_model_name,
|
||||
"display_name": "GPT-5.6-Sol",
|
||||
"description": "Frontier coding model",
|
||||
"default_reasoning_level": "low",
|
||||
"supported_reasoning_levels": [
|
||||
{"effort": "low", "description": "Low"},
|
||||
{"effort": "medium", "description": "Medium"},
|
||||
{"effort": "high", "description": "High"},
|
||||
{"effort": "xhigh", "description": "XHigh"},
|
||||
{"effort": "max", "description": "Max"},
|
||||
{"effort": "ultra", "description": "Ultra"}
|
||||
],
|
||||
"shell_type": "shell_command",
|
||||
"visibility": "list",
|
||||
"supported_in_api": true,
|
||||
"priority": 1,
|
||||
"availability_nux": null,
|
||||
"upgrade": null,
|
||||
"base_instructions": "Use the current Codex instructions.",
|
||||
"model_messages": null,
|
||||
"support_verbosity": true,
|
||||
"default_verbosity": "low",
|
||||
"apply_patch_tool_type": "freeform",
|
||||
"truncation_policy": {"mode": "tokens", "limit": 10000},
|
||||
"supports_parallel_tool_calls": true,
|
||||
"experimental_supported_tools": [],
|
||||
"minimal_client_version": "0.144.0",
|
||||
"future_capability": {"enabled": true}
|
||||
})
|
||||
}
|
||||
|
||||
fn gemini_operation_status_label(status: VideoTaskStatus) -> &'static str {
|
||||
match status {
|
||||
VideoTaskStatus::Pending => "Pending",
|
||||
@@ -360,6 +448,121 @@ async fn gateway_handles_public_openai_models_without_hitting_fallback_probe() {
|
||||
fallback_probe_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_serves_codex_model_cards_for_versioned_models_requests() {
|
||||
let codex_row =
|
||||
sample_codex_models_candidate_row("provider-codex-models", "frontier-sol", "gpt-5.6-sol");
|
||||
let incomplete_codex_row = sample_codex_models_candidate_row(
|
||||
"provider-codex-incomplete",
|
||||
"broken-luna",
|
||||
"gpt-5.6-luna",
|
||||
);
|
||||
let candidate_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
codex_row.clone(),
|
||||
incomplete_codex_row.clone(),
|
||||
sample_models_candidate_row(
|
||||
"provider-openai-responses",
|
||||
"openai",
|
||||
"openai:responses",
|
||||
"custom-responses-model",
|
||||
20,
|
||||
),
|
||||
]));
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![
|
||||
(
|
||||
Some(hash_api_key("sk-codex-models")),
|
||||
codex_models_snapshot("key-codex-models", "user-codex-models"),
|
||||
),
|
||||
(
|
||||
Some(hash_api_key("sk-standard-models")),
|
||||
unrestricted_models_snapshot("key-standard-models", "user-standard-models"),
|
||||
),
|
||||
]));
|
||||
let state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_minimal_candidate_selection_and_auth_for_tests(
|
||||
candidate_repository,
|
||||
auth_repository,
|
||||
),
|
||||
);
|
||||
state
|
||||
.runtime_kv_setex(
|
||||
&format!(
|
||||
"upstream_models:{}:{}",
|
||||
codex_row.provider_id, codex_row.key_id
|
||||
),
|
||||
&serde_json::to_string(&vec![complete_codex_model_card("gpt-5.6-sol")])
|
||||
.expect("model cache should serialize"),
|
||||
60,
|
||||
)
|
||||
.await
|
||||
.expect("model cache should seed");
|
||||
state
|
||||
.runtime_kv_setex(
|
||||
&format!(
|
||||
"upstream_models:{}:{}",
|
||||
incomplete_codex_row.provider_id, incomplete_codex_row.key_id
|
||||
),
|
||||
&serde_json::to_string(&vec![json!({
|
||||
"id": "gpt-5.6-luna",
|
||||
"slug": "gpt-5.6-luna",
|
||||
"display_name": "GPT-5.6-Luna"
|
||||
})])
|
||||
.expect("incomplete model cache should serialize"),
|
||||
60,
|
||||
)
|
||||
.await
|
||||
.expect("incomplete model cache should seed");
|
||||
|
||||
let gateway = build_router_with_state(state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
let codex_response = client
|
||||
.get(format!("{gateway_url}/v1/models?client_version=0.144.1"))
|
||||
.header("authorization", "Bearer sk-codex-models")
|
||||
.send()
|
||||
.await
|
||||
.expect("Codex models request should succeed");
|
||||
assert_eq!(codex_response.status(), StatusCode::OK);
|
||||
let codex_payload: serde_json::Value = codex_response
|
||||
.json()
|
||||
.await
|
||||
.expect("Codex models body should parse");
|
||||
assert_eq!(codex_payload["models"].as_array().map(Vec::len), Some(1));
|
||||
assert_eq!(codex_payload["models"][0]["slug"], "frontier-sol");
|
||||
assert_eq!(
|
||||
codex_payload["models"][0]["supported_reasoning_levels"][5]["effort"],
|
||||
"ultra"
|
||||
);
|
||||
assert_eq!(
|
||||
codex_payload["models"][0]["future_capability"],
|
||||
json!({"enabled": true})
|
||||
);
|
||||
assert!(codex_payload["models"][0].get("id").is_none());
|
||||
assert!(codex_payload["models"][0].get("api_formats").is_none());
|
||||
assert!(codex_payload.get("object").is_none());
|
||||
|
||||
let standard_response = client
|
||||
.get(format!("{gateway_url}/v1/models"))
|
||||
.header("authorization", "Bearer sk-standard-models")
|
||||
.send()
|
||||
.await
|
||||
.expect("standard models request should succeed");
|
||||
assert_eq!(standard_response.status(), StatusCode::OK);
|
||||
let standard_payload: serde_json::Value = standard_response
|
||||
.json()
|
||||
.await
|
||||
.expect("standard models body should parse");
|
||||
assert_eq!(standard_payload["object"], "list");
|
||||
assert!(standard_payload["data"].is_array());
|
||||
assert!(standard_payload.get("models").is_none());
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_openai_models_list_drops_disabled_global_model_after_cache_invalidation() {
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
@@ -1212,7 +1415,7 @@ async fn gateway_does_not_locally_reject_image_model_name_on_chat_completions()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_image_request_with_n_greater_than_four_without_hitting_fallback_probe() {
|
||||
async fn gateway_rejects_image_request_above_gateway_limit_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(
|
||||
@@ -1247,7 +1450,7 @@ async fn gateway_rejects_image_request_with_n_greater_than_four_without_hitting_
|
||||
serde_json::to_vec(&json!({
|
||||
"model": "grok-imagine-image-lite",
|
||||
"prompt": "draw",
|
||||
"n": 5,
|
||||
"n": openai_image_gateway_max_generation_count() + 1,
|
||||
"response_format": "b64_json"
|
||||
}))
|
||||
.expect("request body should encode"),
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
use super::{
|
||||
hash_api_key, sample_endpoint, sample_key, sample_models_candidate_row, sample_provider,
|
||||
unrestricted_models_snapshot, InMemoryAuthApiKeySnapshotRepository,
|
||||
InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository,
|
||||
InMemoryRequestCandidateRepository, DEVELOPMENT_ENCRYPTION_KEY,
|
||||
hash_api_key, run_frontdoor_async_test, sample_endpoint, sample_key,
|
||||
sample_models_candidate_row, sample_provider, unrestricted_models_snapshot,
|
||||
InMemoryAuthApiKeySnapshotRepository, InMemoryMinimalCandidateSelectionReadRepository,
|
||||
InMemoryProviderCatalogReadRepository, InMemoryRequestCandidateRepository,
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
};
|
||||
use crate::tests::{
|
||||
any, build_router, build_router_with_state, build_state_with_execution_runtime_override, json,
|
||||
@@ -160,8 +161,15 @@ async fn gateway_returns_internal_gateway_plan_sync_proxy_public_action_without_
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_internal_gateway_execute_sync_locally() {
|
||||
#[test]
|
||||
fn gateway_handles_internal_gateway_execute_sync_locally() {
|
||||
run_frontdoor_async_test(
|
||||
"gateway_handles_internal_gateway_execute_sync_locally",
|
||||
gateway_handles_internal_gateway_execute_sync_locally_impl(),
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_handles_internal_gateway_execute_sync_locally_impl() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let fallback_probe = Router::new().route(
|
||||
|
||||
@@ -70,6 +70,7 @@ async fn gateway_exposes_frontdoor_manifest_without_proxying_upstream() {
|
||||
assert!(owned_routes
|
||||
.iter()
|
||||
.any(|value| value == "/v1/responses/compact"));
|
||||
assert!(owned_routes.iter().any(|value| value == "/v1/alpha/search"));
|
||||
assert!(owned_routes.iter().any(|value| value == "/health"));
|
||||
assert!(owned_routes.iter().any(|value| value == "/v1/health"));
|
||||
assert!(owned_routes.iter().any(|value| value == "/v1/providers"));
|
||||
|
||||
@@ -5735,7 +5735,7 @@ async fn gateway_handles_users_me_usage_locally_without_proxying_upstream() {
|
||||
payload["records"][0]["cache_creation_ephemeral_5m_input_tokens"],
|
||||
4
|
||||
);
|
||||
assert_eq!(payload["records"][0]["effective_input_tokens"], 105);
|
||||
assert_eq!(payload["records"][0]["effective_input_tokens"], 95);
|
||||
assert_eq!(
|
||||
payload["records"][0]["cache_creation_ephemeral_1h_input_tokens"],
|
||||
6
|
||||
@@ -5762,10 +5762,7 @@ async fn gateway_handles_users_me_usage_locally_without_proxying_upstream() {
|
||||
payload["summary_by_model"][0]["cache_creation_ephemeral_1h_tokens"],
|
||||
6
|
||||
);
|
||||
assert_eq!(
|
||||
payload["summary_by_model"][0]["effective_input_tokens"],
|
||||
105
|
||||
);
|
||||
assert_eq!(payload["summary_by_model"][0]["effective_input_tokens"], 95);
|
||||
assert_eq!(payload["summary_by_model"][0]["total_input_context"], 120);
|
||||
assert!(payload.get("summary_by_provider").is_none());
|
||||
assert_eq!(payload["billing"]["id"], "wallet-auth-1");
|
||||
|
||||
@@ -191,7 +191,7 @@ async fn gateway_handles_dashboard_stats_locally_without_proxying_upstream() {
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["today"]["requests"], 1);
|
||||
assert_eq!(payload["today"]["tokens"], 160);
|
||||
assert_eq!(payload["today"]["tokens"], 150);
|
||||
assert_eq!(payload["api_keys"]["total"], 2);
|
||||
assert_eq!(payload["api_keys"]["active"], 1);
|
||||
assert_eq!(payload["stats"][3]["subValue"], json!("输入 240 / 输出 60"));
|
||||
@@ -646,7 +646,7 @@ async fn gateway_handles_admin_dashboard_stats_locally_without_proxying_upstream
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["today"]["requests"], 2);
|
||||
assert_eq!(payload["today"]["tokens"], 17_450);
|
||||
assert_eq!(payload["today"]["tokens"], 16_250);
|
||||
assert_eq!(payload["today"]["cost"], json!(2.5));
|
||||
assert_eq!(payload["cost_stats"]["cost_savings"], json!(0.025));
|
||||
let stats = payload["stats"].as_array().expect("stats should be array");
|
||||
@@ -664,10 +664,10 @@ async fn gateway_handles_admin_dashboard_stats_locally_without_proxying_upstream
|
||||
.iter()
|
||||
.find(|item| item["name"] == json!("今日 Token"))
|
||||
.expect("today token stats card should exist");
|
||||
assert_eq!(today_token_stats["value"], json!("17.4K"));
|
||||
assert_eq!(today_token_stats["value"], json!("16.2K"));
|
||||
assert_eq!(
|
||||
today_token_stats["subValue"],
|
||||
json!("输入 12.1K / 输出 3.1K · 写缓存 1.25K / 读缓存 1K")
|
||||
json!("输入 10.9K / 输出 3.1K · 写缓存 1.25K / 读缓存 1K")
|
||||
);
|
||||
assert_eq!(payload["users"]["total"], 2);
|
||||
assert_eq!(payload["users"]["active"], 1);
|
||||
|
||||
@@ -89,6 +89,13 @@ fn sample_cli_auth_snapshot(
|
||||
}
|
||||
|
||||
fn sample_provider(provider_id: &str) -> StoredProviderCatalogProvider {
|
||||
sample_provider_with_request_timeout(provider_id, None)
|
||||
}
|
||||
|
||||
fn sample_provider_with_request_timeout(
|
||||
provider_id: &str,
|
||||
request_timeout_secs: Option<f64>,
|
||||
) -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
provider_id.to_string(),
|
||||
provider_id.to_string(),
|
||||
@@ -96,7 +103,17 @@ fn sample_provider(provider_id: &str) -> StoredProviderCatalogProvider {
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(true, false, false, None, None, None, None, None, None)
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
request_timeout_secs,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_endpoint(endpoint_id: &str, provider_id: &str) -> StoredProviderCatalogEndpoint {
|
||||
@@ -435,6 +452,7 @@ async fn gateway_forwards_public_request_to_remote_tunnel_owner_before_fallback_
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
tokio::time::sleep(Duration::from_millis(40)).await;
|
||||
*seen_owner_inner.lock().expect("mutex should lock") = Some(SeenOwnerRequest {
|
||||
path: parts
|
||||
.uri
|
||||
@@ -511,7 +529,10 @@ async fn gateway_forwards_public_request_to_remote_tunnel_owner_before_fallback_
|
||||
let (owner_url, owner_handle) = start_server(owner).await;
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-owner")],
|
||||
vec![sample_provider_with_request_timeout(
|
||||
"provider-owner",
|
||||
Some(0.1),
|
||||
)],
|
||||
vec![sample_endpoint("endpoint-owner", "provider-owner")],
|
||||
vec![sample_key("key-owner", "provider-owner", "node-owner")],
|
||||
));
|
||||
@@ -540,6 +561,12 @@ async fn gateway_forwards_public_request_to_remote_tunnel_owner_before_fallback_
|
||||
state = state
|
||||
.with_data_state_for_tests(data_state)
|
||||
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a:8080"));
|
||||
let short_timeout_client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_millis(10))
|
||||
.build()
|
||||
.expect("test client should build");
|
||||
state.client = short_timeout_client.clone();
|
||||
state.owner_forward_client = short_timeout_client;
|
||||
state.remember_scheduler_affinity_target(
|
||||
"scheduler_affinity:api-key-affinity-1:openai:chat:gpt-4.1",
|
||||
crate::cache::SchedulerAffinityTarget {
|
||||
@@ -922,32 +949,39 @@ async fn gateway_streamifies_sync_json_from_remote_tunnel_owner_before_returning
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
let encoded_response = serde_json::to_vec(&json!({
|
||||
"id": "resp-codex-affinity-stream-123",
|
||||
"object": "response",
|
||||
"model": "gpt-5.4",
|
||||
"status": "completed",
|
||||
"output": [{
|
||||
"type": "message",
|
||||
"id": "msg-codex-affinity-stream-123",
|
||||
"role": "assistant",
|
||||
"content": [{
|
||||
"type": "output_text",
|
||||
"text": "Hello from affinity sync json",
|
||||
"annotations": []
|
||||
}]
|
||||
}],
|
||||
"usage": {
|
||||
"input_tokens": 1,
|
||||
"output_tokens": 2,
|
||||
"total_tokens": 3
|
||||
}
|
||||
}))
|
||||
.expect("body should encode");
|
||||
let split_at = encoded_response.len() / 2;
|
||||
let first = axum::body::Bytes::copy_from_slice(&encoded_response[..split_at]);
|
||||
let second = axum::body::Bytes::copy_from_slice(&encoded_response[split_at..]);
|
||||
let response_body = Body::from_stream(async_stream::stream! {
|
||||
yield Ok::<_, std::io::Error>(first);
|
||||
tokio::time::sleep(Duration::from_millis(40)).await;
|
||||
yield Ok::<_, std::io::Error>(second);
|
||||
});
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from(
|
||||
serde_json::to_vec(&json!({
|
||||
"id": "resp-codex-affinity-stream-123",
|
||||
"object": "response",
|
||||
"model": "gpt-5.4",
|
||||
"status": "completed",
|
||||
"output": [{
|
||||
"type": "message",
|
||||
"id": "msg-codex-affinity-stream-123",
|
||||
"role": "assistant",
|
||||
"content": [{
|
||||
"type": "output_text",
|
||||
"text": "Hello from affinity sync json",
|
||||
"annotations": []
|
||||
}]
|
||||
}],
|
||||
"usage": {
|
||||
"input_tokens": 1,
|
||||
"output_tokens": 2,
|
||||
"total_tokens": 3
|
||||
}
|
||||
}))
|
||||
.expect("body should encode"),
|
||||
))
|
||||
.body(response_body)
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
@@ -1001,6 +1035,10 @@ async fn gateway_streamifies_sync_json_from_remote_tunnel_owner_before_returning
|
||||
state = state
|
||||
.with_data_state_for_tests(data_state)
|
||||
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a:8080"));
|
||||
state.client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_millis(10))
|
||||
.build()
|
||||
.expect("short shared client should build");
|
||||
state.remember_scheduler_affinity_target(
|
||||
"scheduler_affinity:api-key-affinity-cli-1:openai:responses:gpt-5.4",
|
||||
crate::cache::SchedulerAffinityTarget {
|
||||
|
||||
@@ -804,7 +804,7 @@ fn gateway_records_openai_sync_usage_and_pricing_with_cache_tokens() {
|
||||
async fn gateway_records_openai_sync_usage_and_pricing_with_cache_tokens_impl() {
|
||||
let expected = ExpectedUsagePricing {
|
||||
input_tokens: 120,
|
||||
billed_input_tokens: 100,
|
||||
billed_input_tokens: 20,
|
||||
output_tokens: 40,
|
||||
cache_creation_tokens: 80,
|
||||
cache_creation_ephemeral_5m_tokens: 0,
|
||||
@@ -898,7 +898,7 @@ fn gateway_records_openai_stream_usage_and_pricing_with_cache_tokens() {
|
||||
async fn gateway_records_openai_stream_usage_and_pricing_with_cache_tokens_impl() {
|
||||
let expected = ExpectedUsagePricing {
|
||||
input_tokens: 240,
|
||||
billed_input_tokens: 200,
|
||||
billed_input_tokens: 120,
|
||||
output_tokens: 60,
|
||||
cache_creation_tokens: 80,
|
||||
cache_creation_ephemeral_5m_tokens: 0,
|
||||
|
||||
Reference in New Issue
Block a user