Merge remote-tracking branch 'origin/pr/613'

# Conflicts:
#	apps/aether-gateway/src/tests/usage/direct.rs
This commit is contained in:
elky
2026-06-02 19:27:50 +08:00
6 changed files with 623 additions and 169 deletions
@@ -3357,7 +3357,13 @@ async fn execute_stream_from_frame_stream(
"gateway skipped client stream flush after downstream disconnect" "gateway skipped client stream flush after downstream disconnect"
); );
} }
if let Some(normalizer) = private_stream_normalizer.as_mut() { // Buffered stream state is partial after a terminal failure; normal
// finish paths may synthesize successful terminal events.
let should_finish_stream_rewriters = terminal_failure.is_none();
if let Some(normalizer) = private_stream_normalizer
.as_mut()
.filter(|_| should_finish_stream_rewriters)
{
match normalizer.finish() { match normalizer.finish() {
Ok(normalized_chunk) if !normalized_chunk.is_empty() => { Ok(normalized_chunk) if !normalized_chunk.is_empty() => {
let provider_private_error_body_json = let provider_private_error_body_json =
@@ -3391,13 +3397,12 @@ async fn execute_stream_from_frame_stream(
error = ?err, error = ?err,
"gateway failed to rewrite normalized private stream chunk during flush" "gateway failed to rewrite normalized private stream chunk during flush"
); );
terminal_failure.get_or_insert_with(|| { let failure = build_stream_failure_report(
build_stream_failure_report( "execution_runtime_stream_rewrite_flush_error",
"execution_runtime_stream_rewrite_flush_error", format!("failed to rewrite normalized private stream chunk during flush: {err:?}"),
format!("failed to rewrite normalized private stream chunk during flush: {err:?}"), 502,
502, );
) terminal_failure.get_or_insert(failure);
});
Vec::new() Vec::new()
} }
} }
@@ -3472,7 +3477,7 @@ async fn execute_stream_from_frame_stream(
} }
} }
} }
if !downstream_dropped { if !downstream_dropped && terminal_failure.is_none() {
if let Some(rewriter) = local_stream_rewriter.as_mut() { if let Some(rewriter) = local_stream_rewriter.as_mut() {
match rewriter.finish() { match rewriter.finish() {
Ok(flushed_chunk) if !flushed_chunk.is_empty() => { Ok(flushed_chunk) if !flushed_chunk.is_empty() => {
@@ -3898,8 +3903,9 @@ mod tests {
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use aether_contracts::{ use aether_contracts::{
ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTimeouts, RequestBody, ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan,
StandardizedUsage, ExecutionStreamTerminalSummary, ExecutionTimeouts, RequestBody, StandardizedUsage,
StreamFrame, StreamFramePayload, StreamFrameType,
}; };
use aether_data::repository::candidates::InMemoryRequestCandidateRepository; use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::usage::InMemoryUsageReadRepository; use aether_data::repository::usage::InMemoryUsageReadRepository;
@@ -3996,6 +4002,12 @@ mod tests {
out out
} }
fn ndjson_frame(frame: StreamFrame) -> Bytes {
let mut bytes = serde_json::to_vec(&frame).expect("stream frame should serialize");
bytes.push(b'\n');
Bytes::from(bytes)
}
#[test] #[test]
fn merge_stream_terminal_summary_prefers_more_complete_observed_usage() { fn merge_stream_terminal_summary_prefers_more_complete_observed_usage() {
let mut runtime_usage = StandardizedUsage::new(); let mut runtime_usage = StandardizedUsage::new();
@@ -5017,6 +5029,154 @@ mod tests {
assert_eq!(first.as_ref(), b": aether-keepalive\n\n"); assert_eq!(first.as_ref(), b": aether-keepalive\n\n");
} }
#[tokio::test]
async fn execute_stream_from_frame_stream_does_not_finalize_rewritten_tool_call_after_midstream_error(
) {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = AppState::new()
.expect("app 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()
});
let plan = ExecutionPlan {
request_id: "req-responses-tool-midstream-error".into(),
candidate_id: Some("cand-responses-tool-midstream-error".into()),
provider_name: Some("openai".into()),
provider_id: "provider-openai-responses".into(),
endpoint_id: "endpoint-openai-responses".into(),
key_id: "key-openai-responses".into(),
method: "POST".into(),
url: "https://api.openai.com/v1/responses".into(),
headers: BTreeMap::from([
("content-type".into(), "application/json".into()),
("accept".into(), "text/event-stream".into()),
]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gpt-5.5",
"input": [],
"stream": true
})),
stream: true,
client_api_format: "claude:messages".into(),
provider_api_format: "openai:responses".into(),
model_name: Some("gpt-5.5".into()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let upstream_chunk = concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_midstream_error\",\"model\":\"gpt-5.5\",\"status\":\"in_progress\"}}\n\n",
"event: response.output_item.added\n",
"data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"id\":\"fc_1\",\"call_id\":\"call_1\",\"name\":\"lookup\",\"arguments\":\"\",\"status\":\"in_progress\"}}\n\n",
"event: response.function_call_arguments.delta\n",
"data: {\"type\":\"response.function_call_arguments.delta\",\"output_index\":0,\"item_id\":\"fc_1\",\"call_id\":\"call_1\",\"delta\":\"{\\\"query\\\":\\\"abc\"}\n\n"
);
let frame_stream = stream! {
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
frame_type: StreamFrameType::Headers,
payload: StreamFramePayload::Headers {
status_code: 200,
headers: BTreeMap::from([(
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
frame_type: StreamFrameType::Data,
payload: StreamFramePayload::Data {
chunk_b64: None,
text: Some(upstream_chunk.to_string()),
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
frame_type: StreamFrameType::Error,
payload: StreamFramePayload::Error {
error: ExecutionError {
kind: ExecutionErrorKind::Internal,
phase: ExecutionPhase::StreamRead,
message: "error reading a body from connection: stream error received: unexpected internal error encountered".to_string(),
upstream_status: Some(200),
retryable: false,
failover_recommended: false,
},
},
}));
}
.boxed();
let response = execute_stream_from_frame_stream(
&state,
plan,
"trace-responses-tool-midstream-error",
&test_decision(),
"openai_responses_stream",
Some("openai_responses_stream_success".to_string()),
Some(json!({
"request_id": "req-responses-tool-midstream-error",
"candidate_id": "cand-responses-tool-midstream-error",
"candidate_index": 0,
"retry_index": 0,
"provider_api_format": "openai:responses",
"client_api_format": "claude:messages",
"needs_conversion": true,
})),
crate::clock::current_unix_ms(),
Instant::now(),
frame_stream,
None,
)
.await
.expect("execution should succeed")
.expect("execution should return a client response");
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body should read");
let body_text = String::from_utf8(body.to_vec()).expect("body should be utf8");
assert!(body_text.contains("event: content_block_start"));
assert!(body_text.contains("event: content_block_delta"));
assert!(body_text.contains("\"type\":\"tool_use\""));
assert!(!body_text.contains("event: content_block_stop"));
assert!(!body_text.contains("event: message_delta"));
assert!(!body_text.contains("event: message_stop"));
assert!(!body_text.contains("\"stop_reason\":\"tool_use\""));
assert!(body_text.contains("\"error\""));
assert!(body_text.contains("unexpected internal error encountered"));
assert!(body_text.contains("data: [DONE]"));
let candidates = tokio::time::timeout(Duration::from_secs(1), async {
loop {
let candidates = request_candidate_repository
.list_by_request_id("req-responses-tool-midstream-error")
.await
.expect("request candidates should read");
if candidates
.first()
.is_some_and(|candidate| candidate.status == RequestCandidateStatus::Failed)
{
break candidates;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("candidate should be marked failed");
assert_eq!(candidates[0].status_code, Some(200));
assert_eq!(candidates[0].error_type.as_deref(), Some("internal"));
}
#[tokio::test] #[tokio::test]
async fn openai_image_stream_ignores_plan_total_timeout() { async fn openai_image_stream_ignores_plan_total_timeout() {
let state = AppState::new().expect("app state should build"); let state = AppState::new().expect("app state should build");
+128 -125
View File
@@ -154,74 +154,79 @@ async fn gateway_records_usage_for_execution_runtime_sync_when_runtime_enabled()
#[test] #[test]
fn gateway_records_pending_usage_before_execution_runtime_sync_result_arrives() { fn gateway_records_pending_usage_before_execution_runtime_sync_result_arrives() {
run_async_test_on_large_stack("pending-usage-sync-before-runtime-result", async move { run_async_test_on_large_stack(
let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); "gateway_records_pending_usage_before_execution_runtime_sync_result_arrives",
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); gateway_records_pending_usage_before_execution_runtime_sync_result_arrives_impl(),
let execution_request_started = Arc::new(tokio::sync::Notify::new()); );
let allow_execution_response = Arc::new(tokio::sync::Notify::new()); }
let upstream = Router::new().route( async fn gateway_records_pending_usage_before_execution_runtime_sync_result_arrives_impl() {
"/api/internal/gateway/report-sync", let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
any(|_request: Request| async move { Json(json!({"ok": true})) }), let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
); let execution_request_started = Arc::new(tokio::sync::Notify::new());
let allow_execution_response = Arc::new(tokio::sync::Notify::new());
let execution_runtime = Router::new().route( let upstream = Router::new().route(
"/v1/execute/sync", "/api/internal/gateway/report-sync",
any({ any(|_request: Request| async move { Json(json!({"ok": true})) }),
);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any({
let execution_request_started = Arc::clone(&execution_request_started);
let allow_execution_response = Arc::clone(&allow_execution_response);
move |_request: Request| {
let execution_request_started = Arc::clone(&execution_request_started); let execution_request_started = Arc::clone(&execution_request_started);
let allow_execution_response = Arc::clone(&allow_execution_response); let allow_execution_response = Arc::clone(&allow_execution_response);
move |_request: Request| { async move {
let execution_request_started = Arc::clone(&execution_request_started); execution_request_started.notify_one();
let allow_execution_response = Arc::clone(&allow_execution_response); allow_execution_response.notified().await;
async move { Json(json!({
execution_request_started.notify_one(); "request_id": "req-usage-sync-pending-123",
allow_execution_response.notified().await; "status_code": 200,
Json(json!({ "headers": {
"request_id": "req-usage-sync-pending-123", "content-type": "application/json"
"status_code": 200, },
"headers": { "body": {
"content-type": "application/json" "json_body": {
}, "id": "chatcmpl-usage-sync-pending-123",
"body": { "usage": {
"json_body": { "input_tokens": 3,
"id": "chatcmpl-usage-sync-pending-123", "output_tokens": 5,
"usage": { "total_tokens": 8
"input_tokens": 3,
"output_tokens": 5,
"total_tokens": 8
}
} }
},
"telemetry": {
"elapsed_ms": 45
} }
})) },
} "telemetry": {
"elapsed_ms": 45
}
}))
} }
}), }
); }),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-usage-sync-pending")), Some(hash_api_key("sk-client-openai-usage-sync-pending")),
sample_local_openai_auth_snapshot( sample_local_openai_auth_snapshot(
"api-key-usage-sync-pending-123", "api-key-usage-sync-pending-123",
"user-usage-sync-pending-123", "user-usage-sync-pending-123",
), ),
)])); )]));
let candidate_selection_repository = let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_local_openai_candidate_row(), sample_local_openai_candidate_row(),
])); ]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_local_openai_provider()], vec![sample_local_openai_provider()],
vec![sample_local_openai_endpoint()], vec![sample_local_openai_endpoint()],
vec![sample_local_openai_key()], vec![sample_local_openai_key()],
)); ));
let (upstream_url, upstream_handle) = start_server(upstream).await; let (upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
start_server(execution_runtime).await; let gateway_state =
let gateway_state =
build_state_with_execution_runtime_override(execution_runtime_url) build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests( .with_data_state_for_tests(
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests( GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
@@ -237,76 +242,74 @@ fn gateway_records_pending_usage_before_execution_runtime_sync_result_arrives()
enabled: true, enabled: true,
..UsageRuntimeConfig::default() ..UsageRuntimeConfig::default()
}); });
let gateway = build_router_with_state(gateway_state); let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await; let (gateway_url, gateway_handle) = start_server(gateway).await;
let request_task = tokio::spawn({ let request_task = tokio::spawn({
let gateway_url = gateway_url.clone(); let gateway_url = gateway_url.clone();
async move { async move {
let response = reqwest::Client::new() let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions")) .post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json") .header(http::header::CONTENT_TYPE, "application/json")
.header( .header(
http::header::AUTHORIZATION, http::header::AUTHORIZATION,
"Bearer sk-client-openai-usage-sync-pending", "Bearer sk-client-openai-usage-sync-pending",
) )
.header(TRACE_ID_HEADER, "req-usage-sync-pending-123") .header(TRACE_ID_HEADER, "req-usage-sync-pending-123")
.body("{\"model\":\"gpt-5\",\"messages\":[]}") .body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send() .send()
.await
.expect("request should succeed");
let status = response.status();
let body = response.text().await.expect("body should read");
(status, body)
}
});
execution_request_started.notified().await;
let mut pending = None;
for _ in 0..50 {
pending = usage_repository
.find_by_request_id("req-usage-sync-pending-123")
.await .await
.expect("usage lookup should succeed"); .expect("request should succeed");
if pending let status = response.status();
.as_ref() let body = response.text().await.expect("body should read");
.is_some_and(|stored| stored.status == "pending") (status, body)
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
} }
let pending =
pending.expect("pending usage should be recorded before sync result resolves");
assert_eq!(pending.status, "pending");
assert_eq!(pending.billing_status, "pending");
assert_eq!(pending.response_time_ms, None);
allow_execution_response.notify_one();
let (status, _body) = request_task.await.expect("request task should join");
assert_eq!(status, StatusCode::OK);
let mut stored = None;
for _ in 0..50 {
stored = usage_repository
.find_by_request_id("req-usage-sync-pending-123")
.await
.expect("usage lookup should succeed");
if stored.as_ref().is_some_and(|row| row.status == "completed") {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let stored = stored.expect("usage should be finalized");
assert_eq!(stored.status, "completed");
assert_eq!(stored.response_time_ms, Some(45));
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}); });
execution_request_started.notified().await;
let mut pending = None;
for _ in 0..50 {
pending = usage_repository
.find_by_request_id("req-usage-sync-pending-123")
.await
.expect("usage lookup should succeed");
if pending
.as_ref()
.is_some_and(|stored| stored.status == "pending")
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let pending = pending.expect("pending usage should be recorded before sync result resolves");
assert_eq!(pending.status, "pending");
assert_eq!(pending.billing_status, "pending");
assert_eq!(pending.response_time_ms, None);
allow_execution_response.notify_one();
let (status, _body) = request_task.await.expect("request task should join");
assert_eq!(status, StatusCode::OK);
let mut stored = None;
for _ in 0..50 {
stored = usage_repository
.find_by_request_id("req-usage-sync-pending-123")
.await
.expect("usage lookup should succeed");
if stored.as_ref().is_some_and(|row| row.status == "completed") {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let stored = stored.expect("usage should be finalized");
assert_eq!(stored.status, "completed");
assert_eq!(stored.response_time_ms, Some(45));
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
} }
#[tokio::test] #[tokio::test]
+9 -2
View File
@@ -559,8 +559,15 @@ async fn gateway_applies_system_max_request_body_size_to_local_openai_chat_sync_
upstream_handle.abort(); upstream_handle.abort();
} }
#[tokio::test] #[test]
async fn gateway_strips_request_and_response_bodies_when_request_record_level_is_base() { fn gateway_strips_request_and_response_bodies_when_request_record_level_is_base() {
run_async_test_on_large_stack(
"gateway_strips_request_and_response_bodies_when_request_record_level_is_base",
gateway_strips_request_and_response_bodies_when_request_record_level_is_base_impl(),
);
}
async fn gateway_strips_request_and_response_bodies_when_request_record_level_is_base_impl() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
+30 -2
View File
@@ -11,8 +11,36 @@ use super::{
}; };
use aether_data::repository::settlement::InMemorySettlementRepository; use aether_data::repository::settlement::InMemorySettlementRepository;
#[tokio::test] fn run_async_test_on_large_stack<F>(name: &'static str, future: F)
async fn gateway_settles_wallet_for_completed_execution_runtime_sync_usage() { where
F: std::future::Future<Output = ()> + Send + 'static,
{
let handle = std::thread::Builder::new()
.name(name.to_string())
.stack_size(8 * 1024 * 1024)
.spawn(move || {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("tokio runtime should build")
.block_on(future);
})
.expect("large-stack test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
#[test]
fn gateway_settles_wallet_for_completed_execution_runtime_sync_usage() {
run_async_test_on_large_stack(
"gateway_settles_wallet_for_completed_execution_runtime_sync_usage",
gateway_settles_wallet_for_completed_execution_runtime_sync_usage_impl(),
);
}
async fn gateway_settles_wallet_for_completed_execution_runtime_sync_usage_impl() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let billing_repository = Arc::new(InMemoryBillingReadRepository::seed(vec![ let billing_repository = Arc::new(InMemoryBillingReadRepository::seed(vec![
@@ -45,7 +45,7 @@ pub struct OpenAIResponsesProviderState {
model: Option<String>, model: Option<String>,
started: bool, started: bool,
finished: bool, finished: bool,
text: String, text_parts: BTreeMap<String, String>,
reasoning: String, reasoning: String,
reasoning_parts: BTreeMap<usize, String>, reasoning_parts: BTreeMap<usize, String>,
tool_calls: BTreeMap<usize, OpenAIResponsesProviderToolState>, tool_calls: BTreeMap<usize, OpenAIResponsesProviderToolState>,
@@ -420,24 +420,87 @@ impl OpenAIResponsesProviderState {
index index
} }
fn text_part_key_from_event(value: &Value) -> String {
let item_key = value
.get("output_index")
.and_then(Value::as_u64)
.map(|value| format!("output:{value}"))
.or_else(|| {
value
.get("item_id")
.or_else(|| value.get("id"))
.and_then(Value::as_str)
.map(|value| format!("item:{value}"))
})
.unwrap_or_else(|| "output:default".to_string());
let content_index = value
.get("content_index")
.and_then(Value::as_u64)
.unwrap_or(0);
format!("{item_key}:content:{content_index}")
}
fn text_part_key_from_message_item(
output_index: Option<usize>,
item: &Map<String, Value>,
content_index: usize,
) -> String {
let item_key = output_index
.map(|value| format!("output:{value}"))
.or_else(|| {
item.get("id")
.and_then(Value::as_str)
.map(|value| format!("item:{value}"))
})
.unwrap_or_else(|| "output:default".to_string());
format!("{item_key}:content:{content_index}")
}
fn emit_text_delta(
&mut self,
report_context: &Value,
out: &mut Vec<CanonicalStreamFrame>,
key: String,
text: &str,
) {
if text.is_empty() {
return;
}
self.text_parts.entry(key).or_default().push_str(text);
self.ensure_started(report_context, out);
let (id, model) = self.identity(report_context);
out.push(CanonicalStreamFrame {
id,
model,
event: CanonicalStreamEvent::TextDelta(text.to_string()),
});
}
fn emit_missing_text( fn emit_missing_text(
&mut self, &mut self,
report_context: &Value, report_context: &Value,
out: &mut Vec<CanonicalStreamFrame>, out: &mut Vec<CanonicalStreamFrame>,
key: String,
text: &str, text: &str,
) { ) {
let missing = if text.starts_with(&self.text) { let missing = {
text[self.text.len()..].to_string() let current = self.text_parts.entry(key).or_default();
} else if self.text == text || self.text.starts_with(text) { let missing = if text.starts_with(current.as_str()) {
String::new() text[current.len()..].to_string()
} else { } else if current.as_str() == text || current.starts_with(text) {
text.to_string() String::new()
} else {
text.to_string()
};
if !missing.is_empty() {
current.push_str(&missing);
}
missing
}; };
if missing.is_empty() { if missing.is_empty() {
return; return;
} }
self.ensure_started(report_context, out); self.ensure_started(report_context, out);
self.text.push_str(&missing);
let (id, model) = self.identity(report_context); let (id, model) = self.identity(report_context);
out.push(CanonicalStreamFrame { out.push(CanonicalStreamFrame {
id, id,
@@ -695,29 +758,34 @@ impl OpenAIResponsesProviderState {
report_context: &Value, report_context: &Value,
out: &mut Vec<CanonicalStreamFrame>, out: &mut Vec<CanonicalStreamFrame>,
item: &Map<String, Value>, item: &Map<String, Value>,
output_index: Option<usize>,
) { ) {
if item.get("type").and_then(Value::as_str) != Some("message") { if item.get("type").and_then(Value::as_str) != Some("message") {
return; return;
} }
let mut completed_text = String::new(); for (content_index, raw_content) in item
for raw_content in item
.get("content") .get("content")
.and_then(Value::as_array) .and_then(Value::as_array)
.into_iter() .into_iter()
.flatten() .flatten()
.enumerate()
{ {
let Some(content) = raw_content.as_object() else { let Some(content) = raw_content.as_object() else {
continue; continue;
}; };
if content.get("type").and_then(Value::as_str) == Some("output_text") { if content.get("type").and_then(Value::as_str) == Some("output_text") {
if let Some(text) = content.get("text").and_then(Value::as_str) { if let Some(text) = content.get("text").and_then(Value::as_str) {
completed_text.push_str(text); if !text.is_empty() {
let key = Self::text_part_key_from_message_item(
output_index,
item,
content_index,
);
self.emit_missing_text(report_context, out, key, text);
}
} }
} }
} }
if !completed_text.is_empty() {
self.emit_missing_text(report_context, out, &completed_text);
}
} }
fn emit_reasoning_item( fn emit_reasoning_item(
@@ -831,18 +899,13 @@ impl OpenAIResponsesProviderState {
} }
"response.output_text.delta" | "response.outtext.delta" => match value.get("delta") { "response.output_text.delta" | "response.outtext.delta" => match value.get("delta") {
Some(Value::String(piece)) if !piece.is_empty() => { Some(Value::String(piece)) if !piece.is_empty() => {
self.ensure_started(report_context, &mut out); let key = Self::text_part_key_from_event(&value);
self.text.push_str(piece); self.emit_text_delta(report_context, &mut out, key, piece);
let (id, model) = self.identity(report_context);
out.push(CanonicalStreamFrame {
id,
model,
event: CanonicalStreamEvent::TextDelta(piece.clone()),
});
} }
Some(Value::Object(delta)) => { Some(Value::Object(delta)) => {
if let Some(text) = delta.get("text").and_then(Value::as_str) { if let Some(text) = delta.get("text").and_then(Value::as_str) {
self.emit_missing_text(report_context, &mut out, text); let key = Self::text_part_key_from_event(&value);
self.emit_missing_text(report_context, &mut out, key, text);
} }
} }
_ => {} _ => {}
@@ -852,7 +915,8 @@ impl OpenAIResponsesProviderState {
if part.get("type").and_then(Value::as_str) == Some("output_text") { if part.get("type").and_then(Value::as_str) == Some("output_text") {
if let Some(text) = part.get("text").and_then(Value::as_str) { if let Some(text) = part.get("text").and_then(Value::as_str) {
if !text.is_empty() { if !text.is_empty() {
self.emit_missing_text(report_context, &mut out, text); let key = Self::text_part_key_from_event(&value);
self.emit_missing_text(report_context, &mut out, key, text);
} }
} }
} }
@@ -890,7 +954,8 @@ impl OpenAIResponsesProviderState {
}) })
.unwrap_or_default(); .unwrap_or_default();
if !text.is_empty() { if !text.is_empty() {
self.emit_missing_text(report_context, &mut out, text); let key = Self::text_part_key_from_event(&value);
self.emit_missing_text(report_context, &mut out, key, text);
} }
} }
"response.reasoning_summary_text.delta" => { "response.reasoning_summary_text.delta" => {
@@ -967,7 +1032,7 @@ impl OpenAIResponsesProviderState {
self.emit_tool_result_item(report_context, &mut out, item, output_index); self.emit_tool_result_item(report_context, &mut out, item, output_index);
} }
"message" => { "message" => {
self.emit_message_item(report_context, &mut out, item); self.emit_message_item(report_context, &mut out, item, output_index);
} }
"reasoning" => { "reasoning" => {
self.ensure_started(report_context, &mut out); self.ensure_started(report_context, &mut out);
@@ -1139,7 +1204,7 @@ impl OpenAIResponsesProviderState {
self.emit_tool_result_item(report_context, &mut out, item, output_index); self.emit_tool_result_item(report_context, &mut out, item, output_index);
} }
"message" => { "message" => {
self.emit_message_item(report_context, &mut out, item); self.emit_message_item(report_context, &mut out, item, output_index);
} }
"reasoning" => { "reasoning" => {
self.emit_reasoning_item(report_context, &mut out, item); self.emit_reasoning_item(report_context, &mut out, item);
@@ -1194,7 +1259,12 @@ impl OpenAIResponsesProviderState {
}; };
match item.get("type").and_then(Value::as_str).unwrap_or_default() { match item.get("type").and_then(Value::as_str).unwrap_or_default() {
"message" => { "message" => {
self.emit_message_item(report_context, &mut out, item); self.emit_message_item(
report_context,
&mut out,
item,
Some(output_index),
);
} }
"function_call" => { "function_call" => {
self.emit_tool_call_item( self.emit_tool_call_item(
@@ -3235,6 +3305,102 @@ mod tests {
assert_eq!(text, "Hello world"); assert_eq!(text, "Hello world");
} }
#[test]
fn openai_responses_provider_state_dedupes_text_snapshots_per_output_item() {
let mut state = OpenAIResponsesProviderState::default();
let report_context = json!({});
let mut frames = Vec::new();
for event in [
json!({
"type": "response.output_text.delta",
"response_id": "resp_multi_message",
"output_index": 0,
"item_id": "msg_1",
"content_index": 0,
"delta": "First message.",
}),
json!({
"type": "response.output_text.done",
"response_id": "resp_multi_message",
"output_index": 0,
"item_id": "msg_1",
"content_index": 0,
"text": "First message.",
}),
json!({
"type": "response.output_item.done",
"response_id": "resp_multi_message",
"output_index": 0,
"item": {
"type": "message",
"id": "msg_1",
"status": "completed",
"content": [{
"type": "output_text",
"text": "First message.",
}],
},
}),
json!({
"type": "response.output_text.delta",
"response_id": "resp_multi_message",
"output_index": 1,
"item_id": "msg_2",
"content_index": 0,
"delta": "Second message.",
}),
json!({
"type": "response.output_text.done",
"response_id": "resp_multi_message",
"output_index": 1,
"item_id": "msg_2",
"content_index": 0,
"text": "Second message.",
}),
json!({
"type": "response.content_part.done",
"response_id": "resp_multi_message",
"output_index": 1,
"item_id": "msg_2",
"content_index": 0,
"part": {
"type": "output_text",
"text": "Second message.",
},
}),
json!({
"type": "response.output_item.done",
"response_id": "resp_multi_message",
"output_index": 1,
"item": {
"type": "message",
"id": "msg_2",
"status": "completed",
"content": [{
"type": "output_text",
"text": "Second message.",
}],
},
}),
] {
frames.extend(
state
.push_line(&report_context, data_line(event))
.expect("responses text event should parse"),
);
}
let text = frames
.iter()
.filter_map(|frame| match &frame.event {
CanonicalStreamEvent::TextDelta(text) => Some(text.as_str()),
_ => None,
})
.collect::<String>();
assert_eq!(text, "First message.Second message.");
}
#[test] #[test]
fn openai_responses_provider_state_delays_arguments_until_tool_name_is_known() { fn openai_responses_provider_state_delays_arguments_until_tool_name_is_known() {
let mut state = OpenAIResponsesProviderState::default(); let mut state = OpenAIResponsesProviderState::default();
+91 -1
View File
@@ -735,7 +735,10 @@ pub fn build_sync_terminal_usage_payload_seed(
.and_then(|context| context.get(UPSTREAM_IS_STREAM_KEY)) .and_then(|context| context.get(UPSTREAM_IS_STREAM_KEY))
.and_then(Value::as_bool) .and_then(Value::as_bool)
.unwrap_or(false); .unwrap_or(false);
let provider_response_full = if upstream_is_stream && payload.body_base64.is_some() { let provider_response_full = if upstream_is_stream
&& payload.body_base64.is_some()
&& !body_json_has_terminal_error(payload.body_json.as_ref())
{
decode_body_for_storage(payload.body_base64.as_deref()) decode_body_for_storage(payload.body_base64.as_deref())
.or_else(|| payload.body_json.as_ref().cloned()) .or_else(|| payload.body_json.as_ref().cloned())
} else { } else {
@@ -786,6 +789,12 @@ pub fn build_sync_terminal_usage_payload_seed(
} }
} }
fn body_json_has_terminal_error(body_json: Option<&Value>) -> bool {
body_json
.and_then(|value| value.get("error"))
.is_some_and(|error| !error.is_null())
}
pub fn build_stream_terminal_usage_payload_seed( pub fn build_stream_terminal_usage_payload_seed(
payload: &GatewayStreamReportRequest, payload: &GatewayStreamReportRequest,
) -> StreamTerminalUsagePayloadSeed { ) -> StreamTerminalUsagePayloadSeed {
@@ -5108,6 +5117,87 @@ mod tests {
); );
} }
#[test]
fn sync_terminal_usage_prefers_error_body_over_partial_upstream_stream_body() {
let partial_sse_body = concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_partial_123\",\"object\":\"response\",\"model\":\"gpt-5.5\",\"status\":\"in_progress\",\"output\":[]}}\n\n",
"event: response.output_item.added\n",
"data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"name\":\"exec_command\"}}\n\n",
"event: response.function_call_arguments.delta\n",
"data: {\"type\":\"response.function_call_arguments.delta\",\"delta\":\"{\\\"cmd\\\":\"}\n\n",
);
let plan = ExecutionPlan {
request_id: "req-sync-upstream-stream-error-1".to_string(),
candidate_id: Some("cand-sync-upstream-stream-error-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: None,
content_encoding: None,
body: RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
},
stream: false,
client_api_format: "claude:messages".to_string(),
provider_api_format: "openai:responses".to_string(),
model_name: Some("gpt-5.5".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let payload = GatewaySyncReportRequest {
trace_id: "trace-sync-upstream-stream-error-1".to_string(),
report_kind: "openai_responses_sync_error".to_string(),
report_context: Some(json!({
"client_api_format": "claude:messages",
"provider_api_format": "openai:responses",
"upstream_is_stream": true,
"needs_conversion": true
})),
status_code: 200,
headers: BTreeMap::from([(
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
body_json: Some(json!({
"error": {
"type": "internal",
"message": "error decoding response body: stream error received"
}
})),
client_body_json: None,
body_base64: Some(base64::engine::general_purpose::STANDARD.encode(partial_sse_body)),
telemetry: None,
};
let event =
build_sync_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload)
.expect("usage event should build");
assert_eq!(event.event_type, UsageEventType::Failed);
assert_eq!(event.data.status_code, Some(200));
assert_eq!(
event.data.error_message.as_deref(),
Some("error decoding response body: stream error received")
);
assert_eq!(
event
.data
.response_body
.as_ref()
.and_then(|value| value.pointer("/error/type"))
.and_then(Value::as_str),
Some("internal")
);
}
#[test] #[test]
fn sync_terminal_seed_path_matches_legacy_wrapper_event() { fn sync_terminal_seed_path_matches_legacy_wrapper_event() {
let plan = ExecutionPlan { let plan = ExecutionPlan {