roo provider: update session token on every request (#8923)

* roo provider: update session token on every request

* Cleanup: remove unused imports

* Also refresh token before completePrompt()
This commit is contained in:
John Richmond 2025-10-29 22:35:51 -07:00 committed by GitHub
parent ed45d1c081
commit 59ab1f3f3c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 49 additions and 33 deletions

View file

@ -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(

View file

@ -1,7 +1,7 @@
import { Anthropic } from "@anthropic-ai/sdk"
import OpenAI from "openai"
import { AuthState, rooDefaultModelId, type ModelInfo } from "@roo-code/types"
import { rooDefaultModelId } from "@roo-code/types"
import { CloudService } from "@roo-code/cloud"
import type { ApiHandlerOptions, ModelRecord } from "../../shared/api"
@ -12,9 +12,8 @@ import type { RooReasoningParams } from "../transform/reasoning"
import { getRooReasoning } from "../transform/reasoning"
import type { ApiHandlerCreateMessageMetadata } from "../index"
import { DEFAULT_HEADERS } from "./constants"
import { BaseOpenAiCompatibleProvider } from "./base-openai-compatible-provider"
import { getModels, flushModels, getModelsFromCache } from "../providers/fetchers/modelCache"
import { getModels, getModelsFromCache } from "../providers/fetchers/modelCache"
import { handleOpenAIError } from "./utils/openai-error-handler"
// Extend OpenAI's CompletionUsage to include Roo specific fields
@ -28,16 +27,16 @@ type RooChatCompletionParams = OpenAI.Chat.ChatCompletionCreateParamsStreaming &
reasoning?: RooReasoningParams
}
function getSessionToken(): string {
const token = CloudService.hasInstance() ? CloudService.instance.authService?.getSessionToken() : undefined
return token ?? "unauthenticated"
}
export class RooHandler extends BaseOpenAiCompatibleProvider<string> {
private authStateListener?: (state: { state: AuthState }) => void
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"
@ -52,7 +51,7 @@ export class RooHandler extends BaseOpenAiCompatibleProvider<string> {
...options,
providerName: "Roo Code Cloud",
baseURL, // Already has /v1 suffix
apiKey: sessionToken || "unauthenticated", // Use a placeholder if no token.
apiKey: sessionToken,
defaultProviderModelId: rooDefaultModelId,
providerModels: {},
defaultTemperature: 0.7,
@ -63,29 +62,6 @@ export class RooHandler extends BaseOpenAiCompatibleProvider<string> {
this.loadDynamicModels(this.fetcherBaseURL, sessionToken).catch((error) => {
console.error("[RooHandler] Failed to load dynamic models:", error)
})
if (CloudService.hasInstance()) {
const cloudService = CloudService.instance
this.authStateListener = (state: { state: AuthState }) => {
// Update OpenAI client with current auth token
// Note: Model cache flush/reload is handled by extension.ts authStateChangedHandler
const newToken = cloudService.authService?.getSessionToken()
this.client = new OpenAI({
baseURL: this.baseURL,
apiKey: newToken ?? "unauthenticated",
defaultHeaders: DEFAULT_HEADERS,
})
}
cloudService.on("auth-state-changed", this.authStateListener)
}
}
dispose() {
if (this.authStateListener && CloudService.hasInstance()) {
CloudService.instance.off("auth-state-changed", this.authStateListener)
}
}
protected override createStream(
@ -127,6 +103,7 @@ export class RooHandler extends BaseOpenAiCompatibleProvider<string> {
}
try {
this.client.apiKey = getSessionToken()
return this.client.chat.completions.create(rooParams, requestOptions)
} catch (error) {
throw handleOpenAIError(error, this.providerName)
@ -195,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 {