Skip to content

Commit cf3871a

Browse files
jrjr
committed
Add a RCC credit balance display
1 parent 18c4d1a commit cf3871a

File tree

10 files changed

+441
-2
lines changed

10 files changed

+441
-2
lines changed

‎packages/cloud/src/CloudAPI.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,4 +134,14 @@ export class CloudAPI {
134134
.parse(data),
135135
})
136136
}
137+
138+
async creditBalance(): Promise<number> {
139+
return this.request("/api/extension/credit-balance", {
140+
method: "GET",
141+
parseResponse: (data) => {
142+
const result = z.object({ balance: z.number() }).parse(data)
143+
return result.balance
144+
},
145+
})
146+
}
137147
}

‎packages/cloud/src/__tests__/CloudAPI.creditBalance.spec.ts‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
import { describe, it, expect, vi, beforeEach, type Mock } from "vitest"
2+
import { CloudAPI } from "../CloudAPI.js"
3+
import { AuthenticationError, CloudAPIError } from "../errors.js"
4+
import type { AuthService } from "@roo-code/types"
5+
6+
// Mock the config module
7+
vi.mock("../config.js", () => ({
8+
getRooCodeApiUrl: () => "https://api.test.com",
9+
}))
10+
11+
// Mock the utils module
12+
vi.mock("../utils.js", () => ({
13+
getUserAgent: () => "test-user-agent",
14+
}))
15+
16+
describe("CloudAPI.creditBalance", () => {
17+
let mockAuthService: {
18+
getSessionToken: Mock<() => string | undefined>
19+
}
20+
let cloudAPI: CloudAPI
21+
22+
beforeEach(() => {
23+
mockAuthService = {
24+
getSessionToken: vi.fn(),
25+
}
26+
cloudAPI = new CloudAPI(mockAuthService as unknown as AuthService)
27+
28+
// Reset fetch mock
29+
global.fetch = vi.fn()
30+
})
31+
32+
it("should fetch credit balance successfully", async () => {
33+
const mockBalance = 12.34
34+
mockAuthService.getSessionToken.mockReturnValue("test-session-token")
35+
36+
global.fetch = vi.fn().mockResolvedValue({
37+
ok: true,
38+
json: async () => ({ balance: mockBalance }),
39+
})
40+
41+
const balance = await cloudAPI.creditBalance()
42+
43+
expect(balance).toBe(mockBalance)
44+
expect(global.fetch).toHaveBeenCalledWith(
45+
"https://api.test.com/api/extension/credit-balance",
46+
expect.objectContaining({
47+
method: "GET",
48+
headers: expect.objectContaining({
49+
Authorization: "Bearer test-session-token",
50+
"Content-Type": "application/json",
51+
"User-Agent": "test-user-agent",
52+
}),
53+
}),
54+
)
55+
})
56+
57+
it("should throw AuthenticationError when session token is missing", async () => {
58+
mockAuthService.getSessionToken.mockReturnValue(undefined)
59+
60+
await expect(cloudAPI.creditBalance()).rejects.toThrow(AuthenticationError)
61+
})
62+
63+
it("should handle API errors", async () => {
64+
mockAuthService.getSessionToken.mockReturnValue("test-session-token")
65+
66+
global.fetch = vi.fn().mockResolvedValue({
67+
ok: false,
68+
status: 500,
69+
statusText: "Internal Server Error",
70+
json: async () => ({ error: "Server error" }),
71+
})
72+
73+
await expect(cloudAPI.creditBalance()).rejects.toThrow(CloudAPIError)
74+
})
75+
76+
it("should handle network errors", async () => {
77+
mockAuthService.getSessionToken.mockReturnValue("test-session-token")
78+
79+
global.fetch = vi.fn().mockRejectedValue(new TypeError("fetch failed"))
80+
81+
await expect(cloudAPI.creditBalance()).rejects.toThrow(
82+
"Network error while calling /api/extension/credit-balance",
83+
)
84+
})
85+
86+
it("should handle invalid response format", async () => {
87+
mockAuthService.getSessionToken.mockReturnValue("test-session-token")
88+
89+
global.fetch = vi.fn().mockResolvedValue({
90+
ok: true,
91+
json: async () => ({ invalid: "response" }),
92+
})
93+
94+
await expect(cloudAPI.creditBalance()).rejects.toThrow()
95+
})
96+
})

