admin/CC-Switch
1
0
Fork 0
mirror of https://github.com/farion1231/cc-switch.git synced 2026-08-04 11:43:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(pi): localize UI and isolate credential health

Browse Source
This commit is contained in:
SaladDay
2026-08-03 08:07:19 +00:00
parent a8246da3d6
commit a2fa86bea1
4 changed files with 614 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src-tauri/src/proxy/pi_handler.rs
+242 -77
View File
@@ -111,8 +111,7 @@ pub(crate) async fn handle_pi_native(
continue;
}
let mut saw_health_failure = false;
let mut reached_network = false;
let mut provider_health_failure = None;
for candidate in attempts[index..provider_end].iter().cloned() {
if !network_budget.has_remaining() {
break;
@@ -155,16 +154,12 @@ pub(crate) async fn handle_pi_native(
}
if let Some(pending) = pending_retryable.take() {
if pending.provider_id != provider_id {
record_provider_result(
settle_provider_health(
&state,
route.catalog_epoch,
&pending.provider_id,
pending.used_half_open_permit,
false,
Some(format!(
"Pi upstream returned retryable status {}",
pending.response.status()
)),
pending.provider_health.clone(),
)
.await;
} else {
@@ -182,7 +177,6 @@ pub(crate) async fn handle_pi_native(
.headers(outgoing_headers)
.body(body.clone())
.send();
reached_network = true;
let response = match if timeout_seconds > 0 {
tokio::time::timeout(Duration::from_secs(u64::from(timeout_seconds)), send)
.await
@@ -192,25 +186,31 @@ pub(crate) async fn handle_pi_native(
} {
Ok(Ok(response)) => response,
Ok(Err(error)) => {
saw_health_failure = true;
last_error = Some(if error.is_timeout() {
let error = if error.is_timeout() {
"Pi upstream request timed out".to_string()
} else {
"Pi upstream request failed before response".to_string()
});
};
provider_health_failure = Some(error.clone());
last_error = Some(error);
continue;
}
Err(()) => {
saw_health_failure = true;
last_error = Some("Pi upstream response-header timeout".to_string());
let error = "Pi upstream response-header timeout".to_string();
provider_health_failure = Some(error.clone());
last_error = Some(error);
continue;
}
};
let status = response.status();
if retryable_status(status) && network_budget.has_remaining() {
saw_health_failure = true;
last_error = Some(format!("Pi upstream returned retryable status {status}"));
let error = format!("Pi upstream returned retryable status {status}");
let status_health = provider_health_disposition(status);
if status_health == ProviderHealthDisposition::Unhealthy {
provider_health_failure = Some(error.clone());
}
last_error = Some(error);
let selected_is_failover = materialized.is_failover;
pending_retryable = Some(PendingRetryableResponse {
response,
@@ -218,10 +218,16 @@ pub(crate) async fn handle_pi_native(
provider_id: provider_id.clone(),
used_half_open_permit: permit.used_half_open_permit,
selected_is_failover,
provider_health: ProviderHealthOutcome::from_status(
status,
provider_health_failure.as_deref(),
),
});
continue;
}
let selected_is_failover = materialized.is_failover;
let provider_health =
ProviderHealthOutcome::from_status(status, provider_health_failure.as_deref());
match prepare_response(
state.clone(),
response,
@@ -236,19 +242,18 @@ pub(crate) async fn handle_pi_native(
route.app_config.non_streaming_timeout,
permit.used_half_open_permit,
selected_is_failover,
provider_health.clone(),
)
.await
{
Ok(prepared) => {
if !prepared.finalization_deferred {
let health_success = !retryable_status(status);
record_provider_result(
settle_provider_health(
&state,
route.catalog_epoch,
&provider_id,
permit.used_half_open_permit,
health_success,
(!health_success).then(|| format!("Pi upstream returned {status}")),
provider_health,
)
.await;
record_request_finish(
@@ -263,7 +268,7 @@ pub(crate) async fn handle_pi_native(
return Ok(prepared.response);
}
Err(ProxyError::ForwardFailed(error)) | Err(ProxyError::Timeout(error)) => {
saw_health_failure = true;
provider_health_failure = Some(error.clone());
last_error = Some(error);
continue;
}
@@ -273,9 +278,7 @@ pub(crate) async fn handle_pi_native(
route.catalog_epoch,
&provider_id,
permit.used_half_open_permit,
reached_network,
saw_health_failure,
Some(error.to_string()),
provider_health_failure.clone(),
)
.await;
record_request_finish(&state, false, false, Some(error.to_string())).await;
@@ -293,9 +296,7 @@ pub(crate) async fn handle_pi_native(
route.catalog_epoch,
&provider_id,
permit.used_half_open_permit,
reached_network,
saw_health_failure,
last_error.clone(),
provider_health_failure,
)
.await;
}
@@ -307,6 +308,7 @@ pub(crate) async fn handle_pi_native(
let provider_id = pending.provider_id.clone();
let used_half_open_permit = pending.used_half_open_permit;
let selected_is_failover = pending.selected_is_failover;
let provider_health = pending.provider_health;
match prepare_response(
state.clone(),
pending.response,
@@ -321,18 +323,18 @@ pub(crate) async fn handle_pi_native(
route.app_config.non_streaming_timeout,
used_half_open_permit,
selected_is_failover,
provider_health.clone(),
)
.await
{
Ok(prepared) => {
if !prepared.finalization_deferred {
record_provider_result(
settle_provider_health(
&state,
route.catalog_epoch,
&provider_id,
used_half_open_permit,
false,
Some(format!("Pi upstream returned {status}")),
provider_health,
)
.await;
record_request_finish(
@@ -374,6 +376,7 @@ struct PendingRetryableResponse {
provider_id: String,
used_half_open_permit: bool,
selected_is_failover: bool,
provider_health: ProviderHealthOutcome,
}
#[derive(Debug)]
@@ -497,24 +500,114 @@ async fn record_provider_result(
}
}
#[allow(clippy::too_many_arguments)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ProviderHealthDisposition {
Healthy,
Unhealthy,
Neutral,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct ProviderHealthOutcome {
disposition: ProviderHealthDisposition,
error: Option<String>,
}
impl ProviderHealthOutcome {
fn from_status(status: StatusCode, prior_failure: Option<&str>) -> Self {
match provider_health_disposition(status) {
ProviderHealthDisposition::Healthy => Self {
disposition: ProviderHealthDisposition::Healthy,
error: None,
},
ProviderHealthDisposition::Unhealthy => Self {
disposition: ProviderHealthDisposition::Unhealthy,
error: Some(format!("Pi upstream returned {status}")),
},
ProviderHealthDisposition::Neutral => prior_failure.map_or_else(
|| Self {
disposition: ProviderHealthDisposition::Neutral,
error: None,
},
|error| Self {
// A credential rejection is neutral by itself, but it
// cannot erase a real failure from an earlier endpoint
// covered by the same provider-level permit.
disposition: ProviderHealthDisposition::Unhealthy,
error: Some(error.to_string()),
},
),
}
}
}
fn provider_health_disposition(status: StatusCode) -> ProviderHealthDisposition {
if matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN) {
// Authentication/authorization belongs to the candidate credential,
// not endpoint availability. Failover may still try another
// credential, but this provider's health and circuit stay unchanged.
ProviderHealthDisposition::Neutral
} else if retryable_status(status) {
ProviderHealthDisposition::Unhealthy
} else {
ProviderHealthDisposition::Healthy
}
}
async fn settle_provider_health(
state: &ProxyState,
catalog_epoch: u64,
provider_id: &str,
used_half_open_permit: bool,
outcome: ProviderHealthOutcome,
) {
match outcome.disposition {
ProviderHealthDisposition::Healthy => {
record_provider_result(
state,
catalog_epoch,
provider_id,
used_half_open_permit,
true,
None,
)
.await;
}
ProviderHealthDisposition::Unhealthy => {
record_provider_result(
state,
catalog_epoch,
provider_id,
used_half_open_permit,
false,
outcome.error,
)
.await;
}
ProviderHealthDisposition::Neutral => {
state
.provider_router
.release_permit_neutral(provider_id, "pi", used_half_open_permit)
.await;
}
}
}
async fn release_or_record_provider(
state: &ProxyState,
catalog_epoch: u64,
provider_id: &str,
used_half_open_permit: bool,
reached_network: bool,
saw_health_failure: bool,
error: Option<String>,
provider_health_failure: Option<String>,
) {
if reached_network && saw_health_failure {
if provider_health_failure.is_some() {
record_provider_result(
state,
catalog_epoch,
provider_id,
used_half_open_permit,
false,
error,
provider_health_failure,
)
.await;
} else {
@@ -714,6 +807,7 @@ struct PiStreamFinalization {
content_is_sse: bool,
used_half_open_permit: bool,
selected_is_failover: bool,
complete_provider_health: ProviderHealthOutcome,
}
enum PiStreamTermination {
@@ -723,7 +817,7 @@ enum PiStreamTermination {
}
struct PiStreamDisposition {
provider_success: Option<bool>,
provider_health: ProviderHealthDisposition,
provider_error: Option<String>,
request_success: bool,
request_error: Option<String>,
@@ -731,29 +825,25 @@ struct PiStreamDisposition {
fn pi_stream_disposition(
status: StatusCode,
complete_provider_health: &ProviderHealthOutcome,
termination: &PiStreamTermination,
) -> PiStreamDisposition {
match termination {
PiStreamTermination::Complete { .. } => {
let provider_success = !retryable_status(status);
PiStreamDisposition {
provider_success: Some(provider_success),
provider_error: (!provider_success)
.then(|| format!("Pi upstream returned {status}")),
request_success: status.is_success(),
request_error: (!status.is_success())
.then(|| format!("Pi upstream returned {status}")),
}
}
PiStreamTermination::Complete { .. } => PiStreamDisposition {
provider_health: complete_provider_health.disposition,
provider_error: complete_provider_health.error.clone(),
request_success: status.is_success(),
request_error: (!status.is_success()).then(|| format!("Pi upstream returned {status}")),
},
PiStreamTermination::UpstreamFailure { message } => PiStreamDisposition {
provider_success: Some(false),
provider_health: ProviderHealthDisposition::Unhealthy,
provider_error: Some(message.clone()),
request_success: false,
request_error: Some(message.clone()),
},
PiStreamTermination::DownstreamDropped => PiStreamDisposition {
// A downstream cancellation says nothing about upstream health.
provider_success: None,
provider_health: ProviderHealthDisposition::Neutral,
provider_error: None,
request_success: false,
request_error: Some(
@@ -801,27 +891,19 @@ impl PiStreamFinalization {
}
async fn apply(self, termination: PiStreamTermination) {
let disposition = pi_stream_disposition(self.status, &termination);
if let Some(provider_success) = disposition.provider_success {
record_provider_result(
&self.state,
self.catalog_epoch,
&self.candidate.provider_id,
self.used_half_open_permit,
provider_success,
disposition.provider_error,
)
.await;
} else {
self.state
.provider_router
.release_permit_neutral(
&self.candidate.provider_id,
"pi",
self.used_half_open_permit,
)
.await;
}
let disposition =
pi_stream_disposition(self.status, &self.complete_provider_health, &termination);
settle_provider_health(
&self.state,
self.catalog_epoch,
&self.candidate.provider_id,
self.used_half_open_permit,
ProviderHealthOutcome {
disposition: disposition.provider_health,
error: disposition.provider_error,
},
)
.await;
record_request_finish(
&self.state,
disposition.request_success,
@@ -921,6 +1003,7 @@ async fn prepare_response(
non_streaming_timeout_seconds: u32,
used_half_open_permit: bool,
selected_is_failover: bool,
complete_provider_health: ProviderHealthOutcome,
) -> Result<PreparedPiResponse, ProxyError> {
let status = response.status();
let headers = filtered_response_headers(response.headers());
@@ -1019,6 +1102,7 @@ async fn prepare_response(
streaming_idle_timeout_seconds,
used_half_open_permit,
selected_is_failover,
complete_provider_health,
);
let mut builder = Response::builder().status(status);
*builder
@@ -1138,6 +1222,7 @@ fn logged_body_stream(
streaming_idle_timeout_seconds: u32,
used_half_open_permit: bool,
selected_is_failover: bool,
complete_provider_health: ProviderHealthOutcome,
) -> impl futures::Stream<Item = Result<Bytes, std::io::Error>> + Send + 'static {
// Construct the guard before the generator is polled. Axum may drop a
// response body without ever polling it when the client disconnects after
@@ -1154,6 +1239,7 @@ fn logged_body_stream(
content_is_sse,
used_half_open_permit,
selected_is_failover,
complete_provider_health,
});
async_stream::stream! {
let mut captured = Vec::new();
@@ -1358,6 +1444,57 @@ mod tests {
}
}
#[tokio::test]
async fn credential_rejections_remain_retryable_but_health_neutral() {
let app = axum::Router::new()
.route(
"/unauthorized",
axum::routing::get(|| async { StatusCode::UNAUTHORIZED }),
)
.route(
"/forbidden",
axum::routing::get(|| async { StatusCode::FORBIDDEN }),
);
let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0))
.await
.expect("bind local Pi capture endpoint");
let address = listener.local_addr().expect("local capture address");
let server = tokio::spawn(async move {
axum::serve(listener, app)
.await
.expect("serve local Pi capture endpoint");
});
for (path, expected) in [
("unauthorized", StatusCode::UNAUTHORIZED),
("forbidden", StatusCode::FORBIDDEN),
] {
let response = reqwest::get(format!("http://{address}/{path}"))
.await
.expect("request local Pi capture endpoint");
assert_eq!(response.status(), expected);
assert!(retryable_status(response.status()));
assert_eq!(
ProviderHealthOutcome::from_status(response.status(), None),
ProviderHealthOutcome {
disposition: ProviderHealthDisposition::Neutral,
error: None,
}
);
}
assert_eq!(
ProviderHealthOutcome::from_status(
StatusCode::UNAUTHORIZED,
Some("earlier endpoint failed"),
),
ProviderHealthOutcome {
disposition: ProviderHealthDisposition::Unhealthy,
error: Some("earlier endpoint failed".to_string()),
}
);
server.abort();
}
#[test]
fn sse_comments_do_not_cross_the_commit_fence() {
assert!(!contains_semantic_sse_event(
@@ -1497,33 +1634,61 @@ mod tests {
fn streaming_health_waits_for_the_terminal_outcome() {
let complete = pi_stream_disposition(
StatusCode::OK,
&ProviderHealthOutcome::from_status(StatusCode::OK, None),
&PiStreamTermination::Complete {
captured: Some(Vec::new()),
},
);
assert_eq!(complete.provider_success, Some(true));
assert_eq!(complete.provider_health, ProviderHealthDisposition::Healthy);
assert!(complete.request_success);
let truncated = pi_stream_disposition(
StatusCode::OK,
&ProviderHealthOutcome::from_status(StatusCode::OK, None),
&PiStreamTermination::UpstreamFailure {
message: "truncated".to_string(),
},
);
assert_eq!(truncated.provider_success, Some(false));
assert_eq!(
truncated.provider_health,
ProviderHealthDisposition::Unhealthy
);
assert!(!truncated.request_success);
assert_eq!(truncated.provider_error.as_deref(), Some("truncated"));
let client_drop =
pi_stream_disposition(StatusCode::OK, &PiStreamTermination::DownstreamDropped);
assert_eq!(client_drop.provider_success, None);
let client_drop = pi_stream_disposition(
StatusCode::OK,
&ProviderHealthOutcome::from_status(StatusCode::OK, None),
&PiStreamTermination::DownstreamDropped,
);
assert_eq!(
client_drop.provider_health,
ProviderHealthDisposition::Neutral
);
assert!(!client_drop.request_success);
let non_retryable_error = pi_stream_disposition(
StatusCode::BAD_REQUEST,
&ProviderHealthOutcome::from_status(StatusCode::BAD_REQUEST, None),
&PiStreamTermination::Complete { captured: None },
);
assert_eq!(non_retryable_error.provider_success, Some(true));
assert_eq!(
non_retryable_error.provider_health,
ProviderHealthDisposition::Healthy
);
assert!(!non_retryable_error.request_success);
for status in [StatusCode::UNAUTHORIZED, StatusCode::FORBIDDEN] {
let credential_rejection = pi_stream_disposition(
status,
&ProviderHealthOutcome::from_status(status, None),
&PiStreamTermination::Complete { captured: None },
);
assert_eq!(
credential_rejection.provider_health,
ProviderHealthDisposition::Neutral
);
assert!(!credential_rejection.request_success);
}
}
}