mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-13 08:06:41 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
789609dd53 | ||
|
|
10e25ce25f | ||
|
|
6b362a5029 | ||
|
|
13374e2a70 | ||
|
|
19ef484969 | ||
|
|
242ffb2324 | ||
|
|
53a09bbfe9 | ||
|
|
239def9a5b | ||
|
|
5eaca14a64 | ||
|
|
78e943f406 | ||
|
|
0663df14df | ||
|
|
39aa20a13a | ||
|
|
621b1faceb | ||
|
|
28eb4708b5 | ||
|
|
aa003f7d90 | ||
|
|
4c57a0405d | ||
|
|
49cee2397f | ||
|
|
1915806bd1 | ||
|
|
5100f6333e | ||
|
|
187af36306 | ||
|
|
be35f8a6ee | ||
|
|
317c05c74a | ||
|
|
12c32fcdec | ||
|
|
26de671226 | ||
|
|
a92b95c48c | ||
|
|
e7c24781fb | ||
|
|
3b4b547e9c | ||
|
|
2016bda031 | ||
|
|
790f3147e0 |
@@ -15,6 +15,8 @@ description: 准备、创建、更新或审查 GitHub PR,含回复讨论与处
|
||||
|
||||
本仓库目标分支为 `main`,其他仓库使用其可信基线规则与模板,PR 内修改的规则不能降低审查要求
|
||||
|
||||
为本仓库拟写或更新 PR 标题时,完整标题使用英文,沿用提交信息的 Conventional Commits 格式,例如 `fix(store): preserve client key ledgers after deletion`
|
||||
|
||||
涉及上游行为的判断按[问题与方案依据](../../../CONTRIBUTING.md#问题与方案依据)执行
|
||||
|
||||
## 按模式读取
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
# SQLx migration checksums are byte-sensitive; keep migration inputs portable across hosts.
|
||||
backend/migrations/*.sql text eol=lf
|
||||
backend/migrations/.frozen-sha256 text eol=lf
|
||||
@@ -146,6 +146,10 @@ async fn execute(
|
||||
providers.admin,
|
||||
);
|
||||
or_shutdown!(plugin_runtime, services.register_usage_data(&usage_data));
|
||||
or_shutdown!(
|
||||
plugin_runtime,
|
||||
plugin_runtime.bind_usage_queries(&usage_data)
|
||||
);
|
||||
let services = std::sync::Arc::new(services);
|
||||
or_shutdown!(plugin_runtime, plugin_runtime.bind_services(&services));
|
||||
let commands = or_shutdown!(plugin_runtime, plugin_runtime.prepare_command_line().await);
|
||||
|
||||
@@ -160,6 +160,10 @@ pub async fn run() -> Result<(), BootstrapError> {
|
||||
plugin_runtime,
|
||||
plugin_runtime.bind_services(&admin.services().public_services())
|
||||
);
|
||||
or_shutdown!(
|
||||
plugin_runtime,
|
||||
plugin_runtime.bind_usage_queries(&admin.services().usage_data())
|
||||
);
|
||||
|
||||
let mut plan = store.take_worker_contributions();
|
||||
plan.push(analytics_worker);
|
||||
|
||||
@@ -197,6 +197,7 @@ async fn initialize_inner(
|
||||
auth,
|
||||
accounts: accounts.clone(),
|
||||
account_groups: Arc::new(DefaultAccountGroupService::new(
|
||||
upstream_info.clone(),
|
||||
store.account_groups(),
|
||||
store.account_runtime(),
|
||||
snapshot.clone(),
|
||||
|
||||
@@ -7,6 +7,7 @@ use std::time::{Duration, SystemTime};
|
||||
use chrono::{DateTime, Utc};
|
||||
use futures::StreamExt as _;
|
||||
use gateway_core::account::{AccountErrorProtection, ProviderAccountId};
|
||||
use gateway_core::engine::probe::AccountProbeTarget;
|
||||
use gateway_core::routing::UpstreamModelId;
|
||||
use gateway_core::task::{
|
||||
ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerId, WorkerKind,
|
||||
@@ -161,7 +162,7 @@ impl FreezeRecoveryTask {
|
||||
let probe_succeeded = match self
|
||||
.deps
|
||||
.accounts
|
||||
.test_connection(account.clone(), model)
|
||||
.test_connection(AccountProbeTarget::FreezeRecovery(account.clone()), model)
|
||||
.await
|
||||
{
|
||||
Ok(events) => drain_probe(events).await,
|
||||
|
||||
@@ -288,8 +288,8 @@ impl AdminServices {
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn usage_data(&self) -> &dyn UsageDataService {
|
||||
self.usage_data.as_ref()
|
||||
pub fn usage_data(&self) -> Arc<dyn UsageDataService> {
|
||||
self.usage_data.clone()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
|
||||
@@ -73,6 +73,7 @@ pub struct AccountGroupUsage {
|
||||
/// 当前页账号组 membership 对应的持久账号事实;运行态在 Admin query service 合并
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountGroupMemberFact {
|
||||
pub inherited_version: Option<gateway_core::account::QuotaObservationVersion>,
|
||||
pub group_id: AccountGroupId,
|
||||
pub account_id: String,
|
||||
pub status: AccountStatusFacts,
|
||||
|
||||
@@ -219,7 +219,7 @@ pub struct AccountPageItem {
|
||||
pub struct AccountCapacity {
|
||||
/// 实时租约读取失败时为 `None`,不能当作空闲
|
||||
pub used_slots: Option<u64>,
|
||||
/// 应用账号覆盖或全局默认值后的上限;`None` 表示不限
|
||||
/// 应用账号覆盖或继承默认值后的普通并发上限;`None` 表示不限
|
||||
pub total_slots: Option<u64>,
|
||||
}
|
||||
|
||||
|
||||
@@ -684,8 +684,10 @@ pub struct DashboardAccountUsage {
|
||||
}
|
||||
|
||||
/// Dashboard 当前账号池的可重建运行时槽位事实
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct DashboardRuntimeSlots {
|
||||
/// 仅未设置本地覆盖的账号可能继承有效互联观测
|
||||
pub inherited_versions: Vec<gateway_core::account::QuotaObservationVersion>,
|
||||
/// 继承运行参数默认并发上限的可调度账号数
|
||||
pub inherited_accounts: u64,
|
||||
/// 已配置账号级并发上限的可调度账号槽位总数
|
||||
@@ -724,6 +726,11 @@ pub struct UsageListRecord {
|
||||
pub outcome: RequestOutcome,
|
||||
pub error_code: Option<String>,
|
||||
pub error_message: Option<String>,
|
||||
pub error_kind: Option<String>,
|
||||
pub websocket_close_code: Option<u16>,
|
||||
pub continuation_unavailable_reason: Option<String>,
|
||||
pub upstream_connection_exit_reason: Option<String>,
|
||||
pub recovered: bool,
|
||||
pub client_status_code: Option<u16>,
|
||||
pub client_api_key_name: Option<String>,
|
||||
pub id: String,
|
||||
@@ -782,6 +789,23 @@ pub struct UsageListRecord {
|
||||
pub started_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// 受控请求详情中的上游错误字段,不参与错误分类或普通日志输出
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct UsageUpstreamError {
|
||||
pub code: Option<String>,
|
||||
pub error_type: Option<String>,
|
||||
pub message: Option<String>,
|
||||
pub param: Option<String>,
|
||||
pub close_reason: Option<String>,
|
||||
pub last_event_type: Option<String>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for UsageUpstreamError {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter.write_str("UsageUpstreamError(<restricted>)")
|
||||
}
|
||||
}
|
||||
|
||||
/// 一次完整模型请求的公共观测记录
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct UsageRecord {
|
||||
@@ -820,6 +844,11 @@ pub struct UsageRecord {
|
||||
pub error_kind: Option<String>,
|
||||
pub provider_error_code: Option<String>,
|
||||
pub error_message: Option<String>,
|
||||
pub upstream_error: Option<UsageUpstreamError>,
|
||||
pub websocket_close_code: Option<u16>,
|
||||
pub continuation_unavailable_reason: Option<String>,
|
||||
pub upstream_connection_exit_reason: Option<String>,
|
||||
pub recovered: bool,
|
||||
pub retry_after_ms: Option<u64>,
|
||||
pub input_tokens: Option<u64>,
|
||||
pub output_tokens: Option<u64>,
|
||||
|
||||
@@ -17,6 +17,7 @@ pub struct PeerLinkLimits {
|
||||
pub struct PeerLinkExport {
|
||||
pub observed_at: DateTime<Utc>,
|
||||
pub expires_at: DateTime<Utc>,
|
||||
pub key_limits: gateway_core::policy::RateLimits,
|
||||
pub budget: ClientBudgetStatus,
|
||||
pub budget_access: ClientBudgetAccess,
|
||||
pub budget_observed_at: DateTime<Utc>,
|
||||
|
||||
@@ -806,11 +806,12 @@ pub struct ProviderQuotaWindow {
|
||||
pub provider_data: Option<ProviderDocument>,
|
||||
}
|
||||
|
||||
/// 账号用量统计周期,按周优先、月次之选择;短期限流窗口不参与统计面板
|
||||
/// 账号用量统计周期,完整周/月额度优先,随后使用 Provider 提供的本地滚动日统计
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
||||
pub enum AccountUsagePeriod {
|
||||
Weekly,
|
||||
Monthly,
|
||||
Daily,
|
||||
}
|
||||
|
||||
/// Provider 解释 quota 所需的公共请求事实
|
||||
@@ -998,16 +999,27 @@ impl ProviderQuota {
|
||||
}
|
||||
}
|
||||
|
||||
/// 选择账号用量统计的周/月窗口,同时返回展示所需的周期事实
|
||||
/// 选择账号用量统计窗口,同时返回展示所需的周期事实
|
||||
///
|
||||
/// 只使用边界完整、可归属到整个账号的窗口;周优先于月,同周期保持
|
||||
/// Provider 投影顺序
|
||||
/// 没有符合条件的窗口时不回退到短期或历史累计
|
||||
/// 周/月额度使用边界完整、可归属到整个账号的窗口
|
||||
/// 周优先于月,同周期保持 Provider 投影顺序
|
||||
/// Provider 提供的无固定重置点的滚动日统计可独立展示
|
||||
#[must_use]
|
||||
pub fn usage_window(&self) -> Option<(&ProviderQuotaWindow, AccountUsagePeriod)> {
|
||||
self.usage_windows()
|
||||
.filter(|(window, _)| window.local_usage.is_some())
|
||||
.min_by_key(|(_, period)| *period)
|
||||
.or_else(|| {
|
||||
self.windows
|
||||
.iter()
|
||||
.find(|window| {
|
||||
window.local_usage_attribution == QuotaLocalUsageAttribution::AccountWide
|
||||
&& window.window_seconds == Some(86_400)
|
||||
&& window.reset_at.is_none()
|
||||
&& window.local_usage.is_some()
|
||||
})
|
||||
.map(|window| (window, AccountUsagePeriod::Daily))
|
||||
})
|
||||
}
|
||||
|
||||
/// 互联账号按最长的 Key 消费周期统计本地用量,池内账号窗口不代表此 Key 的周期
|
||||
@@ -1025,7 +1037,7 @@ impl ProviderQuota {
|
||||
.max_by_key(|window| window.window_seconds)
|
||||
}
|
||||
|
||||
/// 统计与预测共用窗口归属规则,避免把模型专属桶或短期限流当作账号容量
|
||||
/// 周/月额度的归属规则,供额度周期统计与容量预测共用
|
||||
pub(crate) fn usage_windows(
|
||||
&self,
|
||||
) -> impl Iterator<Item = (&ProviderQuotaWindow, AccountUsagePeriod)> {
|
||||
|
||||
@@ -62,14 +62,15 @@ pub fn account_quota_forecasts(
|
||||
now: DateTime<Utc>,
|
||||
samples: &[QuotaForecastSample],
|
||||
) -> [AccountQuotaForecast; 2] {
|
||||
[AccountUsagePeriod::Weekly, AccountUsagePeriod::Monthly].map(|period| {
|
||||
[
|
||||
(AccountUsagePeriod::Weekly, 7 * DAY_SECONDS),
|
||||
(AccountUsagePeriod::Monthly, 30 * DAY_SECONDS),
|
||||
]
|
||||
.map(|(period, target_seconds)| {
|
||||
let selected = quota_forecast_source_window(quota, period);
|
||||
let mut forecast = AccountQuotaForecast {
|
||||
period,
|
||||
target_seconds: match period {
|
||||
AccountUsagePeriod::Weekly => 7 * DAY_SECONDS,
|
||||
AccountUsagePeriod::Monthly => 30 * DAY_SECONDS,
|
||||
},
|
||||
target_seconds,
|
||||
extrapolated: false,
|
||||
source: None,
|
||||
unavailable_reason: Some("缺少周/月额度窗口,请刷新账号额度"),
|
||||
|
||||
@@ -4,11 +4,14 @@ use chrono::{DateTime, Utc};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderUpstreamInfo {
|
||||
/// 当前连接曾成功识别互联,失权时保留身份但清除观测内容
|
||||
pub discovered: bool,
|
||||
pub observed_at: Option<DateTime<Utc>>,
|
||||
pub stale: bool,
|
||||
pub connection_revision: String,
|
||||
pub last_error: Option<String>,
|
||||
pub max_concurrency: Option<u64>,
|
||||
pub requests_per_minute: Option<u64>,
|
||||
pub pool: Option<UpstreamPool>,
|
||||
}
|
||||
|
||||
|
||||
@@ -10,6 +10,8 @@ pub enum UsageErrorKind {
|
||||
Conflict,
|
||||
CapacityExceeded,
|
||||
Unavailable,
|
||||
Cancelled,
|
||||
Timeout,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
|
||||
@@ -20,3 +20,6 @@ pub use actor::*;
|
||||
|
||||
mod aggregate;
|
||||
pub use aggregate::*;
|
||||
|
||||
mod read;
|
||||
pub use read::*;
|
||||
|
||||
@@ -0,0 +1,162 @@
|
||||
//! 历史用量的窄事实、分组统计与调用内分批读取合同
|
||||
|
||||
use super::{UsageCount, UsageMetric, UsageRecordData, UsageRecordFilter};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
pub const USAGE_READ_MAX_BATCH_ROWS: u16 = 1024;
|
||||
pub const USAGE_READ_MAX_BATCH_BYTES: usize = 2 * 1024 * 1024;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct UsageReadFilter {
|
||||
pub records: UsageRecordFilter,
|
||||
pub routing_group_ref: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ScanUsageRecords {
|
||||
pub filter: UsageReadFilter,
|
||||
pub batch_size: Option<u16>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum UsageGrouping {
|
||||
Model,
|
||||
Key,
|
||||
KeyModel,
|
||||
RoutingGroup,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct QueryUsageStatistics {
|
||||
pub filter: UsageReadFilter,
|
||||
pub group_by: UsageGrouping,
|
||||
pub batch_size: Option<u16>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||
pub enum UsageReadQuery {
|
||||
Records(ScanUsageRecords),
|
||||
Statistics(QueryUsageStatistics),
|
||||
}
|
||||
|
||||
impl UsageReadQuery {
|
||||
pub fn filter(&self) -> &UsageReadFilter {
|
||||
match self {
|
||||
Self::Records(query) => &query.filter,
|
||||
Self::Statistics(query) => &query.filter,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn batch_size(&self) -> u16 {
|
||||
match self {
|
||||
Self::Records(query) => query.batch_size,
|
||||
Self::Statistics(query) => query.batch_size,
|
||||
}
|
||||
.unwrap_or(USAGE_READ_MAX_BATCH_ROWS)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct UsageRecordFact {
|
||||
pub id: String,
|
||||
pub data: UsageRecordData,
|
||||
pub user_agent: Option<String>,
|
||||
pub reasoning_effort: Option<String>,
|
||||
pub subagent_kind: Option<String>,
|
||||
pub request_kind: Option<String>,
|
||||
pub routing_scope: Option<String>,
|
||||
pub routing_group_refs: Vec<String>,
|
||||
pub recovered: bool,
|
||||
/// 由宿主现行规则派生,原始审计中的恢复请求仍可读取
|
||||
pub request_counted: bool,
|
||||
/// 包含完整交付、预热排除及恢复关联规则
|
||||
pub usage_counted: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct UsageStatisticsGroup {
|
||||
pub client_key_ref: Option<String>,
|
||||
pub model: Option<String>,
|
||||
pub routing_group_ref: Option<String>,
|
||||
pub unassigned_contribution: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct UsageStatisticsCost {
|
||||
pub currency: String,
|
||||
pub amount: String,
|
||||
pub native_amount: Option<String>,
|
||||
/// 外部修正允许负值,原生费用来源覆盖单独返回
|
||||
pub contribution_amount: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct UsageStatisticsCostCoverage {
|
||||
pub provider_reported: UsageCount,
|
||||
pub calculated: UsageCount,
|
||||
pub unavailable: UsageCount,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct UsageStatisticsMetrics {
|
||||
/// 缺失的 Token 指标表示未知,计数指标始终返回
|
||||
pub values: BTreeMap<UsageMetric, UsageCount>,
|
||||
pub native_request_count: UsageCount,
|
||||
pub contribution_count: UsageCount,
|
||||
pub native_token_samples: BTreeMap<UsageMetric, UsageCount>,
|
||||
/// 数量表示提供该指标的贡献记录数,不能用作请求覆盖率的分母
|
||||
pub contribution_token_samples: BTreeMap<UsageMetric, UsageCount>,
|
||||
pub cost_coverage: UsageStatisticsCostCoverage,
|
||||
pub costs: Vec<UsageStatisticsCost>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct UsageStatisticsRow {
|
||||
/// null 表示全局总计,分组重叠不会重复计入该行
|
||||
pub group: Option<UsageStatisticsGroup>,
|
||||
pub metrics: UsageStatisticsMetrics,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(tag = "kind", content = "value", rename_all = "snake_case")]
|
||||
pub enum UsageReadItem {
|
||||
Record(Box<UsageRecordFact>),
|
||||
Statistics(Box<UsageStatisticsRow>),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct UsageReadAnalytics {
|
||||
pub coverage_start_ms: Option<i64>,
|
||||
pub coverage_end_ms: Option<i64>,
|
||||
pub lag_seconds: u64,
|
||||
pub state: String,
|
||||
pub generation: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct UsageReadObservation {
|
||||
pub generated_at_ms: i64,
|
||||
pub analytics: Option<UsageReadAnalytics>,
|
||||
pub total_groups: Option<UsageCount>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct UsageReadBatch {
|
||||
pub observation: UsageReadObservation,
|
||||
pub items: Vec<UsageReadItem>,
|
||||
pub complete: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct UsageReadOpened {
|
||||
pub reader: Option<String>,
|
||||
pub batch: UsageReadBatch,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct UsageReadHandle {
|
||||
pub reader: String,
|
||||
}
|
||||
@@ -38,6 +38,13 @@ pub trait PeerTransport: Send + Sync {
|
||||
/// 账号用例读取公共互联观测,显式刷新等待额度查询完成
|
||||
#[async_trait]
|
||||
pub trait PeerUpstreamService: gateway_core::account::UpstreamInfoSource {
|
||||
/// 仅在指定持久观测与当前调度缓存一致时估算容量
|
||||
fn observed_concurrency_limit(
|
||||
&self,
|
||||
version: &gateway_core::account::QuotaObservationVersion,
|
||||
now: std::time::SystemTime,
|
||||
) -> Option<gateway_core::account::AccountConcurrencyLimit>;
|
||||
|
||||
async fn account_changed(
|
||||
&self,
|
||||
id: &gateway_core::account::ProviderAccountId,
|
||||
|
||||
@@ -2,6 +2,14 @@
|
||||
|
||||
use crate::model::{auth::AdminAuditEvent, usage_data::*};
|
||||
use async_trait::async_trait;
|
||||
use gateway_core::lifecycle::CancellationToken;
|
||||
|
||||
#[async_trait]
|
||||
pub trait UsageReadSession: Send {
|
||||
async fn next_batch(&mut self) -> UsageDataResult<UsageReadBatch>;
|
||||
fn is_finished(&self) -> bool;
|
||||
async fn close(&mut self);
|
||||
}
|
||||
|
||||
pub struct StoredUsageWrite {
|
||||
pub changes: Vec<(UsageRecordChange, AdminAuditEvent)>,
|
||||
@@ -13,6 +21,11 @@ pub struct StoredAggregateWrite {
|
||||
|
||||
#[async_trait]
|
||||
pub trait UsageDataStore: Send + Sync {
|
||||
async fn open_read(
|
||||
&self,
|
||||
query: UsageReadQuery,
|
||||
cancellation: CancellationToken,
|
||||
) -> UsageDataResult<Box<dyn UsageReadSession>>;
|
||||
async fn list_aggregates(
|
||||
&self,
|
||||
request: ListUsageAggregates,
|
||||
|
||||
@@ -60,6 +60,7 @@ pub trait AccountGroupService: Send + Sync {
|
||||
}
|
||||
|
||||
pub(crate) struct DefaultAccountGroupService {
|
||||
upstream_info: Option<Arc<dyn crate::ports::peer::PeerUpstreamService>>,
|
||||
store: Arc<dyn AccountGroupStore>,
|
||||
runtime: Arc<dyn AccountRuntimeStore>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
@@ -68,11 +69,13 @@ pub(crate) struct DefaultAccountGroupService {
|
||||
impl DefaultAccountGroupService {
|
||||
#[must_use]
|
||||
pub(crate) fn new(
|
||||
upstream_info: Option<Arc<dyn crate::ports::peer::PeerUpstreamService>>,
|
||||
store: Arc<dyn AccountGroupStore>,
|
||||
runtime: Arc<dyn AccountRuntimeStore>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
) -> Self {
|
||||
Self {
|
||||
upstream_info,
|
||||
store,
|
||||
runtime,
|
||||
snapshot,
|
||||
@@ -105,7 +108,7 @@ impl DefaultAccountGroupService {
|
||||
Default::default()
|
||||
}
|
||||
};
|
||||
project_group_runtime(records, &members, &runtime);
|
||||
project_group_runtime(records, &members, &runtime, self.upstream_info.as_deref());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -215,6 +218,7 @@ fn project_group_runtime(
|
||||
records: &mut [AccountGroupRecord],
|
||||
members: &[AccountGroupMemberFact],
|
||||
runtime: &crate::model::accounts::AccountRuntimeSnapshot,
|
||||
upstream: Option<&dyn crate::ports::peer::PeerUpstreamService>,
|
||||
) {
|
||||
let now = SystemTime::now();
|
||||
let mut status_by_account = BTreeMap::new();
|
||||
@@ -228,7 +232,14 @@ fn project_group_runtime(
|
||||
.or_insert_with(|| resolve_account_status(&status, now).status);
|
||||
slots_by_account
|
||||
.entry(member.account_id.as_str())
|
||||
.or_insert(member.total_slots);
|
||||
.or_insert_with(|| {
|
||||
member
|
||||
.inherited_version
|
||||
.as_ref()
|
||||
.and_then(|version| upstream?.observed_concurrency_limit(version, now))
|
||||
.map(|limit| u64::from(limit.get()))
|
||||
.or(member.total_slots)
|
||||
});
|
||||
accounts_by_group
|
||||
.entry(member.group_id.as_str())
|
||||
.or_default()
|
||||
|
||||
@@ -12,7 +12,7 @@ use chrono::{Duration, Utc};
|
||||
use futures::StreamExt as _;
|
||||
use gateway_core::{
|
||||
account::ProviderAccountId,
|
||||
engine::probe::{AccountProbe, AccountProbeRequest},
|
||||
engine::probe::{AccountProbe, AccountProbeRequest, AccountProbeTarget},
|
||||
routing::{ProviderKind, UpstreamModelId},
|
||||
runtime::SnapshotControl,
|
||||
};
|
||||
@@ -178,7 +178,7 @@ pub trait AccountsService: Send + Sync {
|
||||
|
||||
async fn test_connection(
|
||||
&self,
|
||||
account_id: ProviderAccountId,
|
||||
target: AccountProbeTarget,
|
||||
upstream_model: UpstreamModelId,
|
||||
) -> Result<AccountConnectionTestEventStream, AdminError>;
|
||||
}
|
||||
@@ -504,7 +504,11 @@ impl DefaultAccountsService {
|
||||
});
|
||||
Ok(AccountDirectoryItem {
|
||||
protection_rules: stored.protection_rules,
|
||||
capacity: stored.capacity,
|
||||
capacity: super::peer_link::upstream::account_capacity(
|
||||
&stored.account,
|
||||
stored.capacity,
|
||||
"a,
|
||||
),
|
||||
capabilities: provider
|
||||
.account_capabilities(account_id, &stored.account.authentication_kind),
|
||||
plan_type_display: self.providers.resolve_account_plan(
|
||||
@@ -618,7 +622,11 @@ impl AccountsService for DefaultAccountsService {
|
||||
used_slots: in_flight
|
||||
.as_ref()
|
||||
.map(|counts| counts.get(&item.account.id).copied().unwrap_or(0)),
|
||||
..item.capacity
|
||||
..super::peer_link::upstream::account_capacity(
|
||||
&item.account,
|
||||
item.capacity,
|
||||
"a,
|
||||
)
|
||||
},
|
||||
capabilities,
|
||||
plan_type_display: providers.resolve_account_plan(
|
||||
@@ -1123,10 +1131,10 @@ impl AccountsService for DefaultAccountsService {
|
||||
|
||||
async fn test_connection(
|
||||
&self,
|
||||
account_id: ProviderAccountId,
|
||||
target: AccountProbeTarget,
|
||||
upstream_model: UpstreamModelId,
|
||||
) -> Result<AccountConnectionTestEventStream, AdminError> {
|
||||
let (stored, provider) = self.provider_for_account(&account_id).await?;
|
||||
let (stored, provider) = self.provider_for_account(target.account_id()).await?;
|
||||
let account = stored.account;
|
||||
let model = upstream_model.as_str().to_owned();
|
||||
let operation = provider
|
||||
@@ -1149,7 +1157,7 @@ impl AccountsService for DefaultAccountsService {
|
||||
let result = probe
|
||||
.probe(
|
||||
AccountProbeRequest {
|
||||
account_id,
|
||||
target,
|
||||
provider_kind: account.provider_kind,
|
||||
upstream_model,
|
||||
operation,
|
||||
|
||||
@@ -133,7 +133,7 @@ impl DefaultObservabilityService {
|
||||
self.settings.load_runtime_settings(),
|
||||
)
|
||||
.map_err(|error| map_store_error(error, "dashboard"))?;
|
||||
let runtime_slots = observation.runtime_slots;
|
||||
let runtime_slots = observation.runtime_slots.as_ref();
|
||||
self.enrich_list_records(&mut observation.recent_requests)
|
||||
.await;
|
||||
self.enrich_dashboard_quotas(&mut observation.account_usage)
|
||||
@@ -180,11 +180,19 @@ impl DefaultObservabilityService {
|
||||
.providers
|
||||
.dashboard_wire_profiles(&settings.request_profiles);
|
||||
let max_concurrent_per_account = u64::from(settings.values.max_concurrent_per_account);
|
||||
let (inherited_accounts, overridden_slots) = runtime_slots
|
||||
.as_ref()
|
||||
let (mut inherited_accounts, mut overridden_slots) = runtime_slots
|
||||
.map_or((observation.provider_accounts.normal, 0), |slots| {
|
||||
(slots.inherited_accounts, slots.overridden_slots)
|
||||
});
|
||||
if let (Some(peer), Some(slots)) = (&self.upstream_info, runtime_slots) {
|
||||
let now = std::time::SystemTime::now();
|
||||
for version in &slots.inherited_versions {
|
||||
if let Some(limit) = peer.observed_concurrency_limit(version, now) {
|
||||
inherited_accounts -= 1;
|
||||
overridden_slots = overridden_slots.saturating_add(u64::from(limit.get()));
|
||||
}
|
||||
}
|
||||
}
|
||||
let total_slots = (max_concurrent_per_account > 0 || inherited_accounts == 0).then(|| {
|
||||
inherited_accounts
|
||||
.saturating_mul(max_concurrent_per_account)
|
||||
|
||||
@@ -117,6 +117,7 @@ impl PeerLinkService for DefaultPeerLinkService {
|
||||
.budget
|
||||
.limits
|
||||
.access(key.budget.daily_used_usd, key.budget.weekly_used_usd),
|
||||
key_limits: key.limits,
|
||||
budget: key.budget,
|
||||
budget_observed_at,
|
||||
pool: None,
|
||||
|
||||
@@ -57,6 +57,7 @@ impl PeerInfoService {
|
||||
scheduling: None,
|
||||
summary: None,
|
||||
unsupported: false,
|
||||
unavailable: None,
|
||||
},
|
||||
));
|
||||
}
|
||||
@@ -92,6 +93,7 @@ impl PeerInfoService {
|
||||
(
|
||||
connection.clone(),
|
||||
cached.and_then(|cached| cached.summary.clone()),
|
||||
cached.and_then(|cached| cached.unavailable),
|
||||
cached.is_some_and(|cached| cached.unsupported),
|
||||
)
|
||||
})
|
||||
@@ -99,10 +101,11 @@ impl PeerInfoService {
|
||||
};
|
||||
Ok(summaries
|
||||
.into_iter()
|
||||
.map(|(connection, peer, unsupported)| {
|
||||
.map(|(connection, peer, unavailable, unsupported)| {
|
||||
let quota = Self::project_peer(
|
||||
&connection,
|
||||
peer.as_deref(),
|
||||
unavailable,
|
||||
self.last_error(&connection)
|
||||
.or(unsupported.then_some(UpstreamInfoFailure::Unsupported)),
|
||||
now,
|
||||
@@ -137,6 +140,8 @@ impl PeerInfoService {
|
||||
summary,
|
||||
unsupported: observation
|
||||
.is_some_and(|value| snapshot::is_unsupported(&value.quota, connection)),
|
||||
unavailable: observation
|
||||
.and_then(|value| snapshot::unavailable_reason(&value.quota, connection)),
|
||||
};
|
||||
let previous = self
|
||||
.observations
|
||||
|
||||
@@ -25,6 +25,7 @@ use gateway_core::{
|
||||
error::StoreError,
|
||||
task::{WorkerContribution, WorkerDefinitionError},
|
||||
};
|
||||
pub(crate) use projection::account_capacity;
|
||||
use service::PeerInfoService;
|
||||
|
||||
/// 组合根独立注册公共互联 Worker,不属于任何具体 Provider 的任务集合
|
||||
@@ -62,6 +63,24 @@ pub fn initialize(
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl PeerUpstreamService for PeerInfoService {
|
||||
fn observed_concurrency_limit(
|
||||
&self,
|
||||
version: &gateway_core::account::QuotaObservationVersion,
|
||||
now: std::time::SystemTime,
|
||||
) -> Option<gateway_core::account::AccountConcurrencyLimit> {
|
||||
let observations = self
|
||||
.observations
|
||||
.read()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
let cached = observations.get(&version.account_id)?;
|
||||
if cached.observed_at != Some(version.observed_at)
|
||||
|| cached.connection.credential_revision != version.credential_revision
|
||||
{
|
||||
return None;
|
||||
}
|
||||
cached.scheduling.as_ref()?.concurrency_limit(now)
|
||||
}
|
||||
|
||||
async fn account_changed(&self, id: &ProviderAccountId) -> Result<(), StoreError> {
|
||||
PeerInfoService::account_changed(self, id).await
|
||||
}
|
||||
|
||||
@@ -16,10 +16,31 @@ use gateway_protocol::peer::{PeerComponentStatus, PeerLink};
|
||||
|
||||
use super::{PeerInfoService, quota, snapshot};
|
||||
|
||||
pub(crate) fn account_capacity(
|
||||
account: &crate::model::accounts::AccountRecord,
|
||||
mut capacity: crate::model::accounts::AccountCapacity,
|
||||
quota: &ProviderQuota,
|
||||
) -> crate::model::accounts::AccountCapacity {
|
||||
if account.concurrency_limit.is_none()
|
||||
&& let Some(info) = "a.upstream
|
||||
&& !info.stale
|
||||
&& info
|
||||
.observed_at
|
||||
.is_some_and(|at| snapshot::fresh_at(at.into(), SystemTime::now()))
|
||||
&& let Some(limit) = info
|
||||
.max_concurrency
|
||||
.and_then(gateway_core::account::AccountConcurrencyLimit::from_key_limit)
|
||||
{
|
||||
capacity.total_slots = Some(u64::from(limit.get()));
|
||||
}
|
||||
capacity
|
||||
}
|
||||
|
||||
impl PeerInfoService {
|
||||
pub(super) fn project_peer(
|
||||
connection: &UpstreamConnection,
|
||||
peer: Option<&PeerLink>,
|
||||
unavailable: Option<UpstreamInfoFailure>,
|
||||
last_error: Option<UpstreamInfoFailure>,
|
||||
now: SystemTime,
|
||||
) -> ProviderQuota {
|
||||
@@ -29,12 +50,17 @@ impl PeerInfoService {
|
||||
let stale = peer.as_ref().is_none_or(|peer| {
|
||||
!snapshot::component_fresh(Some(&peer.observed_at), Some(&peer.expires_at), now)
|
||||
});
|
||||
let limits = peer.and_then(|peer| peer.key_limits);
|
||||
let info = ProviderUpstreamInfo {
|
||||
discovered: peer.is_some(),
|
||||
discovered: peer.is_some() || unavailable.is_some(),
|
||||
observed_at,
|
||||
stale,
|
||||
connection_revision: connection.revision.clone(),
|
||||
last_error: last_error.map(|error| error.as_str().to_owned()),
|
||||
last_error: last_error
|
||||
.or(unavailable)
|
||||
.map(|error| error.as_str().to_owned()),
|
||||
max_concurrency: limits.and_then(|limits| limits.max_concurrency),
|
||||
requests_per_minute: limits.and_then(|limits| limits.requests_per_minute),
|
||||
pool: peer
|
||||
.as_ref()
|
||||
.and_then(|peer| peer.pool.as_ref())
|
||||
|
||||
@@ -135,21 +135,30 @@ impl PeerInfoService {
|
||||
}
|
||||
.max(retry_after.unwrap_or_default());
|
||||
state.next_attempt = Instant::now() + delay;
|
||||
state.manual_after = Instant::now()
|
||||
+ (if failure.is_some_and(|error| error != UpstreamInfoFailure::Unsupported) {
|
||||
delay
|
||||
} else {
|
||||
Duration::from_secs(15)
|
||||
})
|
||||
.max(retry_after.unwrap_or_default());
|
||||
// 管理员修复授权后可主动复核,不继承自动重试的长退避,仍遵守上游限流
|
||||
state.manual_after =
|
||||
Instant::now() + Duration::from_secs(15).max(retry_after.unwrap_or_default());
|
||||
state.last_error = failure;
|
||||
let document = match result {
|
||||
Ok(document) => Some(document),
|
||||
// 空观测清除已撤销内容,并用共有观察时间阻止更早请求重新写回
|
||||
Err(UpstreamInfoFailure::Unauthorized) => {
|
||||
Some(OpaqueProviderData::new(Default::default()))
|
||||
Err(error @ (UpstreamInfoFailure::Unauthorized | UpstreamInfoFailure::Unsupported)) => {
|
||||
let discovered = self
|
||||
.read_summary(id)
|
||||
.await
|
||||
.map_err(store_error)?
|
||||
.and_then(|quota| quota.upstream)
|
||||
.is_some_and(|info| {
|
||||
info.connection_revision == connection.revision && info.discovered
|
||||
});
|
||||
// 首次探测失败不能证明对方是互联节点;已确认的连接只留下无共享内容的识别记录
|
||||
Some(if discovered {
|
||||
snapshot::unavailable(&connection, error)
|
||||
} else if error == UpstreamInfoFailure::Unsupported {
|
||||
snapshot::unsupported(&connection)
|
||||
} else {
|
||||
OpaqueProviderData::new(Default::default())
|
||||
})
|
||||
}
|
||||
Err(UpstreamInfoFailure::Unsupported) => Some(snapshot::unsupported(&connection)),
|
||||
Err(_) => None,
|
||||
};
|
||||
if let Some(document) = document {
|
||||
|
||||
@@ -63,6 +63,7 @@ struct ModelObservation {
|
||||
|
||||
pub(super) struct SchedulingObservation {
|
||||
validity: Validity,
|
||||
concurrency_limit: Option<gateway_core::account::AccountConcurrencyLimit>,
|
||||
exhausted_budgets: Vec<Validity>,
|
||||
models: HashMap<String, Option<ModelObservation>>,
|
||||
}
|
||||
@@ -166,43 +167,52 @@ impl SchedulingObservation {
|
||||
}
|
||||
Some(Self {
|
||||
validity,
|
||||
concurrency_limit: peer
|
||||
.key_limits
|
||||
.and_then(|limits| limits.max_concurrency)
|
||||
.and_then(gateway_core::account::AccountConcurrencyLimit::from_key_limit),
|
||||
exhausted_budgets,
|
||||
models,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn fresh(&self, now: SystemTime) -> bool {
|
||||
self.validity.fresh(now)
|
||||
&& now
|
||||
.duration_since(self.validity.observed_at)
|
||||
.unwrap_or_default()
|
||||
<= snapshot::FRESH_FOR
|
||||
self.validity.fresh(now) && snapshot::fresh_at(self.validity.observed_at, now)
|
||||
}
|
||||
|
||||
pub(super) fn signal(&self, model: &str, now: SystemTime) -> Option<UpstreamSchedulingSignal> {
|
||||
pub(super) fn concurrency_limit(
|
||||
&self,
|
||||
now: SystemTime,
|
||||
) -> Option<gateway_core::account::AccountConcurrencyLimit> {
|
||||
self.concurrency_limit.filter(|_| self.fresh(now))
|
||||
}
|
||||
|
||||
pub(super) fn signal(
|
||||
&self,
|
||||
model: Option<&str>,
|
||||
now: SystemTime,
|
||||
) -> Option<UpstreamSchedulingSignal> {
|
||||
if !self.fresh(now) {
|
||||
return None;
|
||||
}
|
||||
if self
|
||||
let budget_exhausted = self
|
||||
.exhausted_budgets
|
||||
.iter()
|
||||
.any(|validity| validity.fresh(now))
|
||||
{
|
||||
return Some(UpstreamSchedulingSignal {
|
||||
exhausted: true,
|
||||
remaining_percent: None,
|
||||
});
|
||||
}
|
||||
let model = self.models.get(model)?.as_ref()?;
|
||||
if !model.validity.fresh(now) {
|
||||
.any(|validity| validity.fresh(now));
|
||||
let model = model
|
||||
.and_then(|id| self.models.get(id)?.as_ref())
|
||||
.filter(|model| model.validity.fresh(now));
|
||||
if !budget_exhausted && model.is_none() && self.concurrency_limit.is_none() {
|
||||
return None;
|
||||
}
|
||||
Some(UpstreamSchedulingSignal {
|
||||
exhausted: model.exhausted,
|
||||
remaining_percent: model
|
||||
.account
|
||||
.as_ref()
|
||||
.and_then(|account| account.remaining_percent(now)),
|
||||
concurrency_limit: self.concurrency_limit,
|
||||
exhausted: budget_exhausted || model.is_some_and(|model| model.exhausted),
|
||||
remaining_percent: if budget_exhausted {
|
||||
None
|
||||
} else {
|
||||
model.and_then(|model| model.account.as_ref()?.remaining_percent(now))
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -35,6 +35,7 @@ pub(super) struct CachedObservation {
|
||||
pub scheduling: Option<Arc<SchedulingObservation>>,
|
||||
pub summary: Option<Arc<PeerLink>>,
|
||||
pub unsupported: bool,
|
||||
pub unavailable: Option<UpstreamInfoFailure>,
|
||||
}
|
||||
|
||||
pub(crate) struct PeerInfoService {
|
||||
@@ -175,12 +176,16 @@ impl PeerInfoService {
|
||||
let unsupported = observation
|
||||
.as_ref()
|
||||
.is_some_and(|value| snapshot::is_unsupported(&value.quota, &connection));
|
||||
let unavailable = observation
|
||||
.as_ref()
|
||||
.and_then(|value| snapshot::unavailable_reason(&value.quota, &connection));
|
||||
if self.current_connection(id).await?.as_ref() != Some(&connection) {
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(Some(Self::project_peer(
|
||||
&connection,
|
||||
peer.as_ref(),
|
||||
unavailable,
|
||||
self.last_error(&connection)
|
||||
.or(unsupported.then_some(UpstreamInfoFailure::Unsupported)),
|
||||
SystemTime::now(),
|
||||
@@ -285,7 +290,7 @@ impl UpstreamInfoSource for PeerInfoService {
|
||||
fn signal(
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
model: &str,
|
||||
model: Option<&str>,
|
||||
now: SystemTime,
|
||||
) -> Option<UpstreamSchedulingSignal> {
|
||||
if !account.enabled() {
|
||||
|
||||
@@ -30,7 +30,7 @@ pub(super) fn unsupported(connection: &UpstreamConnection) -> OpaqueProviderData
|
||||
|
||||
pub(super) fn is_unsupported(data: &OpaqueProviderData, connection: &UpstreamConnection) -> bool {
|
||||
let fields = data.expose_to_provider();
|
||||
fields.get("object").and_then(serde_json::Value::as_str) == Some("cpr.peer.unsupported")
|
||||
(fields.get("object").and_then(serde_json::Value::as_str) == Some("cpr.peer.unsupported")
|
||||
&& fields
|
||||
.get("schema_version")
|
||||
.and_then(serde_json::Value::as_u64)
|
||||
@@ -38,7 +38,42 @@ pub(super) fn is_unsupported(data: &OpaqueProviderData, connection: &UpstreamCon
|
||||
&& fields
|
||||
.get("connection_revision")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
== Some(connection.revision.as_str())
|
||||
== Some(connection.revision.as_str()))
|
||||
|| unavailable_reason(data, connection) == Some(UpstreamInfoFailure::Unsupported)
|
||||
}
|
||||
|
||||
// 失权后只保留同一连接已确认互联的事实,不保留已撤销的池、预算与模型数据
|
||||
pub(super) fn unavailable(
|
||||
connection: &UpstreamConnection,
|
||||
reason: UpstreamInfoFailure,
|
||||
) -> OpaqueProviderData {
|
||||
OpaqueProviderData::new(serde_json::Map::from_iter([
|
||||
("object".into(), "cpr.peer.unavailable".into()),
|
||||
("schema_version".into(), 1.into()),
|
||||
(
|
||||
"connection_revision".into(),
|
||||
connection.revision.clone().into(),
|
||||
),
|
||||
("reason".into(), reason.as_str().into()),
|
||||
]))
|
||||
}
|
||||
|
||||
pub(super) fn unavailable_reason(
|
||||
data: &OpaqueProviderData,
|
||||
connection: &UpstreamConnection,
|
||||
) -> Option<UpstreamInfoFailure> {
|
||||
let fields = data.expose_to_provider();
|
||||
if fields.get("object")?.as_str()? != "cpr.peer.unavailable"
|
||||
|| fields.get("schema_version")?.as_u64()? != 1
|
||||
|| fields.get("connection_revision")?.as_str()? != connection.revision
|
||||
{
|
||||
return None;
|
||||
}
|
||||
match fields.get("reason")?.as_str()? {
|
||||
"unauthorized" => Some(UpstreamInfoFailure::Unauthorized),
|
||||
"unsupported" => Some(UpstreamInfoFailure::Unsupported),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn timestamp(value: &str) -> Option<SystemTime> {
|
||||
@@ -54,10 +89,12 @@ pub(super) fn datetime(value: &str) -> Option<DateTime<Utc>> {
|
||||
}
|
||||
|
||||
pub(super) fn fresh(observed: &str, now: SystemTime) -> bool {
|
||||
timestamp(observed).is_some_and(|observed| {
|
||||
observed <= now + Duration::from_secs(60)
|
||||
&& now.duration_since(observed).unwrap_or_default() <= FRESH_FOR
|
||||
})
|
||||
timestamp(observed).is_some_and(|observed| fresh_at(observed, now))
|
||||
}
|
||||
|
||||
pub(super) fn fresh_at(observed: SystemTime, now: SystemTime) -> bool {
|
||||
observed <= now + Duration::from_secs(60)
|
||||
&& now.duration_since(observed).unwrap_or_default() <= FRESH_FOR
|
||||
}
|
||||
|
||||
pub(super) fn component_fresh(
|
||||
@@ -127,6 +164,9 @@ pub(super) fn parse(
|
||||
}
|
||||
// 未授权能力不能借同名数据块越过逐层公开边界
|
||||
let disclosed = |capability| peer.capabilities.iter().any(|value| value == capability);
|
||||
if !disclosed("key_limits") {
|
||||
peer.key_limits = None;
|
||||
}
|
||||
if !disclosed("key_usage") {
|
||||
peer.usage = None;
|
||||
}
|
||||
|
||||
@@ -15,6 +15,11 @@ use std::sync::Arc;
|
||||
|
||||
#[async_trait]
|
||||
pub trait UsageDataService: Send + Sync {
|
||||
async fn open_read(
|
||||
&self,
|
||||
query: UsageReadQuery,
|
||||
cancellation: gateway_core::lifecycle::CancellationToken,
|
||||
) -> UsageDataResult<Box<dyn UsageReadSession>>;
|
||||
async fn list_aggregates(
|
||||
&self,
|
||||
request: ListUsageAggregates,
|
||||
@@ -58,6 +63,21 @@ impl DefaultUsageDataService {
|
||||
|
||||
#[async_trait]
|
||||
impl UsageDataService for DefaultUsageDataService {
|
||||
async fn open_read(
|
||||
&self,
|
||||
query: UsageReadQuery,
|
||||
cancellation: gateway_core::lifecycle::CancellationToken,
|
||||
) -> UsageDataResult<Box<dyn UsageReadSession>> {
|
||||
let filter = query.filter();
|
||||
validation::query(&filter.records, None, None)?;
|
||||
if let Some(group) = &filter.routing_group_ref {
|
||||
validation::text(group, 128)?;
|
||||
}
|
||||
if !(1..=USAGE_READ_MAX_BATCH_ROWS).contains(&query.batch_size()) {
|
||||
return Err(UsageDataError::invalid("读取批次条数应为 1 到 1024"));
|
||||
}
|
||||
self.store.open_read(query, cancellation).await
|
||||
}
|
||||
async fn list_aggregates(
|
||||
&self,
|
||||
request: ListUsageAggregates,
|
||||
|
||||
@@ -7,7 +7,6 @@ use gateway_core::account::OpaqueProviderData;
|
||||
|
||||
use gateway_admin::model::{
|
||||
PageSize, Revision,
|
||||
accounts::AccountUsage,
|
||||
auth::LoginCommand,
|
||||
client_keys::ClientKeyPageSize,
|
||||
observability::{DecimalAmount, RequestOutcome, TimeRange},
|
||||
@@ -146,57 +145,6 @@ fn representative_quota_should_prefer_account_wide_window_over_model_specific_wi
|
||||
assert_eq!(quota.representative_used_percent(), Some(5.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_window_should_not_use_daily_rolling_usage_for_weekly_statistics() {
|
||||
let usage = AccountUsage {
|
||||
account_id: "acct_xai".to_owned(),
|
||||
request_count: 5,
|
||||
success_count: 5,
|
||||
input_tokens: Some(100),
|
||||
output_tokens: Some(20),
|
||||
cached_input_tokens: Some(80),
|
||||
cache_write_input_tokens: Some(0),
|
||||
reasoning_output_tokens: Some(0),
|
||||
image_input_tokens: Some(0),
|
||||
image_output_tokens: Some(0),
|
||||
image_request_count: 0,
|
||||
image_request_failed_count: 0,
|
||||
total_tokens: Some(120),
|
||||
cost_coverage: Default::default(),
|
||||
costs: Vec::new(),
|
||||
last_used_at: Some(Utc::now()),
|
||||
request_buckets: Vec::new(),
|
||||
models: Vec::new(),
|
||||
};
|
||||
let quota = ProviderQuota {
|
||||
upstream: None,
|
||||
credits: None,
|
||||
plan_type: None,
|
||||
observed_at: None,
|
||||
refresh_token_expires_at: None,
|
||||
windows: vec![ProviderQuotaWindow {
|
||||
upstream: None,
|
||||
key: "free-rolling-24h".to_owned(),
|
||||
group: "shortTerm".to_owned(),
|
||||
label: "日限额".to_owned(),
|
||||
limit_id: None,
|
||||
limit_name: None,
|
||||
role: None,
|
||||
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
|
||||
window_seconds: Some(86_400),
|
||||
used_percent: None,
|
||||
reset_at: None,
|
||||
limit_reached: false,
|
||||
local_usage: Some(usage),
|
||||
provider_data: None,
|
||||
}],
|
||||
limit_reached: false,
|
||||
provider_data: None,
|
||||
};
|
||||
|
||||
assert_eq!(quota.usage_window(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exhausted_quota_should_project_full_usage_to_only_the_representative_window() {
|
||||
let mut quota = ProviderQuota {
|
||||
|
||||
@@ -65,6 +65,36 @@ fn usage_window_should_not_fallback_to_short_daily_or_incomplete_windows() {
|
||||
assert_eq!(quota.usage_window(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_window_should_preserve_provider_rolling_daily_statistics() {
|
||||
let mut rolling = usage_window("rolling-day", "shortTerm", 86_400);
|
||||
rolling.reset_at = None;
|
||||
rolling.used_percent = None;
|
||||
let mut quota = usage_quota(vec![rolling.clone()]);
|
||||
assert_eq!(
|
||||
selected_usage_window("a),
|
||||
Some(("rolling-day", AccountUsagePeriod::Daily))
|
||||
);
|
||||
|
||||
quota.limit_reached = true;
|
||||
quota.apply_limit_reached_display();
|
||||
assert_eq!(
|
||||
selected_usage_window("a),
|
||||
Some(("rolling-day", AccountUsagePeriod::Daily))
|
||||
);
|
||||
|
||||
quota
|
||||
.windows
|
||||
.push(usage_window("week", "shortTerm", 7 * 86_400));
|
||||
assert_eq!(
|
||||
selected_usage_window("a),
|
||||
Some(("week", AccountUsagePeriod::Weekly))
|
||||
);
|
||||
|
||||
rolling.local_usage_attribution = QuotaLocalUsageAttribution::Unavailable;
|
||||
assert_eq!(usage_quota(vec![rolling]).usage_window(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_window_should_preserve_provider_order_for_equal_periods() {
|
||||
let quota = usage_quota(vec![
|
||||
|
||||
@@ -162,7 +162,10 @@ fn source_selection_keeps_provider_order_with_the_maximum_window_count() {
|
||||
fn short_or_model_specific_windows_are_not_account_capacity() {
|
||||
let mut model = window("model", 7);
|
||||
model.local_usage_attribution = QuotaLocalUsageAttribution::Unavailable;
|
||||
for item in forecast("a(vec![window("day", 1), model])) {
|
||||
let mut rolling = window("rolling-day", 1);
|
||||
rolling.reset_at = None;
|
||||
rolling.used_percent = None;
|
||||
for item in forecast("a(vec![window("day", 1), rolling, model])) {
|
||||
assert!(item.source.is_none());
|
||||
assert!(item.unavailable_reason.is_some());
|
||||
assert!(item.estimated_tokens.is_none());
|
||||
|
||||
@@ -373,6 +373,7 @@ fn group_record() -> AccountGroupRecord {
|
||||
|
||||
fn member(account_id: &str, total_slots: u64) -> AccountGroupMemberFact {
|
||||
AccountGroupMemberFact {
|
||||
inherited_version: None,
|
||||
group_id: group_id(),
|
||||
account_id: account_id.to_owned(),
|
||||
status: AccountStatusFacts {
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
mod reset_credit_automation;
|
||||
|
||||
use gateway_core::engine::probe::AccountProbeTarget;
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
sync::{Arc, Mutex},
|
||||
@@ -1335,7 +1336,9 @@ async fn connection_test_should_probe_unavailable_account() {
|
||||
let events = services
|
||||
.accounts()
|
||||
.test_connection(
|
||||
ProviderAccountId::new("acct_test").expect("account ID"),
|
||||
AccountProbeTarget::ConnectionTest(
|
||||
ProviderAccountId::new("acct_test").expect("account ID"),
|
||||
),
|
||||
gateway_core::routing::UpstreamModelId::new("grok-4.5").expect("model"),
|
||||
)
|
||||
.await
|
||||
@@ -1372,7 +1375,9 @@ async fn connection_test_rate_limited_probe_returns_provider_failure() {
|
||||
let events = services
|
||||
.accounts()
|
||||
.test_connection(
|
||||
ProviderAccountId::new("acct_test").expect("account ID"),
|
||||
AccountProbeTarget::ConnectionTest(
|
||||
ProviderAccountId::new("acct_test").expect("account ID"),
|
||||
),
|
||||
gateway_core::routing::UpstreamModelId::new("grok-4.5").expect("model"),
|
||||
)
|
||||
.await
|
||||
@@ -1405,7 +1410,9 @@ async fn connection_test_should_preserve_disabled_account_status() {
|
||||
let events = services
|
||||
.accounts()
|
||||
.test_connection(
|
||||
ProviderAccountId::new("acct_test").expect("account ID"),
|
||||
AccountProbeTarget::ConnectionTest(
|
||||
ProviderAccountId::new("acct_test").expect("account ID"),
|
||||
),
|
||||
gateway_core::routing::UpstreamModelId::new("grok-4.5").expect("model"),
|
||||
)
|
||||
.await
|
||||
@@ -2686,12 +2693,15 @@ async fn peer_account_usage_should_follow_budget_windows_without_borrowing_pool_
|
||||
for longest in [86_400, 604_800] {
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
let mut account = account_record("openai");
|
||||
account.concurrency_limit = None;
|
||||
account.authentication_kind = "api_key".to_owned();
|
||||
account.created_at = Utc::now() - TimeDelta::days(60);
|
||||
let store = FakeAccountStore::with_account(account, events());
|
||||
let reset_at = Utc::now() + TimeDelta::hours(1);
|
||||
let mut quota = empty_quota();
|
||||
quota.upstream = Some(ProviderUpstreamInfo {
|
||||
max_concurrency: Some(7),
|
||||
requests_per_minute: Some(60),
|
||||
discovered: true,
|
||||
observed_at: Some(Utc::now()),
|
||||
stale: false,
|
||||
@@ -2762,6 +2772,8 @@ async fn peer_account_usage_should_follow_budget_windows_without_borrowing_pool_
|
||||
page.items[0].usage.as_ref().unwrap().total_tokens,
|
||||
Some(1_234)
|
||||
);
|
||||
assert_eq!(page.items[0].capacity.total_slots, Some(7));
|
||||
assert_eq!(page.items[0].account.concurrency_limit, None);
|
||||
let queries = store.quota_window_queries();
|
||||
assert_eq!(queries.len(), window_usage.len());
|
||||
for (query, expected) in queries.iter().zip(&window_usage) {
|
||||
|
||||
@@ -143,7 +143,7 @@ impl UpstreamInfoSource for PendingPeer {
|
||||
fn signal(
|
||||
&self,
|
||||
_: &ProviderAccount,
|
||||
_: &str,
|
||||
_: Option<&str>,
|
||||
_: std::time::SystemTime,
|
||||
) -> Option<UpstreamSchedulingSignal> {
|
||||
None
|
||||
@@ -152,6 +152,14 @@ impl UpstreamInfoSource for PendingPeer {
|
||||
|
||||
#[async_trait]
|
||||
impl PeerUpstreamService for PendingPeer {
|
||||
fn observed_concurrency_limit(
|
||||
&self,
|
||||
_: &gateway_core::account::QuotaObservationVersion,
|
||||
_: std::time::SystemTime,
|
||||
) -> Option<gateway_core::account::AccountConcurrencyLimit> {
|
||||
None
|
||||
}
|
||||
|
||||
async fn account_changed(&self, id: &ProviderAccountId) -> Result<(), StoreError> {
|
||||
self.changed.lock().unwrap().push(id.as_str().into());
|
||||
Ok(())
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
//! 账号保护恢复编排:到期探测、快照参数与失败顺延
|
||||
|
||||
use gateway_core::engine::probe::AccountProbeTarget;
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
@@ -242,6 +243,10 @@ impl AccountProbe for FailingProbe {
|
||||
request: AccountProbeRequest,
|
||||
_: Option<Arc<gateway_core::routing::RuntimeSnapshot>>,
|
||||
) -> futures::future::BoxFuture<'_, Result<AccountProbeResult, AccountProbeError>> {
|
||||
assert!(matches!(
|
||||
request.target,
|
||||
AccountProbeTarget::FreezeRecovery(_)
|
||||
));
|
||||
assert_eq!(request.upstream_model.as_str(), "snapshot-model");
|
||||
Box::pin(async {
|
||||
Err(AccountProbeError::new(
|
||||
|
||||
@@ -448,6 +448,7 @@ async fn dashboard_summary_should_project_rebuildable_runtime_slots() {
|
||||
let now = Utc::now();
|
||||
let store = Arc::new(FixtureObservabilityStore::new(observation_range(now)));
|
||||
store.replace_runtime_slots(Some(DashboardRuntimeSlots {
|
||||
inherited_versions: Vec::new(),
|
||||
inherited_accounts: 2,
|
||||
overridden_slots: 5,
|
||||
used_slots: Some(2),
|
||||
@@ -476,6 +477,7 @@ async fn dashboard_capacity_distinguishes_unlimited_inheritance_finite_overrides
|
||||
let now = Utc::now();
|
||||
let store = Arc::new(FixtureObservabilityStore::new(observation_range(now)));
|
||||
store.replace_runtime_slots(Some(DashboardRuntimeSlots {
|
||||
inherited_versions: Vec::new(),
|
||||
inherited_accounts,
|
||||
overridden_slots,
|
||||
used_slots: Some(2),
|
||||
@@ -1212,7 +1214,7 @@ impl ObservabilityStore for FixtureObservabilityStore {
|
||||
range: query.range,
|
||||
totals: Default::default(),
|
||||
provider_accounts: AccountPoolMetrics::default(),
|
||||
runtime_slots: *self.runtime_slots.lock().expect("runtime slots"),
|
||||
runtime_slots: self.runtime_slots.lock().expect("runtime slots").clone(),
|
||||
trend: self.trend.lock().expect("trend").clone(),
|
||||
account_usage: self.account_usage.lock().expect("account usage").clone(),
|
||||
recent_requests: Vec::new(),
|
||||
@@ -1482,6 +1484,11 @@ fn total_record(
|
||||
provider_account_upstream_peer: None,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
error_kind: None,
|
||||
websocket_close_code: None,
|
||||
continuation_unavailable_reason: None,
|
||||
upstream_connection_exit_reason: None,
|
||||
recovered: false,
|
||||
client_status_code: None,
|
||||
outcome: gateway_admin::model::observability::RequestOutcome::Succeeded,
|
||||
client_api_key_name: Some("Production".to_owned()),
|
||||
|
||||
@@ -241,7 +241,7 @@ async fn discovery_uses_provider_capabilities_and_publishes_neutral_signals() {
|
||||
let reader = UpstreamInfoReader::default();
|
||||
assert!(
|
||||
reader
|
||||
.signal(&initial_accounts[0], "model", SystemTime::now())
|
||||
.signal(&initial_accounts[0], Some("model"), SystemTime::now())
|
||||
.is_none()
|
||||
);
|
||||
let source = bundle.scheduling_source();
|
||||
@@ -273,7 +273,11 @@ async fn discovery_uses_provider_capabilities_and_publishes_neutral_signals() {
|
||||
);
|
||||
for account in &initial_accounts[..2] {
|
||||
if Some(account.provider().as_str()) == failed_provider {
|
||||
assert!(reader.signal(account, "model", SystemTime::now()).is_none());
|
||||
assert!(
|
||||
reader
|
||||
.signal(account, Some("model"), SystemTime::now())
|
||||
.is_none()
|
||||
);
|
||||
continue;
|
||||
}
|
||||
let observed = bundle
|
||||
@@ -288,21 +292,21 @@ async fn discovery_uses_provider_capabilities_and_publishes_neutral_signals() {
|
||||
);
|
||||
assert!(
|
||||
reader
|
||||
.signal(account, "model", SystemTime::now())
|
||||
.signal(account, Some("model"), SystemTime::now())
|
||||
.unwrap()
|
||||
.exhausted
|
||||
);
|
||||
}
|
||||
assert!(
|
||||
reader
|
||||
.signal(&initial_accounts[2], "model", SystemTime::now())
|
||||
.signal(&initial_accounts[2], Some("model"), SystemTime::now())
|
||||
.is_none()
|
||||
);
|
||||
drop(worker);
|
||||
drop(bundle);
|
||||
assert!(
|
||||
reader
|
||||
.signal(&initial_accounts[0], "model", SystemTime::now())
|
||||
.signal(&initial_accounts[0], Some("model"), SystemTime::now())
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
@@ -332,7 +336,7 @@ async fn account_edits_preserve_observation_and_refresh_budget_until_connection_
|
||||
assert!(
|
||||
bundle
|
||||
.scheduling_source()
|
||||
.signal(&edited, "model", SystemTime::now())
|
||||
.signal(&edited, Some("model"), SystemTime::now())
|
||||
.is_some()
|
||||
);
|
||||
let after = service.refresh_quota(edited.id()).await.unwrap();
|
||||
@@ -347,7 +351,7 @@ async fn account_edits_preserve_observation_and_refresh_budget_until_connection_
|
||||
assert!(
|
||||
bundle
|
||||
.scheduling_source()
|
||||
.signal(&reconnected, "model", SystemTime::now())
|
||||
.signal(&reconnected, Some("model"), SystemTime::now())
|
||||
.is_none()
|
||||
);
|
||||
service.refresh_quota(reconnected.id()).await.unwrap();
|
||||
@@ -392,7 +396,7 @@ impl PeerTransport for TestTransport {
|
||||
.push(account.id().as_str().to_owned());
|
||||
let now = Utc::now();
|
||||
Ok(serde_json::to_vec(&serde_json::json!({
|
||||
"object":"cpr.peer", "schema_version":2,
|
||||
"object":"cpr.peer", "schema_version":3,
|
||||
"observed_at":now.to_rfc3339(),
|
||||
"expires_at":(now + chrono::Duration::minutes(5)).to_rfc3339(),
|
||||
"capabilities":["key_usage"],
|
||||
|
||||
@@ -175,6 +175,174 @@ async fn uncertain_discovery_failures_retry_without_becoming_unsupported() {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn known_peer_failures_keep_identity_and_only_revocation_discards_shared_data() {
|
||||
for failure in [
|
||||
UpstreamInfoFailure::Unauthorized,
|
||||
UpstreamInfoFailure::Unsupported,
|
||||
UpstreamInfoFailure::Unavailable,
|
||||
UpstreamInfoFailure::RateLimited,
|
||||
UpstreamInfoFailure::InvalidResponse,
|
||||
] {
|
||||
let (bundle, accounts, transport, id) = fixture(failure);
|
||||
*transport.failure.lock().unwrap() = None;
|
||||
let before = bundle.service().refresh_quota(&id).await.unwrap();
|
||||
assert!(!before.windows.is_empty());
|
||||
tokio::time::advance(Duration::from_secs(15)).await;
|
||||
*transport.failure.lock().unwrap() = Some(failure.into());
|
||||
let result = bundle.service().refresh_quota(&id).await;
|
||||
assert_eq!(result.is_ok(), failure == UpstreamInfoFailure::Unsupported);
|
||||
let clears_data = matches!(
|
||||
failure,
|
||||
UpstreamInfoFailure::Unauthorized | UpstreamInfoFailure::Unsupported
|
||||
);
|
||||
let quota = bundle.service().read_quota(&id).await.unwrap().unwrap();
|
||||
let info = quota.upstream.as_ref().unwrap();
|
||||
assert!(info.discovered);
|
||||
assert_eq!(info.last_error.as_deref(), Some(failure.as_str()));
|
||||
assert_eq!(quota.windows.is_empty(), clears_data);
|
||||
assert_eq!(info.observed_at.is_none(), clears_data);
|
||||
assert_eq!(
|
||||
bundle
|
||||
.service()
|
||||
.read_summary(&id)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.upstream,
|
||||
quota.upstream
|
||||
);
|
||||
let account = accounts.accounts.lock().unwrap()[0].clone();
|
||||
assert_eq!(
|
||||
bundle
|
||||
.scheduling_source()
|
||||
.signal(&account, Some("model"), SystemTime::now())
|
||||
.is_none(),
|
||||
clears_data
|
||||
);
|
||||
|
||||
drop(bundle);
|
||||
let restarted = initialize(accounts, transport);
|
||||
let restored = restarted.service().read_quota(&id).await.unwrap().unwrap();
|
||||
assert!(restored.upstream.as_ref().unwrap().discovered);
|
||||
assert_eq!(restored.windows.is_empty(), clears_data);
|
||||
if clears_data {
|
||||
assert_eq!(
|
||||
restored.upstream.as_ref().unwrap().last_error.as_deref(),
|
||||
Some(failure.as_str())
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
restarted
|
||||
.service()
|
||||
.read_summary(&id)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.upstream,
|
||||
restored.upstream
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn known_peer_identity_is_scoped_to_the_connection_after_revocation() {
|
||||
for failure in [
|
||||
UpstreamInfoFailure::Unauthorized,
|
||||
UpstreamInfoFailure::Unsupported,
|
||||
] {
|
||||
let (bundle, accounts, transport, id) = fixture(failure);
|
||||
*transport.failure.lock().unwrap() = None;
|
||||
bundle.service().refresh_quota(&id).await.unwrap();
|
||||
tokio::time::advance(Duration::from_secs(15)).await;
|
||||
*transport.failure.lock().unwrap() = Some(failure.into());
|
||||
let _result = bundle.service().refresh_quota(&id).await;
|
||||
let initial = accounts.accounts.lock().unwrap()[0].clone();
|
||||
let edited = initial.clone().with_notes(Some("保留连接身份".into()));
|
||||
*accounts.accounts.lock().unwrap() = vec![edited.clone()];
|
||||
bundle.service().account_changed(&id).await.unwrap();
|
||||
assert!(
|
||||
bundle
|
||||
.service()
|
||||
.read_summary(&id)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.upstream
|
||||
.unwrap()
|
||||
.discovered
|
||||
);
|
||||
|
||||
// 地址和密钥修改推进凭据版本,出口变化则改变同一版本的连接身份
|
||||
for replacement in [
|
||||
edited.with_outbound_proxy(Some(
|
||||
OutboundProxy::parse("http://other.invalid:8080").unwrap(),
|
||||
)),
|
||||
ProviderAccount::new(
|
||||
id.clone(),
|
||||
initial.provider().clone(),
|
||||
"替换连接".into(),
|
||||
None,
|
||||
"peer_test".into(),
|
||||
CredentialRevision::new(2).unwrap(),
|
||||
None,
|
||||
)
|
||||
.with_account_facts(
|
||||
true,
|
||||
CredentialState::Ready,
|
||||
QuotaState::unknown(),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
] {
|
||||
*accounts.accounts.lock().unwrap() = vec![replacement];
|
||||
bundle.service().account_changed(&id).await.unwrap();
|
||||
let detail = bundle.service().read_quota(&id).await.unwrap().unwrap();
|
||||
assert!(!detail.upstream.as_ref().unwrap().discovered);
|
||||
assert!(detail.upstream.as_ref().unwrap().last_error.is_none());
|
||||
assert_eq!(
|
||||
bundle
|
||||
.service()
|
||||
.read_summary(&id)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.upstream,
|
||||
detail.upstream
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn manual_rediscovery_does_not_wait_for_authorization_backoff_but_respects_retry_after() {
|
||||
for retry_after in [None, Some(Duration::from_secs(90))] {
|
||||
let (mut bundle, _, transport, id) = fixture(UpstreamInfoFailure::Unauthorized);
|
||||
*transport.failure.lock().unwrap() = None;
|
||||
bundle.service().refresh_quota(&id).await.unwrap();
|
||||
tokio::time::advance(Duration::from_secs(15)).await;
|
||||
*transport.failure.lock().unwrap() = Some(PeerFetchFailure {
|
||||
kind: UpstreamInfoFailure::Unauthorized,
|
||||
retry_after,
|
||||
});
|
||||
assert!(bundle.service().refresh_quota(&id).await.is_err());
|
||||
let worker = worker(&mut bundle);
|
||||
*transport.failure.lock().unwrap() = None;
|
||||
let delay = retry_after.unwrap_or(Duration::from_secs(15));
|
||||
tokio::time::advance(delay - Duration::from_secs(1)).await;
|
||||
assert!(bundle.service().refresh_quota(&id).await.is_err());
|
||||
assert_eq!(transport.requests.load(Ordering::SeqCst), 2);
|
||||
tokio::time::advance(Duration::from_secs(1)).await;
|
||||
cycle(&worker).await;
|
||||
assert_eq!(transport.requests.load(Ordering::SeqCst), 2);
|
||||
let recovered = bundle.service().refresh_quota(&id).await.unwrap();
|
||||
assert_eq!(transport.requests.load(Ordering::SeqCst), 3);
|
||||
assert!(recovered.upstream.as_ref().unwrap().discovered);
|
||||
assert!(recovered.upstream.as_ref().unwrap().last_error.is_none());
|
||||
assert!(!recovered.windows.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
fn fixture(
|
||||
failure: UpstreamInfoFailure,
|
||||
) -> (
|
||||
|
||||
@@ -3,6 +3,14 @@ use super::TestPluginPorts;
|
||||
use gateway_admin::{model::usage_data::*, ports::usage_data::*};
|
||||
#[async_trait::async_trait]
|
||||
impl UsageDataStore for TestPluginPorts {
|
||||
async fn open_read(
|
||||
&self,
|
||||
_: UsageReadQuery,
|
||||
_: gateway_core::lifecycle::CancellationToken,
|
||||
) -> UsageDataResult<Box<dyn UsageReadSession>> {
|
||||
Err(UsageDataError::unavailable("test usage reader unavailable"))
|
||||
}
|
||||
|
||||
async fn get_record(&self, _: GetUsageRecord) -> UsageDataResult<UsageRecord> {
|
||||
unreachable!("unexpected usage call")
|
||||
}
|
||||
|
||||
@@ -64,6 +64,7 @@ fn sdk_wire_types_match_host_contracts() {
|
||||
"value.rs",
|
||||
"record.rs",
|
||||
"query.rs",
|
||||
"read.rs",
|
||||
"aggregate.rs",
|
||||
"pricing.rs",
|
||||
"error.rs",
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
//! 账号管理路由与 HTTP handler 编排
|
||||
|
||||
use gateway_core::engine::probe::AccountProbeTarget;
|
||||
|
||||
use crate::auth::SessionState;
|
||||
|
||||
use super::*;
|
||||
@@ -626,7 +628,10 @@ where
|
||||
let stream = state
|
||||
.admin_services()
|
||||
.accounts()
|
||||
.test_connection(account_id, upstream_model)
|
||||
.test_connection(
|
||||
AccountProbeTarget::ConnectionTest(account_id),
|
||||
upstream_model,
|
||||
)
|
||||
.await
|
||||
.map_err(map_service_error)?
|
||||
.map(move |event| {
|
||||
|
||||
@@ -193,6 +193,7 @@ fn quota_forecast_view(
|
||||
period: match forecast.period {
|
||||
AccountUsagePeriod::Weekly => "weekly",
|
||||
AccountUsagePeriod::Monthly => "monthly",
|
||||
AccountUsagePeriod::Daily => "daily",
|
||||
},
|
||||
target_days: forecast.target_seconds as f64 / 86_400.0,
|
||||
extrapolated: forecast.extrapolated,
|
||||
@@ -389,6 +390,7 @@ pub(super) fn account_usage_view(
|
||||
};
|
||||
AccountUsageView {
|
||||
window_kind: match period {
|
||||
Some(AccountUsagePeriod::Daily) => AccountUsageWindowKind::Daily,
|
||||
Some(AccountUsagePeriod::Weekly) => AccountUsageWindowKind::Weekly,
|
||||
Some(AccountUsagePeriod::Monthly) => AccountUsageWindowKind::Monthly,
|
||||
None => AccountUsageWindowKind::Quota,
|
||||
|
||||
@@ -14,6 +14,8 @@ pub struct AccountUpstreamInfoView {
|
||||
pub stale: bool,
|
||||
pub connection_revision: String,
|
||||
pub last_error: Option<String>,
|
||||
pub max_concurrency: Option<u64>,
|
||||
pub requests_per_minute: Option<u64>,
|
||||
pub pool: Option<UpstreamPoolView>,
|
||||
}
|
||||
|
||||
@@ -78,6 +80,8 @@ pub(crate) fn view(
|
||||
stale: info.stale,
|
||||
connection_revision: info.connection_revision,
|
||||
last_error: info.last_error,
|
||||
max_concurrency: info.max_concurrency,
|
||||
requests_per_minute: info.requests_per_minute,
|
||||
pool: info.pool.map(|pool| UpstreamPoolView {
|
||||
account_count: pool.account_count,
|
||||
states: UpstreamPoolStatesView {
|
||||
|
||||
@@ -260,6 +260,11 @@ pub(crate) fn usage_list_record_view(
|
||||
UsageListRecordView {
|
||||
error_code: record.error_code,
|
||||
error_message: record.error_message,
|
||||
error_kind: record.error_kind,
|
||||
websocket_close_code: record.websocket_close_code,
|
||||
continuation_unavailable_reason: record.continuation_unavailable_reason,
|
||||
upstream_connection_exit_reason: record.upstream_connection_exit_reason,
|
||||
recovered: record.recovered,
|
||||
logical_outcome: record.outcome.as_str().to_owned(),
|
||||
client_api_key_name: record.client_api_key_name,
|
||||
id: record.id,
|
||||
@@ -357,6 +362,23 @@ pub(crate) fn usage_record_view(
|
||||
.or_else(|| record.requested_model_id.clone());
|
||||
let metadata = provider_metadata_fields(record.provider_metadata_json.as_deref());
|
||||
UsageRecordView {
|
||||
upstream_error: record
|
||||
.upstream_error
|
||||
.map(|error| super::wire::UsageUpstreamErrorView {
|
||||
code: error.code,
|
||||
error_type: error.error_type,
|
||||
message: error.message,
|
||||
param: error.param,
|
||||
close_reason: error.close_reason,
|
||||
last_event_type: error.last_event_type,
|
||||
}),
|
||||
retry_after_ms: record.retry_after_ms,
|
||||
error_code: record.provider_error_code,
|
||||
error_kind: record.error_kind,
|
||||
websocket_close_code: record.websocket_close_code,
|
||||
continuation_unavailable_reason: record.continuation_unavailable_reason,
|
||||
upstream_connection_exit_reason: record.upstream_connection_exit_reason,
|
||||
recovered: record.recovered,
|
||||
id: record.id.clone(),
|
||||
request_id: record.id,
|
||||
client_api_key_id: Some(record.client_api_key_ref),
|
||||
|
||||
@@ -88,6 +88,11 @@ pub struct ImageBillingView {
|
||||
pub struct UsageListRecordView {
|
||||
pub error_code: Option<String>,
|
||||
pub error_message: Option<String>,
|
||||
pub error_kind: Option<String>,
|
||||
pub websocket_close_code: Option<u16>,
|
||||
pub continuation_unavailable_reason: Option<String>,
|
||||
pub upstream_connection_exit_reason: Option<String>,
|
||||
pub recovered: bool,
|
||||
pub logical_outcome: String,
|
||||
pub client_api_key_name: Option<String>,
|
||||
pub id: String,
|
||||
@@ -122,10 +127,30 @@ pub struct UsageListRecordView {
|
||||
pub user_agent: Option<String>,
|
||||
}
|
||||
|
||||
/// 受控详情中的原始上游错误字段
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct UsageUpstreamErrorView {
|
||||
pub code: Option<String>,
|
||||
pub error_type: Option<String>,
|
||||
pub message: Option<String>,
|
||||
pub param: Option<String>,
|
||||
pub close_reason: Option<String>,
|
||||
pub last_event_type: Option<String>,
|
||||
}
|
||||
|
||||
/// 单条逻辑请求详情展示
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct UsageRecordView {
|
||||
pub upstream_error: Option<UsageUpstreamErrorView>,
|
||||
pub retry_after_ms: Option<u64>,
|
||||
pub error_code: Option<String>,
|
||||
pub error_kind: Option<String>,
|
||||
pub websocket_close_code: Option<u16>,
|
||||
pub continuation_unavailable_reason: Option<String>,
|
||||
pub upstream_connection_exit_reason: Option<String>,
|
||||
pub recovered: bool,
|
||||
pub id: String,
|
||||
pub request_id: String,
|
||||
pub client_api_key_id: Option<String>,
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//! `/healthz` 对 Core、Store 与 Host 关键 worker 健康事实的聚合
|
||||
//! 请求就绪检查与综合健康诊断
|
||||
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
@@ -67,22 +67,7 @@ impl HealthStatus {
|
||||
),
|
||||
})
|
||||
.collect();
|
||||
// 逐个探针共享截止时间,保留已完成结果并指出具体超时组件
|
||||
let deadline = tokio::time::Instant::now() + HEALTH_CHECK_TIMEOUT;
|
||||
let checks = self.probes.iter().map(|probe| async move {
|
||||
let (state, reason) = match tokio::time::timeout_at(deadline, probe.check()).await {
|
||||
Ok(HealthState::Healthy) => return None,
|
||||
Ok(HealthState::Degraded(reason)) => ("degraded", reason),
|
||||
Ok(HealthState::Unhealthy(reason)) => ("unhealthy", reason),
|
||||
Err(_) => ("timed_out", "health probe exceeded 2s deadline".to_owned()),
|
||||
};
|
||||
Some(HealthFailure {
|
||||
component: format!("probe:{}", probe.name()),
|
||||
state,
|
||||
reason,
|
||||
})
|
||||
});
|
||||
failures.extend(join_all(checks).await.into_iter().flatten());
|
||||
failures.extend(probe_failures(self.probes.iter()).await);
|
||||
let healthy = failures.is_empty();
|
||||
let changed = {
|
||||
let mut previous = self
|
||||
@@ -128,8 +113,42 @@ impl HealthStatus {
|
||||
}
|
||||
}
|
||||
|
||||
async fn probe_failures<'a>(
|
||||
probes: impl Iterator<Item = &'a Arc<dyn HealthProbe>>,
|
||||
) -> Vec<HealthFailure> {
|
||||
// 探针并行共用截止时间,保留已完成结果并指出具体超时组件
|
||||
let deadline = tokio::time::Instant::now() + HEALTH_CHECK_TIMEOUT;
|
||||
let checks = probes.map(|probe| async move {
|
||||
let (state, reason) = match tokio::time::timeout_at(deadline, probe.check()).await {
|
||||
Ok(HealthState::Healthy) => return None,
|
||||
Ok(HealthState::Degraded(reason)) => ("degraded", reason),
|
||||
Ok(HealthState::Unhealthy(reason)) => ("unhealthy", reason),
|
||||
Err(_) => ("timed_out", "health probe exceeded 2s deadline".to_owned()),
|
||||
};
|
||||
Some(HealthFailure {
|
||||
component: format!("probe:{}", probe.name()),
|
||||
state,
|
||||
reason,
|
||||
})
|
||||
});
|
||||
join_all(checks).await.into_iter().flatten().collect()
|
||||
}
|
||||
|
||||
pub(crate) async fn readyz(State(state): State<ApiState>) -> StatusCode {
|
||||
let required = state
|
||||
.health()
|
||||
.probes
|
||||
.iter()
|
||||
.filter(|probe| probe.required_for_readiness());
|
||||
status_code(probe_failures(required).await.is_empty())
|
||||
}
|
||||
|
||||
pub(crate) async fn healthz(State(state): State<ApiState>) -> StatusCode {
|
||||
if state.health().healthy().await {
|
||||
status_code(state.health().healthy().await)
|
||||
}
|
||||
|
||||
fn status_code(healthy: bool) -> StatusCode {
|
||||
if healthy {
|
||||
StatusCode::NO_CONTENT
|
||||
} else {
|
||||
StatusCode::SERVICE_UNAVAILABLE
|
||||
|
||||
@@ -216,6 +216,7 @@ pub fn initialize(
|
||||
};
|
||||
let index = config.asset_directory.join("index.html");
|
||||
let mut router = Router::new()
|
||||
.route("/readyz", get(health::readyz))
|
||||
.route("/healthz", get(health::healthz))
|
||||
.merge(openai::router::router())
|
||||
.merge(provider::router())
|
||||
|
||||
@@ -27,12 +27,35 @@ pub(in crate::openai) fn convert_request(
|
||||
("stream".into(), stream.into()),
|
||||
("store".into(), false.into()),
|
||||
]);
|
||||
if let Some(value) = optional(&mut chat, "tools") {
|
||||
body.insert("tools".into(), function_tools(value)?.into());
|
||||
}
|
||||
if let Some(value) = optional(&mut chat, "tool_choice") {
|
||||
let choice = tool_choice(value, body.get("tools"))?;
|
||||
body.insert("tool_choice".into(), choice);
|
||||
if let Some(options) = optional(&mut chat, "web_search_options") {
|
||||
let tool = web_search_tool(options)?;
|
||||
// Chat 搜索要求先检索,不能用强制搜索覆盖调用方的函数工具或选择
|
||||
if let Some(tools) = optional(&mut chat, "tools")
|
||||
&& !array(tools, "tools")?.is_empty()
|
||||
{
|
||||
return Err(RequestDecodeError::UnsupportedCombination {
|
||||
field: "tools",
|
||||
other: "web_search_options",
|
||||
});
|
||||
}
|
||||
if let Some(choice) = optional(&mut chat, "tool_choice")
|
||||
&& choice != "auto"
|
||||
{
|
||||
return Err(RequestDecodeError::UnsupportedCombination {
|
||||
field: "tool_choice",
|
||||
other: "web_search_options",
|
||||
});
|
||||
}
|
||||
body.insert("tools".into(), Value::Array(vec![tool]));
|
||||
body.insert("tool_choice".into(), json!({"type": "web_search"}));
|
||||
} else {
|
||||
if let Some(value) = optional(&mut chat, "tools") {
|
||||
body.insert("tools".into(), function_tools(value)?.into());
|
||||
}
|
||||
if let Some(value) = optional(&mut chat, "tool_choice") {
|
||||
let choice = tool_choice(value, body.get("tools"))?;
|
||||
body.insert("tool_choice".into(), choice);
|
||||
}
|
||||
}
|
||||
if let Some(value) = optional_bool(&mut chat, "parallel_tool_calls", "parallel_tool_calls")? {
|
||||
body.insert("parallel_tool_calls".into(), value.into());
|
||||
@@ -112,7 +135,6 @@ fn validate_options(chat: &Map<String, Value>) -> Result<(), RequestDecodeError>
|
||||
| "prediction"
|
||||
| "seed"
|
||||
| "function_call"
|
||||
| "web_search_options"
|
||||
| "reasoning_content"
|
||||
| "input"
|
||||
| "previous_response_id"
|
||||
@@ -132,6 +154,33 @@ fn validate_options(chat: &Map<String, Value>) -> Result<(), RequestDecodeError>
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn web_search_tool(value: Value) -> Result<Value, RequestDecodeError> {
|
||||
let mut options = object(value, "web_search_options")?;
|
||||
let mut tool = json!({"type": "web_search"});
|
||||
if let Some(size) = optional(&mut options, "search_context_size") {
|
||||
string(&size, "web_search_options.search_context_size", false)?;
|
||||
tool["search_context_size"] = size;
|
||||
}
|
||||
if let Some(location) = optional(&mut options, "user_location") {
|
||||
let mut location = object(location, "web_search_options.user_location")?;
|
||||
let path = "web_search_options.user_location.type";
|
||||
if required_string(&mut location, "type", path, false)? != "approximate" {
|
||||
return Err(unsupported(path));
|
||||
}
|
||||
let path = "web_search_options.user_location.approximate";
|
||||
let mut approximate = object(required(&mut location, "approximate", path)?, path)?;
|
||||
let mut flat = json!({"type": "approximate"});
|
||||
for field in ["city", "country", "region", "timezone"] {
|
||||
if let Some(value) = optional(&mut approximate, field) {
|
||||
string(&value, &format!("{path}.{field}"), true)?;
|
||||
flat[field] = value;
|
||||
}
|
||||
}
|
||||
tool["user_location"] = flat;
|
||||
}
|
||||
Ok(tool)
|
||||
}
|
||||
|
||||
fn stream_options(value: Option<Value>, stream: bool) -> Result<bool, RequestDecodeError> {
|
||||
let Some(value) = value else {
|
||||
return Ok(false);
|
||||
@@ -176,10 +225,13 @@ fn messages(value: Value) -> Result<Vec<Value>, RequestDecodeError> {
|
||||
) {
|
||||
return Err(unsupported(&format!("{path}.role")));
|
||||
}
|
||||
if let Some(name) = optional(&mut message, "name")
|
||||
&& !string(&name, &format!("{path}.name"), true)?.is_empty()
|
||||
{
|
||||
return Err(unsupported(&format!("{path}.name")));
|
||||
let name = optional(&mut message, "name").filter(|name| name != "");
|
||||
if let Some(name) = &name {
|
||||
let name_path = format!("{path}.name");
|
||||
string(name, &name_path, true)?;
|
||||
if role != "tool" {
|
||||
return Err(unsupported(&name_path));
|
||||
}
|
||||
}
|
||||
for field in ["function_call", "reasoning_content", "reasoning", "audio"] {
|
||||
if optional(&mut message, field).is_some() {
|
||||
@@ -209,9 +261,12 @@ fn messages(value: Value) -> Result<Vec<Value>, RequestDecodeError> {
|
||||
} else {
|
||||
content_parts(content, &role, &content_path, false)?.into()
|
||||
};
|
||||
input.push(
|
||||
json!({"type": "function_call_output", "call_id": call_id, "output": output}),
|
||||
);
|
||||
let mut item =
|
||||
json!({"type": "function_call_output", "call_id": call_id, "output": output});
|
||||
if let Some(name) = name {
|
||||
item["name"] = name;
|
||||
}
|
||||
input.push(item);
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -275,6 +330,10 @@ fn content_parts(
|
||||
required(&mut part, "image_url", &format!("{path}.image_url"))?,
|
||||
&format!("{path}.image_url"),
|
||||
)?,
|
||||
"file" if role == "user" => {
|
||||
let file_path = format!("{path}.file");
|
||||
inline_pdf_part(required(&mut part, "file", &file_path)?, &file_path)?
|
||||
}
|
||||
"refusal" if role == "assistant" => json!({
|
||||
"type": "refusal",
|
||||
"refusal": required_string(&mut part, "refusal", &format!("{path}.refusal"), true)?,
|
||||
@@ -301,6 +360,25 @@ fn text_part(text: String, is_assistant: bool) -> Value {
|
||||
}
|
||||
}
|
||||
|
||||
fn inline_pdf_part(value: Value, path: &str) -> Result<Value, RequestDecodeError> {
|
||||
let mut file = object(value, path)?;
|
||||
if optional(&mut file, "file_id").is_some() {
|
||||
return Err(unsupported(&format!("{path}.file_id")));
|
||||
}
|
||||
let data_path = format!("{path}.file_data");
|
||||
let data = required_string(&mut file, "file_data", &data_path, false)?;
|
||||
if !data
|
||||
.strip_prefix("data:application/pdf;base64,")
|
||||
.is_some_and(|payload| !payload.trim().is_empty())
|
||||
{
|
||||
return Err(invalid(&data_path));
|
||||
}
|
||||
let name = required_string(&mut file, "filename", &format!("{path}.filename"), false)?;
|
||||
let mut converted = json!({"type": "input_file", "filename": name});
|
||||
converted["file_data"] = data.into();
|
||||
Ok(converted)
|
||||
}
|
||||
|
||||
fn image_part(value: Value, path: &str) -> Result<Value, RequestDecodeError> {
|
||||
let mut image = object(value, path)?;
|
||||
let url = required_string(&mut image, "url", &format!("{path}.url"), false)?;
|
||||
|
||||
@@ -46,13 +46,21 @@ pub(in crate::openai) fn project_response(
|
||||
let mut content = None;
|
||||
let mut refusal = None;
|
||||
let mut tools = Vec::new();
|
||||
let mut annotations = Vec::new();
|
||||
let mut content_chars = 0;
|
||||
for item in output {
|
||||
match required_str(item, "type")? {
|
||||
"message" => {
|
||||
// 两侧索引都按消息中的字符计数,合并多条 assistant 消息时平移引用
|
||||
let message_offset = content_chars;
|
||||
for part in message_parts(item)? {
|
||||
let (kind, text) = part_text(part)?;
|
||||
let target = match kind {
|
||||
PartKind::Text => &mut content,
|
||||
PartKind::Text => {
|
||||
append_url_citations(part, message_offset, &mut annotations)?;
|
||||
content_chars += text.chars().count() as u64;
|
||||
&mut content
|
||||
}
|
||||
PartKind::Refusal => &mut refusal,
|
||||
};
|
||||
target.get_or_insert_with(String::new).push_str(text);
|
||||
@@ -65,7 +73,7 @@ pub(in crate::openai) fn project_response(
|
||||
"function": {"name": name, "arguments": arguments}
|
||||
}));
|
||||
}
|
||||
"reasoning" => {}
|
||||
"reasoning" | "web_search_call" => {}
|
||||
_ => return Err(invalid_response()),
|
||||
}
|
||||
}
|
||||
@@ -74,6 +82,9 @@ pub(in crate::openai) fn project_response(
|
||||
if !tools.is_empty() {
|
||||
message["tool_calls"] = Value::Array(tools);
|
||||
}
|
||||
if !annotations.is_empty() {
|
||||
message["annotations"] = Value::Array(annotations);
|
||||
}
|
||||
serde_json::to_value(Envelope {
|
||||
identity,
|
||||
object: "chat.completion",
|
||||
@@ -117,6 +128,7 @@ enum Output {
|
||||
Message(BTreeMap<u64, Part>),
|
||||
Function(Tool),
|
||||
Reasoning,
|
||||
WebSearch,
|
||||
}
|
||||
|
||||
struct OutputItem {
|
||||
@@ -357,6 +369,7 @@ impl ChatStreamEncoder {
|
||||
let output = match kind {
|
||||
"message" => Output::Message(BTreeMap::new()),
|
||||
"reasoning" => Output::Reasoning,
|
||||
"web_search_call" => Output::WebSearch,
|
||||
"function_call" => {
|
||||
let (call_id, name, _) = function_fields(snapshot)?;
|
||||
let tool_index = self.tool_count();
|
||||
@@ -399,7 +412,7 @@ impl ChatStreamEncoder {
|
||||
frames.push(delta_frame(&self.identity, self.include_usage, json!({"tool_calls": [{"index": tool.index, "function": {"arguments": suffix}}]}))?);
|
||||
}
|
||||
}
|
||||
(Output::Reasoning, "reasoning") => {}
|
||||
(Output::Reasoning, "reasoning") | (Output::WebSearch, "web_search_call") => {}
|
||||
_ => return Err(invalid_response()),
|
||||
}
|
||||
Ok(())
|
||||
@@ -502,6 +515,33 @@ fn part_text(part: &Value) -> Result<(PartKind, &str), ProtocolErrorBody> {
|
||||
}
|
||||
}
|
||||
|
||||
fn append_url_citations(
|
||||
part: &Value,
|
||||
message_offset: u64,
|
||||
annotations: &mut Vec<Value>,
|
||||
) -> Result<(), ProtocolErrorBody> {
|
||||
let Some(values) = part.get("annotations").filter(|value| !value.is_null()) else {
|
||||
return Ok(());
|
||||
};
|
||||
for value in values.as_array().ok_or_else(invalid_response)? {
|
||||
if value["type"] != "url_citation" {
|
||||
continue;
|
||||
}
|
||||
let start = required_index(value, "start_index")?;
|
||||
let end = required_index(value, "end_index")?;
|
||||
if start > end {
|
||||
return Err(invalid_response());
|
||||
}
|
||||
annotations.push(json!({"type": "url_citation", "url_citation": {
|
||||
"url": required_str(value, "url")?,
|
||||
"title": required_str(value, "title")?,
|
||||
"start_index": message_offset.checked_add(start).ok_or_else(invalid_response)?,
|
||||
"end_index": message_offset.checked_add(end).ok_or_else(invalid_response)?,
|
||||
}}));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn function_fields(item: &Value) -> Result<(&str, &str, &str), ProtocolErrorBody> {
|
||||
Ok((
|
||||
required_nonempty_str(item, "call_id")?,
|
||||
|
||||
@@ -4,14 +4,14 @@ use chrono::{DateTime, Utc};
|
||||
use gateway_admin::model::peer_link::{PeerLinkExport, PeerModelState as ModelState};
|
||||
use gateway_core::metering::Decimal;
|
||||
use gateway_protocol::peer::{
|
||||
OBJECT, PeerBudgetWindow, PeerComponentStatus, PeerKeyUsage, PeerLink, PeerModelAvailability,
|
||||
PeerModelCoverage, PeerModelState, PeerPlanCount, PeerPool, PeerPoolStates, PeerQuota,
|
||||
PeerQuotaAccount, PeerQuotaWindow, PeerUsage, SCHEMA_VERSION,
|
||||
OBJECT, PeerBudgetWindow, PeerComponentStatus, PeerKeyLimits, PeerKeyUsage, PeerLink,
|
||||
PeerModelAvailability, PeerModelCoverage, PeerModelState, PeerPlanCount, PeerPool,
|
||||
PeerPoolStates, PeerQuota, PeerQuotaAccount, PeerQuotaWindow, PeerUsage, SCHEMA_VERSION,
|
||||
};
|
||||
use std::time::SystemTime;
|
||||
|
||||
pub(super) fn response(export: PeerLinkExport) -> PeerLink {
|
||||
let mut capabilities = vec!["key_usage".to_owned()];
|
||||
let mut capabilities = vec!["key_usage".to_owned(), "key_limits".to_owned()];
|
||||
let pool = export.pool.map(|pool| {
|
||||
capabilities.push("pool_summary".to_owned());
|
||||
PeerPool {
|
||||
@@ -74,6 +74,10 @@ pub(super) fn response(export: PeerLinkExport) -> PeerLink {
|
||||
capabilities,
|
||||
observed_at: export.observed_at.to_rfc3339(),
|
||||
expires_at: export.expires_at.to_rfc3339(),
|
||||
key_limits: Some(PeerKeyLimits {
|
||||
max_concurrency: Some(export.key_limits.max_concurrency),
|
||||
requests_per_minute: Some(export.key_limits.requests_per_minute),
|
||||
}),
|
||||
usage: Some(PeerUsage {
|
||||
status: PeerComponentStatus::Available,
|
||||
observed_at: Some(export.budget_observed_at.to_rfc3339()),
|
||||
|
||||
@@ -62,6 +62,14 @@ pub enum RequestDecodeError {
|
||||
/// 字段路径
|
||||
field: String,
|
||||
},
|
||||
/// 两个参数的要求无法同时满足
|
||||
#[error("requested `{field}` is not supported with `{other}`")]
|
||||
UnsupportedCombination {
|
||||
/// 返回给客户端的参数路径
|
||||
field: &'static str,
|
||||
/// 与之冲突的参数路径
|
||||
other: &'static str,
|
||||
},
|
||||
/// Core operation 拒绝规范化后的字段
|
||||
#[error("field `{field}` violates the canonical operation contract")]
|
||||
CanonicalContract {
|
||||
@@ -127,6 +135,11 @@ impl RequestDecodeError {
|
||||
format!("Parameter `{field}` is not supported by this gateway."),
|
||||
Some(field.clone()),
|
||||
),
|
||||
Self::UnsupportedCombination { field, other } => (
|
||||
"unsupported_parameter",
|
||||
format!("The requested `{field}` is not supported with `{other}`."),
|
||||
Some((*field).to_owned()),
|
||||
),
|
||||
};
|
||||
ProtocolErrorBody {
|
||||
error: ProtocolError {
|
||||
|
||||
@@ -618,8 +618,6 @@ fn passthrough_header_name(name: &str, connection_headers: &[String]) -> bool {
|
||||
| "chatgpt-project-id"
|
||||
| "openai-organization"
|
||||
| "openai-project"
|
||||
| "x-openai-account-routing-override"
|
||||
| "x-openai-fedramp"
|
||||
// 上游指纹必须由运行时画像统一生成,客户端 originator/User-Agent/version
|
||||
// 不能作为不透明头透传覆盖,避免不同下游客户端暴露不一致的设备指纹
|
||||
| "originator"
|
||||
|
||||
@@ -275,6 +275,11 @@ impl ResponsesDeliveryValidator {
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn completed_response_id(&self) -> Option<&str> {
|
||||
self.response_id()
|
||||
.filter(|_| self.state.terminal == Some(TerminalKind::Completed))
|
||||
}
|
||||
|
||||
pub(super) fn finish_websocket_delivery(&self) -> Result<(), ResponseValidationError> {
|
||||
if !self.active || self.state.terminal.is_some() {
|
||||
Ok(())
|
||||
|
||||
@@ -36,6 +36,7 @@ pub struct ResponsesWebSocketConnection {
|
||||
commands: Option<mpsc::Sender<ConnectionCommand>>,
|
||||
incoming: mpsc::Receiver<PendingConnectionEvent>,
|
||||
deferred: VecDeque<PendingConnectionEvent>,
|
||||
input_gate: Arc<tokio::sync::Mutex<()>>,
|
||||
exited: oneshot::Receiver<PumpExitReason>,
|
||||
pump_task: Option<JoinHandle<()>>,
|
||||
stats: Arc<ConnectionStats>,
|
||||
@@ -172,6 +173,26 @@ impl ResponsesWebSocketConnection {
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(in super::super) async fn close_for_replay_if_idle(&mut self) -> bool {
|
||||
// 原子划分入站帧与正常关闭,正在处理或已排队的帧都先按原合同执行
|
||||
let Ok(_guard) = self.input_gate.clone().try_lock_owned() else {
|
||||
return false;
|
||||
};
|
||||
if !self.deferred.is_empty() || !self.incoming.is_empty() {
|
||||
return false;
|
||||
}
|
||||
self.close(
|
||||
CloseFrame {
|
||||
code: close_code::NORMAL,
|
||||
reason: "Reconnect with full input.".into(),
|
||||
},
|
||||
WriteContext::connection(FramePhase::Close),
|
||||
PumpExitReason::ServerClose,
|
||||
)
|
||||
.await;
|
||||
true
|
||||
}
|
||||
|
||||
pub(in super::super) async fn close_for_connection_limit(&mut self, reason: &'static str) {
|
||||
self.close(
|
||||
CloseFrame {
|
||||
@@ -344,11 +365,13 @@ where
|
||||
let extensions = crate::middleware::current()
|
||||
.map(|context| context.extensions)
|
||||
.unwrap_or_default();
|
||||
let input_gate = Arc::new(tokio::sync::Mutex::new(()));
|
||||
let pump_task = tokio::spawn(run_pump(
|
||||
socket,
|
||||
activity.clone(),
|
||||
command_rx,
|
||||
incoming_tx,
|
||||
Arc::clone(&input_gate),
|
||||
exit_tx,
|
||||
Arc::clone(&connection_id),
|
||||
cancellation.clone(),
|
||||
@@ -368,6 +391,7 @@ where
|
||||
commands: Some(command_tx),
|
||||
incoming: incoming_rx,
|
||||
deferred: VecDeque::new(),
|
||||
input_gate,
|
||||
exited: exit_rx,
|
||||
pump_task: Some(pump_task),
|
||||
stats,
|
||||
|
||||
@@ -29,6 +29,7 @@ pub(super) async fn run_pump<S, E>(
|
||||
activity: watch::Sender<Instant>,
|
||||
commands: mpsc::Receiver<ConnectionCommand>,
|
||||
incoming: mpsc::Sender<PendingConnectionEvent>,
|
||||
input_gate: Arc<tokio::sync::Mutex<()>>,
|
||||
exited: oneshot::Sender<PumpExitReason>,
|
||||
connection_id: Arc<str>,
|
||||
cancellation: CancellationToken,
|
||||
@@ -63,7 +64,7 @@ pub(super) async fn run_pump<S, E>(
|
||||
// 心跳独立于业务写入等待,但共享唯一 writer,避免长写入或排队阻止失联检测
|
||||
let reason = tokio::select! {
|
||||
biased;
|
||||
reason = read_connection(&mut reader, &incoming, sender.clone(), &activity, &context) => reason,
|
||||
reason = read_connection(&mut reader, &incoming, &input_gate, sender.clone(), &activity, &context) => reason,
|
||||
reason = write_connection(writer.clone(), commands, &context) => reason,
|
||||
reason = heartbeat_connection(writer.clone(), received_activity, &context) => reason,
|
||||
};
|
||||
@@ -81,6 +82,7 @@ pub(super) async fn run_pump<S, E>(
|
||||
async fn read_connection<S, E>(
|
||||
socket: &mut S,
|
||||
incoming: &mpsc::Sender<PendingConnectionEvent>,
|
||||
input_gate: &tokio::sync::Mutex<()>,
|
||||
sender: Arc<dyn contract::Sender>,
|
||||
activity: &watch::Sender<Instant>,
|
||||
context: &PumpContext,
|
||||
@@ -107,6 +109,7 @@ where
|
||||
}
|
||||
}
|
||||
message = socket.next() => {
|
||||
let _guard = input_gate.lock().await;
|
||||
if matches!(&message, Some(Ok(_))) {
|
||||
// 链路活动以真实入站为准,插件丢弃消息不应抹掉对端已经响应的事实
|
||||
activity.send_replace(Instant::now());
|
||||
|
||||
@@ -23,7 +23,9 @@ use crate::openai::responses::validation::{ResponseValidationFacts, ResponsesDel
|
||||
|
||||
use super::{
|
||||
super::{DecodedResponsesRequest, OpenAiResponsesEncoder, ProtocolErrorBody},
|
||||
connection::{ConnectionEvent, FramePhase, ResponsesWebSocketConnection, WriteContext},
|
||||
connection::{
|
||||
ConnectionEvent, FramePhase, ResponsesWebSocketConnection, WriteContext, WriteOutcome,
|
||||
},
|
||||
protocol::{
|
||||
error_event, initial_engine_error_event, is_response_create, response_metadata_event,
|
||||
},
|
||||
@@ -32,6 +34,7 @@ use super::{
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
pub(super) enum ForwardOutcome {
|
||||
Continue,
|
||||
ReconnectIfIdle,
|
||||
Disconnect,
|
||||
}
|
||||
|
||||
@@ -55,9 +58,25 @@ impl ConnectionReplaySnapshot {
|
||||
}
|
||||
}
|
||||
|
||||
fn commit(&mut self, response_id: String, provider_state: Option<ProviderSessionState>) {
|
||||
self.last_response_id = Some(response_id);
|
||||
self.provider_state = provider_state;
|
||||
fn commit(&mut self, capture: &ReplayCaptureHandle, delivered: Option<&str>) -> ForwardOutcome {
|
||||
let mut capture = capture
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
let Some(id) = capture.response_id.take() else {
|
||||
return ForwardOutcome::Continue;
|
||||
};
|
||||
self.provider_state = capture.provider_state.take();
|
||||
let reconnect = delivered == Some(&id)
|
||||
&& self
|
||||
.provider_state
|
||||
.as_ref()
|
||||
.is_some_and(ProviderSessionState::prefers_websocket_reconnect);
|
||||
self.last_response_id = Some(id);
|
||||
if reconnect {
|
||||
ForwardOutcome::ReconnectIfIdle
|
||||
} else {
|
||||
ForwardOutcome::Continue
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -199,7 +218,7 @@ pub(super) async fn forward_response(
|
||||
|
||||
let mut first_frame_written = false;
|
||||
let mut current = Some(first);
|
||||
let mut terminal_seen = false;
|
||||
let mut terminal_write = None;
|
||||
let mut first_protocol_validated = true;
|
||||
loop {
|
||||
let frame = match current.take() {
|
||||
@@ -211,12 +230,14 @@ pub(super) async fn forward_response(
|
||||
response_control,
|
||||
&mut interrupt,
|
||||
&request_id,
|
||||
terminal_seen,
|
||||
terminal_write.is_some(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
BodyInput::Frame(Ok(Some(frame))) => frame,
|
||||
BodyInput::Frame(Ok(None)) if terminal_seen && body.is_finalized() => {
|
||||
BodyInput::Frame(Ok(None))
|
||||
if terminal_write.is_some() && body.is_finalized() =>
|
||||
{
|
||||
if delivery_validator.finish_websocket_delivery().is_err() {
|
||||
return fail_body(
|
||||
connection,
|
||||
@@ -226,16 +247,12 @@ pub(super) async fn forward_response(
|
||||
)
|
||||
.await;
|
||||
}
|
||||
{
|
||||
let captured = capture
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if let Some(response_id) = captured.response_id.clone() {
|
||||
replay.commit(response_id, captured.provider_state.clone());
|
||||
}
|
||||
}
|
||||
let delivered = delivery_validator
|
||||
.completed_response_id()
|
||||
.filter(|_| terminal_write == Some(WriteOutcome::Written));
|
||||
let outcome = replay.commit(&capture, delivered);
|
||||
body.close().await;
|
||||
return ForwardOutcome::Continue;
|
||||
return outcome;
|
||||
}
|
||||
BodyInput::Frame(Ok(None)) => {
|
||||
return fail_body(
|
||||
@@ -256,7 +273,7 @@ pub(super) async fn forward_response(
|
||||
}
|
||||
}
|
||||
};
|
||||
if terminal_seen || validate_frame(&frame).is_err() {
|
||||
if terminal_write.is_some() || validate_frame(&frame).is_err() {
|
||||
return fail_body(connection, body, MiddlewareError::InvalidState, &request_id).await;
|
||||
}
|
||||
if first_protocol_validated {
|
||||
@@ -278,30 +295,27 @@ pub(super) async fn forward_response(
|
||||
(true, false) => FramePhase::Data,
|
||||
};
|
||||
let created_id = super::interrupt::ResponseInterrupt::created_id(&message);
|
||||
match connection
|
||||
let Ok(write) = connection
|
||||
.send_text(message, WriteContext::request(&request_id, phase))
|
||||
.await
|
||||
{
|
||||
Ok(super::connection::WriteOutcome::Written) => {
|
||||
first_frame_written = true;
|
||||
if let Some(id) = created_id {
|
||||
interrupt.observe_created(id, response_control);
|
||||
tokio::select! {
|
||||
_ = connection.wait_for_exit() => {
|
||||
detach_body(body);
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
() = interrupt.forward_pending(response_control) => {}
|
||||
else {
|
||||
detach_body(body);
|
||||
return ForwardOutcome::Disconnect;
|
||||
};
|
||||
if write == WriteOutcome::Written {
|
||||
first_frame_written = true;
|
||||
if let Some(id) = created_id {
|
||||
interrupt.observe_created(id, response_control);
|
||||
tokio::select! {
|
||||
_ = connection.wait_for_exit() => {
|
||||
detach_body(body);
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
() = interrupt.forward_pending(response_control) => {}
|
||||
}
|
||||
}
|
||||
Ok(super::connection::WriteOutcome::Suppressed) => {}
|
||||
Err(_) => {
|
||||
detach_body(body);
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
}
|
||||
terminal_seen = terminal;
|
||||
terminal_write = terminal.then_some(write);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -179,6 +179,10 @@ async fn serve_responses_websocket(socket: Socket, session: ResponsesWebSocketSe
|
||||
crate::middleware::websocket::CancelOnDrop(Some(cancellation.clone()));
|
||||
let request_headers =
|
||||
request_headers.with_downstream_websocket_connection_id(connection_id.clone());
|
||||
// 帧中间件可能改写终态,此时不能据原始响应关闭下游连接
|
||||
let reconnect_for_replay = !middleware
|
||||
.as_ref()
|
||||
.is_some_and(FrozenMiddlewarePlan::has_websocket);
|
||||
let mut connection = ResponsesWebSocketConnection::new(
|
||||
socket,
|
||||
connection_id,
|
||||
@@ -401,7 +405,7 @@ async fn serve_responses_websocket(socket: Socket, session: ResponsesWebSocketSe
|
||||
continue;
|
||||
}
|
||||
};
|
||||
if forward_response(
|
||||
let outcome = forward_response(
|
||||
&mut connection,
|
||||
response,
|
||||
request_id,
|
||||
@@ -410,8 +414,12 @@ async fn serve_responses_websocket(socket: Socket, session: ResponsesWebSocketSe
|
||||
validation,
|
||||
&response_control,
|
||||
)
|
||||
.await
|
||||
== ForwardOutcome::Disconnect
|
||||
.await;
|
||||
if outcome == ForwardOutcome::Disconnect
|
||||
|| (outcome == ForwardOutcome::ReconnectIfIdle
|
||||
&& reconnect_for_replay
|
||||
&& !connection.is_expired()
|
||||
&& connection.close_for_replay_if_idle().await)
|
||||
{
|
||||
break;
|
||||
}
|
||||
|
||||
@@ -444,9 +444,12 @@ fn account_page_item(
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_usage_preserves_cache_read_and_write_counts_and_missing_values() {
|
||||
async fn account_usage_preserves_periods_cache_counts_and_missing_values() {
|
||||
use chrono::Utc;
|
||||
use gateway_admin::model::accounts::AccountUsage;
|
||||
use gateway_admin::model::provider_credentials::{
|
||||
ProviderQuota, ProviderQuotaWindow, QuotaLocalUsageAttribution,
|
||||
};
|
||||
use gateway_core::account::{AccountStatusFacts, CredentialState, QuotaState};
|
||||
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
@@ -461,9 +464,19 @@ async fn account_usage_preserves_cache_read_and_write_counts_and_missing_values(
|
||||
last_error_reason: None,
|
||||
last_error_message: None,
|
||||
};
|
||||
*fixture.account.lock().unwrap() = Some(account_page_item(facts, now, "api_key"));
|
||||
for (read, write) in [(Some(13), Some(7)), (None, None)] {
|
||||
*fixture.account_usage.lock().unwrap() = vec![AccountUsage {
|
||||
for (authentication_kind, read, write) in [
|
||||
("api_key", Some(13), Some(7)),
|
||||
("api_key", None, None),
|
||||
("oauth", Some(13), Some(7)),
|
||||
("oauth", None, None),
|
||||
] {
|
||||
let mut account = account_page_item(facts.clone(), now, authentication_kind);
|
||||
if authentication_kind == "oauth" {
|
||||
account.account.provider_kind =
|
||||
gateway_core::routing::ProviderKind::new("xai").unwrap();
|
||||
}
|
||||
*fixture.account.lock().unwrap() = Some(account);
|
||||
let local_usage = AccountUsage {
|
||||
account_id: "acct_cooldown".to_owned(),
|
||||
request_count: 1,
|
||||
success_count: 1,
|
||||
@@ -480,9 +493,33 @@ async fn account_usage_preserves_cache_read_and_write_counts_and_missing_values(
|
||||
cost_coverage: Default::default(),
|
||||
costs: Vec::new(),
|
||||
last_used_at: Some(now),
|
||||
request_buckets: Vec::new(),
|
||||
request_buckets: vec![gateway_admin::model::accounts::AccountRequestBucket {
|
||||
bucket_start: now,
|
||||
request_count: 1,
|
||||
}],
|
||||
models: Vec::new(),
|
||||
}];
|
||||
};
|
||||
*fixture.account_usage.lock().unwrap() = vec![local_usage.clone()];
|
||||
*fixture.provider_quota.lock().unwrap() =
|
||||
(authentication_kind == "oauth").then(|| ProviderQuota {
|
||||
windows: vec![ProviderQuotaWindow {
|
||||
upstream: None,
|
||||
key: "rolling-day".to_owned(),
|
||||
group: "shortTerm".to_owned(),
|
||||
label: "日限额".to_owned(),
|
||||
limit_id: None,
|
||||
limit_name: None,
|
||||
role: None,
|
||||
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
|
||||
window_seconds: Some(86_400),
|
||||
used_percent: None,
|
||||
reset_at: None,
|
||||
limit_reached: false,
|
||||
local_usage: Some(local_usage),
|
||||
provider_data: None,
|
||||
}],
|
||||
..Default::default()
|
||||
});
|
||||
let response = admin::accounts::router::<AdminTestState>()
|
||||
.with_state(fixture.state())
|
||||
.oneshot(
|
||||
@@ -500,6 +537,27 @@ async fn account_usage_preserves_cache_read_and_write_counts_and_missing_values(
|
||||
let value: serde_json::Value = serde_json::from_slice(&body).unwrap();
|
||||
assert_eq!(status, StatusCode::OK, "{value}");
|
||||
let usage = &value["data"]["items"][0]["usage"];
|
||||
assert_eq!(
|
||||
usage["windowKind"],
|
||||
if authentication_kind == "oauth" {
|
||||
"daily"
|
||||
} else {
|
||||
"general"
|
||||
}
|
||||
);
|
||||
assert_eq!(usage["requestCount"], 1);
|
||||
assert_eq!(usage["totalTokens"], 120);
|
||||
if authentication_kind == "oauth" {
|
||||
let window = &value["data"]["items"][0]["quota"]["windows"][0];
|
||||
assert_eq!(window["localUsage"]["requestCount"], 1);
|
||||
assert_eq!(
|
||||
window["localUsage"]["requestBuckets"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
}
|
||||
assert_eq!(usage["cachedInputTokens"], serde_json::json!(read));
|
||||
assert_eq!(usage["cacheWriteInputTokens"], serde_json::json!(write));
|
||||
assert_eq!(
|
||||
|
||||
@@ -884,6 +884,7 @@ impl AccountGroupStore for MemoryAccountGroupStore {
|
||||
}
|
||||
Ok(vec![
|
||||
AccountGroupMemberFact {
|
||||
inherited_version: None,
|
||||
group_id: group_id(PRIMARY_GROUP_ID),
|
||||
account_id: "acct_group_ready".to_owned(),
|
||||
status: AccountStatusFacts {
|
||||
@@ -898,6 +899,7 @@ impl AccountGroupStore for MemoryAccountGroupStore {
|
||||
total_slots: Some(1),
|
||||
},
|
||||
AccountGroupMemberFact {
|
||||
inherited_version: None,
|
||||
group_id: group_id(PRIMARY_GROUP_ID),
|
||||
account_id: "acct_group_disabled".to_owned(),
|
||||
status: AccountStatusFacts {
|
||||
|
||||
@@ -387,6 +387,21 @@ async fn usage_detail_should_keep_attempt_snapshot_contract() {
|
||||
record.upstream_response_ms = Some(1_000);
|
||||
record.upstream_api_overhead_ms = Some(120.25);
|
||||
record.upstream_engine_ms = Some(6400.0);
|
||||
record.outcome = gateway_admin::model::observability::RequestOutcome::Failed;
|
||||
record.error_kind = Some("continuation_recovery_required".to_owned());
|
||||
record.provider_error_code = Some("previous_response_not_found".to_owned());
|
||||
record.continuation_unavailable_reason = Some("reused_connection_lost".to_owned());
|
||||
record.upstream_connection_exit_reason = Some("max_age_expired".to_owned());
|
||||
record.recovered = true;
|
||||
record.upstream_error = Some(gateway_admin::model::observability::UsageUpstreamError {
|
||||
code: Some("previous_response_not_found".to_owned()),
|
||||
error_type: Some("invalid_request_error".to_owned()),
|
||||
message: Some("Previous response is unavailable.".to_owned()),
|
||||
param: Some("previous_response_id".to_owned()),
|
||||
close_reason: None,
|
||||
last_event_type: None,
|
||||
});
|
||||
record.retry_after_ms = Some(206_000);
|
||||
fixture
|
||||
.usage_detail
|
||||
.lock()
|
||||
@@ -447,6 +462,31 @@ async fn usage_detail_should_keep_attempt_snapshot_contract() {
|
||||
.await
|
||||
.expect("usage detail body");
|
||||
let value: serde_json::Value = serde_json::from_slice(&body).expect("usage detail JSON");
|
||||
assert_eq!(value["data"]["errorKind"], "continuation_recovery_required");
|
||||
assert_eq!(value["data"]["errorCode"], "previous_response_not_found");
|
||||
assert_eq!(
|
||||
value["data"]["upstreamError"]["message"],
|
||||
"Previous response is unavailable."
|
||||
);
|
||||
assert_eq!(
|
||||
value["data"]["upstreamError"]["errorType"],
|
||||
"invalid_request_error"
|
||||
);
|
||||
assert_eq!(
|
||||
value["data"]["upstreamError"]["param"],
|
||||
"previous_response_id"
|
||||
);
|
||||
assert_eq!(value["data"]["retryAfterMs"], 206_000);
|
||||
assert_eq!(
|
||||
value["data"]["continuationUnavailableReason"],
|
||||
"reused_connection_lost"
|
||||
);
|
||||
assert_eq!(
|
||||
value["data"]["upstreamConnectionExitReason"],
|
||||
"max_age_expired"
|
||||
);
|
||||
assert_eq!(value["data"]["websocketCloseCode"], serde_json::Value::Null);
|
||||
assert_eq!(value["data"]["recovered"], true);
|
||||
assert_eq!(value["data"]["firstTokenLatencyMs"], 20);
|
||||
assert_eq!(value["data"]["latencyDetails"]["upstreamResponseMs"], 1_000);
|
||||
assert_eq!(
|
||||
@@ -662,6 +702,8 @@ fn peer_summary(
|
||||
) -> gateway_admin::model::provider_credentials::ProviderQuota {
|
||||
gateway_admin::model::provider_credentials::ProviderQuota {
|
||||
upstream: Some(gateway_admin::model::upstream_info::ProviderUpstreamInfo {
|
||||
max_concurrency: None,
|
||||
requests_per_minute: None,
|
||||
discovered,
|
||||
observed_at: observed.then(chrono::Utc::now),
|
||||
stale: true,
|
||||
@@ -786,6 +828,7 @@ pub(super) fn usage_record_with_account(
|
||||
use gateway_admin::model::observability::{RequestOutcome, UsageRecord};
|
||||
|
||||
UsageRecord {
|
||||
upstream_error: None,
|
||||
id: id.to_owned(),
|
||||
client_api_key_ref: "key_detail".to_owned(),
|
||||
config_revision: 1,
|
||||
@@ -817,6 +860,10 @@ pub(super) fn usage_record_with_account(
|
||||
client_response_id: None,
|
||||
upstream_request_id: None,
|
||||
upstream_response_id: None,
|
||||
websocket_close_code: None,
|
||||
continuation_unavailable_reason: None,
|
||||
upstream_connection_exit_reason: None,
|
||||
recovered: false,
|
||||
error_kind: None,
|
||||
provider_error_code: None,
|
||||
error_message: None,
|
||||
@@ -934,6 +981,11 @@ async fn usage_route_projects_request_facts_and_billing() {
|
||||
provider_account_upstream_peer: None,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
error_kind: None,
|
||||
websocket_close_code: None,
|
||||
continuation_unavailable_reason: None,
|
||||
upstream_connection_exit_reason: None,
|
||||
recovered: false,
|
||||
client_status_code: None,
|
||||
outcome: gateway_admin::model::observability::RequestOutcome::Succeeded,
|
||||
client_api_key_name: Some("Production".to_owned()),
|
||||
@@ -1027,6 +1079,14 @@ async fn usage_route_projects_request_facts_and_billing() {
|
||||
failed.error_code = Some("server_error".to_owned());
|
||||
failed.error_message = Some("The upstream service is temporarily unavailable.".to_owned());
|
||||
records.push(failed);
|
||||
let mut closed = records[0].clone();
|
||||
closed.id = "request_ws_closed".to_owned();
|
||||
closed.outcome = gateway_admin::model::observability::RequestOutcome::Incomplete;
|
||||
closed.error_kind = Some("upstream_unavailable".to_owned());
|
||||
closed.websocket_close_code = Some(1012);
|
||||
closed.error_message =
|
||||
Some("OpenAI WebSocket closed before a terminal response".to_owned());
|
||||
records.push(closed);
|
||||
}
|
||||
let response = observability::router::<AdminTestState>()
|
||||
.with_state(fixture.state())
|
||||
@@ -1049,6 +1109,17 @@ async fn usage_route_projects_request_facts_and_billing() {
|
||||
assert!(value["data"]["items"][0]["errorCode"].is_null());
|
||||
assert!(value["data"]["items"][0]["errorMessage"].is_null());
|
||||
assert_eq!(value["data"]["items"][3]["errorCode"], "server_error");
|
||||
assert_eq!(value["data"]["items"][4]["logicalOutcome"], "incomplete");
|
||||
assert_eq!(
|
||||
value["data"]["items"][4]["errorKind"],
|
||||
"upstream_unavailable"
|
||||
);
|
||||
assert_eq!(value["data"]["items"][4]["websocketCloseCode"], 1012);
|
||||
assert_eq!(
|
||||
value["data"]["items"][4]["errorCode"],
|
||||
serde_json::Value::Null
|
||||
);
|
||||
assert_eq!(value["data"]["items"][4]["recovered"], false);
|
||||
assert_eq!(
|
||||
value["data"]["items"][3]["errorMessage"],
|
||||
"The upstream service is temporarily unavailable."
|
||||
|
||||
@@ -3,6 +3,14 @@ use super::TestPluginPorts;
|
||||
use gateway_admin::{model::usage_data::*, ports::usage_data::*};
|
||||
#[async_trait::async_trait]
|
||||
impl UsageDataStore for TestPluginPorts {
|
||||
async fn open_read(
|
||||
&self,
|
||||
_: UsageReadQuery,
|
||||
_: gateway_core::lifecycle::CancellationToken,
|
||||
) -> UsageDataResult<Box<dyn UsageReadSession>> {
|
||||
Err(UsageDataError::unavailable("test usage reader unavailable"))
|
||||
}
|
||||
|
||||
async fn get_record(&self, _: GetUsageRecord) -> UsageDataResult<UsageRecord> {
|
||||
unreachable!("unexpected usage call")
|
||||
}
|
||||
|
||||
@@ -13,7 +13,7 @@ impl UpstreamInfoSource for TestPeerInfo {
|
||||
fn signal(
|
||||
&self,
|
||||
_: &ProviderAccount,
|
||||
_: &str,
|
||||
_: Option<&str>,
|
||||
_: SystemTime,
|
||||
) -> Option<UpstreamSchedulingSignal> {
|
||||
None
|
||||
@@ -22,6 +22,14 @@ impl UpstreamInfoSource for TestPeerInfo {
|
||||
|
||||
#[async_trait]
|
||||
impl PeerUpstreamService for TestPeerInfo {
|
||||
fn observed_concurrency_limit(
|
||||
&self,
|
||||
_: &gateway_core::account::QuotaObservationVersion,
|
||||
_: std::time::SystemTime,
|
||||
) -> Option<gateway_core::account::AccountConcurrencyLimit> {
|
||||
None
|
||||
}
|
||||
|
||||
async fn account_changed(&self, _: &ProviderAccountId) -> Result<(), StoreError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//! 验证健康接口按基础设施与关键 Worker 状态判定可用性
|
||||
//! 验证请求就绪、综合健康状态与故障诊断
|
||||
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
@@ -86,6 +86,54 @@ async fn healthz_should_reject_critical_worker_failure() {
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn readiness_ignores_freeze_recovery_and_logging_failures_without_clearing_alerts() {
|
||||
let diagnostics = Arc::new(RecordingDiagnostics::default());
|
||||
let router = health_router(
|
||||
vec![Arc::new(MutableProbe(Mutex::new(HealthState::Degraded(
|
||||
"file logging maintenance failed".to_owned(),
|
||||
))))],
|
||||
vec![worker_snapshot(
|
||||
WorkerKind::AccountFreezeRecovery,
|
||||
WorkerRuntimeState::BackingOff,
|
||||
)],
|
||||
diagnostics.clone(),
|
||||
)
|
||||
.await;
|
||||
|
||||
for (path, expected) in [
|
||||
("/healthz", StatusCode::SERVICE_UNAVAILABLE),
|
||||
("/readyz", StatusCode::NO_CONTENT),
|
||||
("/healthz", StatusCode::SERVICE_UNAVAILABLE),
|
||||
] {
|
||||
assert_eq!(
|
||||
health_response(router.clone(), path).await.status(),
|
||||
expected
|
||||
);
|
||||
}
|
||||
assert_eq!(diagnostics.0.lock().unwrap().len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn readiness_rejects_an_unavailable_runtime_snapshot() {
|
||||
let router = health_router(
|
||||
vec![Arc::new(
|
||||
gateway_core::runtime::RuntimeSnapshotHandle::default(),
|
||||
)],
|
||||
Vec::new(),
|
||||
Arc::new(RecordingDiagnostics::default()),
|
||||
)
|
||||
.await;
|
||||
let response = health_response(router, "/readyz").await;
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert!(
|
||||
axum::body::to_bytes(response.into_body(), 1024)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_empty()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn healthz_records_probe_changes_once_and_records_recurrence_after_recovery() {
|
||||
let probe = Arc::new(MutableProbe(Mutex::new(HealthState::Unhealthy(
|
||||
@@ -95,7 +143,8 @@ async fn healthz_records_probe_changes_once_and_records_recurrence_after_recover
|
||||
let router = health_router(vec![probe.clone()], Vec::new(), diagnostics.clone()).await;
|
||||
|
||||
let responses =
|
||||
futures::future::join_all((0..8).map(|_| health_response(router.clone()))).await;
|
||||
futures::future::join_all((0..8).map(|_| health_response(router.clone(), "/healthz")))
|
||||
.await;
|
||||
for response in responses {
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert!(
|
||||
@@ -126,18 +175,18 @@ async fn healthz_records_probe_changes_once_and_records_recurrence_after_recover
|
||||
|
||||
*probe.0.lock().unwrap() = HealthState::Degraded("file logging maintenance failed".to_owned());
|
||||
assert_eq!(
|
||||
health_response(router.clone()).await.status(),
|
||||
health_response(router.clone(), "/healthz").await.status(),
|
||||
StatusCode::SERVICE_UNAVAILABLE
|
||||
);
|
||||
assert_eq!(diagnostics.0.lock().unwrap().len(), 2);
|
||||
*probe.0.lock().unwrap() = HealthState::Healthy;
|
||||
assert_eq!(
|
||||
health_response(router.clone()).await.status(),
|
||||
health_response(router.clone(), "/healthz").await.status(),
|
||||
StatusCode::NO_CONTENT
|
||||
);
|
||||
*probe.0.lock().unwrap() = HealthState::Degraded("file logging maintenance failed".to_owned());
|
||||
assert_eq!(
|
||||
health_response(router).await.status(),
|
||||
health_response(router, "/healthz").await.status(),
|
||||
StatusCode::SERVICE_UNAVAILABLE
|
||||
);
|
||||
assert_eq!(diagnostics.0.lock().unwrap().len(), 3);
|
||||
@@ -168,7 +217,11 @@ async fn healthz_retains_worker_and_completed_probe_failures_when_another_probe_
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
health_response(router).await.status(),
|
||||
health_response(router.clone(), "/readyz").await.status(),
|
||||
StatusCode::SERVICE_UNAVAILABLE
|
||||
);
|
||||
assert_eq!(
|
||||
health_response(router, "/healthz").await.status(),
|
||||
StatusCode::SERVICE_UNAVAILABLE
|
||||
);
|
||||
let events = diagnostics.0.lock().unwrap();
|
||||
@@ -202,12 +255,12 @@ async fn healthz_retries_failed_diagnostic_admission_without_changing_health_res
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
health_response(router.clone()).await.status(),
|
||||
health_response(router.clone(), "/healthz").await.status(),
|
||||
StatusCode::SERVICE_UNAVAILABLE
|
||||
);
|
||||
assert!(diagnostics.recorded.0.lock().unwrap().is_empty());
|
||||
assert_eq!(
|
||||
health_response(router).await.status(),
|
||||
health_response(router, "/healthz").await.status(),
|
||||
StatusCode::SERVICE_UNAVAILABLE
|
||||
);
|
||||
assert_eq!(diagnostics.recorded.0.lock().unwrap().len(), 1);
|
||||
@@ -258,14 +311,9 @@ async fn health_router(
|
||||
.router()
|
||||
}
|
||||
|
||||
async fn health_response(router: axum::Router) -> axum::response::Response {
|
||||
async fn health_response(router: axum::Router, path: &str) -> axum::response::Response {
|
||||
router
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.uri("/healthz")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.oneshot(Request::builder().uri(path).body(Body::empty()).unwrap())
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
@@ -289,6 +337,10 @@ impl HealthProbe for PendingProbe {
|
||||
"pending"
|
||||
}
|
||||
|
||||
fn required_for_readiness(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn check(&self) -> BoxFuture<'_, HealthState> {
|
||||
Box::pin(std::future::pending())
|
||||
}
|
||||
|
||||
@@ -95,6 +95,11 @@ fn usage_record() -> UsageListRecord {
|
||||
provider_account_upstream_peer: None,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
error_kind: None,
|
||||
websocket_close_code: None,
|
||||
continuation_unavailable_reason: None,
|
||||
upstream_connection_exit_reason: None,
|
||||
recovered: false,
|
||||
client_status_code: None,
|
||||
outcome: gateway_admin::model::observability::RequestOutcome::Succeeded,
|
||||
client_api_key_name: Some("Production".to_owned()),
|
||||
|
||||
@@ -28,7 +28,7 @@ async fn preserves_ordered_roles_images_and_modern_tool_round_trip() {
|
||||
{"type":"image_url","image_url":{"url":"data:image/png;base64,aGVsbG8=","detail":"high"}}]},
|
||||
{"role":"assistant","content":"checking","tool_calls":[
|
||||
{"id":"call-weather","type":"function","function":{"name":"weather","arguments":arguments}}]},
|
||||
{"role":"tool","tool_call_id":"call-weather","content":"{\"temperature\":25}"},
|
||||
{"role":"tool","tool_call_id":"call-weather","name":"weather","content":"{\"temperature\":25}"},
|
||||
{"role":"user","content":"continue"}
|
||||
]})).await;
|
||||
let input = converted["input"].as_array().unwrap();
|
||||
@@ -53,11 +53,78 @@ async fn preserves_ordered_roles_images_and_modern_tool_round_trip() {
|
||||
);
|
||||
assert_eq!(
|
||||
input[5],
|
||||
json!({"type":"function_call_output","call_id":"call-weather","output":"{\"temperature\":25}"})
|
||||
json!({"type":"function_call_output","call_id":"call-weather","name":"weather","output":"{\"temperature\":25}"})
|
||||
);
|
||||
assert!(converted.get("previous_response_id").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn web_search_options_require_a_hosted_search_and_flatten_location() {
|
||||
for (options, expected) in [
|
||||
(json!({}), json!({"type":"web_search"})),
|
||||
(
|
||||
json!({"search_context_size":"low","user_location":{"type":"approximate",
|
||||
"approximate":{"country":"US","city":"Seattle","region":"WA","timezone":"America/Los_Angeles"}}}),
|
||||
json!({"type":"web_search","search_context_size":"low","user_location":{"type":"approximate",
|
||||
"country":"US","city":"Seattle","region":"WA","timezone":"America/Los_Angeles"}}),
|
||||
),
|
||||
] {
|
||||
let mut request = chat_request(false);
|
||||
request["web_search_options"] = options;
|
||||
let body = convert(request).await;
|
||||
assert_eq!(body["tools"], json!([expected]));
|
||||
assert_eq!(body["tool_choice"], json!({"type":"web_search"}));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn invalid_search_options_and_conflicting_tool_policies_report_the_field() {
|
||||
for (fields, param) in [
|
||||
(
|
||||
json!({"web_search_options":{"search_context_size":7}}),
|
||||
"web_search_options.search_context_size",
|
||||
),
|
||||
(
|
||||
json!({"web_search_options":{"user_location":{"type":"approximate"}}}),
|
||||
"web_search_options.user_location.approximate",
|
||||
),
|
||||
(
|
||||
json!({"web_search_options":{}, "tool_choice":"none"}),
|
||||
"tool_choice",
|
||||
),
|
||||
(
|
||||
json!({"web_search_options":{},"tools":[{"type":"function","function":{"name":"lookup"}}]}),
|
||||
"tools",
|
||||
),
|
||||
] {
|
||||
let mut request = chat_request(false);
|
||||
request
|
||||
.as_object_mut()
|
||||
.unwrap()
|
||||
.extend(fields.as_object().unwrap().clone());
|
||||
let fixture = Fixture::new(Session::new(Vec::new())).await;
|
||||
let response = fixture.post(request).await;
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
assert_eq!(json_body(response).await["error"]["param"], param);
|
||||
assert!(fixture.execution.captured.lock().unwrap().is_none());
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn maps_inline_user_pdf_without_reencoding_its_data() {
|
||||
let data = "data:application/pdf;base64,JVBERi0xLjQ=";
|
||||
let converted = convert(json!({"model":"model-a","messages":[
|
||||
{"role":"user","content":[{"type":"text","text":"Read the attachment"},
|
||||
{"type":"file","file":{"filename":"document.pdf","file_data":data}}]}
|
||||
]}))
|
||||
.await;
|
||||
assert_eq!(
|
||||
converted["input"][0]["content"],
|
||||
json!([{"type":"input_text","text":"Read the attachment"},
|
||||
{"type":"input_file","filename":"document.pdf","file_data":data}])
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn assistant_history_uses_output_content_for_text_and_refusals() {
|
||||
for content in [
|
||||
@@ -194,7 +261,12 @@ async fn refuses_known_unrepresentable_features_before_starting_execution() {
|
||||
async fn rejects_unmapped_message_content_and_identity_fields() {
|
||||
for message in [
|
||||
json!({"role":"user","name":"speaker","content":"hello"}),
|
||||
json!({"role":"tool","tool_call_id":"call-lookup","name":42,"content":"result"}),
|
||||
json!({"role":"user","content":[{"type":"file","file":{"file_id":"file-example"}}]}),
|
||||
json!({"role":"user","content":[{"type":"file","file":{"file_data":"data:application/pdf;base64,JVBERg=="}}]}),
|
||||
json!({"role":"user","content":[{"type":"file","file":{"filename":"a.pdf","file_data":"data:application/pdf;base64,"}}]}),
|
||||
json!({"role":"user","content":[{"type":"file","file":{"filename":"a.txt","file_data":"data:text/plain;base64,SGk="}}]}),
|
||||
json!({"role":"assistant","content":[{"type":"file","file":{"filename":"a.pdf","file_data":"data:application/pdf;base64,JVBERg=="}}]}),
|
||||
json!({"role":"assistant","content":"hello","reasoning_content":"private"}),
|
||||
json!({"role":"assistant","content":[{"type":"text","text":"history","prompt_cache_breakpoint":{"type":"default"}}]}),
|
||||
json!({"role":"tool","content":"unassociated output"}),
|
||||
@@ -232,8 +304,8 @@ async fn assistant_tool_history_keeps_arguments_opaque_and_cache_annotations_str
|
||||
);
|
||||
assert_eq!(converted["input"][1]["arguments"], arguments);
|
||||
assert_eq!(
|
||||
converted["input"][2]["output"],
|
||||
json!([{"type":"input_text","text":"result"}])
|
||||
converted["input"][2],
|
||||
json!({"type":"function_call_output","call_id":"call-raw","output":[{"type":"input_text","text":"result"}]})
|
||||
);
|
||||
assert_eq!(
|
||||
converted["prompt_cache_options"],
|
||||
|
||||
@@ -61,6 +61,72 @@ async fn json_projects_public_identity_text_and_actual_usage() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn web_search_projects_citations_in_json_and_streams_the_original_links() {
|
||||
let linked = "[IANA](https://www.iana.org/)";
|
||||
let text = format!("见 {linked}");
|
||||
let citation = json!({"type":"url_citation","url":"https://www.iana.org/","title":"IANA",
|
||||
"start_index":2,"end_index":text.chars().count()});
|
||||
let mut prefix = message("🌍 ");
|
||||
prefix["id"] = json!("msg_intro");
|
||||
let search = json!({"type":"web_search_call","id":"ws_lookup","status":"completed",
|
||||
"action":{"type":"search","query":"IANA"}});
|
||||
let mut answer = message(&text);
|
||||
answer["content"][0]["annotations"] = json!([citation]);
|
||||
for stream in [false, true] {
|
||||
let fixture = Fixture::new(Session::new(vec![
|
||||
created(),
|
||||
wire(
|
||||
"response.output_item.done",
|
||||
json!({"output_index":0,"item":prefix}),
|
||||
),
|
||||
wire(
|
||||
"response.output_item.done",
|
||||
json!({"output_index":1,"item":search}),
|
||||
),
|
||||
wire(
|
||||
"response.output_text.delta",
|
||||
json!({"output_index":2,"content_index":0,
|
||||
"item_id":"msg_chat_test","delta":text}),
|
||||
),
|
||||
wire(
|
||||
"response.output_text.annotation.added",
|
||||
json!({"output_index":2,"content_index":0,
|
||||
"item_id":"msg_chat_test","annotation_index":0,"annotation":citation}),
|
||||
),
|
||||
wire(
|
||||
"response.output_item.done",
|
||||
json!({"output_index":2,"item":answer}),
|
||||
),
|
||||
completed(json!([])),
|
||||
]))
|
||||
.await;
|
||||
let mut request = chat_request(stream);
|
||||
request["web_search_options"] = json!({});
|
||||
let response = fixture.post(request).await;
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let expected = format!("🌍 {text}");
|
||||
if stream {
|
||||
let (_, chunks) = sse_body(response).await;
|
||||
assert_eq!(combined(&chunks, "content"), expected);
|
||||
assert_eq!(
|
||||
chunks.last().unwrap()["choices"][0]["finish_reason"],
|
||||
"stop"
|
||||
);
|
||||
} else {
|
||||
let body = json_body(response).await;
|
||||
let choice = &body["choices"][0];
|
||||
assert_eq!(choice["finish_reason"], "stop");
|
||||
assert_eq!(choice["message"]["content"], expected);
|
||||
assert_eq!(
|
||||
choice["message"]["annotations"],
|
||||
json!([{"type":"url_citation","url_citation":{
|
||||
"url":"https://www.iana.org/","title":"IANA","start_index":4,"end_index":expected.chars().count()}}])
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn json_tool_arguments_and_refusal_do_not_become_assistant_text() {
|
||||
for output in [
|
||||
|
||||
@@ -69,6 +69,10 @@ async fn fixture(sharing: ClientPeerSharing, plaintext: &str) -> (AdminTestFixtu
|
||||
.await
|
||||
.unwrap()
|
||||
.record;
|
||||
key.limits = RateLimits {
|
||||
max_concurrency: 5,
|
||||
requests_per_minute: 60,
|
||||
};
|
||||
key.peer_sharing = sharing;
|
||||
key.groups = vec![group.clone()];
|
||||
key.name = "private-key-sentinel".into();
|
||||
@@ -243,7 +247,11 @@ async fn exports_scoped_observations_even_when_the_key_budget_is_exhausted_witho
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = response_json(response).await;
|
||||
assert_eq!(body["object"], "cpr.peer");
|
||||
assert_eq!(body["schema_version"], 2);
|
||||
assert_eq!(body["schema_version"], 3);
|
||||
assert_eq!(
|
||||
body["key_limits"],
|
||||
json!({"max_concurrency":5,"requests_per_minute":60})
|
||||
);
|
||||
assert_eq!(body["usage"]["data"]["windows"][0]["remaining"], "0");
|
||||
assert_eq!(body["usage"]["data"]["allowed"], false);
|
||||
assert_eq!(body["usage"]["data"]["windows"][0]["limit_reached"], true);
|
||||
@@ -358,7 +366,10 @@ async fn plan_only_summary_limits_exported_data_and_preserves_unknown_on_quota_r
|
||||
body["pool"]["plans"],
|
||||
json!([{ "name": "Free", "count": 1 }])
|
||||
);
|
||||
assert_eq!(body["capabilities"], json!(["key_usage", "pool_summary"]));
|
||||
assert_eq!(
|
||||
body["capabilities"],
|
||||
json!(["key_usage", "key_limits", "pool_summary"])
|
||||
);
|
||||
assert!(body.get("quota").is_none());
|
||||
|
||||
*fixture.provider_quota.lock().unwrap() = None;
|
||||
|
||||
@@ -431,6 +431,9 @@ fn decoder_should_preserve_ordinary_request_headers_as_opaque_multivalues() {
|
||||
values("x-codex-installation-id"),
|
||||
vec![b"client-installation".to_vec()]
|
||||
);
|
||||
for name in ["x-openai-account-routing-override", "x-openai-fedramp"] {
|
||||
assert_eq!(values(name), vec![b"first".to_vec(), b"second".to_vec()]);
|
||||
}
|
||||
for preserved in [
|
||||
"accept",
|
||||
"content-type",
|
||||
@@ -450,8 +453,6 @@ fn decoder_should_preserve_ordinary_request_headers_as_opaque_multivalues() {
|
||||
"cookie",
|
||||
"chatgpt-account-id",
|
||||
"chatgpt-project-id",
|
||||
"x-openai-account-routing-override",
|
||||
"x-openai-fedramp",
|
||||
// 上游指纹由运行时画像统一生成,客户端不得覆盖
|
||||
"user-agent",
|
||||
"originator",
|
||||
|
||||
+13
-2
@@ -356,7 +356,7 @@ async fn rewrite_upgrade_and_control_active_response_keep_http_context_until_con
|
||||
),
|
||||
);
|
||||
let (trace, mut socket, server) =
|
||||
start_active_response_with_middleware(Some(frozen), "/custom/websocket").await;
|
||||
start_active_response_with(Some(frozen), "/custom/websocket", "codex-tui/0.166.0").await;
|
||||
// 升级后的默认会话仍持有入口上下文;101 正文结束不再提前取消它
|
||||
assert!(
|
||||
!plan
|
||||
@@ -388,6 +388,7 @@ async fn rewrite_upgrade_and_control_active_response_keep_http_context_until_con
|
||||
assert!(matches!(reply, ClientMessage::Binary(bytes) if bytes.as_ref() == [0,255,1]));
|
||||
assert_eq!(trace.starts.load(Ordering::SeqCst), 1);
|
||||
assert!(!trace.cancelled.load(Ordering::SeqCst));
|
||||
trace.reconnect_for_replay.store(true, Ordering::Release);
|
||||
trace.release_terminal.notify_one();
|
||||
let reply = tokio::time::timeout(Duration::from_secs(2), socket.next())
|
||||
.await
|
||||
@@ -400,6 +401,16 @@ async fn rewrite_upgrade_and_control_active_response_keep_http_context_until_con
|
||||
let event: Value = serde_json::from_str(&text).unwrap();
|
||||
assert_eq!(event["type"], "response.completed");
|
||||
assert_eq!(event["plugin"], true);
|
||||
socket
|
||||
.send(ClientMessage::Ping(vec![9].into()))
|
||||
.await
|
||||
.unwrap();
|
||||
let reply = tokio::time::timeout(Duration::from_secs(2), socket.next())
|
||||
.await
|
||||
.expect("Pong after middleware rewrote the terminal")
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(matches!(reply, ClientMessage::Pong(_)));
|
||||
assert!(plan.incoming.load(Ordering::SeqCst) >= 2);
|
||||
assert!(plan.outgoing.load(Ordering::SeqCst) >= 3);
|
||||
socket.close(None).await.unwrap();
|
||||
@@ -441,7 +452,7 @@ async fn input_plugin_runs_while_a_real_socket_write_is_blocked() {
|
||||
),
|
||||
);
|
||||
let (trace, mut socket, server) =
|
||||
start_active_response_with_middleware(Some(frozen), "/custom/websocket").await;
|
||||
start_active_response_with(Some(frozen), "/custom/websocket", "codex-tui/0.166.0").await;
|
||||
trace.release_terminal.notify_one();
|
||||
// 仅窥视帧头以确认写入开始,不消费 64 MiB 输出,保留真实 TCP 接收窗口背压
|
||||
let MaybeTlsStream::Plain(transport) = socket.get_ref() else {
|
||||
|
||||
@@ -49,6 +49,7 @@ use crate::openai::{api_router, authenticated_client};
|
||||
|
||||
#[derive(Default)]
|
||||
struct CrossingLimitTrace {
|
||||
reconnect_for_replay: AtomicBool,
|
||||
release_terminal: tokio::sync::Notify,
|
||||
starts: AtomicUsize,
|
||||
inputs: Mutex<Vec<Value>>,
|
||||
@@ -74,7 +75,9 @@ impl ExecutionSession for CrossingLimitSession {
|
||||
if !self.terminal_emitted {
|
||||
self.trace.release_terminal.notified().await;
|
||||
self.terminal_emitted = true;
|
||||
return Ok(Some(crossing_limit_terminal_batch()));
|
||||
return Ok(Some(crossing_limit_terminal_batch(
|
||||
self.trace.reconnect_for_replay.load(Ordering::Acquire),
|
||||
)));
|
||||
}
|
||||
self.finalized = true;
|
||||
self.trace.finalized.store(true, Ordering::Release);
|
||||
@@ -197,8 +200,8 @@ fn crossing_limit_started_batch() -> CoordinatedEvent {
|
||||
.expect("started WebSocket batch")
|
||||
}
|
||||
|
||||
fn crossing_limit_terminal_batch() -> CoordinatedEvent {
|
||||
let completed = ProviderEvent::wire(
|
||||
fn crossing_limit_terminal_batch(reconnect: bool) -> CoordinatedEvent {
|
||||
let mut completed = ProviderEvent::wire(
|
||||
ProtocolWireEvent::json(
|
||||
"openai",
|
||||
Some("response.completed".to_owned()),
|
||||
@@ -207,12 +210,20 @@ fn crossing_limit_terminal_batch() -> CoordinatedEvent {
|
||||
"response": {
|
||||
"id": "resp_ws_crossing_limit",
|
||||
"model": "model-a",
|
||||
"output": [],
|
||||
"status": "completed"
|
||||
}
|
||||
}),
|
||||
)
|
||||
.expect("completed wire"),
|
||||
);
|
||||
if reconnect {
|
||||
completed.attach_session_update(
|
||||
gateway_core::operation::ProviderSessionState::new("openai", serde_json::Map::new())
|
||||
.unwrap()
|
||||
.with_websocket_reconnect(),
|
||||
);
|
||||
}
|
||||
CoordinatedEvent::try_batch(vec![completed], CommitRequirement::AlreadyCommitted)
|
||||
.expect("terminal WebSocket batch")
|
||||
}
|
||||
@@ -816,12 +827,13 @@ async fn start_active_response() -> (
|
||||
TestClientSocket,
|
||||
tokio::task::JoinHandle<()>,
|
||||
) {
|
||||
start_active_response_with_middleware(None, "/v1/responses").await
|
||||
start_active_response_with(None, "/v1/responses", "integration-test").await
|
||||
}
|
||||
|
||||
async fn start_active_response_with_middleware(
|
||||
async fn start_active_response_with(
|
||||
plan: Option<gateway_core::engine::middleware::FrozenMiddlewarePlan>,
|
||||
path: &str,
|
||||
user_agent: &str,
|
||||
) -> (
|
||||
Arc<CrossingLimitTrace>,
|
||||
TestClientSocket,
|
||||
@@ -872,6 +884,9 @@ async fn start_active_response_with_middleware(
|
||||
.parse()
|
||||
.expect("authorization"),
|
||||
);
|
||||
request
|
||||
.headers_mut()
|
||||
.insert("user-agent", user_agent.parse().unwrap());
|
||||
let (mut socket, _) = connect_async(request).await.expect("upgrade WebSocket");
|
||||
let response_create = ClientMessage::Text(
|
||||
json!({
|
||||
@@ -905,8 +920,9 @@ async fn start_active_response_with_middleware(
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn requests_received_during_an_active_response_should_execute_in_order() {
|
||||
async fn replay_reconnect_waits_for_queued_requests_and_writes_the_last_terminal() {
|
||||
let (trace, mut socket, server) = start_active_response().await;
|
||||
trace.reconnect_for_replay.store(true, Ordering::Release);
|
||||
for input in ["second", "third"] {
|
||||
socket
|
||||
.send(ClientMessage::Text(
|
||||
@@ -964,7 +980,13 @@ async fn requests_received_during_an_active_response_should_execute_in_order() {
|
||||
vec![json!("hello"), json!("second"), json!("third")]
|
||||
);
|
||||
assert!(!trace.cancelled.load(Ordering::Acquire));
|
||||
socket.close(None).await.expect("close WebSocket");
|
||||
let next = tokio::time::timeout(Duration::from_secs(2), socket.next())
|
||||
.await
|
||||
.expect("replay close timeout")
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(matches!(next, ClientMessage::Close(Some(frame))
|
||||
if frame.code == CloseCode::Normal && frame.reason == "Reconnect with full input."));
|
||||
server.abort();
|
||||
}
|
||||
|
||||
|
||||
@@ -39,6 +39,12 @@ impl ProviderAccountId {
|
||||
pub struct AccountConcurrencyLimit(NonZeroU32);
|
||||
|
||||
impl AccountConcurrencyLimit {
|
||||
/// Key 的正数上限映射到账号可表达的范围,零沿用本机默认值
|
||||
#[must_use]
|
||||
pub fn from_key_limit(limit: u64) -> Option<Self> {
|
||||
Self::new(u32::try_from(limit).unwrap_or(u32::MAX))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn new(value: u32) -> Option<Self> {
|
||||
match NonZeroU32::new(value) {
|
||||
|
||||
@@ -470,6 +470,8 @@ impl AccountQuotaSignals {
|
||||
/// 账号持久事实与可重建运行信号的请求级组合
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountCandidate {
|
||||
/// 本次选号已确认的上游默认上限,与本地可编辑覆盖独立
|
||||
pub inherited_concurrency: Option<super::AccountConcurrencyLimit>,
|
||||
pub account: ProviderAccount,
|
||||
pub signals: AccountRuntimeSignals,
|
||||
}
|
||||
@@ -525,11 +527,16 @@ pub struct AccountSelectionContext {
|
||||
impl AccountSelectionContext {
|
||||
/// 本请求在该账号上可使用的并发上限;资格判断、租约与策略投影必须共用这一口径
|
||||
#[must_use]
|
||||
pub fn concurrency_limit(&self, account: &ProviderAccount) -> AccountConcurrency {
|
||||
pub fn concurrency_limit(&self, candidate: &AccountCandidate) -> AccountConcurrency {
|
||||
if self.reserved_concurrency > 0 {
|
||||
AccountConcurrency::new(self.reserved_concurrency)
|
||||
} else {
|
||||
account.effective_concurrency(self.policy.max_concurrent_per_account())
|
||||
let default = candidate
|
||||
.inherited_concurrency
|
||||
.map_or(self.policy.max_concurrent_per_account(), |limit| {
|
||||
limit.into_non_zero().into()
|
||||
});
|
||||
candidate.account.effective_concurrency(default)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -544,12 +551,23 @@ pub enum AccountEligibilityPolicy {
|
||||
#[default]
|
||||
Enforce,
|
||||
BypassForDiagnostic,
|
||||
FreezeRecovery,
|
||||
}
|
||||
|
||||
impl AccountEligibilityPolicy {
|
||||
#[must_use]
|
||||
pub const fn bypasses_local_eligibility(self) -> bool {
|
||||
matches!(self, Self::BypassForDiagnostic)
|
||||
!matches!(self, Self::Enforce)
|
||||
}
|
||||
|
||||
/// 恢复探测只验证上游结果,不能通过回写 Cookie 或刷新凭据提前解除冻结
|
||||
#[must_use]
|
||||
pub const fn allows_account_state_mutation(self, enabled: bool) -> bool {
|
||||
match self {
|
||||
Self::Enforce => true,
|
||||
Self::BypassForDiagnostic => enabled,
|
||||
Self::FreezeRecovery => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -671,8 +689,7 @@ impl AccountSelector {
|
||||
)
|
||||
})
|
||||
.try_fold((0_u64, 0_u64), |(used, total), candidate| {
|
||||
let capacity =
|
||||
u64::from(context.concurrency_limit(&candidate.account).limit()?.get());
|
||||
let capacity = u64::from(context.concurrency_limit(candidate).limit()?.get());
|
||||
Some((
|
||||
used.saturating_add(u64::from(candidate.signals.in_flight)),
|
||||
total.saturating_add(capacity),
|
||||
@@ -868,7 +885,7 @@ impl AccountSelector {
|
||||
return Some(AccountSchedulingBlocker::ConcurrencyLimit);
|
||||
}
|
||||
if context
|
||||
.concurrency_limit(&candidate.account)
|
||||
.concurrency_limit(candidate)
|
||||
.limit()
|
||||
.is_some_and(|limit| candidate.signals.in_flight >= limit.get())
|
||||
{
|
||||
@@ -885,7 +902,7 @@ const SMART_RESET_HALF_SCORE_SECONDS: f64 = 3_600.0;
|
||||
|
||||
fn capacity_utilization(candidate: &AccountCandidate, context: &AccountSelectionContext) -> f64 {
|
||||
context
|
||||
.concurrency_limit(&candidate.account)
|
||||
.concurrency_limit(candidate)
|
||||
.limit()
|
||||
.map_or(0.0, |limit| {
|
||||
f64::from(candidate.signals.in_flight) / f64::from(limit.get())
|
||||
|
||||
@@ -76,6 +76,7 @@ impl UpstreamInfoFailure {
|
||||
/// 调度侧只消费已发布的中立观测,不读取协议、凭据或持久存储
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub struct UpstreamSchedulingSignal {
|
||||
pub concurrency_limit: Option<super::AccountConcurrencyLimit>,
|
||||
pub exhausted: bool,
|
||||
pub remaining_percent: Option<f64>,
|
||||
}
|
||||
@@ -84,7 +85,7 @@ pub trait UpstreamInfoSource: Send + Sync {
|
||||
fn signal(
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
model: &str,
|
||||
model: Option<&str>,
|
||||
now: SystemTime,
|
||||
) -> Option<UpstreamSchedulingSignal>;
|
||||
}
|
||||
@@ -110,7 +111,7 @@ impl UpstreamInfoSource for UpstreamInfoReader {
|
||||
fn signal(
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
model: &str,
|
||||
model: Option<&str>,
|
||||
now: SystemTime,
|
||||
) -> Option<UpstreamSchedulingSignal> {
|
||||
self.source.get()?.upgrade()?.signal(account, model, now)
|
||||
|
||||
@@ -41,7 +41,7 @@ impl TraceContext {
|
||||
"blocker": AccountSelector.scheduling_blocker(candidate, context)
|
||||
.map(|blocker| format!("{blocker:?}")),
|
||||
"inFlight": signals.in_flight,
|
||||
"concurrencyLimit": context.concurrency_limit(&candidate.account).get(),
|
||||
"concurrencyLimit": context.concurrency_limit(candidate).get(),
|
||||
"quotaRemainingBasisPoints": signals.quota_remaining_rank,
|
||||
"quotaResetAtUnixMs": signals.quota_reset_at.and_then(|reset|
|
||||
reset.duration_since(std::time::UNIX_EPOCH).ok()
|
||||
|
||||
@@ -18,6 +18,7 @@ use crate::engine::continuation::{
|
||||
ContinuationBinding, NativeContinuationPin, NativeContinuationScope, PreviousResponseId,
|
||||
};
|
||||
|
||||
use crate::engine::probe::AccountProbeTarget;
|
||||
use crate::engine::provider::{Provider, ProviderCallMetadata, ProviderRequest, ProviderStream};
|
||||
use crate::engine::{
|
||||
AccountAttemptContext, AttemptContext, AttemptRecord, AttemptTransport, AttemptTrigger,
|
||||
@@ -49,7 +50,7 @@ pub struct AttemptCoordinator<S: ?Sized> {
|
||||
#[derive(Debug, Clone)]
|
||||
enum AccountSelection {
|
||||
Scheduled(Option<crate::account::ProviderAccountId>),
|
||||
Diagnostic(crate::account::ProviderAccountId),
|
||||
Diagnostic(AccountProbeTarget),
|
||||
}
|
||||
|
||||
pub(super) struct CoordinationExtensions {
|
||||
@@ -153,7 +154,7 @@ impl AccountSelection {
|
||||
fn required_account(&self) -> Option<&crate::account::ProviderAccountId> {
|
||||
match self {
|
||||
Self::Scheduled(account) => account.as_ref(),
|
||||
Self::Diagnostic(account) => Some(account),
|
||||
Self::Diagnostic(target) => Some(target.account_id()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -226,7 +227,7 @@ where
|
||||
request: NewModelRequest,
|
||||
operation: Operation,
|
||||
plan: RoutingPlan,
|
||||
required_account: crate::account::ProviderAccountId,
|
||||
target: AccountProbeTarget,
|
||||
continuation: Option<ContinuationBinding>,
|
||||
cancellation: CancellationToken,
|
||||
) -> Result<ResponseExecutionSession<S>, EngineError> {
|
||||
@@ -234,7 +235,7 @@ where
|
||||
request,
|
||||
operation,
|
||||
plan,
|
||||
AccountSelection::Diagnostic(required_account),
|
||||
AccountSelection::Diagnostic(target),
|
||||
CoordinationExtensions::new(continuation, None),
|
||||
cancellation,
|
||||
)
|
||||
@@ -1086,8 +1087,8 @@ where
|
||||
(Some(recovery.account), recovery.transport)
|
||||
} else {
|
||||
match &self.account_selection {
|
||||
AccountSelection::Diagnostic(account) => {
|
||||
(Some(account.clone()), AttemptTransport::Default)
|
||||
AccountSelection::Diagnostic(target) => {
|
||||
(Some(target.account_id().clone()), AttemptTransport::Default)
|
||||
}
|
||||
AccountSelection::Scheduled(_) => (
|
||||
self.recovery_account
|
||||
@@ -1100,9 +1101,9 @@ where
|
||||
}
|
||||
};
|
||||
let account_context = match &self.account_selection {
|
||||
AccountSelection::Diagnostic(account) => AccountAttemptContext::diagnostic(
|
||||
AccountSelection::Diagnostic(target) => AccountAttemptContext::diagnostic(
|
||||
self.excluded_accounts.clone(),
|
||||
account.clone(),
|
||||
target.clone(),
|
||||
self.account_state_owner.clone(),
|
||||
),
|
||||
AccountSelection::Scheduled(_) => AccountAttemptContext::new(
|
||||
|
||||
@@ -62,15 +62,16 @@ use crate::{
|
||||
policy::{ClientApiKeyId, ClientPolicy},
|
||||
provider_ports::{ProviderSessionAffinityPort, ProviderStoreErrorKind},
|
||||
routing::{
|
||||
FrozenAccountScope, ProviderCatalogUnavailable, PublicModelDescriptor, PublicModelId,
|
||||
RoutingContext, RuntimeSnapshot, UpstreamModelId, request_settings::RequestSettings,
|
||||
FrozenAccountScope, ProviderCatalogUnavailable, ProviderModelContent,
|
||||
PublicModelDescriptor, PublicModelId, RoutingContext, RuntimeSnapshot, UpstreamModelId,
|
||||
request_settings::RequestSettings,
|
||||
},
|
||||
runtime::{RuntimeSnapshotHandle, RuntimeSnapshotPublisher},
|
||||
};
|
||||
use futures::{FutureExt as _, future::BoxFuture, pin_mut, select_biased};
|
||||
use futures_timer::Delay;
|
||||
use std::{
|
||||
collections::BTreeSet,
|
||||
collections::{BTreeMap, BTreeSet},
|
||||
sync::{Arc, Weak},
|
||||
time::{Instant, SystemTime},
|
||||
};
|
||||
@@ -1333,14 +1334,14 @@ impl DefaultExecutionService {
|
||||
frozen: Option<Arc<crate::routing::RuntimeSnapshot>>,
|
||||
) -> Result<AccountProbeResult, AccountProbeError> {
|
||||
let AccountProbeRequest {
|
||||
account_id,
|
||||
target,
|
||||
provider_kind,
|
||||
upstream_model,
|
||||
operation,
|
||||
} = request;
|
||||
let observed = ProbeObservation {
|
||||
provider_kind: provider_kind.clone(),
|
||||
account_id: account_id.clone(),
|
||||
account_id: target.account_id().clone(),
|
||||
upstream_model: upstream_model.clone(),
|
||||
};
|
||||
let snapshot = frozen
|
||||
@@ -1401,7 +1402,7 @@ impl DefaultExecutionService {
|
||||
new_request,
|
||||
operation,
|
||||
plan,
|
||||
account_id,
|
||||
target,
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
)
|
||||
@@ -1575,11 +1576,27 @@ impl ExecutionService for DefaultExecutionService {
|
||||
}
|
||||
continue;
|
||||
};
|
||||
let contributions: Vec<_> = client
|
||||
.snapshot
|
||||
.contributed_models_for_scope(kind, scope)
|
||||
.collect();
|
||||
let mut contributed_by_model = BTreeMap::new();
|
||||
if !contributions.is_empty()
|
||||
&& models
|
||||
.iter()
|
||||
.any(|entry| matches!(entry.content, ProviderModelContent::Adapted(_)))
|
||||
{
|
||||
contributed_by_model.extend(
|
||||
contributions
|
||||
.iter()
|
||||
.map(|model| (&model.id, &model.presentation)),
|
||||
);
|
||||
}
|
||||
// 保留上游顺序;映射只选择完整的目标对象,不能跨账号/模型混拼字段
|
||||
let by_id = models
|
||||
.iter()
|
||||
.map(|entry| (entry.model.as_str(), entry))
|
||||
.collect::<std::collections::BTreeMap<_, _>>();
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
for entry in &models {
|
||||
let target = client.snapshot.mapped_model(entry.model.as_str());
|
||||
if !scope.allows_provider_model(kind, &target) {
|
||||
@@ -1598,7 +1615,11 @@ impl ExecutionService for DefaultExecutionService {
|
||||
continue;
|
||||
}
|
||||
if seen.insert(model.clone()) {
|
||||
result.push(source.content.for_public_model(model));
|
||||
result.push(
|
||||
source
|
||||
.content
|
||||
.for_public_model(model, &contributed_by_model),
|
||||
);
|
||||
}
|
||||
}
|
||||
for model in client.snapshot.public_models_for_provider(kind) {
|
||||
@@ -1617,13 +1638,17 @@ impl ExecutionService for DefaultExecutionService {
|
||||
}
|
||||
if let Some(source) = by_id.get(target.as_str()) {
|
||||
seen.insert(model.clone());
|
||||
result.push(source.content.for_public_model(model));
|
||||
result.push(
|
||||
source
|
||||
.content
|
||||
.for_public_model(model, &contributed_by_model),
|
||||
);
|
||||
}
|
||||
}
|
||||
// 原生完整对象优先;只补入当前授权账号实际提供的独立适配画像
|
||||
for profile in client.snapshot.contributed_profiles_for_scope(kind, scope) {
|
||||
if seen.insert(profile.model().clone()) {
|
||||
result.push(PublicModelDescriptor::Adapted(profile));
|
||||
// 没有 Provider 目录条目的适配模型仍由当前授权账号贡献
|
||||
for model in contributions {
|
||||
if seen.insert(model.id.clone()) {
|
||||
result.push(PublicModelDescriptor::Adapted(model.public_profile()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -35,8 +35,8 @@ use async_trait::async_trait;
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::account::{
|
||||
AccountCandidate, AccountSelection, AccountSelectionContext, AccountSelectionPolicy, FastMode,
|
||||
ProviderAccountId,
|
||||
AccountCandidate, AccountEligibilityPolicy, AccountSelection, AccountSelectionContext,
|
||||
AccountSelectionPolicy, FastMode, ProviderAccountId,
|
||||
};
|
||||
use crate::engine::continuation::{ContinuationBinding, NativeContinuationPin};
|
||||
use crate::error::{
|
||||
@@ -128,7 +128,7 @@ pub struct AccountAttemptContext {
|
||||
required_account: Option<ProviderAccountId>,
|
||||
state_owner: Option<ProviderAccountStateOwner>,
|
||||
credential_recovery_attempted: bool,
|
||||
diagnostic_required_account: bool,
|
||||
eligibility: AccountEligibilityPolicy,
|
||||
account_scope: Option<Arc<crate::account::scope::FrozenAccountScope>>,
|
||||
}
|
||||
|
||||
@@ -144,7 +144,7 @@ impl AccountAttemptContext {
|
||||
required_account,
|
||||
state_owner,
|
||||
credential_recovery_attempted: false,
|
||||
diagnostic_required_account: false,
|
||||
eligibility: AccountEligibilityPolicy::Enforce,
|
||||
account_scope: None,
|
||||
}
|
||||
}
|
||||
@@ -153,17 +153,18 @@ impl AccountAttemptContext {
|
||||
///
|
||||
/// 诊断只绕过本地的被动可用性投影,不能换号,并继续受账号租约保护
|
||||
#[must_use]
|
||||
pub const fn diagnostic(
|
||||
pub fn diagnostic(
|
||||
excluded_accounts: BTreeSet<ProviderAccountId>,
|
||||
required_account: ProviderAccountId,
|
||||
target: probe::AccountProbeTarget,
|
||||
state_owner: Option<ProviderAccountStateOwner>,
|
||||
) -> Self {
|
||||
let (required_account, eligibility) = target.into_selection();
|
||||
Self {
|
||||
excluded_accounts,
|
||||
required_account: Some(required_account),
|
||||
state_owner,
|
||||
credential_recovery_attempted: false,
|
||||
diagnostic_required_account: true,
|
||||
eligibility,
|
||||
account_scope: None,
|
||||
}
|
||||
}
|
||||
@@ -205,12 +206,6 @@ impl AccountAttemptContext {
|
||||
self.credential_recovery_attempted
|
||||
}
|
||||
|
||||
/// 此尝试是否要由真实上游而非本地投影确认固定账号的可用性
|
||||
#[must_use]
|
||||
pub const fn is_diagnostic_required_account(&self) -> bool {
|
||||
self.diagnostic_required_account
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn account_scope(&self) -> Option<&Arc<crate::account::scope::FrozenAccountScope>> {
|
||||
self.account_scope.as_ref()
|
||||
@@ -921,7 +916,12 @@ impl AttemptContext {
|
||||
/// 管理端固定账号诊断会跳过本地可用性投影,保留全部租约约束
|
||||
#[must_use]
|
||||
pub const fn is_diagnostic_required_account(&self) -> bool {
|
||||
self.account.is_diagnostic_required_account()
|
||||
self.account_eligibility().bypasses_local_eligibility()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn account_eligibility(&self) -> AccountEligibilityPolicy {
|
||||
self.account.eligibility
|
||||
}
|
||||
|
||||
/// 当前路由候选的冻结账号范围;管理端诊断为 `None`
|
||||
|
||||
@@ -483,7 +483,7 @@ impl RequestPolicyContext {
|
||||
.is_none(),
|
||||
weight: candidate.account.weight().get(),
|
||||
in_flight: candidate.signals.in_flight,
|
||||
maximum_concurrency: context.concurrency_limit(&candidate.account).get(),
|
||||
maximum_concurrency: context.concurrency_limit(candidate).get(),
|
||||
last_started_at: candidate.signals.last_started_at,
|
||||
quota_reset_at: candidate.signals.quota_reset_at,
|
||||
quota_remaining_rank: candidate.signals.quota_remaining_rank,
|
||||
|
||||
@@ -8,11 +8,39 @@ use futures::future::BoxFuture;
|
||||
use crate::error::{ClientVisibleUpstreamResponse, GatewayError, GatewayErrorKind};
|
||||
use crate::identity::ProviderKind;
|
||||
use crate::routing::UpstreamModelId;
|
||||
use crate::{account::ProviderAccountId, operation::Operation, upstream::UpstreamSendState};
|
||||
use crate::{
|
||||
account::{AccountEligibilityPolicy, ProviderAccountId},
|
||||
operation::Operation,
|
||||
upstream::UpstreamSendState,
|
||||
};
|
||||
|
||||
/// 探测始终固定账号,恢复探测的账号状态由解冻任务持有
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum AccountProbeTarget {
|
||||
ConnectionTest(ProviderAccountId),
|
||||
FreezeRecovery(ProviderAccountId),
|
||||
}
|
||||
|
||||
impl AccountProbeTarget {
|
||||
pub fn account_id(&self) -> &ProviderAccountId {
|
||||
match self {
|
||||
Self::ConnectionTest(account) | Self::FreezeRecovery(account) => account,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn into_selection(self) -> (ProviderAccountId, AccountEligibilityPolicy) {
|
||||
match self {
|
||||
Self::ConnectionTest(account) => {
|
||||
(account, AccountEligibilityPolicy::BypassForDiagnostic)
|
||||
}
|
||||
Self::FreezeRecovery(account) => (account, AccountEligibilityPolicy::FreezeRecovery),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct AccountProbeRequest {
|
||||
pub account_id: ProviderAccountId,
|
||||
pub target: AccountProbeTarget,
|
||||
pub provider_kind: ProviderKind,
|
||||
pub upstream_model: UpstreamModelId,
|
||||
pub operation: Operation,
|
||||
|
||||
@@ -18,6 +18,11 @@ pub enum HealthState {
|
||||
pub trait HealthProbe: Send + Sync {
|
||||
fn name(&self) -> &'static str;
|
||||
|
||||
/// 仅接收请求必需的依赖参与就绪检查
|
||||
fn required_for_readiness(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
fn check(&self) -> BoxFuture<'_, HealthState>;
|
||||
}
|
||||
|
||||
|
||||
@@ -123,6 +123,9 @@ pub struct ProviderSessionState {
|
||||
payload: Map<String, Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
extension_owner: Option<Box<ExtensionSessionOwner>>,
|
||||
// 仅供本次终态交付后处置客户端连接,不持久化或传给外部适配器
|
||||
#[serde(skip)]
|
||||
reconnect_websocket: bool,
|
||||
}
|
||||
|
||||
/// 受管扩展会话的宿主归属,不能把插件私有状态交给原生 Provider 解码
|
||||
@@ -155,6 +158,7 @@ impl ProviderSessionState {
|
||||
provider,
|
||||
payload,
|
||||
extension_owner: None,
|
||||
reconnect_websocket: false,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -168,6 +172,18 @@ impl ProviderSessionState {
|
||||
&self.payload
|
||||
}
|
||||
|
||||
/// 提示能重连重放的客户端在本次成功交付后重建连接
|
||||
#[must_use]
|
||||
pub const fn with_websocket_reconnect(mut self) -> Self {
|
||||
self.reconnect_websocket = true;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn prefers_websocket_reconnect(&self) -> bool {
|
||||
self.reconnect_websocket
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_extension_owner(mut self, owner: ExtensionSessionOwner) -> Self {
|
||||
self.extension_owner = Some(Box::new(owner));
|
||||
|
||||
@@ -32,6 +32,12 @@ pub struct ContributedModel {
|
||||
pub presentation: ModelPresentation,
|
||||
}
|
||||
|
||||
impl ContributedModel {
|
||||
pub(crate) fn public_profile(&self) -> PublicModelProfile {
|
||||
PublicModelProfile::new(self.id.clone(), self.presentation.clone())
|
||||
}
|
||||
}
|
||||
|
||||
/// Provider 客户端目录条目;缺少原生协议正文时使用已编译的通用画像
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderModelDescriptor {
|
||||
@@ -56,13 +62,18 @@ pub enum PublicModelDescriptor {
|
||||
}
|
||||
|
||||
impl ProviderModelContent {
|
||||
pub(crate) fn for_public_model(&self, model: PublicModelId) -> PublicModelDescriptor {
|
||||
pub(crate) fn for_public_model(
|
||||
&self,
|
||||
model: PublicModelId,
|
||||
contributed: &BTreeMap<&PublicModelId, &ModelPresentation>,
|
||||
) -> PublicModelDescriptor {
|
||||
match self {
|
||||
Self::Native(payload) => PublicModelDescriptor::Native {
|
||||
model,
|
||||
payload: payload.clone(),
|
||||
},
|
||||
Self::Adapted(presentation) => {
|
||||
let presentation = contributed.get(&model).copied().unwrap_or(presentation);
|
||||
PublicModelDescriptor::Adapted(PublicModelProfile::new(model, presentation.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -23,7 +23,7 @@ use super::{
|
||||
RuntimeAccountDirectory, UpstreamModelId,
|
||||
};
|
||||
|
||||
const MAXIMUM_CATALOG_STABILITY_ATTEMPTS: usize = 4;
|
||||
const MAXIMUM_SNAPSHOT_STABILITY_ATTEMPTS: usize = 4;
|
||||
|
||||
type ModelCatalogAccounts = BTreeMap<ProviderKind, BTreeMap<String, BTreeSet<ProviderAccountId>>>;
|
||||
|
||||
@@ -328,14 +328,17 @@ impl RuntimeSnapshotCompiler {
|
||||
ConfigRevision,
|
||||
crate::routing::extensions::ExtensionSetReference,
|
||||
)> = None;
|
||||
for _ in 0..MAXIMUM_CATALOG_STABILITY_ATTEMPTS {
|
||||
let mut instability = RuntimeSnapshotCompileError::CatalogChanged;
|
||||
for _ in 0..MAXIMUM_SNAPSHOT_STABILITY_ATTEMPTS {
|
||||
let facts = self
|
||||
.store
|
||||
.load_snapshot_facts()
|
||||
.await
|
||||
.map_err(|_| RuntimeSnapshotCompileError::StoreUnavailable)?;
|
||||
if facts.config_revision != facts.observed_current_revision {
|
||||
return Err(RuntimeSnapshotCompileError::RevisionChanged);
|
||||
// 并发提交使本轮事实过期时重读完整快照,不能混用不同版本的权限
|
||||
instability = RuntimeSnapshotCompileError::RevisionChanged;
|
||||
continue;
|
||||
}
|
||||
let revision = facts.config_revision;
|
||||
let extensions = match &self.extensions {
|
||||
@@ -347,7 +350,7 @@ impl RuntimeSnapshotCompiler {
|
||||
*prepared_revision != revision || !prepared.is_ready()
|
||||
});
|
||||
if must_prepare {
|
||||
let prepared = preparer
|
||||
let prepared = match preparer
|
||||
.prepare_with_request_body_limit(
|
||||
revision,
|
||||
usize::try_from(
|
||||
@@ -356,7 +359,21 @@ impl RuntimeSnapshotCompiler {
|
||||
.unwrap_or(usize::MAX),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| RuntimeSnapshotCompileError::ExtensionsUnavailable)?;
|
||||
{
|
||||
Ok(prepared) => prepared,
|
||||
Err(_) => {
|
||||
// 扩展读取也可能遇到并发提交;只有版本已变才重读事实
|
||||
let observed =
|
||||
self.store.current_config_revision().await.map_err(|_| {
|
||||
RuntimeSnapshotCompileError::StoreUnavailable
|
||||
})?;
|
||||
if observed == revision {
|
||||
return Err(RuntimeSnapshotCompileError::ExtensionsUnavailable);
|
||||
}
|
||||
instability = RuntimeSnapshotCompileError::RevisionChanged;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
prepared_extensions = Some((revision, prepared));
|
||||
}
|
||||
prepared_extensions
|
||||
@@ -390,7 +407,8 @@ impl RuntimeSnapshotCompiler {
|
||||
.map_err(|_| RuntimeSnapshotCompileError::StoreUnavailable)?
|
||||
!= revision
|
||||
{
|
||||
return Err(RuntimeSnapshotCompileError::RevisionChanged);
|
||||
instability = RuntimeSnapshotCompileError::RevisionChanged;
|
||||
continue;
|
||||
}
|
||||
let snapshot = snapshot
|
||||
.with_provider_catalog_generations(if cached.is_some() && !refresh_catalog {
|
||||
@@ -401,8 +419,9 @@ impl RuntimeSnapshotCompiler {
|
||||
.with_extensions(extensions);
|
||||
return Ok(snapshot);
|
||||
}
|
||||
instability = RuntimeSnapshotCompileError::CatalogChanged;
|
||||
}
|
||||
Err(RuntimeSnapshotCompileError::CatalogChanged)
|
||||
Err(instability)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -682,24 +701,18 @@ impl RuntimeSnapshot {
|
||||
self.extensions.as_ref().map_or(&[], |set| set.models())
|
||||
}
|
||||
|
||||
pub(crate) fn contributed_profiles_for_scope(
|
||||
pub(crate) fn contributed_models_for_scope(
|
||||
&self,
|
||||
provider: &ProviderKind,
|
||||
scope: &FrozenAccountScope,
|
||||
) -> Vec<super::PublicModelProfile> {
|
||||
self.contributed_models()
|
||||
.iter()
|
||||
.filter(|model| {
|
||||
&model.provider == provider
|
||||
&& model.accounts.iter().any(|account| {
|
||||
scope.account_provider(account) == Some(provider)
|
||||
&& scope.allows_model(account, &self.mapped_model(model.id.as_str()))
|
||||
})
|
||||
})
|
||||
.map(|model| {
|
||||
super::PublicModelProfile::new(model.id.clone(), model.presentation.clone())
|
||||
})
|
||||
.collect()
|
||||
) -> impl Iterator<Item = &super::ContributedModel> {
|
||||
self.contributed_models().iter().filter(move |model| {
|
||||
&model.provider == provider
|
||||
&& model.accounts.iter().any(|account| {
|
||||
scope.account_provider(account) == Some(provider)
|
||||
&& scope.allows_model(account, &self.mapped_model(model.id.as_str()))
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn model_alias(&self, model: &str) -> Option<&super::ContributedModelAlias> {
|
||||
@@ -995,13 +1008,21 @@ impl RuntimeSnapshot {
|
||||
provider: &ProviderKind,
|
||||
scope: &FrozenAccountScope,
|
||||
) -> Vec<super::PublicModelProfile> {
|
||||
self.public_model_profiles_for_provider(provider)
|
||||
.into_iter()
|
||||
.filter(|profile| {
|
||||
self.native_catalog_model_allowed_for_scope(provider, profile.model(), scope)
|
||||
})
|
||||
.chain(self.contributed_profiles_for_scope(provider, scope))
|
||||
.collect()
|
||||
let mut profiles = self.public_model_profiles_for_provider(provider);
|
||||
profiles.retain(|profile| {
|
||||
self.native_catalog_model_allowed_for_scope(provider, profile.model(), scope)
|
||||
});
|
||||
// Provider 画像按公开 ID 排序,追加的插件画像保留注册顺序,不参与二分查找
|
||||
let provider_count = profiles.len();
|
||||
for model in self.contributed_models_for_scope(provider, scope) {
|
||||
match profiles[..provider_count]
|
||||
.binary_search_by(|profile| profile.model().cmp(&model.id))
|
||||
{
|
||||
Ok(index) => profiles[index].presentation.clone_from(&model.presentation),
|
||||
Err(_) => profiles.push(model.public_profile()),
|
||||
}
|
||||
}
|
||||
profiles
|
||||
}
|
||||
|
||||
/// 已知模型只属于明确声明它的 Provider;未知模型仍可由非完整目录透传
|
||||
|
||||
@@ -115,6 +115,10 @@ impl HealthProbe for RuntimeSnapshotHandle {
|
||||
"runtime_snapshot"
|
||||
}
|
||||
|
||||
fn required_for_readiness(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn check(&self) -> BoxFuture<'_, HealthState> {
|
||||
Box::pin(async move {
|
||||
if self.acquire().is_ok() {
|
||||
|
||||
@@ -50,6 +50,7 @@ fn account(id: &str) -> ProviderAccount {
|
||||
|
||||
fn candidate(id: &str, in_flight: u32, remaining: Option<u64>) -> AccountCandidate {
|
||||
AccountCandidate {
|
||||
inherited_concurrency: None,
|
||||
account: account(id),
|
||||
signals: AccountRuntimeSignals {
|
||||
in_flight,
|
||||
@@ -215,6 +216,7 @@ fn inconclusive_quota_observation_cannot_erase_confirmed_access() {
|
||||
fn diagnostic_selection_bypasses_all_local_account_eligibility() {
|
||||
let observed_at = SystemTime::now();
|
||||
let exhausted = AccountCandidate {
|
||||
inherited_concurrency: None,
|
||||
account: account("acct_exhausted").with_account_facts(
|
||||
true,
|
||||
CredentialState::Ready,
|
||||
@@ -233,6 +235,7 @@ fn diagnostic_selection_bypasses_all_local_account_eligibility() {
|
||||
},
|
||||
};
|
||||
let disabled = AccountCandidate {
|
||||
inherited_concurrency: None,
|
||||
account: account("acct_disabled").with_account_facts(
|
||||
false,
|
||||
CredentialState::Ready,
|
||||
@@ -1113,11 +1116,8 @@ fn reserved_concurrency_uses_its_own_limit_without_reducing_normal_capacity() {
|
||||
|
||||
// 独立额度可以大于账号普通上限,极值不能回绕成无限并发
|
||||
reserved.reserved_concurrency = u32::MAX;
|
||||
assert_eq!(
|
||||
reserved.concurrency_limit(&candidates[0].account).get(),
|
||||
u32::MAX
|
||||
);
|
||||
assert_eq!(normal.concurrency_limit(&candidates[0].account).get(), 5);
|
||||
assert_eq!(reserved.concurrency_limit(&candidates[0]).get(), u32::MAX);
|
||||
assert_eq!(normal.concurrency_limit(&candidates[0]).get(), 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
//! 单行 `model_requests`、账号重试与下游提交屏障测试
|
||||
|
||||
use gateway_core::engine::probe::AccountProbeTarget;
|
||||
use std::collections::{BTreeMap, BTreeSet, VecDeque};
|
||||
use std::num::NonZeroU32;
|
||||
use std::sync::{
|
||||
@@ -2753,7 +2754,7 @@ fn continuation_history_failure_without_replay_proof_is_delivered_without_hidden
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn diagnostic_account_reaches_provider_and_matching_metadata_succeeds() {
|
||||
fn recovery_probe_reaches_provider_with_read_only_account() {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let (coordinator, _, provider) = coordinator(vec![Script::Stream {
|
||||
@@ -2765,7 +2766,7 @@ fn diagnostic_account_reaches_provider_and_matching_metadata_succeeds() {
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation,
|
||||
route_plan,
|
||||
required.clone(),
|
||||
AccountProbeTarget::FreezeRecovery(required.clone()),
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
))
|
||||
@@ -2777,6 +2778,11 @@ fn diagnostic_account_reaches_provider_and_matching_metadata_succeeds() {
|
||||
assert_eq!(contexts.len(), 1);
|
||||
assert_eq!(contexts[0].required_account(), Some(&required));
|
||||
assert!(contexts[0].is_diagnostic_required_account());
|
||||
assert!(
|
||||
!contexts[0]
|
||||
.account_eligibility()
|
||||
.allows_account_state_mutation(true)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -7397,7 +7403,9 @@ fn error_protection_tracks_real_retry_attempts_and_skips_diagnostic_runs() {
|
||||
request,
|
||||
operation,
|
||||
route_plan,
|
||||
ProviderAccountId::new("acct_first").unwrap(),
|
||||
AccountProbeTarget::ConnectionTest(
|
||||
ProviderAccountId::new("acct_first").unwrap(),
|
||||
),
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
)
|
||||
|
||||
@@ -16,11 +16,8 @@ fn identity_queries_keep_frozen_permissions_without_using_execution_resources()
|
||||
let service = DefaultExecutionService::new(
|
||||
snapshots.clone(),
|
||||
store.clone(),
|
||||
ProviderRegistry::new([Arc::new(NativeCatalogProvider {
|
||||
fail: false,
|
||||
empty: false,
|
||||
}) as Arc<dyn Provider>])
|
||||
.unwrap(),
|
||||
ProviderRegistry::new([Arc::new(CatalogProvider::Native) as Arc<dyn Provider>])
|
||||
.unwrap(),
|
||||
Arc::new(UnusedAdmissions),
|
||||
Arc::new(UnusedContinuation),
|
||||
usage.clone(),
|
||||
@@ -91,11 +88,8 @@ fn execution_scope_limits_apply_only_when_starting_a_model() {
|
||||
let service = DefaultExecutionService::new(
|
||||
RuntimeSnapshotHandle::new(start_snapshot()),
|
||||
store.clone(),
|
||||
ProviderRegistry::new([Arc::new(NativeCatalogProvider {
|
||||
fail: false,
|
||||
empty: false,
|
||||
}) as Arc<dyn Provider>])
|
||||
.unwrap(),
|
||||
ProviderRegistry::new([Arc::new(CatalogProvider::Native) as Arc<dyn Provider>])
|
||||
.unwrap(),
|
||||
Arc::new(UnusedAdmissions),
|
||||
Arc::new(UnusedContinuation),
|
||||
usage.clone(),
|
||||
|
||||
@@ -2,7 +2,11 @@
|
||||
|
||||
mod identity;
|
||||
mod outbound;
|
||||
use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference};
|
||||
use gateway_core::engine::probe::AccountProbeTarget;
|
||||
use gateway_core::routing::{
|
||||
ContributedModel, ContributedModelAlias,
|
||||
extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference},
|
||||
};
|
||||
|
||||
use std::{
|
||||
collections::{BTreeMap, BTreeSet, VecDeque},
|
||||
@@ -2111,6 +2115,7 @@ impl Provider for PolicySelectingProvider {
|
||||
)
|
||||
.with_scheduling(None, AccountWeight::new(100).unwrap());
|
||||
let candidates = [AccountCandidate {
|
||||
inherited_concurrency: None,
|
||||
account,
|
||||
signals: AccountRuntimeSignals {
|
||||
in_flight: 0,
|
||||
@@ -3131,7 +3136,9 @@ fn account_probe_should_not_write_to_the_persistent_execution_store() {
|
||||
|
||||
let error = block_on(service.probe(
|
||||
AccountProbeRequest {
|
||||
account_id: ProviderAccountId::new("acct_probe").expect("account ID"),
|
||||
target: AccountProbeTarget::ConnectionTest(
|
||||
ProviderAccountId::new("acct_probe").expect("account ID"),
|
||||
),
|
||||
provider_kind: ProviderKind::new("openai").expect("provider kind"),
|
||||
upstream_model: UpstreamModelId::new("gpt-probe").expect("model ID"),
|
||||
operation: probe_operation(),
|
||||
@@ -3163,7 +3170,9 @@ fn probe_failures_should_be_observable_without_a_model_request_row() {
|
||||
|
||||
let error = block_on(service.probe(
|
||||
AccountProbeRequest {
|
||||
account_id: ProviderAccountId::new("acct_probe").expect("account ID"),
|
||||
target: AccountProbeTarget::ConnectionTest(
|
||||
ProviderAccountId::new("acct_probe").expect("account ID"),
|
||||
),
|
||||
provider_kind: ProviderKind::new("openai").expect("provider kind"),
|
||||
upstream_model: UpstreamModelId::new("gpt-probe").expect("model ID"),
|
||||
operation: probe_operation(),
|
||||
@@ -3209,7 +3218,9 @@ fn provider_local_probe_failure_should_remain_distinct_from_upstream() {
|
||||
|
||||
let error = block_on(service.probe(
|
||||
AccountProbeRequest {
|
||||
account_id: ProviderAccountId::new("acct_probe").expect("account ID"),
|
||||
target: AccountProbeTarget::ConnectionTest(
|
||||
ProviderAccountId::new("acct_probe").expect("account ID"),
|
||||
),
|
||||
provider_kind: ProviderKind::new("openai").expect("provider kind"),
|
||||
upstream_model: UpstreamModelId::new("gpt-probe").expect("model ID"),
|
||||
operation: probe_operation(),
|
||||
@@ -3252,7 +3263,7 @@ fn diagnostic_probe_does_not_apply_data_plane_account_model_policy() {
|
||||
|
||||
let error = block_on(service.probe(
|
||||
AccountProbeRequest {
|
||||
account_id,
|
||||
target: AccountProbeTarget::ConnectionTest(account_id),
|
||||
provider_kind: provider,
|
||||
upstream_model: UpstreamModelId::new("gpt-probe").expect("model"),
|
||||
operation: probe_operation(),
|
||||
@@ -3284,7 +3295,9 @@ fn probe_observation_store_failure_preserves_the_provider_error() {
|
||||
|
||||
let error = block_on(service.probe(
|
||||
AccountProbeRequest {
|
||||
account_id: ProviderAccountId::new("acct_probe").expect("account ID"),
|
||||
target: AccountProbeTarget::ConnectionTest(
|
||||
ProviderAccountId::new("acct_probe").expect("account ID"),
|
||||
),
|
||||
provider_kind: ProviderKind::new("openai").expect("provider kind"),
|
||||
upstream_model: UpstreamModelId::new("gpt-probe").expect("model ID"),
|
||||
operation: probe_operation(),
|
||||
@@ -3300,15 +3313,22 @@ fn probe_observation_store_failure_preserves_the_provider_error() {
|
||||
|
||||
struct FailingProvider;
|
||||
|
||||
struct NativeCatalogProvider {
|
||||
fail: bool,
|
||||
empty: bool,
|
||||
#[derive(Clone, Copy)]
|
||||
enum CatalogProvider {
|
||||
Native,
|
||||
Empty,
|
||||
Unavailable,
|
||||
Generic,
|
||||
Xai,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for NativeCatalogProvider {
|
||||
impl Provider for CatalogProvider {
|
||||
fn name(&self) -> &'static str {
|
||||
"openai"
|
||||
match self {
|
||||
Self::Xai => "xai",
|
||||
_ => "openai",
|
||||
}
|
||||
}
|
||||
fn catalog_generation(&self) -> ProviderCatalogGeneration {
|
||||
ProviderCatalogGeneration::default()
|
||||
@@ -3330,29 +3350,41 @@ impl Provider for NativeCatalogProvider {
|
||||
assert!(scope.allows(&ProviderAccountId::new("acct_start").expect("account")));
|
||||
assert!(!scope.allows(&ProviderAccountId::new("acct_other").expect("account")));
|
||||
assert_eq!((protocol, version), ("codex", "0.154.0"));
|
||||
if self.fail {
|
||||
return Err(gateway_core::routing::ProviderCatalogUnavailable);
|
||||
match self {
|
||||
Self::Unavailable => return Err(gateway_core::routing::ProviderCatalogUnavailable),
|
||||
Self::Empty => return Ok(Some(Vec::new())),
|
||||
Self::Xai => return Ok(None),
|
||||
Self::Native | Self::Generic => {}
|
||||
}
|
||||
if self.empty {
|
||||
return Ok(Some(Vec::new()));
|
||||
}
|
||||
Ok(Some(
|
||||
["gpt-new", "gpt-start"]
|
||||
.into_iter()
|
||||
.map(|id| gateway_core::routing::ProviderModelDescriptor {
|
||||
model: UpstreamModelId::new(id).expect("id"),
|
||||
content: gateway_core::routing::ProviderModelContent::Native(
|
||||
RawJsonPayload::new(
|
||||
"codex",
|
||||
serde_json::to_vec(&json!({"slug":id,"future":null}))
|
||||
.expect("JSON")
|
||||
.into(),
|
||||
)
|
||||
.expect("payload"),
|
||||
let mut models = ["gpt-new", "gpt-start"]
|
||||
.into_iter()
|
||||
.map(|id| gateway_core::routing::ProviderModelDescriptor {
|
||||
model: UpstreamModelId::new(id).expect("id"),
|
||||
content: gateway_core::routing::ProviderModelContent::Native(
|
||||
RawJsonPayload::new(
|
||||
"codex",
|
||||
serde_json::to_vec(&json!({"slug":id,"future":null}))
|
||||
.expect("JSON")
|
||||
.into(),
|
||||
)
|
||||
.expect("payload"),
|
||||
),
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
if matches!(self, Self::Generic) {
|
||||
models.extend(["shared-model", "z-last"].map(|id| {
|
||||
gateway_core::routing::ProviderModelDescriptor {
|
||||
model: UpstreamModelId::new(id).unwrap(),
|
||||
content: gateway_core::routing::ProviderModelContent::Adapted(
|
||||
gateway_core::routing::ModelPresentation::new(
|
||||
Some(format!("Provider {id}")),
|
||||
None,
|
||||
),
|
||||
),
|
||||
})
|
||||
.collect(),
|
||||
))
|
||||
}
|
||||
}));
|
||||
}
|
||||
Ok(Some(models))
|
||||
}
|
||||
async fn execute(
|
||||
self: Arc<Self>,
|
||||
@@ -3363,32 +3395,37 @@ impl Provider for NativeCatalogProvider {
|
||||
}
|
||||
}
|
||||
|
||||
struct CatalogContributions {
|
||||
aliases: Vec<ContributedModelAlias>,
|
||||
models: Vec<ContributedModel>,
|
||||
}
|
||||
|
||||
impl ExtensionSetLease for CatalogContributions {
|
||||
fn is_ready(&self) -> bool {
|
||||
true
|
||||
}
|
||||
fn model_aliases(&self) -> &[ContributedModelAlias] {
|
||||
&self.aliases
|
||||
}
|
||||
fn models(&self) -> &[ContributedModel] {
|
||||
&self.models
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_catalog_forwards_scope_maps_whole_objects_and_omits_unroutable_models() {
|
||||
#[derive(Debug)]
|
||||
struct Aliases(
|
||||
Vec<gateway_core::routing::ContributedModelAlias>,
|
||||
Vec<gateway_core::routing::ContributedModel>,
|
||||
);
|
||||
impl gateway_core::routing::extensions::ExtensionSetLease for Aliases {
|
||||
fn is_ready(&self) -> bool {
|
||||
true
|
||||
}
|
||||
fn model_aliases(&self) -> &[gateway_core::routing::ContributedModelAlias] {
|
||||
&self.0
|
||||
}
|
||||
fn models(&self) -> &[gateway_core::routing::ContributedModel] {
|
||||
&self.1
|
||||
}
|
||||
}
|
||||
for (fail, empty) in [(false, false), (false, true), (true, false)] {
|
||||
for provider in [
|
||||
CatalogProvider::Native,
|
||||
CatalogProvider::Empty,
|
||||
CatalogProvider::Unavailable,
|
||||
] {
|
||||
let contributed = [
|
||||
("gpt-start", "acct_start"),
|
||||
("adapted", "acct_start"),
|
||||
("private-model", "acct_other"),
|
||||
]
|
||||
.into_iter()
|
||||
.map(|(id, account)| gateway_core::routing::ContributedModel {
|
||||
.map(|(id, account)| ContributedModel {
|
||||
owner: "catalog-plugin".into(),
|
||||
id: PublicModelId::new(id).unwrap(),
|
||||
provider: ProviderKind::new("openai").unwrap(),
|
||||
@@ -3412,25 +3449,21 @@ fn client_catalog_forwards_scope_maps_whole_objects_and_omits_unroutable_models(
|
||||
.with_settings(&settings)
|
||||
.unwrap()
|
||||
.with_extensions(Some(ExtensionSetReference::new(
|
||||
gateway_core::routing::extensions::ExtensionSetId::new("catalog-aliases".into())
|
||||
.unwrap(),
|
||||
Arc::new(Aliases(
|
||||
vec![gateway_core::routing::ContributedModelAlias {
|
||||
ExtensionSetId::new("catalog-aliases".into()).unwrap(),
|
||||
Arc::new(CatalogContributions {
|
||||
aliases: vec![ContributedModelAlias {
|
||||
owner: "catalog-plugin".into(),
|
||||
id: PublicModelId::new("plugin-alias").unwrap(),
|
||||
provider: ProviderKind::new("openai").unwrap(),
|
||||
target: UpstreamModelId::new("gpt-start").unwrap(),
|
||||
}],
|
||||
contributed,
|
||||
)),
|
||||
models: contributed,
|
||||
}),
|
||||
)));
|
||||
let service = DefaultExecutionService::new(
|
||||
RuntimeSnapshotHandle::new(snapshot),
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
ProviderRegistry::new([
|
||||
Arc::new(NativeCatalogProvider { fail, empty }) as Arc<dyn Provider>
|
||||
])
|
||||
.expect("registry"),
|
||||
ProviderRegistry::new([Arc::new(provider) as Arc<dyn Provider>]).expect("registry"),
|
||||
Arc::new(UnusedAdmissions),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
@@ -3438,7 +3471,7 @@ fn client_catalog_forwards_scope_maps_whole_objects_and_omits_unroutable_models(
|
||||
);
|
||||
let client = service.authenticate("sk_start_test").expect("authenticate");
|
||||
let result = block_on(service.client_model_catalog(&client, "codex", "0.154.0"));
|
||||
if fail {
|
||||
if matches!(provider, CatalogProvider::Unavailable) {
|
||||
assert!(
|
||||
result.is_err(),
|
||||
"native failure cannot use the global synthesized catalog"
|
||||
@@ -3457,13 +3490,13 @@ fn client_catalog_forwards_scope_maps_whole_objects_and_omits_unroutable_models(
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
adapted,
|
||||
if empty {
|
||||
if matches!(provider, CatalogProvider::Empty) {
|
||||
vec!["gpt-start", "adapted"]
|
||||
} else {
|
||||
vec!["adapted"]
|
||||
}
|
||||
);
|
||||
if empty {
|
||||
if matches!(provider, CatalogProvider::Empty) {
|
||||
assert_eq!(models.len(), 2);
|
||||
continue;
|
||||
}
|
||||
@@ -3495,6 +3528,95 @@ fn client_catalog_forwards_scope_maps_whole_objects_and_omits_unroutable_models(
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_catalog_prefers_scoped_profiles_in_provider_order() {
|
||||
use gateway_core::routing::{ModelPresentation, PublicModelDescriptor};
|
||||
|
||||
for provider in [CatalogProvider::Generic, CatalogProvider::Xai] {
|
||||
let kind = ProviderKind::new(provider.name()).unwrap();
|
||||
let capabilities = ModelCapabilities::new([OperationKind::Generate].into(), None);
|
||||
let presentation =
|
||||
|source: &str, id: &str| ModelPresentation::new(Some(format!("{source} {id}")), None);
|
||||
let provider_models = ["shared-model", "z-last"].map(|id| {
|
||||
ProviderModel::new(
|
||||
kind.clone(),
|
||||
UpstreamModelId::new(id).unwrap(),
|
||||
capabilities.clone(),
|
||||
)
|
||||
.with_presentation(presentation("Provider", id))
|
||||
});
|
||||
let contributions = [
|
||||
("extra", "acct_start"),
|
||||
("z-last", "acct_other"),
|
||||
("shared-model", "acct_start"),
|
||||
]
|
||||
.map(|(id, account)| ContributedModel {
|
||||
owner: "catalog-plugin".into(),
|
||||
id: PublicModelId::new(id).unwrap(),
|
||||
provider: kind.clone(),
|
||||
accounts: [ProviderAccountId::new(account).unwrap()].into(),
|
||||
capabilities: capabilities.clone(),
|
||||
presentation: presentation("Plugin", id).with_agent_tools(true, true),
|
||||
});
|
||||
let policy = ClientPolicy::new(
|
||||
ClientApiKeyId::new("key_start_test").unwrap(),
|
||||
PlaintextClientApiKey::new("sk_start_test").unwrap(),
|
||||
account_scope(&kind, "acct_start"),
|
||||
true,
|
||||
RateLimits::unlimited(),
|
||||
);
|
||||
let snapshot = RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).unwrap(),
|
||||
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
|
||||
vec![kind],
|
||||
provider_models.into(),
|
||||
vec![policy],
|
||||
)
|
||||
.unwrap()
|
||||
.with_extensions(Some(ExtensionSetReference::new(
|
||||
ExtensionSetId::new("catalog-precedence".into()).unwrap(),
|
||||
Arc::new(CatalogContributions {
|
||||
aliases: vec![],
|
||||
models: contributions.into(),
|
||||
}),
|
||||
)));
|
||||
let service = DefaultExecutionService::new(
|
||||
RuntimeSnapshotHandle::new(snapshot),
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
ProviderRegistry::new([Arc::new(provider) as Arc<dyn Provider>]).unwrap(),
|
||||
Arc::new(UnusedAdmissions),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
Arc::new(crate::RecordingDiagnostics::default()),
|
||||
);
|
||||
let client = service.authenticate("sk_start_test").unwrap();
|
||||
let models = block_on(service.client_model_catalog(&client, "codex", "0.154.0")).unwrap();
|
||||
let profiles = models
|
||||
.iter()
|
||||
.map(|entry| {
|
||||
let PublicModelDescriptor::Adapted(profile) = entry else {
|
||||
panic!("expected adapted catalog")
|
||||
};
|
||||
(
|
||||
profile.model().as_str(),
|
||||
profile.presentation().display_name().unwrap(),
|
||||
profile.presentation().parallel_tool_calls(),
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
profiles,
|
||||
[
|
||||
("shared-model", "Plugin shared-model", true),
|
||||
("z-last", "Provider z-last", false),
|
||||
("extra", "Plugin extra", true),
|
||||
],
|
||||
"{}",
|
||||
provider.name()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for FailingProvider {
|
||||
fn name(&self) -> &'static str {
|
||||
|
||||
@@ -55,6 +55,7 @@ fn account_id(id: &str) -> ProviderAccountId {
|
||||
|
||||
fn candidate(id: &str) -> AccountCandidate {
|
||||
AccountCandidate {
|
||||
inherited_concurrency: None,
|
||||
account: ProviderAccount::new(
|
||||
account_id(id),
|
||||
ProviderKind::new("openai").unwrap(),
|
||||
|
||||
@@ -730,12 +730,116 @@ impl ProviderCatalogPort for PublishingCatalog {
|
||||
|
||||
#[test]
|
||||
fn compiler_should_reject_revision_changed_during_consistent_read() {
|
||||
let facts = facts(1, 2);
|
||||
let compiler = compiler(Arc::new(TestSnapshotStore::new(Ok(facts))));
|
||||
let store = Arc::new(SequencedSnapshotStore {
|
||||
facts: vec![facts(1, 2)],
|
||||
loads: AtomicUsize::new(0),
|
||||
current_revision: revision(2),
|
||||
});
|
||||
let compiler = compiler(store.clone());
|
||||
|
||||
let error = block_on(compiler.compile()).expect_err("revision drift must fail closed");
|
||||
|
||||
assert_eq!(error, RuntimeSnapshotCompileError::RevisionChanged);
|
||||
assert_eq!(store.loads.load(Ordering::SeqCst), 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compiler_reloads_complete_facts_after_revision_changes_during_read() {
|
||||
let store = Arc::new(SequencedSnapshotStore {
|
||||
facts: vec![facts(1, 2), facts(2, 2)],
|
||||
loads: AtomicUsize::new(0),
|
||||
current_revision: revision(2),
|
||||
});
|
||||
|
||||
let snapshot = block_on(compiler(store.clone()).compile()).expect("latest consistent facts");
|
||||
|
||||
assert_eq!(snapshot.revision(), revision(2));
|
||||
assert_eq!(store.loads.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compiler_reloads_revision_changed_after_extension_preparation() {
|
||||
let drops = Arc::new(AtomicUsize::new(0));
|
||||
let extensions = Arc::new(CountingExtensionPreparation {
|
||||
prepares: AtomicUsize::new(0),
|
||||
drops: Arc::clone(&drops),
|
||||
});
|
||||
let compiler = RuntimeSnapshotCompiler::new(
|
||||
Arc::new(SequencedSnapshotStore {
|
||||
facts: vec![facts(1, 1), facts(2, 2)],
|
||||
loads: AtomicUsize::new(0),
|
||||
current_revision: revision(2),
|
||||
}),
|
||||
Arc::new(TestCatalog::Empty),
|
||||
)
|
||||
.with_extensions(extensions.clone());
|
||||
|
||||
let snapshot = block_on(compiler.compile()).expect("latest prepared snapshot");
|
||||
|
||||
assert_eq!(snapshot.revision(), revision(2));
|
||||
assert_eq!(snapshot.extensions().unwrap().id().as_str(), "prepared-2");
|
||||
assert_eq!(extensions.prepares.load(Ordering::SeqCst), 2);
|
||||
assert_eq!(drops.load(Ordering::SeqCst), 1);
|
||||
drop(snapshot);
|
||||
assert_eq!(drops.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compiler_rechecks_revision_when_extension_preparation_rejects_stale_facts() {
|
||||
struct RacingPreparation {
|
||||
store: Arc<TestSnapshotStore>,
|
||||
changes: usize,
|
||||
calls: AtomicUsize,
|
||||
}
|
||||
|
||||
impl ExtensionPreparationPort for RacingPreparation {
|
||||
fn prepare(
|
||||
&self,
|
||||
_: ConfigRevision,
|
||||
) -> BoxFuture<'_, Result<ExtensionSetReference, ExtensionPreparationError>> {
|
||||
Box::pin(async move {
|
||||
let call = self.calls.fetch_add(1, Ordering::SeqCst);
|
||||
if self.changes == 0 || call < self.changes {
|
||||
if self.changes != 0 {
|
||||
let next = u64::try_from(call).unwrap() + 2;
|
||||
*self.store.facts.lock().unwrap() = Ok(facts(next, next));
|
||||
*self.store.current_revision.lock().unwrap() = Ok(revision(next));
|
||||
}
|
||||
return Err(ExtensionPreparationError);
|
||||
}
|
||||
Ok(ExtensionSetReference::new(
|
||||
ExtensionSetId::new("ready".to_owned()).unwrap(),
|
||||
Arc::new(AliasLease(vec![])),
|
||||
))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
for (changes, expected_calls) in [(0, 1), (1, 2), (4, 4)] {
|
||||
let store = Arc::new(TestSnapshotStore::new(Ok(facts(1, 1))));
|
||||
let extensions = Arc::new(RacingPreparation {
|
||||
store: store.clone(),
|
||||
changes,
|
||||
calls: AtomicUsize::new(0),
|
||||
});
|
||||
let compiler = RuntimeSnapshotCompiler::new(store, Arc::new(TestCatalog::Empty))
|
||||
.with_extensions(extensions.clone());
|
||||
|
||||
let result = block_on(compiler.compile());
|
||||
|
||||
match changes {
|
||||
0 => assert_eq!(
|
||||
result.unwrap_err(),
|
||||
RuntimeSnapshotCompileError::ExtensionsUnavailable
|
||||
),
|
||||
1 => assert_eq!(result.unwrap().revision(), revision(2)),
|
||||
_ => assert_eq!(
|
||||
result.unwrap_err(),
|
||||
RuntimeSnapshotCompileError::RevisionChanged
|
||||
),
|
||||
}
|
||||
assert_eq!(extensions.calls.load(Ordering::SeqCst), expected_calls);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -326,6 +326,39 @@ fn publisher_should_suspend_but_still_notify_after_committed_refresh_failure() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn committed_refresh_reloads_raced_facts_without_losing_revoked_scope() {
|
||||
block_on(async {
|
||||
let store = Arc::new(TestSnapshotStore::new(Ok(scoped_facts(10, true))));
|
||||
let initial = compiler(store.clone())
|
||||
.compile()
|
||||
.await
|
||||
.expect("initial scope");
|
||||
let handle = RuntimeSnapshotHandle::new(initial);
|
||||
let frozen = handle.acquire().expect("initial request");
|
||||
assert!(removed_account_allowed(&frozen));
|
||||
let publisher = RuntimeSnapshotPublisher::new(
|
||||
Arc::new(compiler(store.clone())),
|
||||
handle.clone(),
|
||||
Arc::new(TestSnapshotSubscriptions::default()),
|
||||
);
|
||||
let (release, gate) = oneshot::channel();
|
||||
*store.next_load.lock().expect("load gate lock") = Some(gate);
|
||||
let mut committed = publisher.publish_committed(revision(11));
|
||||
assert!(futures::poll!(committed.as_mut()).is_pending());
|
||||
*store.facts.lock().expect("facts lock") = Ok(scoped_facts(11, false));
|
||||
*store.current_revision.lock().expect("revision lock") = Ok(revision(11));
|
||||
release.send(Ok(facts(10, 11))).expect("release raced read");
|
||||
|
||||
committed.await;
|
||||
|
||||
let current = handle.acquire().expect("latest snapshot remains available");
|
||||
assert_eq!(current.revision(), revision(11));
|
||||
assert!(!removed_account_allowed(¤t));
|
||||
assert!(removed_account_allowed(&frozen));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn publisher_should_contribute_reconciliation_and_subscription_workers() {
|
||||
let store = Arc::new(TestSnapshotStore::new(Ok(facts(1, 1))));
|
||||
|
||||
@@ -125,6 +125,17 @@ impl HttpProxyProbe {
|
||||
}
|
||||
|
||||
impl HttpProxyProbe {
|
||||
async fn timed_exit_ip_at(
|
||||
&self,
|
||||
proxy: Option<&OutboundProxy>,
|
||||
endpoint: &str,
|
||||
) -> (Result<IpAddr, &'static str>, u64) {
|
||||
let started = Instant::now();
|
||||
let result = self.exit_ip_at(proxy, endpoint).await;
|
||||
let latency_ms = u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX);
|
||||
(result, latency_ms)
|
||||
}
|
||||
|
||||
async fn test_connection(&self, proxy: Option<&OutboundProxy>) -> ProxyTestResult {
|
||||
let started = Instant::now();
|
||||
let timeout_limit = Duration::from_secs(15);
|
||||
@@ -169,15 +180,15 @@ impl HttpProxyProbe {
|
||||
} => {
|
||||
let probe_dual = async {
|
||||
tokio::join!(
|
||||
self.exit_ip_at(proxy, ipv4_endpoint),
|
||||
self.exit_ip_at(proxy, ipv6_endpoint),
|
||||
self.timed_exit_ip_at(proxy, ipv4_endpoint),
|
||||
self.timed_exit_ip_at(proxy, ipv6_endpoint),
|
||||
)
|
||||
};
|
||||
let result = tokio::time::timeout(timeout_limit, probe_dual).await;
|
||||
let latency_ms = u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX);
|
||||
|
||||
match result {
|
||||
Ok((res_v4, res_v6)) => {
|
||||
Ok(((res_v4, latency_v4), (res_v6, latency_v6))) => {
|
||||
let exit_ipv4: Option<Ipv4Addr> = match res_v4 {
|
||||
Ok(IpAddr::V4(v4)) => Some(v4),
|
||||
_ => None,
|
||||
@@ -186,6 +197,14 @@ impl HttpProxyProbe {
|
||||
Ok(IpAddr::V6(v6)) => Some(v6),
|
||||
_ => None,
|
||||
};
|
||||
// 延迟对应首选的可用出口,另一地址族失败等待不代表该出口延迟
|
||||
let latency_ms = if exit_ipv4.is_some() {
|
||||
latency_v4
|
||||
} else if exit_ipv6.is_some() {
|
||||
latency_v6
|
||||
} else {
|
||||
latency_ms
|
||||
};
|
||||
|
||||
if exit_ipv4.is_some() && exit_ipv6.is_some() {
|
||||
ProxyTestResult {
|
||||
|
||||
@@ -86,11 +86,7 @@ pub(crate) fn extract_release(
|
||||
.unpack(&target)
|
||||
.map_err(|error| internal("无法解压官方插件文件").with_source(error))?;
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt as _;
|
||||
fs::set_permissions(&target, fs::Permissions::from_mode(0o444))
|
||||
.map_err(|error| internal("无法设置官方插件文件权限").with_source(error))?;
|
||||
}
|
||||
ensure_official_plugin_file_readonly(&target)?;
|
||||
found_official_manifest |= relative == Path::new("plugin-release-manifest.json");
|
||||
continue;
|
||||
}
|
||||
@@ -128,6 +124,24 @@ pub(crate) fn extract_release(
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn ensure_official_plugin_file_readonly(path: &Path) -> Result<(), OperationError> {
|
||||
use std::os::unix::fs::PermissionsExt as _;
|
||||
|
||||
// tar 解包已对打开的文件描述符套用归档模式
|
||||
// Docker Desktop 把 macOS 目录挂进容器后,0444 文件再次 chmod 会返回 EPERM,权限已满足时不必再改
|
||||
let current_mode = fs::symlink_metadata(path)
|
||||
.map_err(|error| internal("无法设置官方插件文件权限").with_source(error))?
|
||||
.permissions()
|
||||
.mode()
|
||||
& 0o777;
|
||||
if current_mode == 0o444 {
|
||||
return Ok(());
|
||||
}
|
||||
fs::set_permissions(path, fs::Permissions::from_mode(0o444))
|
||||
.map_err(|error| internal("无法设置官方插件文件权限").with_source(error))
|
||||
}
|
||||
|
||||
fn unsafe_archive_path(path: &Path) -> bool {
|
||||
path.is_absolute()
|
||||
|| path.components().any(|component| match component {
|
||||
|
||||
@@ -155,6 +155,39 @@ async fn dual_stack_proxy_probe_reports_single_stack_when_only_one_succeeds() {
|
||||
assert!(result.message.contains("仅 IPv4"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dual_stack_latency_measures_the_selected_exit_without_the_other_family_wait() {
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
for working_family in ["v4", "v6"] {
|
||||
let server = MockServer::start().await;
|
||||
for (family, ip) in [("v4", "203.0.113.8"), ("v6", "2001:db8::8")] {
|
||||
let response = if family == working_family {
|
||||
ResponseTemplate::new(200).set_body_json(json!({"ip": ip}))
|
||||
} else {
|
||||
ResponseTemplate::new(503).set_delay(Duration::from_millis(500))
|
||||
};
|
||||
Mock::given(path(format!("/{family}")))
|
||||
.respond_with(response)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
}
|
||||
let started = Instant::now();
|
||||
let result = HttpProxyProbe::new_dual(
|
||||
format!("{}/v4", server.uri()),
|
||||
format!("{}/v6", server.uri()),
|
||||
)
|
||||
.test(None, false)
|
||||
.await;
|
||||
assert!(result.success);
|
||||
assert_eq!(result.exit_ipv4.is_some(), working_family == "v4");
|
||||
assert_eq!(result.exit_ipv6.is_some(), working_family == "v6");
|
||||
assert!(started.elapsed() >= Duration::from_millis(500));
|
||||
assert!(u128::from(result.latency_ms) + 400 < started.elapsed().as_millis());
|
||||
}
|
||||
}
|
||||
|
||||
fn location_response(ip: &str, timezone: &str) -> serde_json::Value {
|
||||
json!({"success": true, "ip": ip, "country_code": "JP", "region": "Tokyo", "city": "Tokyo", "timezone": {"id": timezone}})
|
||||
}
|
||||
|
||||
@@ -419,13 +419,14 @@ enum ChecksumKind {
|
||||
fn release_archive(kind: ArchiveKind) -> Vec<u8> {
|
||||
let encoder = GzEncoder::new(Vec::new(), Compression::default());
|
||||
let mut tar = Builder::new(encoder);
|
||||
append_file(&mut tar, BINARY_NAME, b"new-binary", false);
|
||||
append_file(&mut tar, "web/dist/index.html", b"new-web", false);
|
||||
append_file(&mut tar, BINARY_NAME, b"new-binary", 0o755, false);
|
||||
append_file(&mut tar, "web/dist/index.html", b"new-web", 0o755, false);
|
||||
if !matches!(kind, ArchiveKind::MissingOfficialManifest) {
|
||||
append_file(
|
||||
&mut tar,
|
||||
"plugins/official/plugin-release-manifest.json",
|
||||
b"new-manifest",
|
||||
0o444,
|
||||
false,
|
||||
);
|
||||
}
|
||||
@@ -433,16 +434,18 @@ fn release_archive(kind: ArchiveKind) -> Vec<u8> {
|
||||
&mut tar,
|
||||
"plugins/official/new-plugin.tar.gz",
|
||||
b"new-plugin",
|
||||
0o644,
|
||||
false,
|
||||
);
|
||||
if let ArchiveKind::UnsafePath(name) = kind {
|
||||
append_file(&mut tar, name, b"escape", true);
|
||||
append_file(&mut tar, name, b"escape", 0o755, true);
|
||||
}
|
||||
if matches!(kind, ArchiveKind::NestedOfficialAsset) {
|
||||
append_file(
|
||||
&mut tar,
|
||||
"plugins/official/nested/plugin.tar.gz",
|
||||
b"nested",
|
||||
0o444,
|
||||
false,
|
||||
);
|
||||
}
|
||||
@@ -468,14 +471,16 @@ fn append_symlink(tar: &mut Builder<GzEncoder<Vec<u8>>>, name: &str, target: &st
|
||||
tar.append(&header, std::io::empty()).expect("append link");
|
||||
}
|
||||
|
||||
fn append_file(tar: &mut Builder<GzEncoder<Vec<u8>>>, name: &str, data: &[u8], unsafe_path: bool) {
|
||||
fn append_file(
|
||||
tar: &mut Builder<GzEncoder<Vec<u8>>>,
|
||||
name: &str,
|
||||
data: &[u8],
|
||||
mode: u32,
|
||||
unsafe_path: bool,
|
||||
) {
|
||||
let mut header = Header::new_gnu();
|
||||
header.set_entry_type(EntryType::Regular);
|
||||
header.set_mode(if name.starts_with("plugins/official/") {
|
||||
0o444
|
||||
} else {
|
||||
0o755
|
||||
});
|
||||
header.set_mode(mode);
|
||||
header.set_size(u64::try_from(data.len()).expect("size"));
|
||||
header
|
||||
.set_path(if unsafe_path { "safe" } else { name })
|
||||
|
||||
@@ -471,6 +471,17 @@ async fn rollback_should_restore_binary_web_and_version_state() {
|
||||
& 0o777,
|
||||
0o555
|
||||
);
|
||||
for name in ["plugin-release-manifest.json", "new-plugin.tar.gz"] {
|
||||
assert_eq!(
|
||||
fs::metadata(fixture.official().join(name))
|
||||
.expect("official plugin file")
|
||||
.permissions()
|
||||
.mode()
|
||||
& 0o777,
|
||||
0o444,
|
||||
"{name}"
|
||||
);
|
||||
}
|
||||
}
|
||||
service
|
||||
.rollback(Arc::new(AllowingUpdatePreflight))
|
||||
|
||||
@@ -110,6 +110,10 @@ impl CallResources {
|
||||
stream.close();
|
||||
}
|
||||
state.model_streams.clear();
|
||||
for reader in state.usage_readers.values() {
|
||||
reader.cancel();
|
||||
}
|
||||
state.usage_readers.clear();
|
||||
}
|
||||
|
||||
pub(super) fn start_upstream(
|
||||
@@ -179,6 +183,7 @@ impl CallResources {
|
||||
|
||||
#[derive(Default)]
|
||||
pub(super) struct CallState {
|
||||
pub(super) usage_readers: BTreeMap<String, Arc<super::usage::UsageReader>>,
|
||||
pub(super) closed: bool,
|
||||
pub(super) middleware: Option<Arc<dyn MiddlewareCallback>>,
|
||||
pub(super) streams: BTreeMap<String, Arc<HttpStream>>,
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user