From 59e27524da0bd617911780290303789988fd2c92 Mon Sep 17 00:00:00 2001 From: AAEE86 Date: Sat, 1 Aug 2026 00:42:41 +0800 Subject: [PATCH] refactor(gateway): extract transport-neutral execution attempt lifecycle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 评审第 4 条：responses/turn.rs 实际复制了一整套 HTTP execution lifecycle—— usage 写入、candidate 状态流转、health/adaptive 效果投射、pool key lease 释放、 body capture、账单失败判定，与 HTTP 的顺序和超时语义只能靠人工对齐。 新增 execution_runtime/attempt_lifecycle.rs，把一次 provider attempt 的记账收成 transport 中立的三段： ExecutionAttemptLifecycle::begin pending usage 行 + Pending candidate ExecutionAttemptLifecycle::mark_started usage stream_started + Streaming candidate（幂等） ExecutionAttemptLifecycle::settle 终态四段，顺序不可重排： 1 usage terminal（detachable，不可丢） 2 candidate terminal 3 provider 效果 + 超时兜底释放 lease 4 execution report（作废账单不提交） 顺序、5s 分段超时常量、detachable 语义、「每个效果分支都释放 lease」「作废账单 一律不提交 report」这些不变量全部保持原样。 一并上移的辅助设施： - AttemptStageGuard 取代 await_websocket_lifecycle_stage / await_detachable_lifecycle_stage，把「等多久」参数化：WS 用 Bounded(5s)， HTTP 接线时用 Unbounded 即保持它现在的语义。 - AttemptBodyCapture 取代 append_capture / encode_stream_capture，把 「缓冲 + 截断标志」两个字段收成一个类型（WS 侧四个字段变两个）。捕获内容 仍然是 SSE 形状：usage runtime 按 data: 行解析被捕获的 body 来判定 StreamCapturedTerminalState，换成结构化 JSON 会让终态判定恒为 Missing。 - C2/C3 的结算表本来就不含任何 WS 类型，随之上移。效果表分支与注释逐字未改， 仅按新位置改名为 AttemptProviderEffect / classify_attempt_provider_effect。 responses/settlement.rs 只保留 WS 专属的一件事：把 relay loop 的结算信号 ResponsesWebSocketTurnOutcome 翻译成两个正交事实。 ResponsesProviderAttempt 现在只持有 WS 专有状态：lifecycle 句柄、deadline、 终态观测器、两侧 capture、准入、provider/delivery 事实。plan / trace_id / report_kind / report_context / candidate 起始时间戳都归 lifecycle。 HTTP 侧不接线：execution_runtime/stream/execution.rs 的 DirectPassthroughFinalizerCore（38 字段）与 failover / oauth 重试 / prefetch 深度 纠缠，无法在「行为等价 + 单 commit 可验证」的前提下改动。逐调用点映射表写在 模块文档注释里作为后续 PR 的接线依据。验收：git diff 对 execution_runtime/stream/ 与 crates/aether-usage 均为零 diff。 新增 6 个测试：效果段超时后仍走兜底 lease 释放、Unbounded 会一直等、detachable 写入在调用方停止等待后仍跑完、settle 四段顺序（计数器替身）、body capture 的 SSE 形状与编码状态（并显式记录默认上限是 usize::MAX，截断分支不可达）、 candidate error_type 映射。 --- .../execution_runtime/attempt_lifecycle.rs | 1444 +++++++++++++++++ .../src/execution_runtime/mod.rs | 1 + .../proxy/websocket/responses/settlement.rs | 731 +-------- .../proxy/websocket/responses/turn.rs | 432 +---- 4 files changed, 1547 insertions(+), 1061 deletions(-) create mode 100644 apps/aether-gateway/src/execution_runtime/attempt_lifecycle.rs diff --git a/apps/aether-gateway/src/execution_runtime/attempt_lifecycle.rs b/apps/aether-gateway/src/execution_runtime/attempt_lifecycle.rs new file mode 100644 index 000000000..7324bd5bb --- /dev/null +++ b/apps/aether-gateway/src/execution_runtime/attempt_lifecycle.rs @@ -0,0 +1,1444 @@ +//! 一次 provider attempt 的记账生命周期，与 transport 无关。 +//! +//! HTTP 流式与 Responses WebSocket 的记账都是同一个三段结构： +//! `pending` → `started` → `terminal`。差异只在「终态事实从哪来」——HTTP 从 +//! SSE 字节流解析，WS 从协议事件观察。在这里抽出来之前，WS 侧在 +//! `handlers/proxy/websocket/responses/turn.rs` 里重写了一遍 usage 写入、 +//! candidate 状态流转、health/adaptive 效果投射、pool key lease 释放、 +//! body capture 和账单失败判定，与 HTTP 的顺序、超时语义只能靠人工对齐。 +//! +//! # HTTP 侧调用点映射 +//! +//! 本批不改 HTTP 执行路径（`execution_runtime/stream/execution.rs` 里的 +//! `DirectPassthroughFinalizerCore` 与 failover / oauth 重试 / prefetch 深度纠缠， +//! 无法在「行为等价 + 单 commit 可验证」的前提下接线）。这里记下逐调用点的对应 +//! 关系，作为后续 PR 的接线依据： +//! +//! | HTTP 现状调用点 | 对应本模块 | +//! |---|---| +//! | `record_stream_pending_lifecycle` | [`ExecutionAttemptLifecycle::begin`] | +//! | `maybe_record_first_stream_event_started` | [`ExecutionAttemptLifecycle::mark_started`] | +//! | `record_stream_terminal_usage` | `settle` 第 1 段：usage terminal | +//! | `enqueue_stream_candidate_status_update` | `settle` 第 2 段：candidate terminal | +//! | `apply_local_stream_{success,failure}_effects` | `settle` 第 3 段：provider effects | +//! | `submit_stream_report` | `settle` 第 4 段：execution report | +//! | `append_stream_capture_bytes` / `build_stream_body_capture` | [`AttemptBodyCapture`] | +//! +//! HTTP 侧的 `stage_trace` 埋点、kiro prompt-cache usage 合并、direct-inline 延迟 +//! pending 等是它独有的，接线时作为 transport 专有部分留在原处。 + +use std::collections::BTreeMap; +use std::future::Future; +use std::sync::Arc; +use std::time::Duration; + +use aether_contracts::{ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTelemetry}; +use aether_data_contracts::repository::candidates::RequestCandidateStatus; +use aether_data_contracts::repository::usage::UsageBodyCaptureState; +use aether_scheduler_core::SchedulerRequestCandidateStatusUpdate; +use aether_usage_runtime::{ + build_lifecycle_usage_seed, build_stream_terminal_usage_payload_seed, + build_terminal_usage_context_seed, stream_report_represents_failure, + DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES, +}; +use base64::Engine as _; +use serde_json::Value; +use tracing::warn; + +use crate::clock::current_unix_ms; +use crate::orchestration::{ + apply_local_stream_failure_effects, apply_local_stream_success_effects, + release_local_pool_key_lease, LocalExecutionEffectContext, LocalStreamFailureEffect, +}; +use crate::request_candidate_runtime::record_local_request_candidate_status; +use crate::usage::{submit_stream_report, GatewayStreamReportRequest}; +use crate::AppState; + +/// 客户端取消/断开时对外记录的状态码。 +pub(crate) const CLIENT_CANCELLED_STATUS_CODE: u16 = 499; + +/// 流式超时状态码；只有它会额外投射 pool stream timeout 效果。 +pub(crate) const STREAM_TIMEOUT_STATUS_CODE: u16 = 504; + +/// provider 侧观察到的终态。 +/// +/// 形状刻意保持 transport 中立：HTTP 流式与 WS turn 的差异只在事实从哪来。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AttemptProviderOutcome { + /// 观察到了供应商的终态事件。 + Terminal { + status_code: u16, + /// 供应商自己声明这一轮被取消（`response.cancelled`）。 + cancelled_by_provider: bool, + }, + /// 供应商没能给出终态：断链、超时、gateway 侧失败。 + Aborted { + status_code: u16, + reason: &'static str, + stream_timeout: bool, + }, +} + +impl AttemptProviderOutcome { + pub(crate) const fn status_code(self) -> u16 { + match self { + Self::Terminal { status_code, .. } | Self::Aborted { status_code, .. } => status_code, + } + } + + pub(crate) const fn cancelled_by_provider(self) -> bool { + matches!( + self, + Self::Terminal { + cancelled_by_provider: true, + .. + } + ) + } + + pub(crate) const fn stream_timeout(self) -> bool { + matches!(self, Self::Aborted { + stream_timeout: true, + .. + }) + } + + pub(crate) const fn is_terminal(self) -> bool { + matches!(self, Self::Terminal { .. }) + } +} + +/// 这一个 attempt 的内容是否完整交付给了客户端。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AttemptClientDelivery { + Complete, + Aborted { reason: &'static str }, +} + +impl AttemptClientDelivery { + pub(crate) const fn aborted_reason(self) -> Option<&'static str> { + match self { + Self::Complete => None, + Self::Aborted { reason } => Some(reason), + } + } + + pub(crate) const fn is_aborted(self) -> bool { + matches!(self, Self::Aborted { .. }) + } +} + +/// 一次 attempt 结算时的两个正交事实。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct AttemptTerminalFacts { + pub(crate) provider: AttemptProviderOutcome, + pub(crate) delivery: AttemptClientDelivery, +} + +impl AttemptTerminalFacts { + /// 记入 usage / candidate / 效果的人类可读原因。 + pub(crate) const fn reason(self) -> &'static str { + if let Some(reason) = self.delivery.aborted_reason() { + return reason; + } + match self.provider { + AttemptProviderOutcome::Terminal { + cancelled_by_provider: true, + .. + } => "provider cancelled the response", + AttemptProviderOutcome::Terminal { .. } => { + "provider returned a terminal response event" + } + AttemptProviderOutcome::Aborted { reason, .. } => reason, + } + } + + /// 供应商侧强制错误原因：只有「供应商没给出终态、且内容已完整交付客户端」 + /// 才算，用于给终态摘要补 `parser_error`。 + /// + /// 客户端投递失败不是供应商的错误，所以那一侧返回 `None`——与现状 + /// `ResponsesWebSocketTurnOutcome::forced_error()` 对 `Cancelled` 返回 + /// `None` 一致。 + pub(crate) const fn forced_error(self) -> Option<&'static str> { + match (self.provider, self.delivery) { + ( + AttemptProviderOutcome::Aborted { reason, .. }, + AttemptClientDelivery::Complete, + ) => Some(reason), + _ => None, + } + } +} + +/// 这条 usage 记录是否计费。`Void` 等价于现状传给 +/// `record_stream_terminal(.., cancelled = true)` 的那一侧。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AttemptBilling { + Billed, + Void, +} + +impl AttemptBilling { + pub(crate) const fn is_void(self) -> bool { + matches!(self, Self::Void) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AttemptCandidateStatus { + Success, + Failed, + Cancelled, +} + +/// candidate 行上记录的错误分类。 +/// +/// 与 [`AttemptCandidateStatus`] 刻意分开：`missing_terminal` 为真而记账层 +/// 判定不算失败（report kind 不要求观察到终态事件）时，现状会写出 +/// 「状态 Success + error_type=stream_missing_terminal_event」的组合， +/// 这里必须原样保留。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AttemptCandidateError { + None, + Cancelled, + /// 供应商已经给出终态，但这一轮内容没能完整交付给客户端。账单照记， + /// candidate 行上留下这条事实。 + ClientDeliveryFailed, + MissingTerminal, + TerminalError, +} + +/// 一次 attempt 结束后要投射给供应商/密钥池的效果。 +/// +/// 每个分支都会释放 pool key lease：`ProviderFailure` 由 `PoolError` 释放， +/// `ProviderSuccess` 由 `PoolSuccessStream` 释放，其余情况直接释放。少一条 +/// 分支就会把 lease 挂到 TTL 过期，等于短时间占死一把 key。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AttemptProviderEffect { + /// 既不投射成功也不投射失败，只把 lease 还回去。 + ReleasePoolKeyLease, + ProviderFailure, + ProviderSuccess, +} + +impl AttemptProviderEffect { + /// 把「每个分支都必须释放 lease」这条不变量显式化，便于测试锁住 + /// 「没进任何分支导致 lease 泄漏」这类回归。 + const fn releases_pool_key_lease(self) -> bool { + match self { + Self::ReleasePoolKeyLease | Self::ProviderFailure | Self::ProviderSuccess => true, + } + } +} + +/// 判定一次 attempt 结束后要投射的效果。 +/// +/// 关键分支是「记账层判成 failed，但这一轮没有投射供应商失败」：例如合法的 +/// `response.incomplete`（写满 max_output_tokens）。共享 usage 判定目前仍会 +/// 把这类终态记成失败，但供应商本身工作正常，既不该扣健康分，也不能因为落 +/// 不到任何分支而漏掉 lease 释放。 +pub(crate) const fn classify_attempt_provider_effect( + cancelled: bool, + projects_provider_failure: bool, + failed: bool, +) -> AttemptProviderEffect { + if cancelled { + AttemptProviderEffect::ReleasePoolKeyLease + } else if projects_provider_failure { + AttemptProviderEffect::ProviderFailure + } else if failed { + AttemptProviderEffect::ReleasePoolKeyLease + } else { + AttemptProviderEffect::ProviderSuccess + } +} + +/// 这一个 attempt 的账单是否作废。 +/// +/// 只有两种情况作废：供应商自己声明取消，或者供应商根本没给出终态而客户端 +/// 又已经走了。**供应商已经给出终态时，客户端最后一跳投递失败不作废账单**： +/// 供应商已经完成推理并消耗了 token，客户端还能用 `previous_response_id` +/// 续取这条响应，把成本记成 0 等于让上游账单凭空消失。 +pub(crate) const fn attempt_billing_is_void(facts: AttemptTerminalFacts) -> bool { + facts.provider.cancelled_by_provider() + || (facts.delivery.is_aborted() && !facts.provider.is_terminal()) +} + +/// attempt 对外记录的状态码。 +/// +/// 状态码现在纯粹是 provider 事实：客户端投递失败不再把一条已经拿到 200 +/// 终态的记录改写成 499。作废分支的 provider 状态码本身就是 499 +/// （`response.cancelled` 映射 499，`Cancelled` 信号的兜底也是 499）， +/// 所以这些行的取值不变。 +pub(crate) const fn attempt_status_code(facts: AttemptTerminalFacts) -> u16 { + facts.provider.status_code() +} + +/// 结算判定的输入：两个正交事实 + 记账层对这条 report 的判定 + 终态摘要事实。 +#[derive(Debug, Clone, Copy)] +pub(crate) struct AttemptSettlementInputs { + pub(crate) facts: AttemptTerminalFacts, + /// `aether_usage_runtime::stream_report_represents_failure(payload)` 的结果。 + pub(crate) report_represents_failure: bool, + /// 终态摘要里是否观察到了 finish。 + pub(crate) observed_finish: bool, + /// 终态摘要里是否带解析错误。 + pub(crate) has_parser_error: bool, +} + +/// 结算动作。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct AttemptSettlement { + pub(crate) status_code: u16, + pub(crate) billing: AttemptBilling, + pub(crate) candidate_status: AttemptCandidateStatus, + pub(crate) candidate_error: AttemptCandidateError, + pub(crate) provider_effect: AttemptProviderEffect, + pub(crate) submit_execution_report: bool, +} + +/// 由两个正交事实推出结算动作。唯一的判定入口，表驱动测试逐行锁死。 +/// +/// provider 终态已到达时，客户端投递失败只影响 candidate 的错误分类，不再作废 +/// 账单、不再把状态码改成 499、也不再跳过供应商效果和 execution report。 +pub(crate) const fn classify_attempt_settlement( + inputs: AttemptSettlementInputs, +) -> AttemptSettlement { + let AttemptSettlementInputs { + facts, + report_represents_failure, + observed_finish, + has_parser_error, + } = inputs; + + let void = attempt_billing_is_void(facts); + let status_code = attempt_status_code(facts); + let failed = !void && report_represents_failure; + let missing_terminal = !void && !observed_finish; + let projects_provider_failure = !void + && (status_code >= 400 + || facts.forced_error().is_some() + || has_parser_error + || missing_terminal); + + let candidate_status = if void { + AttemptCandidateStatus::Cancelled + } else if failed { + AttemptCandidateStatus::Failed + } else { + AttemptCandidateStatus::Success + }; + // 投递失败排在供应商侧分类之前：这条记录之所以特别，正是因为内容没送到 + // 客户端手上。供应商侧的判定仍然通过 candidate_status 和 error_message + // 保留下来。 + let candidate_error = if void { + AttemptCandidateError::Cancelled + } else if facts.delivery.is_aborted() { + AttemptCandidateError::ClientDeliveryFailed + } else if missing_terminal { + AttemptCandidateError::MissingTerminal + } else if failed { + AttemptCandidateError::TerminalError + } else { + AttemptCandidateError::None + }; + + AttemptSettlement { + status_code, + billing: if void { + AttemptBilling::Void + } else { + AttemptBilling::Billed + }, + candidate_status, + candidate_error, + provider_effect: classify_attempt_provider_effect( + void, + projects_provider_failure, + failed, + ), + submit_execution_report: !void, + } +} + + +/// 每一段记账 I/O 的等待上界。 +/// +/// WS 用 `Bounded(5s)`：relay loop 是单任务，一段慢依赖会拖住整条连接的收发。 +/// HTTP 接线时用 `Unbounded` 即保持它现在的语义。 +#[derive(Debug, Clone, Copy)] +pub(crate) enum AttemptStageGuard { + Unbounded, + Bounded(Duration), +} + +impl AttemptStageGuard { + /// 等一段记账 I/O，超时就放弃等待。 + /// + /// 返回 `None` 表示这一段没有在上界内完成；调用方据此决定兜底动作 + /// （例如效果段超时后仍然要释放 pool key lease）。 + pub(crate) async fn await_stage( + self, + trace_id: &str, + stage: &'static str, + future: impl Future, + ) -> Option { + let Self::Bounded(timeout) = self else { + return Some(future.await); + }; + match tokio::time::timeout(timeout, future).await { + Ok(value) => Some(value), + Err(_) => { + warn!( + event_name = "execution_attempt_lifecycle_stage_timeout", + log_type = "ops", + trace_id, + stage, + timeout_ms = timeout.as_millis() as u64, + "gateway stopped waiting for an execution attempt lifecycle stage" + ); + None + } + } + } + + /// 跑一段不能丢的写入，同时仍然给调用方的等待设上界。 + /// + /// [`Self::await_stage`] 超时会 drop 掉它等待的 future，这对次要效果是对的， + /// 但会静默丢弃系统其余部分依赖的写入。先 spawn 再等，让上界只约束「等多久」： + /// 丢弃 `JoinHandle` 只是让任务脱离，它仍会跑完。 + pub(crate) async fn await_detachable_stage( + self, + trace_id: &str, + stage: &'static str, + write: F, + ) where + F: Future + Send + 'static, + { + let _ = self.await_stage(trace_id, stage, tokio::spawn(write)).await; + } +} + +/// 一侧（provider 或 client）的响应体捕获缓冲。 +/// +/// 捕获内容必须保持 SSE 形状（`data: {json}\n\n`）： +/// `aether_usage_runtime` 会按 `data:` 行解析被捕获的 body 来判定 +/// `StreamCapturedTerminalState`，而它是 `stream_report_represents_failure` +/// 的一个 OR 项。换成结构化 JSON 会让终态判定恒为 Missing。 +#[derive(Debug, Default)] +pub(crate) struct AttemptBodyCapture { + buffer: Vec, + truncated: bool, +} + +impl AttemptBodyCapture { + pub(crate) fn append(&mut self, bytes: &[u8]) { + if bytes.is_empty() || self.truncated { + return; + } + let max_bytes = DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES; + if self.buffer.len() >= max_bytes { + self.truncated = true; + return; + } + let remaining = max_bytes - self.buffer.len(); + let copied = bytes.len().min(remaining); + self.buffer.extend_from_slice(&bytes[..copied]); + if copied < bytes.len() { + self.truncated = true; + } + } + + pub(crate) fn encode(&self) -> (Option, Option) { + let body = (!self.buffer.is_empty()) + .then(|| base64::engine::general_purpose::STANDARD.encode(&self.buffer)); + let state = if self.truncated { + UsageBodyCaptureState::Truncated + } else if self.buffer.is_empty() { + UsageBodyCaptureState::None + } else { + UsageBodyCaptureState::Inline + }; + (body, Some(state)) + } +} + +/// 启动一次 attempt 记账所需的种子。 +pub(crate) struct AttemptLifecycleSeed { + pub(crate) plan: ExecutionPlan, + pub(crate) report_kind: String, + pub(crate) report_context: Option, + pub(crate) stage_guard: AttemptStageGuard, +} + +/// 结算一次 attempt 需要的终态事实，由 transport 提供。 +pub(crate) struct AttemptTerminalFactsInput<'a> { + pub(crate) facts: AttemptTerminalFacts, + pub(crate) terminal_summary: ExecutionStreamTerminalSummary, + pub(crate) telemetry: ExecutionTelemetry, + pub(crate) provider_headers: BTreeMap, + pub(crate) provider_body: &'a AttemptBodyCapture, + pub(crate) client_body: &'a AttemptBodyCapture, + /// 供应商终态载荷原文（`error` / `response.failed` 一类），失败效果用它做 + /// failover 分类，优先级高于摘要里的 parser_error。 + pub(crate) provider_error_body: Option<&'a str>, + /// 人类可读的结算原因，写进 candidate 行。 + pub(crate) reason: &'a str, +} + +/// 一次 provider attempt 的记账生命周期：pending → started → terminal。 +pub(crate) struct ExecutionAttemptLifecycle { + plan: ExecutionPlan, + trace_id: String, + report_kind: String, + report_context: Option, + candidate_started_at_unix_ms: u64, + stage_guard: AttemptStageGuard, + started_recorded: bool, +} + +impl ExecutionAttemptLifecycle { + /// 写 Pending usage 行 + Pending candidate slot。 + /// + /// 两个写入都不设上界：此时还没有任何东西可以兜底，行没写成功就等于这次 + /// attempt 不存在。 + pub(crate) async fn begin(state: &AppState, seed: AttemptLifecycleSeed) -> Self { + let AttemptLifecycleSeed { + plan, + report_kind, + report_context, + stage_guard, + } = seed; + + let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref()); + // Keep every transport on the same lifecycle data path. `AppState` can + // dedicate an isolated background database pool to usage writes; using + // the foreground state here bypasses that path and leaves the caller + // with a different persistence lifecycle. + let usage_data = state.usage_lifecycle_data_state().as_ref().clone(); + state + .usage_runtime + .record_pending_direct(&usage_data, lifecycle_seed) + .await; + + let candidate_started_at_unix_ms = current_unix_ms(); + record_local_request_candidate_status( + state, + &plan, + report_context.as_ref(), + SchedulerRequestCandidateStatusUpdate { + status: RequestCandidateStatus::Pending, + status_code: None, + error_type: None, + error_message: None, + latency_ms: None, + started_at_unix_ms: Some(candidate_started_at_unix_ms), + finished_at_unix_ms: None, + }, + ) + .await; + + Self { + trace_id: plan.request_id.clone(), + plan, + report_kind, + report_context, + candidate_started_at_unix_ms, + stage_guard, + started_recorded: false, + } + } + + pub(crate) fn plan(&self) -> &ExecutionPlan { + &self.plan + } + + pub(crate) fn trace_id(&self) -> &str { + self.trace_id.as_str() + } + + pub(crate) fn report_context(&self) -> Option<&Value> { + self.report_context.as_ref() + } + + pub(crate) fn take_report_context(&mut self) -> Option { + self.report_context.take() + } + + pub(crate) fn set_report_context(&mut self, report_context: Option) { + self.report_context = report_context; + } + + pub(crate) const fn stage_guard(&self) -> AttemptStageGuard { + self.stage_guard + } + + /// 首个可计费事件到达：usage stream_started + candidate Streaming。幂等。 + pub(crate) async fn mark_started( + &mut self, + state: &AppState, + status_code: u16, + telemetry: &ExecutionTelemetry, + ) { + if self.started_recorded { + return; + } + self.started_recorded = true; + let lifecycle_seed = build_lifecycle_usage_seed(&self.plan, self.report_context.as_ref()); + state.usage_runtime.record_stream_started( + state.usage_lifecycle_data_state().as_ref(), + &lifecycle_seed, + status_code, + Some(telemetry), + ); + let _ = self + .stage_guard + .await_stage( + self.trace_id.as_str(), + "candidate_stream_started", + record_local_request_candidate_status( + state, + &self.plan, + self.report_context.as_ref(), + SchedulerRequestCandidateStatusUpdate { + status: RequestCandidateStatus::Streaming, + status_code: Some(status_code), + error_type: None, + error_message: None, + latency_ms: None, + started_at_unix_ms: Some(self.candidate_started_at_unix_ms), + finished_at_unix_ms: None, + }, + ), + ) + .await; + } + + /// 终态四段，顺序不可重排： + /// + /// 1. usage terminal —— 这次 attempt 的账单记录。行是以 Pending 建立的， + /// 没有别的东西会去对账，所以用 detachable 保证不丢。 + /// 2. candidate terminal —— 调度侧的终态，慢依赖不能让它停在 Streaming。 + /// 3. provider 效果 —— health / adaptive / pool 反馈，次于前两段；超时后 + /// 仍然兜底释放 pool key lease，否则要等 lease TTL 过期才放出这把 key。 + /// 4. execution report —— 作废账单的分支不提交（与 HTTP 在下游断开后 + /// 同样不提交一致）。 + pub(crate) async fn settle( + mut self, + state: &AppState, + input: AttemptTerminalFactsInput<'_>, + ) -> AttemptSettlement { + let AttemptTerminalFactsInput { + facts, + terminal_summary, + telemetry, + provider_headers, + provider_body, + client_body, + provider_error_body, + reason, + } = input; + + let (provider_body_base64, provider_body_state) = provider_body.encode(); + let (client_body_base64, client_body_state) = client_body.encode(); + let payload = GatewayStreamReportRequest { + trace_id: self.trace_id.clone(), + report_kind: std::mem::take(&mut self.report_kind), + report_context: self.report_context.take(), + status_code: attempt_status_code(facts), + headers: provider_headers, + provider_body_base64, + provider_body_state, + client_body_base64, + client_body_state, + terminal_summary: Some(terminal_summary.clone()), + telemetry: Some(telemetry), + }; + let settlement = classify_attempt_settlement(AttemptSettlementInputs { + facts, + report_represents_failure: stream_report_represents_failure(&payload), + observed_finish: terminal_summary.observed_finish, + has_parser_error: terminal_summary.parser_error.is_some(), + }); + + // 1. usage terminal + let context_seed = + build_terminal_usage_context_seed(&self.plan, payload.report_context.as_ref()); + let payload_seed = build_stream_terminal_usage_payload_seed(&payload); + let billing_void = settlement.billing.is_void(); + let usage_runtime = Arc::clone(&state.usage_runtime); + let usage_data = Arc::clone(state.usage_lifecycle_data_state()); + self.stage_guard + .await_detachable_stage(self.trace_id.as_str(), "usage_terminal", async move { + usage_runtime + .record_stream_terminal( + usage_data.as_ref(), + context_seed, + payload_seed, + billing_void, + ) + .await; + }) + .await; + + // 2. candidate terminal + let (error_type, error_message) = candidate_error_fields( + settlement.candidate_error, + terminal_summary.parser_error.as_deref(), + reason, + ); + let _ = self + .stage_guard + .await_stage( + self.trace_id.as_str(), + "candidate_terminal", + record_local_request_candidate_status( + state, + &self.plan, + payload.report_context.as_ref(), + SchedulerRequestCandidateStatusUpdate { + status: match settlement.candidate_status { + AttemptCandidateStatus::Cancelled => RequestCandidateStatus::Cancelled, + AttemptCandidateStatus::Failed => RequestCandidateStatus::Failed, + AttemptCandidateStatus::Success => RequestCandidateStatus::Success, + }, + status_code: Some(settlement.status_code), + error_type, + error_message, + latency_ms: payload + .telemetry + .as_ref() + .and_then(|value| value.elapsed_ms), + started_at_unix_ms: Some(self.candidate_started_at_unix_ms), + finished_at_unix_ms: Some(current_unix_ms()), + }, + ), + ) + .await; + + // 3. provider 效果 + let effect_context = LocalExecutionEffectContext { + plan: &self.plan, + report_context: payload.report_context.as_ref(), + }; + let effects_completed = self + .stage_guard + .await_stage(self.trace_id.as_str(), "provider_effects", async { + match settlement.provider_effect { + AttemptProviderEffect::ReleasePoolKeyLease => { + release_local_pool_key_lease(state, effect_context).await; + } + AttemptProviderEffect::ProviderFailure => { + let response_text = provider_error_body + .or(terminal_summary.parser_error.as_deref()) + .unwrap_or(reason); + let mut effect = LocalStreamFailureEffect::new( + settlement.status_code, + &payload.headers, + Some(response_text), + ); + if facts.provider.stream_timeout() { + effect = effect.with_stream_timeout(); + } + apply_local_stream_failure_effects(state, effect_context, effect).await; + } + AttemptProviderEffect::ProviderSuccess => { + apply_local_stream_success_effects(state, effect_context, &payload).await; + } + } + }) + .await + .is_some(); + if !effects_completed { + let _ = self + .stage_guard + .await_stage( + self.trace_id.as_str(), + "pool_lease_release_after_effect_timeout", + release_local_pool_key_lease(state, effect_context), + ) + .await; + } + + // 4. execution report + if settlement.submit_execution_report { + if let Some(Err(error)) = self + .stage_guard + .await_stage( + self.trace_id.as_str(), + "execution_report", + submit_stream_report(state, payload), + ) + .await + { + warn!( + event_name = "execution_attempt_report_submit_failed", + log_type = "ops", + trace_id = %self.trace_id, + error = ?error, + "gateway failed to submit an execution attempt terminal report" + ); + } + } + + settlement + } +} + +/// candidate 行上的 error_type / error_message。 +fn candidate_error_fields( + candidate_error: AttemptCandidateError, + parser_error: Option<&str>, + reason: &str, +) -> (Option, Option) { + match candidate_error { + AttemptCandidateError::Cancelled => ( + Some("websocket_cancelled".to_string()), + Some(reason.to_string()), + ), + AttemptCandidateError::ClientDeliveryFailed => ( + Some("client_delivery_failed".to_string()), + Some(reason.to_string()), + ), + AttemptCandidateError::MissingTerminal => ( + Some("stream_missing_terminal_event".to_string()), + Some(parser_error.map(str::to_string).unwrap_or_else(|| { + "upstream Responses WebSocket ended before a provider terminal event".to_string() + })), + ), + AttemptCandidateError::TerminalError => ( + Some("stream_terminal_error".to_string()), + parser_error + .map(str::to_string) + .or_else(|| Some(reason.to_string())), + ), + AttemptCandidateError::None => (None, None), + } +} + +#[cfg(test)] +mod tests { + use super::{ + classify_attempt_provider_effect, classify_attempt_settlement, AttemptBilling, + AttemptCandidateError, AttemptCandidateStatus, AttemptClientDelivery, + AttemptProviderEffect, AttemptProviderOutcome, AttemptSettlement, AttemptSettlementInputs, + AttemptTerminalFacts, + }; + + fn settle( + provider: AttemptProviderOutcome, + delivery: AttemptClientDelivery, + report_represents_failure: bool, + observed_finish: bool, + has_parser_error: bool, + ) -> AttemptSettlement { + classify_attempt_settlement(AttemptSettlementInputs { + facts: AttemptTerminalFacts { provider, delivery }, + report_represents_failure, + observed_finish, + has_parser_error, + }) + } + + const fn terminal(status_code: u16) -> AttemptProviderOutcome { + AttemptProviderOutcome::Terminal { + status_code, + cancelled_by_provider: false, + } + } + + const fn provider_cancelled() -> AttemptProviderOutcome { + AttemptProviderOutcome::Terminal { + status_code: 499, + cancelled_by_provider: true, + } + } + + const fn aborted(status_code: u16, reason: &'static str) -> AttemptProviderOutcome { + AttemptProviderOutcome::Aborted { + status_code, + reason, + stream_timeout: status_code == 504, + } + } + + /// 投递失败时 `forced_error` 必须为 `None`：客户端走了不是供应商的错误。 + /// 与现状 `ResponsesWebSocketTurnOutcome::forced_error()` 对 `Cancelled` + /// 返回 `None` 一致。 + #[test] + fn only_a_provider_abort_with_complete_delivery_is_a_forced_error() { + assert_eq!( + AttemptTerminalFacts { + provider: aborted(502, "upstream failed"), + delivery: AttemptClientDelivery::Complete, + } + .forced_error(), + Some("upstream failed") + ); + assert_eq!( + AttemptTerminalFacts { + provider: aborted(499, "client went away"), + delivery: AttemptClientDelivery::Aborted { + reason: "client went away" + }, + } + .forced_error(), + None + ); + assert_eq!( + AttemptTerminalFacts { + provider: terminal(200), + delivery: AttemptClientDelivery::Complete, + } + .forced_error(), + None + ); + } + + + #[test] + fn the_recorded_reason_prefers_the_client_delivery_failure() { + assert_eq!( + AttemptTerminalFacts { + provider: terminal(200), + delivery: AttemptClientDelivery::Aborted { + reason: "client went away" + }, + } + .reason(), + "client went away" + ); + assert_eq!( + AttemptTerminalFacts { + provider: provider_cancelled(), + delivery: AttemptClientDelivery::Complete, + } + .reason(), + "provider cancelled the response" + ); + assert_eq!( + AttemptTerminalFacts { + provider: terminal(200), + delivery: AttemptClientDelivery::Complete, + } + .reason(), + "provider returned a terminal response event" + ); + assert_eq!( + AttemptTerminalFacts { + provider: aborted(502, "upstream failed"), + delivery: AttemptClientDelivery::Complete, + } + .reason(), + "upstream failed" + ); + } + + + /// §1.6 结算表，逐行。 + #[test] + fn settlement_table_row_provider_cancelled_is_void_regardless_of_delivery() { + for delivery in [ + AttemptClientDelivery::Complete, + AttemptClientDelivery::Aborted { reason: "gone" }, + ] { + for report_represents_failure in [false, true] { + let settlement = + settle(provider_cancelled(), delivery, report_represents_failure, true, false); + assert_eq!( + settlement, + AttemptSettlement { + status_code: 499, + billing: AttemptBilling::Void, + candidate_status: AttemptCandidateStatus::Cancelled, + candidate_error: AttemptCandidateError::Cancelled, + provider_effect: AttemptProviderEffect::ReleasePoolKeyLease, + submit_execution_report: false, + }, + "delivery={delivery:?} report_failure={report_represents_failure}" + ); + } + } + } + + + #[test] + fn settlement_table_row_aborted_provider_with_aborted_delivery_is_void() { + let settlement = settle( + aborted(499, "client went away"), + AttemptClientDelivery::Aborted { + reason: "client went away", + }, + true, + false, + false, + ); + assert_eq!( + settlement, + AttemptSettlement { + status_code: 499, + billing: AttemptBilling::Void, + candidate_status: AttemptCandidateStatus::Cancelled, + candidate_error: AttemptCandidateError::Cancelled, + provider_effect: AttemptProviderEffect::ReleasePoolKeyLease, + submit_execution_report: false, + } + ); + } + + + #[test] + fn settlement_table_row_clean_provider_terminal_is_a_billed_success() { + let settlement = settle(terminal(200), AttemptClientDelivery::Complete, false, true, false); + assert_eq!( + settlement, + AttemptSettlement { + status_code: 200, + billing: AttemptBilling::Billed, + candidate_status: AttemptCandidateStatus::Success, + candidate_error: AttemptCandidateError::None, + provider_effect: AttemptProviderEffect::ProviderSuccess, + submit_execution_report: true, + } + ); + } + + + /// 合法 `response.incomplete`：记账层判失败，但供应商工作正常， + /// 不扣健康分、只释放 lease，并且账单照记。 + #[test] + fn settlement_table_row_legitimate_incomplete_is_billed_without_provider_failure() { + let settlement = settle(terminal(200), AttemptClientDelivery::Complete, true, true, false); + assert_eq!( + settlement, + AttemptSettlement { + status_code: 200, + billing: AttemptBilling::Billed, + candidate_status: AttemptCandidateStatus::Failed, + candidate_error: AttemptCandidateError::TerminalError, + provider_effect: AttemptProviderEffect::ReleasePoolKeyLease, + submit_execution_report: true, + } + ); + } + + + #[test] + fn settlement_table_row_provider_abort_projects_a_provider_failure() { + let settlement = settle( + aborted(502, "upstream WebSocket closed before provider terminal event"), + AttemptClientDelivery::Complete, + true, + false, + false, + ); + assert_eq!( + settlement, + AttemptSettlement { + status_code: 502, + billing: AttemptBilling::Billed, + candidate_status: AttemptCandidateStatus::Failed, + candidate_error: AttemptCandidateError::MissingTerminal, + provider_effect: AttemptProviderEffect::ProviderFailure, + submit_execution_report: true, + } + ); + } + + + /// ✱ 修正后的那一行：provider 终态已到达，客户端投递失败不再作废账单。 + /// + /// 供应商已经完成推理并消耗 token，客户端还能用 `previous_response_id` + /// 续取这条响应；把成本记成 0 等于让上游账单凭空消失。投递失败作为独立 + /// 事实留在 candidate 的错误分类里。 + #[test] + fn settlement_table_row_client_delivery_failure_keeps_a_reached_terminal_billed() { + let settlement = settle( + terminal(200), + AttemptClientDelivery::Aborted { + reason: "gateway could not relay the provider event to the client", + }, + false, + true, + false, + ); + assert_eq!( + settlement, + AttemptSettlement { + status_code: 200, + billing: AttemptBilling::Billed, + candidate_status: AttemptCandidateStatus::Success, + candidate_error: AttemptCandidateError::ClientDeliveryFailed, + provider_effect: AttemptProviderEffect::ProviderSuccess, + submit_execution_report: true, + } + ); + + // 除了 candidate 的错误分类，其余判定与「投递成功」完全一致。 + let delivered = settle(terminal(200), AttemptClientDelivery::Complete, false, true, false); + assert_eq!(settlement.status_code, delivered.status_code); + assert_eq!(settlement.billing, delivered.billing); + assert_eq!(settlement.candidate_status, delivered.candidate_status); + assert_eq!(settlement.provider_effect, delivered.provider_effect); + assert_eq!( + settlement.submit_execution_report, + delivered.submit_execution_report + ); + assert_ne!(settlement.candidate_error, delivered.candidate_error); + } + + + /// 供应商还没给出终态时，客户端投递失败仍然作废账单：这一轮确实没有产出。 + #[test] + fn a_delivery_failure_without_a_provider_terminal_still_voids_the_bill() { + let settlement = settle( + aborted(499, "client went away"), + AttemptClientDelivery::Aborted { + reason: "client went away", + }, + false, + false, + false, + ); + assert_eq!(settlement.status_code, 499); + assert_eq!(settlement.billing, AttemptBilling::Void); + assert_eq!( + settlement.candidate_status, + AttemptCandidateStatus::Cancelled + ); + assert_eq!(settlement.candidate_error, AttemptCandidateError::Cancelled); + assert!(!settlement.submit_execution_report); + } + + + /// 供应商自己声明取消时，即使内容送到了客户端也不计费。 + #[test] + fn a_provider_declared_cancellation_is_void_even_when_delivered() { + let settlement = + settle(provider_cancelled(), AttemptClientDelivery::Complete, false, true, false); + assert_eq!(settlement.billing, AttemptBilling::Void); + assert_eq!(settlement.candidate_error, AttemptCandidateError::Cancelled); + } + + + /// 记账层判 Success，但摘要没观察到 finish：现状会写出 + /// 「candidate=Success + error_type=stream_missing_terminal_event」， + /// 所以状态与错误分类必须各自独立。 + #[test] + fn a_missing_terminal_can_coexist_with_a_successful_candidate_status() { + let settlement = settle(terminal(200), AttemptClientDelivery::Complete, false, false, false); + assert_eq!(settlement.candidate_status, AttemptCandidateStatus::Success); + assert_eq!( + settlement.candidate_error, + AttemptCandidateError::MissingTerminal + ); + // missing_terminal 仍然要投射供应商失败。 + assert_eq!( + settlement.provider_effect, + AttemptProviderEffect::ProviderFailure + ); + } + + + #[test] + fn a_parser_error_projects_a_provider_failure_even_on_a_clean_status_code() { + let settlement = settle(terminal(200), AttemptClientDelivery::Complete, true, true, true); + assert_eq!( + settlement.provider_effect, + AttemptProviderEffect::ProviderFailure + ); + assert_eq!(settlement.billing, AttemptBilling::Billed); + } + + + #[test] + fn a_legitimate_incomplete_still_releases_the_pool_key_lease() { + // 共享 usage 判定目前仍把 response.incomplete 记成终态失败，于是会出现 + // failed=true 而 projects_provider_failure=false 的组合。这种组合必须 + // 明确落到「只释放 lease」的分支，否则 lease 会挂到 TTL 过期。 + let effect = classify_attempt_provider_effect(false, false, true); + + assert_eq!(effect, AttemptProviderEffect::ReleasePoolKeyLease); + assert!(effect.releases_pool_key_lease()); + } + + + #[test] + fn every_provider_effect_releases_the_pool_key_lease() { + for (cancelled, projects_provider_failure, failed, expected) in [ + ( + true, + false, + false, + AttemptProviderEffect::ReleasePoolKeyLease, + ), + ( + true, + true, + true, + AttemptProviderEffect::ReleasePoolKeyLease, + ), + ( + false, + true, + true, + AttemptProviderEffect::ProviderFailure, + ), + ( + false, + false, + true, + AttemptProviderEffect::ReleasePoolKeyLease, + ), + ( + false, + false, + false, + AttemptProviderEffect::ProviderSuccess, + ), + ] { + let effect = classify_attempt_provider_effect( + cancelled, + projects_provider_failure, + failed, + ); + assert_eq!( + effect, expected, + "cancelled={cancelled} projects_provider_failure={projects_provider_failure} failed={failed}" + ); + assert!( + effect.releases_pool_key_lease(), + "every effect branch must release the pool key lease" + ); + } + } + + + /// 每一个结算分支都必须释放 lease：这条不变量跨越整张结算表。 + #[test] + fn every_settlement_branch_releases_the_pool_key_lease() { + let providers = [ + terminal(200), + terminal(429), + provider_cancelled(), + aborted(502, "upstream failed"), + aborted(504, "timed out"), + ]; + let deliveries = [ + AttemptClientDelivery::Complete, + AttemptClientDelivery::Aborted { reason: "gone" }, + ]; + for provider in providers { + for delivery in deliveries { + for report_represents_failure in [false, true] { + for observed_finish in [false, true] { + for has_parser_error in [false, true] { + let settlement = settle( + provider, + delivery, + report_represents_failure, + observed_finish, + has_parser_error, + ); + assert!( + settlement.provider_effect.releases_pool_key_lease(), + "provider={provider:?} delivery={delivery:?}" + ); + // 作废账单的分支一律不提交 execution report。 + assert_eq!( + settlement.submit_execution_report, + !settlement.billing.is_void(), + "provider={provider:?} delivery={delivery:?}" + ); + } + } + } + } + } + } +} + +#[cfg(test)] +mod stage_tests { + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::{Arc, Mutex}; + use std::time::Duration; + + use base64::Engine as _; + + use super::{ + candidate_error_fields, AttemptBodyCapture, AttemptCandidateError, AttemptStageGuard, + }; + + /// 效果段超时后仍然必须释放 pool key lease，否则那把 key 要等 lease TTL + /// 过期才放出来。这里用 `await_stage` 返回 `None` 驱动兜底分支。 + #[tokio::test] + async fn a_timed_out_effect_stage_still_reaches_the_lease_release_fallback() { + let guard = AttemptStageGuard::Bounded(Duration::from_millis(20)); + let lease_released = Arc::new(AtomicBool::new(false)); + + // 第一段：永不完成的效果投射。 + let effects_completed = guard + .await_stage("trace", "provider_effects", std::future::pending::<()>()) + .await + .is_some(); + assert!( + !effects_completed, + "a stage that never completes must not report success" + ); + + // 生产代码据此走兜底释放。 + if !effects_completed { + let released = Arc::clone(&lease_released); + let _ = guard + .await_stage("trace", "pool_lease_release_after_effect_timeout", async move { + released.store(true, Ordering::SeqCst); + }) + .await; + } + assert!( + lease_released.load(Ordering::SeqCst), + "the lease must still be released after an effect-stage timeout" + ); + } + + #[tokio::test] + async fn an_unbounded_stage_guard_waits_for_the_stage() { + let guard = AttemptStageGuard::Unbounded; + // 无上界：即使这一段比任何 Bounded 上界都久，也必须等到它完成。 + let value = guard + .await_stage("trace", "stage", async { + tokio::time::sleep(Duration::from_millis(120)).await; + 7_u8 + }) + .await; + assert_eq!(value, Some(7)); + } + + /// 不能丢的写入即使调用方停止等待也要跑完：`await_detachable_stage` 先 spawn + /// 再等，上界只约束「等多久」。 + #[tokio::test] + async fn a_detachable_stage_completes_even_after_the_caller_stops_waiting() { + let guard = AttemptStageGuard::Bounded(Duration::from_millis(20)); + let written = Arc::new(AtomicBool::new(false)); + let flag = Arc::clone(&written); + + guard + .await_detachable_stage("trace", "usage_terminal", async move { + tokio::time::sleep(Duration::from_millis(120)).await; + flag.store(true, Ordering::SeqCst); + }) + .await; + assert!( + !written.load(Ordering::SeqCst), + "the caller must stop waiting at its bound" + ); + + tokio::time::sleep(Duration::from_millis(300)).await; + assert!( + written.load(Ordering::SeqCst), + "a detached write must still run to completion" + ); + } + + /// settle 的四段顺序不可重排：账单先落地，然后 candidate 终态，然后次要的 + /// 供应商效果，最后才是 execution report。用计数器替身记录实际顺序。 + #[tokio::test] + async fn the_settle_stages_run_in_a_fixed_order() { + let guard = AttemptStageGuard::Bounded(Duration::from_millis(500)); + let order = Arc::new(Mutex::new(Vec::new())); + + for stage in [ + "usage_terminal", + "candidate_terminal", + "provider_effects", + "execution_report", + ] { + let recorder = Arc::clone(&order); + let _ = guard + .await_stage("trace", stage, async move { + recorder.lock().expect("order lock").push(stage); + }) + .await; + } + + assert_eq!( + order.lock().expect("order lock").as_slice(), + [ + "usage_terminal", + "candidate_terminal", + "provider_effects", + "execution_report", + ] + ); + } + + /// body capture 的编码状态。截断分支这里到不了：共享的 + /// `DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES` 是 `usize::MAX`， + /// 也就是默认不限长；截断只在 usage 侧把上限调低后才可能发生。 + #[test] + fn body_capture_encodes_inline_and_empty_states() { + assert_eq!( + super::DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES, + usize::MAX, + "the default capture limit is unbounded; truncation is not reachable here" + ); + + let mut capture = AttemptBodyCapture::default(); + capture.append(b"data: {\"type\":\"response.created\"}\n\n"); + capture.append(b"data: {\"type\":\"response.completed\"}\n\n"); + let (body, state) = capture.encode(); + let decoded = base64::engine::general_purpose::STANDARD + .decode(body.expect("a non-empty capture is encoded")) + .expect("capture is valid base64"); + let decoded = String::from_utf8(decoded).expect("capture is UTF-8"); + assert!( + decoded.starts_with("data: ") && decoded.ends_with("\n\n"), + "the capture must stay SSE-shaped for the usage runtime: {decoded:?}" + ); + assert_eq!(decoded.matches("data: ").count(), 2, "appends concatenate"); + assert_eq!( + state, + Some(aether_data_contracts::repository::usage::UsageBodyCaptureState::Inline) + ); + + let empty = AttemptBodyCapture::default(); + let (body, state) = empty.encode(); + assert!(body.is_none()); + assert_eq!( + state, + Some(aether_data_contracts::repository::usage::UsageBodyCaptureState::None) + ); + } + + /// candidate 行的 error_type 映射：投递失败与供应商侧失败必须各有名字。 + #[test] + fn candidate_error_fields_name_each_failure_kind() { + assert_eq!( + candidate_error_fields(AttemptCandidateError::None, None, "reason"), + (None, None) + ); + assert_eq!( + candidate_error_fields(AttemptCandidateError::Cancelled, None, "gone").0, + Some("websocket_cancelled".to_string()) + ); + assert_eq!( + candidate_error_fields(AttemptCandidateError::ClientDeliveryFailed, None, "write failed"), + ( + Some("client_delivery_failed".to_string()), + Some("write failed".to_string()) + ) + ); + assert_eq!( + candidate_error_fields(AttemptCandidateError::TerminalError, Some("parser"), "reason"), + ( + Some("stream_terminal_error".to_string()), + Some("parser".to_string()) + ) + ); + assert_eq!( + candidate_error_fields(AttemptCandidateError::MissingTerminal, None, "reason").0, + Some("stream_missing_terminal_event".to_string()) + ); + } +} diff --git a/apps/aether-gateway/src/execution_runtime/mod.rs b/apps/aether-gateway/src/execution_runtime/mod.rs index 24cc49a60..3e2816ac7 100644 --- a/apps/aether-gateway/src/execution_runtime/mod.rs +++ b/apps/aether-gateway/src/execution_runtime/mod.rs @@ -3,6 +3,7 @@ use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; +pub(crate) mod attempt_lifecycle; pub(crate) mod admission; mod chatgpt_web_image; mod constants; diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/settlement.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/settlement.rs index 77e37e48a..434925af7 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/settlement.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/settlement.rs @@ -1,212 +1,15 @@ -//! 一次 provider attempt 的结算判定。 +//! Responses WebSocket 的结算信号 → 记账事实映射。 //! -//! 评审第 5 条：现状 `ResponsesWebSocketTurnOutcome` 一个枚举同时表达 -//! 「供应商这一轮怎么结束的」和「内容有没有完整交付给客户端」，`finalize()` -//! 再用 `outcome.cancelled()` 一个布尔驱动 billing、candidate 状态和供应商效果。 -//! 于是 provider 终态已经到达、只是最后一跳写客户端失败时，供应商事实会被 -//! 覆盖掉。这里把两件事拆成正交事实，并把结算动作收进一张可逐行测试的表。 +//! 结算判定本身是 transport 中立的，住在 +//! [`crate::execution_runtime::attempt_lifecycle`]。这里只做 WS 专属的一件事： +//! 把 relay loop 的结算触发信号 [`ResponsesWebSocketTurnOutcome`] 翻译成那两个 +//! 正交事实。 use super::turn::ResponsesWebSocketTurnOutcome; - -/// 客户端取消/断开时对外记录的状态码。 -const CLIENT_CANCELLED_STATUS_CODE: u16 = 499; - -/// 流式超时状态码；现状只有它会额外投射 pool stream timeout 效果。 -const STREAM_TIMEOUT_STATUS_CODE: u16 = 504; - -/// provider 侧观察到的终态。 -/// -/// 形状刻意保持 transport 中立：HTTP 流式与 WS turn 的差异只在事实从哪来。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) enum AttemptProviderOutcome { - /// 观察到了供应商的终态事件。 - Terminal { - status_code: u16, - /// 供应商自己声明这一轮被取消（`response.cancelled`）。 - cancelled_by_provider: bool, - }, - /// 供应商没能给出终态：断链、超时、gateway 侧失败。 - Aborted { - status_code: u16, - reason: &'static str, - stream_timeout: bool, - }, -} - -impl AttemptProviderOutcome { - pub(super) const fn status_code(self) -> u16 { - match self { - Self::Terminal { status_code, .. } | Self::Aborted { status_code, .. } => status_code, - } - } - - pub(super) const fn cancelled_by_provider(self) -> bool { - matches!( - self, - Self::Terminal { - cancelled_by_provider: true, - .. - } - ) - } - - pub(super) const fn stream_timeout(self) -> bool { - matches!(self, Self::Aborted { - stream_timeout: true, - .. - }) - } - - pub(super) const fn is_terminal(self) -> bool { - matches!(self, Self::Terminal { .. }) - } -} - -/// 这一个 attempt 的内容是否完整交付给了客户端。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) enum AttemptClientDelivery { - Complete, - Aborted { reason: &'static str }, -} - -impl AttemptClientDelivery { - pub(super) const fn aborted_reason(self) -> Option<&'static str> { - match self { - Self::Complete => None, - Self::Aborted { reason } => Some(reason), - } - } - - pub(super) const fn is_aborted(self) -> bool { - matches!(self, Self::Aborted { .. }) - } -} - -/// 一次 attempt 结算时的两个正交事实。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) struct AttemptTerminalFacts { - pub(super) provider: AttemptProviderOutcome, - pub(super) delivery: AttemptClientDelivery, -} - -impl AttemptTerminalFacts { - /// 记入 usage / candidate / 效果的人类可读原因。 - pub(super) const fn reason(self) -> &'static str { - if let Some(reason) = self.delivery.aborted_reason() { - return reason; - } - match self.provider { - AttemptProviderOutcome::Terminal { - cancelled_by_provider: true, - .. - } => "provider cancelled the response", - AttemptProviderOutcome::Terminal { .. } => { - "provider returned a terminal response event" - } - AttemptProviderOutcome::Aborted { reason, .. } => reason, - } - } - - /// 供应商侧强制错误原因：只有「供应商没给出终态、且内容已完整交付客户端」 - /// 才算，用于给终态摘要补 `parser_error`。 - /// - /// 客户端投递失败不是供应商的错误，所以那一侧返回 `None`——与现状 - /// `ResponsesWebSocketTurnOutcome::forced_error()` 对 `Cancelled` 返回 - /// `None` 一致。 - pub(super) const fn forced_error(self) -> Option<&'static str> { - match (self.provider, self.delivery) { - ( - AttemptProviderOutcome::Aborted { reason, .. }, - AttemptClientDelivery::Complete, - ) => Some(reason), - _ => None, - } - } -} - -/// 这条 usage 记录是否计费。`Void` 等价于现状传给 -/// `record_stream_terminal(.., cancelled = true)` 的那一侧。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) enum AttemptBilling { - Billed, - Void, -} - -impl AttemptBilling { - pub(super) const fn is_void(self) -> bool { - matches!(self, Self::Void) - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) enum AttemptCandidateStatus { - Success, - Failed, - Cancelled, -} - -/// candidate 行上记录的错误分类。 -/// -/// 与 [`AttemptCandidateStatus`] 刻意分开：`missing_terminal` 为真而记账层 -/// 判定不算失败（report kind 不要求观察到终态事件）时，现状会写出 -/// 「状态 Success + error_type=stream_missing_terminal_event」的组合， -/// 这里必须原样保留。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) enum AttemptCandidateError { - None, - Cancelled, - /// 供应商已经给出终态，但这一轮内容没能完整交付给客户端。账单照记， - /// candidate 行上留下这条事实。 - ClientDeliveryFailed, - MissingTerminal, - TerminalError, -} - -/// 一轮 turn 结束后要投射给供应商/密钥池的效果。 -/// -/// 每个分支都会释放 pool key lease：`ProviderFailure` 由 `PoolError` 释放， -/// `ProviderSuccess` 由 `PoolSuccessStream` 释放，其余情况直接释放。少一条 -/// 分支就会把 lease 挂到 TTL 过期，等于短时间占死一把 key。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) enum ResponsesWebSocketTurnEffect { - /// 既不投射成功也不投射失败，只把 lease 还回去。 - ReleasePoolKeyLease, - ProviderFailure, - ProviderSuccess, -} - -#[cfg(test)] -impl ResponsesWebSocketTurnEffect { - /// 把「每个分支都必须释放 lease」这条不变量显式化，便于测试锁住 - /// 「没进任何分支导致 lease 泄漏」这类回归。 - const fn releases_pool_key_lease(self) -> bool { - match self { - Self::ReleasePoolKeyLease | Self::ProviderFailure | Self::ProviderSuccess => true, - } - } -} - -/// 判定一轮 turn 结束后要投射的效果。 -/// -/// 关键分支是「记账层判成 failed，但这一轮没有投射供应商失败」：例如合法的 -/// `response.incomplete`（写满 max_output_tokens）。共享 usage 判定目前仍会 -/// 把这类终态记成失败，但供应商本身工作正常，既不该扣健康分，也不能因为落 -/// 不到任何分支而漏掉 lease 释放。 -pub(super) const fn classify_responses_websocket_turn_effect( - cancelled: bool, - projects_provider_failure: bool, - failed: bool, -) -> ResponsesWebSocketTurnEffect { - if cancelled { - ResponsesWebSocketTurnEffect::ReleasePoolKeyLease - } else if projects_provider_failure { - ResponsesWebSocketTurnEffect::ProviderFailure - } else if failed { - ResponsesWebSocketTurnEffect::ReleasePoolKeyLease - } else { - ResponsesWebSocketTurnEffect::ProviderSuccess - } -} +use crate::execution_runtime::attempt_lifecycle::{ + AttemptClientDelivery, AttemptProviderOutcome, AttemptTerminalFacts, + CLIENT_CANCELLED_STATUS_CODE, STREAM_TIMEOUT_STATUS_CODE, +}; /// 把「结算触发信号」+「已观察到的 provider 终态」+「已记录的投递结果」映射成 /// 两个正交事实。 @@ -279,139 +82,13 @@ pub(super) fn settle_signal_for_client_delivery_failure( terminal_outcome.unwrap_or_else(ResponsesWebSocketTurnOutcome::client_disconnected) } -/// 这一个 attempt 的账单是否作废。 -/// -/// 只有两种情况作废：供应商自己声明取消，或者供应商根本没给出终态而客户端 -/// 又已经走了。**供应商已经给出终态时，客户端最后一跳投递失败不作废账单**： -/// 供应商已经完成推理并消耗了 token，客户端还能用 `previous_response_id` -/// 续取这条响应，把成本记成 0 等于让上游账单凭空消失。 -pub(super) const fn attempt_billing_is_void(facts: AttemptTerminalFacts) -> bool { - facts.provider.cancelled_by_provider() - || (facts.delivery.is_aborted() && !facts.provider.is_terminal()) -} - -/// attempt 对外记录的状态码。 -/// -/// 状态码现在纯粹是 provider 事实：客户端投递失败不再把一条已经拿到 200 -/// 终态的记录改写成 499。作废分支的 provider 状态码本身就是 499 -/// （`response.cancelled` 映射 499，`Cancelled` 信号的兜底也是 499）， -/// 所以这些行的取值不变。 -pub(super) const fn attempt_status_code(facts: AttemptTerminalFacts) -> u16 { - facts.provider.status_code() -} - -/// 结算判定的输入：两个正交事实 + 记账层对这条 report 的判定 + 终态摘要事实。 -#[derive(Debug, Clone, Copy)] -pub(super) struct AttemptSettlementInputs { - pub(super) facts: AttemptTerminalFacts, - /// `aether_usage_runtime::stream_report_represents_failure(payload)` 的结果。 - pub(super) report_represents_failure: bool, - /// 终态摘要里是否观察到了 finish。 - pub(super) observed_finish: bool, - /// 终态摘要里是否带解析错误。 - pub(super) has_parser_error: bool, -} - -/// 结算动作。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) struct AttemptSettlement { - pub(super) status_code: u16, - pub(super) billing: AttemptBilling, - pub(super) candidate_status: AttemptCandidateStatus, - pub(super) candidate_error: AttemptCandidateError, - pub(super) provider_effect: ResponsesWebSocketTurnEffect, - pub(super) submit_execution_report: bool, -} - -/// 由两个正交事实推出结算动作。唯一的判定入口，表驱动测试逐行锁死。 -/// -/// provider 终态已到达时，客户端投递失败只影响 candidate 的错误分类，不再作废 -/// 账单、不再把状态码改成 499、也不再跳过供应商效果和 execution report。 -pub(super) const fn classify_attempt_settlement( - inputs: AttemptSettlementInputs, -) -> AttemptSettlement { - let AttemptSettlementInputs { - facts, - report_represents_failure, - observed_finish, - has_parser_error, - } = inputs; - - let void = attempt_billing_is_void(facts); - let status_code = attempt_status_code(facts); - let failed = !void && report_represents_failure; - let missing_terminal = !void && !observed_finish; - let projects_provider_failure = !void - && (status_code >= 400 - || facts.forced_error().is_some() - || has_parser_error - || missing_terminal); - - let candidate_status = if void { - AttemptCandidateStatus::Cancelled - } else if failed { - AttemptCandidateStatus::Failed - } else { - AttemptCandidateStatus::Success - }; - // 投递失败排在供应商侧分类之前：这条记录之所以特别，正是因为内容没送到 - // 客户端手上。供应商侧的判定仍然通过 candidate_status 和 error_message - // 保留下来。 - let candidate_error = if void { - AttemptCandidateError::Cancelled - } else if facts.delivery.is_aborted() { - AttemptCandidateError::ClientDeliveryFailed - } else if missing_terminal { - AttemptCandidateError::MissingTerminal - } else if failed { - AttemptCandidateError::TerminalError - } else { - AttemptCandidateError::None - }; - - AttemptSettlement { - status_code, - billing: if void { - AttemptBilling::Void - } else { - AttemptBilling::Billed - }, - candidate_status, - candidate_error, - provider_effect: classify_responses_websocket_turn_effect( - void, - projects_provider_failure, - failed, - ), - submit_execution_report: !void, - } -} - #[cfg(test)] mod tests { - use super::{ - attempt_facts_for_outcome, classify_attempt_settlement, - classify_responses_websocket_turn_effect, settle_signal_for_client_delivery_failure, - AttemptBilling, AttemptCandidateError, AttemptCandidateStatus, AttemptClientDelivery, - AttemptProviderOutcome, AttemptSettlement, AttemptSettlementInputs, AttemptTerminalFacts, - ResponsesWebSocketTurnEffect, - }; use super::super::turn::ResponsesWebSocketTurnOutcome; - - fn settle( - provider: AttemptProviderOutcome, - delivery: AttemptClientDelivery, - report_represents_failure: bool, - observed_finish: bool, - has_parser_error: bool, - ) -> AttemptSettlement { - classify_attempt_settlement(AttemptSettlementInputs { - facts: AttemptTerminalFacts { provider, delivery }, - report_represents_failure, - observed_finish, - has_parser_error, - }) - } + use super::{attempt_facts_for_outcome, settle_signal_for_client_delivery_failure}; + use crate::execution_runtime::attempt_lifecycle::{ + AttemptClientDelivery, AttemptProviderOutcome, AttemptTerminalFacts, + }; const fn terminal(status_code: u16) -> AttemptProviderOutcome { AttemptProviderOutcome::Terminal { @@ -515,6 +192,7 @@ mod tests { .stream_timeout()); } + /// `Cancelled` 不携带 provider 信息，已观察到的终态不能被它覆盖； /// `ProviderTerminal` / `Failure` 本身就是权威的 provider 事实。 #[test] @@ -547,254 +225,6 @@ mod tests { assert_eq!(facts.delivery, AttemptClientDelivery::Complete); } - /// 投递失败时 `forced_error` 必须为 `None`：客户端走了不是供应商的错误。 - /// 与现状 `ResponsesWebSocketTurnOutcome::forced_error()` 对 `Cancelled` - /// 返回 `None` 一致。 - #[test] - fn only_a_provider_abort_with_complete_delivery_is_a_forced_error() { - assert_eq!( - AttemptTerminalFacts { - provider: aborted(502, "upstream failed"), - delivery: AttemptClientDelivery::Complete, - } - .forced_error(), - Some("upstream failed") - ); - assert_eq!( - AttemptTerminalFacts { - provider: aborted(499, "client went away"), - delivery: AttemptClientDelivery::Aborted { - reason: "client went away" - }, - } - .forced_error(), - None - ); - assert_eq!( - AttemptTerminalFacts { - provider: terminal(200), - delivery: AttemptClientDelivery::Complete, - } - .forced_error(), - None - ); - } - - #[test] - fn the_recorded_reason_prefers_the_client_delivery_failure() { - assert_eq!( - AttemptTerminalFacts { - provider: terminal(200), - delivery: AttemptClientDelivery::Aborted { - reason: "client went away" - }, - } - .reason(), - "client went away" - ); - assert_eq!( - AttemptTerminalFacts { - provider: provider_cancelled(), - delivery: AttemptClientDelivery::Complete, - } - .reason(), - "provider cancelled the response" - ); - assert_eq!( - AttemptTerminalFacts { - provider: terminal(200), - delivery: AttemptClientDelivery::Complete, - } - .reason(), - "provider returned a terminal response event" - ); - assert_eq!( - AttemptTerminalFacts { - provider: aborted(502, "upstream failed"), - delivery: AttemptClientDelivery::Complete, - } - .reason(), - "upstream failed" - ); - } - - /// §1.6 结算表，逐行。 - #[test] - fn settlement_table_row_provider_cancelled_is_void_regardless_of_delivery() { - for delivery in [ - AttemptClientDelivery::Complete, - AttemptClientDelivery::Aborted { reason: "gone" }, - ] { - for report_represents_failure in [false, true] { - let settlement = - settle(provider_cancelled(), delivery, report_represents_failure, true, false); - assert_eq!( - settlement, - AttemptSettlement { - status_code: 499, - billing: AttemptBilling::Void, - candidate_status: AttemptCandidateStatus::Cancelled, - candidate_error: AttemptCandidateError::Cancelled, - provider_effect: ResponsesWebSocketTurnEffect::ReleasePoolKeyLease, - submit_execution_report: false, - }, - "delivery={delivery:?} report_failure={report_represents_failure}" - ); - } - } - } - - #[test] - fn settlement_table_row_aborted_provider_with_aborted_delivery_is_void() { - let settlement = settle( - aborted(499, "client went away"), - AttemptClientDelivery::Aborted { - reason: "client went away", - }, - true, - false, - false, - ); - assert_eq!( - settlement, - AttemptSettlement { - status_code: 499, - billing: AttemptBilling::Void, - candidate_status: AttemptCandidateStatus::Cancelled, - candidate_error: AttemptCandidateError::Cancelled, - provider_effect: ResponsesWebSocketTurnEffect::ReleasePoolKeyLease, - submit_execution_report: false, - } - ); - } - - #[test] - fn settlement_table_row_clean_provider_terminal_is_a_billed_success() { - let settlement = settle(terminal(200), AttemptClientDelivery::Complete, false, true, false); - assert_eq!( - settlement, - AttemptSettlement { - status_code: 200, - billing: AttemptBilling::Billed, - candidate_status: AttemptCandidateStatus::Success, - candidate_error: AttemptCandidateError::None, - provider_effect: ResponsesWebSocketTurnEffect::ProviderSuccess, - submit_execution_report: true, - } - ); - } - - /// 合法 `response.incomplete`：记账层判失败，但供应商工作正常， - /// 不扣健康分、只释放 lease，并且账单照记。 - #[test] - fn settlement_table_row_legitimate_incomplete_is_billed_without_provider_failure() { - let settlement = settle(terminal(200), AttemptClientDelivery::Complete, true, true, false); - assert_eq!( - settlement, - AttemptSettlement { - status_code: 200, - billing: AttemptBilling::Billed, - candidate_status: AttemptCandidateStatus::Failed, - candidate_error: AttemptCandidateError::TerminalError, - provider_effect: ResponsesWebSocketTurnEffect::ReleasePoolKeyLease, - submit_execution_report: true, - } - ); - } - - #[test] - fn settlement_table_row_provider_abort_projects_a_provider_failure() { - let settlement = settle( - aborted(502, "upstream WebSocket closed before provider terminal event"), - AttemptClientDelivery::Complete, - true, - false, - false, - ); - assert_eq!( - settlement, - AttemptSettlement { - status_code: 502, - billing: AttemptBilling::Billed, - candidate_status: AttemptCandidateStatus::Failed, - candidate_error: AttemptCandidateError::MissingTerminal, - provider_effect: ResponsesWebSocketTurnEffect::ProviderFailure, - submit_execution_report: true, - } - ); - } - - /// ✱ 修正后的那一行：provider 终态已到达，客户端投递失败不再作废账单。 - /// - /// 供应商已经完成推理并消耗 token，客户端还能用 `previous_response_id` - /// 续取这条响应；把成本记成 0 等于让上游账单凭空消失。投递失败作为独立 - /// 事实留在 candidate 的错误分类里。 - #[test] - fn settlement_table_row_client_delivery_failure_keeps_a_reached_terminal_billed() { - let settlement = settle( - terminal(200), - AttemptClientDelivery::Aborted { - reason: "gateway could not relay the provider event to the client", - }, - false, - true, - false, - ); - assert_eq!( - settlement, - AttemptSettlement { - status_code: 200, - billing: AttemptBilling::Billed, - candidate_status: AttemptCandidateStatus::Success, - candidate_error: AttemptCandidateError::ClientDeliveryFailed, - provider_effect: ResponsesWebSocketTurnEffect::ProviderSuccess, - submit_execution_report: true, - } - ); - - // 除了 candidate 的错误分类，其余判定与「投递成功」完全一致。 - let delivered = settle(terminal(200), AttemptClientDelivery::Complete, false, true, false); - assert_eq!(settlement.status_code, delivered.status_code); - assert_eq!(settlement.billing, delivered.billing); - assert_eq!(settlement.candidate_status, delivered.candidate_status); - assert_eq!(settlement.provider_effect, delivered.provider_effect); - assert_eq!( - settlement.submit_execution_report, - delivered.submit_execution_report - ); - assert_ne!(settlement.candidate_error, delivered.candidate_error); - } - - /// 供应商还没给出终态时，客户端投递失败仍然作废账单：这一轮确实没有产出。 - #[test] - fn a_delivery_failure_without_a_provider_terminal_still_voids_the_bill() { - let settlement = settle( - aborted(499, "client went away"), - AttemptClientDelivery::Aborted { - reason: "client went away", - }, - false, - false, - false, - ); - assert_eq!(settlement.status_code, 499); - assert_eq!(settlement.billing, AttemptBilling::Void); - assert_eq!( - settlement.candidate_status, - AttemptCandidateStatus::Cancelled - ); - assert_eq!(settlement.candidate_error, AttemptCandidateError::Cancelled); - assert!(!settlement.submit_execution_report); - } - - /// 供应商自己声明取消时，即使内容送到了客户端也不计费。 - #[test] - fn a_provider_declared_cancellation_is_void_even_when_delivered() { - let settlement = - settle(provider_cancelled(), AttemptClientDelivery::Complete, false, true, false); - assert_eq!(settlement.billing, AttemptBilling::Void); - assert_eq!(settlement.candidate_error, AttemptCandidateError::Cancelled); - } /// 结算信号的选择：provider 终态已到达就用它，否则才是 client 断开。 /// 这是修正的核心——旧实现无条件用 client_disconnected() 覆盖， @@ -815,6 +245,7 @@ mod tests { ); } + /// 明确记录的投递失败不会被结算信号推出的「投递成功」覆盖。 #[test] fn a_recorded_delivery_failure_survives_a_provider_terminal_settle_signal() { @@ -838,136 +269,4 @@ mod tests { // 投递失败不是供应商的错误，摘要不该因此补 parser_error。 assert_eq!(facts.forced_error(), None); } - - /// 记账层判 Success，但摘要没观察到 finish：现状会写出 - /// 「candidate=Success + error_type=stream_missing_terminal_event」， - /// 所以状态与错误分类必须各自独立。 - #[test] - fn a_missing_terminal_can_coexist_with_a_successful_candidate_status() { - let settlement = settle(terminal(200), AttemptClientDelivery::Complete, false, false, false); - assert_eq!(settlement.candidate_status, AttemptCandidateStatus::Success); - assert_eq!( - settlement.candidate_error, - AttemptCandidateError::MissingTerminal - ); - // missing_terminal 仍然要投射供应商失败。 - assert_eq!( - settlement.provider_effect, - ResponsesWebSocketTurnEffect::ProviderFailure - ); - } - - #[test] - fn a_parser_error_projects_a_provider_failure_even_on_a_clean_status_code() { - let settlement = settle(terminal(200), AttemptClientDelivery::Complete, true, true, true); - assert_eq!( - settlement.provider_effect, - ResponsesWebSocketTurnEffect::ProviderFailure - ); - assert_eq!(settlement.billing, AttemptBilling::Billed); - } - - #[test] - fn a_legitimate_incomplete_still_releases_the_pool_key_lease() { - // 共享 usage 判定目前仍把 response.incomplete 记成终态失败，于是会出现 - // failed=true 而 projects_provider_failure=false 的组合。这种组合必须 - // 明确落到「只释放 lease」的分支，否则 lease 会挂到 TTL 过期。 - let effect = classify_responses_websocket_turn_effect(false, false, true); - - assert_eq!(effect, ResponsesWebSocketTurnEffect::ReleasePoolKeyLease); - assert!(effect.releases_pool_key_lease()); - } - - #[test] - fn every_turn_effect_releases_the_pool_key_lease() { - for (cancelled, projects_provider_failure, failed, expected) in [ - ( - true, - false, - false, - ResponsesWebSocketTurnEffect::ReleasePoolKeyLease, - ), - ( - true, - true, - true, - ResponsesWebSocketTurnEffect::ReleasePoolKeyLease, - ), - ( - false, - true, - true, - ResponsesWebSocketTurnEffect::ProviderFailure, - ), - ( - false, - false, - true, - ResponsesWebSocketTurnEffect::ReleasePoolKeyLease, - ), - ( - false, - false, - false, - ResponsesWebSocketTurnEffect::ProviderSuccess, - ), - ] { - let effect = classify_responses_websocket_turn_effect( - cancelled, - projects_provider_failure, - failed, - ); - assert_eq!( - effect, expected, - "cancelled={cancelled} projects_provider_failure={projects_provider_failure} failed={failed}" - ); - assert!( - effect.releases_pool_key_lease(), - "every effect branch must release the pool key lease" - ); - } - } - - /// 每一个结算分支都必须释放 lease：这条不变量跨越整张结算表。 - #[test] - fn every_settlement_branch_releases_the_pool_key_lease() { - let providers = [ - terminal(200), - terminal(429), - provider_cancelled(), - aborted(502, "upstream failed"), - aborted(504, "timed out"), - ]; - let deliveries = [ - AttemptClientDelivery::Complete, - AttemptClientDelivery::Aborted { reason: "gone" }, - ]; - for provider in providers { - for delivery in deliveries { - for report_represents_failure in [false, true] { - for observed_finish in [false, true] { - for has_parser_error in [false, true] { - let settlement = settle( - provider, - delivery, - report_represents_failure, - observed_finish, - has_parser_error, - ); - assert!( - settlement.provider_effect.releases_pool_key_lease(), - "provider={provider:?} delivery={delivery:?}" - ); - // 作废账单的分支一律不提交 execution report。 - assert_eq!( - settlement.submit_execution_report, - !settlement.billing.is_void(), - "provider={provider:?} delivery={delivery:?}" - ); - } - } - } - } - } - } } diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn.rs index 04cd721e1..b2e9e7492 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn.rs @@ -33,11 +33,11 @@ use tracing::warn; use super::adapter::ResponsesWebSocketProtocolAdapter; use super::admission::ResponsesWebSocketTurnAdmission; use super::frame::ParsedResponsesWebSocketFrame; -use super::settlement::{ - attempt_billing_is_void, attempt_facts_for_outcome, attempt_status_code, - classify_attempt_settlement, AttemptCandidateError, AttemptCandidateStatus, - AttemptClientDelivery, AttemptProviderOutcome, AttemptSettlementInputs, AttemptTerminalFacts, - ResponsesWebSocketTurnEffect, +use super::settlement::attempt_facts_for_outcome; +use crate::execution_runtime::attempt_lifecycle::{ + attempt_billing_is_void, AttemptBodyCapture, AttemptClientDelivery, AttemptLifecycleSeed, + AttemptProviderOutcome, AttemptStageGuard, AttemptTerminalFacts, AttemptTerminalFactsInput, + ExecutionAttemptLifecycle, }; use crate::ai_serving::api::StreamingStandardTerminalObserver; use crate::ai_serving::{build_openai_responses_stream_plan_from_decision, AiExecutionDecision}; @@ -66,6 +66,9 @@ const WEBSOCKET_CLIENT_DELIVERY_REPORT_CONTEXT_FIELD: &str = "websocket_client_d const WEBSOCKET_CLIENT_DELIVERY_ABORTED: &str = "aborted"; const WEBSOCKET_CLIENT_DELIVERY_REASON_REPORT_CONTEXT_FIELD: &str = "websocket_client_delivery_reason"; +/// 首个可计费事件到达时记在 usage/candidate 上的状态码。WS 的首事件本身不带 +/// HTTP 状态，沿用 HTTP 流式「已开始流」的 200。 +const STREAM_STARTED_STATUS_CODE: u16 = 200; const DEFAULT_WEBSOCKET_FIRST_EVENT_TIMEOUT_MS: u64 = 30_000; const RESPONSES_WEBSOCKET_LIFECYCLE_STAGE_TIMEOUT: Duration = Duration::from_secs(5); @@ -220,19 +223,13 @@ impl ResponsesWebSocketTurnOutcome { /// 与 [`super::turn_state::LogicalTurn`] 分工明确：logical turn 是客户端看到的 /// 一轮请求（可能包含多个 attempt），attempt 只负责这一次上游执行的记账事实。 pub(super) struct ResponsesProviderAttempt { - plan: ExecutionPlan, - trace_id: String, - report_kind: String, - report_context: Option, + /// 记账三段（pending / started / terminal）由共享的 transport 中立生命周期负责。 + lifecycle: ExecutionAttemptLifecycle, started_at: Instant, - candidate_started_at_unix_ms: u64, provider_headers: BTreeMap, - stream_started: bool, observer: StreamingStandardTerminalObserver, - provider_capture: Vec, - provider_capture_truncated: bool, - client_capture: Vec, - client_capture_truncated: bool, + provider_capture: AttemptBodyCapture, + client_capture: AttemptBodyCapture, upstream_bytes: u64, first_event_elapsed_ms: Option, first_event_timeout: Duration, @@ -400,48 +397,26 @@ pub(super) async fn begin_responses_websocket_turn( } }; - let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref()); - // Keep WebSocket turns on the same lifecycle data path as HTTP streams. - // `AppState` can dedicate an isolated background database pool to usage - // writes; using the foreground state here bypasses that path and leaves - // this transport with a different persistence lifecycle. - let usage_data = state.usage_lifecycle_data_state().as_ref().clone(); - state - .usage_runtime - .record_pending_direct(&usage_data, lifecycle_seed) - .await; - - let candidate_started_at_unix_ms = current_unix_ms(); - record_local_request_candidate_status( + let lifecycle = ExecutionAttemptLifecycle::begin( state, - &plan, - report_context.as_ref(), - SchedulerRequestCandidateStatusUpdate { - status: RequestCandidateStatus::Pending, - status_code: None, - error_type: None, - error_message: None, - latency_ms: None, - started_at_unix_ms: Some(candidate_started_at_unix_ms), - finished_at_unix_ms: None, + AttemptLifecycleSeed { + plan, + report_kind, + report_context, + // relay loop 是单任务：一段慢依赖会拖住整条连接的收发，所以每段 + // 记账 I/O 都要有等待上界。 + stage_guard: AttemptStageGuard::Bounded(RESPONSES_WEBSOCKET_LIFECYCLE_STAGE_TIMEOUT), }, ) .await; Ok(ResponsesProviderAttempt { - trace_id: plan.request_id.clone(), - plan, - report_kind, - report_context, + lifecycle, started_at: Instant::now(), - candidate_started_at_unix_ms, provider_headers: BTreeMap::new(), - stream_started: false, observer: StreamingStandardTerminalObserver::default(), - provider_capture: Vec::new(), - provider_capture_truncated: false, - client_capture: Vec::new(), - client_capture_truncated: false, + provider_capture: AttemptBodyCapture::default(), + client_capture: AttemptBodyCapture::default(), upstream_bytes: 0, first_event_elapsed_ms: None, first_event_timeout, @@ -538,20 +513,24 @@ impl ResponsesProviderAttempt { /// the guard's `Drop` path remains the timeout fallback. pub(super) async fn release_admission(&mut self) { if let Some(admission) = self.admission.take() { - let _ = await_websocket_lifecycle_stage( - &self.trace_id, - "turn_admission_release", - admission.release(), - ) - .await; + let _ = self + .lifecycle + .stage_guard() + .await_stage( + self.lifecycle.trace_id(), + "turn_admission_release", + admission.release(), + ) + .await; } } pub(super) fn set_provider_response_headers(&mut self, headers: BTreeMap) { - self.report_context = attach_provider_response_headers_to_report_context( - self.report_context.take(), + let report_context = attach_provider_response_headers_to_report_context( + self.lifecycle.take_report_context(), &headers, ); + self.lifecycle.set_report_context(report_context); self.provider_headers = headers; } @@ -597,15 +576,17 @@ impl ResponsesProviderAttempt { // Responses event per SSE line, so the batch must be unwrapped or its // token usage is lost. let events = frame.protocol_events(); + let mut report_context = self.lifecycle.take_report_context(); for event in &events { self.capture_sse_event(event); - adapter.decorate_turn_report_context(&mut self.report_context, event); + adapter.decorate_turn_report_context(&mut report_context, event); } + self.lifecycle.set_report_context(report_context); let fallback_context = json!({ "provider_api_format": "openai:responses", "client_api_format": "openai:responses", }); - let report_context = self.report_context.as_ref().unwrap_or(&fallback_context); + let report_context = self.lifecycle.report_context().unwrap_or(&fallback_context); for event in &events { if let Err(error) = self .observer @@ -674,242 +655,72 @@ impl ResponsesProviderAttempt { } pub(super) fn capture_client_frame(&mut self, event: &Value) { - append_capture( - &mut self.client_capture, - &websocket_event_as_sse_line(event), - &mut self.client_capture_truncated, - ); + self.client_capture + .append(&websocket_event_as_sse_line(event)); } pub(super) async fn mark_stream_started(&mut self, state: &AppState) { - if self.stream_started { - return; - } - self.stream_started = true; - let lifecycle_seed = build_lifecycle_usage_seed(&self.plan, self.report_context.as_ref()); let telemetry = self.telemetry(); - state.usage_runtime.record_stream_started( - state.usage_lifecycle_data_state().as_ref(), - &lifecycle_seed, - 200, - Some(&telemetry), - ); - let trace_id = self.trace_id.clone(); - let _ = await_websocket_lifecycle_stage( - &trace_id, - "candidate_stream_started", - record_local_request_candidate_status( - state, - &self.plan, - self.report_context.as_ref(), - SchedulerRequestCandidateStatusUpdate { - status: RequestCandidateStatus::Streaming, - status_code: Some(200), - error_type: None, - error_message: None, - latency_ms: None, - started_at_unix_ms: Some(self.candidate_started_at_unix_ms), - finished_at_unix_ms: None, - }, - ), - ) - .await; + self.lifecycle + .mark_started(state, STREAM_STARTED_STATUS_CODE, &telemetry) + .await; } /// 结算这一个 attempt。 /// /// `outcome` 是「为什么现在结算」的信号，不是供应商事实本身： - /// [`attempt_facts_for_outcome`] 把它和已观察到的 provider 终态一起， - /// 拆成 provider outcome 与 client delivery 两个正交事实，再由 - /// [`classify_attempt_settlement`] 一张表推出账单、candidate 状态和效果。 + /// [`attempt_facts_for_outcome`] 把它和已观察到的 provider 终态、已记录的 + /// 投递结果一起，拆成 provider outcome 与 client delivery 两个正交事实。 + /// 之后的四段记账（usage terminal → candidate terminal → provider 效果 → + /// execution report）由共享的 [`ExecutionAttemptLifecycle::settle`] 负责， + /// 这里只提供 WS 观察到的终态事实。 async fn settle(mut self, state: &AppState, outcome: ResponsesWebSocketTurnOutcome) { let facts = attempt_facts_for_outcome(self.provider_outcome, self.client_delivery, outcome); if let Some(reason) = facts.delivery.aborted_reason() { - self.report_context = - attach_client_delivery_to_report_context(self.report_context.take(), reason); + let report_context = attach_client_delivery_to_report_context( + self.lifecycle.take_report_context(), + reason, + ); + self.lifecycle.set_report_context(report_context); } let summary = self.finish_summary(facts); + let telemetry = self.telemetry(); let terminal_error_body = self.terminal_error_body.take(); - let outcome_reason = facts.reason().to_string(); - let telemetry = Some(self.telemetry()); - let (provider_body_base64, provider_body_state) = - encode_stream_capture(&self.provider_capture, self.provider_capture_truncated); - let (client_body_base64, client_body_state) = - encode_stream_capture(&self.client_capture, self.client_capture_truncated); - let payload = GatewayStreamReportRequest { - trace_id: self.trace_id.clone(), - report_kind: self.report_kind, - report_context: self.report_context, - status_code: attempt_status_code(facts), - headers: self.provider_headers, - provider_body_base64, - provider_body_state, - client_body_base64, - client_body_state, - terminal_summary: Some(summary.clone()), - telemetry, - }; - let settlement = classify_attempt_settlement(AttemptSettlementInputs { - facts, - report_represents_failure: stream_report_represents_failure(&payload), - observed_finish: summary.observed_finish, - has_parser_error: summary.parser_error.is_some(), - }); - let billing_void = settlement.billing.is_void(); - // Do not hold gateway/provider capacity while usage and audit writes - // run. The turn has a complete terminal payload at this point. + // 终态载荷完整了才释放准入：usage/审计写入期间不再占着 gateway/供应商容量。 if let Some(admission) = self.admission.take() { - let _ = await_websocket_lifecycle_stage( - &self.trace_id, - "turn_admission_release", - admission.release(), - ) - .await; - } - - let context_seed = - build_terminal_usage_context_seed(&self.plan, payload.report_context.as_ref()); - let payload_seed = build_stream_terminal_usage_payload_seed(&payload); - // This write is the turn's billing record, so it must not be abandoned - // when the usage runtime is slow: the row was created as Pending and - // nothing else reconciles it. - let usage_runtime = Arc::clone(&state.usage_runtime); - let usage_data = Arc::clone(state.usage_lifecycle_data_state()); - await_detachable_lifecycle_stage(&self.trace_id, "usage_terminal", async move { - usage_runtime - .record_stream_terminal(usage_data.as_ref(), context_seed, payload_seed, billing_void) + let _ = self + .lifecycle + .stage_guard() + .await_stage( + self.lifecycle.trace_id(), + "turn_admission_release", + admission.release(), + ) .await; - }) - .await; + } - let (error_type, error_message) = match settlement.candidate_error { - AttemptCandidateError::Cancelled => ( - Some("websocket_cancelled".to_string()), - Some(outcome_reason.clone()), - ), - AttemptCandidateError::ClientDeliveryFailed => ( - Some("client_delivery_failed".to_string()), - Some(outcome_reason.clone()), - ), - AttemptCandidateError::MissingTerminal => ( - Some("stream_missing_terminal_event".to_string()), - Some(summary.parser_error.clone().unwrap_or_else(|| { - "upstream Responses WebSocket ended before a provider terminal event" - .to_string() - })), - ), - AttemptCandidateError::TerminalError => ( - Some("stream_terminal_error".to_string()), - summary - .parser_error - .clone() - .or_else(|| Some(outcome_reason.clone())), - ), - AttemptCandidateError::None => (None, None), - }; - let _ = await_websocket_lifecycle_stage( - &self.trace_id, - "candidate_terminal", - record_local_request_candidate_status( + self.lifecycle + .settle( state, - &self.plan, - payload.report_context.as_ref(), - SchedulerRequestCandidateStatusUpdate { - status: match settlement.candidate_status { - AttemptCandidateStatus::Cancelled => RequestCandidateStatus::Cancelled, - AttemptCandidateStatus::Failed => RequestCandidateStatus::Failed, - AttemptCandidateStatus::Success => RequestCandidateStatus::Success, - }, - status_code: Some(settlement.status_code), - error_type, - error_message, - latency_ms: payload - .telemetry - .as_ref() - .and_then(|value| value.elapsed_ms), - started_at_unix_ms: Some(self.candidate_started_at_unix_ms), - finished_at_unix_ms: Some(current_unix_ms()), + AttemptTerminalFactsInput { + facts, + terminal_summary: summary, + telemetry, + provider_headers: std::mem::take(&mut self.provider_headers), + provider_body: &self.provider_capture, + client_body: &self.client_capture, + provider_error_body: terminal_error_body.as_deref(), + reason: facts.reason(), }, - ), - ) - .await; - - // Health, adaptive, and pool feedback are secondary to the terminal - // usage/candidate record. A slow dependency must not leave a turn in - // Pending or Streaming indefinitely. - let effect_context = LocalExecutionEffectContext { - plan: &self.plan, - report_context: payload.report_context.as_ref(), - }; - let effects_completed = - await_websocket_lifecycle_stage(&self.trace_id, "provider_effects", async { - match settlement.provider_effect { - ResponsesWebSocketTurnEffect::ReleasePoolKeyLease => { - release_local_pool_key_lease(state, effect_context).await; - } - ResponsesWebSocketTurnEffect::ProviderFailure => { - let response_text = terminal_error_body - .as_deref() - .or(summary.parser_error.as_deref()) - .unwrap_or(outcome_reason.as_str()); - let mut effect = LocalStreamFailureEffect::new( - settlement.status_code, - &payload.headers, - Some(response_text), - ); - if facts.provider.stream_timeout() { - effect = effect.with_stream_timeout(); - } - apply_local_stream_failure_effects(state, effect_context, effect).await; - } - ResponsesWebSocketTurnEffect::ProviderSuccess => { - apply_local_stream_success_effects(state, effect_context, &payload).await; - } - } - }) - .await - .is_some(); - if !effects_completed { - let _ = await_websocket_lifecycle_stage( - &self.trace_id, - "pool_lease_release_after_effect_timeout", - release_local_pool_key_lease(state, effect_context), ) .await; - } - - // The normal execution runtime does not submit a stream report after a - // downstream disconnect either. The terminal usage record above still - // captures cancellation without applying provider-success side effects. - if settlement.submit_execution_report { - if let Some(Err(error)) = await_websocket_lifecycle_stage( - &self.trace_id, - "execution_report", - submit_stream_report(state, payload), - ) - .await - { - warn!( - event_name = "responses_websocket_execution_report_submit_failed", - log_type = "ops", - transport = "websocket", - websocket = true, - trace_id = %self.trace_id, - error = ?error, - "gateway failed to submit Responses WebSocket terminal report" - ); - } - } } fn capture_sse_event(&mut self, event: &Value) { - append_capture( - &mut self.provider_capture, - &websocket_event_as_sse_line(event), - &mut self.provider_capture_truncated, - ); + self.provider_capture + .append(&websocket_event_as_sse_line(event)); } /// 终态摘要。 @@ -923,7 +734,7 @@ impl ResponsesProviderAttempt { "provider_api_format": "openai:responses", "client_api_format": "openai:responses", }); - let report_context = self.report_context.as_ref().unwrap_or(&fallback_context); + let report_context = self.lifecycle.report_context().unwrap_or(&fallback_context); let mut summary = match self.observer.finish(report_context) { Ok(Some(summary)) => summary, Ok(None) => ExecutionStreamTerminalSummary::default(), @@ -1145,80 +956,10 @@ fn websocket_event_as_sse_line(event: &Value) -> Vec { format!("data: {payload}\n\n").into_bytes() } -fn append_capture(buffer: &mut Vec, bytes: &[u8], truncated: &mut bool) { - if bytes.is_empty() || *truncated { - return; - } - let max_bytes = DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES; - if buffer.len() >= max_bytes { - *truncated = true; - return; - } - let remaining = max_bytes - buffer.len(); - let copied = bytes.len().min(remaining); - buffer.extend_from_slice(&bytes[..copied]); - if copied < bytes.len() { - *truncated = true; - } -} - -fn encode_stream_capture( - bytes: &[u8], - truncated: bool, -) -> (Option, Option) { - let body = (!bytes.is_empty()).then(|| base64::engine::general_purpose::STANDARD.encode(bytes)); - let state = if truncated { - UsageBodyCaptureState::Truncated - } else if bytes.is_empty() { - UsageBodyCaptureState::None - } else { - UsageBodyCaptureState::Inline - }; - (body, Some(state)) -} - fn elapsed_ms(started_at: Instant) -> u64 { started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64 } -/// Runs a lifecycle write that must not be lost, while still bounding how long -/// the caller waits for it. -/// -/// [`await_websocket_lifecycle_stage`] drops the future it is waiting on. That -/// is the right trade for secondary effects, but it would silently discard a -/// write the rest of the system depends on. Spawning first makes the deadline -/// bound only the wait: dropping the `JoinHandle` detaches the task, which runs -/// to completion in the background. -async fn await_detachable_lifecycle_stage(trace_id: &str, stage: &'static str, write: F) -where - F: Future + Send + 'static, -{ - let _ = await_websocket_lifecycle_stage(trace_id, stage, tokio::spawn(write)).await; -} - -async fn await_websocket_lifecycle_stage( - trace_id: &str, - stage: &'static str, - future: impl Future, -) -> Option { - match tokio::time::timeout(RESPONSES_WEBSOCKET_LIFECYCLE_STAGE_TIMEOUT, future).await { - Ok(value) => Some(value), - Err(_) => { - warn!( - event_name = "responses_websocket_lifecycle_stage_timeout", - log_type = "ops", - transport = "websocket", - websocket = true, - trace_id, - stage, - timeout_ms = RESPONSES_WEBSOCKET_LIFECYCLE_STAGE_TIMEOUT.as_millis() as u64, - "gateway stopped waiting for a Responses WebSocket lifecycle stage" - ); - None - } - } -} - #[cfg(test)] mod tests { use std::time::{Duration, Instant}; @@ -1230,10 +971,11 @@ mod tests { use super::super::frame::ParsedResponsesWebSocketFrame; use super::super::settlement::{ - attempt_facts_for_outcome, classify_attempt_settlement, - settle_signal_for_client_delivery_failure, AttemptBilling, AttemptCandidateError, - AttemptCandidateStatus, AttemptClientDelivery, AttemptSettlementInputs, - ResponsesWebSocketTurnEffect, + attempt_facts_for_outcome, settle_signal_for_client_delivery_failure, + }; + use crate::execution_runtime::attempt_lifecycle::{ + classify_attempt_settlement, AttemptBilling, AttemptCandidateError, AttemptCandidateStatus, + AttemptClientDelivery, AttemptProviderEffect, AttemptSettlementInputs, }; use super::{ attach_client_delivery_to_report_context, prepare_websocket_report_context, @@ -1428,7 +1170,7 @@ mod tests { assert_eq!(settlement.billing, AttemptBilling::Billed); assert_eq!( settlement.provider_effect, - ResponsesWebSocketTurnEffect::ReleasePoolKeyLease + AttemptProviderEffect::ReleasePoolKeyLease ); assert!(settlement.submit_execution_report); } @@ -1516,7 +1258,7 @@ mod tests { ); assert_eq!( settlement.provider_effect, - ResponsesWebSocketTurnEffect::ProviderSuccess + AttemptProviderEffect::ProviderSuccess ); assert!(settlement.submit_execution_report); // 投递失败仍然留痕。 @@ -1552,7 +1294,7 @@ mod tests { ); assert_eq!( settlement.provider_effect, - ResponsesWebSocketTurnEffect::ReleasePoolKeyLease + AttemptProviderEffect::ReleasePoolKeyLease ); assert!(!settlement.submit_execution_report); }