Stream token usage on every agent message (#5342)

crates/goose-server/src/openapi.rs

@@ -19,8 +19,8 @@ use goose::config::declarative_providers::{
 };
 use goose::conversation::message::{
     FrontendToolRequest, Message, MessageContent, MessageMetadata, RedactedThinkingContent,
-    SystemNotificationContent, SystemNotificationType, ThinkingContent, ToolConfirmationRequest,
-    ToolRequest, ToolResponse,
+    SystemNotificationContent, SystemNotificationType, ThinkingContent, TokenState,
+    ToolConfirmationRequest, ToolRequest, ToolResponse,
 };
 
 use crate::routes::reply::MessageEvent;
@@ -404,6 +404,7 @@ derive_utoipa!(Icon as IconSchema);
         Message,
         MessageContent,
         MessageMetadata,
+        TokenState,
         ContentSchema,
         EmbeddedResourceSchema,
         ImageContentSchema,

crates/goose-server/src/routes/reply.rs

@@ -8,7 +8,7 @@ use axum::{
 };
 use bytes::Bytes;
 use futures::{stream::StreamExt, Stream};
-use goose::conversation::message::{Message, MessageContent};
+use goose::conversation::message::{Message, MessageContent, TokenState};
 use goose::conversation::Conversation;
 use goose::permission::{Permission, PermissionConfirmation};
 use goose::session::SessionManager;
@@ -126,6 +126,7 @@ impl IntoResponse for SseResponse {
 pub enum MessageEvent {
     Message {
         message: Message,
+        token_state: TokenState,
     },
     Error {
         error: String,
@@ -159,6 +160,7 @@ async fn stream_event(
             e
         )
     });
+
     if tx.send(format!("data: {}\n\n", json)).await.is_err() {
         tracing::info!("client hung up");
         cancel_token.cancel();
@@ -305,7 +307,32 @@ pub async fn reply(
                             }
 
                             all_messages.push(message.clone());
-                            stream_event(MessageEvent::Message { message }, &tx, &cancel_token).await;
+
+                            let token_state = match SessionManager::get_session(&session_id, false).await {
+                                Ok(session) => {
+                                    TokenState {
+                                        input_tokens: session.input_tokens.unwrap_or(0),
+                                        output_tokens: session.output_tokens.unwrap_or(0),
+                                        total_tokens: session.total_tokens.unwrap_or(0),
+                                        accumulated_input_tokens: session.accumulated_input_tokens.unwrap_or(0),
+                                        accumulated_output_tokens: session.accumulated_output_tokens.unwrap_or(0),
+                                        accumulated_total_tokens: session.accumulated_total_tokens.unwrap_or(0),
+                                    }
+                                },
+                                Err(e) => {
+                                    tracing::warn!("Failed to fetch session for token state: {}", e);
+                                    TokenState {
+                                        input_tokens: 0,
+                                        output_tokens: 0,
+                                        total_tokens: 0,
+                                        accumulated_input_tokens: 0,
+                                        accumulated_output_tokens: 0,
+                                        accumulated_total_tokens: 0,
+                                    }
+                                }
+                            };
+
+                            stream_event(MessageEvent::Message { message, token_state }, &tx, &cancel_token).await;
                         }
                         Ok(Some(Ok(AgentEvent::HistoryReplaced(new_messages)))) => {
                             all_messages = new_messages.clone();