Replace the provider docs with the balance when logged in

Author: mrubens

webview-ui/src/components/settings/ApiOptions.tsx

@@ -115,6 +115,7 @@ import { TemperatureControl } from "./TemperatureControl"
 import { RateLimitSecondsControl } from "./RateLimitSecondsControl"
 import { ConsecutiveMistakeLimitControl } from "./ConsecutiveMistakeLimitControl"
 import { BedrockCustomArn } from "./providers/BedrockCustomArn"
+import { RooBalanceDisplay } from "./providers/RooBalanceDisplay"
 import { buildDocLink } from "@src/utils/docLinks"
 
 export interface ApiOptionsProps {
@@ -460,12 +461,16 @@ const ApiOptions = ({
 			<div className="flex flex-col gap-1 relative">
 				<div className="flex justify-between items-center">
 					<label className="block font-medium mb-1">{t("settings:providers.apiProvider")}</label>
-					{docs && (
-						<div className="text-xs text-vscode-descriptionForeground">
-							<VSCodeLink href={docs.url} className="hover:text-vscode-foreground" target="_blank">
-								{t("settings:providers.providerDocumentation", { provider: docs.name })}
-							</VSCodeLink>
-						</div>
+					{selectedProvider === "roo" && cloudIsAuthenticated ? (
+						<RooBalanceDisplay />
+					) : (
+						docs && (
+							<div className="text-xs text-vscode-descriptionForeground">
+								<VSCodeLink href={docs.url} className="hover:text-vscode-foreground" target="_blank">
+									{t("settings:providers.providerDocumentation", { provider: docs.name })}
+								</VSCodeLink>
+							</div>
+						)
 					)}
 				</div>
 				<SearchableSelect

webview-ui/src/components/settings/__tests__/ApiOptions.spec.tsx

@@ -5,7 +5,8 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"
 
 import { type ModelInfo, type ProviderSettings, openAiModelInfoSaneDefaults } from "@roo-code/types"
 
-import { ExtensionStateContextProvider } from "@src/context/ExtensionStateContext"
+import * as ExtensionStateContext from "@src/context/ExtensionStateContext"
+const { ExtensionStateContextProvider } = ExtensionStateContext
 
 import ApiOptions, { ApiOptionsProps } from "../ApiOptions"
 
@@ -238,6 +239,18 @@ vi.mock("../providers/LiteLLM", () => ({
 	),
 }))
 
+// Mock Roo provider for tests
+vi.mock("../providers/Roo", () => ({
+	Roo: ({ cloudIsAuthenticated }: any) => (
+		<div data-testid="roo-provider">{cloudIsAuthenticated ? "Authenticated" : "Not Authenticated"}</div>
+	),
+}))
+
+// Mock RooBalanceDisplay for tests
+vi.mock("../providers/RooBalanceDisplay", () => ({
+	RooBalanceDisplay: () => <div data-testid="roo-balance-display">Balance: $10.00</div>,
+}))
+
 vi.mock("@src/components/ui/hooks/useSelectedModel", () => ({
 	useSelectedModel: vi.fn((apiConfiguration: ProviderSettings) => {
 		if (apiConfiguration.apiModelId?.includes("thinking")) {
@@ -563,4 +576,40 @@ describe("ApiOptions", () => {
 			expect(screen.queryByTestId("litellm-provider")).not.toBeInTheDocument()
 		})
 	})
+
+	describe("Roo provider tests", () => {
+		it("shows balance display when authenticated", () => {
+			// Mock useExtensionState to return authenticated state
+			const useExtensionStateMock = vi.spyOn(ExtensionStateContext, "useExtensionState")
+			useExtensionStateMock.mockReturnValue({
+				cloudIsAuthenticated: true,
+				organizationAllowList: { providers: {} },
+			} as any)
+
+			renderApiOptions({
+				apiConfiguration: {
+					apiProvider: "roo",
+				},
+			})
+
+			expect(screen.getByTestId("roo-balance-display")).toBeInTheDocument()
+		})
+
+		it("does not show balance display when not authenticated", () => {
+			// Mock useExtensionState to return unauthenticated state
+			const useExtensionStateMock = vi.spyOn(ExtensionStateContext, "useExtensionState")
+			useExtensionStateMock.mockReturnValue({
+				cloudIsAuthenticated: false,
+				organizationAllowList: { providers: {} },
+			} as any)
+
+			renderApiOptions({
+				apiConfiguration: {
+					apiProvider: "roo",
+				},
+			})
+
+			expect(screen.queryByTestId("roo-balance-display")).not.toBeInTheDocument()
+		})
+	})
 })

webview-ui/src/components/settings/providers/Roo.tsx

@@ -7,7 +7,6 @@ import { vscode } from "@src/utils/vscode"
 import { Button } from "@src/components/ui"
 
 import { ModelPicker } from "../ModelPicker"
-import { RooBalanceDisplay } from "./RooBalanceDisplay"
 
 type RooProps = {
 	apiConfiguration: ProviderSettings
@@ -35,7 +34,6 @@ export const Roo = ({
 					<div className="text-sm text-vscode-descriptionForeground">
 						{t("settings:providers.roo.authenticatedMessage")}
 					</div>
-					<RooBalanceDisplay />
 				</div>
 			) : (
 				<div className="flex flex-col gap-2">