feat(acp): emit session usage updates for context window tracking

Signed-off-by: Taylor Ho <taylorkmho@gmail.com>

crates/goose-acp/src/server.rs

@@ -40,7 +40,7 @@ use sacp::schema::{
     SetSessionConfigOptionResponse, SetSessionModeRequest, SetSessionModeResponse,
     SetSessionModelRequest, SetSessionModelResponse, StopReason, TextContent, TextResourceContents,
     ToolCall, ToolCallContent, ToolCallId, ToolCallLocation, ToolCallStatus, ToolCallUpdate,
-    ToolCallUpdateFields, ToolKind,
+    ToolCallUpdateFields, ToolKind, Usage, UsageUpdate,
 };
 use sacp::util::MatchDispatchFrom;
 use sacp::{
@@ -595,6 +595,27 @@ fn build_config_options(
     ]
 }
 
+fn to_nonnegative_u64(value: Option<i32>) -> Option<u64> {
+    value.and_then(|v| u64::try_from(v).ok())
+}
+
+fn build_prompt_usage(session: &Session) -> Option<Usage> {
+    let total = to_nonnegative_u64(session.accumulated_total_tokens)
+        .or_else(|| to_nonnegative_u64(session.total_tokens))?;
+    let input = to_nonnegative_u64(session.accumulated_input_tokens)
+        .or_else(|| to_nonnegative_u64(session.input_tokens))
+        .unwrap_or(0);
+    let output = to_nonnegative_u64(session.accumulated_output_tokens)
+        .or_else(|| to_nonnegative_u64(session.output_tokens))
+        .unwrap_or(0);
+    Some(Usage::new(total, input, output))
+}
+
+fn build_usage_update(session: &Session, context_limit: usize) -> UsageUpdate {
+    let used = session.total_tokens.unwrap_or(0).max(0) as u64;
+    UsageUpdate::new(used, context_limit as u64)
+}
+
 impl GooseAcpAgent {
     pub fn permission_manager(&self) -> Arc<PermissionManager> {
         Arc::clone(&self.permission_manager)
@@ -1324,27 +1345,38 @@ impl GooseAcpAgent {
         // Resolve provider + model from config so we can include the current
         // model in the response without waiting for the full agent setup.
         let resolved = resolve_provider_and_model(&self.config_dir, &goose_session).await;
+        let initial_usage_update = resolved
+            .as_ref()
+            .ok()
+            .map(|(_, mc)| build_usage_update(&goose_session, mc.context_limit()));
         let (model_state, config_options) =
             build_eager_config(&resolved, &mode_state, &goose_session).await;
+        let session_id = SessionId::new(thread_id.clone());
 
         self.spawn_agent_setup(
             cx,
             agent_tx,
             AgentSetupRequest {
-                session_id: SessionId::new(thread_id.clone()),
+                session_id: session_id.clone(),
                 goose_session,
                 mcp_servers: args.mcp_servers,
-                resolved_provider: resolved.ok(),
+                resolved_provider: resolved.as_ref().ok().cloned(),
             },
         );
 
-        let mut response = NewSessionResponse::new(SessionId::new(thread_id)).modes(mode_state);
+        let mut response = NewSessionResponse::new(session_id.clone()).modes(mode_state);
         if let Some(ms) = model_state {
             response = response.models(ms);
         }
         if let Some(co) = config_options {
             response = response.config_options(co);
         }
+        if let Some(usage_update) = initial_usage_update {
+            cx.send_notification(SessionNotification::new(
+                session_id,
+                SessionUpdate::UsageUpdate(usage_update),
+            ))?;
+        }
         Ok(response)
     }
 
@@ -1653,6 +1685,16 @@ impl GooseAcpAgent {
         let mode_state = build_mode_state(loaded_mode)?;
 
         let resolved = resolve_provider_and_model(&self.config_dir, &goose_session).await;
+        let initial_usage_update = resolved
+            .as_ref()
+            .ok()
+            .map(|(_, mc)| build_usage_update(&goose_session, mc.context_limit()))
+            .or_else(|| {
+                goose_session
+                    .model_config
+                    .as_ref()
+                    .map(|mc| build_usage_update(&goose_session, mc.context_limit()))
+            });
         let (model_state, config_options) =
             build_eager_config(&resolved, &mode_state, &goose_session).await;
 
@@ -1674,6 +1716,12 @@ impl GooseAcpAgent {
         if let Some(co) = config_options {
             response = response.config_options(co);
         }
+        if let Some(usage_update) = initial_usage_update {
+            cx.send_notification(SessionNotification::new(
+                args.session_id.clone(),
+                SessionUpdate::UsageUpdate(usage_update),
+            ))?;
+        }
         Ok(response)
     }
 
@@ -1759,15 +1807,41 @@ impl GooseAcpAgent {
             }
         }
 
-        let mut sessions = self.sessions.lock().await;
-        if let Some(session) = sessions.get_mut(&thread_id) {
-            session.cancel_token = None;
+        {
+            let mut sessions = self.sessions.lock().await;
+            if let Some(session) = sessions.get_mut(&thread_id) {
+                session.cancel_token = None;
+            }
         }
-        Ok(PromptResponse::new(if was_cancelled {
+
+        let session = self
+            .session_manager
+            .get_session(&internal_session_id, false)
+            .await
+            .map_err(|e| {
+                sacp::Error::internal_error().data(format!("Failed to load session: {}", e))
+            })?;
+        let provider = agent.provider().await.map_err(|e| {
+            sacp::Error::internal_error().data(format!("Failed to get provider: {}", e))
+        })?;
+        let usage_update =
+            build_usage_update(&session, provider.get_model_config().context_limit());
+        cx.send_notification(SessionNotification::new(
+            args.session_id.clone(),
+            SessionUpdate::UsageUpdate(usage_update),
+        ))?;
+
+        let stop_reason = if was_cancelled {
             StopReason::Cancelled
         } else {
             StopReason::EndTurn
-        }))
+        };
+
+        let mut response = PromptResponse::new(stop_reason);
+        if let Some(usage) = build_prompt_usage(&session) {
+            response = response.usage(usage);
+        }
+        Ok(response)
     }
 
     async fn on_cancel(&self, args: CancelNotification) -> Result<(), sacp::Error> {
@@ -2909,6 +2983,7 @@ pub async fn run(builtins: Vec<String>) -> Result<()> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use chrono::Utc;
     use goose::conversation::message::{ToolRequest, ToolResponse};
     use goose::providers::errors::ProviderError;
     use rmcp::model::{CallToolRequestParams, Content as RmcpContent};
@@ -3294,6 +3369,80 @@ print(\"hello, world\")
             .map(|locs| locs.into_iter().map(|loc| (loc.path, loc.line)).collect())
     }
 
+    fn make_session_with_usage(
+        total_tokens: Option<i32>,
+        input_tokens: Option<i32>,
+        output_tokens: Option<i32>,
+        accumulated_total_tokens: Option<i32>,
+        accumulated_input_tokens: Option<i32>,
+        accumulated_output_tokens: Option<i32>,
+    ) -> Session {
+        Session {
+            id: "session-1".to_string(),
+            working_dir: PathBuf::from("/tmp"),
+            name: "ACP Session".to_string(),
+            user_set_name: false,
+            session_type: SessionType::Acp,
+            created_at: Utc::now(),
+            updated_at: Utc::now(),
+            extension_data: goose::session::ExtensionData::default(),
+            total_tokens,
+            input_tokens,
+            output_tokens,
+            accumulated_total_tokens,
+            accumulated_input_tokens,
+            accumulated_output_tokens,
+            schedule_id: None,
+            recipe: None,
+            user_recipe_values: None,
+            conversation: None,
+            message_count: 0,
+            provider_name: None,
+            model_config: None,
+            goose_mode: GooseMode::default(),
+            thread_id: None,
+        }
+    }
+
+    #[test]
+    fn test_build_prompt_usage_prefers_accumulated_tokens() {
+        let session = make_session_with_usage(
+            Some(120),
+            Some(80),
+            Some(40),
+            Some(360),
+            Some(210),
+            Some(150),
+        );
+        let usage = build_prompt_usage(&session).expect("usage should be present");
+        assert_eq!(usage.total_tokens, 360);
+        assert_eq!(usage.input_tokens, 210);
+        assert_eq!(usage.output_tokens, 150);
+    }
+
+    #[test]
+    fn test_build_prompt_usage_falls_back_to_current_tokens() {
+        let session = make_session_with_usage(Some(120), Some(80), Some(40), None, None, None);
+        let usage = build_prompt_usage(&session).expect("usage should be present");
+        assert_eq!(usage.total_tokens, 120);
+        assert_eq!(usage.input_tokens, 80);
+        assert_eq!(usage.output_tokens, 40);
+    }
+
+    #[test]
+    fn test_build_prompt_usage_requires_total_tokens() {
+        let session = make_session_with_usage(None, Some(80), Some(40), None, None, None);
+        assert!(build_prompt_usage(&session).is_none());
+    }
+
+    #[test]
+    fn test_build_usage_update_clamps_negative_used_to_zero() {
+        let session = make_session_with_usage(Some(-7), Some(0), Some(0), None, None, None);
+        let usage = build_usage_update(&session, 258_000);
+        assert_eq!(usage.used, 0);
+        assert_eq!(usage.size, 258_000);
+    }
+
     #[test_case(
         GooseMode::Auto
         => Ok(SessionModeState::new(