mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-09 04:30:20 +08:00
fix provider pool quota status handling
This commit is contained in:
@@ -9,7 +9,35 @@ permissions:
|
||||
contents: write
|
||||
|
||||
jobs:
|
||||
preflight:
|
||||
runs-on: ubuntu-latest
|
||||
if: startsWith(github.ref, 'refs/tags/')
|
||||
steps:
|
||||
- uses: actions/checkout@v5
|
||||
|
||||
- name: Ensure proxy tag matches Cargo version
|
||||
shell: bash
|
||||
run: |
|
||||
TAG="${GITHUB_REF_NAME}"
|
||||
EXPECTED="${TAG#proxy-v}"
|
||||
ACTUAL="$(cargo metadata --manifest-path apps/aether-proxy/Cargo.toml --locked --no-deps --format-version 1 | jq -r '.packages[] | select(.name == "aether-proxy") | .version')"
|
||||
|
||||
echo "tag version: ${EXPECTED}"
|
||||
echo "cargo version: ${ACTUAL}"
|
||||
|
||||
if [ -z "${ACTUAL}" ]; then
|
||||
echo "Could not resolve aether-proxy package version" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [ "${EXPECTED}" != "${ACTUAL}" ]; then
|
||||
echo "proxy tag ${TAG} does not match apps/aether-proxy/Cargo.toml version ${ACTUAL}" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
build:
|
||||
needs: preflight
|
||||
if: always() && (needs.preflight.result == 'success' || needs.preflight.result == 'skipped')
|
||||
name: ${{ matrix.name }}
|
||||
runs-on: ${{ matrix.os }}
|
||||
strategy:
|
||||
|
||||
Generated
+1
-1
@@ -294,7 +294,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "aether-proxy"
|
||||
version = "0.3.3"
|
||||
version = "0.3.5"
|
||||
dependencies = [
|
||||
"aether-contracts",
|
||||
"aether-gateway",
|
||||
|
||||
@@ -85,13 +85,13 @@ git pull
|
||||
curl -fsSL https://raw.githubusercontent.com/fawney19/Aether/aether-rust-pioneer/install.sh | sudo bash
|
||||
```
|
||||
|
||||
运行后按提示输入版本和部署方式。固定安装某个 tag 时,版本选择选 `2`,再输入类似 `v0.7.0-rc22` 的 tag。
|
||||
运行后按提示输入版本和部署方式。固定安装某个 tag 时,版本选择选 `2`,再输入类似 `v0.7.0-rc23` 的 tag。默认会安装最新预发布版本;Docker Compose 模式默认使用 `pre` 镜像通道。
|
||||
如果安装目录里已经有配置,脚本会优先复用:Docker Compose 保留已有 `.env`,systemd 保留已有 `/etc/aether/aether-gateway.env`。只有首次生成新配置时才会提示输入管理员密码。
|
||||
|
||||
```text
|
||||
Choose Aether version:
|
||||
1) Latest rc tag
|
||||
2) Exact tag, for example v0.7.0-rc22
|
||||
1) Latest pre release
|
||||
2) Exact tag, for example v0.7.0-rc23
|
||||
|
||||
Enter choice [1]:
|
||||
|
||||
|
||||
@@ -35,6 +35,8 @@ use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_conf
|
||||
use crate::orchestration::{local_attempt_slot_count, ExecutionAttemptIdentity};
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
const POOL_KEY_RETRY_INDEX_STRIDE: u32 = 100;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct LocalExecutionCandidateAttempt {
|
||||
pub(crate) eligible: EligibleLocalExecutionCandidate,
|
||||
@@ -129,6 +131,16 @@ impl LocalExecutionCandidateAttempt {
|
||||
}
|
||||
}
|
||||
|
||||
fn effective_retry_index(retry_index: u32, pool_key_index: Option<u32>) -> u32 {
|
||||
pool_key_index
|
||||
.and_then(|index| {
|
||||
index
|
||||
.checked_mul(POOL_KEY_RETRY_INDEX_STRIDE)
|
||||
.and_then(|base| base.checked_add(retry_index))
|
||||
})
|
||||
.unwrap_or(retry_index)
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub(crate) struct LocalAvailableCandidatePersistenceContext<'a> {
|
||||
pub(crate) user_id: &'a str,
|
||||
@@ -326,7 +338,7 @@ where
|
||||
self.api_key_id,
|
||||
&candidate.candidate,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
effective_retry_index(retry_index, candidate.orchestration.pool_key_index),
|
||||
generated_candidate_id,
|
||||
self.required_capabilities,
|
||||
extra_data,
|
||||
@@ -343,6 +355,8 @@ where
|
||||
retry_index: u32,
|
||||
candidate_id: String,
|
||||
) -> Self::Attempt {
|
||||
let retry_index =
|
||||
effective_retry_index(retry_index, candidate.orchestration.pool_key_index);
|
||||
LocalExecutionCandidateAttempt {
|
||||
eligible: candidate,
|
||||
candidate_index,
|
||||
@@ -908,7 +922,7 @@ where
|
||||
context.api_key_id,
|
||||
&candidate_ref.candidate,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
effective_retry_index(retry_index, candidate_ref.orchestration.pool_key_index),
|
||||
generated_candidate_id.as_str(),
|
||||
context.required_capabilities,
|
||||
extra_data.clone(),
|
||||
@@ -927,6 +941,8 @@ where
|
||||
} else {
|
||||
candidate_ref.clone()
|
||||
};
|
||||
let retry_index =
|
||||
effective_retry_index(retry_index, candidate.orchestration.pool_key_index);
|
||||
attempts.push(LocalExecutionCandidateAttempt {
|
||||
eligible: candidate,
|
||||
candidate_index,
|
||||
@@ -957,6 +973,8 @@ fn build_unpersisted_local_execution_candidate_attempts(
|
||||
.expect("candidate should remain available until final retry")
|
||||
.clone()
|
||||
};
|
||||
let retry_index =
|
||||
effective_retry_index(retry_index, candidate.orchestration.pool_key_index);
|
||||
attempts.push_back(LocalExecutionCandidateAttempt {
|
||||
eligible: candidate,
|
||||
candidate_index,
|
||||
@@ -1290,6 +1308,29 @@ mod tests {
|
||||
assert_eq!(stored[0].candidate_index, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pool_key_attempts_use_distinct_effective_retry_indices() {
|
||||
let first = build_unpersisted_local_execution_candidate_attempts(
|
||||
sample_eligible("pool-key-1", Some(0)),
|
||||
0,
|
||||
)
|
||||
.pop_front()
|
||||
.expect("first pool key attempt");
|
||||
let second = build_unpersisted_local_execution_candidate_attempts(
|
||||
sample_eligible("pool-key-2", Some(1)),
|
||||
0,
|
||||
)
|
||||
.pop_front()
|
||||
.expect("second pool key attempt");
|
||||
|
||||
assert_eq!(first.retry_index, 0);
|
||||
assert_eq!(second.retry_index, 100);
|
||||
assert_eq!(first.attempt_identity().retry_index, 0);
|
||||
assert_eq!(second.attempt_identity().retry_index, 100);
|
||||
assert_eq!(first.attempt_identity().pool_key_index, Some(0));
|
||||
assert_eq!(second.attempt_identity().pool_key_index, Some(1));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn available_candidates_persist_ranking_metadata_in_extra_data() {
|
||||
let repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
|
||||
@@ -279,7 +279,7 @@ pub(crate) async fn record_failed_usage_for_exhausted_request(
|
||||
|
||||
state
|
||||
.usage_runtime
|
||||
.record_terminal_event(
|
||||
.record_terminal_event_direct(
|
||||
state.data.as_ref(),
|
||||
UsageEvent::new(UsageEventType::Failed, request_id, data),
|
||||
)
|
||||
@@ -418,7 +418,7 @@ pub(crate) async fn record_failed_usage_for_runtime_miss_request(
|
||||
|
||||
state
|
||||
.usage_runtime
|
||||
.record_terminal_event(
|
||||
.record_terminal_event_direct(
|
||||
state.data.as_ref(),
|
||||
UsageEvent::new(UsageEventType::Failed, request_id, data),
|
||||
)
|
||||
|
||||
+14
-3
@@ -111,7 +111,14 @@ fn reset_codex_cycle_usage_windows(status_snapshot: &mut Value, now_unix_secs: u
|
||||
}
|
||||
|
||||
window.insert("usage_reset_at".to_string(), json!(now_unix_secs));
|
||||
window.remove("usage");
|
||||
window.insert(
|
||||
"usage".to_string(),
|
||||
json!({
|
||||
"request_count": 0,
|
||||
"total_tokens": 0,
|
||||
"total_cost_usd": "0.00000000",
|
||||
}),
|
||||
);
|
||||
reset_count += 1;
|
||||
}
|
||||
|
||||
@@ -173,9 +180,13 @@ mod tests {
|
||||
assert_eq!(reset_codex_cycle_usage_windows(&mut snapshot, 1_234), 2);
|
||||
let windows = snapshot["quota"]["windows"].as_array().expect("windows");
|
||||
assert_eq!(windows[0]["usage_reset_at"], json!(1_234));
|
||||
assert!(windows[0].get("usage").is_none());
|
||||
assert_eq!(windows[0]["usage"]["request_count"], json!(0));
|
||||
assert_eq!(windows[0]["usage"]["total_tokens"], json!(0));
|
||||
assert_eq!(windows[0]["usage"]["total_cost_usd"], json!("0.00000000"));
|
||||
assert_eq!(windows[1]["usage_reset_at"], json!(1_234));
|
||||
assert!(windows[1].get("usage").is_none());
|
||||
assert_eq!(windows[1]["usage"]["request_count"], json!(0));
|
||||
assert_eq!(windows[1]["usage"]["total_tokens"], json!(0));
|
||||
assert_eq!(windows[1]["usage"]["total_cost_usd"], json!("0.00000000"));
|
||||
assert!(windows[2].get("usage_reset_at").is_none());
|
||||
assert!(windows[2].get("usage").is_some());
|
||||
}
|
||||
|
||||
@@ -147,7 +147,7 @@ pub(crate) fn merge_provider_oauth_refresh_failure_reason(
|
||||
return Some(refresh_reason.to_string());
|
||||
}
|
||||
if current_reason.starts_with(OAUTH_EXPIRED_PREFIX) {
|
||||
return Some(refresh_reason.to_string());
|
||||
return Some(current_reason.to_string());
|
||||
}
|
||||
if oauth_invalid_reason_is_account_level_block(Some(current_reason)) {
|
||||
return None;
|
||||
@@ -190,13 +190,13 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn refresh_failure_replaces_access_token_expired_marker() {
|
||||
fn refresh_failure_does_not_replace_access_token_expired_marker() {
|
||||
assert_eq!(
|
||||
merge_provider_oauth_refresh_failure_reason(
|
||||
Some("[OAUTH_EXPIRED] access token invalid"),
|
||||
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效",
|
||||
),
|
||||
Some("[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效".to_string()),
|
||||
Some("[OAUTH_EXPIRED] access token invalid".to_string()),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -9,11 +9,8 @@ use aether_admin::provider::quota as admin_provider_quota_pure;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
};
|
||||
use aether_data_contracts::repository::usage::{
|
||||
ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyWindowUsageSummary,
|
||||
};
|
||||
use serde_json::json;
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
fn admin_pool_string_list(value: Option<&serde_json::Value>) -> Option<Vec<String>> {
|
||||
let values = value
|
||||
@@ -293,128 +290,6 @@ fn admin_pool_quota_window<'a>(
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) type AdminPoolCodexWindowUsageByKey =
|
||||
BTreeMap<(String, String), StoredProviderApiKeyWindowUsageSummary>;
|
||||
|
||||
fn admin_pool_provider_type_is_codex(provider_type: &str) -> bool {
|
||||
provider_type.trim().eq_ignore_ascii_case("codex")
|
||||
}
|
||||
|
||||
fn admin_pool_codex_window_usage_code(
|
||||
window: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> Option<&'static str> {
|
||||
let code = window
|
||||
.get("code")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)?;
|
||||
if code.eq_ignore_ascii_case("5h") {
|
||||
Some("5h")
|
||||
} else if code.eq_ignore_ascii_case("weekly") {
|
||||
Some("weekly")
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn admin_pool_codex_window_usage_bounds(
|
||||
window: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> Option<(u64, u64)> {
|
||||
let reset_at = admin_pool_json_to_u64(window.get("reset_at"))?;
|
||||
let window_minutes = admin_pool_json_to_u64(window.get("window_minutes"))?;
|
||||
let window_seconds = window_minutes.checked_mul(60)?;
|
||||
let window_start = reset_at.checked_sub(window_seconds)?;
|
||||
let usage_reset_at = admin_pool_json_to_u64(window.get("usage_reset_at"))
|
||||
.filter(|reset_at_override| *reset_at_override < reset_at)
|
||||
.unwrap_or(window_start);
|
||||
let start = window_start.max(usage_reset_at);
|
||||
(start < reset_at).then_some((start, reset_at))
|
||||
}
|
||||
|
||||
pub(super) fn build_admin_pool_codex_window_usage_requests(
|
||||
provider_type: &str,
|
||||
keys: &[StoredProviderCatalogKey],
|
||||
) -> Vec<ProviderApiKeyWindowUsageRequest> {
|
||||
if !admin_pool_provider_type_is_codex(provider_type) {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let mut requests = Vec::new();
|
||||
for key in keys {
|
||||
let status_snapshot = provider_key_status_snapshot_payload(key, provider_type);
|
||||
let Some(windows) = status_snapshot
|
||||
.get("quota")
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.and_then(|quota| quota.get("windows"))
|
||||
.and_then(serde_json::Value::as_array)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
for window in windows.iter().filter_map(serde_json::Value::as_object) {
|
||||
let Some(window_code) = admin_pool_codex_window_usage_code(window) else {
|
||||
continue;
|
||||
};
|
||||
let Some((start_unix_secs, end_unix_secs)) =
|
||||
admin_pool_codex_window_usage_bounds(window)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
requests.push(ProviderApiKeyWindowUsageRequest {
|
||||
provider_api_key_id: key.id.clone(),
|
||||
window_code: window_code.to_string(),
|
||||
start_unix_secs,
|
||||
end_unix_secs,
|
||||
});
|
||||
}
|
||||
}
|
||||
requests
|
||||
}
|
||||
|
||||
fn admin_pool_codex_window_usage_payload(
|
||||
usage: &StoredProviderApiKeyWindowUsageSummary,
|
||||
) -> serde_json::Value {
|
||||
json!({
|
||||
"request_count": usage.request_count,
|
||||
"total_tokens": usage.total_tokens,
|
||||
"total_cost_usd": format!("{:.8}", usage.total_cost_usd),
|
||||
})
|
||||
}
|
||||
|
||||
fn admin_pool_attach_codex_window_usage(
|
||||
status_snapshot: &mut serde_json::Value,
|
||||
key_id: &str,
|
||||
usage_by_key: &AdminPoolCodexWindowUsageByKey,
|
||||
) {
|
||||
if usage_by_key.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(windows) = status_snapshot
|
||||
.get_mut("quota")
|
||||
.and_then(serde_json::Value::as_object_mut)
|
||||
.and_then(|quota| quota.get_mut("windows"))
|
||||
.and_then(serde_json::Value::as_array_mut)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
|
||||
for window in windows
|
||||
.iter_mut()
|
||||
.filter_map(serde_json::Value::as_object_mut)
|
||||
{
|
||||
let Some(window_code) = admin_pool_codex_window_usage_code(window) else {
|
||||
continue;
|
||||
};
|
||||
let lookup_key = (key_id.to_string(), window_code.to_string());
|
||||
if let Some(usage) = usage_by_key.get(&lookup_key) {
|
||||
window.insert(
|
||||
"usage".to_string(),
|
||||
admin_pool_codex_window_usage_payload(usage),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn admin_pool_quota_window_used_percent(
|
||||
window: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> Option<f64> {
|
||||
@@ -528,6 +403,54 @@ fn admin_pool_build_codex_account_quota_from_snapshot(
|
||||
None
|
||||
}
|
||||
|
||||
fn admin_pool_current_unix_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn admin_pool_prune_expired_codex_window_usage(status_snapshot: &mut serde_json::Value) {
|
||||
let now_unix_secs = admin_pool_current_unix_secs();
|
||||
let Some(windows) = status_snapshot
|
||||
.get_mut("quota")
|
||||
.and_then(serde_json::Value::as_object_mut)
|
||||
.and_then(|quota| quota.get_mut("windows"))
|
||||
.and_then(serde_json::Value::as_array_mut)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
|
||||
for window in windows
|
||||
.iter_mut()
|
||||
.filter_map(serde_json::Value::as_object_mut)
|
||||
{
|
||||
let code = window
|
||||
.get("code")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
if !code.eq_ignore_ascii_case("5h") && !code.eq_ignore_ascii_case("weekly") {
|
||||
continue;
|
||||
}
|
||||
let Some(reset_at) = admin_pool_json_to_u64(window.get("reset_at")) else {
|
||||
continue;
|
||||
};
|
||||
if now_unix_secs < reset_at {
|
||||
continue;
|
||||
}
|
||||
window.insert(
|
||||
"usage".to_string(),
|
||||
json!({
|
||||
"request_count": 0,
|
||||
"total_tokens": 0,
|
||||
"total_cost_usd": "0.00000000",
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn admin_pool_quota_windows<'a>(
|
||||
quota_snapshot: &'a serde_json::Map<String, serde_json::Value>,
|
||||
) -> Vec<&'a serde_json::Map<String, serde_json::Value>> {
|
||||
@@ -930,7 +853,6 @@ pub(super) fn build_admin_pool_key_payload(
|
||||
key: &StoredProviderCatalogKey,
|
||||
runtime: &AdminProviderPoolRuntimeState,
|
||||
pool_config: Option<AdminProviderPoolConfig>,
|
||||
codex_window_usage_by_key: &AdminPoolCodexWindowUsageByKey,
|
||||
) -> serde_json::Value {
|
||||
let cooldown_reason = runtime.cooldown_reason_by_key.get(&key.id).cloned();
|
||||
let cooldown_ttl_seconds = cooldown_reason
|
||||
@@ -949,12 +871,8 @@ pub(super) fn build_admin_pool_key_payload(
|
||||
let oauth_plan_type =
|
||||
admin_pool_derive_oauth_plan_type(key, provider_type, auth_config.as_ref());
|
||||
let mut status_snapshot = provider_key_status_snapshot_payload(key, provider_type);
|
||||
if admin_pool_provider_type_is_codex(provider_type) {
|
||||
admin_pool_attach_codex_window_usage(
|
||||
&mut status_snapshot,
|
||||
&key.id,
|
||||
codex_window_usage_by_key,
|
||||
);
|
||||
if provider_type.trim().eq_ignore_ascii_case("codex") {
|
||||
admin_pool_prune_expired_codex_window_usage(&mut status_snapshot);
|
||||
}
|
||||
let account_snapshot = status_snapshot
|
||||
.get("account")
|
||||
@@ -1237,34 +1155,54 @@ mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn codex_window_usage_bounds_honor_manual_usage_reset_at() {
|
||||
let window = json!({
|
||||
"code": "5h",
|
||||
"reset_at": 20_000,
|
||||
"window_minutes": 300,
|
||||
"usage_reset_at": 9_000,
|
||||
fn expired_codex_window_usage_is_zeroed_for_display() {
|
||||
let mut snapshot = json!({
|
||||
"quota": {
|
||||
"windows": [
|
||||
{
|
||||
"code": "5h",
|
||||
"reset_at": 1,
|
||||
"usage": {
|
||||
"request_count": 7,
|
||||
"total_tokens": 700,
|
||||
"total_cost_usd": "7.00000000"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
});
|
||||
let window = window.as_object().expect("window object");
|
||||
|
||||
assert_eq!(
|
||||
admin_pool_codex_window_usage_bounds(window),
|
||||
Some((9_000, 20_000))
|
||||
);
|
||||
admin_pool_prune_expired_codex_window_usage(&mut snapshot);
|
||||
|
||||
let usage = &snapshot["quota"]["windows"][0]["usage"];
|
||||
assert_eq!(usage["request_count"], json!(0));
|
||||
assert_eq!(usage["total_tokens"], json!(0));
|
||||
assert_eq!(usage["total_cost_usd"], json!("0.00000000"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_window_usage_bounds_ignore_reset_before_window_start() {
|
||||
let window = json!({
|
||||
"code": "weekly",
|
||||
"reset_at": 700_000,
|
||||
"window_minutes": 10_080,
|
||||
"usage_reset_at": 1,
|
||||
fn active_codex_window_usage_is_preserved_for_display() {
|
||||
let mut snapshot = json!({
|
||||
"quota": {
|
||||
"windows": [
|
||||
{
|
||||
"code": "weekly",
|
||||
"reset_at": 4_102_444_800u64,
|
||||
"usage": {
|
||||
"request_count": 3,
|
||||
"total_tokens": 375,
|
||||
"total_cost_usd": "0.60000000"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
});
|
||||
let window = window.as_object().expect("window object");
|
||||
|
||||
assert_eq!(
|
||||
admin_pool_codex_window_usage_bounds(window),
|
||||
Some((95_200, 700_000))
|
||||
);
|
||||
admin_pool_prune_expired_codex_window_usage(&mut snapshot);
|
||||
|
||||
let usage = &snapshot["quota"]["windows"][0]["usage"];
|
||||
assert_eq!(usage["request_count"], json!(3));
|
||||
assert_eq!(usage["total_tokens"], json!(375));
|
||||
assert_eq!(usage["total_cost_usd"], json!("0.60000000"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -253,36 +253,6 @@ pub(super) async fn build_admin_pool_list_keys_response(
|
||||
}
|
||||
_ => AdminProviderPoolRuntimeState::default(),
|
||||
};
|
||||
let codex_window_usage_requests =
|
||||
pool_payloads::build_admin_pool_codex_window_usage_requests(&provider.provider_type, &keys);
|
||||
let mut codex_window_usage_by_key = codex_window_usage_requests
|
||||
.iter()
|
||||
.map(|request| {
|
||||
(
|
||||
(
|
||||
request.provider_api_key_id.clone(),
|
||||
request.window_code.clone(),
|
||||
),
|
||||
aether_data_contracts::repository::usage::StoredProviderApiKeyWindowUsageSummary {
|
||||
provider_api_key_id: request.provider_api_key_id.clone(),
|
||||
window_code: request.window_code.clone(),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect::<pool_payloads::AdminPoolCodexWindowUsageByKey>();
|
||||
if !codex_window_usage_requests.is_empty() {
|
||||
for usage in state
|
||||
.app()
|
||||
.summarize_usage_by_provider_api_key_windows(&codex_window_usage_requests)
|
||||
.await?
|
||||
{
|
||||
codex_window_usage_by_key.insert(
|
||||
(usage.provider_api_key_id.clone(), usage.window_code.clone()),
|
||||
usage,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let items = keys
|
||||
.into_iter()
|
||||
@@ -294,7 +264,6 @@ pub(super) async fn build_admin_pool_list_keys_response(
|
||||
&key,
|
||||
&runtime,
|
||||
pool_config.clone(),
|
||||
&codex_window_usage_by_key,
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
@@ -240,6 +240,10 @@ fn tagged_oauth_invalid_reason(reason: Option<&str>, prefix: &str) -> Option<Str
|
||||
})
|
||||
}
|
||||
|
||||
fn oauth_access_token_expired(expires_at_unix_secs: Option<u64>, now_unix_secs: u64) -> bool {
|
||||
expires_at_unix_secs.is_none_or(|expires_at| expires_at == 0 || expires_at <= now_unix_secs)
|
||||
}
|
||||
|
||||
fn build_provider_key_oauth_status_snapshot(key: &StoredProviderCatalogKey) -> Value {
|
||||
if !key.auth_type.trim().eq_ignore_ascii_case("oauth") {
|
||||
return default_oauth_status_snapshot_value();
|
||||
@@ -271,14 +275,16 @@ fn build_provider_key_oauth_status_snapshot(key: &StoredProviderCatalogKey) -> V
|
||||
if let Some(reason) =
|
||||
tagged_oauth_invalid_reason(invalid_reason.as_deref(), OAUTH_REFRESH_FAILED_PREFIX)
|
||||
{
|
||||
let access_token_expired = oauth_access_token_expired(expires_at_unix_secs, now_unix_secs);
|
||||
return json!({
|
||||
"code": "invalid",
|
||||
"label": "已失效",
|
||||
"code": if access_token_expired { "invalid" } else { "reauth_required" },
|
||||
"label": if access_token_expired { "已失效" } else { "续期失败" },
|
||||
"reason": reason,
|
||||
"expires_at": expires_at_unix_secs,
|
||||
"invalid_at": invalid_at_unix_secs,
|
||||
"source": "oauth_refresh",
|
||||
"requires_reauth": true,
|
||||
"usable_until_expiry": !access_token_expired,
|
||||
"expiring_soon": false,
|
||||
});
|
||||
}
|
||||
@@ -320,11 +326,11 @@ fn build_provider_key_oauth_status_snapshot(key: &StoredProviderCatalogKey) -> V
|
||||
return json!({
|
||||
"code": "expired",
|
||||
"label": "已过期",
|
||||
"reason": "Token 已过期,请重新授权",
|
||||
"reason": "Access Token 已过期,等待自动续期",
|
||||
"expires_at": expires_at_unix_secs,
|
||||
"invalid_at": Value::Null,
|
||||
"source": "expires_at",
|
||||
"requires_reauth": true,
|
||||
"requires_reauth": false,
|
||||
"expiring_soon": false,
|
||||
});
|
||||
}
|
||||
@@ -556,10 +562,7 @@ fn quota_windows_all_exhausted(windows: &[Value]) -> bool {
|
||||
total > 0 && exhausted == total
|
||||
}
|
||||
|
||||
fn preserve_quota_window_usage_reset_at(
|
||||
current_status_snapshot: Option<&Value>,
|
||||
quota: &mut Value,
|
||||
) {
|
||||
fn preserve_quota_window_usage_state(current_status_snapshot: Option<&Value>, quota: &mut Value) {
|
||||
let Some(current_windows) = current_status_snapshot
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|snapshot| snapshot.get("quota"))
|
||||
@@ -582,23 +585,39 @@ fn preserve_quota_window_usage_reset_at(
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let Some(usage_reset_at) = current_windows
|
||||
.iter()
|
||||
.filter_map(Value::as_object)
|
||||
.find(|current_window| {
|
||||
current_window
|
||||
.get("code")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|current_code| current_code.eq_ignore_ascii_case(code))
|
||||
})
|
||||
.and_then(|current_window| current_window.get("usage_reset_at"))
|
||||
.and_then(admin_provider_quota_pure::coerce_json_u64)
|
||||
else {
|
||||
let current_window =
|
||||
current_windows
|
||||
.iter()
|
||||
.filter_map(Value::as_object)
|
||||
.find(|current_window| {
|
||||
current_window
|
||||
.get("code")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|current_code| current_code.eq_ignore_ascii_case(code))
|
||||
});
|
||||
let Some(current_window) = current_window else {
|
||||
continue;
|
||||
};
|
||||
let current_reset_at = current_window
|
||||
.get("reset_at")
|
||||
.and_then(admin_provider_quota_pure::coerce_json_u64);
|
||||
let next_reset_at = next_window
|
||||
.get("reset_at")
|
||||
.and_then(admin_provider_quota_pure::coerce_json_u64);
|
||||
if current_reset_at.is_none() || current_reset_at != next_reset_at {
|
||||
continue;
|
||||
}
|
||||
|
||||
next_window.insert("usage_reset_at".to_string(), json!(usage_reset_at));
|
||||
if let Some(usage_reset_at) = current_window
|
||||
.get("usage_reset_at")
|
||||
.and_then(admin_provider_quota_pure::coerce_json_u64)
|
||||
{
|
||||
next_window.insert("usage_reset_at".to_string(), json!(usage_reset_at));
|
||||
}
|
||||
if let Some(usage) = current_window.get("usage") {
|
||||
next_window.insert("usage".to_string(), usage.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1162,7 +1181,7 @@ pub(crate) fn sync_provider_key_quota_status_snapshot(
|
||||
_ => None,
|
||||
}?;
|
||||
if normalized_provider_type == "codex" {
|
||||
preserve_quota_window_usage_reset_at(status_snapshot, &mut quota);
|
||||
preserve_quota_window_usage_state(status_snapshot, &mut quota);
|
||||
}
|
||||
|
||||
let default_snapshot = default_provider_key_status_snapshot();
|
||||
@@ -2062,7 +2081,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_provider_key_quota_status_snapshot_preserves_codex_usage_reset_at() {
|
||||
fn sync_provider_key_quota_status_snapshot_preserves_codex_usage_state() {
|
||||
let current_status_snapshot = json!({
|
||||
"quota": {
|
||||
"version": 2,
|
||||
@@ -2071,12 +2090,22 @@ mod tests {
|
||||
{
|
||||
"code": "weekly",
|
||||
"usage_reset_at": 1_775_600_000u64,
|
||||
"usage": {
|
||||
"request_count": 3,
|
||||
"total_tokens": 375,
|
||||
"total_cost_usd": "0.60000000"
|
||||
},
|
||||
"reset_at": 1_900_000_000u64,
|
||||
"window_minutes": 10_080u64
|
||||
},
|
||||
{
|
||||
"code": "5h",
|
||||
"usage_reset_at": 1_775_700_000u64,
|
||||
"usage": {
|
||||
"request_count": 2,
|
||||
"total_tokens": 225,
|
||||
"total_cost_usd": "0.30000000"
|
||||
},
|
||||
"reset_at": 1_900_500_000u64,
|
||||
"window_minutes": 300u64
|
||||
}
|
||||
@@ -2088,10 +2117,10 @@ mod tests {
|
||||
"updated_at": 1_775_800_000u64,
|
||||
"plan_type": "plus",
|
||||
"primary_used_percent": 5.0,
|
||||
"primary_reset_at": 1_901_000_000u64,
|
||||
"primary_reset_at": 1_900_000_000u64,
|
||||
"primary_window_minutes": 10_080u64,
|
||||
"secondary_used_percent": 1.0,
|
||||
"secondary_reset_at": 1_901_500_000u64,
|
||||
"secondary_reset_at": 1_900_500_000u64,
|
||||
"secondary_window_minutes": 300u64
|
||||
}
|
||||
});
|
||||
@@ -2121,7 +2150,94 @@ mod tests {
|
||||
.expect("5h window should exist");
|
||||
|
||||
assert_eq!(weekly.get("usage_reset_at"), Some(&json!(1_775_600_000u64)));
|
||||
assert_eq!(
|
||||
weekly
|
||||
.get("usage")
|
||||
.and_then(|usage| usage.get("request_count")),
|
||||
Some(&json!(3))
|
||||
);
|
||||
assert_eq!(
|
||||
weekly
|
||||
.get("usage")
|
||||
.and_then(|usage| usage.get("total_tokens")),
|
||||
Some(&json!(375))
|
||||
);
|
||||
assert_eq!(
|
||||
weekly
|
||||
.get("usage")
|
||||
.and_then(|usage| usage.get("total_cost_usd")),
|
||||
Some(&json!("0.60000000"))
|
||||
);
|
||||
assert_eq!(five_h.get("usage_reset_at"), Some(&json!(1_775_700_000u64)));
|
||||
assert_eq!(
|
||||
five_h
|
||||
.get("usage")
|
||||
.and_then(|usage| usage.get("request_count")),
|
||||
Some(&json!(2))
|
||||
);
|
||||
assert_eq!(
|
||||
five_h
|
||||
.get("usage")
|
||||
.and_then(|usage| usage.get("total_tokens")),
|
||||
Some(&json!(225))
|
||||
);
|
||||
assert_eq!(
|
||||
five_h
|
||||
.get("usage")
|
||||
.and_then(|usage| usage.get("total_cost_usd")),
|
||||
Some(&json!("0.30000000"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_provider_key_quota_status_snapshot_drops_codex_usage_state_when_window_resets() {
|
||||
let current_status_snapshot = json!({
|
||||
"quota": {
|
||||
"version": 2,
|
||||
"provider_type": "codex",
|
||||
"windows": [
|
||||
{
|
||||
"code": "weekly",
|
||||
"usage_reset_at": 1_775_600_000u64,
|
||||
"usage": {
|
||||
"request_count": 3,
|
||||
"total_tokens": 375,
|
||||
"total_cost_usd": "0.60000000"
|
||||
},
|
||||
"reset_at": 1_900_000_000u64,
|
||||
"window_minutes": 10_080u64
|
||||
}
|
||||
]
|
||||
}
|
||||
});
|
||||
let upstream_metadata = json!({
|
||||
"codex": {
|
||||
"updated_at": 1_900_000_100u64,
|
||||
"plan_type": "plus",
|
||||
"primary_used_percent": 0.0,
|
||||
"primary_reset_at": 1_960_480_100u64,
|
||||
"primary_window_minutes": 10_080u64
|
||||
}
|
||||
});
|
||||
|
||||
let payload = sync_provider_key_quota_status_snapshot(
|
||||
Some(¤t_status_snapshot),
|
||||
"codex",
|
||||
Some(&upstream_metadata),
|
||||
"refresh_api",
|
||||
)
|
||||
.expect("quota snapshot should sync");
|
||||
let weekly = payload["quota"]["windows"]
|
||||
.as_array()
|
||||
.expect("quota windows should exist")
|
||||
.iter()
|
||||
.filter_map(Value::as_object)
|
||||
.find(|window| window.get("code") == Some(&json!("weekly")))
|
||||
.expect("weekly window should exist");
|
||||
|
||||
assert_eq!(weekly.get("reset_at"), Some(&json!(1_960_480_100u64)));
|
||||
assert!(weekly.get("usage_reset_at").is_none());
|
||||
assert!(weekly.get("usage").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -10,6 +10,7 @@ use aether_scheduler_core::{
|
||||
candidate_is_selectable_with_runtime_state, candidate_runtime_skip_reason_with_state,
|
||||
CandidateRuntimeSelectabilityInput,
|
||||
};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use crate::data::auth::GatewayAuthApiKeySnapshot;
|
||||
use crate::GatewayError;
|
||||
@@ -343,7 +344,7 @@ fn read_key_oauth_invalid_map(
|
||||
fn key_requires_oauth_reauth(
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
_now_unix_secs: u64,
|
||||
now_unix_secs: u64,
|
||||
) -> bool {
|
||||
if !key.auth_type.trim().eq_ignore_ascii_case("oauth") {
|
||||
return false;
|
||||
@@ -355,34 +356,69 @@ fn key_requires_oauth_reauth(
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
if !invalid_reason.is_empty() {
|
||||
return oauth_invalid_reason_is_hard_account_block(key, provider_type, invalid_reason);
|
||||
return oauth_invalid_reason_blocks_scheduling(
|
||||
key,
|
||||
provider_type,
|
||||
invalid_reason,
|
||||
now_unix_secs,
|
||||
);
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
|
||||
fn oauth_invalid_reason_is_hard_account_block(
|
||||
fn oauth_invalid_reason_blocks_scheduling(
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
invalid_reason: &str,
|
||||
now_unix_secs: u64,
|
||||
) -> bool {
|
||||
if provider_type.trim().eq_ignore_ascii_case("kiro")
|
||||
&& invalid_reason.trim().starts_with("[REFRESH_FAILED] ")
|
||||
{
|
||||
let trimmed_reason = invalid_reason.trim();
|
||||
if oauth_invalid_reason_has_tag(trimmed_reason, "[OAUTH_EXPIRED]") {
|
||||
return true;
|
||||
}
|
||||
|
||||
let account_state = admin_provider_status_pure::resolve_pool_account_state(
|
||||
Some(provider_type),
|
||||
key.upstream_metadata.as_ref(),
|
||||
Some(invalid_reason),
|
||||
Some(trimmed_reason),
|
||||
);
|
||||
account_state.blocked
|
||||
if account_state.blocked
|
||||
&& !account_state.recoverable
|
||||
&& account_state
|
||||
.code
|
||||
.as_deref()
|
||||
.is_some_and(oauth_account_state_code_is_hard_block)
|
||||
{
|
||||
return true;
|
||||
}
|
||||
|
||||
if oauth_invalid_reason_has_tag(trimmed_reason, "[REFRESH_FAILED]") {
|
||||
return oauth_access_token_expired(key, now_unix_secs);
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
|
||||
fn oauth_invalid_reason_has_tag(reason: &str, tag: &str) -> bool {
|
||||
reason
|
||||
.lines()
|
||||
.map(str::trim)
|
||||
.any(|line| line.starts_with(tag))
|
||||
}
|
||||
|
||||
fn oauth_access_token_expired(key: &StoredProviderCatalogKey, now_unix_secs: u64) -> bool {
|
||||
let now_unix_secs = if now_unix_secs == 0 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or(0)
|
||||
} else {
|
||||
now_unix_secs
|
||||
};
|
||||
key.expires_at_unix_secs
|
||||
.is_none_or(|expires_at| expires_at == 0 || expires_at <= now_unix_secs)
|
||||
}
|
||||
|
||||
fn oauth_account_state_code_is_hard_block(code: &str) -> bool {
|
||||
|
||||
@@ -1692,6 +1692,7 @@ async fn keeps_refresh_failed_oauth_candidate_selectable_before_local_auth_resol
|
||||
{
|
||||
let mut key = sample_key("key-codex", "provider-codex", Some(10));
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.expires_at_unix_secs = Some(1_710_000_200);
|
||||
key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
|
||||
key.oauth_invalid_reason = Some(
|
||||
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 已被使用并轮换,请重新登录授权"
|
||||
@@ -1733,6 +1734,136 @@ async fn keeps_refresh_failed_oauth_candidate_selectable_before_local_auth_resol
|
||||
assert!(skipped.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn skips_refresh_failed_oauth_candidate_after_access_token_expiry() {
|
||||
let mut row = sample_row();
|
||||
row.provider_id = "provider-codex".to_string();
|
||||
row.provider_name = "codex".to_string();
|
||||
row.provider_type = "codex".to_string();
|
||||
row.endpoint_id = "endpoint-codex".to_string();
|
||||
row.endpoint_api_format = "openai:responses".to_string();
|
||||
row.key_id = "key-codex".to_string();
|
||||
row.key_name = "codex-refresh-failed-expired".to_string();
|
||||
row.key_auth_type = "oauth".to_string();
|
||||
row.key_api_formats = Some(vec!["openai:responses".to_string()]);
|
||||
row.key_global_priority_by_format = Some(serde_json::json!({"openai:responses": 1}));
|
||||
|
||||
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
row,
|
||||
]));
|
||||
let mut provider = sample_provider("provider-codex", None);
|
||||
provider.provider_type = "codex".to_string();
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
Vec::new(),
|
||||
vec![{
|
||||
let mut key = sample_key("key-codex", "provider-codex", Some(10));
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.expires_at_unix_secs = Some(1_710_000_000);
|
||||
key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
|
||||
key.oauth_invalid_reason = Some(
|
||||
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 已被使用并轮换,请重新登录授权"
|
||||
.to_string(),
|
||||
);
|
||||
key
|
||||
}],
|
||||
));
|
||||
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![]));
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_candidate_selection_provider_catalog_quota_and_request_candidates_for_tests(
|
||||
candidates,
|
||||
provider_catalog,
|
||||
quotas,
|
||||
request_candidates,
|
||||
),
|
||||
);
|
||||
|
||||
let (selected, skipped) = collect_selectable_candidates_with_skip_reasons(
|
||||
state.data.as_ref(),
|
||||
&state,
|
||||
"openai:responses",
|
||||
"gpt-4.1",
|
||||
false,
|
||||
None,
|
||||
1_710_000_100,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
|
||||
assert!(selected.is_empty());
|
||||
assert_eq!(skipped.len(), 1);
|
||||
assert_eq!(skipped[0].candidate.key_id, "key-codex");
|
||||
assert_eq!(skipped[0].skip_reason, "oauth_invalid");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn skips_oauth_candidate_with_account_block_even_when_refresh_failed_is_present() {
|
||||
let mut row = sample_row();
|
||||
row.provider_id = "provider-codex".to_string();
|
||||
row.provider_name = "codex".to_string();
|
||||
row.provider_type = "codex".to_string();
|
||||
row.endpoint_id = "endpoint-codex".to_string();
|
||||
row.endpoint_api_format = "openai:responses".to_string();
|
||||
row.key_id = "key-codex".to_string();
|
||||
row.key_name = "codex-account-blocked-refresh-failed".to_string();
|
||||
row.key_auth_type = "oauth".to_string();
|
||||
row.key_api_formats = Some(vec!["openai:responses".to_string()]);
|
||||
row.key_global_priority_by_format = Some(serde_json::json!({"openai:responses": 1}));
|
||||
|
||||
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
row,
|
||||
]));
|
||||
let mut provider = sample_provider("provider-codex", None);
|
||||
provider.provider_type = "codex".to_string();
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
Vec::new(),
|
||||
vec![{
|
||||
let mut key = sample_key("key-codex", "provider-codex", Some(10));
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.expires_at_unix_secs = Some(1_710_000_200);
|
||||
key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
|
||||
key.oauth_invalid_reason = Some(
|
||||
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 已被使用并轮换,请重新登录授权\n[ACCOUNT_BLOCK] account has been deactivated"
|
||||
.to_string(),
|
||||
);
|
||||
key
|
||||
}],
|
||||
));
|
||||
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![]));
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_candidate_selection_provider_catalog_quota_and_request_candidates_for_tests(
|
||||
candidates,
|
||||
provider_catalog,
|
||||
quotas,
|
||||
request_candidates,
|
||||
),
|
||||
);
|
||||
|
||||
let (selected, skipped) = collect_selectable_candidates_with_skip_reasons(
|
||||
state.data.as_ref(),
|
||||
&state,
|
||||
"openai:responses",
|
||||
"gpt-4.1",
|
||||
false,
|
||||
None,
|
||||
1_710_000_100,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
|
||||
assert!(selected.is_empty());
|
||||
assert_eq!(skipped.len(), 1);
|
||||
assert_eq!(skipped[0].candidate.key_id, "key-codex");
|
||||
assert_eq!(skipped[0].skip_reason, "oauth_invalid");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn keeps_request_failed_oauth_candidate_selectable() {
|
||||
let mut row = sample_row();
|
||||
@@ -1971,7 +2102,7 @@ async fn keeps_refreshable_kiro_candidate_selectable_when_oauth_token_expired()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn skips_kiro_candidate_after_refresh_token_failure() {
|
||||
async fn keeps_kiro_candidate_selectable_after_refresh_token_failure_until_access_token_expiry() {
|
||||
let mut row = sample_row();
|
||||
row.provider_id = "provider-kiro".to_string();
|
||||
row.provider_name = "kiro".to_string();
|
||||
@@ -1995,6 +2126,71 @@ async fn skips_kiro_candidate_after_refresh_token_failure() {
|
||||
let mut key = sample_key("key-kiro", "provider-kiro", Some(10));
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.encrypted_auth_config = Some("encrypted-refreshable-session".to_string());
|
||||
key.expires_at_unix_secs = Some(1_710_000_200);
|
||||
key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
|
||||
key.oauth_invalid_reason = Some(
|
||||
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效、已过期或已撤销,请重新登录授权"
|
||||
.to_string(),
|
||||
);
|
||||
key
|
||||
}],
|
||||
));
|
||||
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![]));
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_candidate_selection_provider_catalog_quota_and_request_candidates_for_tests(
|
||||
candidates,
|
||||
provider_catalog,
|
||||
quotas,
|
||||
request_candidates,
|
||||
),
|
||||
);
|
||||
|
||||
let (selected, skipped) = collect_selectable_candidates_with_skip_reasons(
|
||||
state.data.as_ref(),
|
||||
&state,
|
||||
"claude:messages",
|
||||
"gpt-4.1",
|
||||
false,
|
||||
None,
|
||||
1_710_000_100,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
|
||||
assert_eq!(selected.len(), 1);
|
||||
assert_eq!(selected[0].provider_id, "provider-kiro");
|
||||
assert!(skipped.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn skips_kiro_candidate_after_refresh_token_failure_and_access_token_expiry() {
|
||||
let mut row = sample_row();
|
||||
row.provider_id = "provider-kiro".to_string();
|
||||
row.provider_name = "kiro".to_string();
|
||||
row.provider_type = "kiro".to_string();
|
||||
row.endpoint_id = "endpoint-kiro".to_string();
|
||||
row.endpoint_api_format = "claude:messages".to_string();
|
||||
row.key_id = "key-kiro".to_string();
|
||||
row.key_name = "kiro-refresh-failed-expired".to_string();
|
||||
row.key_auth_type = "oauth".to_string();
|
||||
row.key_api_formats = Some(vec!["claude:messages".to_string()]);
|
||||
|
||||
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
row,
|
||||
]));
|
||||
let mut provider = sample_provider("provider-kiro", None);
|
||||
provider.provider_type = "kiro".to_string();
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
Vec::new(),
|
||||
vec![{
|
||||
let mut key = sample_key("key-kiro", "provider-kiro", Some(10));
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.encrypted_auth_config = Some("encrypted-refreshable-session".to_string());
|
||||
key.expires_at_unix_secs = Some(1_710_000_000);
|
||||
key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
|
||||
key.oauth_invalid_reason = Some(
|
||||
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效、已过期或已撤销,请重新登录授权"
|
||||
|
||||
@@ -53,6 +53,10 @@ fn tagged_reason(reason: Option<&str>, prefix: &str) -> Option<String> {
|
||||
})
|
||||
}
|
||||
|
||||
fn oauth_access_token_expired(expires_at_unix_secs: Option<u64>, now_unix_secs: u64) -> bool {
|
||||
expires_at_unix_secs.is_none_or(|expires_at| expires_at == 0 || expires_at <= now_unix_secs)
|
||||
}
|
||||
|
||||
fn oauth_auth_config_refresh_token_fingerprint(auth_config: Option<&str>) -> Option<String> {
|
||||
let parsed = auth_config
|
||||
.map(str::trim)
|
||||
@@ -275,7 +279,7 @@ fn merge_local_oauth_refresh_failure_reason(
|
||||
return Some(refresh_reason.to_string());
|
||||
}
|
||||
if current_reason.starts_with(OAUTH_EXPIRED_PREFIX) {
|
||||
return Some(refresh_reason.to_string());
|
||||
return Some(current_reason.to_string());
|
||||
}
|
||||
if oauth_invalid_reason_is_account_block(Some(current_reason)) {
|
||||
return None;
|
||||
@@ -345,14 +349,16 @@ fn build_oauth_status_snapshot_value(key: &StoredProviderCatalogKey) -> Value {
|
||||
});
|
||||
}
|
||||
if let Some(reason) = tagged_reason(invalid_reason.as_deref(), OAUTH_REFRESH_FAILED_PREFIX) {
|
||||
let access_token_expired = oauth_access_token_expired(expires_at_unix_secs, now_unix_secs);
|
||||
return json!({
|
||||
"code": "invalid",
|
||||
"label": "已失效",
|
||||
"code": if access_token_expired { "invalid" } else { "reauth_required" },
|
||||
"label": if access_token_expired { "已失效" } else { "续期失败" },
|
||||
"reason": reason,
|
||||
"expires_at": expires_at_unix_secs,
|
||||
"invalid_at": invalid_at_unix_secs,
|
||||
"source": "oauth_refresh",
|
||||
"requires_reauth": true,
|
||||
"usable_until_expiry": !access_token_expired,
|
||||
"expiring_soon": false,
|
||||
});
|
||||
}
|
||||
@@ -392,11 +398,11 @@ fn build_oauth_status_snapshot_value(key: &StoredProviderCatalogKey) -> Value {
|
||||
return json!({
|
||||
"code": "expired",
|
||||
"label": "已过期",
|
||||
"reason": "Token 已过期,请重新授权",
|
||||
"reason": "Access Token 已过期,等待自动续期",
|
||||
"expires_at": expires_at_unix_secs,
|
||||
"invalid_at": Value::Null,
|
||||
"source": "expires_at",
|
||||
"requires_reauth": true,
|
||||
"requires_reauth": false,
|
||||
"expiring_soon": false,
|
||||
});
|
||||
}
|
||||
@@ -1688,13 +1694,13 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_refresh_failure_replaces_access_token_expired_marker() {
|
||||
fn local_refresh_failure_does_not_replace_access_token_expired_marker() {
|
||||
assert_eq!(
|
||||
super::merge_local_oauth_refresh_failure_reason(
|
||||
Some("[OAUTH_EXPIRED] access token invalid"),
|
||||
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效",
|
||||
),
|
||||
Some("[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效".to_string()),
|
||||
Some("[OAUTH_EXPIRED] access token invalid".to_string()),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4463,11 +4463,11 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
|
||||
"oauth": {
|
||||
"code": "expired",
|
||||
"label": "已过期",
|
||||
"reason": "Token 已过期,请重新授权",
|
||||
"reason": "Access Token 已过期,等待自动续期",
|
||||
"expires_at": 1u64,
|
||||
"invalid_at": 1_700_000_000u64,
|
||||
"source": "expires_at",
|
||||
"requires_reauth": true,
|
||||
"requires_reauth": false,
|
||||
"expiring_soon": false
|
||||
},
|
||||
"account": {
|
||||
@@ -4754,13 +4754,13 @@ async fn gateway_marks_manual_oauth_refresh_failures_as_invalid_in_pool_payload(
|
||||
"stale-codex-access-token",
|
||||
);
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.expires_at_unix_secs = Some(1_900_000_000);
|
||||
key.expires_at_unix_secs = Some(4_102_444_800);
|
||||
key.status_snapshot = Some(json!({
|
||||
"oauth": {
|
||||
"code": "valid",
|
||||
"label": "有效",
|
||||
"reason": serde_json::Value::Null,
|
||||
"expires_at": 1_900_000_000u64,
|
||||
"expires_at": 4_102_444_800u64,
|
||||
"invalid_at": serde_json::Value::Null,
|
||||
"source": "expires_at",
|
||||
"requires_reauth": false,
|
||||
@@ -4788,7 +4788,7 @@ async fn gateway_marks_manual_oauth_refresh_failures_as_invalid_in_pool_payload(
|
||||
key.encrypted_auth_config = Some(
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"provider_type":"codex","refresh_token":"used-refresh-token","email":"alice@example.com","account_id":"acct-codex-123","plan_type":"plus","expires_at":1900000000}"#,
|
||||
r#"{"provider_type":"codex","refresh_token":"used-refresh-token","email":"alice@example.com","account_id":"acct-codex-123","plan_type":"plus","expires_at":4102444800}"#,
|
||||
)
|
||||
.expect("auth config ciphertext should build"),
|
||||
);
|
||||
@@ -4873,7 +4873,7 @@ async fn gateway_marks_manual_oauth_refresh_failures_as_invalid_in_pool_payload(
|
||||
);
|
||||
assert_eq!(
|
||||
keys[0]["status_snapshot"]["oauth"]["code"],
|
||||
json!("invalid")
|
||||
json!("reauth_required")
|
||||
);
|
||||
assert_eq!(
|
||||
keys[0]["status_snapshot"]["oauth"]["reason"],
|
||||
@@ -4887,6 +4887,10 @@ async fn gateway_marks_manual_oauth_refresh_failures_as_invalid_in_pool_payload(
|
||||
keys[0]["status_snapshot"]["oauth"]["requires_reauth"],
|
||||
json!(true)
|
||||
);
|
||||
assert_eq!(
|
||||
keys[0]["status_snapshot"]["oauth"]["usable_until_expiry"],
|
||||
json!(true)
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
token_handle.abort();
|
||||
|
||||
@@ -2,9 +2,7 @@ use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data::repository::usage::InMemoryUsageReadRepository;
|
||||
use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository;
|
||||
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
|
||||
use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::routing::{any, get, post};
|
||||
use axum::{extract::Request, Router};
|
||||
@@ -23,54 +21,6 @@ use crate::constants::{
|
||||
use crate::control::resolve_public_request_context;
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
fn sample_pool_usage_row(
|
||||
request_id: &str,
|
||||
provider_api_key_id: &str,
|
||||
created_at_unix_secs: i64,
|
||||
total_tokens: i32,
|
||||
total_cost_usd: f64,
|
||||
) -> StoredRequestUsageAudit {
|
||||
StoredRequestUsageAudit::new(
|
||||
format!("usage-{request_id}"),
|
||||
request_id.to_string(),
|
||||
Some("user-codex".to_string()),
|
||||
Some("api-key-codex".to_string()),
|
||||
Some("codex-user".to_string()),
|
||||
Some("codex-api-key".to_string()),
|
||||
"codex".to_string(),
|
||||
"gpt-5-codex".to_string(),
|
||||
None,
|
||||
Some("provider-codex".to_string()),
|
||||
Some("endpoint-codex".to_string()),
|
||||
Some(provider_api_key_id.to_string()),
|
||||
Some("responses".to_string()),
|
||||
Some("openai:responses".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("responses".to_string()),
|
||||
Some("openai:responses".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("responses".to_string()),
|
||||
false,
|
||||
false,
|
||||
total_tokens,
|
||||
0,
|
||||
total_tokens,
|
||||
total_cost_usd,
|
||||
total_cost_usd,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(240),
|
||||
Some(80),
|
||||
"completed".to_string(),
|
||||
"settled".to_string(),
|
||||
created_at_unix_secs,
|
||||
created_at_unix_secs + 1,
|
||||
Some(created_at_unix_secs + 2),
|
||||
)
|
||||
.expect("usage row should build")
|
||||
}
|
||||
|
||||
fn trusted_admin_headers() -> HeaderMap {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(GATEWAY_HEADER, HeaderValue::from_static("rust-phase3b"));
|
||||
@@ -866,8 +816,8 @@ async fn gateway_sorts_admin_pool_keys_by_imported_and_last_used_time() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_pool_list_adds_codex_cycle_usage_to_quota_windows() {
|
||||
const RESET_AT: u64 = 1_711_000_000;
|
||||
async fn gateway_pool_list_reads_materialized_codex_cycle_usage_from_quota_windows() {
|
||||
const RESET_AT: u64 = 4_102_444_800;
|
||||
|
||||
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(
|
||||
true,
|
||||
@@ -898,21 +848,21 @@ async fn gateway_pool_list_adds_codex_cycle_usage_to_quota_windows() {
|
||||
usage_key.total_tokens = 999;
|
||||
usage_key.total_cost_usd = 9.99;
|
||||
usage_key.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"version": 2,
|
||||
"provider_type": "codex",
|
||||
"code": "ok",
|
||||
"label": serde_json::Value::Null,
|
||||
"reason": serde_json::Value::Null,
|
||||
"freshness": "fresh",
|
||||
"source": "response_headers",
|
||||
"observed_at": RESET_AT,
|
||||
"exhausted": false,
|
||||
"usage_ratio": 0.0,
|
||||
"updated_at": RESET_AT,
|
||||
"reset_seconds": serde_json::Value::Null,
|
||||
"plan_type": "plus",
|
||||
"windows": [
|
||||
"quota": {
|
||||
"version": 2,
|
||||
"provider_type": "codex",
|
||||
"code": "ok",
|
||||
"label": serde_json::Value::Null,
|
||||
"reason": serde_json::Value::Null,
|
||||
"freshness": "fresh",
|
||||
"source": "response_headers",
|
||||
"observed_at": RESET_AT,
|
||||
"exhausted": false,
|
||||
"usage_ratio": 0.0,
|
||||
"updated_at": RESET_AT,
|
||||
"reset_seconds": serde_json::Value::Null,
|
||||
"plan_type": "plus",
|
||||
"windows": [
|
||||
{
|
||||
"code": "weekly",
|
||||
"label": "周",
|
||||
@@ -922,7 +872,12 @@ async fn gateway_pool_list_adds_codex_cycle_usage_to_quota_windows() {
|
||||
"remaining_ratio": 1.0,
|
||||
"reset_at": RESET_AT,
|
||||
"reset_seconds": 604_800,
|
||||
"window_minutes": 10_080
|
||||
"window_minutes": 10_080,
|
||||
"usage": {
|
||||
"request_count": 3,
|
||||
"total_tokens": 375,
|
||||
"total_cost_usd": "0.60000000"
|
||||
}
|
||||
},
|
||||
{
|
||||
"code": "5h",
|
||||
@@ -933,7 +888,12 @@ async fn gateway_pool_list_adds_codex_cycle_usage_to_quota_windows() {
|
||||
"remaining_ratio": 1.0,
|
||||
"reset_at": RESET_AT,
|
||||
"reset_seconds": 18_000,
|
||||
"window_minutes": 300
|
||||
"window_minutes": 300,
|
||||
"usage": {
|
||||
"request_count": 2,
|
||||
"total_tokens": 225,
|
||||
"total_cost_usd": "0.30000000"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -947,7 +907,37 @@ async fn gateway_pool_list_adds_codex_cycle_usage_to_quota_windows() {
|
||||
);
|
||||
zero_key.name = "codex zero usage".to_string();
|
||||
zero_key.auth_type = "oauth".to_string();
|
||||
zero_key.status_snapshot = usage_key.status_snapshot.clone();
|
||||
zero_key.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"version": 2,
|
||||
"provider_type": "codex",
|
||||
"code": "ok",
|
||||
"windows": [
|
||||
{
|
||||
"code": "weekly",
|
||||
"label": "周",
|
||||
"reset_at": RESET_AT,
|
||||
"window_minutes": 10_080,
|
||||
"usage": {
|
||||
"request_count": 0,
|
||||
"total_tokens": 0,
|
||||
"total_cost_usd": "0.00000000"
|
||||
}
|
||||
},
|
||||
{
|
||||
"code": "5h",
|
||||
"label": "5H",
|
||||
"reset_at": RESET_AT,
|
||||
"window_minutes": 300,
|
||||
"usage": {
|
||||
"request_count": 0,
|
||||
"total_tokens": 0,
|
||||
"total_cost_usd": "0.00000000"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}));
|
||||
|
||||
let mut invalid_key = sample_key(
|
||||
"key-codex-invalid",
|
||||
@@ -984,44 +974,11 @@ async fn gateway_pool_list_adds_codex_cycle_usage_to_quota_windows() {
|
||||
Vec::new(),
|
||||
vec![usage_key, zero_key, invalid_key],
|
||||
));
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||
sample_pool_usage_row(
|
||||
"req-5h-a",
|
||||
"key-codex-cycle",
|
||||
RESET_AT as i64 - 60,
|
||||
100,
|
||||
0.10,
|
||||
),
|
||||
sample_pool_usage_row(
|
||||
"req-5h-b",
|
||||
"key-codex-cycle",
|
||||
RESET_AT as i64 - 17_999,
|
||||
125,
|
||||
0.20,
|
||||
),
|
||||
sample_pool_usage_row(
|
||||
"req-weekly-only",
|
||||
"key-codex-cycle",
|
||||
RESET_AT as i64 - 18_001,
|
||||
150,
|
||||
0.30,
|
||||
),
|
||||
sample_pool_usage_row(
|
||||
"req-before-weekly",
|
||||
"key-codex-cycle",
|
||||
RESET_AT as i64 - 604_801,
|
||||
200,
|
||||
0.40,
|
||||
),
|
||||
]));
|
||||
let state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_and_usage_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
usage_repository,
|
||||
),
|
||||
);
|
||||
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
));
|
||||
|
||||
let response = local_admin_pool_response(
|
||||
&state,
|
||||
@@ -1302,11 +1259,11 @@ async fn gateway_handles_admin_pool_list_keys_with_quota_compatibility_fields()
|
||||
"oauth": {
|
||||
"code": "expired",
|
||||
"label": "已过期",
|
||||
"reason": "Token 已过期,请重新授权",
|
||||
"reason": "Access Token 已过期,等待自动续期",
|
||||
"expires_at": 1775556730u64,
|
||||
"invalid_at": null,
|
||||
"source": "expires_at",
|
||||
"requires_reauth": true,
|
||||
"requires_reauth": false,
|
||||
"expiring_soon": false
|
||||
},
|
||||
"account": {
|
||||
@@ -1441,11 +1398,11 @@ async fn gateway_includes_pool_quota_and_compat_fields_in_list_keys_response() {
|
||||
"oauth": {
|
||||
"code": "expired",
|
||||
"label": "已过期",
|
||||
"reason": "Token 已过期,请重新授权",
|
||||
"reason": "Access Token 已过期,等待自动续期",
|
||||
"expires_at": 1_775_556_730u64,
|
||||
"invalid_at": serde_json::Value::Null,
|
||||
"source": "expires_at",
|
||||
"requires_reauth": true,
|
||||
"requires_reauth": false,
|
||||
"expiring_soon": false
|
||||
},
|
||||
"account": {
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "aether-proxy"
|
||||
version = "0.3.3"
|
||||
version = "0.3.5"
|
||||
edition = "2021"
|
||||
description = "Tunnel proxy for Aether"
|
||||
|
||||
|
||||
@@ -395,17 +395,38 @@ pub fn admin_pool_is_oauth_invalid(key: &StoredProviderCatalogKey, now_unix_secs
|
||||
if key.auth_type.trim() != "oauth" {
|
||||
return false;
|
||||
}
|
||||
if key
|
||||
.oauth_invalid_reason
|
||||
.as_deref()
|
||||
.is_some_and(|value| !value.trim().is_empty())
|
||||
{
|
||||
return true;
|
||||
if let Some(reason) = key.oauth_invalid_reason.as_deref().map(str::trim) {
|
||||
let account_state = provider_status::resolve_pool_account_state(
|
||||
None,
|
||||
key.upstream_metadata.as_ref(),
|
||||
Some(reason),
|
||||
);
|
||||
if account_state.blocked && !account_state.recoverable {
|
||||
return true;
|
||||
}
|
||||
if admin_pool_reason_has_tag(reason, "[REFRESH_FAILED]") {
|
||||
return key
|
||||
.expires_at_unix_secs
|
||||
.is_none_or(|value| value == 0 || value <= now_unix_secs);
|
||||
}
|
||||
if admin_pool_reason_has_tag(reason, "[REQUEST_FAILED]") {
|
||||
return false;
|
||||
}
|
||||
if !reason.is_empty() {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
key.expires_at_unix_secs
|
||||
.is_some_and(|value| value > 0 && value <= now_unix_secs)
|
||||
}
|
||||
|
||||
fn admin_pool_reason_has_tag(reason: &str, tag: &str) -> bool {
|
||||
reason
|
||||
.lines()
|
||||
.map(str::trim)
|
||||
.any(|line| line.starts_with(tag))
|
||||
}
|
||||
|
||||
pub fn admin_pool_matches_quick_selector(
|
||||
key: &StoredProviderCatalogKey,
|
||||
selector: &str,
|
||||
|
||||
@@ -396,7 +396,8 @@ fn codex_merge_invalid_reason(current: &str, candidate_reason: &str) -> String {
|
||||
return current.to_string();
|
||||
}
|
||||
if current.starts_with(OAUTH_EXPIRED_PREFIX)
|
||||
&& candidate_reason.starts_with(OAUTH_REQUEST_FAILED_PREFIX)
|
||||
&& (candidate_reason.starts_with(OAUTH_REQUEST_FAILED_PREFIX)
|
||||
|| candidate_reason.starts_with(OAUTH_REFRESH_FAILED_PREFIX))
|
||||
{
|
||||
return current.to_string();
|
||||
}
|
||||
@@ -872,9 +873,11 @@ pub fn parse_chatgpt_web_conversation_init_response(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
codex_runtime_invalid_reason, parse_chatgpt_web_conversation_init_response,
|
||||
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX,
|
||||
codex_build_invalid_state, codex_runtime_invalid_reason,
|
||||
parse_chatgpt_web_conversation_init_response, OAUTH_ACCOUNT_BLOCK_PREFIX,
|
||||
OAUTH_EXPIRED_PREFIX, OAUTH_REFRESH_FAILED_PREFIX, OAUTH_REQUEST_FAILED_PREFIX,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
@@ -900,6 +903,73 @@ mod tests {
|
||||
assert_eq!(codex_runtime_invalid_reason(403, Some("forbidden")), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_invalid_state_keeps_oauth_expired_over_refresh_failure() {
|
||||
let mut key = StoredProviderCatalogKey::new(
|
||||
"key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"key-1".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build");
|
||||
key.oauth_invalid_at_unix_secs = Some(100);
|
||||
key.oauth_invalid_reason = Some(format!("{OAUTH_EXPIRED_PREFIX}session expired"));
|
||||
|
||||
assert_eq!(
|
||||
codex_build_invalid_state(
|
||||
&key,
|
||||
format!("{OAUTH_REFRESH_FAILED_PREFIX}Token 续期失败"),
|
||||
200,
|
||||
),
|
||||
(
|
||||
Some(100),
|
||||
Some(format!("{OAUTH_EXPIRED_PREFIX}session expired"))
|
||||
)
|
||||
);
|
||||
assert_eq!(
|
||||
codex_build_invalid_state(
|
||||
&key,
|
||||
format!("{OAUTH_REQUEST_FAILED_PREFIX}账号状态检查失败"),
|
||||
200,
|
||||
),
|
||||
(
|
||||
Some(100),
|
||||
Some(format!("{OAUTH_EXPIRED_PREFIX}session expired"))
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_invalid_state_allows_account_block_to_override_oauth_expired() {
|
||||
let mut key = StoredProviderCatalogKey::new(
|
||||
"key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"key-1".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build");
|
||||
key.oauth_invalid_at_unix_secs = Some(100);
|
||||
key.oauth_invalid_reason = Some(format!("{OAUTH_EXPIRED_PREFIX}session expired"));
|
||||
|
||||
assert_eq!(
|
||||
codex_build_invalid_state(
|
||||
&key,
|
||||
format!("{OAUTH_ACCOUNT_BLOCK_PREFIX}account has been deactivated"),
|
||||
200,
|
||||
),
|
||||
(
|
||||
Some(200),
|
||||
Some(format!(
|
||||
"{OAUTH_ACCOUNT_BLOCK_PREFIX}account has been deactivated"
|
||||
))
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_chatgpt_web_image_quota_from_conversation_init() {
|
||||
let parsed = parse_chatgpt_web_conversation_init_response(
|
||||
|
||||
@@ -1,11 +1,6 @@
|
||||
use serde_json::Value;
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
const OAUTH_ACCOUNT_BLOCK_PREFIX: &str = "[ACCOUNT_BLOCK] ";
|
||||
const OAUTH_REFRESH_FAILED_PREFIX: &str = "[REFRESH_FAILED] ";
|
||||
const OAUTH_EXPIRED_PREFIX: &str = "[OAUTH_EXPIRED] ";
|
||||
const OAUTH_REQUEST_FAILED_PREFIX: &str = "[REQUEST_FAILED] ";
|
||||
|
||||
const ACCOUNT_BLOCK_REASON_KEYWORDS: &[&str] = &[
|
||||
"suspended",
|
||||
"banned",
|
||||
@@ -326,55 +321,55 @@ fn resolve_from_metadata(
|
||||
fn resolve_from_oauth_invalid_reason(reason: Option<&str>) -> Option<PoolAccountState> {
|
||||
let text = clean_text(reason)?;
|
||||
|
||||
if text.starts_with(OAUTH_ACCOUNT_BLOCK_PREFIX) {
|
||||
let cleaned = clean_text(Some(text.trim_start_matches(OAUTH_ACCOUNT_BLOCK_PREFIX)))
|
||||
.unwrap_or_else(|| "账号异常".to_string());
|
||||
let (code, label) = classify_block_reason(&cleaned);
|
||||
if let Some(cleaned) = tagged_reason(&text, "ACCOUNT_BLOCK") {
|
||||
let reason = if cleaned.is_empty() {
|
||||
"账号异常".to_string()
|
||||
} else {
|
||||
cleaned
|
||||
};
|
||||
let (code, label) = classify_block_reason(&reason);
|
||||
return Some(PoolAccountState {
|
||||
blocked: true,
|
||||
code: Some(code.to_string()),
|
||||
label: Some(label.to_string()),
|
||||
reason: Some(cleaned),
|
||||
reason: Some(reason),
|
||||
source: Some("oauth_invalid".to_string()),
|
||||
recoverable: false,
|
||||
});
|
||||
}
|
||||
if text.starts_with(OAUTH_EXPIRED_PREFIX) {
|
||||
let cleaned = clean_text(Some(text.trim_start_matches(OAUTH_EXPIRED_PREFIX)))
|
||||
.unwrap_or_else(|| "OAuth Token 已过期且无法续期".to_string());
|
||||
if let Some(cleaned) = tagged_reason(&text, "OAUTH_EXPIRED") {
|
||||
let reason = if cleaned.is_empty() {
|
||||
"OAuth Token 已过期且无法续期".to_string()
|
||||
} else {
|
||||
cleaned
|
||||
};
|
||||
return Some(PoolAccountState {
|
||||
blocked: true,
|
||||
code: Some("oauth_token_invalid".to_string()),
|
||||
label: Some("Token 失效".to_string()),
|
||||
reason: Some(cleaned),
|
||||
reason: Some(reason),
|
||||
source: Some("oauth_invalid".to_string()),
|
||||
recoverable: false,
|
||||
});
|
||||
}
|
||||
if text.starts_with(OAUTH_REFRESH_FAILED_PREFIX) {
|
||||
let cleaned = clean_text(Some(text.trim_start_matches(OAUTH_REFRESH_FAILED_PREFIX)))
|
||||
.unwrap_or_else(|| "OAuth Token 续期失败".to_string());
|
||||
return Some(PoolAccountState {
|
||||
blocked: true,
|
||||
code: Some("oauth_token_invalid".to_string()),
|
||||
label: Some("Token 失效".to_string()),
|
||||
reason: Some(cleaned),
|
||||
source: Some("oauth_refresh".to_string()),
|
||||
recoverable: false,
|
||||
});
|
||||
}
|
||||
if text.starts_with(OAUTH_REQUEST_FAILED_PREFIX) {
|
||||
let cleaned = clean_text(Some(text.trim_start_matches(OAUTH_REQUEST_FAILED_PREFIX)))
|
||||
.unwrap_or_else(|| "账号状态检查失败".to_string());
|
||||
if let Some(cleaned) = tagged_reason(&text, "REQUEST_FAILED") {
|
||||
let reason = if cleaned.is_empty() {
|
||||
"账号状态检查失败".to_string()
|
||||
} else {
|
||||
cleaned
|
||||
};
|
||||
return Some(PoolAccountState {
|
||||
blocked: false,
|
||||
code: Some("oauth_request_failed".to_string()),
|
||||
label: Some("请求失败".to_string()),
|
||||
reason: Some(cleaned),
|
||||
reason: Some(reason),
|
||||
source: Some("oauth_request".to_string()),
|
||||
recoverable: true,
|
||||
});
|
||||
}
|
||||
if tagged_reason(&text, "REFRESH_FAILED").is_some() {
|
||||
return None;
|
||||
}
|
||||
if text.starts_with('[') {
|
||||
return None;
|
||||
}
|
||||
@@ -462,20 +457,8 @@ pub fn resolve_account_status_snapshot(
|
||||
};
|
||||
}
|
||||
|
||||
if let Some(cleaned) = tagged_reason(&text, "REFRESH_FAILED") {
|
||||
let reason = if cleaned.is_empty() {
|
||||
"OAuth Token 续期失败".to_string()
|
||||
} else {
|
||||
cleaned
|
||||
};
|
||||
return AccountStatusSnapshot {
|
||||
code: "oauth_token_invalid".to_string(),
|
||||
label: Some("Token 失效".to_string()),
|
||||
reason: Some(reason),
|
||||
blocked: true,
|
||||
source: Some("oauth_refresh".to_string()),
|
||||
recoverable: false,
|
||||
};
|
||||
if tagged_reason(&text, "REFRESH_FAILED").is_some() {
|
||||
return AccountStatusSnapshot::default();
|
||||
}
|
||||
|
||||
if text.starts_with('[') {
|
||||
@@ -578,31 +561,30 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_refresh_failed_as_token_invalid_pool_state() {
|
||||
fn ignores_refresh_failed_as_pool_account_block() {
|
||||
let state = resolve_pool_account_state(
|
||||
Some("codex"),
|
||||
None,
|
||||
Some("[REFRESH_FAILED] Token 续期失败 (401): refresh_token 已失效"),
|
||||
);
|
||||
|
||||
assert!(state.blocked);
|
||||
assert!(!state.recoverable);
|
||||
assert_eq!(state.code.as_deref(), Some("oauth_token_invalid"));
|
||||
assert_eq!(state.label.as_deref(), Some("Token 失效"));
|
||||
assert!(!state.blocked);
|
||||
assert_eq!(state.code.as_deref(), None);
|
||||
assert_eq!(state.label.as_deref(), None);
|
||||
assert!(!should_auto_remove_account_state(&state));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn account_snapshot_marks_refresh_failed_as_token_invalid() {
|
||||
fn account_snapshot_ignores_refresh_failed_as_account_block() {
|
||||
let snapshot = resolve_account_status_snapshot(
|
||||
Some("codex"),
|
||||
None,
|
||||
Some("[REFRESH_FAILED] Token 续期失败 (401): refresh_token 已失效"),
|
||||
);
|
||||
|
||||
assert_eq!(snapshot.code, "oauth_token_invalid");
|
||||
assert_eq!(snapshot.label.as_deref(), Some("Token 失效"));
|
||||
assert!(snapshot.blocked);
|
||||
assert_eq!(snapshot.code, "ok");
|
||||
assert_eq!(snapshot.label.as_deref(), None);
|
||||
assert!(!snapshot.blocked);
|
||||
assert!(!snapshot.recoverable);
|
||||
}
|
||||
|
||||
|
||||
-13
@@ -1,19 +1,6 @@
|
||||
ALTER TABLE IF EXISTS public.usage
|
||||
ADD COLUMN IF NOT EXISTS upstream_is_stream boolean;
|
||||
|
||||
UPDATE public.usage
|
||||
SET upstream_is_stream = COALESCE(
|
||||
CASE
|
||||
WHEN (request_metadata->>'upstream_is_stream') IN ('true', 'false')
|
||||
THEN (request_metadata->>'upstream_is_stream')::boolean
|
||||
ELSE NULL
|
||||
END,
|
||||
COALESCE(is_stream, FALSE)
|
||||
)
|
||||
WHERE upstream_is_stream IS NULL;
|
||||
|
||||
ANALYZE public.usage;
|
||||
|
||||
COMMENT ON COLUMN public.usage.upstream_is_stream IS
|
||||
'Resolved upstream stream mode from request_metadata.upstream_is_stream, falling back to is_stream for legacy rows.';
|
||||
|
||||
|
||||
@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use super::{
|
||||
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
|
||||
@@ -88,6 +89,8 @@ impl InMemoryProviderCatalogReadRepository {
|
||||
{
|
||||
key.last_used_at_unix_secs = recomputed_last_used_at_unix_secs;
|
||||
}
|
||||
|
||||
apply_codex_window_usage_stats_delta(&mut key.status_snapshot, delta);
|
||||
}
|
||||
|
||||
pub(crate) fn rebuild_usage_stats(
|
||||
@@ -156,6 +159,156 @@ fn apply_f64_delta(current: f64, delta: f64) -> f64 {
|
||||
}
|
||||
}
|
||||
|
||||
fn json_u64(value: Option<&Value>) -> Option<u64> {
|
||||
value.and_then(|value| {
|
||||
value.as_u64().or_else(|| {
|
||||
value
|
||||
.as_str()
|
||||
.and_then(|text| text.trim().parse::<u64>().ok())
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn json_i64(value: Option<&Value>) -> Option<i64> {
|
||||
value.and_then(|value| {
|
||||
value.as_i64().or_else(|| {
|
||||
value
|
||||
.as_str()
|
||||
.and_then(|text| text.trim().parse::<i64>().ok())
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn json_f64(value: Option<&Value>) -> Option<f64> {
|
||||
value.and_then(|value| {
|
||||
value.as_f64().or_else(|| {
|
||||
value
|
||||
.as_str()
|
||||
.and_then(|text| text.trim().parse::<f64>().ok())
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn apply_i64_delta_to_json_u64(current: u64, delta: i64) -> u64 {
|
||||
if delta >= 0 {
|
||||
current.saturating_add(delta as u64)
|
||||
} else {
|
||||
current.saturating_sub(delta.unsigned_abs())
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_f64_delta_to_json_cost(current: f64, delta: f64) -> f64 {
|
||||
let current = if current.is_finite() { current } else { 0.0 };
|
||||
let delta = if delta.is_finite() { delta } else { 0.0 };
|
||||
(current + delta).max(0.0)
|
||||
}
|
||||
|
||||
fn codex_window_matches_usage_time(window: &Map<String, Value>, usage_created_at: u64) -> bool {
|
||||
let code = window
|
||||
.get("code")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
if !code.eq_ignore_ascii_case("5h") && !code.eq_ignore_ascii_case("weekly") {
|
||||
return false;
|
||||
}
|
||||
|
||||
let Some(reset_at) = json_u64(window.get("reset_at")) else {
|
||||
return false;
|
||||
};
|
||||
let Some(window_minutes) = json_u64(window.get("window_minutes")) else {
|
||||
return false;
|
||||
};
|
||||
let Some(window_seconds) = window_minutes.checked_mul(60) else {
|
||||
return false;
|
||||
};
|
||||
let Some(window_start) = reset_at.checked_sub(window_seconds) else {
|
||||
return false;
|
||||
};
|
||||
let usage_reset_at = json_u64(window.get("usage_reset_at")).unwrap_or(0);
|
||||
let start = window_start.max(usage_reset_at);
|
||||
usage_created_at >= start && usage_created_at < reset_at
|
||||
}
|
||||
|
||||
fn apply_codex_window_usage_stats_delta(
|
||||
status_snapshot: &mut Option<Value>,
|
||||
delta: &ProviderApiKeyUsageDelta,
|
||||
) {
|
||||
let Some(usage_created_at) = delta.usage_created_at_unix_secs else {
|
||||
return;
|
||||
};
|
||||
if delta.request_count == 0 && delta.total_tokens == 0 && delta.total_cost_usd == 0.0 {
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(quota) = status_snapshot
|
||||
.as_mut()
|
||||
.and_then(Value::as_object_mut)
|
||||
.and_then(|snapshot| snapshot.get_mut("quota"))
|
||||
.and_then(Value::as_object_mut)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let quota_provider_type = quota
|
||||
.get("provider_type")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
if !quota_provider_type.eq_ignore_ascii_case("codex") {
|
||||
return;
|
||||
}
|
||||
let Some(windows) = quota.get_mut("windows").and_then(Value::as_array_mut) else {
|
||||
return;
|
||||
};
|
||||
|
||||
for window in windows.iter_mut().filter_map(Value::as_object_mut) {
|
||||
if !codex_window_matches_usage_time(window, usage_created_at) {
|
||||
continue;
|
||||
}
|
||||
|
||||
let usage = window
|
||||
.entry("usage".to_string())
|
||||
.or_insert_with(|| json!({}))
|
||||
.as_object_mut();
|
||||
let Some(usage) = usage else {
|
||||
window.insert("usage".to_string(), json!({}));
|
||||
let Some(usage) = window.get_mut("usage").and_then(Value::as_object_mut) else {
|
||||
continue;
|
||||
};
|
||||
let request_count = apply_i64_delta_to_json_u64(0, delta.request_count);
|
||||
let total_tokens = apply_i64_delta_to_json_u64(0, delta.total_tokens);
|
||||
let total_cost_usd = apply_f64_delta_to_json_cost(0.0, delta.total_cost_usd);
|
||||
usage.insert("request_count".to_string(), json!(request_count));
|
||||
usage.insert("total_tokens".to_string(), json!(total_tokens));
|
||||
usage.insert(
|
||||
"total_cost_usd".to_string(),
|
||||
json!(format!("{total_cost_usd:.8}")),
|
||||
);
|
||||
continue;
|
||||
};
|
||||
|
||||
let request_count = apply_i64_delta_to_json_u64(
|
||||
json_i64(usage.get("request_count")).unwrap_or(0).max(0) as u64,
|
||||
delta.request_count,
|
||||
);
|
||||
let total_tokens = apply_i64_delta_to_json_u64(
|
||||
json_i64(usage.get("total_tokens")).unwrap_or(0).max(0) as u64,
|
||||
delta.total_tokens,
|
||||
);
|
||||
let total_cost_usd = apply_f64_delta_to_json_cost(
|
||||
json_f64(usage.get("total_cost_usd")).unwrap_or(0.0),
|
||||
delta.total_cost_usd,
|
||||
);
|
||||
|
||||
usage.insert("request_count".to_string(), json!(request_count));
|
||||
usage.insert("total_tokens".to_string(), json!(total_tokens));
|
||||
usage.insert(
|
||||
"total_cost_usd".to_string(),
|
||||
json!(format!("{total_cost_usd:.8}")),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderCatalogReadRepository for InMemoryProviderCatalogReadRepository {
|
||||
async fn list_providers(
|
||||
@@ -573,6 +726,8 @@ mod tests {
|
||||
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use crate::repository::usage::ProviderApiKeyUsageDelta;
|
||||
use serde_json::json;
|
||||
|
||||
fn sample_provider(id: &str) -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
@@ -716,6 +871,73 @@ mod tests {
|
||||
assert_eq!(stored[0].expires_at_unix_secs, Some(4_102_444_800));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn materializes_codex_window_usage_stats_delta_in_memory() {
|
||||
let mut key = sample_key("key-1", "provider-1");
|
||||
key.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"provider_type": "codex",
|
||||
"windows": [
|
||||
{
|
||||
"code": "5h",
|
||||
"reset_at": 120_000u64,
|
||||
"window_minutes": 300u64,
|
||||
"usage": {
|
||||
"request_count": 1,
|
||||
"total_tokens": 10,
|
||||
"total_cost_usd": "0.10000000"
|
||||
}
|
||||
},
|
||||
{
|
||||
"code": "weekly",
|
||||
"reset_at": 700_000u64,
|
||||
"window_minutes": 10_080u64
|
||||
}
|
||||
]
|
||||
}
|
||||
}));
|
||||
let repository = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1")],
|
||||
vec![],
|
||||
vec![key],
|
||||
);
|
||||
|
||||
repository.apply_usage_stats_delta(
|
||||
"key-1",
|
||||
&ProviderApiKeyUsageDelta {
|
||||
request_count: 2,
|
||||
total_tokens: 25,
|
||||
total_cost_usd: 0.25,
|
||||
usage_created_at_unix_secs: Some(110_000),
|
||||
..ProviderApiKeyUsageDelta::default()
|
||||
},
|
||||
None,
|
||||
);
|
||||
|
||||
let stored = repository
|
||||
.list_keys_by_ids(&["key-1".to_string()])
|
||||
.await
|
||||
.expect("keys should read");
|
||||
let windows = stored[0].status_snapshot.as_ref().expect("snapshot")["quota"]["windows"]
|
||||
.as_array()
|
||||
.expect("windows");
|
||||
let five_h = windows
|
||||
.iter()
|
||||
.find(|window| window["code"] == json!("5h"))
|
||||
.expect("5h window should exist");
|
||||
let weekly = windows
|
||||
.iter()
|
||||
.find(|window| window["code"] == json!("weekly"))
|
||||
.expect("weekly window should exist");
|
||||
|
||||
assert_eq!(five_h["usage"]["request_count"], json!(3));
|
||||
assert_eq!(five_h["usage"]["total_tokens"], json!(35));
|
||||
assert_eq!(five_h["usage"]["total_cost_usd"], json!("0.35000000"));
|
||||
assert_eq!(weekly["usage"]["request_count"], json!(2));
|
||||
assert_eq!(weekly["usage"]["total_tokens"], json!(25));
|
||||
assert_eq!(weekly["usage"]["total_cost_usd"], json!("0.25000000"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn paginates_provider_keys_with_search_and_active_filter() {
|
||||
let mut alpha = sample_key("key-1", "provider-1");
|
||||
|
||||
@@ -501,6 +501,7 @@ pub(crate) struct ProviderApiKeyUsageContribution {
|
||||
pub total_cost_usd: f64,
|
||||
pub total_response_time_ms: i64,
|
||||
pub last_used_at_unix_secs: Option<u64>,
|
||||
pub usage_created_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Default)]
|
||||
@@ -513,6 +514,7 @@ pub(crate) struct ProviderApiKeyUsageDelta {
|
||||
pub total_response_time_ms: i64,
|
||||
pub candidate_last_used_at_unix_secs: Option<u64>,
|
||||
pub removed_last_used_at_unix_secs: Option<u64>,
|
||||
pub usage_created_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl ProviderApiKeyUsageDelta {
|
||||
@@ -529,6 +531,7 @@ impl ProviderApiKeyUsageDelta {
|
||||
total_response_time_ms: after.total_response_time_ms - before.total_response_time_ms,
|
||||
candidate_last_used_at_unix_secs: after.last_used_at_unix_secs,
|
||||
removed_last_used_at_unix_secs: None,
|
||||
usage_created_at_unix_secs: after.usage_created_at_unix_secs,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -542,6 +545,7 @@ impl ProviderApiKeyUsageDelta {
|
||||
total_response_time_ms: after.total_response_time_ms,
|
||||
candidate_last_used_at_unix_secs: after.last_used_at_unix_secs,
|
||||
removed_last_used_at_unix_secs: None,
|
||||
usage_created_at_unix_secs: after.usage_created_at_unix_secs,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -555,6 +559,7 @@ impl ProviderApiKeyUsageDelta {
|
||||
total_response_time_ms: -before.total_response_time_ms,
|
||||
candidate_last_used_at_unix_secs: None,
|
||||
removed_last_used_at_unix_secs: before.last_used_at_unix_secs,
|
||||
usage_created_at_unix_secs: before.usage_created_at_unix_secs,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -664,6 +669,7 @@ pub(crate) fn provider_api_key_usage_contribution(
|
||||
0
|
||||
},
|
||||
last_used_at_unix_secs: Some(usage.created_at_unix_ms),
|
||||
usage_created_at_unix_secs: Some(usage.created_at_unix_ms),
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -1333,15 +1333,20 @@ const REBUILD_API_KEY_USAGE_STATS_SQL: &str =
|
||||
const APPLY_PROVIDER_API_KEY_USAGE_DELTA_SQL: &str =
|
||||
include_str!("queries/apply_provider_api_key_usage_delta_sql.sql");
|
||||
|
||||
const APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL: &str =
|
||||
include_str!("queries/apply_provider_api_key_codex_window_usage_delta_sql.sql");
|
||||
|
||||
const RESET_PROVIDER_API_KEY_USAGE_STATS_SQL: &str =
|
||||
include_str!("queries/reset_provider_api_key_usage_stats_sql.sql");
|
||||
|
||||
const REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL: &str =
|
||||
include_str!("queries/rebuild_provider_api_key_usage_stats_sql.sql");
|
||||
|
||||
const REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL: &str =
|
||||
include_str!("queries/rebuild_provider_api_key_codex_window_usage_stats_sql.sql");
|
||||
|
||||
const LIST_USAGE_AUDITS_PREFIX: &str = include_str!("queries/list_usage_audits_prefix.sql");
|
||||
const USAGE_RESERVED_PROVIDER_LABELS_FILTER_SQL: &str =
|
||||
" AND BTRIM(COALESCE(\"usage\".provider_name, '')) <> '' AND lower(BTRIM(COALESCE(\"usage\".provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')";
|
||||
const USAGE_RESERVED_PROVIDER_LABELS_FILTER_SQL: &str = " AND BTRIM(COALESCE(\"usage\".provider_name, '')) <> '' AND lower(BTRIM(COALESCE(\"usage\".provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')";
|
||||
|
||||
struct UsageAuditAggregationSqlFragments {
|
||||
filtered_extra_where: &'static str,
|
||||
@@ -6395,33 +6400,38 @@ WHERE stats_daily_api_key.date >=
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let (table_name, group_column, display_name_expr, avg_response_time_expr, success_count_expr) =
|
||||
match group_by {
|
||||
UsageAuditAggregationGroupBy::Model => (
|
||||
"stats_user_daily_model",
|
||||
"model",
|
||||
"NULL::varchar",
|
||||
"NULL::DOUBLE PRECISION",
|
||||
"NULL::BIGINT",
|
||||
),
|
||||
UsageAuditAggregationGroupBy::Provider => (
|
||||
"stats_user_daily_provider",
|
||||
"provider_name",
|
||||
"provider_name",
|
||||
"CASE WHEN COALESCE(SUM(response_time_samples), 0) > 0 THEN COALESCE(SUM(response_time_sum_ms), 0) / COALESCE(SUM(response_time_samples), 0) ELSE NULL END",
|
||||
"COALESCE(SUM(success_requests), 0)::BIGINT",
|
||||
),
|
||||
UsageAuditAggregationGroupBy::ApiFormat => (
|
||||
"stats_user_daily_api_format",
|
||||
"api_format",
|
||||
"NULL::varchar",
|
||||
"CASE WHEN COALESCE(SUM(response_time_samples), 0) > 0 THEN COALESCE(SUM(response_time_sum_ms), 0) / COALESCE(SUM(response_time_samples), 0) ELSE NULL END",
|
||||
"NULL::BIGINT",
|
||||
),
|
||||
UsageAuditAggregationGroupBy::User => {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
};
|
||||
let (
|
||||
table_name,
|
||||
group_column,
|
||||
display_name_expr,
|
||||
avg_response_time_expr,
|
||||
success_count_expr,
|
||||
) = match group_by {
|
||||
UsageAuditAggregationGroupBy::Model => (
|
||||
"stats_user_daily_model",
|
||||
"model",
|
||||
"NULL::varchar",
|
||||
"NULL::DOUBLE PRECISION",
|
||||
"NULL::BIGINT",
|
||||
),
|
||||
UsageAuditAggregationGroupBy::Provider => (
|
||||
"stats_user_daily_provider",
|
||||
"provider_name",
|
||||
"provider_name",
|
||||
"CASE WHEN COALESCE(SUM(response_time_samples), 0) > 0 THEN COALESCE(SUM(response_time_sum_ms), 0) / COALESCE(SUM(response_time_samples), 0) ELSE NULL END",
|
||||
"COALESCE(SUM(success_requests), 0)::BIGINT",
|
||||
),
|
||||
UsageAuditAggregationGroupBy::ApiFormat => (
|
||||
"stats_user_daily_api_format",
|
||||
"api_format",
|
||||
"NULL::varchar",
|
||||
"CASE WHEN COALESCE(SUM(response_time_samples), 0) > 0 THEN COALESCE(SUM(response_time_sum_ms), 0) / COALESCE(SUM(response_time_samples), 0) ELSE NULL END",
|
||||
"NULL::BIGINT",
|
||||
),
|
||||
UsageAuditAggregationGroupBy::User => {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
};
|
||||
|
||||
let provider_extra_where = if matches!(group_by, UsageAuditAggregationGroupBy::Provider) {
|
||||
" AND BTRIM(COALESCE(provider_name, '')) <> '' AND lower(BTRIM(COALESCE(provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')"
|
||||
@@ -7925,6 +7935,10 @@ ORDER BY "usage".user_id ASC
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
.rows_affected();
|
||||
sqlx::query(REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
Ok(rows_affected)
|
||||
}) as BoxFuture<'_, Result<u64, DataLayerError>>
|
||||
})
|
||||
@@ -8371,6 +8385,40 @@ async fn apply_provider_api_key_usage_delta_in_tx(
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
apply_provider_api_key_codex_window_usage_delta_in_tx(tx, key_id, delta).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn apply_provider_api_key_codex_window_usage_delta_in_tx(
|
||||
tx: &mut sqlx::Transaction<'_, Postgres>,
|
||||
key_id: &str,
|
||||
delta: &ProviderApiKeyUsageDelta,
|
||||
) -> Result<(), DataLayerError> {
|
||||
let Some(usage_created_at_unix_secs) = delta.usage_created_at_unix_secs else {
|
||||
return Ok(());
|
||||
};
|
||||
if delta.request_count == 0 && delta.total_tokens == 0 && delta.total_cost_usd == 0.0 {
|
||||
return Ok(());
|
||||
}
|
||||
let total_cost_usd_delta = if delta.total_cost_usd.is_finite() {
|
||||
delta.total_cost_usd
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
sqlx::query(APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL)
|
||||
.bind(key_id)
|
||||
.bind(i64::try_from(usage_created_at_unix_secs).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"provider api key window usage timestamp exceeds i64: {usage_created_at_unix_secs}"
|
||||
))
|
||||
})?)
|
||||
.bind(delta.request_count)
|
||||
.bind(delta.total_tokens)
|
||||
.bind(total_cost_usd_delta)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
+166
@@ -0,0 +1,166 @@
|
||||
WITH target_key AS (
|
||||
SELECT
|
||||
id,
|
||||
COALESCE(status_snapshot::jsonb, '{}'::jsonb) AS snapshot
|
||||
FROM provider_api_keys
|
||||
WHERE id = $1
|
||||
AND jsonb_typeof((status_snapshot::jsonb) -> 'quota' -> 'windows') = 'array'
|
||||
AND lower(BTRIM(COALESCE((status_snapshot::jsonb) -> 'quota' ->> 'provider_type', ''))) = 'codex'
|
||||
FOR UPDATE
|
||||
),
|
||||
window_items AS (
|
||||
SELECT
|
||||
target_key.id,
|
||||
window_rows.window_item,
|
||||
window_rows.window_ordinality
|
||||
FROM target_key
|
||||
CROSS JOIN LATERAL jsonb_array_elements(target_key.snapshot -> 'quota' -> 'windows')
|
||||
WITH ORDINALITY AS window_rows(window_item, window_ordinality)
|
||||
),
|
||||
parsed_windows AS (
|
||||
SELECT
|
||||
window_items.id,
|
||||
window_items.window_item,
|
||||
window_items.window_ordinality,
|
||||
lower(BTRIM(COALESCE(window_items.window_item ->> 'code', ''))) AS window_code,
|
||||
CASE
|
||||
WHEN text_values.reset_at_text ~ '^[0-9]+$'
|
||||
AND (
|
||||
length(text_values.reset_at_text) < 19
|
||||
OR (
|
||||
length(text_values.reset_at_text) = 19
|
||||
AND text_values.reset_at_text <= '9223372036854775807'
|
||||
)
|
||||
)
|
||||
THEN text_values.reset_at_text::BIGINT
|
||||
ELSE NULL
|
||||
END AS reset_at,
|
||||
CASE
|
||||
WHEN text_values.window_minutes_text ~ '^[0-9]+$'
|
||||
AND (
|
||||
length(text_values.window_minutes_text) < 19
|
||||
OR (
|
||||
length(text_values.window_minutes_text) = 19
|
||||
AND text_values.window_minutes_text <= '9223372036854775807'
|
||||
)
|
||||
)
|
||||
THEN text_values.window_minutes_text::BIGINT
|
||||
ELSE NULL
|
||||
END AS window_minutes,
|
||||
CASE
|
||||
WHEN text_values.usage_reset_at_text ~ '^[0-9]+$'
|
||||
AND (
|
||||
length(text_values.usage_reset_at_text) < 19
|
||||
OR (
|
||||
length(text_values.usage_reset_at_text) = 19
|
||||
AND text_values.usage_reset_at_text <= '9223372036854775807'
|
||||
)
|
||||
)
|
||||
THEN text_values.usage_reset_at_text::BIGINT
|
||||
ELSE NULL
|
||||
END AS usage_reset_at,
|
||||
CASE
|
||||
WHEN text_values.request_count_text ~ '^[0-9]+$'
|
||||
AND (
|
||||
length(text_values.request_count_text) < 19
|
||||
OR (
|
||||
length(text_values.request_count_text) = 19
|
||||
AND text_values.request_count_text <= '9223372036854775807'
|
||||
)
|
||||
)
|
||||
THEN text_values.request_count_text::BIGINT
|
||||
ELSE 0
|
||||
END AS current_request_count,
|
||||
CASE
|
||||
WHEN text_values.total_tokens_text ~ '^[0-9]+$'
|
||||
AND (
|
||||
length(text_values.total_tokens_text) < 19
|
||||
OR (
|
||||
length(text_values.total_tokens_text) = 19
|
||||
AND text_values.total_tokens_text <= '9223372036854775807'
|
||||
)
|
||||
)
|
||||
THEN text_values.total_tokens_text::BIGINT
|
||||
ELSE 0
|
||||
END AS current_total_tokens,
|
||||
CASE
|
||||
WHEN text_values.total_cost_usd_text ~ '^[-+]?[0-9]+([.][0-9]+)?$'
|
||||
THEN text_values.total_cost_usd_text::DOUBLE PRECISION
|
||||
ELSE 0
|
||||
END AS current_total_cost_usd
|
||||
FROM window_items
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT
|
||||
BTRIM(COALESCE(window_items.window_item ->> 'reset_at', '')) AS reset_at_text,
|
||||
BTRIM(COALESCE(window_items.window_item ->> 'window_minutes', '')) AS window_minutes_text,
|
||||
BTRIM(COALESCE(window_items.window_item ->> 'usage_reset_at', '')) AS usage_reset_at_text,
|
||||
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'request_count', '')) AS request_count_text,
|
||||
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'total_tokens', '')) AS total_tokens_text,
|
||||
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'total_cost_usd', '')) AS total_cost_usd_text
|
||||
) AS text_values
|
||||
),
|
||||
window_usage AS (
|
||||
SELECT
|
||||
parsed_windows.*,
|
||||
CASE
|
||||
WHEN parsed_windows.window_minutes BETWEEN 0 AND 153722867280912930
|
||||
THEN parsed_windows.window_minutes * 60
|
||||
ELSE NULL
|
||||
END AS window_seconds
|
||||
FROM parsed_windows
|
||||
),
|
||||
updated_windows AS (
|
||||
SELECT
|
||||
window_usage.id,
|
||||
jsonb_agg(
|
||||
CASE
|
||||
WHEN window_usage.window_code IN ('5h', 'weekly')
|
||||
AND window_usage.reset_at IS NOT NULL
|
||||
AND window_usage.window_seconds IS NOT NULL
|
||||
AND window_usage.reset_at >= window_usage.window_seconds
|
||||
AND window_usage.reset_at > $2
|
||||
AND GREATEST(
|
||||
window_usage.reset_at - window_usage.window_seconds,
|
||||
COALESCE(window_usage.usage_reset_at, 0)
|
||||
) <= $2
|
||||
THEN jsonb_set(
|
||||
window_usage.window_item,
|
||||
'{usage}',
|
||||
jsonb_build_object(
|
||||
'request_count',
|
||||
LEAST(
|
||||
GREATEST(window_usage.current_request_count::NUMERIC + $3::NUMERIC, 0),
|
||||
9223372036854775807
|
||||
)::BIGINT,
|
||||
'total_tokens',
|
||||
LEAST(
|
||||
GREATEST(window_usage.current_total_tokens::NUMERIC + $4::NUMERIC, 0),
|
||||
9223372036854775807
|
||||
)::BIGINT,
|
||||
'total_cost_usd',
|
||||
to_char(
|
||||
GREATEST(COALESCE(window_usage.current_total_cost_usd, 0) + $5, 0),
|
||||
'FM999999999999999990.00000000'
|
||||
)
|
||||
),
|
||||
true
|
||||
)
|
||||
ELSE window_usage.window_item
|
||||
END
|
||||
ORDER BY window_usage.window_ordinality
|
||||
) AS windows
|
||||
FROM window_usage
|
||||
GROUP BY window_usage.id
|
||||
)
|
||||
UPDATE provider_api_keys AS keys
|
||||
SET
|
||||
status_snapshot = jsonb_set(
|
||||
target_key.snapshot,
|
||||
'{quota,windows}',
|
||||
updated_windows.windows,
|
||||
true
|
||||
)::json,
|
||||
updated_at = NOW()
|
||||
FROM target_key
|
||||
JOIN updated_windows ON updated_windows.id = target_key.id
|
||||
WHERE keys.id = target_key.id
|
||||
+164
@@ -0,0 +1,164 @@
|
||||
WITH target_keys AS (
|
||||
SELECT
|
||||
id,
|
||||
COALESCE(status_snapshot::jsonb, '{}'::jsonb) AS snapshot
|
||||
FROM provider_api_keys
|
||||
WHERE jsonb_typeof((status_snapshot::jsonb) -> 'quota' -> 'windows') = 'array'
|
||||
AND lower(BTRIM(COALESCE((status_snapshot::jsonb) -> 'quota' ->> 'provider_type', ''))) = 'codex'
|
||||
),
|
||||
window_items AS (
|
||||
SELECT
|
||||
target_keys.id,
|
||||
window_rows.window_item,
|
||||
window_rows.window_ordinality
|
||||
FROM target_keys
|
||||
CROSS JOIN LATERAL jsonb_array_elements(target_keys.snapshot -> 'quota' -> 'windows')
|
||||
WITH ORDINALITY AS window_rows(window_item, window_ordinality)
|
||||
),
|
||||
parsed_windows AS (
|
||||
SELECT
|
||||
window_items.id,
|
||||
window_items.window_item,
|
||||
window_items.window_ordinality,
|
||||
lower(BTRIM(COALESCE(window_items.window_item ->> 'code', ''))) AS window_code,
|
||||
CASE
|
||||
WHEN text_values.reset_at_text ~ '^[0-9]+$'
|
||||
AND (
|
||||
length(text_values.reset_at_text) < 19
|
||||
OR (
|
||||
length(text_values.reset_at_text) = 19
|
||||
AND text_values.reset_at_text <= '9223372036854775807'
|
||||
)
|
||||
)
|
||||
THEN text_values.reset_at_text::BIGINT
|
||||
ELSE NULL
|
||||
END AS reset_at,
|
||||
CASE
|
||||
WHEN text_values.window_minutes_text ~ '^[0-9]+$'
|
||||
AND (
|
||||
length(text_values.window_minutes_text) < 19
|
||||
OR (
|
||||
length(text_values.window_minutes_text) = 19
|
||||
AND text_values.window_minutes_text <= '9223372036854775807'
|
||||
)
|
||||
)
|
||||
THEN text_values.window_minutes_text::BIGINT
|
||||
ELSE NULL
|
||||
END AS window_minutes,
|
||||
CASE
|
||||
WHEN text_values.usage_reset_at_text ~ '^[0-9]+$'
|
||||
AND (
|
||||
length(text_values.usage_reset_at_text) < 19
|
||||
OR (
|
||||
length(text_values.usage_reset_at_text) = 19
|
||||
AND text_values.usage_reset_at_text <= '9223372036854775807'
|
||||
)
|
||||
)
|
||||
THEN text_values.usage_reset_at_text::BIGINT
|
||||
ELSE NULL
|
||||
END AS usage_reset_at
|
||||
FROM window_items
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT
|
||||
BTRIM(COALESCE(window_items.window_item ->> 'reset_at', '')) AS reset_at_text,
|
||||
BTRIM(COALESCE(window_items.window_item ->> 'window_minutes', '')) AS window_minutes_text,
|
||||
BTRIM(COALESCE(window_items.window_item ->> 'usage_reset_at', '')) AS usage_reset_at_text
|
||||
) AS text_values
|
||||
),
|
||||
window_usage AS (
|
||||
SELECT
|
||||
parsed_windows.*,
|
||||
CASE
|
||||
WHEN parsed_windows.window_minutes BETWEEN 0 AND 153722867280912930
|
||||
THEN parsed_windows.window_minutes * 60
|
||||
ELSE NULL
|
||||
END AS window_seconds
|
||||
FROM parsed_windows
|
||||
),
|
||||
window_bounds AS (
|
||||
SELECT
|
||||
window_usage.*,
|
||||
CASE
|
||||
WHEN window_usage.window_code IN ('5h', 'weekly')
|
||||
AND window_usage.reset_at IS NOT NULL
|
||||
AND window_usage.window_seconds IS NOT NULL
|
||||
AND window_usage.reset_at >= window_usage.window_seconds
|
||||
THEN GREATEST(
|
||||
window_usage.reset_at - window_usage.window_seconds,
|
||||
COALESCE(window_usage.usage_reset_at, 0)
|
||||
)
|
||||
ELSE NULL
|
||||
END AS window_start,
|
||||
CASE
|
||||
WHEN window_usage.window_code IN ('5h', 'weekly')
|
||||
AND window_usage.reset_at IS NOT NULL
|
||||
AND window_usage.window_seconds IS NOT NULL
|
||||
AND window_usage.reset_at >= window_usage.window_seconds
|
||||
THEN window_usage.reset_at
|
||||
ELSE NULL
|
||||
END AS window_end
|
||||
FROM window_usage
|
||||
),
|
||||
aggregated AS (
|
||||
SELECT
|
||||
window_bounds.id,
|
||||
window_bounds.window_ordinality,
|
||||
COUNT("usage".id)::BIGINT AS request_count,
|
||||
COALESCE(SUM(GREATEST(COALESCE("usage".total_tokens, 0), 0)::BIGINT), 0)::BIGINT AS total_tokens,
|
||||
CAST(COALESCE(SUM(COALESCE("usage".total_cost_usd, 0)), 0) AS DOUBLE PRECISION) AS total_cost_usd
|
||||
FROM window_bounds
|
||||
LEFT JOIN usage_billing_facts AS "usage"
|
||||
ON window_bounds.window_start IS NOT NULL
|
||||
AND window_bounds.window_end IS NOT NULL
|
||||
AND "usage".provider_api_key_id = window_bounds.id
|
||||
AND "usage".created_at >= to_timestamp(window_bounds.window_start::DOUBLE PRECISION)
|
||||
AND "usage".created_at < to_timestamp(window_bounds.window_end::DOUBLE PRECISION)
|
||||
GROUP BY
|
||||
window_bounds.id,
|
||||
window_bounds.window_ordinality
|
||||
),
|
||||
updated_windows AS (
|
||||
SELECT
|
||||
window_bounds.id,
|
||||
jsonb_agg(
|
||||
CASE
|
||||
WHEN window_bounds.window_start IS NOT NULL
|
||||
AND window_bounds.window_end IS NOT NULL
|
||||
THEN jsonb_set(
|
||||
window_bounds.window_item,
|
||||
'{usage}',
|
||||
jsonb_build_object(
|
||||
'request_count',
|
||||
COALESCE(aggregated.request_count, 0),
|
||||
'total_tokens',
|
||||
COALESCE(aggregated.total_tokens, 0),
|
||||
'total_cost_usd',
|
||||
to_char(
|
||||
GREATEST(COALESCE(aggregated.total_cost_usd, 0), 0),
|
||||
'FM999999999999999990.00000000'
|
||||
)
|
||||
),
|
||||
true
|
||||
)
|
||||
ELSE window_bounds.window_item
|
||||
END
|
||||
ORDER BY window_bounds.window_ordinality
|
||||
) AS windows
|
||||
FROM window_bounds
|
||||
LEFT JOIN aggregated
|
||||
ON aggregated.id = window_bounds.id
|
||||
AND aggregated.window_ordinality = window_bounds.window_ordinality
|
||||
GROUP BY window_bounds.id
|
||||
)
|
||||
UPDATE provider_api_keys AS keys
|
||||
SET
|
||||
status_snapshot = jsonb_set(
|
||||
target_keys.snapshot,
|
||||
'{quota,windows}',
|
||||
updated_windows.windows,
|
||||
true
|
||||
)::json,
|
||||
updated_at = NOW()
|
||||
FROM target_keys
|
||||
JOIN updated_windows ON updated_windows.id = target_keys.id
|
||||
WHERE keys.id = target_keys.id
|
||||
@@ -235,18 +235,31 @@ fn usage_sql_summarizes_usage_by_provider_api_key_ids_in_database() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_sql_summarizes_provider_key_window_usage_from_billing_facts() {
|
||||
assert!(super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL.contains("UNNEST"));
|
||||
assert!(super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL
|
||||
.contains("LEFT JOIN usage_billing_facts AS \"usage\""));
|
||||
fn usage_sql_materializes_provider_key_window_usage_in_status_snapshot() {
|
||||
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL
|
||||
.contains("UPDATE provider_api_keys AS keys"));
|
||||
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("jsonb_set"));
|
||||
assert!(
|
||||
super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL.contains("created_at >= to_timestamp")
|
||||
super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("'{quota,windows}'")
|
||||
);
|
||||
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("'usage'"));
|
||||
assert!(
|
||||
super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL.contains("created_at < to_timestamp")
|
||||
!super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("usage_billing_facts")
|
||||
);
|
||||
assert!(super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL
|
||||
.contains("COUNT(\"usage\".id)::BIGINT AS request_count"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_sql_rebuilds_provider_key_window_usage_into_status_snapshot() {
|
||||
assert!(super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL
|
||||
.contains("UPDATE provider_api_keys AS keys"));
|
||||
assert!(super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL
|
||||
.contains("usage_billing_facts"));
|
||||
assert!(super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL
|
||||
.contains("provider_type', ''))) = 'codex'"));
|
||||
assert!(
|
||||
super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL.contains("'{quota,windows}'")
|
||||
);
|
||||
assert!(super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL.contains("'{usage}'"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -470,7 +483,7 @@ fn usage_sql_raw_aggregates_use_canonical_billing_facts() {
|
||||
.contains("FROM usage_billing_facts AS \"usage\""));
|
||||
assert!(super::SUMMARIZE_USAGE_TOTALS_BY_USER_IDS_SQL
|
||||
.contains("FROM usage_billing_facts AS \"usage\""));
|
||||
assert!(super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL
|
||||
assert!(!super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL
|
||||
.contains("usage_billing_facts AS \"usage\""));
|
||||
}
|
||||
|
||||
@@ -498,7 +511,10 @@ fn usage_billing_facts_projects_upstream_stream_mode() {
|
||||
assert!(migration.contains("COALESCE(usage_rows.upstream_is_stream"));
|
||||
assert!(migration.contains("COALESCE(usage_rows.is_stream, FALSE)"));
|
||||
assert!(migration.contains("ADD COLUMN IF NOT EXISTS upstream_is_stream boolean"));
|
||||
assert!(migration.contains("request_metadata->>'upstream_is_stream'"));
|
||||
assert!(
|
||||
!migration.contains("request_metadata->>'upstream_is_stream'"),
|
||||
"migration should avoid backfilling historical usage rows from request metadata"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -24,6 +24,8 @@ pub struct SchedulerRequestCandidateReportContext {
|
||||
pub body_rules: Option<Value>,
|
||||
pub proxy: Option<Value>,
|
||||
pub error_flow: Option<Value>,
|
||||
pub candidate_group_id: Option<String>,
|
||||
pub pool_key_index: Option<u32>,
|
||||
pub ranking_mode: Option<String>,
|
||||
pub priority_mode: Option<String>,
|
||||
pub ranking_index: Option<u32>,
|
||||
@@ -65,6 +67,8 @@ struct ReportCandidateExtraDataInput {
|
||||
body_rules: Option<Value>,
|
||||
proxy: Option<Value>,
|
||||
error_flow: Option<Value>,
|
||||
candidate_group_id: Option<String>,
|
||||
pool_key_index: Option<u32>,
|
||||
ranking_mode: Option<String>,
|
||||
priority_mode: Option<String>,
|
||||
ranking_index: Option<u32>,
|
||||
@@ -153,6 +157,8 @@ pub fn parse_request_candidate_report_context(
|
||||
.get("error_flow")
|
||||
.cloned()
|
||||
.filter(|value| !value.is_null()),
|
||||
candidate_group_id: string_field(report_context, "candidate_group_id"),
|
||||
pool_key_index: u32_field(report_context, "pool_key_index"),
|
||||
ranking_mode: string_field(report_context, "ranking_mode"),
|
||||
priority_mode: string_field(report_context, "priority_mode"),
|
||||
ranking_index: u32_field(report_context, "ranking_index"),
|
||||
@@ -188,6 +194,8 @@ pub fn resolve_report_request_candidate_slot(
|
||||
body_rules,
|
||||
proxy,
|
||||
error_flow,
|
||||
candidate_group_id,
|
||||
pool_key_index,
|
||||
ranking_mode,
|
||||
priority_mode,
|
||||
ranking_index,
|
||||
@@ -206,6 +214,8 @@ pub fn resolve_report_request_candidate_slot(
|
||||
body_rules,
|
||||
proxy,
|
||||
error_flow,
|
||||
candidate_group_id,
|
||||
pool_key_index,
|
||||
ranking_mode,
|
||||
priority_mode,
|
||||
ranking_index,
|
||||
@@ -323,6 +333,8 @@ pub fn build_execution_request_candidate_seed(
|
||||
body_rules: metadata.body_rules,
|
||||
proxy: metadata.proxy,
|
||||
error_flow: metadata.error_flow,
|
||||
candidate_group_id: metadata.candidate_group_id,
|
||||
pool_key_index: metadata.pool_key_index,
|
||||
ranking_mode: metadata.ranking_mode,
|
||||
priority_mode: metadata.priority_mode,
|
||||
ranking_index: metadata.ranking_index,
|
||||
@@ -426,6 +438,8 @@ pub fn build_local_request_candidate_status_record(
|
||||
body_rules: metadata.body_rules.clone(),
|
||||
proxy: metadata.proxy.clone(),
|
||||
error_flow: metadata.error_flow.clone(),
|
||||
candidate_group_id: metadata.candidate_group_id.clone(),
|
||||
pool_key_index: metadata.pool_key_index,
|
||||
ranking_mode: metadata.ranking_mode.clone(),
|
||||
priority_mode: metadata.priority_mode.clone(),
|
||||
ranking_index: metadata.ranking_index,
|
||||
@@ -660,6 +674,8 @@ fn build_report_candidate_extra_data(input: ReportCandidateExtraDataInput) -> Op
|
||||
body_rules,
|
||||
proxy,
|
||||
error_flow,
|
||||
candidate_group_id,
|
||||
pool_key_index,
|
||||
ranking_mode,
|
||||
priority_mode,
|
||||
ranking_index,
|
||||
@@ -703,6 +719,22 @@ fn build_report_candidate_extra_data(input: ReportCandidateExtraDataInput) -> Op
|
||||
if let Some(error_flow) = error_flow {
|
||||
extra_data.insert("error_flow".to_string(), error_flow);
|
||||
}
|
||||
if let Some(candidate_group_id) = candidate_group_id {
|
||||
extra_data.insert(
|
||||
"candidate_group_id".to_string(),
|
||||
Value::String(candidate_group_id.clone()),
|
||||
);
|
||||
extra_data.insert(
|
||||
"pool_group_id".to_string(),
|
||||
Value::String(candidate_group_id),
|
||||
);
|
||||
}
|
||||
if let Some(pool_key_index) = pool_key_index {
|
||||
extra_data.insert(
|
||||
"pool_key_index".to_string(),
|
||||
Value::Number(pool_key_index.into()),
|
||||
);
|
||||
}
|
||||
if let Some(ranking_mode) = ranking_mode {
|
||||
extra_data.insert("ranking_mode".to_string(), Value::String(ranking_mode));
|
||||
}
|
||||
@@ -880,7 +912,9 @@ mod tests {
|
||||
"classification": "retry_upstream_failure",
|
||||
"decision": "retry_next_candidate",
|
||||
"propagation": "suppressed"
|
||||
}
|
||||
},
|
||||
"candidate_group_id": "pool-group-1",
|
||||
"pool_key_index": 2
|
||||
})))
|
||||
.expect("metadata");
|
||||
|
||||
@@ -935,6 +969,24 @@ mod tests {
|
||||
.and_then(|value| value.get("propagation")),
|
||||
Some(&json!("suppressed"))
|
||||
);
|
||||
assert_eq!(
|
||||
slot.extra_data
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("candidate_group_id")),
|
||||
Some(&json!("pool-group-1"))
|
||||
);
|
||||
assert_eq!(
|
||||
slot.extra_data
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("pool_group_id")),
|
||||
Some(&json!("pool-group-1"))
|
||||
);
|
||||
assert_eq!(
|
||||
slot.extra_data
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("pool_key_index")),
|
||||
Some(&json!(2))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -328,6 +328,26 @@ impl UsageRuntime {
|
||||
self.enqueue_or_write_terminal(data, event).await;
|
||||
}
|
||||
|
||||
pub async fn record_terminal_event_direct<T>(&self, data: &T, mut event: UsageEvent)
|
||||
where
|
||||
T: UsageRuntimeAccess,
|
||||
{
|
||||
if !self.is_enabled() {
|
||||
return;
|
||||
}
|
||||
apply_body_capture_policy_from_data(data, &mut event).await;
|
||||
if let Err(err) = data.enrich_usage_event(&mut event).await {
|
||||
warn!(
|
||||
event_name = "usage_terminal_billing_enrichment_failed",
|
||||
log_type = "event",
|
||||
request_id = %event.request_id,
|
||||
error = %err,
|
||||
"usage runtime failed to enrich terminal usage event with billing"
|
||||
);
|
||||
}
|
||||
self.write_terminal_direct(data, &event).await;
|
||||
}
|
||||
|
||||
async fn enqueue_or_write_terminal<T>(&self, data: &T, event: UsageEvent)
|
||||
where
|
||||
T: UsageRuntimeAccess,
|
||||
@@ -360,7 +380,14 @@ impl UsageRuntime {
|
||||
}
|
||||
}
|
||||
|
||||
match build_upsert_usage_record_from_event(&event) {
|
||||
self.write_terminal_direct(data, &event).await;
|
||||
}
|
||||
|
||||
async fn write_terminal_direct<T>(&self, data: &T, event: &UsageEvent)
|
||||
where
|
||||
T: UsageRuntimeAccess,
|
||||
{
|
||||
match build_upsert_usage_record_from_event(event) {
|
||||
Ok(record) => match data.upsert_usage_record(record).await {
|
||||
Ok(Some(stored)) => {
|
||||
if let Err(err) = settle_usage_if_needed(data, &stored).await {
|
||||
@@ -518,7 +545,9 @@ fn now_unix_secs() -> u64 {
|
||||
mod tests {
|
||||
use std::sync::Mutex;
|
||||
|
||||
use aether_data::driver::redis::RedisStreamRunner;
|
||||
use aether_data::driver::redis::{
|
||||
RedisClientConfig, RedisClientFactory, RedisStreamRunner, RedisStreamRunnerConfig,
|
||||
};
|
||||
use aether_data_contracts::repository::settlement::{
|
||||
StoredUsageSettlement, UsageSettlementInput,
|
||||
};
|
||||
@@ -542,6 +571,29 @@ mod tests {
|
||||
records: Mutex<Vec<UpsertUsageRecord>>,
|
||||
}
|
||||
|
||||
struct RedisConfiguredUsageStore {
|
||||
inner: NoRedisUsageStore,
|
||||
runner: RedisStreamRunner,
|
||||
}
|
||||
|
||||
fn sample_runner() -> RedisStreamRunner {
|
||||
let config = RedisClientConfig {
|
||||
url: "redis://127.0.0.1/0".to_string(),
|
||||
key_prefix: Some("aether".to_string()),
|
||||
};
|
||||
let client = RedisClientFactory::new(config.clone())
|
||||
.expect("factory should build")
|
||||
.connect_lazy()
|
||||
.expect("client should build");
|
||||
|
||||
RedisStreamRunner::new(
|
||||
client,
|
||||
config.keyspace(),
|
||||
RedisStreamRunnerConfig::default(),
|
||||
)
|
||||
.expect("runner should build")
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageRecordWriter for NoRedisUsageStore {
|
||||
async fn upsert_usage_record(
|
||||
@@ -601,6 +653,64 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageRecordWriter for RedisConfiguredUsageStore {
|
||||
async fn upsert_usage_record(
|
||||
&self,
|
||||
record: UpsertUsageRecord,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
self.inner.upsert_usage_record(record).await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageSettlementWriter for RedisConfiguredUsageStore {
|
||||
fn has_usage_settlement_writer(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
async fn settle_usage(
|
||||
&self,
|
||||
_input: UsageSettlementInput,
|
||||
) -> Result<Option<StoredUsageSettlement>, DataLayerError> {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageBillingEventEnricher for RedisConfiguredUsageStore {
|
||||
async fn enrich_usage_event(&self, _event: &mut UsageEvent) -> Result<(), DataLayerError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ManualProxyNodeCounter for RedisConfiguredUsageStore {
|
||||
async fn increment_manual_proxy_node_requests(
|
||||
&self,
|
||||
_node_id: &str,
|
||||
_total_delta: i64,
|
||||
_failed_delta: i64,
|
||||
_latency_ms: Option<i64>,
|
||||
) -> Result<(), DataLayerError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl UsageRuntimeAccess for RedisConfiguredUsageStore {
|
||||
fn has_usage_writer(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn has_usage_worker_runner(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn usage_worker_runner(&self) -> Option<RedisStreamRunner> {
|
||||
Some(self.runner.clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn terminal_usage_without_redis_writes_directly_to_usage_repository() {
|
||||
let runtime = UsageRuntime::new(UsageRuntimeConfig {
|
||||
@@ -633,6 +743,41 @@ mod tests {
|
||||
assert_eq!(records[0].total_tokens, Some(12));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn direct_terminal_usage_bypasses_redis_queue_and_writes_repository() {
|
||||
let runtime = UsageRuntime::new(UsageRuntimeConfig {
|
||||
enabled: true,
|
||||
..UsageRuntimeConfig::default()
|
||||
})
|
||||
.expect("usage runtime should build");
|
||||
let runner = sample_runner();
|
||||
let store = RedisConfiguredUsageStore {
|
||||
inner: NoRedisUsageStore::default(),
|
||||
runner,
|
||||
};
|
||||
let event = UsageEvent::new(
|
||||
UsageEventType::Failed,
|
||||
"req-direct-terminal-1",
|
||||
UsageEventData {
|
||||
user_id: Some("user-direct-terminal-1".to_string()),
|
||||
provider_name: "openai".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
status_code: Some(503),
|
||||
error_message: Some("upstream failed".to_string()),
|
||||
..UsageEventData::default()
|
||||
},
|
||||
);
|
||||
|
||||
runtime.record_terminal_event_direct(&store, event).await;
|
||||
|
||||
let records = store.inner.records.lock().expect("records lock");
|
||||
assert_eq!(records.len(), 1);
|
||||
assert_eq!(records[0].request_id, "req-direct-terminal-1");
|
||||
assert_eq!(records[0].status, "failed");
|
||||
assert_eq!(records[0].billing_status, "void");
|
||||
assert_eq!(records[0].status_code, Some(503));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn basic_request_record_level_strips_body_capture_but_preserves_derived_fields() {
|
||||
let mut event = UsageEvent::new(
|
||||
|
||||
@@ -1,11 +1,12 @@
|
||||
export interface OAuthStatusSnapshot {
|
||||
code: 'none' | 'valid' | 'expiring' | 'expired' | 'invalid' | 'check_failed'
|
||||
code: 'none' | 'valid' | 'expiring' | 'expired' | 'invalid' | 'reauth_required' | 'check_failed'
|
||||
label?: string | null
|
||||
reason?: string | null
|
||||
expires_at?: number | null
|
||||
invalid_at?: number | null
|
||||
source?: string | null
|
||||
requires_reauth?: boolean
|
||||
usable_until_expiry?: boolean
|
||||
expiring_soon?: boolean
|
||||
}
|
||||
|
||||
|
||||
@@ -157,6 +157,8 @@ export interface OAuthStatusInfo {
|
||||
isExpiringSoon: boolean
|
||||
isInvalid: boolean // Token 已失效(账号被封、授权撤销等)
|
||||
invalidReason?: string // 失效原因
|
||||
requiresReauth?: boolean
|
||||
usableUntilExpiry?: boolean
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -486,14 +486,27 @@ function normalizeAuthTypeLabel(key: PoolKeyDetail | PoolKeySelectionItem): stri
|
||||
|
||||
function getStatusBadgeLabel(key: PoolKeyDetail): string | null {
|
||||
const account = getAccountStatusDisplay(key)
|
||||
if (account.blocked && account.label) return account.label
|
||||
if (account.blocked && account.label) return compactStatusBadgeLabel(account.label)
|
||||
|
||||
const oauth = getOAuthStatusDisplay(key, 0)
|
||||
if (oauth?.isInvalid) return 'Token 失效'
|
||||
if (oauth?.isExpired) return 'Token 过期'
|
||||
if (oauth?.requiresReauth) return '续期失败'
|
||||
if (oauth?.isInvalid) return '已失效'
|
||||
if (oauth?.isExpired) return '已过期'
|
||||
return null
|
||||
}
|
||||
|
||||
function compactStatusBadgeLabel(label: string): string {
|
||||
const normalized = label.trim()
|
||||
const mapped: Record<string, string> = {
|
||||
'Token 失效': '已失效',
|
||||
'Token 过期': '已过期',
|
||||
账号已封禁: '账号封禁',
|
||||
工作区已停用: '工作区停用',
|
||||
账号访问受限: '访问受限',
|
||||
}
|
||||
return Array.from(mapped[normalized] || normalized).slice(0, 5).join('')
|
||||
}
|
||||
|
||||
function getStatusBadgeTitle(key: PoolKeyDetail): string {
|
||||
const label = getStatusBadgeLabel(key)
|
||||
if (!label) return ''
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
import { describe, expect, it } from 'vitest'
|
||||
|
||||
import { cleanAccountBlockReason, isRefreshFailedReason } from '@/utils/accountBlock'
|
||||
import {
|
||||
cleanAccountBlockReason,
|
||||
isAccountLevelBlockReason,
|
||||
isRefreshFailedReason,
|
||||
} from '@/utils/accountBlock'
|
||||
|
||||
describe('accountBlock helpers', () => {
|
||||
it('detects refresh failure markers even when account block is also present', () => {
|
||||
@@ -18,4 +22,12 @@ describe('accountBlock helpers', () => {
|
||||
),
|
||||
).toBe('工作区已停用 (deactivated_workspace)')
|
||||
})
|
||||
|
||||
it('does not treat refresh failure text as an account block by itself', () => {
|
||||
expect(
|
||||
isAccountLevelBlockReason(
|
||||
'[REFRESH_FAILED] Token 续期失败 (401): token has been invalidated',
|
||||
),
|
||||
).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
@@ -87,11 +87,59 @@ describe('providerKeyStatus', () => {
|
||||
)
|
||||
|
||||
expect(status).toEqual({
|
||||
text: '已失效',
|
||||
text: expect.stringMatching(/^续期失败 /),
|
||||
isExpired: false,
|
||||
isExpiringSoon: false,
|
||||
isInvalid: true,
|
||||
isExpiringSoon: expect.any(Boolean),
|
||||
isInvalid: false,
|
||||
invalidReason: '[REFRESH_FAILED] Token 续期失败 (401): refresh_token_reused',
|
||||
requiresReauth: true,
|
||||
usableUntilExpiry: true,
|
||||
})
|
||||
expect(status?.requiresReauth).toBe(true)
|
||||
expect(getOAuthStatusTitle({
|
||||
auth_type: 'oauth',
|
||||
oauth_expires_at: future,
|
||||
oauth_invalid_reason: '[REFRESH_FAILED] Token 续期失败 (401): refresh_token_reused',
|
||||
}, 0)).toContain('当前 Access Token 未到期仍可使用')
|
||||
expect(getOAuthRefreshButtonTitle({
|
||||
auth_type: 'oauth',
|
||||
oauth_expires_at: future,
|
||||
oauth_invalid_reason: '[REFRESH_FAILED] Token 续期失败 (401): refresh_token_reused',
|
||||
}, 0)).toBe('重新授权')
|
||||
})
|
||||
|
||||
it('shows refresh failure as reauth required while access token is still usable', () => {
|
||||
const future = Math.floor(Date.now() / 1000) + 2 * 24 * 3600
|
||||
const status = getOAuthStatusDisplay(
|
||||
{
|
||||
auth_type: 'oauth',
|
||||
oauth_expires_at: future,
|
||||
status_snapshot: {
|
||||
oauth: {
|
||||
code: 'reauth_required',
|
||||
reason: 'Token 续期失败 (400): refresh_token_reused',
|
||||
expires_at: future,
|
||||
requires_reauth: true,
|
||||
usable_until_expiry: true,
|
||||
},
|
||||
account: {
|
||||
code: 'ok',
|
||||
blocked: false,
|
||||
},
|
||||
quota: { code: 'ok', exhausted: false },
|
||||
},
|
||||
},
|
||||
0,
|
||||
)
|
||||
|
||||
expect(status).toEqual({
|
||||
text: expect.stringMatching(/^续期失败 /),
|
||||
isExpired: false,
|
||||
isExpiringSoon: expect.any(Boolean),
|
||||
isInvalid: false,
|
||||
invalidReason: 'Token 续期失败 (400): refresh_token_reused',
|
||||
requiresReauth: true,
|
||||
usableUntilExpiry: true,
|
||||
})
|
||||
})
|
||||
|
||||
|
||||
@@ -56,6 +56,7 @@ export function isAccountLevelBlockReason(reason: string | null | undefined): bo
|
||||
if (!text) return false
|
||||
if (text.startsWith('[ACCOUNT_BLOCK]')) return true
|
||||
if (text.startsWith('[OAUTH_EXPIRED]')) return true
|
||||
if (text.startsWith('[REFRESH_FAILED]')) return false
|
||||
const lowered = text.toLowerCase()
|
||||
return ACCOUNT_BLOCK_REASON_KEYWORDS.some(keyword => lowered.includes(keyword))
|
||||
}
|
||||
|
||||
@@ -84,6 +84,19 @@ function getSnapshotOAuthState(
|
||||
const expiresAt = oauth.expires_at ?? input.oauth_expires_at ?? null
|
||||
const reason = normalizeText(oauth.reason)
|
||||
|
||||
if (code === 'reauth_required') {
|
||||
const countdown = expiresAt == null ? null : getOAuthExpiresCountdown(expiresAt, tick, null, null)
|
||||
return {
|
||||
text: countdown?.text ? `续期失败 ${countdown.text}` : '续期失败',
|
||||
isExpired: false,
|
||||
isExpiringSoon: countdown?.isExpiringSoon ?? false,
|
||||
isInvalid: false,
|
||||
invalidReason: reason || undefined,
|
||||
requiresReauth: true,
|
||||
usableUntilExpiry: true,
|
||||
}
|
||||
}
|
||||
|
||||
if (code === 'invalid') {
|
||||
return {
|
||||
text: '已失效',
|
||||
@@ -107,6 +120,27 @@ function getSnapshotOAuthState(
|
||||
return getOAuthExpiresCountdown(expiresAt, tick, null, null)
|
||||
}
|
||||
|
||||
function refreshFailureAccessTokenStillUsable(expiresAt: number | null | undefined): boolean {
|
||||
return typeof expiresAt === 'number' && expiresAt > Math.floor(Date.now() / 1000)
|
||||
}
|
||||
|
||||
function getReauthRequiredOAuthState(
|
||||
expiresAt: number | null | undefined,
|
||||
tick: number,
|
||||
reason: string,
|
||||
): OAuthStatusInfo {
|
||||
const countdown = expiresAt == null ? null : getOAuthExpiresCountdown(expiresAt, tick, null, null)
|
||||
return {
|
||||
text: countdown?.text ? `续期失败 ${countdown.text}` : '续期失败',
|
||||
isExpired: false,
|
||||
isExpiringSoon: countdown?.isExpiringSoon ?? false,
|
||||
isInvalid: false,
|
||||
invalidReason: reason,
|
||||
requiresReauth: true,
|
||||
usableUntilExpiry: true,
|
||||
}
|
||||
}
|
||||
|
||||
function getLegacyOAuthState(
|
||||
input: ProviderKeyStatusCarrier,
|
||||
tick: number,
|
||||
@@ -115,6 +149,14 @@ function getLegacyOAuthState(
|
||||
if (!input.oauth_expires_at && !input.oauth_invalid_at && !input.oauth_invalid_reason) return null
|
||||
|
||||
const rawReason = normalizeText(input.oauth_invalid_reason)
|
||||
if (
|
||||
rawReason
|
||||
&& isRefreshFailedReason(rawReason)
|
||||
&& refreshFailureAccessTokenStillUsable(input.oauth_expires_at)
|
||||
) {
|
||||
return getReauthRequiredOAuthState(input.oauth_expires_at, tick, rawReason)
|
||||
}
|
||||
|
||||
if (rawReason && isAccountLevelBlockReason(rawReason) && !isRefreshFailedReason(rawReason)) {
|
||||
if (!input.oauth_expires_at) return null
|
||||
return getOAuthExpiresCountdown(input.oauth_expires_at, tick, null, null)
|
||||
@@ -132,6 +174,7 @@ function getOAuthStatusSeverity(status: OAuthStatusInfo | null): number {
|
||||
if (!status) return 0
|
||||
if (status.isInvalid) return 3
|
||||
if (status.isExpired) return 2
|
||||
if (status.requiresReauth) return 2
|
||||
return 1
|
||||
}
|
||||
|
||||
@@ -217,8 +260,15 @@ export function getOAuthStatusTitle(
|
||||
const reason = normalizeText(status.invalidReason)
|
||||
return reason ? `Token 已失效: ${reason}` : 'Token 已失效'
|
||||
}
|
||||
const snapshotCode = normalizeText(input.status_snapshot?.oauth?.code)
|
||||
if (snapshotCode === 'reauth_required' || status.requiresReauth) {
|
||||
const reason = normalizeText(status.invalidReason)
|
||||
return reason
|
||||
? `Refresh Token 续期失败,当前 Access Token 未到期仍可使用: ${reason}`
|
||||
: 'Refresh Token 续期失败,当前 Access Token 未到期仍可使用'
|
||||
}
|
||||
if (status.isExpired) {
|
||||
return 'Token 已过期,请重新授权'
|
||||
return 'Access Token 已过期,等待自动续期'
|
||||
}
|
||||
return `Token 剩余有效期: ${status.text}`
|
||||
}
|
||||
@@ -235,7 +285,7 @@ export function getOAuthRefreshButtonTitle(
|
||||
}
|
||||
|
||||
const status = getOAuthStatusDisplay(input, tick)
|
||||
if (status?.isInvalid || status?.isExpired) {
|
||||
if (status?.isInvalid || status?.isExpired || status?.requiresReauth) {
|
||||
return '重新授权'
|
||||
}
|
||||
return '刷新 Token'
|
||||
|
||||
@@ -2873,9 +2873,38 @@ function getSchedulingStatus(key: PoolKeyDetail): 'available' | 'degraded' | 'bl
|
||||
return 'available'
|
||||
}
|
||||
|
||||
function compactPoolStatusLabel(label: string | null | undefined): string | null {
|
||||
const normalized = String(label || '').trim()
|
||||
if (!normalized) return null
|
||||
|
||||
const mapped: Record<string, string> = {
|
||||
'Token 失效': '已失效',
|
||||
'Token 过期': '已过期',
|
||||
Token失效: '已失效',
|
||||
Token过期: '已过期',
|
||||
账号已封禁: '账号封禁',
|
||||
工作区已停用: '工作区停用',
|
||||
账号访问受限: '访问受限',
|
||||
健康度较低: '健康低',
|
||||
}
|
||||
const labelText = mapped[normalized] || normalized
|
||||
return Array.from(labelText).slice(0, 5).join('')
|
||||
}
|
||||
|
||||
function getOAuthStatusBadgeLabel(status: ReturnType<typeof getVisibleOAuthState>): string | null {
|
||||
if (!status) return null
|
||||
if (status.requiresReauth) return '续期失败'
|
||||
if (status.isInvalid) return '已失效'
|
||||
if (status.isExpired) return '已过期'
|
||||
if (status.text === '未添加') return '未添加'
|
||||
if (status.text === '有效期未知') return '未知'
|
||||
if (status.isExpiringSoon) return '将过期'
|
||||
return '有效'
|
||||
}
|
||||
|
||||
function getSchedulingBadgeLabel(key: PoolKeyDetail): string {
|
||||
const accountAlert = getAccountAlertLabel(key)
|
||||
if (accountAlert) return accountAlert
|
||||
if (accountAlert) return compactPoolStatusLabel(accountAlert) || accountAlert
|
||||
|
||||
const rawLabel = String(key.scheduling_label || '').trim()
|
||||
if (
|
||||
@@ -2884,7 +2913,7 @@ function getSchedulingBadgeLabel(key: PoolKeyDetail): string {
|
||||
&& !isHealthDerivedSchedulingLabel(rawLabel)
|
||||
) {
|
||||
if (rawLabel === '禁用' || rawLabel === '停用') return '禁用'
|
||||
return rawLabel
|
||||
return compactPoolStatusLabel(rawLabel) || rawLabel
|
||||
}
|
||||
|
||||
if (!key.is_active) return '禁用'
|
||||
@@ -2961,9 +2990,9 @@ function getMobileTagItems(key: PoolKeyDetail): PoolMobileTagItem[] {
|
||||
const orgBadge = getOAuthOrgBadge(key)
|
||||
|
||||
return buildPoolMobileTagItems({
|
||||
accountStatusLabel: accountAlert,
|
||||
accountStatusLabel: compactPoolStatusLabel(accountAlert),
|
||||
accountStatusTone: accountAlert ? 'danger' : null,
|
||||
oauthStatusLabel: oauthState?.text ?? null,
|
||||
oauthStatusLabel: getOAuthStatusBadgeLabel(oauthState),
|
||||
oauthStatusTone: getMobileOAuthTone(key),
|
||||
priorityLabel: `P${key.internal_priority ?? 50}`,
|
||||
authLabel: getAuthTypeChipLabel(key),
|
||||
|
||||
+17
-10
@@ -4,7 +4,7 @@ set -euo pipefail
|
||||
REPO="${AETHER_REPO:-fawney19/Aether}"
|
||||
SOURCE_REF="${AETHER_SOURCE_REF:-aether-rust-pioneer}"
|
||||
VERSION="${AETHER_VERSION:-}"
|
||||
CHANNEL="${AETHER_CHANNEL:-rc}"
|
||||
CHANNEL="${AETHER_CHANNEL:-pre}"
|
||||
CHANNEL_EXPLICIT="false"
|
||||
if [[ -n "${AETHER_CHANNEL:-}" ]]; then
|
||||
CHANNEL_EXPLICIT="true"
|
||||
@@ -38,9 +38,10 @@ Options:
|
||||
compose: Docker Compose app + Postgres + Redis
|
||||
single: systemd service with SQLite + in-process runtime
|
||||
cluster: systemd service connected to shared database + Redis
|
||||
--channel CHANNEL Release channel to resolve when --version is omitted: rc
|
||||
(default: rc)
|
||||
--version VERSION Exact release tag to install, for example v0.7.0-rc22
|
||||
--channel CHANNEL Release channel to resolve when --version is omitted: pre or rc
|
||||
pre resolves the latest semver prerelease tag (default)
|
||||
rc restricts resolution to tags like v0.7.0-rc23
|
||||
--version VERSION Exact release tag to install, for example v0.7.0-rc23
|
||||
--repo OWNER/REPO GitHub repository to download from (default: fawney19/Aether)
|
||||
--source-ref REF Source branch/tag used for compose templates (default: aether-rust-pioneer)
|
||||
--archive PATH Install from a local release tarball instead of downloading
|
||||
@@ -177,8 +178,8 @@ select_version() {
|
||||
cat >/dev/tty <<'EOF'
|
||||
|
||||
Choose Aether version:
|
||||
1) Latest rc tag
|
||||
2) Exact tag, for example v0.7.0-rc22
|
||||
1) Latest pre release
|
||||
2) Exact tag, for example v0.7.0-rc23
|
||||
|
||||
Enter choice [1]:
|
||||
EOF
|
||||
@@ -186,7 +187,7 @@ EOF
|
||||
IFS= read -r choice </dev/tty || choice=""
|
||||
case "${choice:-1}" in
|
||||
1)
|
||||
CHANNEL="rc"
|
||||
CHANNEL="pre"
|
||||
;;
|
||||
2)
|
||||
cat >/dev/tty <<'EOF'
|
||||
@@ -356,6 +357,12 @@ resolve_version() {
|
||||
|
||||
local tag=""
|
||||
case "${CHANNEL}" in
|
||||
pre)
|
||||
tag="$(download_stdout "https://api.github.com/repos/${REPO}/releases?per_page=50" |
|
||||
sed -n 's/.*"tag_name"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p' |
|
||||
grep -E '^v[0-9]+(\.[0-9]+)*-[0-9A-Za-z][0-9A-Za-z.-]*$' |
|
||||
head -n1 || true)"
|
||||
;;
|
||||
rc)
|
||||
tag="$(download_stdout "https://api.github.com/repos/${REPO}/releases?per_page=50" |
|
||||
sed -n 's/.*"tag_name"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p' |
|
||||
@@ -363,7 +370,7 @@ resolve_version() {
|
||||
head -n1 || true)"
|
||||
;;
|
||||
*)
|
||||
die "unsupported release channel: ${CHANNEL}; expected rc"
|
||||
die "unsupported release channel: ${CHANNEL}; expected pre or rc"
|
||||
;;
|
||||
esac
|
||||
echo "${tag}"
|
||||
@@ -634,8 +641,8 @@ compose_image() {
|
||||
tag="${VERSION#v}"
|
||||
else
|
||||
case "${CHANNEL}" in
|
||||
rc)
|
||||
tag="rc"
|
||||
pre|rc)
|
||||
tag="${CHANNEL}"
|
||||
;;
|
||||
*)
|
||||
tag="${CHANNEL}"
|
||||
|
||||
Reference in New Issue
Block a user