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:
@@ -30,6 +30,7 @@ pub(crate) use crate::ai_serving::{
|
||||
maybe_build_stream_decision_payload, maybe_build_stream_plan_payload,
|
||||
maybe_build_sync_decision_payload, maybe_build_sync_plan_payload,
|
||||
set_local_openai_chat_execution_exhausted_diagnostic,
|
||||
set_local_openai_image_execution_exhausted_diagnostic,
|
||||
};
|
||||
pub(crate) use crate::ai_serving::{
|
||||
maybe_bridge_standard_sync_json_to_stream, maybe_build_provider_private_stream_normalizer,
|
||||
|
||||
@@ -49,7 +49,8 @@ pub(crate) use self::planner::{
|
||||
extract_pool_sticky_session_token, maybe_build_stream_decision_payload,
|
||||
maybe_build_stream_plan_payload, maybe_build_sync_decision_payload,
|
||||
maybe_build_sync_plan_payload, planner_is_matching_stream_request,
|
||||
set_local_openai_chat_execution_exhausted_diagnostic, CandidateFailureDiagnostic,
|
||||
set_local_openai_chat_execution_exhausted_diagnostic,
|
||||
set_local_openai_image_execution_exhausted_diagnostic, CandidateFailureDiagnostic,
|
||||
CandidateFailureDiagnosticKind, GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot,
|
||||
LocalExecutionAttemptSource, LocalResolvedOAuthRequestAuth, PlannerAppState,
|
||||
};
|
||||
|
||||
@@ -48,6 +48,7 @@ pub(crate) use self::specialized::{
|
||||
build_local_image_sync_plan_and_reports_for_kind,
|
||||
build_local_video_sync_attempt_source_for_kind,
|
||||
build_local_video_sync_plan_and_reports_for_kind,
|
||||
set_local_openai_image_execution_exhausted_diagnostic,
|
||||
};
|
||||
pub(crate) use self::standard::{
|
||||
build_local_openai_chat_stream_attempt_source_for_kind,
|
||||
|
||||
@@ -11,6 +11,7 @@ use crate::ai_serving::planner::plan_builders::{
|
||||
build_passthrough_sync_plan_from_decision, build_standard_stream_plan_from_decision,
|
||||
AiStreamAttempt, AiSyncAttempt,
|
||||
};
|
||||
use crate::ai_serving::planner::runtime_miss::set_local_runtime_execution_exhausted_diagnostic;
|
||||
use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata;
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::ai_serving::{
|
||||
@@ -50,6 +51,35 @@ pub(crate) struct LocalOpenAiImageStreamAttemptSource<'a> {
|
||||
candidates: LocalOpenAiImageCandidateAttemptSource<'a>,
|
||||
}
|
||||
|
||||
pub(crate) fn set_local_openai_image_execution_exhausted_diagnostic(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
body_json: &serde_json::Value,
|
||||
candidate_count: usize,
|
||||
) {
|
||||
warn!(
|
||||
event_name = "local_openai_image_candidates_exhausted",
|
||||
log_type = "event",
|
||||
trace_id = %trace_id,
|
||||
plan_kind,
|
||||
route_class = decision.route_class.as_deref().unwrap_or("passthrough"),
|
||||
route_family = decision.route_family.as_deref().unwrap_or("unknown"),
|
||||
candidate_count,
|
||||
model = body_json.get("model").and_then(|value| value.as_str()).unwrap_or(""),
|
||||
"gateway local openai image execution exhausted all candidates"
|
||||
);
|
||||
set_local_runtime_execution_exhausted_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
body_json.get("model").and_then(|value| value.as_str()),
|
||||
candidate_count,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_image_sync_plan_and_reports_for_kind(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
|
||||
@@ -18,6 +18,7 @@ pub(crate) use self::image::{
|
||||
build_local_image_sync_attempt_source_for_kind,
|
||||
build_local_image_sync_plan_and_reports_for_kind,
|
||||
maybe_build_stream_local_image_decision_payload, maybe_build_sync_local_image_decision_payload,
|
||||
set_local_openai_image_execution_exhausted_diagnostic,
|
||||
};
|
||||
pub(crate) use self::video::{
|
||||
build_local_video_sync_attempt_source_for_kind,
|
||||
|
||||
@@ -271,6 +271,8 @@ pub(crate) fn build_local_auth_rejection_response(
|
||||
control_decision: Option<&GatewayControlDecision>,
|
||||
rejection: &GatewayLocalAuthRejection,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
const ACCESS_POLICY_SUBJECT: &str = "当前用户、用户组或密钥的访问控制策略";
|
||||
|
||||
match rejection {
|
||||
GatewayLocalAuthRejection::InvalidApiKey => build_local_http_error_response(
|
||||
trace_id,
|
||||
@@ -298,7 +300,7 @@ pub(crate) fn build_local_auth_rejection_response(
|
||||
trace_id,
|
||||
control_decision,
|
||||
StatusCode::FORBIDDEN,
|
||||
&format!("当前密钥不允许访问 {provider} 提供商"),
|
||||
&format!("{ACCESS_POLICY_SUBJECT}不允许访问 {provider} 提供商"),
|
||||
)
|
||||
}
|
||||
GatewayLocalAuthRejection::ApiFormatNotAllowed { api_format } => {
|
||||
@@ -306,14 +308,14 @@ pub(crate) fn build_local_auth_rejection_response(
|
||||
trace_id,
|
||||
control_decision,
|
||||
StatusCode::FORBIDDEN,
|
||||
&format!("当前密钥不允许访问 {api_format} 格式"),
|
||||
&format!("{ACCESS_POLICY_SUBJECT}不允许访问 {api_format} 格式"),
|
||||
)
|
||||
}
|
||||
GatewayLocalAuthRejection::ModelNotAllowed { model } => build_local_http_error_response(
|
||||
trace_id,
|
||||
control_decision,
|
||||
StatusCode::FORBIDDEN,
|
||||
&format!("当前密钥不允许访问模型 {model}"),
|
||||
&format!("{ACCESS_POLICY_SUBJECT}不允许访问模型 {model}"),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -61,6 +61,8 @@ pub(crate) const TRUSTED_ADMIN_SESSION_ID_HEADER: &str = "x-aether-admin-session
|
||||
pub(crate) const TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER: &str =
|
||||
"x-aether-admin-management-token-id";
|
||||
pub(crate) const TRUSTED_RATE_LIMIT_PREFLIGHT_HEADER: &str = "x-aether-rate-limit-preflight";
|
||||
pub(crate) const DEFAULT_USER_GROUP_CONFIG_KEY: &str = "default_user_group_id";
|
||||
pub(crate) const BUILTIN_DEFAULT_USER_GROUP_ID: &str = "00000000-0000-0000-0000-000000000001";
|
||||
|
||||
pub(crate) const FRONTDOOR_REPLACEABLE_ROUTE_GROUPS: &[&str] = &["frontdoor_compat_router"];
|
||||
pub(crate) const FRONTDOOR_REPLACEABLE_MIDDLEWARE_GROUPS: &[&str] = &["cors"];
|
||||
|
||||
@@ -575,6 +575,95 @@ pub(super) fn classify_admin_operations_family_route(
|
||||
"admin:wallets",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::GET
|
||||
&& matches!(
|
||||
normalized_path,
|
||||
"/api/admin/user-groups" | "/api/admin/user-groups/"
|
||||
)
|
||||
{
|
||||
Some(classified(
|
||||
"admin_proxy",
|
||||
"users_manage",
|
||||
"list_user_groups",
|
||||
"admin:users",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::POST
|
||||
&& matches!(
|
||||
normalized_path,
|
||||
"/api/admin/user-groups" | "/api/admin/user-groups/"
|
||||
)
|
||||
{
|
||||
Some(classified(
|
||||
"admin_proxy",
|
||||
"users_manage",
|
||||
"create_user_group",
|
||||
"admin:users",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::PUT
|
||||
&& matches!(
|
||||
normalized_path,
|
||||
"/api/admin/user-groups/default" | "/api/admin/user-groups/default/"
|
||||
)
|
||||
{
|
||||
Some(classified(
|
||||
"admin_proxy",
|
||||
"users_manage",
|
||||
"set_default_user_group",
|
||||
"admin:users",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::GET
|
||||
&& normalized_path.starts_with("/api/admin/user-groups/")
|
||||
&& normalized_path.ends_with("/members")
|
||||
&& normalized_path.matches('/').count() == 5
|
||||
{
|
||||
Some(classified(
|
||||
"admin_proxy",
|
||||
"users_manage",
|
||||
"list_user_group_members",
|
||||
"admin:users",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::PUT
|
||||
&& normalized_path.starts_with("/api/admin/user-groups/")
|
||||
&& normalized_path.ends_with("/members")
|
||||
&& normalized_path.matches('/').count() == 5
|
||||
{
|
||||
Some(classified(
|
||||
"admin_proxy",
|
||||
"users_manage",
|
||||
"replace_user_group_members",
|
||||
"admin:users",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::PUT
|
||||
&& normalized_path.starts_with("/api/admin/user-groups/")
|
||||
&& normalized_path.matches('/').count() == 4
|
||||
&& !normalized_path.ends_with("/default")
|
||||
&& !normalized_path.ends_with("/members")
|
||||
{
|
||||
Some(classified(
|
||||
"admin_proxy",
|
||||
"users_manage",
|
||||
"update_user_group",
|
||||
"admin:users",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::DELETE
|
||||
&& normalized_path.starts_with("/api/admin/user-groups/")
|
||||
&& normalized_path.matches('/').count() == 4
|
||||
&& !normalized_path.ends_with("/default")
|
||||
&& !normalized_path.ends_with("/members")
|
||||
{
|
||||
Some(classified(
|
||||
"admin_proxy",
|
||||
"users_manage",
|
||||
"delete_user_group",
|
||||
"admin:users",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::GET
|
||||
&& matches!(normalized_path, "/api/admin/users" | "/api/admin/users/")
|
||||
{
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
use http::Uri;
|
||||
|
||||
use super::{classify_control_route, headers};
|
||||
use crate::handlers::shared::local_proxy_route_requires_buffered_body;
|
||||
|
||||
use super::{classify_control_route, headers, GatewayPublicRequestContext};
|
||||
|
||||
#[test]
|
||||
fn classifies_admin_users_list_as_admin_proxy_route() {
|
||||
@@ -68,6 +70,86 @@ fn classifies_admin_user_batch_routes_as_admin_proxy_route() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classifies_admin_user_group_routes_as_admin_proxy_route() {
|
||||
let headers = headers(&[]);
|
||||
|
||||
let list_uri: Uri = "/api/admin/user-groups".parse().expect("uri should parse");
|
||||
let list = classify_control_route(&http::Method::GET, &list_uri, &headers)
|
||||
.expect("route should classify");
|
||||
assert_eq!(list.route_family.as_deref(), Some("users_manage"));
|
||||
assert_eq!(list.route_kind.as_deref(), Some("list_user_groups"));
|
||||
|
||||
let create_uri: Uri = "/api/admin/user-groups".parse().expect("uri should parse");
|
||||
let create = classify_control_route(&http::Method::POST, &create_uri, &headers)
|
||||
.expect("route should classify");
|
||||
assert_eq!(create.route_family.as_deref(), Some("users_manage"));
|
||||
assert_eq!(create.route_kind.as_deref(), Some("create_user_group"));
|
||||
|
||||
let update_uri: Uri = "/api/admin/user-groups/group-1"
|
||||
.parse()
|
||||
.expect("uri should parse");
|
||||
let update = classify_control_route(&http::Method::PUT, &update_uri, &headers)
|
||||
.expect("route should classify");
|
||||
assert_eq!(update.route_family.as_deref(), Some("users_manage"));
|
||||
assert_eq!(update.route_kind.as_deref(), Some("update_user_group"));
|
||||
|
||||
let members_uri: Uri = "/api/admin/user-groups/group-1/members"
|
||||
.parse()
|
||||
.expect("uri should parse");
|
||||
let members = classify_control_route(&http::Method::PUT, &members_uri, &headers)
|
||||
.expect("route should classify");
|
||||
assert_eq!(members.route_family.as_deref(), Some("users_manage"));
|
||||
assert_eq!(
|
||||
members.route_kind.as_deref(),
|
||||
Some("replace_user_group_members")
|
||||
);
|
||||
|
||||
let default_uri: Uri = "/api/admin/user-groups/default"
|
||||
.parse()
|
||||
.expect("uri should parse");
|
||||
let default = classify_control_route(&http::Method::PUT, &default_uri, &headers)
|
||||
.expect("route should classify");
|
||||
assert_eq!(default.route_family.as_deref(), Some("users_manage"));
|
||||
assert_eq!(
|
||||
default.route_kind.as_deref(),
|
||||
Some("set_default_user_group")
|
||||
);
|
||||
assert_eq!(
|
||||
default.auth_endpoint_signature.as_deref(),
|
||||
Some("admin:users")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn admin_user_group_write_routes_buffer_request_body() {
|
||||
let headers = headers(&[]);
|
||||
let routes = [
|
||||
(http::Method::POST, "/api/admin/user-groups"),
|
||||
(http::Method::PUT, "/api/admin/user-groups/group-1"),
|
||||
(http::Method::PUT, "/api/admin/user-groups/group-1/members"),
|
||||
(http::Method::PUT, "/api/admin/user-groups/default"),
|
||||
];
|
||||
|
||||
for (method, path) in routes {
|
||||
let uri: Uri = path.parse().expect("uri should parse");
|
||||
let decision =
|
||||
classify_control_route(&method, &uri, &headers).expect("route should classify");
|
||||
let context = GatewayPublicRequestContext::from_request_parts(
|
||||
"trace-user-group-write",
|
||||
&method,
|
||||
&uri,
|
||||
&headers,
|
||||
Some(decision),
|
||||
);
|
||||
|
||||
assert!(
|
||||
local_proxy_route_requires_buffered_body(&context),
|
||||
"{method} {path} should buffer request body"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classifies_admin_user_detail_routes_as_admin_proxy_route() {
|
||||
let headers = headers(&[]);
|
||||
|
||||
@@ -86,6 +86,139 @@ impl GatewayDataState {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_groups(
|
||||
&self,
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroup>, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => repository.list_user_groups().await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn find_user_group_by_id(
|
||||
&self,
|
||||
group_id: &str,
|
||||
) -> Result<Option<aether_data::repository::users::StoredUserGroup>, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => repository.find_user_group_by_id(group_id).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_groups_by_ids(
|
||||
&self,
|
||||
group_ids: &[String],
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroup>, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => repository.list_user_groups_by_ids(group_ids).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn create_user_group(
|
||||
&self,
|
||||
record: aether_data::repository::users::UpsertUserGroupRecord,
|
||||
) -> Result<Option<aether_data::repository::users::StoredUserGroup>, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => repository.create_user_group(record).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn update_user_group(
|
||||
&self,
|
||||
group_id: &str,
|
||||
record: aether_data::repository::users::UpsertUserGroupRecord,
|
||||
) -> Result<Option<aether_data::repository::users::StoredUserGroup>, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => repository.update_user_group(group_id, record).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_user_group(&self, group_id: &str) -> Result<bool, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => repository.delete_user_group(group_id).await,
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_group_members(
|
||||
&self,
|
||||
group_id: &str,
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroupMember>, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => repository.list_user_group_members(group_id).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn replace_user_group_members(
|
||||
&self,
|
||||
group_id: &str,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroupMember>, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.replace_user_group_members(group_id, user_ids)
|
||||
.await
|
||||
}
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_groups_for_user(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroup>, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => repository.list_user_groups_for_user(user_id).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_group_memberships_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroupMembership>, DataLayerError>
|
||||
{
|
||||
match &self.user_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.list_user_group_memberships_by_user_ids(user_ids)
|
||||
.await
|
||||
}
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn replace_user_groups_for_user(
|
||||
&self,
|
||||
user_id: &str,
|
||||
group_ids: &[String],
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroup>, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.replace_user_groups_for_user(user_id, group_ids)
|
||||
.await
|
||||
}
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn add_user_to_group(
|
||||
&self,
|
||||
group_id: &str,
|
||||
user_id: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
match &self.user_reader {
|
||||
Some(repository) => repository.add_user_to_group(group_id, user_id).await,
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_oauth_links(
|
||||
&self,
|
||||
user_id: &str,
|
||||
@@ -424,6 +557,28 @@ impl GatewayDataState {
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn update_local_auth_user_policy_modes(
|
||||
&self,
|
||||
user_id: &str,
|
||||
allowed_providers_mode: Option<String>,
|
||||
allowed_api_formats_mode: Option<String>,
|
||||
allowed_models_mode: Option<String>,
|
||||
rate_limit_mode: Option<String>,
|
||||
) -> Result<Option<StoredUserAuthRecord>, DataLayerError> {
|
||||
let Some(repository) = self.user_reader.as_ref() else {
|
||||
return Ok(None);
|
||||
};
|
||||
repository
|
||||
.update_local_auth_user_policy_modes(
|
||||
user_id,
|
||||
allowed_providers_mode,
|
||||
allowed_api_formats_mode,
|
||||
allowed_models_mode,
|
||||
rate_limit_mode,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn touch_auth_user_last_login(
|
||||
&self,
|
||||
user_id: &str,
|
||||
@@ -1470,13 +1625,14 @@ impl GatewayDataState {
|
||||
api_key_id: &str,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Option<GatewayAuthApiKeySnapshot>, DataLayerError> {
|
||||
read_resolved_auth_api_key_snapshot_by_user_api_key_ids(
|
||||
let snapshot = read_resolved_auth_api_key_snapshot_by_user_api_key_ids(
|
||||
self,
|
||||
user_id,
|
||||
api_key_id,
|
||||
now_unix_secs,
|
||||
)
|
||||
.await
|
||||
.await?;
|
||||
self.apply_user_group_effective_policies(snapshot).await
|
||||
}
|
||||
|
||||
pub(crate) async fn read_auth_api_key_snapshot_by_key_hash(
|
||||
@@ -1484,18 +1640,242 @@ impl GatewayDataState {
|
||||
key_hash: &str,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Option<GatewayAuthApiKeySnapshot>, DataLayerError> {
|
||||
read_resolved_auth_api_key_snapshot_by_key_hash(self, key_hash, now_unix_secs).await
|
||||
let snapshot =
|
||||
read_resolved_auth_api_key_snapshot_by_key_hash(self, key_hash, now_unix_secs).await?;
|
||||
self.apply_user_group_effective_policies(snapshot).await
|
||||
}
|
||||
|
||||
async fn apply_user_group_effective_policies(
|
||||
&self,
|
||||
snapshot: Option<GatewayAuthApiKeySnapshot>,
|
||||
) -> Result<Option<GatewayAuthApiKeySnapshot>, DataLayerError> {
|
||||
let Some(mut snapshot) = snapshot else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(repository) = self.user_reader.as_ref() else {
|
||||
return Ok(Some(snapshot));
|
||||
};
|
||||
let Some(user) = repository.find_user_auth_by_id(&snapshot.user_id).await? else {
|
||||
return Ok(Some(snapshot));
|
||||
};
|
||||
let export_row = repository.find_export_user_by_id(&snapshot.user_id).await?;
|
||||
let mut groups = repository
|
||||
.list_user_groups_for_user(&snapshot.user_id)
|
||||
.await?;
|
||||
groups.sort_by(|left, right| {
|
||||
left.name
|
||||
.cmp(&right.name)
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
|
||||
let mut allowed_providers = resolve_effective_list_policy(
|
||||
user.allowed_providers,
|
||||
&user.allowed_providers_mode,
|
||||
&groups,
|
||||
|group| {
|
||||
(
|
||||
&group.allowed_providers_mode,
|
||||
group.allowed_providers.clone(),
|
||||
)
|
||||
},
|
||||
);
|
||||
let mut allowed_api_formats = resolve_effective_list_policy(
|
||||
user.allowed_api_formats,
|
||||
&user.allowed_api_formats_mode,
|
||||
&groups,
|
||||
|group| {
|
||||
(
|
||||
&group.allowed_api_formats_mode,
|
||||
group.allowed_api_formats.clone(),
|
||||
)
|
||||
},
|
||||
);
|
||||
let mut allowed_models = resolve_effective_list_policy(
|
||||
user.allowed_models,
|
||||
&user.allowed_models_mode,
|
||||
&groups,
|
||||
|group| (&group.allowed_models_mode, group.allowed_models.clone()),
|
||||
);
|
||||
let snapshot_user_rate_limit = snapshot.user_rate_limit;
|
||||
let export_user_rate_limit = export_row.as_ref().and_then(|row| row.rate_limit);
|
||||
let user_rate_limit_mode = match export_row.as_ref() {
|
||||
Some(row)
|
||||
if row.rate_limit.is_none()
|
||||
&& row.rate_limit_mode == "system"
|
||||
&& snapshot_user_rate_limit.is_some() =>
|
||||
{
|
||||
"custom"
|
||||
}
|
||||
Some(row) => row.rate_limit_mode.as_str(),
|
||||
None if snapshot_user_rate_limit.is_some() => "custom",
|
||||
None => "system",
|
||||
};
|
||||
let user_rate_limit = resolve_effective_rate_limit_policy(
|
||||
export_user_rate_limit.or(snapshot_user_rate_limit),
|
||||
user_rate_limit_mode,
|
||||
&groups,
|
||||
);
|
||||
if !snapshot.api_key_is_standalone {
|
||||
constrain_api_key_list_policy_to_user_policy(
|
||||
&mut allowed_providers,
|
||||
&mut snapshot.api_key_allowed_providers,
|
||||
);
|
||||
constrain_api_key_list_policy_to_user_policy(
|
||||
&mut allowed_api_formats,
|
||||
&mut snapshot.api_key_allowed_api_formats,
|
||||
);
|
||||
constrain_api_key_list_policy_to_user_policy(
|
||||
&mut allowed_models,
|
||||
&mut snapshot.api_key_allowed_models,
|
||||
);
|
||||
}
|
||||
snapshot.apply_user_policy(
|
||||
allowed_providers,
|
||||
allowed_api_formats,
|
||||
allowed_models,
|
||||
user_rate_limit,
|
||||
);
|
||||
Ok(Some(snapshot))
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_effective_list_policy(
|
||||
user_values: Option<Vec<String>>,
|
||||
user_mode: &str,
|
||||
groups: &[aether_data::repository::users::StoredUserGroup],
|
||||
group_field: impl Fn(
|
||||
&aether_data::repository::users::StoredUserGroup,
|
||||
) -> (&str, Option<Vec<String>>),
|
||||
) -> Option<Vec<String>> {
|
||||
let group_policy = groups.iter().fold(None, |effective, group| {
|
||||
let (mode, values) = group_field(group);
|
||||
intersect_list_policies(effective, list_restriction_from_mode(mode, values))
|
||||
});
|
||||
let user_policy = list_restriction_from_mode(user_mode, user_values);
|
||||
intersect_list_policies(group_policy, user_policy)
|
||||
}
|
||||
|
||||
fn list_restriction_from_mode(mode: &str, values: Option<Vec<String>>) -> Option<Vec<String>> {
|
||||
match mode {
|
||||
"specific" => Some(values.unwrap_or_default()),
|
||||
"deny_all" => Some(Vec::new()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_effective_rate_limit_policy(
|
||||
user_rate_limit: Option<i32>,
|
||||
user_mode: &str,
|
||||
groups: &[aether_data::repository::users::StoredUserGroup],
|
||||
) -> Option<i32> {
|
||||
let group_policy = groups.iter().fold(None, |effective, group| {
|
||||
intersect_rate_limit_policies(
|
||||
effective,
|
||||
rate_limit_restriction_from_mode(&group.rate_limit_mode, group.rate_limit),
|
||||
)
|
||||
});
|
||||
let user_policy = rate_limit_restriction_from_mode(user_mode, user_rate_limit);
|
||||
rate_limit_policy_value(intersect_rate_limit_policies(group_policy, user_policy))
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
enum RateLimitRestriction {
|
||||
Unlimited,
|
||||
Limited(i32),
|
||||
}
|
||||
|
||||
fn rate_limit_restriction_from_mode(
|
||||
mode: &str,
|
||||
rate_limit: Option<i32>,
|
||||
) -> Option<RateLimitRestriction> {
|
||||
match mode {
|
||||
"custom" => {
|
||||
let rate_limit = rate_limit.unwrap_or(0).max(0);
|
||||
if rate_limit == 0 {
|
||||
Some(RateLimitRestriction::Unlimited)
|
||||
} else {
|
||||
Some(RateLimitRestriction::Limited(rate_limit))
|
||||
}
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn intersect_list_policies(
|
||||
left: Option<Vec<String>>,
|
||||
right: Option<Vec<String>>,
|
||||
) -> Option<Vec<String>> {
|
||||
match (left, right) {
|
||||
(None, None) => None,
|
||||
(Some(values), None) | (None, Some(values)) => Some(values),
|
||||
(Some(left_values), Some(right_values)) => {
|
||||
let right_values = right_values
|
||||
.into_iter()
|
||||
.collect::<std::collections::BTreeSet<_>>();
|
||||
Some(
|
||||
left_values
|
||||
.into_iter()
|
||||
.filter(|value| right_values.contains(value))
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn intersect_rate_limit_policies(
|
||||
left: Option<RateLimitRestriction>,
|
||||
right: Option<RateLimitRestriction>,
|
||||
) -> Option<RateLimitRestriction> {
|
||||
match (left, right) {
|
||||
(None, None) => None,
|
||||
(Some(value), None) | (None, Some(value)) => Some(value),
|
||||
(Some(RateLimitRestriction::Unlimited), Some(RateLimitRestriction::Unlimited)) => {
|
||||
Some(RateLimitRestriction::Unlimited)
|
||||
}
|
||||
(Some(RateLimitRestriction::Limited(value)), Some(RateLimitRestriction::Unlimited))
|
||||
| (Some(RateLimitRestriction::Unlimited), Some(RateLimitRestriction::Limited(value))) => {
|
||||
Some(RateLimitRestriction::Limited(value))
|
||||
}
|
||||
(Some(RateLimitRestriction::Limited(left)), Some(RateLimitRestriction::Limited(right))) => {
|
||||
Some(RateLimitRestriction::Limited(left.min(right)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn rate_limit_policy_value(policy: Option<RateLimitRestriction>) -> Option<i32> {
|
||||
match policy {
|
||||
None => None,
|
||||
Some(RateLimitRestriction::Unlimited) => Some(0),
|
||||
Some(RateLimitRestriction::Limited(value)) => Some(value),
|
||||
}
|
||||
}
|
||||
|
||||
fn constrain_api_key_list_policy_to_user_policy(
|
||||
user_policy: &mut Option<Vec<String>>,
|
||||
api_key_policy: &mut Option<Vec<String>>,
|
||||
) {
|
||||
let Some(api_key_values) = api_key_policy.as_ref().filter(|values| !values.is_empty()) else {
|
||||
return;
|
||||
};
|
||||
let Some(user_values) = user_policy.clone() else {
|
||||
return;
|
||||
};
|
||||
let effective = intersect_list_policies(Some(api_key_values.to_vec()), Some(user_values))
|
||||
.unwrap_or_default();
|
||||
*user_policy = Some(effective.clone());
|
||||
*api_key_policy = Some(effective);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::*;
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeyExportRecord,
|
||||
StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::users::StoredUserGroup;
|
||||
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
@@ -1526,6 +1906,125 @@ mod tests {
|
||||
.expect("snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_group(
|
||||
id: &str,
|
||||
priority: i32,
|
||||
allowed_models: Option<Vec<&str>>,
|
||||
allowed_models_mode: &str,
|
||||
rate_limit: Option<i32>,
|
||||
rate_limit_mode: &str,
|
||||
) -> StoredUserGroup {
|
||||
StoredUserGroup {
|
||||
id: id.to_string(),
|
||||
name: id.to_string(),
|
||||
normalized_name: id.to_string(),
|
||||
description: None,
|
||||
priority,
|
||||
allowed_providers: None,
|
||||
allowed_providers_mode: "unrestricted".to_string(),
|
||||
allowed_api_formats: None,
|
||||
allowed_api_formats_mode: "unrestricted".to_string(),
|
||||
allowed_models: allowed_models.map(|values| {
|
||||
values
|
||||
.into_iter()
|
||||
.map(ToOwned::to_owned)
|
||||
.collect::<Vec<_>>()
|
||||
}),
|
||||
allowed_models_mode: allowed_models_mode.to_string(),
|
||||
rate_limit,
|
||||
rate_limit_mode: rate_limit_mode.to_string(),
|
||||
created_at: None,
|
||||
updated_at: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_policy_intersects_group_and_user_restrictions() {
|
||||
let groups = vec![
|
||||
sample_group("default", 0, None, "unrestricted", None, "system"),
|
||||
sample_group(
|
||||
"restricted",
|
||||
10,
|
||||
Some(vec!["gpt-5", "gpt-4.1"]),
|
||||
"specific",
|
||||
None,
|
||||
"system",
|
||||
),
|
||||
];
|
||||
|
||||
let policy = resolve_effective_list_policy(
|
||||
Some(vec!["gpt-4.1".to_string(), "gemini-2.5-pro".to_string()]),
|
||||
"specific",
|
||||
&groups,
|
||||
|group| (&group.allowed_models_mode, group.allowed_models.clone()),
|
||||
);
|
||||
|
||||
assert_eq!(policy, Some(vec!["gpt-4.1".to_string()]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_unrestricted_does_not_bypass_group_restrictions() {
|
||||
let groups = vec![sample_group(
|
||||
"restricted",
|
||||
10,
|
||||
Some(vec!["gpt-5"]),
|
||||
"specific",
|
||||
None,
|
||||
"system",
|
||||
)];
|
||||
|
||||
let policy = resolve_effective_list_policy(None, "unrestricted", &groups, |group| {
|
||||
(&group.allowed_models_mode, group.allowed_models.clone())
|
||||
});
|
||||
|
||||
assert_eq!(policy, Some(vec!["gpt-5".to_string()]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rate_limit_policy_uses_most_restrictive_custom_limit() {
|
||||
let groups = vec![sample_group(
|
||||
"restricted",
|
||||
10,
|
||||
None,
|
||||
"unrestricted",
|
||||
Some(60),
|
||||
"custom",
|
||||
)];
|
||||
|
||||
assert_eq!(
|
||||
resolve_effective_rate_limit_policy(Some(120), "custom", &groups),
|
||||
Some(60)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rate_limit_unlimited_does_not_bypass_limited_group() {
|
||||
let groups = vec![sample_group(
|
||||
"restricted",
|
||||
10,
|
||||
None,
|
||||
"unrestricted",
|
||||
Some(60),
|
||||
"custom",
|
||||
)];
|
||||
|
||||
assert_eq!(
|
||||
resolve_effective_rate_limit_policy(Some(0), "custom", &groups),
|
||||
Some(60)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn api_key_specific_policy_cannot_expand_user_policy() {
|
||||
let mut user_policy = Some(vec!["gpt-5".to_string()]);
|
||||
let mut api_key_policy = Some(vec!["gpt-4.1".to_string()]);
|
||||
|
||||
constrain_api_key_list_policy_to_user_policy(&mut user_policy, &mut api_key_policy);
|
||||
|
||||
assert_eq!(user_policy, Some(Vec::<String>::new()));
|
||||
assert_eq!(api_key_policy, Some(Vec::<String>::new()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn data_state_lists_auth_api_key_export_records() {
|
||||
let repository = Arc::new(
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
use std::collections::{BTreeMap, VecDeque};
|
||||
use std::io::Error as IoError;
|
||||
use std::time::Instant;
|
||||
use std::sync::{
|
||||
atomic::{AtomicBool, AtomicU64, Ordering},
|
||||
Arc,
|
||||
};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use aether_contracts::{
|
||||
ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTelemetry, StreamFrame,
|
||||
@@ -25,6 +29,7 @@ use futures_util::stream::BoxStream;
|
||||
use futures_util::{StreamExt, TryStreamExt};
|
||||
use serde_json::{json, Value};
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::time::MissedTickBehavior;
|
||||
use tokio_util::codec::{FramedRead, LinesCodec};
|
||||
use tokio_util::io::StreamReader;
|
||||
use tracing::{debug, info, warn};
|
||||
@@ -94,6 +99,14 @@ use crate::{
|
||||
AppState, GatewayError, GEMINI_FILES_DOWNLOAD_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
};
|
||||
|
||||
const OPENAI_IMAGE_STREAM_PLAN_KIND: &str = "openai_image_stream";
|
||||
const SSE_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15);
|
||||
const SSE_KEEPALIVE_BYTES: &[u8] = b": aether-keepalive\n\n";
|
||||
const STREAM_IDLE_LOG_INTERVAL: Duration = Duration::from_secs(60);
|
||||
const STREAM_IDLE_LOG_INTERVAL_MS: u64 = 60_000;
|
||||
const REWRITTEN_STREAM_PREFETCH_TIMEOUT: Duration = Duration::from_millis(750);
|
||||
const OPENAI_IMAGE_STREAM_DEFAULT_TOTAL_TIMEOUT_MS: u64 = 900_000;
|
||||
|
||||
fn record_sync_terminal_usage(
|
||||
state: &AppState,
|
||||
plan: &ExecutionPlan,
|
||||
@@ -851,6 +864,114 @@ fn encode_terminal_sse_error_event(failure: &StreamFailureReport) -> Result<Byte
|
||||
Ok(Bytes::from(event))
|
||||
}
|
||||
|
||||
fn image_stream_failed_event_name(report_context: Option<&Value>) -> &'static str {
|
||||
let operation = report_context
|
||||
.and_then(|value| value.get("image_request"))
|
||||
.and_then(|value| value.get("operation"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
if operation == "edit" {
|
||||
"image_edit.failed"
|
||||
} else {
|
||||
"image_generation.failed"
|
||||
}
|
||||
}
|
||||
|
||||
fn encode_openai_image_failed_event(
|
||||
report_context: Option<&Value>,
|
||||
failure: &StreamFailureReport,
|
||||
) -> Result<Bytes, std::io::Error> {
|
||||
let event_name = image_stream_failed_event_name(report_context);
|
||||
let failure_body = failure
|
||||
.to_json_string()
|
||||
.map_err(|err| IoError::other(err.to_string()))?;
|
||||
let failure_json: Value =
|
||||
serde_json::from_str(&failure_body).map_err(|err| IoError::other(err.to_string()))?;
|
||||
let error = failure_json.get("error").cloned().unwrap_or_else(|| {
|
||||
serde_json::json!({
|
||||
"type": failure.error_type.as_str(),
|
||||
"message": failure.error_message.as_str(),
|
||||
"code": failure.status_code,
|
||||
})
|
||||
});
|
||||
let payload = serde_json::json!({
|
||||
"type": event_name,
|
||||
"error": error,
|
||||
});
|
||||
let payload = serde_json::to_string(&payload).map_err(|err| IoError::other(err.to_string()))?;
|
||||
let mut event = format!("event: {event_name}\n");
|
||||
for line in payload.lines() {
|
||||
event.push_str("data: ");
|
||||
event.push_str(line);
|
||||
event.push('\n');
|
||||
}
|
||||
event.push('\n');
|
||||
Ok(Bytes::from(event))
|
||||
}
|
||||
|
||||
fn resolve_openai_image_stream_total_timeout_ms(
|
||||
plan_kind: &str,
|
||||
plan: &ExecutionPlan,
|
||||
) -> Option<u64> {
|
||||
if plan_kind != OPENAI_IMAGE_STREAM_PLAN_KIND {
|
||||
return None;
|
||||
}
|
||||
Some(
|
||||
plan.timeouts
|
||||
.as_ref()
|
||||
.and_then(|timeouts| timeouts.total_ms)
|
||||
.unwrap_or(OPENAI_IMAGE_STREAM_DEFAULT_TOTAL_TIMEOUT_MS)
|
||||
.max(1),
|
||||
)
|
||||
}
|
||||
|
||||
fn should_limit_direct_finalize_prefetch(plan_kind: &str, has_local_stream_rewriter: bool) -> bool {
|
||||
plan_kind == OPENAI_IMAGE_STREAM_PLAN_KIND || has_local_stream_rewriter
|
||||
}
|
||||
|
||||
fn build_sse_body_stream(
|
||||
prefetched_chunks_for_body: Vec<Bytes>,
|
||||
mut rx: mpsc::Receiver<Result<Bytes, IoError>>,
|
||||
emit_keepalive: bool,
|
||||
keepalive_interval: Duration,
|
||||
) -> impl futures_util::Stream<Item = Result<Bytes, IoError>> + Send + 'static {
|
||||
stream! {
|
||||
let mut sent_prefetched_chunk = false;
|
||||
for chunk in prefetched_chunks_for_body {
|
||||
sent_prefetched_chunk = true;
|
||||
yield Ok(chunk);
|
||||
}
|
||||
|
||||
if emit_keepalive {
|
||||
if !sent_prefetched_chunk {
|
||||
yield Ok(Bytes::from_static(SSE_KEEPALIVE_BYTES));
|
||||
}
|
||||
let mut keepalive = tokio::time::interval(keepalive_interval);
|
||||
keepalive.set_missed_tick_behavior(MissedTickBehavior::Delay);
|
||||
keepalive.tick().await;
|
||||
loop {
|
||||
tokio::select! {
|
||||
biased;
|
||||
item = rx.recv() => {
|
||||
let Some(item) = item else {
|
||||
break;
|
||||
};
|
||||
yield item;
|
||||
}
|
||||
_ = keepalive.tick() => {
|
||||
yield Ok(Bytes::from_static(SSE_KEEPALIVE_BYTES));
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
while let Some(item) = rx.recv().await {
|
||||
yield item;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn next_stream_frame<R>(
|
||||
buffered_frames: &mut VecDeque<StreamFrame>,
|
||||
lines: &mut FramedRead<R, LinesCodec>,
|
||||
@@ -1347,6 +1468,8 @@ async fn execute_stream_from_frame_stream(
|
||||
private_stream_normalizer.is_some(),
|
||||
local_stream_rewriter.is_some(),
|
||||
);
|
||||
let limit_direct_finalize_prefetch =
|
||||
should_limit_direct_finalize_prefetch(plan_kind, local_stream_rewriter.is_some());
|
||||
let mut prefetched_chunks: Vec<Bytes> = Vec::new();
|
||||
let mut provider_prefetched_body = Vec::new();
|
||||
let mut prefetched_body = Vec::new();
|
||||
@@ -1380,7 +1503,38 @@ async fn execute_stream_from_frame_stream(
|
||||
while prefetched_chunks.len() < MAX_STREAM_PREFETCH_FRAMES
|
||||
&& prefetched_inspection_body.len() < MAX_STREAM_PREFETCH_BYTES
|
||||
{
|
||||
let Some(frame) = (match next_stream_frame(&mut buffered_frames, &mut lines).await {
|
||||
let next_frame_result = if limit_direct_finalize_prefetch {
|
||||
match tokio::time::timeout(
|
||||
REWRITTEN_STREAM_PREFETCH_TIMEOUT,
|
||||
next_stream_frame(&mut buffered_frames, &mut lines),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(result) => result,
|
||||
Err(_) => {
|
||||
debug!(
|
||||
event_name = "execution_runtime_stream_prefetch_limited",
|
||||
log_type = "debug",
|
||||
trace_id = %trace_id,
|
||||
request_id = %request_id_for_log,
|
||||
candidate_id = ?candidate_id,
|
||||
plan_kind,
|
||||
report_kind,
|
||||
provider_name,
|
||||
endpoint_id = %plan.endpoint_id,
|
||||
key_id = %plan.key_id,
|
||||
model_name,
|
||||
candidate_index = candidate_index.as_str(),
|
||||
timeout_ms = REWRITTEN_STREAM_PREFETCH_TIMEOUT.as_millis() as u64,
|
||||
"gateway stopped rewritten stream prefetch before client-visible body"
|
||||
);
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
next_stream_frame(&mut buffered_frames, &mut lines).await
|
||||
};
|
||||
let Some(frame) = (match next_frame_result {
|
||||
Ok(frame) => frame,
|
||||
Err(err) => {
|
||||
let failure = build_stream_failure_report(
|
||||
@@ -1705,7 +1859,6 @@ async fn execute_stream_from_frame_stream(
|
||||
let candidate_id = candidate_id.map(ToOwned::to_owned);
|
||||
let (tx, mut rx) = mpsc::channel::<Result<Bytes, IoError>>(16);
|
||||
let state_for_report = state.clone();
|
||||
let plan_for_report = plan;
|
||||
let trace_id_owned = trace_id.to_string();
|
||||
let headers_for_report = headers.clone();
|
||||
let report_kind_owned = report_kind;
|
||||
@@ -1723,8 +1876,14 @@ async fn execute_stream_from_frame_stream(
|
||||
let request_id_for_report = request_id.clone();
|
||||
let request_id_for_report_log = short_request_id(&request_id);
|
||||
let candidate_id_for_report = candidate_id.clone();
|
||||
let emit_passthrough_sse_terminal_error =
|
||||
skip_direct_finalize_prefetch && response_headers_indicate_sse(&upstream_headers);
|
||||
let candidate_index_for_report = candidate_index.clone();
|
||||
let is_openai_image_stream_for_report = plan_kind == OPENAI_IMAGE_STREAM_PLAN_KIND;
|
||||
let openai_image_stream_total_timeout_ms =
|
||||
resolve_openai_image_stream_total_timeout_ms(plan_kind, &plan);
|
||||
let plan_for_report = plan;
|
||||
let emit_passthrough_sse_terminal_error = skip_direct_finalize_prefetch
|
||||
&& response_headers_indicate_sse(&upstream_headers)
|
||||
&& !is_openai_image_stream_for_report;
|
||||
let body_capture_policy = match UsageRuntimeAccess::body_capture_policy(state.data.as_ref())
|
||||
.await
|
||||
{
|
||||
@@ -1753,6 +1912,8 @@ async fn execute_stream_from_frame_stream(
|
||||
.max_response_body_bytes
|
||||
.unwrap_or(DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES)
|
||||
};
|
||||
let plan_kind_for_report = plan_kind.to_string();
|
||||
let stream_started_at_for_report = stream_started_at;
|
||||
tokio::spawn(async move {
|
||||
let mut provider_buffered_body = Vec::new();
|
||||
let mut buffered_body = Vec::new();
|
||||
@@ -1797,6 +1958,115 @@ async fn execute_stream_from_frame_stream(
|
||||
let reached_eof = initial_reached_eof;
|
||||
let mut downstream_dropped = false;
|
||||
let mut terminal_failure: Option<StreamFailureReport> = None;
|
||||
let initial_elapsed_ms = stream_started_at_for_report
|
||||
.elapsed()
|
||||
.as_millis()
|
||||
.min(u128::from(u64::MAX)) as u64;
|
||||
let last_upstream_frame_elapsed_ms = Arc::new(AtomicU64::new(initial_elapsed_ms));
|
||||
let last_client_chunk_elapsed_ms =
|
||||
Arc::new(AtomicU64::new(if prefetched_body_for_report.is_empty() {
|
||||
0
|
||||
} else {
|
||||
initial_elapsed_ms
|
||||
}));
|
||||
let provider_stream_bytes = Arc::new(AtomicU64::new(
|
||||
u64::try_from(provider_prefetched_body_for_report.len()).unwrap_or(u64::MAX),
|
||||
));
|
||||
let client_stream_bytes = Arc::new(AtomicU64::new(
|
||||
u64::try_from(prefetched_body_for_report.len()).unwrap_or(u64::MAX),
|
||||
));
|
||||
let idle_monitor_done = Arc::new(AtomicBool::new(false));
|
||||
let idle_monitor_handle = {
|
||||
let done = Arc::clone(&idle_monitor_done);
|
||||
let last_upstream = Arc::clone(&last_upstream_frame_elapsed_ms);
|
||||
let last_client = Arc::clone(&last_client_chunk_elapsed_ms);
|
||||
let provider_bytes = Arc::clone(&provider_stream_bytes);
|
||||
let client_bytes = Arc::clone(&client_stream_bytes);
|
||||
let trace_id_for_idle = trace_id_owned.clone();
|
||||
let request_id_for_idle = request_id_for_report_log.clone();
|
||||
let candidate_id_for_idle = candidate_id_for_report.clone();
|
||||
let candidate_index_for_idle = candidate_index_for_report.clone();
|
||||
let plan_kind_for_idle = plan_kind_for_report.clone();
|
||||
let provider_name_for_idle = plan_for_report
|
||||
.provider_name
|
||||
.clone()
|
||||
.unwrap_or_else(|| "-".to_string());
|
||||
let endpoint_id_for_idle = plan_for_report.endpoint_id.clone();
|
||||
let key_id_for_idle = plan_for_report.key_id.clone();
|
||||
let model_name_for_idle = plan_for_report
|
||||
.model_name
|
||||
.clone()
|
||||
.unwrap_or_else(|| "-".to_string());
|
||||
let has_local_stream_rewriter_for_idle = local_stream_rewriter.is_some();
|
||||
tokio::spawn(async move {
|
||||
let mut interval = tokio::time::interval(STREAM_IDLE_LOG_INTERVAL);
|
||||
interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
|
||||
interval.tick().await;
|
||||
loop {
|
||||
interval.tick().await;
|
||||
if done.load(Ordering::Relaxed) {
|
||||
break;
|
||||
}
|
||||
let elapsed_ms = stream_started_at_for_report
|
||||
.elapsed()
|
||||
.as_millis()
|
||||
.min(u128::from(u64::MAX)) as u64;
|
||||
let last_upstream_frame_elapsed_ms = last_upstream.load(Ordering::Relaxed);
|
||||
let last_client_chunk_elapsed_ms = last_client.load(Ordering::Relaxed);
|
||||
let upstream_idle_ms =
|
||||
elapsed_ms.saturating_sub(last_upstream_frame_elapsed_ms);
|
||||
let client_idle_ms = if last_client_chunk_elapsed_ms == 0 {
|
||||
elapsed_ms
|
||||
} else {
|
||||
elapsed_ms.saturating_sub(last_client_chunk_elapsed_ms)
|
||||
};
|
||||
if upstream_idle_ms >= STREAM_IDLE_LOG_INTERVAL_MS {
|
||||
warn!(
|
||||
event_name = "stream_execution_upstream_idle",
|
||||
log_type = "ops",
|
||||
trace_id = %trace_id_for_idle,
|
||||
request_id = %request_id_for_idle,
|
||||
candidate_id = ?candidate_id_for_idle.as_deref(),
|
||||
candidate_index = candidate_index_for_idle.as_str(),
|
||||
plan_kind = plan_kind_for_idle.as_str(),
|
||||
provider_name = provider_name_for_idle.as_str(),
|
||||
endpoint_id = %endpoint_id_for_idle,
|
||||
key_id = %key_id_for_idle,
|
||||
model_name = model_name_for_idle.as_str(),
|
||||
elapsed_ms,
|
||||
provider_bytes = provider_bytes.load(Ordering::Relaxed),
|
||||
client_bytes = client_bytes.load(Ordering::Relaxed),
|
||||
last_upstream_frame_elapsed_ms,
|
||||
last_client_chunk_elapsed_ms,
|
||||
"gateway stream has not received an upstream frame within the idle window"
|
||||
);
|
||||
} else if client_idle_ms >= STREAM_IDLE_LOG_INTERVAL_MS
|
||||
&& last_upstream_frame_elapsed_ms >= last_client_chunk_elapsed_ms
|
||||
{
|
||||
warn!(
|
||||
event_name = "stream_execution_client_visible_idle",
|
||||
log_type = "ops",
|
||||
trace_id = %trace_id_for_idle,
|
||||
request_id = %request_id_for_idle,
|
||||
candidate_id = ?candidate_id_for_idle.as_deref(),
|
||||
candidate_index = candidate_index_for_idle.as_str(),
|
||||
plan_kind = plan_kind_for_idle.as_str(),
|
||||
provider_name = provider_name_for_idle.as_str(),
|
||||
endpoint_id = %endpoint_id_for_idle,
|
||||
key_id = %key_id_for_idle,
|
||||
model_name = model_name_for_idle.as_str(),
|
||||
elapsed_ms,
|
||||
provider_bytes = provider_bytes.load(Ordering::Relaxed),
|
||||
client_bytes = client_bytes.load(Ordering::Relaxed),
|
||||
last_upstream_frame_elapsed_ms,
|
||||
last_client_chunk_elapsed_ms,
|
||||
local_stream_rewriter = has_local_stream_rewriter_for_idle,
|
||||
"gateway stream received upstream frames but has no recent client-visible chunk"
|
||||
);
|
||||
}
|
||||
}
|
||||
})
|
||||
};
|
||||
if !provider_prefetched_body_for_report.is_empty() {
|
||||
let normalized_prefetched_chunk = if let Some(normalizer) =
|
||||
private_stream_normalizer.as_mut()
|
||||
@@ -1865,8 +2135,67 @@ async fn execute_stream_from_frame_stream(
|
||||
}
|
||||
|
||||
if terminal_failure.is_none() && !reached_eof {
|
||||
let mut image_stream_total_timeout = openai_image_stream_total_timeout_ms
|
||||
.map(|timeout_ms| Box::pin(tokio::time::sleep(Duration::from_millis(timeout_ms))));
|
||||
loop {
|
||||
let next_frame = match next_stream_frame(&mut buffered_frames, &mut lines).await {
|
||||
let next_frame_result = if let Some(timeout_sleep) =
|
||||
image_stream_total_timeout.as_mut()
|
||||
{
|
||||
tokio::select! {
|
||||
result = next_stream_frame(&mut buffered_frames, &mut lines) => result,
|
||||
_ = timeout_sleep.as_mut() => {
|
||||
let timeout_ms = openai_image_stream_total_timeout_ms
|
||||
.unwrap_or(OPENAI_IMAGE_STREAM_DEFAULT_TOTAL_TIMEOUT_MS);
|
||||
let elapsed_ms = stream_started_at_for_report
|
||||
.elapsed()
|
||||
.as_millis()
|
||||
.min(u128::from(u64::MAX)) as u64;
|
||||
warn!(
|
||||
event_name = "openai_image_stream_total_timeout",
|
||||
log_type = "ops",
|
||||
trace_id = %trace_id_owned,
|
||||
request_id = %request_id_for_report_log,
|
||||
candidate_id = ?candidate_id_for_report.as_deref(),
|
||||
candidate_index = candidate_index_for_report.as_str(),
|
||||
plan_kind = plan_kind_for_report.as_str(),
|
||||
provider_name = plan_for_report.provider_name.as_deref().unwrap_or("-"),
|
||||
endpoint_id = %plan_for_report.endpoint_id,
|
||||
key_id = %plan_for_report.key_id,
|
||||
model_name = plan_for_report.model_name.as_deref().unwrap_or("-"),
|
||||
elapsed_ms,
|
||||
timeout_ms,
|
||||
provider_bytes = provider_stream_bytes.load(Ordering::Relaxed),
|
||||
client_bytes = client_stream_bytes.load(Ordering::Relaxed),
|
||||
last_upstream_frame_elapsed_ms = last_upstream_frame_elapsed_ms.load(Ordering::Relaxed),
|
||||
last_client_chunk_elapsed_ms = last_client_chunk_elapsed_ms.load(Ordering::Relaxed),
|
||||
"gateway OpenAI image stream exceeded total timeout"
|
||||
);
|
||||
telemetry = Some(ExecutionTelemetry {
|
||||
ttfb_ms: telemetry
|
||||
.as_ref()
|
||||
.and_then(|telemetry| telemetry.ttfb_ms)
|
||||
.or_else(|| {
|
||||
usage_stream_telemetry
|
||||
.as_ref()
|
||||
.and_then(|telemetry| telemetry.ttfb_ms)
|
||||
}),
|
||||
elapsed_ms: Some(elapsed_ms),
|
||||
upstream_bytes: Some(provider_stream_bytes.load(Ordering::Relaxed)),
|
||||
});
|
||||
terminal_failure = Some(build_stream_failure_report(
|
||||
"image_stream_total_timeout",
|
||||
format!(
|
||||
"OpenAI image stream exceeded total timeout of {timeout_ms}ms"
|
||||
),
|
||||
504,
|
||||
));
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
next_stream_frame(&mut buffered_frames, &mut lines).await
|
||||
};
|
||||
let next_frame = match next_frame_result {
|
||||
Ok(frame) => frame,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
@@ -1889,6 +2218,11 @@ async fn execute_stream_from_frame_stream(
|
||||
let Some(frame) = next_frame else {
|
||||
break;
|
||||
};
|
||||
let frame_elapsed_ms = stream_started_at_for_report
|
||||
.elapsed()
|
||||
.as_millis()
|
||||
.min(u128::from(u64::MAX)) as u64;
|
||||
last_upstream_frame_elapsed_ms.store(frame_elapsed_ms, Ordering::Relaxed);
|
||||
match frame.payload {
|
||||
StreamFramePayload::Data { chunk_b64, text } => {
|
||||
if sync_json_stream_bridge_active_for_report {
|
||||
@@ -1922,6 +2256,10 @@ async fn execute_stream_from_frame_stream(
|
||||
continue;
|
||||
}
|
||||
|
||||
provider_stream_bytes.fetch_add(
|
||||
u64::try_from(chunk.len()).unwrap_or(u64::MAX),
|
||||
Ordering::Relaxed,
|
||||
);
|
||||
append_stream_capture_bytes(
|
||||
&mut provider_buffered_body,
|
||||
&chunk,
|
||||
@@ -2000,7 +2338,7 @@ async fn execute_stream_from_frame_stream(
|
||||
.and_then(|telemetry| telemetry.ttfb_ms)
|
||||
.is_none()
|
||||
{
|
||||
let first_data_elapsed_ms = stream_started_at
|
||||
let first_data_elapsed_ms = stream_started_at_for_report
|
||||
.elapsed()
|
||||
.as_millis()
|
||||
.min(u128::from(u64::MAX))
|
||||
@@ -2027,6 +2365,8 @@ async fn execute_stream_from_frame_stream(
|
||||
max_stream_body_buffer_bytes,
|
||||
&mut client_body_truncated,
|
||||
);
|
||||
let rewritten_chunk_len =
|
||||
u64::try_from(rewritten_chunk.len()).unwrap_or(u64::MAX);
|
||||
if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
|
||||
warn!(
|
||||
event_name = "stream_execution_downstream_disconnected",
|
||||
@@ -2038,6 +2378,16 @@ async fn execute_stream_from_frame_stream(
|
||||
);
|
||||
downstream_dropped = true;
|
||||
break;
|
||||
} else {
|
||||
client_stream_bytes.fetch_add(rewritten_chunk_len, Ordering::Relaxed);
|
||||
last_client_chunk_elapsed_ms.store(
|
||||
stream_started_at_for_report
|
||||
.elapsed()
|
||||
.as_millis()
|
||||
.min(u128::from(u64::MAX))
|
||||
as u64,
|
||||
Ordering::Relaxed,
|
||||
);
|
||||
}
|
||||
}
|
||||
StreamFramePayload::Telemetry {
|
||||
@@ -2140,16 +2490,29 @@ async fn execute_stream_from_frame_stream(
|
||||
max_stream_body_buffer_bytes,
|
||||
&mut client_body_truncated,
|
||||
);
|
||||
let rewritten_chunk_len =
|
||||
u64::try_from(rewritten_chunk.len()).unwrap_or(u64::MAX);
|
||||
if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
|
||||
warn!(
|
||||
event_name = "stream_execution_downstream_flush_disconnected",
|
||||
log_type = "ops",
|
||||
trace_id = %trace_id_owned,
|
||||
request_id = %request_id_for_report_log,
|
||||
candidate_id = ?candidate_id_for_report.as_deref(),
|
||||
"gateway stream downstream dropped while flushing private stream normalization"
|
||||
candidate_id = ?candidate_id_for_report.as_deref(),
|
||||
"gateway stream downstream dropped while flushing private stream normalization"
|
||||
);
|
||||
downstream_dropped = true;
|
||||
} else {
|
||||
client_stream_bytes
|
||||
.fetch_add(rewritten_chunk_len, Ordering::Relaxed);
|
||||
last_client_chunk_elapsed_ms.store(
|
||||
stream_started_at_for_report
|
||||
.elapsed()
|
||||
.as_millis()
|
||||
.min(u128::from(u64::MAX))
|
||||
as u64,
|
||||
Ordering::Relaxed,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -2184,16 +2547,28 @@ async fn execute_stream_from_frame_stream(
|
||||
max_stream_body_buffer_bytes,
|
||||
&mut client_body_truncated,
|
||||
);
|
||||
let flushed_chunk_len =
|
||||
u64::try_from(flushed_chunk.len()).unwrap_or(u64::MAX);
|
||||
if tx.send(Ok(Bytes::from(flushed_chunk))).await.is_err() {
|
||||
warn!(
|
||||
event_name = "stream_execution_downstream_rewrite_flush_disconnected",
|
||||
log_type = "ops",
|
||||
trace_id = %trace_id_owned,
|
||||
request_id = %request_id_for_report_log,
|
||||
candidate_id = ?candidate_id_for_report.as_deref(),
|
||||
"gateway stream downstream dropped while flushing local stream rewrite"
|
||||
candidate_id = ?candidate_id_for_report.as_deref(),
|
||||
"gateway stream downstream dropped while flushing local stream rewrite"
|
||||
);
|
||||
downstream_dropped = true;
|
||||
} else {
|
||||
client_stream_bytes.fetch_add(flushed_chunk_len, Ordering::Relaxed);
|
||||
last_client_chunk_elapsed_ms.store(
|
||||
stream_started_at_for_report
|
||||
.elapsed()
|
||||
.as_millis()
|
||||
.min(u128::from(u64::MAX))
|
||||
as u64,
|
||||
Ordering::Relaxed,
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(_) => {}
|
||||
@@ -2220,44 +2595,70 @@ async fn execute_stream_from_frame_stream(
|
||||
}
|
||||
}
|
||||
|
||||
if !downstream_dropped && emit_passthrough_sse_terminal_error {
|
||||
if !downstream_dropped {
|
||||
if let Some(failure) = terminal_failure.as_ref() {
|
||||
match encode_terminal_sse_error_event(failure) {
|
||||
Ok(error_event) => {
|
||||
append_stream_capture_bytes(
|
||||
&mut buffered_body,
|
||||
error_event.as_ref(),
|
||||
max_stream_body_buffer_bytes,
|
||||
&mut client_body_truncated,
|
||||
);
|
||||
if tx.send(Ok(error_event)).await.is_err() {
|
||||
warn!(
|
||||
let terminal_event = if is_openai_image_stream_for_report {
|
||||
Some(encode_openai_image_failed_event(
|
||||
report_context_owned.as_ref(),
|
||||
failure,
|
||||
))
|
||||
} else if emit_passthrough_sse_terminal_error {
|
||||
Some(encode_terminal_sse_error_event(failure))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if let Some(terminal_event) = terminal_event {
|
||||
match terminal_event {
|
||||
Ok(error_event) => {
|
||||
let error_event_len =
|
||||
u64::try_from(error_event.len()).unwrap_or(u64::MAX);
|
||||
append_stream_capture_bytes(
|
||||
&mut buffered_body,
|
||||
error_event.as_ref(),
|
||||
max_stream_body_buffer_bytes,
|
||||
&mut client_body_truncated,
|
||||
);
|
||||
if tx.send(Ok(error_event)).await.is_err() {
|
||||
warn!(
|
||||
event_name = "stream_execution_downstream_terminal_error_disconnected",
|
||||
log_type = "ops",
|
||||
trace_id = %trace_id_owned,
|
||||
request_id = %request_id_for_report_log,
|
||||
candidate_id = ?candidate_id_for_report.as_deref(),
|
||||
"gateway stream downstream dropped while sending terminal SSE error event"
|
||||
);
|
||||
downstream_dropped = true;
|
||||
);
|
||||
downstream_dropped = true;
|
||||
} else {
|
||||
client_stream_bytes.fetch_add(error_event_len, Ordering::Relaxed);
|
||||
last_client_chunk_elapsed_ms.store(
|
||||
stream_started_at_for_report
|
||||
.elapsed()
|
||||
.as_millis()
|
||||
.min(u128::from(u64::MAX))
|
||||
as u64,
|
||||
Ordering::Relaxed,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "stream_execution_terminal_error_event_encode_failed",
|
||||
log_type = "ops",
|
||||
trace_id = %trace_id_owned,
|
||||
request_id = %request_id_for_report_log,
|
||||
candidate_id = ?candidate_id_for_report.as_deref(),
|
||||
error = ?err,
|
||||
"gateway failed to encode terminal SSE error event"
|
||||
);
|
||||
"gateway failed to encode terminal SSE error event"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
drop(tx);
|
||||
idle_monitor_done.store(true, Ordering::Relaxed);
|
||||
idle_monitor_handle.abort();
|
||||
|
||||
stream_terminal_summary = merge_stream_terminal_summary(
|
||||
stream_terminal_summary,
|
||||
@@ -2422,15 +2823,6 @@ async fn execute_stream_from_frame_stream(
|
||||
}
|
||||
});
|
||||
|
||||
let body_stream = stream! {
|
||||
for chunk in prefetched_chunks_for_body {
|
||||
yield Ok(chunk);
|
||||
}
|
||||
while let Some(item) = rx.recv().await {
|
||||
yield item;
|
||||
}
|
||||
};
|
||||
|
||||
headers.insert(CONTROL_REQUEST_ID_HEADER.to_string(), request_id.clone());
|
||||
|
||||
if let Some(candidate_id) = candidate_id
|
||||
@@ -2444,6 +2836,17 @@ async fn execute_stream_from_frame_stream(
|
||||
);
|
||||
}
|
||||
|
||||
let emit_sse_keepalive = response_headers_indicate_sse(&headers);
|
||||
if emit_sse_keepalive {
|
||||
headers.remove("content-length");
|
||||
}
|
||||
let body_stream = build_sse_body_stream(
|
||||
prefetched_chunks_for_body,
|
||||
rx,
|
||||
emit_sse_keepalive,
|
||||
SSE_KEEPALIVE_INTERVAL,
|
||||
);
|
||||
|
||||
Ok(Some(build_client_response_from_parts(
|
||||
status_code,
|
||||
&headers,
|
||||
@@ -2467,7 +2870,7 @@ mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
use std::convert::Infallible;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use aether_contracts::{
|
||||
ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTimeouts, RequestBody,
|
||||
@@ -2484,11 +2887,13 @@ mod tests {
|
||||
use axum::extract::Request;
|
||||
use axum::routing::any;
|
||||
use axum::{http::header, http::HeaderValue, Router};
|
||||
use futures_util::StreamExt as _;
|
||||
use serde_json::{json, Value};
|
||||
use tokio::sync::{watch, Notify};
|
||||
use tokio::sync::{mpsc, watch, Notify};
|
||||
|
||||
use super::{
|
||||
execute_execution_runtime_stream, merge_stream_terminal_summary,
|
||||
build_sse_body_stream, execute_execution_runtime_stream, execute_stream_from_frame_stream,
|
||||
merge_stream_terminal_summary, should_limit_direct_finalize_prefetch,
|
||||
should_probe_success_failover_before_stream, should_skip_direct_finalize_prefetch,
|
||||
};
|
||||
use crate::control::GatewayControlDecision;
|
||||
@@ -2614,6 +3019,132 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn limits_prefetch_for_openai_image_and_rewritten_streams() {
|
||||
assert!(should_limit_direct_finalize_prefetch(
|
||||
"openai_image_stream",
|
||||
false
|
||||
));
|
||||
assert!(should_limit_direct_finalize_prefetch(
|
||||
"openai_chat_stream",
|
||||
true
|
||||
));
|
||||
assert!(!should_limit_direct_finalize_prefetch(
|
||||
"openai_chat_stream",
|
||||
false
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sse_body_stream_emits_initial_and_periodic_keepalive_without_business_chunks() {
|
||||
let (_tx, rx) = mpsc::channel::<Result<Bytes, std::io::Error>>(1);
|
||||
let mut body_stream = Box::pin(build_sse_body_stream(
|
||||
Vec::new(),
|
||||
rx,
|
||||
true,
|
||||
Duration::from_millis(10),
|
||||
));
|
||||
|
||||
let first = tokio::time::timeout(Duration::from_millis(50), body_stream.next())
|
||||
.await
|
||||
.expect("initial keepalive should be immediate")
|
||||
.expect("stream should yield initial keepalive")
|
||||
.expect("initial keepalive should be ok");
|
||||
assert_eq!(first.as_ref(), b": aether-keepalive\n\n");
|
||||
|
||||
let second = tokio::time::timeout(Duration::from_millis(100), body_stream.next())
|
||||
.await
|
||||
.expect("periodic keepalive should arrive")
|
||||
.expect("stream should yield periodic keepalive")
|
||||
.expect("periodic keepalive should be ok");
|
||||
assert_eq!(second.as_ref(), b": aether-keepalive\n\n");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_image_stream_total_timeout_emits_image_failed_event() {
|
||||
let state = AppState::new().expect("app state should build");
|
||||
let plan = ExecutionPlan {
|
||||
request_id: "req-image-stream-timeout".into(),
|
||||
candidate_id: Some("cand-image-stream-timeout".into()),
|
||||
provider_name: Some("codex".into()),
|
||||
provider_id: "prov-1".into(),
|
||||
endpoint_id: "ep-1".into(),
|
||||
key_id: "key-1".into(),
|
||||
method: "POST".into(),
|
||||
url: "https://chatgpt.com/backend-api/codex/responses".into(),
|
||||
headers: BTreeMap::from([
|
||||
("content-type".into(), "application/json".into()),
|
||||
("accept".into(), "text/event-stream".into()),
|
||||
]),
|
||||
content_type: Some("application/json".into()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({
|
||||
"model": "gpt-image-1",
|
||||
"prompt": "hello",
|
||||
"stream": true
|
||||
})),
|
||||
stream: true,
|
||||
client_api_format: "openai:image".into(),
|
||||
provider_api_format: "openai:image".into(),
|
||||
model_name: Some("gpt-image-1".into()),
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: Some(ExecutionTimeouts {
|
||||
total_ms: Some(25),
|
||||
..ExecutionTimeouts::default()
|
||||
}),
|
||||
};
|
||||
let decision = GatewayControlDecision::synthetic(
|
||||
"/v1/images/generations",
|
||||
Some("ai_public".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("image".to_string()),
|
||||
Some("openai:image".to_string()),
|
||||
)
|
||||
.with_execution_runtime_candidate(true);
|
||||
let frame_stream = stream! {
|
||||
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
|
||||
b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
|
||||
));
|
||||
std::future::pending::<()>().await;
|
||||
}
|
||||
.boxed();
|
||||
|
||||
let response = execute_stream_from_frame_stream(
|
||||
&state,
|
||||
plan,
|
||||
"trace-image-stream-timeout",
|
||||
&decision,
|
||||
"openai_image_stream",
|
||||
None,
|
||||
Some(json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:image",
|
||||
"image_request": {
|
||||
"operation": "generate"
|
||||
}
|
||||
})),
|
||||
crate::clock::current_unix_ms(),
|
||||
Instant::now(),
|
||||
frame_stream,
|
||||
)
|
||||
.await
|
||||
.expect("execution should succeed")
|
||||
.expect("execution should return a client response");
|
||||
|
||||
let body = tokio::time::timeout(
|
||||
Duration::from_secs(2),
|
||||
to_bytes(response.into_body(), usize::MAX),
|
||||
)
|
||||
.await
|
||||
.expect("timeout failure should close the response body")
|
||||
.expect("response body should read");
|
||||
let text = String::from_utf8(body.to_vec()).expect("response body should be utf8");
|
||||
assert!(text.contains(": aether-keepalive\n\n"));
|
||||
assert!(text.contains("event: image_generation.failed"));
|
||||
assert!(text.contains("\"type\":\"image_stream_total_timeout\""));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_execution_runtime_stream_records_first_data_as_streaming_before_terminal_telemetry(
|
||||
) {
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,6 +1,8 @@
|
||||
mod execution;
|
||||
|
||||
pub(crate) use execution::execute_execution_runtime_sync;
|
||||
pub(crate) use execution::{
|
||||
build_openai_image_sync_json_whitespace_heartbeat_stream, execute_execution_runtime_sync,
|
||||
};
|
||||
|
||||
#[allow(unused_imports)]
|
||||
pub(crate) use execution::{
|
||||
|
||||
@@ -551,7 +551,7 @@ fn build_direct_tunnel_request_meta(
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_request(
|
||||
pub(crate) async fn send_request(
|
||||
plan: &ExecutionPlan,
|
||||
body_bytes: Vec<u8>,
|
||||
) -> Result<reqwest::Response, ExecutionRuntimeTransportError> {
|
||||
@@ -728,7 +728,9 @@ async fn send_via_tunnel_relay(
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
fn build_request_body(plan: &ExecutionPlan) -> Result<Vec<u8>, ExecutionRuntimeTransportError> {
|
||||
pub(crate) fn build_request_body(
|
||||
plan: &ExecutionPlan,
|
||||
) -> Result<Vec<u8>, ExecutionRuntimeTransportError> {
|
||||
let mut body_bytes = if let Some(json_body) = plan.body.json_body.clone() {
|
||||
serde_json::to_vec(&json_body).map_err(ExecutionRuntimeTransportError::BodyEncode)?
|
||||
} else if let Some(body_b64) = plan.body.body_bytes_b64.as_deref() {
|
||||
@@ -1102,7 +1104,7 @@ fn is_hop_by_hop_header(name: &str) -> bool {
|
||||
)
|
||||
}
|
||||
|
||||
fn collect_response_headers(headers: &HeaderMap) -> BTreeMap<String, String> {
|
||||
pub(crate) fn collect_response_headers(headers: &HeaderMap) -> BTreeMap<String, String> {
|
||||
header_map_to_string_map(headers)
|
||||
}
|
||||
|
||||
@@ -1130,7 +1132,7 @@ fn execution_log_url_host(url: &str) -> String {
|
||||
.unwrap_or_else(|| "-".to_string())
|
||||
}
|
||||
|
||||
fn decode_response_body_bytes(
|
||||
pub(crate) fn decode_response_body_bytes(
|
||||
headers: &BTreeMap<String, String>,
|
||||
body_bytes: &[u8],
|
||||
) -> Option<Vec<u8>> {
|
||||
@@ -1157,7 +1159,7 @@ fn decode_response_body_bytes(
|
||||
}
|
||||
}
|
||||
|
||||
fn response_body_is_json(headers: &BTreeMap<String, String>, body_bytes: &[u8]) -> bool {
|
||||
pub(crate) fn response_body_is_json(headers: &BTreeMap<String, String>, body_bytes: &[u8]) -> bool {
|
||||
if headers
|
||||
.get("content-type")
|
||||
.map(|value| value.to_ascii_lowercase())
|
||||
|
||||
@@ -158,6 +158,19 @@ where
|
||||
last_plan: aether_contracts::ExecutionPlan,
|
||||
last_report_context: Option<serde_json::Value>,
|
||||
) -> Result<Self::Exhaustion, Self::Error> {
|
||||
warn!(
|
||||
event_name = "candidate_loop_exhausted",
|
||||
log_type = "ops",
|
||||
trace_id = %self.trace_id,
|
||||
plan_kind = self.plan_kind,
|
||||
request_id = %short_request_id(last_plan.request_id.as_str()),
|
||||
candidate_id = ?last_plan.candidate_id,
|
||||
provider_name = last_plan.provider_name.as_deref().unwrap_or("-"),
|
||||
endpoint_id = %last_plan.endpoint_id,
|
||||
key_id = %last_plan.key_id,
|
||||
model_name = last_plan.model_name.as_deref().unwrap_or("-"),
|
||||
"candidate loop exhausted local sync candidates"
|
||||
);
|
||||
Ok(
|
||||
build_local_execution_exhaustion(self.state, &last_plan, last_report_context.as_ref())
|
||||
.await,
|
||||
|
||||
@@ -1,3 +1,13 @@
|
||||
use std::collections::{BTreeMap, VecDeque};
|
||||
use std::io::Error as IoError;
|
||||
use std::time::Instant;
|
||||
|
||||
use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::http::header::{CACHE_CONTROL, CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE};
|
||||
use axum::http::{HeaderName, HeaderValue, Response, StatusCode};
|
||||
use serde_json::{json, Value};
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::ai_serving::api::{
|
||||
build_local_gemini_files_stream_attempt_source_for_kind,
|
||||
build_local_gemini_files_sync_attempt_source_for_kind,
|
||||
@@ -18,16 +28,35 @@ use crate::ai_serving::api::{
|
||||
resolve_claude_stream_spec, resolve_claude_sync_spec, resolve_gemini_stream_spec,
|
||||
resolve_gemini_sync_spec, resolve_local_same_format_stream_spec,
|
||||
resolve_local_same_format_sync_spec, set_local_openai_chat_execution_exhausted_diagnostic,
|
||||
AiStreamAttempt, AiSyncAttempt, LocalStandardSpec, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
|
||||
set_local_openai_image_execution_exhausted_diagnostic, AiStreamAttempt, AiSyncAttempt,
|
||||
LocalStandardSpec, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
|
||||
};
|
||||
use crate::ai_serving::LocalExecutionAttemptSource;
|
||||
use crate::api::response::{
|
||||
attach_control_metadata_headers, build_client_response_from_parts_with_mutator,
|
||||
};
|
||||
use crate::constants::EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS;
|
||||
use crate::control::GatewayControlDecision;
|
||||
use crate::execution_runtime::sync::{
|
||||
build_openai_image_sync_json_whitespace_heartbeat_stream, execute_execution_runtime_sync,
|
||||
};
|
||||
use crate::executor::candidate_loop::{
|
||||
execute_stream_attempt_source, execute_sync_attempt_source, execute_sync_plan_and_reports,
|
||||
mark_unused_local_candidates,
|
||||
};
|
||||
use crate::executor::LocalExecutionRequestOutcome;
|
||||
use crate::executor::{
|
||||
build_local_execution_exhaustion, record_failed_usage_for_exhausted_request,
|
||||
LocalExecutionRequestOutcome,
|
||||
};
|
||||
use crate::handlers::shared::system_config_bool;
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
const ENABLE_OPENAI_IMAGE_SYNC_HEARTBEAT_CONFIG_KEY: &str = "enable_openai_image_sync_heartbeat";
|
||||
const OPENAI_IMAGE_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS: u16 = 502;
|
||||
const OPENAI_IMAGE_SYNC_HEARTBEAT_EXHAUSTED_STATUS: u16 = 503;
|
||||
const OPENAI_IMAGE_SYNC_HEARTBEAT_ERROR_MESSAGE_LIMIT: usize = 4096;
|
||||
|
||||
pub(crate) async fn maybe_execute_sync_local_path(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
@@ -448,6 +477,226 @@ pub(crate) async fn maybe_execute_sync_via_local_gemini_files_decision(
|
||||
.await
|
||||
}
|
||||
|
||||
async fn openai_image_sync_heartbeat_enabled(state: &AppState) -> bool {
|
||||
match state
|
||||
.read_system_config_json_value(ENABLE_OPENAI_IMAGE_SYNC_HEARTBEAT_CONFIG_KEY)
|
||||
.await
|
||||
{
|
||||
Ok(value) => system_config_bool(value.as_ref(), false),
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
event_name = "openai_image_sync_heartbeat_config_read_failed",
|
||||
log_type = "ops",
|
||||
error = ?err,
|
||||
"gateway failed to read sync image heartbeat config; defaulting disabled"
|
||||
);
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn build_openai_image_sync_heartbeat_shell_response(
|
||||
state: AppState,
|
||||
request_path: String,
|
||||
trace_id: String,
|
||||
decision: GatewayControlDecision,
|
||||
plan_kind: String,
|
||||
attempts: Vec<AiSyncAttempt>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let request_id = attempts
|
||||
.first()
|
||||
.map(|attempt| attempt.plan.request_id.clone())
|
||||
.filter(|value| !value.trim().is_empty());
|
||||
let trace_id_for_response = trace_id.clone();
|
||||
let decision_for_response = decision.clone();
|
||||
let started_at = Instant::now();
|
||||
let (tx, rx) = mpsc::channel::<Result<Bytes, IoError>>(1);
|
||||
|
||||
tokio::spawn(async move {
|
||||
let bytes = openai_image_sync_heartbeat_final_bytes(
|
||||
execute_openai_image_sync_heartbeat_attempts(
|
||||
state,
|
||||
request_path,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
attempts,
|
||||
started_at,
|
||||
)
|
||||
.await,
|
||||
)
|
||||
.await;
|
||||
let _ = tx.send(Ok(Bytes::from(bytes))).await;
|
||||
});
|
||||
|
||||
let headers = BTreeMap::from([(
|
||||
CONTENT_TYPE.as_str().to_string(),
|
||||
"application/json".to_string(),
|
||||
)]);
|
||||
let response = build_client_response_from_parts_with_mutator(
|
||||
StatusCode::OK.as_u16(),
|
||||
&headers,
|
||||
Body::from_stream(build_openai_image_sync_json_whitespace_heartbeat_stream(rx)),
|
||||
trace_id_for_response.as_str(),
|
||||
Some(&decision_for_response),
|
||||
|headers| {
|
||||
headers.remove(CONTENT_LENGTH);
|
||||
headers.remove(CONTENT_ENCODING);
|
||||
headers.insert(
|
||||
CACHE_CONTROL,
|
||||
HeaderValue::from_static("no-cache, no-transform"),
|
||||
);
|
||||
headers.insert(
|
||||
HeaderName::from_static("x-accel-buffering"),
|
||||
HeaderValue::from_static("no"),
|
||||
);
|
||||
Ok(())
|
||||
},
|
||||
)?;
|
||||
attach_control_metadata_headers(response, request_id.as_deref(), None)
|
||||
}
|
||||
|
||||
async fn execute_openai_image_sync_heartbeat_attempts(
|
||||
state: AppState,
|
||||
request_path: String,
|
||||
trace_id: String,
|
||||
decision: GatewayControlDecision,
|
||||
plan_kind: String,
|
||||
attempts: Vec<AiSyncAttempt>,
|
||||
started_at: Instant,
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let mut attempts = VecDeque::from(attempts);
|
||||
let mut last_attempted = None;
|
||||
|
||||
while let Some(attempt) = attempts.pop_front() {
|
||||
let plan = attempt.plan;
|
||||
let report_kind = attempt.report_kind;
|
||||
let report_context = attempt.report_context;
|
||||
last_attempted = Some((plan.clone(), report_context.clone()));
|
||||
match execute_execution_runtime_sync(
|
||||
&state,
|
||||
request_path.as_str(),
|
||||
plan,
|
||||
trace_id.as_str(),
|
||||
&decision,
|
||||
plan_kind.as_str(),
|
||||
report_kind,
|
||||
report_context,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Some(response) => {
|
||||
mark_unused_local_candidates(&state, attempts.into_iter().collect()).await;
|
||||
return Ok(LocalExecutionRequestOutcome::responded(response));
|
||||
}
|
||||
None => continue,
|
||||
}
|
||||
}
|
||||
|
||||
let Some((last_plan, last_report_context)) = last_attempted else {
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
let exhaustion =
|
||||
build_local_execution_exhaustion(&state, &last_plan, last_report_context.as_ref()).await;
|
||||
record_failed_usage_for_exhausted_request(
|
||||
&state,
|
||||
exhaustion,
|
||||
&started_at,
|
||||
"OpenAI image sync heartbeat exhausted all local candidates",
|
||||
EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
Ok(LocalExecutionRequestOutcome::NoPath)
|
||||
}
|
||||
|
||||
async fn openai_image_sync_heartbeat_final_bytes(
|
||||
result: Result<LocalExecutionRequestOutcome, GatewayError>,
|
||||
) -> Vec<u8> {
|
||||
match result {
|
||||
Ok(LocalExecutionRequestOutcome::Responded(response)) => {
|
||||
openai_image_sync_heartbeat_response_body_bytes(response).await
|
||||
}
|
||||
Ok(LocalExecutionRequestOutcome::Exhausted(_))
|
||||
| Ok(LocalExecutionRequestOutcome::NoPath) => openai_image_sync_heartbeat_error_body(
|
||||
OPENAI_IMAGE_SYNC_HEARTBEAT_EXHAUSTED_STATUS,
|
||||
"OpenAI image sync exhausted all local candidates",
|
||||
),
|
||||
Err(err) => openai_image_sync_heartbeat_error_body(
|
||||
OPENAI_IMAGE_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS,
|
||||
&format!("{err:?}"),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
async fn openai_image_sync_heartbeat_response_body_bytes(response: Response<Body>) -> Vec<u8> {
|
||||
let status_code = response.status().as_u16();
|
||||
match to_bytes(response.into_body(), usize::MAX).await {
|
||||
Ok(bytes) if status_code < 400 && !bytes.is_empty() => bytes.to_vec(),
|
||||
Ok(bytes) if status_code >= 400 => {
|
||||
openai_image_sync_heartbeat_error_body_from_response(status_code, bytes.as_ref())
|
||||
}
|
||||
Ok(_) => openai_image_sync_heartbeat_error_body(
|
||||
OPENAI_IMAGE_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS,
|
||||
"empty sync image response",
|
||||
),
|
||||
Err(err) => openai_image_sync_heartbeat_error_body(
|
||||
OPENAI_IMAGE_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS,
|
||||
&err.to_string(),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn openai_image_sync_heartbeat_error_body_from_response(status_code: u16, body: &[u8]) -> Vec<u8> {
|
||||
if let Ok(mut value) = serde_json::from_slice::<Value>(body) {
|
||||
if let Some(error) = value.get_mut("error").and_then(Value::as_object_mut) {
|
||||
error.insert("upstream_status".to_string(), Value::from(status_code));
|
||||
error
|
||||
.entry("type".to_string())
|
||||
.or_insert_with(|| Value::String("upstream_error".to_string()));
|
||||
error.entry("message".to_string()).or_insert_with(|| {
|
||||
Value::String(format!("upstream returned status {status_code}"))
|
||||
});
|
||||
return serde_json::to_vec(&value).unwrap_or_else(|_| {
|
||||
openai_image_sync_heartbeat_error_body(
|
||||
status_code,
|
||||
&format!("upstream returned status {status_code}"),
|
||||
)
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let message = openai_image_sync_heartbeat_error_message_from_body(status_code, body);
|
||||
openai_image_sync_heartbeat_error_body(status_code, message.as_str())
|
||||
}
|
||||
|
||||
fn openai_image_sync_heartbeat_error_message_from_body(status_code: u16, body: &[u8]) -> String {
|
||||
let text = String::from_utf8_lossy(body).trim().to_string();
|
||||
if text.is_empty() {
|
||||
return format!("upstream returned status {status_code}");
|
||||
}
|
||||
text.chars()
|
||||
.take(OPENAI_IMAGE_SYNC_HEARTBEAT_ERROR_MESSAGE_LIMIT)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn openai_image_sync_heartbeat_error_body(status_code: u16, message: &str) -> Vec<u8> {
|
||||
serde_json::to_vec(&json!({
|
||||
"error": {
|
||||
"type": "upstream_error",
|
||||
"message": message,
|
||||
"code": status_code,
|
||||
"upstream_status": status_code,
|
||||
}
|
||||
}))
|
||||
.unwrap_or_else(|_| {
|
||||
format!(
|
||||
"{{\"error\":{{\"type\":\"upstream_error\",\"code\":{status_code},\"upstream_status\":{status_code}}}}}"
|
||||
)
|
||||
.into_bytes()
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_execute_sync_via_local_image_decision(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
@@ -457,21 +706,39 @@ pub(crate) async fn maybe_execute_sync_via_local_image_decision(
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let Some((attempt_source, _candidate_count)) = build_local_image_sync_attempt_source_for_kind(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
)
|
||||
.await?
|
||||
let Some((mut attempt_source, candidate_count)) =
|
||||
build_local_image_sync_attempt_source_for_kind(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
)
|
||||
.await?
|
||||
else {
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
execute_sync_attempt_source::<AiSyncAttempt, _>(
|
||||
if openai_image_sync_heartbeat_enabled(state).await {
|
||||
let mut attempts = Vec::new();
|
||||
while let Some(attempt) = attempt_source.next_execution_attempt().await? {
|
||||
attempts.push(attempt);
|
||||
}
|
||||
return Ok(LocalExecutionRequestOutcome::responded(
|
||||
build_openai_image_sync_heartbeat_shell_response(
|
||||
state.clone(),
|
||||
parts.uri.path().to_string(),
|
||||
trace_id.to_string(),
|
||||
decision.clone(),
|
||||
plan_kind.to_string(),
|
||||
attempts,
|
||||
)?,
|
||||
));
|
||||
}
|
||||
|
||||
let outcome = execute_sync_attempt_source::<AiSyncAttempt, _>(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
@@ -479,7 +746,20 @@ pub(crate) async fn maybe_execute_sync_via_local_image_decision(
|
||||
plan_kind,
|
||||
attempt_source,
|
||||
)
|
||||
.await
|
||||
.await?;
|
||||
|
||||
if let LocalExecutionRequestOutcome::Exhausted(_) = &outcome {
|
||||
set_local_openai_image_execution_exhausted_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
body_json,
|
||||
candidate_count,
|
||||
);
|
||||
}
|
||||
|
||||
Ok(outcome)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_execute_stream_via_local_gemini_files_decision(
|
||||
@@ -517,29 +797,41 @@ pub(crate) async fn maybe_execute_stream_via_local_image_decision(
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let Some((attempt_source, _candidate_count)) =
|
||||
build_local_image_stream_attempt_source_for_kind(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
)
|
||||
.await?
|
||||
let Some((attempt_source, candidate_count)) = build_local_image_stream_attempt_source_for_kind(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
)
|
||||
.await?
|
||||
else {
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
execute_stream_attempt_source::<AiStreamAttempt, _>(
|
||||
let outcome = execute_stream_attempt_source::<AiStreamAttempt, _>(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
attempt_source,
|
||||
)
|
||||
.await
|
||||
.await?;
|
||||
|
||||
if let LocalExecutionRequestOutcome::Exhausted(_) = &outcome {
|
||||
set_local_openai_image_execution_exhausted_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
body_json,
|
||||
candidate_count,
|
||||
);
|
||||
}
|
||||
|
||||
Ok(outcome)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_execute_sync_via_local_video_decision(
|
||||
@@ -648,3 +940,197 @@ pub(crate) fn parse_local_request_body(
|
||||
pub(crate) fn decision_payload_is_direct_execution(payload: &AiExecutionDecision) -> bool {
|
||||
planner_decision_action(payload.action.as_str())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
const TEST_OPENAI_IMAGE_SYNC_PLAN_KIND: &str = "openai_image_sync";
|
||||
|
||||
fn test_openai_image_heartbeat_decision() -> GatewayControlDecision {
|
||||
GatewayControlDecision::synthetic(
|
||||
"/v1/images/generations",
|
||||
Some("ai_public".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("image".to_string()),
|
||||
Some("openai:image".to_string()),
|
||||
)
|
||||
.with_execution_runtime_candidate(true)
|
||||
}
|
||||
|
||||
fn test_openai_image_heartbeat_plan(
|
||||
endpoint_id: &str,
|
||||
candidate_id: &str,
|
||||
) -> aether_contracts::ExecutionPlan {
|
||||
aether_contracts::ExecutionPlan {
|
||||
request_id: "trace-image-heartbeat-retry".to_string(),
|
||||
candidate_id: Some(candidate_id.to_string()),
|
||||
provider_name: Some("OpenAI".to_string()),
|
||||
provider_id: "provider-openai".to_string(),
|
||||
endpoint_id: endpoint_id.to_string(),
|
||||
key_id: "key-openai".to_string(),
|
||||
method: "POST".to_string(),
|
||||
url: "https://api.openai.com/v1/images/generations".to_string(),
|
||||
headers: BTreeMap::new(),
|
||||
content_type: Some("application/json".to_string()),
|
||||
content_encoding: None,
|
||||
body: aether_contracts::RequestBody::from_json(json!({"prompt": "test"})),
|
||||
stream: false,
|
||||
client_api_format: "openai:image".to_string(),
|
||||
provider_api_format: "openai:image".to_string(),
|
||||
model_name: Some("gpt-image-1".to_string()),
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn test_openai_image_heartbeat_attempt(
|
||||
candidate_index: u32,
|
||||
endpoint_id: &str,
|
||||
candidate_id: &str,
|
||||
) -> AiSyncAttempt {
|
||||
AiSyncAttempt {
|
||||
plan: test_openai_image_heartbeat_plan(endpoint_id, candidate_id),
|
||||
report_kind: None,
|
||||
report_context: Some(json!({
|
||||
"candidate_index": candidate_index,
|
||||
"retry_index": 0,
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
fn test_openai_image_execution_result(
|
||||
plan: &aether_contracts::ExecutionPlan,
|
||||
status_code: u16,
|
||||
body_json: Value,
|
||||
) -> aether_contracts::ExecutionResult {
|
||||
aether_contracts::ExecutionResult {
|
||||
request_id: plan.request_id.clone(),
|
||||
candidate_id: plan.candidate_id.clone(),
|
||||
status_code,
|
||||
headers: BTreeMap::from([(
|
||||
CONTENT_TYPE.as_str().to_string(),
|
||||
"application/json".to_string(),
|
||||
)]),
|
||||
body: Some(aether_contracts::ResponseBody {
|
||||
json_body: Some(body_json),
|
||||
body_bytes_b64: None,
|
||||
}),
|
||||
telemetry: Some(aether_contracts::ExecutionTelemetry {
|
||||
ttfb_ms: None,
|
||||
elapsed_ms: Some(10),
|
||||
upstream_bytes: None,
|
||||
}),
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_image_sync_heartbeat_success_body_is_unchanged() {
|
||||
let response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from(r#"{"data":[{"b64_json":"x"}]}"#))
|
||||
.expect("response should build");
|
||||
|
||||
let bytes = openai_image_sync_heartbeat_response_body_bytes(response).await;
|
||||
let body: Value = serde_json::from_slice(&bytes).expect("body should decode");
|
||||
|
||||
assert_eq!(body, json!({"data": [{"b64_json": "x"}]}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_image_sync_heartbeat_missing_config_defaults_disabled() {
|
||||
let state = AppState::new().expect("state should build");
|
||||
|
||||
assert!(!openai_image_sync_heartbeat_enabled(&state).await);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_image_sync_heartbeat_error_body_includes_upstream_status() {
|
||||
let response = Response::builder()
|
||||
.status(StatusCode::TOO_MANY_REQUESTS)
|
||||
.body(Body::from(
|
||||
r#"{"error":{"type":"rate_limit","message":"slow down"}}"#,
|
||||
))
|
||||
.expect("response should build");
|
||||
|
||||
let bytes = openai_image_sync_heartbeat_response_body_bytes(response).await;
|
||||
let body: Value = serde_json::from_slice(&bytes).expect("body should decode");
|
||||
|
||||
assert_eq!(body["error"]["type"], json!("rate_limit"));
|
||||
assert_eq!(body["error"]["message"], json!("slow down"));
|
||||
assert_eq!(body["error"]["upstream_status"], json!(429));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_image_sync_heartbeat_non_json_error_body_is_wrapped() {
|
||||
let bytes =
|
||||
openai_image_sync_heartbeat_error_body_from_response(502, b"bad gateway from upstream");
|
||||
let body: Value = serde_json::from_slice(&bytes).expect("body should decode");
|
||||
|
||||
assert_eq!(body["error"]["type"], json!("upstream_error"));
|
||||
assert_eq!(body["error"]["message"], json!("bad gateway from upstream"));
|
||||
assert_eq!(body["error"]["upstream_status"], json!(502));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_image_sync_heartbeat_no_path_returns_json_error_body() {
|
||||
let bytes =
|
||||
openai_image_sync_heartbeat_final_bytes(Ok(LocalExecutionRequestOutcome::NoPath)).await;
|
||||
let body: Value = serde_json::from_slice(&bytes).expect("body should decode");
|
||||
|
||||
assert_eq!(body["error"]["type"], json!("upstream_error"));
|
||||
assert_eq!(body["error"]["upstream_status"], json!(503));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_image_sync_heartbeat_attempts_retry_first_candidate_then_return_second() {
|
||||
let call_count = Arc::new(AtomicUsize::new(0));
|
||||
let call_count_for_override = Arc::clone(&call_count);
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_execution_runtime_sync_override_for_tests(move |plan| {
|
||||
call_count_for_override.fetch_add(1, Ordering::SeqCst);
|
||||
if plan.endpoint_id == "endpoint-retry" {
|
||||
Ok(test_openai_image_execution_result(
|
||||
plan,
|
||||
StatusCode::TOO_MANY_REQUESTS.as_u16(),
|
||||
json!({"error": {"message": "retry this candidate"}}),
|
||||
))
|
||||
} else {
|
||||
Ok(test_openai_image_execution_result(
|
||||
plan,
|
||||
StatusCode::OK.as_u16(),
|
||||
json!({"data": [{"b64_json": "second-candidate"}]}),
|
||||
))
|
||||
}
|
||||
});
|
||||
let attempts = vec![
|
||||
test_openai_image_heartbeat_attempt(0, "endpoint-retry", "candidate-retry"),
|
||||
test_openai_image_heartbeat_attempt(1, "endpoint-success", "candidate-success"),
|
||||
];
|
||||
|
||||
let outcome = execute_openai_image_sync_heartbeat_attempts(
|
||||
state,
|
||||
"/v1/images/generations".to_string(),
|
||||
"trace-image-heartbeat-retry".to_string(),
|
||||
test_openai_image_heartbeat_decision(),
|
||||
TEST_OPENAI_IMAGE_SYNC_PLAN_KIND.to_string(),
|
||||
attempts,
|
||||
Instant::now(),
|
||||
)
|
||||
.await
|
||||
.expect("heartbeat attempts should execute");
|
||||
let LocalExecutionRequestOutcome::Responded(response) = outcome else {
|
||||
panic!("second candidate should return a response");
|
||||
};
|
||||
let bytes = openai_image_sync_heartbeat_response_body_bytes(response).await;
|
||||
let body: Value = serde_json::from_slice(&bytes).expect("body should decode");
|
||||
|
||||
assert_eq!(call_count.load(Ordering::SeqCst), 2);
|
||||
assert_eq!(body, json!({"data": [{"b64_json": "second-candidate"}]}));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -194,6 +194,57 @@ async fn resolve_admin_usage_attempt_flags_by_usage_id(
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn resolve_admin_usage_image_progress_by_request_id(
|
||||
state: &AdminAppState<'_>,
|
||||
items: &[StoredRequestUsageAudit],
|
||||
) -> Result<BTreeMap<String, serde_json::Value>, GatewayError> {
|
||||
if !state.has_request_candidate_data_reader() || items.is_empty() {
|
||||
return Ok(BTreeMap::new());
|
||||
}
|
||||
|
||||
let request_ids = items
|
||||
.iter()
|
||||
.map(|item| item.request_id.clone())
|
||||
.collect::<BTreeSet<_>>();
|
||||
let mut progress_by_request_id = BTreeMap::new();
|
||||
for request_id in request_ids {
|
||||
let candidates = state
|
||||
.app()
|
||||
.read_request_candidates_by_request_id(&request_id)
|
||||
.await?;
|
||||
if let Some(progress) = latest_admin_usage_image_progress(&candidates) {
|
||||
progress_by_request_id.insert(request_id, progress);
|
||||
}
|
||||
}
|
||||
Ok(progress_by_request_id)
|
||||
}
|
||||
|
||||
fn latest_admin_usage_image_progress(
|
||||
candidates: &[StoredRequestCandidate],
|
||||
) -> Option<serde_json::Value> {
|
||||
candidates
|
||||
.iter()
|
||||
.filter_map(|candidate| {
|
||||
let progress = candidate
|
||||
.extra_data
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("image_progress"))?
|
||||
.clone();
|
||||
Some((
|
||||
candidate
|
||||
.started_at_unix_ms
|
||||
.unwrap_or(candidate.created_at_unix_ms),
|
||||
candidate.candidate_index,
|
||||
candidate.retry_index,
|
||||
progress,
|
||||
))
|
||||
})
|
||||
.max_by_key(|(started_at, candidate_index, retry_index, _)| {
|
||||
(*started_at, *candidate_index, *retry_index)
|
||||
})
|
||||
.map(|(_, _, _, progress)| progress)
|
||||
}
|
||||
|
||||
fn admin_usage_matches_attempt_status(
|
||||
item: &StoredRequestUsageAudit,
|
||||
status: &str,
|
||||
@@ -490,6 +541,7 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
|
||||
&BTreeMap::new(),
|
||||
state.has_auth_api_key_data_reader(),
|
||||
&BTreeMap::new(),
|
||||
&BTreeMap::new(),
|
||||
)));
|
||||
};
|
||||
state
|
||||
@@ -513,12 +565,15 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
|
||||
};
|
||||
let api_key_names = admin_usage_api_key_names(state, &items).await?;
|
||||
let provider_key_names = admin_usage_provider_key_names(state, &items).await?;
|
||||
let image_progress_by_request_id =
|
||||
resolve_admin_usage_image_progress_by_request_id(state, &items).await?;
|
||||
|
||||
return Ok(Some(build_admin_usage_active_requests_response(
|
||||
&items,
|
||||
&api_key_names,
|
||||
state.has_auth_api_key_data_reader(),
|
||||
&provider_key_names,
|
||||
&image_progress_by_request_id,
|
||||
)));
|
||||
}
|
||||
Some("records")
|
||||
|
||||
@@ -172,6 +172,10 @@ impl<'a> AdminAppState<'a> {
|
||||
let user_api_keys = self
|
||||
.list_auth_api_key_export_records_by_user_ids(&user_ids)
|
||||
.await?;
|
||||
let groups = self.list_user_groups().await?;
|
||||
let memberships = self
|
||||
.list_user_group_memberships_by_user_ids(&user_ids)
|
||||
.await?;
|
||||
let standalone_api_keys = self.list_auth_api_key_export_standalone_records().await?;
|
||||
let standalone_api_key_ids = standalone_api_keys
|
||||
.iter()
|
||||
@@ -205,12 +209,49 @@ impl<'a> AdminAppState<'a> {
|
||||
.or_default()
|
||||
.push(key);
|
||||
}
|
||||
let mut memberships_by_user_id = BTreeMap::<
|
||||
String,
|
||||
Vec<aether_data::repository::users::StoredUserGroupMembership>,
|
||||
>::new();
|
||||
for membership in memberships {
|
||||
memberships_by_user_id
|
||||
.entry(membership.user_id.clone())
|
||||
.or_default()
|
||||
.push(membership);
|
||||
}
|
||||
let user_groups_data = groups
|
||||
.iter()
|
||||
.map(|group| {
|
||||
json!({
|
||||
"id": group.id.clone(),
|
||||
"name": group.name.clone(),
|
||||
"description": group.description.clone(),
|
||||
"allowed_providers": group.allowed_providers.clone(),
|
||||
"allowed_providers_mode": group.allowed_providers_mode.clone(),
|
||||
"allowed_api_formats": group.allowed_api_formats.clone(),
|
||||
"allowed_api_formats_mode": group.allowed_api_formats_mode.clone(),
|
||||
"allowed_models": group.allowed_models.clone(),
|
||||
"allowed_models_mode": group.allowed_models_mode.clone(),
|
||||
"rate_limit": group.rate_limit,
|
||||
"rate_limit_mode": group.rate_limit_mode.clone(),
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let users_data = users
|
||||
.iter()
|
||||
.map(|user| {
|
||||
let wallet = wallets_by_user_id.get(&user.id);
|
||||
let wallet_payload = serialize_admin_system_users_export_wallet(wallet);
|
||||
let memberships = memberships_by_user_id.remove(&user.id).unwrap_or_default();
|
||||
let group_ids = memberships
|
||||
.iter()
|
||||
.map(|membership| membership.group_id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let group_names = memberships
|
||||
.iter()
|
||||
.map(|membership| membership.group_name.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let api_keys = api_keys_by_user_id.remove(&user.id).unwrap_or_default();
|
||||
let api_keys_payload = api_keys
|
||||
.iter()
|
||||
@@ -226,10 +267,16 @@ impl<'a> AdminAppState<'a> {
|
||||
"password_hash": user.password_hash.clone(),
|
||||
"role": user.role.clone(),
|
||||
"allowed_providers": user.allowed_providers.clone(),
|
||||
"allowed_providers_mode": user.allowed_providers_mode.clone(),
|
||||
"allowed_api_formats": user.allowed_api_formats.clone(),
|
||||
"allowed_api_formats_mode": user.allowed_api_formats_mode.clone(),
|
||||
"allowed_models": user.allowed_models.clone(),
|
||||
"allowed_models_mode": user.allowed_models_mode.clone(),
|
||||
"rate_limit": user.rate_limit,
|
||||
"rate_limit_mode": user.rate_limit_mode.clone(),
|
||||
"model_capability_settings": user.model_capability_settings.clone(),
|
||||
"group_ids": group_ids,
|
||||
"group_names": group_names,
|
||||
"unlimited": wallet
|
||||
.map(|entry| entry.limit_mode.eq_ignore_ascii_case("unlimited"))
|
||||
.unwrap_or(false),
|
||||
@@ -254,6 +301,7 @@ impl<'a> AdminAppState<'a> {
|
||||
Ok(json!({
|
||||
"version": ADMIN_SYSTEM_USERS_EXPORT_VERSION,
|
||||
"exported_at": Utc::now().to_rfc3339(),
|
||||
"user_groups": user_groups_data,
|
||||
"users": users_data,
|
||||
"standalone_keys": standalone_keys_data,
|
||||
}))
|
||||
|
||||
@@ -10,7 +10,9 @@ use crate::handlers::admin::shared::{
|
||||
};
|
||||
use crate::handlers::admin::system::shared::configs::apply_admin_system_config_update;
|
||||
use crate::handlers::admin::users::{
|
||||
hash_admin_user_api_key, normalize_admin_user_api_formats, normalize_admin_user_string_list,
|
||||
hash_admin_user_api_key, normalize_admin_list_policy_mode,
|
||||
normalize_admin_rate_limit_policy_mode, normalize_admin_user_api_formats,
|
||||
normalize_admin_user_string_list,
|
||||
};
|
||||
use crate::handlers::public::normalize_admin_base_url;
|
||||
use crate::GatewayError;
|
||||
@@ -397,6 +399,7 @@ fn build_import_provider_model_record(
|
||||
|
||||
#[derive(Debug, Clone, Default, serde::Serialize)]
|
||||
struct AdminSystemUsersImportStats {
|
||||
user_groups: AdminSystemConfigImportCounter,
|
||||
users: AdminSystemConfigImportCounter,
|
||||
api_keys: AdminSystemConfigImportCounter,
|
||||
standalone_keys: AdminSystemConfigImportCounter,
|
||||
@@ -542,6 +545,78 @@ fn imported_optional_value(value: Option<&Value>) -> Option<Value> {
|
||||
value.cloned().filter(|value| !value.is_null())
|
||||
}
|
||||
|
||||
fn imported_optional_list_policy_mode(
|
||||
value: Option<&Value>,
|
||||
field_name: &str,
|
||||
) -> Result<Option<String>, String> {
|
||||
let Some(value) = imported_optional_string(value)? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let value = value.to_ascii_lowercase();
|
||||
normalize_admin_list_policy_mode(&value)
|
||||
.map(Some)
|
||||
.map_err(|_| format!("{field_name} 不合法"))
|
||||
}
|
||||
|
||||
fn imported_optional_rate_limit_policy_mode(
|
||||
value: Option<&Value>,
|
||||
field_name: &str,
|
||||
) -> Result<Option<String>, String> {
|
||||
let Some(value) = imported_optional_string(value)? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let value = value.to_ascii_lowercase();
|
||||
normalize_admin_rate_limit_policy_mode(&value)
|
||||
.map(Some)
|
||||
.map_err(|_| format!("{field_name} 不合法"))
|
||||
}
|
||||
|
||||
fn legacy_imported_list_policy_mode(values: &Option<Vec<String>>) -> String {
|
||||
if values.is_some() {
|
||||
"specific".to_string()
|
||||
} else {
|
||||
"unrestricted".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn legacy_imported_rate_limit_policy_mode(value: Option<i32>) -> String {
|
||||
if value.is_some() {
|
||||
"custom".to_string()
|
||||
} else {
|
||||
"system".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn imported_user_list_policy_mode(
|
||||
object: &Map<String, Value>,
|
||||
mode_field: &str,
|
||||
value_field: &str,
|
||||
values: &Option<Vec<String>>,
|
||||
) -> Result<Option<String>, String> {
|
||||
imported_optional_list_policy_mode(object.get(mode_field), mode_field).map(|mode| {
|
||||
mode.or_else(|| {
|
||||
object
|
||||
.contains_key(value_field)
|
||||
.then(|| legacy_imported_list_policy_mode(values))
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn imported_user_rate_limit_policy_mode(
|
||||
object: &Map<String, Value>,
|
||||
mode_field: &str,
|
||||
value_field: &str,
|
||||
value: Option<i32>,
|
||||
) -> Result<Option<String>, String> {
|
||||
imported_optional_rate_limit_policy_mode(object.get(mode_field), mode_field).map(|mode| {
|
||||
mode.or_else(|| {
|
||||
object
|
||||
.contains_key(value_field)
|
||||
.then(|| legacy_imported_rate_limit_policy_mode(value))
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn imported_rfc3339_to_unix_secs(
|
||||
value: Option<&Value>,
|
||||
field_name: &str,
|
||||
@@ -601,6 +676,130 @@ fn normalize_imported_user_api_formats(
|
||||
)?)
|
||||
}
|
||||
|
||||
fn build_imported_user_group_record(
|
||||
group: &Map<String, Value>,
|
||||
field_name: &str,
|
||||
) -> Result<
|
||||
(
|
||||
Option<String>,
|
||||
String,
|
||||
aether_data::repository::users::UpsertUserGroupRecord,
|
||||
),
|
||||
String,
|
||||
> {
|
||||
let export_id = imported_optional_string(group.get("id"))?;
|
||||
let name = imported_optional_string(group.get("name"))?
|
||||
.ok_or_else(|| format!("{field_name}.name 不能为空"))?;
|
||||
let name = aether_data::repository::users::normalize_user_group_name(&name);
|
||||
if name.is_empty() {
|
||||
return Err(format!("{field_name}.name 不能为空"));
|
||||
}
|
||||
let description = imported_optional_string(group.get("description"))?;
|
||||
let allowed_providers = normalize_imported_user_string_list(group, "allowed_providers")?;
|
||||
let allowed_api_formats = normalize_imported_user_api_formats(group, "allowed_api_formats")?;
|
||||
let allowed_models = normalize_imported_user_string_list(group, "allowed_models")?;
|
||||
let rate_limit = imported_optional_i32(group.get("rate_limit"), "rate_limit")?;
|
||||
|
||||
let allowed_providers_mode = imported_optional_list_policy_mode(
|
||||
group.get("allowed_providers_mode"),
|
||||
"allowed_providers_mode",
|
||||
)?
|
||||
.unwrap_or_else(|| {
|
||||
if group.contains_key("allowed_providers") {
|
||||
legacy_imported_list_policy_mode(&allowed_providers)
|
||||
} else {
|
||||
"inherit".to_string()
|
||||
}
|
||||
});
|
||||
let allowed_api_formats_mode = imported_optional_list_policy_mode(
|
||||
group.get("allowed_api_formats_mode"),
|
||||
"allowed_api_formats_mode",
|
||||
)?
|
||||
.unwrap_or_else(|| {
|
||||
if group.contains_key("allowed_api_formats") {
|
||||
legacy_imported_list_policy_mode(&allowed_api_formats)
|
||||
} else {
|
||||
"inherit".to_string()
|
||||
}
|
||||
});
|
||||
let allowed_models_mode = imported_optional_list_policy_mode(
|
||||
group.get("allowed_models_mode"),
|
||||
"allowed_models_mode",
|
||||
)?
|
||||
.unwrap_or_else(|| {
|
||||
if group.contains_key("allowed_models") {
|
||||
legacy_imported_list_policy_mode(&allowed_models)
|
||||
} else {
|
||||
"inherit".to_string()
|
||||
}
|
||||
});
|
||||
let rate_limit_mode =
|
||||
imported_optional_rate_limit_policy_mode(group.get("rate_limit_mode"), "rate_limit_mode")?
|
||||
.unwrap_or_else(|| {
|
||||
if group.contains_key("rate_limit") {
|
||||
legacy_imported_rate_limit_policy_mode(rate_limit)
|
||||
} else {
|
||||
"inherit".to_string()
|
||||
}
|
||||
});
|
||||
|
||||
let normalized_name = name.to_ascii_lowercase();
|
||||
|
||||
Ok((
|
||||
export_id,
|
||||
normalized_name,
|
||||
aether_data::repository::users::UpsertUserGroupRecord {
|
||||
name,
|
||||
description,
|
||||
priority: 0,
|
||||
allowed_providers,
|
||||
allowed_providers_mode,
|
||||
allowed_api_formats,
|
||||
allowed_api_formats_mode,
|
||||
allowed_models,
|
||||
allowed_models_mode,
|
||||
rate_limit,
|
||||
rate_limit_mode,
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
fn resolve_imported_user_group_ids(
|
||||
user: &Map<String, Value>,
|
||||
imported_group_id_map: &BTreeMap<String, String>,
|
||||
imported_group_name_map: &BTreeMap<String, String>,
|
||||
groups_by_name: &BTreeMap<String, aether_data::repository::users::StoredUserGroup>,
|
||||
) -> Result<Vec<String>, String> {
|
||||
let raw_group_ids =
|
||||
imported_string_list_from_value(user.get("group_ids"), "group_ids")?.unwrap_or_default();
|
||||
let raw_group_names = imported_string_list_from_value(user.get("group_names"), "group_names")?
|
||||
.unwrap_or_default();
|
||||
let mut group_ids = BTreeSet::new();
|
||||
for raw_group_id in raw_group_ids {
|
||||
if let Some(group_id) = imported_group_id_map.get(&raw_group_id) {
|
||||
group_ids.insert(group_id.clone());
|
||||
continue;
|
||||
}
|
||||
group_ids.insert(raw_group_id);
|
||||
}
|
||||
for raw_group_name in raw_group_names {
|
||||
let normalized_name =
|
||||
aether_data::repository::users::normalize_user_group_name(&raw_group_name)
|
||||
.to_ascii_lowercase();
|
||||
if normalized_name.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if let Some(group_id) = imported_group_name_map.get(&normalized_name) {
|
||||
group_ids.insert(group_id.clone());
|
||||
continue;
|
||||
}
|
||||
if let Some(group) = groups_by_name.get(&normalized_name) {
|
||||
group_ids.insert(group.id.clone());
|
||||
}
|
||||
}
|
||||
Ok(group_ids.into_iter().collect())
|
||||
}
|
||||
|
||||
fn normalize_imported_wallet_target(
|
||||
wallet: Option<&Map<String, Value>>,
|
||||
unlimited: bool,
|
||||
@@ -1686,6 +1885,11 @@ impl<'a> AdminAppState<'a> {
|
||||
Some(_) => return Ok(Err(invalid_request("standalone_keys 必须是数组"))),
|
||||
None => &empty,
|
||||
};
|
||||
let imported_user_groups = match root.get("user_groups") {
|
||||
Some(Value::Array(items)) => items,
|
||||
Some(_) => return Ok(Err(invalid_request("user_groups 必须是数组"))),
|
||||
None => &empty,
|
||||
};
|
||||
|
||||
let standalone_owner_id = match operator_id {
|
||||
Some(candidate) => match self.find_user_auth_by_id(candidate).await? {
|
||||
@@ -1709,6 +1913,86 @@ impl<'a> AdminAppState<'a> {
|
||||
));
|
||||
|
||||
let mut stats = AdminSystemUsersImportStats::default();
|
||||
let default_group_id = self.effective_default_user_group_id().await?;
|
||||
let existing_groups = self.list_user_groups().await?;
|
||||
let mut groups_by_name = existing_groups
|
||||
.into_iter()
|
||||
.map(|group| {
|
||||
(
|
||||
aether_data::repository::users::normalize_user_group_name(&group.name)
|
||||
.to_ascii_lowercase(),
|
||||
group,
|
||||
)
|
||||
})
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let mut imported_group_id_map = BTreeMap::<String, String>::new();
|
||||
let mut imported_group_name_map = BTreeMap::<String, String>::new();
|
||||
|
||||
for (index, raw_group) in imported_user_groups.iter().enumerate() {
|
||||
let group = match imported_object_field(raw_group, &format!("user_groups[{index}]")) {
|
||||
Ok(value) => value,
|
||||
Err(detail) => return Ok(Err(invalid_request(detail))),
|
||||
};
|
||||
let (export_id, normalized_name, record) = invalid_value!(
|
||||
build_imported_user_group_record(group, &format!("user_groups[{index}]"))
|
||||
);
|
||||
if default_group_id
|
||||
.as_deref()
|
||||
.is_some_and(|group_id| export_id.as_deref() == Some(group_id))
|
||||
|| normalized_name == "default"
|
||||
{
|
||||
if let Some(default_group_id) = default_group_id.as_ref() {
|
||||
if let Some(export_id) = export_id {
|
||||
imported_group_id_map.insert(export_id, default_group_id.clone());
|
||||
}
|
||||
imported_group_name_map.insert(normalized_name, default_group_id.clone());
|
||||
}
|
||||
stats.user_groups.skipped += 1;
|
||||
continue;
|
||||
}
|
||||
if let Some(existing) = groups_by_name.get(&normalized_name).cloned() {
|
||||
if let Some(export_id) = export_id {
|
||||
imported_group_id_map.insert(export_id, existing.id.clone());
|
||||
}
|
||||
imported_group_name_map.insert(normalized_name.clone(), existing.id.clone());
|
||||
match merge_mode {
|
||||
AdminImportMergeMode::Skip => {
|
||||
stats.user_groups.skipped += 1;
|
||||
}
|
||||
AdminImportMergeMode::Error => {
|
||||
return Ok(Err(invalid_request(format!(
|
||||
"用户组 '{}' 已存在",
|
||||
existing.name
|
||||
))));
|
||||
}
|
||||
AdminImportMergeMode::Overwrite => {
|
||||
let Some(updated) = self.update_user_group(&existing.id, record).await?
|
||||
else {
|
||||
return Ok(Err((
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
json!({ "detail": "Admin system data unavailable" }),
|
||||
)));
|
||||
};
|
||||
groups_by_name.insert(normalized_name, updated);
|
||||
stats.user_groups.updated += 1;
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
let Some(created) = self.create_user_group(record).await? else {
|
||||
return Ok(Err((
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
json!({ "detail": "Admin system data unavailable" }),
|
||||
)));
|
||||
};
|
||||
if let Some(export_id) = export_id {
|
||||
imported_group_id_map.insert(export_id, created.id.clone());
|
||||
}
|
||||
imported_group_name_map.insert(normalized_name.clone(), created.id.clone());
|
||||
groups_by_name.insert(normalized_name, created);
|
||||
stats.user_groups.created += 1;
|
||||
}
|
||||
|
||||
for (index, raw_user) in users.iter().enumerate() {
|
||||
let user = match imported_object_field(raw_user, &format!("users[{index}]")) {
|
||||
@@ -1760,6 +2044,53 @@ impl<'a> AdminAppState<'a> {
|
||||
invalid_value!(normalize_imported_user_string_list(user, "allowed_models"));
|
||||
let rate_limit =
|
||||
invalid_value!(imported_optional_i32(user.get("rate_limit"), "rate_limit"));
|
||||
let allowed_providers_mode = invalid_value!(imported_user_list_policy_mode(
|
||||
user,
|
||||
"allowed_providers_mode",
|
||||
"allowed_providers",
|
||||
&allowed_providers,
|
||||
));
|
||||
let allowed_api_formats_mode = invalid_value!(imported_user_list_policy_mode(
|
||||
user,
|
||||
"allowed_api_formats_mode",
|
||||
"allowed_api_formats",
|
||||
&allowed_api_formats,
|
||||
));
|
||||
let allowed_models_mode = invalid_value!(imported_user_list_policy_mode(
|
||||
user,
|
||||
"allowed_models_mode",
|
||||
"allowed_models",
|
||||
&allowed_models,
|
||||
));
|
||||
let rate_limit_mode = invalid_value!(imported_user_rate_limit_policy_mode(
|
||||
user,
|
||||
"rate_limit_mode",
|
||||
"rate_limit",
|
||||
rate_limit,
|
||||
));
|
||||
let imported_user_group_ids = invalid_value!(resolve_imported_user_group_ids(
|
||||
user,
|
||||
&imported_group_id_map,
|
||||
&imported_group_name_map,
|
||||
&groups_by_name,
|
||||
));
|
||||
let group_ids = if user.contains_key("group_ids") || user.contains_key("group_names") {
|
||||
let group_ids = self
|
||||
.include_default_user_group_ids(&imported_user_group_ids)
|
||||
.await?;
|
||||
if !group_ids.is_empty() {
|
||||
let existing_groups = self.list_user_groups_by_ids(&group_ids).await?;
|
||||
if existing_groups.len() != group_ids.len() {
|
||||
return Ok(Err(invalid_request(format!(
|
||||
"用户 '{}' 的用户组不存在",
|
||||
email.clone().unwrap_or(username.clone())
|
||||
))));
|
||||
}
|
||||
}
|
||||
Some(group_ids)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let is_active =
|
||||
invalid_value!(imported_optional_bool(user.get("is_active"))).unwrap_or(true);
|
||||
let model_capability_settings = invalid_value!(imported_optional_json_object(
|
||||
@@ -1883,6 +2214,31 @@ impl<'a> AdminAppState<'a> {
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
if allowed_providers_mode.is_some()
|
||||
|| allowed_api_formats_mode.is_some()
|
||||
|| allowed_models_mode.is_some()
|
||||
|| rate_limit_mode.is_some()
|
||||
{
|
||||
let updated_policy_modes = self
|
||||
.update_local_auth_user_policy_modes(
|
||||
&existing.id,
|
||||
allowed_providers_mode.clone(),
|
||||
allowed_api_formats_mode.clone(),
|
||||
allowed_models_mode.clone(),
|
||||
rate_limit_mode.clone(),
|
||||
)
|
||||
.await?;
|
||||
if updated_policy_modes.is_none() {
|
||||
return Ok(Err((
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
json!({ "detail": "Admin system data unavailable" }),
|
||||
)));
|
||||
}
|
||||
}
|
||||
if let Some(group_ids) = group_ids.as_ref() {
|
||||
self.replace_user_groups_for_user(&existing.id, group_ids)
|
||||
.await?;
|
||||
}
|
||||
self.sync_imported_user_wallet(
|
||||
&existing.id,
|
||||
&wallet_target,
|
||||
@@ -1921,6 +2277,34 @@ impl<'a> AdminAppState<'a> {
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
let created = if allowed_providers_mode.is_some()
|
||||
|| allowed_api_formats_mode.is_some()
|
||||
|| allowed_models_mode.is_some()
|
||||
|| rate_limit_mode.is_some()
|
||||
{
|
||||
let Some(updated_policy_modes) = self
|
||||
.update_local_auth_user_policy_modes(
|
||||
&created.id,
|
||||
allowed_providers_mode.clone(),
|
||||
allowed_api_formats_mode.clone(),
|
||||
allowed_models_mode.clone(),
|
||||
rate_limit_mode.clone(),
|
||||
)
|
||||
.await?
|
||||
else {
|
||||
return Ok(Err((
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
json!({ "detail": "Admin system data unavailable" }),
|
||||
)));
|
||||
};
|
||||
updated_policy_modes
|
||||
} else {
|
||||
created
|
||||
};
|
||||
if let Some(group_ids) = group_ids.as_ref() {
|
||||
self.replace_user_groups_for_user(&created.id, group_ids)
|
||||
.await?;
|
||||
}
|
||||
self.sync_imported_user_wallet(
|
||||
&created.id,
|
||||
&wallet_target,
|
||||
@@ -2454,9 +2838,10 @@ mod tests {
|
||||
#[test]
|
||||
fn users_import_requires_supported_export_version() {
|
||||
assert!(validate_imported_system_users_export_version(Some(&json!("1.3"))).is_ok());
|
||||
assert!(validate_imported_system_users_export_version(Some(&json!("1.4"))).is_ok());
|
||||
assert_eq!(
|
||||
validate_imported_system_users_export_version(Some(&json!("2.2"))).unwrap_err(),
|
||||
"不支持的用户数据版本: 2.2,支持的版本: 1.3"
|
||||
"不支持的用户数据版本: 2.2,支持的版本: 1.3, 1.4"
|
||||
);
|
||||
assert_eq!(
|
||||
validate_imported_system_users_export_version(Some(&json!(null))).unwrap_err(),
|
||||
|
||||
@@ -63,6 +63,113 @@ impl<'a> AdminAppState<'a> {
|
||||
self.app.find_user_auth_by_identifier(identifier).await
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_groups(
|
||||
&self,
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroup>, GatewayError> {
|
||||
self.app.list_user_groups().await
|
||||
}
|
||||
|
||||
pub(crate) async fn find_user_group_by_id(
|
||||
&self,
|
||||
group_id: &str,
|
||||
) -> Result<Option<aether_data::repository::users::StoredUserGroup>, GatewayError> {
|
||||
self.app.find_user_group_by_id(group_id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_groups_by_ids(
|
||||
&self,
|
||||
group_ids: &[String],
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroup>, GatewayError> {
|
||||
self.app.list_user_groups_by_ids(group_ids).await
|
||||
}
|
||||
|
||||
pub(crate) async fn create_user_group(
|
||||
&self,
|
||||
record: aether_data::repository::users::UpsertUserGroupRecord,
|
||||
) -> Result<Option<aether_data::repository::users::StoredUserGroup>, GatewayError> {
|
||||
self.app.create_user_group(record).await
|
||||
}
|
||||
|
||||
pub(crate) async fn update_user_group(
|
||||
&self,
|
||||
group_id: &str,
|
||||
record: aether_data::repository::users::UpsertUserGroupRecord,
|
||||
) -> Result<Option<aether_data::repository::users::StoredUserGroup>, GatewayError> {
|
||||
self.app.update_user_group(group_id, record).await
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_user_group(&self, group_id: &str) -> Result<bool, GatewayError> {
|
||||
self.app.delete_user_group(group_id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_group_members(
|
||||
&self,
|
||||
group_id: &str,
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroupMember>, GatewayError> {
|
||||
self.app.list_user_group_members(group_id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn replace_user_group_members(
|
||||
&self,
|
||||
group_id: &str,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroupMember>, GatewayError> {
|
||||
self.app
|
||||
.replace_user_group_members(group_id, user_ids)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_groups_for_user(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroup>, GatewayError> {
|
||||
self.app.list_user_groups_for_user(user_id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_group_memberships_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroupMembership>, GatewayError> {
|
||||
self.app
|
||||
.list_user_group_memberships_by_user_ids(user_ids)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn replace_user_groups_for_user(
|
||||
&self,
|
||||
user_id: &str,
|
||||
group_ids: &[String],
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroup>, GatewayError> {
|
||||
self.app
|
||||
.replace_user_groups_for_user(user_id, group_ids)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn include_default_user_group_ids(
|
||||
&self,
|
||||
group_ids: &[String],
|
||||
) -> Result<Vec<String>, GatewayError> {
|
||||
self.app.include_default_user_group_ids(group_ids).await
|
||||
}
|
||||
|
||||
pub(crate) async fn effective_default_user_group_id(
|
||||
&self,
|
||||
) -> Result<Option<String>, GatewayError> {
|
||||
self.app.effective_default_user_group_id().await
|
||||
}
|
||||
|
||||
pub(crate) async fn add_user_to_group(
|
||||
&self,
|
||||
group_id: &str,
|
||||
user_id: &str,
|
||||
) -> Result<bool, GatewayError> {
|
||||
self.app.add_user_to_group(group_id, user_id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn add_all_users_to_group(&self, group_id: &str) -> Result<(), GatewayError> {
|
||||
self.app.add_all_users_to_group(group_id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn is_other_user_auth_email_taken(
|
||||
&self,
|
||||
email: &str,
|
||||
@@ -187,6 +294,25 @@ impl<'a> AdminAppState<'a> {
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn update_local_auth_user_policy_modes(
|
||||
&self,
|
||||
user_id: &str,
|
||||
allowed_providers_mode: Option<String>,
|
||||
allowed_api_formats_mode: Option<String>,
|
||||
allowed_models_mode: Option<String>,
|
||||
rate_limit_mode: Option<String>,
|
||||
) -> Result<Option<aether_data::repository::users::StoredUserAuthRecord>, GatewayError> {
|
||||
self.app
|
||||
.update_local_auth_user_policy_modes(
|
||||
user_id,
|
||||
allowed_providers_mode,
|
||||
allowed_api_formats_mode,
|
||||
allowed_models_mode,
|
||||
rate_limit_mode,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn update_auth_user_wallet_limit_mode(
|
||||
&self,
|
||||
user_id: &str,
|
||||
|
||||
@@ -22,11 +22,14 @@ struct AdminUserSelectionFilters {
|
||||
role: Option<String>,
|
||||
#[serde(default)]
|
||||
is_active: Option<bool>,
|
||||
#[serde(default)]
|
||||
group_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
struct AdminUserSelectionRequest {
|
||||
user_ids: Vec<String>,
|
||||
group_ids: Vec<String>,
|
||||
filters: Option<AdminUserSelectionFilters>,
|
||||
filters_scope_present: bool,
|
||||
}
|
||||
@@ -51,6 +54,7 @@ struct NormalizedAdminUserSelectionFilters {
|
||||
search: Option<String>,
|
||||
role: Option<String>,
|
||||
is_active: Option<bool>,
|
||||
group_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
@@ -60,12 +64,22 @@ struct AdminUserSelectionItem {
|
||||
email: Option<String>,
|
||||
role: String,
|
||||
is_active: bool,
|
||||
matched_by: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
struct AdminUserSelectionWarning {
|
||||
#[serde(rename = "type")]
|
||||
warning_type: String,
|
||||
group_id: Option<String>,
|
||||
message: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
struct ResolvedAdminUserSelection {
|
||||
items: Vec<AdminUserSelectionItem>,
|
||||
missing_user_ids: Vec<String>,
|
||||
warnings: Vec<AdminUserSelectionWarning>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
@@ -112,6 +126,7 @@ pub(in super::super) async fn build_admin_resolve_user_selection_response(
|
||||
Ok(Json(json!({
|
||||
"total": resolved.items.len(),
|
||||
"items": resolved.items,
|
||||
"warnings": resolved.warnings,
|
||||
}))
|
||||
.into_response())
|
||||
}
|
||||
@@ -229,6 +244,7 @@ pub(in super::super) async fn build_admin_user_batch_action_response(
|
||||
"success": success,
|
||||
"failed": failed,
|
||||
"failures": failures,
|
||||
"warnings": resolved.warnings,
|
||||
"action": request.action.trim().to_ascii_lowercase(),
|
||||
"modified_fields": mutation.modified_fields,
|
||||
}))
|
||||
@@ -284,6 +300,11 @@ fn parse_selection_request_value(value: Value) -> Result<AdminUserSelectionReque
|
||||
Some(value) => serde_json::from_value::<Vec<String>>(value.clone())
|
||||
.map_err(|_| "user_ids 必须是字符串数组".to_string())?,
|
||||
};
|
||||
let group_ids = match map.get("group_ids") {
|
||||
None | Some(Value::Null) => Vec::new(),
|
||||
Some(value) => serde_json::from_value::<Vec<String>>(value.clone())
|
||||
.map_err(|_| "group_ids 必须是字符串数组".to_string())?,
|
||||
};
|
||||
|
||||
let (filters_scope_present, filters) = match map.get("filters") {
|
||||
Some(Value::Object(_)) => {
|
||||
@@ -299,6 +320,7 @@ fn parse_selection_request_value(value: Value) -> Result<AdminUserSelectionReque
|
||||
|
||||
Ok(AdminUserSelectionRequest {
|
||||
user_ids,
|
||||
group_ids,
|
||||
filters,
|
||||
filters_scope_present,
|
||||
})
|
||||
@@ -310,12 +332,36 @@ async fn resolve_admin_user_selection(
|
||||
) -> Result<ResolvedAdminUserSelection, String> {
|
||||
let filters = normalize_selection_filters(selection.filters)?;
|
||||
let explicit_user_ids = normalize_user_ids(selection.user_ids);
|
||||
if explicit_user_ids.is_empty() && !selection.filters_scope_present {
|
||||
return Err("至少需要选择一个用户或明确提供筛选条件".to_string());
|
||||
let explicit_group_ids = normalize_user_ids(selection.group_ids);
|
||||
if explicit_user_ids.is_empty()
|
||||
&& explicit_group_ids.is_empty()
|
||||
&& !selection.filters_scope_present
|
||||
{
|
||||
return Err("至少需要选择一个用户、用户组或明确提供筛选条件".to_string());
|
||||
}
|
||||
let should_resolve_filters = selection.filters_scope_present;
|
||||
let mut items_by_id = BTreeMap::new();
|
||||
let mut missing_user_ids = Vec::new();
|
||||
let mut warnings = Vec::new();
|
||||
|
||||
if !explicit_group_ids.is_empty() {
|
||||
let groups = state
|
||||
.list_user_groups_by_ids(&explicit_group_ids)
|
||||
.await
|
||||
.map_err(|_| "用户分组数据不可用".to_string())?;
|
||||
let found_group_ids = groups
|
||||
.iter()
|
||||
.map(|group| group.id.clone())
|
||||
.collect::<BTreeSet<_>>();
|
||||
let missing_group_ids = explicit_group_ids
|
||||
.iter()
|
||||
.filter(|group_id| !found_group_ids.contains(*group_id))
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
if !missing_group_ids.is_empty() {
|
||||
return Err(format!("用户分组不存在: {}", missing_group_ids.join(", ")));
|
||||
}
|
||||
}
|
||||
|
||||
if !explicit_user_ids.is_empty() {
|
||||
let users = state
|
||||
@@ -325,15 +371,14 @@ async fn resolve_admin_user_selection(
|
||||
for user_id in explicit_user_ids {
|
||||
match users.get(&user_id).filter(|user| !user.is_deleted) {
|
||||
Some(user) => {
|
||||
items_by_id.insert(
|
||||
insert_or_update_selection_item(
|
||||
&mut items_by_id,
|
||||
user.id.clone(),
|
||||
AdminUserSelectionItem {
|
||||
user_id: user.id.clone(),
|
||||
username: user.username.clone(),
|
||||
email: user.email.clone(),
|
||||
role: user.role.clone(),
|
||||
is_active: user.is_active,
|
||||
},
|
||||
user.username.clone(),
|
||||
user.email.clone(),
|
||||
user.role.clone(),
|
||||
user.is_active,
|
||||
"direct".to_string(),
|
||||
);
|
||||
}
|
||||
None => missing_user_ids.push(user_id),
|
||||
@@ -341,24 +386,71 @@ async fn resolve_admin_user_selection(
|
||||
}
|
||||
}
|
||||
|
||||
if should_resolve_filters {
|
||||
let users = state
|
||||
.list_export_users()
|
||||
for group_id in &explicit_group_ids {
|
||||
let members = state
|
||||
.list_user_group_members(group_id)
|
||||
.await
|
||||
.map_err(|_| "用户数据不可用".to_string())?;
|
||||
.map_err(|_| "用户分组成员数据不可用".to_string())?;
|
||||
let mut matched_count = 0usize;
|
||||
for member in members.into_iter().filter(|member| !member.is_deleted) {
|
||||
matched_count += 1;
|
||||
insert_or_update_selection_item(
|
||||
&mut items_by_id,
|
||||
member.user_id,
|
||||
member.username,
|
||||
member.email,
|
||||
member.role,
|
||||
member.is_active,
|
||||
format!("group:{group_id}"),
|
||||
);
|
||||
}
|
||||
if matched_count == 0 {
|
||||
warnings.push(AdminUserSelectionWarning {
|
||||
warning_type: "empty_group".to_string(),
|
||||
group_id: Some(group_id.clone()),
|
||||
message: "分组内没有可操作用户".to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if should_resolve_filters {
|
||||
let users = if filters.as_ref().is_some_and(|filters| {
|
||||
filters.search.is_some()
|
||||
|| filters.role.is_some()
|
||||
|| filters.is_active.is_some()
|
||||
|| filters.group_id.is_some()
|
||||
}) {
|
||||
state
|
||||
.list_export_users_page(&aether_data::repository::users::UserExportListQuery {
|
||||
skip: 0,
|
||||
limit: 100_000,
|
||||
role: filters.as_ref().and_then(|filters| filters.role.clone()),
|
||||
is_active: filters.as_ref().and_then(|filters| filters.is_active),
|
||||
search: filters.as_ref().and_then(|filters| filters.search.clone()),
|
||||
group_id: filters
|
||||
.as_ref()
|
||||
.and_then(|filters| filters.group_id.clone()),
|
||||
})
|
||||
.await
|
||||
.map_err(|_| "用户数据不可用".to_string())?
|
||||
} else {
|
||||
state
|
||||
.list_export_users()
|
||||
.await
|
||||
.map_err(|_| "用户数据不可用".to_string())?
|
||||
};
|
||||
for user in users
|
||||
.into_iter()
|
||||
.filter(|user| admin_user_matches_filters(user, filters.as_ref()))
|
||||
{
|
||||
items_by_id.insert(
|
||||
user.id.clone(),
|
||||
AdminUserSelectionItem {
|
||||
user_id: user.id,
|
||||
username: user.username,
|
||||
email: user.email,
|
||||
role: user.role,
|
||||
is_active: user.is_active,
|
||||
},
|
||||
insert_or_update_selection_item(
|
||||
&mut items_by_id,
|
||||
user.id,
|
||||
user.username,
|
||||
user.email,
|
||||
user.role,
|
||||
user.is_active,
|
||||
"filter".to_string(),
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -374,9 +466,41 @@ async fn resolve_admin_user_selection(
|
||||
Ok(ResolvedAdminUserSelection {
|
||||
items,
|
||||
missing_user_ids,
|
||||
warnings,
|
||||
})
|
||||
}
|
||||
|
||||
fn insert_or_update_selection_item(
|
||||
items_by_id: &mut BTreeMap<String, AdminUserSelectionItem>,
|
||||
user_id: String,
|
||||
username: String,
|
||||
email: Option<String>,
|
||||
role: String,
|
||||
is_active: bool,
|
||||
matched_by: String,
|
||||
) {
|
||||
match items_by_id.get_mut(&user_id) {
|
||||
Some(item) => {
|
||||
if !item.matched_by.iter().any(|value| value == &matched_by) {
|
||||
item.matched_by.push(matched_by);
|
||||
}
|
||||
}
|
||||
None => {
|
||||
items_by_id.insert(
|
||||
user_id.clone(),
|
||||
AdminUserSelectionItem {
|
||||
user_id,
|
||||
username,
|
||||
email,
|
||||
role,
|
||||
is_active,
|
||||
matched_by: vec![matched_by],
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_selection_filters(
|
||||
filters: Option<AdminUserSelectionFilters>,
|
||||
) -> Result<Option<NormalizedAdminUserSelectionFilters>, String> {
|
||||
@@ -401,6 +525,10 @@ fn normalize_selection_filters(
|
||||
search,
|
||||
role,
|
||||
is_active: filters.is_active,
|
||||
group_id: filters
|
||||
.group_id
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty()),
|
||||
}))
|
||||
}
|
||||
|
||||
|
||||
491
apps/aether-gateway/src/handlers/admin/users/groups.rs
Normal file
491
apps/aether-gateway/src/handlers/admin/users/groups.rs
Normal file
@@ -0,0 +1,491 @@
|
||||
use super::{
|
||||
build_admin_users_bad_request_response, build_admin_users_read_only_response,
|
||||
format_optional_datetime_iso8601, normalize_admin_user_api_formats,
|
||||
normalize_admin_user_string_list,
|
||||
};
|
||||
use crate::constants::DEFAULT_USER_GROUP_CONFIG_KEY;
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::handlers::admin::shared::attach_admin_audit_response;
|
||||
use crate::GatewayError;
|
||||
use axum::{
|
||||
body::Body,
|
||||
http,
|
||||
response::{IntoResponse, Response},
|
||||
Json,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
struct AdminUserGroupPayload {
|
||||
name: String,
|
||||
#[serde(default)]
|
||||
description: Option<String>,
|
||||
#[serde(default)]
|
||||
allowed_providers: Option<Vec<String>>,
|
||||
#[serde(default = "default_list_mode")]
|
||||
allowed_providers_mode: String,
|
||||
#[serde(default)]
|
||||
allowed_api_formats: Option<Vec<String>>,
|
||||
#[serde(default = "default_list_mode")]
|
||||
allowed_api_formats_mode: String,
|
||||
#[serde(default)]
|
||||
allowed_models: Option<Vec<String>>,
|
||||
#[serde(default = "default_list_mode")]
|
||||
allowed_models_mode: String,
|
||||
#[serde(default)]
|
||||
rate_limit: Option<i32>,
|
||||
#[serde(default = "default_rate_limit_mode")]
|
||||
rate_limit_mode: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
struct AdminUserGroupMembersPayload {
|
||||
user_ids: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
struct AdminDefaultUserGroupPayload {
|
||||
#[serde(default)]
|
||||
group_id: Option<String>,
|
||||
}
|
||||
|
||||
pub(in super::super) async fn build_admin_list_user_groups_response(
|
||||
state: &AdminAppState<'_>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let default_group_id = read_default_user_group_id(state).await?;
|
||||
let items = state
|
||||
.list_user_groups()
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|group| user_group_payload(group, default_group_id.as_deref()))
|
||||
.collect::<Vec<_>>();
|
||||
Ok(Json(json!({
|
||||
"items": items,
|
||||
"default_group_id": default_group_id,
|
||||
}))
|
||||
.into_response())
|
||||
}
|
||||
|
||||
pub(in super::super) async fn build_admin_create_user_group_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_body: Option<&axum::body::Bytes>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
if !state.has_auth_user_write_capability() {
|
||||
return Ok(build_admin_users_read_only_response(
|
||||
"当前为只读模式,无法创建用户分组",
|
||||
));
|
||||
}
|
||||
let record = match parse_group_record(request_body) {
|
||||
Ok(value) => value,
|
||||
Err(detail) => return Ok(bad_request_owned(detail)),
|
||||
};
|
||||
let group = match state.create_user_group(record).await {
|
||||
Ok(Some(group)) => group,
|
||||
Ok(None) => {
|
||||
return Ok(build_admin_users_read_only_response(
|
||||
"当前为只读模式,无法创建用户分组",
|
||||
))
|
||||
}
|
||||
Err(err) if is_duplicate_group_name_error(&err) => {
|
||||
return Ok(bad_request_owned("用户分组名称已存在".to_string()))
|
||||
}
|
||||
Err(err) => return Err(err),
|
||||
};
|
||||
let default_group_id = read_default_user_group_id(state).await?;
|
||||
Ok(attach_admin_audit_response(
|
||||
Json(user_group_payload(group, default_group_id.as_deref())).into_response(),
|
||||
"admin_user_group_created",
|
||||
"create_user_group",
|
||||
"user_group",
|
||||
"user_groups",
|
||||
))
|
||||
}
|
||||
|
||||
pub(in super::super) async fn build_admin_update_user_group_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
request_body: Option<&axum::body::Bytes>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
if !state.has_auth_user_write_capability() {
|
||||
return Ok(build_admin_users_read_only_response(
|
||||
"当前为只读模式,无法更新用户分组",
|
||||
));
|
||||
}
|
||||
let Some(group_id) = user_group_id_from_path(request_context.path()) else {
|
||||
return Ok(build_admin_users_bad_request_response("缺少 group_id"));
|
||||
};
|
||||
let record = match parse_group_record(request_body) {
|
||||
Ok(value) => value,
|
||||
Err(detail) => return Ok(bad_request_owned(detail)),
|
||||
};
|
||||
if read_default_user_group_id(state).await?.as_deref() == Some(group_id.as_str())
|
||||
&& !is_unrestricted_default_group_record(&record)
|
||||
{
|
||||
return Ok(bad_request_owned("默认用户组不能配置访问限制".to_string()));
|
||||
}
|
||||
let group = match state.update_user_group(&group_id, record).await {
|
||||
Ok(Some(group)) => group,
|
||||
Ok(None) => return Ok(not_found("用户分组不存在")),
|
||||
Err(err) if is_duplicate_group_name_error(&err) => {
|
||||
return Ok(bad_request_owned("用户分组名称已存在".to_string()))
|
||||
}
|
||||
Err(err) => return Err(err),
|
||||
};
|
||||
let default_group_id = read_default_user_group_id(state).await?;
|
||||
Ok(attach_admin_audit_response(
|
||||
Json(user_group_payload(group, default_group_id.as_deref())).into_response(),
|
||||
"admin_user_group_updated",
|
||||
"update_user_group",
|
||||
"user_group",
|
||||
&group_id,
|
||||
))
|
||||
}
|
||||
|
||||
pub(in super::super) async fn build_admin_delete_user_group_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
if !state.has_auth_user_write_capability() {
|
||||
return Ok(build_admin_users_read_only_response(
|
||||
"当前为只读模式,无法删除用户分组",
|
||||
));
|
||||
}
|
||||
let Some(group_id) = user_group_id_from_path(request_context.path()) else {
|
||||
return Ok(build_admin_users_bad_request_response("缺少 group_id"));
|
||||
};
|
||||
if read_default_user_group_id(state).await?.as_deref() == Some(group_id.as_str()) {
|
||||
return Ok(bad_request_owned("默认用户组不能删除".to_string()));
|
||||
}
|
||||
if !state.delete_user_group(&group_id).await? {
|
||||
return Ok(not_found("用户分组不存在"));
|
||||
}
|
||||
Ok(attach_admin_audit_response(
|
||||
Json(json!({ "deleted": true })).into_response(),
|
||||
"admin_user_group_deleted",
|
||||
"delete_user_group",
|
||||
"user_group",
|
||||
&group_id,
|
||||
))
|
||||
}
|
||||
|
||||
pub(in super::super) async fn build_admin_list_user_group_members_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let Some(group_id) = user_group_member_group_id_from_path(request_context.path()) else {
|
||||
return Ok(build_admin_users_bad_request_response("缺少 group_id"));
|
||||
};
|
||||
if state.find_user_group_by_id(&group_id).await?.is_none() {
|
||||
return Ok(not_found("用户分组不存在"));
|
||||
}
|
||||
let items = state
|
||||
.list_user_group_members(&group_id)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|member| {
|
||||
json!({
|
||||
"group_id": member.group_id,
|
||||
"user_id": member.user_id,
|
||||
"username": member.username,
|
||||
"email": member.email,
|
||||
"role": member.role,
|
||||
"is_active": member.is_active,
|
||||
"is_deleted": member.is_deleted,
|
||||
"created_at": format_optional_datetime_iso8601(member.created_at),
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
Ok(Json(json!({ "items": items })).into_response())
|
||||
}
|
||||
|
||||
pub(in super::super) async fn build_admin_replace_user_group_members_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
request_body: Option<&axum::body::Bytes>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
if !state.has_auth_user_write_capability() {
|
||||
return Ok(build_admin_users_read_only_response(
|
||||
"当前为只读模式,无法更新分组成员",
|
||||
));
|
||||
}
|
||||
let Some(group_id) = user_group_member_group_id_from_path(request_context.path()) else {
|
||||
return Ok(build_admin_users_bad_request_response("缺少 group_id"));
|
||||
};
|
||||
if state.find_user_group_by_id(&group_id).await?.is_none() {
|
||||
return Ok(not_found("用户分组不存在"));
|
||||
}
|
||||
if read_default_user_group_id(state).await?.as_deref() == Some(group_id.as_str()) {
|
||||
return Ok(bad_request_owned("默认用户组成员由系统维护".to_string()));
|
||||
}
|
||||
let payload = match parse_members_payload(request_body) {
|
||||
Ok(value) => value,
|
||||
Err(detail) => return Ok(bad_request_owned(detail)),
|
||||
};
|
||||
let user_ids = normalize_ids(payload.user_ids);
|
||||
let known_users = state.resolve_auth_user_summaries_by_ids(&user_ids).await?;
|
||||
if known_users.len() != user_ids.len() {
|
||||
return Ok(bad_request_owned("成员包含不存在的用户".to_string()));
|
||||
}
|
||||
let items = state
|
||||
.replace_user_group_members(&group_id, &user_ids)
|
||||
.await?;
|
||||
Ok(attach_admin_audit_response(
|
||||
Json(json!({
|
||||
"items": items.into_iter().map(|member| json!({
|
||||
"group_id": member.group_id,
|
||||
"user_id": member.user_id,
|
||||
"username": member.username,
|
||||
"email": member.email,
|
||||
"role": member.role,
|
||||
"is_active": member.is_active,
|
||||
"is_deleted": member.is_deleted,
|
||||
"created_at": format_optional_datetime_iso8601(member.created_at),
|
||||
})).collect::<Vec<_>>()
|
||||
}))
|
||||
.into_response(),
|
||||
"admin_user_group_members_updated",
|
||||
"update_user_group_members",
|
||||
"user_group",
|
||||
&group_id,
|
||||
))
|
||||
}
|
||||
|
||||
pub(in super::super) async fn build_admin_set_default_user_group_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_body: Option<&axum::body::Bytes>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
if !state.has_auth_user_write_capability() {
|
||||
return Ok(build_admin_users_read_only_response(
|
||||
"当前为只读模式,无法设置默认用户组",
|
||||
));
|
||||
}
|
||||
let payload = match request_body {
|
||||
Some(body) if !body.is_empty() => {
|
||||
serde_json::from_slice::<AdminDefaultUserGroupPayload>(body)
|
||||
.map_err(|_| "请求数据验证失败".to_string())
|
||||
}
|
||||
_ => Err("请求数据验证失败".to_string()),
|
||||
};
|
||||
let payload = match payload {
|
||||
Ok(value) => value,
|
||||
Err(detail) => return Ok(bad_request_owned(detail)),
|
||||
};
|
||||
let group_id = payload
|
||||
.group_id
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty());
|
||||
if let Some(group_id) = group_id.as_deref() {
|
||||
let Some(group) = state.find_user_group_by_id(group_id).await? else {
|
||||
return Ok(bad_request_owned("默认用户组不存在".to_string()));
|
||||
};
|
||||
if !is_unrestricted_default_group(&group) {
|
||||
return Ok(bad_request_owned("默认用户组不能配置访问限制".to_string()));
|
||||
}
|
||||
state
|
||||
.upsert_system_config_json_value(
|
||||
DEFAULT_USER_GROUP_CONFIG_KEY,
|
||||
&json!(group_id),
|
||||
Some("Default group for self-registered users"),
|
||||
)
|
||||
.await?;
|
||||
} else {
|
||||
state
|
||||
.delete_system_config_value(DEFAULT_USER_GROUP_CONFIG_KEY)
|
||||
.await?;
|
||||
}
|
||||
let effective_group_id = read_default_user_group_id(state).await?;
|
||||
if let Some(group_id) = effective_group_id.as_deref() {
|
||||
state.add_all_users_to_group(group_id).await?;
|
||||
}
|
||||
Ok(attach_admin_audit_response(
|
||||
Json(json!({ "default_group_id": effective_group_id })).into_response(),
|
||||
"admin_default_user_group_set",
|
||||
"set_default_user_group",
|
||||
"user_group",
|
||||
group_id.as_deref().unwrap_or("default_user_group"),
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) async fn read_default_user_group_id(
|
||||
state: &AdminAppState<'_>,
|
||||
) -> Result<Option<String>, GatewayError> {
|
||||
state.effective_default_user_group_id().await
|
||||
}
|
||||
|
||||
fn parse_group_record(
|
||||
request_body: Option<&axum::body::Bytes>,
|
||||
) -> Result<aether_data::repository::users::UpsertUserGroupRecord, String> {
|
||||
let Some(body) = request_body.filter(|body| !body.is_empty()) else {
|
||||
return Err("请求数据验证失败".to_string());
|
||||
};
|
||||
let payload = serde_json::from_slice::<AdminUserGroupPayload>(body)
|
||||
.map_err(|_| "请求数据验证失败".to_string())?;
|
||||
let name = aether_data::repository::users::normalize_user_group_name(&payload.name);
|
||||
if name.is_empty() {
|
||||
return Err("分组名称不能为空".to_string());
|
||||
}
|
||||
if payload.rate_limit.is_some_and(|value| value < 0) {
|
||||
return Err("rate_limit 必须大于等于 0".to_string());
|
||||
}
|
||||
let allowed_providers =
|
||||
normalize_admin_user_string_list(payload.allowed_providers, "allowed_providers")?;
|
||||
let allowed_api_formats = normalize_admin_user_api_formats(payload.allowed_api_formats)?;
|
||||
let allowed_models =
|
||||
normalize_admin_user_string_list(payload.allowed_models, "allowed_models")?;
|
||||
Ok(aether_data::repository::users::UpsertUserGroupRecord {
|
||||
name,
|
||||
description: payload
|
||||
.description
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty()),
|
||||
priority: 0,
|
||||
allowed_providers,
|
||||
allowed_providers_mode: normalize_list_mode(&payload.allowed_providers_mode)?,
|
||||
allowed_api_formats,
|
||||
allowed_api_formats_mode: normalize_list_mode(&payload.allowed_api_formats_mode)?,
|
||||
allowed_models,
|
||||
allowed_models_mode: normalize_list_mode(&payload.allowed_models_mode)?,
|
||||
rate_limit: payload.rate_limit,
|
||||
rate_limit_mode: normalize_rate_mode(&payload.rate_limit_mode)?,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_members_payload(
|
||||
request_body: Option<&axum::body::Bytes>,
|
||||
) -> Result<AdminUserGroupMembersPayload, String> {
|
||||
let Some(body) = request_body.filter(|body| !body.is_empty()) else {
|
||||
return Err("请求数据验证失败".to_string());
|
||||
};
|
||||
serde_json::from_slice::<AdminUserGroupMembersPayload>(body)
|
||||
.map_err(|_| "请求数据验证失败".to_string())
|
||||
}
|
||||
|
||||
fn user_group_payload(
|
||||
group: aether_data::repository::users::StoredUserGroup,
|
||||
default_group_id: Option<&str>,
|
||||
) -> serde_json::Value {
|
||||
json!({
|
||||
"id": group.id,
|
||||
"name": group.name,
|
||||
"normalized_name": group.normalized_name,
|
||||
"description": group.description,
|
||||
"allowed_providers": group.allowed_providers,
|
||||
"allowed_providers_mode": group.allowed_providers_mode,
|
||||
"allowed_api_formats": group.allowed_api_formats,
|
||||
"allowed_api_formats_mode": group.allowed_api_formats_mode,
|
||||
"allowed_models": group.allowed_models,
|
||||
"allowed_models_mode": group.allowed_models_mode,
|
||||
"rate_limit": group.rate_limit,
|
||||
"rate_limit_mode": group.rate_limit_mode,
|
||||
"is_default": default_group_id == Some(group.id.as_str()),
|
||||
"created_at": format_optional_datetime_iso8601(group.created_at),
|
||||
"updated_at": format_optional_datetime_iso8601(group.updated_at),
|
||||
})
|
||||
}
|
||||
|
||||
fn normalize_list_mode(value: &str) -> Result<String, String> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"inherit" | "unrestricted" | "specific" | "deny_all" => {
|
||||
Ok(value.trim().to_ascii_lowercase())
|
||||
}
|
||||
_ => Err("权限列表模式不合法".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_rate_mode(value: &str) -> Result<String, String> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"inherit" | "system" | "custom" => Ok(value.trim().to_ascii_lowercase()),
|
||||
_ => Err("限速模式不合法".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn default_list_mode() -> String {
|
||||
"inherit".to_string()
|
||||
}
|
||||
|
||||
fn default_rate_limit_mode() -> String {
|
||||
"inherit".to_string()
|
||||
}
|
||||
|
||||
fn normalize_ids(values: Vec<String>) -> Vec<String> {
|
||||
values
|
||||
.into_iter()
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
.collect::<std::collections::BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn is_unrestricted_default_group(group: &aether_data::repository::users::StoredUserGroup) -> bool {
|
||||
list_mode_has_no_restriction(&group.allowed_providers_mode)
|
||||
&& list_mode_has_no_restriction(&group.allowed_api_formats_mode)
|
||||
&& list_mode_has_no_restriction(&group.allowed_models_mode)
|
||||
&& rate_mode_has_no_restriction(&group.rate_limit_mode)
|
||||
}
|
||||
|
||||
fn is_unrestricted_default_group_record(
|
||||
record: &aether_data::repository::users::UpsertUserGroupRecord,
|
||||
) -> bool {
|
||||
list_mode_has_no_restriction(&record.allowed_providers_mode)
|
||||
&& list_mode_has_no_restriction(&record.allowed_api_formats_mode)
|
||||
&& list_mode_has_no_restriction(&record.allowed_models_mode)
|
||||
&& rate_mode_has_no_restriction(&record.rate_limit_mode)
|
||||
}
|
||||
|
||||
fn list_mode_has_no_restriction(mode: &str) -> bool {
|
||||
matches!(mode, "inherit" | "unrestricted")
|
||||
}
|
||||
|
||||
fn rate_mode_has_no_restriction(mode: &str) -> bool {
|
||||
matches!(mode, "inherit" | "system")
|
||||
}
|
||||
|
||||
fn user_group_id_from_path(request_path: &str) -> Option<String> {
|
||||
let value = request_path
|
||||
.strip_prefix("/api/admin/user-groups/")?
|
||||
.trim()
|
||||
.trim_matches('/')
|
||||
.to_string();
|
||||
if value.is_empty() || value.contains('/') || value == "default" {
|
||||
None
|
||||
} else {
|
||||
Some(value)
|
||||
}
|
||||
}
|
||||
|
||||
fn user_group_member_group_id_from_path(request_path: &str) -> Option<String> {
|
||||
let value = request_path
|
||||
.strip_prefix("/api/admin/user-groups/")?
|
||||
.trim()
|
||||
.trim_matches('/');
|
||||
let group_id = value.strip_suffix("/members")?.trim_matches('/');
|
||||
if group_id.is_empty() || group_id.contains('/') {
|
||||
None
|
||||
} else {
|
||||
Some(group_id.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
fn bad_request_owned(detail: String) -> Response<Body> {
|
||||
(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": detail })),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
fn not_found(detail: &'static str) -> Response<Body> {
|
||||
(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
Json(json!({ "detail": detail })),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
fn is_duplicate_group_name_error(err: &GatewayError) -> bool {
|
||||
match err {
|
||||
GatewayError::Internal(message) => message.contains("duplicate user group name"),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
@@ -1,10 +1,12 @@
|
||||
use super::super::{
|
||||
admin_default_user_initial_gift, build_admin_users_read_only_response,
|
||||
normalize_admin_optional_user_email, normalize_admin_user_api_formats,
|
||||
normalize_admin_user_role, normalize_admin_user_string_list, normalize_admin_username,
|
||||
validate_admin_user_password, AdminCreateUserRequest,
|
||||
legacy_admin_list_policy_mode, legacy_admin_rate_limit_policy_mode,
|
||||
normalize_admin_list_policy_mode, normalize_admin_optional_user_email,
|
||||
normalize_admin_rate_limit_policy_mode, normalize_admin_user_api_formats,
|
||||
normalize_admin_user_group_ids, normalize_admin_user_role, normalize_admin_user_string_list,
|
||||
normalize_admin_username, validate_admin_user_password, AdminCreateUserRequest,
|
||||
};
|
||||
use super::support::{admin_user_password_policy, build_admin_user_payload};
|
||||
use super::support::{admin_user_password_policy, build_admin_user_payload_with_groups};
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::handlers::admin::shared::attach_admin_audit_response;
|
||||
use crate::GatewayError;
|
||||
@@ -136,6 +138,75 @@ pub(in super::super) async fn build_admin_create_user_response(
|
||||
.into_response())
|
||||
}
|
||||
};
|
||||
let allowed_providers_mode = match payload.allowed_providers_mode.as_deref() {
|
||||
Some(value) => match normalize_admin_list_policy_mode(value) {
|
||||
Ok(value) => value,
|
||||
Err(detail) => {
|
||||
return Ok((
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": detail })),
|
||||
)
|
||||
.into_response())
|
||||
}
|
||||
},
|
||||
None => legacy_admin_list_policy_mode(&allowed_providers),
|
||||
};
|
||||
let allowed_api_formats_mode = match payload.allowed_api_formats_mode.as_deref() {
|
||||
Some(value) => match normalize_admin_list_policy_mode(value) {
|
||||
Ok(value) => value,
|
||||
Err(detail) => {
|
||||
return Ok((
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": detail })),
|
||||
)
|
||||
.into_response())
|
||||
}
|
||||
},
|
||||
None => legacy_admin_list_policy_mode(&allowed_api_formats),
|
||||
};
|
||||
let allowed_models_mode = match payload.allowed_models_mode.as_deref() {
|
||||
Some(value) => match normalize_admin_list_policy_mode(value) {
|
||||
Ok(value) => value,
|
||||
Err(detail) => {
|
||||
return Ok((
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": detail })),
|
||||
)
|
||||
.into_response())
|
||||
}
|
||||
},
|
||||
None => legacy_admin_list_policy_mode(&allowed_models),
|
||||
};
|
||||
let rate_limit_mode = match payload.rate_limit_mode.as_deref() {
|
||||
Some(value) => match normalize_admin_rate_limit_policy_mode(value) {
|
||||
Ok(value) => value,
|
||||
Err(detail) => {
|
||||
return Ok((
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": detail })),
|
||||
)
|
||||
.into_response())
|
||||
}
|
||||
},
|
||||
None => legacy_admin_rate_limit_policy_mode(payload.rate_limit),
|
||||
};
|
||||
let requested_group_ids = normalize_admin_user_group_ids(payload.group_ids);
|
||||
let group_ids = state
|
||||
.include_default_user_group_ids(&requested_group_ids)
|
||||
.await?;
|
||||
let groups = if group_ids.is_empty() {
|
||||
Vec::new()
|
||||
} else {
|
||||
let groups = state.list_user_groups_by_ids(&group_ids).await?;
|
||||
if groups.len() != group_ids.len() {
|
||||
return Ok((
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": "用户分组不存在" })),
|
||||
)
|
||||
.into_response());
|
||||
}
|
||||
groups
|
||||
};
|
||||
|
||||
if let Some(email) = email.as_deref() {
|
||||
if state.find_user_auth_by_identifier(email).await?.is_some() {
|
||||
@@ -209,12 +280,33 @@ pub(in super::super) async fn build_admin_create_user_response(
|
||||
"当前为只读模式,无法初始化用户钱包",
|
||||
));
|
||||
}
|
||||
let Some(user) = state
|
||||
.update_local_auth_user_policy_modes(
|
||||
&user.id,
|
||||
Some(allowed_providers_mode.clone()),
|
||||
Some(allowed_api_formats_mode.clone()),
|
||||
Some(allowed_models_mode.clone()),
|
||||
Some(rate_limit_mode.clone()),
|
||||
)
|
||||
.await?
|
||||
else {
|
||||
return Ok(build_admin_users_read_only_response(
|
||||
"当前为只读模式,无法创建用户",
|
||||
));
|
||||
};
|
||||
if !group_ids.is_empty() {
|
||||
state
|
||||
.replace_user_groups_for_user(&user.id, &group_ids)
|
||||
.await?;
|
||||
}
|
||||
|
||||
Ok(attach_admin_audit_response(
|
||||
Json(build_admin_user_payload(
|
||||
Json(build_admin_user_payload_with_groups(
|
||||
&user,
|
||||
payload.rate_limit,
|
||||
Some(rate_limit_mode.as_str()),
|
||||
payload.unlimited,
|
||||
&groups,
|
||||
))
|
||||
.into_response(),
|
||||
"admin_user_created",
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use super::super::{build_admin_users_bad_request_response, format_optional_datetime_iso8601};
|
||||
use super::support::{
|
||||
admin_user_id_from_detail_path, build_admin_user_payload, find_admin_export_user,
|
||||
admin_user_id_from_detail_path, build_admin_user_export_payload,
|
||||
build_admin_user_payload_with_groups, find_admin_export_user,
|
||||
};
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::handlers::admin::shared::{query_param_optional_bool, query_param_value};
|
||||
@@ -32,6 +33,9 @@ pub(in super::super) async fn build_admin_list_users_response(
|
||||
let search = query_param_value(request_context.query_string(), "search")
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty());
|
||||
let group_id = query_param_value(request_context.query_string(), "group_id")
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty());
|
||||
|
||||
let paged_rows = state
|
||||
.list_export_users_page(&aether_data::repository::users::UserExportListQuery {
|
||||
@@ -40,16 +44,25 @@ pub(in super::super) async fn build_admin_list_users_response(
|
||||
role: role.clone(),
|
||||
is_active,
|
||||
search,
|
||||
group_id,
|
||||
})
|
||||
.await?;
|
||||
let user_ids = paged_rows
|
||||
.iter()
|
||||
.map(|row| row.id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let (auth_rows_result, wallet_rows_result, usage_totals_result) = tokio::join!(
|
||||
let (
|
||||
auth_rows_result,
|
||||
wallet_rows_result,
|
||||
usage_totals_result,
|
||||
memberships_result,
|
||||
groups_result,
|
||||
) = tokio::join!(
|
||||
state.list_user_auth_by_ids(&user_ids),
|
||||
state.list_wallet_snapshots_by_user_ids(&user_ids),
|
||||
state.summarize_usage_totals_by_user_ids(&user_ids),
|
||||
state.list_user_group_memberships_by_user_ids(&user_ids),
|
||||
state.list_user_groups(),
|
||||
);
|
||||
let auth_by_user_id = auth_rows_result?
|
||||
.into_iter()
|
||||
@@ -63,6 +76,17 @@ pub(in super::super) async fn build_admin_list_users_response(
|
||||
.into_iter()
|
||||
.map(|item| (item.user_id.clone(), item))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let groups_by_id = groups_result?
|
||||
.into_iter()
|
||||
.map(|group| (group.id.clone(), group))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let mut group_ids_by_user_id = BTreeMap::<String, Vec<String>>::new();
|
||||
for membership in memberships_result? {
|
||||
group_ids_by_user_id
|
||||
.entry(membership.user_id)
|
||||
.or_default()
|
||||
.push(membership.group_id);
|
||||
}
|
||||
|
||||
let mut payload = Vec::with_capacity(paged_rows.len());
|
||||
for row in paged_rows {
|
||||
@@ -71,25 +95,25 @@ pub(in super::super) async fn build_admin_list_users_response(
|
||||
.get(&row.id)
|
||||
.is_some_and(|wallet| wallet.limit_mode.eq_ignore_ascii_case("unlimited"));
|
||||
let usage_totals = usage_totals_by_user_id.get(&row.id);
|
||||
payload.push(json!({
|
||||
"id": row.id,
|
||||
"email": row.email,
|
||||
"username": row.username,
|
||||
"role": row.role,
|
||||
"allowed_providers": row.allowed_providers,
|
||||
"allowed_api_formats": row.allowed_api_formats,
|
||||
"allowed_models": row.allowed_models,
|
||||
"rate_limit": row.rate_limit,
|
||||
"unlimited": unlimited,
|
||||
"is_active": row.is_active,
|
||||
"created_at": format_optional_datetime_iso8601(auth.as_ref().and_then(|user| user.created_at)),
|
||||
"updated_at": serde_json::Value::Null,
|
||||
"last_login_at": format_optional_datetime_iso8601(
|
||||
auth.as_ref().and_then(|user| user.last_login_at),
|
||||
),
|
||||
"request_count": usage_totals.map(|item| item.request_count).unwrap_or_default(),
|
||||
"total_tokens": usage_totals.map(|item| item.total_tokens).unwrap_or_default(),
|
||||
}));
|
||||
let groups = group_ids_by_user_id
|
||||
.get(&row.id)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(|group_id| groups_by_id.get(group_id).cloned())
|
||||
.collect::<Vec<_>>();
|
||||
payload.push(build_admin_user_export_payload(
|
||||
&row,
|
||||
unlimited,
|
||||
auth.as_ref().and_then(|user| user.created_at),
|
||||
auth.as_ref().and_then(|user| user.last_login_at),
|
||||
usage_totals
|
||||
.map(|item| item.request_count)
|
||||
.unwrap_or_default(),
|
||||
usage_totals
|
||||
.map(|item| item.total_tokens)
|
||||
.unwrap_or_default(),
|
||||
&groups,
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Json(payload).into_response())
|
||||
@@ -116,13 +140,16 @@ pub(in super::super) async fn build_admin_get_user_response(
|
||||
))
|
||||
.await?;
|
||||
let export_row = find_admin_export_user(state, &user_id).await?;
|
||||
let groups = state.list_user_groups_for_user(&user_id).await?;
|
||||
let unlimited = wallet
|
||||
.as_ref()
|
||||
.is_some_and(|wallet| wallet.limit_mode.eq_ignore_ascii_case("unlimited"));
|
||||
Ok(Json(build_admin_user_payload(
|
||||
Ok(Json(build_admin_user_payload_with_groups(
|
||||
&user,
|
||||
export_row.as_ref().and_then(|row| row.rate_limit),
|
||||
export_row.as_ref().map(|row| row.rate_limit_mode.as_str()),
|
||||
unlimited,
|
||||
&groups,
|
||||
))
|
||||
.into_response())
|
||||
}
|
||||
|
||||
@@ -36,6 +36,16 @@ pub(super) fn build_admin_user_payload(
|
||||
user: &aether_data::repository::users::StoredUserAuthRecord,
|
||||
rate_limit: Option<i32>,
|
||||
unlimited: bool,
|
||||
) -> serde_json::Value {
|
||||
build_admin_user_payload_with_groups(user, rate_limit, None, unlimited, &[])
|
||||
}
|
||||
|
||||
pub(super) fn build_admin_user_payload_with_groups(
|
||||
user: &aether_data::repository::users::StoredUserAuthRecord,
|
||||
rate_limit: Option<i32>,
|
||||
rate_limit_mode: Option<&str>,
|
||||
unlimited: bool,
|
||||
groups: &[aether_data::repository::users::StoredUserGroup],
|
||||
) -> serde_json::Value {
|
||||
json!({
|
||||
"id": user.id,
|
||||
@@ -43,17 +53,299 @@ pub(super) fn build_admin_user_payload(
|
||||
"username": user.username,
|
||||
"role": user.role,
|
||||
"allowed_providers": user.allowed_providers,
|
||||
"allowed_providers_mode": user.allowed_providers_mode,
|
||||
"allowed_api_formats": user.allowed_api_formats,
|
||||
"allowed_api_formats_mode": user.allowed_api_formats_mode,
|
||||
"allowed_models": user.allowed_models,
|
||||
"allowed_models_mode": user.allowed_models_mode,
|
||||
"rate_limit": rate_limit,
|
||||
"rate_limit_mode": rate_limit_mode.unwrap_or("system"),
|
||||
"unlimited": unlimited,
|
||||
"is_active": user.is_active,
|
||||
"created_at": format_optional_datetime_iso8601(user.created_at),
|
||||
"updated_at": serde_json::Value::Null,
|
||||
"last_login_at": format_optional_datetime_iso8601(user.last_login_at),
|
||||
"groups": groups.iter().map(user_group_badge_payload).collect::<Vec<_>>(),
|
||||
"effective_policy": effective_policy_payload(
|
||||
user.allowed_providers.as_ref(),
|
||||
&user.allowed_providers_mode,
|
||||
user.allowed_api_formats.as_ref(),
|
||||
&user.allowed_api_formats_mode,
|
||||
user.allowed_models.as_ref(),
|
||||
&user.allowed_models_mode,
|
||||
rate_limit,
|
||||
rate_limit_mode.unwrap_or("system"),
|
||||
groups,
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) fn build_admin_user_export_payload(
|
||||
row: &aether_data::repository::users::StoredUserExportRow,
|
||||
unlimited: bool,
|
||||
created_at: Option<chrono::DateTime<chrono::Utc>>,
|
||||
last_login_at: Option<chrono::DateTime<chrono::Utc>>,
|
||||
request_count: u64,
|
||||
total_tokens: u64,
|
||||
groups: &[aether_data::repository::users::StoredUserGroup],
|
||||
) -> serde_json::Value {
|
||||
json!({
|
||||
"id": row.id,
|
||||
"email": row.email,
|
||||
"username": row.username,
|
||||
"role": row.role,
|
||||
"allowed_providers": row.allowed_providers,
|
||||
"allowed_providers_mode": row.allowed_providers_mode,
|
||||
"allowed_api_formats": row.allowed_api_formats,
|
||||
"allowed_api_formats_mode": row.allowed_api_formats_mode,
|
||||
"allowed_models": row.allowed_models,
|
||||
"allowed_models_mode": row.allowed_models_mode,
|
||||
"rate_limit": row.rate_limit,
|
||||
"rate_limit_mode": row.rate_limit_mode,
|
||||
"unlimited": unlimited,
|
||||
"is_active": row.is_active,
|
||||
"created_at": format_optional_datetime_iso8601(created_at),
|
||||
"updated_at": serde_json::Value::Null,
|
||||
"last_login_at": format_optional_datetime_iso8601(last_login_at),
|
||||
"request_count": request_count,
|
||||
"total_tokens": total_tokens,
|
||||
"groups": groups.iter().map(user_group_badge_payload).collect::<Vec<_>>(),
|
||||
"effective_policy": effective_policy_payload(
|
||||
row.allowed_providers.as_ref(),
|
||||
&row.allowed_providers_mode,
|
||||
row.allowed_api_formats.as_ref(),
|
||||
&row.allowed_api_formats_mode,
|
||||
row.allowed_models.as_ref(),
|
||||
&row.allowed_models_mode,
|
||||
row.rate_limit,
|
||||
&row.rate_limit_mode,
|
||||
groups,
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn user_group_badge_payload(
|
||||
group: &aether_data::repository::users::StoredUserGroup,
|
||||
) -> serde_json::Value {
|
||||
json!({
|
||||
"id": group.id,
|
||||
"name": group.name,
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn effective_policy_payload(
|
||||
allowed_providers: Option<&Vec<String>>,
|
||||
allowed_providers_mode: &str,
|
||||
allowed_api_formats: Option<&Vec<String>>,
|
||||
allowed_api_formats_mode: &str,
|
||||
allowed_models: Option<&Vec<String>>,
|
||||
allowed_models_mode: &str,
|
||||
rate_limit: Option<i32>,
|
||||
rate_limit_mode: &str,
|
||||
groups: &[aether_data::repository::users::StoredUserGroup],
|
||||
) -> serde_json::Value {
|
||||
let mut sorted_groups = groups.to_vec();
|
||||
sorted_groups.sort_by(|left, right| {
|
||||
left.name
|
||||
.cmp(&right.name)
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
json!({
|
||||
"allowed_providers": effective_list_policy_payload(
|
||||
allowed_providers,
|
||||
allowed_providers_mode,
|
||||
&sorted_groups,
|
||||
|group| (&group.allowed_providers_mode, group.allowed_providers.as_ref()),
|
||||
),
|
||||
"allowed_api_formats": effective_list_policy_payload(
|
||||
allowed_api_formats,
|
||||
allowed_api_formats_mode,
|
||||
&sorted_groups,
|
||||
|group| (&group.allowed_api_formats_mode, group.allowed_api_formats.as_ref()),
|
||||
),
|
||||
"allowed_models": effective_list_policy_payload(
|
||||
allowed_models,
|
||||
allowed_models_mode,
|
||||
&sorted_groups,
|
||||
|group| (&group.allowed_models_mode, group.allowed_models.as_ref()),
|
||||
),
|
||||
"rate_limit": effective_rate_limit_policy_payload(rate_limit, rate_limit_mode, &sorted_groups),
|
||||
})
|
||||
}
|
||||
|
||||
fn effective_list_policy_payload(
|
||||
user_values: Option<&Vec<String>>,
|
||||
user_mode: &str,
|
||||
groups: &[aether_data::repository::users::StoredUserGroup],
|
||||
group_field: impl Fn(
|
||||
&aether_data::repository::users::StoredUserGroup,
|
||||
) -> (&String, Option<&Vec<String>>),
|
||||
) -> serde_json::Value {
|
||||
let mut effective = None;
|
||||
let mut group_sources = Vec::new();
|
||||
for group in groups {
|
||||
let (mode, values) = group_field(group);
|
||||
if let Some(restriction) = list_restriction_from_mode(mode, values.cloned()) {
|
||||
effective = intersect_list_policies(effective, Some(restriction));
|
||||
group_sources.push(group);
|
||||
}
|
||||
}
|
||||
let mut has_user_source = false;
|
||||
if let Some(restriction) = list_restriction_from_mode(user_mode, user_values.cloned()) {
|
||||
effective = intersect_list_policies(effective, Some(restriction));
|
||||
has_user_source = true;
|
||||
}
|
||||
|
||||
let (mode, value) = match effective {
|
||||
Some(values) if values.is_empty() => ("deny_all", json!(Vec::<String>::new())),
|
||||
Some(values) => ("specific", json!(values)),
|
||||
None => ("unrestricted", serde_json::Value::Null),
|
||||
};
|
||||
let source = combined_policy_source(has_user_source, group_sources.len(), "fallback");
|
||||
policy_payload(mode, value, source, group_sources.as_slice())
|
||||
}
|
||||
|
||||
fn effective_rate_limit_policy_payload(
|
||||
user_rate_limit: Option<i32>,
|
||||
user_mode: &str,
|
||||
groups: &[aether_data::repository::users::StoredUserGroup],
|
||||
) -> serde_json::Value {
|
||||
let mut effective = None;
|
||||
let mut group_sources = Vec::new();
|
||||
for group in groups {
|
||||
if let Some(restriction) =
|
||||
rate_limit_restriction_from_mode(&group.rate_limit_mode, group.rate_limit)
|
||||
{
|
||||
effective = intersect_rate_limit_policies(effective, Some(restriction));
|
||||
group_sources.push(group);
|
||||
}
|
||||
}
|
||||
let mut has_user_source = false;
|
||||
if let Some(restriction) = rate_limit_restriction_from_mode(user_mode, user_rate_limit) {
|
||||
effective = intersect_rate_limit_policies(effective, Some(restriction));
|
||||
has_user_source = true;
|
||||
}
|
||||
|
||||
let source = combined_policy_source(has_user_source, group_sources.len(), "fallback");
|
||||
match rate_limit_policy_value(effective) {
|
||||
Some(rate_limit) => policy_payload("custom", json!(rate_limit), source, &group_sources),
|
||||
None => policy_payload("system", serde_json::Value::Null, source, &group_sources),
|
||||
}
|
||||
}
|
||||
|
||||
fn policy_payload(
|
||||
mode: &str,
|
||||
value: serde_json::Value,
|
||||
source: &str,
|
||||
groups: &[&aether_data::repository::users::StoredUserGroup],
|
||||
) -> serde_json::Value {
|
||||
let single_group = groups.first().copied().filter(|_| groups.len() == 1);
|
||||
json!({
|
||||
"mode": mode,
|
||||
"value": value,
|
||||
"source": source,
|
||||
"group_id": single_group.map(|group| group.id.as_str()),
|
||||
"group_name": single_group.map(|group| group.name.as_str()),
|
||||
"group_ids": groups.iter().map(|group| group.id.as_str()).collect::<Vec<_>>(),
|
||||
"group_names": groups.iter().map(|group| group.name.as_str()).collect::<Vec<_>>(),
|
||||
})
|
||||
}
|
||||
|
||||
fn list_restriction_from_mode(mode: &str, values: Option<Vec<String>>) -> Option<Vec<String>> {
|
||||
match mode {
|
||||
"specific" => Some(values.unwrap_or_default()),
|
||||
"deny_all" => Some(Vec::new()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
enum RateLimitRestriction {
|
||||
Unlimited,
|
||||
Limited(i32),
|
||||
}
|
||||
|
||||
fn rate_limit_restriction_from_mode(
|
||||
mode: &str,
|
||||
rate_limit: Option<i32>,
|
||||
) -> Option<RateLimitRestriction> {
|
||||
match mode {
|
||||
"custom" => {
|
||||
let rate_limit = rate_limit.unwrap_or(0).max(0);
|
||||
if rate_limit == 0 {
|
||||
Some(RateLimitRestriction::Unlimited)
|
||||
} else {
|
||||
Some(RateLimitRestriction::Limited(rate_limit))
|
||||
}
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn intersect_list_policies(
|
||||
left: Option<Vec<String>>,
|
||||
right: Option<Vec<String>>,
|
||||
) -> Option<Vec<String>> {
|
||||
match (left, right) {
|
||||
(None, None) => None,
|
||||
(Some(values), None) | (None, Some(values)) => Some(values),
|
||||
(Some(left_values), Some(right_values)) => {
|
||||
let right_values = right_values
|
||||
.into_iter()
|
||||
.collect::<std::collections::BTreeSet<_>>();
|
||||
Some(
|
||||
left_values
|
||||
.into_iter()
|
||||
.filter(|value| right_values.contains(value))
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn intersect_rate_limit_policies(
|
||||
left: Option<RateLimitRestriction>,
|
||||
right: Option<RateLimitRestriction>,
|
||||
) -> Option<RateLimitRestriction> {
|
||||
match (left, right) {
|
||||
(None, None) => None,
|
||||
(Some(value), None) | (None, Some(value)) => Some(value),
|
||||
(Some(RateLimitRestriction::Unlimited), Some(RateLimitRestriction::Unlimited)) => {
|
||||
Some(RateLimitRestriction::Unlimited)
|
||||
}
|
||||
(Some(RateLimitRestriction::Limited(value)), Some(RateLimitRestriction::Unlimited))
|
||||
| (Some(RateLimitRestriction::Unlimited), Some(RateLimitRestriction::Limited(value))) => {
|
||||
Some(RateLimitRestriction::Limited(value))
|
||||
}
|
||||
(Some(RateLimitRestriction::Limited(left)), Some(RateLimitRestriction::Limited(right))) => {
|
||||
Some(RateLimitRestriction::Limited(left.min(right)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn rate_limit_policy_value(policy: Option<RateLimitRestriction>) -> Option<i32> {
|
||||
match policy {
|
||||
None => None,
|
||||
Some(RateLimitRestriction::Unlimited) => Some(0),
|
||||
Some(RateLimitRestriction::Limited(value)) => Some(value),
|
||||
}
|
||||
}
|
||||
|
||||
fn combined_policy_source(
|
||||
has_user_source: bool,
|
||||
group_source_count: usize,
|
||||
fallback_source: &'static str,
|
||||
) -> &'static str {
|
||||
match (has_user_source, group_source_count) {
|
||||
(true, 0) => "user",
|
||||
(false, 1) => "group",
|
||||
(false, 0) => fallback_source,
|
||||
_ => "combined",
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn admin_user_id_from_detail_path(request_path: &str) -> Option<String> {
|
||||
let value = request_path
|
||||
.strip_prefix("/api/admin/users/")?
|
||||
|
||||
@@ -1,12 +1,14 @@
|
||||
use super::super::{
|
||||
build_admin_users_bad_request_response, build_admin_users_data_unavailable_response,
|
||||
build_admin_users_read_only_response, normalize_admin_optional_user_email,
|
||||
normalize_admin_user_api_formats, normalize_admin_user_role, normalize_admin_user_string_list,
|
||||
normalize_admin_username, validate_admin_user_password, AdminUpdateUserPatch,
|
||||
build_admin_users_read_only_response, normalize_admin_list_policy_mode,
|
||||
normalize_admin_optional_user_email, normalize_admin_rate_limit_policy_mode,
|
||||
normalize_admin_user_api_formats, normalize_admin_user_group_ids, normalize_admin_user_role,
|
||||
normalize_admin_user_string_list, normalize_admin_username, validate_admin_user_password,
|
||||
AdminUpdateUserPatch,
|
||||
};
|
||||
use super::support::{
|
||||
admin_user_id_from_detail_path, admin_user_password_policy, build_admin_user_payload,
|
||||
find_admin_export_user,
|
||||
admin_user_id_from_detail_path, admin_user_password_policy,
|
||||
build_admin_user_payload_with_groups, find_admin_export_user,
|
||||
};
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::handlers::admin::shared::attach_admin_audit_response;
|
||||
@@ -177,15 +179,110 @@ pub(in super::super) async fn build_admin_update_user_response(
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let allowed_providers_mode = if field_presence.contains("allowed_providers_mode") {
|
||||
match payload.allowed_providers_mode.as_deref() {
|
||||
Some(value) => match normalize_admin_list_policy_mode(value) {
|
||||
Ok(value) => Some(value),
|
||||
Err(detail) => {
|
||||
return Ok((
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": detail })),
|
||||
)
|
||||
.into_response())
|
||||
}
|
||||
},
|
||||
None => None,
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let allowed_api_formats_mode = if field_presence.contains("allowed_api_formats_mode") {
|
||||
match payload.allowed_api_formats_mode.as_deref() {
|
||||
Some(value) => match normalize_admin_list_policy_mode(value) {
|
||||
Ok(value) => Some(value),
|
||||
Err(detail) => {
|
||||
return Ok((
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": detail })),
|
||||
)
|
||||
.into_response())
|
||||
}
|
||||
},
|
||||
None => None,
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let allowed_models_mode = if field_presence.contains("allowed_models_mode") {
|
||||
match payload.allowed_models_mode.as_deref() {
|
||||
Some(value) => match normalize_admin_list_policy_mode(value) {
|
||||
Ok(value) => Some(value),
|
||||
Err(detail) => {
|
||||
return Ok((
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": detail })),
|
||||
)
|
||||
.into_response())
|
||||
}
|
||||
},
|
||||
None => None,
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let rate_limit_mode = if field_presence.contains("rate_limit_mode") {
|
||||
match payload.rate_limit_mode.as_deref() {
|
||||
Some(value) => match normalize_admin_rate_limit_policy_mode(value) {
|
||||
Ok(value) => Some(value),
|
||||
Err(detail) => {
|
||||
return Ok((
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": detail })),
|
||||
)
|
||||
.into_response())
|
||||
}
|
||||
},
|
||||
None => None,
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let group_ids = if field_presence.contains("group_ids") {
|
||||
let requested_group_ids = normalize_admin_user_group_ids(payload.group_ids);
|
||||
Some(
|
||||
state
|
||||
.include_default_user_group_ids(&requested_group_ids)
|
||||
.await?,
|
||||
)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if let Some(group_ids) = group_ids.as_ref() {
|
||||
if !group_ids.is_empty() {
|
||||
let groups = state.list_user_groups_by_ids(group_ids).await?;
|
||||
if groups.len() != group_ids.len() {
|
||||
return Ok((
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": "用户分组不存在" })),
|
||||
)
|
||||
.into_response());
|
||||
}
|
||||
}
|
||||
}
|
||||
let needs_auth_user_write = email.is_some()
|
||||
|| username.is_some()
|
||||
|| payload.password.is_some()
|
||||
|| role.is_some()
|
||||
|| field_presence.contains("allowed_providers")
|
||||
|| allowed_providers_mode.is_some()
|
||||
|| field_presence.contains("allowed_api_formats")
|
||||
|| allowed_api_formats_mode.is_some()
|
||||
|| field_presence.contains("allowed_models")
|
||||
|| allowed_models_mode.is_some()
|
||||
|| field_presence.contains("rate_limit")
|
||||
|| payload.is_active.is_some();
|
||||
|| rate_limit_mode.is_some()
|
||||
|| payload.is_active.is_some()
|
||||
|| group_ids.is_some();
|
||||
if needs_auth_user_write && !state.has_auth_user_write_capability() {
|
||||
return Ok(build_admin_users_read_only_response(
|
||||
"当前为只读模式,无法更新用户",
|
||||
@@ -210,6 +307,11 @@ pub(in super::super) async fn build_admin_update_user_response(
|
||||
.into_response());
|
||||
}
|
||||
}
|
||||
if let Some(group_ids) = group_ids.as_ref() {
|
||||
state
|
||||
.replace_user_groups_for_user(&user_id, group_ids)
|
||||
.await?;
|
||||
}
|
||||
|
||||
if let Some(password) = payload.password.as_deref() {
|
||||
let password_policy = admin_user_password_policy(state).await?;
|
||||
@@ -274,6 +376,29 @@ pub(in super::super) async fn build_admin_update_user_response(
|
||||
.into_response());
|
||||
}
|
||||
}
|
||||
if allowed_providers_mode.is_some()
|
||||
|| allowed_api_formats_mode.is_some()
|
||||
|| allowed_models_mode.is_some()
|
||||
|| rate_limit_mode.is_some()
|
||||
{
|
||||
if state
|
||||
.update_local_auth_user_policy_modes(
|
||||
&user_id,
|
||||
allowed_providers_mode,
|
||||
allowed_api_formats_mode,
|
||||
allowed_models_mode,
|
||||
rate_limit_mode,
|
||||
)
|
||||
.await?
|
||||
.is_none()
|
||||
{
|
||||
return Ok((
|
||||
http::StatusCode::NOT_FOUND,
|
||||
Json(json!({ "detail": "用户不存在" })),
|
||||
)
|
||||
.into_response());
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(unlimited) = payload.unlimited {
|
||||
match state
|
||||
@@ -322,13 +447,21 @@ pub(in super::super) async fn build_admin_update_user_response(
|
||||
.as_ref()
|
||||
.is_some_and(|wallet| wallet.limit_mode.eq_ignore_ascii_case("unlimited"));
|
||||
let export_row = find_admin_export_user(state, &user_id).await?;
|
||||
let groups = state.list_user_groups_for_user(&user_id).await?;
|
||||
let rate_limit = export_row
|
||||
.as_ref()
|
||||
.and_then(|row| row.rate_limit)
|
||||
.or(payload.rate_limit);
|
||||
|
||||
Ok(attach_admin_audit_response(
|
||||
Json(build_admin_user_payload(&user, rate_limit, unlimited)).into_response(),
|
||||
Json(build_admin_user_payload_with_groups(
|
||||
&user,
|
||||
rate_limit,
|
||||
export_row.as_ref().map(|row| row.rate_limit_mode.as_str()),
|
||||
unlimited,
|
||||
&groups,
|
||||
))
|
||||
.into_response(),
|
||||
"admin_user_updated",
|
||||
"update_user",
|
||||
"user",
|
||||
|
||||
@@ -4,6 +4,7 @@ const ADMIN_USERS_DATA_UNAVAILABLE_DETAIL: &str = "Admin user management data un
|
||||
|
||||
mod api_keys;
|
||||
mod batch;
|
||||
mod groups;
|
||||
mod lifecycle;
|
||||
mod route_seam;
|
||||
mod routes;
|
||||
@@ -23,6 +24,12 @@ pub(crate) use self::api_keys::{
|
||||
use self::batch::{
|
||||
build_admin_resolve_user_selection_response, build_admin_user_batch_action_response,
|
||||
};
|
||||
use self::groups::{
|
||||
build_admin_create_user_group_response, build_admin_delete_user_group_response,
|
||||
build_admin_list_user_group_members_response, build_admin_list_user_groups_response,
|
||||
build_admin_replace_user_group_members_response, build_admin_set_default_user_group_response,
|
||||
build_admin_update_user_group_response,
|
||||
};
|
||||
use self::lifecycle::{
|
||||
build_admin_create_user_response, build_admin_delete_user_response,
|
||||
build_admin_get_user_response, build_admin_list_users_response,
|
||||
@@ -36,12 +43,16 @@ use self::shared::AdminUpdateUserPatch;
|
||||
use self::shared::{
|
||||
admin_default_user_initial_gift, build_admin_users_bad_request_response,
|
||||
build_admin_users_data_unavailable_response, build_admin_users_read_only_response,
|
||||
format_optional_datetime_iso8601, normalize_admin_optional_user_email,
|
||||
normalize_admin_user_role, normalize_admin_username, validate_admin_user_password,
|
||||
AdminCreateUserApiKeyRequest, AdminCreateUserRequest, AdminToggleUserApiKeyLockRequest,
|
||||
AdminUpdateUserApiKeyRequest,
|
||||
format_optional_datetime_iso8601, legacy_admin_list_policy_mode,
|
||||
legacy_admin_rate_limit_policy_mode, normalize_admin_optional_user_email,
|
||||
normalize_admin_user_group_ids, normalize_admin_user_role, normalize_admin_username,
|
||||
validate_admin_user_password, AdminCreateUserApiKeyRequest, AdminCreateUserRequest,
|
||||
AdminToggleUserApiKeyLockRequest, AdminUpdateUserApiKeyRequest,
|
||||
};
|
||||
pub(crate) use self::shared::{
|
||||
normalize_admin_list_policy_mode, normalize_admin_rate_limit_policy_mode,
|
||||
normalize_admin_user_api_formats, normalize_admin_user_string_list,
|
||||
};
|
||||
pub(crate) use self::shared::{normalize_admin_user_api_formats, normalize_admin_user_string_list};
|
||||
|
||||
pub(crate) async fn maybe_build_local_admin_users_response(
|
||||
request: AdminRouteRequest<'_>,
|
||||
|
||||
@@ -1,13 +1,16 @@
|
||||
use super::{
|
||||
build_admin_create_user_api_key_response, build_admin_create_user_response,
|
||||
build_admin_delete_user_api_key_response, build_admin_delete_user_response,
|
||||
build_admin_create_user_api_key_response, build_admin_create_user_group_response,
|
||||
build_admin_create_user_response, build_admin_delete_user_api_key_response,
|
||||
build_admin_delete_user_group_response, build_admin_delete_user_response,
|
||||
build_admin_delete_user_session_response, build_admin_delete_user_sessions_response,
|
||||
build_admin_get_user_response, build_admin_list_user_api_keys_response,
|
||||
build_admin_list_user_group_members_response, build_admin_list_user_groups_response,
|
||||
build_admin_list_user_sessions_response, build_admin_list_users_response,
|
||||
build_admin_resolve_user_selection_response, build_admin_reveal_user_api_key_response,
|
||||
build_admin_replace_user_group_members_response, build_admin_resolve_user_selection_response,
|
||||
build_admin_reveal_user_api_key_response, build_admin_set_default_user_group_response,
|
||||
build_admin_toggle_user_api_key_lock_response, build_admin_update_user_api_key_response,
|
||||
build_admin_update_user_response, build_admin_user_batch_action_response,
|
||||
build_admin_users_data_unavailable_response,
|
||||
build_admin_update_user_group_response, build_admin_update_user_response,
|
||||
build_admin_user_batch_action_response, build_admin_users_data_unavailable_response,
|
||||
};
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::GatewayError;
|
||||
@@ -15,8 +18,27 @@ use axum::{body::Body, http, response::Response};
|
||||
|
||||
fn is_admin_users_route(request_context: &AdminRequestContext<'_>) -> bool {
|
||||
let path = request_context.path();
|
||||
(request_context.method() == http::Method::GET
|
||||
&& matches!(path, "/api/admin/users" | "/api/admin/users/"))
|
||||
((request_context.method() == http::Method::GET
|
||||
|| request_context.method() == http::Method::POST)
|
||||
&& matches!(path, "/api/admin/user-groups" | "/api/admin/user-groups/"))
|
||||
|| (request_context.method() == http::Method::PUT
|
||||
&& matches!(
|
||||
path,
|
||||
"/api/admin/user-groups/default" | "/api/admin/user-groups/default/"
|
||||
))
|
||||
|| ((request_context.method() == http::Method::PUT
|
||||
|| request_context.method() == http::Method::DELETE)
|
||||
&& path.starts_with("/api/admin/user-groups/")
|
||||
&& path.matches('/').count() == 4
|
||||
&& !path.ends_with("/members")
|
||||
&& !path.ends_with("/default"))
|
||||
|| ((request_context.method() == http::Method::GET
|
||||
|| request_context.method() == http::Method::PUT)
|
||||
&& path.starts_with("/api/admin/user-groups/")
|
||||
&& path.ends_with("/members")
|
||||
&& path.matches('/').count() == 5)
|
||||
|| (request_context.method() == http::Method::GET
|
||||
&& matches!(path, "/api/admin/users" | "/api/admin/users/"))
|
||||
|| (request_context.method() == http::Method::POST
|
||||
&& matches!(path, "/api/admin/users" | "/api/admin/users/"))
|
||||
|| (request_context.method() == http::Method::POST
|
||||
@@ -84,6 +106,26 @@ pub(super) async fn maybe_build_local_admin_users_routes_response(
|
||||
}
|
||||
|
||||
match decision.route_kind.as_deref() {
|
||||
Some("list_user_groups") => Ok(Some(build_admin_list_user_groups_response(state).await?)),
|
||||
Some("create_user_group") => Ok(Some(
|
||||
build_admin_create_user_group_response(state, request_body).await?,
|
||||
)),
|
||||
Some("update_user_group") => Ok(Some(
|
||||
build_admin_update_user_group_response(state, request_context, request_body).await?,
|
||||
)),
|
||||
Some("delete_user_group") => Ok(Some(
|
||||
build_admin_delete_user_group_response(state, request_context).await?,
|
||||
)),
|
||||
Some("list_user_group_members") => Ok(Some(
|
||||
build_admin_list_user_group_members_response(state, request_context).await?,
|
||||
)),
|
||||
Some("replace_user_group_members") => Ok(Some(
|
||||
build_admin_replace_user_group_members_response(state, request_context, request_body)
|
||||
.await?,
|
||||
)),
|
||||
Some("set_default_user_group") => Ok(Some(
|
||||
build_admin_set_default_user_group_response(state, request_body).await?,
|
||||
)),
|
||||
Some("create_user") => Ok(Some(
|
||||
build_admin_create_user_response(state, request_context, request_body).await?,
|
||||
)),
|
||||
|
||||
@@ -68,11 +68,21 @@ pub(super) struct AdminCreateUserRequest {
|
||||
#[serde(default)]
|
||||
pub(super) allowed_providers: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
pub(super) allowed_providers_mode: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(super) allowed_api_formats: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
pub(super) allowed_api_formats_mode: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(super) allowed_models: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
pub(super) allowed_models_mode: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(super) rate_limit: Option<i32>,
|
||||
#[serde(default)]
|
||||
pub(super) rate_limit_mode: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(super) group_ids: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
@@ -90,12 +100,22 @@ pub(super) struct AdminUpdateUserRequest {
|
||||
#[serde(default)]
|
||||
pub(super) allowed_providers: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
pub(super) allowed_providers_mode: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(super) allowed_api_formats: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
pub(super) allowed_api_formats_mode: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(super) allowed_models: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
pub(super) allowed_models_mode: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(super) rate_limit: Option<i32>,
|
||||
#[serde(default)]
|
||||
pub(super) rate_limit_mode: Option<String>,
|
||||
#[serde(default)]
|
||||
pub(super) group_ids: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub(super) is_active: Option<bool>,
|
||||
}
|
||||
|
||||
@@ -263,6 +283,48 @@ pub(crate) fn normalize_admin_user_api_formats(
|
||||
Ok(Some(normalized))
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_admin_list_policy_mode(value: &str) -> Result<String, String> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"inherit" | "unrestricted" | "specific" | "deny_all" => {
|
||||
Ok(value.trim().to_ascii_lowercase())
|
||||
}
|
||||
_ => Err("权限列表模式不合法".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_admin_rate_limit_policy_mode(value: &str) -> Result<String, String> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"inherit" | "system" | "custom" => Ok(value.trim().to_ascii_lowercase()),
|
||||
_ => Err("限速模式不合法".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn legacy_admin_list_policy_mode(values: &Option<Vec<String>>) -> String {
|
||||
if values.is_some() {
|
||||
"specific".to_string()
|
||||
} else {
|
||||
"unrestricted".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn legacy_admin_rate_limit_policy_mode(value: Option<i32>) -> String {
|
||||
if value.is_some() {
|
||||
"custom".to_string()
|
||||
} else {
|
||||
"system".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn normalize_admin_user_group_ids(values: Vec<String>) -> Vec<String> {
|
||||
values
|
||||
.into_iter()
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
.collect::<std::collections::BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(super) fn admin_default_user_initial_gift(value: Option<&serde_json::Value>) -> f64 {
|
||||
match value {
|
||||
Some(serde_json::Value::Number(number)) => number.as_f64().unwrap_or(10.0),
|
||||
|
||||
@@ -513,6 +513,17 @@ pub(super) async fn handle_auth_register(
|
||||
false,
|
||||
);
|
||||
};
|
||||
if let Err(err) = state
|
||||
.assign_default_group_to_self_registered_user(&user.id)
|
||||
.await
|
||||
{
|
||||
let _ = state.delete_local_auth_user(&user.id).await;
|
||||
return build_auth_error_response(
|
||||
http::StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("auth default user group assignment failed: {err:?}"),
|
||||
false,
|
||||
);
|
||||
}
|
||||
|
||||
if require_verification {
|
||||
if let Some(email) = email.as_deref() {
|
||||
|
||||
@@ -325,6 +325,13 @@ pub(crate) fn admin_proxy_local_requires_buffered_body(
|
||||
| (Some("users_manage"), http::Method::POST, Some("resolve_user_selection"))
|
||||
| (Some("users_manage"), http::Method::POST, Some("batch_action_users"))
|
||||
| (Some("users_manage"), http::Method::PUT, Some("update_user"))
|
||||
| (Some("users_manage"), http::Method::POST, Some("create_user_group"))
|
||||
| (Some("users_manage"), http::Method::PUT, Some("update_user_group"))
|
||||
| (
|
||||
Some("users_manage"),
|
||||
http::Method::PUT,
|
||||
Some("replace_user_group_members" | "set_default_user_group"),
|
||||
)
|
||||
| (Some("users_manage"), http::Method::POST, Some("create_user_api_key"))
|
||||
| (Some("users_manage"), http::Method::PUT, Some("update_user_api_key"))
|
||||
| (Some("users_manage"), http::Method::PATCH, Some("lock_user_api_key"))
|
||||
|
||||
@@ -211,6 +211,13 @@ pub(crate) async fn resolve_identity_oauth_login_user(
|
||||
return Err(IdentityOAuthAccountError::Storage(format!("{err:?}")));
|
||||
}
|
||||
}
|
||||
if let Err(err) = state
|
||||
.assign_default_group_to_self_registered_user(&user.id)
|
||||
.await
|
||||
{
|
||||
let _ = state.delete_local_auth_user(&user.id).await;
|
||||
return Err(IdentityOAuthAccountError::Storage(format!("{err:?}")));
|
||||
}
|
||||
if let Err(err) = upsert_oauth_link(state, &user.id, claims, now).await {
|
||||
let _ = state.delete_local_auth_user(&user.id).await;
|
||||
return Err(err);
|
||||
|
||||
@@ -274,6 +274,47 @@ pub(crate) async fn record_local_request_candidate_status(
|
||||
persist_local_request_candidate_status_record(state, record).await;
|
||||
}
|
||||
|
||||
pub(crate) async fn record_local_request_candidate_extra_data(
|
||||
state: &(impl RequestCandidateRuntimeWriter + ?Sized),
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&Value>,
|
||||
status: RequestCandidateStatus,
|
||||
status_code: Option<u16>,
|
||||
latency_ms: Option<u64>,
|
||||
extra_data: Value,
|
||||
) {
|
||||
let Some(snapshot) = snapshot_local_request_candidate_status(plan, report_context) else {
|
||||
return;
|
||||
};
|
||||
let record = UpsertRequestCandidateRecord {
|
||||
id: snapshot.candidate_id.clone(),
|
||||
request_id: snapshot.request_id.clone(),
|
||||
user_id: snapshot.user_id.clone(),
|
||||
api_key_id: snapshot.api_key_id.clone(),
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
candidate_index: snapshot.candidate_index,
|
||||
retry_index: snapshot.retry_index,
|
||||
provider_id: Some(snapshot.provider_id.clone()),
|
||||
endpoint_id: Some(snapshot.endpoint_id.clone()),
|
||||
key_id: Some(snapshot.key_id.clone()),
|
||||
status,
|
||||
skip_reason: None,
|
||||
is_cached: None,
|
||||
status_code,
|
||||
error_type: None,
|
||||
error_message: None,
|
||||
latency_ms,
|
||||
concurrent_requests: None,
|
||||
extra_data: Some(extra_data),
|
||||
required_capabilities: None,
|
||||
created_at_unix_ms: None,
|
||||
started_at_unix_ms: None,
|
||||
finished_at_unix_ms: None,
|
||||
};
|
||||
persist_local_request_candidate_status_record(state, record).await;
|
||||
}
|
||||
|
||||
pub(crate) async fn record_local_request_candidate_status_snapshot(
|
||||
state: &(impl RequestCandidateRuntimeWriter + ?Sized),
|
||||
snapshot: &LocalRequestCandidateStatusSnapshot,
|
||||
|
||||
@@ -1,8 +1,80 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use crate::constants::{BUILTIN_DEFAULT_USER_GROUP_ID, DEFAULT_USER_GROUP_CONFIG_KEY};
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
impl AppState {
|
||||
pub(crate) async fn assign_default_group_to_self_registered_user(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<(), GatewayError> {
|
||||
let group_id = self.effective_default_user_group_id().await?;
|
||||
let Some(group_id) = group_id else {
|
||||
return Ok(());
|
||||
};
|
||||
if !self.add_user_to_group(&group_id, user_id).await? {
|
||||
return Err(GatewayError::Internal(format!(
|
||||
"failed to add user {user_id} to default group {group_id}"
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn configured_default_user_group_id(
|
||||
&self,
|
||||
) -> Result<Option<String>, GatewayError> {
|
||||
Ok(self
|
||||
.read_system_config_json_value(DEFAULT_USER_GROUP_CONFIG_KEY)
|
||||
.await?
|
||||
.and_then(|value| value.as_str().map(str::to_string))
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty()))
|
||||
}
|
||||
|
||||
pub(crate) async fn effective_default_user_group_id(
|
||||
&self,
|
||||
) -> Result<Option<String>, GatewayError> {
|
||||
if let Some(group_id) = self.configured_default_user_group_id().await? {
|
||||
if self.find_user_group_by_id(&group_id).await?.is_none() {
|
||||
return Err(GatewayError::Internal(format!(
|
||||
"{DEFAULT_USER_GROUP_CONFIG_KEY} points to missing group: {group_id}"
|
||||
)));
|
||||
}
|
||||
return Ok(Some(group_id));
|
||||
}
|
||||
if self
|
||||
.find_user_group_by_id(BUILTIN_DEFAULT_USER_GROUP_ID)
|
||||
.await?
|
||||
.is_some()
|
||||
{
|
||||
return Ok(Some(BUILTIN_DEFAULT_USER_GROUP_ID.to_string()));
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub(crate) async fn include_default_user_group_ids(
|
||||
&self,
|
||||
group_ids: &[String],
|
||||
) -> Result<Vec<String>, GatewayError> {
|
||||
let mut group_ids = group_ids
|
||||
.iter()
|
||||
.map(|value| value.trim())
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.collect::<BTreeSet<_>>();
|
||||
if let Some(default_group_id) = self.effective_default_user_group_id().await? {
|
||||
group_ids.insert(default_group_id);
|
||||
}
|
||||
Ok(group_ids.into_iter().collect())
|
||||
}
|
||||
|
||||
pub(crate) async fn add_all_users_to_group(&self, group_id: &str) -> Result<(), GatewayError> {
|
||||
for user in self.list_export_users().await? {
|
||||
self.add_user_to_group(group_id, &user.id).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_auth_user_summaries_by_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
@@ -132,6 +204,126 @@ impl AppState {
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_groups(
|
||||
&self,
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroup>, GatewayError> {
|
||||
self.data
|
||||
.list_user_groups()
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn find_user_group_by_id(
|
||||
&self,
|
||||
group_id: &str,
|
||||
) -> Result<Option<aether_data::repository::users::StoredUserGroup>, GatewayError> {
|
||||
self.data
|
||||
.find_user_group_by_id(group_id)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_groups_by_ids(
|
||||
&self,
|
||||
group_ids: &[String],
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroup>, GatewayError> {
|
||||
self.data
|
||||
.list_user_groups_by_ids(group_ids)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn create_user_group(
|
||||
&self,
|
||||
record: aether_data::repository::users::UpsertUserGroupRecord,
|
||||
) -> Result<Option<aether_data::repository::users::StoredUserGroup>, GatewayError> {
|
||||
self.data
|
||||
.create_user_group(record)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn update_user_group(
|
||||
&self,
|
||||
group_id: &str,
|
||||
record: aether_data::repository::users::UpsertUserGroupRecord,
|
||||
) -> Result<Option<aether_data::repository::users::StoredUserGroup>, GatewayError> {
|
||||
self.data
|
||||
.update_user_group(group_id, record)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_user_group(&self, group_id: &str) -> Result<bool, GatewayError> {
|
||||
self.data
|
||||
.delete_user_group(group_id)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_group_members(
|
||||
&self,
|
||||
group_id: &str,
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroupMember>, GatewayError> {
|
||||
self.data
|
||||
.list_user_group_members(group_id)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn replace_user_group_members(
|
||||
&self,
|
||||
group_id: &str,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroupMember>, GatewayError> {
|
||||
self.data
|
||||
.replace_user_group_members(group_id, user_ids)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_groups_for_user(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroup>, GatewayError> {
|
||||
self.data
|
||||
.list_user_groups_for_user(user_id)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_user_group_memberships_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroupMembership>, GatewayError> {
|
||||
self.data
|
||||
.list_user_group_memberships_by_user_ids(user_ids)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn replace_user_groups_for_user(
|
||||
&self,
|
||||
user_id: &str,
|
||||
group_ids: &[String],
|
||||
) -> Result<Vec<aether_data::repository::users::StoredUserGroup>, GatewayError> {
|
||||
self.data
|
||||
.replace_user_groups_for_user(user_id, group_ids)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn add_user_to_group(
|
||||
&self,
|
||||
group_id: &str,
|
||||
user_id: &str,
|
||||
) -> Result<bool, GatewayError> {
|
||||
self.data
|
||||
.add_user_to_group(group_id, user_id)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn is_other_user_auth_email_taken(
|
||||
&self,
|
||||
email: &str,
|
||||
@@ -289,6 +481,12 @@ impl AppState {
|
||||
Some(now),
|
||||
None,
|
||||
)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?
|
||||
.with_policy_modes(
|
||||
"inherit".to_string(),
|
||||
"inherit".to_string(),
|
||||
"inherit".to_string(),
|
||||
)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
store
|
||||
.lock()
|
||||
@@ -418,6 +616,45 @@ impl AppState {
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn update_local_auth_user_policy_modes(
|
||||
&self,
|
||||
user_id: &str,
|
||||
allowed_providers_mode: Option<String>,
|
||||
allowed_api_formats_mode: Option<String>,
|
||||
allowed_models_mode: Option<String>,
|
||||
rate_limit_mode: Option<String>,
|
||||
) -> Result<Option<aether_data::repository::users::StoredUserAuthRecord>, GatewayError> {
|
||||
#[cfg(test)]
|
||||
if let Some(store) = self.auth_user_store.as_ref() {
|
||||
let mut guard = store.lock().expect("auth user store should lock");
|
||||
let Some(user) = guard.get_mut(user_id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if let Some(mode) = allowed_providers_mode {
|
||||
user.allowed_providers_mode = mode;
|
||||
}
|
||||
if let Some(mode) = allowed_api_formats_mode {
|
||||
user.allowed_api_formats_mode = mode;
|
||||
}
|
||||
if let Some(mode) = allowed_models_mode {
|
||||
user.allowed_models_mode = mode;
|
||||
}
|
||||
let _ = rate_limit_mode;
|
||||
return Ok(Some(user.clone()));
|
||||
}
|
||||
|
||||
self.data
|
||||
.update_local_auth_user_policy_modes(
|
||||
user_id,
|
||||
allowed_providers_mode,
|
||||
allowed_api_formats_mode,
|
||||
allowed_models_mode,
|
||||
rate_limit_mode,
|
||||
)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn touch_auth_user_last_login(
|
||||
&self,
|
||||
user_id: &str,
|
||||
@@ -508,6 +745,12 @@ impl AppState {
|
||||
Some(now),
|
||||
None,
|
||||
)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?
|
||||
.with_policy_modes(
|
||||
"inherit".to_string(),
|
||||
"inherit".to_string(),
|
||||
"inherit".to_string(),
|
||||
)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let gift_balance = if unlimited {
|
||||
0.0
|
||||
|
||||
@@ -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