From 2a9a02e27e2f44eca3b80314b63090ab05ab8d85 Mon Sep 17 00:00:00 2001 From: LAP16603 Date: Sun, 9 Aug 2026 06:30:58 +0700 Subject: [PATCH] fix(agent): forward AionRS provider token usage --- crates/aionui-ai-agent/src/agent_runtime.rs | 39 +++++- .../src/capability/backend_output_sink.rs | 21 ++-- .../src/manager/acp/agent_session_flow.rs | 1 + .../src/manager/aionrs/agent.rs | 39 +++++- .../src/manager/aionrs/agent_test.rs | 113 +++++++++++++++++- .../src/protocol/events/mod.rs | 5 + crates/aionui-ai-agent/src/session_agent.rs | 1 + crates/aionui-channel/src/message_service.rs | 5 +- .../aionui-channel/tests/stream_relay_test.rs | 20 +++- .../aionui-conversation/src/stream_relay.rs | 9 +- crates/aionui-team/tests/e2e_team_flow.rs | 7 +- 11 files changed, 234 insertions(+), 26 deletions(-) diff --git a/crates/aionui-ai-agent/src/agent_runtime.rs b/crates/aionui-ai-agent/src/agent_runtime.rs index e5f3f4d1f..b628d2edd 100644 --- a/crates/aionui-ai-agent/src/agent_runtime.rs +++ b/crates/aionui-ai-agent/src/agent_runtime.rs @@ -109,6 +109,22 @@ impl AgentRuntime { /// Atomic: set status ← Finished AND broadcast `Finish(session_id)`. /// Idempotent in the Finished absorbing state (no-op). pub fn emit_finish(&self, session_id: Option) { + self.emit_finish_data(FinishEventData { + session_id, + ..Default::default() + }); + } + + /// Atomic finish carrying provider-reported consumption for one completed turn. + pub fn emit_finish_with_usage(&self, session_id: Option, input_tokens: u64, output_tokens: u64) { + self.emit_finish_data(FinishEventData { + session_id, + input_tokens: Some(input_tokens), + output_tokens: Some(output_tokens), + }); + } + + fn emit_finish_data(&self, data: FinishEventData) { let already_finished = { let mut guard = self.status.write().unwrap_or_else(|e| e.into_inner()); let was_finished = matches!(*guard, Some(ConversationStatus::Finished)); @@ -120,9 +136,7 @@ impl AgentRuntime { if already_finished { return; } - let _ = self - .event_tx - .send(AgentStreamEvent::Finish(FinishEventData { session_id })); + let _ = self.event_tx.send(AgentStreamEvent::Finish(data)); } /// Atomic: set status ← Finished AND broadcast `Error { message }`. @@ -193,6 +207,25 @@ mod tests { } } + #[tokio::test] + async fn emit_finish_with_usage_transitions_and_broadcasts_provider_totals() { + let rt = runtime(); + let mut rx = rt.subscribe(); + + rt.emit_finish_with_usage(Some("sess-usage".into()), 1_200, 345); + + assert_eq!(rt.status(), Some(ConversationStatus::Finished)); + let ev = rx.recv().await.expect("finish event"); + match ev { + AgentStreamEvent::Finish(data) => { + assert_eq!(data.session_id.as_deref(), Some("sess-usage")); + assert_eq!(data.input_tokens, Some(1_200)); + assert_eq!(data.output_tokens, Some(345)); + } + other => panic!("expected Finish, got {other:?}"), + } + } + #[tokio::test] async fn emit_error_transitions_and_broadcasts() { let rt = runtime(); diff --git a/crates/aionui-ai-agent/src/capability/backend_output_sink.rs b/crates/aionui-ai-agent/src/capability/backend_output_sink.rs index 087295aff..c82af5921 100644 --- a/crates/aionui-ai-agent/src/capability/backend_output_sink.rs +++ b/crates/aionui-ai-agent/src/capability/backend_output_sink.rs @@ -128,14 +128,16 @@ impl OutputSink for BackendOutputSink { &self, _msg_id: &str, _turns: usize, - _input_tokens: u64, - _output_tokens: u64, + input_tokens: u64, + output_tokens: u64, _cache_creation_tokens: u64, _cache_read_tokens: u64, ) { - let _ = self - .event_tx - .send(AgentStreamEvent::Finish(FinishEventData { session_id: None })); + let _ = self.event_tx.send(AgentStreamEvent::Finish(FinishEventData { + session_id: None, + input_tokens: Some(input_tokens), + output_tokens: Some(output_tokens), + })); } fn emit_error(&self, msg: &str) { @@ -284,10 +286,11 @@ mod tests { let (sink, mut rx) = make_sink(); sink.emit_stream_end("msg-1", 3, 1000, 500, 100, 200); let event = rx.try_recv().unwrap(); - match event { - AgentStreamEvent::Finish(_) => {} - other => panic!("Expected Finish, got {:?}", other), - } + let value = serde_json::to_value(event).unwrap(); + + assert_eq!(value["type"], "finish"); + assert_eq!(value["data"]["input_tokens"], 1000); + assert_eq!(value["data"]["output_tokens"], 500); } #[test] diff --git a/crates/aionui-ai-agent/src/manager/acp/agent_session_flow.rs b/crates/aionui-ai-agent/src/manager/acp/agent_session_flow.rs index dbe6dcc29..e67ccc0aa 100644 --- a/crates/aionui-ai-agent/src/manager/acp/agent_session_flow.rs +++ b/crates/aionui-ai-agent/src/manager/acp/agent_session_flow.rs @@ -1203,6 +1203,7 @@ mod tests { .unwrap(); tx.send(AgentStreamEvent::Finish(FinishEventData { session_id: Some("s1".into()), + ..Default::default() })) .unwrap(); diff --git a/crates/aionui-ai-agent/src/manager/aionrs/agent.rs b/crates/aionui-ai-agent/src/manager/aionrs/agent.rs index 8e5c9677a..5ccceccd4 100644 --- a/crates/aionui-ai-agent/src/manager/aionrs/agent.rs +++ b/crates/aionui-ai-agent/src/manager/aionrs/agent.rs @@ -14,6 +14,7 @@ use aion_config::config::{CliArgs, Config, McpServerConfig, ProviderType}; use aion_mcp::manager::McpManager; use aion_protocol::commands::{ApprovalScope, SessionMode}; use aion_protocol::{ToolApprovalManager, ToolApprovalResult}; +use aion_types::message::TokenUsage; use aionui_api_types::{ AcpConfigOptionDto, AcpConfigSelectOptionDto, AgentModeResponse, ConfigOptionConfirmation, GetConfigOptionsResponse, SetConfigOptionResponse, SlashCommandItem, @@ -104,6 +105,8 @@ fn build_aionrs_final_input_dump_value( pub struct AionrsAgentManager { runtime: AgentRuntime, engine: Mutex, + /// Last cumulative engine usage already projected onto per-turn finish events. + usage_watermark: Mutex>, /// Static slash command metadata captured at bootstrap so UI lookups do /// not wait behind an active `engine.run()` turn. slash_commands: Vec, @@ -143,6 +146,10 @@ impl AionrsAgentManager { ) -> Result { let runtime = AgentRuntime::new(conversation_id.clone(), workspace.clone(), 128); let sink: Arc = Arc::new(BackendOutputSink::new(runtime.event_sender())); + let usage_watermark = resume_session + .as_ref() + .map(|session| session.total_usage.clone()) + .unwrap_or_default(); let runtime_env = config_extra.runtime_env.clone(); let image_input_override = config_extra.compat_overrides.image_input; let image_input_capability = image_input_override.unwrap_or_else(|| { @@ -275,6 +282,7 @@ impl AionrsAgentManager { Ok(Self { runtime, engine: Mutex::new(engine), + usage_watermark: Mutex::new(Some(usage_watermark)), slash_commands, mcp_managers: result.mcp_managers, approval_manager, @@ -285,6 +293,31 @@ impl AionrsAgentManager { }) } + async fn emit_successful_finish(&self, cumulative_usage: TokenUsage) { + let mut usage_watermark = self.usage_watermark.lock().await; + let usage_delta = usage_watermark + .as_ref() + .and_then(|watermark| { + cumulative_usage + .input_tokens + .checked_sub(watermark.input_tokens) + .zip(cumulative_usage.output_tokens.checked_sub(watermark.output_tokens)) + }) + .filter(|(input_tokens, output_tokens)| *input_tokens > 0 || *output_tokens > 0); + *usage_watermark = Some(cumulative_usage); + drop(usage_watermark); + + if let Some((input_tokens, output_tokens)) = usage_delta { + self.runtime.emit_finish_with_usage(None, input_tokens, output_tokens); + } else { + self.runtime.emit_finish(None); + } + } + + async fn invalidate_usage_watermark(&self) { + *self.usage_watermark.lock().await = None; + } + fn request_stop(&self, reason: Option, operation: &'static str) -> bool { let was_running = self.runtime.status() == Some(ConversationStatus::Running); @@ -422,16 +455,17 @@ impl IAgentTask for AionrsAgentManager { self.runtime.bump_activity(); let send_result = match result { - Some(Ok(_)) => { + Some(Ok(run_result)) => { info!( conversation_id = %self.runtime.conversation_id(), elapsed_ms, "Aionrs engine.run() completed, emitting Finish" ); - self.runtime.emit_finish(None); + self.emit_successful_finish(run_result.usage).await; Ok(()) } Some(Err(e)) => { + self.invalidate_usage_watermark().await; let summary = aionrs_runtime_error_summary(&e); error!( conversation_id = %self.runtime.conversation_id(), @@ -458,6 +492,7 @@ impl IAgentTask for AionrsAgentManager { Err(send_error) } None => { + self.invalidate_usage_watermark().await; self.runtime.emit_finish(None); Ok(()) } diff --git a/crates/aionui-ai-agent/src/manager/aionrs/agent_test.rs b/crates/aionui-ai-agent/src/manager/aionrs/agent_test.rs index 9310f1e0b..c1f23f38a 100644 --- a/crates/aionui-ai-agent/src/manager/aionrs/agent_test.rs +++ b/crates/aionui-ai-agent/src/manager/aionrs/agent_test.rs @@ -5,6 +5,7 @@ use std::path::PathBuf; use std::time::Duration; use aion_config::config::{McpServerConfig, TransportType}; +use aion_types::message::TokenUsage; use tokio::sync::broadcast::error::TryRecvError; use tokio::time::timeout; @@ -330,9 +331,7 @@ async fn runtime_can_emit_error_and_finish() { agent.runtime.emit_error("test error"); // emit_error sets status to Finished, so emit_finish is a no-op here. // We emit directly for the Finish broadcast path test: - agent - .runtime - .emit(AgentStreamEvent::Finish(FinishEventData { session_id: None })); + agent.runtime.emit(AgentStreamEvent::Finish(FinishEventData::default())); match rx.try_recv().unwrap() { AgentStreamEvent::Error(data) => assert_eq!(data.message, "test error"), @@ -343,3 +342,111 @@ async fn runtime_can_emit_error_and_finish() { other => panic!("Expected Finish, got {:?}", other), } } + +#[tokio::test] +async fn successful_finish_emits_only_usage_added_since_the_previous_run() { + let agent = AionrsAgentManager::new("conv-usage".into(), "/project".into(), make_test_config(), None) + .await + .unwrap(); + *agent.usage_watermark.lock().await = Some(TokenUsage { + input_tokens: 1_000, + output_tokens: 400, + ..Default::default() + }); + agent.runtime.reset_for_new_turn(ConversationStatus::Running); + let mut rx = agent.subscribe(); + + agent + .emit_successful_finish(TokenUsage { + input_tokens: 1_250, + output_tokens: 475, + ..Default::default() + }) + .await; + + match rx.recv().await.expect("finish event") { + AgentStreamEvent::Finish(data) => { + assert_eq!(data.input_tokens, Some(250)); + assert_eq!(data.output_tokens, Some(75)); + } + other => panic!("expected Finish, got {other:?}"), + } +} + +#[tokio::test] +async fn successful_finish_omits_usage_when_the_provider_total_did_not_advance() { + let agent = AionrsAgentManager::new("conv-no-usage".into(), "/project".into(), make_test_config(), None) + .await + .unwrap(); + *agent.usage_watermark.lock().await = Some(TokenUsage { + input_tokens: 1_000, + output_tokens: 400, + ..Default::default() + }); + agent.runtime.reset_for_new_turn(ConversationStatus::Running); + let mut rx = agent.subscribe(); + + agent + .emit_successful_finish(TokenUsage { + input_tokens: 1_000, + output_tokens: 400, + ..Default::default() + }) + .await; + + match rx.recv().await.expect("finish event") { + AgentStreamEvent::Finish(data) => { + assert_eq!(data.input_tokens, None); + assert_eq!(data.output_tokens, None); + } + other => panic!("expected Finish, got {other:?}"), + } +} + +#[tokio::test] +async fn failed_or_cancelled_turn_rebaselines_usage_before_reporting_again() { + let agent = AionrsAgentManager::new("conv-rebaseline".into(), "/project".into(), make_test_config(), None) + .await + .unwrap(); + *agent.usage_watermark.lock().await = Some(TokenUsage { + input_tokens: 1_000, + output_tokens: 400, + ..Default::default() + }); + agent.invalidate_usage_watermark().await; + agent.runtime.reset_for_new_turn(ConversationStatus::Running); + let mut rx = agent.subscribe(); + + agent + .emit_successful_finish(TokenUsage { + input_tokens: 1_250, + output_tokens: 475, + ..Default::default() + }) + .await; + + match rx.recv().await.expect("rebaseline finish event") { + AgentStreamEvent::Finish(data) => { + assert_eq!(data.input_tokens, None); + assert_eq!(data.output_tokens, None); + } + other => panic!("expected Finish, got {other:?}"), + } + + agent.runtime.reset_for_new_turn(ConversationStatus::Running); + agent + .emit_successful_finish(TokenUsage { + input_tokens: 1_400, + output_tokens: 525, + ..Default::default() + }) + .await; + + match rx.recv().await.expect("post-rebaseline finish event") { + AgentStreamEvent::Finish(data) => { + assert_eq!(data.input_tokens, Some(150)); + assert_eq!(data.output_tokens, Some(50)); + } + other => panic!("expected Finish, got {other:?}"), + } +} diff --git a/crates/aionui-ai-agent/src/protocol/events/mod.rs b/crates/aionui-ai-agent/src/protocol/events/mod.rs index 4a10492a5..db01689c0 100644 --- a/crates/aionui-ai-agent/src/protocol/events/mod.rs +++ b/crates/aionui-ai-agent/src/protocol/events/mod.rs @@ -159,6 +159,10 @@ pub enum TipType { pub struct FinishEventData { #[serde(default)] pub session_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub input_tokens: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub output_tokens: Option, } /// Kind of CodeBuddy ACP dialect signal absorbed by the tolerant transport @@ -328,6 +332,7 @@ mod tests { fn finish_event_roundtrip() { let event = AgentStreamEvent::Finish(FinishEventData { session_id: Some("sess-abc".into()), + ..Default::default() }); let json = serde_json::to_value(&event).unwrap(); assert_eq!(json["type"], "finish"); diff --git a/crates/aionui-ai-agent/src/session_agent.rs b/crates/aionui-ai-agent/src/session_agent.rs index b41d5e7c6..acd71753b 100644 --- a/crates/aionui-ai-agent/src/session_agent.rs +++ b/crates/aionui-ai-agent/src/session_agent.rs @@ -171,6 +171,7 @@ impl SessionRuntime { } let _ = self.tx.send(AgentStreamEvent::Finish(FinishEventData { session_id: self.session_id(), + ..Default::default() })); } } diff --git a/crates/aionui-channel/src/message_service.rs b/crates/aionui-channel/src/message_service.rs index f09e37edb..065231221 100644 --- a/crates/aionui-channel/src/message_service.rs +++ b/crates/aionui-channel/src/message_service.rs @@ -523,7 +523,10 @@ mod tests { #[test] fn finish_event_produces_finish() { - let event = AgentStreamEvent::Finish(FinishEventData { session_id: None }); + let event = AgentStreamEvent::Finish(FinishEventData { + session_id: None, + ..Default::default() + }); let action = ChannelMessageService::process_stream_event(&event); assert!(matches!(action, Some(StreamAction::Finish))); } diff --git a/crates/aionui-channel/tests/stream_relay_test.rs b/crates/aionui-channel/tests/stream_relay_test.rs index 42135555f..9f38f8b35 100644 --- a/crates/aionui-channel/tests/stream_relay_test.rs +++ b/crates/aionui-channel/tests/stream_relay_test.rs @@ -52,7 +52,10 @@ async fn relay_sends_thinking_then_final_message() { })) .unwrap(); event_tx - .send(AgentStreamEvent::Finish(FinishEventData { session_id: None })) + .send(AgentStreamEvent::Finish(FinishEventData { + session_id: None, + ..Default::default() + })) .unwrap(); relay.run(rx).await; @@ -131,7 +134,10 @@ async fn weixin_flushes_pending_text_before_tool_call() { })) .unwrap(); event_tx - .send(AgentStreamEvent::Finish(FinishEventData { session_id: None })) + .send(AgentStreamEvent::Finish(FinishEventData { + session_id: None, + ..Default::default() + })) .unwrap(); relay.run(rx).await; @@ -184,7 +190,10 @@ async fn telegram_does_not_flush_text_before_tool_call() { })) .unwrap(); event_tx - .send(AgentStreamEvent::Finish(FinishEventData { session_id: None })) + .send(AgentStreamEvent::Finish(FinishEventData { + session_id: None, + ..Default::default() + })) .unwrap(); relay.run(rx).await; @@ -222,7 +231,10 @@ async fn weixin_skips_flush_when_buffer_is_empty() { })) .unwrap(); event_tx - .send(AgentStreamEvent::Finish(FinishEventData { session_id: None })) + .send(AgentStreamEvent::Finish(FinishEventData { + session_id: None, + ..Default::default() + })) .unwrap(); relay.run(rx).await; diff --git a/crates/aionui-conversation/src/stream_relay.rs b/crates/aionui-conversation/src/stream_relay.rs index 943b26b5b..dae682a81 100644 --- a/crates/aionui-conversation/src/stream_relay.rs +++ b/crates/aionui-conversation/src/stream_relay.rs @@ -2336,7 +2336,12 @@ mod tests { let mut ws_rx = bus.subscribe(); let rx = tx.subscribe(); - tx.send(AgentStreamEvent::Finish(FinishEventData::default())).unwrap(); + tx.send(AgentStreamEvent::Finish(FinishEventData { + input_tokens: Some(1_200), + output_tokens: Some(345), + ..Default::default() + })) + .unwrap(); let outcome = relay.consume(rx).await; assert!(outcome.system_responses.is_empty()); @@ -2362,6 +2367,8 @@ mod tests { .find(|e| e.name == "message.stream") .expect("finish should be forwarded as message.stream"); assert_eq!(stream_event.data["turn_id"], "turn-1"); + assert_eq!(stream_event.data["data"]["input_tokens"], 1_200); + assert_eq!(stream_event.data["data"]["output_tokens"], 345); } // ── Tool persistence tests ──────────────────────────────────── diff --git a/crates/aionui-team/tests/e2e_team_flow.rs b/crates/aionui-team/tests/e2e_team_flow.rs index 66cbf4ff9..8d766d1a1 100644 --- a/crates/aionui-team/tests/e2e_team_flow.rs +++ b/crates/aionui-team/tests/e2e_team_flow.rs @@ -390,9 +390,10 @@ impl RecordingAgent { /// Fire a Finish event on the agent's stream (simulates agent completing a turn). #[allow(dead_code)] fn fire_finish(&self) { - let _ = self - .event_tx - .send(AgentStreamEvent::Finish(FinishEventData { session_id: None })); + let _ = self.event_tx.send(AgentStreamEvent::Finish(FinishEventData { + session_id: None, + ..Default::default() + })); } }