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

fix(gateway): drain downstream-disconnected streams and stop inferring cancelled usage

Browse Source
This commit is contained in:
zhefox
2026-05-21 20:30:44 +08:00
parent 71f9afc526
commit 77c2d91eb0
4 changed files with 186 additions and 466 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
apps/aether-gateway/src/ai_serving/planner/standard/codex/tests.rs
View File

@@ -85,7 +85,7 @@ fn injects_stable_prompt_cache_key_for_codex_requests() {
assert_eq!(
body["prompt_cache_key"],
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
"53363264-dbb0-5f9d-b9c7-3e92c45c5bdf"
);
}

6
apps/aether-gateway/src/ai_serving/planner/standard/normalize/tests.rs
View File

@@ -209,7 +209,7 @@ fn local_openai_responses_compact_wrapper_strips_include_for_codex_requests() {
assert_eq!(provider_request_body["instructions"], "");
assert_eq!(
provider_request_body["prompt_cache_key"],
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
"3d2e2842-74cb-55dd-803a-b8940b3500c2"
);
}
@@ -355,7 +355,7 @@ fn injects_codex_prompt_cache_key_for_openai_responses_cross_format_requests() {
assert_eq!(
provider_request_body["prompt_cache_key"],
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
"b4dfeb75-b105-544c-a706-39b92f0bddb0"
);
}
@@ -385,6 +385,6 @@ fn injects_codex_prompt_cache_key_for_openai_chat_cross_format_requests() {
assert_eq!(
provider_request_body["prompt_cache_key"],
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
"4ee6ea6e-3ac6-5a18-8cb8-1f8b956419e5"
);
}

228
apps/aether-gateway/src/execution_runtime/stream/execution.rs
View File

