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:
Entropy.Xu
2026-05-10 18:36:53 +08:00
109 changed files with 11022 additions and 948 deletions

View File

@@ -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,

View File

@@ -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,
};

View File

@@ -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,

View File

@@ -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,

View File

@@ -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,

View File

@@ -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}"),
),
}
}

View File

@@ -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"];

View File

@@ -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/")
{

View File

@@ -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(&[]);

View File

@@ -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(

View File

@@ -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

View File

@@ -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::{

View File

@@ -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())

View File

@@ -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,

View File

@@ -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"}]}));
}
}

View File

@@ -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")

View File

@@ -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,
}))

View File

@@ -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(),

View File

@@ -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,

View File

@@ -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()),
}))
}

View 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,
}
}

View File

@@ -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",

View File

@@ -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())
}

View File

@@ -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/")?

View File

@@ -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",

View File

@@ -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<'_>,

View File

@@ -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?,
)),

View File

@@ -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),

View File

@@ -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() {

View File

@@ -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"))

View File

@@ -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);

View File

@@ -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,

View File

@@ -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

View File

@@ -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")

View File

@@ -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;

View File

@@ -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"
);

View File

@@ -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;

View File

@@ -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"
);

View File

@@ -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;

View File

@@ -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"
);

View File

@@ -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"
);

View File

@@ -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: ")

View File

@@ -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::{

View File

@@ -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"],

View File

@@ -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

View File

@@ -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,

View File

@@ -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();

View File

@@ -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);

View File

@@ -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();

View File

@@ -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);

View File

@@ -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;
}
}
}

View File

@@ -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;

View File

@@ -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"

View File

@@ -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()
);