mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-04 16:37:46 +08:00
Merge branch 'merge-pr-483'
This commit is contained in:
fawney19
@@ -13,6 +13,7 @@ use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadReposi
|
||||
use aether_data_contracts::repository::candidate_selection::{
|
||||
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
|
||||
};
|
||||
use aether_data_contracts::repository::candidates::RequestCandidateReadRepository;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
@@ -741,3 +742,451 @@ async fn gateway_bridges_codex_image_sync_json_to_streaming_image_sse() {
|
||||
execution_runtime_handle.abort();
|
||||
refresh_handle.abort();
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenImageBridgeExecutionPlan {
|
||||
trace_id: String,
|
||||
client_api_format: String,
|
||||
provider_api_format: String,
|
||||
url: String,
|
||||
plan_stream: bool,
|
||||
auth_header: String,
|
||||
chatgpt_web_marker: String,
|
||||
body_json: serde_json::Value,
|
||||
}
|
||||
|
||||
fn image_bridge_hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn image_bridge_auth_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,
|
||||
None,
|
||||
Some(serde_json::json!([
|
||||
"openai:chat",
|
||||
"openai:responses",
|
||||
"openai:image"
|
||||
])),
|
||||
Some(serde_json::json!(["gpt-image-2"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800_i64),
|
||||
None,
|
||||
Some(serde_json::json!([
|
||||
"openai:chat",
|
||||
"openai:responses",
|
||||
"openai:image"
|
||||
])),
|
||||
Some(serde_json::json!(["gpt-image-2"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn image_bridge_candidate_row(
|
||||
prefix: &str,
|
||||
provider_name: &str,
|
||||
provider_type: &str,
|
||||
) -> StoredMinimalCandidateSelectionRow {
|
||||
let key_auth_type = if provider_type == "chatgpt_web" {
|
||||
"bearer"
|
||||
} else {
|
||||
"api_key"
|
||||
};
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: format!("provider-{prefix}"),
|
||||
provider_name: provider_name.to_string(),
|
||||
provider_type: provider_type.to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: format!("endpoint-{prefix}"),
|
||||
endpoint_api_format: "openai:image".to_string(),
|
||||
endpoint_api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("image".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: format!("key-{prefix}"),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: key_auth_type.to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["openai:image".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"openai:image": 1})),
|
||||
model_id: format!("model-{prefix}"),
|
||||
global_model_id: format!("global-model-{prefix}"),
|
||||
global_model_name: "gpt-image-2".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(false),
|
||||
model_provider_model_name: "gpt-image-2".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-image-2".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["openai:image".to_string()]),
|
||||
endpoint_ids: None,
|
||||
}]),
|
||||
model_supports_streaming: Some(false),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn image_bridge_provider_catalog_provider(
|
||||
prefix: &str,
|
||||
provider_name: &str,
|
||||
provider_type: &str,
|
||||
base_url: &str,
|
||||
) -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
format!("provider-{prefix}"),
|
||||
provider_name.to_string(),
|
||||
Some(base_url.to_string()),
|
||||
provider_type.to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn image_bridge_provider_catalog_endpoint(
|
||||
prefix: &str,
|
||||
base_url: &str,
|
||||
) -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
format!("endpoint-{prefix}"),
|
||||
format!("provider-{prefix}"),
|
||||
"openai:image".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("image".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
base_url.to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn image_bridge_provider_catalog_key(
|
||||
prefix: &str,
|
||||
provider_type: &str,
|
||||
) -> StoredProviderCatalogKey {
|
||||
let auth_type = if provider_type == "chatgpt_web" {
|
||||
"bearer"
|
||||
} else {
|
||||
"api_key"
|
||||
};
|
||||
StoredProviderCatalogKey::new(
|
||||
format!("key-{prefix}"),
|
||||
format!("provider-{prefix}"),
|
||||
"prod".to_string(),
|
||||
auth_type.to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["openai:image"])),
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-image-bridge")
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"openai:image": 1})),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
async fn start_image_bridge_gateway(
|
||||
prefix: &str,
|
||||
provider_name: &str,
|
||||
provider_type: &str,
|
||||
base_url: &str,
|
||||
execution_runtime_url: String,
|
||||
) -> (
|
||||
String,
|
||||
tokio::task::JoinHandle<()>,
|
||||
String,
|
||||
Arc<InMemoryRequestCandidateRepository>,
|
||||
) {
|
||||
let client_api_key = format!("sk-client-{prefix}");
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(image_bridge_hash_api_key(&client_api_key)),
|
||||
image_bridge_auth_snapshot(&format!("api-key-{prefix}"), &format!("user-{prefix}")),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
image_bridge_candidate_row(prefix, provider_name, provider_type),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![image_bridge_provider_catalog_provider(
|
||||
prefix,
|
||||
provider_name,
|
||||
provider_type,
|
||||
base_url,
|
||||
)],
|
||||
vec![image_bridge_provider_catalog_endpoint(prefix, base_url)],
|
||||
vec![image_bridge_provider_catalog_key(prefix, provider_type)],
|
||||
));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
(
|
||||
gateway_url,
|
||||
gateway_handle,
|
||||
client_api_key,
|
||||
request_candidate_repository,
|
||||
)
|
||||
}
|
||||
|
||||
fn capture_image_bridge_execution_plan(
|
||||
parts: http::request::Parts,
|
||||
payload: serde_json::Value,
|
||||
) -> SeenImageBridgeExecutionPlan {
|
||||
SeenImageBridgeExecutionPlan {
|
||||
trace_id: parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
client_api_format: payload
|
||||
.get("client_api_format")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
provider_api_format: payload
|
||||
.get("provider_api_format")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
url: payload
|
||||
.get("url")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
plan_stream: payload
|
||||
.get("stream")
|
||||
.and_then(|value| value.as_bool())
|
||||
.unwrap_or(false),
|
||||
auth_header: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("authorization"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
chatgpt_web_marker: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-aether-chatgpt-web-image"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
body_json: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.cloned()
|
||||
.unwrap_or(serde_json::Value::Null),
|
||||
}
|
||||
}
|
||||
|
||||
fn image_bridge_execution_runtime(
|
||||
seen_execution_plan: Arc<Mutex<Option<SeenImageBridgeExecutionPlan>>>,
|
||||
) -> Router {
|
||||
Router::new().route(
|
||||
"/v1/execute/stream",
|
||||
any(move |request: Request| {
|
||||
let seen_execution_plan_inner = Arc::clone(&seen_execution_plan);
|
||||
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_plan_inner.lock().expect("mutex should lock") =
|
||||
Some(capture_image_bridge_execution_plan(parts, payload));
|
||||
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_bridge_123\\\",\\\"type\\\":\\\"image_generation_call\\\",\\\"result\\\":\\\"aGVsbG8=\\\",\\\"output_format\\\":\\\"png\\\"}}\\n\\n\"}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp_bridge_123\\\",\\\"object\\\":\\\"response\\\",\\\"model\\\":\\\"gpt-image-2\\\",\\\"status\\\":\\\"completed\\\",\\\"output\\\":[]}}\\n\\n\"}}\n",
|
||||
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
|
||||
);
|
||||
let mut response = http::Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from(frames))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
http::HeaderValue::from_static("application/x-ndjson"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_routes_openai_chat_stream_image_intent_to_openai_image_plan_without_streaming_support(
|
||||
) {
|
||||
let seen_execution_plan = Arc::new(Mutex::new(None::<SeenImageBridgeExecutionPlan>));
|
||||
let execution_runtime = image_bridge_execution_runtime(Arc::clone(&seen_execution_plan));
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let (gateway_url, gateway_handle, client_api_key, request_candidate_repository) =
|
||||
start_image_bridge_gateway(
|
||||
"chat-stream-image-bridge",
|
||||
"image-provider",
|
||||
"custom",
|
||||
"https://images.example.com",
|
||||
execution_runtime_url,
|
||||
)
|
||||
.await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(http::header::AUTHORIZATION, format!("Bearer {client_api_key}"))
|
||||
.header(TRACE_ID_HEADER, "trace-chat-stream-image-bridge-123")
|
||||
.body(
|
||||
r#"{"model":"gpt-image-2","messages":[{"role":"user","content":"Draw a city made of glass"}],"stream":true,"size":"1024x1024"}"#,
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let status = response.status();
|
||||
let response_text = response.text().await.expect("body should read");
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-chat-stream-image-bridge-123")
|
||||
.await
|
||||
.expect("request candidates should read");
|
||||
assert_eq!(
|
||||
status,
|
||||
StatusCode::OK,
|
||||
"{response_text}\n{stored_candidates:#?}"
|
||||
);
|
||||
assert!(response_text.contains("\"object\":\"chat.completion.chunk\""));
|
||||
assert!(response_text.contains(""));
|
||||
assert!(response_text.contains("data: [DONE]"));
|
||||
assert!(!response_text.contains("image_generation.completed"));
|
||||
|
||||
let seen_plan = seen_execution_plan
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution plan should be captured");
|
||||
assert_eq!(seen_plan.trace_id, "trace-chat-stream-image-bridge-123");
|
||||
assert_eq!(seen_plan.client_api_format, "openai:chat");
|
||||
assert_eq!(seen_plan.provider_api_format, "openai:image");
|
||||
assert_eq!(seen_plan.url, "https://images.example.com/v1/responses");
|
||||
assert!(seen_plan.plan_stream);
|
||||
assert_eq!(seen_plan.auth_header, "Bearer sk-upstream-image-bridge");
|
||||
assert_eq!(seen_plan.chatgpt_web_marker, "");
|
||||
assert_eq!(seen_plan.body_json["model"], "gpt-image-2");
|
||||
assert_eq!(seen_plan.body_json["stream"], true);
|
||||
assert_eq!(
|
||||
seen_plan.body_json["input"][0]["content"],
|
||||
"Draw a city made of glass"
|
||||
);
|
||||
assert_eq!(seen_plan.body_json["tools"][0]["type"], "image_generation");
|
||||
assert_eq!(seen_plan.body_json["tools"][0]["size"], "1024x1024");
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_routes_openai_responses_stream_image_intent_to_openai_image_plan_without_streaming_support(
|
||||
) {
|
||||
let seen_execution_plan = Arc::new(Mutex::new(None::<SeenImageBridgeExecutionPlan>));
|
||||
let execution_runtime = image_bridge_execution_runtime(Arc::clone(&seen_execution_plan));
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let (gateway_url, gateway_handle, client_api_key, _request_candidate_repository) =
|
||||
start_image_bridge_gateway(
|
||||
"responses-stream-image-bridge",
|
||||
"image-provider",
|
||||
"custom",
|
||||
"https://images.example.com",
|
||||
execution_runtime_url,
|
||||
)
|
||||
.await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/responses"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(http::header::AUTHORIZATION, format!("Bearer {client_api_key}"))
|
||||
.header(TRACE_ID_HEADER, "trace-responses-stream-image-bridge-123")
|
||||
.body(
|
||||
r#"{"model":"gpt-image-2","input":"Draw a mountain observatory","tools":[{"type":"image_generation","size":"1024x1024"}],"stream":true}"#,
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let status = response.status();
|
||||
let response_text = response.text().await.expect("body should read");
|
||||
assert_eq!(status, StatusCode::OK, "{response_text}");
|
||||
assert!(response_text.contains("response.output_item.done"));
|
||||
assert!(response_text.contains("image_generation_call"));
|
||||
|
||||
let seen_plan = seen_execution_plan
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution plan should be captured");
|
||||
assert_eq!(
|
||||
seen_plan.trace_id,
|
||||
"trace-responses-stream-image-bridge-123"
|
||||
);
|
||||
assert_eq!(seen_plan.client_api_format, "openai:responses");
|
||||
assert_eq!(seen_plan.provider_api_format, "openai:image");
|
||||
assert_eq!(seen_plan.url, "https://images.example.com/v1/responses");
|
||||
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");
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
}
|
||||
|
||||
@@ -148,6 +148,120 @@ async fn gateway_handles_admin_provider_endpoints_locally_with_trusted_admin_pri
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_counts_fixed_provider_oauth_keys_for_inherited_endpoint_formats() {
|
||||
let mut codex_provider = sample_provider("provider-codex", "codex", 10);
|
||||
codex_provider.provider_type = "codex".to_string();
|
||||
let mut chatgpt_web_provider = sample_provider("provider-chatgpt-web", "chatgpt_web", 20);
|
||||
chatgpt_web_provider.provider_type = "chatgpt_web".to_string();
|
||||
|
||||
let mut codex_key = sample_key(
|
||||
"key-codex-oauth",
|
||||
"provider-codex",
|
||||
"openai:responses:compact",
|
||||
"oauth-token",
|
||||
);
|
||||
codex_key.auth_type = "oauth".to_string();
|
||||
codex_key.api_formats = Some(json!(["legacy:mismatch"]));
|
||||
|
||||
let mut chatgpt_web_key = sample_key(
|
||||
"key-chatgpt-web-oauth",
|
||||
"provider-chatgpt-web",
|
||||
"openai:image",
|
||||
"oauth-token",
|
||||
);
|
||||
chatgpt_web_key.auth_type = "oauth".to_string();
|
||||
chatgpt_web_key.api_formats = Some(json!(["legacy:mismatch"]));
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![codex_provider, chatgpt_web_provider],
|
||||
vec![
|
||||
sample_endpoint(
|
||||
"endpoint-codex-compact",
|
||||
"provider-codex",
|
||||
"openai:responses:compact",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
),
|
||||
sample_endpoint(
|
||||
"endpoint-codex-image",
|
||||
"provider-codex",
|
||||
"openai:image",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
),
|
||||
sample_endpoint(
|
||||
"endpoint-chatgpt-web-image",
|
||||
"provider-chatgpt-web",
|
||||
"openai:image",
|
||||
"https://chatgpt.com",
|
||||
),
|
||||
],
|
||||
vec![codex_key, chatgpt_web_key],
|
||||
));
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let codex_response = client
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-codex/endpoints?skip=0&limit=50"
|
||||
))
|
||||
.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")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(codex_response.status(), StatusCode::OK);
|
||||
let codex_payload: serde_json::Value = codex_response.json().await.expect("json should parse");
|
||||
let codex_items = codex_payload
|
||||
.as_array()
|
||||
.expect("payload should be an array");
|
||||
for api_format in ["openai:responses:compact", "openai:image"] {
|
||||
let endpoint = codex_items
|
||||
.iter()
|
||||
.find(|item| item["api_format"] == api_format)
|
||||
.expect("endpoint should exist");
|
||||
assert_eq!(endpoint["total_keys"], 1);
|
||||
assert_eq!(endpoint["active_keys"], 1);
|
||||
}
|
||||
|
||||
let chatgpt_web_response = client
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-chatgpt-web/endpoints?skip=0&limit=50"
|
||||
))
|
||||
.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")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(chatgpt_web_response.status(), StatusCode::OK);
|
||||
let chatgpt_web_payload: serde_json::Value = chatgpt_web_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json should parse");
|
||||
let chatgpt_web_items = chatgpt_web_payload
|
||||
.as_array()
|
||||
.expect("payload should be an array");
|
||||
let chatgpt_web_image = chatgpt_web_items
|
||||
.iter()
|
||||
.find(|item| item["api_format"] == "openai:image")
|
||||
.expect("image endpoint should exist");
|
||||
assert_eq!(chatgpt_web_image["total_keys"], 1);
|
||||
assert_eq!(chatgpt_web_image["active_keys"], 1);
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_counts_keys_with_null_api_formats_for_each_fixed_provider_endpoint() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
@@ -195,7 +309,7 @@ async fn gateway_counts_keys_with_null_api_formats_for_each_fixed_provider_endpo
|
||||
vec![inherited_key],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (_, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
|
||||
@@ -936,6 +936,142 @@ async fn gateway_handles_admin_global_model_routing_locally_with_trusted_admin_p
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_global_model_routing_counts_image_provider_keys_by_provider_model_name() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/models/global/global-gpt-image/routing",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut image_provider = sample_provider("provider-image", "image", 10);
|
||||
image_provider.provider_type = "chatgpt_web".to_string();
|
||||
let grok_provider = sample_provider("provider-grok", "grok2api", 20);
|
||||
|
||||
let mut image_key = sample_key(
|
||||
"key-image-routing",
|
||||
"provider-image",
|
||||
"legacy:mismatch",
|
||||
"sk-image-routing-1234",
|
||||
);
|
||||
image_key.name = "image-account".to_string();
|
||||
image_key.auth_type = "oauth".to_string();
|
||||
image_key.allowed_models = Some(json!(["gpt-image-2"]));
|
||||
|
||||
let mut grok_key = sample_key(
|
||||
"key-grok-routing",
|
||||
"provider-grok",
|
||||
"openai:chat",
|
||||
"sk-grok-routing-5678",
|
||||
);
|
||||
grok_key.name = "all".to_string();
|
||||
grok_key.allowed_models = Some(json!(["gpt-image-2"]));
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![image_provider, grok_provider],
|
||||
vec![
|
||||
sample_endpoint(
|
||||
"endpoint-image",
|
||||
"provider-image",
|
||||
"openai:image",
|
||||
"https://chatgpt.example",
|
||||
),
|
||||
sample_endpoint(
|
||||
"endpoint-grok-chat",
|
||||
"provider-grok",
|
||||
"openai:chat",
|
||||
"https://grok.example",
|
||||
),
|
||||
],
|
||||
vec![image_key, grok_key],
|
||||
));
|
||||
let global_model_repository = Arc::new(
|
||||
InMemoryGlobalModelReadRepository::seed(Vec::new())
|
||||
.with_admin_global_models(vec![sample_admin_global_model(
|
||||
"global-gpt-image",
|
||||
"GPT-Image-2",
|
||||
"GPT-Image-2",
|
||||
)])
|
||||
.with_admin_provider_models(vec![
|
||||
sample_admin_provider_model(
|
||||
"model-image-gpt-image",
|
||||
"provider-image",
|
||||
"global-gpt-image",
|
||||
"gpt-image-2",
|
||||
),
|
||||
sample_admin_provider_model(
|
||||
"model-grok-gpt-image",
|
||||
"provider-grok",
|
||||
"global-gpt-image",
|
||||
"gpt-image-2",
|
||||
),
|
||||
]),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)
|
||||
.with_global_model_repository_for_tests(global_model_repository),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/models/global/global-gpt-image/routing"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["global_model_name"], "GPT-Image-2");
|
||||
assert_eq!(payload["total_providers"], 2);
|
||||
assert_eq!(payload["active_providers"], 2);
|
||||
|
||||
let providers = payload["providers"].as_array().expect("providers array");
|
||||
assert_eq!(providers.len(), 2);
|
||||
|
||||
let image_endpoints = providers[0]["endpoints"]
|
||||
.as_array()
|
||||
.expect("image endpoints array");
|
||||
assert_eq!(providers[0]["id"], "provider-image");
|
||||
assert_eq!(image_endpoints[0]["api_format"], "openai:image");
|
||||
assert_eq!(image_endpoints[0]["total_keys"], 1);
|
||||
assert_eq!(image_endpoints[0]["active_keys"], 1);
|
||||
assert_eq!(image_endpoints[0]["keys"][0]["name"], "image-account");
|
||||
|
||||
let grok_endpoints = providers[1]["endpoints"]
|
||||
.as_array()
|
||||
.expect("grok endpoints array");
|
||||
assert_eq!(providers[1]["id"], "provider-grok");
|
||||
assert_eq!(grok_endpoints[0]["api_format"], "openai:chat");
|
||||
assert_eq!(grok_endpoints[0]["total_keys"], 1);
|
||||
assert_eq!(grok_endpoints[0]["active_keys"], 1);
|
||||
assert_eq!(grok_endpoints[0]["keys"][0]["name"], "all");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_creates_admin_global_model_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
|
||||
@@ -113,10 +113,7 @@ struct ExpectedUsagePricing {
|
||||
|
||||
impl ExpectedUsagePricing {
|
||||
fn total_tokens(self) -> u64 {
|
||||
self.input_tokens
|
||||
.saturating_add(self.output_tokens)
|
||||
.saturating_add(self.cache_creation_tokens)
|
||||
.saturating_add(self.cache_read_tokens)
|
||||
self.input_tokens.saturating_add(self.output_tokens)
|
||||
}
|
||||
|
||||
fn cache_creation_uncategorized_tokens(self) -> u64 {
|
||||
|
||||
Reference in New Issue
Block a user