diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index 4ba2996c91..811fd35d4a 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -551,12 +551,9 @@ export class Task extends EventEmitter implements TaskLike { this.emit(RooCodeEventName.QueuedMessagesUpdated, this.taskId, this.messageQueueService.messages) void this.providerRef .deref() - ?.postStateToWebviewWithoutTaskHistory() + ?.postStateToWebviewThrottled() .catch((error) => { - console.error( - "[Task#messageQueueStateChangedHandler] postStateToWebviewWithoutTaskHistory failed:", - error, - ) + console.error("[Task#messageQueueStateChangedHandler] postStateToWebviewThrottled failed:", error) }) } @@ -1047,9 +1044,13 @@ export class Task extends EventEmitter implements TaskLike { private async addToClineMessages(message: ClineMessage) { this.clineMessages.push(message) const provider = this.providerRef.deref() - // Avoid resending large, mostly-static fields (notably taskHistory) on every chat message update. - // taskHistory is maintained in-memory in the webview and updated via taskHistoryItemUpdated. - await provider?.postStateToWebviewWithoutTaskHistory() + // Unanswered asks must reach the webview before Message listeners can respond against its state. + const requiresImmediateState = + message.partial === true || (message.type === "ask" && message.isAnswered !== true) + await provider?.postStateToWebviewThrottled() + if (requiresImmediateState) { + await provider?.flushPostStateToWebviewThrottled() + } this.emit(RooCodeEventName.Message, { action: "created", message }) await this.saveClineMessages() @@ -2250,6 +2251,8 @@ export class Task extends EventEmitter implements TaskLike { // Force final token usage update before abort event this.emitFinalTokenUsageUpdate() + await this.providerRef.deref()?.flushPostStateToWebviewThrottled() + this.emit(RooCodeEventName.TaskAborted) try { diff --git a/src/core/task/__tests__/Task.spec.ts b/src/core/task/__tests__/Task.spec.ts index 330241e221..6e4d55e3a1 100644 --- a/src/core/task/__tests__/Task.spec.ts +++ b/src/core/task/__tests__/Task.spec.ts @@ -8,6 +8,7 @@ import { Anthropic } from "@anthropic-ai/sdk" import { providerIdentifiers, + RooCodeEventName, type GlobalState, type ProviderSettings, type ModelInfo, @@ -27,9 +28,12 @@ import type { ApiMessage } from "../../task-persistence" type TaskTestAccess = { getSystemPrompt: () => Promise + getEnabledMcpToolsCount: () => Promise<{ enabledToolCount: number; enabledServerCount: number }> + initiateTaskLoop: (userContent: Anthropic.Messages.ContentBlockParam[]) => Promise startTask: (task?: string, images?: string[]) => Promise resumeTaskFromHistory: () => Promise presentAssistantMessageSafe: () => void + addToClineMessages: (message: import("@roo-code/types").ClineMessage) => Promise updateClineMessage: (message: import("@roo-code/types").ClineMessage) => Promise saveClineMessages: () => Promise safeEnsureModelFetched: () => Promise @@ -338,6 +342,8 @@ describe("Cline", () => { mockProvider.postMessageToWebview = vi.fn().mockResolvedValue(undefined) mockProvider.postStateToWebview = vi.fn().mockResolvedValue(undefined) mockProvider.postStateToWebviewWithoutTaskHistory = vi.fn().mockResolvedValue(undefined) + mockProvider.postStateToWebviewThrottled = vi.fn().mockResolvedValue(undefined) + mockProvider.flushPostStateToWebviewThrottled = vi.fn().mockResolvedValue(undefined) mockProvider.getTaskWithId = vi.fn().mockImplementation(async (id) => ({ historyItem: { id, @@ -1029,6 +1035,8 @@ describe("Cline", () => { say: vi.fn(), postStateToWebview: vi.fn().mockResolvedValue(undefined), postStateToWebviewWithoutTaskHistory: vi.fn().mockResolvedValue(undefined), + postStateToWebviewThrottled: vi.fn().mockResolvedValue(undefined), + flushPostStateToWebviewThrottled: vi.fn().mockResolvedValue(undefined), postMessageToWebview: vi.fn().mockResolvedValue(undefined), updateTaskHistory: vi.fn().mockResolvedValue(undefined), } @@ -1663,6 +1671,137 @@ describe("Cline", () => { }) }) + describe("webview state throttling", () => { + it("schedules a complete new message without forcing an immediate state push", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + startTask: false, + }) + vi.spyOn(getTaskTestAccess(task), "saveClineMessages").mockResolvedValue(true) + const message = { + ts: Date.now(), + type: "say" as const, + say: "text" as const, + text: "message", + } + + await getTaskTestAccess(task).addToClineMessages(message) + + expect(mockProvider.postStateToWebviewThrottled).toHaveBeenCalledOnce() + expect(mockProvider.postStateToWebviewThrottled).toHaveBeenCalledWith() + expect(mockProvider.flushPostStateToWebviewThrottled).not.toHaveBeenCalled() + expect(mockProvider.postStateToWebviewWithoutTaskHistory).not.toHaveBeenCalled() + }) + + it("waits for an unanswered ask flush before emitting the message", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + startTask: false, + }) + const taskAccess = getTaskTestAccess(task) + vi.spyOn(taskAccess, "saveClineMessages").mockResolvedValue(true) + let releaseFlush!: () => void + const pendingFlush = new Promise((resolve) => { + releaseFlush = resolve + }) + const flushSpy = vi.mocked(mockProvider.flushPostStateToWebviewThrottled).mockReturnValueOnce(pendingFlush) + const messageListener = vi.fn() + task.on(RooCodeEventName.Message, messageListener) + const message = { + ts: 1, + type: "ask" as const, + ask: "resume_task" as const, + } + + const addPromise = taskAccess.addToClineMessages(message) + + await Promise.resolve() + expect(mockProvider.postStateToWebviewThrottled).toHaveBeenCalledOnce() + expect(mockProvider.postStateToWebviewThrottled).toHaveBeenCalledWith() + expect(flushSpy).toHaveBeenCalledOnce() + expect(flushSpy).toHaveBeenCalledWith() + expect(messageListener).not.toHaveBeenCalled() + + releaseFlush() + await addPromise + + expect(flushSpy.mock.invocationCallOrder[0]).toBeLessThan(messageListener.mock.invocationCallOrder[0]) + expect(messageListener).toHaveBeenCalledWith({ action: "created", message }) + }) + + it("keeps an already answered ask on the throttled path", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + startTask: false, + }) + vi.spyOn(getTaskTestAccess(task), "saveClineMessages").mockResolvedValue(true) + + await getTaskTestAccess(task).addToClineMessages({ + ts: 1, + type: "ask", + ask: "tool", + isAnswered: true, + }) + + expect(mockProvider.postStateToWebviewThrottled).toHaveBeenCalledOnce() + expect(mockProvider.postStateToWebviewThrottled).toHaveBeenCalledWith() + expect(mockProvider.flushPostStateToWebviewThrottled).not.toHaveBeenCalled() + }) + + it("waits for a new partial message flush before a following message update", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + startTask: false, + }) + const taskAccess = getTaskTestAccess(task) + vi.spyOn(taskAccess, "saveClineMessages").mockResolvedValue(true) + let releaseFlush!: () => void + const pendingFlush = new Promise((resolve) => { + releaseFlush = resolve + }) + const flushSpy = vi.mocked(mockProvider.flushPostStateToWebviewThrottled).mockReturnValueOnce(pendingFlush) + const updatePostSpy = vi.mocked(mockProvider.postMessageToWebview) + const partialMessage = { + ts: 1, + type: "say" as const, + say: "text" as const, + text: "partial message", + partial: true, + } + let partialAddSettled = false + const addThenUpdate = taskAccess.addToClineMessages(partialMessage).then(async () => { + partialAddSettled = true + await taskAccess.updateClineMessage({ ...partialMessage, text: "updated partial" }) + }) + + await Promise.resolve() + expect(mockProvider.postStateToWebviewThrottled).toHaveBeenCalledWith() + expect(flushSpy).toHaveBeenCalledWith() + expect(partialAddSettled).toBe(false) + expect(updatePostSpy).not.toHaveBeenCalled() + + releaseFlush() + await addThenUpdate + + expect(flushSpy.mock.invocationCallOrder[0]).toBeLessThan(updatePostSpy.mock.invocationCallOrder[0]) + expect(updatePostSpy).toHaveBeenCalledWith({ + type: "messageUpdated", + clineMessage: { + ...partialMessage, + text: "updated partial", + }, + }) + }) + }) + describe("abortTask", () => { it("should set abort flag and emit TaskAborted event", async () => { const task = new Task({ @@ -1707,6 +1846,37 @@ describe("Cline", () => { expect(disposeSpy).toHaveBeenCalled() }) + it("flushes pending state before TaskAborted and disposal while queue state is intact", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + startTask: false, + }) + let queuedMessagesAtFlush = -1 + const flushSpy = vi.mocked(mockProvider.flushPostStateToWebviewThrottled).mockImplementation(async () => { + queuedMessagesAtFlush = task.messageQueueService.messages.length + }) + const emitSpy = vi.spyOn(task, "emit") + const disposeSpy = vi.spyOn(task, "dispose").mockImplementation(() => {}) + + task.messageQueueService.addMessage("queued text") + await task.abortTask() + + const taskAbortedCallIndex = (emitSpy.mock.calls as unknown[][]).findIndex( + ([event]) => event === "taskAborted", + ) + expect(taskAbortedCallIndex).toBeGreaterThanOrEqual(0) + expect(queuedMessagesAtFlush).toBe(1) + expect(flushSpy).toHaveBeenCalledWith() + expect(flushSpy.mock.invocationCallOrder[0]).toBeLessThan( + emitSpy.mock.invocationCallOrder[taskAbortedCallIndex], + ) + expect(emitSpy.mock.invocationCallOrder[taskAbortedCallIndex]).toBeLessThan( + disposeSpy.mock.invocationCallOrder[0], + ) + }) + it("should work with TaskLike interface", async () => { const task = new Task({ provider: mockProvider, @@ -2741,6 +2911,50 @@ describe("Cline", () => { }) }) + describe("startTask", () => { + it("posts a clean state immediately before adding the first task message", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "new task", + startTask: false, + }) + const taskAccess = getTaskTestAccess(task) + + task.clineMessages = [{ ts: 1, type: "say", say: "text", text: "stale message" }] + + let resolvePostState: (() => void) | undefined + const pendingPostState = new Promise((resolve) => { + resolvePostState = resolve + }) + const postStateSpy = vi + .mocked(mockProvider.postStateToWebviewWithoutTaskHistory) + .mockImplementationOnce(async () => { + expect(task.clineMessages).toEqual([]) + await pendingPostState + }) + const saySpy = vi.spyOn(task, "say").mockResolvedValue(undefined) + vi.spyOn(taskAccess, "getEnabledMcpToolsCount").mockResolvedValue({ + enabledToolCount: 0, + enabledServerCount: 0, + }) + const initiateTaskLoopSpy = vi.spyOn(taskAccess, "initiateTaskLoop").mockResolvedValue(undefined) + + const startPromise = taskAccess.startTask("new task") + + expect(postStateSpy).toHaveBeenCalledTimes(1) + expect(mockProvider.postStateToWebviewThrottled).not.toHaveBeenCalled() + expect(saySpy).not.toHaveBeenCalled() + + resolvePostState?.() + await startPromise + + expect(saySpy).toHaveBeenCalledOnce() + expect(saySpy).toHaveBeenCalledWith("text", "new task", undefined) + expect(initiateTaskLoopSpy).toHaveBeenCalledOnce() + }) + }) + describe("start()", () => { it("should be a no-op if the task was already started in the constructor", () => { const task = new Task({ @@ -2859,9 +3073,9 @@ describe("Cline", () => { resumeSpy.mockRestore() }) - it("logs (instead of crashing) when postStateToWebviewWithoutTaskHistory rejects from the queue handler", async () => { + it("logs (instead of crashing) when postStateToWebviewThrottled rejects from the queue handler", async () => { const boom = new Error("postState boom") - mockProvider.postStateToWebviewWithoutTaskHistory = vi.fn().mockRejectedValue(boom) + mockProvider.postStateToWebviewThrottled = vi.fn().mockRejectedValue(boom) const task = new Task({ provider: mockProvider, @@ -2870,13 +3084,14 @@ describe("Cline", () => { startTask: false, }) - // Triggers messageQueueStateChangedHandler -> void postStateToWebviewWithoutTaskHistory() + // Triggers messageQueueStateChangedHandler -> void postStateToWebviewThrottled() task.messageQueueService.addMessage("queued text") await flushMicrotasks() - expect(mockProvider.postStateToWebviewWithoutTaskHistory).toHaveBeenCalled() + expect(mockProvider.postStateToWebviewThrottled).toHaveBeenCalledWith() + expect(mockProvider.postStateToWebviewWithoutTaskHistory).not.toHaveBeenCalled() expect(consoleErrorSpy).toHaveBeenCalledWith( - "[Task#messageQueueStateChangedHandler] postStateToWebviewWithoutTaskHistory failed:", + "[Task#messageQueueStateChangedHandler] postStateToWebviewThrottled failed:", boom, ) }) diff --git a/src/core/task/__tests__/Task.throttle.test.ts b/src/core/task/__tests__/Task.throttle.test.ts index 34d78a4ef9..0eac687e64 100644 --- a/src/core/task/__tests__/Task.throttle.test.ts +++ b/src/core/task/__tests__/Task.throttle.test.ts @@ -79,6 +79,8 @@ describe("Task token usage throttling", () => { log: vi.fn(), postStateToWebview: vi.fn().mockResolvedValue(undefined), postStateToWebviewWithoutTaskHistory: vi.fn().mockResolvedValue(undefined), + postStateToWebviewThrottled: vi.fn().mockResolvedValue(undefined), + flushPostStateToWebviewThrottled: vi.fn().mockResolvedValue(undefined), updateTaskHistory: vi.fn().mockResolvedValue(undefined), } diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index 2ee92edebc..3a6344b9b9 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -6,6 +6,7 @@ import EventEmitter from "events" import { Anthropic } from "@anthropic-ai/sdk" import delay from "delay" import axios from "axios" +import debounce from "lodash.debounce" import pWaitFor from "p-wait-for" import * as vscode from "vscode" @@ -188,6 +189,21 @@ export class ClineProvider private taskEventListeners: WeakMap void>> = new WeakMap() private currentWorkspacePath: string | undefined private _disposed = false + private readonly _postStateToWebviewThrottled = debounce( + async () => { + try { + await this.postStateToWebviewWithoutTaskHistory() + } catch (error) { + this.log( + `[ClineProvider#postStateToWebviewThrottled] Failed to post state: ${ + error instanceof Error ? error.message : String(error) + }`, + ) + } + }, + 500, + { leading: true, trailing: true, maxWait: 1000 }, + ) private readonly rateLimitClock: RateLimitClock = createRateLimitClock() private recentTasksCache?: string[] @@ -689,6 +705,7 @@ export class ClineProvider } this._disposed = true + this._postStateToWebviewThrottled.cancel() this.log("Disposing ClineProvider...") // Reject any tasks still waiting for a scheduler permit so they don't @@ -2178,6 +2195,22 @@ export class ClineProvider await this.postMessageToWebview({ type: "state", state: rest }) } + async postStateToWebviewThrottled(): Promise { + if (this._disposed) { + return + } + + await this._postStateToWebviewThrottled() + } + + async flushPostStateToWebviewThrottled(): Promise { + if (this._disposed) { + return + } + + await this._postStateToWebviewThrottled.flush() + } + /** * Like postStateToWebview but intentionally omits both clineMessages and taskHistory. * diff --git a/src/core/webview/__tests__/ClineProvider.spec.ts b/src/core/webview/__tests__/ClineProvider.spec.ts index b99e502f61..231d4dd3f7 100644 --- a/src/core/webview/__tests__/ClineProvider.spec.ts +++ b/src/core/webview/__tests__/ClineProvider.spec.ts @@ -794,6 +794,135 @@ describe("ClineProvider", () => { expect(postMessageSpy).not.toHaveBeenCalledWith(expect.objectContaining({ type: "action" })) }) + test("postStateToWebviewWithoutTaskHistory waits for the webview post boundary", async () => { + let releasePost!: () => void + const pendingPost = new Promise((resolve) => { + releasePost = resolve + }) + let statePostSettled = false + + vi.spyOn(provider, "getStateToPostToWebview").mockResolvedValue({ + taskHistory: [], + } as unknown as ExtensionState) + const postMessageSpy = vi.spyOn(provider, "postMessageToWebview").mockReturnValue(pendingPost) + + const statePost = provider.postStateToWebviewWithoutTaskHistory() + void statePost.then(() => { + statePostSettled = true + }) + await Promise.resolve() + + expect(postMessageSpy).toHaveBeenCalledOnce() + expect(statePostSettled).toBe(false) + + releasePost() + await statePost + expect(statePostSettled).toBe(true) + }) + + describe("postStateToWebviewThrottled", () => { + beforeEach(() => { + vi.useFakeTimers() + }) + + afterEach(async () => { + await provider.dispose() + vi.useRealTimers() + }) + + test("posts on the leading edge and coalesces a burst into one trailing post", async () => { + const postStateSpy = vi.spyOn(provider, "postStateToWebviewWithoutTaskHistory").mockResolvedValue(undefined) + + await provider.postStateToWebviewThrottled() + await provider.postStateToWebviewThrottled() + await provider.postStateToWebviewThrottled() + + expect(postStateSpy).toHaveBeenCalledTimes(1) + + await vi.advanceTimersByTimeAsync(499) + expect(postStateSpy).toHaveBeenCalledTimes(1) + + await vi.advanceTimersByTimeAsync(1) + expect(postStateSpy).toHaveBeenCalledTimes(2) + }) + + test("does not starve state posts during continuous updates", async () => { + const postStateSpy = vi.spyOn(provider, "postStateToWebviewWithoutTaskHistory").mockResolvedValue(undefined) + + await provider.postStateToWebviewThrottled() + await vi.advanceTimersByTimeAsync(400) + await provider.postStateToWebviewThrottled() + await vi.advanceTimersByTimeAsync(400) + await provider.postStateToWebviewThrottled() + await vi.advanceTimersByTimeAsync(199) + + expect(postStateSpy).toHaveBeenCalledTimes(1) + + await vi.advanceTimersByTimeAsync(1) + expect(postStateSpy).toHaveBeenCalledTimes(2) + }) + + test("flushes a pending trailing post exactly once", async () => { + const postStateSpy = vi.spyOn(provider, "postStateToWebviewWithoutTaskHistory").mockResolvedValue(undefined) + + await provider.postStateToWebviewThrottled() + await provider.postStateToWebviewThrottled() + expect(postStateSpy).toHaveBeenCalledTimes(1) + + await provider.flushPostStateToWebviewThrottled() + expect(postStateSpy).toHaveBeenCalledTimes(2) + + await vi.advanceTimersByTimeAsync(1000) + expect(postStateSpy).toHaveBeenCalledTimes(2) + }) + + test("does not duplicate an idle leading post when flushed", async () => { + const postStateSpy = vi.spyOn(provider, "postStateToWebviewWithoutTaskHistory").mockResolvedValue(undefined) + + await provider.postStateToWebviewThrottled() + await provider.flushPostStateToWebviewThrottled() + await vi.advanceTimersByTimeAsync(1000) + + expect(postStateSpy).toHaveBeenCalledOnce() + }) + + test("handles state post failures inside the debounced callback", async () => { + const error = new Error("state post failed") + const logSpy = vi.spyOn(provider, "log").mockImplementation(() => {}) + vi.spyOn(provider, "postStateToWebviewWithoutTaskHistory").mockRejectedValue(error) + + await expect(provider.postStateToWebviewThrottled()).resolves.toBeUndefined() + expect(logSpy).toHaveBeenCalledWith( + "[ClineProvider#postStateToWebviewThrottled] Failed to post state: state post failed", + ) + }) + + test("stringifies non-Error state post failures inside the debounced callback", async () => { + const logSpy = vi.spyOn(provider, "log").mockImplementation(() => {}) + vi.spyOn(provider, "postStateToWebviewWithoutTaskHistory").mockRejectedValue("state post failed") + + await expect(provider.postStateToWebviewThrottled()).resolves.toBeUndefined() + expect(logSpy).toHaveBeenCalledWith( + "[ClineProvider#postStateToWebviewThrottled] Failed to post state: state post failed", + ) + }) + + test("cancels pending work on dispose and ignores later schedule or flush calls", async () => { + const postStateSpy = vi.spyOn(provider, "postStateToWebviewWithoutTaskHistory").mockResolvedValue(undefined) + + await provider.postStateToWebviewThrottled() + await provider.postStateToWebviewThrottled() + expect(postStateSpy).toHaveBeenCalledTimes(1) + + await provider.dispose() + await vi.advanceTimersByTimeAsync(1000) + await provider.postStateToWebviewThrottled() + await provider.flushPostStateToWebviewThrottled() + + expect(postStateSpy).toHaveBeenCalledTimes(1) + }) + }) + test("postMessageToWebview skips postMessage after dispose", async () => { await provider.resolveWebviewView(mockWebviewView) diff --git a/webview-ui/src/components/chat/ChatView.tsx b/webview-ui/src/components/chat/ChatView.tsx index 0410336609..9390a5d483 100644 --- a/webview-ui/src/components/chat/ChatView.tsx +++ b/webview-ui/src/components/chat/ChatView.tsx @@ -1,4 +1,13 @@ -import React, { forwardRef, useCallback, useEffect, useImperativeHandle, useMemo, useRef, useState } from "react" +import React, { + forwardRef, + useCallback, + useEffect, + useImperativeHandle, + useLayoutEffect, + useMemo, + useRef, + useState, +} from "react" import { useDeepCompareEffect, useEvent } from "react-use" import { Virtuoso, type VirtuosoHandle } from "react-virtuoso" import removeMd from "remove-markdown" @@ -76,6 +85,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction { + currentTaskIdRef.current = currentTaskId + }, [currentTaskId]) useEffect(() => { messagesRef.current = messages @@ -513,13 +528,14 @@ const ChatViewComponent: React.ForwardRefRenderFunction { - const newMap = new Map(prev) - newMap.set(message.text!, message.aggregatedCosts!) - return newMap - }) + if (message.text && message.text === currentTaskIdRef.current && message.aggregatedCosts) { + setAggregatedCostsMap(new Map([[message.text, message.aggregatedCosts]])) } break } @@ -1612,6 +1624,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction 0 - ) - } + aggregatedCost={currentTaskAggregatedCosts?.totalCost} + hasSubtasks={(currentTaskAggregatedCosts?.childrenCost ?? 0) > 0} parentTaskId={currentTaskItem?.parentTaskId} costBreakdown={ - currentTaskItem?.id && aggregatedCostsMap.has(currentTaskItem.id) - ? getCostBreakdownIfNeeded(aggregatedCostsMap.get(currentTaskItem.id)!, { + currentTaskAggregatedCosts + ? getCostBreakdownIfNeeded(currentTaskAggregatedCosts, { own: t("common:costs.own"), subtasks: t("common:costs.subtasks"), }) diff --git a/webview-ui/src/components/chat/__tests__/ChatView.spec.tsx b/webview-ui/src/components/chat/__tests__/ChatView.spec.tsx index 66169e0bac..6e76c2a5b3 100644 --- a/webview-ui/src/components/chat/__tests__/ChatView.spec.tsx +++ b/webview-ui/src/components/chat/__tests__/ChatView.spec.tsx @@ -86,6 +86,25 @@ vi.mock("../ChatRow", () => ({ }, })) +const mockTaskHeaderState = vi.hoisted(() => ({ + renders: [] as Array<{ taskId?: string; aggregatedCost?: number }>, +})) + +vi.mock("../TaskHeader", () => ({ + default: function MockTaskHeader({ task, aggregatedCost }: { task: ClineMessage; aggregatedCost?: number }) { + mockTaskHeaderState.renders.push({ taskId: task.text, aggregatedCost }) + + return ( +
+ ) + }, +})) + vi.mock("../AutoApproveMenu", () => ({ default: () => null, })) @@ -331,6 +350,54 @@ const mockPostMessage = (state: Partial) => { ) } +const dispatchExtensionMessage = async (data: Record) => { + await act(async () => { + window.dispatchEvent(new MessageEvent("message", { data })) + }) +} + +const dispatchTaskState = async (id: string, taskTs: number, childIds: string[] = []) => { + await dispatchExtensionMessage({ + type: "state", + state: { + version: "1.0.0", + clineMessages: [ + { + type: "say", + say: "task", + ts: taskTs, + text: id, + }, + ], + currentTaskId: id, + currentTaskItem: { + id, + ts: taskTs, + task: id, + childIds, + }, + taskHistory: [], + shouldShowAnnouncement: false, + allowedCommands: [], + alwaysAllowExecute: false, + cloudIsAuthenticated: false, + telemetrySetting: "enabled", + }, + }) +} + +const dispatchAggregatedCosts = async (taskId: string, totalCost: number) => { + await dispatchExtensionMessage({ + type: "taskWithAggregatedCosts", + text: taskId, + aggregatedCosts: { + totalCost, + ownCost: 1, + childrenCost: totalCost - 1, + }, + }) +} + const defaultProps: ChatViewProps = { isHidden: false, showAnnouncement: false, @@ -349,6 +416,65 @@ const renderChatView = (props: Partial = {}) => { ) } +describe("ChatView - Aggregated Costs Lifecycle", () => { + beforeEach(() => { + vi.clearAllMocks() + mockTaskHeaderState.renders.length = 0 + }) + + it("clears cached aggregated costs when switching tasks", async () => { + const { getByTestId } = renderChatView() + + await dispatchTaskState("task-a", 1_000, ["child-a"]) + await dispatchAggregatedCosts("task-a", 9) + + await waitFor(() => { + expect(getByTestId("task-header")).toHaveAttribute("data-aggregated-cost", "9") + }) + + // Use the same message timestamp to prove task identity, rather than task.ts, + // drives the reset. + await dispatchTaskState("task-b", 1_000) + await waitFor(() => { + expect(getByTestId("task-header")).toHaveAttribute("data-task-id", "task-b") + expect(getByTestId("task-header")).toHaveAttribute("data-aggregated-cost", "") + }) + + await dispatchTaskState("task-a", 1_000, ["child-a"]) + await waitFor(() => { + expect(getByTestId("task-header")).toHaveAttribute("data-task-id", "task-a") + expect(getByTestId("task-header")).toHaveAttribute("data-aggregated-cost", "") + }) + }) + + it("rejects a delayed aggregated-cost response from the previous task", async () => { + const { getByTestId } = renderChatView() + + await dispatchTaskState("task-a", 1_001, ["child-a"]) + await dispatchTaskState("task-b", 2_001) + await dispatchAggregatedCosts("task-a", 13) + + await waitFor(() => { + expect(getByTestId("task-header")).toHaveAttribute("data-task-id", "task-b") + expect(getByTestId("task-header")).toHaveAttribute("data-aggregated-cost", "") + }) + + mockTaskHeaderState.renders.length = 0 + await dispatchTaskState("task-a", 1_001, ["child-a"]) + + await waitFor(() => { + expect(getByTestId("task-header")).toHaveAttribute("data-task-id", "task-a") + expect(getByTestId("task-header")).toHaveAttribute("data-aggregated-cost", "") + }) + + expect( + mockTaskHeaderState.renders.some( + ({ taskId, aggregatedCost }) => taskId === "task-a" && aggregatedCost === 13, + ), + ).toBe(false) + }) +}) + describe("ChatView - Sound Playing Tests", () => { beforeEach(() => vi.clearAllMocks())