mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 18:07:47 +08:00
Merge origin/main into fix/gemini-cli-v1internal
This commit is contained in:
@@ -23,6 +23,17 @@ use super::{
|
||||
LocalProviderDeleteTaskState, ProviderTransportSnapshotCacheKey,
|
||||
};
|
||||
|
||||
const DEFAULT_REQUEST_BODY_READ_TIMEOUT_MS: u64 = 120_000;
|
||||
const MIN_REQUEST_BODY_READ_TIMEOUT_MS: u64 = 1_000;
|
||||
const MAX_REQUEST_BODY_READ_TIMEOUT_MS: u64 = 600_000;
|
||||
const REQUEST_BODY_READ_TIMEOUT_MS_ENV: &str = "AETHER_GATEWAY_REQUEST_BODY_READ_TIMEOUT_MS";
|
||||
|
||||
const DEFAULT_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS: u64 = 30_000;
|
||||
const MIN_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS: u64 = 500;
|
||||
const MAX_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS: u64 = 120_000;
|
||||
const LOCAL_EXECUTION_PLANNING_TIMEOUT_MS_ENV: &str =
|
||||
"AETHER_GATEWAY_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS";
|
||||
|
||||
#[cfg(test)]
|
||||
type TestExecutionRuntimeSyncOverrideFn = dyn Fn(
|
||||
&aether_contracts::ExecutionPlan,
|
||||
@@ -43,6 +54,52 @@ impl std::fmt::Debug for TestExecutionRuntimeSyncOverride {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct FrontdoorRuntimeGuardConfig {
|
||||
pub(crate) request_body_read_timeout: Duration,
|
||||
pub(crate) local_execution_planning_timeout: Duration,
|
||||
}
|
||||
|
||||
impl FrontdoorRuntimeGuardConfig {
|
||||
pub(crate) fn from_env() -> Self {
|
||||
Self {
|
||||
request_body_read_timeout: env_duration_ms(
|
||||
REQUEST_BODY_READ_TIMEOUT_MS_ENV,
|
||||
DEFAULT_REQUEST_BODY_READ_TIMEOUT_MS,
|
||||
MIN_REQUEST_BODY_READ_TIMEOUT_MS,
|
||||
MAX_REQUEST_BODY_READ_TIMEOUT_MS,
|
||||
),
|
||||
local_execution_planning_timeout: env_duration_ms(
|
||||
LOCAL_EXECUTION_PLANNING_TIMEOUT_MS_ENV,
|
||||
DEFAULT_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS,
|
||||
MIN_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS,
|
||||
MAX_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn for_tests(
|
||||
request_body_read_timeout: Duration,
|
||||
local_execution_planning_timeout: Duration,
|
||||
) -> Self {
|
||||
Self {
|
||||
request_body_read_timeout,
|
||||
local_execution_planning_timeout,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn env_duration_ms(key: &str, default_ms: u64, min_ms: u64, max_ms: u64) -> Duration {
|
||||
let ms = std::env::var(key)
|
||||
.ok()
|
||||
.and_then(|value| value.trim().parse::<u64>().ok())
|
||||
.filter(|value| *value > 0)
|
||||
.unwrap_or(default_ms)
|
||||
.clamp(min_ms, max_ms);
|
||||
Duration::from_millis(ms)
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AppState {
|
||||
#[cfg(test)]
|
||||
@@ -54,6 +111,7 @@ pub struct AppState {
|
||||
pub(crate) usage_runtime: Arc<usage::UsageRuntime>,
|
||||
pub(crate) video_tasks: Arc<VideoTaskService>,
|
||||
pub(crate) video_task_poller: Option<VideoTaskPollerConfig>,
|
||||
pub(crate) frontdoor_runtime_guards: Arc<FrontdoorRuntimeGuardConfig>,
|
||||
pub(crate) request_gate: Option<Arc<ConcurrencyGate>>,
|
||||
pub(crate) distributed_request_gate: Option<Arc<RuntimeSemaphore>>,
|
||||
pub(crate) client: reqwest::Client,
|
||||
|
||||
@@ -429,6 +429,17 @@ impl AppState {
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_key_maintenance_summaries_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<provider_catalog::StoredProviderCatalogKeyMaintenanceSummary>, GatewayError>
|
||||
{
|
||||
self.data
|
||||
.list_provider_catalog_key_maintenance_summaries_by_provider_ids(provider_ids)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_provider_catalog_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
|
||||
@@ -21,7 +21,9 @@ use aether_runtime_state::{
|
||||
};
|
||||
use aether_scheduler_core::PROVIDER_KEY_RPM_WINDOW_SECS;
|
||||
|
||||
use super::{AppState, FrontdoorCorsConfig, LocalExecutionRuntimeMissDiagnostic};
|
||||
use super::{
|
||||
AppState, FrontdoorCorsConfig, FrontdoorRuntimeGuardConfig, LocalExecutionRuntimeMissDiagnostic,
|
||||
};
|
||||
|
||||
use super::super::async_task::{
|
||||
spawn_video_task_poller, VideoTaskPollerConfig, VideoTaskService, VideoTaskTruthSourceMode,
|
||||
@@ -67,8 +69,10 @@ const SCHEDULER_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] = &[
|
||||
"provider_priority_mode",
|
||||
"scheduling_mode",
|
||||
];
|
||||
const AUTH_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] =
|
||||
&[crate::constants::DEFAULT_USER_GROUP_CONFIG_KEY];
|
||||
const AUTH_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] = &[
|
||||
crate::constants::DEFAULT_USER_GROUP_CONFIG_KEY,
|
||||
crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY,
|
||||
];
|
||||
const FRONTDOOR_RPM_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] = &["rate_limit_per_minute"];
|
||||
|
||||
fn system_config_key_affects_scheduler(key: &str) -> bool {
|
||||
@@ -227,6 +231,7 @@ impl AppState {
|
||||
VideoTaskTruthSourceMode::PythonSyncReport,
|
||||
)),
|
||||
video_task_poller: None,
|
||||
frontdoor_runtime_guards: Arc::new(FrontdoorRuntimeGuardConfig::from_env()),
|
||||
request_gate: None,
|
||||
distributed_request_gate: None,
|
||||
client,
|
||||
@@ -568,12 +573,14 @@ impl AppState {
|
||||
|
||||
pub(crate) fn invalidate_provider_routing_caches(&self) {
|
||||
self.data.clear_minimal_candidate_selection_cache();
|
||||
self.data.clear_provider_catalog_cache();
|
||||
self.clear_provider_transport_snapshot_cache();
|
||||
self.invalidate_scheduler_affinity_cache();
|
||||
}
|
||||
|
||||
pub(crate) fn invalidate_provider_health_routing_caches(&self) {
|
||||
self.data.clear_minimal_candidate_selection_cache();
|
||||
self.data.clear_provider_catalog_cache();
|
||||
self.clear_provider_transport_snapshot_cache();
|
||||
}
|
||||
|
||||
@@ -626,6 +633,36 @@ impl AppState {
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
pub(crate) async fn export_admin_system_usage_aggregates(
|
||||
&self,
|
||||
) -> Result<aether_data::repository::system::AdminSystemUsageAggregateSnapshot, GatewayError>
|
||||
{
|
||||
self.data
|
||||
.export_admin_system_usage_aggregates()
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn import_admin_system_usage_aggregates(
|
||||
&self,
|
||||
snapshot: &aether_data::repository::system::AdminSystemUsageAggregateSnapshot,
|
||||
user_id_map: &std::collections::BTreeMap<String, String>,
|
||||
api_key_id_map: &std::collections::BTreeMap<String, String>,
|
||||
mode: aether_data::repository::system::AdminSystemUsageAggregateImportMode,
|
||||
) -> Result<aether_data::repository::system::AdminSystemUsageAggregateImportSummary, GatewayError>
|
||||
{
|
||||
self.data
|
||||
.import_admin_system_usage_aggregates(snapshot, user_id_map, api_key_id_map, mode)
|
||||
.await
|
||||
.map_err(|err| match err {
|
||||
aether_data::DataLayerError::InvalidInput(detail) => GatewayError::Client {
|
||||
status: http::StatusCode::BAD_REQUEST,
|
||||
message: detail,
|
||||
},
|
||||
other => GatewayError::Internal(other.to_string()),
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn run_admin_system_cleanup_once(
|
||||
&self,
|
||||
) -> Result<crate::maintenance::AdminSystemCleanupSummary, GatewayError> {
|
||||
@@ -1233,6 +1270,10 @@ impl AppState {
|
||||
crate::task_runtime::TASK_KEY_VIDEO_TASK_POLLER,
|
||||
spawn_video_task_poller(self.clone()),
|
||||
);
|
||||
supervise_worker(
|
||||
crate::backup::worker::S3_BACKUP_WORKER_TASK_KEY,
|
||||
crate::backup::worker::spawn_s3_backup_worker(self.clone()),
|
||||
);
|
||||
|
||||
supervisor
|
||||
}
|
||||
|
||||
@@ -118,12 +118,7 @@ impl provider_transport::VideoTaskTransportSnapshotLookup for AppState {
|
||||
) -> Result<Option<GatewayProviderTransportSnapshot>, String> {
|
||||
self.read_provider_transport_snapshot(provider_id, endpoint_id, key_id)
|
||||
.await
|
||||
.map_err(|err| match err {
|
||||
GatewayError::UpstreamUnavailable { message, .. }
|
||||
| GatewayError::ControlUnavailable { message, .. }
|
||||
| GatewayError::Client { message, .. }
|
||||
| GatewayError::Internal(message) => message,
|
||||
})
|
||||
.map_err(GatewayError::into_message)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -135,12 +130,7 @@ impl ModelFetchTransportRuntime for AppState {
|
||||
) -> Result<Option<LocalResolvedOAuthRequestAuth>, String> {
|
||||
AppState::resolve_local_oauth_request_auth(self, transport)
|
||||
.await
|
||||
.map_err(|err| match err {
|
||||
GatewayError::UpstreamUnavailable { message, .. }
|
||||
| GatewayError::ControlUnavailable { message, .. }
|
||||
| GatewayError::Client { message, .. }
|
||||
| GatewayError::Internal(message) => message,
|
||||
})
|
||||
.map_err(GatewayError::into_message)
|
||||
}
|
||||
|
||||
async fn resolve_model_fetch_proxy(
|
||||
@@ -157,12 +147,7 @@ impl ModelFetchTransportRuntime for AppState {
|
||||
) -> Result<ExecutionResult, String> {
|
||||
execution_runtime::execute_execution_runtime_sync_plan(self, None, plan)
|
||||
.await
|
||||
.map_err(|err| match err {
|
||||
GatewayError::UpstreamUnavailable { message, .. }
|
||||
| GatewayError::ControlUnavailable { message, .. }
|
||||
| GatewayError::Client { message, .. }
|
||||
| GatewayError::Internal(message) => message,
|
||||
})
|
||||
.map_err(GatewayError::into_message)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -27,6 +27,7 @@ pub(crate) use self::admin_types::{
|
||||
UserDailyQuotaAvailabilityRecord, UserPlanEntitlementRecord,
|
||||
};
|
||||
pub use self::app::AppState;
|
||||
pub(crate) use self::app::FrontdoorRuntimeGuardConfig;
|
||||
pub(crate) use self::cache::{
|
||||
CachedProviderTransportSnapshot, AUTH_API_KEY_LAST_USED_MAX_ENTRIES,
|
||||
AUTH_API_KEY_LAST_USED_TTL, PROVIDER_TRANSPORT_SNAPSHOT_CACHE_MAX_ENTRIES,
|
||||
|
||||
@@ -1166,7 +1166,6 @@ impl AppState {
|
||||
|
||||
latest_key.encrypted_api_key = Some(encrypted_api_key);
|
||||
latest_key.encrypted_auth_config = encrypted_auth_config;
|
||||
latest_key.is_active = true;
|
||||
latest_key.expires_at_unix_secs = entry.expires_at_unix_secs;
|
||||
let (oauth_invalid_at_unix_secs, oauth_invalid_reason) =
|
||||
local_oauth_refresh_success_invalid_state(&latest_key);
|
||||
@@ -1448,12 +1447,7 @@ impl AppState {
|
||||
.map_err(
|
||||
|err| provider_transport::LocalOAuthRefreshError::InvalidResponse {
|
||||
provider_type,
|
||||
message: match err {
|
||||
GatewayError::UpstreamUnavailable { message, .. }
|
||||
| GatewayError::ControlUnavailable { message, .. }
|
||||
| GatewayError::Client { message, .. }
|
||||
| GatewayError::Internal(message) => message,
|
||||
},
|
||||
message: err.into_message(),
|
||||
},
|
||||
)?;
|
||||
let response_body_text = local_oauth_execution_body_text(&result);
|
||||
|
||||
@@ -369,6 +369,25 @@ impl AppState {
|
||||
Ok(api_key)
|
||||
}
|
||||
|
||||
pub(crate) async fn set_api_key_usage_totals(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
total_requests: u64,
|
||||
total_tokens: u64,
|
||||
total_cost_usd: f64,
|
||||
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
|
||||
{
|
||||
let api_key = self
|
||||
.data
|
||||
.set_api_key_usage_totals(api_key_id, total_requests, total_tokens, total_cost_usd)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if api_key.is_some() {
|
||||
self.invalidate_auth_context_cache();
|
||||
}
|
||||
Ok(api_key)
|
||||
}
|
||||
|
||||
pub(crate) async fn set_standalone_api_key_feature_settings(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
|
||||
@@ -11,7 +11,7 @@ use aether_data_contracts::repository::video_tasks::{
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
use super::{AppState, GatewayDataState};
|
||||
use super::{AppState, FrontdoorRuntimeGuardConfig, GatewayDataState};
|
||||
use crate::{provider_transport, usage};
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -31,6 +31,14 @@ impl AppState {
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) fn with_frontdoor_runtime_guard_config_for_tests(
|
||||
mut self,
|
||||
config: FrontdoorRuntimeGuardConfig,
|
||||
) -> Self {
|
||||
self.frontdoor_runtime_guards = Arc::new(config);
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) fn with_tunnel_identity_for_tests(
|
||||
mut self,
|
||||
instance_id: &str,
|
||||
|
||||
Reference in New Issue
Block a user