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

fix(kiro): 隐藏 OAuth 刷新重试流程 (#339)

Browse Source
This commit is contained in:
Entropy.Xu
2026-04-25 21:29:45 +08:00
committed by GitHub
parent 429fdb47e6
commit d784c540b6
10 changed files with 317 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
apps/aether-gateway/src/execution_runtime/mod.rs
View File

@@ -6,6 +6,7 @@ use serde_json::{Map, Value};
mod constants;
mod fallback;
pub(crate) mod ndjson;
mod oauth_retry;
#[cfg(test)]
pub(crate) mod remote_compat;
mod server;

104
apps/aether-gateway/src/execution_runtime/oauth_retry.rs Normal file
View File

@@ -0,0 +1,104 @@
use aether_contracts::ExecutionPlan;
use tracing::warn;
use crate::AppState;
pub(crate) async fn refresh_oauth_plan_auth_for_retry(
state: &AppState,
plan: &mut ExecutionPlan,
status_code: u16,
response_text: Option<&str>,
trace_id: &str,
) -> bool {
if !status_may_be_oauth_invalid(status_code, response_text) {
return false;
}
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 false,
Err(err) => {
warn!(
event_name = "local_oauth_retry_transport_read_failed",
log_type = "ops",
trace_id = %trace_id,
provider_id = %plan.provider_id,
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
error = ?err,
"gateway failed to read transport before oauth retry refresh"
);
return false;
}
};
if transport.key.decrypted_auth_config.is_none()
&& !transport.key.auth_type.trim().eq_ignore_ascii_case("oauth")
{
return false;
}
match state.force_local_oauth_refresh_entry(&transport).await {
Ok(Some(entry)) => {
let header_name = entry.auth_header_name.trim().to_ascii_lowercase();
let header_value = entry.auth_header_value.trim();
if header_name.is_empty() || header_value.is_empty() {
return false;
}
plan.headers.insert(header_name, header_value.to_string());
true
}
Ok(None) => false,
Err(err) => {
warn!(
event_name = "local_oauth_retry_refresh_failed",
log_type = "ops",
trace_id = %trace_id,
provider_id = %plan.provider_id,
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
status_code,
error = %err,
"gateway oauth retry refresh failed"
);
false
}
}
}
fn status_may_be_oauth_invalid(status_code: u16, response_text: Option<&str>) -> bool {
if status_code == 401 {
return true;
}
if status_code != 403 {
return false;
}
let Some(response_text) = response_text else {
return true;
};
let response_text = response_text.to_ascii_lowercase();
["oauth", "token", "auth", "credential", "expired"]
.iter()
.any(|needle| response_text.contains(needle))
}
#[cfg(test)]
mod tests {
use super::status_may_be_oauth_invalid;
#[test]
fn recognizes_oauth_invalid_statuses() {
assert!(status_may_be_oauth_invalid(401, None));
assert!(status_may_be_oauth_invalid(
403,
Some("The security token included in the request is expired")
));
assert!(status_may_be_oauth_invalid(403, None));
assert!(!status_may_be_oauth_invalid(403, Some("quota exceeded")));
assert!(!status_may_be_oauth_invalid(429, Some("token bucket")));
}
}

58
apps/aether-gateway/src/execution_runtime/stream/execution.rs
View File

