diff --git a/apps/aether-gateway/src/ai_serving/finalize/internal/stream_rewrite.rs b/apps/aether-gateway/src/ai_serving/finalize/internal/stream_rewrite.rs index c1b8c0425..4355437ea 100644 --- a/apps/aether-gateway/src/ai_serving/finalize/internal/stream_rewrite.rs +++ b/apps/aether-gateway/src/ai_serving/finalize/internal/stream_rewrite.rs @@ -18,6 +18,12 @@ pub(crate) fn maybe_build_local_stream_rewriter<'a>( } impl LocalStreamRewriter<'_> { + pub(crate) fn into_owned(self) -> LocalStreamRewriter<'static> { + LocalStreamRewriter { + inner: self.inner.into_owned(), + } + } + pub(crate) fn push_chunk(&mut self, chunk: &[u8]) -> Result, GatewayError> { self.inner.push_chunk(chunk).map_err(map_surface_error) } diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index 8b7948968..9506d32e9 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -6462,6 +6462,21 @@ async fn execute_stream_from_frame_stream_with_retry_scope( let normalized_stream_report_context = normalize_provider_private_report_context(report_context.as_ref()); + // Observers follow the live protocol stream across prefetch and transfer. + // Diagnostic capture limits must never determine parser state. + let stream_usage_report_context = normalized_stream_report_context.clone().or_else(|| { + Some(json!({ + "provider_api_format": plan.provider_api_format.as_str(), + "client_api_format": plan.client_api_format.as_str(), + })) + }); + let mut stream_usage_observer = stream_usage_report_context + .as_ref() + .map(|_| StreamingStandardTerminalObserver::default()); + let mut stream_usage_observer_buffered = + StreamUsageObservationBuffer::new(max_stream_body_buffer_bytes); + let mut provider_error_inspection = ProviderStreamErrorInspection::default(); + let mut prefetched_provider_error = None; let upstream_headers = headers.clone(); let mut private_stream_normalizer = maybe_build_provider_private_stream_normalizer(report_context.as_ref()); @@ -6562,7 +6577,8 @@ async fn execute_stream_from_frame_stream_with_retry_scope( stream_commit_gate.commit(); } let mut prefetched_chunks: Vec = Vec::new(); - let mut provider_prefetched_body = Vec::new(); + let mut provider_prefetched_body = StreamBodyCapture::default(); + let mut provider_prefetched_bytes = 0_u64; let mut provider_prefetched_body_truncated = false; let mut prefetched_body = Vec::new(); let mut prefetched_inspection_body = Vec::new(); @@ -6815,10 +6831,12 @@ async fn execute_stream_from_frame_stream_with_retry_scope( } } - append_stream_capture_bytes( + provider_prefetched_bytes = + provider_prefetched_bytes.saturating_add(chunk.len() as u64); + append_budgeted_stream_capture_bytes( &mut provider_prefetched_body, &chunk, - MAX_STREAM_PREFETCH_BYTES, + max_stream_body_buffer_bytes, &mut provider_prefetched_body_truncated, ); append_stream_capture_bytes( @@ -7063,6 +7081,22 @@ async fn execute_stream_from_frame_stream_with_retry_scope( } else { chunk }; + if let Some(error) = provider_error_inspection + .observe(stream_usage_report_context.as_ref(), &normalized_chunk) + { + prefetched_provider_error.get_or_insert(error); + } + if let (Some(observer), Some(context)) = ( + stream_usage_observer.as_mut(), + stream_usage_report_context.as_ref(), + ) { + observe_stream_usage_bytes( + observer, + context, + &mut stream_usage_observer_buffered, + &normalized_chunk, + ); + } let rewritten_chunk = if let Some(rewriter) = local_stream_rewriter.as_mut() { match rewriter.push_chunk(&normalized_chunk) { Ok(rewritten_chunk) => rewritten_chunk, @@ -7193,17 +7227,21 @@ async fn execute_stream_from_frame_stream_with_retry_scope( if stream_commit_gate.is_uncommitted() { stream_commit_gate.commit(); } - let prefetched_response_history_persisted = if let Some(record) = local_stream_rewriter + if let Some(record) = local_stream_rewriter .as_mut() .and_then(|rewriter| rewriter.take_response_history_record()) { crate::ai_serving::persist_response_history_record(state, record).await; - true - } else { - false - }; - drop(private_stream_normalizer); - drop(local_stream_rewriter); + } + // Keep partial records and conversion state; replaying the bounded + // inspection/capture prefix loses any bytes consumed beyond that prefix. + let mut private_stream_normalizer = private_stream_normalizer.map(|parser| parser.into_owned()); + let mut local_stream_rewriter = local_stream_rewriter.map(|parser| parser.into_owned()); + if sync_json_stream_bridge_active { + private_stream_normalizer = None; + local_stream_rewriter = None; + stream_usage_observer = None; + } let initial_usage_telemetry = prefetched_usage_telemetry.clone().or_else(|| { prefetched_telemetry @@ -7246,7 +7284,6 @@ async fn execute_stream_from_frame_stream_with_retry_scope( let headers_for_report = headers.clone(); let report_kind_owned = report_kind; let report_context_owned = report_context; - let normalized_stream_report_context_owned = normalized_stream_report_context; let lifecycle_seed_for_report = lifecycle_seed; let provider_prefetched_body_for_report = provider_prefetched_body; let prefetched_body_for_report = prefetched_body; @@ -7288,40 +7325,10 @@ async fn execute_stream_from_frame_stream_with_retry_scope( let _stream_total_guard = StageElapsedGuard::from_started_at("stream_total", stream_started_at_for_report); let _provider_pool_in_flight_guard = provider_pool_in_flight_guard_for_report; - let mut provider_buffered_body = StreamBodyCapture::default(); + let mut provider_buffered_body = provider_prefetched_body_for_report; let mut buffered_body = StreamBodyCapture::default(); - let mut provider_body_truncated = false; + let mut provider_body_truncated = provider_prefetched_body_truncated; let mut client_body_truncated = false; - let mut private_stream_normalizer = if sync_json_stream_bridge_active_for_report { - None - } else { - maybe_build_provider_private_stream_normalizer(report_context_owned.as_ref()) - }; - let mut local_stream_rewriter = if sync_json_stream_bridge_active_for_report { - None - } else { - maybe_build_stream_response_rewriter(normalized_stream_report_context_owned.as_ref()) - }; - let stream_usage_report_context = - normalized_stream_report_context_owned.clone().or_else(|| { - Some(serde_json::json!({ - "provider_api_format": plan_for_report.provider_api_format.as_str(), - "client_api_format": plan_for_report.client_api_format.as_str(), - })) - }); - let mut stream_usage_observer = stream_usage_report_context - .as_ref() - .filter(|_| !sync_json_stream_bridge_active_for_report) - .map(|_| StreamingStandardTerminalObserver::default()); - let mut stream_usage_observer_buffered = - StreamUsageObservationBuffer::new(max_stream_body_buffer_bytes); - let mut provider_error_inspection = ProviderStreamErrorInspection::default(); - append_budgeted_stream_capture_bytes( - &mut provider_buffered_body, - &provider_prefetched_body_for_report, - max_stream_body_buffer_bytes, - &mut provider_body_truncated, - ); append_budgeted_stream_capture_bytes( &mut buffered_body, &prefetched_body_for_report, @@ -7365,9 +7372,7 @@ async fn execute_stream_from_frame_stream_with_retry_scope( } else { initial_elapsed_ms })); - let provider_stream_bytes = Arc::new(AtomicU64::new( - u64::try_from(provider_prefetched_body_for_report.len()).unwrap_or(u64::MAX), - )); + let provider_stream_bytes = Arc::new(AtomicU64::new(provider_prefetched_bytes)); let client_stream_bytes = Arc::new(AtomicU64::new( u64::try_from(prefetched_body_for_report.len()).unwrap_or(u64::MAX), )); @@ -7463,96 +7468,20 @@ async fn execute_stream_from_frame_stream_with_retry_scope( } }) }; - if !provider_prefetched_body_for_report.is_empty() { - let normalized_prefetched_chunk = if let Some(normalizer) = - private_stream_normalizer.as_mut() - { - match normalizer.push_chunk(&provider_prefetched_body_for_report) { - Ok(normalized_chunk) => Some(normalized_chunk), - Err(err) => { - warn!( - event_name = "stream_execution_prefetch_normalize_restore_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_category = "stream_normalization_restore_failed", - "gateway failed to restore private stream normalization state after prefetch" - ); - terminal_failure = Some(build_stream_failure_report( - "execution_runtime_stream_rewrite_error", - format!( - "failed to restore private stream normalization state after prefetch: {err:?}" - ), - 502, - )); - None - } - } - } else { - None - }; - let replay_chunk = normalized_prefetched_chunk - .as_deref() - .unwrap_or(provider_prefetched_body_for_report.as_slice()); - if let Some(error_body_json) = provider_error_inspection - .observe(stream_usage_report_context.as_ref(), replay_chunk) - { - provider_error_forwarded_to_client = !prefetched_body_for_report.is_empty(); - let error_status_code = resolve_provider_stream_error_status_code( - plan_for_report.provider_api_format.as_str(), - status_code, - &error_body_json, - ); - terminal_failure = Some(build_stream_failure_from_provider_error_body( - error_status_code, - &error_body_json, - )); - } - 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, - replay_chunk, - ); - } - if terminal_failure.is_none() { - if let Some(rewriter) = local_stream_rewriter.as_mut() { - if let Err(err) = rewriter.push_chunk(replay_chunk) { - warn!( - event_name = "stream_execution_prefetch_rewrite_restore_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_category = "stream_rewrite_restore_failed", - "gateway failed to restore local stream rewrite state after prefetch" - ); - terminal_failure = Some(build_stream_failure_report( - "execution_runtime_stream_rewrite_error", - format!( - "failed to restore local stream rewrite state after prefetch: {err:?}" - ), - 502, - )); - } - } - } - if prefetched_response_history_persisted { - if let Some(rewriter) = local_stream_rewriter.as_mut() { - let _ = rewriter.take_response_history_record(); - } - } + if let Some(error_body_json) = prefetched_provider_error { + provider_error_forwarded_to_client = !prefetched_body_for_report.is_empty(); + let error_status_code = resolve_provider_stream_error_status_code( + plan_for_report.provider_api_format.as_str(), + status_code, + &error_body_json, + ); + terminal_failure = Some(build_stream_failure_from_provider_error_body( + error_status_code, + &error_body_json, + )); } - - // These buffers restore parser/rewriter state above. Audit capture owns - // its budgeted copies; retaining semantic prefetch duplicates for the - // rest of the stream would bypass the capture memory limit. - drop(provider_prefetched_body_for_report); + // Parser state is already current and capture owns its budgeted bytes. + // This output prefix is needed only to initialize client-side trackers. drop(prefetched_body_for_report); if terminal_failure.is_none() && !reached_eof { @@ -9312,6 +9241,188 @@ mod tests { .unwrap() } + #[tokio::test] + async fn prefetch_handoff_preserves_large_responses_setup_event() { + let event = format!( + "event: response.created\ndata: {}\n\n", + json!({"type":"response.created", "response": { + "id":"resp-large-setup", "status":"in_progress", "output":[], + "tools":[{"name":"write", "description":"x".repeat(64 * 1024)}] + }}) + ); + let done = "event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-large-setup\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":7,\"output_tokens\":2}}}\n\n"; + // Include the two observed transport boundaries, exact/near budget + // boundaries, and multiple prefetch chunks crossing the budget. + for cuts in [ + vec![16_383], + vec![16_384], + vec![17_735], + vec![17_741], + vec![8_192, 17_735], + ] { + let mut chunks = Vec::new(); + let mut start = 0; + for end in cuts { + chunks.push(&event[start..end]); + start = end; + } + chunks.push(&event[start..]); + chunks.push(done); + let response = execute_generic_sse_precommit(chunks, json!({}), None, false) + .await + .expect("large setup should commit at the bounded prefetch limit"); + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + let body = String::from_utf8(body.to_vec()).unwrap(); + assert!( + body.starts_with(&event), + "setup bytes lost or duplicated at split {start}" + ); + let events: Vec = body + .lines() + .filter_map(|line| line.strip_prefix("data: ")) + .filter(|payload| *payload != "[DONE]") + .map(|payload| { + serde_json::from_str(payload).expect("every SSE payload must be valid JSON") + }) + .collect(); + assert_eq!(events.len(), 2, "events must be forwarded exactly once"); + assert_eq!(events[1]["type"], "response.completed"); + } + } + + #[tokio::test] + async fn prefetch_handoff_keeps_audit_usage_and_private_conversion() { + for private in [false, true] { + let request_id = format!("handoff-audit-{}", uuid::Uuid::new_v4()); + let mut plan = if private { + antigravity_gemini_stream_plan(&request_id) + } else { + native_anthropic_stream_plan(&request_id) + }; + if !private { + plan.provider_api_format = "openai:responses".into(); + plan.client_api_format = "openai:responses".into(); + } + let context = json!({ + "request_id": request_id, "candidate_id": plan.candidate_id, + "candidate_index":0, "retry_index":0, + "provider_api_format": plan.provider_api_format, + "client_api_format": plan.client_api_format, + "needs_conversion": private, "has_envelope": private, + "envelope_name": if private { "antigravity:v1internal" } else { "" }, + }); + let repository = Arc::new(InMemoryUsageReadRepository::default()); + let catalog = provider_catalog_for_plan(&plan, None); + let state = AppState::new() + .unwrap() + .with_data_state_for_tests( + crate::data::GatewayDataState::with_usage_repository_for_tests(Arc::clone( + &repository, + )) + .with_provider_catalog_reader(Arc::new(catalog)) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY) + .with_system_config_values_for_tests([( + "request_record_level".into(), + json!("full"), + )]), + ) + .with_usage_runtime_for_tests(UsageRuntimeConfig { + enabled: true, + ..Default::default() + }); + let text = "hello".repeat(12_000); + let payload = if private { + json!({"response":{"candidates":[{"content":{"role":"model","parts":[{"text":text}]}, + "finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":1234,"candidatesTokenCount":567}, + "modelVersion":"gemini-3.7-flash-tiered"}}) + } else { + json!({"type":"response.completed","response":{"id":"resp-handoff-usage","status":"completed", + "output":[{"type":"message","id":"msg-handoff","role":"assistant","status":"completed", + "content":[{"type":"output_text","text":text,"annotations":[]}]}], + "usage":{"input_tokens":1234,"output_tokens":567,"total_tokens":1801}}}) + }; + let input = format!("data: {payload}\n\n"); + // One complete large chunk exercises an already-emitted prefetch + // result; the private path exercises incomplete normalization too. + let chunks = if private { + vec![input[..17_735].to_string(), input[17_735..].to_string()] + } else { + vec![input.clone()] + }; + let frames = stream! { + yield Ok::(ndjson_frame(StreamFrame { + frame_type:StreamFrameType::Headers, + payload:StreamFramePayload::Headers { status_code:200, + headers:BTreeMap::from([("content-type".into(),"text/event-stream".into())]), + response_observation:None }, + })); + for chunk in chunks { + yield Ok(ndjson_frame(StreamFrame { frame_type:StreamFrameType::Data, + payload:StreamFramePayload::Data { text:Some(chunk),chunk_b64:None } })); + } + yield Ok(ndjson_frame(StreamFrame::eof())); + }.boxed(); + let response = execute_stream_from_frame_stream( + &state, + plan, + "trace-handoff-audit", + &test_decision(), + OPENAI_RESPONSES_STREAM_PLAN_KIND, + Some("openai_responses_stream_success".into()), + Some(context), + crate::clock::current_unix_ms(), + Instant::now(), + RequestStageTrace::from_env(), + false, + frames, + None, + ) + .await + .unwrap() + .unwrap(); + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + let body = String::from_utf8(body.to_vec()).unwrap(); + let events: Vec = body + .lines() + .filter_map(|l| l.strip_prefix("data: ")) + .filter(|p| *p != "[DONE]") + .map(|p| serde_json::from_str(p).unwrap()) + .collect(); + assert_eq!( + events + .iter() + .filter(|e| e["type"] == "response.completed") + .count(), + 1 + ); + assert!(body.contains(&text)); + let usage = tokio::time::timeout(Duration::from_secs(3), async { + loop { + if let Some(u) = repository + .find_by_request_id(&request_id) + .await + .unwrap() + .filter(|u| u.status == "completed" || u.status == "failed") + { + break u; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("usage should finalize"); + assert_eq!(usage.status, "completed", "{:?}", usage.error_message); + assert_eq!(usage.input_tokens, 1234); + assert_eq!(usage.output_tokens, 567); + let captured = usage.response_body.as_ref().expect("provider capture"); + assert!( + captured["metadata"].get("dropped_chunks").is_none(), + "{captured}" + ); + assert_eq!(captured["chunks"].as_array().unwrap(), &vec![payload]); + } + } + #[tokio::test] async fn generic_stream_success_regex_matches_fragmented_plain_body() { for chunks in [ @@ -9850,7 +9961,7 @@ mod tests { let mut buffer = super::StreamUsageObservationBuffer::new(32 * 1024); let mut rewriter = super::maybe_build_stream_response_rewriter(Some(&context)).unwrap(); let mut delivered = Vec::new(); - for chunk in chunks { + for (index, chunk) in chunks.into_iter().enumerate() { provider.append(chunk, 32 * 1024, &mut provider_truncated); super::observe_stream_usage_bytes( observer.as_mut().unwrap(), @@ -9861,6 +9972,10 @@ mod tests { let output = rewriter.push_chunk(chunk).unwrap(); client.append(&output, 32 * 1024, &mut client_truncated); delivered.extend(output); + if index == 0 { + // Task handoff must also work when audit admits no bytes. + rewriter = rewriter.into_owned(); + } } let tail = rewriter.finish().unwrap(); client.append(&tail, 32 * 1024, &mut client_truncated); diff --git a/crates/aether-ai/formats/src/formats/shared/stream_rewrite.rs b/crates/aether-ai/formats/src/formats/shared/stream_rewrite.rs index d766c5aa2..10a9f3808 100644 --- a/crates/aether-ai/formats/src/formats/shared/stream_rewrite.rs +++ b/crates/aether-ai/formats/src/formats/shared/stream_rewrite.rs @@ -1,3 +1,4 @@ +use std::borrow::Cow; use std::collections::BTreeMap; use serde_json::{json, Map, Value}; @@ -226,7 +227,7 @@ enum AiSurfaceStreamRewriteState { } pub struct AiSurfaceStreamRewriter<'a> { - report_context: &'a Value, + report_context: Cow<'a, Value>, buffered: Vec, state: AiSurfaceStreamRewriteState, } @@ -271,30 +272,39 @@ pub fn maybe_build_ai_surface_stream_rewriter<'a>( }; Some(AiSurfaceStreamRewriter { - report_context, + report_context: Cow::Borrowed(report_context), buffered: Vec::new(), state, }) } impl AiSurfaceStreamRewriter<'_> { + /// Move parser state across task boundaries without replaying captured bytes. + pub fn into_owned(self) -> AiSurfaceStreamRewriter<'static> { + AiSurfaceStreamRewriter { + report_context: Cow::Owned(self.report_context.into_owned()), + buffered: self.buffered, + state: self.state, + } + } + pub fn push_chunk(&mut self, chunk: &[u8]) -> Result, AiSurfaceFinalizeError> { match &mut self.state { AiSurfaceStreamRewriteState::OpenAiImage(state) => { - state.push_chunk(self.report_context, chunk) + state.push_chunk(self.report_context.as_ref(), chunk) } AiSurfaceStreamRewriteState::OpenAiImageToOpenAiChat(state) => { - state.push_chunk(self.report_context, chunk) + state.push_chunk(self.report_context.as_ref(), chunk) } AiSurfaceStreamRewriteState::ClaudeReadToolSanitize(state) => { - state.push_chunk(self.report_context, chunk) + state.push_chunk(self.report_context.as_ref(), chunk) } AiSurfaceStreamRewriteState::KiroToClaudeCli(state) => { - state.push_chunk(self.report_context, chunk) + state.push_chunk(self.report_context.as_ref(), chunk) } AiSurfaceStreamRewriteState::KiroToClaudeCliThenStandard { kiro, standard } => { - let claude_bytes = kiro.push_chunk(self.report_context, chunk)?; - transform_standard_bytes(standard, self.report_context, claude_bytes) + let claude_bytes = kiro.push_chunk(self.report_context.as_ref(), chunk)?; + transform_standard_bytes(standard, self.report_context.as_ref(), claude_bytes) } AiSurfaceStreamRewriteState::EnvelopeUnwrap | AiSurfaceStreamRewriteState::ModelDirectiveDisplay @@ -313,23 +323,25 @@ impl AiSurfaceStreamRewriter<'_> { pub fn finish(&mut self) -> Result, AiSurfaceFinalizeError> { match &mut self.state { - AiSurfaceStreamRewriteState::OpenAiImage(state) => state.finish(self.report_context), + AiSurfaceStreamRewriteState::OpenAiImage(state) => { + state.finish(self.report_context.as_ref()) + } AiSurfaceStreamRewriteState::OpenAiImageToOpenAiChat(state) => { - state.finish(self.report_context) + state.finish(self.report_context.as_ref()) } AiSurfaceStreamRewriteState::ClaudeReadToolSanitize(state) => { - state.finish(self.report_context) + state.finish(self.report_context.as_ref()) } AiSurfaceStreamRewriteState::KiroToClaudeCli(state) => { - state.finish(self.report_context) + state.finish(self.report_context.as_ref()) } AiSurfaceStreamRewriteState::KiroToClaudeCliThenStandard { kiro, standard } => { let mut output = transform_standard_bytes( standard, - self.report_context, - kiro.finish(self.report_context)?, + self.report_context.as_ref(), + kiro.finish(self.report_context.as_ref())?, )?; - output.extend(standard.finish(self.report_context)?); + output.extend(standard.finish(self.report_context.as_ref())?); Ok(output) } AiSurfaceStreamRewriteState::EnvelopeUnwrap @@ -338,14 +350,14 @@ impl AiSurfaceStreamRewriter<'_> { | AiSurfaceStreamRewriteState::Standard(_) => { if self.buffered.is_empty() { if let AiSurfaceStreamRewriteState::Standard(state) = &mut self.state { - return state.finish(self.report_context); + return state.finish(self.report_context.as_ref()); } return Ok(Vec::new()); } let line = std::mem::take(&mut self.buffered); let mut output = self.transform_line(line)?; if let AiSurfaceStreamRewriteState::Standard(state) = &mut self.state { - output.extend(state.finish(self.report_context)?); + output.extend(state.finish(self.report_context.as_ref())?); } Ok(output) } @@ -365,18 +377,19 @@ impl AiSurfaceStreamRewriter<'_> { fn transform_line(&mut self, line: Vec) -> Result, AiSurfaceFinalizeError> { match &mut self.state { AiSurfaceStreamRewriteState::EnvelopeUnwrap => { - let output = transform_provider_private_stream_line(self.report_context, line) - .map_err(AiSurfaceFinalizeError::from)?; - rewrite_model_directive_stream_line(self.report_context, output) + let output = + transform_provider_private_stream_line(self.report_context.as_ref(), line) + .map_err(AiSurfaceFinalizeError::from)?; + rewrite_model_directive_stream_line(self.report_context.as_ref(), output) } AiSurfaceStreamRewriteState::ModelDirectiveDisplay => { - rewrite_model_directive_stream_line(self.report_context, line) + rewrite_model_directive_stream_line(self.report_context.as_ref(), line) } AiSurfaceStreamRewriteState::OpenAiResponsesCompat => { - rewrite_openai_responses_compat_stream_line(self.report_context, line) + rewrite_openai_responses_compat_stream_line(self.report_context.as_ref(), line) } AiSurfaceStreamRewriteState::Standard(state) => { - transform_standard_line(state, self.report_context, line) + transform_standard_line(state, self.report_context.as_ref(), line) } AiSurfaceStreamRewriteState::OpenAiImage(_) | AiSurfaceStreamRewriteState::OpenAiImageToOpenAiChat(_) @@ -892,7 +905,7 @@ fn is_standard_cli_client_api_format(api_format: &str) -> bool { #[cfg(test)] mod tests { - use serde_json::json; + use serde_json::{json, Value}; use super::{ maybe_build_ai_surface_stream_rewriter, resolve_finalize_stream_rewrite_mode, @@ -1067,6 +1080,50 @@ data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_123\",\"object\ assert!(!output.contains("\"model\":\"gpt-5.5\"")); } + #[test] + fn owned_handoff_preserves_partial_utf8_and_conversion_state() { + for client in ["openai:responses", "openai:chat"] { + let text = "界".repeat(12_000); + let delta = format!( + "data: {}\n\n", + json!({ + "type":"response.output_text.delta", "response_id":"resp_handoff", + "item_id":"msg_handoff", "output_index":0, "content_index":0, "delta":text, + }) + ); + let split = delta.find('界').unwrap() + 17_002; + assert!(!delta.is_char_boundary(split)); + let (mut owned, mut output) = { + let context = json!({"provider_api_format":"openai:responses", + "client_api_format":client, "needs_conversion":client == "openai:chat"}); + let mut parser = maybe_build_ai_surface_stream_rewriter(Some(&context)).unwrap(); + let output = parser.push_chunk(&delta.as_bytes()[..split]).unwrap(); + (parser.into_owned(), output) + }; + output.extend(owned.push_chunk(&delta.as_bytes()[split..]).unwrap()); + output.extend(owned.finish().unwrap()); + let output = String::from_utf8(output).unwrap(); + let events: Vec = output + .lines() + .filter_map(|l| l.strip_prefix("data: ")) + .filter(|p| *p != "[DONE]") + .map(|p| serde_json::from_str(p).unwrap()) + .collect(); + let recovered: String = events + .iter() + .filter_map(|e| { + if client == "openai:responses" { + e["delta"].as_str() + } else { + e.pointer("/choices/0/delta/content") + .and_then(Value::as_str) + } + }) + .collect(); + assert_eq!(recovered, text); + } + } + #[test] fn standard_rewriter_converts_openai_responses_reasoning_delta_to_chat() { let report_context = json!({ diff --git a/crates/aether-ai/formats/src/provider_compat/private_envelope.rs b/crates/aether-ai/formats/src/provider_compat/private_envelope.rs index d63f0e858..4adbd1e2f 100644 --- a/crates/aether-ai/formats/src/provider_compat/private_envelope.rs +++ b/crates/aether-ai/formats/src/provider_compat/private_envelope.rs @@ -1,3 +1,4 @@ +use std::borrow::Cow; use std::collections::BTreeMap; use serde_json::Value; @@ -354,7 +355,7 @@ enum ProviderPrivateStreamNormalizeMode { } pub struct ProviderPrivateStreamNormalizer<'a> { - report_context: &'a Value, + report_context: Cow<'a, Value>, buffered: Vec, current_event_type: Option, mode: ProviderPrivateStreamNormalizeMode, @@ -401,7 +402,7 @@ pub fn maybe_build_provider_private_stream_normalizer<'a>( return None; }; Some(ProviderPrivateStreamNormalizer { - report_context, + report_context: Cow::Borrowed(report_context), buffered: Vec::new(), current_event_type: None, mode, @@ -422,10 +423,20 @@ pub fn extract_provider_private_stream_error_body( } impl ProviderPrivateStreamNormalizer<'_> { + /// Move parser state across task boundaries without replaying captured bytes. + pub fn into_owned(self) -> ProviderPrivateStreamNormalizer<'static> { + ProviderPrivateStreamNormalizer { + report_context: Cow::Owned(self.report_context.into_owned()), + buffered: self.buffered, + current_event_type: self.current_event_type, + mode: self.mode, + } + } + pub fn push_chunk(&mut self, chunk: &[u8]) -> Result, AiSurfaceFinalizeError> { match &mut self.mode { ProviderPrivateStreamNormalizeMode::KiroToClaudeCli(state) => { - state.push_chunk(self.report_context, chunk) + state.push_chunk(self.report_context.as_ref(), chunk) } ProviderPrivateStreamNormalizeMode::EnvelopeUnwrap => { let next_len = self @@ -441,7 +452,7 @@ impl ProviderPrivateStreamNormalizer<'_> { ))); } self.buffered.extend_from_slice(chunk); - if report_context_is_windsurf_envelope(self.report_context) + if report_context_is_windsurf_envelope(self.report_context.as_ref()) && buffer_looks_like_connect_frame(&self.buffered) { return drain_windsurf_connect_json_frames(&mut self.buffered); @@ -451,7 +462,7 @@ impl ProviderPrivateStreamNormalizer<'_> { let line = self.buffered.drain(..=line_end).collect::>(); output.extend( transform_provider_private_stream_line_with_event_state( - self.report_context, + self.report_context.as_ref(), line, &mut self.current_event_type, ) @@ -466,20 +477,20 @@ impl ProviderPrivateStreamNormalizer<'_> { pub fn finish(&mut self) -> Result, AiSurfaceFinalizeError> { match &mut self.mode { ProviderPrivateStreamNormalizeMode::KiroToClaudeCli(state) => { - state.finish(self.report_context) + state.finish(self.report_context.as_ref()) } ProviderPrivateStreamNormalizeMode::EnvelopeUnwrap => { if self.buffered.is_empty() { return Ok(Vec::new()); } - if report_context_is_windsurf_envelope(self.report_context) + if report_context_is_windsurf_envelope(self.report_context.as_ref()) && buffer_looks_like_connect_frame(&self.buffered) { return drain_windsurf_connect_json_frames(&mut self.buffered); } let line = std::mem::take(&mut self.buffered); transform_provider_private_stream_line_with_event_state( - self.report_context, + self.report_context.as_ref(), line, &mut self.current_event_type, ) @@ -939,7 +950,7 @@ fn postprocess_private_response_value(data: &mut Value, report_context: &Value) #[cfg(test)] mod tests { - use serde_json::json; + use serde_json::{json, Value}; use super::{ extract_provider_private_stream_error_body, maybe_build_provider_private_stream_normalizer, @@ -1116,6 +1127,44 @@ mod tests { assert!(text.contains(r#""content":"chunk""#)); } + #[test] + fn owned_handoff_preserves_private_binary_frame() { + let text = "frame".repeat(10_000); + let framed = connect_json_frame( + 0, + &serde_json::to_vec(&json!({ + "responseId":"ws-handoff", "response":{"text":text} + })) + .unwrap(), + ); + let split = 17_735; + let mut normalizer = { + let context = json!({"has_envelope":true, + "envelope_name":"windsurf:GetChatMessage", "provider_api_format":"openai:chat"}); + let mut normalizer = + maybe_build_provider_private_stream_normalizer(Some(&context)).unwrap(); + assert!(normalizer.push_chunk(&framed[..split]).unwrap().is_empty()); + normalizer.into_owned() + }; + let mut output = normalizer.push_chunk(&framed[split..]).unwrap(); + output.extend(normalizer.finish().unwrap()); + let output = String::from_utf8(output).unwrap(); + let events: Vec = output + .lines() + .filter_map(|l| l.strip_prefix("data: ")) + .filter(|p| *p != "[DONE]") + .map(|p| serde_json::from_str(p).unwrap()) + .collect(); + let recovered: String = events + .iter() + .filter_map(|e| { + e.pointer("/choices/0/delta/content") + .and_then(Value::as_str) + }) + .collect(); + assert_eq!(recovered, text); + } + #[test] fn unwraps_windsurf_connect_json_stream_frames() { let report_context = json!({ diff --git a/crates/aether-usage/runtime/src/write.rs b/crates/aether-usage/runtime/src/write.rs index 1fb82a2be..5f4659930 100644 --- a/crates/aether-usage/runtime/src/write.rs +++ b/crates/aether-usage/runtime/src/write.rs @@ -3052,14 +3052,26 @@ fn parse_sse_body_for_storage(text: &str) -> Option { let mut chunks = Vec::new(); let mut total_chunks = 0_u64; let mut saw_done = false; + let mut first_parse_error = None; for_each_sse_payload(text, |payload| { if payload == "[DONE]" { saw_done = true; return; } total_chunks += 1; - if let Ok(json_body) = serde_json::from_str::(payload) { - chunks.push(json_body); + match serde_json::from_str::(payload) { + Ok(json_body) => chunks.push(json_body), + Err(error) if first_parse_error.is_none() => { + // A later valid event must not hide an earlier broken one. + // Store diagnostics only, without duplicating raw user content. + first_parse_error = Some(json!({ + "chunk_index": total_chunks - 1, + "line": error.line(), + "column": error.column(), + "message": error.to_string(), + })); + } + Err(_) => {} } }); if total_chunks == 0 && !saw_done { @@ -3076,6 +3088,11 @@ fn parse_sse_body_for_storage(text: &str) -> Option { if saw_done { metadata.insert("has_completion".to_string(), Value::Bool(true)); } + if let Some(error) = first_parse_error { + // Capture truncation can also cause a parse error; this describes the + // captured payload, not an assertion that the provider sent bad JSON. + metadata.insert("first_parse_error".to_string(), error); + } if stored_chunks < total_chunks { metadata.insert( "dropped_chunks".to_string(), @@ -7104,6 +7121,20 @@ mod tests { ); } + #[test] + fn parse_sse_body_for_storage_reports_bad_event_before_valid_terminal() { + let body = concat!( + "data: {\"tools\":[}\n\n", + "data: {\"type\":\"response.completed\"}\n\n", + ); + let parsed = parse_sse_body_for_storage(body).unwrap(); + assert_eq!(parsed["metadata"]["dropped_chunks"], 1); + assert_eq!(parsed["metadata"]["first_parse_error"]["chunk_index"], 0); + assert!(parsed["metadata"]["first_parse_error"]["message"].is_string()); + assert_eq!(parsed["chunks"][0]["type"], "response.completed"); + assert!(parsed.get("raw_response").is_none()); + } + #[test] fn extract_token_counts_from_value_handles_crlf_and_cr_sse_text() { let sse_body = concat!(