Stream token usage on every agent message

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;
@@ -402,6 +402,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();

crates/goose/src/agents/agent.rs

@@ -834,9 +834,11 @@ impl Agent {
                     }
                 }
                 Err(e) => {
-                    yield AgentEvent::Message(Message::assistant().with_text(
-                        format!("Ran into this error trying to compact: {e}.\n\nPlease try again or create a new session")
-                    ));
+                    yield AgentEvent::Message(
+                        Message::assistant().with_text(
+                            format!("Ran into this error trying to compact: {e}.\n\nPlease try again or create a new session")
+                        )
+                    );
                 }
             }
         }))
@@ -926,7 +928,7 @@ impl Agent {
                 if let Some(final_output_tool) = self.final_output_tool.lock().await.as_ref() {
                     if final_output_tool.final_output.is_some() {
                         let final_event = AgentEvent::Message(
-                            Message::assistant().with_text(final_output_tool.final_output.clone().unwrap()),
+                            Message::assistant().with_text(final_output_tool.final_output.clone().unwrap())
                         );
                         yield final_event;
                         break;
@@ -935,9 +937,11 @@ impl Agent {
 
                 turns_taken += 1;
                 if turns_taken > max_turns {
-                    yield AgentEvent::Message(Message::assistant().with_text(
-                        "I've reached the maximum number of actions I can do without user input. Would you like me to continue?"
-                    ));
+                    yield AgentEvent::Message(
+                        Message::assistant().with_text(
+                            "I've reached the maximum number of actions I can do without user input. Would you like me to continue?"
+                        )
+                    );
                     break;
                 }
 
@@ -1187,18 +1191,22 @@ impl Agent {
                                 }
                                 Err(e) => {
                                     error!("Error: {}", e);
-                                    yield AgentEvent::Message(Message::assistant().with_text(
+                                    yield AgentEvent::Message(
+                                        Message::assistant().with_text(
                                             format!("Ran into this error trying to compact: {e}.\n\nPlease retry if you think this is a transient or recoverable error.")
-                                        ));
+                                        )
+                                    );
                                     break;
                                 }
                             }
                         }
                         Err(e) => {
                             error!("Error: {}", e);
-                            yield AgentEvent::Message(Message::assistant().with_text(
+                            yield AgentEvent::Message(
+                                Message::assistant().with_text(
                                     format!("Ran into this error: {e}.\n\nPlease retry if you think this is a transient or recoverable error.")
-                                ));
+                                )
+                            );
                             break;
                         }
                     }
@@ -1233,9 +1241,11 @@ impl Agent {
                             }
                             Err(e) => {
                                 error!("Retry logic failed: {}", e);
-                                yield AgentEvent::Message(Message::assistant().with_text(
-                                    format!("Retry logic encountered an error: {}", e)
-                                ));
+                                yield AgentEvent::Message(
+                                    Message::assistant().with_text(
+                                        format!("Retry logic encountered an error: {}", e)
+                                    )
+                                );
                                 exit_chat = true;
                             }
                         }

crates/goose/src/conversation/message.rs

@@ -711,6 +711,17 @@ impl Message {
     }
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
+#[serde(rename_all = "camelCase")]
+pub struct TokenState {
+    pub input_tokens: i32,
+    pub output_tokens: i32,
+    pub total_tokens: i32,
+    pub accumulated_input_tokens: i32,
+    pub accumulated_output_tokens: i32,
+    pub accumulated_total_tokens: i32,
+}
+
 #[cfg(test)]
 mod tests {
     use crate::conversation::message::{Message, MessageContent, MessageMetadata};

crates/goose/src/providers/base.rs

@@ -278,11 +278,11 @@ impl Add for Usage {
     type Output = Self;
 
     fn add(self, other: Self) -> Self {
-        Self {
-            input_tokens: sum_optionals(self.input_tokens, other.input_tokens),
-            output_tokens: sum_optionals(self.output_tokens, other.output_tokens),
-            total_tokens: sum_optionals(self.total_tokens, other.total_tokens),
-        }
+        Self::new(
+            sum_optionals(self.input_tokens, other.input_tokens),
+            sum_optionals(self.output_tokens, other.output_tokens),
+            sum_optionals(self.total_tokens, other.total_tokens),
+        )
     }
 }
 
@@ -298,10 +298,21 @@ impl Usage {
         output_tokens: Option<i32>,
         total_tokens: Option<i32>,
     ) -> Self {
+        let calculated_total = if total_tokens.is_none() {
+            match (input_tokens, output_tokens) {
+                (Some(input), Some(output)) => Some(input + output),
+                (Some(input), None) => Some(input),
+                (None, Some(output)) => Some(output),
+                (None, None) => None,
+            }
+        } else {
+            total_tokens
+        };
+
         Self {
             input_tokens,
             output_tokens,
-            total_tokens,
+            total_tokens: calculated_total,
         }
     }
 }

crates/goose/src/providers/formats/bedrock.rs

