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

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

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();
+  });
+});

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. */

ui/goose2/src/shared/api/acpApi.ts

@@ -132,11 +132,23 @@ export async function cancelSession(sessionId: string): Promise<void> {
 
 export async function newSession(
   workingDir: string,
+  providerId?: string,
 ): Promise<NewSessionResponse> {
   const tClient = performance.now();
   const client = await getClient();
+  const request: Parameters<typeof client.newSession>[0] & {
+    meta?: Record<string, string>;
+  } = {
+    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`,

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<string>(),
+      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();
+  });
+});

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<string, string>();
+const pendingUsageUpdates = new Map<string, SessionUpdate[]>();
+
+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<void> {
   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);
   }
 }
 

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<string, PreparedSession>();
 const gooseToLocal = new Map<string, string>();
+const registrationListeners = new Set<SessionRegistrationListener>();
+
+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);
+  }
 }

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...",

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...",

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",