i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-01 17:00:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix OAuth import proxy handling

Browse Source
This commit is contained in:
fawney19
2026-04-27 20:05:43 +08:00
parent 4c16b11cb4
commit a7a3c9f023
11 changed files with 707 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/execution.rs
View File

@@ -4,11 +4,15 @@ use super::parse::{
parse_admin_provider_oauth_batch_import_entries, AdminProviderOAuthBatchImportEntry,
AdminProviderOAuthBatchImportOutcome,
};
use super::progress::{
maybe_report_admin_provider_oauth_batch_import_progress,
AdminProviderOAuthBatchProgressReporter,
};
use crate::handlers::admin::provider::oauth::duplicates::find_duplicate_provider_oauth_key;
use crate::handlers::admin::provider::oauth::provisioning::build_provider_oauth_auth_config_from_token_payload;
use crate::handlers::admin::provider::oauth::provisioning::{
create_provider_oauth_catalog_key, provider_oauth_active_api_formats,
update_existing_provider_oauth_catalog_key,
provider_oauth_key_proxy_value, update_existing_provider_oauth_catalog_key,
};
use crate::handlers::admin::provider::oauth::runtime::{
provider_oauth_runtime_endpoint_for_provider,
@@ -41,6 +45,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import_for_provider_type(
provider_type: &str,
raw_credentials: &str,
proxy_node_id: Option<&str>,
progress: Option<&mut dyn AdminProviderOAuthBatchProgressReporter>,
) -> Result<AdminProviderOAuthBatchImportOutcome, GatewayError> {
if provider_type.eq_ignore_ascii_case("kiro") {
execute_admin_provider_oauth_kiro_batch_import(
@@ -48,6 +53,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import_for_provider_type(
provider_id,
raw_credentials,
proxy_node_id,
progress,
)
.await
} else {
@@ -58,6 +64,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import_for_provider_type(
provider_type,
&entries,
proxy_node_id,
progress,
)
.await
}
@@ -69,6 +76,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
provider_type: &str,
entries: &[AdminProviderOAuthBatchImportEntry],
proxy_node_id: Option<&str>,
mut progress: Option<&mut dyn AdminProviderOAuthBatchProgressReporter>,
) -> Result<AdminProviderOAuthBatchImportOutcome, GatewayError> {
let Some(provider) = state
.read_provider_catalog_providers_by_ids(&[provider_id.to_string()])
@@ -131,6 +139,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
],
)
.await;
let key_proxy = provider_oauth_key_proxy_value(proxy_node_id);
let mut results = Vec::with_capacity(entries.len());
let mut success = 0usize;
let mut failed = 0usize;
@@ -156,6 +165,14 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
),
"replaced": false,
}));
maybe_report_admin_provider_oauth_batch_import_progress(
&mut progress,
entries.len(),
success,
failed,
&results,
)
.await;
continue;
}
};
@@ -170,6 +187,14 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
"error": "Token 刷新返回缺少 access_token",
"replaced": false,
}));
maybe_report_admin_provider_oauth_batch_import_progress(
&mut progress,
entries.len(),
success,
failed,
&results,
)
.await;
continue;
};
@@ -192,6 +217,14 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
"error": detail,
"replaced": false,
}));
maybe_report_admin_provider_oauth_batch_import_progress(
&mut progress,
entries.len(),
success,
failed,
&results,
)
.await;
continue;
}
};
@@ -205,7 +238,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
&access_token,
&auth_config,
&api_formats,
None,
key_proxy.clone(),
expires_at,
)
.await?
@@ -219,6 +252,14 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
"error": "provider oauth write unavailable",
"replaced": true,
}));
maybe_report_admin_provider_oauth_batch_import_progress(
&mut progress,
entries.len(),
success,
failed,
&results,
)
.await;
continue;
}
}
@@ -249,7 +290,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
&access_token,
&auth_config,
&api_formats,
None,
key_proxy.clone(),
expires_at,
)
.await?
@@ -263,6 +304,14 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
"error": "provider oauth write unavailable",
"replaced": false,
}));
maybe_report_admin_provider_oauth_batch_import_progress(
&mut progress,
entries.len(),
success,
failed,
&results,
)
.await;
continue;
}
}
@@ -284,6 +333,14 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
"error": serde_json::Value::Null,
"replaced": replaced,
}));
maybe_report_admin_provider_oauth_batch_import_progress(
&mut progress,
entries.len(),
success,
failed,
&results,
)
.await;
}
Ok(AdminProviderOAuthBatchImportOutcome {

76
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/kiro_import.rs
View File

@@ -3,10 +3,14 @@ use super::super::kiro::{
refresh_admin_provider_oauth_kiro_auth_config,
};
use super::parse::AdminProviderOAuthBatchImportOutcome;
use super::progress::{
maybe_report_admin_provider_oauth_batch_import_progress,
AdminProviderOAuthBatchProgressReporter,
};
use crate::handlers::admin::provider::oauth::duplicates::find_duplicate_provider_oauth_key;
use crate::handlers::admin::provider::oauth::provisioning::{
create_provider_oauth_catalog_key, provider_oauth_active_api_formats,
update_existing_provider_oauth_catalog_key,
provider_oauth_key_proxy_value, update_existing_provider_oauth_catalog_key,
};
use crate::handlers::admin::provider::oauth::runtime::{
provider_oauth_runtime_endpoint_for_provider,
@@ -27,6 +31,7 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
provider_id: &str,
raw_credentials: &str,
proxy_node_id: Option<&str>,
mut progress: Option<&mut dyn AdminProviderOAuthBatchProgressReporter>,
) -> Result<AdminProviderOAuthBatchImportOutcome, GatewayError> {
let entries = parse_admin_provider_oauth_kiro_batch_import_entries(raw_credentials);
let Some(provider) = state
@@ -70,6 +75,7 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
],
)
.await;
let key_proxy = provider_oauth_key_proxy_value(proxy_node_id);
let social_refresh_base_url =
admin_provider_oauth_kiro_refresh_base_url_override(state, "kiro_social_refresh");
let idc_refresh_base_url =
@@ -87,6 +93,14 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
"error": "未找到有效的凭据数据",
"replaced": false,
}));
maybe_report_admin_provider_oauth_batch_import_progress(
&mut progress,
entries.len(),
success,
failed,
&results,
)
.await;
continue;
};
@@ -103,6 +117,14 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
"error": "缺少可用的 Kiro refresh 凭据",
"replaced": false,
}));
maybe_report_admin_provider_oauth_batch_import_progress(
&mut progress,
entries.len(),
success,
failed,
&results,
)
.await;
continue;
}
@@ -124,6 +146,14 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
"error": format!("Token 验证失败: {err}"),
"replaced": false,
}));
maybe_report_admin_provider_oauth_batch_import_progress(
&mut progress,
entries.len(),
success,
failed,
&results,
)
.await;
continue;
}
};
@@ -166,6 +196,14 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
"error": detail,
"replaced": false,
}));
maybe_report_admin_provider_oauth_batch_import_progress(
&mut progress,
entries.len(),
success,
failed,
&results,
)
.await;
continue;
}
};
@@ -184,6 +222,14 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
"error": "Token 验证失败: accessToken 为空",
"replaced": false,
}));
maybe_report_admin_provider_oauth_batch_import_progress(
&mut progress,
entries.len(),
success,
failed,
&results,
)
.await;
continue;
};
@@ -196,7 +242,7 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
&access_token,
&auth_config,
&api_formats,
None,
key_proxy.clone(),
refreshed_auth_config.expires_at,
)
.await?
@@ -210,6 +256,14 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
"error": "provider oauth write unavailable",
"replaced": true,
}));
maybe_report_admin_provider_oauth_batch_import_progress(
&mut progress,
entries.len(),
success,
failed,
&results,
)
.await;
continue;
}
}
@@ -231,7 +285,7 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
&access_token,
&auth_config,
&api_formats,
None,
key_proxy.clone(),
refreshed_auth_config.expires_at,
)
.await?
@@ -245,6 +299,14 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
"error": "provider oauth write unavailable",
"replaced": false,
}));
maybe_report_admin_provider_oauth_batch_import_progress(
&mut progress,
entries.len(),
success,
failed,
&results,
)
.await;
continue;
}
}
@@ -271,6 +333,14 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
"error": serde_json::Value::Null,
"replaced": replaced,
}));
maybe_report_admin_provider_oauth_batch_import_progress(
&mut progress,
entries.len(),
success,
failed,
&results,
)
.await;
}
Ok(AdminProviderOAuthBatchImportOutcome {

1
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/mod.rs
View File

@@ -2,6 +2,7 @@ mod execution;
mod kiro_import;
mod orchestration;
mod parse;
mod progress;
mod task;
pub(super) use orchestration::handle_admin_provider_oauth_batch_import;

1
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/orchestration.rs
View File

@@ -81,6 +81,7 @@ pub(in super::super) async fn handle_admin_provider_oauth_batch_import(
&provider_type,
payload.credentials.as_str(),
payload.proxy_node_id.as_deref(),
None,
)
.await?;
Ok(build_admin_provider_oauth_batch_import_response(&outcome).into_response())

47
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/progress.rs Normal file
View File

@@ -0,0 +1,47 @@
use serde_json::Value;
#[derive(Debug, Clone)]
pub(super) struct AdminProviderOAuthBatchImportProgress {
pub total: usize,
pub processed: usize,
pub success: usize,
pub failed: usize,
pub created_count: usize,
pub replaced_count: usize,
pub latest_result: Option<Value>,
}
#[async_trait::async_trait]
pub(super) trait AdminProviderOAuthBatchProgressReporter: Send {
async fn report(&mut self, progress: AdminProviderOAuthBatchImportProgress);
}
pub(super) async fn maybe_report_admin_provider_oauth_batch_import_progress(
reporter: &mut Option<&mut dyn AdminProviderOAuthBatchProgressReporter>,
total: usize,
success: usize,
failed: usize,
results: &[Value],
) {
let Some(reporter) = reporter.as_deref_mut() else {
return;
};
let replaced_count = results
.iter()
.filter(|item| {
item.get("status").and_then(Value::as_str) == Some("success")
&& item.get("replaced").and_then(Value::as_bool) == Some(true)
})
.count();
reporter
.report(AdminProviderOAuthBatchImportProgress {
total,
processed: success.saturating_add(failed).min(total),
success,
failed,
created_count: success.saturating_sub(replaced_count),
replaced_count,
latest_result: results.last().cloned(),
})
.await;
}

62
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/task.rs
View File

@@ -5,6 +5,9 @@ use super::execution::{
use super::parse::{
build_admin_provider_oauth_batch_task_state, parse_admin_provider_oauth_batch_import_request,
};
use super::progress::{
AdminProviderOAuthBatchImportProgress, AdminProviderOAuthBatchProgressReporter,
};
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
use crate::handlers::admin::provider::oauth::state::{
admin_provider_oauth_template, build_admin_provider_oauth_backend_unavailable_response,
@@ -26,6 +29,55 @@ use uuid::Uuid;
const PROVIDER_OAUTH_BATCH_TASK_MAX_ERROR_SAMPLES: usize = 20;
struct BatchTaskProgressReporter {
app: crate::AppState,
task_id: String,
provider_id: String,
provider_type: String,
created_at: u64,
started_at: u64,
error_samples: Vec<serde_json::Value>,
}
#[async_trait::async_trait]
impl AdminProviderOAuthBatchProgressReporter for BatchTaskProgressReporter {
async fn report(&mut self, progress: AdminProviderOAuthBatchImportProgress) {
if let Some(latest_result) = progress.latest_result.as_ref() {
if latest_result
.get("status")
.and_then(serde_json::Value::as_str)
== Some("error")
&& self.error_samples.len() < PROVIDER_OAUTH_BATCH_TASK_MAX_ERROR_SAMPLES
{
self.error_samples.push(latest_result.clone());
}
}
let message = format!("处理中 {}/{}", progress.processed, progress.total);
let progress_state = build_admin_provider_oauth_batch_task_state(
&self.task_id,
&self.provider_id,
&self.provider_type,
"processing",
progress.total,
progress.processed,
progress.success,
progress.failed,
progress.created_count,
progress.replaced_count,
Some(message.as_str()),
None,
self.error_samples.clone(),
self.created_at,
Some(self.started_at),
None,
);
let _ = AdminAppState::new(&self.app)
.save_provider_oauth_batch_task_payload(&self.task_id, &progress_state)
.await;
}
}
pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_task(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
@@ -149,12 +201,22 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
.save_provider_oauth_batch_task_payload(&task_id_for_worker, &processing_state)
.await;
let mut progress_reporter = BatchTaskProgressReporter {
app: task_state.clone(),
task_id: task_id_for_worker.clone(),
provider_id: provider_id_for_worker.clone(),
provider_type: provider_type_for_worker.clone(),
created_at,
started_at,
error_samples: Vec::new(),
};
match execute_admin_provider_oauth_batch_import_for_provider_type(
&AdminAppState::new(&task_state),
&provider_id_for_worker,
&provider_type_for_worker,
raw_credentials.as_str(),
proxy_node_id.as_deref(),
Some(&mut progress_reporter),
)
.await
{

8
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/provider.rs
View File

@@ -2,7 +2,8 @@ use super::super::super::duplicates::find_duplicate_provider_oauth_key;
use super::super::super::errors::build_internal_control_error_response;
use super::super::super::provisioning::{
build_provider_oauth_auth_config_from_token_payload, create_provider_oauth_catalog_key,
provider_oauth_active_api_formats, update_existing_provider_oauth_catalog_key,
provider_oauth_active_api_formats, provider_oauth_key_proxy_value,
update_existing_provider_oauth_catalog_key,
};
use super::super::super::runtime::{
provider_oauth_runtime_endpoint_for_provider,
@@ -131,6 +132,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
],
)
.await;
let key_proxy = provider_oauth_key_proxy_value(payload.proxy_node_id.as_deref());
let token_payload = match state
.exchange_admin_provider_oauth_code(
@@ -178,7 +180,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
&access_token,
&auth_config,
&api_formats,
None,
key_proxy.clone(),
expires_at,
)
.await?
@@ -220,7 +222,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
&access_token,
&auth_config,
&api_formats,
None,
key_proxy.clone(),
expires_at,
)
.await?

8
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/import.rs
View File

@@ -2,7 +2,8 @@ use super::super::duplicates::find_duplicate_provider_oauth_key;
use super::super::errors::build_internal_control_error_response;
use super::super::provisioning::{
build_provider_oauth_auth_config_from_token_payload, create_provider_oauth_catalog_key,
provider_oauth_active_api_formats, update_existing_provider_oauth_catalog_key,
provider_oauth_active_api_formats, provider_oauth_key_proxy_value,
update_existing_provider_oauth_catalog_key,
};
use super::super::runtime::{
provider_oauth_runtime_endpoint_for_provider,
@@ -119,6 +120,7 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
],
)
.await;
let key_proxy = provider_oauth_key_proxy_value(proxy_node_id.as_deref());
let token_payload = match state
.exchange_admin_provider_oauth_refresh_token(
@@ -170,7 +172,7 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
&access_token,
&auth_config,
&api_formats,
None,
key_proxy.clone(),
expires_at,
)
.await?
@@ -211,7 +213,7 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
&access_token,
&auth_config,
&api_formats,
None,
key_proxy.clone(),
expires_at,
)
.await?

147
apps/aether-gateway/src/state/proxy.rs
View File

@@ -33,7 +33,12 @@ impl AppState {
}
if node.tunnel_mode && node.tunnel_connected {
let mut extra = Map::new();
if let Ok(Some(owner)) = self.lookup_tunnel_attachment_owner(node_id).await {
let owner = self
.lookup_tunnel_attachment_owner(node_id)
.await
.ok()
.flatten();
if let Some(owner) = owner {
extra.insert(
TUNNEL_BASE_URL_EXTRA_KEY.to_string(),
Value::String(owner.relay_base_url),
@@ -46,6 +51,8 @@ impl AppState {
TUNNEL_OWNER_OBSERVED_AT_EXTRA_KEY.to_string(),
json!(owner.observed_at_unix_secs),
);
} else if !self.tunnel.has_local_proxy(node_id) {
return None;
}
return Some(ProxySnapshot {
enabled: Some(true),
@@ -119,6 +126,13 @@ impl AppState {
if let Some(snapshot) = self.resolve_proxy_node_snapshot(node_id.as_deref()).await {
return Some(snapshot);
}
if let Some(node_id) = node_id.as_deref() {
if !proxy_object_has_inline_url(object)
&& self.find_proxy_node(node_id).await.ok().flatten().is_some()
{
return None;
}
}
proxy_snapshot_from_object(object)
}
@@ -197,6 +211,10 @@ fn proxy_snapshot_from_object(object: &Map<String, Value>) -> Option<ProxySnapsh
})
}
fn proxy_object_has_inline_url(object: &Map<String, Value>) -> bool {
json_string_field(object, "url").is_some() || json_string_field(object, "proxy_url").is_some()
}
fn json_string_field(object: &Map<String, Value>, key: &str) -> Option<String> {
object
.get(key)
@@ -241,7 +259,13 @@ fn proxy_url_with_node_auth(
#[cfg(test)]
mod tests {
use std::sync::Arc;
use aether_data::repository::proxy_nodes::{InMemoryProxyNodeRepository, StoredProxyNode};
use serde_json::json;
use super::proxy_url_with_node_auth;
use crate::{data::GatewayDataState, AppState};
#[test]
fn proxy_url_with_node_auth_omits_empty_password_separator() {
@@ -250,4 +274,125 @@ mod tests {
Some("socks5://alice@proxy.example:1080")
);
}
#[tokio::test]
async fn resolve_proxy_node_snapshot_rejects_unroutable_tunnel_node() {
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![sample_tunnel_node(
"proxy-node-stale",
)]));
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
repository,
));
let snapshot = state
.resolve_proxy_node_snapshot(Some("proxy-node-stale"))
.await;
assert_eq!(snapshot, None);
}
#[tokio::test]
async fn resolve_proxy_node_snapshot_keeps_tunnel_node_with_owner_hint() {
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![sample_tunnel_node(
"proxy-node-owned",
)]));
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::with_proxy_node_repository_for_tests(repository)
.with_system_config_values_for_tests(vec![(
"tunnel.attachments.proxy-node-owned".to_string(),
json!({
"gateway_instance_id": "gateway-owner",
"relay_base_url": "http://gateway-owner.internal",
"conn_count": 1,
"observed_at_unix_secs": 4_102_444_800u64,
}),
)]),
);
let snapshot = state
.resolve_proxy_node_snapshot(Some("proxy-node-owned"))
.await
.expect("owned tunnel snapshot should resolve");
assert_eq!(snapshot.mode.as_deref(), Some("tunnel"));
assert_eq!(snapshot.node_id.as_deref(), Some("proxy-node-owned"));
assert_eq!(
snapshot
.extra
.as_ref()
.and_then(|extra| extra.get("tunnel_base_url"))
.and_then(serde_json::Value::as_str),
Some("http://gateway-owner.internal")
);
}
#[tokio::test]
async fn resolve_configured_proxy_snapshot_rejects_unroutable_stored_tunnel_reference() {
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![sample_tunnel_node(
"proxy-node-stale",
)]));
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
repository,
));
let snapshot = state
.resolve_configured_proxy_snapshot_with_tunnel_affinity(Some(&json!({
"node_id": "proxy-node-stale",
"enabled": true,
})))
.await;
assert_eq!(snapshot, None);
}
#[tokio::test]
async fn resolve_configured_proxy_snapshot_keeps_inline_url_when_stored_node_is_unroutable() {
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![sample_tunnel_node(
"proxy-node-stale",
)]));
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
repository,
));
let snapshot = state
.resolve_configured_proxy_snapshot_with_tunnel_affinity(Some(&json!({
"node_id": "proxy-node-stale",
"url": "http://proxy.example:8080",
"enabled": true,
})))
.await
.expect("inline proxy URL should still resolve");
assert_eq!(snapshot.node_id.as_deref(), Some("proxy-node-stale"));
assert_eq!(snapshot.url.as_deref(), Some("http://proxy.example:8080"));
}
fn sample_tunnel_node(id: &str) -> StoredProxyNode {
StoredProxyNode::new(
id.to_string(),
id.to_string(),
"127.0.0.1".to_string(),
0,
false,
"online".to_string(),
15,
1,
0,
0,
0,
0,
true,
true,
1,
)
.expect("sample tunnel node should build")
}
}

