From 61d43876f7cda2ce98e268e1befed79bd048e7cb Mon Sep 17 00:00:00 2001 From: Piotr Wilkin Date: Thu, 17 Jul 2025 15:07:41 +0200 Subject: [PATCH] fix: proper context window loading for LMStudio (fixes #5075) --- src/api/providers/fetchers/lm-studio.ts | 35 +++++++ src/api/providers/fetchers/modelCache.ts | 1 + src/api/providers/lm-studio.ts | 50 ++++++---- src/core/webview/webviewMessageHandler.ts | 92 ++++++++----------- src/shared/WebviewMessage.ts | 1 - .../src/components/settings/ApiOptions.tsx | 6 +- .../settings/providers/LMStudio.tsx | 6 +- 7 files changed, 108 insertions(+), 83 deletions(-) create mode 100644 src/api/providers/fetchers/lm-studio.ts diff --git a/src/api/providers/fetchers/lm-studio.ts b/src/api/providers/fetchers/lm-studio.ts new file mode 100644 index 0000000000..4b5f43bf7c --- /dev/null +++ b/src/api/providers/fetchers/lm-studio.ts @@ -0,0 +1,35 @@ +import axios from "axios" +import { ModelRecord } from "../../../shared/api" +import { openAiModelInfoSaneDefaults } from "@roo-code/types" + +export async function getLmStudioModels(baseUrl = "http://localhost:1234"): Promise { + try { + if (!URL.canParse(baseUrl)) { + return {} + } + + const response = await axios.get(`${baseUrl}/api/v0/models`) + return response.data?.data?.reduce((acc: ModelRecord, model: any) => { + acc[model.id] = { + maxTokens: + model.loaded_context_length || + model.max_context_length || + openAiModelInfoSaneDefaults.contextWindow, + contextWindow: + model.loaded_context_length || + model.max_context_length || + openAiModelInfoSaneDefaults.contextWindow, + supportsImages: false, + supportsPromptCache: false, + supportsComputerUse: false, + inputPrice: 0, + outputPrice: 0, + cacheWritesPrice: 0, + cacheReadsPrice: 0, + } + return acc + }, {}) + } catch (error) { + return {} + } +} diff --git a/src/api/providers/fetchers/modelCache.ts b/src/api/providers/fetchers/modelCache.ts index fef700268d..df63c83a36 100644 --- a/src/api/providers/fetchers/modelCache.ts +++ b/src/api/providers/fetchers/modelCache.ts @@ -14,6 +14,7 @@ import { getRequestyModels } from "./requesty" import { getGlamaModels } from "./glama" import { getUnboundModels } from "./unbound" import { getLiteLLMModels } from "./litellm" +import { getLmStudioModels } from "./lm-studio" import { GetModelsOptions } from "../../../shared/api" import { getOllamaModels } from "./ollama" import { getLMStudioModels } from "./lmstudio" diff --git a/src/api/providers/lm-studio.ts b/src/api/providers/lm-studio.ts index f032e2d560..4eb014caf0 100644 --- a/src/api/providers/lm-studio.ts +++ b/src/api/providers/lm-studio.ts @@ -1,8 +1,7 @@ import { Anthropic } from "@anthropic-ai/sdk" import OpenAI from "openai" -import axios from "axios" -import { type ModelInfo, openAiModelInfoSaneDefaults, LMSTUDIO_DEFAULT_TEMPERATURE } from "@roo-code/types" +import { LMSTUDIO_DEFAULT_TEMPERATURE, type ModelInfo, openAiModelInfoSaneDefaults } from "@roo-code/types" import type { ApiHandlerOptions } from "../../shared/api" @@ -11,22 +10,32 @@ import { XmlMatcher } from "../../utils/xml-matcher" import { convertToOpenAiMessages } from "../transform/openai-format" import { ApiStream } from "../transform/stream" +import type { ApiHandlerCreateMessageMetadata, SingleCompletionHandler } from "../index" import { BaseProvider } from "./base-provider" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import { flushModels, getModels } from "./fetchers/modelCache" export class LmStudioHandler extends BaseProvider implements SingleCompletionHandler { protected options: ApiHandlerOptions private client: OpenAI + private cachedModelInfo: ModelInfo = openAiModelInfoSaneDefaults constructor(options: ApiHandlerOptions) { super() this.options = options this.client = new OpenAI({ - baseURL: (this.options.lmStudioBaseUrl || "http://localhost:1234") + "/v1", + baseURL: this.getBaseUrl() + "/v1", apiKey: "noop", }) } + private getBaseUrl(): string { + if (this.options.lmStudioBaseUrl && this.options.lmStudioBaseUrl.trim() !== "") { + return this.options.lmStudioBaseUrl.trim() + } else { + return "http://localhost:1234" + } + } + override async *createMessage( systemPrompt: string, messages: Anthropic.Messages.MessageParam[], @@ -118,6 +127,23 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan outputTokens = 0 } + if (this.cachedModelInfo === openAiModelInfoSaneDefaults) { + // We need to fetch the model info every time we open a new session + // to ensure we have the latest context window and other details + // since LM Studio models can chance their context windows on reload + await flushModels("lmstudio") + const models = await getModels({ provider: "lmstudio", baseUrl: this.getBaseUrl() }) + if (models && models[this.getModel().id]) { + this.cachedModelInfo = models[this.getModel().id] + } else { + // If model info is not found, use sane defaults + this.cachedModelInfo = { + ...openAiModelInfoSaneDefaults, + description: "Fake description to avoid recache", + } + } + } + yield { type: "usage", inputTokens, @@ -133,7 +159,7 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan override getModel(): { id: string; info: ModelInfo } { return { id: this.options.lmStudioModelId || "", - info: openAiModelInfoSaneDefaults, + info: this.cachedModelInfo, } } @@ -161,17 +187,3 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan } } } - -export async function getLmStudioModels(baseUrl = "http://localhost:1234") { - try { - if (!URL.canParse(baseUrl)) { - return [] - } - - const response = await axios.get(`${baseUrl}/v1/models`) - const modelsArray = response.data?.data?.map((model: any) => model.id) || [] - return [...new Set(modelsArray)] - } catch (error) { - return [] - } -} diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index e70b39df8f..c623542ee1 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -1,58 +1,56 @@ -import { safeWriteJson } from "../../utils/safeWriteJson" -import * as path from "path" -import * as os from "os" import * as fs from "fs/promises" +import * as os from "os" import pWaitFor from "p-wait-for" +import * as path from "path" import * as vscode from "vscode" -import * as yaml from "yaml" +import { safeWriteJson } from "../../utils/safeWriteJson" -import { - type Language, - type ProviderSettings, - type GlobalState, - type ClineMessage, - TelemetryEventName, -} from "@roo-code/types" import { CloudService } from "@roo-code/cloud" import { TelemetryService } from "@roo-code/telemetry" +import { + type ClineMessage, + type GlobalState, + type Language, + type ProviderSettings, + TelemetryEventName, +} from "@roo-code/types" import { type ApiMessage } from "../task-persistence/apiMessages" -import { ClineProvider } from "./ClineProvider" import { changeLanguage, t } from "../../i18n" +import { ModelRecord, RouterName, toRouterName } from "../../shared/api" import { Package } from "../../shared/package" -import { RouterName, toRouterName, ModelRecord } from "../../shared/api" import { supportPrompt } from "../../shared/support-prompt" +import { ClineProvider } from "./ClineProvider" -import { checkoutDiffPayloadSchema, checkoutRestorePayloadSchema, WebviewMessage } from "../../shared/WebviewMessage" -import { checkExistKey } from "../../shared/checkExistApiConfig" -import { experimentDefault } from "../../shared/experiments" -import { Terminal } from "../../integrations/terminal/Terminal" -import { openFile } from "../../integrations/misc/open-file" +import { flushModels, getModels } from "../../api/providers/fetchers/modelCache" +import { getOpenAiModels } from "../../api/providers/openai" +import { getVsCodeLmModels } from "../../api/providers/vscode-lm" import { openImage, saveImage } from "../../integrations/misc/image-handler" +import { openFile } from "../../integrations/misc/open-file" import { selectImages } from "../../integrations/misc/process-images" +import { Terminal } from "../../integrations/terminal/Terminal" import { getTheme } from "../../integrations/theme/getTheme" import { discoverChromeHostUrl, tryChromeHostUrl } from "../../services/browser/browserDiscovery" import { searchWorkspaceFiles } from "../../services/search/file-search" -import { fileExistsAtPath } from "../../utils/fs" -import { playTts, setTtsEnabled, setTtsSpeed, stopTts } from "../../utils/tts" -import { singleCompletionHandler } from "../../utils/single-completion-handler" -import { searchCommits } from "../../utils/git" -import { exportSettings, importSettingsWithFeedback } from "../config/importExport" -import { getOpenAiModels } from "../../api/providers/openai" -import { getVsCodeLmModels } from "../../api/providers/vscode-lm" -import { openMention } from "../mentions" import { TelemetrySetting } from "../../shared/TelemetrySetting" -import { getWorkspacePath } from "../../utils/path" -import { ensureSettingsDirectoryExists } from "../../utils/globalContext" -import { Mode, defaultModeSlug } from "../../shared/modes" -import { getModels, flushModels } from "../../api/providers/fetchers/modelCache" +import { checkoutDiffPayloadSchema, checkoutRestorePayloadSchema, WebviewMessage } from "../../shared/WebviewMessage" import { GetModelsOptions } from "../../shared/api" -import { generateSystemPrompt } from "./generateSystemPrompt" +import { checkExistKey } from "../../shared/checkExistApiConfig" +import { experimentDefault } from "../../shared/experiments" +import { defaultModeSlug, Mode } from "../../shared/modes" import { getCommand } from "../../utils/commands" +import { fileExistsAtPath } from "../../utils/fs" +import { searchCommits } from "../../utils/git" +import { getWorkspacePath } from "../../utils/path" +import { singleCompletionHandler } from "../../utils/single-completion-handler" +import { playTts, setTtsEnabled, setTtsSpeed, stopTts } from "../../utils/tts" +import { exportSettings, importSettingsWithFeedback } from "../config/importExport" +import { openMention } from "../mentions" +import { generateSystemPrompt } from "./generateSystemPrompt" const ALLOWED_VSCODE_SETTINGS = new Set(["terminal.integrated.inheritEnv"]) -import { MarketplaceManager, MarketplaceItemType } from "../../services/marketplace" +import { MarketplaceItemType, MarketplaceManager } from "../../services/marketplace" import { setPendingTodoList } from "../tools/updateTodoListTool" export const webviewMessageHandler = async ( @@ -592,6 +590,12 @@ export const webviewMessageHandler = async ( }) } + const lmStudioBaseUrl = apiConfiguration.lmStudioBaseUrl || message?.values?.lmStudioBaseUrl + modelFetchPromises.push({ + key: "lmstudio", + options: { provider: "lmstudio", baseUrl: lmStudioBaseUrl }, + }) + const results = await Promise.allSettled( modelFetchPromises.map(async ({ key, options }) => { const models = await safeGetModels(options) @@ -670,30 +674,6 @@ export const webviewMessageHandler = async ( } break } - case "requestLmStudioModels": { - // Specific handler for LM Studio models only - const { apiConfiguration: lmStudioApiConfig } = await provider.getState() - try { - // Flush cache first to ensure fresh models - await flushModels("lmstudio") - - const lmStudioModels = await getModels({ - provider: "lmstudio", - baseUrl: lmStudioApiConfig.lmStudioBaseUrl, - }) - - if (Object.keys(lmStudioModels).length > 0) { - provider.postMessageToWebview({ - type: "lmStudioModels", - lmStudioModels: Object.keys(lmStudioModels), - }) - } - } catch (error) { - // Silently fail - user hasn't configured LM Studio yet - console.debug("LM Studio models fetch failed:", error) - } - break - } case "requestOpenAiModels": if (message?.values?.baseUrl && message?.values?.apiKey) { const openAiModels = await getOpenAiModels( diff --git a/src/shared/WebviewMessage.ts b/src/shared/WebviewMessage.ts index d5dc3f8c28..1411e3ed31 100644 --- a/src/shared/WebviewMessage.ts +++ b/src/shared/WebviewMessage.ts @@ -65,7 +65,6 @@ export interface WebviewMessage { | "requestRouterModels" | "requestOpenAiModels" | "requestOllamaModels" - | "requestLmStudioModels" | "requestVsCodeLmModels" | "openImage" | "saveImage" diff --git a/webview-ui/src/components/settings/ApiOptions.tsx b/webview-ui/src/components/settings/ApiOptions.tsx index 06994b16b9..245346aedf 100644 --- a/webview-ui/src/components/settings/ApiOptions.tsx +++ b/webview-ui/src/components/settings/ApiOptions.tsx @@ -190,12 +190,10 @@ const ApiOptions = ({ }, }) } else if (selectedProvider === "ollama") { - vscode.postMessage({ type: "requestOllamaModels" }) - } else if (selectedProvider === "lmstudio") { - vscode.postMessage({ type: "requestLmStudioModels" }) + vscode.postMessage({ type: "requestOllamaModels", text: apiConfiguration?.ollamaBaseUrl }) } else if (selectedProvider === "vscode-lm") { vscode.postMessage({ type: "requestVsCodeLmModels" }) - } else if (selectedProvider === "litellm") { + } else if (selectedProvider === "litellm" || selectedProvider === "lmstudio") { vscode.postMessage({ type: "requestRouterModels" }) } }, diff --git a/webview-ui/src/components/settings/providers/LMStudio.tsx b/webview-ui/src/components/settings/providers/LMStudio.tsx index a907e43e1b..74e1f1b193 100644 --- a/webview-ui/src/components/settings/providers/LMStudio.tsx +++ b/webview-ui/src/components/settings/providers/LMStudio.tsx @@ -39,9 +39,9 @@ export const LMStudio = ({ apiConfiguration, setApiConfigurationField }: LMStudi const message: ExtensionMessage = event.data switch (message.type) { - case "lmStudioModels": + case "routerModels": { - const newModels = message.lmStudioModels ?? [] + const newModels = Object.keys(message.routerModels?.lmstudio || {}) setLmStudioModels(newModels) } break @@ -53,7 +53,7 @@ export const LMStudio = ({ apiConfiguration, setApiConfigurationField }: LMStudi // Refresh models on mount useEffect(() => { // Request fresh models - the handler now flushes cache automatically - vscode.postMessage({ type: "requestLmStudioModels" }) + vscode.postMessage({ type: "requestRouterModels" }) }, []) // Check if the selected model exists in the fetched models