fix: address PR review feedback for mode display indicators

- Changed from findModeBySlug to getModeBySlug for better mode resolution
- Removed unnecessary 'items-center' CSS class from ModeBadge
- Simplified mock setup in TaskHeader.spec.tsx for better maintainability
- Improved tooltip test clarity in ModeBadge.spec.tsx
- All tests passing successfully

Author: MuriloFP

webview-ui/src/components/chat/__tests__/TaskHeader.spec.tsx

@@ -33,15 +33,9 @@ vi.mock("@vscode/webview-ui-toolkit/react", () => ({
 }))
 
 // Mock the ExtensionStateContext
-const { mockGetCurrentTaskItem, mockSetCurrentTaskItem } = vi.hoisted(() => {
-	let currentTaskItem: any = { id: "test-task-id", mode: "code" }
-	return {
-		mockGetCurrentTaskItem: () => currentTaskItem,
-		mockSetCurrentTaskItem: (newItem: any) => {
-			currentTaskItem = newItem
-		},
-	}
-})
+const mockCurrentTaskItem = vi.hoisted(() => ({
+	value: { id: "test-task-id", mode: "code" } as any,
+}))
 
 vi.mock("@src/context/ExtensionStateContext", () => ({
 	useExtensionState: () => ({
@@ -50,7 +44,7 @@ vi.mock("@src/context/ExtensionStateContext", () => ({
 			apiKey: "test-api-key",
 			apiModelId: "claude-3-opus-20240229",
 		} as ProviderSettings,
-		currentTaskItem: mockGetCurrentTaskItem(),
+		currentTaskItem: mockCurrentTaskItem.value,
 		customModes: [],
 	}),
 }))
@@ -147,22 +141,16 @@ describe("TaskHeader", () => {
 
 	it("should not display mode badge when currentTaskItem has no mode", () => {
 		// Override the mock for this test
-		const originalTaskItem = mockGetCurrentTaskItem()
-		mockSetCurrentTaskItem({
+		const originalTaskItem = mockCurrentTaskItem.value
+		mockCurrentTaskItem.value = {
 			id: "test-task-id",
-			number: 1,
-			ts: Date.now(),
-			task: "Test task",
-			tokensIn: 100,
-			tokensOut: 50,
-			totalCost: 0.05,
 			// No mode property
-		})
+		}
 
 		renderTaskHeader()
 		expect(screen.queryByTestId("mode-badge")).not.toBeInTheDocument()
 
 		// Restore original mock
-		mockSetCurrentTaskItem(originalTaskItem)
+		mockCurrentTaskItem.value = originalTaskItem
 	})
 })

webview-ui/src/components/common/ModeBadge.tsx

@@ -1,5 +1,5 @@
 import React from "react"
-import { findModeBySlug, getAllModes } from "@roo/modes"
+import { getModeBySlug } from "@roo/modes"
 import { useExtensionState } from "@/context/ExtensionStateContext"
 import { StandardTooltip } from "@/components/ui"
 import { cn } from "@/lib/utils"
@@ -16,16 +16,15 @@ export const ModeBadge: React.FC<ModeBadgeProps> = ({ modeSlug, className }) =>
 		return null
 	}
 
-	// Get all modes (built-in + custom)
-	const allModes = getAllModes(customModes)
-	const mode = findModeBySlug(modeSlug, allModes)
+	// Get mode using getModeBySlug which checks custom modes first, then built-in
+	const mode = getModeBySlug(modeSlug, customModes)
 	const displayName = mode?.name || modeSlug // Fallback to slug if mode deleted
 
 	return (
 		<StandardTooltip content={displayName}>
 			<span
 				className={cn(
-					"inline-flex items-center px-2 py-0.5 rounded text-xs font-medium",
+					"inline-flex px-2 py-0.5 rounded text-xs font-medium",
 					"bg-vscode-badge-background text-vscode-badge-foreground",
 					"max-w-[120px] truncate",
 					className,

webview-ui/src/components/common/__tests__/ModeBadge.spec.tsx

@@ -3,13 +3,12 @@ import { ModeBadge } from "../ModeBadge"
 
 // Mock the shared modes module
 vi.mock("@roo/modes", () => ({
-	findModeBySlug: vi.fn((slug, _modes) => {
+	getModeBySlug: vi.fn((slug, _customModes) => {
 		if (slug === "code") return { slug: "code", name: "💻 Code" }
 		if (slug === "architect") return { slug: "architect", name: "🏗️ Architect" }
 		if (slug === "custom") return { slug: "custom", name: "Very Long Custom Mode Name That Should Be Truncated" }
 		return undefined
 	}),
-	getAllModes: vi.fn(() => []),
 }))
 
 // Mock ExtensionStateContext
@@ -42,10 +41,11 @@ describe("ModeBadge", () => {
 		expect(badge).toHaveClass("max-w-[120px]")
 	})
 
-	it("shows full name in tooltip", async () => {
+	it("shows full name in tooltip", () => {
 		render(<ModeBadge modeSlug="custom" />)
-		// Tooltip content would be tested with user interaction
-		// This is a simplified test
-		expect(screen.getByText(/Very Long Custom Mode Name/)).toBeInTheDocument()
+		// The tooltip content is set via the StandardTooltip's content prop
+		// We verify the text is rendered in the badge itself
+		const badge = screen.getByText(/Very Long Custom Mode Name/)
+		expect(badge).toBeInTheDocument()
 	})
 })