mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
Merge remote-tracking branch 'upstream/aether-rust-pioneer' into fix-management-token-oauth-jsonb
# Conflicts: # crates/aether-data/src/lifecycle/bootstrap/postgres.rs # crates/aether-data/src/lifecycle/migrate/tests.rs
This commit is contained in:
@@ -20,6 +20,7 @@ use aether_data_contracts::repository::provider_catalog::{
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::data::GatewayDataState;
|
||||
use crate::tests::next_non_keepalive_chunk;
|
||||
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
@@ -370,11 +371,7 @@ async fn gateway_stops_execution_runtime_stream_when_client_disconnects() {
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let mut response = response;
|
||||
let first_chunk = response
|
||||
.chunk()
|
||||
.await
|
||||
.expect("first chunk should read")
|
||||
.expect("first chunk should exist");
|
||||
let first_chunk = next_non_keepalive_chunk(&mut response).await;
|
||||
assert_eq!(
|
||||
first_chunk,
|
||||
Bytes::from_static(b"data: {\"id\":\"chatcmpl-first\"}\n\n")
|
||||
|
||||
@@ -18,9 +18,10 @@ use crate::constants::{
|
||||
|
||||
use super::{
|
||||
build_router, build_router_with_execution_runtime_override, build_router_with_state,
|
||||
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
|
||||
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
|
||||
UsageRuntimeConfig, VideoTaskTruthSourceMode,
|
||||
build_state_with_execution_runtime_override, next_non_keepalive_chunk, start_server,
|
||||
strip_sse_keepalive_comments, wait_until, AppState, FrontdoorCorsConfig,
|
||||
FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason, UsageRuntimeConfig,
|
||||
VideoTaskTruthSourceMode,
|
||||
};
|
||||
|
||||
mod control_execute;
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
|
||||
to_bytes, AppState, Arc, Body, Bytes, HeaderName, HeaderValue, Json, Mutex, Request, Response,
|
||||
Router, StatusCode, EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_HEADER,
|
||||
TRACE_ID_HEADER,
|
||||
strip_sse_keepalive_comments, to_bytes, AppState, Arc, Body, Bytes, HeaderName, HeaderValue,
|
||||
Json, Mutex, Request, Response, Router, StatusCode, EXECUTION_PATH_EXECUTION_RUNTIME_STREAM,
|
||||
EXECUTION_PATH_HEADER, TRACE_ID_HEADER,
|
||||
};
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
@@ -391,7 +391,7 @@ async fn gateway_executes_openai_chat_stream_via_local_decision_gate_without_exe
|
||||
Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM)
|
||||
);
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
|
||||
"data: {\"id\":\"chatcmpl-local-123\"}\n\ndata: [DONE]\n\n"
|
||||
);
|
||||
|
||||
@@ -1774,7 +1774,7 @@ async fn gateway_executes_openai_chat_stream_with_custom_path_via_local_decision
|
||||
Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM)
|
||||
);
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
|
||||
"data: {\"id\":\"chatcmpl-local-custom-path-123\"}\n\ndata: [DONE]\n\n"
|
||||
);
|
||||
|
||||
@@ -2285,7 +2285,7 @@ async fn gateway_retries_next_local_openai_chat_stream_candidate_after_retryable
|
||||
Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM)
|
||||
);
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
|
||||
"data: {\"id\":\"chatcmpl-local-stream-failover-123\"}\n\ndata: [DONE]\n\n"
|
||||
);
|
||||
|
||||
|
||||
@@ -18,9 +18,10 @@ use crate::constants::{
|
||||
|
||||
use super::{
|
||||
build_router, build_router_with_execution_runtime_override, build_router_with_state,
|
||||
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
|
||||
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
|
||||
UsageRuntimeConfig, VideoTaskTruthSourceMode,
|
||||
build_state_with_execution_runtime_override, next_non_keepalive_chunk, start_server,
|
||||
strip_sse_keepalive_comments, wait_until, AppState, FrontdoorCorsConfig,
|
||||
FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason, UsageRuntimeConfig,
|
||||
VideoTaskTruthSourceMode,
|
||||
};
|
||||
|
||||
mod decision;
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
|
||||
to_bytes, Arc, Body, Bytes, HeaderName, HeaderValue, Infallible, Json, Mutex, Request,
|
||||
Response, Router, StatusCode, UsageRuntimeConfig, TRACE_ID_HEADER,
|
||||
strip_sse_keepalive_comments, to_bytes, Arc, Body, Bytes, HeaderName, HeaderValue, Infallible,
|
||||
Json, Mutex, Request, Response, Router, StatusCode, UsageRuntimeConfig, TRACE_ID_HEADER,
|
||||
};
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
@@ -469,7 +469,7 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
|
||||
"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"
|
||||
);
|
||||
|
||||
|
||||
@@ -18,9 +18,9 @@ use crate::constants::{
|
||||
|
||||
use super::{
|
||||
build_router, build_router_with_execution_runtime_override, build_router_with_state,
|
||||
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
|
||||
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
|
||||
UsageRuntimeConfig, VideoTaskTruthSourceMode,
|
||||
build_state_with_execution_runtime_override, start_server, strip_sse_keepalive_comments,
|
||||
wait_until, AppState, FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind,
|
||||
GatewayFallbackReason, UsageRuntimeConfig, VideoTaskTruthSourceMode,
|
||||
};
|
||||
|
||||
mod compact;
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
|
||||
to_bytes, Arc, Body, Bytes, HeaderName, HeaderValue, Json, Mutex, Request, Response, Router,
|
||||
StatusCode, TRACE_ID_HEADER,
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override, json,
|
||||
next_non_keepalive_chunk, start_server, strip_sse_keepalive_comments, to_bytes, Arc, Body,
|
||||
Bytes, HeaderName, HeaderValue, Json, Mutex, Request, Response, Router, StatusCode,
|
||||
TRACE_ID_HEADER,
|
||||
};
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
@@ -952,11 +953,12 @@ async fn gateway_executes_claude_cli_stream_via_local_decision_gate_without_wait
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
tokio::time::timeout(std::time::Duration::from_millis(100), response.chunk())
|
||||
.await
|
||||
.expect("same-format passthrough should yield first chunk before eof")
|
||||
.expect("first chunk should read")
|
||||
.expect("first chunk should exist"),
|
||||
tokio::time::timeout(
|
||||
std::time::Duration::from_millis(100),
|
||||
next_non_keepalive_chunk(&mut response),
|
||||
)
|
||||
.await
|
||||
.expect("same-format passthrough should yield first chunk before eof"),
|
||||
Bytes::from_static(b"event: message_start\ndata: {\"type\":\"message_start\"}\n\n")
|
||||
);
|
||||
assert_eq!(
|
||||
@@ -1470,7 +1472,7 @@ async fn gateway_executes_claude_code_cli_stream_via_local_decision_gate_with_lo
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
|
||||
"event: message_start\ndata: {\"type\":\"message_start\"}\n\n"
|
||||
);
|
||||
|
||||
@@ -1924,7 +1926,7 @@ async fn gateway_executes_claude_chat_stream_via_local_decision_gate_with_local_
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
|
||||
"event: message_start\ndata: {\"type\":\"message_start\"}\n\n"
|
||||
);
|
||||
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override,
|
||||
encrypt_python_fernet_plaintext, json, start_server, to_bytes, Arc, Body, Bytes, Digest,
|
||||
HeaderName, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
|
||||
encrypt_python_fernet_plaintext, json, start_server, strip_sse_keepalive_comments, to_bytes,
|
||||
Arc, Body, Bytes, Digest, HeaderName, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
|
||||
InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository,
|
||||
InMemoryRequestCandidateRepository, Json, Mutex, Request, RequestCandidateReadRepository,
|
||||
RequestCandidateStatus, Response, Router, Sha256, StatusCode, StoredAuthApiKeySnapshot,
|
||||
@@ -404,7 +404,7 @@ async fn gateway_executes_gemini_chat_stream_via_local_decision_gate_with_local_
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
|
||||
"data: {\"candidates\":[]}\n\n"
|
||||
);
|
||||
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override,
|
||||
encrypt_python_fernet_plaintext, json, start_server, to_bytes, Arc, Body, Bytes, Digest,
|
||||
HeaderName, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
|
||||
encrypt_python_fernet_plaintext, json, start_server, strip_sse_keepalive_comments, to_bytes,
|
||||
Arc, Body, Bytes, Digest, HeaderName, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
|
||||
InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository,
|
||||
InMemoryRequestCandidateRepository, Json, Mutex, Request, RequestCandidateReadRepository,
|
||||
RequestCandidateStatus, Response, Router, Sha256, StatusCode, StoredAuthApiKeySnapshot,
|
||||
@@ -381,7 +381,7 @@ async fn gateway_executes_gemini_cli_stream_via_local_decision_gate_with_local_s
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
|
||||
"data: {\"candidates\":[]}\n\n"
|
||||
);
|
||||
|
||||
@@ -872,7 +872,7 @@ async fn gateway_executes_gemini_cli_stream_via_local_decision_gate_after_oauth_
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
|
||||
"data: {\"candidates\":[]}\n\n"
|
||||
);
|
||||
|
||||
@@ -1339,7 +1339,7 @@ async fn gateway_executes_vertex_ai_gemini_cli_stream_via_local_decision_gate_wi
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
|
||||
"data: {\"candidates\":[]}\n\n"
|
||||
);
|
||||
|
||||
@@ -1849,7 +1849,8 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let response_text = response.text().await.expect("body should read");
|
||||
let response_text =
|
||||
strip_sse_keepalive_comments(&response.text().await.expect("body should read"));
|
||||
let payload = response_text
|
||||
.trim()
|
||||
.strip_prefix("data: ")
|
||||
|
||||
@@ -18,9 +18,9 @@ use crate::constants::{
|
||||
|
||||
use super::{
|
||||
build_router, build_router_with_execution_runtime_override, build_router_with_state,
|
||||
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
|
||||
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
|
||||
UsageRuntimeConfig, VideoTaskTruthSourceMode,
|
||||
build_state_with_execution_runtime_override, start_server, strip_sse_keepalive_comments,
|
||||
wait_until, AppState, FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind,
|
||||
GatewayFallbackReason, UsageRuntimeConfig, VideoTaskTruthSourceMode,
|
||||
};
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
|
||||
@@ -13,7 +13,9 @@ use aether_data::repository::global_models::InMemoryGlobalModelReadRepository;
|
||||
use aether_data::repository::oauth_providers::InMemoryOAuthProviderRepository;
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository;
|
||||
use aether_data::repository::users::{InMemoryUserReadRepository, StoredUserExportRow};
|
||||
use aether_data::repository::users::{
|
||||
InMemoryUserReadRepository, StoredUserAuthRecord, UpsertUserGroupRecord, UserReadRepository,
|
||||
};
|
||||
use aether_data::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
|
||||
use aether_data_contracts::repository::global_models::StoredPublicGlobalModel;
|
||||
use axum::body::Body;
|
||||
@@ -596,8 +598,8 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
|
||||
}),
|
||||
);
|
||||
|
||||
let user_repository = Arc::new(InMemoryUserReadRepository::seed_export_users(vec![
|
||||
StoredUserExportRow::new(
|
||||
let user_repository = Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![
|
||||
StoredUserAuthRecord::new(
|
||||
"user-1".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
true,
|
||||
@@ -608,12 +610,40 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
|
||||
Some(json!(["openai"])),
|
||||
Some(json!(["openai:chat"])),
|
||||
Some(json!(["gpt-5"])),
|
||||
Some(120),
|
||||
Some(json!({"gpt-5": {"cache_1h": true}})),
|
||||
true,
|
||||
false,
|
||||
Some(chrono::Utc::now()),
|
||||
None,
|
||||
)
|
||||
.expect("user export row should build"),
|
||||
.expect("user export row should build")
|
||||
.with_policy_modes(
|
||||
"specific".to_string(),
|
||||
"specific".to_string(),
|
||||
"specific".to_string(),
|
||||
)
|
||||
.expect("user policy modes should build"),
|
||||
]));
|
||||
let user_group = user_repository
|
||||
.create_user_group(UpsertUserGroupRecord {
|
||||
name: "Restricted GPT".to_string(),
|
||||
description: Some("GPT-only users".to_string()),
|
||||
priority: 10,
|
||||
allowed_providers: Some(vec!["openai".to_string()]),
|
||||
allowed_providers_mode: "specific".to_string(),
|
||||
allowed_api_formats: Some(vec!["openai:chat".to_string()]),
|
||||
allowed_api_formats_mode: "specific".to_string(),
|
||||
allowed_models: Some(vec!["gpt-5".to_string()]),
|
||||
allowed_models_mode: "specific".to_string(),
|
||||
rate_limit: Some(60),
|
||||
rate_limit_mode: "custom".to_string(),
|
||||
})
|
||||
.await
|
||||
.expect("user group should create")
|
||||
.expect("user group should exist");
|
||||
user_repository
|
||||
.replace_user_groups_for_user("user-1", std::slice::from_ref(&user_group.id))
|
||||
.await
|
||||
.expect("user group membership should create");
|
||||
let auth_repository = Arc::new(
|
||||
InMemoryAuthApiKeySnapshotRepository::default().with_export_records(vec![
|
||||
StoredAuthApiKeyExportRecord::new(
|
||||
@@ -728,9 +758,24 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["version"], "1.3");
|
||||
assert_eq!(payload["version"], "1.4");
|
||||
assert!(payload["exported_at"].as_str().is_some());
|
||||
assert_eq!(payload["user_groups"][0]["name"], "Restricted GPT");
|
||||
assert!(payload["user_groups"][0].get("priority").is_none());
|
||||
assert_eq!(
|
||||
payload["user_groups"][0]["allowed_models"],
|
||||
json!(["gpt-5"])
|
||||
);
|
||||
assert_eq!(payload["users"][0]["email"], "alice@example.com");
|
||||
assert_eq!(
|
||||
payload["users"][0]["allowed_models_mode"],
|
||||
json!("specific")
|
||||
);
|
||||
assert_eq!(payload["users"][0]["rate_limit_mode"], json!("system"));
|
||||
assert_eq!(
|
||||
payload["users"][0]["group_names"],
|
||||
json!(["Restricted GPT"])
|
||||
);
|
||||
assert_eq!(payload["users"][0]["wallet"]["balance"], json!(12.5));
|
||||
assert_eq!(
|
||||
payload["users"][0]["wallet"]["recharge_balance"],
|
||||
|
||||
@@ -12,7 +12,7 @@ use aether_data::repository::oauth_providers::{
|
||||
InMemoryOAuthProviderRepository, OAuthProviderReadRepository, StoredOAuthProviderConfig,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data::repository::users::StoredUserAuthRecord;
|
||||
use aether_data::repository::users::{StoredUserAuthRecord, UserReadRepository};
|
||||
use aether_data::repository::wallet::{StoredWalletSnapshot, WalletLookupKey};
|
||||
use aether_data_contracts::repository::global_models::{
|
||||
AdminGlobalModelListQuery, AdminProviderModelListQuery, GlobalModelReadRepository,
|
||||
@@ -556,11 +556,14 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
|
||||
}));
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::default());
|
||||
let user_repository =
|
||||
Arc::new(aether_data::repository::users::InMemoryUserReadRepository::default());
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_auth_api_key_repository_for_tests(Arc::clone(&auth_repository))
|
||||
.with_user_reader(user_repository)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
)
|
||||
.with_auth_users_for_tests([sample_import_admin_user("admin-user-123")])
|
||||
@@ -575,8 +578,21 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"version": "1.3",
|
||||
"version": "1.4",
|
||||
"merge_mode": "overwrite",
|
||||
"user_groups": [{
|
||||
"id": "source-group-1",
|
||||
"name": "GPT Import",
|
||||
"description": "Imported group",
|
||||
"allowed_providers": ["openai"],
|
||||
"allowed_providers_mode": "specific",
|
||||
"allowed_api_formats": ["openai:chat"],
|
||||
"allowed_api_formats_mode": "specific",
|
||||
"allowed_models": ["gpt-5"],
|
||||
"allowed_models_mode": "specific",
|
||||
"rate_limit": 44,
|
||||
"rate_limit_mode": "custom"
|
||||
}],
|
||||
"users": [{
|
||||
"email": "alice@example.com",
|
||||
"email_verified": true,
|
||||
@@ -587,6 +603,10 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
|
||||
"allowed_api_formats": ["openai:chat"],
|
||||
"allowed_models": ["gpt-5"],
|
||||
"rate_limit": 77,
|
||||
"allowed_models_mode": "specific",
|
||||
"rate_limit_mode": "custom",
|
||||
"group_ids": ["source-group-1"],
|
||||
"group_names": ["GPT Import"],
|
||||
"is_active": true,
|
||||
"wallet": {
|
||||
"balance": 20.0,
|
||||
@@ -654,6 +674,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
|
||||
let payload: Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(status, StatusCode::OK, "payload={payload}");
|
||||
assert_eq!(payload["message"], "用户数据导入成功");
|
||||
assert_eq!(payload["stats"]["user_groups"]["created"], json!(1));
|
||||
assert_eq!(payload["stats"]["users"]["created"], json!(1));
|
||||
assert_eq!(payload["stats"]["api_keys"]["created"], json!(1));
|
||||
assert_eq!(payload["stats"]["standalone_keys"]["created"], json!(1));
|
||||
@@ -683,8 +704,22 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
|
||||
imported_user.allowed_models,
|
||||
Some(vec!["gpt-5".to_string()])
|
||||
);
|
||||
assert_eq!(imported_user.allowed_models_mode, "specific");
|
||||
assert!(imported_user.is_active);
|
||||
|
||||
let imported_groups = state
|
||||
.list_user_groups_for_user(&imported_user.id)
|
||||
.await
|
||||
.expect("user groups should load");
|
||||
assert_eq!(imported_groups.len(), 1);
|
||||
assert_eq!(imported_groups[0].name, "GPT Import");
|
||||
assert_eq!(imported_groups[0].allowed_models_mode, "specific");
|
||||
assert_eq!(
|
||||
imported_groups[0].allowed_models,
|
||||
Some(vec!["gpt-5".to_string()])
|
||||
);
|
||||
assert_eq!(imported_groups[0].rate_limit, Some(44));
|
||||
|
||||
let user_wallet = state
|
||||
.find_wallet(WalletLookupKey::UserId(&imported_user.id))
|
||||
.await
|
||||
|
||||
@@ -2927,8 +2927,11 @@ async fn gateway_handles_admin_usage_cache_affinity_interval_timeline_with_legac
|
||||
role: "user".to_string(),
|
||||
auth_source: "local".to_string(),
|
||||
allowed_providers: None,
|
||||
allowed_providers_mode: "unrestricted".to_string(),
|
||||
allowed_api_formats: None,
|
||||
allowed_api_formats_mode: "unrestricted".to_string(),
|
||||
allowed_models: None,
|
||||
allowed_models_mode: "unrestricted".to_string(),
|
||||
is_active: true,
|
||||
is_deleted: false,
|
||||
created_at: None,
|
||||
|
||||
@@ -377,7 +377,10 @@ async fn embeddings_route_rejects_chat_only_model() {
|
||||
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(payload["error"]["message"], "当前密钥不允许访问模型 gpt-5");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"当前用户、用户组或密钥的访问控制策略不允许访问模型 gpt-5"
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
@@ -416,7 +419,7 @@ async fn embeddings_route_rejects_chat_only_api_format() {
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"当前密钥不允许访问 openai:embedding 格式"
|
||||
"当前用户、用户组或密钥的访问控制策略不允许访问 openai:embedding 格式"
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
|
||||
@@ -501,7 +501,7 @@ async fn gateway_locally_denies_disallowed_claude_api_format_without_hitting_con
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"当前密钥不允许访问 claude:messages 格式"
|
||||
"当前用户、用户组或密钥的访问控制策略不允许访问 claude:messages 格式"
|
||||
);
|
||||
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
@@ -584,7 +584,7 @@ async fn gateway_locally_denies_disallowed_provider_without_hitting_control_or_u
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"当前密钥不允许访问 claude 提供商"
|
||||
"当前用户、用户组或密钥的访问控制策略不允许访问 claude 提供商"
|
||||
);
|
||||
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
@@ -656,7 +656,7 @@ async fn gateway_locally_denies_disallowed_gemini_model_without_hitting_control_
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"当前密钥不允许访问模型 gemini-2.5-pro"
|
||||
"当前用户、用户组或密钥的访问控制策略不允许访问模型 gemini-2.5-pro"
|
||||
);
|
||||
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
@@ -806,7 +806,10 @@ async fn gateway_locally_denies_disallowed_openai_model_without_hitting_control_
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(payload["error"]["message"], "当前密钥不允许访问模型 gpt-5");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"当前用户、用户组或密钥的访问控制策略不允许访问模型 gpt-5"
|
||||
);
|
||||
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
|
||||
@@ -319,7 +319,7 @@ async fn rerank_route_rejects_chat_only_api_format() {
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"当前密钥不允许访问 openai:rerank 格式"
|
||||
"当前用户、用户组或密钥的访问控制策略不允许访问 openai:rerank 格式"
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
|
||||
@@ -6,8 +6,8 @@ use super::{
|
||||
};
|
||||
use crate::tests::{
|
||||
any, build_router, build_router_with_state, build_state_with_execution_runtime_override, json,
|
||||
start_server, AppState, Arc, Body, HeaderValue, Json, Mutex, Request, Response, Router,
|
||||
StatusCode, CONTROL_ACTION_PROXY_PUBLIC, CONTROL_EXECUTED_HEADER,
|
||||
start_server, strip_sse_keepalive_comments, AppState, Arc, Body, HeaderValue, Json, Mutex,
|
||||
Request, Response, Router, StatusCode, CONTROL_ACTION_PROXY_PUBLIC, CONTROL_EXECUTED_HEADER,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||
EXECUTION_PATH_HEADER,
|
||||
};
|
||||
@@ -458,7 +458,7 @@ async fn gateway_handles_internal_gateway_execute_stream_locally() {
|
||||
Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM)
|
||||
);
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
|
||||
"data: one\n\ndata: [DONE]\n\n"
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
@@ -88,3 +88,20 @@ pub(super) async fn wait_until(timeout_ms: u64, mut predicate: impl FnMut() -> b
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn strip_sse_keepalive_comments(body: &str) -> String {
|
||||
body.replace(": aether-keepalive\n\n", "")
|
||||
}
|
||||
|
||||
pub(crate) async fn next_non_keepalive_chunk(response: &mut reqwest::Response) -> Bytes {
|
||||
loop {
|
||||
let chunk = response
|
||||
.chunk()
|
||||
.await
|
||||
.expect("chunk should read")
|
||||
.expect("chunk should exist");
|
||||
if chunk.as_ref() != b": aether-keepalive\n\n" {
|
||||
return chunk;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -28,8 +28,8 @@ use sha2::{Digest, Sha256};
|
||||
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override, send_request,
|
||||
start_server, Body, HeaderValue, Json, Mutex, Request, Response, Router, StatusCode,
|
||||
UsageRuntimeConfig, TRACE_ID_HEADER,
|
||||
start_server, strip_sse_keepalive_comments, Body, HeaderValue, Json, Mutex, Request, Response,
|
||||
Router, StatusCode, UsageRuntimeConfig, TRACE_ID_HEADER,
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
|
||||
@@ -2,8 +2,8 @@ use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override,
|
||||
encrypt_python_fernet_plaintext, hash_api_key, json, sample_local_openai_auth_snapshot,
|
||||
sample_local_openai_candidate_row, sample_local_openai_endpoint, sample_local_openai_key,
|
||||
sample_local_openai_provider, send_request, start_server, Arc, Body, GatewayDataState,
|
||||
HeaderValue, InMemoryAuthApiKeySnapshotRepository,
|
||||
sample_local_openai_provider, send_request, start_server, strip_sse_keepalive_comments, Arc,
|
||||
Body, GatewayDataState, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
|
||||
InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository,
|
||||
InMemoryRequestCandidateRepository, InMemoryUsageReadRepository, Json, Mutex, Request,
|
||||
RequestCandidateReadRepository, RequestCandidateStatus, Response, Router, StatusCode,
|
||||
@@ -840,6 +840,104 @@ async fn gateway_records_failed_usage_when_all_local_openai_chat_candidates_exha
|
||||
assert_eq!(stored_candidates[0].status_code, Some(503));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_records_failed_usage_when_sync_runtime_transport_is_unavailable_without_plan_fallback(
|
||||
) {
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let execution_hits = Arc::new(Mutex::new(0usize));
|
||||
let execution_hits_clone = Arc::clone(&execution_hits);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-openai-local-transport-unavailable")),
|
||||
sample_local_openai_auth_snapshot(
|
||||
"api-key-openai-usage-local-transport-unavailable-1",
|
||||
"user-openai-usage-local-transport-unavailable-1",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_local_openai_candidate_row(),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_local_openai_provider()],
|
||||
vec![sample_local_openai_endpoint()],
|
||||
vec![sample_local_openai_key()],
|
||||
));
|
||||
|
||||
let gateway_state = crate::AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_execution_runtime_sync_override_for_tests(move |_plan| {
|
||||
*execution_hits_clone.lock().expect("mutex should lock") += 1;
|
||||
Err(crate::GatewayError::Internal(
|
||||
"simulated transport unavailable".to_string(),
|
||||
))
|
||||
})
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
Arc::clone(&usage_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
)
|
||||
.with_usage_runtime_for_tests(UsageRuntimeConfig {
|
||||
enabled: true,
|
||||
..UsageRuntimeConfig::default()
|
||||
});
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let request = Request::builder()
|
||||
.method(http::Method::POST)
|
||||
.uri("/v1/chat/completions")
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-openai-local-transport-unavailable",
|
||||
)
|
||||
.header(
|
||||
TRACE_ID_HEADER,
|
||||
"trace-openai-chat-local-transport-unavailable-123",
|
||||
)
|
||||
.body(Body::from("{\"model\":\"gpt-5\",\"messages\":[]}"))
|
||||
.expect("request should build");
|
||||
let response = send_request(gateway, request).await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(*execution_hits.lock().expect("mutex should lock"), 1);
|
||||
|
||||
let stored_usage = wait_for_usage_status(
|
||||
usage_repository.as_ref(),
|
||||
"trace-openai-chat-local-transport-unavailable-123",
|
||||
"failed",
|
||||
)
|
||||
.await;
|
||||
assert_eq!(stored_usage.status, "failed");
|
||||
assert_eq!(stored_usage.billing_status, "void");
|
||||
assert_eq!(stored_usage.status_code, Some(503));
|
||||
assert_eq!(
|
||||
stored_usage
|
||||
.response_body
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("error"))
|
||||
.and_then(|value| value.get("type"))
|
||||
.and_then(|value| value.as_str()),
|
||||
Some("execution_runtime_unavailable")
|
||||
);
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-openai-chat-local-transport-unavailable-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Failed);
|
||||
assert_eq!(
|
||||
stored_candidates[0].error_type.as_deref(),
|
||||
Some("execution_runtime_unavailable")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_records_failed_usage_for_claude_runtime_miss_without_execution_exhaustion() {
|
||||
run_async_test_on_large_stack(
|
||||
@@ -1184,7 +1282,8 @@ async fn gateway_handles_local_openai_chat_stream_report_with_local_reporting_wh
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body_text = response.text().await.expect("stream body should read");
|
||||
let body_text =
|
||||
strip_sse_keepalive_comments(&response.text().await.expect("stream body should read"));
|
||||
assert_eq!(
|
||||
body_text,
|
||||
"data: {\"id\":\"chatcmpl-local-report-stream-123\",\"usage\":{\"input_tokens\":2,\"output_tokens\":4,\"total_tokens\":6}}\n\ndata: [DONE]\n\n"
|
||||
|
||||
@@ -1,13 +1,14 @@
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override,
|
||||
encrypt_python_fernet_plaintext, hash_api_key, json, start_server, Arc, Body, GatewayDataState,
|
||||
HeaderValue, InMemoryAuthApiKeySnapshotRepository,
|
||||
InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository,
|
||||
InMemoryRequestCandidateRepository, InMemoryUsageReadRepository, Json, Request,
|
||||
RequestCandidateReadRepository, RequestCandidateStatus, Response, Router, StatusCode,
|
||||
StoredAuthApiKeySnapshot, StoredMinimalCandidateSelectionRow, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider, StoredProviderModelMapping,
|
||||
UsageReadRepository, UsageRuntimeConfig, DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
|
||||
encrypt_python_fernet_plaintext, hash_api_key, json, start_server,
|
||||
strip_sse_keepalive_comments, Arc, Body, GatewayDataState, HeaderValue,
|
||||
InMemoryAuthApiKeySnapshotRepository, InMemoryMinimalCandidateSelectionReadRepository,
|
||||
InMemoryProviderCatalogReadRepository, InMemoryRequestCandidateRepository,
|
||||
InMemoryUsageReadRepository, Json, Request, RequestCandidateReadRepository,
|
||||
RequestCandidateStatus, Response, Router, StatusCode, StoredAuthApiKeySnapshot,
|
||||
StoredMinimalCandidateSelectionRow, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider, StoredProviderModelMapping, UsageReadRepository,
|
||||
UsageRuntimeConfig, DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
|
||||
};
|
||||
use aether_data::repository::billing::InMemoryBillingReadRepository;
|
||||
use aether_data::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
|
||||
@@ -939,7 +940,7 @@ async fn gateway_records_openai_stream_usage_and_pricing_with_cache_tokens_impl(
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.text().await.expect("stream body should read"),
|
||||
strip_sse_keepalive_comments(&response.text().await.expect("stream body should read")),
|
||||
stream_body.concat()
|
||||
);
|
||||
|
||||
@@ -1132,7 +1133,7 @@ async fn gateway_records_claude_stream_usage_and_pricing_with_cache_breakdown_im
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.text().await.expect("stream body should read"),
|
||||
strip_sse_keepalive_comments(&response.text().await.expect("stream body should read")),
|
||||
stream_body.concat()
|
||||
);
|
||||
|
||||
@@ -1310,7 +1311,7 @@ async fn gateway_records_gemini_stream_usage_and_pricing_with_cache_read_tokens_
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.text().await.expect("stream body should read"),
|
||||
strip_sse_keepalive_comments(&response.text().await.expect("stream body should read")),
|
||||
stream_body.concat()
|
||||
);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user