diff --git a/webview-ui/src/__tests__/ContextWindowProgress.test.tsx b/webview-ui/src/__tests__/ContextWindowProgress.test.tsx
index bf0af4598d..431cb6136e 100644
--- a/webview-ui/src/__tests__/ContextWindowProgress.test.tsx
+++ b/webview-ui/src/__tests__/ContextWindowProgress.test.tsx
@@ -9,6 +9,11 @@ jest.mock("@/utils/format", () => ({
 	formatLargeNumber: jest.fn((num) => num.toString()),
 }))
 
+// Mock VSCodeBadge component for all tests
+jest.mock("@vscode/webview-ui-toolkit/react", () => ({
+	VSCodeBadge: ({ children }: { children: React.ReactNode }) => <div data-testid="vscode-badge">{children}</div>,
+}))
+
 // Mock ExtensionStateContext since we use useExtensionState
 jest.mock("../context/ExtensionStateContext", () => ({
 	useExtensionState: jest.fn(() => ({
diff --git a/webview-ui/src/components/chat/TaskHeader.tsx b/webview-ui/src/components/chat/TaskHeader.tsx
index 3053099b10..982b76bc0d 100644
--- a/webview-ui/src/components/chat/TaskHeader.tsx
+++ b/webview-ui/src/components/chat/TaskHeader.tsx
@@ -6,7 +6,7 @@ import { useTranslation } from "react-i18next"
 import { vscode } from "@/utils/vscode"
 import { formatLargeNumber } from "@/utils/format"
 import { calculateTokenDistribution, getMaxTokensForModel } from "@/utils/model-utils"
-import { Button, Badge } from "@/components/ui"
+import { Button } from "@/components/ui"
 
 import { ClineMessage } from "../../../../src/shared/ExtensionMessage"
 import { mentionRegexGlobal } from "../../../../src/shared/context-mentions"
@@ -17,6 +17,7 @@ import Thumbnails from "../common/Thumbnails"
 import { normalizeApiConfiguration } from "../settings/ApiOptions"
 import { DeleteTaskDialog } from "../history/DeleteTaskDialog"
 import { cn } from "@/lib/utils"
+import { VSCodeBadge } from "@vscode/webview-ui-toolkit/react"
 
 interface TaskHeaderProps {
 	task: ClineMessage
@@ -95,7 +96,7 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 							contextTokens={contextTokens || 0}
 							maxTokens={getMaxTokensForModel(selectedModelInfo, apiConfiguration)}
 						/>
-						{!!totalCost && <Badge>${totalCost.toFixed(2)}</Badge>}
+						{!!totalCost && <VSCodeBadge>${totalCost.toFixed(2)}</VSCodeBadge>}
 					</div>
 				)}
 				{/* Expanded state: Show task text and images */}
@@ -278,7 +279,7 @@ const ContextWindowProgress = ({ contextWindow, contextTokens, maxTokens }: Cont
 					/>
 
 					{/* Main progress bar container */}
-					<div className="flex items-center h-1 rounded-[2px] overflow-hidden w-full bg-[color-mix(in_srgb,var(--vscode-badge-foreground)_20%,transparent)]">
+					<div className="flex items-center h-1 rounded-[2px] overflow-hidden w-full bg-[color-mix(in_srgb,var(--vscode-foreground)_20%,transparent)]">
 						{/* Current tokens container */}
 						<div className="relative h-full" style={{ width: `${currentPercent}%` }}>
 							{/* Invisible overlay for current tokens section */}
@@ -291,7 +292,7 @@ const ContextWindowProgress = ({ contextWindow, contextTokens, maxTokens }: Cont
 								data-testid="context-tokens-used"
 							/>
 							{/* Current tokens used - darkest */}
-							<div className="h-full w-full bg-[var(--vscode-badge-foreground)] transition-width duration-300 ease-out" />
+							<div className="h-full w-full bg-[var(--vscode-foreground)] transition-width duration-300 ease-out" />
 						</div>
 
 						{/* Container for reserved tokens */}
@@ -305,7 +306,7 @@ const ContextWindowProgress = ({ contextWindow, contextTokens, maxTokens }: Cont
 								data-testid="context-reserved-tokens"
 							/>
 							{/* Reserved for output section - medium gray */}
-							<div className="h-full w-full bg-[color-mix(in_srgb,var(--vscode-badge-foreground)_30%,transparent)] transition-width duration-300 ease-out" />
+							<div className="h-full w-full bg-[color-mix(in_srgb,var(--vscode-foreground)_30%,transparent)] transition-width duration-300 ease-out" />
 						</div>
 
 						{/* Empty section (if any) */}
diff --git a/webview-ui/src/components/chat/__tests__/TaskHeader.test.tsx b/webview-ui/src/components/chat/__tests__/TaskHeader.test.tsx
index 6a81a9b540..6880de29d7 100644
--- a/webview-ui/src/components/chat/__tests__/TaskHeader.test.tsx
+++ b/webview-ui/src/components/chat/__tests__/TaskHeader.test.tsx
@@ -12,6 +12,11 @@ jest.mock("@/utils/vscode", () => ({
 	},
 }))
 
+// Mock the VSCodeBadge component
+jest.mock("@vscode/webview-ui-toolkit/react", () => ({
+	VSCodeBadge: ({ children }: { children: React.ReactNode }) => <div data-testid="vscode-badge">{children}</div>,
+}))
+
 // Mock the ExtensionStateContext
 jest.mock("../../../context/ExtensionStateContext", () => ({
 	useExtensionState: () => ({
diff --git a/webview-ui/src/i18n/__tests__/TranslationContext.test.tsx b/webview-ui/src/i18n/__tests__/TranslationContext.test.tsx
index 0b4dbf8238..78fbfc23d4 100644
--- a/webview-ui/src/i18n/__tests__/TranslationContext.test.tsx
+++ b/webview-ui/src/i18n/__tests__/TranslationContext.test.tsx
@@ -2,7 +2,6 @@ import React from "react"
 import { render } from "@testing-library/react"
 import "@testing-library/jest-dom"
 import TranslationProvider, { useAppTranslation } from "../TranslationContext"
-import { setupI18nForTests } from "../test-utils"
 
 // Mock the useExtensionState hook
 jest.mock("@/context/ExtensionStateContext", () => ({
@@ -23,11 +22,6 @@ const TestComponent = () => {
 }
 
 describe("TranslationContext", () => {
-	beforeAll(() => {
-		// Initialize i18next with test translations
-		setupI18nForTests()
-	})
-
 	it("should provide translations via context", () => {
 		const { getByTestId } = render(
 			<TranslationProvider>
diff --git a/webview-ui/src/setupTests.ts b/webview-ui/src/setupTests.ts
index 298fa25f8f..01034af37f 100644
--- a/webview-ui/src/setupTests.ts
+++ b/webview-ui/src/setupTests.ts
@@ -1,4 +1,8 @@
 import "@testing-library/jest-dom"
+import { setupI18nForTests } from "./i18n/test-utils"
+
+// Set up i18n for all tests
+setupI18nForTests()
 
 // Mock crypto.getRandomValues
 Object.defineProperty(window, "crypto", {