Merge pull request #2893 from ivanarifin/virtual-quota-provider

fix(virtual-quota): display active model in UI for the frontend

Author: chrarnoldus

.changeset/tender-files-leave.md

@@ -0,0 +1,10 @@
+---
+"kilo-code": patch
+---
+
+fix(virtual-quota): display active model in UI for the frontend
+
+When the backend switches the model, it now sends out a "model has changed" signal by emitting event.
+The main application logic catches this signal and immediately tells the user interface to refresh itself.
+The user interface then updates the display to show the name of the new, currently active model.
+This will also keep the backend and the frontend active model in sync

packages/types/src/task.ts

@@ -149,6 +149,9 @@ export type TaskEvents = {
 	// Task Execution
 	[RooCodeEventName.Message]: [{ action: "created" | "updated"; message: ClineMessage }]
 	[RooCodeEventName.TaskModeSwitched]: [taskId: string, mode: string]
+
+	modelChanged: [] // kilocode_change: Add modelChanged event for virtual quota fallback
+
 	[RooCodeEventName.TaskAskResponded]: []
 	[RooCodeEventName.TaskUserMessage]: [taskId: string]
 

src/api/index.ts

@@ -109,6 +109,8 @@ export interface ApiHandler {
 	 * @returns A promise resolving to the token count
 	 */
 	countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number>
+
+	contextWindow?: number // kilocode_change: Add contextWindow property for virtual quota fallback
 }
 
 export function buildApiHandler(configuration: ProviderSettings): ApiHandler {

src/api/providers/__tests__/virtual-quota-fallback-provider.spec.ts

@@ -524,10 +524,10 @@ describe("VirtualQuotaFallbackProvider", () => {
 				;(handler as any).activeHandler = undefined
 				const result = handler.getModel()
 				expect(result).toEqual({
-					id: "unknown",
+					id: "",
 					info: {
-						maxTokens: 100000,
-						contextWindow: 100000,
+						maxTokens: 1,
+						contextWindow: 1,
 						supportsPromptCache: false,
 					},
 				})

src/api/providers/virtual-quota-fallback.ts

@@ -2,6 +2,7 @@
 import { Anthropic } from "@anthropic-ai/sdk"
 import { z } from "zod"
 import * as vscode from "vscode"
+import EventEmitter from "events"
 import type { ModelInfo, ProviderSettings } from "@roo-code/types"
 import { ProviderSettingsManager } from "../../core/config/ProviderSettingsManager"
 import { ContextProxy } from "../../core/config/ContextProxy"
@@ -27,7 +28,7 @@ interface HandlerConfig {
  * Virtual Quota Fallback Provider API processor.
  * This handler is designed to call other API handlers with automatic fallback when quota limits are reached.
  */
-export class VirtualQuotaFallbackHandler implements ApiHandler {
+export class VirtualQuotaFallbackHandler extends EventEmitter implements ApiHandler {
 	private settingsManager: ProviderSettingsManager
 	private settings: ProviderSettings
 
@@ -38,6 +39,7 @@ export class VirtualQuotaFallbackHandler implements ApiHandler {
 	private isInitialized: boolean = false
 
 	constructor(options: ProviderSettings) {
+		super()
 		this.settings = options
 		this.settingsManager = new ProviderSettingsManager(ContextProxy.instance.rawContext)
 		this.usage = UsageTracker.getInstance()
@@ -121,19 +123,29 @@ export class VirtualQuotaFallbackHandler implements ApiHandler {
 	}
 
 	getModel(): { id: string; info: ModelInfo } {
+		// This is a synchronous method, so we can't await adjustActiveHandler here.
+		// The handler should be adjusted before this method is called.
 		if (!this.activeHandler) {
 			return {
-				id: "unknown",
+				id: "",
 				info: {
-					maxTokens: 100000,
-					contextWindow: 100000,
+					maxTokens: 1,
+					contextWindow: 1,
 					supportsPromptCache: false,
 				},
 			}
 		}
 		return this.activeHandler.getModel()
 	}
 
+	get contextWindow(): number {
+		if (!this.activeHandler) {
+			return 1 // Default fallback
+		}
+		const model = this.activeHandler.getModel()
+		return model.info.contextWindow
+	}
+
 	private async loadConfiguredProfiles(): Promise<void> {
 		this.handlerConfigs = []
 
@@ -234,6 +246,7 @@ export class VirtualQuotaFallbackHandler implements ApiHandler {
 			}
 			this.activeHandler = handler
 			this.activeProfileId = profileId
+			this.emit("handlerChanged", this.activeHandler)
 			return
 		}
 
@@ -243,6 +256,7 @@ export class VirtualQuotaFallbackHandler implements ApiHandler {
 		}
 		this.activeHandler = undefined
 		this.activeProfileId = undefined
+		this.emit("handlerChanged", this.activeHandler)
 	}
 
 	private async notifyHandlerSwitch(newProfileId: string | undefined, reason?: string): Promise<void> {

src/core/task/Task.ts

@@ -44,6 +44,7 @@ import { CloudService, BridgeOrchestrator } from "@roo-code/cloud"
 import { ApiHandler, ApiHandlerCreateMessageMetadata, buildApiHandler } from "../../api"
 import { ApiStream, GroundingSource } from "../../api/transform/stream"
 import { maybeRemoveImageBlocks } from "../../api/transform/image-cleaning"
+import { VirtualQuotaFallbackHandler } from "../../api/providers/virtual-quota-fallback" // kilocode_change: Import VirtualQuotaFallbackHandler for model change notifications
 
 // shared
 import { findLastIndex } from "../../shared/array"
@@ -373,6 +374,13 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 
 		this.apiConfiguration = apiConfiguration
 		this.api = buildApiHandler(apiConfiguration)
+		// kilocode_change start: Listen for model changes in virtual quota fallback
+		if (this.api instanceof VirtualQuotaFallbackHandler) {
+			this.api.on("handlerChanged", () => {
+				this.emit("modelChanged")
+			})
+		}
+		// kilocode_change end
 		this.autoApprovalHandler = new AutoApprovalHandler()
 
 		this.urlContentFetcher = new UrlContentFetcher(provider.context)
@@ -411,6 +419,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 		this.messageQueueStateChangedHandler = () => {
 			this.emit(RooCodeEventName.TaskUserMessage, this.taskId)
 			this.providerRef.deref()?.postStateToWebview()
+			this.emit("modelChanged") // kilocode_change: Emit modelChanged for virtual quota fallback UI updates
 		}
 
 		this.messageQueueService.on("stateChanged", this.messageQueueStateChangedHandler)
@@ -2740,6 +2749,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 		const { profileThresholds = {} } = state ?? {}
 
 		const { contextTokens } = this.getTokenUsage()
+		// kilocode_change start: Initialize virtual quota fallback handler
+		if (this.api instanceof VirtualQuotaFallbackHandler) {
+			await this.api.initialize()
+		}
+		// kilocode_change end
 		const modelInfo = this.api.getModel().info
 
 		const maxTokens = getModelMaxOutputTokens({
@@ -2748,7 +2762,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 			settings: this.apiConfiguration,
 		})
 
-		const contextWindow = modelInfo.contextWindow
+		const contextWindow = this.api.contextWindow ?? modelInfo.contextWindow // kilocode_change: Use contextWindow from API handler if available
 
 		// Get the current profile ID using the helper method
 		const currentProfileId = this.getCurrentProfileId(state)
@@ -2872,6 +2886,12 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 		const { contextTokens } = this.getTokenUsage()
 
 		if (contextTokens) {
+			// kilocode_change start: Initialize and adjust virtual quota fallback handler
+			if (this.api instanceof VirtualQuotaFallbackHandler) {
+				await this.api.initialize()
+				await this.api.adjustActiveHandler("Pre-Request Adjustment")
+			}
+			// kilocode_change end
 			const modelInfo = this.api.getModel().info
 
 			const maxTokens = getModelMaxOutputTokens({
@@ -2880,7 +2900,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 				settings: this.apiConfiguration,
 			})
 
-			const contextWindow = modelInfo.contextWindow
+			const contextWindow = this.api.contextWindow ?? modelInfo.contextWindow // kilocode_change
 
 			// Get the current profile ID using the helper method
 			const currentProfileId = this.getCurrentProfileId(state)

src/core/webview/ClineProvider.ts

@@ -257,6 +257,7 @@ export class ClineProvider
 			const onTaskUserMessage = (taskId: string) => this.emit(RooCodeEventName.TaskUserMessage, taskId)
 			const onTaskTokenUsageUpdated = (taskId: string, tokenUsage: TokenUsage) =>
 				this.emit(RooCodeEventName.TaskTokenUsageUpdated, taskId, tokenUsage)
+			const onModelChanged = () => this.postStateToWebview() // kilocode_change: Listen for model changes in virtual quota fallback
 
 			// Attach the listeners.
 			instance.on(RooCodeEventName.TaskStarted, onTaskStarted)
@@ -273,6 +274,7 @@ export class ClineProvider
 			instance.on(RooCodeEventName.TaskSpawned, onTaskSpawned)
 			instance.on(RooCodeEventName.TaskUserMessage, onTaskUserMessage)
 			instance.on(RooCodeEventName.TaskTokenUsageUpdated, onTaskTokenUsageUpdated)
+			instance.on("modelChanged", onModelChanged) // kilocode_change: Listen for model changes in virtual quota fallback
 
 			// Store the cleanup functions for later removal.
 			this.taskEventListeners.set(instance, [
@@ -290,6 +292,7 @@ export class ClineProvider
 				() => instance.off(RooCodeEventName.TaskUnpaused, onTaskUnpaused),
 				() => instance.off(RooCodeEventName.TaskSpawned, onTaskSpawned),
 				() => instance.off(RooCodeEventName.TaskTokenUsageUpdated, onTaskTokenUsageUpdated),
+				() => instance.off("modelChanged", onModelChanged), // kilocode_change: Clean up model change listener
 			])
 		}
 
@@ -2014,6 +2017,13 @@ ${prompt}
 			yoloMode, // kilocode_change
 		} = await this.getState()
 
+		// kilocode_change start: Get active model for virtual quota fallback UI display
+		const virtualQuotaActiveModel =
+			apiConfiguration?.apiProvider === "virtual-quota-fallback" && this.getCurrentTask()
+				? this.getCurrentTask()!.api.getModel()
+				: undefined
+		// kilocode_change end
+
 		let cloudOrganizations: CloudOrganizationMembership[] = []
 
 		try {
@@ -2206,6 +2216,7 @@ ${prompt}
 			openRouterImageGenerationSelectedModel,
 			openRouterUseMiddleOutTransform,
 			featureRoomoteControlEnabled,
+			virtualQuotaActiveModel, // kilocode_change: Include virtual quota active model in state
 		}
 	}
 

src/shared/ExtensionMessage.ts

@@ -2,6 +2,7 @@ import type {
 	GlobalSettings,
 	ProviderSettingsEntry,
 	ProviderSettings,
+	ModelInfo, // kilocode_change
 	HistoryItem,
 	ModeConfig,
 	TelemetrySetting,
@@ -454,6 +455,7 @@ export type ExtensionState = Pick<
 	remoteControlEnabled: boolean
 	taskSyncEnabled: boolean
 	featureRoomoteControlEnabled: boolean
+	virtualQuotaActiveModel?: { id: string; info: ModelInfo } // kilocode_change: Add virtual quota active model for UI display
 	showTimestamps?: boolean // kilocode_change: Show timestamps in chat messages
 }
 

webview-ui/src/components/kilocode/BottomApiConfig.tsx

@@ -3,7 +3,7 @@ import { useExtensionState } from "@/context/ExtensionStateContext"
 import { useSelectedModel } from "../ui/hooks/useSelectedModel"
 
 export const BottomApiConfig = () => {
-	const { currentApiConfigName, apiConfiguration } = useExtensionState()
+	const { currentApiConfigName, apiConfiguration, virtualQuotaActiveModel } = useExtensionState() // kilocode_change: Get virtual quota active model for UI display
 	const { id: selectedModelId, provider: selectedProvider } = useSelectedModel(apiConfiguration)
 
 	if (!apiConfiguration) {
@@ -18,6 +18,12 @@ export const BottomApiConfig = () => {
 					currentApiConfigName={currentApiConfigName}
 					apiConfiguration={apiConfiguration}
 					fallbackText={`${selectedProvider}:${selectedModelId}`}
+					//kilocode_change: Pass virtual quota active model to ModelSelector
+					virtualQuotaActiveModel={
+						virtualQuotaActiveModel
+							? { id: virtualQuotaActiveModel.id, name: virtualQuotaActiveModel.id }
+							: undefined
+					}
 				/>
 			</div>
 		</>

webview-ui/src/components/kilocode/chat/ModelSelector.tsx

@@ -13,9 +13,15 @@ interface ModelSelectorProps {
 	currentApiConfigName?: string
 	apiConfiguration: ProviderSettings
 	fallbackText: string
+	virtualQuotaActiveModel?: { id: string; name: string } // kilocode_change: Add virtual quota active model for UI display
 }
 
-export const ModelSelector = ({ currentApiConfigName, apiConfiguration, fallbackText }: ModelSelectorProps) => {
+export const ModelSelector = ({
+	currentApiConfigName,
+	apiConfiguration,
+	fallbackText,
+	virtualQuotaActiveModel, //kilocode_change
+}: ModelSelectorProps) => {
 	const { t } = useAppTranslation()
 	const { provider, providerModels, providerDefaultModel, isLoading, isError } = useProviderModels(apiConfiguration)
 	const selectedModelId = getSelectedModelId({
@@ -60,6 +66,16 @@ export const ModelSelector = ({ currentApiConfigName, apiConfiguration, fallback
 		return null
 	}
 
+	// kilocode_change start: Display active model for virtual quota fallback
+	if (provider === "virtual-quota-fallback" && virtualQuotaActiveModel) {
+		return (
+			<span className="text-xs text-vscode-descriptionForeground opacity-70 truncate">
+				{prettyModelName(virtualQuotaActiveModel.id)}
+			</span>
+		)
+	}
+	// kilocode_change end
+
 	if (isError || options.length <= 0) {
 		return <span className="text-xs text-vscode-descriptionForeground opacity-70 truncate">{fallbackText}</span>
 	}