Also refresh token before completePrompt()

Author: jr

src/api/providers/__tests__/roo.spec.ts

@@ -182,6 +182,21 @@ describe("RooHandler", () => {
 			handler = new RooHandler(mockOptions)
 		})
 
+		it("should update API key before making request", async () => {
+			// Set up a fresh token that will be returned when createMessage is called
+			const freshToken = "fresh-session-token"
+			mockGetSessionTokenFn.mockReturnValue(freshToken)
+
+			const stream = handler.createMessage(systemPrompt, messages)
+			// Consume the stream to trigger the API call
+			for await (const _chunk of stream) {
+				// Just consume
+			}
+
+			// Verify getSessionToken was called to get the fresh token
+			expect(mockGetSessionTokenFn).toHaveBeenCalled()
+		})
+
 		it("should handle streaming responses", async () => {
 			const stream = handler.createMessage(systemPrompt, messages)
 			const chunks: any[] = []
@@ -290,6 +305,25 @@ describe("RooHandler", () => {
 			})
 		})
 
+		it("should update API key before making request", async () => {
+			// Set up a fresh token that will be returned when completePrompt is called
+			const freshToken = "fresh-session-token"
+			mockGetSessionTokenFn.mockReturnValue(freshToken)
+
+			// Access the client's apiKey property to verify it gets updated
+			const clientApiKeyGetter = vitest.fn()
+			Object.defineProperty(handler["client"], "apiKey", {
+				get: clientApiKeyGetter,
+				set: vitest.fn(),
+				configurable: true,
+			})
+
+			await handler.completePrompt("Test prompt")
+
+			// Verify getSessionToken was called to get the fresh token
+			expect(mockGetSessionTokenFn).toHaveBeenCalled()
+		})
+
 		it("should handle API errors", async () => {
 			mockCreate.mockRejectedValueOnce(new Error("API Error"))
 			await expect(handler.completePrompt("Test prompt")).rejects.toThrow(

src/api/providers/roo.ts

@@ -36,11 +36,7 @@ export class RooHandler extends BaseOpenAiCompatibleProvider<string> {
 	private fetcherBaseURL: string
 
 	constructor(options: ApiHandlerOptions) {
-		let sessionToken: string | undefined = undefined
-
-		if (CloudService.hasInstance()) {
-			sessionToken = CloudService.instance.authService?.getSessionToken()
-		}
+		const sessionToken = getSessionToken()
 
 		let baseURL = process.env.ROO_CODE_PROVIDER_URL ?? "https://api.roocode.com/proxy"
 
@@ -55,7 +51,7 @@ export class RooHandler extends BaseOpenAiCompatibleProvider<string> {
 			...options,
 			providerName: "Roo Code Cloud",
 			baseURL, // Already has /v1 suffix
-			apiKey: getSessionToken(),
+			apiKey: sessionToken,
 			defaultProviderModelId: rooDefaultModelId,
 			providerModels: {},
 			defaultTemperature: 0.7,
@@ -176,6 +172,11 @@ export class RooHandler extends BaseOpenAiCompatibleProvider<string> {
 			}
 		}
 	}
+	override async completePrompt(prompt: string): Promise<string> {
+		// Update API key before making request to ensure we use the latest session token
+		this.client.apiKey = getSessionToken()
+		return super.completePrompt(prompt)
+	}
 
 	private async loadDynamicModels(baseURL: string, apiKey?: string): Promise<void> {
 		try {