diff --git a/ui/goose2/src/app/AppShell.tsx b/ui/goose2/src/app/AppShell.tsx index 726e59d847..ee9d58328c 100644 --- a/ui/goose2/src/app/AppShell.tsx +++ b/ui/goose2/src/app/AppShell.tsx @@ -134,9 +134,6 @@ export function AppShell({ children }: { children?: React.ReactNode }) { : undefined; const modelName = activeSession?.modelName; const tokenCount = chatStore.tokenState.totalTokens; - const connectionStatus = chatStore.isConnected - ? ("connected" as const) - : ("disconnected" as const); const [pendingInitialMessage, setPendingInitialMessage] = useState< string | undefined @@ -507,7 +504,6 @@ export function AppShell({ children }: { children?: React.ReactNode }) { modelName={modelName} sessionId={activeTabId ?? undefined} tokenCount={tokenCount} - status={connectionStatus} /> diff --git a/ui/goose2/src/features/chat/hooks/useAcpStream.ts b/ui/goose2/src/features/chat/hooks/useAcpStream.ts index cd2edef051..caa6e05098 100644 --- a/ui/goose2/src/features/chat/hooks/useAcpStream.ts +++ b/ui/goose2/src/features/chat/hooks/useAcpStream.ts @@ -47,6 +47,12 @@ interface AcpModelStatePayload { currentModelName?: string; } +interface AcpUsageUpdatePayload { + sessionId: string; + used: number; + size: number; +} + function findLatestUnpairedToolRequest( content: MessageContent[], ): ToolRequestContent | null { @@ -219,6 +225,18 @@ export function useAcpStream(sessionId: string, enabled: boolean): void { }), ); + // acp:usage_update — update context window token usage + unlisteners.push( + listen("acp:usage_update", (event) => { + if (!active) return; + if (event.payload.sessionId !== sessionIdRef.current) return; + useChatStore.getState().updateTokenState({ + accumulatedTotal: event.payload.used, + contextLimit: event.payload.size, + }); + }), + ); + // Cleanup: mark inactive immediately to prevent stale listeners from // firing during the async unlisten, then tear down actual subscriptions. return () => { diff --git a/ui/goose2/src/features/chat/stores/chatStore.ts b/ui/goose2/src/features/chat/stores/chatStore.ts index ea7e2064e9..d6bab1cad5 100644 --- a/ui/goose2/src/features/chat/stores/chatStore.ts +++ b/ui/goose2/src/features/chat/stores/chatStore.ts @@ -206,22 +206,23 @@ export const useChatStore = create((set, get) => ({ const current = state.tokenState; const inputTokens = partial.inputTokens ?? current.inputTokens; const outputTokens = partial.outputTokens ?? current.outputTokens; + const accumulatedInput = + partial.accumulatedInput ?? + current.accumulatedInput + (partial.inputTokens ?? 0); + const accumulatedOutput = + partial.accumulatedOutput ?? + current.accumulatedOutput + (partial.outputTokens ?? 0); + const accumulatedTotal = + partial.accumulatedTotal ?? accumulatedInput + accumulatedOutput; return { tokenState: { inputTokens, outputTokens, totalTokens: inputTokens + outputTokens, - accumulatedInput: - partial.accumulatedInput ?? - current.accumulatedInput + (partial.inputTokens ?? 0), - accumulatedOutput: - partial.accumulatedOutput ?? - current.accumulatedOutput + (partial.outputTokens ?? 0), - accumulatedTotal: - (partial.accumulatedInput ?? - current.accumulatedInput + (partial.inputTokens ?? 0)) + - (partial.accumulatedOutput ?? - current.accumulatedOutput + (partial.outputTokens ?? 0)), + accumulatedInput, + accumulatedOutput, + accumulatedTotal, + contextLimit: partial.contextLimit ?? current.contextLimit, }, }; }), diff --git a/ui/goose2/src/features/chat/ui/ChatView.tsx b/ui/goose2/src/features/chat/ui/ChatView.tsx index 014b4b82ea..25e3d5cd19 100644 --- a/ui/goose2/src/features/chat/ui/ChatView.tsx +++ b/ui/goose2/src/features/chat/ui/ChatView.tsx @@ -243,6 +243,7 @@ export function ChatView({ // Wrap sendMessage to handle @ mentioned persona overrides const chatStore = useChatStore(); + const tokenState = useChatStore((s) => s.tokenState); const handleSend = useCallback( (text: string, personaId?: string) => { if (personaId && personaId !== selectedPersonaId) { @@ -362,6 +363,8 @@ export function ChatView({ }, }) } + contextTokens={tokenState.accumulatedTotal} + contextLimit={tokenState.contextLimit} /> ); diff --git a/ui/goose2/src/features/status/ui/StatusBar.tsx b/ui/goose2/src/features/status/ui/StatusBar.tsx index 93ff024b42..6dd1a2035e 100644 --- a/ui/goose2/src/features/status/ui/StatusBar.tsx +++ b/ui/goose2/src/features/status/ui/StatusBar.tsx @@ -1,25 +1,17 @@ import { useState } from "react"; -import { Copy, Check } from "lucide-react"; +import { Bot, Copy, Check } from "lucide-react"; import { cn } from "@/shared/lib/cn"; interface StatusBarProps { modelName?: string; sessionId?: string; tokenCount?: number; - status?: "connected" | "disconnected" | "loading"; } -const statusColor = { - connected: "bg-green-500", - disconnected: "bg-red-500", - loading: "bg-yellow-500", -} as const; - export function StatusBar({ modelName, sessionId, tokenCount = 0, - status = "disconnected", }: StatusBarProps) { const [copied, setCopied] = useState(false); @@ -37,28 +29,32 @@ export function StatusBar({ "bg-background/80 px-3 text-xs text-foreground-secondary", )} > - {modelName ? {modelName} : } +
+
+ + + {modelName ?? "No model"} + +
+
{sessionId && ( -
- {sessionId} - -
+ + )} + {tokenCount > 0 && ( + + {tokenCount.toLocaleString()} tokens + )} - {tokenCount > 0 && {tokenCount.toLocaleString()} tokens} -
); diff --git a/ui/goose2/src/shared/types/chat.ts b/ui/goose2/src/shared/types/chat.ts index 7b8d11aafa..ec3cabbfd8 100644 --- a/ui/goose2/src/shared/types/chat.ts +++ b/ui/goose2/src/shared/types/chat.ts @@ -18,6 +18,7 @@ export interface TokenState { accumulatedInput: number; accumulatedOutput: number; accumulatedTotal: number; + contextLimit: number; } export const INITIAL_TOKEN_STATE: TokenState = { @@ -27,6 +28,7 @@ export const INITIAL_TOKEN_STATE: TokenState = { accumulatedInput: 0, accumulatedOutput: 0, accumulatedTotal: 0, + contextLimit: 0, }; // Session