feat: goose2 context window usage in chat input (aaif-goose#8613)

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

Author: tellaho

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::{
@@ -603,6 +603,22 @@ 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.total_tokens)?;
+    let input = to_nonnegative_u64(session.input_tokens).unwrap_or(0);
+    let output = 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)
@@ -1383,27 +1399,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),
+            ))?;
+        }
         debug!(
             target: "perf",
             sid = %sid,
@@ -1748,6 +1775,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;
 
@@ -1769,6 +1806,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),
+            ))?;
+        }
         debug!(
             target: "perf",
             sid = %sid,
@@ -1883,10 +1926,30 @@ 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;
+            }
         }
+
+        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),
+        ))?;
+
         debug!(
             target: "perf",
             sid = %sid,
@@ -1895,11 +1958,17 @@ impl GooseAcpAgent {
             cancelled = was_cancelled,
             "perf: prompt done"
         );
-        Ok(PromptResponse::new(if was_cancelled {
+        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> {
@@ -3487,6 +3556,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: Default::default(),
+            updated_at: Default::default(),
+            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_uses_current_turn_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, 120);
+        assert_eq!(usage.input_tokens, 80);
+        assert_eq!(usage.output_tokens, 40);
+    }
+
+    #[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(

ui/goose2/justfile

@@ -88,10 +88,24 @@ dev:
     # Override with e.g. RUST_LOG=info just dev to disable.
     export RUST_LOG="${RUST_LOG:-perf=debug,info}"
     PROJECT_DIR=$(pwd)
-    GOOSE_BIN="${PROJECT_DIR}/../../target/debug/goose"
-    export GOOSE_BIN
+    REPO_ROOT=$(cd ../.. && pwd)
+    LOCAL_GOOSE_DEBUG="${REPO_ROOT}/target/debug/goose"
+    LOCAL_GOOSE_RELEASE="${REPO_ROOT}/target/release/goose"
+    if [[ -x "${LOCAL_GOOSE_DEBUG}" ]]; then
+        export GOOSE_BIN="${LOCAL_GOOSE_DEBUG}"
+    elif [[ -x "${LOCAL_GOOSE_RELEASE}" ]]; then
+        export GOOSE_BIN="${LOCAL_GOOSE_RELEASE}"
+    else
+        unset GOOSE_BIN
+    fi
     EXTRA_CONFIG_ARGS=(--config "{\"build\":{\"devUrl\":\"http://localhost:${VITE_PORT}\",\"beforeDevCommand\":{\"script\":\"cd ${PROJECT_DIR} && exec pnpm exec vite --port ${VITE_PORT} --strictPort\",\"cwd\":\".\",\"wait\":false}}}")
 
+    if [[ -n "${GOOSE_BIN:-}" ]]; then
+        echo "Using local goose binary: ${GOOSE_BIN}"
+    else
+        echo "No local goose binary found under ${REPO_ROOT}/target; falling back to PATH"
+    fi
+
     # In worktrees, generate a labeled icon so you can tell instances apart
     if git rev-parse --is-inside-work-tree &>/dev/null; then
         GIT_DIR=$(git rev-parse --git-dir)
@@ -117,14 +131,29 @@ dev-debug:
     #!/usr/bin/env bash
     set -euo pipefail
 
+    VITE_PORT={{ vite_port }}
+    export VITE_PORT
     # Enable perf logs in the child `goose serve` process by default.
     # Override with e.g. RUST_LOG=info just dev-debug to disable.
     export RUST_LOG="${RUST_LOG:-perf=debug,info}"
-    PROJECT_DIR=$(pwd)
-    GOOSE_BIN="${PROJECT_DIR}/../../target/debug/goose"
-    export GOOSE_BIN
+    REPO_ROOT=$(cd ../.. && pwd)
+    LOCAL_GOOSE_DEBUG="${REPO_ROOT}/target/debug/goose"
+    LOCAL_GOOSE_RELEASE="${REPO_ROOT}/target/release/goose"
+    if [[ -x "${LOCAL_GOOSE_DEBUG}" ]]; then
+        export GOOSE_BIN="${LOCAL_GOOSE_DEBUG}"
+    elif [[ -x "${LOCAL_GOOSE_RELEASE}" ]]; then
+        export GOOSE_BIN="${LOCAL_GOOSE_RELEASE}"
+    else
+        unset GOOSE_BIN
+    fi
     EXTRA_CONFIG_ARGS=(--config "{\"build\":{\"devUrl\":\"http://localhost:${VITE_PORT}\",\"beforeDevCommand\":{\"script\":\"exec ./node_modules/.bin/vite --port ${VITE_PORT} --strictPort\",\"cwd\":\"..\",\"wait\":false}}}")
 
+    if [[ -n "${GOOSE_BIN:-}" ]]; then
+        echo "Using local goose binary: ${GOOSE_BIN}"
+    else
+        echo "No local goose binary found under ${REPO_ROOT}/target; falling back to PATH"
+    fi
+
     # In worktrees, generate a labeled icon so you can tell instances apart
     if git rev-parse --is-inside-work-tree &>/dev/null; then
         GIT_DIR=$(git rev-parse --git-dir)

ui/goose2/scripts/check-file-sizes.mjs

@@ -13,7 +13,17 @@ const EXCEPTIONS = {
   "src/features/chat/ui/ChatView.tsx": {
     limit: 570,
     justification:
-      "ACP prewarm guards, project-aware working dir selection, working context sync, and chat bootstrapping still live together here. Includes gated [perf:chatview] logging via perfLog (dev-only by default).",
+      "ACP prewarm guards, project-aware working dir selection, working context sync, chat bootstrapping, context-ring compaction wiring, and gated [perf:chatview] logging via perfLog (dev-only by default).",
+  },
+  "src/features/chat/hooks/useChat.ts": {
+    limit: 510,
+    justification:
+      "Session preparation, provider/model handoff, persona-aware sends, cancellation, and compaction replay still live in one chat lifecycle hook.",
+  },
+  "src/shared/api/acpNotificationHandler.ts": {
+    limit: 550,
+    justification:
+      "ACP replay/live update handling, pending session buffering, model/config propagation, and streaming perf tracking still share one notification entrypoint.",
   },
   "src/features/chat/ui/__tests__/ContextPanel.test.tsx": {
     limit: 550,