diff --git a/src/components/ai-edition/LeftPanel.tsx b/src/components/ai-edition/LeftPanel.tsx index 1942d1da2..3695a938d 100644 --- a/src/components/ai-edition/LeftPanel.tsx +++ b/src/components/ai-edition/LeftPanel.tsx @@ -27,11 +27,11 @@ import { } from "../../../electron/ai-edition/provider-registry"; import { ChatWelcome } from "./ChatWelcome"; import { canSendChat } from "./chatAvailability"; -import { computeBudget } from "./chatBudget"; import { ChatHistoryModal, SourceTranscriptModal } from "./Modals"; import styles from "./NewEditorShell.module.css"; import { ProviderSettings } from "./ProviderSettings"; import { TranscriptionStatusDot } from "./TranscriptionStatus"; +import { useChatBudget } from "./useChatBudget"; export type LeftTab = "chat" | "media"; @@ -1141,10 +1141,10 @@ function ChatStripPanel() { }); }, [llmConfig]); - // Real context usage — feeds the badge in the chat strip and gates the - // auto-compact heuristic on the main side. Recomputed on every messages - // change so the % tracks the live history. - const budget = computeBudget(messages); + // Prefer the main process's model-message budget so manual compaction can + // shrink this meter while the complete transcript remains visible. The hook + // falls back to a renderer estimate in browser/shim or bridge-failure cases. + const budget = useChatBudget({ projectId, sessionId: activeSessionId, messages }); const [compactNowPending, setCompactNowPending] = useState(false); const compactNow = useCallback(async () => { diff --git a/src/components/ai-edition/chatBudget.ts b/src/components/ai-edition/chatBudget.ts index 912c91877..9f76190d9 100644 --- a/src/components/ai-edition/chatBudget.ts +++ b/src/components/ai-edition/chatBudget.ts @@ -1,12 +1,10 @@ // Renderer-side budget helper. Mirrors `electron/ai-edition/chat-compaction.ts` // but inline so we don't drag electron/ into the renderer bundle. // -// This feeds the context pill and NOTHING else — no code decides anything from -// it. `DEFAULT_CHAT_BUDGET_TOKENS` is a made-up denominator (the app has no way -// to ask a provider how big its context window is), which is exactly why the -// automatic compaction that used to branch on the main-process twin is gone. -// Read the pill as "the conversation is about this big", never as "you are this -// close to a limit", and do not let this number regain a decision. +// This is the renderer fallback for the context pill while native usage is +// loading or unavailable. Desktop builds replace it with the main process's +// model-message estimate, which understands compaction; no code makes an +// automatic compaction decision from either value. const CHARS_PER_TOKEN = 4; @@ -18,12 +16,12 @@ export interface ChatBudget { const DEFAULT_CHAT_BUDGET_TOKENS = 80_000; -interface RenderableChatMessage { +export interface RenderableChatMessage { content: string; toolCalls?: Array<{ name?: string; summary?: string }>; } -function estimateTokens(messages: RenderableChatMessage[]): number { +function estimateTokens(messages: readonly RenderableChatMessage[]): number { let chars = 0; for (const m of messages) { chars += m.content.length; @@ -35,7 +33,7 @@ function estimateTokens(messages: RenderableChatMessage[]): number { } export function computeBudget( - messages: RenderableChatMessage[], + messages: readonly RenderableChatMessage[], budgetTokens: number = DEFAULT_CHAT_BUDGET_TOKENS, ): ChatBudget { const used = estimateTokens(messages); diff --git a/src/components/ai-edition/useChatBudget.test.ts b/src/components/ai-edition/useChatBudget.test.ts new file mode 100644 index 000000000..33718d22e --- /dev/null +++ b/src/components/ai-edition/useChatBudget.test.ts @@ -0,0 +1,107 @@ +// @vitest-environment jsdom +import { act, renderHook, waitFor } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { useChatBudget } from "./useChatBudget"; + +const chatBudgetMock = vi.hoisted(() => vi.fn()); + +vi.mock("@/native/client", () => ({ + nativeBridgeClient: { + aiEdition: { chatBudget: chatBudgetMock }, + }, +})); + +function deferred() { + let resolve!: (value: T) => void; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +const message = (content: string) => [{ content }]; + +describe("useChatBudget", () => { + beforeEach(() => chatBudgetMock.mockReset()); + + it("uses the transcript estimate until native model-context usage arrives", async () => { + const native = deferred<{ + usedTokens: number; + budgetTokens: number; + ratio: number; + fillPercent: number; + }>(); + chatBudgetMock.mockReturnValue(native.promise); + const visibleMessages = message("x".repeat(400)); + + const { result } = renderHook(() => + useChatBudget({ projectId: "project_1", sessionId: "session_1", messages: visibleMessages }), + ); + expect(result.current.usedTokens).toBe(100); + + act(() => + native.resolve({ usedTokens: 12, budgetTokens: 80_000, ratio: 0.00015, fillPercent: 0.015 }), + ); + await waitFor(() => expect(result.current.usedTokens).toBe(12)); + }); + + it("ignores a late response from the previously selected session", async () => { + const first = deferred<{ + usedTokens: number; + budgetTokens: number; + ratio: number; + fillPercent: number; + }>(); + const second = deferred<{ + usedTokens: number; + budgetTokens: number; + ratio: number; + fillPercent: number; + }>(); + chatBudgetMock.mockReturnValueOnce(first.promise).mockReturnValueOnce(second.promise); + const visibleMessages = message("visible transcript"); + + const { result, rerender } = renderHook( + ({ sessionId }) => + useChatBudget({ projectId: "project_1", sessionId, messages: visibleMessages }), + { initialProps: { sessionId: "session_1" } }, + ); + rerender({ sessionId: "session_2" }); + act(() => + second.resolve({ usedTokens: 20, budgetTokens: 80_000, ratio: 0.00025, fillPercent: 0.025 }), + ); + await waitFor(() => expect(result.current.usedTokens).toBe(20)); + + act(() => + first.resolve({ usedTokens: 999, budgetTokens: 80_000, ratio: 0.012, fillPercent: 1.2 }), + ); + await act(async () => Promise.resolve()); + expect(result.current.usedTokens).toBe(20); + }); + + it("refreshes native usage when compaction returns a new transcript array", async () => { + chatBudgetMock + .mockResolvedValueOnce({ + usedTokens: 500, + budgetTokens: 80_000, + ratio: 0.00625, + fillPercent: 0.625, + }) + .mockResolvedValueOnce({ + usedTokens: 40, + budgetTokens: 80_000, + ratio: 0.0005, + fillPercent: 0.05, + }); + const visibleMessages = message("the transcript remains visible"); + const { result, rerender } = renderHook( + ({ messages }) => useChatBudget({ projectId: "project_1", sessionId: "session_1", messages }), + { initialProps: { messages: visibleMessages } }, + ); + await waitFor(() => expect(result.current.usedTokens).toBe(500)); + + rerender({ messages: [...visibleMessages] }); + await waitFor(() => expect(result.current.usedTokens).toBe(40)); + expect(chatBudgetMock).toHaveBeenCalledTimes(2); + }); +}); diff --git a/src/components/ai-edition/useChatBudget.ts b/src/components/ai-edition/useChatBudget.ts new file mode 100644 index 000000000..01c0ae9ef --- /dev/null +++ b/src/components/ai-edition/useChatBudget.ts @@ -0,0 +1,45 @@ +import { useEffect, useMemo, useRef, useState } from "react"; +import { nativeBridgeClient } from "@/native/client"; +import { type ChatBudget, computeBudget, type RenderableChatMessage } from "./chatBudget"; + +interface NativeBudgetState { + sessionKey: string; + budget: ChatBudget; +} + +export function useChatBudget(options: { + projectId: string | null; + sessionId: string | null; + messages: readonly RenderableChatMessage[]; +}): ChatBudget { + const { projectId, sessionId, messages } = options; + const fallback = useMemo(() => computeBudget(messages), [messages]); + const sessionKey = projectId && sessionId ? `${projectId}\0${sessionId}` : null; + const [nativeState, setNativeState] = useState(null); + const requestIdRef = useRef(0); + + useEffect(() => { + const requestId = ++requestIdRef.current; + if (!projectId || !sessionId || !sessionKey) { + setNativeState(null); + return; + } + + void nativeBridgeClient.aiEdition + .chatBudget(projectId, sessionId) + .then((budget) => { + if (requestIdRef.current !== requestId) return; + setNativeState({ sessionKey, budget: budget ?? fallback }); + }) + .catch(() => { + if (requestIdRef.current !== requestId) return; + setNativeState({ sessionKey, budget: fallback }); + }); + + return () => { + if (requestIdRef.current === requestId) requestIdRef.current++; + }; + }, [projectId, sessionId, sessionKey, fallback]); + + return nativeState?.sessionKey === sessionKey ? nativeState.budget : fallback; +}