i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-08 20:20:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #798 from zhefox/codex/antigravity-streaming-finish-reason

Browse Source 
fix: harden stream lifecycle, Antigravity reasoning, and Gemini replay
This commit is contained in:
ZheFox
2026-09-04 18:19:38 +08:00
committed by GitHub
parent 18d78dd6c9 86f7cc0d58
commit 6b1074cfcd
14 changed files with 1402 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/data/state/routing_group_cache.rs
+1 -1
View File
@@ -684,7 +684,7 @@ mod tests {
.unwrap_or_default();
// One Arc is retained by the map and every active request
// owns one through its leader guard or follower state.
if participant_count >= participants + 1 {
if participant_count > participants {
break;
}
tokio::task::yield_now().await;
apps/aether-gateway/src/execution_runtime/attempt_cancellation.rs
+568
View File
@@ -0,0 +1,568 @@
//! Terminal settlement for a local stream attempt whose future is dropped
//! mid-flight.
//!
//! A local stream attempt writes its `usage` row and its `request_candidates`
//! slot as `pending` before it dispatches to the provider, then keeps running
//! inside the downstream request future. When the client disconnects, axum drops
//! that future: the remaining `.await`s never resume and nothing settles either
//! row. They stay `pending` until the maintenance sweeper rewrites them as a 504
//! timeout roughly ten minutes later, which loses the real outcome and the real
//! latency.
//!
//! The stream transport therefore keeps a guard alive across the window between
//! the `pending` write and terminal settlement, and settles the attempt from
//! `Drop` when that window is left by cancellation instead of by a terminal
//! state.
use std::sync::Arc;
use std::time::Instant;
use aether_contracts::ExecutionPlan;
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_scheduler_core::SchedulerRequestCandidateStatusUpdate;
use aether_usage_runtime::{
build_usage_event_data_seed_describing_request_bodies, UsageEvent, UsageEventData,
UsageEventType,
};
use serde_json::{json, Value};
use tracing::warn;
use crate::clock::current_unix_ms as current_request_candidate_unix_ms;
use crate::execution_runtime::attempt_lifecycle::CLIENT_CANCELLED_STATUS_CODE;
use crate::execution_runtime::transport_failure::StreamCandidateWatchdogProgress;
use crate::log_ids::short_request_id;
use crate::request_candidate_runtime::{
record_local_request_candidate_status_snapshot, LocalRequestCandidateStatusSnapshot,
};
use crate::request_diagnostics::{
attach_request_diagnostics_to_report_context, current_request_diagnostics, RequestDiagnostics,
};
use crate::AppState;
fn elapsed_ms_since(started_at: Instant) -> u64 {
started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64
}
/// The facts the guard needs to settle the attempt it is watching.
///
/// This is held for the whole attempt, so it is deliberately free of request
/// bodies. A request body can be megabytes, and holding one per in-flight
/// attempt would cost far more than the row it settles: the usage seed is built
/// with [`build_usage_event_data_seed_describing_request_bodies`], which derives
/// every capture state, body reference and derived request fact from the real
/// plan and report context but keeps neither body. The terminal write it
/// produces therefore preserves the capture the `pending` write recorded instead
/// of clearing it.
struct ArmedAttempt {
request_id: String,
candidate_id: Option<String>,
candidate: Option<LocalRequestCandidateStatusSnapshot>,
// Boxed: the guard lives inside the stream request future, which is already
// very large, and `UsageEventData` is a wide struct.
usage_seed: Option<Box<UsageEventData>>,
request_diagnostics: Option<Arc<RequestDiagnostics>>,
candidate_started_unix_ms: u64,
candidate_started_at: Instant,
}
/// Settles an attempt as cancelled when its future is dropped before the
/// transport reaches a terminal state.
///
/// The guard is created disarmed and stays inert until [`Self::arm`] is called,
/// so an attempt that is dropped before it owns any `pending` row does not grow
/// a settlement row it never had. The owner disarms it as soon as the attempt
/// completes, whichever way it completes: from that point terminal settlement
/// belongs to the transport (for streams, to the stream finalizer that lives in
/// the response body), and the guard must not write a second terminal state.
///
/// A stream candidate also runs under a first-byte watchdog that drops the
/// attempt future when it gives up. That drop is not a client disconnect and the
/// watchdog settles the attempt itself, so the guard stands down for it.
pub(crate) struct AttemptCancellationGuard {
state: AppState,
error_type: &'static str,
error_message: &'static str,
watchdog: Option<Arc<StreamCandidateWatchdogProgress>>,
armed: Option<ArmedAttempt>,
}
impl AttemptCancellationGuard {
pub(crate) fn disarmed(
state: &AppState,
error_type: &'static str,
error_message: &'static str,
) -> Self {
Self {
state: state.clone(),
error_type,
error_message,
watchdog: StreamCandidateWatchdogProgress::current(),
armed: None,
}
}
/// Takes ownership of the attempt's settlement until it is disarmed.
pub(crate) fn arm(
&mut self,
plan: &ExecutionPlan,
report_context: Option<&Value>,
candidate: Option<&LocalRequestCandidateStatusSnapshot>,
candidate_started_unix_ms: u64,
candidate_started_at: Instant,
) {
let usage_seed = self.state.usage_runtime.is_enabled().then(|| {
Box::new(build_usage_event_data_seed_describing_request_bodies(
plan,
report_context,
))
});
self.armed = Some(ArmedAttempt {
request_id: plan.request_id.clone(),
candidate_id: plan.candidate_id.clone(),
candidate: candidate.cloned(),
usage_seed,
request_diagnostics: current_request_diagnostics(),
candidate_started_unix_ms,
candidate_started_at,
});
}
pub(crate) fn disarm(&mut self) {
self.armed = None;
}
}
/// Writes the candidate terminal row and the terminal usage event for an attempt
/// that never reached its own terminal path.
async fn settle_cancelled_attempt(
state: AppState,
armed: ArmedAttempt,
error_type: &'static str,
error_message: &'static str,
) {
let ArmedAttempt {
request_id,
candidate_id: _,
candidate,
usage_seed,
request_diagnostics,
candidate_started_unix_ms,
candidate_started_at,
} = armed;
let terminal_unix_ms = current_request_candidate_unix_ms();
let latency_ms = elapsed_ms_since(candidate_started_at);
if let Some(candidate) = candidate.as_ref() {
record_local_request_candidate_status_snapshot(
&state,
candidate,
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Cancelled,
status_code: Some(CLIENT_CANCELLED_STATUS_CODE),
error_type: Some(error_type.to_string()),
error_message: Some(error_message.to_string()),
latency_ms: Some(latency_ms),
started_at_unix_ms: Some(candidate_started_unix_ms),
finished_at_unix_ms: Some(terminal_unix_ms),
},
)
.await;
}
let Some(usage_data) = usage_seed else {
return;
};
let mut usage_data = *usage_data;
// The seed was built when the attempt was armed, so it predates the
// diagnostics it should carry. Attaching them to the seed's metadata is the
// same write the report context would have carried into a seed built here:
// both land the same keys in the same object.
usage_data.request_metadata = attach_request_diagnostics_to_report_context(
usage_data.request_metadata.take(),
request_diagnostics.as_ref(),
);
usage_data.status_code = Some(CLIENT_CANCELLED_STATUS_CODE);
usage_data.error_message = Some(error_message.to_string());
usage_data.error_category = Some("cancelled".to_string());
usage_data.response_time_ms = Some(latency_ms);
let error_body = json!({
"error": {
"type": error_type,
"message": error_message,
"code": CLIENT_CANCELLED_STATUS_CODE
}
});
usage_data.response_headers = Some(json!({"content-type": "application/json"}));
usage_data.response_body = Some(error_body.clone());
usage_data.client_response_headers = Some(json!({"content-type": "application/json"}));
usage_data.client_response_body = Some(error_body);
state
.usage_runtime
.record_terminal_event_direct(
state.usage_lifecycle_data_state().as_ref(),
UsageEvent::new(UsageEventType::Cancelled, request_id, usage_data),
)
.await;
}
impl Drop for AttemptCancellationGuard {
fn drop(&mut self) {
let Some(armed) = self.armed.take() else {
return;
};
if self
.watchdog
.as_ref()
.is_some_and(|watchdog| watchdog.abandoned())
{
return;
}
let state = self.state.clone();
let error_type = self.error_type;
let error_message = self.error_message;
// `Drop` cannot await, and the settlement writes touch the database.
// Hand them to the runtime so they survive the dropped request future.
let Ok(handle) = tokio::runtime::Handle::try_current() else {
warn!(
event_name = "local_attempt_cancellation_guard_no_runtime",
log_type = "ops",
request_id = %short_request_id(armed.request_id.as_str()),
candidate_id = ?armed.candidate_id,
error_type,
"gateway could not settle dropped local attempt because no Tokio runtime is available"
);
return;
};
handle.spawn(async move {
settle_cancelled_attempt(state, armed, error_type, error_message).await;
});
}
}
#[cfg(test)]
mod tests {
use super::*;
use aether_contracts::RequestBody;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::usage::InMemoryUsageReadRepository;
use aether_data_contracts::repository::candidates::RequestCandidateReadRepository;
use aether_data_contracts::repository::usage::{
StoredRequestUsageAudit, UsageBodyCaptureState, UsageReadRepository, UsageWriteRepository,
};
use aether_usage_runtime::{
build_lifecycle_usage_seed, build_pending_usage_record, UsageRuntimeConfig,
};
use std::collections::BTreeMap;
use std::time::Duration;
use crate::request_candidate_runtime::{
ensure_execution_request_candidate_slot, snapshot_local_request_candidate_status,
};
const TEST_ERROR_TYPE: &str = "local_stream_attempt_cancelled";
const TEST_ERROR_MESSAGE: &str =
"Local stream attempt was dropped before terminal finalization.";
fn test_stream_plan(request_id: &str) -> ExecutionPlan {
ExecutionPlan {
request_id: request_id.to_string(),
candidate_id: None,
provider_name: Some("Anthropic".to_string()),
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "key-1".to_string(),
method: "POST".to_string(),
url: "https://example.test/v1/messages".to_string(),
headers: BTreeMap::new(),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(json!({"stream": true, "service_tier": "priority"})),
stream: true,
client_api_format: "claude:messages".to_string(),
provider_api_format: "claude:messages".to_string(),
model_name: Some("claude-sonnet-4-5".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
}
}
fn test_report_context() -> Option<Value> {
Some(json!({
"candidate_index": 0,
"retry_index": 0,
"user_id": "user-cancel",
"api_key_id": "api-key-cancel",
"client_api_format": "claude:messages",
"provider_api_format": "claude:messages",
"request_path": "/v1/messages",
"request_path_and_query": "/v1/messages?beta=true",
"upstream_url": "https://example.test/v1/messages",
"mapped_model": "claude-sonnet-4-5",
"original_request_body": {"stream": true, "messages": []},
}))
}
fn test_state(
usage_repository: &Arc<InMemoryUsageReadRepository>,
request_candidate_repository: &Arc<InMemoryRequestCandidateRepository>,
) -> AppState {
AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
Arc::clone(request_candidate_repository),
Arc::clone(usage_repository),
),
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
})
}
/// Writes the `pending` rows the same way a stream attempt does before it
/// dispatches to the provider, and returns the candidate slot snapshot the
/// attempt owns from that point on.
async fn record_pending_attempt(
state: &AppState,
plan: &mut ExecutionPlan,
report_context: &mut Option<Value>,
candidate_started_unix_ms: u64,
) -> LocalRequestCandidateStatusSnapshot {
ensure_execution_request_candidate_slot(state, plan, report_context).await;
state.usage_runtime.record_pending(
state.usage_lifecycle_data_state().as_ref(),
build_lifecycle_usage_seed(plan, report_context.as_ref()),
);
let snapshot = snapshot_local_request_candidate_status(plan, report_context.as_ref())
.expect("attempt should own a candidate slot");
record_local_request_candidate_status_snapshot(
state,
&snapshot,
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Pending,
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
started_at_unix_ms: Some(candidate_started_unix_ms),
finished_at_unix_ms: None,
},
)
.await;
snapshot
}
async fn wait_for_usage_status(
usage_repository: &InMemoryUsageReadRepository,
request_id: &str,
status: &str,
) -> Option<StoredRequestUsageAudit> {
for _ in 0..50 {
if let Some(usage) = usage_repository
.find_by_request_id(request_id)
.await
.expect("usage should read")
{
if usage.status == status {
return Some(usage);
}
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
None
}
#[tokio::test]
async fn armed_guard_settles_a_dropped_attempt_as_cancelled() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = test_state(&usage_repository, &request_candidate_repository);
let mut plan = test_stream_plan("stream-cancel-guard-request");
let mut report_context = test_report_context();
let candidate_started_unix_ms = current_request_candidate_unix_ms();
let snapshot = record_pending_attempt(
&state,
&mut plan,
&mut report_context,
candidate_started_unix_ms,
)
.await;
{
let mut guard =
AttemptCancellationGuard::disarmed(&state, TEST_ERROR_TYPE, TEST_ERROR_MESSAGE);
guard.arm(
&plan,
report_context.as_ref(),
Some(&snapshot),
candidate_started_unix_ms,
Instant::now(),
);
}
let usage = wait_for_usage_status(
usage_repository.as_ref(),
"stream-cancel-guard-request",
"cancelled",
)
.await
.expect("cancelled usage should be recorded");
assert_eq!(usage.billing_status, "void");
assert_eq!(usage.status_code, Some(CLIENT_CANCELLED_STATUS_CODE));
assert_eq!(usage.error_category.as_deref(), Some("cancelled"));
assert!(usage.response_time_ms.is_some());
let candidates = request_candidate_repository
.list_by_request_id("stream-cancel-guard-request")
.await
.expect("candidates should read");
let candidate = candidates.first().expect("candidate row should exist");
assert_eq!(candidate.status, RequestCandidateStatus::Cancelled);
assert_eq!(candidate.status_code, Some(CLIENT_CANCELLED_STATUS_CODE));
assert_eq!(candidate.error_type.as_deref(), Some(TEST_ERROR_TYPE));
assert!(candidate.finished_at_unix_ms.is_some());
}
/// The guard holds no request body, so its settlement write must describe the
/// capture rather than deny it: a typed `none` capture state would clear the
/// stored request body instead of leaving it alone.
#[tokio::test]
async fn settling_a_dropped_attempt_leaves_the_captured_request_body_alone() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = test_state(&usage_repository, &request_candidate_repository);
let mut plan = test_stream_plan("stream-cancel-guard-capture");
let mut report_context = test_report_context();
let candidate_started_unix_ms = current_request_candidate_unix_ms();
let snapshot = record_pending_attempt(
&state,
&mut plan,
&mut report_context,
candidate_started_unix_ms,
)
.await;
// Stand in for a write that already captured this request's body.
let captured_body = json!({"stream": true, "service_tier": "priority"});
let mut capture = build_pending_usage_record(
&plan,
report_context.as_ref(),
current_request_candidate_unix_ms() / 1_000,
)
.expect("pending usage record should build");
capture.provider_request_body = Some(captured_body.clone());
capture.provider_request_body_state = Some(UsageBodyCaptureState::Inline);
usage_repository
.upsert(capture)
.await
.expect("captured request body should upsert");
{
let mut guard =
AttemptCancellationGuard::disarmed(&state, TEST_ERROR_TYPE, TEST_ERROR_MESSAGE);
guard.arm(
&plan,
report_context.as_ref(),
Some(&snapshot),
candidate_started_unix_ms,
Instant::now(),
);
}
let usage = wait_for_usage_status(
usage_repository.as_ref(),
"stream-cancel-guard-capture",
"cancelled",
)
.await
.expect("cancelled usage should be recorded");
assert_eq!(usage.provider_request_body, Some(captured_body));
assert_ne!(
usage.provider_request_body_state,
Some(UsageBodyCaptureState::None)
);
}
#[tokio::test]
async fn guard_stands_down_when_the_watchdog_abandons_the_attempt() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = test_state(&usage_repository, &request_candidate_repository);
let mut plan = test_stream_plan("stream-watchdog-guard-request");
let mut report_context = test_report_context();
let candidate_started_unix_ms = current_request_candidate_unix_ms();
let snapshot = record_pending_attempt(
&state,
&mut plan,
&mut report_context,
candidate_started_unix_ms,
)
.await;
let watchdog = StreamCandidateWatchdogProgress::shared();
Arc::clone(&watchdog)
.scope(async {
let mut guard =
AttemptCancellationGuard::disarmed(&state, TEST_ERROR_TYPE, TEST_ERROR_MESSAGE);
guard.arm(
&plan,
report_context.as_ref(),
Some(&snapshot),
candidate_started_unix_ms,
Instant::now(),
);
// The watchdog gives up and takes over settlement before the
// abandoned attempt is dropped.
watchdog.mark_abandoned();
})
.await;
assert!(wait_for_usage_status(
usage_repository.as_ref(),
"stream-watchdog-guard-request",
"cancelled",
)
.await
.is_none());
}
#[tokio::test]
async fn disarmed_guard_leaves_the_attempt_pending() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = test_state(&usage_repository, &request_candidate_repository);
let mut plan = test_stream_plan("stream-disarmed-guard-request");
let mut report_context = test_report_context();
let candidate_started_unix_ms = current_request_candidate_unix_ms();
let snapshot = record_pending_attempt(
&state,
&mut plan,
&mut report_context,
candidate_started_unix_ms,
)
.await;
{
let mut guard =
AttemptCancellationGuard::disarmed(&state, TEST_ERROR_TYPE, TEST_ERROR_MESSAGE);
guard.arm(
&plan,
report_context.as_ref(),
Some(&snapshot),
candidate_started_unix_ms,
Instant::now(),
);
guard.disarm();
}
assert!(wait_for_usage_status(
usage_repository.as_ref(),
"stream-disarmed-guard-request",
"cancelled",
)
.await
.is_none());
}
}
apps/aether-gateway/src/execution_runtime/mod.rs
+1
View File
@@ -4,6 +4,7 @@ use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
pub(crate) mod admission;
pub(crate) mod attempt_cancellation;
pub(crate) mod attempt_lifecycle;
mod chatgpt_web_image;
mod constants;
apps/aether-gateway/src/execution_runtime/stream/commit_policy.rs
+6 -8
View File
@@ -455,7 +455,10 @@ fn gemini_part_is_client_semantic(part: &Value) -> bool {
return true;
}
if part.get("thought").and_then(Value::as_bool) == Some(true) {
return false;
return part
.get("text")
.and_then(Value::as_str)
.is_some_and(|text| !text.is_empty());
}
if part.keys().all(|key| key == "thoughtSignature") {
return false;
@@ -584,17 +587,12 @@ mod tests {
}
#[test]
fn gemini_gate_waits_through_thought_and_commits_on_text() {
fn gemini_gate_commits_on_first_nonempty_thought() {
let mut gate = StreamCommitGate::new(gemini_policy());
let thought = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thought\":true,\"text\":\"checking\"}]}}]}}\n\n";
let text = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"answer\"}]}}]}}\n\n";
assert_eq!(
gate.observe_provider_bytes(thought),
StreamPrecommitObservation::Pending
);
assert_eq!(
gate.observe_provider_bytes(text),
StreamPrecommitObservation::Commit
);
assert_eq!(gate.state(), StreamCommitState::Committed);
@@ -615,7 +613,7 @@ mod tests {
#[test]
fn gemini_gate_rejects_malformed_function_call_before_commit() {
let mut gate = StreamCommitGate::new(gemini_policy());
let thought = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thought\":true,\"text\":\"calling\"}]}}]}}\n\n";
let thought = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thoughtSignature\":\"signature\",\"text\":\"\"}]}}]}}\n\n";
let malformed = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thoughtSignature\":\"signature\",\"text\":\"\"}]},\"finishReason\":\"MALFORMED_FUNCTION_CALL\",\"finishMessage\":\"Malformed function call: Function call is empty - no input to parse.\"}]}}\n\n";
assert_eq!(
apps/aether-gateway/src/execution_runtime/stream/execution.rs
+83 -17
View File
@@ -78,6 +78,7 @@ use crate::api::response::{
use crate::clock::current_unix_ms as current_request_candidate_unix_ms;
use crate::constants::{CONTROL_CANDIDATE_ID_HEADER, CONTROL_REQUEST_ID_HEADER};
use crate::control::GatewayControlDecision;
use crate::execution_runtime::attempt_cancellation::AttemptCancellationGuard;
use crate::execution_runtime::build_direct_execution_frame_stream;
use crate::execution_runtime::chatgpt_web_image::maybe_execute_chatgpt_web_image_stream;
use crate::execution_runtime::grok::maybe_execute_grok_stream;
@@ -149,6 +150,11 @@ use crate::{
AppState, GatewayError, GEMINI_FILES_DOWNLOAD_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
};
/// Settlement labels for a stream attempt whose future is dropped before the
/// transport reaches a terminal state.
const STREAM_ATTEMPT_CANCELLED_ERROR_TYPE: &str = "local_stream_attempt_cancelled";
const STREAM_ATTEMPT_CANCELLED_ERROR_MESSAGE: &str = "Local stream attempt was dropped before terminal finalization, usually because the client disconnected or the request task was cancelled.";
const SSE_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15);
const SSE_KEEPALIVE_BYTES: &[u8] = b": aether-keepalive\n\n";
const SSE_CONTROL_FILTER_MAX_BUFFER_BYTES: usize = 1024 * 1024;
@@ -3655,17 +3661,30 @@ pub(crate) fn execute_execution_runtime_stream<'a>(
report_kind: Option<String>,
report_context: Option<serde_json::Value>,
) -> Pin<Box<dyn Future<Output = Result<Option<Response<Body>>, GatewayError>> + Send + 'a>> {
Box::pin(execute_execution_runtime_stream_inner(
state,
plan,
trace_id,
decision,
plan_kind,
report_kind,
report_context,
None,
None,
))
Box::pin(async move {
let mut cancellation_guard = AttemptCancellationGuard::disarmed(
state,
STREAM_ATTEMPT_CANCELLED_ERROR_TYPE,
STREAM_ATTEMPT_CANCELLED_ERROR_MESSAGE,
);
let result = execute_execution_runtime_stream_inner(
state,
plan,
trace_id,
decision,
plan_kind,
report_kind,
report_context,
None,
None,
&mut cancellation_guard,
)
.await;
// The attempt reached its own terminal path, or handed settlement to the
// stream finalizer that now lives in the response body.
cancellation_guard.disarm();
result
})
}
#[allow(clippy::too_many_arguments)]
@@ -3687,7 +3706,12 @@ pub(crate) fn execute_execution_runtime_stream_with_retry_scope<'a>(
Box::pin(async move {
let mut retry_scope = AiAttemptRetryScope::Candidate;
let mut fallback_response = None;
let response = execute_execution_runtime_stream_inner(
let mut cancellation_guard = AttemptCancellationGuard::disarmed(
state,
STREAM_ATTEMPT_CANCELLED_ERROR_TYPE,
STREAM_ATTEMPT_CANCELLED_ERROR_MESSAGE,
);
let result = execute_execution_runtime_stream_inner(
state,
plan,
trace_id,
@@ -3697,8 +3721,13 @@ pub(crate) fn execute_execution_runtime_stream_with_retry_scope<'a>(
report_context,
Some(&mut retry_scope),
Some(&mut fallback_response),
&mut cancellation_guard,
)
.await?;
.await;
// The attempt reached its own terminal path, or handed settlement to the
// stream finalizer that now lives in the response body.
cancellation_guard.disarm();
let response = result?;
Ok(match response {
Some(response) => AiAttemptExecutionOutcome::Responded(response),
None => AiAttemptExecutionOutcome::Retry {
@@ -3744,6 +3773,7 @@ async fn maybe_build_stream_transport_error_stop_response(
.map(Some)
}
#[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection
async fn execute_execution_runtime_stream_inner(
state: &AppState,
mut plan: ExecutionPlan,
@@ -3754,6 +3784,7 @@ async fn execute_execution_runtime_stream_inner(
mut report_context: Option<serde_json::Value>,
mut retry_scope_out: Option<&mut AiAttemptRetryScope>,
mut retry_fallback_out: Option<&mut Option<Response<Body>>>,
cancellation_guard: &mut AttemptCancellationGuard,
) -> Result<Option<Response<Body>>, GatewayError> {
let stream_started_at = Instant::now();
let mut stage_trace = RequestStageTrace::from_env();
@@ -3837,6 +3868,16 @@ async fn execute_execution_runtime_stream_inner(
)
.await;
}
// From here the attempt owns non-terminal rows, and everything that could
// settle them runs inside the downstream request future. Arm the guard so a
// client disconnect before the stream finalizer exists still settles them.
cancellation_guard.arm(
&plan,
report_context.as_ref(),
request_candidate_status_snapshot.as_ref(),
candidate_started_unix_secs,
stream_started_at,
);
let plan_request_id_for_log = short_request_id(plan.request_id.as_str());
let provider_name = plan
.provider_name
@@ -10561,7 +10602,7 @@ mod tests {
}
#[tokio::test]
async fn malformed_antigravity_function_call_retries_before_stream_commit() {
async fn malformed_antigravity_function_call_streams_thought_then_fails_in_band() {
let request_id = "req-antigravity-malformed-function-call";
let plan = antigravity_gemini_stream_plan(request_id);
let provider_catalog = provider_catalog_for_plan(
@@ -10640,10 +10681,35 @@ mod tests {
None,
)
.await
.expect("malformed Antigravity stream should resolve through failover");
.expect("malformed Antigravity stream should return a client stream")
.expect("the first reasoning delta should commit the selected candidate");
assert!(response.is_none());
assert_eq!(retry_scope, AiAttemptRetryScope::Candidate);
assert_eq!(response.status(), StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body should read");
let body = String::from_utf8(body.to_vec()).expect("response body should be utf8");
assert!(
body.contains("event: response.reasoning_summary_text.delta\n"),
"{body}"
);
assert!(
body.contains("\"delta\":\"Validating the document.\""),
"{body}"
);
assert!(body.contains("event: response.failed\n"), "{body}");
assert!(
body.contains("\"code\":\"MALFORMED_FUNCTION_CALL\""),
"{body}"
);
assert!(
body.contains(
"\"message\":\"Malformed function call: Function call is empty - no input to parse.\""
),
"{body}"
);
assert!(!body.contains("unsupported_finish_reason"), "{body}");
assert_eq!(retry_scope, AiAttemptRetryScope::Provider);
}
fn tunnel_proxy_snapshot(base_url: String) -> aether_contracts::ProxySnapshot {
apps/aether-gateway/src/execution_runtime/transport_failure.rs
+19
View File
@@ -22,6 +22,7 @@ const TRANSPORT_ERROR_CLIENT_MESSAGE: &str =
#[derive(Debug, Default)]
pub(crate) struct StreamCandidateWatchdogProgress {
terminal_started: AtomicBool,
abandoned: AtomicBool,
}
tokio::task_local! {
@@ -37,6 +38,24 @@ impl StreamCandidateWatchdogProgress {
self.terminal_started.load(Ordering::Acquire)
}
/// The watchdog gave up waiting and settles this attempt itself.
///
/// The attempt future is dropped once the watchdog returns, so its own
/// cancellation guard must stay out of the way instead of racing the
/// watchdog's terminal rows with a cancellation.
pub(crate) fn mark_abandoned(&self) {
self.abandoned.store(true, Ordering::Release);
}
pub(crate) fn abandoned(&self) -> bool {
self.abandoned.load(Ordering::Acquire)
}
/// The watchdog watching the attempt on this task, if it runs under one.
pub(crate) fn current() -> Option<Arc<Self>> {
STREAM_CANDIDATE_WATCHDOG_PROGRESS.try_with(Arc::clone).ok()
}
pub(crate) async fn scope<F>(self: Arc<Self>, future: F) -> F::Output
where
F: Future,
apps/aether-gateway/src/executor/candidate_loop.rs
+84 -2
View File
@@ -412,6 +412,7 @@ where
decision,
plan_kind,
transfer_tracker,
request_first_byte_started_at: Instant::now(),
};
match run_ai_attempt_loop(&port, plan_and_reports).await? {
AiAttemptLoopOutcome::Responded(response) => {
@@ -482,6 +483,7 @@ where
decision,
plan_kind,
transfer_tracker,
request_first_byte_started_at: Instant::now(),
};
run_dynamic_attempt_loop(
&port,
@@ -946,6 +948,10 @@ struct StreamAttemptLoopPort<'a> {
decision: &'a GatewayControlDecision,
plan_kind: &'a str,
transfer_tracker: &'a ProviderTransferTracker,
/// All candidates in one downstream stream request share this origin.
/// Without it every retry receives a fresh full first-byte timeout and a
/// 30-second provider timeout can accumulate into a 60-120 second stall.
request_first_byte_started_at: Instant,
}
#[async_trait]
@@ -1058,6 +1064,7 @@ where
self.plan_kind,
plan,
watchdog_report_context,
self.request_first_byte_started_at,
stop_on_transport_errors,
move || async move {
execute_execution_runtime_stream_with_retry_scope(
@@ -1085,7 +1092,7 @@ where
http::StatusCode::GATEWAY_TIMEOUT.as_u16(),
"local_stream_candidate_watchdog_timeout",
stream_candidate_watchdog_timeout_message(),
watchdog_started_at.elapsed().as_millis() as u64,
self.request_first_byte_started_at.elapsed().as_millis() as u64,
)
.await?,
)
@@ -1368,6 +1375,7 @@ async fn execute_stream_candidate_with_watchdog<Fut>(
plan_kind: &str,
plan: &aether_contracts::ExecutionPlan,
report_context: Option<&serde_json::Value>,
request_first_byte_started_at: Instant,
stop_on_transport_errors: bool,
execute: impl FnOnce() -> Fut,
) -> Result<StreamCandidateWatchdogOutcome, GatewayError>
@@ -1377,6 +1385,7 @@ where
> + Send,
{
let timeout_duration = resolve_stream_candidate_watchdog_timeout(plan, report_context);
let request_first_byte_deadline = request_first_byte_started_at + timeout_duration;
let candidate_started_at = std::time::Instant::now();
let candidate_started_unix_ms = current_unix_ms();
let permit = match acquire_upstream_execution_gate(state, trace_id).await {
@@ -1402,7 +1411,14 @@ where
let watchdog_progress = StreamCandidateWatchdogProgress::shared();
let execution = watchdog_progress.clone().scope(execute());
tokio::pin!(execution);
let deadline = tokio::time::sleep(timeout_duration);
// This is an absolute request-level deadline, not a new timeout for this
// candidate. Retries therefore consume only the budget left by earlier
// candidates instead of resetting the full provider timeout.
let candidate_budget_ms = request_first_byte_deadline
.saturating_duration_since(Instant::now())
.as_millis()
.min(u128::from(u64::MAX)) as u64;
let deadline = tokio::time::sleep_until(request_first_byte_deadline);
tokio::pin!(deadline);
let execution_result = tokio::select! {
biased;
@@ -1418,6 +1434,10 @@ where
let outcome = match execution_result {
Some(result) => result.map(StreamCandidateWatchdogOutcome::Executed),
None => {
// The abandoned attempt is dropped when this function returns.
// Claim its settlement before that so its cancellation guard does
// not race the watchdog rows written just below.
watchdog_progress.mark_abandoned();
let finished_at_unix_ms = current_unix_ms();
let request_id = short_request_id(plan.request_id.as_str());
let provider_name = plan.provider_name.as_deref().unwrap_or("-");
@@ -1427,6 +1447,10 @@ where
.map(|value| value.to_string())
.unwrap_or_else(|| "-".to_string());
let timeout_ms = u64::try_from(timeout_duration.as_millis()).unwrap_or(u64::MAX);
let request_elapsed_ms = request_first_byte_started_at
.elapsed()
.as_millis()
.min(u128::from(u64::MAX)) as u64;
record_local_request_candidate_status(
state,
plan,
@@ -1455,6 +1479,8 @@ where
model_name,
candidate_index = candidate_index.as_str(),
timeout_ms,
candidate_budget_ms,
request_elapsed_ms,
"gateway local stream candidate watchdog timed out"
);
if stop_on_transport_errors {
@@ -2385,6 +2411,7 @@ mod tests {
"claude_cli_stream",
&plan,
Some(&report_context),
Instant::now(),
false,
|| {
std::future::pending::<
@@ -2423,6 +2450,56 @@ mod tests {
assert_eq!(record.candidate_index, 2);
}
#[tokio::test]
async fn stream_candidate_retry_does_not_reset_an_expired_request_first_byte_budget() {
let writer = Arc::new(TestRequestCandidateWriter::default());
let plan = test_plan(Some(ExecutionTimeouts {
first_byte_ms: Some(250),
..ExecutionTimeouts::default()
}));
let report_context = test_report_context();
// Stand in for earlier candidates having already consumed the request's
// complete first-byte budget. A per-candidate watchdog would wait a new
// 250 ms here; the shared absolute deadline must settle immediately.
let request_first_byte_started_at = Instant::now() - Duration::from_millis(300);
let result = tokio::time::timeout(
Duration::from_millis(100),
execute_stream_candidate_with_watchdog(
writer.as_ref(),
"trace_watchdog_shared_budget",
"claude_cli_stream",
&plan,
Some(&report_context),
request_first_byte_started_at,
false,
|| {
std::future::pending::<
Result<AiAttemptExecutionOutcome<Response<Body>>, GatewayError>,
>()
},
),
)
.await
.expect("an expired request-level first-byte budget must not restart per candidate");
assert!(matches!(
result,
Ok(StreamCandidateWatchdogOutcome::Executed(
AiAttemptExecutionOutcome::Retry {
scope: AiAttemptRetryScope::Candidate,
fallback_response: None,
}
))
));
let records = writer.records.lock().await;
assert_eq!(records.len(), 1);
assert_eq!(
records[0].error_type.as_deref(),
Some("local_stream_candidate_watchdog_timeout")
);
}
#[tokio::test]
async fn stream_candidate_watchdog_can_stop_on_transport_error() {
let writer = Arc::new(TestRequestCandidateWriter::default());
@@ -2438,6 +2515,7 @@ mod tests {
"claude_cli_stream",
&plan,
Some(&report_context),
Instant::now(),
true,
|| {
std::future::pending::<
@@ -2475,6 +2553,7 @@ mod tests {
"claude_cli_stream",
&plan,
Some(&report_context),
Instant::now(),
true,
|| async {
mark_stream_candidate_watchdog_terminal_started();
@@ -2507,6 +2586,7 @@ mod tests {
"claude_cli_stream",
&plan,
Some(&report_context),
Instant::now(),
true,
|| async {
Err(GatewayError::UpstreamUnavailable {
@@ -2546,6 +2626,7 @@ mod tests {
"claude_cli_stream",
&plan,
Some(&report_context),
Instant::now(),
false,
|| async {
panic!("execute future should not run while upstream execution gate is saturated")
@@ -2593,6 +2674,7 @@ mod tests {
"claude_cli_stream",
&plan,
Some(&report_context),
Instant::now(),
false,
|| async {
Err(GatewayError::AdmissionTimeout {
apps/aether-gateway/src/tests/ai_execute/lifecycle.rs
+109
View File
@@ -14,6 +14,9 @@ use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadReposi
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::repository::candidates::{
RequestCandidateReadRepository, RequestCandidateStatus,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -427,6 +430,112 @@ async fn gateway_stops_execution_runtime_stream_when_client_disconnects_impl() {
upstream_handle.abort();
}
#[test]
fn gateway_settles_stream_attempt_when_client_disconnects_before_first_byte() {
run_lifecycle_test(
"gateway_settles_stream_attempt_when_client_disconnects_before_first_byte",
gateway_settles_stream_attempt_when_client_disconnects_before_first_byte_impl,
);
}
async fn gateway_settles_stream_attempt_when_client_disconnects_before_first_byte_impl() {
// The execution runtime accepts the plan and then goes quiet, so the attempt
// is parked between its `pending` rows and the first upstream byte.
let execution_runtime = Router::new().route(
"/v1/execute/stream",
any(|_request: Request| async move {
tokio::time::sleep(std::time::Duration::from_secs(30)).await;
StatusCode::OK
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-stream-precommit-disconnect")),
sample_local_openai_auth_snapshot(
"api-key-openai-lifecycle-local-1",
"user-openai-lifecycle-local-1",
),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_local_openai_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_local_openai_provider()],
vec![sample_local_openai_endpoint()],
vec![sample_local_openai_key()],
));
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let request = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-openai-stream-precommit-disconnect",
)
.header(
TRACE_ID_HEADER,
"trace-openai-chat-stream-precommit-disconnect-123",
)
.body("{\"model\":\"gpt-5\",\"messages\":[],\"stream\":true}")
.send();
// Drop the in-flight request the way a downstream client does when its own
// first-byte timeout fires, before any response header exists.
assert!(
tokio::time::timeout(std::time::Duration::from_millis(750), request)
.await
.is_err(),
"the execution runtime should not have answered before the client gave up"
);
let mut stored_candidates = Vec::new();
for _ in 0..200 {
stored_candidates = request_candidate_repository
.list_by_request_id("trace-openai-chat-stream-precommit-disconnect-123")
.await
.expect("request candidate trace should read");
if stored_candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Cancelled)
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let cancelled = stored_candidates
.iter()
.find(|candidate| candidate.status == RequestCandidateStatus::Cancelled)
.unwrap_or_else(|| {
panic!("dropped stream attempt should settle as cancelled: {stored_candidates:?}")
});
assert_eq!(cancelled.status_code, Some(499));
assert_eq!(
cancelled.error_type.as_deref(),
Some("local_stream_attempt_cancelled")
);
assert!(cancelled.finished_at_unix_ms.is_some());
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[test]
fn gateway_returns_error_body_when_prefetch_detects_embedded_stream_error() {
run_lifecycle_test(
crates/aether-ai/formats/src/formats/gemini/generate_content/stream.rs
+131
View File
@@ -104,10 +104,25 @@ impl GeminiProviderState {
let Some(candidate_object) = candidate.as_object() else {
continue;
};
let (response_id, response_model) = self.identity(report_context);
let terminal_error = gemini_stream_terminal_error_payload(
candidate_object,
response_id.as_str(),
response_model.as_str(),
event_object.get("usageMetadata"),
);
let Some(content) = candidate_object.get("content").and_then(Value::as_object) else {
if let Some(payload) = terminal_error {
out.push(self.unknown_frame(report_context, payload));
self.finished = true;
}
continue;
};
let Some(parts) = content.get("parts").and_then(Value::as_array) else {
if let Some(payload) = terminal_error {
out.push(self.unknown_frame(report_context, payload));
self.finished = true;
}
continue;
};
if !parts.is_empty() {
@@ -306,6 +321,11 @@ impl GeminiProviderState {
});
}
}
if let Some(payload) = terminal_error {
out.push(self.unknown_frame(report_context, payload));
self.finished = true;
continue;
}
if let Some(finish_reason) =
candidate_object.get("finishReason").and_then(Value::as_str)
{
@@ -350,6 +370,60 @@ impl GeminiProviderState {
}
}
fn gemini_stream_terminal_error_payload(
candidate: &Map<String, Value>,
response_id: &str,
model: &str,
usage_metadata: Option<&Value>,
) -> Option<Value> {
let finish_reason = candidate
.get("finishReason")
.or_else(|| candidate.get("finish_reason"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| {
matches!(
*value,
"MALFORMED_FUNCTION_CALL"
| "UNEXPECTED_TOOL_CALL"
| "TOO_MANY_TOOL_CALLS"
| "MISSING_THOUGHT_SIGNATURE"
| "MALFORMED_RESPONSE"
)
})?;
let message = candidate
.get("finishMessage")
.or_else(|| candidate.get("finish_message"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.unwrap_or_else(|| format!("Gemini stream ended with {finish_reason}"));
let mut response = json!({
"id": response_id,
"object": "response",
"model": model,
"status": "failed",
"error": {
"type": "upstream_gemini_finish_error",
"code": finish_reason,
"message": message,
"upstream_status": 200
}
});
if let Some(usage) = canonical_usage_from_gemini_usage(usage_metadata)
.map(|usage| openai_responses_usage_from_usage(&usage))
{
response["usage"] = usage;
}
Some(json!({
"type": "response.failed",
"response": response
}))
}
fn map_gemini_stream_finish_reason(value: &str) -> Option<&str> {
match value {
"STOP" => Some("stop"),
@@ -1011,6 +1085,63 @@ mod tests {
)));
}
#[test]
fn gemini_provider_state_emits_terminal_error_for_malformed_function_call() {
let mut state = GeminiProviderState::default();
let report_context = json!({});
let frames = state
.push_line(
&report_context,
data_line(json!({
"response": {
"responseId": "resp_malformed_tool_call",
"modelVersion": "gemini-3.7-flash-tiered",
"candidates": [{
"index": 0,
"content": {
"role": "model",
"parts": [{
"text": "",
"thoughtSignature": "opaque-thought-signature"
}]
},
"finishReason": "MALFORMED_FUNCTION_CALL",
"finishMessage": "Malformed function call: Function call is empty - no input to parse."
}],
"usageMetadata": {
"promptTokenCount": 206744,
"cachedContentTokenCount": 203947,
"thoughtsTokenCount": 1130,
"totalTokenCount": 207874
}
}
})),
)
.expect("malformed function call terminal should parse");
assert!(frames.iter().any(|frame| matches!(
&frame.event,
CanonicalStreamEvent::UnknownEvent(payload)
if payload["type"] == "response.failed"
&& payload["response"]["status"] == "failed"
&& payload["response"]["id"] == "resp_malformed_tool_call"
&& payload["response"]["model"] == "gemini-3.7-flash-tiered"
&& payload["response"]["error"]["code"] == "MALFORMED_FUNCTION_CALL"
&& payload["response"]["error"]["message"]
== "Malformed function call: Function call is empty - no input to parse."
&& payload["response"]["usage"]["input_tokens"] == 206744
&& payload["response"]["usage"]["output_tokens"] == 1130
&& payload["response"]["usage"]["total_tokens"] == 207874
)));
assert!(!frames
.iter()
.any(|frame| matches!(frame.event, CanonicalStreamEvent::Finish { .. })));
assert!(state
.finish(&report_context)
.expect("finished error stream should not synthesize success")
.is_empty());
}
#[test]
fn gemini_provider_state_parses_function_response_as_tool_result() {
let mut state = GeminiProviderState::default();
crates/aether-ai/formats/src/formats/openai/request_contract.rs
+11
View File
@@ -381,10 +381,21 @@ mod tests {
#[test]
fn finalization_strips_non_replayable_responses_reasoning_history() {
let gemini_carrier =
crate::formats::openai::responses::encode_gemini_tool_signature_carrier(
"opaque-gemini-thought-signature",
)
.expect("Gemini signature carrier");
let mut body = json!({
"model": "gpt-5.4",
"input": [
{"type": "reasoning", "id": "rs_provider_123", "summary": []},
{
"type": "reasoning",
"id": "rs_aether_55070860f6d45c6b8f6fa11efd9dff8a",
"summary": [],
"encrypted_content": gemini_carrier
},
{
"type": "reasoning",
"id": "item_72d3bd8d367d01977ace23f1",
crates/aether-ai/formats/src/formats/openai/responses/mod.rs
+48 -3
View File
@@ -167,9 +167,10 @@ pub fn normalize_openai_responses_message_item_ids(body: &mut Value) -> usize {
/// Removes reasoning history items that cannot be replayed against an OpenAI Responses backend.
///
/// Reasoning IDs are opaque provider references and must never be repaired by changing their
/// prefix. Foreign IDs (for example `item_...`) are therefore removed. Aether-synthesized
/// reasoning summaries are also removed unless they carry encrypted reasoning state that can be
/// replayed statelessly.
/// prefix. Foreign IDs (for example `item_...`) are therefore removed. Aether's Gemini signature
/// carriers are also removed: they are intentionally transported through the Responses
/// `encrypted_content` field so they can be restored on a later Gemini tool turn, but they are not
/// OpenAI ciphertext and must never be replayed to an OpenAI/Codex backend.
pub fn strip_incompatible_openai_responses_reasoning_items(
body: &mut Value,
provider_api_format: &str,
@@ -221,6 +222,13 @@ fn openai_responses_reasoning_item_is_replayable(
if object.get("type").and_then(Value::as_str) != Some("reasoning") {
return true;
}
if object
.get("encrypted_content")
.and_then(Value::as_str)
.is_some_and(|value| value.starts_with(GEMINI_TOOL_SIGNATURE_CARRIER_PREFIX))
{
return false;
}
if policy == OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque
&& deepseek_opaque_reasoning_item_is_replayable(object)
{
@@ -475,6 +483,43 @@ mod tests {
assert_eq!(input[2]["id"], "item_message_123");
}
#[test]
fn strips_gemini_signature_carriers_before_openai_replay() {
let gemini_item_id = openai_responses_synthetic_reasoning_item_id("resp_gemini", 0);
let openai_item_id = openai_responses_synthetic_reasoning_item_id("resp_openai", 0);
let carrier = encode_gemini_tool_signature_carrier_with_direction(
"opaque-gemini-thought-signature",
GeminiToolSignatureCarrierDirection::Next,
)
.expect("Gemini signature carrier");
let mut body = json!({
"input": [
{
"type": "reasoning",
"id": gemini_item_id,
"summary": [],
"encrypted_content": carrier
},
{
"type": "reasoning",
"id": openai_item_id,
"summary": [],
"encrypted_content": "provider-encrypted-state"
},
{"type": "reasoning", "id": "rs_provider_123", "summary": []}
]
});
assert_eq!(
strip_incompatible_openai_responses_reasoning_items(&mut body, "openai:responses"),
1
);
let input = body["input"].as_array().expect("input array");
assert_eq!(input.len(), 2);
assert_eq!(input[0]["encrypted_content"], "provider-encrypted-state");
assert_eq!(input[1]["id"], "rs_provider_123");
}
#[test]
fn reasoning_item_sanitizer_is_scoped_to_responses_targets() {
let mut body = json!({
crates/aether-ai/formats/src/formats/shared/stream_core/format_matrix.rs
+153 -9
View File
@@ -17,8 +17,9 @@ use crate::formats::shared::error_body::{
};
use crate::formats::shared::sse::encode_json_sse;
use crate::formats::shared::stream_core::common::{
decode_json_data_line, openai_stream_terminal_error_body, openai_stream_terminal_error_message,
unsupported_stream_event_message, CanonicalStreamEvent, CanonicalStreamFrame, CanonicalUsage,
canonical_usage_from_openai_usage, decode_json_data_line, openai_stream_terminal_error_body,
openai_stream_terminal_error_message, unsupported_stream_event_message, CanonicalStreamEvent,
CanonicalStreamFrame, CanonicalUsage,
};
use crate::formats::shared::AiSurfaceFinalizeError;
@@ -134,7 +135,11 @@ impl StreamingStandardFormatMatrix {
}
if let CanonicalStreamEvent::UnknownEvent(payload) = &frame.event {
self.terminated = true;
out.extend(client.emit_unknown_event(payload)?);
if openai_stream_terminal_error_body(payload).is_some() {
out.extend(client.emit_terminal_error_frame(frame)?);
} else {
out.extend(client.emit_unknown_event(payload)?);
}
break;
}
if let CanonicalStreamEvent::OpenAiResponsesOutputItem { raw_event, .. } = &frame.event
@@ -368,6 +373,10 @@ impl StreamingStandardTerminalObserver {
summary.observed_finish = true;
summary.finish_reason = Some("error".to_string());
summary.parser_error = openai_stream_terminal_error_message(&payload);
summary.standardized_usage = payload
.pointer("/response/usage")
.and_then(|usage| canonical_usage_from_openai_usage(Some(usage)))
.map(standardized_usage_from_canonical);
}
CanonicalStreamEvent::UnknownEvent(_) => {
summary.unknown_event_count = summary.unknown_event_count.saturating_add(1);
@@ -613,6 +622,44 @@ impl ClientStreamEmitter {
self.emit_error(error_body)
}
fn emit_terminal_error_frame(
&mut self,
frame: CanonicalStreamFrame,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if matches!(
self,
ClientStreamEmitter::OpenAIChat(_) | ClientStreamEmitter::OpenAIResponses(_)
) {
return self.emit(frame);
}
let CanonicalStreamEvent::UnknownEvent(payload) = frame.event else {
return self.emit(frame);
};
let Some(source_error_body) = openai_stream_terminal_error_body(&payload) else {
return self.emit_unknown_event(&payload);
};
let Some(error) = source_error_body.get("error") else {
return self.emit_unknown_event(&payload);
};
let message = error
.get("message")
.and_then(Value::as_str)
.unwrap_or("Upstream stream ended with an error");
let code = error.get("code").and_then(|value| match value {
Value::String(value) => Some(value.as_str()),
_ => None,
});
let Some(error_body) = build_core_error_body_for_client_format(
self.api_format(),
message,
code,
LocalCoreSyncErrorKind::ServerError,
) else {
return Ok(Vec::new());
};
self.emit_error(error_body)
}
fn emit_unsupported_finish_reason(
&mut self,
finish_reason: &str,
@@ -811,7 +858,13 @@ mod tests {
},
"finishReason": "MALFORMED_FUNCTION_CALL",
"finishMessage": "Malformed function call: Function call is empty - no input to parse."
}]
}],
"usageMetadata": {
"promptTokenCount": 206744,
"cachedContentTokenCount": 203947,
"thoughtsTokenCount": 1130,
"totalTokenCount": 207874
}
},
"responseId": "resp_malformed_tool_call"
})),
@@ -824,14 +877,105 @@ mod tests {
.expect("Gemini terminal frame should produce a summary");
assert!(summary.observed_finish);
assert_eq!(
summary.finish_reason.as_deref(),
Some("MALFORMED_FUNCTION_CALL")
);
assert_eq!(summary.finish_reason.as_deref(), Some("error"));
assert_eq!(
summary.parser_error.as_deref(),
Some("unsupported provider stream finish reason: MALFORMED_FUNCTION_CALL")
Some("Malformed function call: Function call is empty - no input to parse.")
);
let usage = summary
.standardized_usage
.expect("failed Gemini terminal should preserve usage");
assert_eq!(usage.input_tokens, 206744);
assert_eq!(usage.output_tokens, 1130);
assert_eq!(usage.cache_read_tokens, 203947);
}
#[test]
fn streams_gemini_thought_text_to_openai_responses_immediately() {
let context = report_context("gemini:generate_content", "openai:responses");
let mut matrix = StreamingStandardFormatMatrix::default();
let output = matrix
.transform_line(
&context,
data_line(json!({
"response": {
"responseId": "resp_reasoning_123",
"modelVersion": "gemini-3.7-flash-tiered",
"candidates": [{
"index": 0,
"content": {
"role": "model",
"parts": [{"thought": true, "text": "checking"}]
}
}]
}
})),
)
.expect("first Gemini thought chunk should transform");
let sse = String::from_utf8(output).expect("reasoning SSE should be utf8");
assert!(
sse.contains("event: response.reasoning_summary_text.delta\n"),
"{sse}"
);
assert!(sse.contains("\"delta\":\"checking\""), "{sse}");
}
#[test]
fn transforms_malformed_gemini_function_call_to_responses_failed() {
let context = report_context("gemini:generate_content", "openai:responses");
let mut matrix = StreamingStandardFormatMatrix::default();
let output = matrix
.transform_line(
&context,
data_line(json!({
"response": {
"responseId": "resp_malformed_tool_call",
"modelVersion": "gemini-3.7-flash-tiered",
"candidates": [{
"index": 0,
"content": {
"role": "model",
"parts": [{
"text": "",
"thoughtSignature": "opaque-thought-signature"
}]
},
"finishReason": "MALFORMED_FUNCTION_CALL",
"finishMessage": "Malformed function call: Function call is empty - no input to parse."
}],
"usageMetadata": {
"promptTokenCount": 206744,
"cachedContentTokenCount": 203947,
"thoughtsTokenCount": 1130,
"totalTokenCount": 207874
}
}
})),
)
.expect("malformed Gemini terminal should transform to a stream error");
let sse = String::from_utf8(output).expect("failed response SSE should be utf8");
assert!(sse.contains("event: response.failed\n"), "{sse}");
assert!(sse.contains("\"type\":\"response.failed\""), "{sse}");
assert!(
sse.contains("\"code\":\"MALFORMED_FUNCTION_CALL\""),
"{sse}"
);
assert!(
sse.contains(
"\"message\":\"Malformed function call: Function call is empty - no input to parse.\""
),
"{sse}"
);
assert!(sse.contains("\"input_tokens\":206744"), "{sse}");
assert!(sse.contains("\"output_tokens\":1130"), "{sse}");
assert!(sse.contains("\"cached_tokens\":203947"), "{sse}");
assert!(!sse.contains("unsupported_finish_reason"), "{sse}");
assert!(matrix
.finish(&context)
.expect("failed matrix should stay terminated")
.is_empty());
}
#[test]
crates/aether-usage/runtime/src/lib.rs
+2 -1
View File
@@ -61,7 +61,8 @@ pub use write::{
build_sync_terminal_usage_event, build_sync_terminal_usage_outcome,
build_sync_terminal_usage_payload_seed, build_sync_terminal_usage_seed,
build_terminal_usage_context_seed, build_terminal_usage_event_from_outcome,
build_terminal_usage_event_from_seed, build_usage_event_data_seed, LifecycleUsageSeed,
build_terminal_usage_event_from_seed, build_usage_event_data_seed,
build_usage_event_data_seed_describing_request_bodies, LifecycleUsageSeed,
StreamTerminalUsagePayloadSeed, SyncTerminalUsagePayloadSeed, TerminalUsageContextSeed,
TerminalUsageOutcome, TerminalUsageSeed, UsageTerminalState,
};
crates/aether-usage/runtime/src/write.rs
+186 -12
View File
@@ -72,6 +72,20 @@ struct RuntimeRequestCaptureSeed {
provider_request: Option<Value>,
provider_request_body_ref: Option<String>,
body_states: UsageBodyStatesSeed,
request_has_inline_body: bool,
provider_request_has_inline_body: bool,
}
/// Whether a seed keeps the request bodies it describes, or only describes them.
///
/// A holder that has to outlive the request itself pays for every byte it keeps,
/// and a request body can be megabytes. [`RequestBodyCapture::Describe`] computes
/// every capture state, reference and derived fact from the real plan and report
/// context, and leaves out only the body values themselves.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RequestBodyCapture {
Keep,
Describe,
}
#[derive(Debug, Clone, PartialEq)]
@@ -804,7 +818,8 @@ pub fn build_terminal_usage_context_seed(
report_context: Option<&Value>,
) -> TerminalUsageContextSeed {
let context = report_context.and_then(Value::as_object);
let request_capture = build_runtime_request_capture_seed(plan, context);
let request_capture =
build_runtime_request_capture_seed(plan, context, RequestBodyCapture::Keep);
let client_contract = context_string(context, "client_contract")
.or_else(|| context_string(context, "client_api_format"))
.or_else(|| non_empty_str(Some(plan.client_api_format.as_str())))
@@ -1691,16 +1706,34 @@ pub fn build_usage_event_data_seed(
plan: &ExecutionPlan,
report_context: Option<&Value>,
) -> UsageEventData {
build_usage_event_data_seed_with_detail(plan, report_context)
build_usage_event_data_seed_with_detail(plan, report_context, RequestBodyCapture::Keep)
}
/// Builds the same seed as [`build_usage_event_data_seed`] without keeping the
/// request bodies.
///
/// This is for a caller that has to hold a seed for the whole life of an attempt
/// so it can still write a terminal row if the attempt is dropped: a request body
/// can be megabytes, and holding one per in-flight attempt is far more expensive
/// than the row it would eventually capture. Every capture state, body reference
/// and derived request fact is still computed from the real plan and report
/// context, so the resulting terminal write preserves the capture an earlier
/// non-terminal write recorded rather than clearing it.
pub fn build_usage_event_data_seed_describing_request_bodies(
plan: &ExecutionPlan,
report_context: Option<&Value>,
) -> UsageEventData {
build_usage_event_data_seed_with_detail(plan, report_context, RequestBodyCapture::Describe)
}
fn build_usage_event_data_seed_with_detail(
plan: &ExecutionPlan,
report_context: Option<&Value>,
capture: RequestBodyCapture,
) -> UsageEventData {
let context = report_context.and_then(Value::as_object);
let routing = build_runtime_routing_seed(plan, context);
let request_capture = build_runtime_request_capture_seed(plan, context);
let request_capture = build_runtime_request_capture_seed(plan, context, capture);
let api_format = context_string(context, "client_api_format")
.or_else(|| non_empty_str(Some(plan.client_api_format.as_str())));
let endpoint_api_format = context_string(context, "provider_api_format")
@@ -1714,7 +1747,10 @@ fn build_usage_event_data_seed_with_detail(
let request_type = Some(infer_request_type_from_contracts(
api_format.as_deref(),
endpoint_api_format.as_deref(),
request_capture.provider_request.as_ref(),
request_capture
.provider_request
.as_ref()
.or_else(|| provider_request_body_ref_for_inference(plan, context)),
));
let api_family = api_format
.as_deref()
@@ -1737,9 +1773,9 @@ fn build_usage_event_data_seed_with_detail(
build_runtime_request_metadata_seed_from_parts(
plan,
context,
request_capture.request_body.is_some(),
request_capture.request_has_inline_body,
request_capture.request_body_ref.as_deref(),
request_capture.provider_request.is_some(),
request_capture.provider_request_has_inline_body,
request_capture.provider_request_body_ref.as_deref(),
plan.body.body_bytes_b64.as_deref(),
),
@@ -1999,17 +2035,29 @@ fn plan_has_inline_json_body_for_usage(plan: &ExecutionPlan) -> bool {
fn build_runtime_request_capture_seed(
plan: &ExecutionPlan,
context: Option<&Map<String, Value>>,
capture: RequestBodyCapture,
) -> RuntimeRequestCaptureSeed {
let request_body = context_body_value(context, "original_request_body");
// Presence, not the value, is what every capture state and derived fact is
// built from, so both capture modes agree on all of them.
let request_has_inline_body = context_has_inline_body(context, "original_request_body");
let provider_request_has_inline_body =
context_has_inline_body(context, "provider_request_body")
|| plan_has_inline_json_body_for_usage(plan);
let (request_body, provider_request) = match capture {
RequestBodyCapture::Keep => (
context_body_value(context, "original_request_body"),
context_body_value(context, "provider_request_body")
.or_else(|| plan_json_body_capture_for_usage(plan)),
),
RequestBodyCapture::Describe => (None, None),
};
let request_body_ref = context_string(context, "request_body_ref");
let provider_request = context_body_value(context, "provider_request_body")
.or_else(|| plan_json_body_capture_for_usage(plan));
let provider_request_body_ref = context_string(context, "provider_request_body_ref")
.or_else(|| non_empty_str(plan.body.body_ref.as_deref()));
let body_states = build_runtime_body_states_seed_from_parts(
request_body.is_some(),
request_has_inline_body,
request_body_ref.as_deref(),
provider_request.is_some(),
provider_request_has_inline_body,
provider_request_body_ref.as_deref(),
plan.body.body_bytes_b64.is_some(),
);
@@ -2020,9 +2068,26 @@ fn build_runtime_request_capture_seed(
provider_request,
provider_request_body_ref,
body_states,
request_has_inline_body,
provider_request_has_inline_body,
}
}
/// Borrows the provider request body that request-type inference reads, without
/// cloning it.
fn provider_request_body_ref_for_inference<'a>(
plan: &'a ExecutionPlan,
context: Option<&'a Map<String, Value>>,
) -> Option<&'a Value> {
context_value_ref(context, "provider_request_body")
.filter(|value| !value.is_null())
.or_else(|| {
plan_has_inline_json_body_for_usage(plan)
.then_some(plan.body.json_body.as_ref())
.flatten()
})
}
fn build_runtime_request_metadata_seed(
plan: &ExecutionPlan,
context: Option<&Map<String, Value>>,
@@ -3493,7 +3558,8 @@ mod tests {
build_streaming_usage_event_from_owned_seed, build_streaming_usage_record,
build_sync_terminal_usage_event, build_sync_terminal_usage_payload_seed,
build_sync_terminal_usage_seed, build_terminal_usage_context_seed,
build_terminal_usage_event_from_seed, build_usage_event_data_seed, decode_body_for_storage,
build_terminal_usage_event_from_seed, build_usage_event_data_seed,
build_usage_event_data_seed_describing_request_bodies, decode_body_for_storage,
extract_token_counts_from_json, extract_token_counts_from_value, headers_to_json,
mask_header_value, mask_sensitive_body_fields, mask_sensitive_headers_in_json_value,
parse_sse_body_for_storage, resolve_error_message, trim_owned_non_empty_string,
@@ -6842,6 +6908,114 @@ mod tests {
assert!(body_size.get("provider_request_body").is_some());
}
#[test]
fn describing_request_bodies_matches_the_capturing_seed_apart_from_the_bodies() {
let plan = ExecutionPlan {
request_id: "req-seed-describe-1".to_string(),
candidate_id: Some("cand-seed-describe-1".to_string()),
provider_name: Some("OpenAI".to_string()),
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "key-1".to_string(),
method: "POST".to_string(),
url: "https://example.com/v1/chat/completions".to_string(),
headers: BTreeMap::new(),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gpt-5",
"service_tier": "priority",
"reasoning": {"effort": "high"}
})),
stream: false,
client_api_format: "openai:chat".to_string(),
provider_api_format: "openai:chat".to_string(),
model_name: Some("gpt-5".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let report_context = json!({
"client_api_format": "openai:chat",
"provider_api_format": "openai:chat",
"original_request_body": {"model": "gpt-5", "messages": []},
"original_headers": {"accept": "application/json"}
});
let captured = build_usage_event_data_seed(&plan, Some(&report_context));
let described =
build_usage_event_data_seed_describing_request_bodies(&plan, Some(&report_context));
// Only the two heavy values differ.
assert!(captured.request_body.is_some());
assert!(captured.provider_request_body.is_some());
assert_eq!(described.request_body, None);
assert_eq!(described.provider_request_body, None);
// Everything a terminal write reads to decide what to do with the stored
// capture is identical, so the described seed preserves it rather than
// clearing it.
assert_eq!(
described.request_body_state,
Some(UsageBodyCaptureState::Inline)
);
assert_eq!(
described.provider_request_body_state,
Some(UsageBodyCaptureState::Inline)
);
assert_eq!(described.request_body_state, captured.request_body_state);
assert_eq!(
described.provider_request_body_state,
captured.provider_request_body_state
);
assert_eq!(described.request_body_ref, captured.request_body_ref);
assert_eq!(
described.provider_request_body_ref,
captured.provider_request_body_ref
);
assert_eq!(described.request_type, captured.request_type);
assert_eq!(described.request_metadata, captured.request_metadata);
assert_eq!(
described.provider_request_headers,
captured.provider_request_headers
);
assert_eq!(described.request_headers, captured.request_headers);
}
#[test]
fn describing_request_bodies_keeps_the_unavailable_marker_for_raw_bodies() {
let mut plan = ExecutionPlan {
request_id: "req-seed-describe-2".to_string(),
candidate_id: None,
provider_name: Some("OpenAI".to_string()),
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "key-1".to_string(),
method: "POST".to_string(),
url: "https://example.com/v1/chat/completions".to_string(),
headers: BTreeMap::new(),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(json!({"model": "gpt-5"})),
stream: false,
client_api_format: "openai:chat".to_string(),
provider_api_format: "openai:chat".to_string(),
model_name: Some("gpt-5".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
};
plan.body.json_body = None;
plan.body.body_bytes_b64 = Some("eyJtb2RlbCI6ICJncHQtNSJ9".to_string());
let described = build_usage_event_data_seed_describing_request_bodies(&plan, None);
assert_eq!(
described.provider_request_body_state,
Some(UsageBodyCaptureState::Unavailable)
);
}
#[test]
fn masks_known_sensitive_header_values() {
let token = "Bearer eyJhbGciOiJSUzI1NiJ9.payload-here.signature-tail";