fix: proper context window loading for LMStudio (fixes #5075)

This commit is contained in:
Piotr Wilkin 2025-07-17 15:07:41 +02:00
parent 6cf376f832
commit 61d43876f7
7 changed files with 108 additions and 83 deletions

View file

@ -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<ModelRecord> {
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 {}
}
}

View file

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

View file

@ -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<string>(modelsArray)]
} catch (error) {
return []
}
}

View file

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

View file

@ -65,7 +65,6 @@ export interface WebviewMessage {
| "requestRouterModels"
| "requestOpenAiModels"
| "requestOllamaModels"
| "requestLmStudioModels"
| "requestVsCodeLmModels"
| "openImage"
| "saveImage"

View file

@ -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" })
}
},

View file

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