@@ -1546,24 +1546,6 @@ where
read_next_frame(lines).await
}
async fn next_stream_frame_until_downstream_closed<R>(
buffered_frames: &mut VecDeque<StreamFrame>,
lines: &mut FramedRead<R, LinesCodec>,
tx: &mpsc::Sender<Result<Bytes, IoError>>,
) -> Result<Option<StreamFrame>, GatewayError>
where
R: tokio::io::AsyncRead + Unpin,
{
if let Some(frame) = buffered_frames.pop_front() {
return Ok(Some(frame));
}
tokio::select! {
frame = read_next_frame(lines) => frame,
() = tx.closed() => Ok(None),
}
}
fn should_refresh_stream_usage_telemetry(
previous: Option<&ExecutionTelemetry>,
next: &ExecutionTelemetry,
@@ -2764,15 +2746,7 @@ async fn execute_stream_from_frame_stream(
image_stream_total_timeout.as_mut()
{
tokio::select! {
result = next_stream_frame_until_downstream_closed(
&mut buffered_frames,
&mut lines,
&tx,
) => result,
() = tx.closed() => {
downstream_dropped = true;
break;
}
result = next_stream_frame(&mut buffered_frames, &mut lines) => result,
_ = timeout_sleep.as_mut() => {
let timeout_ms = openai_image_stream_total_timeout_ms
.unwrap_or(OPENAI_IMAGE_STREAM_DEFAULT_TOTAL_TIMEOUT_MS);
@@ -2823,8 +2797,7 @@ async fn execute_stream_from_frame_stream(
}
}
} else {
next_stream_frame_until_downstream_closed(&mut buffered_frames, &mut lines, &tx)
.await
next_stream_frame(&mut buffered_frames, &mut lines).await
};
let next_frame = match next_frame_result {
Ok(frame) => frame,
@@ -3003,6 +2976,9 @@ async fn execute_stream_from_frame_stream(
u64::try_from(rewritten_chunk.len()).unwrap_or(u64::MAX);
let chunk_completed_stream =
stream_chunk_contains_sse_done(&rewritten_chunk);
if downstream_dropped {
continue;
}
if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_disconnected",
@@ -3010,10 +2986,9 @@ async fn execute_stream_from_frame_stream(
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
"gateway stream downstream dropped; stopping execution runtime stream forwarding"
"gateway stream downstream dropped; continuing to drain execution runtime stream"
);
downstream_dropped = true;
break;
} else {
client_visible_stream_completed |= chunk_completed_stream;
client_stream_bytes.fetch_add(rewritten_chunk_len, Ordering::Relaxed);
@@ -3070,30 +3045,30 @@ async fn execute_stream_from_frame_stream(
}
if downstream_dropped {
drop(lines);
debug!(
event_name = "execution_runtime_stream_flush_skipped",
event_name = "execution_runtime_stream_client_flush_skipped",
log_type = "debug",
debug_context = "redacted",
stream_status = "downstream_disconnected",
trace_id = %trace_id_owned,
"gateway skipped local stream flush after downstream disconnect"
"gateway skipped client stream flush after downstream disconnect"
);
} else {
if let Some(normalizer) = private_stream_normalizer.as_mut() {
match normalizer.finish() {
Ok(normalized_chunk) if !normalized_chunk.is_empty() => {
if let (Some(observer), Some(report_context)) = (
stream_usage_observer.as_mut(),
stream_usage_report_context.as_ref(),
) {
observe_stream_usage_bytes(
observer,
report_context,
&mut stream_usage_observer_buffered,
&normalized_chunk,
);
}
}
if let Some(normalizer) = private_stream_normalizer.as_mut() {
match normalizer.finish() {
Ok(normalized_chunk) if !normalized_chunk.is_empty() => {
if let (Some(observer), Some(report_context)) = (
stream_usage_observer.as_mut(),
stream_usage_report_context.as_ref(),
) {
observe_stream_usage_bytes(
observer,
report_context,
&mut stream_usage_observer_buffered,
&normalized_chunk,
);
}
if !downstream_dropped {
let rewritten_chunk = if let Some(rewriter) = local_stream_rewriter.as_mut()
{
match rewriter.push_chunk(&normalized_chunk) {
@@ -3157,86 +3132,85 @@ async fn execute_stream_from_frame_stream(
}
}
}
}
Ok(_) => {}
Err(err) => {
warn!(
event_name = "stream_execution_normalization_flush_failed",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
error = ?err,
"gateway failed to flush private stream normalization"
);
terminal_failure.get_or_insert_with(|| {
build_stream_failure_report(
"execution_runtime_stream_rewrite_flush_error",
format!("failed to flush private stream normalization: {err:?}"),
502,
)
});
}
}
}
if !downstream_dropped {
if let Some(rewriter) = local_stream_rewriter.as_mut() {
match rewriter.finish() {
Ok(flushed_chunk) if !flushed_chunk.is_empty() => {
append_stream_capture_bytes(
&mut buffered_body,
&flushed_chunk,
max_stream_body_buffer_bytes,
&mut client_body_truncated,
);
let flushed_chunk_len =
u64::try_from(flushed_chunk.len()).unwrap_or(u64::MAX);
let chunk_completed_stream = stream_chunk_contains_sse_done(&flushed_chunk);
if tx.send(Ok(Bytes::from(flushed_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_rewrite_flush_disconnected",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
"gateway stream downstream dropped while flushing local stream rewrite"
);
downstream_dropped = true;
} else {
client_visible_stream_completed |= chunk_completed_stream;
client_stream_bytes.fetch_add(flushed_chunk_len, Ordering::Relaxed);
last_client_chunk_elapsed_ms.store(
stream_started_at_for_report
.elapsed()
.as_millis()
.min(u128::from(u64::MAX))
as u64,
Ordering::Relaxed,
);
}
}
Ok(_) => {}
Err(err) => {
warn!(
event_name = "stream_execution_normalization_flush_failed",
event_name = "stream_execution_rewrite_flush_failed",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
error = ?err,
"gateway failed to flush private stream normalization"
"gateway failed to flush local stream rewrite"
);
terminal_failure.get_or_insert_with(|| {
build_stream_failure_report(
"execution_runtime_stream_rewrite_flush_error",
format!("failed to flush private stream normalization: {err:?}"),
format!("failed to flush local stream rewrite: {err:?}"),
502,
)
});
}
}
}
if !downstream_dropped {
if let Some(rewriter) = local_stream_rewriter.as_mut() {
match rewriter.finish() {
Ok(flushed_chunk) if !flushed_chunk.is_empty() => {
append_stream_capture_bytes(
&mut buffered_body,
&flushed_chunk,
max_stream_body_buffer_bytes,
&mut client_body_truncated,
);
let flushed_chunk_len =
u64::try_from(flushed_chunk.len()).unwrap_or(u64::MAX);
let chunk_completed_stream =
stream_chunk_contains_sse_done(&flushed_chunk);
if tx.send(Ok(Bytes::from(flushed_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_rewrite_flush_disconnected",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
"gateway stream downstream dropped while flushing local stream rewrite"
);
downstream_dropped = true;
} else {
client_visible_stream_completed |= chunk_completed_stream;
client_stream_bytes.fetch_add(flushed_chunk_len, Ordering::Relaxed);
last_client_chunk_elapsed_ms.store(
stream_started_at_for_report
.elapsed()
.as_millis()
.min(u128::from(u64::MAX))
as u64,
Ordering::Relaxed,
);
}
}
Ok(_) => {}
Err(err) => {
warn!(
event_name = "stream_execution_rewrite_flush_failed",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
error = ?err,
"gateway failed to flush local stream rewrite"
);
terminal_failure.get_or_insert_with(|| {
build_stream_failure_report(
"execution_runtime_stream_rewrite_flush_error",
format!("failed to flush local stream rewrite: {err:?}"),
502,
)
});
}
}
}
}
}
if !downstream_dropped {
@@ -3527,7 +3501,7 @@ async fn execute_stream_from_frame_stream(
None
},
error_message: stream_failed
.then(|| stream_terminal_error_message)
.then_some(stream_terminal_error_message)
.flatten(),
latency_ms: usage_payload
.telemetry
@@ -4655,7 +4629,7 @@ mod tests {
}
#[tokio::test]
async fn execute_stream_from_frame_stream_stops_upstream_when_client_drops_body() {
async fn execute_stream_from_frame_stream_drains_upstream_when_client_drops_body() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = AppState::new()
@@ -4698,24 +4672,22 @@ mod tests {
transport_profile: None,
timeouts: None,
};
let frame_stream_dropped = Arc::new(Notify::new());
let frame_stream_dropped_for_stream = Arc::clone(&frame_stream_dropped);
let release_terminal = Arc::new(Notify::new());
let terminal_frame_drained = Arc::new(Notify::new());
let release_terminal_for_stream = Arc::clone(&release_terminal);
let terminal_frame_drained_for_stream = Arc::clone(&terminal_frame_drained);
let frame_stream = stream! {
struct NotifyOnDrop(Arc<Notify>);
impl Drop for NotifyOnDrop {
fn drop(&mut self) {
self.0.notify_waiters();
}
}
let _drop_guard = NotifyOnDrop(frame_stream_dropped_for_stream);
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
));
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"first\\\"}\\n\\n\"}}\n",
));
std::future::pending::<()>().await;
release_terminal_for_stream.notified().await;
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"terminal\\\",\\\"object\\\":\\\"chat.completion.chunk\\\",\\\"model\\\":\\\"gpt-5.4\\\",\\\"choices\\\":[{\\\"index\\\":0,\\\"delta\\\":{},\\\"finish_reason\\\":\\\"stop\\\"}],\\\"usage\\\":{\\\"prompt_tokens\\\":7,\\\"completion_tokens\\\":11,\\\"total_tokens\\\":18}}\\n\\ndata: [DONE]\\n\\n\"}}\n",
));
terminal_frame_drained_for_stream.notify_one();
}
.boxed();
@@ -4761,10 +4733,11 @@ mod tests {
assert_eq!(first.as_ref(), b"data: {\"id\":\"first\"}\n\n");
tokio::time::sleep(Duration::from_millis(30)).await;
drop(body_stream);
release_terminal.notify_one();
tokio::time::timeout(Duration::from_secs(1), frame_stream_dropped.notified())
tokio::time::timeout(Duration::from_secs(1), terminal_frame_drained.notified())
.await
.expect("upstream frame stream should be dropped after client disconnect");
.expect("upstream frame stream should be drained after client disconnect");
let candidates = tokio::time::timeout(Duration::from_secs(1), async {
loop {
let candidates = request_candidate_repository
@@ -4807,6 +4780,9 @@ mod tests {
.expect("usage should be marked cancelled");
assert_eq!(stored_usage.billing_status, "pending");
assert_eq!(stored_usage.status_code, Some(499));
assert_eq!(stored_usage.input_tokens, 7);
assert_eq!(stored_usage.output_tokens, 11);
assert_eq!(stored_usage.total_tokens, 18);
let first_byte_time_ms = stored_usage
.first_byte_time_ms
.expect("cancelled stream should retain first byte time");

416
crates/aether-usage-runtime/src/write.rs
View File

@@ -628,10 +628,6 @@ fn build_terminal_usage_event_from_seed_impl(
apply_completed_image_usage_estimate(&mut data);
}
if matches!(event_type, UsageEventType::Cancelled) {
apply_cancelled_usage_estimate(&mut data);
}
let data = if trusted_request_metadata {
sanitize_usage_event_capture_fields_trusted(data)
} else {
@@ -2321,48 +2317,6 @@ fn extract_token_counts_from_value(value: &Value) -> Option<(u64, u64, u64)> {
}
}
fn apply_cancelled_usage_estimate(data: &mut UsageEventData) {
let provider_usage_available = data
.response_body
.as_ref()
.and_then(extract_token_counts_from_value)
.is_some();
let request_usage = data
.provider_request_body
.as_ref()
.or(data.request_body.as_ref())
.and_then(estimate_request_usage);
if positive_tokens(data.input_tokens) == 0 {
if let Some(usage) = request_usage.as_ref() {
data.input_tokens = Some(usage.input_tokens);
}
}
if !provider_usage_available {
apply_cancelled_request_cache_estimate(data, request_usage.as_ref());
}
if positive_tokens(data.output_tokens) == 0 {
if let Some(output_tokens) = data
.response_body
.as_ref()
.or(data.client_response_body.as_ref())
.and_then(estimate_response_output_tokens)
{
data.output_tokens = Some(output_tokens);
}
}
if positive_tokens(data.total_tokens) == 0 {
let total_tokens =
positive_tokens(data.input_tokens).saturating_add(positive_tokens(data.output_tokens));
if total_tokens > 0 {
data.total_tokens = Some(total_tokens);
}
}
}
fn apply_completed_image_usage_estimate(data: &mut UsageEventData) {
if !usage_event_data_is_image(data) {
return;
@@ -2387,7 +2341,7 @@ fn apply_completed_image_usage_estimate(data: &mut UsageEventData) {
data.input_tokens = Some(usage.input_tokens);
}
}
apply_cancelled_request_cache_estimate(data, request_usage.as_ref());
apply_request_cache_usage_estimate(data, request_usage.as_ref());
if positive_tokens(data.total_tokens) == 0 {
let total_tokens =
positive_tokens(data.input_tokens).saturating_add(positive_tokens(data.output_tokens));
@@ -2525,7 +2479,7 @@ fn usage_event_data_is_image(data: &UsageEventData) -> bool {
.is_some_and(|value| value.eq_ignore_ascii_case("image"))
}
fn apply_cancelled_request_cache_estimate(
fn apply_request_cache_usage_estimate(
data: &mut UsageEventData,
request_usage: Option<&EstimatedRequestUsage>,
) {
@@ -2692,280 +2646,6 @@ fn estimate_text_tokens(text: &str) -> u64 {
}
}
#[derive(Default)]
struct StreamOutputEstimate {
text: String,
saw_delta: bool,
}
impl StreamOutputEstimate {
fn push_delta(&mut self, text: &str) {
if text.is_empty() {
return;
}
self.saw_delta = true;
self.text.push_str(text);
}
fn push_done(&mut self, text: &str) {
if text.is_empty() || self.saw_delta {
return;
}
self.text.push_str(text);
}
}
fn estimate_response_output_tokens(value: &Value) -> Option<u64> {
let mut estimate = StreamOutputEstimate::default();
collect_stream_output_text(value, &mut estimate);
let tokens = estimate_text_tokens(estimate.text.as_str());
(tokens > 0).then_some(tokens)
}
fn collect_stream_output_text(value: &Value, estimate: &mut StreamOutputEstimate) {
match value {
Value::String(text) => {
for_each_sse_payload(text, |payload| {
if payload == "[DONE]" {
return;
}
if let Ok(json_body) = serde_json::from_str::<Value>(payload) {
collect_stream_output_text(&json_body, estimate);
}
});
}
Value::Array(items) => {
for item in items {
collect_stream_output_text(item, estimate);
}
}
Value::Object(object) => {
if let Some(chunks) = object.get("chunks").and_then(Value::as_array) {
for chunk in chunks {
collect_stream_output_text(chunk, estimate);
}
return;
}
collect_openai_responses_output_text(object, estimate);
collect_openai_chat_output_text(object, estimate);
collect_claude_output_text(object, estimate);
collect_gemini_output_text(object, estimate);
}
_ => {}
}
}
fn collect_openai_responses_output_text(
object: &Map<String, Value>,
estimate: &mut StreamOutputEstimate,
) {
match object
.get("type")
.and_then(Value::as_str)
.unwrap_or_default()
{
"response.output_text.delta" | "response.outtext.delta" => {
if let Some(text) = openai_delta_text(object.get("delta")) {
estimate.push_delta(text.as_str());
}
}
"response.reasoning_summary_text.delta" | "response.function_call_arguments.delta" => {
if let Some(text) = object.get("delta").and_then(Value::as_str) {
estimate.push_delta(text);
}
}
"response.output_text.done" | "response.reasoning_summary_text.done" => {
if let Some(text) = object
.get("text")
.and_then(Value::as_str)
.or_else(|| part_text(object.get("part")))
{
estimate.push_done(text);
}
}
"response.function_call_arguments.done" => {
if let Some(text) = object.get("arguments").and_then(Value::as_str) {
estimate.push_done(text);
}
}
"response.output_item.done" => {
if let Some(item) = object.get("item").and_then(Value::as_object) {
collect_openai_responses_output_item_text(item, estimate);
}
}
"response.completed" => {
if let Some(response) = object.get("response").and_then(Value::as_object) {
collect_openai_responses_completed_text(response, estimate);
}
}
_ => {}
}
}
fn collect_openai_responses_completed_text(
response: &Map<String, Value>,
estimate: &mut StreamOutputEstimate,
) {
for item in response
.get("output")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
{
collect_openai_responses_output_item_text(item, estimate);
}
}
fn collect_openai_responses_output_item_text(
item: &Map<String, Value>,
estimate: &mut StreamOutputEstimate,
) {
match item.get("type").and_then(Value::as_str).unwrap_or_default() {
"message" => {
for content in item
.get("content")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
{
if content.get("type").and_then(Value::as_str) == Some("output_text") {
if let Some(text) = content.get("text").and_then(Value::as_str) {
estimate.push_done(text);
}
}
}
}
"reasoning" => {
for summary in item
.get("summary")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
{
if let Some(text) = summary.get("text").and_then(Value::as_str) {
estimate.push_done(text);
}
}
}
"function_call" => {
if let Some(arguments) = item.get("arguments").and_then(Value::as_str) {
estimate.push_done(arguments);
}
}
_ => {}
}
}
fn collect_openai_chat_output_text(
object: &Map<String, Value>,
estimate: &mut StreamOutputEstimate,
) {
for choice in object
.get("choices")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
{
if let Some(delta) = choice.get("delta").and_then(Value::as_object) {
if let Some(content) = delta.get("content").and_then(Value::as_str) {
estimate.push_delta(content);
}
if let Some(reasoning_content) = delta.get("reasoning_content").and_then(Value::as_str)
{
estimate.push_delta(reasoning_content);
}
for tool_call in delta
.get("tool_calls")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
{
if let Some(arguments) = tool_call
.get("function")
.and_then(Value::as_object)
.and_then(|function| function.get("arguments"))
.and_then(Value::as_str)
{
estimate.push_delta(arguments);
}
}
}
}
}
fn collect_claude_output_text(object: &Map<String, Value>, estimate: &mut StreamOutputEstimate) {
if object.get("type").and_then(Value::as_str) != Some("content_block_delta") {
return;
}
let Some(delta) = object.get("delta").and_then(Value::as_object) else {
return;
};
match delta
.get("type")
.and_then(Value::as_str)
.unwrap_or_default()
{
"text_delta" => {
if let Some(text) = delta.get("text").and_then(Value::as_str) {
estimate.push_delta(text);
}
}
"thinking_delta" => {
if let Some(text) = delta.get("thinking").and_then(Value::as_str) {
estimate.push_delta(text);
}
}
"input_json_delta" => {
if let Some(text) = delta.get("partial_json").and_then(Value::as_str) {
estimate.push_delta(text);
}
}
_ => {}
}
}
fn collect_gemini_output_text(object: &Map<String, Value>, estimate: &mut StreamOutputEstimate) {
for part in object
.get("candidates")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(|candidate| candidate.get("content"))
.filter_map(Value::as_object)
.filter_map(|content| content.get("parts"))
.filter_map(Value::as_array)
.flatten()
.filter_map(Value::as_object)
{
if let Some(text) = part.get("text").and_then(Value::as_str) {
estimate.push_delta(text);
}
}
}
fn openai_delta_text(value: Option<&Value>) -> Option<String> {
match value {
Some(Value::String(text)) => Some(text.clone()),
Some(Value::Object(object)) => object
.get("text")
.and_then(Value::as_str)
.map(ToOwned::to_owned),
_ => None,
}
}
fn part_text(value: Option<&Value>) -> Option<&str> {
value
.and_then(Value::as_object)
.and_then(|part| part.get("text"))
.and_then(Value::as_str)
}
fn extract_token_counts_from_json(value: &Value) -> Option<(u64, u64, u64)> {
if let Some(usage) = value.get("usage").and_then(Value::as_object) {
let input = usage
@@ -3543,7 +3223,7 @@ mod tests {
}
#[test]
fn cancelled_stream_usage_estimates_tokens_from_request_and_partial_response() {
fn cancelled_stream_usage_does_not_estimate_tokens_from_request_or_partial_response() {
let plan = ExecutionPlan {
request_id: "req-stream-cancelled-estimated-usage-1".to_string(),
candidate_id: Some("cand-stream-cancelled-estimated-usage-1".to_string()),
@@ -3602,16 +3282,14 @@ mod tests {
.expect("usage event should build");
assert_eq!(event.event_type, UsageEventType::Cancelled);
assert!(event.data.input_tokens.unwrap_or_default() > 0);
assert_eq!(event.data.output_tokens, Some(5));
assert_eq!(
event.data.total_tokens,
Some(event.data.input_tokens.unwrap_or_default() + 5)
);
assert_eq!(event.data.input_tokens, None);
assert_eq!(event.data.output_tokens, None);
assert_eq!(event.data.total_tokens, None);
assert_eq!(event.data.cache_read_input_tokens, None);
}
#[test]
fn cancelled_stream_usage_does_not_infer_cache_read_from_prompt_cache_key() {
fn cancelled_stream_usage_does_not_infer_cache_or_token_estimates_from_prompt_cache_key() {
let request_body = json!({
"model": "gpt-5.4",
"input": "Use the cached project context and answer briefly",
@@ -3662,15 +3340,81 @@ mod tests {
let event =
build_stream_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload)
.expect("usage event should build");
let input_tokens = event
.data
.input_tokens
.expect("input estimate should exist");
assert_eq!(event.event_type, UsageEventType::Cancelled);
assert_eq!(event.data.input_tokens, None);
assert_eq!(event.data.output_tokens, None);
assert_eq!(event.data.total_tokens, None);
assert_eq!(event.data.cache_read_input_tokens, None);
assert_eq!(event.data.output_tokens, Some(4));
assert_eq!(event.data.total_tokens, Some(input_tokens + 4));
}
#[test]
fn cancelled_stream_usage_preserves_terminal_summary_usage() {
let plan = ExecutionPlan {
request_id: "req-stream-cancelled-summary-usage-1".to_string(),
candidate_id: Some("cand-stream-cancelled-summary-usage-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/responses".to_string(),
headers: BTreeMap::new(),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gpt-5.4",
"input": "This cancelled request has terminal upstream usage",
"stream": true
})),
stream: true,
client_api_format: "openai:responses".to_string(),
provider_api_format: "openai:responses".to_string(),
model_name: Some("gpt-5.4".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let mut standardized_usage = StandardizedUsage::new();
standardized_usage.input_tokens = 13;
standardized_usage.output_tokens = 21;
standardized_usage.cache_creation_tokens = 2;
standardized_usage.cache_read_tokens = 3;
let payload = GatewayStreamReportRequest {
trace_id: "trace-stream-cancelled-summary-usage-1".to_string(),
report_kind: "openai_responses_stream_cancelled".to_string(),
report_context: Some(json!({
"client_api_format": "openai:responses",
"provider_api_format": "openai:responses"
})),
status_code: 499,
headers: BTreeMap::new(),
provider_body_base64: None,
provider_body_state: Some(UsageBodyCaptureState::None),
client_body_base64: None,
client_body_state: Some(UsageBodyCaptureState::None),
terminal_summary: Some(ExecutionStreamTerminalSummary {
standardized_usage: Some(standardized_usage),
finish_reason: None,
response_id: Some("resp_cancel_summary_1".to_string()),
model: Some("gpt-5.4".to_string()),
observed_finish: true,
unknown_event_count: 0,
parser_error: None,
}),
telemetry: None,
};
let event =
build_stream_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload)
.expect("usage event should build");
assert_eq!(event.event_type, UsageEventType::Cancelled);
assert_eq!(event.data.input_tokens, Some(13));
assert_eq!(event.data.output_tokens, Some(21));
assert_eq!(event.data.total_tokens, Some(34));
assert_eq!(event.data.cache_creation_input_tokens, Some(2));
assert_eq!(event.data.cache_read_input_tokens, Some(3));
}
#[test]