Token state not showing on load, or after message is finished. (#5606)

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

@@ -133,6 +133,7 @@ pub enum MessageEvent {
     },
     Finish {
         reason: String,
+        token_state: TokenState,
     },
     ModelChange {
         model: String,
@@ -149,6 +150,27 @@ pub enum MessageEvent {
     Ping,
 }
 
+async fn get_token_state(session_id: &str) -> TokenState {
+    SessionManager::get_session(session_id, false)
+        .await
+        .map(|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),
+        })
+        .inspect_err(|e| {
+            tracing::warn!(
+                "Failed to fetch session token state for {}: {}",
+                session_id,
+                e
+            );
+        })
+        .unwrap_or_default()
+}
+
 async fn stream_event(
     event: MessageEvent,
     tx: &mpsc::Sender<String>,
@@ -321,29 +343,7 @@ pub async fn reply(
 
                             all_messages.push(message.clone());
 
-                            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,
-                                    }
-                                }
-                            };
+                            let token_state = get_token_state(&session_id).await;
 
                             stream_event(MessageEvent::Message { message, token_state }, &tx, &cancel_token).await;
                         }
@@ -437,9 +437,12 @@ pub async fn reply(
             );
         }
 
+        let final_token_state = get_token_state(&session_id).await;
+
         let _ = stream_event(
             MessageEvent::Finish {
                 reason: "stop".to_string(),
+                token_state: final_token_state,
             },
             &task_tx,
             &cancel_token,

crates/goose/src/conversation/message.rs

@@ -711,7 +711,7 @@ impl Message {
     }
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
+#[derive(Debug, Clone, Default, Serialize, Deserialize, ToSchema)]
 #[serde(rename_all = "camelCase")]
 pub struct TokenState {
     pub input_tokens: i32,