@@ -51,6 +51,7 @@ use crate::clock::current_unix_ms as current_request_candidate_unix_ms;
use crate::constants::{CONTROL_CANDIDATE_ID_HEADER, CONTROL_REQUEST_ID_HEADER};
use crate::control::GatewayControlDecision;
use crate::execution_runtime::build_direct_execution_frame_stream;
use crate::execution_runtime::oauth_retry::refresh_oauth_plan_auth_for_retry;
#[cfg(test)]
use crate::execution_runtime::remote_compat::post_stream_plan_to_remote_execution_runtime;
use crate::execution_runtime::submission::{
@@ -312,6 +313,24 @@ async fn execute_in_process_stream(
DirectSyncExecutionRuntime::new().execute_stream(plan).await
}
async fn execute_in_process_stream_with_oauth_retry(
state: &AppState,
plan: &mut ExecutionPlan,
trace_id: &str,
report_context: Option<&Value>,
) -> Result<DirectUpstreamStreamExecution, ExecutionRuntimeTransportError> {
let mut execution = execute_in_process_stream(state, plan).await?;
apply_stream_summary_report_context(&mut execution, report_context);
if execution.status_code >= 400
&& refresh_oauth_plan_auth_for_retry(state, plan, execution.status_code, None, trace_id)
.await
{
execution = execute_in_process_stream(state, plan).await?;
apply_stream_summary_report_context(&mut execution, report_context);
}
Ok(execution)
}
#[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection
pub(crate) async fn execute_execution_runtime_stream(
state: &AppState,
@@ -351,21 +370,28 @@ pub(crate) async fn execute_execution_runtime_stream(
});
}
let plan_request_id_for_log = short_request_id(plan.request_id.as_str());
let provider_name = plan.provider_name.as_deref().unwrap_or("-");
let endpoint_id = plan.endpoint_id.as_str();
let key_id = plan.key_id.as_str();
let model_name = plan.model_name.as_deref().unwrap_or("-");
let provider_name = plan
.provider_name
.clone()
.unwrap_or_else(|| "-".to_string());
let endpoint_id = plan.endpoint_id.clone();
let key_id = plan.key_id.clone();
let model_name = plan.model_name.clone().unwrap_or_else(|| "-".to_string());
let candidate_index = parse_request_candidate_report_context(report_context.as_ref())
.and_then(|context| context.candidate_index)
.map(|value| value.to_string())
.unwrap_or_else(|| "-".to_string());
#[cfg(not(test))]
{
let execution = match execute_in_process_stream(state, &plan).await {
Ok(mut execution) => {
apply_stream_summary_report_context(&mut execution, report_context.as_ref());
execution
}
let execution = match execute_in_process_stream_with_oauth_retry(
state,
&mut plan,
trace_id,
report_context.as_ref(),
)
.await
{
Ok(execution) => execution,
Err(err) => {
info!(
event_name = "stream_execution_runtime_unavailable",
@@ -421,11 +447,15 @@ pub(crate) async fn execute_execution_runtime_stream(
.execution_runtime_override_base_url()
.unwrap_or_default();
if remote_execution_runtime_base_url.trim().is_empty() {
let execution = match execute_in_process_stream(state, &plan).await {
Ok(mut execution) => {
apply_stream_summary_report_context(&mut execution, report_context.as_ref());
execution
}
let execution = match execute_in_process_stream_with_oauth_retry(
state,
&mut plan,
trace_id,
report_context.as_ref(),
)
.await
{
Ok(execution) => execution,
Err(err) => {
info!(
event_name = "stream_execution_runtime_unavailable",

150
apps/aether-gateway/src/execution_runtime/sync/execution.rs
View File

@@ -24,6 +24,7 @@ use crate::api::response::{
};
use crate::clock::current_unix_ms as current_request_candidate_unix_ms;
use crate::control::GatewayControlDecision;
use crate::execution_runtime::oauth_retry::refresh_oauth_plan_auth_for_retry;
#[cfg(test)]
use crate::execution_runtime::remote_compat::post_sync_plan_to_remote_execution_runtime;
use crate::execution_runtime::submission::submit_local_core_error_or_sync_finalize;
@@ -155,13 +156,16 @@ pub(crate) async fn execute_execution_runtime_sync(
mut report_context: Option<serde_json::Value>,
) -> Result<Option<Response<Body>>, GatewayError> {
ensure_execution_request_candidate_slot(state, &mut plan, &mut report_context).await;
let plan_request_id = plan.request_id.as_str();
let plan_request_id_for_log = short_request_id(plan_request_id);
let plan_candidate_id = plan.candidate_id.as_deref();
let provider_name = plan.provider_name.as_deref().unwrap_or("-");
let endpoint_id = plan.endpoint_id.as_str();
let key_id = plan.key_id.as_str();
let model_name = plan.model_name.as_deref().unwrap_or("-");
let plan_request_id = plan.request_id.clone();
let plan_request_id_for_log = short_request_id(plan_request_id.as_str());
let plan_candidate_id = plan.candidate_id.clone();
let provider_name = plan
.provider_name
.clone()
.unwrap_or_else(|| "-".to_string());
let endpoint_id = plan.endpoint_id.clone();
let key_id = plan.key_id.clone();
let model_name = plan.model_name.clone().unwrap_or_else(|| "-".to_string());
let candidate_index = parse_request_candidate_report_context(report_context.as_ref())
.and_then(|context| context.candidate_index)
.map(|value| value.to_string())
@@ -316,8 +320,8 @@ pub(crate) async fn execute_execution_runtime_sync(
trace_id,
decision,
&plan,
plan_request_id,
plan_candidate_id,
plan_request_id.as_str(),
plan_candidate_id.as_deref(),
report_context.as_ref(),
candidate_started_unix_secs,
)
@@ -329,33 +333,100 @@ pub(crate) async fn execute_execution_runtime_sync(
}
}
};
let result_body_json = result
.body
.as_ref()
.and_then(|body| body.json_body.as_ref());
let (result_error_type, result_error_message) =
execution_error_details(result.error.as_ref(), result_body_json);
let result_latency_ms = result
.telemetry
.as_ref()
.and_then(|telemetry| telemetry.elapsed_ms);
let mut headers = std::mem::take(&mut result.headers);
let (body_bytes, body_json, body_base64) =
decode_execution_result_body(result.body.take(), &mut headers)?;
let local_failover_response_text = local_failover_response_text(
body_json.as_ref(),
&body_bytes,
result.error.as_ref().map(|error| error.message.as_str()),
);
let local_failover_analysis = analyze_local_candidate_failover_sync(
state,
&plan,
plan_kind,
report_context.as_ref(),
&result,
local_failover_response_text.as_deref(),
)
.await;
let mut oauth_retry_attempted = false;
let (
result_error_type,
result_error_message,
result_latency_ms,
headers,
body_bytes,
body_json,
body_base64,
local_failover_response_text,
local_failover_analysis,
) = loop {
let result_body_json = result
.body
.as_ref()
.and_then(|body| body.json_body.as_ref());
let (result_error_type, result_error_message) =
execution_error_details(result.error.as_ref(), result_body_json);
let result_latency_ms = result
.telemetry
.as_ref()
.and_then(|telemetry| telemetry.elapsed_ms);
let mut headers = std::mem::take(&mut result.headers);
let (body_bytes, body_json, body_base64) =
decode_execution_result_body(result.body.take(), &mut headers)?;
let local_failover_response_text = local_failover_response_text(
body_json.as_ref(),
&body_bytes,
result.error.as_ref().map(|error| error.message.as_str()),
);
if result.status_code >= 400
&& !oauth_retry_attempted
&& refresh_oauth_plan_auth_for_retry(
state,
&mut plan,
result.status_code,
local_failover_response_text.as_deref(),
trace_id,
)
.await
{
oauth_retry_attempted = true;
match crate::execution_runtime::execute_execution_runtime_sync_plan(
state,
Some(trace_id),
&plan,
)
.await
{
Ok(retry_result) => {
result = retry_result;
continue;
}
Err(err) => {
warn!(
event_name = "local_sync_oauth_retry_execution_failed",
log_type = "ops",
trace_id = %trace_id,
request_id = %plan_request_id_for_log,
candidate_id = ?plan_candidate_id,
provider_name,
endpoint_id,
key_id,
model_name,
candidate_index = candidate_index.as_str(),
error = ?err,
"gateway oauth retry sync execution failed"
);
}
}
}
let local_failover_analysis = analyze_local_candidate_failover_sync(
state,
&plan,
plan_kind,
report_context.as_ref(),
&result,
local_failover_response_text.as_deref(),
)
.await;
break (
result_error_type,
result_error_message,
result_latency_ms,
headers,
body_bytes,
body_json,
body_base64,
local_failover_response_text,
local_failover_analysis,
);
};
if result.status_code >= 400 {
apply_local_execution_effect(
state,
@@ -565,11 +636,12 @@ pub(crate) async fn execute_execution_runtime_sync(
let candidate_id_owned = result.candidate_id;
let request_id = (!request_id_owned.trim().is_empty())
.then_some(request_id_owned.as_str())
.or(Some(plan_request_id));
.or(Some(plan_request_id.as_str()));
let request_id_for_log = short_request_id(request_id.unwrap_or("-"));
let candidate_id = candidate_id_owned.as_deref().or(plan_candidate_id);
let candidate_id = candidate_id_owned
.as_deref()
.or(plan_candidate_id.as_deref());
let report_context = report_context;
let headers = headers;
let body_json = body_json;
let telemetry = result.telemetry;

1
apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/direct.rs
View File

@@ -664,6 +664,7 @@ async fn gateway_executes_kiro_claude_cli_sync_upstream_stream_via_local_finaliz
let auth_config = serde_json::json!({
"provider_type": "kiro",
"access_token": "cached-kiro-finalize-access-token",
"expires_at": 4102444800_u64,
"refresh_token": "rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
"machine_id": "123e4567-e89b-12d3-a456-426614174000",
"api_region": "us-east-1",

1
apps/aether-gateway/src/tests/ai_execute/stream_provider.rs
View File

@@ -209,6 +209,7 @@ async fn gateway_executes_kiro_claude_cli_stream_via_local_provider_catalog_cand
let auth_config = serde_json::json!({
"provider_type": "kiro",
"access_token": "cached-kiro-access-token",
"expires_at": 4102444800_u64,
"refresh_token": "rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
"machine_id": "123e4567-e89b-12d3-a456-426614174000",
"api_region": "us-east-1",

39
apps/aether-gateway/src/tests/ai_execute/sync/claude/kiro.rs
View File

@@ -198,6 +198,7 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid
let auth_config = serde_json::json!({
"provider_type": "kiro",
"access_token": "cached-kiro-access-token",
"expires_at": 4102444800_u64,
"refresh_token": "rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
"machine_id": "123e4567-e89b-12d3-a456-426614174000",
"api_region": "us-east-1",
@@ -750,6 +751,8 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
let auth_config = serde_json::json!({
"provider_type": "kiro",
"access_token": "stale-kiro-access-token",
"expires_at": 4102444800_u64,
"refresh_token": "rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
"machine_id": "123e4567-e89b-12d3-a456-426614174000",
"api_region": "us-east-1",
@@ -791,6 +794,8 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let execution_hits = Arc::new(Mutex::new(0usize));
let execution_hits_clone = Arc::clone(&execution_hits);
let refresh_hits = Arc::new(Mutex::new(0usize));
let refresh_hits_clone = Arc::clone(&refresh_hits);
@@ -864,11 +869,19 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid
"/v1/execute/sync",
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
let execution_hits_inner = Arc::clone(&execution_hits_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("execution runtime payload should parse");
*execution_hits_inner.lock().expect("mutex should lock") += 1;
let authorization = payload
.get("headers")
.and_then(|value| value.get("authorization"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string();
*seen_execution_runtime_inner.lock().expect("mutex should lock") =
Some(SeenExecutionRuntimeSyncRequest {
trace_id: parts
@@ -877,12 +890,7 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
authorization: payload
.get("headers")
.and_then(|value| value.get("authorization"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
authorization: authorization.clone(),
mapped_model: payload
.get("body")
.and_then(|value| value.get("json_body"))
@@ -905,6 +913,24 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid
.to_string(),
});
if authorization == "Bearer stale-kiro-access-token" {
return Json(json!({
"request_id": "trace-kiro-cli-local-refresh-123",
"status_code": 401,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"message": "The security token included in the request is expired"
}
},
"telemetry": {
"elapsed_ms": 7
}
}));
}
let kiro_frames = [
encode_event_frame(
"event",
@@ -1001,6 +1027,7 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid
);
assert_eq!(*refresh_hits.lock().expect("mutex should lock"), 1);
assert_eq!(*execution_hits.lock().expect("mutex should lock"), 2);
let seen_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")

3
apps/aether-gateway/src/tests/control/admin/provider_query.rs
View File

@@ -1021,6 +1021,7 @@ async fn gateway_handles_admin_provider_query_test_model_for_kiro_locally() {
"provider_type":"kiro",
"auth_method":"idc",
"access_token":"cached-kiro-token",
"expires_at":4102444800,
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
"machine_id":"123e4567-e89b-12d3-a456-426614174000",
"api_region":"us-east-1",
@@ -1155,6 +1156,7 @@ async fn gateway_handles_admin_provider_query_test_model_failover_for_kiro_local
"provider_type":"kiro",
"auth_method":"idc",
"access_token":"cached-kiro-token",
"expires_at":4102444800,
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
"machine_id":"123e4567-e89b-12d3-a456-426614174000",
"api_region":"us-east-1",
@@ -1301,6 +1303,7 @@ async fn gateway_retries_kiro_failover_after_http_error_without_message() {
"provider_type":"kiro",
"auth_method":"idc",
"access_token":"cached-kiro-token",
"expires_at":4102444800,
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
"machine_id":"123e4567-e89b-12d3-a456-426614174000",
"api_region":"us-east-1",