‎src/core/webview/__tests__/webviewMessageHandler.rooBalance.spec.ts‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
import { describe, it, expect, vi, beforeEach } from "vitest"
2+
import { webviewMessageHandler } from "../webviewMessageHandler"
3+
import { CloudService } from "@roo-code/cloud"
4+
5+
vi.mock("@roo-code/cloud", () => ({
6+
CloudService: {
7+
hasInstance: vi.fn(),
8+
instance: {
9+
cloudAPI: {
10+
creditBalance: vi.fn(),
11+
},
12+
},
13+
},
14+
}))
15+
16+
describe("webviewMessageHandler - requestRooCreditBalance", () => {
17+
let mockProvider: any
18+
19+
beforeEach(() => {
20+
mockProvider = {
21+
postMessageToWebview: vi.fn(),
22+
contextProxy: {
23+
getValue: vi.fn(),
24+
setValue: vi.fn(),
25+
},
26+
getCurrentTask: vi.fn(),
27+
cwd: "/test/path",
28+
}
29+
30+
vi.clearAllMocks()
31+
})
32+
33+
it("should handle requestRooCreditBalance and return balance", async () => {
34+
const mockBalance = 42.75
35+
const requestId = "test-request-id"
36+
37+
;(CloudService.hasInstance as any).mockReturnValue(true)
38+
;(CloudService.instance.cloudAPI!.creditBalance as any).mockResolvedValue(mockBalance)
39+
40+
await webviewMessageHandler(
41+
mockProvider as any,
42+
{
43+
type: "requestRooCreditBalance",
44+
requestId,
45+
} as any,
46+
)
47+
48+
expect(mockProvider.postMessageToWebview).toHaveBeenCalledWith({
49+
type: "rooCreditBalance",
50+
requestId,
51+
values: { balance: mockBalance },
52+
})
53+
})
54+
55+
it("should handle CloudAPI errors", async () => {
56+
const requestId = "test-request-id"
57+
const errorMessage = "Failed to fetch balance"
58+
59+
;(CloudService.hasInstance as any).mockReturnValue(true)
60+
;(CloudService.instance.cloudAPI!.creditBalance as any).mockRejectedValue(new Error(errorMessage))
61+
62+
await webviewMessageHandler(
63+
mockProvider as any,
64+
{
65+
type: "requestRooCreditBalance",
66+
requestId,
67+
} as any,
68+
)
69+
70+
expect(mockProvider.postMessageToWebview).toHaveBeenCalledWith({
71+
type: "rooCreditBalance",
72+
requestId,
73+
values: { error: errorMessage },
74+
})
75+
})
76+
77+
it("should handle missing CloudService", async () => {
78+
const requestId = "test-request-id"
79+
80+
;(CloudService.hasInstance as any).mockReturnValue(false)
81+
82+
await webviewMessageHandler(
83+
mockProvider as any,
84+
{
85+
type: "requestRooCreditBalance",
86+
requestId,
87+
} as any,
88+
)
89+
90+
expect(mockProvider.postMessageToWebview).toHaveBeenCalledWith({
91+
type: "rooCreditBalance",
92+
requestId,
93+
values: { error: "Cloud service not available" },
94+
})
95+
})
96+
97+
it("should handle missing cloudAPI", async () => {
98+
const requestId = "test-request-id"
99+
100+
;(CloudService.hasInstance as any).mockReturnValue(true)
101+
;(CloudService.instance as any).cloudAPI = null
102+
103+
await webviewMessageHandler(
104+
mockProvider as any,
105+
{
106+
type: "requestRooCreditBalance",
107+
requestId,
108+
} as any,
109+
)
110+
111+
expect(mockProvider.postMessageToWebview).toHaveBeenCalledWith({
112+
type: "rooCreditBalance",
113+
requestId,
114+
values: { error: "Cloud service not available" },
115+
})
116+
})
117+
})

