feat(stream): 在流式执行中引入 StreamingStandardTerminalObserver 用量采集

- 在 execute_stream_from_frame_stream 中集成 StreamingStandardTerminalObserver,对流式响应逐行观察并在结束时合并终态摘要
- 新增 observe_stream_usage_bytes / finalize_stream_usage_observer / merge_stream_terminal_summary 辅助函数
- 更新 codex-cli 流式集成测试:补充 response.completed 含完整 usage 字段的 mock 数据,并断言用量写入 usage_repository
This commit is contained in:
fawney19
2026-04-24 11:34:41 +08:00
parent 35400b0c2d
commit f29649e3a8
2 changed files with 177 additions and 6 deletions

View File

@@ -41,6 +41,7 @@ use self::execution_failures::{
use crate::ai_pipeline_api::{
maybe_bridge_standard_sync_json_to_stream, maybe_build_provider_private_stream_normalizer,
maybe_build_stream_response_rewriter, normalize_provider_private_report_context,
StreamingStandardTerminalObserver,
};
use crate::api::response::{
attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
@@ -204,6 +205,84 @@ fn append_stream_capture_bytes(
}
}
fn observe_stream_usage_bytes(
observer: &mut StreamingStandardTerminalObserver,
report_context: &Value,
buffered: &mut Vec<u8>,
chunk: &[u8],
) {
if chunk.is_empty() {
return;
}
buffered.extend_from_slice(chunk);
while let Some(line_end) = buffered.iter().position(|byte| *byte == b'\n') {
let line = buffered.drain(..=line_end).collect::<Vec<_>>();
if let Err(err) = observer.push_line(report_context, line) {
observer.disable_with_error(err.to_string());
buffered.clear();
break;
}
}
}
fn finalize_stream_usage_observer(
observer: &mut Option<StreamingStandardTerminalObserver>,
report_context: Option<&Value>,
buffered: &mut Vec<u8>,
) -> Option<ExecutionStreamTerminalSummary> {
let (Some(observer), Some(report_context)) = (observer.as_mut(), report_context) else {
return None;
};
if !buffered.is_empty() {
let line = std::mem::take(buffered);
if let Err(err) = observer.push_line(report_context, line) {
observer.disable_with_error(err.to_string());
}
}
match observer.finish(report_context) {
Ok(summary) => summary,
Err(err) => {
observer.disable_with_error(err.to_string());
observer.latest_summary().cloned()
}
}
}
fn merge_stream_terminal_summary(
mut current: Option<ExecutionStreamTerminalSummary>,
observed: Option<ExecutionStreamTerminalSummary>,
) -> Option<ExecutionStreamTerminalSummary> {
let Some(observed) = observed else {
return current;
};
let Some(current_summary) = current.as_mut() else {
return Some(observed);
};
if current_summary.standardized_usage.is_none() {
current_summary.standardized_usage = observed.standardized_usage;
}
if current_summary.finish_reason.is_none() {
current_summary.finish_reason = observed.finish_reason;
}
if current_summary.response_id.is_none() {
current_summary.response_id = observed.response_id;
}
if current_summary.model.is_none() {
current_summary.model = observed.model;
}
current_summary.observed_finish |= observed.observed_finish;
if current_summary.parser_error.is_none() {
current_summary.parser_error = observed.parser_error;
}
current
}
async fn execute_in_process_stream(
state: &AppState,
plan: &ExecutionPlan,
@@ -1333,6 +1412,18 @@ async fn execute_stream_from_frame_stream(
} 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 = Vec::new();
append_stream_capture_bytes(
&mut provider_buffered_body,
&provider_prefetched_body_for_report,
@@ -1382,6 +1473,17 @@ async fn execute_stream_from_frame_stream(
let replay_chunk = normalized_prefetched_chunk
.as_deref()
.unwrap_or(provider_prefetched_body_for_report.as_slice());
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) {
@@ -1496,6 +1598,17 @@ async fn execute_stream_from_frame_stream(
} else {
chunk
};
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,
);
}
let rewritten_chunk = if let Some(rewriter) = local_stream_rewriter.as_mut()
{
match rewriter.push_chunk(&normalized_chunk) {
@@ -1600,6 +1713,17 @@ async fn execute_stream_from_frame_stream(
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,
);
}
let rewritten_chunk = if let Some(rewriter) = local_stream_rewriter.as_mut()
{
match rewriter.push_chunk(&normalized_chunk) {
@@ -1753,6 +1877,15 @@ async fn execute_stream_from_frame_stream(
drop(tx);
stream_terminal_summary = merge_stream_terminal_summary(
stream_terminal_summary,
finalize_stream_usage_observer(
&mut stream_usage_observer,
stream_usage_report_context.as_ref(),
&mut stream_usage_observer_buffered,
),
);
if downstream_dropped {
debug!(
event_name = "execution_runtime_stream_report_skipped",