mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-13 06:30:20 +08:00
fix(provider): harden Agent Identity OAuth lifecycle
This commit is contained in:
Generated
+1
@@ -535,6 +535,7 @@ version = "0.1.0"
|
||||
dependencies = [
|
||||
"aether-data-contracts",
|
||||
"aether-pool-core",
|
||||
"aether-provider-transport",
|
||||
"serde_json",
|
||||
"url",
|
||||
"uuid",
|
||||
|
||||
@@ -222,6 +222,28 @@ pub(super) fn classify_oauth_route(
|
||||
"admin:provider_oauth",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::POST
|
||||
&& normalized_path.starts_with("/api/admin/provider-oauth/providers/")
|
||||
&& normalized_path.ends_with("/agent-identity-import/tasks")
|
||||
{
|
||||
Some(classified(
|
||||
"admin_proxy",
|
||||
"provider_oauth_manage",
|
||||
"start_agent_identity_import_task",
|
||||
"admin:provider_oauth",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::GET
|
||||
&& normalized_path.starts_with("/api/admin/provider-oauth/providers/")
|
||||
&& normalized_path.contains("/agent-identity-import/tasks/")
|
||||
{
|
||||
Some(classified(
|
||||
"admin_proxy",
|
||||
"provider_oauth_manage",
|
||||
"get_agent_identity_import_task_status",
|
||||
"admin:provider_oauth",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::POST
|
||||
&& normalized_path.starts_with("/api/admin/provider-oauth/providers/")
|
||||
&& normalized_path.ends_with("/batch-import")
|
||||
|
||||
@@ -109,6 +109,20 @@ fn classifies_admin_provider_oauth_maintenance_routes_as_admin_proxy_route() {
|
||||
"admin:provider_oauth",
|
||||
"admin:provider_oauth:write",
|
||||
),
|
||||
(
|
||||
http::Method::POST,
|
||||
"/api/admin/provider-oauth/providers/provider-123/agent-identity-import/tasks",
|
||||
"start_agent_identity_import_task",
|
||||
"admin:provider_oauth",
|
||||
"admin:provider_oauth:write",
|
||||
),
|
||||
(
|
||||
http::Method::GET,
|
||||
"/api/admin/provider-oauth/providers/provider-123/agent-identity-import/tasks/agent-identity-task-123",
|
||||
"get_agent_identity_import_task_status",
|
||||
"admin:provider_oauth",
|
||||
"admin:provider_oauth:read",
|
||||
),
|
||||
(
|
||||
http::Method::POST,
|
||||
"/api/admin/provider-oauth/providers/provider-123/batch-import",
|
||||
|
||||
@@ -2,12 +2,12 @@ use super::{
|
||||
ApiKeyLastUsedDelta, DataLayerError, GatewayDataState, GeminiFileMappingListQuery,
|
||||
GeminiFileMappingStats, ProviderCatalogKeyAdaptiveStateUpdate,
|
||||
ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListQuery,
|
||||
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
|
||||
PublicHealthStatusCount, PublicHealthTimelineBucket, StoredGeminiFileMapping,
|
||||
StoredGeminiFileMappingListPage, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
|
||||
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, StoredRequestCandidate,
|
||||
UpsertGeminiFileMappingRecord, UpsertRequestCandidateRecord,
|
||||
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
|
||||
ProviderCatalogKeyStatusSnapshotUpdate, PublicHealthStatusCount, PublicHealthTimelineBucket,
|
||||
StoredGeminiFileMapping, StoredGeminiFileMappingListPage, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary,
|
||||
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
||||
StoredRequestCandidate, UpsertGeminiFileMappingRecord, UpsertRequestCandidateRecord,
|
||||
};
|
||||
|
||||
impl GatewayDataState {
|
||||
@@ -400,6 +400,24 @@ impl GatewayDataState {
|
||||
Ok(updated)
|
||||
}
|
||||
|
||||
pub(crate) async fn compare_and_update_provider_catalog_key_oauth_runtime_state(
|
||||
&self,
|
||||
update: &ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
let updated = match &self.provider_catalog_writer {
|
||||
Some(repository) => {
|
||||
repository
|
||||
.compare_and_update_key_oauth_runtime_state(update)
|
||||
.await
|
||||
}
|
||||
None => Ok(false),
|
||||
}?;
|
||||
// A false result is a credential CAS conflict. Clear cached snapshots
|
||||
// either way so the next read observes the authoritative row.
|
||||
self.clear_provider_catalog_cache();
|
||||
Ok(updated)
|
||||
}
|
||||
|
||||
pub(crate) async fn create_provider_catalog_key(
|
||||
&self,
|
||||
key: &StoredProviderCatalogKey,
|
||||
|
||||
@@ -122,11 +122,11 @@ use aether_data_contracts::repository::pool_scores::{
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate,
|
||||
ProviderCatalogKeyListQuery, ProviderCatalogKeyRuntimeMetadataUpdate,
|
||||
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository,
|
||||
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
|
||||
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
||||
ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
|
||||
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
|
||||
ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary,
|
||||
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_data_contracts::repository::quota::{
|
||||
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, StoredProviderQuotaSnapshot,
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use aether_contracts::ExecutionPlan;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::state::AgentIdentityAuthConfigFence;
|
||||
use crate::{provider_transport::LocalOAuthRefreshError, AppState};
|
||||
|
||||
pub(crate) async fn refresh_oauth_plan_auth_for_retry(
|
||||
@@ -13,8 +14,11 @@ pub(crate) async fn refresh_oauth_plan_auth_for_retry(
|
||||
if !status_may_be_oauth_invalid(status_code, response_text) {
|
||||
return false;
|
||||
}
|
||||
let access_token_invalid_proven =
|
||||
status_proves_access_token_invalid(status_code, response_text);
|
||||
let request_authorization = execution_plan_authorization(plan);
|
||||
let request_uses_agent_identity = request_authorization
|
||||
.is_some_and(aether_provider_transport::is_codex_agent_identity_authorization);
|
||||
let access_token_invalid_proven = !request_uses_agent_identity
|
||||
&& status_proves_access_token_invalid(status_code, response_text);
|
||||
|
||||
let transport = match state
|
||||
.read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
|
||||
@@ -37,11 +41,44 @@ pub(crate) async fn refresh_oauth_plan_auth_for_retry(
|
||||
}
|
||||
};
|
||||
|
||||
if aether_provider_transport::is_codex_agent_identity_transport(&transport)
|
||||
&& !aether_provider_transport::is_codex_agent_identity_invalid_task_response(
|
||||
status_code,
|
||||
response_text,
|
||||
)
|
||||
let current_uses_agent_identity =
|
||||
aether_provider_transport::is_codex_agent_identity_transport(&transport);
|
||||
if request_uses_agent_identity {
|
||||
if !current_uses_agent_identity
|
||||
|| !aether_provider_transport::is_codex_agent_identity_invalid_task_response(
|
||||
status_code,
|
||||
response_text,
|
||||
)
|
||||
|| !request_authorization.is_some_and(|authorization| {
|
||||
aether_provider_transport::codex_agent_identity_authorization_matches_transport(
|
||||
&transport,
|
||||
authorization,
|
||||
)
|
||||
})
|
||||
{
|
||||
return false;
|
||||
}
|
||||
if !matches!(
|
||||
state
|
||||
.capture_agent_identity_auth_config_fence(&transport)
|
||||
.await,
|
||||
Ok(AgentIdentityAuthConfigFence::Current(_))
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
} else if current_uses_agent_identity {
|
||||
// A bearer-token response cannot authorize refreshing an Agent Identity
|
||||
// installed under the same key id while the request was in flight.
|
||||
return false;
|
||||
} else if transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("codex")
|
||||
&& transport.key.auth_type.trim().eq_ignore_ascii_case("oauth")
|
||||
&& !request_authorization.is_some_and(|authorization| {
|
||||
bearer_authorization_matches_transport(authorization, &transport)
|
||||
})
|
||||
{
|
||||
return false;
|
||||
}
|
||||
@@ -120,6 +157,26 @@ pub(crate) async fn refresh_oauth_plan_auth_for_retry(
|
||||
}
|
||||
}
|
||||
|
||||
fn execution_plan_authorization(plan: &ExecutionPlan) -> Option<&str> {
|
||||
plan.headers
|
||||
.iter()
|
||||
.find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
|
||||
.map(|(_, value)| value.as_str())
|
||||
}
|
||||
|
||||
fn bearer_authorization_matches_transport(
|
||||
authorization: &str,
|
||||
transport: &aether_provider_transport::GatewayProviderTransportSnapshot,
|
||||
) -> bool {
|
||||
let current_token = transport.key.decrypted_api_key.trim();
|
||||
!current_token.is_empty()
|
||||
&& authorization
|
||||
.trim()
|
||||
.strip_prefix("Bearer ")
|
||||
.map(str::trim)
|
||||
.is_some_and(|token| token == current_token)
|
||||
}
|
||||
|
||||
fn status_may_be_oauth_invalid(status_code: u16, response_text: Option<&str>) -> bool {
|
||||
if status_code == 401 {
|
||||
return true;
|
||||
|
||||
@@ -98,7 +98,9 @@ use crate::execution_runtime::{
|
||||
resolve_local_candidate_failover_analysis_stream, should_fallback_to_control_stream,
|
||||
should_retry_next_local_candidate_stream, LocalFailoverDecision,
|
||||
};
|
||||
use crate::execution_runtime::{MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES};
|
||||
use crate::execution_runtime::{
|
||||
MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES,
|
||||
};
|
||||
use crate::log_ids::short_request_id;
|
||||
use crate::orchestration::{
|
||||
apply_local_execution_effect, build_local_error_flow_metadata, cyber_continue_failover_enabled,
|
||||
@@ -1120,9 +1122,22 @@ async fn execute_in_process_stream_with_oauth_retry(
|
||||
) -> Result<DirectUpstreamStreamExecution, InProcessStreamExecutionError> {
|
||||
let mut execution = execute_in_process_stream(state, plan, trace_id).await?;
|
||||
apply_stream_summary_report_context(&mut execution, report_context);
|
||||
let response_text = if execution.status_code == 401
|
||||
&& stream_plan_uses_codex_agent_identity(state, plan).await
|
||||
{
|
||||
prefetch_direct_stream_error_body(&mut execution).await
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if execution.status_code >= 400
|
||||
&& refresh_oauth_plan_auth_for_retry(state, plan, execution.status_code, None, trace_id)
|
||||
.await
|
||||
&& refresh_oauth_plan_auth_for_retry(
|
||||
state,
|
||||
plan,
|
||||
execution.status_code,
|
||||
response_text.as_deref(),
|
||||
trace_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
drop(execution);
|
||||
execution = execute_in_process_stream(state, plan, trace_id).await?;
|
||||
@@ -1131,6 +1146,103 @@ async fn execute_in_process_stream_with_oauth_retry(
|
||||
Ok(execution)
|
||||
}
|
||||
|
||||
async fn stream_plan_uses_codex_agent_identity(state: &AppState, plan: &ExecutionPlan) -> bool {
|
||||
state
|
||||
.read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
.as_ref()
|
||||
.is_some_and(aether_provider_transport::is_codex_agent_identity_transport)
|
||||
}
|
||||
|
||||
async fn next_direct_upstream_response_chunk(
|
||||
response: &mut DirectUpstreamResponse,
|
||||
) -> Result<Option<Bytes>, String> {
|
||||
match response {
|
||||
DirectUpstreamResponse::Reqwest(response) => response
|
||||
.chunk()
|
||||
.await
|
||||
.map_err(|err| format_upstream_request_error(&err)),
|
||||
DirectUpstreamResponse::HyperH2c(response) => loop {
|
||||
let Some(frame) = response.body_mut().frame().await else {
|
||||
return Ok(None);
|
||||
};
|
||||
let frame = frame.map_err(|err| format_hyper_error_chain(&err))?;
|
||||
if let Ok(chunk) = frame.into_data() {
|
||||
return Ok(Some(chunk));
|
||||
}
|
||||
},
|
||||
DirectUpstreamResponse::BrowserWreq(response) => response
|
||||
.chunk()
|
||||
.await
|
||||
.map_err(|err| format_wreq_upstream_request_error(&err)),
|
||||
DirectUpstreamResponse::LocalTunnel(response) => response.next_chunk().await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn prefetch_direct_stream_error_body(
|
||||
execution: &mut DirectUpstreamStreamExecution,
|
||||
) -> Option<String> {
|
||||
let mut inspected = Vec::with_capacity(MAX_ERROR_BODY_BYTES);
|
||||
let mut fully_buffered = false;
|
||||
while inspected.len() < MAX_ERROR_BODY_BYTES {
|
||||
let next_chunk = if execution.prefetched_body.is_empty() {
|
||||
match await_direct_passthrough_first_item(
|
||||
next_direct_upstream_response_chunk(&mut execution.response),
|
||||
execution.started_at,
|
||||
execution.stream_first_byte_timeout,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(item) => item,
|
||||
Err(_) => break,
|
||||
}
|
||||
} else {
|
||||
next_direct_upstream_response_chunk(&mut execution.response).await
|
||||
};
|
||||
let chunk = match next_chunk {
|
||||
Ok(Some(chunk)) => chunk,
|
||||
Ok(None) => {
|
||||
fully_buffered = true;
|
||||
break;
|
||||
}
|
||||
Err(error) => {
|
||||
execution.prefetched_body.push_back(Err(error));
|
||||
break;
|
||||
}
|
||||
};
|
||||
if chunk.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let remaining = MAX_ERROR_BODY_BYTES.saturating_sub(inspected.len());
|
||||
inspected.extend_from_slice(&chunk[..chunk.len().min(remaining)]);
|
||||
execution.prefetched_body.push_back(Ok(chunk));
|
||||
|
||||
let response_text = String::from_utf8_lossy(&inspected);
|
||||
if aether_provider_transport::is_codex_agent_identity_invalid_task_response(
|
||||
execution.status_code,
|
||||
Some(response_text.as_ref()),
|
||||
) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if inspected.is_empty() {
|
||||
return None;
|
||||
}
|
||||
if fully_buffered {
|
||||
let (body_json, _) = decode_stream_error_body(&execution.headers, &inspected);
|
||||
if let Some(body_json) = body_json {
|
||||
if let Ok(response_text) = serde_json::to_string(&body_json) {
|
||||
return Some(response_text);
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(String::from_utf8_lossy(&inspected).into_owned())
|
||||
}
|
||||
|
||||
fn should_use_direct_sse_passthrough(
|
||||
plan: &ExecutionPlan,
|
||||
plan_kind: &str,
|
||||
@@ -1184,9 +1296,10 @@ fn should_use_direct_sse_passthrough(
|
||||
type DirectUpstreamByteStream = BoxStream<'static, Result<Bytes, String>>;
|
||||
|
||||
fn direct_upstream_response_byte_stream(
|
||||
prefetched_body: VecDeque<Result<Bytes, String>>,
|
||||
response: DirectUpstreamResponse,
|
||||
) -> DirectUpstreamByteStream {
|
||||
match response {
|
||||
let response_stream = match response {
|
||||
DirectUpstreamResponse::Reqwest(response) => response
|
||||
.bytes_stream()
|
||||
.map(|item| item.map_err(|err| format_upstream_request_error(&err)))
|
||||
@@ -1213,7 +1326,10 @@ fn direct_upstream_response_byte_stream(
|
||||
}
|
||||
}
|
||||
.boxed(),
|
||||
}
|
||||
};
|
||||
futures_stream::iter(prefetched_body)
|
||||
.chain(response_stream)
|
||||
.boxed()
|
||||
}
|
||||
|
||||
async fn await_direct_passthrough_first_item<T, F>(
|
||||
@@ -1952,12 +2068,14 @@ impl DirectPassthroughFinalizerCore {
|
||||
|
||||
fn build_direct_passthrough_inline_body_stream(
|
||||
finalizer: DirectPassthroughFinalizer,
|
||||
prefetched_body: VecDeque<Result<Bytes, String>>,
|
||||
response: DirectUpstreamResponse,
|
||||
upstream_started_at: Instant,
|
||||
stream_first_byte_timeout: Option<Duration>,
|
||||
) -> impl futures_util::Stream<Item = Result<Bytes, IoError>> + Send + 'static {
|
||||
let state = DirectPassthroughInlineBodyState::new(
|
||||
finalizer,
|
||||
prefetched_body,
|
||||
response,
|
||||
upstream_started_at,
|
||||
stream_first_byte_timeout,
|
||||
@@ -1982,13 +2100,17 @@ struct DirectPassthroughInlineBodyState {
|
||||
impl DirectPassthroughInlineBodyState {
|
||||
fn new(
|
||||
finalizer: DirectPassthroughFinalizer,
|
||||
prefetched_body: VecDeque<Result<Bytes, String>>,
|
||||
response: DirectUpstreamResponse,
|
||||
upstream_started_at: Instant,
|
||||
stream_first_byte_timeout: Option<Duration>,
|
||||
) -> Self {
|
||||
Self {
|
||||
finalizer: Some(finalizer),
|
||||
upstream: Some(direct_upstream_response_byte_stream(response)),
|
||||
upstream: Some(direct_upstream_response_byte_stream(
|
||||
prefetched_body,
|
||||
response,
|
||||
)),
|
||||
upstream_control_filter: Some(SseControlBlockFilter::default()),
|
||||
upstream_started_at,
|
||||
stream_first_byte_timeout,
|
||||
@@ -2271,6 +2393,7 @@ async fn execute_stream_from_direct_passthrough(
|
||||
mut headers,
|
||||
provider_api_format: _,
|
||||
stream_summary_report_context: _,
|
||||
prefetched_body,
|
||||
response,
|
||||
started_at: upstream_started_at,
|
||||
stream_first_byte_timeout,
|
||||
@@ -2405,6 +2528,7 @@ async fn execute_stream_from_direct_passthrough(
|
||||
});
|
||||
let body_stream = build_direct_passthrough_inline_body_stream(
|
||||
finalizer,
|
||||
prefetched_body,
|
||||
response,
|
||||
upstream_started_at,
|
||||
stream_first_byte_timeout,
|
||||
@@ -2480,7 +2604,7 @@ async fn execute_stream_from_direct_passthrough(
|
||||
let mut last_client_chunk_elapsed_ms = 0u64;
|
||||
let mut downstream_dropped = false;
|
||||
let mut terminal_failure: Option<StreamFailureReport> = None;
|
||||
let mut upstream = direct_upstream_response_byte_stream(response);
|
||||
let mut upstream = direct_upstream_response_byte_stream(prefetched_body, response);
|
||||
let mut observed_first_upstream_body = false;
|
||||
let mut observed_first_client_send = false;
|
||||
|
||||
@@ -6333,8 +6457,8 @@ mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
use std::convert::Infallible;
|
||||
use std::sync::{
|
||||
atomic::{AtomicBool, Ordering},
|
||||
Arc,
|
||||
atomic::{AtomicBool, AtomicUsize, Ordering},
|
||||
Arc, Mutex,
|
||||
};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
@@ -6371,7 +6495,9 @@ mod tests {
|
||||
use axum::extract::ws::Message;
|
||||
use axum::extract::Request;
|
||||
use axum::routing::any;
|
||||
use axum::{http::header, http::HeaderValue, Router};
|
||||
use axum::{
|
||||
http::header, http::HeaderValue, http::StatusCode, response::IntoResponse, Json, Router,
|
||||
};
|
||||
use base64::Engine as _;
|
||||
use futures_util::StreamExt as _;
|
||||
use serde_json::{json, Value};
|
||||
@@ -6380,10 +6506,12 @@ mod tests {
|
||||
use super::{
|
||||
build_sse_body_stream, build_stream_sync_payload,
|
||||
client_format_allows_proxy_generated_sse_control_blocks,
|
||||
direct_upstream_response_byte_stream,
|
||||
ensure_stream_terminal_summary_for_missing_observed_finish,
|
||||
execute_execution_runtime_stream, execute_stream_from_frame_stream,
|
||||
maybe_apply_kiro_prompt_cache_usage_to_stream_summary, merge_stream_terminal_summary,
|
||||
parse_direct_passthrough_mode, prefetched_openai_responses_body_has_output_boundary,
|
||||
execute_execution_runtime_stream, execute_in_process_stream_with_oauth_retry,
|
||||
execute_stream_from_frame_stream, maybe_apply_kiro_prompt_cache_usage_to_stream_summary,
|
||||
merge_stream_terminal_summary, parse_direct_passthrough_mode,
|
||||
prefetch_direct_stream_error_body, prefetched_openai_responses_body_has_output_boundary,
|
||||
record_sync_terminal_usage_with_handoff,
|
||||
record_sync_terminal_usage_with_handoff_after_spawn, should_limit_direct_finalize_prefetch,
|
||||
should_probe_success_failover_before_stream, should_skip_direct_finalize_prefetch,
|
||||
@@ -6399,6 +6527,7 @@ mod tests {
|
||||
use crate::stage_metrics::RequestStageTrace;
|
||||
use crate::tunnel::{tunnel_protocol, TunnelProxyConn};
|
||||
use crate::AppState;
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
|
||||
fn provider_catalog_stop_429_for_plan(
|
||||
plan: &ExecutionPlan,
|
||||
@@ -6481,6 +6610,124 @@ mod tests {
|
||||
InMemoryProviderCatalogReadRepository::seed(vec![provider], vec![endpoint], vec![key])
|
||||
}
|
||||
|
||||
fn provider_catalog_for_stream_auth_plan(
|
||||
plan: &ExecutionPlan,
|
||||
provider_type: &str,
|
||||
auth_type: &str,
|
||||
auth_config: Option<Value>,
|
||||
) -> InMemoryProviderCatalogReadRepository {
|
||||
let provider = StoredProviderCatalogProvider::new(
|
||||
plan.provider_id.clone(),
|
||||
plan.provider_id.clone(),
|
||||
Some("https://provider.example".to_string()),
|
||||
provider_type.to_string(),
|
||||
)
|
||||
.expect("provider should build");
|
||||
let endpoint = StoredProviderCatalogEndpoint::new(
|
||||
plan.endpoint_id.clone(),
|
||||
plan.provider_id.clone(),
|
||||
plan.provider_api_format.clone(),
|
||||
None,
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
plan.url.clone(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build");
|
||||
let encrypted_auth_config = auth_config.map(|config| {
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, &config.to_string())
|
||||
.expect("auth config should encrypt")
|
||||
});
|
||||
let key = StoredProviderCatalogKey::new(
|
||||
plan.key_id.clone(),
|
||||
plan.provider_id.clone(),
|
||||
plan.key_id.clone(),
|
||||
auth_type.to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(json!([plan.provider_api_format.clone()])),
|
||||
None,
|
||||
encrypted_auth_config,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build");
|
||||
|
||||
InMemoryProviderCatalogReadRepository::seed(vec![provider], vec![endpoint], vec![key])
|
||||
}
|
||||
|
||||
fn direct_stream_test_plan(request_id: &str, url: String) -> ExecutionPlan {
|
||||
ExecutionPlan {
|
||||
request_id: request_id.to_string(),
|
||||
candidate_id: Some(format!("candidate-{request_id}")),
|
||||
provider_name: Some("codex".to_string()),
|
||||
provider_id: format!("provider-{request_id}"),
|
||||
endpoint_id: format!("endpoint-{request_id}"),
|
||||
key_id: format!("key-{request_id}"),
|
||||
method: "POST".to_string(),
|
||||
url,
|
||||
headers: BTreeMap::from([
|
||||
("content-type".to_string(), "application/json".to_string()),
|
||||
(
|
||||
"authorization".to_string(),
|
||||
"AgentAssertion stale-task".to_string(),
|
||||
),
|
||||
]),
|
||||
content_type: Some("application/json".to_string()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"stream": true})),
|
||||
stream: true,
|
||||
client_api_format: "openai:responses".to_string(),
|
||||
provider_api_format: "openai:responses".to_string(),
|
||||
model_name: Some("gpt-5".to_string()),
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: Some(ExecutionTimeouts {
|
||||
connect_ms: Some(30_000),
|
||||
first_byte_ms: Some(30_000),
|
||||
..ExecutionTimeouts::default()
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn agent_identity_test_auth_config(task_id: &str) -> Value {
|
||||
json!({
|
||||
"provider_type": "codex",
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_runtime_id": "runtime-test",
|
||||
"agent_private_key": "MC4CAQAwBQYDK2VwBCIEIAcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcH",
|
||||
"task_id": task_id
|
||||
})
|
||||
}
|
||||
|
||||
async fn collect_direct_execution_body(
|
||||
mut execution: crate::execution_runtime::DirectUpstreamStreamExecution,
|
||||
) -> Result<Vec<u8>, String> {
|
||||
let prefetched_body = std::mem::take(&mut execution.prefetched_body);
|
||||
let mut stream = direct_upstream_response_byte_stream(prefetched_body, execution.response);
|
||||
let mut body = Vec::new();
|
||||
while let Some(item) = stream.next().await {
|
||||
body.extend_from_slice(&item?);
|
||||
}
|
||||
Ok(body)
|
||||
}
|
||||
|
||||
fn codex_cyber_policy_plan(request_id: &str) -> ExecutionPlan {
|
||||
ExecutionPlan {
|
||||
request_id: request_id.to_string(),
|
||||
@@ -6609,6 +6856,326 @@ mod tests {
|
||||
AppState::new().expect("gateway state should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn agent_identity_stream_error_prefetch_is_bounded_and_replayed() {
|
||||
let upstream_body = format!(
|
||||
"{}{}",
|
||||
"x".repeat(crate::execution_runtime::MAX_ERROR_BODY_BYTES),
|
||||
"body-after-inspection-limit"
|
||||
);
|
||||
let expected_body = upstream_body.clone().into_bytes();
|
||||
let listener = crate::test_support::bind_loopback_listener()
|
||||
.await
|
||||
.expect("listener should bind");
|
||||
let addr = listener.local_addr().expect("address should resolve");
|
||||
let server = tokio::spawn(async move {
|
||||
let app = Router::new().route(
|
||||
"/responses",
|
||||
any(move || {
|
||||
let body = upstream_body.clone();
|
||||
async move { (StatusCode::UNAUTHORIZED, body) }
|
||||
}),
|
||||
);
|
||||
axum::serve(listener, app)
|
||||
.await
|
||||
.expect("server should start");
|
||||
});
|
||||
let plan =
|
||||
direct_stream_test_plan("bounded-agent-error", format!("http://{addr}/responses"));
|
||||
let mut execution = crate::execution_runtime::DirectSyncExecutionRuntime::new()
|
||||
.execute_stream(&plan)
|
||||
.await
|
||||
.expect("stream headers should execute");
|
||||
|
||||
let inspected = prefetch_direct_stream_error_body(&mut execution)
|
||||
.await
|
||||
.expect("error body should be inspected");
|
||||
|
||||
assert_eq!(
|
||||
inspected.len(),
|
||||
crate::execution_runtime::MAX_ERROR_BODY_BYTES
|
||||
);
|
||||
let replayed = collect_direct_execution_body(execution)
|
||||
.await
|
||||
.expect("prefetched response should replay");
|
||||
assert_eq!(replayed, expected_body);
|
||||
server.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn non_agent_stream_401_is_not_prefetched_and_body_passes_through() {
|
||||
let upstream_body = br#"{"error":{"code":"ordinary_unauthorized","message":"sign in"}}"#;
|
||||
let listener = crate::test_support::bind_loopback_listener()
|
||||
.await
|
||||
.expect("listener should bind");
|
||||
let addr = listener.local_addr().expect("address should resolve");
|
||||
let server = tokio::spawn(async move {
|
||||
let app = Router::new().route(
|
||||
"/responses",
|
||||
any(|| async {
|
||||
(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
[(header::CONTENT_TYPE, "application/json")],
|
||||
upstream_body.as_slice(),
|
||||
)
|
||||
}),
|
||||
);
|
||||
axum::serve(listener, app)
|
||||
.await
|
||||
.expect("server should start");
|
||||
});
|
||||
let mut plan = direct_stream_test_plan("non-agent-401", format!("http://{addr}/responses"));
|
||||
plan.provider_name = Some("openai".to_string());
|
||||
let repository = provider_catalog_for_stream_auth_plan(&plan, "openai", "api_key", None);
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_provider_transport_reader_for_tests(
|
||||
Arc::new(repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
|
||||
let execution = execute_in_process_stream_with_oauth_retry(
|
||||
&state,
|
||||
&mut plan,
|
||||
"trace-non-agent-401",
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("stream request should execute");
|
||||
|
||||
assert!(execution.prefetched_body.is_empty());
|
||||
let replayed = collect_direct_execution_body(execution)
|
||||
.await
|
||||
.expect("response body should pass through");
|
||||
assert_eq!(replayed, upstream_body);
|
||||
server.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn agent_identity_stream_non_task_401_replays_original_body_without_refresh() {
|
||||
let upstream_body =
|
||||
br#"{"error":{"code":"account_disabled","message":"account unavailable"}}"#;
|
||||
let listener = crate::test_support::bind_loopback_listener()
|
||||
.await
|
||||
.expect("listener should bind");
|
||||
let addr = listener.local_addr().expect("address should resolve");
|
||||
let task_registration_hits = Arc::new(AtomicUsize::new(0));
|
||||
let task_registration_hits_for_server = Arc::clone(&task_registration_hits);
|
||||
let server = tokio::spawn(async move {
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/responses",
|
||||
any(|| async {
|
||||
(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
[(header::CONTENT_TYPE, "application/json")],
|
||||
upstream_body.as_slice(),
|
||||
)
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/accounts/v1/agent/runtime-test/task/register",
|
||||
any(move || {
|
||||
let hits = Arc::clone(&task_registration_hits_for_server);
|
||||
async move {
|
||||
hits.fetch_add(1, Ordering::SeqCst);
|
||||
Json(json!({"task_id": "unexpected-task"}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
axum::serve(listener, app)
|
||||
.await
|
||||
.expect("server should start");
|
||||
});
|
||||
let mut plan =
|
||||
direct_stream_test_plan("agent-non-task-401", format!("http://{addr}/responses"));
|
||||
let repository = Arc::new(provider_catalog_for_stream_auth_plan(
|
||||
&plan,
|
||||
"codex",
|
||||
"oauth",
|
||||
Some(agent_identity_test_auth_config("task-old")),
|
||||
));
|
||||
let oauth_refresh =
|
||||
aether_provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
|
||||
Arc::new(
|
||||
aether_provider_transport::CodexAgentIdentityRefreshAdapter::default()
|
||||
.with_auth_api_base_url_for_tests(format!("http://{addr}/api/accounts")),
|
||||
),
|
||||
]);
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
repository,
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
)
|
||||
.with_oauth_refresh_coordinator_for_tests(oauth_refresh);
|
||||
|
||||
let execution = execute_in_process_stream_with_oauth_retry(
|
||||
&state,
|
||||
&mut plan,
|
||||
"trace-agent-non-task-401",
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("stream request should execute");
|
||||
|
||||
assert!(!execution.prefetched_body.is_empty());
|
||||
assert_eq!(task_registration_hits.load(Ordering::SeqCst), 0);
|
||||
let replayed = collect_direct_execution_body(execution)
|
||||
.await
|
||||
.expect("response body should replay");
|
||||
assert_eq!(replayed, upstream_body);
|
||||
server.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn agent_identity_stream_invalid_task_refreshes_and_retries_once() {
|
||||
let upstream_hits = Arc::new(AtomicUsize::new(0));
|
||||
let upstream_hits_for_server = Arc::clone(&upstream_hits);
|
||||
let task_registration_hits = Arc::new(AtomicUsize::new(0));
|
||||
let task_registration_hits_for_server = Arc::clone(&task_registration_hits);
|
||||
let observed_authorization = Arc::new(Mutex::new(Vec::<String>::new()));
|
||||
let observed_authorization_for_server = Arc::clone(&observed_authorization);
|
||||
let listener = crate::test_support::bind_loopback_listener()
|
||||
.await
|
||||
.expect("listener should bind");
|
||||
let addr = listener.local_addr().expect("address should resolve");
|
||||
let server = tokio::spawn(async move {
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/responses",
|
||||
any(move |request: Request| {
|
||||
let hits = Arc::clone(&upstream_hits_for_server);
|
||||
let authorizations = Arc::clone(&observed_authorization_for_server);
|
||||
async move {
|
||||
let authorization = request
|
||||
.headers()
|
||||
.get(header::AUTHORIZATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
authorizations
|
||||
.lock()
|
||||
.expect("authorization mutex should lock")
|
||||
.push(authorization);
|
||||
if hits.fetch_add(1, Ordering::SeqCst) == 0 {
|
||||
(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"code": "invalid_task_id",
|
||||
"message": "registered task expired"
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response()
|
||||
} else {
|
||||
(StatusCode::OK, Json(json!({"ok": true}))).into_response()
|
||||
}
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/accounts/v1/agent/runtime-test/task/register",
|
||||
any(move || {
|
||||
let hits = Arc::clone(&task_registration_hits_for_server);
|
||||
async move {
|
||||
hits.fetch_add(1, Ordering::SeqCst);
|
||||
Json(json!({"task_id": "task-new"}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
axum::serve(listener, app)
|
||||
.await
|
||||
.expect("server should start");
|
||||
});
|
||||
let mut plan =
|
||||
direct_stream_test_plan("agent-invalid-task", format!("http://{addr}/responses"));
|
||||
let repository = Arc::new(provider_catalog_for_stream_auth_plan(
|
||||
&plan,
|
||||
"codex",
|
||||
"oauth",
|
||||
Some(agent_identity_test_auth_config("task-old")),
|
||||
));
|
||||
let oauth_refresh =
|
||||
aether_provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
|
||||
Arc::new(
|
||||
aether_provider_transport::CodexAgentIdentityRefreshAdapter::default()
|
||||
.with_auth_api_base_url_for_tests(format!("http://{addr}/api/accounts")),
|
||||
),
|
||||
]);
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
repository,
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
)
|
||||
.with_oauth_refresh_coordinator_for_tests(oauth_refresh);
|
||||
let transport = state
|
||||
.read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
|
||||
.await
|
||||
.expect("transport should load")
|
||||
.expect("transport should exist");
|
||||
let initial_authorization = match state
|
||||
.resolve_local_oauth_request_auth(&transport)
|
||||
.await
|
||||
.expect("initial Agent Identity auth should resolve")
|
||||
.expect("initial Agent Identity auth should exist")
|
||||
{
|
||||
aether_provider_transport::LocalResolvedOAuthRequestAuth::Header { name, value } => {
|
||||
assert_eq!(name, "authorization");
|
||||
value
|
||||
}
|
||||
aether_provider_transport::LocalResolvedOAuthRequestAuth::Kiro(_) => {
|
||||
panic!("Agent Identity should resolve to header auth")
|
||||
}
|
||||
};
|
||||
assert!(
|
||||
aether_provider_transport::codex_agent_identity_authorization_matches_transport(
|
||||
&transport,
|
||||
&initial_authorization,
|
||||
)
|
||||
);
|
||||
plan.headers
|
||||
.insert("authorization".to_string(), initial_authorization.clone());
|
||||
|
||||
let execution = execute_in_process_stream_with_oauth_retry(
|
||||
&state,
|
||||
&mut plan,
|
||||
"trace-agent-invalid-task",
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("stream request should recover");
|
||||
|
||||
assert_eq!(execution.status_code, 200);
|
||||
assert!(execution.prefetched_body.is_empty());
|
||||
assert_eq!(upstream_hits.load(Ordering::SeqCst), 2);
|
||||
assert_eq!(task_registration_hits.load(Ordering::SeqCst), 1);
|
||||
let authorizations = observed_authorization
|
||||
.lock()
|
||||
.expect("authorization mutex should lock");
|
||||
assert_eq!(authorizations.len(), 2);
|
||||
assert_eq!(authorizations[0], initial_authorization);
|
||||
assert!(authorizations[1].starts_with("AgentAssertion "));
|
||||
assert_ne!(authorizations[1], authorizations[0]);
|
||||
drop(authorizations);
|
||||
let replayed = collect_direct_execution_body(execution)
|
||||
.await
|
||||
.expect("retried response body should read");
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<Value>(&replayed).expect("response should be JSON"),
|
||||
json!({"ok": true})
|
||||
);
|
||||
server.abort();
|
||||
}
|
||||
|
||||
struct BlockingStreamingRequestCandidateRepository {
|
||||
inner: InMemoryRequestCandidateRepository,
|
||||
block_streaming: AtomicBool,
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::collections::{BTreeMap, VecDeque};
|
||||
use std::future::Future;
|
||||
use std::io::Error as IoError;
|
||||
use std::time::{Duration, Instant};
|
||||
@@ -38,6 +38,7 @@ pub(crate) fn build_direct_execution_frame_stream(
|
||||
headers,
|
||||
provider_api_format,
|
||||
stream_summary_report_context,
|
||||
prefetched_body,
|
||||
response,
|
||||
started_at,
|
||||
stream_first_byte_timeout,
|
||||
@@ -70,7 +71,14 @@ pub(crate) fn build_direct_execution_frame_stream(
|
||||
|
||||
if should_buffer_non_stream_response(&headers, &observer_context) {
|
||||
let original_headers = headers.clone();
|
||||
match buffer_non_sse_upstream_body(response, started_at, stream_first_byte_timeout).await {
|
||||
match buffer_non_sse_upstream_body(
|
||||
prefetched_body,
|
||||
response,
|
||||
started_at,
|
||||
stream_first_byte_timeout,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(buffered) => {
|
||||
let mut response_headers = original_headers;
|
||||
let mut response_body = Bytes::from(buffered.body_bytes);
|
||||
@@ -192,6 +200,61 @@ pub(crate) fn build_direct_execution_frame_stream(
|
||||
let mut upstream_bytes = 0u64;
|
||||
let mut ttfb_ms = None;
|
||||
let mut first_chunk_telemetry_emitted = false;
|
||||
let mut prefetched_body_failed = false;
|
||||
for item in prefetched_body {
|
||||
match item {
|
||||
Ok(chunk) => {
|
||||
if ttfb_ms.is_none() {
|
||||
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
}
|
||||
if !first_chunk_telemetry_emitted {
|
||||
match encode_telemetry_frame(ttfb_ms, ttfb_ms, upstream_bytes) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(err) => {
|
||||
yield Err(err);
|
||||
return;
|
||||
}
|
||||
}
|
||||
first_chunk_telemetry_emitted = true;
|
||||
}
|
||||
upstream_bytes += chunk.len() as u64;
|
||||
observe_stream_chunk(
|
||||
&mut stream_terminal_observer,
|
||||
&normalized_observer_context,
|
||||
private_stream_normalizer.as_mut(),
|
||||
&mut observer_buffered,
|
||||
chunk.as_ref(),
|
||||
);
|
||||
match encode_data_frame(&chunk) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(err) => {
|
||||
yield Err(err);
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(message) => {
|
||||
warn!(
|
||||
event_name = "stream_pump_body_read_error",
|
||||
log_type = "ops",
|
||||
status_code,
|
||||
upstream_bytes,
|
||||
error = %message,
|
||||
"upstream body stream read error"
|
||||
);
|
||||
match encode_error_frame(status_code, message) {
|
||||
Ok(frame) => yield Ok(frame),
|
||||
Err(encode_err) => {
|
||||
yield Err(encode_err);
|
||||
return;
|
||||
}
|
||||
}
|
||||
prefetched_body_failed = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
if !prefetched_body_failed {
|
||||
match response {
|
||||
DirectUpstreamResponse::Reqwest(response) => {
|
||||
let mut bytes_stream = response.bytes_stream();
|
||||
@@ -516,6 +579,7 @@ pub(crate) fn build_direct_execution_frame_stream(
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let summary = finalize_stream_terminal_summary(
|
||||
&mut stream_terminal_observer,
|
||||
&normalized_observer_context,
|
||||
@@ -691,6 +755,7 @@ fn should_buffer_non_stream_response(
|
||||
}
|
||||
|
||||
async fn buffer_non_sse_upstream_body(
|
||||
mut prefetched_body: VecDeque<Result<Bytes, String>>,
|
||||
response: DirectUpstreamResponse,
|
||||
started_at: Instant,
|
||||
stream_first_byte_timeout: Option<Duration>,
|
||||
@@ -699,6 +764,26 @@ async fn buffer_non_sse_upstream_body(
|
||||
let mut upstream_bytes = 0u64;
|
||||
let mut ttfb_ms = None;
|
||||
|
||||
while let Some(item) = prefetched_body.pop_front() {
|
||||
match item {
|
||||
Ok(chunk) => {
|
||||
if ttfb_ms.is_none() {
|
||||
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
}
|
||||
upstream_bytes += chunk.len() as u64;
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
}
|
||||
Err(message) => {
|
||||
return Err(BufferedUpstreamBodyError {
|
||||
message,
|
||||
ttfb_ms,
|
||||
upstream_bytes,
|
||||
first_byte_timeout: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
match response {
|
||||
DirectUpstreamResponse::Reqwest(response) => {
|
||||
let mut bytes_stream = response.bytes_stream();
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::collections::{BTreeMap, HashMap, HashSet};
|
||||
use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
|
||||
use std::error::Error as _;
|
||||
use std::future::Future;
|
||||
use std::io::Read;
|
||||
@@ -607,6 +607,7 @@ pub(crate) struct DirectUpstreamStreamExecution {
|
||||
pub(crate) headers: BTreeMap<String, String>,
|
||||
pub(crate) provider_api_format: String,
|
||||
pub(crate) stream_summary_report_context: Value,
|
||||
pub(crate) prefetched_body: VecDeque<Result<Bytes, String>>,
|
||||
pub(crate) response: DirectUpstreamResponse,
|
||||
pub(crate) started_at: Instant,
|
||||
pub(crate) stream_first_byte_timeout: Option<Duration>,
|
||||
@@ -711,6 +712,7 @@ impl DirectSyncExecutionRuntime {
|
||||
headers,
|
||||
provider_api_format: plan.provider_api_format.clone(),
|
||||
stream_summary_report_context,
|
||||
prefetched_body: VecDeque::new(),
|
||||
response: response.into_direct_upstream_response(),
|
||||
started_at,
|
||||
stream_first_byte_timeout: resolve_stream_first_byte_timeout(plan),
|
||||
@@ -824,6 +826,7 @@ pub(crate) async fn execute_stream_plan_via_local_tunnel(
|
||||
headers,
|
||||
provider_api_format: plan.provider_api_format.clone(),
|
||||
stream_summary_report_context: build_stream_summary_report_context(plan),
|
||||
prefetched_body: VecDeque::new(),
|
||||
response: DirectUpstreamResponse::LocalTunnel(response),
|
||||
started_at,
|
||||
stream_first_byte_timeout: resolve_stream_first_byte_timeout(plan),
|
||||
|
||||
@@ -190,7 +190,10 @@ pub(crate) async fn build_admin_global_model_routing_payload(
|
||||
let payload = json!({
|
||||
"id": key.id,
|
||||
"name": key.name,
|
||||
"masked_key": state.masked_catalog_api_key(key),
|
||||
"masked_key": state.masked_catalog_api_key_for_provider(
|
||||
key,
|
||||
&provider.provider_type,
|
||||
),
|
||||
"is_active": key.is_active,
|
||||
"is_adaptive": is_adaptive,
|
||||
"effective_rpm": effective_rpm,
|
||||
@@ -252,9 +255,9 @@ pub(crate) async fn build_admin_global_model_routing_payload(
|
||||
.iter()
|
||||
.map(|provider| provider.id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let active_provider_name_by_id = active_providers
|
||||
let active_provider_metadata_by_id = active_providers
|
||||
.into_iter()
|
||||
.map(|provider| (provider.id, provider.name))
|
||||
.map(|provider| (provider.id, (provider.name, provider.provider_type)))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let active_keys = if active_provider_ids.is_empty() {
|
||||
Vec::new()
|
||||
@@ -273,14 +276,14 @@ pub(crate) async fn build_admin_global_model_routing_payload(
|
||||
if allowed_models.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let provider_name = active_provider_name_by_id
|
||||
let (provider_name, provider_type) = active_provider_metadata_by_id
|
||||
.get(&key.provider_id)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
all_keys_whitelist.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"masked_key": state.masked_catalog_api_key(&key),
|
||||
"masked_key": state.masked_catalog_api_key_for_provider(&key, &provider_type),
|
||||
"provider_id": key.provider_id,
|
||||
"provider_name": provider_name,
|
||||
"allowed_models": allowed_models,
|
||||
|
||||
+9
-2
@@ -149,8 +149,15 @@ pub(super) async fn build_admin_monitoring_cache_affinities_response(
|
||||
.map(|item| item.base_url.clone())
|
||||
.filter(|value| !value.trim().is_empty());
|
||||
let key_name = key.map(|item| item.name.clone());
|
||||
let key_prefix =
|
||||
key.and_then(|item| admin_monitoring_masked_provider_key_prefix(state, item));
|
||||
let key_prefix = key.and_then(|item| {
|
||||
admin_monitoring_masked_provider_key_prefix(
|
||||
state,
|
||||
item,
|
||||
provider
|
||||
.map(|provider| provider.provider_type.as_str())
|
||||
.unwrap_or(""),
|
||||
)
|
||||
});
|
||||
let user_id_text = user_id.clone();
|
||||
let username = user.map(|item| item.username.clone());
|
||||
let email = user.and_then(|item| item.email.clone());
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::provider_key_auth::provider_key_auth_config_uses_header_authorization;
|
||||
use crate::provider_key_auth::{
|
||||
provider_key_auth_config_is_agent_identity, provider_key_auth_config_uses_header_authorization,
|
||||
};
|
||||
use aether_crypto::decrypt_python_fernet_ciphertext;
|
||||
#[cfg(test)]
|
||||
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
|
||||
@@ -26,13 +28,15 @@ pub(super) fn admin_monitoring_masked_user_api_key_prefix(
|
||||
pub(super) fn admin_monitoring_masked_provider_key_prefix(
|
||||
state: &AdminAppState<'_>,
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
) -> Option<String> {
|
||||
match key.auth_type.trim() {
|
||||
"service_account" | "vertex_ai" => Some("[Service Account]".to_string()),
|
||||
"oauth" => {
|
||||
if provider_key_auth_config_uses_header_authorization(
|
||||
state.parse_catalog_auth_config_json(key).as_ref(),
|
||||
) {
|
||||
let auth_config = state.parse_catalog_auth_config_json(key);
|
||||
if provider_key_auth_config_is_agent_identity(provider_type, auth_config.as_ref()) {
|
||||
Some("[Agent Identity]".to_string())
|
||||
} else if provider_key_auth_config_uses_header_authorization(auth_config.as_ref()) {
|
||||
Some("[OAuth Header]".to_string())
|
||||
} else {
|
||||
Some("[OAuth Token]".to_string())
|
||||
@@ -115,3 +119,52 @@ pub(super) fn admin_monitoring_cache_affinity_sort_value(value: Option<&serde_js
|
||||
}
|
||||
0.0
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::admin_monitoring_masked_provider_key_prefix;
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::AppState;
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
|
||||
#[test]
|
||||
fn monitoring_labels_agent_identity_instead_of_oauth_token() {
|
||||
let app = AppState::new().expect("gateway should build");
|
||||
let state = AdminAppState::new(&app);
|
||||
let encrypted_placeholder =
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__")
|
||||
.expect("placeholder should encrypt");
|
||||
let encrypted_auth_config = encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"auth_mode":"agentIdentity","agent_runtime_id":"runtime-1","agent_private_key":"base64-private-key","task_id":"task-1"}"#,
|
||||
)
|
||||
.expect("auth config should encrypt");
|
||||
let key = StoredProviderCatalogKey::new(
|
||||
"key-agent".to_string(),
|
||||
"provider-codex".to_string(),
|
||||
"agent".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
None,
|
||||
encrypted_placeholder,
|
||||
Some(encrypted_auth_config),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("transport should build");
|
||||
|
||||
assert_eq!(
|
||||
admin_monitoring_masked_provider_key_prefix(&state, &key, "codex").as_deref(),
|
||||
Some("[Agent Identity]")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,9 +4,11 @@ use crate::handlers::admin::provider::shared::support::{
|
||||
};
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::handlers::admin::shared::{
|
||||
decrypt_catalog_secret_with_fallbacks, json_string_list, take_secret_prefix, take_secret_suffix,
|
||||
decrypt_catalog_secret_with_fallbacks, json_string_list, parse_catalog_auth_config_json,
|
||||
take_secret_prefix, take_secret_suffix,
|
||||
};
|
||||
use crate::handlers::public::matches_model_mapping_for_models;
|
||||
use crate::provider_key_auth::provider_key_auth_config_is_agent_identity;
|
||||
use crate::{GatewayError, LocalProviderDeleteTaskState};
|
||||
use aether_data_contracts::repository::global_models::{
|
||||
AdminGlobalModelListQuery, AdminProviderModelListQuery, PublicGlobalModelQuery,
|
||||
@@ -169,7 +171,12 @@ pub(crate) fn global_model_mapping_patterns_from_config(
|
||||
pub(crate) fn mapping_preview_masked_catalog_api_key(
|
||||
state: &AdminAppState<'_>,
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
) -> String {
|
||||
let auth_config = parse_catalog_auth_config_json(state.as_ref(), key);
|
||||
if provider_key_auth_config_is_agent_identity(provider_type, auth_config.as_ref()) {
|
||||
return "[Agent Identity]".to_string();
|
||||
}
|
||||
let ciphertext = key.encrypted_api_key.as_deref().unwrap_or("").trim();
|
||||
if ciphertext.is_empty() {
|
||||
return "***".to_string();
|
||||
@@ -345,7 +352,11 @@ pub(crate) async fn build_admin_provider_mapping_preview_payload(
|
||||
key_payloads.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"masked_key": mapping_preview_masked_catalog_api_key(state, &key),
|
||||
"masked_key": mapping_preview_masked_catalog_api_key(
|
||||
state,
|
||||
&key,
|
||||
&provider.provider_type,
|
||||
),
|
||||
"is_active": key.is_active,
|
||||
"allowed_models": allowed_models,
|
||||
"matching_global_models": matching_global_models,
|
||||
@@ -363,3 +374,52 @@ pub(crate) async fn build_admin_provider_mapping_preview_payload(
|
||||
"truncated_models": truncated_models,
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::mapping_preview_masked_catalog_api_key;
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::AppState;
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
|
||||
#[test]
|
||||
fn delete_preview_never_masks_internal_agent_identity_placeholder() {
|
||||
let app = AppState::new().expect("gateway should build");
|
||||
let state = AdminAppState::new(&app);
|
||||
let encrypted_placeholder =
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__")
|
||||
.expect("placeholder should encrypt");
|
||||
let encrypted_auth_config = encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"auth_mode":"agentIdentity","agent_runtime_id":"runtime-1","agent_private_key":"base64-private-key","task_id":"task-1"}"#,
|
||||
)
|
||||
.expect("auth config should encrypt");
|
||||
let key = StoredProviderCatalogKey::new(
|
||||
"key-agent".to_string(),
|
||||
"provider-codex".to_string(),
|
||||
"agent".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
None,
|
||||
encrypted_placeholder,
|
||||
Some(encrypted_auth_config),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("transport should build");
|
||||
|
||||
assert_eq!(
|
||||
mapping_preview_masked_catalog_api_key(&state, &key, "codex"),
|
||||
"[Agent Identity]"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
+197
-106
@@ -13,7 +13,10 @@ 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::duplicates::{
|
||||
acquire_codex_oauth_account_locks, find_duplicate_provider_oauth_key,
|
||||
release_codex_oauth_account_locks,
|
||||
};
|
||||
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,
|
||||
@@ -53,38 +56,92 @@ fn sanitize_windsurf_batch_import_error(error: &OAuthError) -> String {
|
||||
}
|
||||
}
|
||||
|
||||
fn copy_codex_agent_identity_field(
|
||||
const CODEX_AGENT_IDENTITY_SAFE_FIELDS: &[(&str, &[&str])] = &[
|
||||
("agent_runtime_id", &["agent_runtime_id", "agentRuntimeId"]),
|
||||
(
|
||||
"agent_private_key",
|
||||
&["agent_private_key", "agentPrivateKey"],
|
||||
),
|
||||
("task_id", &["task_id", "taskId"]),
|
||||
(
|
||||
"account_id",
|
||||
&[
|
||||
"account_id",
|
||||
"accountId",
|
||||
"chatgpt_account_id",
|
||||
"chatgptAccountId",
|
||||
],
|
||||
),
|
||||
(
|
||||
"account_user_id",
|
||||
&[
|
||||
"account_user_id",
|
||||
"accountUserId",
|
||||
"chatgpt_account_user_id",
|
||||
"chatgptAccountUserId",
|
||||
],
|
||||
),
|
||||
(
|
||||
"user_id",
|
||||
&["user_id", "userId", "chatgpt_user_id", "chatgptUserId"],
|
||||
),
|
||||
("email", &["email"]),
|
||||
(
|
||||
"plan_type",
|
||||
&[
|
||||
"plan_type",
|
||||
"planType",
|
||||
"chatgpt_plan_type",
|
||||
"chatgptPlanType",
|
||||
],
|
||||
),
|
||||
("account_name", &["account_name", "accountName"]),
|
||||
(
|
||||
"is_fedramp",
|
||||
&[
|
||||
"is_fedramp",
|
||||
"chatgpt_account_is_fedramp",
|
||||
"chatgptAccountIsFedramp",
|
||||
],
|
||||
),
|
||||
("workspace_id", &["workspace_id", "workspaceId"]),
|
||||
];
|
||||
|
||||
fn copy_codex_agent_identity_safe_fields(
|
||||
auth_config: &mut Map<String, Value>,
|
||||
nested: &Map<String, Value>,
|
||||
canonical_key: &str,
|
||||
aliases: &[&str],
|
||||
preferred: Option<&Map<String, Value>>,
|
||||
fallback: &Map<String, Value>,
|
||||
) {
|
||||
if auth_config.contains_key(canonical_key) {
|
||||
return;
|
||||
}
|
||||
if let Some(value) = aliases.iter().find_map(|key| nested.get(*key)).cloned() {
|
||||
auth_config.insert(canonical_key.to_string(), value);
|
||||
for (canonical_key, aliases) in CODEX_AGENT_IDENTITY_SAFE_FIELDS {
|
||||
if auth_config.contains_key(*canonical_key) {
|
||||
continue;
|
||||
}
|
||||
let value = preferred
|
||||
.and_then(|map| aliases.iter().find_map(|key| map.get(*key)))
|
||||
.or_else(|| aliases.iter().find_map(|key| fallback.get(*key)))
|
||||
.cloned();
|
||||
let Some(value) = value else {
|
||||
continue;
|
||||
};
|
||||
let type_is_allowed = if *canonical_key == "is_fedramp" {
|
||||
value.is_boolean()
|
||||
} else {
|
||||
value.as_str().is_some_and(|text| !text.trim().is_empty())
|
||||
};
|
||||
if !type_is_allowed {
|
||||
continue;
|
||||
}
|
||||
auth_config.insert((*canonical_key).to_string(), value);
|
||||
}
|
||||
}
|
||||
|
||||
fn remove_codex_agent_identity_oauth_tokens(auth_config: &mut Map<String, Value>) {
|
||||
for key in [
|
||||
"access_token",
|
||||
"accessToken",
|
||||
"refresh_token",
|
||||
"refreshToken",
|
||||
"id_token",
|
||||
"idToken",
|
||||
"expires_at",
|
||||
"expiresAt",
|
||||
"expires_in",
|
||||
"expiresIn",
|
||||
] {
|
||||
auth_config.remove(key);
|
||||
}
|
||||
fn sanitize_codex_agent_identity_nested_fields(nested: &Map<String, Value>) -> Map<String, Value> {
|
||||
let mut sanitized = Map::new();
|
||||
copy_codex_agent_identity_safe_fields(&mut sanitized, None, nested);
|
||||
sanitized
|
||||
}
|
||||
|
||||
fn codex_agent_identity_auth_config_from_import(
|
||||
pub(super) fn codex_agent_identity_auth_config_from_import(
|
||||
entry: &AdminProviderOAuthBatchImportEntry,
|
||||
) -> Result<Option<Map<String, Value>>, String> {
|
||||
let Some(raw_credentials) = entry.raw_credentials.as_ref() else {
|
||||
@@ -93,81 +150,24 @@ fn codex_agent_identity_auth_config_from_import(
|
||||
if !aether_provider_transport::is_codex_agent_identity_auth_config_value(raw_credentials) {
|
||||
return Ok(None);
|
||||
}
|
||||
let mut auth_config = raw_credentials
|
||||
let root = raw_credentials
|
||||
.as_object()
|
||||
.cloned()
|
||||
.ok_or_else(|| "Agent Identity 凭据必须是 JSON 对象".to_string())?;
|
||||
remove_codex_agent_identity_oauth_tokens(&mut auth_config);
|
||||
for nested_key in ["agent_identity", "agentIdentity"] {
|
||||
if let Some(nested) = auth_config
|
||||
.get_mut(nested_key)
|
||||
.and_then(Value::as_object_mut)
|
||||
{
|
||||
remove_codex_agent_identity_oauth_tokens(nested);
|
||||
}
|
||||
}
|
||||
let nested = auth_config
|
||||
let nested = root
|
||||
.get("agent_identity")
|
||||
.or_else(|| auth_config.get("agentIdentity"))
|
||||
.or_else(|| root.get("agentIdentity"))
|
||||
.and_then(Value::as_object)
|
||||
.cloned();
|
||||
let root = auth_config.clone();
|
||||
for (canonical_key, aliases) in [
|
||||
(
|
||||
"agent_runtime_id",
|
||||
&["agent_runtime_id", "agentRuntimeId"][..],
|
||||
),
|
||||
(
|
||||
"agent_private_key",
|
||||
&["agent_private_key", "agentPrivateKey"][..],
|
||||
),
|
||||
("task_id", &["task_id", "taskId"][..]),
|
||||
(
|
||||
"account_id",
|
||||
&[
|
||||
"account_id",
|
||||
"accountId",
|
||||
"chatgpt_account_id",
|
||||
"chatgptAccountId",
|
||||
][..],
|
||||
),
|
||||
(
|
||||
"account_user_id",
|
||||
&[
|
||||
"account_user_id",
|
||||
"accountUserId",
|
||||
"chatgpt_account_user_id",
|
||||
"chatgptAccountUserId",
|
||||
][..],
|
||||
),
|
||||
(
|
||||
"user_id",
|
||||
&["user_id", "userId", "chatgpt_user_id", "chatgptUserId"][..],
|
||||
),
|
||||
("email", &["email"][..]),
|
||||
(
|
||||
"plan_type",
|
||||
&[
|
||||
"plan_type",
|
||||
"planType",
|
||||
"chatgpt_plan_type",
|
||||
"chatgptPlanType",
|
||||
][..],
|
||||
),
|
||||
("account_name", &["account_name", "accountName"][..]),
|
||||
(
|
||||
"is_fedramp",
|
||||
&[
|
||||
"is_fedramp",
|
||||
"chatgpt_account_is_fedramp",
|
||||
"chatgptAccountIsFedramp",
|
||||
][..],
|
||||
),
|
||||
] {
|
||||
if let Some(nested) = nested.as_ref() {
|
||||
copy_codex_agent_identity_field(&mut auth_config, nested, canonical_key, aliases);
|
||||
let mut auth_config = Map::new();
|
||||
copy_codex_agent_identity_safe_fields(&mut auth_config, nested.as_ref(), root);
|
||||
if let Some(nested) = nested.as_ref() {
|
||||
let sanitized_nested = sanitize_codex_agent_identity_nested_fields(nested);
|
||||
if !sanitized_nested.is_empty() {
|
||||
auth_config.insert(
|
||||
"agent_identity".to_string(),
|
||||
Value::Object(sanitized_nested),
|
||||
);
|
||||
}
|
||||
copy_codex_agent_identity_field(&mut auth_config, &root, canonical_key, aliases);
|
||||
}
|
||||
auth_config.insert("provider_type".to_string(), json!("codex"));
|
||||
auth_config.insert("auth_mode".to_string(), json!("agentIdentity"));
|
||||
@@ -548,10 +548,49 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
|
||||
}
|
||||
}
|
||||
|
||||
let is_agent_identity = provider_type.eq_ignore_ascii_case("codex")
|
||||
&& aether_provider_transport::is_codex_agent_identity_auth_config_value(
|
||||
&Value::Object(auth_config.clone()),
|
||||
);
|
||||
let codex_oauth_account_leases =
|
||||
if provider_type.eq_ignore_ascii_case("codex") && !is_agent_identity {
|
||||
match acquire_codex_oauth_account_locks(
|
||||
state,
|
||||
provider_id,
|
||||
&auth_config,
|
||||
"batch-import",
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(leases) => leases,
|
||||
Err(error) => {
|
||||
failed += 1;
|
||||
results.push(json!({
|
||||
"index": index,
|
||||
"status": "error",
|
||||
"error": error.detail(),
|
||||
"replaced": false,
|
||||
}));
|
||||
maybe_report_admin_provider_oauth_batch_import_progress(
|
||||
&mut progress,
|
||||
entries.len(),
|
||||
success,
|
||||
failed,
|
||||
&results,
|
||||
)
|
||||
.await;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
let duplicate =
|
||||
match find_duplicate_provider_oauth_key(state, provider_id, &auth_config, None).await {
|
||||
Ok(value) => value,
|
||||
Err(detail) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
failed += 1;
|
||||
results.push(json!({
|
||||
"index": index,
|
||||
@@ -573,7 +612,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
|
||||
|
||||
let replaced = duplicate.is_some();
|
||||
let (persisted_key, key_name) = if let Some(existing_key) = duplicate {
|
||||
match update_existing_provider_oauth_catalog_key(
|
||||
let update_result = update_existing_provider_oauth_catalog_key(
|
||||
state,
|
||||
&existing_key,
|
||||
provider_type,
|
||||
@@ -583,10 +622,15 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
|
||||
key_proxy.clone(),
|
||||
expires_at,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Some(key) => (key, existing_key.name.clone()),
|
||||
None => {
|
||||
.await;
|
||||
match update_result {
|
||||
Err(error) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
return Err(error);
|
||||
}
|
||||
Ok(Some(key)) => (key, existing_key.name.clone()),
|
||||
Ok(None) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
failed += 1;
|
||||
results.push(json!({
|
||||
"index": index,
|
||||
@@ -611,7 +655,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
|
||||
&auth_config,
|
||||
Some(index),
|
||||
);
|
||||
match create_provider_oauth_catalog_key(
|
||||
let create_result = create_provider_oauth_catalog_key(
|
||||
state,
|
||||
provider_id,
|
||||
provider_type,
|
||||
@@ -622,10 +666,15 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
|
||||
key_proxy.clone(),
|
||||
expires_at,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Some(key) => (key, key_name),
|
||||
None => {
|
||||
.await;
|
||||
match create_result {
|
||||
Err(error) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
return Err(error);
|
||||
}
|
||||
Ok(Some(key)) => (key, key_name),
|
||||
Ok(None) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
failed += 1;
|
||||
results.push(json!({
|
||||
"index": index,
|
||||
@@ -645,6 +694,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
|
||||
}
|
||||
}
|
||||
};
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
|
||||
spawn_provider_oauth_account_state_refresh_after_update(
|
||||
state.cloned_app(),
|
||||
@@ -726,13 +776,24 @@ mod tests {
|
||||
"credentials":{
|
||||
"auth_mode":"agentIdentity",
|
||||
"id_token":"stale-id-token",
|
||||
"sessionToken":"stale-session-token",
|
||||
"apiKey":"stale-api-key",
|
||||
"cookie":"stale-cookie",
|
||||
"headers":{
|
||||
"authorization":"Bearer stale-bearer-token",
|
||||
"cookie":"stale-header-cookie",
|
||||
"x-api-key":"stale-header-api-key"
|
||||
},
|
||||
"profile":{"token":"stale-deep-token"},
|
||||
"agent_identity":{
|
||||
"agent_runtime_id":"runtime-1",
|
||||
"agent_private_key":"MC4CAQAwBQYDK2VwBCIEIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
|
||||
"accountId":"account-1",
|
||||
"chatgptUserId":"user-1",
|
||||
"chatgptAccountIsFedramp":true,
|
||||
"access_token":"stale-access-token"
|
||||
"access_token":"stale-access-token",
|
||||
"refreshToken":"stale-refresh-token",
|
||||
"headers":{"authorization":"Bearer stale-nested-bearer"}
|
||||
}
|
||||
}
|
||||
}]
|
||||
@@ -761,5 +822,35 @@ mod tests {
|
||||
.get("agent_identity")
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.is_some_and(|nested| !nested.contains_key("access_token")));
|
||||
let serialized = serde_json::to_string(&auth_config).expect("auth config should serialize");
|
||||
for secret in [
|
||||
"stale-id-token",
|
||||
"stale-session-token",
|
||||
"stale-api-key",
|
||||
"stale-cookie",
|
||||
"stale-bearer-token",
|
||||
"stale-header-cookie",
|
||||
"stale-header-api-key",
|
||||
"stale-deep-token",
|
||||
"stale-access-token",
|
||||
"stale-refresh-token",
|
||||
"stale-nested-bearer",
|
||||
] {
|
||||
assert!(!serialized.contains(secret), "secret leaked: {secret}");
|
||||
}
|
||||
for forbidden_key in [
|
||||
"access_token",
|
||||
"refreshToken",
|
||||
"sessionToken",
|
||||
"apiKey",
|
||||
"cookie",
|
||||
"headers",
|
||||
"profile",
|
||||
] {
|
||||
assert!(
|
||||
!serialized.contains(forbidden_key),
|
||||
"forbidden key persisted: {forbidden_key}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,4 +6,7 @@ mod progress;
|
||||
mod task;
|
||||
|
||||
pub(super) use orchestration::handle_admin_provider_oauth_batch_import;
|
||||
pub(super) use task::handle_admin_provider_oauth_start_batch_import_task;
|
||||
pub(super) use task::{
|
||||
handle_admin_provider_oauth_start_agent_identity_import_task,
|
||||
handle_admin_provider_oauth_start_batch_import_task,
|
||||
};
|
||||
|
||||
@@ -3,6 +3,7 @@ use super::execution::{
|
||||
execute_admin_provider_oauth_batch_import_for_provider_type,
|
||||
};
|
||||
use super::parse::{
|
||||
admin_provider_oauth_batch_contains_agent_identity,
|
||||
build_admin_provider_oauth_batch_import_response,
|
||||
parse_admin_provider_oauth_batch_import_request, AdminProviderOAuthBatchImportRequest,
|
||||
};
|
||||
@@ -41,6 +42,12 @@ pub(in super::super) async fn handle_admin_provider_oauth_batch_import(
|
||||
Ok(payload) => payload,
|
||||
Err(response) => return Ok(response),
|
||||
};
|
||||
if admin_provider_oauth_batch_contains_agent_identity(&payload.credentials) {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"Agent Identity JSON 必须使用专属导入接口",
|
||||
));
|
||||
}
|
||||
|
||||
let Some(provider) = state
|
||||
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
|
||||
|
||||
@@ -78,6 +78,36 @@ pub(super) fn parse_admin_provider_oauth_batch_import_request(
|
||||
}
|
||||
}
|
||||
|
||||
fn json_value_contains_agent_identity(value: &serde_json::Value) -> bool {
|
||||
if value.as_object().is_some_and(|_| {
|
||||
aether_provider_transport::is_codex_agent_identity_auth_config_value(value)
|
||||
}) {
|
||||
return true;
|
||||
}
|
||||
match value {
|
||||
serde_json::Value::Array(items) => items.iter().any(json_value_contains_agent_identity),
|
||||
serde_json::Value::Object(object) => {
|
||||
object.values().any(json_value_contains_agent_identity)
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn admin_provider_oauth_batch_contains_agent_identity(raw_credentials: &str) -> bool {
|
||||
let raw = raw_credentials.trim();
|
||||
if raw.is_empty() {
|
||||
return false;
|
||||
}
|
||||
if let Ok(value) = serde_json::from_str::<serde_json::Value>(raw) {
|
||||
return json_value_contains_agent_identity(&value);
|
||||
}
|
||||
raw.lines()
|
||||
.map(str::trim)
|
||||
.filter(|line| !line.is_empty() && !line.starts_with('#'))
|
||||
.filter_map(|line| serde_json::from_str::<serde_json::Value>(line).ok())
|
||||
.any(|value| json_value_contains_agent_identity(&value))
|
||||
}
|
||||
|
||||
fn coerce_admin_provider_oauth_import_str(value: Option<&serde_json::Value>) -> Option<String> {
|
||||
value
|
||||
.and_then(serde_json::Value::as_str)
|
||||
@@ -623,6 +653,47 @@ pub(super) fn parse_admin_provider_oauth_batch_import_entries(
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(super) fn parse_admin_provider_oauth_agent_identity_import_entries(
|
||||
raw_credentials: &str,
|
||||
) -> Result<Vec<AdminProviderOAuthBatchImportEntry>, String> {
|
||||
let raw = raw_credentials.trim();
|
||||
if raw.is_empty() {
|
||||
return Err("Agent Identity 凭据不能为空".to_string());
|
||||
}
|
||||
let value = serde_json::from_str::<serde_json::Value>(raw)
|
||||
.map_err(|error| format!("Agent Identity JSON 解析失败: {error}"))?;
|
||||
let entries = match &value {
|
||||
serde_json::Value::Array(items) => items
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, item)| {
|
||||
extract_admin_provider_oauth_batch_import_entry("codex", item).unwrap_or_else(
|
||||
|| {
|
||||
parse_error_entry(format!(
|
||||
"第 {} 个条目没有可导入的 Agent Identity 凭据",
|
||||
index + 1
|
||||
))
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect(),
|
||||
serde_json::Value::Object(object) => parse_sub2api_export_accounts("codex", object)
|
||||
.unwrap_or_else(|| {
|
||||
vec![
|
||||
extract_admin_provider_oauth_batch_import_entry("codex", &value)
|
||||
.unwrap_or_else(|| {
|
||||
parse_error_entry("没有可导入的 Agent Identity 凭据".to_string())
|
||||
}),
|
||||
]
|
||||
}),
|
||||
_ => return Err("Agent Identity 凭据必须是 JSON 对象、数组或 sub2api 导出".to_string()),
|
||||
};
|
||||
if entries.is_empty() {
|
||||
return Err("Agent Identity 凭据不能为空".to_string());
|
||||
}
|
||||
Ok(entries)
|
||||
}
|
||||
|
||||
fn parse_error_entry(error: String) -> AdminProviderOAuthBatchImportEntry {
|
||||
AdminProviderOAuthBatchImportEntry {
|
||||
parse_error: Some(error),
|
||||
@@ -811,6 +882,7 @@ pub(super) fn build_admin_provider_oauth_batch_task_state(
|
||||
task_id: &str,
|
||||
provider_id: &str,
|
||||
provider_type: &str,
|
||||
import_kind: &str,
|
||||
status: &str,
|
||||
total: usize,
|
||||
processed: usize,
|
||||
@@ -839,6 +911,7 @@ pub(super) fn build_admin_provider_oauth_batch_task_state(
|
||||
"task_id": task_id,
|
||||
"provider_id": provider_id,
|
||||
"provider_type": provider_type,
|
||||
"import_kind": import_kind,
|
||||
"status": status,
|
||||
"total": total,
|
||||
"processed": processed,
|
||||
@@ -860,6 +933,7 @@ pub(super) fn build_admin_provider_oauth_batch_task_state(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
admin_provider_oauth_batch_contains_agent_identity,
|
||||
apply_admin_provider_oauth_batch_import_hints,
|
||||
parse_admin_provider_oauth_batch_import_entries,
|
||||
};
|
||||
@@ -874,6 +948,37 @@ mod tests {
|
||||
format!("{}.{}.signature", encode(header), encode(payload))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ordinary_batch_guard_detects_agent_identity_in_all_json_shapes() {
|
||||
let single = json!({
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_runtime_id": "runtime-1",
|
||||
"agent_private_key": "private-key"
|
||||
});
|
||||
assert!(admin_provider_oauth_batch_contains_agent_identity(
|
||||
&single.to_string()
|
||||
));
|
||||
assert!(admin_provider_oauth_batch_contains_agent_identity(
|
||||
&json!([{"refresh_token":"ordinary"}, single.clone()]).to_string()
|
||||
));
|
||||
assert!(admin_provider_oauth_batch_contains_agent_identity(
|
||||
&format!("ordinary-token\n{}", single)
|
||||
));
|
||||
assert!(admin_provider_oauth_batch_contains_agent_identity(
|
||||
&json!({
|
||||
"type": "sub2api-data",
|
||||
"accounts": [{"credentials": single.clone()}]
|
||||
})
|
||||
.to_string()
|
||||
));
|
||||
assert!(!admin_provider_oauth_batch_contains_agent_identity(
|
||||
&json!([{"refresh_token":"ordinary"}]).to_string()
|
||||
));
|
||||
assert!(!admin_provider_oauth_batch_contains_agent_identity(
|
||||
"ordinary-token"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_access_token_only_entry() {
|
||||
let entries = parse_admin_provider_oauth_batch_import_entries(
|
||||
|
||||
@@ -1,19 +1,26 @@
|
||||
use super::execution::{
|
||||
estimate_admin_provider_oauth_batch_import_total,
|
||||
codex_agent_identity_auth_config_from_import, estimate_admin_provider_oauth_batch_import_total,
|
||||
execute_admin_provider_oauth_batch_import_for_provider_type,
|
||||
};
|
||||
use super::parse::{
|
||||
build_admin_provider_oauth_batch_task_state, parse_admin_provider_oauth_batch_import_request,
|
||||
admin_provider_oauth_batch_contains_agent_identity,
|
||||
build_admin_provider_oauth_batch_task_state,
|
||||
parse_admin_provider_oauth_agent_identity_import_entries,
|
||||
parse_admin_provider_oauth_batch_import_request,
|
||||
};
|
||||
use super::progress::{
|
||||
AdminProviderOAuthBatchImportProgress, AdminProviderOAuthBatchProgressReporter,
|
||||
};
|
||||
use crate::handlers::admin::provider::oauth::duplicates::codex_agent_identity_account_lock_keys;
|
||||
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,
|
||||
is_fixed_provider_type_for_provider_oauth,
|
||||
};
|
||||
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_batch_import_task_provider_id;
|
||||
use crate::handlers::admin::provider::shared::paths::{
|
||||
admin_provider_oauth_agent_identity_import_task_provider_id,
|
||||
admin_provider_oauth_batch_import_task_provider_id,
|
||||
};
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::task_runtime::{
|
||||
append_event_with_logging, now_unix_secs, task_definition, update_run_status,
|
||||
@@ -23,24 +30,177 @@ use crate::GatewayError;
|
||||
use aether_data_contracts::repository::background_tasks::{
|
||||
BackgroundTaskKind, BackgroundTaskStatus, UpsertBackgroundTaskRun,
|
||||
};
|
||||
use aether_runtime_state::RuntimeLockLease;
|
||||
use axum::{
|
||||
body::Bytes,
|
||||
http,
|
||||
response::{IntoResponse, Response},
|
||||
Json,
|
||||
};
|
||||
use serde_json::json;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use serde_json::{json, Map, Value};
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
use tokio::task;
|
||||
use uuid::Uuid;
|
||||
|
||||
const PROVIDER_OAUTH_BATCH_TASK_MAX_ERROR_SAMPLES: usize = 20;
|
||||
const PROVIDER_OAUTH_BATCH_IMPORT_KIND: &str = "oauth_batch";
|
||||
const PROVIDER_AGENT_IDENTITY_IMPORT_KIND: &str = "agent_identity";
|
||||
const PROVIDER_AGENT_IDENTITY_IMPORT_LOCK_TTL: Duration = Duration::from_secs(180);
|
||||
const PROVIDER_AGENT_IDENTITY_IMPORT_LOCK_RENEW_INTERVAL: Duration = Duration::from_secs(60);
|
||||
|
||||
fn provider_oauth_import_kind(agent_identity_only: bool) -> &'static str {
|
||||
if agent_identity_only {
|
||||
PROVIDER_AGENT_IDENTITY_IMPORT_KIND
|
||||
} else {
|
||||
PROVIDER_OAUTH_BATCH_IMPORT_KIND
|
||||
}
|
||||
}
|
||||
|
||||
fn codex_agent_identity_import_auth_configs(
|
||||
credentials: &str,
|
||||
) -> Result<Vec<Map<String, Value>>, String> {
|
||||
let entries = parse_admin_provider_oauth_agent_identity_import_entries(credentials)?;
|
||||
entries
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, entry)| {
|
||||
if let Some(error) = entry.parse_error.as_deref() {
|
||||
return Err(format!("第 {} 个条目无效: {error}", index + 1));
|
||||
}
|
||||
match codex_agent_identity_auth_config_from_import(entry) {
|
||||
Ok(Some(auth_config)) => Ok(auth_config),
|
||||
Ok(None) => Err(format!("第 {} 个条目不是 Agent Identity", index + 1)),
|
||||
Err(error) => Err(format!("第 {} 个条目无效: {error}", index + 1)),
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn provider_agent_identity_import_lock_key(provider_id: &str, agent_runtime_id: &str) -> String {
|
||||
let digest = Sha256::digest(format!("{provider_id}\0{agent_runtime_id}").as_bytes());
|
||||
format!("provider_oauth_agent_identity_import:{digest:x}")
|
||||
}
|
||||
|
||||
fn provider_agent_identity_import_lock_keys(
|
||||
provider_id: &str,
|
||||
auth_configs: &[Map<String, Value>],
|
||||
) -> Vec<String> {
|
||||
let mut lock_keys = Vec::new();
|
||||
for auth_config in auth_configs {
|
||||
if let Some(agent_runtime_id) = auth_config
|
||||
.get("agent_runtime_id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
lock_keys.push(provider_agent_identity_import_lock_key(
|
||||
provider_id,
|
||||
agent_runtime_id,
|
||||
));
|
||||
}
|
||||
lock_keys.extend(codex_agent_identity_account_lock_keys(
|
||||
provider_id,
|
||||
auth_config,
|
||||
));
|
||||
}
|
||||
lock_keys.sort_unstable();
|
||||
lock_keys.dedup();
|
||||
lock_keys
|
||||
}
|
||||
|
||||
async fn acquire_provider_agent_identity_import_locks(
|
||||
state: &AdminAppState<'_>,
|
||||
provider_id: &str,
|
||||
lock_keys: &[String],
|
||||
task_id: &str,
|
||||
) -> Result<Vec<RuntimeLockLease>, Response> {
|
||||
let owner = format!("aether-gateway-agent-identity-import-{task_id}");
|
||||
let mut leases = Vec::with_capacity(lock_keys.len());
|
||||
for lock_key in lock_keys {
|
||||
match state
|
||||
.runtime_state()
|
||||
.lock_try_acquire(
|
||||
lock_key.as_str(),
|
||||
owner.as_str(),
|
||||
PROVIDER_AGENT_IDENTITY_IMPORT_LOCK_TTL,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(lease)) => leases.push(lease),
|
||||
Ok(None) => {
|
||||
release_provider_agent_identity_import_locks(state, leases).await;
|
||||
return Err(build_internal_control_error_response(
|
||||
http::StatusCode::CONFLICT,
|
||||
"其中一个 Agent Identity 正在导入或创建,请稍后重试",
|
||||
));
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
provider_id = %provider_id,
|
||||
lock_key = %lock_key,
|
||||
error = ?error,
|
||||
"gateway Agent Identity import lock unavailable"
|
||||
);
|
||||
release_provider_agent_identity_import_locks(state, leases).await;
|
||||
return Err(build_internal_control_error_response(
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Agent Identity 导入锁暂不可用,请稍后重试",
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(leases)
|
||||
}
|
||||
|
||||
async fn release_provider_agent_identity_import_locks(
|
||||
state: &AdminAppState<'_>,
|
||||
leases: Vec<RuntimeLockLease>,
|
||||
) {
|
||||
for lease in leases {
|
||||
match state.runtime_state().lock_release(&lease).await {
|
||||
Ok(true) => {}
|
||||
Ok(false) => tracing::warn!(
|
||||
lock_key = %lease.key,
|
||||
"gateway Agent Identity import lock was not owned during release"
|
||||
),
|
||||
Err(error) => tracing::warn!(
|
||||
lock_key = %lease.key,
|
||||
error = ?error,
|
||||
"gateway Agent Identity import lock release failed"
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn renew_provider_agent_identity_import_locks(
|
||||
state: &AdminAppState<'_>,
|
||||
leases: &[RuntimeLockLease],
|
||||
ttl: Duration,
|
||||
) -> Result<(), String> {
|
||||
for lease in leases {
|
||||
match state.runtime_state().lock_renew(lease, ttl).await {
|
||||
Ok(true) => {}
|
||||
Ok(false) => {
|
||||
return Err(format!("Agent Identity 导入锁已失效: {}", lease.key));
|
||||
}
|
||||
Err(error) => {
|
||||
return Err(format!(
|
||||
"Agent Identity 导入锁续租失败 ({}): {error:?}",
|
||||
lease.key
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
struct BatchTaskProgressReporter {
|
||||
app: crate::AppState,
|
||||
task_id: String,
|
||||
provider_id: String,
|
||||
provider_type: String,
|
||||
import_kind: &'static str,
|
||||
created_at: u64,
|
||||
started_at: u64,
|
||||
error_samples: Vec<serde_json::Value>,
|
||||
@@ -65,6 +225,7 @@ impl AdminProviderOAuthBatchProgressReporter for BatchTaskProgressReporter {
|
||||
&self.task_id,
|
||||
&self.provider_id,
|
||||
&self.provider_type,
|
||||
self.import_kind,
|
||||
"processing",
|
||||
progress.total,
|
||||
progress.processed,
|
||||
@@ -89,13 +250,33 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
request_body: Option<&Bytes>,
|
||||
) -> Result<Response, GatewayError> {
|
||||
handle_admin_provider_oauth_start_import_task(state, request_context, request_body, false).await
|
||||
}
|
||||
|
||||
pub(in super::super) async fn handle_admin_provider_oauth_start_agent_identity_import_task(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
request_body: Option<&Bytes>,
|
||||
) -> Result<Response, GatewayError> {
|
||||
handle_admin_provider_oauth_start_import_task(state, request_context, request_body, true).await
|
||||
}
|
||||
|
||||
async fn handle_admin_provider_oauth_start_import_task(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
request_body: Option<&Bytes>,
|
||||
agent_identity_only: bool,
|
||||
) -> Result<Response, GatewayError> {
|
||||
if !state.has_provider_catalog_data_reader() {
|
||||
return Ok(build_admin_provider_oauth_backend_unavailable_response());
|
||||
}
|
||||
let Some(provider_id) =
|
||||
let provider_id = if agent_identity_only {
|
||||
admin_provider_oauth_agent_identity_import_task_provider_id(request_context.path())
|
||||
} else {
|
||||
admin_provider_oauth_batch_import_task_provider_id(request_context.path())
|
||||
else {
|
||||
};
|
||||
let Some(provider_id) = provider_id else {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
"Provider 不存在",
|
||||
@@ -105,6 +286,14 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
Ok(payload) => payload,
|
||||
Err(response) => return Ok(response),
|
||||
};
|
||||
if !agent_identity_only
|
||||
&& admin_provider_oauth_batch_contains_agent_identity(&payload.credentials)
|
||||
{
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"Agent Identity JSON 必须使用专属导入接口",
|
||||
));
|
||||
}
|
||||
|
||||
let Some(provider) = state
|
||||
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
|
||||
@@ -118,6 +307,12 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
));
|
||||
};
|
||||
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
|
||||
if agent_identity_only && provider_type != "codex" {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"仅 Codex Provider 支持导入 Agent Identity",
|
||||
));
|
||||
}
|
||||
if !is_fixed_provider_type_for_provider_oauth(&provider_type) {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
@@ -131,9 +326,27 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
return Ok(build_admin_provider_oauth_backend_unavailable_response());
|
||||
}
|
||||
|
||||
let total = estimate_admin_provider_oauth_batch_import_total(
|
||||
&provider_type,
|
||||
payload.credentials.as_str(),
|
||||
let agent_identity_auth_configs = if agent_identity_only {
|
||||
match codex_agent_identity_import_auth_configs(&payload.credentials) {
|
||||
Ok(auth_configs) => Some(auth_configs),
|
||||
Err(detail) => {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
format!("该接口仅接受有效的 Agent Identity JSON: {detail}"),
|
||||
));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let total = agent_identity_auth_configs.as_ref().map_or_else(
|
||||
|| {
|
||||
estimate_admin_provider_oauth_batch_import_total(
|
||||
&provider_type,
|
||||
payload.credentials.as_str(),
|
||||
)
|
||||
},
|
||||
Vec::len,
|
||||
);
|
||||
if total == 0 {
|
||||
return Ok(build_internal_control_error_response(
|
||||
@@ -142,12 +355,35 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
));
|
||||
}
|
||||
|
||||
let task_id = Uuid::new_v4().to_string();
|
||||
let task_id = if agent_identity_only {
|
||||
format!("agent-identity-{}", Uuid::new_v4())
|
||||
} else {
|
||||
Uuid::new_v4().to_string()
|
||||
};
|
||||
let mut agent_identity_import_leases =
|
||||
if let Some(auth_configs) = agent_identity_auth_configs.as_deref() {
|
||||
let lock_keys = provider_agent_identity_import_lock_keys(&provider_id, auth_configs);
|
||||
match acquire_provider_agent_identity_import_locks(
|
||||
state,
|
||||
&provider_id,
|
||||
&lock_keys,
|
||||
&task_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(leases) => leases,
|
||||
Err(response) => return Ok(response),
|
||||
}
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
let import_kind = provider_oauth_import_kind(agent_identity_only);
|
||||
let created_at = now_unix_secs();
|
||||
let submitted_state = build_admin_provider_oauth_batch_task_state(
|
||||
&task_id,
|
||||
&provider_id,
|
||||
&provider_type,
|
||||
import_kind,
|
||||
"submitted",
|
||||
total,
|
||||
0,
|
||||
@@ -167,6 +403,11 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
release_provider_agent_identity_import_locks(
|
||||
state,
|
||||
std::mem::take(&mut agent_identity_import_leases),
|
||||
)
|
||||
.await;
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
"provider oauth batch task redis unavailable",
|
||||
@@ -191,6 +432,7 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
payload_json: Some(json!({
|
||||
"provider_id": provider_id.clone(),
|
||||
"provider_type": provider_type.clone(),
|
||||
"import_kind": import_kind,
|
||||
"total": total,
|
||||
})),
|
||||
result_json: None,
|
||||
@@ -211,6 +453,7 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
Some(json!({
|
||||
"provider_id": provider_id.clone(),
|
||||
"provider_type": provider_type.clone(),
|
||||
"import_kind": import_kind,
|
||||
"total": total,
|
||||
})),
|
||||
)
|
||||
@@ -223,6 +466,7 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
let provider_type_for_worker = provider_type.clone();
|
||||
let proxy_node_id = payload.proxy_node_id.clone();
|
||||
let raw_credentials = payload.credentials.clone();
|
||||
let agent_identity_import_leases_for_worker = std::mem::take(&mut agent_identity_import_leases);
|
||||
task::spawn(async move {
|
||||
let started_at = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
@@ -233,6 +477,7 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
&task_id_for_worker,
|
||||
&provider_id_for_worker,
|
||||
&provider_type_for_worker,
|
||||
import_kind,
|
||||
"processing",
|
||||
total,
|
||||
0,
|
||||
@@ -277,20 +522,55 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
task_id: task_id_for_worker.clone(),
|
||||
provider_id: provider_id_for_worker.clone(),
|
||||
provider_type: provider_type_for_worker.clone(),
|
||||
import_kind,
|
||||
created_at,
|
||||
started_at,
|
||||
error_samples: Vec::new(),
|
||||
};
|
||||
match execute_admin_provider_oauth_batch_import_for_provider_type(
|
||||
&AdminAppState::new(&task_state),
|
||||
let task_admin_state = AdminAppState::new(&task_state);
|
||||
let execution = execute_admin_provider_oauth_batch_import_for_provider_type(
|
||||
&task_admin_state,
|
||||
&provider_id_for_worker,
|
||||
&provider_type_for_worker,
|
||||
raw_credentials.as_str(),
|
||||
proxy_node_id.as_deref(),
|
||||
Some(&mut progress_reporter),
|
||||
);
|
||||
tokio::pin!(execution);
|
||||
let execution_result = if agent_identity_import_leases_for_worker.is_empty() {
|
||||
execution.await
|
||||
} else {
|
||||
let mut renew_timer =
|
||||
tokio::time::interval(PROVIDER_AGENT_IDENTITY_IMPORT_LOCK_RENEW_INTERVAL);
|
||||
renew_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
|
||||
renew_timer.tick().await;
|
||||
loop {
|
||||
tokio::select! {
|
||||
result = &mut execution => break result,
|
||||
_ = renew_timer.tick() => {
|
||||
if let Err(detail) = renew_provider_agent_identity_import_locks(
|
||||
&task_admin_state,
|
||||
&agent_identity_import_leases_for_worker,
|
||||
PROVIDER_AGENT_IDENTITY_IMPORT_LOCK_TTL,
|
||||
).await {
|
||||
tracing::error!(
|
||||
provider_id = %provider_id_for_worker,
|
||||
task_id = %task_id_for_worker,
|
||||
detail = %detail,
|
||||
"gateway Agent Identity import lease lost"
|
||||
);
|
||||
break Err(GatewayError::Internal(detail));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
release_provider_agent_identity_import_locks(
|
||||
&task_admin_state,
|
||||
agent_identity_import_leases_for_worker,
|
||||
)
|
||||
.await
|
||||
{
|
||||
.await;
|
||||
match execution_result {
|
||||
Ok(outcome) => {
|
||||
let finished_at = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
@@ -324,6 +604,7 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
&task_id_for_worker,
|
||||
&provider_id_for_worker,
|
||||
&provider_type_for_worker,
|
||||
import_kind,
|
||||
"completed",
|
||||
outcome.total,
|
||||
outcome.total,
|
||||
@@ -350,6 +631,7 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
Some(json!({
|
||||
"provider_id": provider_id_for_worker,
|
||||
"provider_type": provider_type_for_worker,
|
||||
"import_kind": import_kind,
|
||||
"total": outcome.total,
|
||||
"success": outcome.success,
|
||||
"failed": outcome.failed,
|
||||
@@ -381,6 +663,7 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
&task_id_for_worker,
|
||||
&provider_id_for_worker,
|
||||
&provider_type_for_worker,
|
||||
import_kind,
|
||||
"failed",
|
||||
total,
|
||||
0,
|
||||
@@ -432,6 +715,7 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
&task_id,
|
||||
&provider_id,
|
||||
&provider_type,
|
||||
import_kind,
|
||||
"submitted",
|
||||
total,
|
||||
0,
|
||||
@@ -448,3 +732,251 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
||||
);
|
||||
Ok(Json(submitted_response).into_response())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
acquire_provider_agent_identity_import_locks, codex_agent_identity_import_auth_configs,
|
||||
provider_agent_identity_import_lock_key, provider_agent_identity_import_lock_keys,
|
||||
release_provider_agent_identity_import_locks, renew_provider_agent_identity_import_locks,
|
||||
};
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::AppState;
|
||||
use serde_json::json;
|
||||
|
||||
fn agent_identity(runtime_id: &str) -> serde_json::Value {
|
||||
json!({
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_runtime_id": runtime_id,
|
||||
"agent_private_key": "MC4CAQAwBQYDK2VwBCIEIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
|
||||
"task_id": format!("task-{runtime_id}"),
|
||||
"account_id": "account-1",
|
||||
"user_id": "user-1",
|
||||
"email": "agent@example.com"
|
||||
})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dedicated_import_accepts_root_array_and_sub2api_agent_identities() {
|
||||
let single = agent_identity("runtime-1");
|
||||
assert_eq!(
|
||||
codex_agent_identity_import_auth_configs(&single.to_string())
|
||||
.expect("single Agent Identity should parse")
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
|
||||
let array = json!([agent_identity("runtime-1"), agent_identity("runtime-2")]);
|
||||
assert_eq!(
|
||||
codex_agent_identity_import_auth_configs(&array.to_string())
|
||||
.expect("Agent Identity array should parse")
|
||||
.len(),
|
||||
2
|
||||
);
|
||||
|
||||
let sub2api = json!({
|
||||
"type": "sub2api-data",
|
||||
"accounts": [
|
||||
{
|
||||
"name": "agent-1@example.com",
|
||||
"platform": "openai",
|
||||
"credentials": agent_identity("runtime-1")
|
||||
},
|
||||
{
|
||||
"name": "agent-2@example.com",
|
||||
"platform": "openai",
|
||||
"credentials": agent_identity("runtime-2")
|
||||
},
|
||||
{
|
||||
"name": "ignored@example.com",
|
||||
"platform": "anthropic",
|
||||
"credentials": { "access_token": "ignored" }
|
||||
}
|
||||
]
|
||||
});
|
||||
assert_eq!(
|
||||
codex_agent_identity_import_auth_configs(&sub2api.to_string())
|
||||
.expect("sub2api Agent Identity export should parse")
|
||||
.len(),
|
||||
2
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dedicated_import_rejects_mixed_and_invalid_entries() {
|
||||
let mixed = json!([
|
||||
agent_identity("runtime-1"),
|
||||
{ "refresh_token": "refresh-token" }
|
||||
]);
|
||||
assert!(codex_agent_identity_import_auth_configs(&mixed.to_string()).is_err());
|
||||
|
||||
let invalid = json!({
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_runtime_id": "runtime-invalid",
|
||||
"agent_private_key": "not-a-pkcs8-key"
|
||||
});
|
||||
assert!(codex_agent_identity_import_auth_configs(&invalid.to_string()).is_err());
|
||||
assert!(codex_agent_identity_import_auth_configs("[]").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_identity_import_lock_key_is_scoped_to_provider_and_runtime() {
|
||||
let first = provider_agent_identity_import_lock_key("provider-a", "runtime-1");
|
||||
assert_eq!(
|
||||
first,
|
||||
provider_agent_identity_import_lock_key("provider-a", "runtime-1")
|
||||
);
|
||||
assert_ne!(
|
||||
first,
|
||||
provider_agent_identity_import_lock_key("provider-b", "runtime-1")
|
||||
);
|
||||
assert_ne!(
|
||||
first,
|
||||
provider_agent_identity_import_lock_key("provider-a", "runtime-2")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_identity_imports_with_distinct_runtimes_share_account_lock_keys() {
|
||||
let first =
|
||||
codex_agent_identity_import_auth_configs(&agent_identity("runtime-1").to_string())
|
||||
.expect("first Agent Identity should parse");
|
||||
let second =
|
||||
codex_agent_identity_import_auth_configs(&agent_identity("runtime-2").to_string())
|
||||
.expect("second Agent Identity should parse");
|
||||
let first_keys = provider_agent_identity_import_lock_keys("provider-a", &first);
|
||||
let second_keys = provider_agent_identity_import_lock_keys("provider-a", &second);
|
||||
|
||||
assert!(first_keys.iter().any(|key| second_keys.contains(key)));
|
||||
assert!(
|
||||
first_keys.contains(&provider_agent_identity_import_lock_key(
|
||||
"provider-a",
|
||||
"runtime-1"
|
||||
))
|
||||
);
|
||||
assert!(
|
||||
second_keys.contains(&provider_agent_identity_import_lock_key(
|
||||
"provider-a",
|
||||
"runtime-2"
|
||||
))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn agent_identity_import_runtime_lock_releases_after_contention() {
|
||||
let app = AppState::new().expect("app state should build");
|
||||
let state = AdminAppState::new(&app);
|
||||
let lock_keys = vec![provider_agent_identity_import_lock_key(
|
||||
"provider-a",
|
||||
"runtime-1",
|
||||
)];
|
||||
let first = acquire_provider_agent_identity_import_locks(
|
||||
&state,
|
||||
"provider-a",
|
||||
&lock_keys,
|
||||
"agent-identity-task-1",
|
||||
)
|
||||
.await
|
||||
.expect("first lock should acquire");
|
||||
let second = acquire_provider_agent_identity_import_locks(
|
||||
&state,
|
||||
"provider-a",
|
||||
&lock_keys,
|
||||
"agent-identity-task-2",
|
||||
)
|
||||
.await
|
||||
.expect_err("second lock should be rejected");
|
||||
assert_eq!(second.status(), axum::http::StatusCode::CONFLICT);
|
||||
release_provider_agent_identity_import_locks(&state, first).await;
|
||||
let third = acquire_provider_agent_identity_import_locks(
|
||||
&state,
|
||||
"provider-a",
|
||||
&lock_keys,
|
||||
"agent-identity-task-3",
|
||||
)
|
||||
.await
|
||||
.expect("lock should be reusable after release");
|
||||
release_provider_agent_identity_import_locks(&state, third).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn agent_identity_import_partial_lock_failure_releases_acquired_leases() {
|
||||
let app = AppState::new().expect("app state should build");
|
||||
let state = AdminAppState::new(&app);
|
||||
let held = state
|
||||
.runtime_state()
|
||||
.lock_try_acquire("z-held", "other-task", std::time::Duration::from_secs(30))
|
||||
.await
|
||||
.expect("runtime lock should be available")
|
||||
.expect("held lock should acquire");
|
||||
let lock_keys = vec!["a-free".to_string(), "z-held".to_string()];
|
||||
|
||||
let response = acquire_provider_agent_identity_import_locks(
|
||||
&state,
|
||||
"provider-a",
|
||||
&lock_keys,
|
||||
"agent-identity-task-partial",
|
||||
)
|
||||
.await
|
||||
.expect_err("second lock should cause contention");
|
||||
assert_eq!(response.status(), axum::http::StatusCode::CONFLICT);
|
||||
|
||||
let free = state
|
||||
.runtime_state()
|
||||
.lock_try_acquire(
|
||||
"a-free",
|
||||
"verification-task",
|
||||
std::time::Duration::from_secs(30),
|
||||
)
|
||||
.await
|
||||
.expect("runtime lock should be available")
|
||||
.expect("partially acquired lock should have been released");
|
||||
assert!(state
|
||||
.runtime_state()
|
||||
.lock_release(&free)
|
||||
.await
|
||||
.expect("free lock should release"));
|
||||
assert!(state
|
||||
.runtime_state()
|
||||
.lock_release(&held)
|
||||
.await
|
||||
.expect("held lock should release"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn agent_identity_import_lock_renewal_extends_all_leases() {
|
||||
let app = AppState::new().expect("app state should build");
|
||||
let state = AdminAppState::new(&app);
|
||||
let lease = state
|
||||
.runtime_state()
|
||||
.lock_try_acquire(
|
||||
"renewed-agent-lock",
|
||||
"agent-identity-task-renew",
|
||||
std::time::Duration::from_secs(1),
|
||||
)
|
||||
.await
|
||||
.expect("runtime lock should be available")
|
||||
.expect("lock should acquire");
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
renew_provider_agent_identity_import_locks(
|
||||
&state,
|
||||
std::slice::from_ref(&lease),
|
||||
std::time::Duration::from_secs(2),
|
||||
)
|
||||
.await
|
||||
.expect("lock should renew");
|
||||
tokio::time::sleep(std::time::Duration::from_millis(1_000)).await;
|
||||
|
||||
let contender = state
|
||||
.runtime_state()
|
||||
.lock_try_acquire(
|
||||
"renewed-agent-lock",
|
||||
"contending-task",
|
||||
std::time::Duration::from_secs(1),
|
||||
)
|
||||
.await
|
||||
.expect("runtime lock should be available");
|
||||
assert!(contender.is_none(), "renewed lease should still be held");
|
||||
release_provider_agent_identity_import_locks(&state, vec![lease]).await;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,3 +1,7 @@
|
||||
use super::super::super::duplicates::{
|
||||
acquire_codex_oauth_account_locks, find_duplicate_provider_oauth_key,
|
||||
release_codex_oauth_account_locks,
|
||||
};
|
||||
use super::super::super::errors::build_internal_control_error_response;
|
||||
use super::super::super::provisioning::{
|
||||
provider_oauth_token_payload_expires_at_unix_secs, seed_provider_oauth_pool_score,
|
||||
@@ -15,8 +19,10 @@ use super::shared::{
|
||||
};
|
||||
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_complete_key_id;
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::handlers::shared::sync_provider_key_oauth_status_snapshot;
|
||||
use crate::provider_key_auth::provider_key_is_oauth_managed;
|
||||
use crate::GatewayError;
|
||||
use aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyOAuthRuntimeStateCasUpdate;
|
||||
use axum::{
|
||||
body::{Body, Bytes},
|
||||
http,
|
||||
@@ -88,6 +94,12 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
|
||||
"state 无效或已过期",
|
||||
));
|
||||
}
|
||||
if state_data.expected_encrypted_auth_config != key.encrypted_auth_config {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::CONFLICT,
|
||||
"授权期间 Key 认证信息已变更,请重新获取授权",
|
||||
));
|
||||
}
|
||||
|
||||
let provider_id = key.provider_id.clone();
|
||||
let provider = state
|
||||
@@ -208,39 +220,164 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
|
||||
"provider oauth encryption unavailable",
|
||||
));
|
||||
};
|
||||
let updated = state
|
||||
.update_provider_catalog_key_oauth_credentials(
|
||||
&key_id,
|
||||
&encrypted_api_key,
|
||||
Some(&encrypted_auth_config),
|
||||
expires_at,
|
||||
let codex_oauth_account_leases = if provider_type == "codex" {
|
||||
match acquire_codex_oauth_account_locks(state, &provider_id, &auth_config, "key-complete")
|
||||
.await
|
||||
{
|
||||
Ok(leases) => leases,
|
||||
Err(error) => {
|
||||
return Ok(build_internal_control_error_response(
|
||||
error.status_code(),
|
||||
error.detail(),
|
||||
));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
if provider_type == "codex" {
|
||||
let duplicate = match state
|
||||
.find_duplicate_provider_oauth_key(&provider_id, &auth_config, Some(&key_id))
|
||||
.await
|
||||
{
|
||||
Ok(duplicate) => duplicate,
|
||||
Err(detail) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::CONFLICT,
|
||||
detail,
|
||||
));
|
||||
}
|
||||
};
|
||||
if let Some(duplicate) = duplicate {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::CONFLICT,
|
||||
format!(
|
||||
"该 ChatGPT 账号已存在于其他 Key(名称: {})",
|
||||
duplicate.name
|
||||
),
|
||||
));
|
||||
}
|
||||
}
|
||||
let mut recovered_key = key.clone();
|
||||
recovered_key.encrypted_api_key = Some(encrypted_api_key.clone());
|
||||
recovered_key.encrypted_auth_config = Some(encrypted_auth_config.clone());
|
||||
recovered_key.expires_at_unix_secs = expires_at;
|
||||
recovered_key.oauth_invalid_at_unix_secs = None;
|
||||
recovered_key.oauth_invalid_reason = None;
|
||||
recovered_key.updated_at_unix_secs = Some(now_unix_secs);
|
||||
recovered_key.status_snapshot = sync_provider_key_oauth_status_snapshot(
|
||||
recovered_key.status_snapshot.as_ref(),
|
||||
&recovered_key,
|
||||
);
|
||||
let oauth_status = recovered_key
|
||||
.status_snapshot
|
||||
.as_ref()
|
||||
.and_then(|snapshot| snapshot.get("oauth"))
|
||||
.cloned()
|
||||
.unwrap_or(serde_json::Value::Null);
|
||||
let persisted_encrypted_auth_config = recovered_key
|
||||
.encrypted_auth_config
|
||||
.clone()
|
||||
.expect("recovered auth config should be present");
|
||||
let updated_result = state
|
||||
.app()
|
||||
.compare_and_update_provider_catalog_key_oauth_runtime_state(
|
||||
&ProviderCatalogKeyOAuthRuntimeStateCasUpdate {
|
||||
key_id: key_id.clone(),
|
||||
expected_encrypted_auth_config: state_data.expected_encrypted_auth_config,
|
||||
encrypted_auth_config: persisted_encrypted_auth_config.clone(),
|
||||
encrypted_api_key_update: Some(encrypted_api_key),
|
||||
expires_at_unix_secs_update: Some(expires_at),
|
||||
oauth_invalid_at_unix_secs: None,
|
||||
oauth_invalid_reason: None,
|
||||
reset_error_count: true,
|
||||
upstream_metadata_patch: None,
|
||||
status_snapshot_patch: json!({ "oauth": oauth_status }),
|
||||
updated_at_unix_secs: Some(now_unix_secs),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
.await;
|
||||
let _ = state
|
||||
.app()
|
||||
.invalidate_local_oauth_refresh_entry(&key_id)
|
||||
.await;
|
||||
let updated = match updated_result {
|
||||
Ok(updated) => updated,
|
||||
Err(error) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
if !updated {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
"Key 不存在",
|
||||
http::StatusCode::CONFLICT,
|
||||
"授权期间 Key 认证信息已变更,请重新获取授权",
|
||||
));
|
||||
}
|
||||
if !state
|
||||
.clear_provider_catalog_key_oauth_invalid_marker(&key_id)
|
||||
.await?
|
||||
let current_after_cas = match state
|
||||
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
|
||||
.await
|
||||
{
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
"Key 不存在",
|
||||
));
|
||||
}
|
||||
let Some(recovered_key) = state
|
||||
.reset_provider_catalog_key_recovery_state(&key_id)
|
||||
.await?
|
||||
else {
|
||||
Ok(keys) => keys.into_iter().next(),
|
||||
Err(error) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
let Some(current_after_cas) = current_after_cas else {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
"Key 不存在",
|
||||
));
|
||||
};
|
||||
if current_after_cas.encrypted_auth_config.as_deref()
|
||||
!= Some(persisted_encrypted_auth_config.as_str())
|
||||
{
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::CONFLICT,
|
||||
"授权期间 Key 认证信息已变更,请重新获取授权",
|
||||
));
|
||||
}
|
||||
let recovered_key = match state
|
||||
.reset_provider_catalog_key_recovery_state_fenced(&key_id, &persisted_encrypted_auth_config)
|
||||
.await
|
||||
{
|
||||
Ok(recovered_key) => recovered_key,
|
||||
Err(error) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
let Some(recovered_key) = recovered_key else {
|
||||
let key_exists = match state
|
||||
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
|
||||
.await
|
||||
{
|
||||
Ok(keys) => !keys.is_empty(),
|
||||
Err(error) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
if !key_exists {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
"Key 不存在",
|
||||
));
|
||||
}
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::CONFLICT,
|
||||
"授权期间 Key 认证信息已变更,请重新获取授权",
|
||||
));
|
||||
};
|
||||
seed_provider_oauth_pool_score(state, &provider.id, &recovered_key, now_unix_secs).await;
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
|
||||
spawn_provider_oauth_account_state_refresh_after_update(
|
||||
state.cloned_app(),
|
||||
|
||||
+51
-12
@@ -1,4 +1,7 @@
|
||||
use super::super::super::duplicates::find_duplicate_provider_oauth_key;
|
||||
use super::super::super::duplicates::{
|
||||
acquire_codex_oauth_account_locks, find_duplicate_provider_oauth_key,
|
||||
release_codex_oauth_account_locks,
|
||||
};
|
||||
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,
|
||||
@@ -158,14 +161,39 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
|
||||
};
|
||||
|
||||
let api_formats = provider_oauth_active_api_formats(&endpoints);
|
||||
let codex_oauth_account_leases = if provider_type == "codex" {
|
||||
match acquire_codex_oauth_account_locks(
|
||||
state,
|
||||
&provider_id,
|
||||
&auth_config,
|
||||
"provider-complete",
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(leases) => leases,
|
||||
Err(error) => {
|
||||
return Ok(build_internal_control_error_response(
|
||||
error.status_code(),
|
||||
error.detail(),
|
||||
));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
let duplicate = match state
|
||||
.find_duplicate_provider_oauth_key(&provider_id, &auth_config, None)
|
||||
.await
|
||||
{
|
||||
Ok(duplicate) => duplicate,
|
||||
Err(detail) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
if provider_type == "codex" {
|
||||
http::StatusCode::CONFLICT
|
||||
} else {
|
||||
http::StatusCode::BAD_REQUEST
|
||||
},
|
||||
detail,
|
||||
));
|
||||
}
|
||||
@@ -173,7 +201,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
|
||||
|
||||
let replaced = duplicate.is_some();
|
||||
let persisted_key = if let Some(existing_key) = duplicate {
|
||||
match state
|
||||
let update_result = state
|
||||
.update_existing_provider_oauth_catalog_key(
|
||||
&existing_key,
|
||||
&provider_type,
|
||||
@@ -183,10 +211,15 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
|
||||
key_proxy.clone(),
|
||||
expires_at,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Some(key) => key,
|
||||
None => {
|
||||
.await;
|
||||
match update_result {
|
||||
Err(error) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
return Err(error);
|
||||
}
|
||||
Ok(Some(key)) => key,
|
||||
Ok(None) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
"provider oauth write unavailable",
|
||||
@@ -214,7 +247,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
|
||||
.unwrap_or(0)
|
||||
)
|
||||
});
|
||||
match state
|
||||
let create_result = state
|
||||
.create_provider_oauth_catalog_key(
|
||||
&provider_id,
|
||||
&provider_type,
|
||||
@@ -225,10 +258,15 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
|
||||
key_proxy.clone(),
|
||||
expires_at,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Some(key) => key,
|
||||
None => {
|
||||
.await;
|
||||
match create_result {
|
||||
Err(error) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
return Err(error);
|
||||
}
|
||||
Ok(Some(key)) => key,
|
||||
Ok(None) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
"provider oauth write unavailable",
|
||||
@@ -236,6 +274,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
|
||||
}
|
||||
}
|
||||
};
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
|
||||
spawn_provider_oauth_account_state_refresh_after_update(
|
||||
state.cloned_app(),
|
||||
|
||||
@@ -6,6 +6,31 @@ use axum::{
|
||||
use serde_json::{Map, Value};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
pub(super) fn admin_provider_oauth_single_import_audit_taxonomy(
|
||||
request_body: Option<&axum::body::Bytes>,
|
||||
) -> (&'static str, &'static str) {
|
||||
let creates_agent_identity = request_body
|
||||
.and_then(|body| serde_json::from_slice::<Value>(body).ok())
|
||||
.and_then(|value| value.as_object().cloned())
|
||||
.is_some_and(|payload| {
|
||||
payload
|
||||
.get("create_agent_identity")
|
||||
.and_then(Value::as_bool)
|
||||
== Some(true)
|
||||
});
|
||||
if creates_agent_identity {
|
||||
(
|
||||
"admin_provider_oauth_agent_identity_created",
|
||||
"create_provider_agent_identity",
|
||||
)
|
||||
} else {
|
||||
(
|
||||
"admin_provider_oauth_refresh_token_imported",
|
||||
"import_provider_oauth_refresh_token",
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn attach_admin_provider_oauth_audit_response(
|
||||
response: Response<Body>,
|
||||
event_name: &'static str,
|
||||
@@ -96,4 +121,53 @@ mod tests {
|
||||
assert!(name.starts_with("codex_"));
|
||||
assert!(name.ends_with("_3"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_import_audit_distinguishes_agent_identity_creation_without_exposing_input() {
|
||||
let body = axum::body::Bytes::from(
|
||||
json!({
|
||||
"create_agent_identity": true,
|
||||
"access_token": "secret-access-token"
|
||||
})
|
||||
.to_string(),
|
||||
);
|
||||
assert_eq!(
|
||||
admin_provider_oauth_single_import_audit_taxonomy(Some(&body)),
|
||||
(
|
||||
"admin_provider_oauth_agent_identity_created",
|
||||
"create_provider_agent_identity",
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_import_audit_rejects_removed_session_token_creation_alias() {
|
||||
let body = axum::body::Bytes::from(
|
||||
json!({
|
||||
"create_agent_identity_from_session_token": true,
|
||||
"access_token": "secret-access-token"
|
||||
})
|
||||
.to_string(),
|
||||
);
|
||||
assert_eq!(
|
||||
admin_provider_oauth_single_import_audit_taxonomy(Some(&body)),
|
||||
(
|
||||
"admin_provider_oauth_refresh_token_imported",
|
||||
"import_provider_oauth_refresh_token",
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_import_audit_keeps_standard_import_taxonomy() {
|
||||
let body =
|
||||
axum::body::Bytes::from(json!({ "refresh_token": "secret-refresh-token" }).to_string());
|
||||
assert_eq!(
|
||||
admin_provider_oauth_single_import_audit_taxonomy(Some(&body)),
|
||||
(
|
||||
"admin_provider_oauth_refresh_token_imported",
|
||||
"import_provider_oauth_refresh_token",
|
||||
)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
use super::super::duplicates::find_duplicate_provider_oauth_key;
|
||||
use super::super::duplicates::{
|
||||
acquire_codex_oauth_account_locks, find_duplicate_provider_oauth_key,
|
||||
release_codex_oauth_account_locks,
|
||||
};
|
||||
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,
|
||||
@@ -27,10 +30,12 @@ use crate::handlers::admin::request::{
|
||||
};
|
||||
use crate::GatewayError;
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use aether_oauth::provider::{
|
||||
ProviderOAuthImportInput, ProviderOAuthService, ProviderOAuthTransportContext,
|
||||
};
|
||||
use aether_oauth::{core::OAuthError, network::OAuthNetworkContext};
|
||||
use aether_runtime_state::RuntimeLockLease;
|
||||
use axum::{
|
||||
body::Body,
|
||||
http,
|
||||
@@ -38,6 +43,15 @@ use axum::{
|
||||
Json,
|
||||
};
|
||||
use serde_json::json;
|
||||
use std::sync::{
|
||||
atomic::{AtomicBool, Ordering},
|
||||
Arc,
|
||||
};
|
||||
use std::time::Duration;
|
||||
use uuid::Uuid;
|
||||
|
||||
const CODEX_AGENT_IDENTITY_ENROLLMENT_LOCK_TTL: Duration = Duration::from_secs(180);
|
||||
const CODEX_AGENT_IDENTITY_ENROLLMENT_LOCK_RENEW_INTERVAL: Duration = Duration::from_secs(60);
|
||||
|
||||
struct AdminProviderOAuthSingleImportTokens {
|
||||
access_token: String,
|
||||
@@ -45,6 +59,19 @@ struct AdminProviderOAuthSingleImportTokens {
|
||||
expires_at: Option<u64>,
|
||||
}
|
||||
|
||||
struct CodexAgentIdentityEnrollment {
|
||||
leases: Vec<RuntimeLockLease>,
|
||||
duplicate: Option<StoredProviderCatalogKey>,
|
||||
lease_lost: Arc<AtomicBool>,
|
||||
heartbeat: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl Drop for CodexAgentIdentityEnrollment {
|
||||
fn drop(&mut self) {
|
||||
self.heartbeat.abort();
|
||||
}
|
||||
}
|
||||
|
||||
fn sanitize_windsurf_import_error(error: &OAuthError) -> String {
|
||||
match error {
|
||||
OAuthError::InvalidRequest(_) => "Windsurf 凭据验证失败: 请求参数无效".to_string(),
|
||||
@@ -115,14 +142,32 @@ fn import_payload_bool(payload: &serde_json::Map<String, serde_json::Value>, key
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn codex_session_token_identity_hints(
|
||||
session_token: &str,
|
||||
fn codex_agent_identity_access_token_input(
|
||||
payload: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> Option<String> {
|
||||
import_payload_string_any(payload, &["access_token", "accessToken"])
|
||||
}
|
||||
|
||||
fn import_payload_requests_agent_identity(
|
||||
payload: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> bool {
|
||||
import_payload_bool(payload, "create_agent_identity")
|
||||
}
|
||||
|
||||
fn import_payload_requests_legacy_agent_identity(
|
||||
payload: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> bool {
|
||||
import_payload_bool(payload, "create_agent_identity_from_session_token")
|
||||
}
|
||||
|
||||
fn codex_access_token_identity_hints(
|
||||
access_token: &str,
|
||||
) -> Result<serde_json::Map<String, serde_json::Value>, &'static str> {
|
||||
let mut hints = serde_json::Map::new();
|
||||
enrich_admin_provider_oauth_auth_config(
|
||||
"codex",
|
||||
&mut hints,
|
||||
&json!({ "access_token": session_token }),
|
||||
&json!({ "access_token": access_token }),
|
||||
);
|
||||
let account_id = hints
|
||||
.get("account_id")
|
||||
@@ -135,22 +180,173 @@ fn codex_session_token_identity_hints(
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
if account_id.is_none() || user_id.is_none() {
|
||||
return Err("ChatGPT Session Token 缺少账号身份字段");
|
||||
return Err("ChatGPT Access Token 缺少账号身份字段");
|
||||
}
|
||||
Ok(hints)
|
||||
}
|
||||
|
||||
async fn resolve_admin_provider_oauth_codex_session_agent_identity_import(
|
||||
async fn prepare_codex_agent_identity_enrollment(
|
||||
state: &AdminAppState<'_>,
|
||||
session_token: &str,
|
||||
provider_id: &str,
|
||||
identity_hints: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> Result<CodexAgentIdentityEnrollment, Response<Body>> {
|
||||
let lock_keys =
|
||||
crate::handlers::admin::provider::oauth::duplicates::codex_agent_identity_account_lock_keys(
|
||||
provider_id,
|
||||
identity_hints,
|
||||
);
|
||||
if lock_keys.is_empty() {
|
||||
return Err(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"ChatGPT Access Token 缺少账号身份字段",
|
||||
));
|
||||
}
|
||||
let owner = format!(
|
||||
"aether-gateway-agent-identity-enrollment-{}",
|
||||
Uuid::new_v4()
|
||||
);
|
||||
let mut leases = Vec::with_capacity(lock_keys.len());
|
||||
for lock_key in lock_keys {
|
||||
match state
|
||||
.runtime_state()
|
||||
.lock_try_acquire(
|
||||
lock_key.as_str(),
|
||||
owner.as_str(),
|
||||
CODEX_AGENT_IDENTITY_ENROLLMENT_LOCK_TTL,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(lease)) => leases.push(lease),
|
||||
Ok(None) => {
|
||||
release_codex_agent_identity_leases(state, leases).await;
|
||||
return Err(build_internal_control_error_response(
|
||||
http::StatusCode::CONFLICT,
|
||||
"该 ChatGPT 账号正在创建 Agent Identity,请稍后重试",
|
||||
));
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
provider_id = %provider_id,
|
||||
error = ?error,
|
||||
"gateway Agent Identity enrollment lock unavailable"
|
||||
);
|
||||
release_codex_agent_identity_leases(state, leases).await;
|
||||
return Err(build_internal_control_error_response(
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Agent Identity 创建锁暂不可用,请稍后重试",
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let lease_lost = Arc::new(AtomicBool::new(false));
|
||||
let heartbeat = spawn_codex_agent_identity_enrollment_heartbeat(
|
||||
state.cloned_app(),
|
||||
leases.clone(),
|
||||
Arc::clone(&lease_lost),
|
||||
);
|
||||
|
||||
let duplicate = match state
|
||||
.find_duplicate_provider_oauth_key(provider_id, identity_hints, None)
|
||||
.await
|
||||
{
|
||||
Ok(duplicate) => duplicate,
|
||||
Err(detail) => {
|
||||
heartbeat.abort();
|
||||
release_codex_agent_identity_leases(state, leases).await;
|
||||
return Err(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
detail,
|
||||
));
|
||||
}
|
||||
};
|
||||
Ok(CodexAgentIdentityEnrollment {
|
||||
leases,
|
||||
duplicate,
|
||||
lease_lost,
|
||||
heartbeat,
|
||||
})
|
||||
}
|
||||
|
||||
fn spawn_codex_agent_identity_enrollment_heartbeat(
|
||||
app: crate::AppState,
|
||||
leases: Vec<RuntimeLockLease>,
|
||||
lease_lost: Arc<AtomicBool>,
|
||||
) -> tokio::task::JoinHandle<()> {
|
||||
tokio::spawn(async move {
|
||||
let mut renew_timer =
|
||||
tokio::time::interval(CODEX_AGENT_IDENTITY_ENROLLMENT_LOCK_RENEW_INTERVAL);
|
||||
renew_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
|
||||
renew_timer.tick().await;
|
||||
loop {
|
||||
renew_timer.tick().await;
|
||||
for lease in &leases {
|
||||
match app
|
||||
.runtime_state()
|
||||
.lock_renew(lease, CODEX_AGENT_IDENTITY_ENROLLMENT_LOCK_TTL)
|
||||
.await
|
||||
{
|
||||
Ok(true) => {}
|
||||
Ok(false) => {
|
||||
lease_lost.store(true, Ordering::Release);
|
||||
tracing::error!(
|
||||
lock_key = %lease.key,
|
||||
"gateway Agent Identity enrollment lock was lost"
|
||||
);
|
||||
return;
|
||||
}
|
||||
Err(error) => {
|
||||
lease_lost.store(true, Ordering::Release);
|
||||
tracing::error!(
|
||||
lock_key = %lease.key,
|
||||
error = ?error,
|
||||
"gateway Agent Identity enrollment lock renewal failed"
|
||||
);
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn release_codex_agent_identity_leases(
|
||||
state: &AdminAppState<'_>,
|
||||
leases: Vec<RuntimeLockLease>,
|
||||
) {
|
||||
for lease in leases {
|
||||
if let Err(error) = state.runtime_state().lock_release(&lease).await {
|
||||
tracing::warn!(
|
||||
lock_key = %lease.key,
|
||||
error = ?error,
|
||||
"gateway Agent Identity enrollment lock release failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn release_codex_agent_identity_enrollment(
|
||||
state: &AdminAppState<'_>,
|
||||
enrollment: Option<CodexAgentIdentityEnrollment>,
|
||||
) {
|
||||
let Some(mut enrollment) = enrollment else {
|
||||
return;
|
||||
};
|
||||
enrollment.heartbeat.abort();
|
||||
release_codex_agent_identity_leases(state, std::mem::take(&mut enrollment.leases)).await;
|
||||
}
|
||||
|
||||
async fn resolve_admin_provider_oauth_codex_access_token_agent_identity_import(
|
||||
state: &AdminAppState<'_>,
|
||||
access_token: &str,
|
||||
identity_hints: serde_json::Map<String, serde_json::Value>,
|
||||
request_proxy: Option<ProxySnapshot>,
|
||||
) -> Result<AdminProviderOAuthSingleImportTokens, Response<Body>> {
|
||||
let executor = crate::oauth::GatewayOAuthHttpExecutor::new(*state);
|
||||
let mut auth_config =
|
||||
aether_provider_transport::create_codex_agent_identity_from_session_token(
|
||||
aether_provider_transport::register_codex_agent_identity_from_access_token(
|
||||
&executor,
|
||||
session_token,
|
||||
access_token,
|
||||
OAuthNetworkContext::provider_operation(request_proxy),
|
||||
)
|
||||
.await
|
||||
@@ -584,18 +780,15 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
||||
));
|
||||
}
|
||||
};
|
||||
let create_agent_identity_from_session_token =
|
||||
import_payload_bool(&raw_payload, "create_agent_identity_from_session_token");
|
||||
let session_token_agent_identity_input = if create_agent_identity_from_session_token {
|
||||
import_payload_string_any(
|
||||
&raw_payload,
|
||||
&[
|
||||
"session_token",
|
||||
"sessionToken",
|
||||
"access_token",
|
||||
"accessToken",
|
||||
],
|
||||
)
|
||||
if import_payload_requests_legacy_agent_identity(&raw_payload) {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"旧版 Agent Identity 创建参数已停用,请使用 create_agent_identity 和 access_token",
|
||||
));
|
||||
}
|
||||
let create_agent_identity = import_payload_requests_agent_identity(&raw_payload);
|
||||
let agent_identity_access_token_input = if create_agent_identity {
|
||||
codex_agent_identity_access_token_input(&raw_payload)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
@@ -644,10 +837,7 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
||||
refresh_token_input.as_deref(),
|
||||
access_token_input.as_deref(),
|
||||
);
|
||||
if !create_agent_identity_from_session_token
|
||||
&& refresh_token_input.is_none()
|
||||
&& access_token_input.is_none()
|
||||
{
|
||||
if !create_agent_identity && refresh_token_input.is_none() && access_token_input.is_none() {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"Refresh Token、Access Token 或 sso_token 不能为空",
|
||||
@@ -665,16 +855,16 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
||||
"Kiro 不支持单条 Refresh Token 导入,请使用批量导入或设备授权。",
|
||||
));
|
||||
}
|
||||
if create_agent_identity_from_session_token && provider_type != "codex" {
|
||||
if create_agent_identity && provider_type != "codex" {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"仅 Codex Provider 支持使用 Session Token 创建 Agent Identity",
|
||||
"仅 Codex Provider 支持使用 Access Token 创建 Agent Identity",
|
||||
));
|
||||
}
|
||||
if create_agent_identity_from_session_token && refresh_token_input.is_some() {
|
||||
if create_agent_identity && refresh_token_input.is_some() {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"使用 Session Token 创建 Agent Identity 时不能同时提交 Refresh Token",
|
||||
"使用 Access Token 创建 Agent Identity 时不能同时提交 Refresh Token",
|
||||
));
|
||||
}
|
||||
let template = admin_provider_oauth_template(&provider_type);
|
||||
@@ -697,15 +887,16 @@ 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 mut agent_identity_enrollment = None;
|
||||
|
||||
let resolved_import = if create_agent_identity_from_session_token {
|
||||
let Some(session_token) = session_token_agent_identity_input.as_deref() else {
|
||||
let resolved_import = if create_agent_identity {
|
||||
let Some(access_token) = agent_identity_access_token_input.as_deref() else {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"ChatGPT Session Token(JWT)不能为空",
|
||||
"ChatGPT Access Token(JWT)不能为空",
|
||||
));
|
||||
};
|
||||
let identity_hints = match codex_session_token_identity_hints(session_token) {
|
||||
let mut identity_hints = match codex_access_token_identity_hints(access_token) {
|
||||
Ok(hints) => hints,
|
||||
Err(detail) => {
|
||||
return Ok(build_internal_control_error_response(
|
||||
@@ -714,16 +905,30 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
||||
));
|
||||
}
|
||||
};
|
||||
match resolve_admin_provider_oauth_codex_session_agent_identity_import(
|
||||
identity_hints.insert("provider_type".to_string(), json!("codex"));
|
||||
let enrollment =
|
||||
match prepare_codex_agent_identity_enrollment(state, &provider_id, &identity_hints)
|
||||
.await
|
||||
{
|
||||
Ok(enrollment) => enrollment,
|
||||
Err(response) => return Ok(response),
|
||||
};
|
||||
match resolve_admin_provider_oauth_codex_access_token_agent_identity_import(
|
||||
state,
|
||||
session_token,
|
||||
access_token,
|
||||
identity_hints,
|
||||
request_proxy.clone(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(value) => value,
|
||||
Err(response) => return Ok(response),
|
||||
Ok(value) => {
|
||||
agent_identity_enrollment = Some(enrollment);
|
||||
value
|
||||
}
|
||||
Err(response) => {
|
||||
release_codex_agent_identity_enrollment(state, Some(enrollment)).await;
|
||||
return Ok(response);
|
||||
}
|
||||
}
|
||||
} else if provider_type == "windsurf" {
|
||||
if !import_payload_has_windsurf_credentials(&raw_payload) {
|
||||
@@ -772,7 +977,7 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
||||
mut auth_config,
|
||||
mut expires_at,
|
||||
} = resolved_import;
|
||||
if !create_agent_identity_from_session_token {
|
||||
if !create_agent_identity {
|
||||
apply_single_import_hints(&provider_type, &raw_payload, &mut auth_config);
|
||||
if let Some(header_access_token) =
|
||||
provider_oauth_import_authorization_bearer_token_from_object(&raw_payload)
|
||||
@@ -794,23 +999,85 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
|
||||
let codex_oauth_account_leases = if !create_agent_identity && provider_type == "codex" {
|
||||
match acquire_codex_oauth_account_locks(state, &provider_id, &auth_config, "single-import")
|
||||
.await
|
||||
{
|
||||
Ok(leases) => leases,
|
||||
Err(error) => {
|
||||
return Ok(build_internal_control_error_response(
|
||||
error.status_code(),
|
||||
error.detail(),
|
||||
));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
let api_formats = provider_oauth_active_api_formats(&endpoints);
|
||||
let duplicate = match state
|
||||
.find_duplicate_provider_oauth_key(&provider_id, &auth_config, None)
|
||||
.await
|
||||
{
|
||||
Ok(duplicate) => duplicate,
|
||||
Err(detail) => {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
detail,
|
||||
));
|
||||
let duplicate = if create_agent_identity {
|
||||
let initial_duplicate_id = agent_identity_enrollment
|
||||
.as_ref()
|
||||
.and_then(|enrollment| enrollment.duplicate.as_ref())
|
||||
.map(|key| key.id.clone());
|
||||
match state
|
||||
.find_duplicate_provider_oauth_key(&provider_id, &auth_config, None)
|
||||
.await
|
||||
{
|
||||
Ok(duplicate) => {
|
||||
let current_duplicate_id = duplicate.as_ref().map(|key| key.id.as_str());
|
||||
if initial_duplicate_id.as_deref() != current_duplicate_id {
|
||||
tracing::info!(
|
||||
provider_id = %provider_id,
|
||||
initial_duplicate_id = ?initial_duplicate_id,
|
||||
current_duplicate_id,
|
||||
"gateway Agent Identity duplicate changed during enrollment"
|
||||
);
|
||||
}
|
||||
duplicate
|
||||
}
|
||||
Err(detail) => {
|
||||
release_codex_agent_identity_enrollment(state, agent_identity_enrollment).await;
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::CONFLICT,
|
||||
detail,
|
||||
));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
match state
|
||||
.find_duplicate_provider_oauth_key(&provider_id, &auth_config, None)
|
||||
.await
|
||||
{
|
||||
Ok(duplicate) => duplicate,
|
||||
Err(detail) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
return Ok(build_internal_control_error_response(
|
||||
if provider_type == "codex" {
|
||||
http::StatusCode::CONFLICT
|
||||
} else {
|
||||
http::StatusCode::BAD_REQUEST
|
||||
},
|
||||
detail,
|
||||
));
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let replaced = duplicate.is_some();
|
||||
let persisted_key = if let Some(existing_key) = duplicate {
|
||||
match state
|
||||
if agent_identity_enrollment
|
||||
.as_ref()
|
||||
.is_some_and(|enrollment| enrollment.lease_lost.load(Ordering::Acquire))
|
||||
{
|
||||
release_codex_agent_identity_enrollment(state, agent_identity_enrollment).await;
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Agent Identity 创建锁已失效,请稍后重试",
|
||||
));
|
||||
}
|
||||
let persisted_key_result = if let Some(existing_key) = duplicate {
|
||||
state
|
||||
.update_existing_provider_oauth_catalog_key(
|
||||
&existing_key,
|
||||
&provider_type,
|
||||
@@ -820,21 +1087,12 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
||||
key_proxy.clone(),
|
||||
expires_at,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Some(key) => key,
|
||||
None => {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
"provider oauth write unavailable",
|
||||
));
|
||||
}
|
||||
}
|
||||
.await
|
||||
} else {
|
||||
let name = name.unwrap_or_else(|| {
|
||||
admin_provider_oauth_key_name_from_auth_config(&provider_type, &auth_config, None)
|
||||
});
|
||||
match state
|
||||
state
|
||||
.create_provider_oauth_catalog_key(
|
||||
&provider_id,
|
||||
&provider_type,
|
||||
@@ -845,17 +1103,68 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
||||
key_proxy.clone(),
|
||||
expires_at,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Some(key) => key,
|
||||
None => {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
"provider oauth write unavailable",
|
||||
));
|
||||
}
|
||||
.await
|
||||
};
|
||||
let persisted_key = match persisted_key_result {
|
||||
Ok(Some(key)) => key,
|
||||
Ok(None) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
release_codex_agent_identity_enrollment(state, agent_identity_enrollment).await;
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
"provider oauth write unavailable",
|
||||
));
|
||||
}
|
||||
Err(error) => {
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
release_codex_agent_identity_enrollment(state, agent_identity_enrollment).await;
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
release_codex_oauth_account_locks(state, codex_oauth_account_leases).await;
|
||||
|
||||
let agent_identity_task_ready = if create_agent_identity {
|
||||
match runtime_endpoint.as_ref() {
|
||||
Some(endpoint) => match state
|
||||
.read_provider_transport_snapshot_uncached(
|
||||
&provider_id,
|
||||
&endpoint.id,
|
||||
&persisted_key.id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(transport)) => {
|
||||
match state.resolve_local_oauth_request_auth(&transport).await {
|
||||
Ok(Some(_)) => true,
|
||||
Ok(None) => false,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
provider_id = %provider_id,
|
||||
key_id = %persisted_key.id,
|
||||
error = ?error,
|
||||
"gateway Agent Identity initial task registration failed"
|
||||
);
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(None) => false,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
provider_id = %provider_id,
|
||||
key_id = %persisted_key.id,
|
||||
error = ?error,
|
||||
"gateway Agent Identity pending transport reload failed"
|
||||
);
|
||||
false
|
||||
}
|
||||
},
|
||||
None => false,
|
||||
}
|
||||
} else {
|
||||
true
|
||||
};
|
||||
release_codex_agent_identity_enrollment(state, agent_identity_enrollment).await;
|
||||
|
||||
spawn_provider_oauth_account_state_refresh_after_update(
|
||||
state.cloned_app(),
|
||||
@@ -864,6 +1173,28 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
||||
request_proxy.clone(),
|
||||
);
|
||||
|
||||
if !agent_identity_task_ready {
|
||||
return Ok((
|
||||
http::StatusCode::ACCEPTED,
|
||||
Json(json!({
|
||||
"detail": "Agent Identity 已安全保存,但 task 初始化暂未完成,系统将自动重试",
|
||||
"key_id": persisted_key.id,
|
||||
"provider_type": provider_type,
|
||||
"expires_at": serde_json::Value::Null,
|
||||
"has_refresh_token": false,
|
||||
"temporary": false,
|
||||
"email": auth_config
|
||||
.get("email")
|
||||
.cloned()
|
||||
.unwrap_or(serde_json::Value::Null),
|
||||
"replaced": replaced,
|
||||
"task_ready": false,
|
||||
"recoverable": true,
|
||||
})),
|
||||
)
|
||||
.into_response());
|
||||
}
|
||||
|
||||
Ok(Json(json!({
|
||||
"key_id": persisted_key.id,
|
||||
"provider_type": provider_type,
|
||||
@@ -882,9 +1213,12 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
apply_single_import_hints, codex_session_token_identity_hints, import_payload_bool,
|
||||
import_payload_string_any, import_payload_u64_any, sanitize_windsurf_import_error,
|
||||
apply_single_import_hints, codex_access_token_identity_hints,
|
||||
codex_agent_identity_access_token_input, import_payload_requests_agent_identity,
|
||||
import_payload_requests_legacy_agent_identity, import_payload_string_any,
|
||||
import_payload_u64_any, sanitize_windsurf_import_error,
|
||||
};
|
||||
use crate::handlers::admin::provider::oauth::duplicates::codex_agent_identity_account_lock_keys;
|
||||
use aether_oauth::core::OAuthError;
|
||||
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
|
||||
use serde_json::json;
|
||||
@@ -911,8 +1245,36 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_token_agent_identity_hints_require_and_extract_chatgpt_identity() {
|
||||
let session_token = unsigned_jwt(json!({
|
||||
fn agent_identity_reads_access_token_and_ignores_session_token_alias() {
|
||||
let payload = json!({
|
||||
"accessToken": "access-token",
|
||||
"sessionToken": "session-token-must-not-win",
|
||||
})
|
||||
.as_object()
|
||||
.cloned()
|
||||
.expect("payload should be an object");
|
||||
|
||||
assert_eq!(
|
||||
codex_agent_identity_access_token_input(&payload).as_deref(),
|
||||
Some("access-token")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_agent_identity_flag_does_not_treat_session_token_as_access_token() {
|
||||
let payload = json!({
|
||||
"sessionToken": "session-token-must-not-be-used",
|
||||
})
|
||||
.as_object()
|
||||
.cloned()
|
||||
.expect("payload should be an object");
|
||||
|
||||
assert_eq!(codex_agent_identity_access_token_input(&payload), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn access_token_agent_identity_hints_require_and_extract_chatgpt_identity() {
|
||||
let access_token = unsigned_jwt(json!({
|
||||
"https://api.openai.com/auth": {
|
||||
"chatgpt_account_id": "account-1",
|
||||
"chatgpt_user_id": "user-1",
|
||||
@@ -923,8 +1285,8 @@ mod tests {
|
||||
}
|
||||
}));
|
||||
|
||||
let hints = codex_session_token_identity_hints(&session_token)
|
||||
.expect("session token identity hints should parse");
|
||||
let hints = codex_access_token_identity_hints(&access_token)
|
||||
.expect("access token identity hints should parse");
|
||||
|
||||
assert_eq!(hints.get("account_id"), Some(&json!("account-1")));
|
||||
assert_eq!(hints.get("user_id"), Some(&json!("user-1")));
|
||||
@@ -935,41 +1297,81 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_token_agent_identity_hints_reject_missing_identity() {
|
||||
let session_token = unsigned_jwt(json!({
|
||||
fn access_token_agent_identity_hints_reject_missing_identity() {
|
||||
let access_token = unsigned_jwt(json!({
|
||||
"https://api.openai.com/auth": {
|
||||
"chatgpt_account_id": "account-1"
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(
|
||||
codex_session_token_identity_hints(&session_token),
|
||||
Err("ChatGPT Session Token 缺少账号身份字段")
|
||||
codex_access_token_identity_hints(&access_token),
|
||||
Err("ChatGPT Access Token 缺少账号身份字段")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_token_agent_identity_flag_is_explicit_boolean_only() {
|
||||
fn agent_identity_enrollment_lock_is_stable_and_account_scoped() {
|
||||
let first = json!({
|
||||
"account_id": "account-1",
|
||||
"user_id": "user-1",
|
||||
})
|
||||
.as_object()
|
||||
.cloned()
|
||||
.expect("identity hints should be an object");
|
||||
let second = json!({
|
||||
"account_id": "account-2",
|
||||
"user_id": "user-1",
|
||||
})
|
||||
.as_object()
|
||||
.cloned()
|
||||
.expect("identity hints should be an object");
|
||||
|
||||
let first_keys = codex_agent_identity_account_lock_keys("provider-1", &first);
|
||||
assert_eq!(
|
||||
first_keys,
|
||||
codex_agent_identity_account_lock_keys("provider-1", &first)
|
||||
);
|
||||
assert_ne!(
|
||||
first_keys,
|
||||
codex_agent_identity_account_lock_keys("provider-1", &second)
|
||||
);
|
||||
assert_ne!(
|
||||
first_keys,
|
||||
codex_agent_identity_account_lock_keys("provider-2", &first)
|
||||
);
|
||||
assert!(first_keys
|
||||
.iter()
|
||||
.all(|key| !key.contains("account-1") && !key.contains("user-1")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_identity_flag_is_explicit_boolean_and_rejects_legacy_alias() {
|
||||
let payload = json!({
|
||||
"create_agent_identity": true,
|
||||
})
|
||||
.as_object()
|
||||
.cloned()
|
||||
.expect("payload should be an object");
|
||||
assert!(import_payload_requests_agent_identity(&payload));
|
||||
|
||||
let string_payload = json!({
|
||||
"create_agent_identity": "true",
|
||||
})
|
||||
.as_object()
|
||||
.cloned()
|
||||
.expect("payload should be an object");
|
||||
assert!(!import_payload_requests_agent_identity(&string_payload));
|
||||
|
||||
let legacy_payload = json!({
|
||||
"create_agent_identity_from_session_token": true,
|
||||
})
|
||||
.as_object()
|
||||
.cloned()
|
||||
.expect("payload should be an object");
|
||||
assert!(import_payload_bool(
|
||||
&payload,
|
||||
"create_agent_identity_from_session_token"
|
||||
));
|
||||
|
||||
let string_payload = json!({
|
||||
"create_agent_identity_from_session_token": "true",
|
||||
})
|
||||
.as_object()
|
||||
.cloned()
|
||||
.expect("payload should be an object");
|
||||
assert!(!import_payload_bool(
|
||||
&string_payload,
|
||||
"create_agent_identity_from_session_token"
|
||||
assert!(!import_payload_requests_agent_identity(&legacy_payload));
|
||||
assert!(import_payload_requests_legacy_agent_identity(
|
||||
&legacy_payload
|
||||
));
|
||||
}
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@ use super::state::{
|
||||
build_admin_provider_oauth_supported_types_payload,
|
||||
};
|
||||
use crate::handlers::admin::provider::shared::paths::{
|
||||
admin_provider_oauth_agent_identity_import_task_provider_id,
|
||||
admin_provider_oauth_batch_import_provider_id,
|
||||
admin_provider_oauth_batch_import_task_provider_id, admin_provider_oauth_complete_key_id,
|
||||
admin_provider_oauth_complete_provider_id, admin_provider_oauth_device_authorize_provider_id,
|
||||
@@ -83,6 +84,16 @@ pub(crate) async fn maybe_build_local_admin_provider_oauth_response(
|
||||
));
|
||||
}
|
||||
|
||||
if route_kind == Some("get_agent_identity_import_task_status") && *method == http::Method::GET {
|
||||
return Ok(Some(
|
||||
tasks::handle_admin_provider_oauth_agent_identity_import_task_status(
|
||||
state,
|
||||
request_context,
|
||||
)
|
||||
.await?,
|
||||
));
|
||||
}
|
||||
|
||||
if route_kind == Some("complete_key_oauth") && *method == http::Method::POST {
|
||||
let response = complete::handle_admin_provider_oauth_complete_key(
|
||||
state,
|
||||
@@ -128,6 +139,8 @@ pub(crate) async fn maybe_build_local_admin_provider_oauth_response(
|
||||
}
|
||||
|
||||
if route_kind == Some("import_refresh_token") && *method == http::Method::POST {
|
||||
let (event_name, action) =
|
||||
helpers::admin_provider_oauth_single_import_audit_taxonomy(request_body);
|
||||
let response = import::handle_admin_provider_oauth_import_refresh_token(
|
||||
state,
|
||||
request_context,
|
||||
@@ -136,8 +149,8 @@ pub(crate) async fn maybe_build_local_admin_provider_oauth_response(
|
||||
.await?;
|
||||
return Ok(Some(helpers::attach_admin_provider_oauth_audit_response(
|
||||
response,
|
||||
"admin_provider_oauth_refresh_token_imported",
|
||||
"import_provider_oauth_refresh_token",
|
||||
event_name,
|
||||
action,
|
||||
"provider",
|
||||
admin_provider_oauth_import_provider_id(request_context.path()),
|
||||
)));
|
||||
@@ -172,6 +185,22 @@ pub(crate) async fn maybe_build_local_admin_provider_oauth_response(
|
||||
)));
|
||||
}
|
||||
|
||||
if route_kind == Some("start_agent_identity_import_task") && *method == http::Method::POST {
|
||||
let response = batch::handle_admin_provider_oauth_start_agent_identity_import_task(
|
||||
state,
|
||||
request_context,
|
||||
request_body,
|
||||
)
|
||||
.await?;
|
||||
return Ok(Some(helpers::attach_admin_provider_oauth_audit_response(
|
||||
response,
|
||||
"admin_provider_oauth_agent_identity_import_started",
|
||||
"start_provider_agent_identity_import",
|
||||
"provider",
|
||||
admin_provider_oauth_agent_identity_import_task_provider_id(request_context.path()),
|
||||
)));
|
||||
}
|
||||
|
||||
if route_kind == Some("device_authorize") && *method == http::Method::POST {
|
||||
let response = device::handle_admin_provider_oauth_device_authorize(
|
||||
state,
|
||||
|
||||
+35
-79
@@ -1,16 +1,7 @@
|
||||
use super::super::super::errors::{
|
||||
merge_provider_oauth_refresh_failure_reason, normalize_provider_oauth_refresh_error_message,
|
||||
};
|
||||
use super::super::super::quota::shared::{
|
||||
persist_provider_quota_refresh_state, provider_auto_remove_banned_keys,
|
||||
should_auto_remove_oauth_invalid_key,
|
||||
};
|
||||
use super::super::super::errors::normalize_provider_oauth_refresh_error_message;
|
||||
use super::super::super::runtime::refresh_provider_oauth_account_state_after_update;
|
||||
use super::helpers::{self, RefreshDispatch, RefreshRequestContext, RefreshSuccessContext};
|
||||
use super::response;
|
||||
use crate::handlers::admin::provider::shared::payloads::{
|
||||
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_REFRESH_FAILED_PREFIX,
|
||||
};
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminLocalOAuthRefreshError};
|
||||
use crate::GatewayError;
|
||||
use axum::http;
|
||||
@@ -62,62 +53,36 @@ pub(super) async fn execute_admin_provider_oauth_refresh(
|
||||
"gateway manual provider oauth refresh failed"
|
||||
);
|
||||
if matches!(status_code, 400 | 401 | 403) {
|
||||
let failure_reason = format!(
|
||||
"{OAUTH_REFRESH_FAILED_PREFIX}Token 续期失败 ({status_code}): {error_reason}"
|
||||
);
|
||||
let merged_reason = merge_provider_oauth_refresh_failure_reason(
|
||||
key.oauth_invalid_reason.as_deref(),
|
||||
&failure_reason,
|
||||
);
|
||||
if let Some(merged_reason) = merged_reason {
|
||||
let _ = persist_provider_quota_refresh_state(
|
||||
state,
|
||||
&key_id,
|
||||
None,
|
||||
Some(helpers::unix_now_secs()),
|
||||
Some(merged_reason),
|
||||
None,
|
||||
let auto_removed = state
|
||||
.app()
|
||||
.persist_local_oauth_refresh_failure_state(
|
||||
&transport,
|
||||
status_code,
|
||||
body_excerpt.as_str(),
|
||||
false,
|
||||
)
|
||||
.await?;
|
||||
if provider_auto_remove_banned_keys(provider.config.as_ref()) {
|
||||
let now_unix_secs = helpers::unix_now_secs();
|
||||
let auto_removed = state
|
||||
.cleanup_provider_catalog_key_if_current(
|
||||
&provider,
|
||||
&key_id,
|
||||
|latest_key| {
|
||||
should_auto_remove_oauth_invalid_key(
|
||||
latest_key,
|
||||
Some(&failure_reason),
|
||||
false,
|
||||
now_unix_secs,
|
||||
)
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
if auto_removed {
|
||||
tracing::info!(
|
||||
trace_id = %trace_id,
|
||||
key_id = %key_id,
|
||||
provider_id = %provider.id,
|
||||
provider_type = %provider_type,
|
||||
event_name = "auto_removed_oauth_refresh_failed",
|
||||
"gateway manual provider oauth refresh auto-removed unusable key"
|
||||
);
|
||||
return Ok(RefreshDispatch::Respond(
|
||||
response::oauth_refresh_auto_removed_response(&error_reason),
|
||||
));
|
||||
}
|
||||
}
|
||||
if auto_removed {
|
||||
tracing::info!(
|
||||
trace_id = %trace_id,
|
||||
key_id = %key_id,
|
||||
provider_id = %provider.id,
|
||||
provider_type = %provider_type,
|
||||
event_name = "refresh_failed_retained",
|
||||
"gateway manual provider oauth refresh failure retained key"
|
||||
event_name = "auto_removed_oauth_refresh_failed",
|
||||
"gateway manual provider oauth refresh auto-removed unusable key"
|
||||
);
|
||||
return Ok(RefreshDispatch::Respond(
|
||||
response::oauth_refresh_auto_removed_response(&error_reason),
|
||||
));
|
||||
}
|
||||
tracing::info!(
|
||||
trace_id = %trace_id,
|
||||
key_id = %key_id,
|
||||
provider_id = %provider.id,
|
||||
provider_type = %provider_type,
|
||||
event_name = "refresh_failed_retained",
|
||||
"gateway manual provider oauth refresh failure retained key"
|
||||
);
|
||||
}
|
||||
return Ok(RefreshDispatch::Respond(
|
||||
response::oauth_refresh_failed_bad_request_response(&error_reason),
|
||||
@@ -164,28 +129,6 @@ pub(super) async fn execute_admin_provider_oauth_refresh(
|
||||
}
|
||||
};
|
||||
|
||||
if !helpers::key_is_account_blocked(&key, OAUTH_ACCOUNT_BLOCK_PREFIX) {
|
||||
let previous_oauth_refresh_issue =
|
||||
key.oauth_invalid_reason.as_deref().is_some_and(|reason| {
|
||||
reason.lines().map(str::trim).any(|line| {
|
||||
line.starts_with("[OAUTH_EXPIRED]") || line.starts_with("[REFRESH_FAILED]")
|
||||
})
|
||||
});
|
||||
let cleared = state
|
||||
.clear_provider_catalog_key_oauth_invalid_marker(&key_id)
|
||||
.await?;
|
||||
if cleared && previous_oauth_refresh_issue {
|
||||
tracing::info!(
|
||||
trace_id = %trace_id,
|
||||
key_id = %key_id,
|
||||
provider_id = %provider.id,
|
||||
provider_type = %provider_type,
|
||||
event_name = "refresh_fixed",
|
||||
"gateway manual provider oauth refresh cleared oauth invalid marker"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let refreshed_key = state
|
||||
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
|
||||
.await?
|
||||
@@ -223,3 +166,16 @@ pub(super) async fn execute_admin_provider_oauth_refresh(
|
||||
account_state_recheck_error,
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
#[test]
|
||||
fn manual_refresh_uses_fenced_state_persistence_without_redundant_clear() {
|
||||
let source = include_str!("execution.rs");
|
||||
assert!(source.contains("persist_local_oauth_refresh_failure_state"));
|
||||
let redundant_clear = ["clear_provider_catalog_key_", "oauth_invalid_marker"].concat();
|
||||
let unfenced_persistence = ["persist_provider_quota_", "refresh_state"].concat();
|
||||
assert!(!source.contains(&redundant_clear));
|
||||
assert!(!source.contains(&unfenced_persistence));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -43,12 +43,6 @@ pub(super) async fn parse_admin_provider_oauth_refresh_request(
|
||||
)));
|
||||
};
|
||||
let parsed_auth_config = helpers::parse_auth_config_object(&decrypted_auth_config);
|
||||
if !helpers::auth_config_has_refresh_token(&parsed_auth_config) {
|
||||
return Ok(RefreshDispatch::Respond(response::control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"缺少 refresh_token,需要重新授权",
|
||||
)));
|
||||
}
|
||||
|
||||
let provider_id = key.provider_id.clone();
|
||||
let Some(provider) = state
|
||||
@@ -63,6 +57,16 @@ pub(super) async fn parse_admin_provider_oauth_refresh_request(
|
||||
)));
|
||||
};
|
||||
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
|
||||
let is_agent_identity = provider_type == "codex"
|
||||
&& crate::provider_transport::is_codex_agent_identity_auth_config_value(
|
||||
&serde_json::Value::Object(parsed_auth_config.clone()),
|
||||
);
|
||||
if !is_agent_identity && !helpers::auth_config_has_refresh_token(&parsed_auth_config) {
|
||||
return Ok(RefreshDispatch::Respond(response::control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
"缺少 refresh_token,需要重新授权",
|
||||
)));
|
||||
}
|
||||
if !provider_key_is_oauth_managed(&key, provider_type.as_str()) {
|
||||
return Ok(RefreshDispatch::Respond(response::control_error_response(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
|
||||
@@ -86,6 +86,7 @@ pub(super) async fn handle_admin_provider_oauth_start_key(
|
||||
&provider_id,
|
||||
&provider_type,
|
||||
pkce_verifier.as_deref(),
|
||||
key.encrypted_auth_config.as_deref(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -158,7 +159,13 @@ pub(super) async fn handle_admin_provider_oauth_start_provider(
|
||||
.then(generate_provider_oauth_pkce_verifier);
|
||||
let code_challenge = pkce_verifier.as_deref().map(provider_oauth_pkce_s256);
|
||||
let nonce = match state
|
||||
.save_provider_oauth_state("", &provider_id, &provider_type, pkce_verifier.as_deref())
|
||||
.save_provider_oauth_state(
|
||||
"",
|
||||
&provider_id,
|
||||
&provider_type,
|
||||
pkce_verifier.as_deref(),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(nonce) => nonce,
|
||||
|
||||
@@ -1,5 +1,8 @@
|
||||
use super::super::errors::build_internal_control_error_response;
|
||||
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_batch_import_task_path;
|
||||
use crate::handlers::admin::provider::shared::paths::{
|
||||
admin_provider_oauth_agent_identity_import_task_path,
|
||||
admin_provider_oauth_batch_import_task_path,
|
||||
};
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::handlers::admin::shared::attach_admin_audit_response;
|
||||
use crate::GatewayError;
|
||||
@@ -10,13 +13,49 @@ use axum::{
|
||||
Json,
|
||||
};
|
||||
|
||||
const PROVIDER_AGENT_IDENTITY_IMPORT_KIND: &str = "agent_identity";
|
||||
|
||||
fn provider_oauth_import_task_matches_route(
|
||||
task_id: &str,
|
||||
payload: &serde_json::Value,
|
||||
agent_identity_only: bool,
|
||||
) -> bool {
|
||||
let has_agent_prefix = task_id.starts_with("agent-identity-");
|
||||
let import_kind = payload
|
||||
.get("import_kind")
|
||||
.and_then(serde_json::Value::as_str);
|
||||
if agent_identity_only {
|
||||
has_agent_prefix && import_kind == Some(PROVIDER_AGENT_IDENTITY_IMPORT_KIND)
|
||||
} else {
|
||||
!has_agent_prefix && import_kind != Some(PROVIDER_AGENT_IDENTITY_IMPORT_KIND)
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn handle_admin_provider_oauth_batch_import_task_status(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let Some((provider_id, task_id)) =
|
||||
handle_admin_provider_oauth_import_task_status(state, request_context, false).await
|
||||
}
|
||||
|
||||
pub(super) async fn handle_admin_provider_oauth_agent_identity_import_task_status(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
handle_admin_provider_oauth_import_task_status(state, request_context, true).await
|
||||
}
|
||||
|
||||
async fn handle_admin_provider_oauth_import_task_status(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
agent_identity_only: bool,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let task_path = if agent_identity_only {
|
||||
admin_provider_oauth_agent_identity_import_task_path(request_context.path())
|
||||
} else {
|
||||
admin_provider_oauth_batch_import_task_path(request_context.path())
|
||||
else {
|
||||
};
|
||||
let Some((provider_id, task_id)) = task_path else {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
"批量导入任务不存在",
|
||||
@@ -40,27 +79,86 @@ pub(super) async fn handle_admin_provider_oauth_batch_import_task_status(
|
||||
));
|
||||
}
|
||||
};
|
||||
if !provider_oauth_import_task_matches_route(&task_id, &payload, agent_identity_only) {
|
||||
return Ok(build_internal_control_error_response(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
"导入任务不存在或已过期",
|
||||
));
|
||||
}
|
||||
let status = payload
|
||||
.get("status")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(ToOwned::to_owned)
|
||||
.unwrap_or_default();
|
||||
let response = Json(payload).into_response();
|
||||
let (completed_event, failed_event, action, target_type) = if agent_identity_only {
|
||||
(
|
||||
"admin_provider_oauth_agent_identity_import_completed_viewed",
|
||||
"admin_provider_oauth_agent_identity_import_failed_viewed",
|
||||
"view_provider_agent_identity_import_terminal_state",
|
||||
"provider_agent_identity_import_task",
|
||||
)
|
||||
} else {
|
||||
(
|
||||
"admin_provider_oauth_batch_task_completed_viewed",
|
||||
"admin_provider_oauth_batch_task_failed_viewed",
|
||||
"view_provider_oauth_batch_task_terminal_state",
|
||||
"provider_oauth_batch_task",
|
||||
)
|
||||
};
|
||||
Ok(match status.as_str() {
|
||||
"completed" => attach_admin_audit_response(
|
||||
response,
|
||||
"admin_provider_oauth_batch_task_completed_viewed",
|
||||
"view_provider_oauth_batch_task_terminal_state",
|
||||
"provider_oauth_batch_task",
|
||||
completed_event,
|
||||
action,
|
||||
target_type,
|
||||
&format!("{provider_id}:{task_id}"),
|
||||
),
|
||||
"failed" => attach_admin_audit_response(
|
||||
response,
|
||||
"admin_provider_oauth_batch_task_failed_viewed",
|
||||
"view_provider_oauth_batch_task_terminal_state",
|
||||
"provider_oauth_batch_task",
|
||||
failed_event,
|
||||
action,
|
||||
target_type,
|
||||
&format!("{provider_id}:{task_id}"),
|
||||
),
|
||||
_ => response,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::provider_oauth_import_task_matches_route;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn import_task_status_routes_are_bidirectionally_isolated() {
|
||||
let agent_payload = json!({ "import_kind": "agent_identity" });
|
||||
let batch_payload = json!({ "import_kind": "oauth_batch" });
|
||||
|
||||
assert!(provider_oauth_import_task_matches_route(
|
||||
"agent-identity-task-1",
|
||||
&agent_payload,
|
||||
true,
|
||||
));
|
||||
assert!(!provider_oauth_import_task_matches_route(
|
||||
"agent-identity-task-1",
|
||||
&agent_payload,
|
||||
false,
|
||||
));
|
||||
assert!(provider_oauth_import_task_matches_route(
|
||||
"batch-task-1",
|
||||
&batch_payload,
|
||||
false,
|
||||
));
|
||||
assert!(!provider_oauth_import_task_matches_route(
|
||||
"batch-task-1",
|
||||
&batch_payload,
|
||||
true,
|
||||
));
|
||||
assert!(provider_oauth_import_task_matches_route(
|
||||
"legacy-batch-task",
|
||||
&json!({}),
|
||||
false,
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,7 +1,38 @@
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::provider_key_auth::provider_key_is_oauth_managed;
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use aether_runtime_state::RuntimeLockLease;
|
||||
use axum::http;
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
use uuid::Uuid;
|
||||
|
||||
const CODEX_OAUTH_ACCOUNT_LOCK_TTL: Duration = Duration::from_secs(180);
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum CodexOAuthAccountLockError {
|
||||
MissingIdentity,
|
||||
Contended,
|
||||
Unavailable,
|
||||
}
|
||||
|
||||
impl CodexOAuthAccountLockError {
|
||||
pub(crate) const fn status_code(self) -> http::StatusCode {
|
||||
match self {
|
||||
Self::MissingIdentity => http::StatusCode::BAD_REQUEST,
|
||||
Self::Contended => http::StatusCode::CONFLICT,
|
||||
Self::Unavailable => http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) const fn detail(self) -> &'static str {
|
||||
match self {
|
||||
Self::MissingIdentity => "Codex 账号身份字段缺失,无法安全写入授权",
|
||||
Self::Contended => "该 ChatGPT 账号正在更新授权,请稍后重试",
|
||||
Self::Unavailable => "Codex 账号授权锁暂不可用,请稍后重试",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_codex_plan_group_for_provider_oauth(
|
||||
plan_type: Option<&serde_json::Value>,
|
||||
@@ -26,6 +57,174 @@ fn normalize_provider_oauth_identity_value(value: Option<&serde_json::Value>) ->
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn normalize_provider_oauth_identity_value_from_keys(
|
||||
auth_config: &serde_json::Map<String, serde_json::Value>,
|
||||
keys: &[&str],
|
||||
) -> Option<String> {
|
||||
keys.iter()
|
||||
.find_map(|key| normalize_provider_oauth_identity_value(auth_config.get(*key)))
|
||||
}
|
||||
|
||||
fn codex_agent_identity_account_lock_key(
|
||||
provider_id: &str,
|
||||
identity_kind: &str,
|
||||
identity_parts: &[&str],
|
||||
) -> String {
|
||||
let mut digest = Sha256::new();
|
||||
digest.update(provider_id.trim().as_bytes());
|
||||
digest.update([0]);
|
||||
digest.update(identity_kind.as_bytes());
|
||||
for part in identity_parts {
|
||||
digest.update([0]);
|
||||
digest.update(part.as_bytes());
|
||||
}
|
||||
format!(
|
||||
"provider_oauth_agent_identity_account:{:x}",
|
||||
digest.finalize()
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn codex_agent_identity_account_lock_keys(
|
||||
provider_id: &str,
|
||||
auth_config: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> Vec<String> {
|
||||
let account_user_id = normalize_provider_oauth_identity_value_from_keys(
|
||||
auth_config,
|
||||
&[
|
||||
"account_user_id",
|
||||
"accountUserId",
|
||||
"chatgpt_account_user_id",
|
||||
"chatgptAccountUserId",
|
||||
],
|
||||
);
|
||||
let account_id = normalize_provider_oauth_identity_value_from_keys(
|
||||
auth_config,
|
||||
&[
|
||||
"account_id",
|
||||
"accountId",
|
||||
"chatgpt_account_id",
|
||||
"chatgptAccountId",
|
||||
],
|
||||
);
|
||||
let user_id = normalize_provider_oauth_identity_value_from_keys(
|
||||
auth_config,
|
||||
&["user_id", "userId", "chatgpt_user_id", "chatgptUserId"],
|
||||
);
|
||||
let email = normalize_provider_oauth_identity_value_from_keys(auth_config, &["email"]);
|
||||
|
||||
let mut keys = Vec::with_capacity(5);
|
||||
if let Some(account_user_id) = account_user_id.as_deref() {
|
||||
keys.push(codex_agent_identity_account_lock_key(
|
||||
provider_id,
|
||||
"account_user_id",
|
||||
&[account_user_id],
|
||||
));
|
||||
}
|
||||
if let (Some(account_id), Some(user_id)) = (account_id.as_deref(), user_id.as_deref()) {
|
||||
keys.push(codex_agent_identity_account_lock_key(
|
||||
provider_id,
|
||||
"account_id_user_id",
|
||||
&[account_id, user_id],
|
||||
));
|
||||
}
|
||||
if let (Some(account_id), Some(email)) = (account_id.as_deref(), email.as_deref()) {
|
||||
keys.push(codex_agent_identity_account_lock_key(
|
||||
provider_id,
|
||||
"account_id_email",
|
||||
&[account_id, email],
|
||||
));
|
||||
}
|
||||
if let Some(user_id) = user_id.as_deref() {
|
||||
keys.push(codex_agent_identity_account_lock_key(
|
||||
provider_id,
|
||||
"user_id",
|
||||
&[user_id],
|
||||
));
|
||||
}
|
||||
if let Some(email) = email.as_deref() {
|
||||
keys.push(codex_agent_identity_account_lock_key(
|
||||
provider_id,
|
||||
"email",
|
||||
&[email],
|
||||
));
|
||||
}
|
||||
keys.sort_unstable();
|
||||
keys.dedup();
|
||||
keys
|
||||
}
|
||||
|
||||
pub(crate) async fn acquire_codex_oauth_account_locks(
|
||||
state: &AdminAppState<'_>,
|
||||
provider_id: &str,
|
||||
auth_config: &serde_json::Map<String, serde_json::Value>,
|
||||
operation: &str,
|
||||
) -> Result<Vec<RuntimeLockLease>, CodexOAuthAccountLockError> {
|
||||
let lock_keys = codex_agent_identity_account_lock_keys(provider_id, auth_config);
|
||||
if lock_keys.is_empty() {
|
||||
return Err(CodexOAuthAccountLockError::MissingIdentity);
|
||||
}
|
||||
|
||||
let owner = format!(
|
||||
"aether-gateway-codex-oauth-{}-{}",
|
||||
operation.trim(),
|
||||
Uuid::new_v4()
|
||||
);
|
||||
let mut leases = Vec::with_capacity(lock_keys.len());
|
||||
for lock_key in lock_keys {
|
||||
match state
|
||||
.runtime_state()
|
||||
.lock_try_acquire(
|
||||
lock_key.as_str(),
|
||||
owner.as_str(),
|
||||
CODEX_OAUTH_ACCOUNT_LOCK_TTL,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(lease)) => leases.push(lease),
|
||||
Ok(None) => {
|
||||
release_codex_oauth_account_locks(state, leases).await;
|
||||
return Err(CodexOAuthAccountLockError::Contended);
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
provider_id = %provider_id,
|
||||
lock_key = %lock_key,
|
||||
operation,
|
||||
error = ?error,
|
||||
"gateway Codex OAuth account lock unavailable"
|
||||
);
|
||||
release_codex_oauth_account_locks(state, leases).await;
|
||||
return Err(CodexOAuthAccountLockError::Unavailable);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The lock is distributed, while the catalog cache is process-local. A
|
||||
// fresh read inside the lease is required to observe the previous holder.
|
||||
state.app().data.clear_provider_catalog_cache();
|
||||
Ok(leases)
|
||||
}
|
||||
|
||||
pub(crate) async fn release_codex_oauth_account_locks(
|
||||
state: &AdminAppState<'_>,
|
||||
leases: Vec<RuntimeLockLease>,
|
||||
) {
|
||||
for lease in leases.into_iter().rev() {
|
||||
match state.runtime_state().lock_release(&lease).await {
|
||||
Ok(true) => {}
|
||||
Ok(false) => tracing::warn!(
|
||||
lock_key = %lease.key,
|
||||
"gateway Codex OAuth account lock was not owned during release"
|
||||
),
|
||||
Err(error) => tracing::warn!(
|
||||
lock_key = %lease.key,
|
||||
error = ?error,
|
||||
"gateway Codex OAuth account lock release failed"
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn is_openai_provider_oauth_provider_type(value: Option<&serde_json::Value>) -> bool {
|
||||
value
|
||||
.and_then(serde_json::Value::as_str)
|
||||
@@ -55,6 +254,24 @@ fn match_codex_provider_oauth_identity(
|
||||
return None;
|
||||
}
|
||||
|
||||
let new_agent_runtime_id = normalize_provider_oauth_identity_value(
|
||||
new_auth_config
|
||||
.get("agent_runtime_id")
|
||||
.or_else(|| new_auth_config.get("agentRuntimeId")),
|
||||
);
|
||||
let existing_agent_runtime_id = normalize_provider_oauth_identity_value(
|
||||
existing_auth_config
|
||||
.get("agent_runtime_id")
|
||||
.or_else(|| existing_auth_config.get("agentRuntimeId")),
|
||||
);
|
||||
if new_agent_runtime_id
|
||||
.as_deref()
|
||||
.zip(existing_agent_runtime_id.as_deref())
|
||||
.is_some_and(|(left, right)| left == right)
|
||||
{
|
||||
return Some(true);
|
||||
}
|
||||
|
||||
let new_account_user_id =
|
||||
normalize_provider_oauth_identity_value(new_auth_config.get("account_user_id"));
|
||||
let existing_account_user_id =
|
||||
@@ -212,6 +429,11 @@ pub(crate) async fn find_duplicate_provider_oauth_key(
|
||||
let new_email = normalize_provider_oauth_identity_value(auth_config.get("email"));
|
||||
let new_user_id = normalize_provider_oauth_identity_value(auth_config.get("user_id"));
|
||||
let new_account_id = normalize_provider_oauth_identity_value(auth_config.get("account_id"));
|
||||
let new_agent_runtime_id = normalize_provider_oauth_identity_value(
|
||||
auth_config
|
||||
.get("agent_runtime_id")
|
||||
.or_else(|| auth_config.get("agentRuntimeId")),
|
||||
);
|
||||
let new_credential_fingerprint =
|
||||
normalize_provider_oauth_identity_value(auth_config.get("credential_fingerprint"));
|
||||
let new_auth_method = normalize_provider_oauth_identity_value(auth_config.get("auth_method"));
|
||||
@@ -220,11 +442,15 @@ pub(crate) async fn find_duplicate_provider_oauth_key(
|
||||
if new_email.is_none()
|
||||
&& new_user_id.is_none()
|
||||
&& new_account_id.is_none()
|
||||
&& new_agent_runtime_id.is_none()
|
||||
&& new_credential_fingerprint.is_none()
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// Duplicate checks are write admission checks. Never let a process-local
|
||||
// read-through cache hide a row committed by the previous lock holder.
|
||||
state.app().data.clear_provider_catalog_cache();
|
||||
let existing_keys = state
|
||||
.list_provider_catalog_keys_by_provider_ids(&[provider_id.to_string()])
|
||||
.await
|
||||
@@ -325,6 +551,7 @@ pub(crate) async fn find_duplicate_provider_oauth_key(
|
||||
let identifier =
|
||||
normalize_provider_oauth_identity_value(auth_config.get("account_user_id"))
|
||||
.or_else(|| normalize_provider_oauth_identity_value(auth_config.get("account_id")))
|
||||
.or_else(|| new_agent_runtime_id.clone())
|
||||
.or_else(|| {
|
||||
normalize_provider_oauth_identity_value(
|
||||
auth_config.get("credential_fingerprint"),
|
||||
@@ -345,7 +572,13 @@ pub(crate) async fn find_duplicate_provider_oauth_key(
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::match_windsurf_provider_oauth_identity;
|
||||
use super::{
|
||||
acquire_codex_oauth_account_locks, codex_agent_identity_account_lock_keys,
|
||||
match_codex_provider_oauth_identity, match_windsurf_provider_oauth_identity,
|
||||
release_codex_oauth_account_locks, CodexOAuthAccountLockError,
|
||||
};
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::AppState;
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
fn auth_config(value: Value) -> Map<String, Value> {
|
||||
@@ -371,6 +604,189 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_agent_identity_matches_runtime_without_account_metadata() {
|
||||
let new_auth_config = auth_config(json!({
|
||||
"provider_type": "codex",
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_runtime_id": "runtime-1",
|
||||
"agent_private_key": "new-private-key"
|
||||
}));
|
||||
let existing_auth_config = auth_config(json!({
|
||||
"provider_type": "codex",
|
||||
"auth_mode": "agentIdentity",
|
||||
"agentRuntimeId": "runtime-1",
|
||||
"agent_private_key": "existing-private-key"
|
||||
}));
|
||||
|
||||
assert_eq!(
|
||||
match_codex_provider_oauth_identity(&new_auth_config, &existing_auth_config),
|
||||
Some(true)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn direct_and_json_agent_identity_imports_share_account_lock_keys() {
|
||||
let direct_identity_hints = auth_config(json!({
|
||||
"provider_type": "codex",
|
||||
"account_id": "account-1",
|
||||
"account_user_id": "account-user-1",
|
||||
"user_id": "user-1",
|
||||
"email": "agent@example.com"
|
||||
}));
|
||||
let imported_auth_config = auth_config(json!({
|
||||
"provider_type": "codex",
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_runtime_id": "runtime-1",
|
||||
"accountId": "account-1",
|
||||
"chatgptAccountUserId": "account-user-1",
|
||||
"chatgptUserId": "user-1",
|
||||
"email": "agent@example.com"
|
||||
}));
|
||||
|
||||
let direct_keys =
|
||||
codex_agent_identity_account_lock_keys("provider-codex", &direct_identity_hints);
|
||||
let imported_keys =
|
||||
codex_agent_identity_account_lock_keys("provider-codex", &imported_auth_config);
|
||||
let shared_keys = direct_keys
|
||||
.iter()
|
||||
.filter(|key| imported_keys.contains(key))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(shared_keys.len(), 5);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ordinary_codex_oauth_and_agent_identity_share_runtime_account_locks() {
|
||||
let app = AppState::new().expect("app state should build");
|
||||
let state = AdminAppState::new(&app);
|
||||
let ordinary = auth_config(json!({
|
||||
"provider_type": "codex",
|
||||
"account_id": "account-1",
|
||||
"account_user_id": "account-user-1",
|
||||
"user_id": "user-1",
|
||||
"email": "agent@example.com"
|
||||
}));
|
||||
let agent = auth_config(json!({
|
||||
"provider_type": "codex",
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_runtime_id": "runtime-1",
|
||||
"account_id": "account-1",
|
||||
"account_user_id": "account-user-1",
|
||||
"user_id": "user-1",
|
||||
"email": "agent@example.com"
|
||||
}));
|
||||
|
||||
let first =
|
||||
acquire_codex_oauth_account_locks(&state, "provider-codex", &ordinary, "ordinary-test")
|
||||
.await
|
||||
.expect("ordinary OAuth lock should acquire");
|
||||
let second =
|
||||
acquire_codex_oauth_account_locks(&state, "provider-codex", &agent, "agent-test")
|
||||
.await
|
||||
.expect_err("Agent Identity must contend on the same account locks");
|
||||
assert_eq!(second, CodexOAuthAccountLockError::Contended);
|
||||
|
||||
release_codex_oauth_account_locks(&state, first).await;
|
||||
let third =
|
||||
acquire_codex_oauth_account_locks(&state, "provider-codex", &agent, "agent-retry-test")
|
||||
.await
|
||||
.expect("account locks should be reusable after release");
|
||||
release_codex_oauth_account_locks(&state, third).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn codex_oauth_account_lock_rejects_identity_free_config() {
|
||||
let app = AppState::new().expect("app state should build");
|
||||
let state = AdminAppState::new(&app);
|
||||
let config = auth_config(json!({"provider_type": "codex"}));
|
||||
|
||||
let error = acquire_codex_oauth_account_locks(
|
||||
&state,
|
||||
"provider-codex",
|
||||
&config,
|
||||
"missing-identity-test",
|
||||
)
|
||||
.await
|
||||
.expect_err("identity-free Codex writes must not proceed unlocked");
|
||||
|
||||
assert_eq!(error, CodexOAuthAccountLockError::MissingIdentity);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn codex_oauth_account_lock_releases_partial_acquisition() {
|
||||
let app = AppState::new().expect("app state should build");
|
||||
let state = AdminAppState::new(&app);
|
||||
let config = auth_config(json!({
|
||||
"provider_type": "codex",
|
||||
"account_id": "account-partial",
|
||||
"account_user_id": "account-user-partial",
|
||||
"user_id": "user-partial",
|
||||
"email": "partial@example.com"
|
||||
}));
|
||||
let keys = codex_agent_identity_account_lock_keys("provider-codex", &config);
|
||||
let held_key = keys.last().expect("account locks should not be empty");
|
||||
let held = state
|
||||
.runtime_state()
|
||||
.lock_try_acquire(held_key, "other-owner", std::time::Duration::from_secs(30))
|
||||
.await
|
||||
.expect("runtime lock should be available")
|
||||
.expect("last account lock should acquire");
|
||||
|
||||
let error =
|
||||
acquire_codex_oauth_account_locks(&state, "provider-codex", &config, "partial-test")
|
||||
.await
|
||||
.expect_err("held final lock should cause contention");
|
||||
assert_eq!(error, CodexOAuthAccountLockError::Contended);
|
||||
|
||||
let first_key = keys.first().expect("account locks should not be empty");
|
||||
let first = state
|
||||
.runtime_state()
|
||||
.lock_try_acquire(
|
||||
first_key,
|
||||
"verification-owner",
|
||||
std::time::Duration::from_secs(30),
|
||||
)
|
||||
.await
|
||||
.expect("runtime lock should be available")
|
||||
.expect("partially acquired account lock should have been released");
|
||||
assert!(state
|
||||
.runtime_state()
|
||||
.lock_release(&first)
|
||||
.await
|
||||
.expect("verification lock should release"));
|
||||
assert!(state
|
||||
.runtime_state()
|
||||
.lock_release(&held)
|
||||
.await
|
||||
.expect("held lock should release"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_identity_account_locks_cover_generic_user_and_email_deduplication() {
|
||||
let first = auth_config(json!({
|
||||
"provider_type": "codex",
|
||||
"agent_runtime_id": "runtime-1",
|
||||
"user_id": "user-1",
|
||||
"email": "agent@example.com"
|
||||
}));
|
||||
let second = auth_config(json!({
|
||||
"provider_type": "codex",
|
||||
"agent_runtime_id": "runtime-2",
|
||||
"user_id": "user-1",
|
||||
"email": "agent@example.com"
|
||||
}));
|
||||
|
||||
let first_keys = codex_agent_identity_account_lock_keys("provider-codex", &first);
|
||||
let second_keys = codex_agent_identity_account_lock_keys("provider-codex", &second);
|
||||
let shared_keys = first_keys
|
||||
.iter()
|
||||
.filter(|key| second_keys.contains(key))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(shared_keys.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn windsurf_identity_rejects_different_account_id() {
|
||||
let new_auth_config = auth_config(json!({
|
||||
|
||||
@@ -19,10 +19,9 @@ use self::plan::{
|
||||
};
|
||||
use super::shared::{
|
||||
build_quota_snapshot_payload, extract_execution_error_message,
|
||||
oauth_refresh_auto_removed_result, persist_provider_quota_refresh_state,
|
||||
provider_auto_remove_banned_keys, provider_auto_remove_quota_exhausted_keys,
|
||||
quota_key_auto_removed, quota_refresh_success_invalid_state,
|
||||
should_auto_remove_structured_reason, ProviderQuotaExecutionOutcome,
|
||||
oauth_refresh_auto_removed_result, persist_fenced_provider_quota_refresh_state,
|
||||
persist_provider_quota_refresh_state, quota_key_auto_removed,
|
||||
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
|
||||
};
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::provider_key_auth::provider_key_is_oauth_managed;
|
||||
@@ -399,9 +398,6 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
proxy_override: Option<ProxySnapshot>,
|
||||
) -> Result<Option<serde_json::Value>, GatewayError> {
|
||||
let auto_remove_abnormal_keys = provider_auto_remove_banned_keys(provider.config.as_ref());
|
||||
let auto_remove_quota_exhausted_keys =
|
||||
provider_auto_remove_quota_exhausted_keys(provider.config.as_ref());
|
||||
let mut results = Vec::new();
|
||||
let mut success_count = 0usize;
|
||||
let mut failed_count = 0usize;
|
||||
@@ -429,8 +425,29 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let is_oauth_managed = provider_key_is_oauth_managed(&key, provider.provider_type.as_str());
|
||||
let quota_auth_config_fence = if is_oauth_managed {
|
||||
match state
|
||||
.app()
|
||||
.capture_provider_transport_auth_config_fence(&transport)
|
||||
.await?
|
||||
{
|
||||
Some(ciphertext) => Some(ciphertext),
|
||||
None => {
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "OAuth credential changed before quota refresh",
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let resolved_oauth_auth = if is_oauth_managed {
|
||||
state.resolve_local_oauth_header_auth(&transport).await?
|
||||
} else {
|
||||
@@ -647,17 +664,27 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
|
||||
}
|
||||
}
|
||||
|
||||
let auto_remove_candidate = auto_remove_abnormal_keys
|
||||
&& should_auto_remove_structured_reason(oauth_invalid_reason.as_deref());
|
||||
let persisted = persist_provider_quota_refresh_state(
|
||||
state,
|
||||
&key.id,
|
||||
metadata_update.as_ref(),
|
||||
oauth_invalid_at_unix_secs,
|
||||
oauth_invalid_reason.clone(),
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
let persisted = if let Some(expected_auth_config) = quota_auth_config_fence.as_deref() {
|
||||
persist_fenced_provider_quota_refresh_state(
|
||||
state,
|
||||
&key.id,
|
||||
expected_auth_config,
|
||||
metadata_update.as_ref(),
|
||||
oauth_invalid_at_unix_secs,
|
||||
oauth_invalid_reason.clone(),
|
||||
)
|
||||
.await?
|
||||
} else {
|
||||
persist_provider_quota_refresh_state(
|
||||
state,
|
||||
&key.id,
|
||||
metadata_update.as_ref(),
|
||||
oauth_invalid_at_unix_secs,
|
||||
oauth_invalid_reason.clone(),
|
||||
None,
|
||||
)
|
||||
.await?
|
||||
};
|
||||
if !persisted {
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
@@ -668,32 +695,15 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
let auto_removed_hard_banned = if auto_remove_candidate {
|
||||
state
|
||||
.cleanup_provider_catalog_key_if_current(provider, &key.id, |latest_key| {
|
||||
should_auto_remove_structured_reason(latest_key.oauth_invalid_reason.as_deref())
|
||||
})
|
||||
.await?
|
||||
} else {
|
||||
false
|
||||
};
|
||||
// Codex quota responses never auto-delete keys. Without a repository
|
||||
// conditional delete, any read-then-delete sequence could remove a
|
||||
// replacement Agent Identity installed while the response was in flight.
|
||||
let auto_removed_hard_banned = false;
|
||||
if auto_removed_hard_banned {
|
||||
auto_removed_count += 1;
|
||||
auto_removed_hard_banned_count += 1;
|
||||
}
|
||||
let auto_removed_quota_exhausted =
|
||||
if !auto_removed_hard_banned && auto_remove_quota_exhausted_keys {
|
||||
state
|
||||
.cleanup_provider_catalog_key_if_current(provider, &key.id, |latest_key| {
|
||||
aether_admin::provider::pool::admin_pool_key_account_quota_exhausted(
|
||||
latest_key,
|
||||
provider.provider_type.as_str(),
|
||||
)
|
||||
})
|
||||
.await?
|
||||
} else {
|
||||
false
|
||||
};
|
||||
let auto_removed_quota_exhausted = false;
|
||||
if auto_removed_quota_exhausted {
|
||||
auto_removed_count += 1;
|
||||
status = "quota_exhausted".to_string();
|
||||
|
||||
@@ -14,8 +14,9 @@ use aether_contracts::{
|
||||
ResolvedTransportProfile, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
|
||||
ProviderCatalogKeyStatusSnapshotUpdate, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey,
|
||||
};
|
||||
use aether_provider_pool::{ProviderPoolQuotaRequestSpec, ProviderPoolService};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
@@ -200,6 +201,10 @@ pub(super) fn extract_execution_error_message(result: &ExecutionResult) -> Optio
|
||||
admin_provider_quota_pure::extract_execution_error_message(result)
|
||||
}
|
||||
|
||||
fn extract_execution_error_detail(result: &ExecutionResult) -> Option<String> {
|
||||
admin_provider_quota_pure::extract_execution_error_detail(result)
|
||||
}
|
||||
|
||||
pub(super) fn quota_refresh_success_invalid_state(
|
||||
key: &StoredProviderCatalogKey,
|
||||
) -> (Option<u64>, Option<String>) {
|
||||
@@ -301,6 +306,86 @@ pub(crate) async fn persist_provider_quota_refresh_state(
|
||||
.await
|
||||
}
|
||||
|
||||
/// Persist a Codex Agent Identity quota response only when the exact encrypted
|
||||
/// auth_config used for the request is still installed. Metadata, OAuth state,
|
||||
/// and their status projection share one repository CAS so a replacement cannot
|
||||
/// receive any portion of an older response.
|
||||
pub(crate) async fn persist_fenced_provider_quota_refresh_state(
|
||||
state: &AdminAppState<'_>,
|
||||
key_id: &str,
|
||||
expected_encrypted_auth_config: &str,
|
||||
metadata_update: Option<&serde_json::Value>,
|
||||
oauth_invalid_at_unix_secs: Option<u64>,
|
||||
oauth_invalid_reason: Option<String>,
|
||||
) -> Result<bool, GatewayError> {
|
||||
let expected_encrypted_auth_config = expected_encrypted_auth_config.trim();
|
||||
if expected_encrypted_auth_config.is_empty() {
|
||||
return Ok(false);
|
||||
}
|
||||
if metadata_update.is_some_and(|value| !value.is_object()) {
|
||||
return Err(GatewayError::Internal(
|
||||
"fenced quota metadata update must be an object".to_string(),
|
||||
));
|
||||
}
|
||||
let Some(mut latest_key) = state
|
||||
.read_provider_catalog_keys_by_ids(&[key_id.to_string()])
|
||||
.await?
|
||||
.into_iter()
|
||||
.next()
|
||||
else {
|
||||
return Ok(false);
|
||||
};
|
||||
if latest_key.encrypted_auth_config.as_deref() != Some(expected_encrypted_auth_config) {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let quota_snapshot_provider_type =
|
||||
metadata_update.and_then(aether_provider_pool::provider_pool_quota_metadata_provider_type);
|
||||
if let Some(metadata_update) = metadata_update {
|
||||
latest_key.upstream_metadata = Some(merge_upstream_metadata(
|
||||
latest_key.upstream_metadata.as_ref(),
|
||||
metadata_update,
|
||||
));
|
||||
}
|
||||
latest_key.oauth_invalid_at_unix_secs = oauth_invalid_at_unix_secs;
|
||||
latest_key.oauth_invalid_reason = oauth_invalid_reason;
|
||||
if let Some(provider_type) = quota_snapshot_provider_type.as_deref() {
|
||||
latest_key.status_snapshot = sync_provider_key_quota_status_snapshot(
|
||||
latest_key.status_snapshot.as_ref(),
|
||||
provider_type,
|
||||
latest_key.upstream_metadata.as_ref(),
|
||||
"refresh_api",
|
||||
);
|
||||
}
|
||||
latest_key.status_snapshot =
|
||||
sync_provider_key_oauth_status_snapshot(latest_key.status_snapshot.as_ref(), &latest_key);
|
||||
latest_key.updated_at_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs());
|
||||
|
||||
state
|
||||
.app()
|
||||
.compare_and_update_provider_catalog_key_oauth_runtime_state(
|
||||
&ProviderCatalogKeyOAuthRuntimeStateCasUpdate {
|
||||
key_id: key_id.to_string(),
|
||||
expected_encrypted_auth_config: Some(expected_encrypted_auth_config.to_string()),
|
||||
encrypted_auth_config: expected_encrypted_auth_config.to_string(),
|
||||
encrypted_api_key_update: None,
|
||||
expires_at_unix_secs_update: None,
|
||||
oauth_invalid_at_unix_secs: latest_key.oauth_invalid_at_unix_secs,
|
||||
oauth_invalid_reason: latest_key.oauth_invalid_reason.clone(),
|
||||
upstream_metadata_patch: metadata_update.cloned(),
|
||||
status_snapshot_patch: provider_quota_refresh_status_patch(
|
||||
latest_key.status_snapshot.as_ref(),
|
||||
),
|
||||
reset_error_count: false,
|
||||
updated_at_unix_secs: latest_key.updated_at_unix_secs,
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn persist_provider_quota_refresh_state_after_read<F>(
|
||||
state: &AdminAppState<'_>,
|
||||
key_id: &str,
|
||||
@@ -452,7 +537,7 @@ pub(super) async fn execute_provider_quota_plan(
|
||||
if !crate::provider_transport::is_codex_agent_identity_transport(transport)
|
||||
|| !crate::provider_transport::is_codex_agent_identity_invalid_task_response(
|
||||
result.status_code,
|
||||
extract_execution_error_message(&result).as_deref(),
|
||||
extract_execution_error_detail(&result).as_deref(),
|
||||
)
|
||||
{
|
||||
return Ok(ProviderQuotaExecutionOutcome::Response(result));
|
||||
|
||||
@@ -5,7 +5,7 @@ use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::handlers::admin::shared::{provider_key_status_snapshot_payload, unix_secs_to_rfc3339};
|
||||
use crate::provider_key_auth::{
|
||||
provider_key_auth_config_is_agent_identity, provider_key_auth_config_uses_header_authorization,
|
||||
provider_key_auth_semantics, provider_key_can_refresh_oauth,
|
||||
provider_key_auth_semantics, provider_key_can_export_oauth, provider_key_can_refresh_oauth,
|
||||
provider_key_effective_api_formats,
|
||||
};
|
||||
use aether_admin::provider::pool as admin_provider_pool_pure;
|
||||
@@ -1228,12 +1228,17 @@ pub(super) fn build_admin_pool_key_payload(
|
||||
"can_refresh_oauth".to_string(),
|
||||
json!(provider_key_can_refresh_oauth(
|
||||
auth_semantics,
|
||||
provider_type,
|
||||
auth_config.as_ref()
|
||||
)),
|
||||
);
|
||||
payload.insert(
|
||||
"can_export_oauth".to_string(),
|
||||
json!(auth_semantics.can_export_oauth()),
|
||||
json!(provider_key_can_export_oauth(
|
||||
auth_semantics,
|
||||
provider_type,
|
||||
auth_config.as_ref()
|
||||
)),
|
||||
);
|
||||
payload.insert(
|
||||
"can_edit_oauth".to_string(),
|
||||
|
||||
+11
-3
@@ -8,7 +8,7 @@ use super::{
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::provider_key_auth::{
|
||||
provider_key_auth_config_is_agent_identity, provider_key_auth_config_uses_header_authorization,
|
||||
provider_key_auth_semantics, provider_key_can_refresh_oauth,
|
||||
provider_key_auth_semantics, provider_key_can_export_oauth, provider_key_can_refresh_oauth,
|
||||
};
|
||||
use crate::GatewayError;
|
||||
use aether_admin::provider::pool as admin_provider_pool_pure;
|
||||
@@ -156,8 +156,16 @@ pub(super) async fn build_admin_pool_resolve_selection_response(
|
||||
&provider_type,
|
||||
auth_config.as_ref(),
|
||||
),
|
||||
"can_refresh_oauth": provider_key_can_refresh_oauth(auth_semantics, auth_config.as_ref()),
|
||||
"can_export_oauth": auth_semantics.can_export_oauth(),
|
||||
"can_refresh_oauth": provider_key_can_refresh_oauth(
|
||||
auth_semantics,
|
||||
&provider_type,
|
||||
auth_config.as_ref(),
|
||||
),
|
||||
"can_export_oauth": provider_key_can_export_oauth(
|
||||
auth_semantics,
|
||||
&provider_type,
|
||||
auth_config.as_ref(),
|
||||
),
|
||||
"can_edit_oauth": auth_semantics.can_edit_oauth(),
|
||||
"oauth_header_auth": auth_semantics.oauth_managed()
|
||||
&& provider_key_auth_config_uses_header_authorization(auth_config.as_ref()),
|
||||
|
||||
@@ -20,6 +20,8 @@ pub(crate) use self::endpoint_keys::{
|
||||
admin_reset_cycle_stats_key_id, admin_reveal_key_id, admin_update_key_id,
|
||||
};
|
||||
pub(crate) use self::oauth::{
|
||||
admin_provider_oauth_agent_identity_import_task_path,
|
||||
admin_provider_oauth_agent_identity_import_task_provider_id,
|
||||
admin_provider_oauth_batch_import_provider_id, admin_provider_oauth_batch_import_task_path,
|
||||
admin_provider_oauth_batch_import_task_provider_id, admin_provider_oauth_complete_key_id,
|
||||
admin_provider_oauth_complete_provider_id, admin_provider_oauth_device_authorize_provider_id,
|
||||
|
||||
@@ -44,6 +44,12 @@ pub(crate) fn admin_provider_oauth_batch_import_task_provider_id(
|
||||
provider_oauth_provider_id_for_suffix(request_path, "/batch-import/tasks")
|
||||
}
|
||||
|
||||
pub(crate) fn admin_provider_oauth_agent_identity_import_task_provider_id(
|
||||
request_path: &str,
|
||||
) -> Option<String> {
|
||||
provider_oauth_provider_id_for_suffix(request_path, "/agent-identity-import/tasks")
|
||||
}
|
||||
|
||||
pub(crate) fn admin_provider_oauth_batch_import_task_path(
|
||||
request_path: &str,
|
||||
) -> Option<(String, String)> {
|
||||
@@ -62,6 +68,25 @@ pub(crate) fn admin_provider_oauth_batch_import_task_path(
|
||||
Some((provider_id.to_string(), task_path.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) fn admin_provider_oauth_agent_identity_import_task_path(
|
||||
request_path: &str,
|
||||
) -> Option<(String, String)> {
|
||||
let suffix = request_path
|
||||
.strip_prefix("/api/admin/provider-oauth/providers/")?
|
||||
.strip_suffix("/")
|
||||
.unwrap_or(request_path.strip_prefix("/api/admin/provider-oauth/providers/")?);
|
||||
let (provider_id, task_path) = suffix.split_once("/agent-identity-import/tasks/")?;
|
||||
if provider_id.is_empty()
|
||||
|| provider_id.contains('/')
|
||||
|| task_path.is_empty()
|
||||
|| task_path.contains('/')
|
||||
|| !task_path.starts_with("agent-identity-")
|
||||
{
|
||||
return None;
|
||||
}
|
||||
Some((provider_id.to_string(), task_path.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) fn admin_provider_oauth_device_authorize_provider_id(
|
||||
request_path: &str,
|
||||
) -> Option<String> {
|
||||
@@ -79,3 +104,39 @@ fn provider_oauth_provider_id_for_suffix(request_path: &str, suffix: &str) -> Op
|
||||
.filter(|provider_id| !provider_id.is_empty() && !provider_id.contains('/'))
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
admin_provider_oauth_agent_identity_import_task_path,
|
||||
admin_provider_oauth_agent_identity_import_task_provider_id,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn parses_dedicated_agent_identity_import_task_paths() {
|
||||
assert_eq!(
|
||||
admin_provider_oauth_agent_identity_import_task_provider_id(
|
||||
"/api/admin/provider-oauth/providers/provider-codex/agent-identity-import/tasks",
|
||||
)
|
||||
.as_deref(),
|
||||
Some("provider-codex")
|
||||
);
|
||||
assert_eq!(
|
||||
admin_provider_oauth_agent_identity_import_task_path(
|
||||
"/api/admin/provider-oauth/providers/provider-codex/agent-identity-import/tasks/agent-identity-task-1",
|
||||
),
|
||||
Some((
|
||||
"provider-codex".to_string(),
|
||||
"agent-identity-task-1".to_string(),
|
||||
))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dedicated_status_path_rejects_generic_batch_task_ids() {
|
||||
assert!(admin_provider_oauth_agent_identity_import_task_path(
|
||||
"/api/admin/provider-oauth/providers/provider-codex/agent-identity-import/tasks/task-1",
|
||||
)
|
||||
.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -52,17 +52,14 @@ pub(crate) async fn build_admin_create_provider_key_record(
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.cloned();
|
||||
|
||||
if auth_type == "oauth"
|
||||
&& provider.provider_type.trim().eq_ignore_ascii_case("codex")
|
||||
&& auth_config
|
||||
.as_ref()
|
||||
.is_some_and(aether_provider_transport::is_codex_agent_identity_auth_config_value)
|
||||
if auth_config
|
||||
.as_ref()
|
||||
.is_some_and(aether_provider_transport::is_codex_agent_identity_auth_config_value)
|
||||
{
|
||||
aether_provider_transport::validate_codex_agent_identity_auth_config(
|
||||
auth_config
|
||||
.as_ref()
|
||||
.expect("Agent Identity auth_config was checked"),
|
||||
)?;
|
||||
return Err(
|
||||
"Agent Identity 凭据必须通过专属创建或导入接口管理,不能通过通用 Key 接口写入"
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
match auth_type.as_str() {
|
||||
|
||||
@@ -78,17 +78,14 @@ pub(crate) fn build_admin_update_provider_key_record_with_existing_keys(
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.cloned();
|
||||
|
||||
if target_auth_type == "oauth"
|
||||
&& provider.provider_type.trim().eq_ignore_ascii_case("codex")
|
||||
&& auth_config
|
||||
.as_ref()
|
||||
.is_some_and(aether_provider_transport::is_codex_agent_identity_auth_config_value)
|
||||
if auth_config
|
||||
.as_ref()
|
||||
.is_some_and(aether_provider_transport::is_codex_agent_identity_auth_config_value)
|
||||
{
|
||||
aether_provider_transport::validate_codex_agent_identity_auth_config(
|
||||
auth_config
|
||||
.as_ref()
|
||||
.expect("Agent Identity auth_config was checked"),
|
||||
)?;
|
||||
return Err(
|
||||
"Agent Identity 凭据必须通过专属创建或导入接口管理,不能通过通用 Key 接口写入"
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
match target_auth_type.as_str() {
|
||||
|
||||
@@ -31,6 +31,16 @@ pub(crate) fn build_admin_reveal_key_payload(
|
||||
key: &StoredProviderCatalogKey,
|
||||
) -> Result<serde_json::Value, String> {
|
||||
let parsed_auth_config = state.parse_catalog_auth_config_json(key);
|
||||
if parsed_auth_config.as_ref().is_some_and(|auth_config| {
|
||||
aether_provider_transport::is_codex_agent_identity_auth_config_value(
|
||||
&serde_json::Value::Object(auth_config.clone()),
|
||||
)
|
||||
}) {
|
||||
return Err(
|
||||
"Agent Identity 凭据不能通过通用 Key 查看接口读取,请使用专属 provider-oauth 管理面"
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
let provider_type = reveal_provider_type_from_auth_config(parsed_auth_config.as_ref());
|
||||
let auth_semantics = provider_key_auth_semantics(key, provider_type.as_str());
|
||||
let auth_type = if auth_semantics.oauth_managed() {
|
||||
@@ -183,6 +193,15 @@ pub(crate) async fn build_admin_export_key_payload(
|
||||
.and_then(|value| value.as_object().cloned())
|
||||
.ok_or_else(|| "无法解密认证配置".to_string())?;
|
||||
|
||||
if aether_provider_transport::is_codex_agent_identity_auth_config_value(
|
||||
&serde_json::Value::Object(auth_config.clone()),
|
||||
) {
|
||||
return Err(
|
||||
"Agent Identity 凭据不能通过通用 Key 导出接口导出,请使用专属 provider-oauth 管理面"
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
let provider_type_from_config = auth_config
|
||||
.get("provider_type")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
|
||||
@@ -69,11 +69,16 @@ impl<'a> AdminAppState<'a> {
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn masked_catalog_api_key(
|
||||
pub(crate) fn masked_catalog_api_key_for_provider(
|
||||
&self,
|
||||
key: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
) -> String {
|
||||
crate::handlers::admin::shared::masked_catalog_api_key(self.app, key)
|
||||
crate::handlers::admin::shared::masked_catalog_api_key_for_provider(
|
||||
self.app,
|
||||
key,
|
||||
provider_type,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) async fn build_admin_provider_keys_payload(
|
||||
|
||||
@@ -227,6 +227,7 @@ impl<'a> AdminAppState<'a> {
|
||||
.compare_and_update_provider_catalog_key_adaptive_state(
|
||||
&ProviderCatalogKeyAdaptiveStateUpdate {
|
||||
key_id: key_id.to_string(),
|
||||
expected_encrypted_auth_config: None,
|
||||
expected,
|
||||
next,
|
||||
status_snapshot_patch: serde_json::json!({
|
||||
@@ -257,6 +258,33 @@ impl<'a> AdminAppState<'a> {
|
||||
) -> Result<
|
||||
Option<aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey>,
|
||||
GatewayError,
|
||||
> {
|
||||
self.reset_provider_catalog_key_recovery_state_inner(key_id, None)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn reset_provider_catalog_key_recovery_state_fenced(
|
||||
&self,
|
||||
key_id: &str,
|
||||
expected_encrypted_auth_config: &str,
|
||||
) -> Result<
|
||||
Option<aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey>,
|
||||
GatewayError,
|
||||
> {
|
||||
self.reset_provider_catalog_key_recovery_state_inner(
|
||||
key_id,
|
||||
Some(expected_encrypted_auth_config),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn reset_provider_catalog_key_recovery_state_inner(
|
||||
&self,
|
||||
key_id: &str,
|
||||
expected_auth_config: Option<&str>,
|
||||
) -> Result<
|
||||
Option<aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey>,
|
||||
GatewayError,
|
||||
> {
|
||||
use aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyHealthStateUpdate;
|
||||
|
||||
@@ -271,6 +299,11 @@ impl<'a> AdminAppState<'a> {
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
if expected_auth_config
|
||||
.is_some_and(|expected| current.encrypted_auth_config.as_deref() != Some(expected))
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
if current.health_by_format.as_ref() == Some(&empty)
|
||||
&& current.circuit_breaker_by_format.as_ref() == Some(&empty)
|
||||
{
|
||||
@@ -282,6 +315,7 @@ impl<'a> AdminAppState<'a> {
|
||||
.compare_and_update_provider_catalog_key_health_state(
|
||||
&ProviderCatalogKeyHealthStateUpdate {
|
||||
key_id: key_id.to_string(),
|
||||
expected_encrypted_auth_config: expected_auth_config.map(ToOwned::to_owned),
|
||||
expected_health_by_format: current.health_by_format,
|
||||
expected_circuit_breaker_by_format: current.circuit_breaker_by_format,
|
||||
health_by_format: Some(empty.clone()),
|
||||
@@ -299,15 +333,24 @@ impl<'a> AdminAppState<'a> {
|
||||
"provider key {key_id} health state changed repeatedly while resetting OAuth recovery state"
|
||||
)));
|
||||
}
|
||||
if !self.reset_provider_catalog_key_error_count(key_id).await? {
|
||||
if expected_auth_config.is_none()
|
||||
&& !self.reset_provider_catalog_key_error_count(key_id).await?
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(self
|
||||
let current = self
|
||||
.read_provider_catalog_keys_by_ids(&[key_id.to_string()])
|
||||
.await?
|
||||
.into_iter()
|
||||
.next())
|
||||
.next();
|
||||
if current.as_ref().is_some_and(|key| {
|
||||
expected_auth_config
|
||||
.is_some_and(|expected| key.encrypted_auth_config.as_deref() != Some(expected))
|
||||
}) {
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(current)
|
||||
}
|
||||
|
||||
pub(crate) async fn update_provider_catalog_key_status_snapshot(
|
||||
|
||||
@@ -90,6 +90,7 @@ impl<'a> AdminAppState<'a> {
|
||||
provider_id: &str,
|
||||
provider_type: &str,
|
||||
pkce_verifier: Option<&str>,
|
||||
expected_encrypted_auth_config: Option<&str>,
|
||||
) -> Result<String, GatewayError> {
|
||||
let nonce = aether_admin::provider::state::generate_provider_oauth_nonce();
|
||||
let payload = json!({
|
||||
@@ -98,6 +99,7 @@ impl<'a> AdminAppState<'a> {
|
||||
"provider_id": provider_id,
|
||||
"provider_type": provider_type,
|
||||
"pkce_verifier": pkce_verifier,
|
||||
"expected_encrypted_auth_config": expected_encrypted_auth_config,
|
||||
"created_at": aether_admin::provider::state::current_unix_secs(),
|
||||
});
|
||||
let key = provider_oauth_state_storage_key(&nonce);
|
||||
|
||||
@@ -16,6 +16,10 @@ use serde_json::{json, Map, Value};
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
fn provider_skips_automatic_key_cleanup(provider: &StoredProviderCatalogProvider) -> bool {
|
||||
provider.provider_type.trim().eq_ignore_ascii_case("codex")
|
||||
}
|
||||
|
||||
impl<'a> AdminAppState<'a> {
|
||||
pub(crate) async fn clear_admin_provider_pool_cooldown(&self, provider_id: &str, key_id: &str) {
|
||||
crate::handlers::admin::provider::pool::runtime::clear_admin_provider_pool_cooldown(
|
||||
@@ -368,6 +372,13 @@ impl<'a> AdminAppState<'a> {
|
||||
) -> Result<usize, GatewayError> {
|
||||
use aether_admin::provider::pool as admin_provider_pool_pure;
|
||||
|
||||
// Codex OAuth credentials can be replaced by a long-lived Agent
|
||||
// Identity under the same key id. Until deletes support an auth_config
|
||||
// CAS, automatic cleanup must retain every Codex key.
|
||||
if provider_skips_automatic_key_cleanup(provider) {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let banned_keys = self
|
||||
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
|
||||
.await?
|
||||
@@ -408,6 +419,10 @@ impl<'a> AdminAppState<'a> {
|
||||
) -> Result<usize, GatewayError> {
|
||||
use aether_admin::provider::pool as admin_provider_pool_pure;
|
||||
|
||||
if provider_skips_automatic_key_cleanup(provider) {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let keys = self
|
||||
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
|
||||
.await?;
|
||||
@@ -849,3 +864,26 @@ impl<'a> AdminAppState<'a> {
|
||||
.into_response())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod automatic_cleanup_tests {
|
||||
use super::provider_skips_automatic_key_cleanup;
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
|
||||
|
||||
fn provider(provider_type: &str) -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
format!("provider-{provider_type}"),
|
||||
provider_type.to_string(),
|
||||
None,
|
||||
provider_type.to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_automatic_cleanup_is_disabled_for_replaceable_agent_credentials() {
|
||||
assert!(provider_skips_automatic_key_cleanup(&provider("codex")));
|
||||
assert!(provider_skips_automatic_key_cleanup(&provider("CoDeX")));
|
||||
assert!(!provider_skips_automatic_key_cleanup(&provider("kiro")));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -153,6 +153,7 @@ impl<'a> AdminAppState<'a> {
|
||||
.compare_and_update_provider_catalog_key_adaptive_state(
|
||||
&ProviderCatalogKeyAdaptiveStateUpdate {
|
||||
key_id: key.id.clone(),
|
||||
expected_encrypted_auth_config: None,
|
||||
expected,
|
||||
next,
|
||||
status_snapshot_patch: json!({
|
||||
|
||||
@@ -11,10 +11,10 @@ pub(crate) use crate::handlers::shared::{
|
||||
attach_admin_audit_response, build_admin_provider_key_response,
|
||||
decrypt_catalog_secret_with_fallbacks, default_provider_key_status_snapshot,
|
||||
effective_catalog_encryption_key, encrypt_catalog_secret_with_fallbacks, json_string_list,
|
||||
masked_catalog_api_key, normalize_json_array, normalize_json_object, normalize_string_list,
|
||||
parse_catalog_auth_config_json, provider_catalog_key_supports_format,
|
||||
provider_key_health_summary, provider_key_health_summary_at,
|
||||
provider_key_status_snapshot_payload, query_param_bool, query_param_optional_bool,
|
||||
query_param_value, take_secret_prefix, take_secret_suffix, unix_secs_to_rfc3339,
|
||||
OFFICIAL_EXTERNAL_MODEL_PROVIDERS,
|
||||
masked_catalog_api_key, masked_catalog_api_key_for_provider, normalize_json_array,
|
||||
normalize_json_object, normalize_string_list, parse_catalog_auth_config_json,
|
||||
provider_catalog_key_supports_format, provider_key_health_summary,
|
||||
provider_key_health_summary_at, provider_key_status_snapshot_payload, query_param_bool,
|
||||
query_param_optional_bool, query_param_value, take_secret_prefix, take_secret_suffix,
|
||||
unix_secs_to_rfc3339, OFFICIAL_EXTERNAL_MODEL_PROVIDERS,
|
||||
};
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
use super::enabled_key_capability_short_names;
|
||||
use crate::handlers::shared::{parse_catalog_auth_config_json, unix_secs_to_rfc3339};
|
||||
use crate::provider_key_auth::{
|
||||
provider_key_auth_config_uses_header_authorization, provider_key_effective_api_formats,
|
||||
provider_key_auth_config_is_agent_identity, provider_key_auth_config_uses_header_authorization,
|
||||
provider_key_effective_api_formats,
|
||||
};
|
||||
use crate::AppState;
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
@@ -10,13 +11,18 @@ use serde_json::json;
|
||||
use std::collections::{BTreeMap, HashMap};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
fn grouped_key_masked_label(state: &AppState, key: &StoredProviderCatalogKey) -> &'static str {
|
||||
fn grouped_key_masked_label(
|
||||
state: &AppState,
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
) -> &'static str {
|
||||
match key.auth_type.trim() {
|
||||
"service_account" | "vertex_ai" => "[Service Account]",
|
||||
"oauth" => {
|
||||
if provider_key_auth_config_uses_header_authorization(
|
||||
parse_catalog_auth_config_json(state, key).as_ref(),
|
||||
) {
|
||||
let auth_config = parse_catalog_auth_config_json(state, key);
|
||||
if provider_key_auth_config_is_agent_identity(provider_type, auth_config.as_ref()) {
|
||||
"[Agent Identity]"
|
||||
} else if provider_key_auth_config_uses_header_authorization(auth_config.as_ref()) {
|
||||
"[OAuth Header]"
|
||||
} else {
|
||||
"[OAuth Token]"
|
||||
@@ -161,7 +167,7 @@ pub(crate) async fn build_admin_keys_grouped_by_format_payload(
|
||||
"provider_id": key.provider_id,
|
||||
"name": key.name,
|
||||
"auth_type": key.auth_type,
|
||||
"api_key_masked": grouped_key_masked_label(state, &key),
|
||||
"api_key_masked": grouped_key_masked_label(state, &key, provider_type),
|
||||
"internal_priority": key.internal_priority,
|
||||
"global_priority_by_format": key.global_priority_by_format,
|
||||
"rate_multipliers": key.rate_multipliers,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use crate::handlers::shared::{json_string_list, unix_secs_to_rfc3339};
|
||||
use crate::provider_key_auth::{
|
||||
provider_key_auth_config_is_agent_identity, provider_key_auth_config_uses_header_authorization,
|
||||
provider_key_auth_semantics, provider_key_can_refresh_oauth,
|
||||
provider_key_auth_semantics, provider_key_can_export_oauth, provider_key_can_refresh_oauth,
|
||||
provider_key_configured_api_formats, provider_key_inherits_provider_api_formats,
|
||||
};
|
||||
use crate::AppState;
|
||||
@@ -168,6 +168,19 @@ pub(crate) fn masked_catalog_api_key(state: &AppState, key: &StoredProviderCatal
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn masked_catalog_api_key_for_provider(
|
||||
state: &AppState,
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
) -> String {
|
||||
let auth_config = parse_catalog_auth_config_json(state, key);
|
||||
if provider_key_auth_config_is_agent_identity(provider_type, auth_config.as_ref()) {
|
||||
"[Agent Identity]".to_string()
|
||||
} else {
|
||||
masked_catalog_api_key(state, key)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn parse_catalog_auth_config_json(
|
||||
state: &AppState,
|
||||
key: &StoredProviderCatalogKey,
|
||||
@@ -2496,11 +2509,11 @@ pub(crate) fn build_admin_provider_key_response(
|
||||
);
|
||||
payload.insert(
|
||||
"api_key_masked".to_string(),
|
||||
json!(if agent_identity {
|
||||
"[Agent Identity]".to_string()
|
||||
} else {
|
||||
masked_catalog_api_key(state, key)
|
||||
}),
|
||||
json!(masked_catalog_api_key_for_provider(
|
||||
state,
|
||||
key,
|
||||
provider_type,
|
||||
)),
|
||||
);
|
||||
payload.insert("api_key_plain".to_string(), serde_json::Value::Null);
|
||||
payload.insert("auth_type".to_string(), json!(key.auth_type));
|
||||
@@ -2529,12 +2542,17 @@ pub(crate) fn build_admin_provider_key_response(
|
||||
"can_refresh_oauth".to_string(),
|
||||
json!(provider_key_can_refresh_oauth(
|
||||
auth_semantics,
|
||||
provider_type,
|
||||
auth_config.as_ref()
|
||||
)),
|
||||
);
|
||||
payload.insert(
|
||||
"can_export_oauth".to_string(),
|
||||
json!(auth_semantics.can_export_oauth()),
|
||||
json!(provider_key_can_export_oauth(
|
||||
auth_semantics,
|
||||
provider_type,
|
||||
auth_config.as_ref()
|
||||
)),
|
||||
);
|
||||
payload.insert(
|
||||
"can_edit_oauth".to_string(),
|
||||
@@ -2874,6 +2892,46 @@ mod tests {
|
||||
assert_ne!(masked, "***ERROR***");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_aware_mask_labels_agent_identity_without_exposing_placeholder() {
|
||||
let state = AppState::new().expect("gateway should build");
|
||||
let encrypted_placeholder =
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__")
|
||||
.expect("placeholder ciphertext should build");
|
||||
let encrypted_auth_config = encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"provider_type":"codex","auth_mode":"agentIdentity","agent_runtime_id":"runtime-1","agent_private_key":"base64-private-key","task_id":"task-1"}"#,
|
||||
)
|
||||
.expect("auth config ciphertext should build");
|
||||
let key = StoredProviderCatalogKey::new(
|
||||
"key-agent".to_string(),
|
||||
"provider-codex".to_string(),
|
||||
"agent".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(json!(["openai:responses"])),
|
||||
encrypted_placeholder,
|
||||
Some(encrypted_auth_config),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build");
|
||||
|
||||
assert_eq!(
|
||||
masked_catalog_api_key_for_provider(&state, &key, "codex"),
|
||||
"[Agent Identity]"
|
||||
);
|
||||
assert!(!masked_catalog_api_key_for_provider(&state, &key, "codex").contains("placeholder"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_backfills_missing_quota_from_upstream_metadata() {
|
||||
let mut key = sample_catalog_key();
|
||||
|
||||
@@ -24,7 +24,8 @@ pub(crate) use self::api_keys::{
|
||||
pub(crate) use self::catalog::{
|
||||
build_admin_provider_key_response, decrypt_catalog_secret_with_fallbacks,
|
||||
default_provider_key_status_snapshot, effective_catalog_encryption_key,
|
||||
encrypt_catalog_secret_with_fallbacks, masked_catalog_api_key, parse_catalog_auth_config_json,
|
||||
encrypt_catalog_secret_with_fallbacks, masked_catalog_api_key,
|
||||
masked_catalog_api_key_for_provider, parse_catalog_auth_config_json,
|
||||
provider_catalog_key_supports_format, provider_key_health_summary,
|
||||
provider_key_health_summary_at, provider_key_status_snapshot_payload,
|
||||
sync_provider_key_oauth_status_snapshot, sync_provider_key_quota_status_snapshot,
|
||||
|
||||
@@ -14,6 +14,7 @@ use crate::{AppState, GatewayError};
|
||||
use super::system_config_bool;
|
||||
|
||||
const OAUTH_TOKEN_REFRESH_LOOKAHEAD_SECS: u64 = 120;
|
||||
const OAUTH_REFRESH_FAILED_PREFIX: &str = "[REFRESH_FAILED] ";
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize)]
|
||||
pub(crate) struct OAuthTokenRefreshRunSummary {
|
||||
@@ -91,13 +92,34 @@ pub(crate) async fn perform_oauth_token_refresh_once(
|
||||
summary.skipped = summary.skipped.saturating_add(1);
|
||||
continue;
|
||||
};
|
||||
if !auth_config_has_refresh_token(transport.key.decrypted_auth_config.as_deref()) {
|
||||
let is_agent_identity =
|
||||
crate::provider_transport::is_codex_agent_identity_transport(&transport);
|
||||
let needs_agent_task_recovery = is_agent_identity
|
||||
&& agent_identity_needs_task_recovery(
|
||||
transport.key.decrypted_auth_config.as_deref(),
|
||||
key.oauth_invalid_reason.as_deref(),
|
||||
);
|
||||
if !needs_agent_task_recovery
|
||||
&& !auth_config_has_refresh_token(transport.key.decrypted_auth_config.as_deref())
|
||||
{
|
||||
summary.skipped = summary.skipped.saturating_add(1);
|
||||
continue;
|
||||
}
|
||||
|
||||
match state.resolve_local_oauth_request_auth(&transport).await {
|
||||
Ok(Some(_auth)) => {
|
||||
let refresh_result = if needs_agent_task_recovery {
|
||||
state
|
||||
.force_local_oauth_refresh_entry(&transport)
|
||||
.await
|
||||
.map(|entry| entry.map(|_| ()))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
} else {
|
||||
state
|
||||
.resolve_local_oauth_request_auth(&transport)
|
||||
.await
|
||||
.map(|auth| auth.map(|_| ()))
|
||||
};
|
||||
match refresh_result {
|
||||
Ok(Some(())) => {
|
||||
summary.resolved = summary.resolved.saturating_add(1);
|
||||
if provider_key_credentials_changed(state, key).await? {
|
||||
summary.refreshed = summary.refreshed.saturating_add(1);
|
||||
@@ -171,19 +193,45 @@ fn oauth_refresh_candidate(
|
||||
key: &StoredProviderCatalogKey,
|
||||
refresh_cutoff_unix_secs: u64,
|
||||
) -> bool {
|
||||
key.is_active
|
||||
&& key.oauth_invalid_at_unix_secs.is_none()
|
||||
&& key
|
||||
.encrypted_auth_config
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty())
|
||||
let has_auth_config = key
|
||||
.encrypted_auth_config
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty());
|
||||
let regular_oauth_candidate = key.oauth_invalid_at_unix_secs.is_none()
|
||||
&& key
|
||||
.expires_at_unix_secs
|
||||
.is_some_and(|expires_at| expires_at <= refresh_cutoff_unix_secs)
|
||||
.is_some_and(|expires_at| expires_at <= refresh_cutoff_unix_secs);
|
||||
// The catalog row is encrypted here, so exact Agent Identity validation is
|
||||
// deferred until the transport snapshot has decrypted auth_config.
|
||||
let possible_agent_candidate = provider.provider_type.trim().eq_ignore_ascii_case("codex")
|
||||
&& key.auth_type.trim().eq_ignore_ascii_case("oauth")
|
||||
&& (key.expires_at_unix_secs.is_none()
|
||||
|| key
|
||||
.oauth_invalid_reason
|
||||
.as_deref()
|
||||
.is_some_and(|reason| reason.contains(OAUTH_REFRESH_FAILED_PREFIX)));
|
||||
key.is_active
|
||||
&& has_auth_config
|
||||
&& (regular_oauth_candidate || possible_agent_candidate)
|
||||
&& provider_key_is_oauth_managed(key, provider.provider_type.as_str())
|
||||
}
|
||||
|
||||
fn agent_identity_needs_task_recovery(
|
||||
auth_config: Option<&str>,
|
||||
oauth_invalid_reason: Option<&str>,
|
||||
) -> bool {
|
||||
if oauth_invalid_reason.is_some_and(|reason| reason.contains(OAUTH_REFRESH_FAILED_PREFIX)) {
|
||||
return true;
|
||||
}
|
||||
auth_config
|
||||
.and_then(|value| serde_json::from_str::<Value>(value).ok())
|
||||
.is_some_and(|config| {
|
||||
crate::provider_transport::is_codex_agent_identity_auth_config_value(&config)
|
||||
&& !crate::provider_transport::codex_agent_identity_auth_config_has_task_id(&config)
|
||||
})
|
||||
}
|
||||
|
||||
async fn provider_key_credentials_changed(
|
||||
state: &AppState,
|
||||
before: &StoredProviderCatalogKey,
|
||||
@@ -222,3 +270,29 @@ fn now_unix_secs() -> u64 {
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::agent_identity_needs_task_recovery;
|
||||
|
||||
#[test]
|
||||
fn pending_agent_identity_without_task_is_recoverable() {
|
||||
let config = serde_json::json!({
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_runtime_id": "runtime-1",
|
||||
"agent_private_key": "private-key-present",
|
||||
});
|
||||
assert!(agent_identity_needs_task_recovery(
|
||||
Some(&config.to_string()),
|
||||
None,
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn refresh_failure_marker_forces_agent_task_recovery() {
|
||||
assert!(agent_identity_needs_task_recovery(
|
||||
Some("{}"),
|
||||
Some("[REFRESH_FAILED] temporary"),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -44,9 +44,7 @@ use crate::handlers::shared::provider_pool::{
|
||||
use crate::orchestration::local_execution_candidate_metadata_from_report_context;
|
||||
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
|
||||
use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerSchedulingMode};
|
||||
use crate::{
|
||||
provider_transport::snapshot::GatewayProviderTransportProvider, AppState, GatewayError,
|
||||
};
|
||||
use crate::AppState;
|
||||
|
||||
const POOL_SCORE_FEEDBACK_GATE_MAX_ENTRIES: usize = 50_000;
|
||||
const HEALTH_SUCCESS_PERSIST_GATE_MAX_ENTRIES: usize = 50_000;
|
||||
@@ -253,6 +251,21 @@ struct PoolFeedbackContext {
|
||||
sticky_session_token: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
enum LocalExecutionAuthConfigFence {
|
||||
Unfenced,
|
||||
Fenced(String),
|
||||
}
|
||||
|
||||
impl LocalExecutionAuthConfigFence {
|
||||
fn encrypted_auth_config(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::Unfenced => None,
|
||||
Self::Fenced(ciphertext) => Some(ciphertext),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const ADAPTIVE_RPM_RECENT_CANDIDATE_LIMIT: usize = 512;
|
||||
const LOCAL_EXECUTION_SCHEDULER_AFFINITY_MAX_ENTRIES: usize = 10_000;
|
||||
|
||||
@@ -405,6 +418,73 @@ async fn local_execution_plan_uses_pool(state: &AppState, plan: &ExecutionPlan)
|
||||
admin_provider_pool_config_from_config_value(transport.provider.config.as_ref()).is_some()
|
||||
}
|
||||
|
||||
async fn capture_local_execution_auth_config_fence(
|
||||
state: &AppState,
|
||||
plan: &ExecutionPlan,
|
||||
) -> Option<LocalExecutionAuthConfigFence> {
|
||||
let transport = match state
|
||||
.read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
|
||||
.await
|
||||
{
|
||||
Ok(Some(transport)) => transport,
|
||||
Ok(None) => return None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
provider_id = %plan.provider_id,
|
||||
endpoint_id = %plan.endpoint_id,
|
||||
key_id = %plan.key_id,
|
||||
error = ?err,
|
||||
"gateway orchestration effects: failed to read transport for credential fencing"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
if !transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("codex")
|
||||
|| !transport.key.auth_type.trim().eq_ignore_ascii_case("oauth")
|
||||
{
|
||||
return Some(LocalExecutionAuthConfigFence::Unfenced);
|
||||
}
|
||||
|
||||
let authorization = execution_plan_authorization(plan)?;
|
||||
let current_uses_agent_identity =
|
||||
crate::provider_transport::is_codex_agent_identity_transport(&transport);
|
||||
let authorization_matches = if current_uses_agent_identity {
|
||||
crate::provider_transport::codex_agent_identity_authorization_matches_transport(
|
||||
&transport,
|
||||
authorization,
|
||||
)
|
||||
} else if crate::provider_transport::is_codex_agent_identity_authorization(authorization) {
|
||||
false
|
||||
} else {
|
||||
execution_plan_bearer_matches_transport(plan, &transport)
|
||||
};
|
||||
if !authorization_matches {
|
||||
return None;
|
||||
}
|
||||
|
||||
match state
|
||||
.capture_provider_transport_auth_config_fence(&transport)
|
||||
.await
|
||||
{
|
||||
Ok(Some(ciphertext)) => Some(LocalExecutionAuthConfigFence::Fenced(ciphertext)),
|
||||
Ok(None) => None,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
provider_id = %plan.provider_id,
|
||||
endpoint_id = %plan.endpoint_id,
|
||||
key_id = %plan.key_id,
|
||||
error = ?err,
|
||||
"gateway orchestration effects: failed to capture credential fence"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn local_scheduler_affinity_matches_failed_target(
|
||||
state: &AppState,
|
||||
plan: &ExecutionPlan,
|
||||
@@ -480,6 +560,7 @@ async fn resolve_pool_feedback_context(
|
||||
context: LocalExecutionEffectContext<'_>,
|
||||
) -> Option<PoolFeedbackContext> {
|
||||
let plan = context.plan;
|
||||
capture_local_execution_auth_config_fence(state, plan).await?;
|
||||
let transport = match state
|
||||
.read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
|
||||
.await
|
||||
@@ -602,6 +683,11 @@ async fn record_adaptive_rate_limit_effect(
|
||||
context: LocalExecutionEffectContext<'_>,
|
||||
effect: LocalAdaptiveRateLimitEffect<'_>,
|
||||
) {
|
||||
let Some(auth_config_fence) =
|
||||
capture_local_execution_auth_config_fence(state, context.plan).await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let effect_lock = PROVIDER_KEY_EFFECT_LOCKS.lock_for(&context.plan.key_id);
|
||||
let _effect_guard = effect_lock.lock().await;
|
||||
let observed_at_unix_secs = current_unix_secs();
|
||||
@@ -626,6 +712,12 @@ async fn record_adaptive_rate_limit_effect(
|
||||
else {
|
||||
return;
|
||||
};
|
||||
if auth_config_fence
|
||||
.encrypted_auth_config()
|
||||
.is_some_and(|expected| current_key.encrypted_auth_config.as_deref() != Some(expected))
|
||||
{
|
||||
return;
|
||||
}
|
||||
let Some(projection) = project_local_adaptive_rate_limit(
|
||||
¤t_key,
|
||||
effect.classification,
|
||||
@@ -648,6 +740,9 @@ async fn record_adaptive_rate_limit_effect(
|
||||
next.last_rpm_peak = projection.last_rpm_peak;
|
||||
let update = ProviderCatalogKeyAdaptiveStateUpdate {
|
||||
key_id: context.plan.key_id.clone(),
|
||||
expected_encrypted_auth_config: auth_config_fence
|
||||
.encrypted_auth_config()
|
||||
.map(ToOwned::to_owned),
|
||||
expected,
|
||||
next,
|
||||
status_snapshot_patch: adaptive_status_snapshot_patch(&projection.status_snapshot),
|
||||
@@ -705,6 +800,11 @@ async fn record_adaptive_success_effect(
|
||||
context: LocalExecutionEffectContext<'_>,
|
||||
_effect: LocalAdaptiveSuccessEffect,
|
||||
) {
|
||||
let Some(auth_config_fence) =
|
||||
capture_local_execution_auth_config_fence(state, context.plan).await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let observed_at_unix_secs = current_unix_secs();
|
||||
let Some(current_key) = state
|
||||
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&context.plan.key_id))
|
||||
@@ -714,6 +814,12 @@ async fn record_adaptive_success_effect(
|
||||
else {
|
||||
return;
|
||||
};
|
||||
if auth_config_fence
|
||||
.encrypted_auth_config()
|
||||
.is_some_and(|expected| current_key.encrypted_auth_config.as_deref() != Some(expected))
|
||||
{
|
||||
return;
|
||||
}
|
||||
if current_key.rpm_limit.is_some()
|
||||
|| current_key
|
||||
.learned_rpm_limit
|
||||
@@ -757,6 +863,12 @@ async fn record_adaptive_success_effect(
|
||||
else {
|
||||
return;
|
||||
};
|
||||
if auth_config_fence
|
||||
.encrypted_auth_config()
|
||||
.is_some_and(|expected| current_key.encrypted_auth_config.as_deref() != Some(expected))
|
||||
{
|
||||
return;
|
||||
}
|
||||
if current_key.rpm_limit.is_some()
|
||||
|| current_key
|
||||
.learned_rpm_limit
|
||||
@@ -778,6 +890,9 @@ async fn record_adaptive_success_effect(
|
||||
next.last_probe_increase_at_unix_secs = projection.last_probe_increase_at_unix_secs;
|
||||
let update = ProviderCatalogKeyAdaptiveStateUpdate {
|
||||
key_id: context.plan.key_id.clone(),
|
||||
expected_encrypted_auth_config: auth_config_fence
|
||||
.encrypted_auth_config()
|
||||
.map(ToOwned::to_owned),
|
||||
expected,
|
||||
next,
|
||||
status_snapshot_patch: adaptive_status_snapshot_patch(&projection.status_snapshot),
|
||||
@@ -853,6 +968,11 @@ async fn record_health_failure_effect(
|
||||
if api_format.is_empty() {
|
||||
return;
|
||||
}
|
||||
let Some(auth_config_fence) =
|
||||
capture_local_execution_auth_config_fence(state, context.plan).await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
|
||||
let effect_lock = PROVIDER_KEY_EFFECT_LOCKS.lock_for(&context.plan.key_id);
|
||||
let _effect_guard = effect_lock.lock().await;
|
||||
@@ -869,6 +989,12 @@ async fn record_health_failure_effect(
|
||||
else {
|
||||
return;
|
||||
};
|
||||
if auth_config_fence
|
||||
.encrypted_auth_config()
|
||||
.is_some_and(|expected| current_key.encrypted_auth_config.as_deref() != Some(expected))
|
||||
{
|
||||
return;
|
||||
}
|
||||
let Some(health_by_format) = project_local_failure_health(
|
||||
current_key.health_by_format.as_ref(),
|
||||
api_format,
|
||||
@@ -897,6 +1023,9 @@ async fn record_health_failure_effect(
|
||||
};
|
||||
let update = ProviderCatalogKeyHealthStateUpdate {
|
||||
key_id: context.plan.key_id.clone(),
|
||||
expected_encrypted_auth_config: auth_config_fence
|
||||
.encrypted_auth_config()
|
||||
.map(ToOwned::to_owned),
|
||||
expected_health_by_format: current_key.health_by_format,
|
||||
expected_circuit_breaker_by_format: current_key.circuit_breaker_by_format,
|
||||
health_by_format: Some(health_by_format),
|
||||
@@ -934,6 +1063,11 @@ async fn record_health_success_effect(
|
||||
if api_format.is_empty() {
|
||||
return;
|
||||
}
|
||||
let Some(auth_config_fence) =
|
||||
capture_local_execution_auth_config_fence(state, context.plan).await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
|
||||
// Health updates replace both JSON snapshots in one write. Serialize the success
|
||||
// read/project/write with failure and circuit-clear effects for this provider key so a
|
||||
@@ -953,6 +1087,12 @@ async fn record_health_success_effect(
|
||||
else {
|
||||
return;
|
||||
};
|
||||
if auth_config_fence
|
||||
.encrypted_auth_config()
|
||||
.is_some_and(|expected| current_key.encrypted_auth_config.as_deref() != Some(expected))
|
||||
{
|
||||
return;
|
||||
}
|
||||
let Some(health_by_format) =
|
||||
project_local_success_health(current_key.health_by_format.as_ref(), api_format)
|
||||
else {
|
||||
@@ -991,6 +1131,9 @@ async fn record_health_success_effect(
|
||||
};
|
||||
let update = ProviderCatalogKeyHealthStateUpdate {
|
||||
key_id: context.plan.key_id.clone(),
|
||||
expected_encrypted_auth_config: auth_config_fence
|
||||
.encrypted_auth_config()
|
||||
.map(ToOwned::to_owned),
|
||||
expected_health_by_format: current_key.health_by_format,
|
||||
expected_circuit_breaker_by_format: current_key.circuit_breaker_by_format,
|
||||
health_by_format: Some(health_by_format),
|
||||
@@ -1104,6 +1247,12 @@ async fn record_pool_error_effect(
|
||||
};
|
||||
|
||||
clear_pool_key_circuit_breaker(state, context).await;
|
||||
if capture_local_execution_auth_config_fence(state, context.plan)
|
||||
.await
|
||||
.is_none()
|
||||
{
|
||||
return;
|
||||
}
|
||||
record_admin_provider_pool_error(
|
||||
state.runtime_state.as_ref(),
|
||||
&context.plan.provider_id,
|
||||
@@ -1135,6 +1284,11 @@ async fn clear_pool_key_circuit_breaker(
|
||||
state: &AppState,
|
||||
context: LocalExecutionEffectContext<'_>,
|
||||
) {
|
||||
let Some(auth_config_fence) =
|
||||
capture_local_execution_auth_config_fence(state, context.plan).await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let effect_lock = PROVIDER_KEY_EFFECT_LOCKS.lock_for(&context.plan.key_id);
|
||||
let _effect_guard = effect_lock.lock().await;
|
||||
|
||||
@@ -1147,11 +1301,20 @@ async fn clear_pool_key_circuit_breaker(
|
||||
else {
|
||||
return;
|
||||
};
|
||||
if auth_config_fence
|
||||
.encrypted_auth_config()
|
||||
.is_some_and(|expected| current_key.encrypted_auth_config.as_deref() != Some(expected))
|
||||
{
|
||||
return;
|
||||
}
|
||||
if current_key.circuit_breaker_by_format.is_none() {
|
||||
return;
|
||||
}
|
||||
let update = ProviderCatalogKeyHealthStateUpdate {
|
||||
key_id: context.plan.key_id.clone(),
|
||||
expected_encrypted_auth_config: auth_config_fence
|
||||
.encrypted_auth_config()
|
||||
.map(ToOwned::to_owned),
|
||||
expected_health_by_format: current_key.health_by_format.clone(),
|
||||
expected_circuit_breaker_by_format: current_key.circuit_breaker_by_format,
|
||||
health_by_format: current_key.health_by_format,
|
||||
@@ -1188,6 +1351,12 @@ async fn record_oauth_invalidation_effect(
|
||||
}
|
||||
|
||||
let plan = context.plan;
|
||||
// Agent assertions are long-lived credential requests whose task can rotate
|
||||
// while the response is in flight. Runtime 401/403 handling must not project
|
||||
// that response onto whichever credential generation is stored later.
|
||||
if execution_plan_uses_codex_agent_identity(plan) {
|
||||
return;
|
||||
}
|
||||
let transport = match state
|
||||
.read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
|
||||
.await
|
||||
@@ -1202,9 +1371,23 @@ async fn record_oauth_invalidation_effect(
|
||||
return;
|
||||
}
|
||||
};
|
||||
// The inverse replacement is equally unsafe: a response sent with an old
|
||||
// bearer token must not invalidate a newly installed Agent Identity.
|
||||
if crate::provider_transport::is_codex_agent_identity_transport(&transport) {
|
||||
return;
|
||||
}
|
||||
if !transport.key.auth_type.trim().eq_ignore_ascii_case("oauth") {
|
||||
return;
|
||||
}
|
||||
if transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("codex")
|
||||
&& !execution_plan_bearer_matches_transport(plan, &transport)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(invalid_reason) = resolve_local_oauth_invalid_reason(
|
||||
transport.provider.provider_type.as_str(),
|
||||
@@ -1214,11 +1397,27 @@ async fn record_oauth_invalidation_effect(
|
||||
return;
|
||||
};
|
||||
|
||||
let expected_auth_config = match state
|
||||
.capture_provider_transport_auth_config_fence(&transport)
|
||||
.await
|
||||
{
|
||||
Ok(Some(ciphertext)) => ciphertext,
|
||||
Ok(None) => return,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
"gateway orchestration effects: failed to capture oauth invalidation fence for provider {} endpoint {} key {}: {:?}",
|
||||
plan.provider_id, plan.endpoint_id, plan.key_id, err
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(err) = state
|
||||
.mark_provider_catalog_key_oauth_invalid(
|
||||
.mark_provider_catalog_key_oauth_invalid_fenced(
|
||||
&plan.key_id,
|
||||
transport.provider.provider_type.as_str(),
|
||||
invalid_reason.as_str(),
|
||||
expected_auth_config.as_str(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -1227,98 +1426,34 @@ async fn record_oauth_invalidation_effect(
|
||||
plan.provider_id, plan.endpoint_id, plan.key_id, err
|
||||
);
|
||||
}
|
||||
record_pool_score_schedule_feedback(
|
||||
state,
|
||||
context,
|
||||
Some(false),
|
||||
Some(PoolMemberHardState::AuthInvalid),
|
||||
Some(-2_000),
|
||||
serde_json::json!({
|
||||
"last_request_feedback": {
|
||||
"source": "oauth_invalidation",
|
||||
"status_code": effect.status_code,
|
||||
"reason": invalid_reason.as_str()
|
||||
}
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
match auto_remove_runtime_oauth_invalid_key(
|
||||
state,
|
||||
&transport.provider,
|
||||
&plan.key_id,
|
||||
invalid_reason.as_str(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(true) => {
|
||||
tracing::info!(
|
||||
provider_id = %plan.provider_id,
|
||||
endpoint_id = %plan.endpoint_id,
|
||||
key_id = %plan.key_id,
|
||||
provider_type = %transport.provider.provider_type,
|
||||
event_name = "auto_removed_oauth_runtime_invalid",
|
||||
"gateway auto-removed runtime invalid oauth key"
|
||||
);
|
||||
}
|
||||
Ok(false) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
"gateway orchestration effects: failed to auto-remove oauth invalid key for provider {} endpoint {} key {}: {:?}",
|
||||
plan.provider_id, plan.endpoint_id, plan.key_id, err
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn auto_remove_runtime_oauth_invalid_key(
|
||||
state: &AppState,
|
||||
provider: &GatewayProviderTransportProvider,
|
||||
key_id: &str,
|
||||
invalid_reason: &str,
|
||||
) -> Result<bool, GatewayError> {
|
||||
if !admin_provider_quota_pure::provider_auto_remove_banned_keys(provider.config.as_ref()) {
|
||||
return Ok(false);
|
||||
}
|
||||
fn execution_plan_uses_codex_agent_identity(plan: &ExecutionPlan) -> bool {
|
||||
execution_plan_authorization(plan)
|
||||
.is_some_and(crate::provider_transport::is_codex_agent_identity_authorization)
|
||||
}
|
||||
|
||||
let key_ids = [key_id.to_string()];
|
||||
let Some(key) = state
|
||||
.read_provider_catalog_keys_by_ids(&key_ids)
|
||||
.await?
|
||||
.into_iter()
|
||||
.next()
|
||||
else {
|
||||
return Ok(false);
|
||||
};
|
||||
if key.provider_id != provider.id {
|
||||
return Ok(false);
|
||||
}
|
||||
fn execution_plan_authorization(plan: &ExecutionPlan) -> Option<&str> {
|
||||
plan.headers
|
||||
.iter()
|
||||
.find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
|
||||
.map(|(_, value)| value.as_str())
|
||||
}
|
||||
|
||||
if !admin_provider_quota_pure::should_auto_remove_oauth_invalid_key(
|
||||
&key,
|
||||
Some(invalid_reason),
|
||||
true,
|
||||
current_unix_secs(),
|
||||
) {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let deleted_key_id = key.id.clone();
|
||||
if !state.delete_provider_catalog_key(&deleted_key_id).await? {
|
||||
return Ok(false);
|
||||
}
|
||||
state
|
||||
.cleanup_deleted_provider_catalog_refs(
|
||||
&provider.id,
|
||||
false,
|
||||
&[],
|
||||
std::slice::from_ref(&deleted_key_id),
|
||||
)
|
||||
.await?;
|
||||
let _ = state
|
||||
.invalidate_local_oauth_refresh_entry(&deleted_key_id)
|
||||
.await;
|
||||
Ok(true)
|
||||
fn execution_plan_bearer_matches_transport(
|
||||
plan: &ExecutionPlan,
|
||||
transport: &crate::provider_transport::GatewayProviderTransportSnapshot,
|
||||
) -> bool {
|
||||
let current_token = transport.key.decrypted_api_key.trim();
|
||||
!current_token.is_empty()
|
||||
&& plan.headers.iter().any(|(name, value)| {
|
||||
name.eq_ignore_ascii_case("authorization")
|
||||
&& value
|
||||
.trim()
|
||||
.strip_prefix("Bearer ")
|
||||
.map(str::trim)
|
||||
.is_some_and(|token| token == current_token)
|
||||
})
|
||||
}
|
||||
|
||||
fn resolve_local_oauth_invalid_reason(
|
||||
@@ -1409,6 +1544,12 @@ async fn record_pool_score_schedule_feedback(
|
||||
if context.plan.provider_id.trim().is_empty() || context.plan.key_id.trim().is_empty() {
|
||||
return;
|
||||
}
|
||||
if capture_local_execution_auth_config_fence(state, context.plan)
|
||||
.await
|
||||
.is_none()
|
||||
{
|
||||
return;
|
||||
}
|
||||
if !pool_score_feedback_gate_allows(context.plan, succeeded, hard_state, score_delta) {
|
||||
return;
|
||||
}
|
||||
@@ -1560,7 +1701,8 @@ mod tests {
|
||||
};
|
||||
use aether_data_contracts::repository::pool_scores::PoolMemberHardState;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
ProviderCatalogKeyAdaptiveState, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_test_support::ManagedRedisServer;
|
||||
use serde_json::{json, Value};
|
||||
@@ -1720,7 +1862,10 @@ mod tests {
|
||||
key_id: "key-codex-cli-local-1".to_string(),
|
||||
method: "POST".to_string(),
|
||||
url: "https://chatgpt.com/backend-api/codex".to_string(),
|
||||
headers: BTreeMap::new(),
|
||||
headers: BTreeMap::from([(
|
||||
"authorization".to_string(),
|
||||
"Bearer __placeholder__".to_string(),
|
||||
)]),
|
||||
content_type: Some("application/json".to_string()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"model":"gpt-5.4"})),
|
||||
@@ -1734,6 +1879,15 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_codex_agent_identity_plan() -> ExecutionPlan {
|
||||
let mut plan = sample_codex_plan();
|
||||
plan.headers.insert(
|
||||
"Authorization".to_string(),
|
||||
"AgentAssertion in-flight-assertion".to_string(),
|
||||
);
|
||||
plan
|
||||
}
|
||||
|
||||
fn sample_codex_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-codex-cli-local-1".to_string(),
|
||||
@@ -1818,6 +1972,19 @@ mod tests {
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
fn sample_codex_agent_identity_key() -> StoredProviderCatalogKey {
|
||||
let mut key = sample_codex_key();
|
||||
key.name = "Agent Identity".to_string();
|
||||
key.encrypted_auth_config = Some(
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"provider_type":"codex","auth_mode":"agentIdentity","agent_runtime_id":"runtime-current","agent_private_key":"MC4CAQAwBQYDK2VwBCIEIAcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcH","task_id":"task-current"}"#,
|
||||
)
|
||||
.expect("Agent Identity auth config should encrypt"),
|
||||
);
|
||||
key
|
||||
}
|
||||
|
||||
fn codex_state() -> AppState {
|
||||
codex_state_with_provider(sample_codex_provider())
|
||||
}
|
||||
@@ -1827,10 +1994,17 @@ mod tests {
|
||||
}
|
||||
|
||||
fn codex_state_with_provider(provider: StoredProviderCatalogProvider) -> AppState {
|
||||
codex_state_with_provider_and_key(provider, sample_codex_key())
|
||||
}
|
||||
|
||||
fn codex_state_with_provider_and_key(
|
||||
provider: StoredProviderCatalogProvider,
|
||||
key: StoredProviderCatalogKey,
|
||||
) -> AppState {
|
||||
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
vec![sample_codex_endpoint()],
|
||||
vec![sample_codex_key()],
|
||||
vec![key],
|
||||
));
|
||||
AppState::new()
|
||||
.expect("gateway state should build")
|
||||
@@ -2744,7 +2918,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn oauth_invalidation_auto_removes_inactive_pat_owner_when_enabled() {
|
||||
async fn oauth_invalidation_retains_inactive_pat_owner_when_auto_remove_is_enabled() {
|
||||
let state = codex_state_with_auto_remove();
|
||||
let plan = sample_codex_plan();
|
||||
|
||||
@@ -2763,13 +2937,16 @@ mod tests {
|
||||
)
|
||||
.await;
|
||||
|
||||
let keys = state
|
||||
let stored_key = state
|
||||
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id))
|
||||
.await
|
||||
.expect("provider catalog keys should load");
|
||||
assert!(
|
||||
keys.is_empty(),
|
||||
"hard-invalid PAT owner should be auto removed"
|
||||
.expect("provider catalog keys should load")
|
||||
.into_iter()
|
||||
.next()
|
||||
.expect("runtime invalidation must not race-delete a replacement key");
|
||||
assert_eq!(
|
||||
stored_key.oauth_invalid_reason.as_deref(),
|
||||
Some("[OAUTH_EXPIRED] Personal access token owner is inactive.")
|
||||
);
|
||||
}
|
||||
|
||||
@@ -2806,6 +2983,253 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn oauth_invalidation_does_not_mutate_replacement_after_agent_request() {
|
||||
let state = codex_state_with_auto_remove();
|
||||
let plan = sample_codex_agent_identity_plan();
|
||||
|
||||
apply_local_execution_effect(
|
||||
&state,
|
||||
LocalExecutionEffectContext {
|
||||
plan: &plan,
|
||||
report_context: None,
|
||||
},
|
||||
LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect {
|
||||
status_code: 403,
|
||||
response_text: Some(
|
||||
r#"{"error":{"code":"biscuit_baker_service_auth_credential_error_status","message":"Personal access token owner is inactive."},"status":403}"#,
|
||||
),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let stored_key = state
|
||||
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id))
|
||||
.await
|
||||
.expect("provider catalog keys should load")
|
||||
.into_iter()
|
||||
.next()
|
||||
.expect("replacement OAuth key should not be removed");
|
||||
assert_eq!(stored_key.oauth_invalid_at_unix_secs, None);
|
||||
assert_eq!(stored_key.oauth_invalid_reason, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn oauth_invalidation_does_not_mutate_agent_replacement_after_bearer_request() {
|
||||
let state = codex_state_with_provider_and_key(
|
||||
sample_codex_provider_with_auto_remove(),
|
||||
sample_codex_agent_identity_key(),
|
||||
);
|
||||
let mut plan = sample_codex_plan();
|
||||
plan.headers.insert(
|
||||
"authorization".to_string(),
|
||||
"Bearer old-access-token".to_string(),
|
||||
);
|
||||
|
||||
apply_local_execution_effect(
|
||||
&state,
|
||||
LocalExecutionEffectContext {
|
||||
plan: &plan,
|
||||
report_context: None,
|
||||
},
|
||||
LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect {
|
||||
status_code: 403,
|
||||
response_text: Some(
|
||||
r#"{"error":{"code":"biscuit_baker_service_auth_credential_error_status","message":"Personal access token owner is inactive."},"status":403}"#,
|
||||
),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let stored_key = state
|
||||
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id))
|
||||
.await
|
||||
.expect("provider catalog keys should load")
|
||||
.into_iter()
|
||||
.next()
|
||||
.expect("Agent Identity replacement should not be removed");
|
||||
assert_eq!(stored_key.oauth_invalid_at_unix_secs, None);
|
||||
assert_eq!(stored_key.oauth_invalid_reason, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn health_failure_updates_codex_key_for_current_bearer_request() {
|
||||
let state = codex_state();
|
||||
let plan = sample_codex_plan();
|
||||
|
||||
apply_local_execution_effect(
|
||||
&state,
|
||||
LocalExecutionEffectContext {
|
||||
plan: &plan,
|
||||
report_context: None,
|
||||
},
|
||||
LocalExecutionEffect::HealthFailure(LocalHealthFailureEffect {
|
||||
status_code: 503,
|
||||
classification: LocalFailoverClassification::RetryUpstreamFailure,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let stored_key = state
|
||||
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id))
|
||||
.await
|
||||
.expect("provider catalog keys should load")
|
||||
.into_iter()
|
||||
.next()
|
||||
.expect("stored key should exist");
|
||||
assert_eq!(
|
||||
stored_key
|
||||
.health_by_format
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("openai:responses"))
|
||||
.and_then(|value| value.get("consecutive_failures"))
|
||||
.and_then(Value::as_u64),
|
||||
Some(1)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn health_failure_does_not_mutate_codex_bearer_replacement() {
|
||||
let mut replacement = sample_codex_key();
|
||||
replacement.encrypted_api_key = Some(
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "replacement-token")
|
||||
.expect("replacement token should encrypt"),
|
||||
);
|
||||
replacement.encrypted_auth_config = Some(
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"provider_type":"codex","refresh_token":"replacement-refresh-token"}"#,
|
||||
)
|
||||
.expect("replacement auth config should encrypt"),
|
||||
);
|
||||
let expected_health = replacement.health_by_format.clone();
|
||||
let expected_circuit = replacement.circuit_breaker_by_format.clone();
|
||||
let state = codex_state_with_provider_and_key(sample_codex_provider(), replacement);
|
||||
let plan = sample_codex_plan();
|
||||
|
||||
apply_local_execution_effect(
|
||||
&state,
|
||||
LocalExecutionEffectContext {
|
||||
plan: &plan,
|
||||
report_context: None,
|
||||
},
|
||||
LocalExecutionEffect::HealthFailure(LocalHealthFailureEffect {
|
||||
status_code: 503,
|
||||
classification: LocalFailoverClassification::RetryUpstreamFailure,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let stored_key = state
|
||||
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id))
|
||||
.await
|
||||
.expect("provider catalog keys should load")
|
||||
.into_iter()
|
||||
.next()
|
||||
.expect("replacement key should exist");
|
||||
assert_eq!(stored_key.health_by_format, expected_health);
|
||||
assert_eq!(stored_key.circuit_breaker_by_format, expected_circuit);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn adaptive_rate_limit_does_not_mutate_codex_bearer_replacement() {
|
||||
let mut replacement = sample_codex_key();
|
||||
replacement.encrypted_api_key = Some(
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "replacement-token")
|
||||
.expect("replacement token should encrypt"),
|
||||
);
|
||||
replacement.encrypted_auth_config = Some(
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"provider_type":"codex","refresh_token":"replacement-refresh-token"}"#,
|
||||
)
|
||||
.expect("replacement auth config should encrypt"),
|
||||
);
|
||||
replacement.learned_rpm_limit = Some(12);
|
||||
replacement.rpm_429_count = Some(1);
|
||||
let expected_adaptive_state = ProviderCatalogKeyAdaptiveState::from(&replacement);
|
||||
let state = codex_state_with_provider_and_key(sample_codex_provider(), replacement);
|
||||
let plan = sample_codex_plan();
|
||||
|
||||
apply_local_execution_effect(
|
||||
&state,
|
||||
LocalExecutionEffectContext {
|
||||
plan: &plan,
|
||||
report_context: None,
|
||||
},
|
||||
LocalExecutionEffect::AdaptiveRateLimit(LocalAdaptiveRateLimitEffect {
|
||||
status_code: 429,
|
||||
classification: LocalFailoverClassification::RetryUpstreamFailure,
|
||||
headers: Some(&BTreeMap::from([(
|
||||
"x-ratelimit-limit-requests".to_string(),
|
||||
"42".to_string(),
|
||||
)])),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let stored_key = state
|
||||
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id))
|
||||
.await
|
||||
.expect("provider catalog keys should load")
|
||||
.into_iter()
|
||||
.next()
|
||||
.expect("replacement key should exist");
|
||||
assert_eq!(
|
||||
ProviderCatalogKeyAdaptiveState::from(&stored_key),
|
||||
expected_adaptive_state
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_error_does_not_clear_codex_bearer_replacement_circuit() {
|
||||
let legacy_circuit = json!({
|
||||
"openai:responses": {
|
||||
"open": true,
|
||||
"reason": "replacement-state"
|
||||
}
|
||||
});
|
||||
let mut replacement = sample_codex_key();
|
||||
replacement.encrypted_api_key = Some(
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "replacement-token")
|
||||
.expect("replacement token should encrypt"),
|
||||
);
|
||||
replacement.encrypted_auth_config = Some(
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"provider_type":"codex","refresh_token":"replacement-refresh-token"}"#,
|
||||
)
|
||||
.expect("replacement auth config should encrypt"),
|
||||
);
|
||||
replacement.circuit_breaker_by_format = Some(legacy_circuit.clone());
|
||||
let state = codex_state_with_provider_and_key(sample_codex_provider(), replacement);
|
||||
let plan = sample_codex_plan();
|
||||
|
||||
apply_local_execution_effect(
|
||||
&state,
|
||||
LocalExecutionEffectContext {
|
||||
plan: &plan,
|
||||
report_context: None,
|
||||
},
|
||||
LocalExecutionEffect::PoolError(LocalPoolErrorEffect {
|
||||
status_code: 401,
|
||||
classification: LocalFailoverClassification::StopErrorPattern,
|
||||
headers: &BTreeMap::new(),
|
||||
error_body: Some(r#"{"error":{"message":"account has been deactivated"}}"#),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let stored_key = state
|
||||
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id))
|
||||
.await
|
||||
.expect("provider catalog keys should load")
|
||||
.into_iter()
|
||||
.next()
|
||||
.expect("replacement key should exist");
|
||||
assert_eq!(stored_key.circuit_breaker_by_format, Some(legacy_circuit));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn health_failure_projection_updates_key_health_for_format() {
|
||||
let state = health_state();
|
||||
|
||||
@@ -85,14 +85,25 @@ impl ProviderKeyAuthSemantics {
|
||||
|
||||
pub(crate) fn provider_key_can_refresh_oauth(
|
||||
auth_semantics: ProviderKeyAuthSemantics,
|
||||
provider_type: &str,
|
||||
auth_config: Option<&Map<String, Value>>,
|
||||
) -> bool {
|
||||
auth_semantics.can_refresh_oauth()
|
||||
&& auth_config
|
||||
.and_then(|config| config.get("refresh_token"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty())
|
||||
&& (provider_key_auth_config_is_agent_identity(provider_type, auth_config)
|
||||
|| auth_config
|
||||
.and_then(|config| config.get("refresh_token"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty()))
|
||||
}
|
||||
|
||||
pub(crate) fn provider_key_can_export_oauth(
|
||||
auth_semantics: ProviderKeyAuthSemantics,
|
||||
provider_type: &str,
|
||||
auth_config: Option<&Map<String, Value>>,
|
||||
) -> bool {
|
||||
auth_semantics.can_export_oauth()
|
||||
&& !provider_key_auth_config_is_agent_identity(provider_type, auth_config)
|
||||
}
|
||||
|
||||
pub(crate) fn provider_key_auth_config_uses_header_authorization(
|
||||
@@ -291,9 +302,10 @@ mod tests {
|
||||
use super::{
|
||||
provider_active_api_formats, provider_key_auth_config_is_agent_identity,
|
||||
provider_key_auth_config_uses_header_authorization, provider_key_auth_semantics,
|
||||
provider_key_can_refresh_oauth, provider_key_configured_api_formats,
|
||||
provider_key_effective_api_formats, provider_key_inherits_provider_api_formats,
|
||||
ProviderKeyCredentialKind, ProviderKeyRuntimeAuthKind,
|
||||
provider_key_can_export_oauth, provider_key_can_refresh_oauth,
|
||||
provider_key_configured_api_formats, provider_key_effective_api_formats,
|
||||
provider_key_inherits_provider_api_formats, ProviderKeyCredentialKind,
|
||||
ProviderKeyRuntimeAuthKind,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
@@ -396,6 +408,7 @@ mod tests {
|
||||
|
||||
assert!(!provider_key_can_refresh_oauth(
|
||||
semantics,
|
||||
"codex",
|
||||
json!({
|
||||
"access_token": "access-token",
|
||||
"access_token_import_temporary": true
|
||||
@@ -404,12 +417,24 @@ mod tests {
|
||||
));
|
||||
assert!(!provider_key_can_refresh_oauth(
|
||||
semantics,
|
||||
"codex",
|
||||
json!({ "refresh_token": " " }).as_object()
|
||||
));
|
||||
assert!(provider_key_can_refresh_oauth(
|
||||
semantics,
|
||||
"codex",
|
||||
json!({ "refresh_token": "refresh-token" }).as_object()
|
||||
));
|
||||
assert!(provider_key_can_refresh_oauth(
|
||||
semantics,
|
||||
"codex",
|
||||
json!({
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_runtime_id": "runtime-1",
|
||||
"agent_private_key": "private-key-present"
|
||||
})
|
||||
.as_object()
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -455,6 +480,28 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_identity_is_not_exportable_through_generic_oauth_export() {
|
||||
let semantics = provider_key_auth_semantics(&sample_key("oauth"), "codex");
|
||||
let agent_identity = json!({
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_runtime_id": "runtime-1",
|
||||
"agent_private_key": "base64-private-key",
|
||||
"task_id": "task-1"
|
||||
});
|
||||
|
||||
assert!(!provider_key_can_export_oauth(
|
||||
semantics,
|
||||
"codex",
|
||||
agent_identity.as_object()
|
||||
));
|
||||
assert!(provider_key_can_export_oauth(
|
||||
semantics,
|
||||
"codex",
|
||||
json!({ "refresh_token": "refresh-token" }).as_object()
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recognizes_legacy_kiro_bearer_key_with_auth_config_as_oauth_managed() {
|
||||
let mut key = sample_key("bearer");
|
||||
|
||||
@@ -1465,6 +1465,7 @@ mod tests {
|
||||
.compare_and_update_provider_catalog_key_health_state(
|
||||
&aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyHealthStateUpdate {
|
||||
key_id: "key-1".to_string(),
|
||||
expected_encrypted_auth_config: None,
|
||||
expected_health_by_format: None,
|
||||
expected_circuit_breaker_by_format: None,
|
||||
health_by_format: Some(health_by_format),
|
||||
|
||||
@@ -38,6 +38,7 @@ pub(crate) use self::cache::{
|
||||
PROVIDER_TRANSPORT_SNAPSHOT_CACHE_STALE_TTL, PROVIDER_TRANSPORT_SNAPSHOT_CACHE_TTL,
|
||||
};
|
||||
pub use self::cors::FrontdoorCorsConfig;
|
||||
pub(crate) use self::oauth::AgentIdentityAuthConfigFence;
|
||||
pub(crate) use self::types::{
|
||||
AdminWalletMutationOutcome, GatewayAdminPaymentCallbackView, GatewayUserPreferenceView,
|
||||
GatewayUserSessionView, LocalExecutionRuntimeMissDiagnostic, LocalMutationOutcome,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -677,6 +677,103 @@ async fn gateway_creates_admin_provider_key_locally_with_trusted_admin_principal
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn generic_key_routes_reject_agent_identity_credential_writes() {
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-codex", "codex", 10)],
|
||||
vec![],
|
||||
vec![sample_key(
|
||||
"key-codex-existing",
|
||||
"provider-codex",
|
||||
"openai:responses",
|
||||
"existing-secret",
|
||||
)],
|
||||
));
|
||||
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.clone(),
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
let agent_identity = json!({
|
||||
"provider_type": "codex",
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_runtime_id": "runtime-bypass",
|
||||
"agent_private_key": "private-key-must-use-dedicated-import"
|
||||
});
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
let create_response = client
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-codex/keys"
|
||||
))
|
||||
.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!({
|
||||
"api_formats": ["openai:responses"],
|
||||
"auth_type": "oauth",
|
||||
"auth_config": agent_identity,
|
||||
"name": "bypass create"
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("create request should complete");
|
||||
assert_eq!(create_response.status(), StatusCode::BAD_REQUEST);
|
||||
let create_payload: serde_json::Value = create_response
|
||||
.json()
|
||||
.await
|
||||
.expect("create error should be JSON");
|
||||
assert!(create_payload["detail"]
|
||||
.as_str()
|
||||
.is_some_and(|detail| detail.contains("专属创建或导入接口")));
|
||||
|
||||
let update_response = client
|
||||
.put(format!(
|
||||
"{gateway_url}/api/admin/endpoints/keys/key-codex-existing"
|
||||
))
|
||||
.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!({
|
||||
"auth_type": "oauth",
|
||||
"auth_config": {
|
||||
"provider_type": "codex",
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_runtime_id": "runtime-bypass-update",
|
||||
"agent_private_key": "private-key-must-use-dedicated-import"
|
||||
}
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("update request should complete");
|
||||
assert_eq!(update_response.status(), StatusCode::BAD_REQUEST);
|
||||
let update_payload: serde_json::Value = update_response
|
||||
.json()
|
||||
.await
|
||||
.expect("update error should be JSON");
|
||||
assert!(update_payload["detail"]
|
||||
.as_str()
|
||||
.is_some_and(|detail| detail.contains("专属创建或导入接口")));
|
||||
|
||||
let keys = provider_catalog_repository
|
||||
.list_keys_by_provider_ids(&["provider-codex".to_string()])
|
||||
.await
|
||||
.expect("keys should read");
|
||||
assert_eq!(keys.len(), 1);
|
||||
assert_eq!(keys[0].id, "key-codex-existing");
|
||||
assert_eq!(keys[0].auth_type, "api_key");
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_key_concurrent_limit_create_and_list_responses() {
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
@@ -1380,6 +1477,121 @@ async fn gateway_export_preserves_distinct_imported_access_token_with_authorizat
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_generic_export_rejects_agent_identity_without_exposing_private_key() {
|
||||
let private_key = "agent-private-key-must-not-leak";
|
||||
let mut provider = sample_provider("provider-codex", "codex", 10);
|
||||
provider.provider_type = "codex".to_string();
|
||||
let mut key = sample_key(
|
||||
"key-codex-agent",
|
||||
"provider-codex",
|
||||
"openai:responses",
|
||||
"__placeholder__",
|
||||
);
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.encrypted_auth_config = Some(
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
&json!({
|
||||
"provider_type": "codex",
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_runtime_id": "runtime-must-not-leak",
|
||||
"agent_private_key": private_key,
|
||||
"task_id": "task-must-not-leak"
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.expect("Agent Identity auth config should encrypt"),
|
||||
);
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
vec![],
|
||||
vec![key],
|
||||
));
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
let reveal_response = client
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/keys/key-codex-agent/reveal"
|
||||
))
|
||||
.header(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("reveal request should complete");
|
||||
|
||||
assert_eq!(reveal_response.status(), StatusCode::BAD_REQUEST);
|
||||
let reveal_body = reveal_response
|
||||
.text()
|
||||
.await
|
||||
.expect("reveal error body should read");
|
||||
assert!(reveal_body.contains("专属 provider-oauth 管理面"));
|
||||
assert!(!reveal_body.contains(private_key));
|
||||
assert!(!reveal_body.contains("runtime-must-not-leak"));
|
||||
assert!(!reveal_body.contains("task-must-not-leak"));
|
||||
|
||||
let export_response = client
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/keys/key-codex-agent/export"
|
||||
))
|
||||
.header(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("export request should complete");
|
||||
|
||||
assert_eq!(export_response.status(), StatusCode::BAD_REQUEST);
|
||||
let export_body = export_response
|
||||
.text()
|
||||
.await
|
||||
.expect("export error body should read");
|
||||
assert!(export_body.contains("专属 provider-oauth 管理面"));
|
||||
assert!(!export_body.contains(private_key));
|
||||
assert!(!export_body.contains("runtime-must-not-leak"));
|
||||
assert!(!export_body.contains("task-must-not-leak"));
|
||||
|
||||
let list_response = client
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-codex/keys"
|
||||
))
|
||||
.header(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("list request should complete");
|
||||
|
||||
assert_eq!(list_response.status(), StatusCode::OK);
|
||||
let list_payload: serde_json::Value = list_response
|
||||
.json()
|
||||
.await
|
||||
.expect("list body should be JSON");
|
||||
let agent = list_payload
|
||||
.as_array()
|
||||
.and_then(|items| items.iter().find(|item| item["id"] == "key-codex-agent"))
|
||||
.expect("Agent Identity key should be listed");
|
||||
assert_eq!(agent["agent_identity"], true);
|
||||
assert_eq!(agent["can_export_oauth"], false);
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_clears_admin_provider_key_oauth_invalid_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
@@ -2601,11 +2813,27 @@ async fn gateway_handles_admin_keys_grouped_by_format_locally_with_trusted_admin
|
||||
key_b.created_at_unix_ms = Some(1_711_100_000);
|
||||
key_b.updated_at_unix_secs = Some(1_711_100_100);
|
||||
|
||||
let mut key_agent = sample_key(
|
||||
"key-codex-agent",
|
||||
"provider-codex",
|
||||
"openai:responses",
|
||||
"__placeholder__",
|
||||
);
|
||||
key_agent.auth_type = "oauth".to_string();
|
||||
key_agent.encrypted_auth_config = Some(
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"provider_type":"codex","auth_mode":"agentIdentity","agent_runtime_id":"runtime-1","agent_private_key":"base64-private-key","task_id":"task-1"}"#,
|
||||
)
|
||||
.expect("Agent Identity auth config should encrypt"),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(SummaryNullingProviderCatalogReadRepository::seed(
|
||||
vec![
|
||||
sample_provider("provider-openai", "openai", 10),
|
||||
sample_provider("provider-claude", "claude", 20)
|
||||
.with_transport_fields(false, false, true, None, None, None, None, None, None),
|
||||
sample_provider("provider-codex", "codex", 30),
|
||||
],
|
||||
vec![
|
||||
sample_endpoint(
|
||||
@@ -2620,8 +2848,14 @@ async fn gateway_handles_admin_keys_grouped_by_format_locally_with_trusted_admin
|
||||
"claude:messages",
|
||||
"https://api.claude.example",
|
||||
),
|
||||
sample_endpoint(
|
||||
"endpoint-codex-responses",
|
||||
"provider-codex",
|
||||
"openai:responses",
|
||||
"https://api.codex.example",
|
||||
),
|
||||
],
|
||||
vec![key_a, key_b],
|
||||
vec![key_a, key_b, key_agent],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
@@ -2665,6 +2899,11 @@ async fn gateway_handles_admin_keys_grouped_by_format_locally_with_trusted_admin
|
||||
);
|
||||
assert_eq!(payload["openai:chat"][0]["internal_priority"], 10);
|
||||
assert_eq!(payload["claude:messages"][0]["provider_active"], false);
|
||||
let agent_item = payload["openai:responses"]
|
||||
.as_array()
|
||||
.and_then(|items| items.iter().find(|item| item["id"] == "key-codex-agent"))
|
||||
.expect("Agent Identity key should be grouped");
|
||||
assert_eq!(agent_item["api_key_masked"], "[Agent Identity]");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
|
||||
@@ -8968,6 +8968,137 @@ async fn gateway_allows_management_token_with_pool_write_for_provider_oauth_batc
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_prevents_pool_write_token_from_importing_agent_identity_via_batch_routes() {
|
||||
run_admin_oauth_test(
|
||||
"gateway_prevents_pool_write_token_from_importing_agent_identity_via_batch_routes",
|
||||
gateway_prevents_pool_write_token_from_importing_agent_identity_via_batch_routes_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_prevents_pool_write_token_from_importing_agent_identity_via_batch_routes_impl() {
|
||||
let raw_token = "ae-provider-oauth-agent-identity-pool-write";
|
||||
let state = AppState::new().expect("gateway should build");
|
||||
let admin_user = state
|
||||
.create_local_auth_user_with_settings(
|
||||
Some("provider-oauth-agent-identity-pool@example.com".to_string()),
|
||||
true,
|
||||
"admin".to_string(),
|
||||
"hash".to_string(),
|
||||
"admin".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("admin user should be created")
|
||||
.expect("admin user should exist");
|
||||
let mut management_token = sample_management_token(
|
||||
"token-provider-oauth-agent-identity-pool",
|
||||
&admin_user.id,
|
||||
"provider-oauth-agent-identity-pool",
|
||||
true,
|
||||
);
|
||||
management_token.token.allowed_ips = None;
|
||||
management_token.token.permissions = Some(json!(["admin:pool:read", "admin:pool:write"]));
|
||||
let management_token_repository =
|
||||
Arc::new(InMemoryManagementTokenRepository::seed_with_hashes(
|
||||
vec![management_token],
|
||||
vec![(
|
||||
hash_management_token(raw_token),
|
||||
"token-provider-oauth-agent-identity-pool".to_string(),
|
||||
)],
|
||||
));
|
||||
|
||||
let mut provider = sample_provider("provider-codex", "codex", 10);
|
||||
provider.provider_type = "codex".to_string();
|
||||
let endpoint = sample_endpoint(
|
||||
"endpoint-codex-agent-identity",
|
||||
"provider-codex",
|
||||
"openai:chat",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
);
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
vec![endpoint],
|
||||
vec![],
|
||||
));
|
||||
let data_state =
|
||||
GatewayDataState::with_management_token_repository_for_tests(management_token_repository)
|
||||
.attach_provider_catalog_repository_for_tests(provider_catalog_repository.clone())
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY);
|
||||
let gateway = build_router_with_state(state.with_data_state_for_tests(data_state));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let credentials = json!({
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_runtime_id": "runtime-rbac-guard",
|
||||
"agent_private_key": "MC4CAQAwBQYDK2VwBCIEIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
|
||||
"task_id": "task-rbac-guard"
|
||||
})
|
||||
.to_string();
|
||||
let client = reqwest::Client::new();
|
||||
for path in [
|
||||
"/api/admin/provider-oauth/providers/provider-codex/batch-import",
|
||||
"/api/admin/provider-oauth/providers/provider-codex/batch-import/tasks",
|
||||
] {
|
||||
let response = client
|
||||
.post(format!("{gateway_url}{path}"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.bearer_auth(raw_token)
|
||||
.json(&json!({ "credentials": credentials }))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let status = response.status();
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(
|
||||
status,
|
||||
StatusCode::BAD_REQUEST,
|
||||
"path={path} payload={payload}"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["detail"],
|
||||
"Agent Identity JSON 必须使用专属导入接口"
|
||||
);
|
||||
}
|
||||
|
||||
let dedicated_response = client
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/agent-identity-import/tasks"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.bearer_auth(raw_token)
|
||||
.json(&json!({ "credentials": credentials }))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
let dedicated_status = dedicated_response.status();
|
||||
let dedicated_payload: serde_json::Value = dedicated_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
assert_eq!(
|
||||
dedicated_status,
|
||||
StatusCode::FORBIDDEN,
|
||||
"payload={dedicated_payload}"
|
||||
);
|
||||
assert_eq!(
|
||||
dedicated_payload["required_permission"],
|
||||
"admin:provider_oauth:write"
|
||||
);
|
||||
|
||||
let keys = provider_catalog_repository
|
||||
.list_keys_by_provider_ids(&["provider-codex".to_string()])
|
||||
.await
|
||||
.expect("keys should load");
|
||||
assert!(keys.is_empty());
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_rejects_management_token_without_pool_write_for_provider_oauth_batch_import() {
|
||||
run_admin_oauth_test(
|
||||
|
||||
@@ -3147,6 +3147,8 @@ async fn gateway_pool_keys_classify_oauth_credentials() {
|
||||
.expect("Agent Identity key should exist");
|
||||
assert_eq!(agent_identity_key["oauth_header_auth"], false);
|
||||
assert_eq!(agent_identity_key["agent_identity"], true);
|
||||
assert_eq!(agent_identity_key["can_refresh_oauth"], true);
|
||||
assert_eq!(agent_identity_key["can_export_oauth"], false);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -200,6 +200,18 @@ pub fn extract_execution_error_message(result: &ExecutionResult) -> Option<Strin
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
/// Keeps the structured upstream error intact for classifiers that depend on fields such as
|
||||
/// `error.code`, while retaining the execution-error fallback used by transport failures.
|
||||
pub fn extract_execution_error_detail(result: &ExecutionResult) -> Option<String> {
|
||||
result
|
||||
.body
|
||||
.as_ref()
|
||||
.and_then(|body| body.json_body.as_ref())
|
||||
.and_then(|body| serde_json::to_string(body).ok())
|
||||
.filter(|value| !value.is_empty())
|
||||
.or_else(|| extract_execution_error_message(result))
|
||||
}
|
||||
|
||||
pub fn quota_refresh_success_invalid_state(
|
||||
key: &StoredProviderCatalogKey,
|
||||
) -> (Option<u64>, Option<String>) {
|
||||
@@ -2109,7 +2121,7 @@ pub fn parse_chatgpt_web_conversation_init_response(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
codex_build_invalid_state, codex_runtime_invalid_reason,
|
||||
codex_build_invalid_state, codex_runtime_invalid_reason, extract_execution_error_detail,
|
||||
normalize_codex_reset_credit_consume_outcome, parse_antigravity_usage_response,
|
||||
parse_chatgpt_web_conversation_init_response, parse_codex_backend_me_response,
|
||||
parse_codex_usage_headers, parse_codex_wham_reset_credits_detail_response,
|
||||
@@ -2120,10 +2132,43 @@ mod tests {
|
||||
should_auto_remove_structured_reason, OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX,
|
||||
OAUTH_REFRESH_FAILED_PREFIX, OAUTH_REQUEST_FAILED_PREFIX,
|
||||
};
|
||||
use aether_contracts::{ExecutionResult, ResponseBody};
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use serde_json::json;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
#[test]
|
||||
fn execution_error_detail_preserves_structured_code_and_message() {
|
||||
let result = ExecutionResult {
|
||||
request_id: "quota-agent-identity".to_string(),
|
||||
candidate_id: None,
|
||||
status_code: 401,
|
||||
headers: BTreeMap::new(),
|
||||
body: Some(ResponseBody {
|
||||
json_body: Some(json!({
|
||||
"error": {
|
||||
"code": "invalid_task_id",
|
||||
"message": "registered task is no longer valid"
|
||||
}
|
||||
})),
|
||||
body_bytes_b64: None,
|
||||
}),
|
||||
telemetry: None,
|
||||
error: None,
|
||||
};
|
||||
|
||||
let detail = extract_execution_error_detail(&result)
|
||||
.expect("structured execution error should be retained");
|
||||
assert!(detail.contains(r#""code":"invalid_task_id""#));
|
||||
assert!(detail.contains(r#""message":"registered task is no longer valid""#));
|
||||
assert!(
|
||||
aether_provider_transport::is_codex_agent_identity_invalid_task_response(
|
||||
result.status_code,
|
||||
Some(&detail),
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_auto_remove_quota_exhausted_keys_defaults_to_false() {
|
||||
assert!(!provider_auto_remove_quota_exhausted_keys(None));
|
||||
|
||||
@@ -9,8 +9,9 @@ use sqlx::{
|
||||
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate,
|
||||
ProviderCatalogKeyListQuery, ProviderCatalogKeyRuntimeMetadataUpdate,
|
||||
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository, ProviderCatalogSnapshot,
|
||||
ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
|
||||
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
|
||||
ProviderCatalogReadRepository, ProviderCatalogSnapshot,
|
||||
ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
|
||||
@@ -1019,6 +1020,89 @@ WHERE id = ?
|
||||
Ok(rows_affected > 0)
|
||||
}
|
||||
|
||||
pub async fn compare_and_update_key_oauth_runtime_state(
|
||||
&self,
|
||||
update: &ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
validate_non_empty(&update.key_id, "provider catalog key_id")?;
|
||||
validate_non_empty(
|
||||
&update.encrypted_auth_config,
|
||||
"provider catalog OAuth auth_config",
|
||||
)?;
|
||||
if update
|
||||
.encrypted_api_key_update
|
||||
.as_deref()
|
||||
.is_some_and(|value| value.trim().is_empty())
|
||||
{
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"provider catalog OAuth api_key update must not be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if !update.status_snapshot_patch.is_object() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"provider catalog status snapshot patch must be an object".to_string(),
|
||||
));
|
||||
}
|
||||
if update
|
||||
.upstream_metadata_patch
|
||||
.as_ref()
|
||||
.is_some_and(|patch| !patch.is_object())
|
||||
{
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"provider catalog upstream metadata patch must be an object".to_string(),
|
||||
));
|
||||
}
|
||||
let mut builder =
|
||||
QueryBuilder::<MySql>::new("UPDATE provider_api_keys SET oauth_invalid_at = ");
|
||||
builder
|
||||
.push_bind(optional_i64_from_u64(
|
||||
update.oauth_invalid_at_unix_secs,
|
||||
"provider_api_keys.oauth_invalid_at",
|
||||
)?)
|
||||
.push(", oauth_invalid_reason = ")
|
||||
.push_bind(update.oauth_invalid_reason.as_deref())
|
||||
.push(", auth_config = ")
|
||||
.push_bind(&update.encrypted_auth_config);
|
||||
if let Some(encrypted_api_key) = update.encrypted_api_key_update.as_deref() {
|
||||
builder.push(", api_key = ").push_bind(encrypted_api_key);
|
||||
}
|
||||
if let Some(expires_at_unix_secs) = update.expires_at_unix_secs_update {
|
||||
builder
|
||||
.push(", expires_at = ")
|
||||
.push_bind(optional_i64_from_u64(
|
||||
expires_at_unix_secs,
|
||||
"provider_api_keys.expires_at",
|
||||
)?);
|
||||
}
|
||||
if let Some(metadata_patch) = update.upstream_metadata_patch.as_ref() {
|
||||
builder.push(", upstream_metadata = ");
|
||||
push_upstream_metadata_shallow_patch(&mut builder, metadata_patch)?;
|
||||
}
|
||||
builder.push(", status_snapshot = ");
|
||||
push_status_snapshot_shallow_patch(&mut builder, &update.status_snapshot_patch)?;
|
||||
if update.reset_error_count {
|
||||
builder.push(", error_count = 0");
|
||||
}
|
||||
builder
|
||||
.push(", updated_at = ")
|
||||
.push_bind(
|
||||
update
|
||||
.updated_at_unix_secs
|
||||
.unwrap_or_else(current_unix_secs) as i64,
|
||||
)
|
||||
.push(" WHERE id = ")
|
||||
.push_bind(&update.key_id)
|
||||
.push(" AND auth_config <=> ")
|
||||
.push_bind(update.expected_encrypted_auth_config.as_deref());
|
||||
let rows_affected = builder
|
||||
.build()
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected();
|
||||
Ok(rows_affected > 0)
|
||||
}
|
||||
|
||||
pub async fn update_key_health_state(
|
||||
&self,
|
||||
key_id: &str,
|
||||
@@ -1153,6 +1237,13 @@ WHERE id = ?
|
||||
.push_bind(optional_i64_from_u32(expected.last_rpm_peak))
|
||||
.push(" AND concurrent_429_count <=> ")
|
||||
.push_bind(optional_i64_from_u32(expected.concurrent_429_count));
|
||||
if let Some(expected_encrypted_auth_config) =
|
||||
update.expected_encrypted_auth_config.as_deref()
|
||||
{
|
||||
builder
|
||||
.push(" AND auth_config <=> ")
|
||||
.push_bind(expected_encrypted_auth_config);
|
||||
}
|
||||
let rows_affected = builder
|
||||
.build()
|
||||
.execute(&self.pool)
|
||||
@@ -1267,6 +1358,7 @@ SET health_by_format = ?, circuit_breaker_by_format = ?, updated_at = ?
|
||||
WHERE id = ?
|
||||
AND JSON_EXTRACT(health_by_format, '$') <=> CAST(? AS JSON)
|
||||
AND JSON_EXTRACT(circuit_breaker_by_format, '$') <=> CAST(? AS JSON)
|
||||
AND (? IS NULL OR auth_config <=> ?)
|
||||
"#,
|
||||
)
|
||||
.bind(optional_json_to_string(
|
||||
@@ -1287,6 +1379,8 @@ WHERE id = ?
|
||||
&update.expected_circuit_breaker_by_format,
|
||||
"provider_api_keys.circuit_breaker_by_format",
|
||||
)?)
|
||||
.bind(update.expected_encrypted_auth_config.as_deref())
|
||||
.bind(update.expected_encrypted_auth_config.as_deref())
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
@@ -1616,6 +1710,13 @@ impl ProviderCatalogWriteRepository for MysqlProviderCatalogReadRepository {
|
||||
.await
|
||||
}
|
||||
|
||||
async fn compare_and_update_key_oauth_runtime_state(
|
||||
&self,
|
||||
update: &ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
Self::compare_and_update_key_oauth_runtime_state(self, update).await
|
||||
}
|
||||
|
||||
async fn update_key_health_state(
|
||||
&self,
|
||||
key_id: &str,
|
||||
@@ -1760,6 +1861,42 @@ fn push_status_snapshot_shallow_patch<'args>(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn push_upstream_metadata_shallow_patch<'args>(
|
||||
builder: &mut QueryBuilder<'args, MySql>,
|
||||
patch: &serde_json::Value,
|
||||
) -> Result<(), DataLayerError> {
|
||||
let object = patch.as_object().ok_or_else(|| {
|
||||
DataLayerError::InvalidInput(
|
||||
"provider catalog upstream metadata patch must be an object".to_string(),
|
||||
)
|
||||
})?;
|
||||
if object.is_empty() {
|
||||
builder.push("COALESCE(NULLIF(upstream_metadata, ''), '{}')");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
builder.push("JSON_SET(COALESCE(NULLIF(upstream_metadata, ''), '{}')");
|
||||
for (field, value) in object {
|
||||
let path = format!(
|
||||
"$.{}",
|
||||
serde_json::to_string(field).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"provider_api_keys.upstream_metadata field is not serializable: {err}"
|
||||
))
|
||||
})?
|
||||
);
|
||||
let value = serde_json::to_string(value).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"provider_api_keys.upstream_metadata value is not serializable: {err}"
|
||||
))
|
||||
})?;
|
||||
builder.push(", ").push_bind(path).push(", CAST(");
|
||||
builder.push_bind(value).push(" AS JSON)");
|
||||
}
|
||||
builder.push(")");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_provider(provider: &StoredProviderCatalogProvider) -> Result<(), DataLayerError> {
|
||||
validate_non_empty(&provider.id, "provider catalog provider.id")?;
|
||||
validate_non_empty(&provider.name, "provider catalog provider.name")?;
|
||||
|
||||
@@ -11,9 +11,10 @@ use sqlx::{
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate,
|
||||
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery,
|
||||
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
|
||||
ProviderCatalogReadRepository, ProviderCatalogUpstreamMetadataNamespaceUpdate,
|
||||
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
|
||||
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository,
|
||||
ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
|
||||
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
||||
};
|
||||
@@ -390,6 +391,7 @@ SET
|
||||
WHERE id = $1
|
||||
AND health_by_format::jsonb IS NOT DISTINCT FROM $4::jsonb
|
||||
AND circuit_breaker_by_format::jsonb IS NOT DISTINCT FROM $5::jsonb
|
||||
AND ($6::text IS NULL OR auth_config IS NOT DISTINCT FROM $6)
|
||||
"#;
|
||||
|
||||
const KEY_RUNTIME_METADATA_CAS_SQL: &str = r#"
|
||||
@@ -880,6 +882,84 @@ WHERE id = $1
|
||||
Ok(rows_affected > 0)
|
||||
}
|
||||
|
||||
pub async fn compare_and_update_key_oauth_runtime_state(
|
||||
&self,
|
||||
update: &ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
if update.key_id.trim().is_empty()
|
||||
|| update.encrypted_auth_config.trim().is_empty()
|
||||
|| update
|
||||
.encrypted_api_key_update
|
||||
.as_deref()
|
||||
.is_some_and(|value| value.trim().is_empty())
|
||||
|| !update.status_snapshot_patch.is_object()
|
||||
|| update
|
||||
.upstream_metadata_patch
|
||||
.as_ref()
|
||||
.is_some_and(|patch| !patch.is_object())
|
||||
{
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"provider catalog OAuth runtime CAS requires key_id, auth_config, and object status patch"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
let rows_affected = sqlx::query(
|
||||
r#"
|
||||
UPDATE provider_api_keys
|
||||
SET
|
||||
oauth_invalid_at = CASE
|
||||
WHEN $2::double precision IS NULL THEN NULL
|
||||
ELSE TO_TIMESTAMP($2::double precision)
|
||||
END,
|
||||
oauth_invalid_reason = $3,
|
||||
auth_config = $4,
|
||||
api_key = CASE
|
||||
WHEN $5::text IS NULL THEN api_key
|
||||
ELSE $5
|
||||
END,
|
||||
expires_at = CASE
|
||||
WHEN $6::boolean IS FALSE THEN expires_at
|
||||
WHEN $7::double precision IS NULL THEN NULL
|
||||
ELSE TO_TIMESTAMP($7::double precision)
|
||||
END,
|
||||
upstream_metadata = CASE
|
||||
WHEN $8::jsonb IS NULL THEN upstream_metadata
|
||||
ELSE COALESCE(upstream_metadata, '{}'::jsonb) || $8::jsonb
|
||||
END,
|
||||
status_snapshot = (COALESCE(status_snapshot::jsonb, '{}'::jsonb) || $9::jsonb)::json,
|
||||
error_count = CASE WHEN $10::boolean THEN 0 ELSE error_count END,
|
||||
updated_at = CASE
|
||||
WHEN $11::double precision IS NULL THEN NOW()
|
||||
ELSE TO_TIMESTAMP($11::double precision)
|
||||
END
|
||||
WHERE id = $1
|
||||
AND auth_config IS NOT DISTINCT FROM $12
|
||||
"#,
|
||||
)
|
||||
.bind(&update.key_id)
|
||||
.bind(update.oauth_invalid_at_unix_secs.map(|value| value as f64))
|
||||
.bind(update.oauth_invalid_reason.as_deref())
|
||||
.bind(&update.encrypted_auth_config)
|
||||
.bind(update.encrypted_api_key_update.as_deref())
|
||||
.bind(update.expires_at_unix_secs_update.is_some())
|
||||
.bind(
|
||||
update
|
||||
.expires_at_unix_secs_update
|
||||
.flatten()
|
||||
.map(|value| value as f64),
|
||||
)
|
||||
.bind(update.upstream_metadata_patch.as_ref())
|
||||
.bind(&update.status_snapshot_patch)
|
||||
.bind(update.reset_error_count)
|
||||
.bind(update.updated_at_unix_secs.map(|value| value as f64))
|
||||
.bind(update.expected_encrypted_auth_config.as_deref())
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
.rows_affected();
|
||||
Ok(rows_affected > 0)
|
||||
}
|
||||
|
||||
pub async fn create_provider(
|
||||
&self,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
@@ -2230,6 +2310,7 @@ WHERE id = $1
|
||||
AND CAST(EXTRACT(EPOCH FROM last_probe_increase_at) AS BIGINT) IS NOT DISTINCT FROM $19
|
||||
AND last_rpm_peak IS NOT DISTINCT FROM $20
|
||||
AND concurrent_429_count IS NOT DISTINCT FROM $21
|
||||
AND ($22::text IS NULL OR auth_config IS NOT DISTINCT FROM $22)
|
||||
"#,
|
||||
)
|
||||
.bind(&update.key_id)
|
||||
@@ -2264,6 +2345,7 @@ WHERE id = $1
|
||||
)
|
||||
.bind(expected.last_rpm_peak.map(|value| value as i32))
|
||||
.bind(expected.concurrent_429_count.map(|value| value as i32))
|
||||
.bind(update.expected_encrypted_auth_config.as_deref())
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
@@ -2336,6 +2418,7 @@ WHERE id = $1
|
||||
.bind(&update.circuit_breaker_by_format)
|
||||
.bind(&update.expected_health_by_format)
|
||||
.bind(&update.expected_circuit_breaker_by_format)
|
||||
.bind(update.expected_encrypted_auth_config.as_deref())
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
@@ -2603,6 +2686,13 @@ impl ProviderCatalogWriteRepository for SqlxProviderCatalogReadRepository {
|
||||
.await
|
||||
}
|
||||
|
||||
async fn compare_and_update_key_oauth_runtime_state(
|
||||
&self,
|
||||
update: &ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
Self::compare_and_update_key_oauth_runtime_state(self, update).await
|
||||
}
|
||||
|
||||
async fn update_key_health_state(
|
||||
&self,
|
||||
key_id: &str,
|
||||
|
||||
@@ -10,9 +10,10 @@ use sqlx::{
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate,
|
||||
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery,
|
||||
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
|
||||
ProviderCatalogReadRepository, ProviderCatalogUpstreamMetadataNamespaceUpdate,
|
||||
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
|
||||
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository,
|
||||
ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
|
||||
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
||||
};
|
||||
@@ -1443,6 +1444,89 @@ WHERE id = ?
|
||||
Ok(rows_affected > 0)
|
||||
}
|
||||
|
||||
pub async fn compare_and_update_key_oauth_runtime_state(
|
||||
&self,
|
||||
update: &ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
validate_non_empty(&update.key_id, "provider catalog key_id")?;
|
||||
validate_non_empty(
|
||||
&update.encrypted_auth_config,
|
||||
"provider catalog OAuth auth_config",
|
||||
)?;
|
||||
if update
|
||||
.encrypted_api_key_update
|
||||
.as_deref()
|
||||
.is_some_and(|value| value.trim().is_empty())
|
||||
{
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"provider catalog OAuth api_key update must not be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if !update.status_snapshot_patch.is_object() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"provider catalog status snapshot patch must be an object".to_string(),
|
||||
));
|
||||
}
|
||||
if update
|
||||
.upstream_metadata_patch
|
||||
.as_ref()
|
||||
.is_some_and(|patch| !patch.is_object())
|
||||
{
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"provider catalog upstream metadata patch must be an object".to_string(),
|
||||
));
|
||||
}
|
||||
let mut builder =
|
||||
QueryBuilder::<Sqlite>::new("UPDATE provider_api_keys SET oauth_invalid_at = ");
|
||||
builder
|
||||
.push_bind(optional_i64_from_u64(
|
||||
update.oauth_invalid_at_unix_secs,
|
||||
"provider_api_keys.oauth_invalid_at",
|
||||
)?)
|
||||
.push(", oauth_invalid_reason = ")
|
||||
.push_bind(update.oauth_invalid_reason.as_deref())
|
||||
.push(", auth_config = ")
|
||||
.push_bind(&update.encrypted_auth_config);
|
||||
if let Some(encrypted_api_key) = update.encrypted_api_key_update.as_deref() {
|
||||
builder.push(", api_key = ").push_bind(encrypted_api_key);
|
||||
}
|
||||
if let Some(expires_at_unix_secs) = update.expires_at_unix_secs_update {
|
||||
builder
|
||||
.push(", expires_at = ")
|
||||
.push_bind(optional_i64_from_u64(
|
||||
expires_at_unix_secs,
|
||||
"provider_api_keys.expires_at",
|
||||
)?);
|
||||
}
|
||||
if let Some(metadata_patch) = update.upstream_metadata_patch.as_ref() {
|
||||
builder.push(", upstream_metadata = ");
|
||||
push_upstream_metadata_shallow_patch(&mut builder, metadata_patch)?;
|
||||
}
|
||||
builder.push(", status_snapshot = ");
|
||||
push_status_snapshot_shallow_patch(&mut builder, &update.status_snapshot_patch)?;
|
||||
if update.reset_error_count {
|
||||
builder.push(", error_count = 0");
|
||||
}
|
||||
builder
|
||||
.push(", updated_at = ")
|
||||
.push_bind(
|
||||
update
|
||||
.updated_at_unix_secs
|
||||
.unwrap_or_else(current_unix_secs) as i64,
|
||||
)
|
||||
.push(" WHERE id = ")
|
||||
.push_bind(&update.key_id)
|
||||
.push(" AND auth_config IS ")
|
||||
.push_bind(update.expected_encrypted_auth_config.as_deref());
|
||||
let rows_affected = builder
|
||||
.build()
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected();
|
||||
Ok(rows_affected > 0)
|
||||
}
|
||||
|
||||
pub async fn update_key_health_state(
|
||||
&self,
|
||||
key_id: &str,
|
||||
@@ -1577,6 +1661,13 @@ WHERE id = ?
|
||||
.push_bind(optional_i64_from_u32(expected.last_rpm_peak))
|
||||
.push(" AND concurrent_429_count IS ")
|
||||
.push_bind(optional_i64_from_u32(expected.concurrent_429_count));
|
||||
if let Some(expected_encrypted_auth_config) =
|
||||
update.expected_encrypted_auth_config.as_deref()
|
||||
{
|
||||
builder
|
||||
.push(" AND auth_config IS ")
|
||||
.push_bind(expected_encrypted_auth_config);
|
||||
}
|
||||
let rows_affected = builder
|
||||
.build()
|
||||
.execute(&self.pool)
|
||||
@@ -1703,6 +1794,7 @@ SET health_by_format = ?, circuit_breaker_by_format = ?, updated_at = ?
|
||||
WHERE id = ?
|
||||
AND json(health_by_format) IS json(?)
|
||||
AND json(circuit_breaker_by_format) IS json(?)
|
||||
AND (? IS NULL OR auth_config IS ?)
|
||||
"#,
|
||||
)
|
||||
.bind(optional_json_to_string(
|
||||
@@ -1723,6 +1815,8 @@ WHERE id = ?
|
||||
&update.expected_circuit_breaker_by_format,
|
||||
"provider_api_keys.circuit_breaker_by_format",
|
||||
)?)
|
||||
.bind(update.expected_encrypted_auth_config.as_deref())
|
||||
.bind(update.expected_encrypted_auth_config.as_deref())
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
@@ -2038,6 +2132,13 @@ impl ProviderCatalogWriteRepository for SqliteProviderCatalogReadRepository {
|
||||
.await
|
||||
}
|
||||
|
||||
async fn compare_and_update_key_oauth_runtime_state(
|
||||
&self,
|
||||
update: &ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
Self::compare_and_update_key_oauth_runtime_state(self, update).await
|
||||
}
|
||||
|
||||
async fn update_key_health_state(
|
||||
&self,
|
||||
key_id: &str,
|
||||
@@ -2237,6 +2338,42 @@ fn push_status_snapshot_shallow_patch<'args>(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn push_upstream_metadata_shallow_patch<'args>(
|
||||
builder: &mut QueryBuilder<'args, Sqlite>,
|
||||
patch: &serde_json::Value,
|
||||
) -> Result<(), DataLayerError> {
|
||||
let object = patch.as_object().ok_or_else(|| {
|
||||
DataLayerError::InvalidInput(
|
||||
"provider catalog upstream metadata patch must be an object".to_string(),
|
||||
)
|
||||
})?;
|
||||
if object.is_empty() {
|
||||
builder.push("COALESCE(NULLIF(upstream_metadata, ''), '{}')");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
builder.push("json_set(COALESCE(NULLIF(upstream_metadata, ''), '{}')");
|
||||
for (field, value) in object {
|
||||
let path = format!(
|
||||
"$.{}",
|
||||
serde_json::to_string(field).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"provider_api_keys.upstream_metadata field is not serializable: {err}"
|
||||
))
|
||||
})?
|
||||
);
|
||||
let value = serde_json::to_string(value).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"provider_api_keys.upstream_metadata value is not serializable: {err}"
|
||||
))
|
||||
})?;
|
||||
builder.push(", ").push_bind(path).push(", json(");
|
||||
builder.push_bind(value).push(")");
|
||||
}
|
||||
builder.push(")");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_provider(provider: &StoredProviderCatalogProvider) -> Result<(), DataLayerError> {
|
||||
validate_non_empty(&provider.id, "provider catalog provider.id")?;
|
||||
validate_non_empty(&provider.name, "provider catalog provider.name")?;
|
||||
@@ -2801,9 +2938,9 @@ mod tests {
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate,
|
||||
ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder,
|
||||
ProviderCatalogKeyListQuery, ProviderCatalogKeyRuntimeMetadataUpdate,
|
||||
ProviderCatalogUpstreamMetadataNamespaceUpdate, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
|
||||
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogUpstreamMetadataNamespaceUpdate,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
@@ -2923,6 +3060,7 @@ mod tests {
|
||||
"observation_count": 1,
|
||||
"known_boundary": "old"
|
||||
}));
|
||||
key.encrypted_auth_config = Some("auth-current".to_string());
|
||||
let mut stale_admin_key = key.clone();
|
||||
repository
|
||||
.create_key(&key)
|
||||
@@ -2931,6 +3069,7 @@ mod tests {
|
||||
|
||||
let health_update = ProviderCatalogKeyHealthStateUpdate {
|
||||
key_id: "runtime-key".to_string(),
|
||||
expected_encrypted_auth_config: None,
|
||||
expected_health_by_format: key.health_by_format.clone(),
|
||||
expected_circuit_breaker_by_format: None,
|
||||
health_by_format: Some(json!({"openai:chat":{"consecutive_failures":2}})),
|
||||
@@ -2955,24 +3094,35 @@ mod tests {
|
||||
let mut next = expected.clone();
|
||||
next.learned_rpm_limit = Some(8);
|
||||
next.rpm_429_count = Some(2);
|
||||
let adaptive_update = ProviderCatalogKeyAdaptiveStateUpdate {
|
||||
key_id: "runtime-key".to_string(),
|
||||
expected_encrypted_auth_config: Some("auth-current".to_string()),
|
||||
expected: expected.clone(),
|
||||
next,
|
||||
status_snapshot_patch: json!({
|
||||
"observation_count": 2,
|
||||
"learning_confidence": 0.5,
|
||||
"known_boundary": null,
|
||||
"quota": {"remaining": 0}
|
||||
}),
|
||||
updated_at_unix_secs: Some(200),
|
||||
};
|
||||
let stale_generation_update = ProviderCatalogKeyAdaptiveStateUpdate {
|
||||
expected_encrypted_auth_config: Some("auth-stale".to_string()),
|
||||
..adaptive_update.clone()
|
||||
};
|
||||
assert!(!repository
|
||||
.compare_and_update_key_adaptive_state(&stale_generation_update)
|
||||
.await
|
||||
.expect("stale auth generation should conflict"));
|
||||
assert!(repository
|
||||
.compare_and_update_key_adaptive_state(&ProviderCatalogKeyAdaptiveStateUpdate {
|
||||
key_id: "runtime-key".to_string(),
|
||||
expected: expected.clone(),
|
||||
next,
|
||||
status_snapshot_patch: json!({
|
||||
"observation_count": 2,
|
||||
"learning_confidence": 0.5,
|
||||
"known_boundary": null,
|
||||
"quota": {"remaining": 0}
|
||||
}),
|
||||
updated_at_unix_secs: Some(200),
|
||||
})
|
||||
.compare_and_update_key_adaptive_state(&adaptive_update)
|
||||
.await
|
||||
.expect("adaptive CAS should succeed"));
|
||||
assert!(!repository
|
||||
.compare_and_update_key_adaptive_state(&ProviderCatalogKeyAdaptiveStateUpdate {
|
||||
key_id: "runtime-key".to_string(),
|
||||
expected_encrypted_auth_config: Some("auth-current".to_string()),
|
||||
expected: expected.clone(),
|
||||
next: expected,
|
||||
status_snapshot_patch: json!({}),
|
||||
@@ -3035,6 +3185,160 @@ mod tests {
|
||||
assert_eq!(status["known_boundary"], serde_json::Value::Null);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sqlite_oauth_runtime_cas_fences_auth_config_and_preserves_admin_fields() {
|
||||
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect("sqlite::memory:")
|
||||
.await
|
||||
.expect("sqlite pool should connect");
|
||||
run_migrations(&pool)
|
||||
.await
|
||||
.expect("sqlite migrations should run");
|
||||
let repository = SqliteProviderCatalogReadRepository::new(pool);
|
||||
repository
|
||||
.create_provider(
|
||||
&StoredProviderCatalogProvider::new(
|
||||
"oauth-cas-provider".to_string(),
|
||||
"OAuth CAS Provider".to_string(),
|
||||
None,
|
||||
"codex".to_string(),
|
||||
)
|
||||
.expect("provider should build"),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("provider should create");
|
||||
|
||||
let mut key = StoredProviderCatalogKey::new(
|
||||
"oauth-cas-key".to_string(),
|
||||
"oauth-cas-provider".to_string(),
|
||||
"Admin Managed Name".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
false,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(json!(["openai:responses"])),
|
||||
Some("encrypted-api-key".to_string()),
|
||||
Some("encrypted-auth-v1".to_string()),
|
||||
None,
|
||||
Some(json!({"openai:responses": 17})),
|
||||
Some(json!(["gpt-5"])),
|
||||
Some(4_102_444_800),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build");
|
||||
key.note = Some("admin note".to_string());
|
||||
key.internal_priority = 23;
|
||||
key.status_snapshot = Some(json!({
|
||||
"oauth": {"invalid": true, "source": "old-task"},
|
||||
"quota": {"remaining": 7},
|
||||
"admin": {"label": "keep"}
|
||||
}));
|
||||
repository
|
||||
.create_key(&key)
|
||||
.await
|
||||
.expect("key should create");
|
||||
|
||||
let update = ProviderCatalogKeyOAuthRuntimeStateCasUpdate {
|
||||
key_id: key.id.clone(),
|
||||
expected_encrypted_auth_config: Some("encrypted-auth-v1".to_string()),
|
||||
encrypted_auth_config: "encrypted-auth-v2".to_string(),
|
||||
encrypted_api_key_update: Some("encrypted-api-v2".to_string()),
|
||||
expires_at_unix_secs_update: Some(Some(4_102_555_900)),
|
||||
oauth_invalid_at_unix_secs: None,
|
||||
oauth_invalid_reason: None,
|
||||
upstream_metadata_patch: Some(json!({"codex": {"remaining": 3}})),
|
||||
status_snapshot_patch: json!({
|
||||
"oauth": {"invalid": false, "task_id": "task-v2"},
|
||||
"runtime": {"generation": 2}
|
||||
}),
|
||||
reset_error_count: false,
|
||||
updated_at_unix_secs: Some(200),
|
||||
};
|
||||
assert!(repository
|
||||
.compare_and_update_key_oauth_runtime_state(&update)
|
||||
.await
|
||||
.expect("matching OAuth runtime CAS should succeed"));
|
||||
|
||||
let stored = repository
|
||||
.list_keys_by_ids(&[key.id.clone()])
|
||||
.await
|
||||
.expect("key should reload")
|
||||
.pop()
|
||||
.expect("key should exist");
|
||||
assert_eq!(
|
||||
stored.encrypted_auth_config.as_deref(),
|
||||
Some("encrypted-auth-v2")
|
||||
);
|
||||
assert_eq!(stored.oauth_invalid_at_unix_secs, None);
|
||||
assert_eq!(stored.oauth_invalid_reason, None);
|
||||
assert_eq!(stored.name, "Admin Managed Name");
|
||||
assert!(!stored.is_active);
|
||||
assert_eq!(stored.note.as_deref(), Some("admin note"));
|
||||
assert_eq!(stored.internal_priority, 23);
|
||||
assert_eq!(
|
||||
stored.global_priority_by_format,
|
||||
Some(json!({"openai:responses": 17}))
|
||||
);
|
||||
assert_eq!(stored.allowed_models, Some(json!(["gpt-5"])));
|
||||
assert_eq!(
|
||||
stored.encrypted_api_key.as_deref(),
|
||||
Some("encrypted-api-v2")
|
||||
);
|
||||
assert_eq!(stored.expires_at_unix_secs, Some(4_102_555_900));
|
||||
assert_eq!(
|
||||
stored.upstream_metadata.as_ref().unwrap()["codex"]["remaining"],
|
||||
3
|
||||
);
|
||||
let status = stored.status_snapshot.expect("status should exist");
|
||||
assert_eq!(
|
||||
status["oauth"],
|
||||
json!({"invalid": false, "task_id": "task-v2"})
|
||||
);
|
||||
assert!(status["oauth"].get("source").is_none());
|
||||
assert_eq!(status["quota"], json!({"remaining": 7}));
|
||||
assert_eq!(status["admin"], json!({"label": "keep"}));
|
||||
assert_eq!(status["runtime"], json!({"generation": 2}));
|
||||
|
||||
let stale_update = ProviderCatalogKeyOAuthRuntimeStateCasUpdate {
|
||||
expected_encrypted_auth_config: Some("encrypted-auth-v1".to_string()),
|
||||
encrypted_auth_config: "encrypted-auth-v3".to_string(),
|
||||
status_snapshot_patch: json!({"quota": {"remaining": 0}}),
|
||||
updated_at_unix_secs: Some(201),
|
||||
..update
|
||||
};
|
||||
assert!(!repository
|
||||
.compare_and_update_key_oauth_runtime_state(&stale_update)
|
||||
.await
|
||||
.expect("stale OAuth runtime CAS should conflict"));
|
||||
|
||||
let stored_after_stale = repository
|
||||
.list_keys_by_ids(&[key.id])
|
||||
.await
|
||||
.expect("key should reload after stale CAS")
|
||||
.pop()
|
||||
.expect("key should exist");
|
||||
assert_eq!(
|
||||
stored_after_stale.encrypted_auth_config.as_deref(),
|
||||
Some("encrypted-auth-v2")
|
||||
);
|
||||
assert_eq!(
|
||||
stored_after_stale
|
||||
.status_snapshot
|
||||
.as_ref()
|
||||
.expect("status should remain")["quota"],
|
||||
json!({"remaining": 7})
|
||||
);
|
||||
assert_eq!(
|
||||
stored_after_stale.upstream_metadata.as_ref().unwrap()["codex"]["remaining"],
|
||||
3
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sqlite_repository_writes_provider_catalog_contract_views() {
|
||||
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
||||
|
||||
@@ -5,9 +5,10 @@ pub use snapshot::ProviderCatalogSnapshot;
|
||||
pub use types::{
|
||||
ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate,
|
||||
ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery,
|
||||
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
|
||||
ProviderCatalogReadRepository, ProviderCatalogUpstreamMetadataNamespaceUpdate,
|
||||
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
|
||||
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository,
|
||||
ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
|
||||
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
||||
};
|
||||
|
||||
@@ -31,6 +31,9 @@ impl ProviderCatalogKeyAdaptiveState {
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct ProviderCatalogKeyAdaptiveStateUpdate {
|
||||
pub key_id: String,
|
||||
/// Optional auth_config fence for request-owned adaptive feedback.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub expected_encrypted_auth_config: Option<String>,
|
||||
pub expected: ProviderCatalogKeyAdaptiveState,
|
||||
pub next: ProviderCatalogKeyAdaptiveState,
|
||||
/// Top-level status fields owned by adaptive rate-limit learning.
|
||||
@@ -64,9 +67,37 @@ pub struct ProviderCatalogKeyStatusSnapshotUpdate {
|
||||
pub updated_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
/// Agent/runtime-owned OAuth state update fenced by the exact encrypted
|
||||
/// auth_config observed before the refresh started. Repositories must update
|
||||
/// only these fields and return `false` when the expected config changed.
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct ProviderCatalogKeyOAuthRuntimeStateCasUpdate {
|
||||
pub key_id: String,
|
||||
pub expected_encrypted_auth_config: Option<String>,
|
||||
pub encrypted_auth_config: String,
|
||||
/// Optional access-token ciphertext replacement owned by refresh success.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub encrypted_api_key_update: Option<String>,
|
||||
/// `None` preserves expiry; `Some(None)` clears it; `Some(Some(_))` replaces it.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub expires_at_unix_secs_update: Option<Option<u64>>,
|
||||
pub oauth_invalid_at_unix_secs: Option<u64>,
|
||||
pub oauth_invalid_reason: Option<String>,
|
||||
/// Top-level runtime metadata namespaces to merge in the same fenced write.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub upstream_metadata_patch: Option<serde_json::Value>,
|
||||
pub status_snapshot_patch: serde_json::Value,
|
||||
#[serde(default)]
|
||||
pub reset_error_count: bool,
|
||||
pub updated_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct ProviderCatalogKeyHealthStateUpdate {
|
||||
pub key_id: String,
|
||||
/// Optional auth_config fence for lifecycle-owned health recovery.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub expected_encrypted_auth_config: Option<String>,
|
||||
pub expected_health_by_format: Option<serde_json::Value>,
|
||||
pub expected_circuit_breaker_by_format: Option<serde_json::Value>,
|
||||
pub health_by_format: Option<serde_json::Value>,
|
||||
@@ -829,6 +860,18 @@ pub trait ProviderCatalogWriteRepository: Send + Sync {
|
||||
updated_at_unix_secs: Option<u64>,
|
||||
) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
/// Compare-and-swap OAuth runtime credentials/status without replacing any
|
||||
/// administrator-owned key fields.
|
||||
async fn compare_and_update_key_oauth_runtime_state(
|
||||
&self,
|
||||
_update: &ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
|
||||
) -> Result<bool, crate::DataLayerError> {
|
||||
Err(crate::DataLayerError::InvalidConfiguration(
|
||||
"provider catalog OAuth runtime CAS updates are not supported by this repository"
|
||||
.to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
async fn update_key_health_state(
|
||||
&self,
|
||||
key_id: &str,
|
||||
|
||||
+9
-15
@@ -8,11 +8,9 @@ BEGIN
|
||||
AND column_name = 'enabled'
|
||||
) THEN
|
||||
UPDATE public.providers
|
||||
SET
|
||||
is_active = enabled,
|
||||
updated_at = NOW()
|
||||
WHERE enabled IS NOT NULL
|
||||
AND is_active IS DISTINCT FROM enabled;
|
||||
SET enabled = is_active
|
||||
WHERE is_active IS NOT NULL
|
||||
AND enabled IS DISTINCT FROM is_active;
|
||||
END IF;
|
||||
|
||||
IF EXISTS (
|
||||
@@ -23,11 +21,9 @@ BEGIN
|
||||
AND column_name = 'enabled'
|
||||
) THEN
|
||||
UPDATE public.provider_endpoints
|
||||
SET
|
||||
is_active = enabled,
|
||||
updated_at = NOW()
|
||||
WHERE enabled IS NOT NULL
|
||||
AND is_active IS DISTINCT FROM enabled;
|
||||
SET enabled = is_active
|
||||
WHERE is_active IS NOT NULL
|
||||
AND enabled IS DISTINCT FROM is_active;
|
||||
END IF;
|
||||
|
||||
IF EXISTS (
|
||||
@@ -38,10 +34,8 @@ BEGIN
|
||||
AND column_name = 'enabled'
|
||||
) THEN
|
||||
UPDATE public.models
|
||||
SET
|
||||
is_active = enabled,
|
||||
updated_at = NOW()
|
||||
WHERE enabled IS NOT NULL
|
||||
AND is_active IS DISTINCT FROM enabled;
|
||||
SET enabled = is_active
|
||||
WHERE is_active IS NOT NULL
|
||||
AND enabled IS DISTINCT FROM is_active;
|
||||
END IF;
|
||||
END $$;
|
||||
|
||||
-59
@@ -1,59 +0,0 @@
|
||||
DO $$
|
||||
BEGIN
|
||||
IF EXISTS (
|
||||
SELECT 1
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = 'public'
|
||||
AND table_name = 'providers'
|
||||
AND column_name = 'enabled'
|
||||
) AND EXISTS (
|
||||
SELECT 1
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = 'public'
|
||||
AND table_name = 'providers'
|
||||
AND column_name = 'is_active'
|
||||
) THEN
|
||||
UPDATE public.providers
|
||||
SET enabled = is_active
|
||||
WHERE is_active IS NOT NULL
|
||||
AND enabled IS DISTINCT FROM is_active;
|
||||
END IF;
|
||||
|
||||
IF EXISTS (
|
||||
SELECT 1
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = 'public'
|
||||
AND table_name = 'provider_endpoints'
|
||||
AND column_name = 'enabled'
|
||||
) AND EXISTS (
|
||||
SELECT 1
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = 'public'
|
||||
AND table_name = 'provider_endpoints'
|
||||
AND column_name = 'is_active'
|
||||
) THEN
|
||||
UPDATE public.provider_endpoints
|
||||
SET enabled = is_active
|
||||
WHERE is_active IS NOT NULL
|
||||
AND enabled IS DISTINCT FROM is_active;
|
||||
END IF;
|
||||
|
||||
IF EXISTS (
|
||||
SELECT 1
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = 'public'
|
||||
AND table_name = 'models'
|
||||
AND column_name = 'enabled'
|
||||
) AND EXISTS (
|
||||
SELECT 1
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = 'public'
|
||||
AND table_name = 'models'
|
||||
AND column_name = 'is_active'
|
||||
) THEN
|
||||
UPDATE public.models
|
||||
SET enabled = is_active
|
||||
WHERE is_active IS NOT NULL
|
||||
AND enabled IS DISTINCT FROM is_active;
|
||||
END IF;
|
||||
END $$;
|
||||
@@ -4,8 +4,7 @@ use std::{
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use sha2::{Digest, Sha384};
|
||||
use sqlx::{query, query_as, query_scalar, Connection, PgConnection, PgPool};
|
||||
use sqlx::{query, query_scalar, Connection, PgConnection, PgPool};
|
||||
|
||||
use super::{
|
||||
pending_backfills, pending_backfills_from_applied, pending_mysql_backfills,
|
||||
@@ -16,54 +15,22 @@ use crate::lifecycle::migrate::prepare_database_for_startup;
|
||||
use crate::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig};
|
||||
|
||||
const LEGACY_SYNC_ENABLED_ACTIVE_FLAGS_VERSION: i64 = 20260517012000;
|
||||
const ACTIVE_FLAG_REPAIR_VERSION: i64 = 20260722140744;
|
||||
const LEGACY_SYNC_ENABLED_ACTIVE_FLAGS_SQL: &str =
|
||||
include_str!("../../../backfills/postgres/20260517012000_sync_legacy_enabled_active_flags.sql");
|
||||
const ACTIVE_FLAG_REPAIR_SQL: &str = include_str!(
|
||||
"../../../backfills/postgres/20260722140744_sync_legacy_enabled_from_is_active.sql"
|
||||
);
|
||||
|
||||
#[test]
|
||||
fn published_legacy_enabled_backfill_preserves_its_original_direction() {
|
||||
fn legacy_enabled_backfill_preserves_canonical_active_flags() {
|
||||
for table in ["providers", "provider_endpoints", "models"] {
|
||||
assert!(
|
||||
LEGACY_SYNC_ENABLED_ACTIVE_FLAGS_SQL.contains(&format!(
|
||||
"UPDATE public.{table}\n SET\n is_active = enabled,"
|
||||
)),
|
||||
"published legacy {table} backfill should remain unchanged"
|
||||
);
|
||||
}
|
||||
assert!(
|
||||
!LEGACY_SYNC_ENABLED_ACTIVE_FLAGS_SQL.contains("enabled = is_active"),
|
||||
"the released legacy backfill must not be rewritten in place"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn active_flag_repair_backfill_treats_is_active_as_authoritative() {
|
||||
for table in ["providers", "provider_endpoints", "models"] {
|
||||
assert!(
|
||||
ACTIVE_FLAG_REPAIR_SQL.contains(&format!(
|
||||
"UPDATE public.{table}\n SET enabled = is_active"
|
||||
)),
|
||||
"active flag repair should copy canonical state for {table}"
|
||||
);
|
||||
assert!(
|
||||
ACTIVE_FLAG_REPAIR_SQL.contains(&format!(
|
||||
"AND table_name = '{table}'\n AND column_name = 'enabled'"
|
||||
)),
|
||||
"active flag repair should check the legacy column for {table}"
|
||||
);
|
||||
assert!(
|
||||
ACTIVE_FLAG_REPAIR_SQL.contains(&format!(
|
||||
"AND table_name = '{table}'\n AND column_name = 'is_active'"
|
||||
)),
|
||||
"active flag repair should check the canonical column for {table}"
|
||||
"legacy {table}.enabled should be initialized from canonical is_active"
|
||||
);
|
||||
}
|
||||
assert!(
|
||||
!ACTIVE_FLAG_REPAIR_SQL.contains("is_active = enabled"),
|
||||
"the repair must not overwrite canonical active state"
|
||||
!LEGACY_SYNC_ENABLED_ACTIVE_FLAGS_SQL.contains("is_active = enabled"),
|
||||
"the corrected script must not enable canonically disabled records"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -81,8 +48,7 @@ fn pending_backfills_from_applied_returns_all_versions_when_none_applied() {
|
||||
20260504120000,
|
||||
20260505120000,
|
||||
20260517012000,
|
||||
20260716010000,
|
||||
ACTIVE_FLAG_REPAIR_VERSION
|
||||
20260716010000
|
||||
]
|
||||
);
|
||||
}
|
||||
@@ -103,14 +69,13 @@ fn pending_backfills_from_applied_skips_versions_already_applied() {
|
||||
20260504120000,
|
||||
20260505120000,
|
||||
20260517012000,
|
||||
20260716010000,
|
||||
ACTIVE_FLAG_REPAIR_VERSION
|
||||
20260716010000
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn active_flag_repair_remains_pending_after_legacy_backfill_was_applied() {
|
||||
fn corrected_legacy_backfill_is_not_requeued_after_application() {
|
||||
let pending_versions = pending_backfills_from_applied(&[AppliedBackfill {
|
||||
version: LEGACY_SYNC_ENABLED_ACTIVE_FLAGS_VERSION,
|
||||
checksum: Vec::new(),
|
||||
@@ -120,8 +85,6 @@ fn active_flag_repair_remains_pending_after_legacy_backfill_was_applied() {
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert!(!pending_versions.contains(&LEGACY_SYNC_ENABLED_ACTIVE_FLAGS_VERSION));
|
||||
assert!(pending_versions.contains(&ACTIVE_FLAG_REPAIR_VERSION));
|
||||
assert_eq!(pending_versions.last(), Some(&ACTIVE_FLAG_REPAIR_VERSION));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -353,108 +316,6 @@ async fn wait_for_postgres(database_url: &str) -> Result<(), Box<dyn std::error:
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn active_flag_repair_runs_after_legacy_backfill_was_applied() {
|
||||
let Some(server) = ManagedPostgresServer::try_start()
|
||||
.await
|
||||
.expect("postgres backfill test should start or skip")
|
||||
else {
|
||||
return;
|
||||
};
|
||||
|
||||
let pool = PgPool::connect(server.database_url())
|
||||
.await
|
||||
.expect("pool should connect");
|
||||
prepare_database_for_startup(&pool)
|
||||
.await
|
||||
.expect("schema should prepare");
|
||||
|
||||
pending_backfills(&pool)
|
||||
.await
|
||||
.expect("backfill ledger should initialize");
|
||||
query(
|
||||
"INSERT INTO public.schema_backfills (version, description, success, checksum, execution_time) VALUES ($1, $2, TRUE, $3, 0)",
|
||||
)
|
||||
.bind(LEGACY_SYNC_ENABLED_ACTIVE_FLAGS_VERSION)
|
||||
.bind("sync_legacy_enabled_active_flags")
|
||||
.bind(Sha384::digest(LEGACY_SYNC_ENABLED_ACTIVE_FLAGS_SQL.as_bytes()).to_vec())
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("legacy backfill ledger fixture should insert");
|
||||
|
||||
let legacy_applied: bool =
|
||||
query_scalar("SELECT EXISTS(SELECT 1 FROM public.schema_backfills WHERE version = $1)")
|
||||
.bind(LEGACY_SYNC_ENABLED_ACTIVE_FLAGS_VERSION)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("legacy backfill ledger state should load");
|
||||
assert!(legacy_applied);
|
||||
|
||||
query(
|
||||
r#"
|
||||
INSERT INTO public.providers (id, name, provider_type, enabled, is_active)
|
||||
VALUES
|
||||
('provider-disabled', 'Disabled Provider', 'custom', TRUE, FALSE),
|
||||
('provider-active', 'Active Provider', 'custom', FALSE, TRUE)
|
||||
"#,
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("provider fixtures should insert");
|
||||
|
||||
let pending_before = pending_backfills(&pool)
|
||||
.await
|
||||
.expect("pending repair should load");
|
||||
let pending_versions = pending_before
|
||||
.iter()
|
||||
.map(|item| item.version)
|
||||
.collect::<Vec<_>>();
|
||||
assert!(!pending_versions.contains(&LEGACY_SYNC_ENABLED_ACTIVE_FLAGS_VERSION));
|
||||
assert!(pending_versions.contains(&ACTIVE_FLAG_REPAIR_VERSION));
|
||||
assert_eq!(pending_versions.last(), Some(&ACTIVE_FLAG_REPAIR_VERSION));
|
||||
|
||||
run_backfills(&pool)
|
||||
.await
|
||||
.expect("active flag repair should apply");
|
||||
|
||||
let pending_after = pending_backfills(&pool)
|
||||
.await
|
||||
.expect("pending backfills should reload");
|
||||
assert!(pending_after.is_empty());
|
||||
|
||||
let states: Vec<(String, bool, bool)> =
|
||||
query_as("SELECT id, enabled, is_active FROM public.providers ORDER BY id ASC")
|
||||
.fetch_all(&pool)
|
||||
.await
|
||||
.expect("provider flags should reload");
|
||||
assert_eq!(
|
||||
states,
|
||||
vec![
|
||||
("provider-active".to_string(), true, true),
|
||||
("provider-disabled".to_string(), false, false),
|
||||
]
|
||||
);
|
||||
|
||||
sqlx::raw_sql(ACTIVE_FLAG_REPAIR_SQL)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("re-running the repair SQL should be idempotent");
|
||||
let states_after_rerun: Vec<(String, bool, bool)> =
|
||||
query_as("SELECT id, enabled, is_active FROM public.providers ORDER BY id ASC")
|
||||
.fetch_all(&pool)
|
||||
.await
|
||||
.expect("provider flags should reload after repair rerun");
|
||||
assert_eq!(states_after_rerun, states);
|
||||
|
||||
let repair_applied: bool =
|
||||
query_scalar("SELECT EXISTS(SELECT 1 FROM public.schema_backfills WHERE version = $1)")
|
||||
.bind(ACTIVE_FLAG_REPAIR_VERSION)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("repair backfill ledger state should load");
|
||||
assert!(repair_applied);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn run_backfills_rebuilds_stats_and_records_execution() {
|
||||
let Some(server) = ManagedPostgresServer::try_start()
|
||||
@@ -591,14 +452,13 @@ async fn run_backfills_rebuilds_stats_and_records_execution() {
|
||||
let pending_before = pending_backfills(&pool)
|
||||
.await
|
||||
.expect("pending backfills should load");
|
||||
assert_eq!(pending_before.len(), 7);
|
||||
assert_eq!(pending_before.len(), 6);
|
||||
assert_eq!(pending_before[0].version, 20260422110000);
|
||||
assert_eq!(pending_before[1].version, 20260422120000);
|
||||
assert_eq!(pending_before[2].version, 20260504120000);
|
||||
assert_eq!(pending_before[3].version, 20260505120000);
|
||||
assert_eq!(pending_before[4].version, 20260517012000);
|
||||
assert_eq!(pending_before[5].version, 20260716010000);
|
||||
assert_eq!(pending_before[6].version, ACTIVE_FLAG_REPAIR_VERSION);
|
||||
|
||||
run_backfills(&pool)
|
||||
.await
|
||||
@@ -622,8 +482,7 @@ async fn run_backfills_rebuilds_stats_and_records_execution() {
|
||||
20260504120000,
|
||||
20260505120000,
|
||||
20260517012000,
|
||||
20260716010000,
|
||||
ACTIVE_FLAG_REPAIR_VERSION
|
||||
20260716010000
|
||||
]
|
||||
);
|
||||
|
||||
@@ -1297,5 +1156,5 @@ ORDER BY total_tokens
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("backfill count should load");
|
||||
assert_eq!(applied_count, 7);
|
||||
assert_eq!(applied_count, 6);
|
||||
}
|
||||
|
||||
@@ -8,8 +8,8 @@ use serde_json::{json, Map, Value};
|
||||
use super::{
|
||||
ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate,
|
||||
ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListQuery,
|
||||
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
|
||||
ProviderCatalogReadRepository, ProviderCatalogSnapshot,
|
||||
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
|
||||
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository, ProviderCatalogSnapshot,
|
||||
ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
|
||||
@@ -770,6 +770,82 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository {
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn compare_and_update_key_oauth_runtime_state(
|
||||
&self,
|
||||
update: &ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
if update.encrypted_auth_config.trim().is_empty()
|
||||
|| update
|
||||
.encrypted_api_key_update
|
||||
.as_deref()
|
||||
.is_some_and(|value| value.trim().is_empty())
|
||||
|| !update.status_snapshot_patch.is_object()
|
||||
|| update
|
||||
.upstream_metadata_patch
|
||||
.as_ref()
|
||||
.is_some_and(|patch| !patch.is_object())
|
||||
{
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"provider catalog OAuth runtime CAS requires auth_config and object status patch"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
let patch = update
|
||||
.status_snapshot_patch
|
||||
.as_object()
|
||||
.cloned()
|
||||
.expect("status patch object was validated");
|
||||
let mut index = self
|
||||
.index
|
||||
.write()
|
||||
.expect("provider catalog repository lock");
|
||||
let Some(key) = index.keys.get_mut(&update.key_id) else {
|
||||
return Ok(false);
|
||||
};
|
||||
if key.encrypted_auth_config.as_deref() != update.expected_encrypted_auth_config.as_deref()
|
||||
{
|
||||
return Ok(false);
|
||||
}
|
||||
if let Some(encrypted_api_key) = update.encrypted_api_key_update.as_ref() {
|
||||
key.encrypted_api_key = Some(encrypted_api_key.clone());
|
||||
}
|
||||
key.encrypted_auth_config = Some(update.encrypted_auth_config.clone());
|
||||
if let Some(expires_at_unix_secs) = update.expires_at_unix_secs_update {
|
||||
key.expires_at_unix_secs = expires_at_unix_secs;
|
||||
}
|
||||
key.oauth_invalid_at_unix_secs = update.oauth_invalid_at_unix_secs;
|
||||
key.oauth_invalid_reason = update.oauth_invalid_reason.clone();
|
||||
if update.reset_error_count {
|
||||
key.error_count = Some(0);
|
||||
}
|
||||
if let Some(metadata_patch) = update
|
||||
.upstream_metadata_patch
|
||||
.as_ref()
|
||||
.and_then(Value::as_object)
|
||||
.cloned()
|
||||
{
|
||||
let upstream_metadata = json_object_for_merge(
|
||||
key.upstream_metadata.as_ref(),
|
||||
"provider catalog upstream metadata",
|
||||
)?;
|
||||
key.upstream_metadata = Some(Value::Object(merge_json_objects(
|
||||
upstream_metadata,
|
||||
metadata_patch,
|
||||
)));
|
||||
}
|
||||
let status_snapshot = json_object_for_merge(
|
||||
key.status_snapshot.as_ref(),
|
||||
"provider catalog status snapshot",
|
||||
)?;
|
||||
key.status_snapshot = Some(Value::Object(merge_json_objects(status_snapshot, patch)));
|
||||
key.updated_at_unix_secs = Some(
|
||||
update
|
||||
.updated_at_unix_secs
|
||||
.unwrap_or_else(current_unix_secs),
|
||||
);
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn update_key_health_state(
|
||||
&self,
|
||||
key_id: &str,
|
||||
@@ -825,7 +901,12 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository {
|
||||
};
|
||||
let expected = update.expected.canonicalized();
|
||||
let next = update.next.canonicalized();
|
||||
if ProviderCatalogKeyAdaptiveState::from(&*key) != expected {
|
||||
if update
|
||||
.expected_encrypted_auth_config
|
||||
.as_deref()
|
||||
.is_some_and(|expected| key.encrypted_auth_config.as_deref() != Some(expected))
|
||||
|| ProviderCatalogKeyAdaptiveState::from(&*key) != expected
|
||||
{
|
||||
return Ok(false);
|
||||
}
|
||||
let status_snapshot = json_object_for_merge(
|
||||
@@ -964,7 +1045,11 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository {
|
||||
let Some(key) = index.keys.get_mut(&update.key_id) else {
|
||||
return Ok(false);
|
||||
};
|
||||
if key.health_by_format != update.expected_health_by_format
|
||||
if update
|
||||
.expected_encrypted_auth_config
|
||||
.as_deref()
|
||||
.is_some_and(|expected| key.encrypted_auth_config.as_deref() != Some(expected))
|
||||
|| key.health_by_format != update.expected_health_by_format
|
||||
|| key.circuit_breaker_by_format != update.expected_circuit_breaker_by_format
|
||||
{
|
||||
return Ok(false);
|
||||
@@ -1085,9 +1170,10 @@ mod tests {
|
||||
use crate::repository::provider_catalog::{
|
||||
ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate,
|
||||
ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder,
|
||||
ProviderCatalogKeyListQuery, ProviderCatalogKeyRuntimeMetadataUpdate,
|
||||
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
|
||||
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogReadRepository,
|
||||
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use crate::repository::usage::ProviderApiKeyUsageDelta;
|
||||
use serde_json::{json, Value};
|
||||
@@ -1236,6 +1322,75 @@ mod tests {
|
||||
assert_eq!(stored[0].expires_at_unix_secs, Some(4_102_444_800));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn oauth_runtime_cas_preserves_admin_fields_and_rejects_stale_config() {
|
||||
let mut key = sample_key("key-1", "provider-1")
|
||||
.with_transport_fields(
|
||||
None,
|
||||
"ciphertext-placeholder".to_string(),
|
||||
Some("ciphertext-auth-1".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build");
|
||||
key.note = Some("admin-owned-note".to_string());
|
||||
key.status_snapshot = Some(json!({"quota":{"remaining":7}}));
|
||||
let repository = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1")],
|
||||
vec![],
|
||||
vec![key],
|
||||
);
|
||||
let update = ProviderCatalogKeyOAuthRuntimeStateCasUpdate {
|
||||
key_id: "key-1".to_string(),
|
||||
expected_encrypted_auth_config: Some("ciphertext-auth-1".to_string()),
|
||||
encrypted_auth_config: "ciphertext-auth-2".to_string(),
|
||||
encrypted_api_key_update: Some("ciphertext-api-2".to_string()),
|
||||
expires_at_unix_secs_update: Some(Some(456)),
|
||||
oauth_invalid_at_unix_secs: None,
|
||||
oauth_invalid_reason: None,
|
||||
upstream_metadata_patch: Some(json!({"codex":{"remaining":3}})),
|
||||
status_snapshot_patch: json!({"oauth":{"code":"none"}}),
|
||||
reset_error_count: false,
|
||||
updated_at_unix_secs: Some(123),
|
||||
};
|
||||
assert!(repository
|
||||
.compare_and_update_key_oauth_runtime_state(&update)
|
||||
.await
|
||||
.expect("CAS should succeed"));
|
||||
assert!(!repository
|
||||
.compare_and_update_key_oauth_runtime_state(&update)
|
||||
.await
|
||||
.expect("stale CAS should not fail"));
|
||||
let stored = repository
|
||||
.list_keys_by_ids(&["key-1".to_string()])
|
||||
.await
|
||||
.expect("key should read")
|
||||
.pop()
|
||||
.expect("key should exist");
|
||||
assert_eq!(stored.note.as_deref(), Some("admin-owned-note"));
|
||||
assert_eq!(
|
||||
stored.encrypted_api_key.as_deref(),
|
||||
Some("ciphertext-api-2")
|
||||
);
|
||||
assert_eq!(stored.expires_at_unix_secs, Some(456));
|
||||
assert_eq!(
|
||||
stored.status_snapshot.as_ref().unwrap()["quota"]["remaining"],
|
||||
7
|
||||
);
|
||||
assert_eq!(
|
||||
stored.status_snapshot.as_ref().unwrap()["oauth"]["code"],
|
||||
"none"
|
||||
);
|
||||
assert_eq!(
|
||||
stored.upstream_metadata.as_ref().unwrap()["codex"]["remaining"],
|
||||
3
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn materializes_codex_window_usage_stats_delta_in_memory() {
|
||||
let mut key = sample_key("key-1", "provider-1");
|
||||
@@ -1586,6 +1741,7 @@ mod tests {
|
||||
);
|
||||
let update = ProviderCatalogKeyHealthStateUpdate {
|
||||
key_id: "key-1".to_string(),
|
||||
expected_encrypted_auth_config: None,
|
||||
expected_health_by_format: Some(json!({"openai:chat":{"consecutive_failures":1}})),
|
||||
expected_circuit_breaker_by_format: None,
|
||||
health_by_format: Some(json!({"openai:chat":{"consecutive_failures":2}})),
|
||||
@@ -1617,6 +1773,7 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn adaptive_cas_detects_conflicts_and_merges_only_owned_status_fields() {
|
||||
let mut key = sample_key("key-1", "provider-1");
|
||||
key.encrypted_auth_config = Some("auth-current".to_string());
|
||||
key.learned_rpm_limit = Some(10);
|
||||
key.rpm_429_count = Some(1);
|
||||
key.status_snapshot = Some(json!({
|
||||
@@ -1635,6 +1792,7 @@ mod tests {
|
||||
);
|
||||
let update = ProviderCatalogKeyAdaptiveStateUpdate {
|
||||
key_id: "key-1".to_string(),
|
||||
expected_encrypted_auth_config: Some("auth-current".to_string()),
|
||||
expected,
|
||||
next,
|
||||
status_snapshot_patch: json!({
|
||||
@@ -1645,6 +1803,14 @@ mod tests {
|
||||
updated_at_unix_secs: Some(10),
|
||||
};
|
||||
|
||||
let stale_generation_update = ProviderCatalogKeyAdaptiveStateUpdate {
|
||||
expected_encrypted_auth_config: Some("auth-stale".to_string()),
|
||||
..update.clone()
|
||||
};
|
||||
assert!(!repository
|
||||
.compare_and_update_key_adaptive_state(&stale_generation_update)
|
||||
.await
|
||||
.expect("stale auth generation should conflict"));
|
||||
assert!(repository
|
||||
.compare_and_update_key_adaptive_state(&update)
|
||||
.await
|
||||
@@ -1887,6 +2053,7 @@ mod tests {
|
||||
repository
|
||||
.compare_and_update_key_health_state(&ProviderCatalogKeyHealthStateUpdate {
|
||||
key_id: key.id.clone(),
|
||||
expected_encrypted_auth_config: None,
|
||||
expected_health_by_format: key.health_by_format.clone(),
|
||||
expected_circuit_breaker_by_format: None,
|
||||
health_by_format: Some(json!({"openai:chat":{"consecutive_failures":2}})),
|
||||
@@ -1901,6 +2068,7 @@ mod tests {
|
||||
repository
|
||||
.compare_and_update_key_adaptive_state(&ProviderCatalogKeyAdaptiveStateUpdate {
|
||||
key_id: key.id.clone(),
|
||||
expected_encrypted_auth_config: None,
|
||||
expected,
|
||||
next,
|
||||
status_snapshot_patch: json!({"observation_count":2}),
|
||||
|
||||
@@ -4,8 +4,8 @@ mod memory;
|
||||
pub(crate) use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate,
|
||||
ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery,
|
||||
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
|
||||
ProviderCatalogReadRepository, ProviderCatalogSnapshot,
|
||||
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
|
||||
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository, ProviderCatalogSnapshot,
|
||||
ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
|
||||
|
||||
@@ -40,6 +40,8 @@ pub struct StoredAdminProviderOAuthState {
|
||||
pub provider_id: String,
|
||||
pub provider_type: String,
|
||||
pub pkce_verifier: Option<String>,
|
||||
#[serde(default)]
|
||||
pub expected_encrypted_auth_config: Option<String>,
|
||||
}
|
||||
|
||||
pub fn provider_oauth_device_session_storage_key(session_id: &str) -> String {
|
||||
|
||||
@@ -9,6 +9,7 @@ description = "Provider-specific pool behavior adapters for Aether"
|
||||
[dependencies]
|
||||
aether-data-contracts.workspace = true
|
||||
aether-pool-core.workspace = true
|
||||
aether-provider-transport.workspace = true
|
||||
serde_json.workspace = true
|
||||
url.workspace = true
|
||||
uuid.workspace = true
|
||||
|
||||
@@ -184,6 +184,33 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_nested_agent_identity_quota_request_prefers_dynamic_assertion() {
|
||||
let spec = build_codex_pool_quota_request(
|
||||
"key-1",
|
||||
Some((
|
||||
"authorization".to_string(),
|
||||
"AgentAssertion signed-at-request-time".to_string(),
|
||||
)),
|
||||
None,
|
||||
Some(&json!({
|
||||
"agent_identity": {
|
||||
"agent_runtime_id": "runtime-1",
|
||||
"agent_private_key": "private-key"
|
||||
},
|
||||
"headers": {
|
||||
"authorization": "Bearer stale-imported-session"
|
||||
}
|
||||
})),
|
||||
)
|
||||
.expect("spec should build");
|
||||
|
||||
assert_eq!(
|
||||
spec.headers.get("authorization").map(String::as_str),
|
||||
Some("AgentAssertion signed-at-request-time")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_cli_quota_request_uses_v1internal_retrieve_user_quota() {
|
||||
let spec = build_gemini_cli_pool_quota_request(
|
||||
|
||||
@@ -93,14 +93,7 @@ fn build_codex_wham_headers(
|
||||
.filter(|value| !value.is_empty());
|
||||
|
||||
let is_agent_identity = auth_config
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|object| {
|
||||
object
|
||||
.get("auth_mode")
|
||||
.or_else(|| object.get("authMode"))
|
||||
.and_then(Value::as_str)
|
||||
})
|
||||
.is_some_and(|value| value.trim().eq_ignore_ascii_case("agentIdentity"));
|
||||
.is_some_and(aether_provider_transport::is_codex_agent_identity_auth_config_value);
|
||||
|
||||
if is_agent_identity {
|
||||
let Some((name, value)) = resolved_oauth_auth else {
|
||||
|
||||
@@ -13,11 +13,11 @@ use crypto_box::{
|
||||
};
|
||||
use ed25519_dalek::{
|
||||
pkcs8::{DecodePrivateKey, EncodePrivateKey},
|
||||
Signer, SigningKey,
|
||||
Signature, Signer, SigningKey, Verifier,
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Map, Value};
|
||||
use sha2::{Digest, Sha512};
|
||||
use sha2::{Digest, Sha256, Sha512};
|
||||
use thiserror::Error;
|
||||
use url::Url;
|
||||
|
||||
@@ -40,9 +40,17 @@ const ASSERTION_PREFIX: &str = "AgentAssertion ";
|
||||
const CODEX_AGENT_IDENTITY_AGENT_HARNESS_ID: &str = "codex-cli";
|
||||
const CODEX_AGENT_IDENTITY_RUNNING_LOCATION: &str = "local";
|
||||
|
||||
/// The AgentAssertion scheme is generated internally after an Agent Identity
|
||||
/// task has been registered. Keep the scheme check separate from envelope
|
||||
/// validation so a malformed in-flight assertion is still treated as an
|
||||
/// Agent-originated request by defensive runtime state writers.
|
||||
pub fn is_codex_agent_identity_authorization(value: &str) -> bool {
|
||||
encoded_agent_identity_assertion(value).is_some()
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Error)]
|
||||
pub enum CodexAgentIdentityEnrollmentError {
|
||||
#[error("ChatGPT Session Token 不能为空")]
|
||||
#[error("ChatGPT Access Token 不能为空")]
|
||||
MissingSessionToken,
|
||||
#[error("Agent Identity 注册请求失败")]
|
||||
RegistrationRequestFailed,
|
||||
@@ -67,6 +75,14 @@ struct AgentIdentityCredentials {
|
||||
task_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct AgentIdentityAssertionEnvelope {
|
||||
agent_runtime_id: String,
|
||||
task_id: String,
|
||||
timestamp: String,
|
||||
signature: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct AgentTaskRegistrationResponse {
|
||||
#[serde(default)]
|
||||
@@ -143,6 +159,90 @@ impl CodexAgentIdentityRefreshAdapter {
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a stable, non-secret fingerprint for the Agent Identity key pair
|
||||
/// and runtime. The task id is deliberately excluded so task rotation does not
|
||||
/// look like a credential replacement.
|
||||
pub fn codex_agent_identity_credential_fingerprint(config: &Value) -> Option<String> {
|
||||
let credentials = agent_identity_credentials(config).ok()?;
|
||||
Some(agent_identity_credential_fingerprint_from_credentials(
|
||||
&credentials,
|
||||
))
|
||||
}
|
||||
|
||||
pub fn codex_agent_identity_transport_credential_fingerprint(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<String> {
|
||||
CodexAgentIdentityRefreshAdapter::config_from_transport(transport)
|
||||
.and_then(|config| codex_agent_identity_credential_fingerprint(&config))
|
||||
}
|
||||
|
||||
/// Returns the fencing generation for a transport/config entry, including the
|
||||
/// current task id. It is safe to log/compare but must never be used as a key.
|
||||
pub fn codex_agent_identity_refresh_fingerprint(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
entry: Option<&CachedOAuthEntry>,
|
||||
) -> Option<String> {
|
||||
let transport_config = CodexAgentIdentityRefreshAdapter::config_from_transport(transport)?;
|
||||
let transport_credentials = agent_identity_credentials(&transport_config).ok()?;
|
||||
let transport_credential_fingerprint =
|
||||
agent_identity_credential_fingerprint_from_credentials(&transport_credentials);
|
||||
let config = entry
|
||||
.filter(|entry| {
|
||||
entry.source_fingerprint.as_deref() == Some(transport_credential_fingerprint.as_str())
|
||||
})
|
||||
.and_then(CodexAgentIdentityRefreshAdapter::config_from_entry)
|
||||
.filter(|entry_config| {
|
||||
agent_identity_credentials(entry_config)
|
||||
.ok()
|
||||
.is_some_and(|entry_credentials| {
|
||||
entry_credentials.task_id == transport_credentials.task_id
|
||||
})
|
||||
})
|
||||
.unwrap_or(transport_config);
|
||||
codex_agent_identity_config_refresh_fingerprint(&config)
|
||||
}
|
||||
|
||||
pub fn codex_agent_identity_config_refresh_fingerprint(config: &Value) -> Option<String> {
|
||||
let credentials = agent_identity_credentials(config).ok()?;
|
||||
let credential_fingerprint =
|
||||
agent_identity_credential_fingerprint_from_credentials(&credentials);
|
||||
let task_id = credentials.task_id.as_deref().unwrap_or_default();
|
||||
let mut digest = Sha256::new();
|
||||
digest.update(credential_fingerprint.as_bytes());
|
||||
digest.update([0]);
|
||||
digest.update(task_id.as_bytes());
|
||||
Some(URL_SAFE_NO_PAD.encode(digest.finalize()))
|
||||
}
|
||||
|
||||
pub fn codex_agent_identity_cached_entry_from_transport(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<CachedOAuthEntry> {
|
||||
let config = CodexAgentIdentityRefreshAdapter::config_from_transport(transport)?;
|
||||
let credentials = agent_identity_credentials(&config).ok()?;
|
||||
let task_id = credentials.task_id.as_deref()?;
|
||||
let auth_header_value = build_agent_assertion(&credentials, task_id, Utc::now()).ok()?;
|
||||
Some(CachedOAuthEntry {
|
||||
provider_type: CODEX_AGENT_IDENTITY_CACHED_ENTRY_PROVIDER_TYPE.to_string(),
|
||||
auth_header_name: AUTHORIZATION_HEADER.to_string(),
|
||||
auth_header_value,
|
||||
expires_at_unix_secs: None,
|
||||
metadata: Some(config),
|
||||
source_fingerprint: Some(agent_identity_credential_fingerprint_from_credentials(
|
||||
&credentials,
|
||||
)),
|
||||
})
|
||||
}
|
||||
|
||||
fn agent_identity_credential_fingerprint_from_credentials(
|
||||
credentials: &AgentIdentityCredentials,
|
||||
) -> String {
|
||||
let mut digest = Sha256::new();
|
||||
digest.update(credentials.runtime_id.as_bytes());
|
||||
digest.update([0]);
|
||||
digest.update(credentials.signing_key.to_bytes());
|
||||
URL_SAFE_NO_PAD.encode(digest.finalize())
|
||||
}
|
||||
|
||||
pub fn is_codex_agent_identity_auth_config_value(config: &Value) -> bool {
|
||||
let Some(root) = config.as_object() else {
|
||||
return false;
|
||||
@@ -168,6 +268,90 @@ pub fn is_codex_agent_identity_transport(transport: &GatewayProviderTransportSna
|
||||
.is_some_and(is_codex_agent_identity_auth_config_value)
|
||||
}
|
||||
|
||||
/// Verifies that an in-flight AgentAssertion was signed by the exact Agent
|
||||
/// Identity credential and task represented by the current transport.
|
||||
pub fn codex_agent_identity_authorization_matches_transport(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
authorization: &str,
|
||||
) -> bool {
|
||||
if !is_codex_agent_identity_transport(transport) {
|
||||
return false;
|
||||
}
|
||||
let Some(encoded) = encoded_agent_identity_assertion(authorization) else {
|
||||
return false;
|
||||
};
|
||||
let Ok(envelope_bytes) = URL_SAFE_NO_PAD.decode(encoded) else {
|
||||
return false;
|
||||
};
|
||||
let Ok(envelope) = serde_json::from_slice::<AgentIdentityAssertionEnvelope>(&envelope_bytes)
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
let Some(config) = CodexAgentIdentityRefreshAdapter::config_from_transport(transport) else {
|
||||
return false;
|
||||
};
|
||||
let Ok(credentials) = agent_identity_credentials(&config) else {
|
||||
return false;
|
||||
};
|
||||
let Some(current_task_id) = credentials.task_id.as_deref() else {
|
||||
return false;
|
||||
};
|
||||
let runtime_id = envelope.agent_runtime_id.trim();
|
||||
let task_id = envelope.task_id.trim();
|
||||
let timestamp = envelope.timestamp.trim();
|
||||
if runtime_id != credentials.runtime_id || task_id != current_task_id || timestamp.is_empty() {
|
||||
return false;
|
||||
}
|
||||
let Ok(signature_bytes) = STANDARD.decode(envelope.signature.trim()) else {
|
||||
return false;
|
||||
};
|
||||
let Ok(signature) = Signature::from_slice(&signature_bytes) else {
|
||||
return false;
|
||||
};
|
||||
let payload = format!("{runtime_id}:{task_id}:{timestamp}");
|
||||
credentials
|
||||
.signing_key
|
||||
.verifying_key()
|
||||
.verify(payload.as_bytes(), &signature)
|
||||
.is_ok()
|
||||
}
|
||||
|
||||
/// Agent task rotation may change only task_id. Any other auth_config change
|
||||
/// means the caller must rebuild the whole request from a fresh transport.
|
||||
pub fn codex_agent_identity_transport_allows_task_rotation_from(
|
||||
initial: &GatewayProviderTransportSnapshot,
|
||||
current: &GatewayProviderTransportSnapshot,
|
||||
) -> bool {
|
||||
let Some(initial_config) = CodexAgentIdentityRefreshAdapter::config_from_transport(initial)
|
||||
.and_then(agent_identity_config_without_task)
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
let Some(current_config) = CodexAgentIdentityRefreshAdapter::config_from_transport(current)
|
||||
.and_then(agent_identity_config_without_task)
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
initial_config == current_config
|
||||
}
|
||||
|
||||
pub fn codex_agent_identity_entry_allows_task_rotation_from(
|
||||
initial: &GatewayProviderTransportSnapshot,
|
||||
entry: &CachedOAuthEntry,
|
||||
) -> bool {
|
||||
let Some(initial_config) = CodexAgentIdentityRefreshAdapter::config_from_transport(initial)
|
||||
.and_then(agent_identity_config_without_task)
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
let Some(entry_config) = CodexAgentIdentityRefreshAdapter::config_from_entry(entry)
|
||||
.and_then(agent_identity_config_without_task)
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
initial_config == entry_config
|
||||
}
|
||||
|
||||
pub fn is_codex_agent_identity_cached_entry(entry: &CachedOAuthEntry) -> bool {
|
||||
entry
|
||||
.provider_type
|
||||
@@ -179,6 +363,18 @@ pub fn validate_codex_agent_identity_auth_config(config: &Value) -> Result<(), S
|
||||
agent_identity_credentials(config).map(|_| ())
|
||||
}
|
||||
|
||||
pub fn codex_agent_identity_auth_config_has_task_id(config: &Value) -> bool {
|
||||
let Some(root) = config.as_object() else {
|
||||
return false;
|
||||
};
|
||||
string_from_maps(
|
||||
root,
|
||||
agent_identity_nested_object(root),
|
||||
&["task_id", "taskId"],
|
||||
)
|
||||
.is_some()
|
||||
}
|
||||
|
||||
/// Returns whether an upstream response proves that the registered Agent Identity task is no
|
||||
/// longer usable. Only this condition should trigger task registration again; an arbitrary 401
|
||||
/// can instead mean that the account itself has lost access.
|
||||
@@ -228,6 +424,22 @@ fn agent_identity_nested_object(root: &Map<String, Value>) -> Option<&Map<String
|
||||
.and_then(Value::as_object)
|
||||
}
|
||||
|
||||
fn agent_identity_config_without_task(mut config: Value) -> Option<Value> {
|
||||
if !is_codex_agent_identity_auth_config_value(&config) {
|
||||
return None;
|
||||
}
|
||||
let root = config.as_object_mut()?;
|
||||
root.remove("task_id");
|
||||
root.remove("taskId");
|
||||
for nested_key in ["agent_identity", "agentIdentity"] {
|
||||
if let Some(nested) = root.get_mut(nested_key).and_then(Value::as_object_mut) {
|
||||
nested.remove("task_id");
|
||||
nested.remove("taskId");
|
||||
}
|
||||
}
|
||||
Some(config)
|
||||
}
|
||||
|
||||
fn string_from_map(map: &Map<String, Value>, keys: &[&str]) -> Option<String> {
|
||||
keys.iter().find_map(|key| {
|
||||
map.get(*key)
|
||||
@@ -279,6 +491,19 @@ fn agent_identity_timestamp(now: DateTime<Utc>) -> String {
|
||||
now.to_rfc3339_opts(SecondsFormat::Secs, true)
|
||||
}
|
||||
|
||||
fn encoded_agent_identity_assertion(value: &str) -> Option<&str> {
|
||||
let mut parts = value.split_ascii_whitespace();
|
||||
let scheme = parts.next()?;
|
||||
let encoded = parts.next()?;
|
||||
if !scheme.eq_ignore_ascii_case(ASSERTION_PREFIX.trim())
|
||||
|| encoded.is_empty()
|
||||
|| parts.next().is_some()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
Some(encoded)
|
||||
}
|
||||
|
||||
fn build_agent_assertion(
|
||||
credentials: &AgentIdentityCredentials,
|
||||
task_id: &str,
|
||||
@@ -373,20 +598,44 @@ fn agent_runtime_id_from_registration_response(body: &str) -> Result<String, ()>
|
||||
.ok_or(())
|
||||
}
|
||||
|
||||
/// Uses a ChatGPT session token once to register a fresh Agent Identity. The returned config
|
||||
/// contains only the generated signing credentials and is deliberately free of the session token.
|
||||
/// Uses a ChatGPT access token once to register a fresh Agent Identity. The returned config
|
||||
/// contains only the generated signing credentials and is deliberately free of the access token.
|
||||
pub async fn register_codex_agent_identity_from_access_token(
|
||||
executor: &dyn OAuthHttpExecutor,
|
||||
access_token: &str,
|
||||
network: OAuthNetworkContext,
|
||||
) -> Result<Map<String, Value>, CodexAgentIdentityEnrollmentError> {
|
||||
register_codex_agent_identity_from_access_token_with_auth_api_base_url(
|
||||
executor,
|
||||
access_token,
|
||||
network,
|
||||
CODEX_AGENT_IDENTITY_AUTH_API_BASE_URL,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Registers an Agent Identity and its initial task for backwards compatibility.
|
||||
pub async fn create_codex_agent_identity_from_access_token(
|
||||
executor: &dyn OAuthHttpExecutor,
|
||||
access_token: &str,
|
||||
network: OAuthNetworkContext,
|
||||
) -> Result<Map<String, Value>, CodexAgentIdentityEnrollmentError> {
|
||||
create_codex_agent_identity_from_session_token_with_auth_api_base_url(
|
||||
executor,
|
||||
access_token,
|
||||
network,
|
||||
CODEX_AGENT_IDENTITY_AUTH_API_BASE_URL,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Backwards-compatible alias for callers using the original, inaccurate name.
|
||||
pub async fn create_codex_agent_identity_from_session_token(
|
||||
executor: &dyn OAuthHttpExecutor,
|
||||
session_token: &str,
|
||||
network: OAuthNetworkContext,
|
||||
) -> Result<Map<String, Value>, CodexAgentIdentityEnrollmentError> {
|
||||
create_codex_agent_identity_from_session_token_with_auth_api_base_url(
|
||||
executor,
|
||||
session_token,
|
||||
network,
|
||||
CODEX_AGENT_IDENTITY_AUTH_API_BASE_URL,
|
||||
)
|
||||
.await
|
||||
create_codex_agent_identity_from_access_token(executor, session_token, network).await
|
||||
}
|
||||
|
||||
async fn create_codex_agent_identity_from_session_token_with_auth_api_base_url(
|
||||
@@ -395,75 +644,13 @@ async fn create_codex_agent_identity_from_session_token_with_auth_api_base_url(
|
||||
network: OAuthNetworkContext,
|
||||
auth_api_base_url: &str,
|
||||
) -> Result<Map<String, Value>, CodexAgentIdentityEnrollmentError> {
|
||||
let session_token = session_token.trim();
|
||||
if session_token.is_empty() {
|
||||
return Err(CodexAgentIdentityEnrollmentError::MissingSessionToken);
|
||||
}
|
||||
|
||||
let signing_key = generate_agent_identity_signing_key();
|
||||
let private_key_der = signing_key
|
||||
.to_pkcs8_der()
|
||||
.map_err(|_| CodexAgentIdentityEnrollmentError::KeyGenerationFailed)?;
|
||||
let agent_private_key = STANDARD.encode(private_key_der.as_bytes());
|
||||
let agent_public_key = agent_identity_ssh_public_key(&signing_key);
|
||||
let registration_url = agent_registration_url(auth_api_base_url)
|
||||
.map_err(|_| CodexAgentIdentityEnrollmentError::RegistrationRequestFailed)?;
|
||||
let registration_response = executor
|
||||
.execute(OAuthHttpRequest {
|
||||
request_id: CODEX_AGENT_IDENTITY_AGENT_REGISTRATION_REQUEST_ID.to_string(),
|
||||
method: reqwest::Method::POST,
|
||||
url: registration_url,
|
||||
headers: BTreeMap::from([
|
||||
("accept".to_string(), "application/json".to_string()),
|
||||
("content-type".to_string(), "application/json".to_string()),
|
||||
(
|
||||
"authorization".to_string(),
|
||||
format!("Bearer {session_token}"),
|
||||
),
|
||||
(
|
||||
"user-agent".to_string(),
|
||||
aether_ai_formats::CODEX_CLIENT_USER_AGENT.to_string(),
|
||||
),
|
||||
(
|
||||
"originator".to_string(),
|
||||
aether_ai_formats::CODEX_CLIENT_ORIGINATOR.to_string(),
|
||||
),
|
||||
]),
|
||||
content_type: Some("application/json".to_string()),
|
||||
json_body: Some(json!({
|
||||
"abom": {
|
||||
"agent_version": aether_ai_formats::CODEX_CLIENT_VERSION,
|
||||
"agent_harness_id": CODEX_AGENT_IDENTITY_AGENT_HARNESS_ID,
|
||||
"running_location": CODEX_AGENT_IDENTITY_RUNNING_LOCATION,
|
||||
},
|
||||
"agent_public_key": agent_public_key,
|
||||
})),
|
||||
body_bytes: None,
|
||||
network: network.clone(),
|
||||
})
|
||||
.await
|
||||
.map_err(|_| CodexAgentIdentityEnrollmentError::RegistrationRequestFailed)?;
|
||||
if !(200..300).contains(®istration_response.status_code) {
|
||||
return Err(CodexAgentIdentityEnrollmentError::RegistrationRejected {
|
||||
status_code: registration_response.status_code,
|
||||
});
|
||||
}
|
||||
let agent_runtime_id =
|
||||
agent_runtime_id_from_registration_response(registration_response.body_text.as_str())
|
||||
.map_err(|_| CodexAgentIdentityEnrollmentError::InvalidRegistrationResponse)?;
|
||||
|
||||
let mut auth_config = Map::from_iter([
|
||||
(
|
||||
"provider_type".to_string(),
|
||||
json!(CODEX_AGENT_IDENTITY_PROVIDER_TYPE),
|
||||
),
|
||||
(
|
||||
"auth_mode".to_string(),
|
||||
json!(CODEX_AGENT_IDENTITY_AUTH_MODE),
|
||||
),
|
||||
("agent_runtime_id".to_string(), json!(agent_runtime_id)),
|
||||
("agent_private_key".to_string(), json!(agent_private_key)),
|
||||
]);
|
||||
let mut auth_config = register_codex_agent_identity_from_access_token_with_auth_api_base_url(
|
||||
executor,
|
||||
session_token,
|
||||
network.clone(),
|
||||
auth_api_base_url,
|
||||
)
|
||||
.await?;
|
||||
let config_value = Value::Object(auth_config.clone());
|
||||
let credentials = agent_identity_credentials(&config_value)
|
||||
.map_err(|_| CodexAgentIdentityEnrollmentError::KeyGenerationFailed)?;
|
||||
@@ -505,6 +692,84 @@ async fn create_codex_agent_identity_from_session_token_with_auth_api_base_url(
|
||||
Ok(auth_config)
|
||||
}
|
||||
|
||||
async fn register_codex_agent_identity_from_access_token_with_auth_api_base_url(
|
||||
executor: &dyn OAuthHttpExecutor,
|
||||
access_token: &str,
|
||||
network: OAuthNetworkContext,
|
||||
auth_api_base_url: &str,
|
||||
) -> Result<Map<String, Value>, CodexAgentIdentityEnrollmentError> {
|
||||
let access_token = access_token.trim();
|
||||
if access_token.is_empty() {
|
||||
return Err(CodexAgentIdentityEnrollmentError::MissingSessionToken);
|
||||
}
|
||||
|
||||
let signing_key = generate_agent_identity_signing_key();
|
||||
let private_key_der = signing_key
|
||||
.to_pkcs8_der()
|
||||
.map_err(|_| CodexAgentIdentityEnrollmentError::KeyGenerationFailed)?;
|
||||
let agent_private_key = STANDARD.encode(private_key_der.as_bytes());
|
||||
let agent_public_key = agent_identity_ssh_public_key(&signing_key);
|
||||
let registration_url = agent_registration_url(auth_api_base_url)
|
||||
.map_err(|_| CodexAgentIdentityEnrollmentError::RegistrationRequestFailed)?;
|
||||
let response = executor
|
||||
.execute(OAuthHttpRequest {
|
||||
request_id: CODEX_AGENT_IDENTITY_AGENT_REGISTRATION_REQUEST_ID.to_string(),
|
||||
method: reqwest::Method::POST,
|
||||
url: registration_url,
|
||||
headers: BTreeMap::from([
|
||||
("accept".to_string(), "application/json".to_string()),
|
||||
("content-type".to_string(), "application/json".to_string()),
|
||||
(
|
||||
"authorization".to_string(),
|
||||
format!("Bearer {access_token}"),
|
||||
),
|
||||
(
|
||||
"user-agent".to_string(),
|
||||
aether_ai_formats::CODEX_CLIENT_USER_AGENT.to_string(),
|
||||
),
|
||||
(
|
||||
"originator".to_string(),
|
||||
aether_ai_formats::CODEX_CLIENT_ORIGINATOR.to_string(),
|
||||
),
|
||||
]),
|
||||
content_type: Some("application/json".to_string()),
|
||||
json_body: Some(json!({
|
||||
"abom": {
|
||||
"agent_version": aether_ai_formats::CODEX_CLIENT_VERSION,
|
||||
"agent_harness_id": CODEX_AGENT_IDENTITY_AGENT_HARNESS_ID,
|
||||
"running_location": CODEX_AGENT_IDENTITY_RUNNING_LOCATION,
|
||||
},
|
||||
"agent_public_key": agent_public_key,
|
||||
})),
|
||||
body_bytes: None,
|
||||
network,
|
||||
})
|
||||
.await
|
||||
.map_err(|_| CodexAgentIdentityEnrollmentError::RegistrationRequestFailed)?;
|
||||
if !(200..300).contains(&response.status_code) {
|
||||
return Err(CodexAgentIdentityEnrollmentError::RegistrationRejected {
|
||||
status_code: response.status_code,
|
||||
});
|
||||
}
|
||||
let agent_runtime_id = agent_runtime_id_from_registration_response(response.body_text.as_str())
|
||||
.map_err(|_| CodexAgentIdentityEnrollmentError::InvalidRegistrationResponse)?;
|
||||
let auth_config = Map::from_iter([
|
||||
(
|
||||
"provider_type".to_string(),
|
||||
json!(CODEX_AGENT_IDENTITY_PROVIDER_TYPE),
|
||||
),
|
||||
(
|
||||
"auth_mode".to_string(),
|
||||
json!(CODEX_AGENT_IDENTITY_AUTH_MODE),
|
||||
),
|
||||
("agent_runtime_id".to_string(), json!(agent_runtime_id)),
|
||||
("agent_private_key".to_string(), json!(agent_private_key)),
|
||||
]);
|
||||
validate_codex_agent_identity_auth_config(&Value::Object(auth_config.clone()))
|
||||
.map_err(|_| CodexAgentIdentityEnrollmentError::KeyGenerationFailed)?;
|
||||
Ok(auth_config)
|
||||
}
|
||||
|
||||
fn task_id_from_registration_response(
|
||||
credentials: &AgentIdentityCredentials,
|
||||
body: &str,
|
||||
@@ -592,9 +857,40 @@ impl LocalOAuthRefreshAdapter for CodexAgentIdentityRefreshAdapter {
|
||||
|
||||
fn resolve_cached(
|
||||
&self,
|
||||
_transport: &GatewayProviderTransportSnapshot,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
entry: &CachedOAuthEntry,
|
||||
) -> Option<LocalResolvedOAuthRequestAuth> {
|
||||
let current_fingerprint = codex_agent_identity_transport_credential_fingerprint(transport)?;
|
||||
if entry.source_fingerprint.as_deref() != Some(current_fingerprint.as_str()) {
|
||||
return None;
|
||||
}
|
||||
let transport_config = Self::config_from_transport(transport)?;
|
||||
let entry_config = Self::config_from_entry(entry)?;
|
||||
let transport_task = agent_identity_credentials(&transport_config).ok()?.task_id;
|
||||
let entry_task = agent_identity_credentials(&entry_config).ok()?.task_id;
|
||||
if transport_task != entry_task {
|
||||
return None;
|
||||
}
|
||||
Self::resolve_from_config(&entry_config)
|
||||
}
|
||||
|
||||
fn resolve_fenced_cached(
|
||||
&self,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
entry: &CachedOAuthEntry,
|
||||
) -> Option<LocalResolvedOAuthRequestAuth> {
|
||||
self.resolve_cached(transport, entry)
|
||||
}
|
||||
|
||||
fn resolve_refreshed(
|
||||
&self,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
entry: &CachedOAuthEntry,
|
||||
) -> Option<LocalResolvedOAuthRequestAuth> {
|
||||
let current_fingerprint = codex_agent_identity_transport_credential_fingerprint(transport)?;
|
||||
if entry.source_fingerprint.as_deref() != Some(current_fingerprint.as_str()) {
|
||||
return None;
|
||||
}
|
||||
Self::config_from_entry(entry).and_then(|config| Self::resolve_from_config(&config))
|
||||
}
|
||||
|
||||
@@ -618,6 +914,36 @@ impl LocalOAuthRefreshAdapter for CodexAgentIdentityRefreshAdapter {
|
||||
.is_some_and(|credentials| credentials.task_id.is_none())
|
||||
}
|
||||
|
||||
fn refresh_fingerprint(
|
||||
&self,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
entry: Option<&CachedOAuthEntry>,
|
||||
) -> Option<String> {
|
||||
codex_agent_identity_refresh_fingerprint(transport, entry)
|
||||
}
|
||||
|
||||
fn cached_entry_from_transport(
|
||||
&self,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<CachedOAuthEntry> {
|
||||
codex_agent_identity_cached_entry_from_transport(transport)
|
||||
}
|
||||
|
||||
fn should_backoff_after_error(&self, error: &LocalOAuthRefreshError) -> bool {
|
||||
match error {
|
||||
LocalOAuthRefreshError::HttpStatus { status_code, .. } => {
|
||||
*status_code == 429 || *status_code >= 500
|
||||
}
|
||||
LocalOAuthRefreshError::Transport { .. }
|
||||
| LocalOAuthRefreshError::TransportMessage { .. }
|
||||
| LocalOAuthRefreshError::InvalidResponse { .. } => true,
|
||||
}
|
||||
}
|
||||
|
||||
fn requires_distributed_refresh_lock(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn refresh(
|
||||
&self,
|
||||
executor: &dyn LocalOAuthHttpExecutor,
|
||||
@@ -676,6 +1002,9 @@ impl LocalOAuthRefreshAdapter for CodexAgentIdentityRefreshAdapter {
|
||||
auth_header_value,
|
||||
expires_at_unix_secs: None,
|
||||
metadata: Some(config),
|
||||
source_fingerprint: Some(agent_identity_credential_fingerprint_from_credentials(
|
||||
&updated_credentials,
|
||||
)),
|
||||
}))
|
||||
}
|
||||
}
|
||||
@@ -693,16 +1022,22 @@ mod tests {
|
||||
|
||||
use super::{
|
||||
agent_identity_credentials, build_agent_assertion,
|
||||
codex_agent_identity_authorization_matches_transport,
|
||||
codex_agent_identity_cached_entry_from_transport,
|
||||
codex_agent_identity_config_refresh_fingerprint, codex_agent_identity_refresh_fingerprint,
|
||||
codex_agent_identity_transport_allows_task_rotation_from,
|
||||
create_codex_agent_identity_from_session_token_with_auth_api_base_url,
|
||||
decrypt_agent_task_id, is_codex_agent_identity_auth_config_value,
|
||||
is_codex_agent_identity_invalid_task_response, task_id_from_registration_response,
|
||||
validate_codex_agent_identity_auth_config, with_agent_identity_task_id,
|
||||
CodexAgentIdentityEnrollmentError, CodexAgentIdentityRefreshAdapter,
|
||||
CODEX_AGENT_IDENTITY_CACHED_ENTRY_PROVIDER_TYPE,
|
||||
is_codex_agent_identity_authorization, is_codex_agent_identity_invalid_task_response,
|
||||
register_codex_agent_identity_from_access_token_with_auth_api_base_url,
|
||||
task_id_from_registration_response, validate_codex_agent_identity_auth_config,
|
||||
with_agent_identity_task_id, CodexAgentIdentityEnrollmentError,
|
||||
CodexAgentIdentityRefreshAdapter, CODEX_AGENT_IDENTITY_CACHED_ENTRY_PROVIDER_TYPE,
|
||||
};
|
||||
use crate::oauth_refresh::{
|
||||
LocalOAuthHttpExecutor, LocalOAuthHttpRequest, LocalOAuthHttpResponse,
|
||||
LocalOAuthRefreshAdapter, LocalOAuthRefreshError, LocalResolvedOAuthRequestAuth,
|
||||
LocalOAuthRefreshAdapter, LocalOAuthRefreshCoordinator, LocalOAuthRefreshError,
|
||||
LocalResolvedOAuthRequestAuth,
|
||||
};
|
||||
use crate::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
@@ -826,6 +1161,79 @@ mod tests {
|
||||
.expect("assertion signature should verify");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recognizes_agent_assertion_authorization_scheme_without_parsing_secrets() {
|
||||
assert!(is_codex_agent_identity_authorization(
|
||||
"AgentAssertion assertion-envelope"
|
||||
));
|
||||
assert!(is_codex_agent_identity_authorization(
|
||||
"agentassertion\tassertion-envelope"
|
||||
));
|
||||
assert!(!is_codex_agent_identity_authorization(
|
||||
"Bearer access-token"
|
||||
));
|
||||
assert!(!is_codex_agent_identity_authorization("AgentAssertion"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assertion_match_binds_runtime_task_and_signing_key_generation() {
|
||||
let config = test_auth_config(Some("task-test"));
|
||||
let credentials = agent_identity_credentials(&config).expect("credentials should parse");
|
||||
let assertion = build_agent_assertion(
|
||||
&credentials,
|
||||
"task-test",
|
||||
Utc.with_ymd_and_hms(2030, 1, 2, 3, 4, 5).unwrap(),
|
||||
)
|
||||
.expect("assertion should build");
|
||||
|
||||
assert!(codex_agent_identity_authorization_matches_transport(
|
||||
&sample_transport(config.clone()),
|
||||
&assertion,
|
||||
));
|
||||
assert!(!codex_agent_identity_authorization_matches_transport(
|
||||
&sample_transport(test_auth_config(Some("task-replaced"))),
|
||||
&assertion,
|
||||
));
|
||||
|
||||
let replacement_key = SigningKey::from_bytes(&[8u8; 32])
|
||||
.to_pkcs8_der()
|
||||
.expect("replacement key should encode");
|
||||
let mut replacement_config = config;
|
||||
replacement_config["agent_private_key"] =
|
||||
json!(STANDARD.encode(replacement_key.as_bytes()));
|
||||
assert!(!codex_agent_identity_authorization_matches_transport(
|
||||
&sample_transport(replacement_config),
|
||||
&assertion,
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn task_rotation_context_rejects_metadata_and_credential_replacement() {
|
||||
let initial = sample_transport(test_auth_config(Some("task-old")));
|
||||
let rotated = sample_transport(test_auth_config(Some("task-new")));
|
||||
assert!(codex_agent_identity_transport_allows_task_rotation_from(
|
||||
&initial, &rotated
|
||||
));
|
||||
|
||||
let mut metadata_rewrite = test_auth_config(Some("task-new"));
|
||||
metadata_rewrite["account_id"] = json!("account-replaced");
|
||||
assert!(!codex_agent_identity_transport_allows_task_rotation_from(
|
||||
&initial,
|
||||
&sample_transport(metadata_rewrite),
|
||||
));
|
||||
|
||||
let replacement_key = SigningKey::from_bytes(&[9u8; 32])
|
||||
.to_pkcs8_der()
|
||||
.expect("replacement key should encode");
|
||||
let mut credential_rewrite = test_auth_config(Some("task-new"));
|
||||
credential_rewrite["agent_private_key"] =
|
||||
json!(STANDARD.encode(replacement_key.as_bytes()));
|
||||
assert!(!codex_agent_identity_transport_allows_task_rotation_from(
|
||||
&initial,
|
||||
&sample_transport(credential_rewrite),
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decrypts_sealed_task_registration_response() {
|
||||
let config = test_auth_config(None);
|
||||
@@ -886,6 +1294,16 @@ mod tests {
|
||||
assert!(updated["agent_identity"].get("taskId").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn refresh_fingerprint_fences_task_generation_not_only_keypair() {
|
||||
let pending = test_auth_config(None);
|
||||
let registered = test_auth_config(Some("task-winner"));
|
||||
assert_ne!(
|
||||
codex_agent_identity_config_refresh_fingerprint(&pending),
|
||||
codex_agent_identity_config_refresh_fingerprint(®istered)
|
||||
);
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct RecordingExecutor {
|
||||
requests: Arc<Mutex<Vec<LocalOAuthHttpRequest>>>,
|
||||
@@ -1009,6 +1427,37 @@ mod tests {
|
||||
assert!(!requests[1].headers.contains_key("authorization"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn register_only_returns_pending_config_without_registering_task() {
|
||||
let requests = Arc::new(Mutex::new(Vec::new()));
|
||||
let executor = RecordingEnrollmentExecutor {
|
||||
requests: Arc::clone(&requests),
|
||||
responses: Arc::new(Mutex::new(vec![OAuthHttpResponse {
|
||||
status_code: 200,
|
||||
body_text: r#"{"agent_runtime_id":"runtime-pending"}"#.to_string(),
|
||||
json_body: None,
|
||||
}])),
|
||||
};
|
||||
let config = register_codex_agent_identity_from_access_token_with_auth_api_base_url(
|
||||
&executor,
|
||||
"access-token-for-test-only",
|
||||
OAuthNetworkContext::direct_identity(),
|
||||
"https://auth.test/api/accounts",
|
||||
)
|
||||
.await
|
||||
.expect("register-only enrollment should succeed");
|
||||
validate_codex_agent_identity_auth_config(&serde_json::Value::Object(config.clone()))
|
||||
.expect("pending config should be valid");
|
||||
assert_eq!(
|
||||
config.get("agent_runtime_id"),
|
||||
Some(&json!("runtime-pending"))
|
||||
);
|
||||
assert!(!config.contains_key("task_id"));
|
||||
let requests = requests.lock().expect("recording lock should hold");
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert!(requests[0].url.ends_with("/v1/agent/register"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn enrollment_error_does_not_echo_session_token_or_response_body() {
|
||||
let executor = RecordingEnrollmentExecutor {
|
||||
@@ -1066,9 +1515,10 @@ mod tests {
|
||||
.and_then(|value| value.get("task_id")),
|
||||
Some(&json!("task-registered"))
|
||||
);
|
||||
assert!(adapter.resolve_cached(&transport, &entry).is_none());
|
||||
let cached_auth = adapter
|
||||
.resolve_cached(&transport, &entry)
|
||||
.expect("cached task should create a new assertion");
|
||||
.resolve_refreshed(&transport, &entry)
|
||||
.expect("new refresh result should create an assertion");
|
||||
assert!(matches!(
|
||||
cached_auth,
|
||||
LocalResolvedOAuthRequestAuth::Header { ref name, ref value }
|
||||
@@ -1093,6 +1543,120 @@ mod tests {
|
||||
.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_cached_task_after_agent_credential_rotation() {
|
||||
let transport = sample_transport(test_auth_config(None));
|
||||
let requests = Arc::new(Mutex::new(Vec::new()));
|
||||
let adapter = CodexAgentIdentityRefreshAdapter::default()
|
||||
.with_auth_api_base_url_for_tests("https://auth.test/api/accounts");
|
||||
let entry = adapter
|
||||
.refresh(
|
||||
&RecordingExecutor {
|
||||
requests: Arc::clone(&requests),
|
||||
},
|
||||
&transport,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("registration should succeed")
|
||||
.expect("registration should return an entry");
|
||||
let mut rotated_config = test_auth_config(Some("task-rotated"));
|
||||
rotated_config["agent_runtime_id"] = json!("runtime-rotated");
|
||||
let rotated_transport = sample_transport(rotated_config);
|
||||
|
||||
assert!(adapter.resolve_cached(&rotated_transport, &entry).is_none());
|
||||
assert!(adapter
|
||||
.resolve_without_refresh(&rotated_transport)
|
||||
.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_cached_task_after_remote_task_rotation() {
|
||||
let transport = sample_transport(test_auth_config(None));
|
||||
let adapter = CodexAgentIdentityRefreshAdapter::default()
|
||||
.with_auth_api_base_url_for_tests("https://auth.test/api/accounts");
|
||||
let entry = adapter
|
||||
.refresh(
|
||||
&RecordingExecutor {
|
||||
requests: Arc::new(Mutex::new(Vec::new())),
|
||||
},
|
||||
&transport,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("registration should succeed")
|
||||
.expect("registration should return an entry");
|
||||
let rotated_transport = sample_transport(test_auth_config(Some("task-new-winner")));
|
||||
|
||||
assert!(adapter.resolve_cached(&rotated_transport, &entry).is_none());
|
||||
assert!(adapter
|
||||
.resolve_without_refresh(&rotated_transport)
|
||||
.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn newer_transport_task_is_authoritative_over_stale_local_cache() {
|
||||
let stale_transport = sample_transport(test_auth_config(Some("task-stale")));
|
||||
let stale_entry = codex_agent_identity_cached_entry_from_transport(&stale_transport)
|
||||
.expect("stale transport should produce a cache entry");
|
||||
let current_config = test_auth_config(Some("task-current"));
|
||||
let current_transport = sample_transport(current_config.clone());
|
||||
|
||||
assert_eq!(
|
||||
codex_agent_identity_refresh_fingerprint(¤t_transport, Some(&stale_entry)),
|
||||
codex_agent_identity_config_refresh_fingerprint(¤t_config)
|
||||
);
|
||||
assert!(CodexAgentIdentityRefreshAdapter::default()
|
||||
.resolve_fenced_cached(¤t_transport, &stale_entry)
|
||||
.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn distributed_waiter_reuses_reloaded_transport_not_stale_cache() {
|
||||
let stale_transport = sample_transport(test_auth_config(Some("task-stale")));
|
||||
let stale_entry = codex_agent_identity_cached_entry_from_transport(&stale_transport)
|
||||
.expect("stale transport should produce a cache entry");
|
||||
let expected = codex_agent_identity_refresh_fingerprint(&stale_transport, None)
|
||||
.expect("stale transport should have a generation");
|
||||
let current_transport = sample_transport(test_auth_config(Some("task-current")));
|
||||
let requests = Arc::new(Mutex::new(Vec::new()));
|
||||
let coordinator = LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![Arc::new(
|
||||
CodexAgentIdentityRefreshAdapter::default()
|
||||
.with_auth_api_base_url_for_tests("https://auth.test/api/accounts"),
|
||||
)]);
|
||||
coordinator
|
||||
.store_cached_entry(current_transport.key.id.as_str(), stale_entry)
|
||||
.await;
|
||||
|
||||
let resolution = coordinator
|
||||
.force_refresh_with_result_fenced(
|
||||
&RecordingExecutor {
|
||||
requests: Arc::clone(&requests),
|
||||
},
|
||||
¤t_transport,
|
||||
None,
|
||||
None,
|
||||
Some(expected.as_str()),
|
||||
)
|
||||
.await
|
||||
.expect("waiter resolution should succeed")
|
||||
.expect("waiter should reuse the DB winner");
|
||||
|
||||
assert!(resolution.reused_refresh);
|
||||
assert_eq!(
|
||||
resolution
|
||||
.refreshed_entry
|
||||
.as_ref()
|
||||
.and_then(|entry| entry.metadata.as_ref())
|
||||
.and_then(|config| config.get("task_id")),
|
||||
Some(&json!("task-current"))
|
||||
);
|
||||
assert!(requests
|
||||
.lock()
|
||||
.expect("recording lock should hold")
|
||||
.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn registration_response_accepts_encrypted_task_aliases() {
|
||||
let config = test_auth_config(Some("task-original"));
|
||||
|
||||
@@ -6,6 +6,7 @@ use aether_oauth::provider::providers::{
|
||||
use aether_oauth::provider::{ProviderOAuthAccount, ProviderOAuthAdapter, ProviderOAuthTokenSet};
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use super::oauth_refresh::{
|
||||
oauth_error_to_local_refresh_error, provider_oauth_transport_context_from_snapshot,
|
||||
@@ -73,6 +74,7 @@ impl GenericOAuthRefreshAdapter {
|
||||
entry
|
||||
.provider_type
|
||||
.eq_ignore_ascii_case(transport.provider.provider_type.as_str())
|
||||
&& generic_oauth_cached_entry_matches_transport(transport, entry)
|
||||
})
|
||||
.cloned()
|
||||
}
|
||||
@@ -147,6 +149,7 @@ impl GenericOAuthRefreshAdapter {
|
||||
|
||||
fn build_cached_entry(
|
||||
provider_type: &'static str,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
refreshed: ProviderOAuthTokenSet,
|
||||
) -> CachedOAuthEntry {
|
||||
CachedOAuthEntry {
|
||||
@@ -155,6 +158,7 @@ impl GenericOAuthRefreshAdapter {
|
||||
auth_header_value: refreshed.token_set.bearer_header_value(),
|
||||
expires_at_unix_secs: refreshed.token_set.expires_at_unix_secs,
|
||||
metadata: Some(refreshed.auth_config),
|
||||
source_fingerprint: Some(generic_oauth_transport_source_fingerprint(transport)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -188,6 +192,9 @@ impl LocalOAuthRefreshAdapter for GenericOAuthRefreshAdapter {
|
||||
value,
|
||||
});
|
||||
}
|
||||
if !generic_oauth_cached_entry_matches_transport(transport, entry) {
|
||||
return None;
|
||||
}
|
||||
if expires_at_requires_refresh(entry.expires_at_unix_secs) {
|
||||
return None;
|
||||
}
|
||||
@@ -293,10 +300,46 @@ impl LocalOAuthRefreshAdapter for GenericOAuthRefreshAdapter {
|
||||
"gateway generic oauth refresh succeeded"
|
||||
);
|
||||
|
||||
Ok(Some(Self::build_cached_entry(provider_type, refreshed)))
|
||||
Ok(Some(Self::build_cached_entry(
|
||||
provider_type,
|
||||
transport,
|
||||
refreshed,
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
fn generic_oauth_transport_source_fingerprint(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> String {
|
||||
let provider_type = transport.provider.provider_type.trim().to_ascii_lowercase();
|
||||
let auth_type = transport.key.auth_type.trim().to_ascii_lowercase();
|
||||
let auth_config = transport
|
||||
.key
|
||||
.decrypted_auth_config
|
||||
.as_deref()
|
||||
.unwrap_or_default();
|
||||
let api_key = transport.key.decrypted_api_key.as_str();
|
||||
let mut digest = Sha256::new();
|
||||
for field in [
|
||||
provider_type.as_bytes(),
|
||||
auth_type.as_bytes(),
|
||||
auth_config.as_bytes(),
|
||||
api_key.as_bytes(),
|
||||
] {
|
||||
digest.update((field.len() as u64).to_be_bytes());
|
||||
digest.update(field);
|
||||
}
|
||||
format!("{:x}", digest.finalize())
|
||||
}
|
||||
|
||||
fn generic_oauth_cached_entry_matches_transport(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
entry: &CachedOAuthEntry,
|
||||
) -> bool {
|
||||
let transport_fingerprint = generic_oauth_transport_source_fingerprint(transport);
|
||||
entry.source_fingerprint.as_deref() == Some(transport_fingerprint.as_str())
|
||||
}
|
||||
|
||||
fn generic_provider_type(provider_type: &str) -> Option<&'static str> {
|
||||
let normalized = provider_type.trim();
|
||||
GENERIC_PROVIDER_OAUTH_TEMPLATES
|
||||
@@ -362,6 +405,7 @@ fn current_access_token(
|
||||
entry: Option<&CachedOAuthEntry>,
|
||||
) -> Option<String> {
|
||||
entry
|
||||
.filter(|entry| generic_oauth_cached_entry_matches_transport(transport, entry))
|
||||
.and_then(|entry| {
|
||||
entry
|
||||
.auth_header_value
|
||||
@@ -388,7 +432,10 @@ mod tests {
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
|
||||
};
|
||||
use super::GenericOAuthRefreshAdapter;
|
||||
use super::{
|
||||
current_access_token, generic_oauth_transport_source_fingerprint,
|
||||
GenericOAuthRefreshAdapter,
|
||||
};
|
||||
|
||||
fn sample_transport() -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
@@ -481,6 +528,7 @@ mod tests {
|
||||
auth_header_value: "Bearer refreshed-access-token".to_string(),
|
||||
expires_at_unix_secs: Some(u64::MAX),
|
||||
metadata: None,
|
||||
source_fingerprint: None,
|
||||
};
|
||||
let auth = adapter
|
||||
.resolve_cached(&sample_transport(), &entry)
|
||||
@@ -494,4 +542,98 @@ mod tests {
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_cached_bearer_and_metadata_from_replaced_credential_generation() {
|
||||
let adapter = GenericOAuthRefreshAdapter::default();
|
||||
let mut original = sample_transport();
|
||||
original.key.decrypted_api_key = "access-a".to_string();
|
||||
original.key.decrypted_auth_config = Some(
|
||||
json!({
|
||||
"provider_type": "codex",
|
||||
"refresh_token": "refresh-a",
|
||||
"expires_at": 1,
|
||||
"updated_at": 100,
|
||||
})
|
||||
.to_string(),
|
||||
);
|
||||
let entry = CachedOAuthEntry {
|
||||
provider_type: "codex".to_string(),
|
||||
auth_header_name: "authorization".to_string(),
|
||||
auth_header_value: "Bearer cached-access-a".to_string(),
|
||||
expires_at_unix_secs: Some(u64::MAX),
|
||||
metadata: Some(json!({
|
||||
"provider_type": "codex",
|
||||
"refresh_token": "rotated-refresh-a",
|
||||
"expires_at": u64::MAX,
|
||||
"updated_at": 200,
|
||||
})),
|
||||
source_fingerprint: Some(generic_oauth_transport_source_fingerprint(&original)),
|
||||
};
|
||||
|
||||
let mut replacement = original.clone();
|
||||
replacement.key.decrypted_api_key = "access-b".to_string();
|
||||
replacement.key.decrypted_auth_config = Some(
|
||||
json!({
|
||||
"provider_type": "codex",
|
||||
"refresh_token": "refresh-b",
|
||||
"expires_at": 1,
|
||||
"updated_at": 300,
|
||||
})
|
||||
.to_string(),
|
||||
);
|
||||
|
||||
assert!(adapter.resolve_cached(&replacement, &entry).is_none());
|
||||
assert_eq!(
|
||||
adapter.base_auth_config(&replacement, Some(&entry)),
|
||||
replacement
|
||||
.key
|
||||
.decrypted_auth_config
|
||||
.as_deref()
|
||||
.and_then(|value| serde_json::from_str(value).ok())
|
||||
);
|
||||
assert_eq!(
|
||||
current_access_token(&replacement, Some(&entry)).as_deref(),
|
||||
Some("access-b")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reuses_cached_bearer_from_matching_credential_generation() {
|
||||
let adapter = GenericOAuthRefreshAdapter::default();
|
||||
let mut transport = sample_transport();
|
||||
transport.key.decrypted_api_key = "access-a".to_string();
|
||||
transport.key.decrypted_auth_config = Some(
|
||||
json!({
|
||||
"provider_type": "codex",
|
||||
"refresh_token": "refresh-a",
|
||||
"expires_at": 1,
|
||||
})
|
||||
.to_string(),
|
||||
);
|
||||
let entry = CachedOAuthEntry {
|
||||
provider_type: "codex".to_string(),
|
||||
auth_header_name: "authorization".to_string(),
|
||||
auth_header_value: "Bearer refreshed-access-a".to_string(),
|
||||
expires_at_unix_secs: Some(u64::MAX),
|
||||
metadata: Some(json!({
|
||||
"provider_type": "codex",
|
||||
"refresh_token": "rotated-refresh-a",
|
||||
"expires_at": u64::MAX,
|
||||
})),
|
||||
source_fingerprint: Some(generic_oauth_transport_source_fingerprint(&transport)),
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
adapter.resolve_cached(&transport, &entry),
|
||||
Some(LocalResolvedOAuthRequestAuth::Header {
|
||||
name: "authorization".to_string(),
|
||||
value: "Bearer refreshed-access-a".to_string(),
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
current_access_token(&transport, Some(&entry)).as_deref(),
|
||||
Some("refreshed-access-a")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -77,6 +77,7 @@ impl KiroOAuthRefreshAdapter {
|
||||
auth_header_value: request_auth.value,
|
||||
expires_at_unix_secs: auth_config.expires_at,
|
||||
metadata: Some(auth_config.to_json_value()),
|
||||
source_fingerprint: None,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -30,13 +30,21 @@ pub mod windsurf;
|
||||
|
||||
pub use aether_oauth as oauth;
|
||||
pub use agent_identity::{
|
||||
create_codex_agent_identity_from_session_token, is_codex_agent_identity_auth_config_value,
|
||||
codex_agent_identity_auth_config_has_task_id,
|
||||
codex_agent_identity_authorization_matches_transport,
|
||||
codex_agent_identity_cached_entry_from_transport,
|
||||
codex_agent_identity_config_refresh_fingerprint, codex_agent_identity_credential_fingerprint,
|
||||
codex_agent_identity_entry_allows_task_rotation_from, codex_agent_identity_refresh_fingerprint,
|
||||
codex_agent_identity_transport_allows_task_rotation_from,
|
||||
codex_agent_identity_transport_credential_fingerprint,
|
||||
create_codex_agent_identity_from_access_token, create_codex_agent_identity_from_session_token,
|
||||
is_codex_agent_identity_auth_config_value, is_codex_agent_identity_authorization,
|
||||
is_codex_agent_identity_cached_entry, is_codex_agent_identity_invalid_task_response,
|
||||
is_codex_agent_identity_transport, validate_codex_agent_identity_auth_config,
|
||||
CodexAgentIdentityEnrollmentError, CodexAgentIdentityRefreshAdapter,
|
||||
CODEX_AGENT_IDENTITY_AGENT_REGISTRATION_REQUEST_ID, CODEX_AGENT_IDENTITY_AUTH_MODE,
|
||||
CODEX_AGENT_IDENTITY_CACHED_ENTRY_PROVIDER_TYPE, CODEX_AGENT_IDENTITY_PROVIDER_TYPE,
|
||||
CODEX_AGENT_IDENTITY_TASK_REGISTRATION_REQUEST_ID,
|
||||
is_codex_agent_identity_transport, register_codex_agent_identity_from_access_token,
|
||||
validate_codex_agent_identity_auth_config, CodexAgentIdentityEnrollmentError,
|
||||
CodexAgentIdentityRefreshAdapter, CODEX_AGENT_IDENTITY_AGENT_REGISTRATION_REQUEST_ID,
|
||||
CODEX_AGENT_IDENTITY_AUTH_MODE, CODEX_AGENT_IDENTITY_CACHED_ENTRY_PROVIDER_TYPE,
|
||||
CODEX_AGENT_IDENTITY_PROVIDER_TYPE, CODEX_AGENT_IDENTITY_TASK_REGISTRATION_REQUEST_ID,
|
||||
};
|
||||
pub use auth::{build_passthrough_headers, ensure_upstream_auth_header};
|
||||
pub use auth_config::apply_local_auth_config_header_overrides;
|
||||
@@ -91,7 +99,8 @@ pub use network::{
|
||||
pub use oauth_refresh::{
|
||||
supports_local_oauth_request_auth_resolution, CachedOAuthEntry, LocalOAuthHttpExecutor,
|
||||
LocalOAuthHttpRequest, LocalOAuthHttpResponse, LocalOAuthRefreshCoordinator,
|
||||
LocalOAuthRefreshError, LocalResolvedOAuthRequestAuth, ReqwestLocalOAuthHttpExecutor,
|
||||
LocalOAuthRefreshError, LocalOAuthResolution, LocalResolvedOAuthRequestAuth,
|
||||
ReqwestLocalOAuthHttpExecutor,
|
||||
};
|
||||
pub use openai_image::{
|
||||
build_openai_image_headers, build_openai_image_upstream_url,
|
||||
|
||||
@@ -1,13 +1,14 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use aether_oauth::core::OAuthError;
|
||||
use aether_oauth::network::{
|
||||
OAuthHttpExecutor, OAuthHttpRequest, OAuthHttpResponse, OAuthNetworkContext,
|
||||
};
|
||||
use aether_oauth::provider::ProviderOAuthTransportContext;
|
||||
use aether_runtime_state::RuntimeState;
|
||||
use aether_runtime_state::{RuntimeLockLease, RuntimeState};
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
use thiserror::Error;
|
||||
@@ -40,6 +41,12 @@ pub struct LocalOAuthResolution {
|
||||
pub auth: Option<LocalResolvedOAuthRequestAuth>,
|
||||
pub refreshed_entry: Option<CachedOAuthEntry>,
|
||||
pub refresh_in_flight: bool,
|
||||
/// Indicates that a forced caller reused a newer completed refresh rather
|
||||
/// than producing a new entry that needs persistence.
|
||||
pub reused_refresh: bool,
|
||||
/// Held until the caller persists `refreshed_entry`. The lease TTL remains
|
||||
/// the cancellation fallback if the caller is dropped.
|
||||
pub distributed_lease: Option<RuntimeLockLease>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
@@ -49,6 +56,8 @@ pub struct CachedOAuthEntry {
|
||||
pub auth_header_value: String,
|
||||
pub expires_at_unix_secs: Option<u64>,
|
||||
pub metadata: Option<Value>,
|
||||
/// Non-secret fingerprint of the credential/configuration that produced it.
|
||||
pub source_fingerprint: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
@@ -297,6 +306,28 @@ pub trait LocalOAuthRefreshAdapter: Send + Sync {
|
||||
entry: &CachedOAuthEntry,
|
||||
) -> Option<LocalResolvedOAuthRequestAuth>;
|
||||
|
||||
/// Resolves a cache entry that is known to have advanced the caller's
|
||||
/// refresh fence. Agent task rotation can safely use the winner even while
|
||||
/// the caller still holds the pre-refresh transport snapshot.
|
||||
fn resolve_fenced_cached(
|
||||
&self,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
entry: &CachedOAuthEntry,
|
||||
) -> Option<LocalResolvedOAuthRequestAuth> {
|
||||
self.resolve_cached(transport, entry)
|
||||
}
|
||||
|
||||
/// Resolves the entry returned by this adapter's immediately preceding
|
||||
/// refresh. Unlike a reusable cache entry, this entry is expected to have
|
||||
/// advanced the transport generation.
|
||||
fn resolve_refreshed(
|
||||
&self,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
entry: &CachedOAuthEntry,
|
||||
) -> Option<LocalResolvedOAuthRequestAuth> {
|
||||
self.resolve_cached(transport, entry)
|
||||
}
|
||||
|
||||
fn resolve_without_refresh(
|
||||
&self,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
@@ -308,6 +339,36 @@ pub trait LocalOAuthRefreshAdapter: Send + Sync {
|
||||
entry: Option<&CachedOAuthEntry>,
|
||||
) -> bool;
|
||||
|
||||
/// Identifies the credential/configuration generation used by a refresh.
|
||||
/// Adapters that support fencing override this method.
|
||||
fn refresh_fingerprint(
|
||||
&self,
|
||||
_transport: &GatewayProviderTransportSnapshot,
|
||||
_entry: Option<&CachedOAuthEntry>,
|
||||
) -> Option<String> {
|
||||
None
|
||||
}
|
||||
|
||||
/// Reconstructs a cache entry from an already-persisted transport after a
|
||||
/// distributed refresh waiter reloads the winner.
|
||||
fn cached_entry_from_transport(
|
||||
&self,
|
||||
_transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<CachedOAuthEntry> {
|
||||
None
|
||||
}
|
||||
|
||||
/// Enables bounded negative backoff for transient refresh failures.
|
||||
fn should_backoff_after_error(&self, _error: &LocalOAuthRefreshError) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// Agent task registration is a non-idempotent external mutation and must
|
||||
/// not continue unlocked when a configured distributed lock is unavailable.
|
||||
fn requires_distributed_refresh_lock(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
async fn refresh(
|
||||
&self,
|
||||
executor: &dyn LocalOAuthHttpExecutor,
|
||||
@@ -320,6 +381,13 @@ pub struct LocalOAuthRefreshCoordinator {
|
||||
adapters: Vec<Arc<dyn LocalOAuthRefreshAdapter>>,
|
||||
cache: Mutex<BTreeMap<String, CachedOAuthEntry>>,
|
||||
key_locks: Mutex<BTreeMap<String, Arc<Mutex<()>>>>,
|
||||
refresh_backoff: Mutex<BTreeMap<String, RefreshBackoffState>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct RefreshBackoffState {
|
||||
failures: u32,
|
||||
retry_after: Instant,
|
||||
}
|
||||
|
||||
impl fmt::Debug for LocalOAuthRefreshCoordinator {
|
||||
@@ -337,7 +405,10 @@ impl Default for LocalOAuthRefreshCoordinator {
|
||||
}
|
||||
|
||||
impl LocalOAuthRefreshCoordinator {
|
||||
const DISTRIBUTED_REFRESH_LOCK_TTL_MS: u64 = 30_000;
|
||||
// Keep the lease alive through the 30s upstream HTTP timeout and the
|
||||
// subsequent encrypted DB CAS/persistence step. Cancellation still relies
|
||||
// on expiry as the last-resort release path.
|
||||
const DISTRIBUTED_REFRESH_LOCK_TTL_MS: u64 = 120_000;
|
||||
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
@@ -349,6 +420,7 @@ impl LocalOAuthRefreshCoordinator {
|
||||
],
|
||||
cache: Mutex::new(BTreeMap::new()),
|
||||
key_locks: Mutex::new(BTreeMap::new()),
|
||||
refresh_backoff: Mutex::new(BTreeMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -373,7 +445,9 @@ impl LocalOAuthRefreshCoordinator {
|
||||
}
|
||||
|
||||
pub async fn invalidate_cached_entry(&self, key_id: &str) -> bool {
|
||||
self.cache.lock().await.remove(key_id).is_some()
|
||||
let removed = self.cache.lock().await.remove(key_id).is_some();
|
||||
self.clear_refresh_backoff(key_id).await;
|
||||
removed
|
||||
}
|
||||
|
||||
pub async fn resolve_with_result(
|
||||
@@ -389,6 +463,7 @@ impl LocalOAuthRefreshCoordinator {
|
||||
distributed_lock,
|
||||
distributed_owner,
|
||||
false,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -399,6 +474,27 @@ impl LocalOAuthRefreshCoordinator {
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
distributed_lock: Option<&RuntimeState>,
|
||||
distributed_owner: Option<&str>,
|
||||
) -> Result<Option<LocalOAuthResolution>, LocalOAuthRefreshError> {
|
||||
self.force_refresh_with_result_fenced(
|
||||
executor,
|
||||
transport,
|
||||
distributed_lock,
|
||||
distributed_owner,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Force a refresh unless another request has already advanced the supplied
|
||||
/// refresh fence. This prevents a distributed waiter from re-registering a
|
||||
/// task after the winner has persisted it.
|
||||
pub async fn force_refresh_with_result_fenced(
|
||||
&self,
|
||||
executor: &dyn LocalOAuthHttpExecutor,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
distributed_lock: Option<&RuntimeState>,
|
||||
distributed_owner: Option<&str>,
|
||||
expected_refresh_fingerprint: Option<&str>,
|
||||
) -> Result<Option<LocalOAuthResolution>, LocalOAuthRefreshError> {
|
||||
self.resolve_with_result_mode(
|
||||
executor,
|
||||
@@ -406,6 +502,7 @@ impl LocalOAuthRefreshCoordinator {
|
||||
distributed_lock,
|
||||
distributed_owner,
|
||||
true,
|
||||
expected_refresh_fingerprint,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -417,6 +514,7 @@ impl LocalOAuthRefreshCoordinator {
|
||||
distributed_lock: Option<&RuntimeState>,
|
||||
distributed_owner: Option<&str>,
|
||||
force_refresh: bool,
|
||||
expected_refresh_fingerprint: Option<&str>,
|
||||
) -> Result<Option<LocalOAuthResolution>, LocalOAuthRefreshError> {
|
||||
let Some(adapter) = self
|
||||
.adapters
|
||||
@@ -450,10 +548,37 @@ impl LocalOAuthRefreshCoordinator {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if force_refresh {
|
||||
if let Some(resolution) = Self::resolve_if_refresh_fence_advanced(
|
||||
adapter.as_ref(),
|
||||
transport,
|
||||
cached_entry.as_ref(),
|
||||
expected_refresh_fingerprint,
|
||||
) {
|
||||
return Ok(Some(resolution));
|
||||
}
|
||||
}
|
||||
if let Some(error) = self.backoff_error(key_id, adapter.provider_type()).await {
|
||||
return Err(error);
|
||||
}
|
||||
|
||||
let key_lock = self.lock_for_key(key_id).await;
|
||||
let _key_guard = key_lock.lock().await;
|
||||
|
||||
let cached_entry = self.cached_entry(key_id).await;
|
||||
if force_refresh {
|
||||
if let Some(resolution) = Self::resolve_if_refresh_fence_advanced(
|
||||
adapter.as_ref(),
|
||||
transport,
|
||||
cached_entry.as_ref(),
|
||||
expected_refresh_fingerprint,
|
||||
) {
|
||||
return Ok(Some(resolution));
|
||||
}
|
||||
}
|
||||
if let Some(error) = self.backoff_error(key_id, adapter.provider_type()).await {
|
||||
return Err(error);
|
||||
}
|
||||
if !force_refresh {
|
||||
if let Some(auth) = cached_entry
|
||||
.as_ref()
|
||||
@@ -488,6 +613,16 @@ impl LocalOAuthRefreshCoordinator {
|
||||
error = ?err,
|
||||
"gateway local oauth refresh distributed lock unavailable"
|
||||
);
|
||||
if adapter.requires_distributed_refresh_lock() {
|
||||
let error = LocalOAuthRefreshError::TransportMessage {
|
||||
provider_type: adapter.provider_type(),
|
||||
message: "distributed refresh lock is unavailable".to_string(),
|
||||
};
|
||||
if adapter.should_backoff_after_error(&error) {
|
||||
self.record_refresh_failure(key_id).await;
|
||||
}
|
||||
return Err(error);
|
||||
}
|
||||
None
|
||||
}
|
||||
}
|
||||
@@ -501,22 +636,147 @@ impl LocalOAuthRefreshCoordinator {
|
||||
// came from the original transport snapshot.
|
||||
let refresh_entry = cached_entry.as_ref();
|
||||
let refresh_result = adapter.refresh(executor, transport, refresh_entry).await;
|
||||
if let (Some(lock), Some(lease)) = (distributed_lock, distributed_lease.as_ref()) {
|
||||
if let Err(err) = lock.lock_release(lease).await {
|
||||
tracing::warn!(
|
||||
key_id = %key_id,
|
||||
provider_type = adapter.provider_type(),
|
||||
error = ?err,
|
||||
"gateway local oauth refresh distributed lock release failed"
|
||||
);
|
||||
let refreshed_entry = match refresh_result {
|
||||
Ok(Some(entry)) => {
|
||||
self.clear_refresh_backoff(key_id).await;
|
||||
entry
|
||||
}
|
||||
Ok(None) => {
|
||||
Self::release_distributed_lease(
|
||||
distributed_lock,
|
||||
distributed_lease.as_ref(),
|
||||
key_id,
|
||||
adapter.provider_type(),
|
||||
)
|
||||
.await;
|
||||
return Ok(None);
|
||||
}
|
||||
Err(error) => {
|
||||
if adapter.should_backoff_after_error(&error) {
|
||||
self.record_refresh_failure(key_id).await;
|
||||
}
|
||||
Self::release_distributed_lease(
|
||||
distributed_lock,
|
||||
distributed_lease.as_ref(),
|
||||
key_id,
|
||||
adapter.provider_type(),
|
||||
)
|
||||
.await;
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
// In production the distributed lease is held through the gateway's
|
||||
// DB CAS. Do not publish a provisional task before that CAS succeeds;
|
||||
// otherwise a waiter could consume an assertion that loses the CAS.
|
||||
// Lock-free/test callers retain the historical in-memory behavior.
|
||||
if distributed_lease.is_none() {
|
||||
self.insert_cached_entry(key_id, refreshed_entry.clone())
|
||||
.await;
|
||||
}
|
||||
let Some(refreshed_entry) = refresh_result? else {
|
||||
let Some(auth) = adapter.resolve_refreshed(transport, &refreshed_entry) else {
|
||||
Self::release_distributed_lease(
|
||||
distributed_lock,
|
||||
distributed_lease.as_ref(),
|
||||
key_id,
|
||||
adapter.provider_type(),
|
||||
)
|
||||
.await;
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(adapter
|
||||
.resolve_cached(transport, &refreshed_entry)
|
||||
.map(|auth| LocalOAuthResolution::resolved(auth, Some(refreshed_entry))))
|
||||
Ok(Some(LocalOAuthResolution::refreshed(
|
||||
auth,
|
||||
refreshed_entry,
|
||||
distributed_lease,
|
||||
)))
|
||||
}
|
||||
|
||||
fn resolve_if_refresh_fence_advanced(
|
||||
adapter: &dyn LocalOAuthRefreshAdapter,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
entry: Option<&CachedOAuthEntry>,
|
||||
expected_refresh_fingerprint: Option<&str>,
|
||||
) -> Option<LocalOAuthResolution> {
|
||||
let expected = expected_refresh_fingerprint?;
|
||||
if adapter.refresh_fingerprint(transport, entry).as_deref() == Some(expected) {
|
||||
return None;
|
||||
}
|
||||
entry
|
||||
.and_then(|entry| adapter.resolve_fenced_cached(transport, entry))
|
||||
.map(|auth| {
|
||||
LocalOAuthResolution::reused(
|
||||
auth,
|
||||
entry.expect("cached auth is required when a refresh fence advanced"),
|
||||
)
|
||||
})
|
||||
.or_else(|| {
|
||||
adapter
|
||||
.cached_entry_from_transport(transport)
|
||||
.and_then(|entry| {
|
||||
adapter
|
||||
.resolve_cached(transport, &entry)
|
||||
.map(|auth| LocalOAuthResolution::reused(auth, &entry))
|
||||
})
|
||||
})
|
||||
.or_else(|| {
|
||||
adapter
|
||||
.resolve_without_refresh(transport)
|
||||
.map(|auth| LocalOAuthResolution::resolved(auth, None))
|
||||
})
|
||||
}
|
||||
|
||||
async fn backoff_error(
|
||||
&self,
|
||||
key_id: &str,
|
||||
provider_type: &'static str,
|
||||
) -> Option<LocalOAuthRefreshError> {
|
||||
let backoff = self.refresh_backoff.lock().await;
|
||||
let state = backoff.get(key_id)?;
|
||||
let remaining = state.retry_after.checked_duration_since(Instant::now())?;
|
||||
Some(LocalOAuthRefreshError::InvalidResponse {
|
||||
provider_type,
|
||||
message: format!(
|
||||
"refresh temporarily backed off after {} failed attempts (retry in {}ms)",
|
||||
state.failures,
|
||||
remaining.as_millis()
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
async fn record_refresh_failure(&self, key_id: &str) {
|
||||
let mut backoff = self.refresh_backoff.lock().await;
|
||||
let state = backoff
|
||||
.entry(key_id.to_string())
|
||||
.or_insert(RefreshBackoffState {
|
||||
failures: 0,
|
||||
retry_after: Instant::now(),
|
||||
});
|
||||
state.failures = state.failures.saturating_add(1);
|
||||
let exponent = state.failures.saturating_sub(1).min(4);
|
||||
let delay = Duration::from_millis(500u64.saturating_mul(1u64 << exponent));
|
||||
state.retry_after = Instant::now() + delay.min(Duration::from_secs(8));
|
||||
}
|
||||
|
||||
async fn clear_refresh_backoff(&self, key_id: &str) {
|
||||
self.refresh_backoff.lock().await.remove(key_id);
|
||||
}
|
||||
|
||||
async fn release_distributed_lease(
|
||||
distributed_lock: Option<&RuntimeState>,
|
||||
lease: Option<&RuntimeLockLease>,
|
||||
key_id: &str,
|
||||
provider_type: &'static str,
|
||||
) {
|
||||
let (Some(lock), Some(lease)) = (distributed_lock, lease) else {
|
||||
return;
|
||||
};
|
||||
if let Err(err) = lock.lock_release(lease).await {
|
||||
tracing::warn!(
|
||||
key_id = %key_id,
|
||||
provider_type,
|
||||
error = ?err,
|
||||
"gateway local oauth refresh distributed lock release failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_adapters_for_tests(adapters: Vec<Arc<dyn LocalOAuthRefreshAdapter>>) -> Self {
|
||||
@@ -524,6 +784,7 @@ impl LocalOAuthRefreshCoordinator {
|
||||
adapters,
|
||||
cache: Mutex::new(BTreeMap::new()),
|
||||
key_locks: Mutex::new(BTreeMap::new()),
|
||||
refresh_backoff: Mutex::new(BTreeMap::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -537,6 +798,37 @@ impl LocalOAuthResolution {
|
||||
auth: Some(auth),
|
||||
refreshed_entry,
|
||||
refresh_in_flight: false,
|
||||
reused_refresh: false,
|
||||
distributed_lease: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn refreshed(
|
||||
auth: LocalResolvedOAuthRequestAuth,
|
||||
refreshed_entry: CachedOAuthEntry,
|
||||
distributed_lease: Option<RuntimeLockLease>,
|
||||
) -> Self {
|
||||
Self {
|
||||
auth: Some(auth),
|
||||
refreshed_entry: Some(refreshed_entry),
|
||||
refresh_in_flight: false,
|
||||
reused_refresh: false,
|
||||
distributed_lease,
|
||||
}
|
||||
}
|
||||
|
||||
fn reused(auth: LocalResolvedOAuthRequestAuth, entry: &CachedOAuthEntry) -> Self {
|
||||
let mut refreshed_entry = entry.clone();
|
||||
if let LocalResolvedOAuthRequestAuth::Header { name, value } = &auth {
|
||||
refreshed_entry.auth_header_name = name.clone();
|
||||
refreshed_entry.auth_header_value = value.clone();
|
||||
}
|
||||
Self {
|
||||
auth: Some(auth),
|
||||
refreshed_entry: Some(refreshed_entry),
|
||||
refresh_in_flight: false,
|
||||
reused_refresh: true,
|
||||
distributed_lease: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -545,6 +837,8 @@ impl LocalOAuthResolution {
|
||||
auth: None,
|
||||
refreshed_entry: None,
|
||||
refresh_in_flight: true,
|
||||
reused_refresh: false,
|
||||
distributed_lease: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -560,7 +854,7 @@ pub fn supports_local_oauth_request_auth_resolution(
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||
|
||||
use super::super::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
@@ -580,6 +874,12 @@ mod tests {
|
||||
refresh_with_entry_hits: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct FencedTestAdapter {
|
||||
refresh_hits: Arc<AtomicUsize>,
|
||||
fail_refresh: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LocalOAuthRefreshAdapter for TestAdapter {
|
||||
fn provider_type(&self) -> &'static str {
|
||||
@@ -634,6 +934,77 @@ mod tests {
|
||||
auth_header_value: "Bearer refreshed-token".to_string(),
|
||||
expires_at_unix_secs: Some(4_102_444_800),
|
||||
metadata: None,
|
||||
source_fingerprint: None,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LocalOAuthRefreshAdapter for FencedTestAdapter {
|
||||
fn provider_type(&self) -> &'static str {
|
||||
"test-oauth"
|
||||
}
|
||||
|
||||
fn resolve_cached(
|
||||
&self,
|
||||
_transport: &GatewayProviderTransportSnapshot,
|
||||
entry: &CachedOAuthEntry,
|
||||
) -> Option<LocalResolvedOAuthRequestAuth> {
|
||||
Some(LocalResolvedOAuthRequestAuth::Header {
|
||||
name: entry.auth_header_name.clone(),
|
||||
value: "fresh-winner-assertion".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn resolve_without_refresh(
|
||||
&self,
|
||||
_transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<LocalResolvedOAuthRequestAuth> {
|
||||
None
|
||||
}
|
||||
|
||||
fn should_refresh(
|
||||
&self,
|
||||
_transport: &GatewayProviderTransportSnapshot,
|
||||
_entry: Option<&CachedOAuthEntry>,
|
||||
) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn refresh_fingerprint(
|
||||
&self,
|
||||
_transport: &GatewayProviderTransportSnapshot,
|
||||
entry: Option<&CachedOAuthEntry>,
|
||||
) -> Option<String> {
|
||||
entry
|
||||
.and_then(|entry| entry.source_fingerprint.clone())
|
||||
.or_else(|| Some("generation-1".to_string()))
|
||||
}
|
||||
|
||||
fn should_backoff_after_error(&self, _error: &LocalOAuthRefreshError) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn refresh(
|
||||
&self,
|
||||
_executor: &dyn LocalOAuthHttpExecutor,
|
||||
_transport: &GatewayProviderTransportSnapshot,
|
||||
_entry: Option<&CachedOAuthEntry>,
|
||||
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError> {
|
||||
self.refresh_hits.fetch_add(1, Ordering::SeqCst);
|
||||
if self.fail_refresh.load(Ordering::SeqCst) {
|
||||
return Err(LocalOAuthRefreshError::TransportMessage {
|
||||
provider_type: "test-oauth",
|
||||
message: "temporary failure".to_string(),
|
||||
});
|
||||
}
|
||||
Ok(Some(CachedOAuthEntry {
|
||||
provider_type: "test-oauth".to_string(),
|
||||
auth_header_name: "authorization".to_string(),
|
||||
auth_header_value: "stale-winner-cache-value".to_string(),
|
||||
expires_at_unix_secs: None,
|
||||
metadata: None,
|
||||
source_fingerprint: Some("generation-2".to_string()),
|
||||
}))
|
||||
}
|
||||
}
|
||||
@@ -739,8 +1110,11 @@ mod tests {
|
||||
auth_header_value: "Bearer refreshed-token".to_string(),
|
||||
expires_at_unix_secs: Some(4_102_444_800),
|
||||
metadata: None,
|
||||
source_fingerprint: None,
|
||||
}),
|
||||
refresh_in_flight: false,
|
||||
reused_refresh: false,
|
||||
distributed_lease: None,
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
@@ -752,6 +1126,8 @@ mod tests {
|
||||
}),
|
||||
refreshed_entry: None,
|
||||
refresh_in_flight: false,
|
||||
reused_refresh: false,
|
||||
distributed_lease: None,
|
||||
})
|
||||
);
|
||||
}
|
||||
@@ -791,4 +1167,86 @@ mod tests {
|
||||
assert_eq!(refresh_hits.load(Ordering::SeqCst), 2);
|
||||
assert_eq!(refresh_with_entry_hits.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fenced_force_refresh_reuses_the_winner() {
|
||||
let refresh_hits = Arc::new(AtomicUsize::new(0));
|
||||
let coordinator = LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![Arc::new(
|
||||
FencedTestAdapter {
|
||||
refresh_hits: Arc::clone(&refresh_hits),
|
||||
fail_refresh: Arc::new(AtomicBool::new(false)),
|
||||
},
|
||||
)]);
|
||||
let transport = sample_transport();
|
||||
let executor = ReqwestLocalOAuthHttpExecutor::new(reqwest::Client::new());
|
||||
|
||||
let first = coordinator
|
||||
.force_refresh_with_result_fenced(
|
||||
&executor,
|
||||
&transport,
|
||||
None,
|
||||
None,
|
||||
Some("generation-1"),
|
||||
)
|
||||
.await
|
||||
.expect("first refresh should succeed")
|
||||
.expect("first refresh should resolve");
|
||||
let waiter = coordinator
|
||||
.force_refresh_with_result_fenced(
|
||||
&executor,
|
||||
&transport,
|
||||
None,
|
||||
None,
|
||||
Some("generation-1"),
|
||||
)
|
||||
.await
|
||||
.expect("waiter should reuse winner")
|
||||
.expect("waiter should resolve");
|
||||
|
||||
assert!(first.refreshed_entry.is_some());
|
||||
assert!(waiter.refreshed_entry.is_some());
|
||||
assert!(waiter.reused_refresh);
|
||||
assert_eq!(
|
||||
waiter
|
||||
.refreshed_entry
|
||||
.as_ref()
|
||||
.expect("reused entry")
|
||||
.auth_header_value,
|
||||
"fresh-winner-assertion"
|
||||
);
|
||||
assert_eq!(refresh_hits.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_failure_enters_bounded_negative_backoff() {
|
||||
let refresh_hits = Arc::new(AtomicUsize::new(0));
|
||||
let fail_refresh = Arc::new(AtomicBool::new(true));
|
||||
let coordinator = LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![Arc::new(
|
||||
FencedTestAdapter {
|
||||
refresh_hits: Arc::clone(&refresh_hits),
|
||||
fail_refresh: Arc::clone(&fail_refresh),
|
||||
},
|
||||
)]);
|
||||
let transport = sample_transport();
|
||||
let executor = ReqwestLocalOAuthHttpExecutor::new(reqwest::Client::new());
|
||||
|
||||
assert!(coordinator
|
||||
.force_refresh_with_result(&executor, &transport, None, None)
|
||||
.await
|
||||
.is_err());
|
||||
let second = coordinator
|
||||
.force_refresh_with_result(&executor, &transport, None, None)
|
||||
.await
|
||||
.expect_err("second refresh should be backed off");
|
||||
assert!(second.to_string().contains("temporarily backed off"));
|
||||
assert_eq!(refresh_hits.load(Ordering::SeqCst), 1);
|
||||
fail_refresh.store(false, Ordering::SeqCst);
|
||||
coordinator.invalidate_cached_entry("key-1").await;
|
||||
assert!(coordinator
|
||||
.force_refresh_with_result(&executor, &transport, None, None)
|
||||
.await
|
||||
.expect("replacement should refresh immediately")
|
||||
.is_some());
|
||||
assert_eq!(refresh_hits.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -259,6 +259,7 @@ impl LocalOAuthRefreshAdapter for VertexServiceAccountRefreshAdapter {
|
||||
"project_id": auth_config.project_id,
|
||||
"client_email": auth_config.client_email,
|
||||
})),
|
||||
source_fingerprint: None,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,140 @@
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
|
||||
const { getMock, postMock } = vi.hoisted(() => ({
|
||||
getMock: vi.fn(),
|
||||
postMock: vi.fn(),
|
||||
}))
|
||||
|
||||
vi.mock('@/api/client', () => ({
|
||||
default: {
|
||||
get: getMock,
|
||||
post: postMock,
|
||||
},
|
||||
}))
|
||||
|
||||
import {
|
||||
getBatchImportOAuthTaskStatus,
|
||||
importProviderRefreshToken,
|
||||
startBatchImportOAuthTask,
|
||||
} from '@/api/endpoints/provider_oauth'
|
||||
|
||||
describe('Agent Identity OAuth management routes', () => {
|
||||
beforeEach(() => {
|
||||
getMock.mockReset()
|
||||
postMock.mockReset()
|
||||
getMock.mockResolvedValue({ data: {} })
|
||||
postMock.mockResolvedValue({ data: {} })
|
||||
})
|
||||
|
||||
it('routes one Agent Identity JSON through the provider-oauth permission surface', async () => {
|
||||
const credentials = JSON.stringify({
|
||||
auth_mode: 'agentIdentity',
|
||||
agent_runtime_id: 'runtime-1',
|
||||
agent_private_key: 'private-key',
|
||||
task_id: 'task-1',
|
||||
})
|
||||
|
||||
await startBatchImportOAuthTask('provider-codex', credentials, 'proxy-1')
|
||||
|
||||
expect(postMock).toHaveBeenCalledWith(
|
||||
'/api/admin/provider-oauth/providers/provider-codex/agent-identity-import/tasks',
|
||||
{ credentials, proxy_node_id: 'proxy-1' },
|
||||
)
|
||||
})
|
||||
|
||||
it('routes Agent Identity arrays through the dedicated import surface', async () => {
|
||||
const credentials = JSON.stringify([
|
||||
{
|
||||
auth_mode: 'agentIdentity',
|
||||
agent_runtime_id: 'runtime-1',
|
||||
agent_private_key: 'private-key-1',
|
||||
},
|
||||
{
|
||||
auth_mode: 'agentIdentity',
|
||||
agent_runtime_id: 'runtime-2',
|
||||
agent_private_key: 'private-key-2',
|
||||
},
|
||||
])
|
||||
|
||||
await startBatchImportOAuthTask('provider-codex', credentials)
|
||||
|
||||
expect(postMock).toHaveBeenCalledWith(
|
||||
'/api/admin/provider-oauth/providers/provider-codex/agent-identity-import/tasks',
|
||||
{ credentials, proxy_node_id: undefined },
|
||||
)
|
||||
})
|
||||
|
||||
it('routes sub2api Agent Identity exports through the dedicated import surface', async () => {
|
||||
const credentials = JSON.stringify({
|
||||
type: 'sub2api-data',
|
||||
accounts: [{
|
||||
platform: 'openai',
|
||||
credentials: {
|
||||
auth_mode: 'agentIdentity',
|
||||
agent_runtime_id: 'runtime-1',
|
||||
agent_private_key: 'private-key',
|
||||
},
|
||||
}],
|
||||
})
|
||||
|
||||
await startBatchImportOAuthTask('provider-codex', credentials)
|
||||
|
||||
expect(postMock).toHaveBeenCalledWith(
|
||||
'/api/admin/provider-oauth/providers/provider-codex/agent-identity-import/tasks',
|
||||
{ credentials, proxy_node_id: undefined },
|
||||
)
|
||||
})
|
||||
|
||||
it('routes mixed Agent Identity credentials to the dedicated surface for rejection', async () => {
|
||||
const credentials = JSON.stringify([
|
||||
{ refresh_token: 'ordinary-refresh-token' },
|
||||
{
|
||||
auth_mode: 'agentIdentity',
|
||||
agent_runtime_id: 'runtime-1',
|
||||
agent_private_key: 'private-key',
|
||||
},
|
||||
])
|
||||
|
||||
await startBatchImportOAuthTask('provider-codex', credentials)
|
||||
|
||||
expect(postMock).toHaveBeenCalledWith(
|
||||
'/api/admin/provider-oauth/providers/provider-codex/agent-identity-import/tasks',
|
||||
{ credentials, proxy_node_id: undefined },
|
||||
)
|
||||
})
|
||||
|
||||
it('keeps ordinary batch credentials on the pool permission surface', async () => {
|
||||
await startBatchImportOAuthTask('provider-codex', 'refresh-token')
|
||||
|
||||
expect(postMock).toHaveBeenCalledWith(
|
||||
'/api/admin/provider-oauth/providers/provider-codex/batch-import/tasks',
|
||||
{ credentials: 'refresh-token', proxy_node_id: undefined },
|
||||
)
|
||||
})
|
||||
|
||||
it('polls Agent Identity tasks through the matching dedicated status route', async () => {
|
||||
await getBatchImportOAuthTaskStatus(
|
||||
'provider-codex',
|
||||
'agent-identity-task-1',
|
||||
)
|
||||
|
||||
expect(getMock).toHaveBeenCalledWith(
|
||||
'/api/admin/provider-oauth/providers/provider-codex/agent-identity-import/tasks/agent-identity-task-1',
|
||||
)
|
||||
})
|
||||
|
||||
it('keeps create requests on the single registration endpoint', async () => {
|
||||
await importProviderRefreshToken('provider-codex', {
|
||||
access_token: 'access-token',
|
||||
create_agent_identity: true,
|
||||
})
|
||||
|
||||
expect(postMock).toHaveBeenCalledWith(
|
||||
'/api/admin/provider-oauth/providers/provider-codex/import-refresh-token',
|
||||
{
|
||||
access_token: 'access-token',
|
||||
create_agent_identity: true,
|
||||
},
|
||||
)
|
||||
})
|
||||
})
|
||||
@@ -31,6 +31,9 @@ export interface ProviderOAuthCompleteResponseWithKey {
|
||||
temporary?: boolean
|
||||
email?: string | null
|
||||
replaced?: boolean
|
||||
task_ready?: boolean
|
||||
recoverable?: boolean
|
||||
detail?: string
|
||||
}
|
||||
|
||||
export interface OAuthBatchImportResultItem {
|
||||
@@ -241,7 +244,7 @@ export async function importProviderRefreshToken(
|
||||
refresh_token?: string
|
||||
access_token?: string
|
||||
session_token?: string
|
||||
create_agent_identity_from_session_token?: boolean
|
||||
create_agent_identity?: boolean
|
||||
password?: string
|
||||
expires_at?: number
|
||||
name?: string
|
||||
@@ -270,7 +273,10 @@ export async function startBatchImportOAuthTask(
|
||||
credentials: string,
|
||||
proxyNodeId?: string
|
||||
): Promise<OAuthBatchImportTaskStartResponse> {
|
||||
const resp = await client.post(`/api/admin/provider-oauth/providers/${providerId}/batch-import/tasks`, {
|
||||
const route = containsAgentIdentityImport(credentials)
|
||||
? 'agent-identity-import/tasks'
|
||||
: 'batch-import/tasks'
|
||||
const resp = await client.post(`/api/admin/provider-oauth/providers/${providerId}/${route}`, {
|
||||
credentials,
|
||||
proxy_node_id: proxyNodeId || undefined,
|
||||
})
|
||||
@@ -281,10 +287,41 @@ export async function getBatchImportOAuthTaskStatus(
|
||||
providerId: string,
|
||||
taskId: string
|
||||
): Promise<OAuthBatchImportTaskStatusResponse> {
|
||||
const resp = await client.get(`/api/admin/provider-oauth/providers/${providerId}/batch-import/tasks/${taskId}`)
|
||||
const route = taskId.startsWith('agent-identity-')
|
||||
? 'agent-identity-import/tasks'
|
||||
: 'batch-import/tasks'
|
||||
const resp = await client.get(`/api/admin/provider-oauth/providers/${providerId}/${route}/${taskId}`)
|
||||
return resp.data
|
||||
}
|
||||
|
||||
function containsAgentIdentityImport(credentials: string): boolean {
|
||||
let parsed: unknown
|
||||
try {
|
||||
parsed = JSON.parse(credentials)
|
||||
} catch {
|
||||
return false
|
||||
}
|
||||
return jsonValueContainsAgentIdentity(parsed)
|
||||
}
|
||||
|
||||
function jsonValueContainsAgentIdentity(value: unknown): boolean {
|
||||
if (Array.isArray(value)) return value.some(jsonValueContainsAgentIdentity)
|
||||
if (typeof value !== 'object' || value === null) return false
|
||||
const root = value as Record<string, unknown>
|
||||
const nestedValue = root.agent_identity ?? root.agentIdentity
|
||||
const nested = typeof nestedValue === 'object' && nestedValue !== null && !Array.isArray(nestedValue)
|
||||
? nestedValue as Record<string, unknown>
|
||||
: undefined
|
||||
const authMode = root.auth_mode ?? root.authMode ?? nested?.auth_mode ?? nested?.authMode
|
||||
if (typeof authMode === 'string' && authMode.trim().toLowerCase() === 'agentidentity') return true
|
||||
const runtimeId = nested?.agent_runtime_id ?? nested?.agentRuntimeId ?? root.agent_runtime_id ?? root.agentRuntimeId
|
||||
const privateKey = nested?.agent_private_key ?? nested?.agentPrivateKey ?? root.agent_private_key ?? root.agentPrivateKey
|
||||
if (typeof runtimeId === 'string' && runtimeId.trim().length > 0
|
||||
&& typeof privateKey === 'string' && privateKey.trim().length > 0
|
||||
) return true
|
||||
return Object.values(root).some(jsonValueContainsAgentIdentity)
|
||||
}
|
||||
|
||||
// Device Authorization (AWS SSO OIDC)
|
||||
|
||||
export interface DeviceAuthorizeRequest {
|
||||
|
||||
@@ -725,19 +725,11 @@
|
||||
class="flex flex-col gap-3 justify-center transition-opacity duration-150"
|
||||
:class="mode === 'agent_identity' ? 'opacity-100' : 'opacity-0 pointer-events-none'"
|
||||
>
|
||||
<div class="space-y-1">
|
||||
<label class="text-xs font-medium">
|
||||
{{ legacyT('ChatGPT Session Token') }}
|
||||
</label>
|
||||
<p class="text-[11px] text-muted-foreground">
|
||||
{{ legacyT('仅用于一次性注册,成功后不会保存 Token。') }}
|
||||
</p>
|
||||
</div>
|
||||
<Textarea
|
||||
v-model="agentIdentitySessionToken"
|
||||
v-model="agentIdentityInput"
|
||||
:disabled="creatingAgentIdentity"
|
||||
:placeholder="legacyT('粘贴 ChatGPT Session Token(JWT)')"
|
||||
class="min-h-[230px] text-xs font-mono break-all !rounded-xl"
|
||||
:placeholder="legacyT('粘贴 AT 或 ChatGPT auth/session JSON')"
|
||||
class="min-h-[200px] text-xs font-mono break-all !rounded-xl"
|
||||
autocomplete="off"
|
||||
spellcheck="false"
|
||||
/>
|
||||
@@ -778,7 +770,7 @@
|
||||
:disabled="!canCreateAgentIdentity"
|
||||
@click="handleCreateAgentIdentity"
|
||||
>
|
||||
{{ creatingAgentIdentity ? legacyT('创建中...') : legacyT('创建并导入 Agent Identity') }}
|
||||
{{ creatingAgentIdentity ? legacyT('创建中...') : legacyT('创建') }}
|
||||
</Button>
|
||||
</template>
|
||||
</Dialog>
|
||||
@@ -833,7 +825,7 @@ const emit = defineEmits<{
|
||||
saved: []
|
||||
}>()
|
||||
|
||||
const { success, error: showError } = useToast()
|
||||
const { success, warning, error: showError } = useToast()
|
||||
const { copyToClipboard } = useClipboard()
|
||||
const { legacyT, locale } = useI18n()
|
||||
const proxyNodesStore = useProxyNodesStore()
|
||||
@@ -987,7 +979,7 @@ const windsurfImportMethod = ref<WindsurfImportMethod>('email_password')
|
||||
const windsurfEmail = ref('')
|
||||
const windsurfPassword = ref('')
|
||||
const windsurfAccountName = ref('')
|
||||
const agentIdentitySessionToken = ref('')
|
||||
const agentIdentityInput = ref('')
|
||||
const creatingAgentIdentity = ref(false)
|
||||
let agentIdentityRequestId = 0
|
||||
|
||||
@@ -1077,7 +1069,7 @@ const canImport = computed(() => {
|
||||
|
||||
const canCreateAgentIdentity = computed(() =>
|
||||
isCodexProvider.value
|
||||
&& agentIdentitySessionToken.value.trim().length > 0
|
||||
&& agentIdentityInput.value.trim().length > 0
|
||||
&& !creatingAgentIdentity.value
|
||||
)
|
||||
|
||||
@@ -1383,7 +1375,7 @@ function resetForm() {
|
||||
windsurfEmail.value = ''
|
||||
windsurfPassword.value = ''
|
||||
windsurfAccountName.value = ''
|
||||
agentIdentitySessionToken.value = ''
|
||||
agentIdentityInput.value = ''
|
||||
creatingAgentIdentity.value = false
|
||||
proxyPopoverOpen.value = false
|
||||
selectedProxyNodeId.value = ''
|
||||
@@ -1760,6 +1752,33 @@ function isObjectRecord(value: unknown): value is Record<string, unknown> {
|
||||
return typeof value === 'object' && value !== null && !Array.isArray(value)
|
||||
}
|
||||
|
||||
type AgentIdentityAccessTokenResolution =
|
||||
| { ok: true, accessToken: string }
|
||||
| { ok: false, message: string }
|
||||
|
||||
function resolveAgentIdentityAccessToken(input: string): AgentIdentityAccessTokenResolution {
|
||||
const normalized = input.trim()
|
||||
if (!normalized.startsWith('{') && !normalized.startsWith('[')) {
|
||||
return { ok: true, accessToken: normalized }
|
||||
}
|
||||
|
||||
let parsed: unknown
|
||||
try {
|
||||
parsed = JSON.parse(normalized)
|
||||
} catch {
|
||||
return { ok: false, message: 'ChatGPT auth/session JSON 格式无效' }
|
||||
}
|
||||
if (!isObjectRecord(parsed)) {
|
||||
return { ok: false, message: 'ChatGPT auth/session JSON 缺少 accessToken' }
|
||||
}
|
||||
|
||||
const accessToken = normalizeStringField(parsed.accessToken)
|
||||
?? normalizeStringField(parsed.access_token)
|
||||
return accessToken
|
||||
? { ok: true, accessToken }
|
||||
: { ok: false, message: 'ChatGPT auth/session JSON 缺少 accessToken' }
|
||||
}
|
||||
|
||||
function isCodexAgentIdentityObject(root: Record<string, unknown>): boolean {
|
||||
const nestedValue = root.agent_identity ?? root.agentIdentity
|
||||
const nested = isObjectRecord(nestedValue) ? nestedValue : null
|
||||
@@ -1908,18 +1927,26 @@ async function handleImport() {
|
||||
async function handleCreateAgentIdentity() {
|
||||
if (!canCreateAgentIdentity.value || !props.providerId) return
|
||||
|
||||
const sessionToken = agentIdentitySessionToken.value.trim()
|
||||
const resolvedInput = resolveAgentIdentityAccessToken(agentIdentityInput.value)
|
||||
if (!resolvedInput.ok) {
|
||||
showError(legacyT(resolvedInput.message), legacyT('格式错误'))
|
||||
return
|
||||
}
|
||||
const requestId = ++agentIdentityRequestId
|
||||
creatingAgentIdentity.value = true
|
||||
try {
|
||||
const result = await importProviderRefreshToken(props.providerId, {
|
||||
session_token: sessionToken,
|
||||
create_agent_identity_from_session_token: true,
|
||||
access_token: resolvedInput.accessToken,
|
||||
create_agent_identity: true,
|
||||
proxy_node_id: selectedProxyNodeId.value || undefined,
|
||||
})
|
||||
if (requestId !== agentIdentityRequestId) return
|
||||
|
||||
success(getOAuthSuccessMessage('创建', result))
|
||||
if (result.task_ready === false) {
|
||||
warning(legacyT('Agent Identity 已保存,任务将在后台初始化'), legacyT('已保存'))
|
||||
} else {
|
||||
success(getOAuthSuccessMessage('创建', result))
|
||||
}
|
||||
emit('saved')
|
||||
handleClose()
|
||||
} catch (err: unknown) {
|
||||
|
||||
@@ -58,7 +58,7 @@
|
||||
<Download class="w-2.5 h-2.5" />
|
||||
</Button>
|
||||
<Button
|
||||
v-else
|
||||
v-else-if="apiKey.agent_identity !== true"
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
class="h-4 w-4 shrink-0"
|
||||
|
||||
+105
-11
@@ -13,6 +13,12 @@ const endpointMocks = vi.hoisted(() => ({
|
||||
getAwsRegions: vi.fn(),
|
||||
}))
|
||||
|
||||
const toastMocks = vi.hoisted(() => ({
|
||||
success: vi.fn(),
|
||||
warning: vi.fn(),
|
||||
error: vi.fn(),
|
||||
}))
|
||||
|
||||
vi.mock('@/api/endpoints', async () => {
|
||||
const actual = await vi.importActual<typeof import('@/api/endpoints/provider_oauth')>(
|
||||
'@/api/endpoints/provider_oauth',
|
||||
@@ -212,10 +218,7 @@ vi.mock('@/stores/proxy-nodes', () => ({
|
||||
}))
|
||||
|
||||
vi.mock('@/composables/useToast', () => ({
|
||||
useToast: () => ({
|
||||
success: vi.fn(),
|
||||
error: vi.fn(),
|
||||
}),
|
||||
useToast: () => toastMocks,
|
||||
}))
|
||||
|
||||
vi.mock('@/composables/useClipboard', () => ({
|
||||
@@ -302,6 +305,9 @@ describe('OAuthAccountDialog Grok import', () => {
|
||||
endpointMocks.startDeviceAuthorize.mockReset()
|
||||
endpointMocks.pollDeviceAuthorize.mockReset()
|
||||
endpointMocks.getAwsRegions.mockReset()
|
||||
toastMocks.success.mockReset()
|
||||
toastMocks.warning.mockReset()
|
||||
toastMocks.error.mockReset()
|
||||
|
||||
endpointMocks.importProviderRefreshToken.mockResolvedValue({
|
||||
provider_type: 'grok',
|
||||
@@ -438,7 +444,7 @@ describe('OAuthAccountDialog Grok import', () => {
|
||||
expect(endpointMocks.importProviderRefreshToken).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('shows a dedicated Codex Agent Identity mode and creates from a Session Token', async () => {
|
||||
it('shows a dedicated Codex Agent Identity mode and creates from an access token', async () => {
|
||||
const root = mountDialog('codex')
|
||||
await settle()
|
||||
|
||||
@@ -450,25 +456,113 @@ describe('OAuthAccountDialog Grok import', () => {
|
||||
await settle()
|
||||
|
||||
const textarea = root.querySelector<HTMLTextAreaElement>(
|
||||
'textarea[placeholder="粘贴 ChatGPT Session Token(JWT)"]',
|
||||
'textarea[placeholder="粘贴 AT 或 ChatGPT auth/session JSON"]',
|
||||
)
|
||||
expect(textarea).toBeTruthy()
|
||||
if (!textarea) throw new Error('Expected Agent Identity Session Token textarea to exist')
|
||||
textarea.value = 'session-token-for-test-only'
|
||||
if (!textarea) throw new Error('Expected Agent Identity access token textarea to exist')
|
||||
expect(textarea.classList.contains('min-h-[200px]')).toBe(true)
|
||||
expect(root.textContent).not.toContain('ChatGPT Session Token')
|
||||
expect(root.textContent).not.toContain('仅用于一次性注册')
|
||||
expect(root.textContent).not.toContain('创建并导入 Agent Identity')
|
||||
textarea.value = 'access-token-for-test-only'
|
||||
textarea.dispatchEvent(new Event('input'))
|
||||
await settle()
|
||||
|
||||
getExactButton(root, '创建并导入 Agent Identity')?.click()
|
||||
getExactButton(root, '创建')?.click()
|
||||
await settle()
|
||||
|
||||
expect(endpointMocks.importProviderRefreshToken).toHaveBeenCalledWith('provider-1', {
|
||||
session_token: 'session-token-for-test-only',
|
||||
create_agent_identity_from_session_token: true,
|
||||
access_token: 'access-token-for-test-only',
|
||||
create_agent_identity: true,
|
||||
proxy_node_id: undefined,
|
||||
})
|
||||
expect(endpointMocks.startBatchImportOAuthTask).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('extracts only accessToken from ChatGPT auth/session JSON for Agent Identity', async () => {
|
||||
const root = mountDialog('codex')
|
||||
await settle()
|
||||
|
||||
getExactButton(root, 'Agent Identity')?.click()
|
||||
await settle()
|
||||
|
||||
const textarea = root.querySelector<HTMLTextAreaElement>(
|
||||
'textarea[placeholder="粘贴 AT 或 ChatGPT auth/session JSON"]',
|
||||
)
|
||||
if (!textarea) throw new Error('Expected Agent Identity access token textarea to exist')
|
||||
textarea.value = JSON.stringify({
|
||||
WARNING_BANNER: 'sensitive session data',
|
||||
accessToken: 'access-token-from-json',
|
||||
sessionToken: 'session-token-must-not-be-used',
|
||||
user: { email: 'private@example.com' },
|
||||
})
|
||||
textarea.dispatchEvent(new Event('input'))
|
||||
await settle()
|
||||
|
||||
getExactButton(root, '创建')?.click()
|
||||
await settle()
|
||||
|
||||
expect(endpointMocks.importProviderRefreshToken).toHaveBeenCalledWith('provider-1', {
|
||||
access_token: 'access-token-from-json',
|
||||
create_agent_identity: true,
|
||||
proxy_node_id: undefined,
|
||||
})
|
||||
expect(endpointMocks.importProviderRefreshToken).not.toHaveBeenCalledWith(
|
||||
'provider-1',
|
||||
expect.objectContaining({ session_token: 'session-token-must-not-be-used' }),
|
||||
)
|
||||
})
|
||||
|
||||
it('does not use sessionToken when ChatGPT auth/session JSON has no accessToken', async () => {
|
||||
const root = mountDialog('codex')
|
||||
await settle()
|
||||
|
||||
getExactButton(root, 'Agent Identity')?.click()
|
||||
await settle()
|
||||
|
||||
const textarea = root.querySelector<HTMLTextAreaElement>(
|
||||
'textarea[placeholder="粘贴 AT 或 ChatGPT auth/session JSON"]',
|
||||
)
|
||||
if (!textarea) throw new Error('Expected Agent Identity access token textarea to exist')
|
||||
textarea.value = JSON.stringify({ sessionToken: 'session-token-must-not-be-used' })
|
||||
textarea.dispatchEvent(new Event('input'))
|
||||
await settle()
|
||||
|
||||
getExactButton(root, '创建')?.click()
|
||||
await settle()
|
||||
|
||||
expect(endpointMocks.importProviderRefreshToken).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('treats a saved Agent Identity with pending task initialization as accepted', async () => {
|
||||
endpointMocks.importProviderRefreshToken.mockResolvedValueOnce({
|
||||
key_id: 'key-agent',
|
||||
provider_type: 'codex',
|
||||
has_refresh_token: false,
|
||||
task_ready: false,
|
||||
recoverable: true,
|
||||
detail: 'pending task',
|
||||
})
|
||||
const root = mountDialog('codex')
|
||||
await settle()
|
||||
|
||||
getExactButton(root, 'Agent Identity')?.click()
|
||||
await settle()
|
||||
const textarea = root.querySelector<HTMLTextAreaElement>(
|
||||
'textarea[placeholder="粘贴 AT 或 ChatGPT auth/session JSON"]',
|
||||
)
|
||||
if (!textarea) throw new Error('Expected Agent Identity access token textarea to exist')
|
||||
textarea.value = 'access-token-for-pending-task'
|
||||
textarea.dispatchEvent(new Event('input'))
|
||||
await settle()
|
||||
|
||||
getExactButton(root, '创建')?.click()
|
||||
await settle()
|
||||
|
||||
expect(toastMocks.warning).toHaveBeenCalled()
|
||||
expect(toastMocks.error).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('keeps Agent Identity creation unavailable for non-Codex providers', async () => {
|
||||
const root = mountDialog('openai')
|
||||
await settle()
|
||||
|
||||
@@ -140,4 +140,18 @@ describe('ProviderKeyIdentityBlock', () => {
|
||||
|
||||
unmount()
|
||||
})
|
||||
|
||||
it('does not offer generic copy or export actions for Agent Identity', () => {
|
||||
const { root, unmount } = mount({
|
||||
apiKey: createProviderKey({ agent_identity: true }),
|
||||
maskedSecretLabel: '[Agent Identity]',
|
||||
canExportCredential: false,
|
||||
})
|
||||
|
||||
expect(root.textContent).toContain('[Agent Identity]')
|
||||
expect(root.querySelector('button[title="下载 OAuth 授权文件"]')).toBeNull()
|
||||
expect(root.querySelector('button[title="复制密钥"]')).toBeNull()
|
||||
|
||||
unmount()
|
||||
})
|
||||
})
|
||||
|
||||
@@ -1583,11 +1583,16 @@ const legacyExactEnglishMessages: Record<string, string> = {
|
||||
'设备授权': 'Device authorization',
|
||||
'导入授权': 'Import authorization',
|
||||
'Agent Identity': 'Agent Identity',
|
||||
'创建并导入 Agent Identity': 'Create and import Agent Identity',
|
||||
'创建 Agent Identity 失败': 'Failed to create Agent Identity',
|
||||
'ChatGPT Session Token': 'ChatGPT Session Token',
|
||||
'粘贴 ChatGPT Session Token(JWT)': 'Paste a ChatGPT Session Token (JWT)',
|
||||
'仅用于一次性注册,成功后不会保存 Token。': 'Used only for one-time registration. The token is not stored after success.',
|
||||
'Agent Identity 已保存,任务将在后台初始化': 'Agent Identity was saved; its task will initialize in the background',
|
||||
'已保存': 'Saved',
|
||||
'粘贴 AT 或 ChatGPT auth/session JSON': 'Paste an access token or ChatGPT auth/session JSON',
|
||||
'ChatGPT auth/session JSON 格式无效': 'Invalid ChatGPT auth/session JSON',
|
||||
'ChatGPT auth/session JSON 缺少 accessToken': 'ChatGPT auth/session JSON is missing accessToken',
|
||||
'ChatGPT Access Token 缺少账号身份字段': 'ChatGPT access token is missing account identity claims',
|
||||
'ChatGPT Access Token(JWT)不能为空': 'ChatGPT access token (JWT) is required',
|
||||
'仅 Codex Provider 支持使用 Access Token 创建 Agent Identity': 'Only Codex providers support creating Agent Identity with an access token',
|
||||
'使用 Access Token 创建 Agent Identity 时不能同时提交 Refresh Token': 'A refresh token cannot be submitted when creating Agent Identity with an access token',
|
||||
'导入账号': 'Import accounts',
|
||||
'授权已过期': 'Authorization expired',
|
||||
'授权失败': 'Authorization failed',
|
||||
|
||||
@@ -261,7 +261,7 @@
|
||||
<Download class="w-2.5 h-2.5" />
|
||||
</Button>
|
||||
<Button
|
||||
v-else
|
||||
v-else-if="key.agent_identity !== true"
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
class="h-4 w-4 shrink-0"
|
||||
@@ -725,7 +725,7 @@
|
||||
<Download class="w-3.5 h-3.5" />
|
||||
</Button>
|
||||
<Button
|
||||
v-else-if="actionId === 'copy_or_download'"
|
||||
v-else-if="actionId === 'copy_or_download' && key.agent_identity !== true"
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
class="h-7 w-7 shrink-0"
|
||||
|
||||
Reference in New Issue
Block a user