‎src/core/webview/webviewMessageHandler.ts‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1006,6 +1006,31 @@ export const webviewMessageHandler = async (
10061006
}
10071007
break
10081008
}
1009+
case "requestRooCreditBalance": {
1010+
// Fetch Roo credit balance using CloudAPI
1011+
const requestId = message.requestId
1012+
try {
1013+
if (!CloudService.hasInstance() || !CloudService.instance.cloudAPI) {
1014+
throw new Error("Cloud service not available")
1015+
}
1016+
1017+
const balance = await CloudService.instance.cloudAPI.creditBalance()
1018+
1019+
provider.postMessageToWebview({
1020+
type: "rooCreditBalance",
1021+
requestId,
1022+
values: { balance },
1023+
})
1024+
} catch (error) {
1025+
const errorMessage = error instanceof Error ? error.message : String(error)
1026+
provider.postMessageToWebview({
1027+
type: "rooCreditBalance",
1028+
requestId,
1029+
values: { error: errorMessage },
1030+
})
1031+
}
1032+
break
1033+
}
10091034
case "requestOpenAiModels":
10101035
if (message?.values?.baseUrl && message?.values?.apiKey) {
10111036
const openAiModels = await getOpenAiModels(

‎src/shared/ExtensionMessage.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,7 @@ export interface ExtensionMessage {
112112
| "authenticatedUser"
113113
| "condenseTaskContextResponse"
114114
| "singleRouterModelFetchResponse"
115+
| "rooCreditBalance"
115116
| "indexingStatusUpdate"
116117
| "indexCleared"
117118
| "codebaseIndexConfig"

‎src/shared/WebviewMessage.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ export interface WebviewMessage {
6060
| "requestOllamaModels"
6161
| "requestLmStudioModels"
6262
| "requestRooModels"
63+
| "requestRooCreditBalance"
6364
| "requestVsCodeLmModels"
6465
| "requestHuggingFaceModels"
6566
| "openImage"

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import { vscode } from "@src/utils/vscode"
77
import { Button } from "@src/components/ui"
88

99
import { ModelPicker } from "../ModelPicker"
10+
import { RooBalanceDisplay } from "./RooBalanceDisplay"
1011

1112
type RooProps = {
1213
apiConfiguration: ProviderSettings
@@ -30,8 +31,11 @@ export const Roo = ({
3031
return (
3132
<>
3233
{cloudIsAuthenticated ? (
33-
<div className="text-sm text-vscode-descriptionForeground">
34-
{t("settings:providers.roo.authenticatedMessage")}
34+
<div className="flex justify-between items-center mb-2">
35+
<div className="text-sm text-vscode-descriptionForeground">
36+
{t("settings:providers.roo.authenticatedMessage")}
37+
</div>
38+
<RooBalanceDisplay />
3539
</div>
3640
) : (
3741
<div className="flex flex-col gap-2">

‎webview-ui/src/components/settings/providers/RooBalanceDisplay.tsx‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
import { VSCodeLink } from "@vscode/webview-ui-toolkit/react"
2+
3+
import { useRooCreditBalance } from "@/components/ui/hooks/useRooCreditBalance"
4+
import { useExtensionState } from "@src/context/ExtensionStateContext"
5+
6+
export const RooBalanceDisplay = () => {
7+
const { data: balance } = useRooCreditBalance()
8+
const { cloudApiUrl } = useExtensionState()
9+
10+
if (balance === null || balance === undefined) {
11+
return null
12+
}
13+
14+
const formattedBalance = balance.toFixed(2)
15+
const billingUrl = cloudApiUrl ? `${cloudApiUrl.replace(/\/$/, "")}/billing` : "https://app.roocode.com/billing"
16+
17+
return (
18+
<VSCodeLink href={billingUrl} className="text-vscode-foreground hover:underline whitespace-nowrap">
19+
${formattedBalance}
20+
</VSCodeLink>
21+
)
22+
}

0 commit comments

Comments
 (0)