From a7d78ee59e6a0eded1f45426c7ad2f9e71ce5710 Mon Sep 17 00:00:00 2001 From: Taylor Ho Date: Sun, 19 Apr 2026 18:53:44 -1000 Subject: [PATCH] feat: goose2 context window usage in chat input (#8613) Signed-off-by: Taylor Ho --- crates/goose-acp/src/server.rs | 161 ++++++++++++- ui/goose2/justfile | 39 ++- ui/goose2/scripts/check-file-sizes.mjs | 12 +- .../__tests__/useChat.compaction.test.ts | 228 ++++++++++++++++++ .../chat/hooks/__tests__/useChat.test.ts | 20 +- ui/goose2/src/features/chat/hooks/useChat.ts | 107 +++++++- ui/goose2/src/features/chat/ui/ChatInput.tsx | 9 + .../src/features/chat/ui/ChatInputToolbar.tsx | 110 ++++++++- ui/goose2/src/features/chat/ui/ChatView.tsx | 10 +- .../src/features/chat/ui/ContextRing.tsx | 23 +- .../src/features/chat/ui/LoadingGoose.tsx | 4 +- .../chat/ui/__tests__/ChatInput.test.tsx | 43 ++++ .../chat/ui/__tests__/LoadingGoose.test.tsx | 9 +- .../src/shared/api/__tests__/acp.test.ts | 57 +++++ ui/goose2/src/shared/api/acp.ts | 17 +- ui/goose2/src/shared/api/acpApi.ts | 14 +- .../shared/api/acpNotificationHandler.test.ts | 76 ++++++ .../src/shared/api/acpNotificationHandler.ts | 70 +++++- ui/goose2/src/shared/api/acpSessionTracker.ts | 85 ++++++- .../src/shared/i18n/locales/en/chat.json | 6 + .../src/shared/i18n/locales/es/chat.json | 6 + ui/goose2/src/shared/ui/button.tsx | 2 +- 22 files changed, 1035 insertions(+), 73 deletions(-) create mode 100644 ui/goose2/src/features/chat/hooks/__tests__/useChat.compaction.test.ts create mode 100644 ui/goose2/src/shared/api/__tests__/acp.test.ts create mode 100644 ui/goose2/src/shared/api/acpNotificationHandler.test.ts diff --git a/crates/goose-acp/src/server.rs b/crates/goose-acp/src/server.rs index eeeae74df0..d1a8212c75 100644 --- a/crates/goose-acp/src/server.rs +++ b/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) -> Option { + value.and_then(|v| u64::try_from(v).ok()) +} + +fn build_prompt_usage(session: &Session) -> Option { + 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 { 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, @@ -1882,10 +1925,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, @@ -1894,11 +1957,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> { @@ -3486,6 +3555,80 @@ print(\"hello, world\") .map(|locs| locs.into_iter().map(|loc| (loc.path, loc.line)).collect()) } + fn make_session_with_usage( + total_tokens: Option, + input_tokens: Option, + output_tokens: Option, + accumulated_total_tokens: Option, + accumulated_input_tokens: Option, + accumulated_output_tokens: Option, + ) -> 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( diff --git a/ui/goose2/justfile b/ui/goose2/justfile index f02b33a8d3..11c2a6ca8d 100644 --- a/ui/goose2/justfile +++ b/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) diff --git a/ui/goose2/scripts/check-file-sizes.mjs b/ui/goose2/scripts/check-file-sizes.mjs index 5d2066a196..07e1d124f2 100644 --- a/ui/goose2/scripts/check-file-sizes.mjs +++ b/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, diff --git a/ui/goose2/src/features/chat/hooks/__tests__/useChat.compaction.test.ts b/ui/goose2/src/features/chat/hooks/__tests__/useChat.compaction.test.ts new file mode 100644 index 0000000000..09b1f6fd2f --- /dev/null +++ b/ui/goose2/src/features/chat/hooks/__tests__/useChat.compaction.test.ts @@ -0,0 +1,228 @@ +import { act, renderHook } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { Message } from "@/shared/types/messages"; +import { useChatStore } from "../../stores/chatStore"; +import { clearReplayBuffer, ensureReplayBuffer } from "../replayBuffer"; + +const mockAcpSendMessage = vi.fn(); +const mockAcpLoadSession = vi.fn(); +const mockGetGooseSessionId = vi.fn(); + +vi.mock("@/shared/api/acp", () => ({ + acpSendMessage: (...args: unknown[]) => mockAcpSendMessage(...args), + acpCancelSession: vi.fn(), + acpLoadSession: (...args: unknown[]) => mockAcpLoadSession(...args), + acpPrepareSession: vi.fn(), + acpSetModel: vi.fn(), +})); + +vi.mock("@/shared/api/acpSessionTracker", () => ({ + getGooseSessionId: (...args: unknown[]) => mockGetGooseSessionId(...args), +})); + +import { useChat } from "../useChat"; + +function createDeferredPromise() { + let resolve!: (value: T) => void; + const promise = new Promise((res) => { + resolve = res; + }); + return { promise, resolve }; +} + +function createTextMessage( + id: string, + role: Message["role"], + text: string, +): Message { + return { + id, + role, + created: 0, + content: [{ type: "text", text }], + metadata: { + userVisible: true, + agentVisible: role !== "system", + }, + }; +} + +describe("useChat compaction", () => { + beforeEach(() => { + mockAcpSendMessage.mockReset(); + mockAcpLoadSession.mockReset(); + mockGetGooseSessionId.mockReset(); + clearReplayBuffer("session-1"); + useChatStore.setState({ + messagesBySession: {}, + sessionStateById: {}, + activeSessionId: null, + isConnected: true, + loadingSessionIds: new Set(), + }); + mockAcpSendMessage.mockResolvedValue(undefined); + mockAcpLoadSession.mockResolvedValue(undefined); + mockGetGooseSessionId.mockReturnValue(null); + }); + + it("reloads compacted history after sending the compact command", async () => { + mockGetGooseSessionId.mockReturnValue("goose-session-1"); + mockAcpLoadSession.mockImplementation(async (sessionId: string) => { + const buffer = ensureReplayBuffer(sessionId); + buffer.push(createTextMessage("user-1", "user", "Before compact")); + buffer.push(createTextMessage("compact-1", "user", "/compact/compact")); + buffer.push( + createTextMessage("assistant-1", "assistant", "After compact"), + ); + }); + + useChatStore + .getState() + .setMessages("session-1", [ + createTextMessage("stale-1", "assistant", "Stale"), + ]); + + const { result } = renderHook(() => useChat("session-1")); + + await act(async () => { + await result.current.compactConversation(); + }); + + expect(mockAcpSendMessage).toHaveBeenCalledWith( + "session-1", + "/compact", + undefined, + ); + expect(mockAcpLoadSession).toHaveBeenCalledWith( + "session-1", + "goose-session-1", + undefined, + ); + + const messages = useChatStore.getState().messagesBySession["session-1"]; + const runtime = useChatStore.getState().getSessionRuntime("session-1"); + + expect(messages).toEqual([ + createTextMessage("user-1", "user", "Before compact"), + createTextMessage("assistant-1", "assistant", "After compact"), + ]); + expect(runtime.chatState).toBe("idle"); + expect(runtime.error).toBeNull(); + expect(useChatStore.getState().loadingSessionIds.has("session-1")).toBe( + false, + ); + }); + + it("blocks new sends while compaction is in flight", async () => { + mockGetGooseSessionId.mockReturnValue("goose-session-1"); + const compactDeferred = createDeferredPromise(); + mockAcpSendMessage.mockImplementation( + (_sessionId: string, prompt: string) => + prompt === "/compact" ? compactDeferred.promise : Promise.resolve(), + ); + + const { result } = renderHook(() => useChat("session-1")); + + let compactPromise!: Promise; + await act(async () => { + compactPromise = result.current.compactConversation(); + await Promise.resolve(); + }); + + expect( + useChatStore.getState().getSessionRuntime("session-1").chatState, + ).toBe("compacting"); + + await act(async () => { + await result.current.sendMessage("Hello during compact"); + }); + + expect(mockAcpSendMessage).toHaveBeenCalledTimes(1); + expect(mockAcpSendMessage).toHaveBeenCalledWith( + "session-1", + "/compact", + undefined, + ); + expect( + useChatStore.getState().messagesBySession["session-1"], + ).toBeUndefined(); + expect( + useChatStore.getState().getSessionRuntime("session-1").chatState, + ).toBe("compacting"); + + compactDeferred.resolve(); + await act(async () => { + await compactPromise; + }); + + expect( + useChatStore.getState().getSessionRuntime("session-1").chatState, + ).toBe("idle"); + }); + + it("ignores a second compact request while the first one is still in flight", async () => { + mockGetGooseSessionId.mockReturnValue("goose-session-1"); + const compactDeferred = createDeferredPromise(); + mockAcpSendMessage.mockImplementation( + (_sessionId: string, prompt: string) => + prompt === "/compact" ? compactDeferred.promise : Promise.resolve(), + ); + + const { result } = renderHook(() => useChat("session-1")); + + let firstCompact!: Promise; + let secondCompact!: Promise; + await act(async () => { + firstCompact = result.current.compactConversation(); + secondCompact = result.current.compactConversation(); + await Promise.resolve(); + }); + + expect(mockAcpSendMessage).toHaveBeenCalledTimes(1); + expect(mockAcpSendMessage).toHaveBeenCalledWith( + "session-1", + "/compact", + undefined, + ); + expect(mockAcpLoadSession).not.toHaveBeenCalled(); + expect( + useChatStore.getState().getSessionRuntime("session-1").chatState, + ).toBe("compacting"); + + compactDeferred.resolve(); + await act(async () => { + await Promise.all([firstCompact, secondCompact]); + }); + + expect(mockAcpLoadSession).toHaveBeenCalledTimes(1); + expect( + useChatStore.getState().getSessionRuntime("session-1").chatState, + ).toBe("idle"); + }); + + it("surfaces an error when compacting before the session is prepared", async () => { + const { result } = renderHook(() => useChat("session-1")); + + await act(async () => { + await result.current.compactConversation(); + }); + + expect(mockAcpSendMessage).not.toHaveBeenCalled(); + expect(mockAcpLoadSession).not.toHaveBeenCalled(); + + const messages = useChatStore.getState().messagesBySession["session-1"]; + const runtime = useChatStore.getState().getSessionRuntime("session-1"); + + expect(messages).toHaveLength(1); + expect(messages[0].content).toEqual([ + { + type: "systemNotification", + notificationType: "error", + text: "Session not prepared. Send a message before compacting.", + }, + ]); + expect(runtime.error).toBe( + "Session not prepared. Send a message before compacting.", + ); + }); +}); diff --git a/ui/goose2/src/features/chat/hooks/__tests__/useChat.test.ts b/ui/goose2/src/features/chat/hooks/__tests__/useChat.test.ts index 4c9509a5af..3dc918e7d6 100644 --- a/ui/goose2/src/features/chat/hooks/__tests__/useChat.test.ts +++ b/ui/goose2/src/features/chat/hooks/__tests__/useChat.test.ts @@ -4,19 +4,27 @@ import { useAgentStore } from "@/features/agents/stores/agentStore"; import { useChatStore } from "../../stores/chatStore"; import { useChatSessionStore } from "../../stores/chatSessionStore"; import type { Message } from "@/shared/types/messages"; +import { clearReplayBuffer } from "../replayBuffer"; const mockAcpSendMessage = vi.fn(); const mockAcpCancelSession = vi.fn(); +const mockAcpLoadSession = vi.fn(); const mockAcpPrepareSession = vi.fn(); const mockAcpSetModel = vi.fn(); +const mockGetGooseSessionId = vi.fn(); vi.mock("@/shared/api/acp", () => ({ acpSendMessage: (...args: unknown[]) => mockAcpSendMessage(...args), acpCancelSession: (...args: unknown[]) => mockAcpCancelSession(...args), + acpLoadSession: (...args: unknown[]) => mockAcpLoadSession(...args), acpPrepareSession: (...args: unknown[]) => mockAcpPrepareSession(...args), acpSetModel: (...args: unknown[]) => mockAcpSetModel(...args), })); +vi.mock("@/shared/api/acpSessionTracker", () => ({ + getGooseSessionId: (...args: unknown[]) => mockGetGooseSessionId(...args), +})); + import { useChat } from "../useChat"; function addStreamingAssistantMessage( @@ -53,7 +61,14 @@ function createDeferredPromise() { describe("useChat", () => { beforeEach(() => { - vi.clearAllMocks(); + mockAcpSendMessage.mockReset(); + mockAcpCancelSession.mockReset(); + mockAcpLoadSession.mockReset(); + mockAcpPrepareSession.mockReset(); + mockAcpSetModel.mockReset(); + mockGetGooseSessionId.mockReset(); + clearReplayBuffer("session-1"); + clearReplayBuffer("session-2"); useChatStore.setState({ messagesBySession: {}, sessionStateById: {}, @@ -96,9 +111,12 @@ describe("useChat", () => { personaEditorOpen: false, editingPersona: null, }); + mockAcpSendMessage.mockResolvedValue(undefined); mockAcpCancelSession.mockResolvedValue(true); + mockAcpLoadSession.mockResolvedValue(undefined); mockAcpPrepareSession.mockResolvedValue(undefined); mockAcpSetModel.mockResolvedValue(undefined); + mockGetGooseSessionId.mockReturnValue(null); }); it("cancels the active override persona instead of the hook default persona", async () => { diff --git a/ui/goose2/src/features/chat/hooks/useChat.ts b/ui/goose2/src/features/chat/hooks/useChat.ts index 3f16395a62..51af77f0dc 100644 --- a/ui/goose2/src/features/chat/hooks/useChat.ts +++ b/ui/goose2/src/features/chat/hooks/useChat.ts @@ -1,18 +1,23 @@ import { useCallback, useRef } from "react"; import { useChatStore } from "../stores/chatStore"; import { useChatSessionStore } from "../stores/chatSessionStore"; +import { clearReplayBuffer, getAndDeleteReplayBuffer } from "./replayBuffer"; import { type ChatAttachmentDraft, + type Message, createSystemNotificationMessage, createUserMessage, + getTextContent, } from "@/shared/types/messages"; import type { ChatState, TokenState } from "@/shared/types/chat"; import { acpSendMessage, acpCancelSession, + acpLoadSession, acpPrepareSession, acpSetModel, } from "@/shared/api/acp"; +import { getGooseSessionId } from "@/shared/api/acpSessionTracker"; import { useAgentStore } from "@/features/agents/stores/agentStore"; import { getSessionTitleFromDraft, @@ -26,6 +31,26 @@ import { buildMessageAttachments, } from "../lib/attachments"; +// TODO: Remove this fallback once goose2 has first-class /-commands. +const MANUAL_COMPACT_TRIGGER = "/compact"; + +function isManualCompactCommandMessage(message: Message): boolean { + if (message.role !== "user") { + return false; + } + + const normalizedText = getTextContent(message).replace(/\s+/g, ""); + if (!normalizedText) { + return false; + } + + return normalizedText.replaceAll(MANUAL_COMPACT_TRIGGER, "").length === 0; +} + +function removeManualCompactCommandMessages(messages: Message[]): Message[] { + return messages.filter((message) => !isManualCompactCommandMessage(message)); +} + function getErrorMessage(error: unknown): string { // Tauri command rejections typically arrive as plain strings, so handle // that shape first before falling back to standard Error objects. @@ -134,10 +159,14 @@ export function useChat( const tSendStart = performance.now(); const images = buildAcpImages(attachments); const hasAttachments = (attachments?.length ?? 0) > 0; + const currentChatState = useChatStore + .getState() + .getSessionRuntime(sessionId).chatState; if ( (!text.trim() && !hasAttachments) || - chatState === "streaming" || - chatState === "thinking" + currentChatState === "streaming" || + currentChatState === "thinking" || + currentChatState === "compacting" ) return; perfLog( @@ -318,7 +347,6 @@ export function useChat( }, [ sessionId, - chatState, store, providerOverride, systemPromptOverride, @@ -387,6 +415,78 @@ export function useChat( store.setPendingAssistantProvider(sessionId, null); }, [sessionId, store]); + const compactConversation = useCallback(async () => { + const currentChatState = useChatStore + .getState() + .getSessionRuntime(sessionId).chatState; + if (currentChatState !== "idle") { + return; + } + + const effectivePersonaInfo = resolvePersonaInfo(); + const gooseSessionId = getGooseSessionId( + sessionId, + effectivePersonaInfo?.id, + ); + + if (!gooseSessionId) { + const errorMessage = + "Session not prepared. Send a message before compacting."; + store.addMessage( + sessionId, + createSystemNotificationMessage(errorMessage, "error"), + ); + store.setError(sessionId, errorMessage); + return; + } + + store.setActiveSession(sessionId); + store.setChatState(sessionId, "compacting"); + store.setStreamingMessageId(sessionId, null); + store.setError(sessionId, null); + store.setSessionLoading(sessionId, true); + clearReplayBuffer(sessionId); + + try { + const sendOptions = effectivePersonaInfo?.id + ? { personaId: effectivePersonaInfo.id } + : undefined; + await acpSendMessage(sessionId, MANUAL_COMPACT_TRIGGER, sendOptions); + + // Command responses are streamed via prompt notifications, but the ACP + // layer does not currently forward history replacement events. Drop those + // transient chunks and refresh the session from replay instead. + clearReplayBuffer(sessionId); + const workingDir = await getWorkingDir?.(); + await acpLoadSession(sessionId, gooseSessionId, workingDir); + + store.setSessionLoading(sessionId, false); + + const buffer = getAndDeleteReplayBuffer(sessionId); + if (buffer) { + store.setMessages( + sessionId, + removeManualCompactCommandMessages(buffer), + ); + } + } catch (err) { + clearReplayBuffer(sessionId); + store.setSessionLoading(sessionId, false); + + const errorMessage = getErrorMessage(err); + store.addMessage( + sessionId, + createSystemNotificationMessage(errorMessage, "error"), + ); + store.setError(sessionId, errorMessage); + } finally { + store.setChatState(sessionId, "idle"); + store.setStreamingMessageId(sessionId, null); + store.setPendingAssistantProvider(sessionId, null); + store.setSessionLoading(sessionId, false); + } + }, [getWorkingDir, resolvePersonaInfo, sessionId, store]); + const stopStreaming = stopGeneration; return { @@ -400,6 +500,7 @@ export function useChat( stopStreaming, retryLastMessage, clearChat, + compactConversation, isStreaming, }; } diff --git a/ui/goose2/src/features/chat/ui/ChatInput.tsx b/ui/goose2/src/features/chat/ui/ChatInput.tsx index fdf2bea7d0..9b40f2b768 100644 --- a/ui/goose2/src/features/chat/ui/ChatInput.tsx +++ b/ui/goose2/src/features/chat/ui/ChatInput.tsx @@ -64,6 +64,9 @@ interface ChatInputProps { }) => void; contextTokens?: number; contextLimit?: number; + onCompactContext?: () => void | Promise; + canCompactContext?: boolean; + isCompactingContext?: boolean; } export function ChatInput({ @@ -94,6 +97,9 @@ export function ChatInput({ onCreateProject, contextTokens = 0, contextLimit = 0, + onCompactContext, + canCompactContext = false, + isCompactingContext = false, }: ChatInputProps) { const { t } = useTranslation("chat"); const [text, setTextRaw] = useState(initialValue); @@ -429,6 +435,9 @@ export function ChatInput({ onCreateProject={onCreateProject} contextTokens={contextTokens} contextLimit={contextLimit} + onCompactContext={onCompactContext} + canCompactContext={canCompactContext} + isCompactingContext={isCompactingContext} canSend={canSend} isStreaming={isStreaming} hasQueuedMessage={hasQueuedMessage} diff --git a/ui/goose2/src/features/chat/ui/ChatInputToolbar.tsx b/ui/goose2/src/features/chat/ui/ChatInputToolbar.tsx index 85cffd3790..3e25b8f084 100644 --- a/ui/goose2/src/features/chat/ui/ChatInputToolbar.tsx +++ b/ui/goose2/src/features/chat/ui/ChatInputToolbar.tsx @@ -1,4 +1,4 @@ -import { useMemo } from "react"; +import { useMemo, useState } from "react"; import { Mic, ArrowUp, @@ -24,6 +24,8 @@ import { DropdownMenuItem, DropdownMenuTrigger, } from "@/shared/ui/dropdown-menu"; +import { Popover, PopoverContent, PopoverTrigger } from "@/shared/ui/popover"; +import { Progress } from "@/shared/ui/progress"; import { Tooltip, TooltipTrigger, TooltipContent } from "@/shared/ui/tooltip"; import { AgentModelPicker } from "./AgentModelPicker"; import type { ModelOption } from "../types"; @@ -77,6 +79,9 @@ interface ChatInputToolbarProps { contextTokens: number; contextLimit: number; // Actions + canCompactContext?: boolean; + isCompactingContext?: boolean; + onCompactContext?: () => void | Promise; canSend: boolean; isStreaming: boolean; hasQueuedMessage: boolean; @@ -108,6 +113,9 @@ export function ChatInputToolbar({ onCreateProject, contextTokens, contextLimit, + canCompactContext = false, + isCompactingContext = false, + onCompactContext, canSend, isStreaming, hasQueuedMessage, @@ -121,6 +129,7 @@ export function ChatInputToolbar({ const { t } = useTranslation("chat"); const { formatNumber } = useLocaleFormatting(); const { readyAgentIds } = useAgentProviderStatus(); + const [isContextPopoverOpen, setIsContextPopoverOpen] = useState(false); const agentProviders = useMemo(() => { const seen = new Set(); @@ -156,6 +165,21 @@ export function ChatInputToolbar({ const projectTitle = selectedProject?.workingDirs.length ? selectedProject.workingDirs.join(", ") : undefined; + const contextProgress = + contextLimit > 0 ? Math.min(contextTokens / contextLimit, 1) : 0; + const contextPercentDigits = + contextProgress > 0 && contextProgress < 0.1 ? 1 : 0; + const usedPercentLabel = formatNumber(contextProgress, { + style: "percent", + minimumFractionDigits: contextPercentDigits, + maximumFractionDigits: contextPercentDigits, + }); + const formatCompactTokenCount = (value: number) => + formatNumber(value, { + notation: "compact", + compactDisplay: "short", + maximumFractionDigits: value < 10_000 ? 1 : 0, + }); const handleProjectValueChange = (value: string) => { if (value === CREATE_PROJECT_VALUE) { @@ -166,6 +190,15 @@ export function ChatInputToolbar({ onProjectChange?.(value === NO_PROJECT_VALUE ? null : value); }; + const handleCompactContext = () => { + if (!canCompactContext || isCompactingContext || !onCompactContext) { + return; + } + + setIsContextPopoverOpen(false); + void onCompactContext(); + }; + return (
{/* Left side: pickers */} @@ -247,19 +280,70 @@ export function ChatInputToolbar({ )} {contextLimit > 0 && ( - + + + + +
+ {t("toolbar.contextWindow")} +
+
+ +
+
+ {t("toolbar.contextTokensBreakdown", { + tokens: formatCompactTokenCount(contextTokens), + limit: formatCompactTokenCount(contextLimit), + })} +
+
{usedPercentLabel}
+
+ +
+ + )} diff --git a/ui/goose2/src/features/chat/ui/ChatView.tsx b/ui/goose2/src/features/chat/ui/ChatView.tsx index f71313a37c..ce0c6184f7 100644 --- a/ui/goose2/src/features/chat/ui/ChatView.tsx +++ b/ui/goose2/src/features/chat/ui/ChatView.tsx @@ -53,7 +53,6 @@ export function ChatView({ const { t } = useTranslation("chat"); const activeSessionId = sessionId; const mountStart = useRef(performance.now()); - // biome-ignore lint/correctness/useExhaustiveDependencies: log once on mount per session useEffect(() => { const ms = (performance.now() - mountStart.current).toFixed(1); perfLog(`[perf:chatview] ${sessionId.slice(0, 8)} mounted in ${ms}ms`); @@ -363,6 +362,7 @@ export function ChatView({ chatState, tokenState, sendMessage, + compactConversation, stopStreaming, streamingMessageId, } = useChat( @@ -453,6 +453,7 @@ export function ChatView({ onInitialMessageConsumed, ]); const isStreaming = chatState === "streaming"; + const isCompacting = chatState === "compacting"; const showIndicator = chatState === "thinking" || chatState === "streaming" || @@ -544,6 +545,13 @@ export function ChatView({ } contextTokens={tokenState.accumulatedTotal} contextLimit={tokenState.contextLimit} + onCompactContext={compactConversation} + canCompactContext={ + chatState === "idle" && + tokenState.accumulatedTotal > 0 && + !projectMetadataPending + } + isCompactingContext={isCompacting} />
diff --git a/ui/goose2/src/features/chat/ui/ContextRing.tsx b/ui/goose2/src/features/chat/ui/ContextRing.tsx index cacb436d71..9f49d28482 100644 --- a/ui/goose2/src/features/chat/ui/ContextRing.tsx +++ b/ui/goose2/src/features/chat/ui/ContextRing.tsx @@ -22,14 +22,6 @@ export function ContextRing({ const offset = circumference - progress * circumference; const percent = formatNumber(Math.round(progress * 100)); - // Color based on usage - const strokeColor = - progress > 0.9 - ? "var(--text-danger)" - : progress > 0.7 - ? "var(--text-warning)" - : "var(--text-muted)"; - return ( - {/* Background track */} - {/* Progress arc */} 0 ? "round" : "butt"} strokeDasharray={circumference} strokeDashoffset={offset} transform={`rotate(-90 ${size / 2} ${size / 2})`} - className="transition-all duration-300" + className="transition-all duration-300 ease-out" /> ); diff --git a/ui/goose2/src/features/chat/ui/LoadingGoose.tsx b/ui/goose2/src/features/chat/ui/LoadingGoose.tsx index e7a2e3f3b6..bd5c885985 100644 --- a/ui/goose2/src/features/chat/ui/LoadingGoose.tsx +++ b/ui/goose2/src/features/chat/ui/LoadingGoose.tsx @@ -21,12 +21,12 @@ const LOADING_SHIMMER_REPEAT_DELAY_S = 0.9; const MESSAGE_KEY_BY_STATE: Record< Exclude, - "thinking" | "responding" + "thinking" | "responding" | "compacting" > = { thinking: "thinking", streaming: "responding", waiting: "responding", - compacting: "responding", + compacting: "compacting", }; export function LoadingGoose({ chatState = "idle" }: LoadingGooseProps) { diff --git a/ui/goose2/src/features/chat/ui/__tests__/ChatInput.test.tsx b/ui/goose2/src/features/chat/ui/__tests__/ChatInput.test.tsx index 68d141de08..0561892dd0 100644 --- a/ui/goose2/src/features/chat/ui/__tests__/ChatInput.test.tsx +++ b/ui/goose2/src/features/chat/ui/__tests__/ChatInput.test.tsx @@ -223,6 +223,49 @@ describe("ChatInput", () => { expect(screen.getByText("goose2")).toBeInTheDocument(); }); + it("opens a context usage popover when token tracking is available", async () => { + const user = userEvent.setup(); + render( + , + ); + + await user.click(screen.getByRole("button", { name: /context usage/i })); + + expect(screen.getByText("Context window")).toBeInTheDocument(); + expect(screen.getByText("1.5K / 8.2K tokens used")).toBeInTheDocument(); + expect(screen.getByText("19%")).toBeInTheDocument(); + }); + + it("runs compaction from the context usage popover", async () => { + const user = userEvent.setup(); + const onCompactContext = vi.fn(); + + render( + , + ); + + await user.click(screen.getByRole("button", { name: /context usage/i })); + await user.click(screen.getByRole("button", { name: "Compact" })); + + expect(onCompactContext).toHaveBeenCalledOnce(); + }); + + it("hides the context usage control when the context limit is unavailable", () => { + render( + , + ); + + expect( + screen.queryByRole("button", { name: /context usage/i }), + ).not.toBeInTheDocument(); + }); + it("shows stop button when streaming", () => { render(); expect( diff --git a/ui/goose2/src/features/chat/ui/__tests__/LoadingGoose.test.tsx b/ui/goose2/src/features/chat/ui/__tests__/LoadingGoose.test.tsx index 939ee3f824..9c34927c35 100644 --- a/ui/goose2/src/features/chat/ui/__tests__/LoadingGoose.test.tsx +++ b/ui/goose2/src/features/chat/ui/__tests__/LoadingGoose.test.tsx @@ -3,7 +3,7 @@ import { describe, expect, it } from "vitest"; import { LoadingGoose } from "../LoadingGoose"; import chat from "@/shared/i18n/locales/en/chat.json"; -const { thinking, responding } = chat.loading; +const { thinking, responding, compacting } = chat.loading; describe("LoadingGoose", () => { it("renders thinking copy for the thinking state", () => { @@ -23,10 +23,13 @@ describe("LoadingGoose", () => { expect( screen.getByRole("status", { name: responding }), ).toBeInTheDocument(); + }); + + it("renders compacting copy for the compacting state", () => { + render(); - rerender(); expect( - screen.getByRole("status", { name: responding }), + screen.getByRole("status", { name: compacting }), ).toBeInTheDocument(); }); diff --git a/ui/goose2/src/shared/api/__tests__/acp.test.ts b/ui/goose2/src/shared/api/__tests__/acp.test.ts new file mode 100644 index 0000000000..880a9fb4fd --- /dev/null +++ b/ui/goose2/src/shared/api/__tests__/acp.test.ts @@ -0,0 +1,57 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const mockLoadSession = vi.fn(); + +vi.mock("../acpApi", () => ({ + listProviders: vi.fn(), + prompt: vi.fn(), + setModel: vi.fn(), + listSessions: vi.fn(), + loadSession: (...args: unknown[]) => mockLoadSession(...args), + exportSession: vi.fn(), + importSession: vi.fn(), + forkSession: vi.fn(), + cancelSession: vi.fn(), +})); + +vi.mock("../acpNotificationHandler", () => ({ + setActiveMessageId: vi.fn(), + clearActiveMessageId: vi.fn(), +})); + +vi.mock("../sessionSearch", () => ({ + searchSessionsViaExports: vi.fn(), +})); + +describe("acpLoadSession", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.resetModules(); + }); + + it("restores the prior session mapping when replay loading fails", async () => { + mockLoadSession.mockRejectedValueOnce(new Error("load failed")); + + const sessionTracker = await import("../acpSessionTracker"); + const { acpLoadSession } = await import("../acp"); + + sessionTracker.registerSession( + "local-session", + "goose-session-1", + "goose", + "/tmp/original", + ); + + await expect( + acpLoadSession("local-session", "goose-session-2", "/tmp/replay"), + ).rejects.toThrow("load failed"); + + expect(sessionTracker.getGooseSessionId("local-session")).toBe( + "goose-session-1", + ); + expect(sessionTracker.getLocalSessionId("goose-session-1")).toBe( + "local-session", + ); + expect(sessionTracker.getLocalSessionId("goose-session-2")).toBeNull(); + }); +}); diff --git a/ui/goose2/src/shared/api/acp.ts b/ui/goose2/src/shared/api/acp.ts index 85fa8dd549..79b1b7c0b9 100644 --- a/ui/goose2/src/shared/api/acp.ts +++ b/ui/goose2/src/shared/api/acp.ts @@ -146,17 +146,22 @@ export async function acpLoadSession( const effectiveWorkingDir = workingDir ?? "~/.goose/artifacts"; const sid = sessionId.slice(0, 8); const t0 = performance.now(); - perfLog(`[perf:load] ${sid} acpLoadSession → client.loadSession`); - await directAcp.loadSession(gooseSessionId, effectiveWorkingDir); - perfLog( - `[perf:load] ${sid} client.loadSession resolved in ${(performance.now() - t0).toFixed(1)}ms`, - ); - sessionTracker.registerSession( + const rollbackSessionRegistration = sessionTracker.registerSession( sessionId, gooseSessionId, "goose", effectiveWorkingDir, ); + try { + perfLog(`[perf:load] ${sid} acpLoadSession → client.loadSession`); + await directAcp.loadSession(gooseSessionId, effectiveWorkingDir); + perfLog( + `[perf:load] ${sid} client.loadSession resolved in ${(performance.now() - t0).toFixed(1)}ms`, + ); + } catch (error) { + rollbackSessionRegistration(); + throw error; + } } /** Export a session as JSON via the goose binary. */ diff --git a/ui/goose2/src/shared/api/acpApi.ts b/ui/goose2/src/shared/api/acpApi.ts index 015e7be18f..f5bf706007 100644 --- a/ui/goose2/src/shared/api/acpApi.ts +++ b/ui/goose2/src/shared/api/acpApi.ts @@ -132,11 +132,23 @@ export async function cancelSession(sessionId: string): Promise { export async function newSession( workingDir: string, + providerId?: string, ): Promise { const tClient = performance.now(); const client = await getClient(); + const request: Parameters[0] & { + meta?: Record; + } = { + cwd: workingDir, + mcpServers: [], + }; + + if (providerId) { + request.meta = { provider: providerId }; + } + const tCall = performance.now(); - const response = await client.newSession({ cwd: workingDir, mcpServers: [] }); + const response = await client.newSession(request); const sid = response.sessionId.slice(0, 8); perfLog( `[perf:api] ${sid} newSession getClient=${(tCall - tClient).toFixed(1)}ms wire=${(performance.now() - tCall).toFixed(1)}ms`, diff --git a/ui/goose2/src/shared/api/acpNotificationHandler.test.ts b/ui/goose2/src/shared/api/acpNotificationHandler.test.ts new file mode 100644 index 0000000000..1a3df9e7e1 --- /dev/null +++ b/ui/goose2/src/shared/api/acpNotificationHandler.test.ts @@ -0,0 +1,76 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import type { SessionNotification } from "@agentclientprotocol/sdk"; +import { useChatStore } from "@/features/chat/stores/chatStore"; +import { + clearReplayBuffer, + getAndDeleteReplayBuffer, +} from "@/features/chat/hooks/replayBuffer"; +import { registerSession } from "./acpSessionTracker"; +import { handleSessionNotification } from "./acpNotificationHandler"; + +describe("acpNotificationHandler", () => { + beforeEach(() => { + clearReplayBuffer("draft-session-1"); + clearReplayBuffer("draft-session-2"); + useChatStore.setState({ + messagesBySession: {}, + sessionStateById: {}, + queuedMessageBySession: {}, + draftsBySession: {}, + activeSessionId: null, + isConnected: false, + loadingSessionIds: new Set(), + scrollTargetMessageBySession: {}, + }); + }); + + it("buffers usage updates until the local session mapping is registered", async () => { + const notification = { + sessionId: "goose-session-1", + update: { + sessionUpdate: "usage_update", + used: 512, + size: 8192, + }, + } as SessionNotification; + + await handleSessionNotification(notification); + + expect( + useChatStore.getState().sessionStateById["draft-session-1"], + ).toBeUndefined(); + expect( + useChatStore.getState().sessionStateById["goose-session-1"], + ).toBeUndefined(); + + registerSession("draft-session-1", "goose-session-1", "goose", "/tmp"); + + const runtime = useChatStore + .getState() + .getSessionRuntime("draft-session-1"); + expect(runtime.tokenState.accumulatedTotal).toBe(512); + expect(runtime.tokenState.contextLimit).toBe(8192); + }); + + it("does not buffer non-usage updates before the local session mapping exists", async () => { + const notification = { + sessionId: "goose-session-2", + update: { + sessionUpdate: "agent_message_chunk", + messageId: "message-1", + content: { + type: "text", + text: "hello from replay", + }, + }, + } as SessionNotification; + + await handleSessionNotification(notification); + registerSession("draft-session-2", "goose-session-2", "goose", "/tmp"); + + expect(getAndDeleteReplayBuffer("draft-session-2")).toBeUndefined(); + expect( + useChatStore.getState().messagesBySession["draft-session-2"], + ).toBeUndefined(); + }); +}); diff --git a/ui/goose2/src/shared/api/acpNotificationHandler.ts b/ui/goose2/src/shared/api/acpNotificationHandler.ts index b0d8d8a578..23b5bf3133 100644 --- a/ui/goose2/src/shared/api/acpNotificationHandler.ts +++ b/ui/goose2/src/shared/api/acpNotificationHandler.ts @@ -14,11 +14,55 @@ import type { ToolResponseContent, } from "@/shared/types/messages"; import type { AcpNotificationHandler } from "./acpConnection"; -import { getLocalSessionId } from "./acpSessionTracker"; +import { + getLocalSessionId, + subscribeToSessionRegistration, +} from "./acpSessionTracker"; import { perfLog } from "@/shared/lib/perfLog"; // Pre-set message ID for the next live stream per goose session const presetMessageIds = new Map(); +const pendingUsageUpdates = new Map(); + +function shouldBufferPendingUpdate(update: SessionUpdate): boolean { + return update.sessionUpdate === "usage_update"; +} + +function queuePendingUsageUpdate( + gooseSessionId: string, + update: SessionUpdate, +): void { + const pending = pendingUsageUpdates.get(gooseSessionId); + if (pending) { + pending.push(update); + return; + } + pendingUsageUpdates.set(gooseSessionId, [update]); +} + +function flushPendingUsageUpdates( + localSessionId: string, + gooseSessionId: string, +): void { + const pending = pendingUsageUpdates.get(gooseSessionId); + if (!pending?.length) { + return; + } + + pendingUsageUpdates.delete(gooseSessionId); + + for (const update of pending) { + if (useChatStore.getState().loadingSessionIds.has(localSessionId)) { + handleReplay(localSessionId, update); + } else { + handleLive(localSessionId, gooseSessionId, update); + } + } +} + +subscribeToSessionRegistration((localSessionId, gooseSessionId) => { + flushPendingUsageUpdates(localSessionId, gooseSessionId); +}); // Per-session perf counters for replay/live streaming. interface ReplayPerf { @@ -67,22 +111,32 @@ export async function handleSessionNotification( notification: SessionNotification, ): Promise { const gooseSessionId = notification.sessionId; - const sessionId = getLocalSessionId(gooseSessionId) ?? gooseSessionId; const { update } = notification; - const isReplay = useChatStore.getState().loadingSessionIds.has(sessionId); + const localSessionId = getLocalSessionId(gooseSessionId); + + if (!localSessionId) { + if (shouldBufferPendingUpdate(update)) { + queuePendingUsageUpdate(gooseSessionId, update); + } + return; + } + + const isReplay = useChatStore + .getState() + .loadingSessionIds.has(localSessionId); if (isReplay) { - const sid = sessionId.slice(0, 8); - let perf = replayPerf.get(sessionId); + const sid = localSessionId.slice(0, 8); + let perf = replayPerf.get(localSessionId); const now = performance.now(); if (!perf) { perf = { firstAt: now, lastAt: now, count: 0 }; - replayPerf.set(sessionId, perf); + replayPerf.set(localSessionId, perf); perfLog(`[perf:replay] ${sid} first notification received`); } perf.lastAt = now; perf.count += 1; - handleReplay(sessionId, update); + handleReplay(localSessionId, update); } else { const perf = livePerf.get(gooseSessionId); if (perf && update.sessionUpdate === "agent_message_chunk") { @@ -95,7 +149,7 @@ export async function handleSessionNotification( ); } } - handleLive(sessionId, gooseSessionId, update); + handleLive(localSessionId, gooseSessionId, update); } } diff --git a/ui/goose2/src/shared/api/acpSessionTracker.ts b/ui/goose2/src/shared/api/acpSessionTracker.ts index 835bcc1900..f521a6293a 100644 --- a/ui/goose2/src/shared/api/acpSessionTracker.ts +++ b/ui/goose2/src/shared/api/acpSessionTracker.ts @@ -7,8 +7,26 @@ interface PreparedSession { workingDir: string; } +type SessionRegistrationListener = ( + localSessionId: string, + gooseSessionId: string, +) => void; + const prepared = new Map(); const gooseToLocal = new Map(); +const registrationListeners = new Set(); + +function restoreGooseRegistration( + gooseSessionId: string, + localSessionId: string | undefined, +): void { + if (localSessionId === undefined) { + gooseToLocal.delete(gooseSessionId); + return; + } + + gooseToLocal.set(gooseSessionId, localSessionId); +} function makeKey(sessionId: string, personaId?: string): string { if (personaId && personaId.length > 0) { @@ -17,6 +35,22 @@ function makeKey(sessionId: string, personaId?: string): string { return sessionId; } +function notifySessionRegistered( + localSessionId: string, + gooseSessionId: string, +): void { + for (const listener of registrationListeners) { + listener(localSessionId, gooseSessionId); + } +} + +export function subscribeToSessionRegistration( + listener: SessionRegistrationListener, +): () => void { + registrationListeners.add(listener); + return () => registrationListeners.delete(listener); +} + export async function prepareSession( sessionId: string, providerId: string, @@ -67,7 +101,7 @@ export async function prepareSession( if (!gooseSessionId) { const tNew = performance.now(); - const response = await acpApi.newSession(workingDir); + const response = await acpApi.newSession(workingDir, providerId); gooseSessionId = response.sessionId; perfLog( `[perf:prepare] ${sid} tracker newSession done in ${(performance.now() - tNew).toFixed(1)}ms (goose_sid=${gooseSessionId.slice(0, 8)})`, @@ -84,6 +118,7 @@ export async function prepareSession( prepared.set(key, { gooseSessionId, providerId, workingDir }); prepared.set(sessionId, { gooseSessionId, providerId, workingDir }); gooseToLocal.set(gooseSessionId, sessionId); + notifySessionRegistered(sessionId, gooseSessionId); return gooseSessionId; } @@ -109,8 +144,54 @@ export function registerSession( gooseSessionId: string, providerId: string, workingDir: string, -): void { +): () => void { + const previousEntry = prepared.get(sessionId); + const previousGooseSessionLocal = gooseToLocal.get(gooseSessionId); + const previousSessionGooseLocal = previousEntry + ? gooseToLocal.get(previousEntry.gooseSessionId) + : undefined; const entry = { gooseSessionId, providerId, workingDir }; + + if ( + previousEntry && + previousEntry.gooseSessionId !== gooseSessionId && + gooseToLocal.get(previousEntry.gooseSessionId) === sessionId + ) { + gooseToLocal.delete(previousEntry.gooseSessionId); + } + prepared.set(sessionId, entry); gooseToLocal.set(gooseSessionId, sessionId); + notifySessionRegistered(sessionId, gooseSessionId); + + return () => { + prepared.delete(sessionId); + if (previousEntry) { + prepared.set(sessionId, previousEntry); + } + + restoreGooseRegistration(gooseSessionId, previousGooseSessionLocal); + if (previousEntry && previousEntry.gooseSessionId !== gooseSessionId) { + restoreGooseRegistration( + previousEntry.gooseSessionId, + previousSessionGooseLocal, + ); + } + }; +} + +export function unregisterSession( + sessionId: string, + gooseSessionId?: string, +): void { + const entry = prepared.get(sessionId); + prepared.delete(sessionId); + + const resolvedGooseSessionId = gooseSessionId ?? entry?.gooseSessionId; + if ( + resolvedGooseSessionId && + gooseToLocal.get(resolvedGooseSessionId) === sessionId + ) { + gooseToLocal.delete(resolvedGooseSessionId); + } } diff --git a/ui/goose2/src/shared/i18n/locales/en/chat.json b/ui/goose2/src/shared/i18n/locales/en/chat.json index 4c9fcb5e38..efe6776e3d 100644 --- a/ui/goose2/src/shared/i18n/locales/en/chat.json +++ b/ui/goose2/src/shared/i18n/locales/en/chat.json @@ -111,6 +111,7 @@ "placeholder": "Message {{agent}}, @ to mention personas" }, "loading": { + "compacting": "Compacting conversation...", "thinking": "Thinking...", "responding": "Responding..." }, @@ -151,9 +152,14 @@ "attachFolder": "Folder", "chooseAgentModel": "Choose agent and model", "chooseProject": "Choose a project", + "compactNow": "Compact", + "compacting": "Compacting...", "chooseProvider": "Choose a provider", + "contextTokensBreakdown": "{{tokens}} / {{limit}} tokens used", "contextUsage": "Context usage", + "contextUsageBreakdown": "{{used}} used ({{left}} left)", "contextUsageTitle": "{{tokens}} / {{limit}} tokens", + "contextWindow": "Context window", "createProject": "Create project", "generalChatWithoutProject": "General chat without project context", "loading": "Loading...", diff --git a/ui/goose2/src/shared/i18n/locales/es/chat.json b/ui/goose2/src/shared/i18n/locales/es/chat.json index e73f7a1fc7..3a5760189e 100644 --- a/ui/goose2/src/shared/i18n/locales/es/chat.json +++ b/ui/goose2/src/shared/i18n/locales/es/chat.json @@ -111,6 +111,7 @@ "placeholder": "Enviar mensaje a {{agent}}, usa @ para mencionar personas" }, "loading": { + "compacting": "Compactando conversación...", "thinking": "Pensando...", "responding": "Respondiendo..." }, @@ -151,9 +152,14 @@ "attachFolder": "Carpeta", "chooseAgentModel": "Elegir agente y modelo", "chooseProject": "Elegir un proyecto", + "compactNow": "Compactar", + "compacting": "Compactando...", "chooseProvider": "Elegir un proveedor", + "contextTokensBreakdown": "{{tokens}} / {{limit}} tokens usados", "contextUsage": "Uso del contexto", + "contextUsageBreakdown": "{{used}} en uso ({{left}} libres)", "contextUsageTitle": "{{tokens}} / {{limit}} tokens", + "contextWindow": "Ventana de contexto", "createProject": "Crear proyecto", "generalChatWithoutProject": "Chat general sin contexto de proyecto", "loading": "Cargando...", diff --git a/ui/goose2/src/shared/ui/button.tsx b/ui/goose2/src/shared/ui/button.tsx index 0f8dfe971e..b1e161d7a4 100644 --- a/ui/goose2/src/shared/ui/button.tsx +++ b/ui/goose2/src/shared/ui/button.tsx @@ -18,7 +18,7 @@ const buttonVariants = cva( "outline-flat": "border border-input bg-background shadow-none hover:bg-accent hover:text-accent-foreground", secondary: - "bg-secondary text-secondary-foreground shadow-sm hover:bg-secondary/80", + "bg-secondary text-secondary-foreground hover:bg-secondary/80", ghost: "hover:bg-accent hover:text-accent-foreground", "ghost-light": "font-normal hover:bg-accent hover:text-accent-foreground",