Expose thinking tokens for roo/sonic (#7212)

Co-authored-by: Roo Code <[email protected]>
Co-authored-by: Matt Rubens <[email protected]>

Author: 3 people

.env.sample

@@ -3,3 +3,4 @@ POSTHOG_API_KEY=key-goes-here
 # Roo Code Cloud / Local Development
 CLERK_BASE_URL=https://epic-chamois-85.clerk.accounts.dev
 ROO_CODE_API_URL=http://localhost:3000
+ROO_CODE_PROVIDER_URL=http://localhost:8080/proxy/v1

packages/types/src/providers/roo.ts

@@ -10,10 +10,10 @@ export const rooModels = {
 		maxTokens: 8192,
 		contextWindow: 262_144,
 		supportsImages: false,
-		supportsPromptCache: false,
+		supportsPromptCache: true,
 		inputPrice: 0,
 		outputPrice: 0,
 		description:
-			"Stealth coding model with 262K context window, accessible for free through Roo Code Cloud for a limited time. (Note: prompts and completions are logged by the model creator and used to improve the model.)",
+			"A stealth reasoning model that is blazing fast and excels at agentic coding, accessible for free through Roo Code Cloud for a limited time. (Note: prompts and completions are logged by the model creator and used to improve the model.)",
 	},
 } as const satisfies Record<string, ModelInfo>

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

@@ -331,7 +331,7 @@ describe("RooHandler", () => {
 			expect(modelInfo.info.maxTokens).toBe(8192)
 			expect(modelInfo.info.contextWindow).toBe(262_144)
 			expect(modelInfo.info.supportsImages).toBe(false)
-			expect(modelInfo.info.supportsPromptCache).toBe(false)
+			expect(modelInfo.info.supportsPromptCache).toBe(true)
 			expect(modelInfo.info.inputPrice).toBe(0)
 			expect(modelInfo.info.outputPrice).toBe(0)
 		})

src/api/providers/base-openai-compatible-provider.ts

@@ -62,11 +62,11 @@ export abstract class BaseOpenAiCompatibleProvider<ModelName extends string>
 		})
 	}
 
-	override async *createMessage(
+	protected createStream(
 		systemPrompt: string,
 		messages: Anthropic.Messages.MessageParam[],
 		metadata?: ApiHandlerCreateMessageMetadata,
-	): ApiStream {
+	) {
 		const {
 			id: model,
 			info: { maxTokens: max_tokens },
@@ -83,7 +83,15 @@ export abstract class BaseOpenAiCompatibleProvider<ModelName extends string>
 			stream_options: { include_usage: true },
 		}
 
-		const stream = await this.client.chat.completions.create(params)
+		return this.client.chat.completions.create(params)
+	}
+
+	override async *createMessage(
+		systemPrompt: string,
+		messages: Anthropic.Messages.MessageParam[],
+		metadata?: ApiHandlerCreateMessageMetadata,
+	): ApiStream {
+		const stream = await this.createStream(systemPrompt, messages, metadata)
 
 		for await (const chunk of stream) {
 			const delta = chunk.choices[0]?.delta

src/api/providers/roo.ts

@@ -1,10 +1,14 @@
+import { Anthropic } from "@anthropic-ai/sdk"
 import { rooDefaultModelId, rooModels, type RooModelId } from "@roo-code/types"
 import { CloudService } from "@roo-code/cloud"
 
 import type { ApiHandlerOptions } from "../../shared/api"
-import { BaseOpenAiCompatibleProvider } from "./base-openai-compatible-provider"
+import { ApiStream } from "../transform/stream"
 import { t } from "../../i18n"
 
+import type { ApiHandlerCreateMessageMetadata } from "../index"
+import { BaseOpenAiCompatibleProvider } from "./base-openai-compatible-provider"
+
 export class RooHandler extends BaseOpenAiCompatibleProvider<RooModelId> {
 	constructor(options: ApiHandlerOptions) {
 		// Check if CloudService is available and get the session token.
@@ -21,14 +25,50 @@ export class RooHandler extends BaseOpenAiCompatibleProvider<RooModelId> {
 		super({
 			...options,
 			providerName: "Roo Code Cloud",
-			baseURL: "https://api.roocode.com/proxy/v1",
+			baseURL: process.env.ROO_CODE_PROVIDER_URL ?? "https://api.roocode.com/proxy/v1",
 			apiKey: sessionToken,
 			defaultProviderModelId: rooDefaultModelId,
 			providerModels: rooModels,
 			defaultTemperature: 0.7,
 		})
 	}
 
+	override async *createMessage(
+		systemPrompt: string,
+		messages: Anthropic.Messages.MessageParam[],
+		metadata?: ApiHandlerCreateMessageMetadata,
+	): ApiStream {
+		const stream = await this.createStream(systemPrompt, messages, metadata)
+
+		for await (const chunk of stream) {
+			const delta = chunk.choices[0]?.delta
+
+			if (delta) {
+				if (delta.content) {
+					yield {
+						type: "text",
+						text: delta.content,
+					}
+				}
+
+				if ("reasoning_content" in delta && typeof delta.reasoning_content === "string") {
+					yield {
+						type: "reasoning",
+						text: delta.reasoning_content,
+					}
+				}
+			}
+
+			if (chunk.usage) {
+				yield {
+					type: "usage",
+					inputTokens: chunk.usage.prompt_tokens || 0,
+					outputTokens: chunk.usage.completion_tokens || 0,
+				}
+			}
+		}
+	}
+
 	override getModel() {
 		const modelId = this.options.apiModelId || rooDefaultModelId
 		const modelInfo = this.providerModels[modelId as RooModelId] ?? this.providerModels[rooDefaultModelId]
@@ -44,7 +84,7 @@ export class RooHandler extends BaseOpenAiCompatibleProvider<RooModelId> {
 				maxTokens: 8192,
 				contextWindow: 262_144,
 				supportsImages: false,
-				supportsPromptCache: false,
+				supportsPromptCache: true,
 				inputPrice: 0,
 				outputPrice: 0,
 			},