@@ -345,11 +345,11 @@ pub fn from_bedrock_role(role: &bedrock::ConversationRole) -> Result<Role> {
 }
 
 pub fn from_bedrock_usage(usage: &bedrock::TokenUsage) -> Usage {
-    Usage {
-        input_tokens: Some(usage.input_tokens),
-        output_tokens: Some(usage.output_tokens),
-        total_tokens: Some(usage.total_tokens),
-    }
+    Usage::new(
+        Some(usage.input_tokens),
+        Some(usage.output_tokens),
+        Some(usage.total_tokens),
+    )
 }
 
 pub fn from_bedrock_json(document: &Document) -> Result<Value> {

crates/goose/src/providers/sagemaker_tgi.rs

@@ -307,11 +307,11 @@ impl Provider for SageMakerTgiProvider {
         let message = self.parse_tgi_response(response)?;
 
         // TGI doesn't provide usage statistics, so we estimate
-        let usage = Usage {
-            input_tokens: Some(0),  // Would need to tokenize input to get accurate count
-            output_tokens: Some(0), // Would need to tokenize output to get accurate count
-            total_tokens: Some(0),
-        };
+        let usage = Usage::new(
+            Some(0), // Would need to tokenize input to get accurate count
+            Some(0), // Would need to tokenize output to get accurate count
+            Some(0),
+        );
 
         // Add debug trace
         let debug_payload = serde_json::json!({

crates/goose/src/providers/venice.rs

@@ -508,11 +508,11 @@ impl Provider for VeniceProvider {
 
         // Extract usage
         let usage_data = &response_json["usage"];
-        let usage = Usage {
-            input_tokens: usage_data["prompt_tokens"].as_i64().map(|v| v as i32),
-            output_tokens: usage_data["completion_tokens"].as_i64().map(|v| v as i32),
-            total_tokens: usage_data["total_tokens"].as_i64().map(|v| v as i32),
-        };
+        let usage = Usage::new(
+            usage_data["prompt_tokens"].as_i64().map(|v| v as i32),
+            usage_data["completion_tokens"].as_i64().map(|v| v as i32),
+            usage_data["total_tokens"].as_i64().map(|v| v as i32),
+        );
 
         Ok((
             Message::new(Role::Assistant, Utc::now().timestamp(), content),

ui/desktop/openapi.json

@@ -3185,12 +3185,16 @@
             "type": "object",
             "required": [
               "message",
+              "token_state",
               "type"
             ],
             "properties": {
               "message": {
                 "$ref": "#/components/schemas/Message"
               },
+              "token_state": {
+                "$ref": "#/components/schemas/TokenState"
+              },
               "type": {
                 "type": "string",
                 "enum": [
@@ -4418,6 +4422,43 @@
           }
         }
       },
+      "TokenState": {
+        "type": "object",
+        "required": [
+          "inputTokens",
+          "outputTokens",
+          "totalTokens",
+          "accumulatedInputTokens",
+          "accumulatedOutputTokens",
+          "accumulatedTotalTokens"
+        ],
+        "properties": {
+          "accumulatedInputTokens": {
+            "type": "integer",
+            "format": "int32"
+          },
+          "accumulatedOutputTokens": {
+            "type": "integer",
+            "format": "int32"
+          },
+          "accumulatedTotalTokens": {
+            "type": "integer",
+            "format": "int32"
+          },
+          "inputTokens": {
+            "type": "integer",
+            "format": "int32"
+          },
+          "outputTokens": {
+            "type": "integer",
+            "format": "int32"
+          },
+          "totalTokens": {
+            "type": "integer",
+            "format": "int32"
+          }
+        }
+      },
       "Tool": {
         "type": "object",
         "required": [

ui/desktop/src/api/types.gen.ts

@@ -362,6 +362,7 @@ export type MessageContent = (TextContent & {
 
 export type MessageEvent = {
     message: Message;
+    token_state: TokenState;
     type: 'Message';
 } | {
     error: string;
@@ -779,6 +780,15 @@ export type ThinkingContent = {
     thinking: string;
 };
 
+export type TokenState = {
+    accumulatedInputTokens: number;
+    accumulatedOutputTokens: number;
+    accumulatedTotalTokens: number;
+    inputTokens: number;
+    outputTokens: number;
+    totalTokens: number;
+};
+
 export type Tool = {
     annotations?: ToolAnnotations | {
         [key: string]: unknown;

ui/desktop/src/components/BaseChat.tsx

@@ -132,6 +132,7 @@ function BaseChatContent({
     sessionOutputTokens,
     localInputTokens,
     localOutputTokens,
+    tokenState,
     commandHistory,
     toolCallNotifications,
     sessionMetadata,
@@ -442,9 +443,13 @@ function BaseChatContent({
             commandHistory={commandHistory}
             initialValue={input || ''}
             setView={setView}
-            numTokens={sessionTokenCount}
-            inputTokens={sessionInputTokens || localInputTokens}
-            outputTokens={sessionOutputTokens || localOutputTokens}
+            totalTokens={tokenState?.totalTokens ?? sessionTokenCount}
+            accumulatedInputTokens={
+              tokenState?.accumulatedInputTokens ?? sessionInputTokens ?? localInputTokens
+            }
+            accumulatedOutputTokens={
+              tokenState?.accumulatedOutputTokens ?? sessionOutputTokens ?? localOutputTokens
+            }
             droppedFiles={droppedFiles}
             onFilesProcessed={() => setDroppedFiles([])} // Clear dropped files after processing
             messages={messages}