diff --git a/webview-ui/src/components/chat/TaskHeader.tsx b/webview-ui/src/components/chat/TaskHeader.tsx index 558d01f9583..2dcf5fd09cf 100644 --- a/webview-ui/src/components/chat/TaskHeader.tsx +++ b/webview-ui/src/components/chat/TaskHeader.tsx @@ -122,13 +122,8 @@ const TaskHeader: React.FC = ({ }, [task.text, windowWidth]) const isCostAvailable = useMemo(() => { - return ( - apiConfiguration?.apiProvider !== "openai" && - apiConfiguration?.apiProvider !== "ollama" && - apiConfiguration?.apiProvider !== "lmstudio" && - apiConfiguration?.apiProvider !== "gemini" - ) - }, [apiConfiguration?.apiProvider]) + return totalCost !== null && totalCost !== undefined && totalCost > 0 && !isNaN(totalCost) + }, [totalCost]) const shouldShowPromptCacheInfo = doesModelSupportPromptCache && apiConfiguration?.apiProvider !== "openrouter" diff --git a/webview-ui/src/components/chat/__tests__/TaskHeader.test.tsx b/webview-ui/src/components/chat/__tests__/TaskHeader.test.tsx new file mode 100644 index 00000000000..41d45bfeac8 --- /dev/null +++ b/webview-ui/src/components/chat/__tests__/TaskHeader.test.tsx @@ -0,0 +1,114 @@ +import React from "react" +import { render, screen } from "@testing-library/react" +import TaskHeader from "../TaskHeader" +import { ApiConfiguration } from "../../../../../src/shared/api" + +// Mock the vscode API +jest.mock("@/utils/vscode", () => ({ + vscode: { + postMessage: jest.fn(), + }, +})) + +// Mock the ExtensionStateContext +jest.mock("../../../context/ExtensionStateContext", () => ({ + useExtensionState: () => ({ + apiConfiguration: { + apiProvider: "anthropic", + apiKey: "test-api-key", // Add relevant fields + apiModelId: "claude-3-opus-20240229", // Add relevant fields + } as ApiConfiguration, // Optional: Add type assertion if ApiConfiguration is imported + currentTaskItem: null, + }), +})) + +describe("TaskHeader", () => { + const defaultProps = { + task: { text: "Test task", images: [] }, + tokensIn: 100, + tokensOut: 50, + doesModelSupportPromptCache: true, + totalCost: 0.05, + contextTokens: 200, + onClose: jest.fn(), + } + + it("should display cost when totalCost is greater than 0", () => { + render( + , + ) + expect(screen.getByText("$0.0500")).toBeInTheDocument() + }) + + it("should not display cost when totalCost is 0", () => { + render( + , + ) + expect(screen.queryByText("$0.0000")).not.toBeInTheDocument() + }) + + it("should not display cost when totalCost is null", () => { + render( + , + ) + expect(screen.queryByText(/\$/)).not.toBeInTheDocument() + }) + + it("should not display cost when totalCost is undefined", () => { + render( + , + ) + expect(screen.queryByText(/\$/)).not.toBeInTheDocument() + }) + + it("should not display cost when totalCost is NaN", () => { + render( + , + ) + expect(screen.queryByText(/\$/)).not.toBeInTheDocument() + }) +})