fix(agent): forward AionRS provider token usage

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<String>) {
+        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<String>, 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();

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]

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();
 

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<AgentEngine>,
+    /// Last cumulative engine usage already projected onto per-turn finish events.
+    usage_watermark: Mutex<Option<TokenUsage>>,
     /// Static slash command metadata captured at bootstrap so UI lookups do
     /// not wait behind an active `engine.run()` turn.
     slash_commands: Vec<SlashCommandItem>,
@@ -143,6 +146,10 @@ impl AionrsAgentManager {
     ) -> Result<Self, AgentError> {
         let runtime = AgentRuntime::new(conversation_id.clone(), workspace.clone(), 128);
         let sink: Arc<dyn OutputSink> = 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<AgentKillReason>, 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(())
             }

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:?}"),
+    }
+}

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<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub input_tokens: Option<u64>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub output_tokens: Option<u64>,
 }
 
 /// 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");

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()
         }));
     }
 }

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)));
     }