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;