198
apps/aether-gateway/src/tests/control/admin/oauth.rs
View File

@@ -1465,7 +1465,10 @@ async fn gateway_batch_imports_admin_provider_oauth_locally_with_trusted_admin_p
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
assert!(persisted.is_active);
assert_eq!(persisted.proxy, None);
assert_eq!(
persisted.proxy,
Some(json!({"node_id": "proxy-node-batch-import", "enabled": true}))
);
let decrypted_api_key =
decrypt_python_fernet_ciphertext(DEVELOPMENT_ENCRYPTION_KEY, &persisted.encrypted_api_key)
.expect("api key should decrypt");
@@ -1616,6 +1619,164 @@ async fn gateway_starts_admin_provider_oauth_batch_import_task_locally_with_trus
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_updates_admin_provider_oauth_batch_import_task_progress() {
let upstream = Router::new().fallback(any(|| async {
(StatusCode::OK, Body::from("quota refresh body"))
}));
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let token_server = Router::new().route(
"/oauth/token",
any(move |_request: Request| {
let token_hits_inner = Arc::clone(&token_hits_clone);
async move {
let hit = {
let mut guard = token_hits_inner.lock().expect("mutex should lock");
*guard += 1;
*guard
};
if hit == 2 {
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
}
Json(json!({
"access_token": format!("progress-codex-access-token-{hit}"),
"refresh_token": format!("progress-codex-refresh-token-{hit}"),
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": format!("progress-{hit}@example.com"),
"account_id": format!("acct-progress-{hit}"),
"account_user_id": format!("acct-user-progress-{hit}"),
"plan_type": "plus",
}))
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-chat",
"provider-codex",
"openai:chat",
"https://chatgpt.com/backend-api/codex",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (token_url, token_handle) = start_server(token_server).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests("codex", format!("{token_url}/oauth/token")),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let submit_response = client
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/batch-import/tasks"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"credentials": "progress-refresh-one\nprogress-refresh-two"
}))
.send()
.await
.expect("submit request should succeed");
assert_eq!(submit_response.status(), StatusCode::OK);
let submit_payload: serde_json::Value = submit_response
.json()
.await
.expect("submit payload should parse");
let task_id = submit_payload["task_id"]
.as_str()
.expect("task id should exist")
.to_string();
let mut progress_payload = serde_json::Value::Null;
for _ in 0..50 {
let response = client
.get(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/batch-import/tasks/{task_id}"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("status request should succeed");
assert_eq!(response.status(), StatusCode::OK);
progress_payload = response.json().await.expect("status payload should parse");
if progress_payload["status"] == "processing" && progress_payload["processed"] == 1 {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
}
assert_eq!(
progress_payload["status"], "processing",
"payload={progress_payload}"
);
assert_eq!(progress_payload["total"], 2);
assert_eq!(progress_payload["processed"], 1);
assert_eq!(progress_payload["success"], 1);
assert_eq!(progress_payload["failed"], 0);
assert_eq!(progress_payload["created_count"], 1);
assert_eq!(progress_payload["progress_percent"], 50);
let mut completed_payload = progress_payload;
for _ in 0..50 {
let response = client
.get(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/batch-import/tasks/{task_id}"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("status request should succeed");
assert_eq!(response.status(), StatusCode::OK);
completed_payload = response.json().await.expect("status payload should parse");
if completed_payload["status"] == "completed" {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
}
assert_eq!(
completed_payload["status"], "completed",
"payload={completed_payload}"
);
assert_eq!(completed_payload["processed"], 2);
assert_eq!(completed_payload["success"], 2);
assert_eq!(completed_payload["progress_percent"], 100);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 2);
gateway_handle.abort();
token_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_principal() {
#[derive(Debug, Clone)]
@@ -1935,7 +2096,10 @@ async fn gateway_completes_admin_provider_oauth_provider_locally_with_trusted_ad
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
assert!(persisted.is_active);
assert_eq!(persisted.proxy, None);
assert_eq!(
persisted.proxy,
Some(json!({"node_id": "proxy-node-codex-oauth", "enabled": true}))
);
let decrypted_api_key =
decrypt_python_fernet_ciphertext(DEVELOPMENT_ENCRYPTION_KEY, &persisted.encrypted_api_key)
.expect("api key should decrypt");
@@ -2103,7 +2267,10 @@ async fn gateway_imports_admin_provider_oauth_refresh_token_locally_with_trusted
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
assert!(persisted.is_active);
assert_eq!(persisted.proxy, None);
assert_eq!(
persisted.proxy,
Some(json!({"node_id": "proxy-node-codex-import", "enabled": true}))
);
let decrypted_api_key =
decrypt_python_fernet_ciphertext(DEVELOPMENT_ENCRYPTION_KEY, &persisted.encrypted_api_key)
.expect("api key should decrypt");
@@ -2275,7 +2442,10 @@ async fn gateway_imports_admin_provider_oauth_refresh_token_over_active_expired_
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
assert!(persisted.is_active);
assert_eq!(persisted.proxy, None);
assert_eq!(
persisted.proxy,
Some(json!({"node_id": "proxy-node-codex-import", "enabled": true}))
);
assert_eq!(persisted.oauth_invalid_at_unix_secs, None);
assert_eq!(persisted.oauth_invalid_reason, None);
let decrypted_api_key =
@@ -2473,7 +2643,10 @@ async fn gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtim
.await
.expect("keys should load");
assert_eq!(keys.len(), 1);
assert_eq!(keys[0].proxy, None);
assert_eq!(
keys[0].proxy,
Some(json!({"node_id": "proxy-node-codex-import", "enabled": true}))
);
let plans = execution_plans.lock().expect("mutex should lock");
assert_eq!(plans.len(), 1);
@@ -2907,7 +3080,10 @@ async fn gateway_batch_imports_admin_provider_oauth_kiro_locally_with_trusted_ad
.next()
.expect("persisted key should exist");
assert!(stored_key.is_active);
assert_eq!(stored_key.proxy, None);
assert_eq!(
stored_key.proxy,
Some(json!({"node_id": "proxy-node-kiro-batch", "enabled": true}))
);
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
stored_key
@@ -3048,7 +3224,10 @@ async fn gateway_batch_imports_admin_provider_oauth_kiro_over_active_expired_dup
.next()
.expect("persisted key should exist");
assert!(stored_key.is_active);
assert_eq!(stored_key.proxy, None);
assert_eq!(
stored_key.proxy,
Some(json!({"node_id": "proxy-node-kiro-batch", "enabled": true}))
);
assert_eq!(stored_key.oauth_invalid_at_unix_secs, None);
assert_eq!(stored_key.oauth_invalid_reason, None);
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
@@ -3244,7 +3423,10 @@ async fn gateway_batch_imports_admin_provider_oauth_kiro_via_execution_runtime_p
.next()
.expect("persisted key should exist");
assert!(stored_key.is_active);
assert_eq!(stored_key.proxy, None);
assert_eq!(
stored_key.proxy,
Some(json!({"node_id": "proxy-node-kiro-batch-runtime", "enabled": true}))
);
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
stored_key