From 02ebc5e8e1d90789f80bb4554c7a915b47f5b03b Mon Sep 17 00:00:00 2001 From: Sage Grigull Date: Tue, 3 Mar 2026 12:26:57 +1000 Subject: [PATCH] add LM Studio support (#259) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: Add LM Studio as a supported model provider with OpenAI-compatible API fetching and discovery. Add LM Studio as a supported AI provider with optional API key, environment variables, and discovery. * feat: Add rustup as a dev dependency. * fix: Refine LM Studio API key handling to conditionally send authorization headers during model discovery based on whether the key is a default local token or a custom key, and add new tests. * feat: Enhance OAuth token and account ID resolution for model providers in the model registry. * rebase for packages/ai/CHANGELOG.md * feat: improve implicit model discovery to independently auto-detect Ollama and LM Studio, and refine LM Studio base URL handling. --------- Co-authored-by: Can Bölük --- README.md | 3 + bun.lock | 3 + docs/environment-variables.md | 2 + docs/models.md | 13 +- package.json | 5 +- packages/ai/CHANGELOG.md | 2 + packages/ai/src/auth-storage.ts | 6 + .../ai/src/provider-models/descriptors.ts | 8 ++ .../ai/src/provider-models/openai-compat.ts | 31 +++++ packages/ai/src/stream.ts | 1 + packages/ai/src/types.ts | 3 +- packages/ai/src/utils/oauth/index.ts | 8 ++ packages/ai/src/utils/oauth/lm-studio.ts | 40 ++++++ packages/ai/src/utils/oauth/types.ts | 1 + packages/coding-agent/CHANGELOG.md | 5 + .../coding-agent/src/config/model-registry.ts | 96 +++++++++++-- .../coding-agent/test/lm-studio-fix.test.ts | 129 ++++++++++++++++++ 17 files changed, 343 insertions(+), 13 deletions(-) create mode 100644 packages/ai/src/utils/oauth/lm-studio.ts create mode 100644 packages/coding-agent/test/lm-studio-fix.test.ts diff --git a/README.md b/README.md index 3885b9f65..52e9fc3df 100644 --- a/README.md +++ b/README.md @@ -490,6 +490,7 @@ return config | Together (`together`) | `TOGETHER_API_KEY` | | Ollama (`ollama`) | `OLLAMA_API_KEY` _(optional)_ | | LiteLLM (`litellm`) | `LITELLM_API_KEY` | +| LM Studio (`lm-studio`) | `LM_STUDIO_API_KEY` _(optional)_ | | Xiaomi MiMo (`xiaomi`) | `XIAOMI_API_KEY` | | Moonshot (`moonshot`) | `MOONSHOT_API_KEY` | | Venice (`venice`) | `VENICE_API_KEY` | @@ -527,6 +528,7 @@ Use `/login` with supported providers: - GitLab Duo (`gitlab-duo`) - Qianfan (`qianfan`) - Ollama (local / self-hosted, `ollama`) +- LM Studio (local / self-hosted, `lm-studio`) - vLLM (local OpenAI-compatible, `vllm`) - Z.AI (GLM Coding Plan) - Synthetic @@ -540,6 +542,7 @@ Use `/login` with supported providers: - Cloudflare AI Gateway (`cloudflare-ai-gateway`) For `ollama`, API key is optional. Leave it unset for local no-auth instances, or set `OLLAMA_API_KEY` for authenticated hosts. +For `lm-studio`, API key is optional. Leave it unset for local no-auth instances, or set `LM_STUDIO_API_KEY` for authenticated hosts. For `vllm`, paste your key in `/login` (or use `VLLM_API_KEY`). For local no-auth servers, any placeholder value works (for example `vllm-local`). For `nanogpt`, `/login nanogpt` opens `https://nano-gpt.com/api` and prompts for your `sk-...` key (or set `NANO_GPT_API_KEY`). Login validates the key via NanoGPT's models endpoint (not a fixed model entitlement). For `cloudflare-ai-gateway`, set provider base URL to diff --git a/bun.lock b/bun.lock index befbc1cb2..deede252f 100644 --- a/bun.lock +++ b/bun.lock @@ -11,6 +11,7 @@ "@typescript/native-preview": "^7.0.0-dev.20260222.1", "lint-staged": "^16.2", "prettier": "^3.8", + "rustup": "1.0.10", }, }, "packages/agent": { @@ -887,6 +888,8 @@ "rimraf": ["rimraf@5.0.10", "", { "dependencies": { "glob": "^10.3.7" }, "bin": { "rimraf": "dist/esm/bin.mjs" } }, "sha512-l0OE8wL34P4nJH/H2ffoaniAokM2qSmrtXHmlpvYr5AVVX8msAyW0l8NVJFDxlSK4u3Uh/f41cQheDVdnYijwQ=="], + "rustup": ["rustup@1.0.10", "", { "bin": { "rustup-init": "rustup-init", "rustfmt": "rustup", "rustup": "rustup", "cargo": "rustup", "cargo-clippy": "rustup", "cargo-fmt": "rustup", "cargo-miri": "rustup", "clippy-driver": "rustup", "rls": "rustup", "rust-gdb": "rustup", "rust-gdbgui": "rustup", "rust-lldb": "rustup", "rustc": "rustup", "rustdoc": "rustup" } }, "sha512-PnUuHNrwRLQwld4Ha2uiO1wAi+8EYHNebXQlCtgI9pUpP+FLtnL79jAYpa3PvtWBDVZi54CK1I9pnUiR7yKYsQ=="], + "safe-buffer": ["safe-buffer@5.2.1", "", {}, "sha512-rp3So07KcdmmKbGvgaNxQSJr7bGVSVk5S9Eq1F+ppbRo70+YeaDxkw5Dd8NPN+GD6bjnYm2VuPuCXmpuYvmCXQ=="], "safe-stable-stringify": ["safe-stable-stringify@2.5.0", "", {}, "sha512-b3rppTKm9T+PsVCBEOUR46GWI7fdOs00VKZ1+9c1EWDaDMvjQc6tUwuFyIprgGgTcWoVHSKrU8H31ZHA2e0RHA=="], diff --git a/docs/environment-variables.md b/docs/environment-variables.md index a22cdaf3f..c5b95e66b 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -46,6 +46,7 @@ These are consumed via `getEnvApiKey()` (`packages/ai/src/stream.ts`) unless not | `NANO_GPT_API_KEY` | NanoGPT auth | Using `nanogpt` provider | | | `VENICE_API_KEY` | Venice auth | Using `venice` provider | | | `LITELLM_API_KEY` | LiteLLM auth | Using `litellm` provider | OpenAI-compatible LiteLLM proxy key | +| `LM_STUDIO_API_KEY` | LM Studio auth (optional) | Using `lm-studio` provider with authenticated hosts | Local LM Studio usually runs without auth; any non-empty token works when a key is required | | `OLLAMA_API_KEY` | Ollama auth (optional) | Using `ollama` provider with authenticated hosts | Local Ollama usually runs without auth; any non-empty token works when a key is required | | `XIAOMI_API_KEY` | Xiaomi MiMo auth | Using `xiaomi` provider | | | `MOONSHOT_API_KEY` | Moonshot auth | Using `moonshot` provider | | @@ -253,6 +254,7 @@ Extra conditional behavior: | `PI_DEBUG_STARTUP` | Enables startup stage debug prints to stderr in multiple startup paths | | `PI_PACKAGE_DIR` | Overrides package asset base dir resolution (docs/examples/changelog path lookup) | | `PI_DISABLE_LSPMUX` | If `1`, disables lspmux detection/integration and forces direct LSP server spawning | +| `LM_STUDIO_BASE_URL` | Default implicit LM Studio discovery base URL override (`http://127.0.0.1:1234/v1` if unset) | | `OLLAMA_BASE_URL` | Default implicit Ollama discovery base URL override (`http://127.0.0.1:11434` if unset) | | `PI_EDIT_VARIANT` | If `hashline`, forces hashline read/grep display mode when edit tool available | | `PI_NO_PTY` | If `1`, disables interactive PTY path for bash tool | diff --git a/docs/models.md b/docs/models.md index f6622ed05..7805ff242 100644 --- a/docs/models.md +++ b/docs/models.md @@ -129,7 +129,7 @@ ModelRegistry pipeline (on refresh): 5. Merge custom `models`: - same `provider + id` replaces existing - otherwise append -6. Apply runtime-discovered models (currently Ollama), then re-apply model overrides. +6. Apply runtime-discovered models (currently Ollama and LM Studio), then re-apply model overrides. Provider defaults vs per-model overrides: @@ -151,6 +151,17 @@ If `ollama` is not explicitly configured, registry adds an implicit discoverable Runtime discovery calls `GET /api/tags` on Ollama and synthesizes model entries with local defaults. +### Implicit LM Studio discovery + +If `lm-studio` is not explicitly configured, registry adds an implicit discoverable provider: + +- provider: `lm-studio` +- api: `openai-completions` +- base URL: `LM_STUDIO_BASE_URL` or `http://127.0.0.1:1234/v1` +- auth mode: keyless (`auth: none` behavior) + +Runtime discovery fetches models (`GET /models`) and synthesizes model entries with local defaults. + ### Explicit provider discovery You can configure discovery yourself: diff --git a/package.json b/package.json index db9a1ad43..f47b6a9fb 100644 --- a/package.json +++ b/package.json @@ -46,9 +46,10 @@ "@types/bun": "^1.3", "@typescript/native-preview": "^7.0.0-dev.20260222.1", "lint-staged": "^16.2", - "prettier": "^3.8" + "prettier": "^3.8", + "rustup": "1.0.10" }, "lint-staged": { "*.{js,ts,jsx,tsx,json,jsonc,css}": "biome check --write --no-errors-on-unmatched" } -} +} \ No newline at end of file diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 1584687c9..7a1004cfd 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -4,6 +4,8 @@ ### Added - Added Anthropic Foundry gateway mode controlled by `CLAUDE_CODE_USE_FOUNDRY`, with support for `FOUNDRY_BASE_URL`, `ANTHROPIC_FOUNDRY_API_KEY`, `ANTHROPIC_CUSTOM_HEADERS`, and optional mTLS material (`CLAUDE_CODE_CLIENT_CERT`, `CLAUDE_CODE_CLIENT_KEY`, `NODE_EXTRA_CA_CERTS`) +- Added LM Studio provider support with OpenAI-compatible model discovery and OAuth login. +- Added support for `LM_STUDIO_API_KEY` and `LM_STUDIO_BASE_URL` environment variables for authentication and custom host configuration. ### Changed diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index f255715c9..055cc0323 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -45,6 +45,7 @@ import { loginHuggingface } from "./utils/oauth/huggingface"; import { loginKilo } from "./utils/oauth/kilo"; import { loginKimi } from "./utils/oauth/kimi"; import { loginLiteLLM } from "./utils/oauth/litellm"; +import { loginLmStudio } from "./utils/oauth/lm-studio"; import { loginMiniMaxCode, loginMiniMaxCodeCn } from "./utils/oauth/minimax-code"; import { loginMoonshot } from "./utils/oauth/moonshot"; import { loginNanoGPT } from "./utils/oauth/nanogpt"; @@ -820,6 +821,11 @@ export class AuthStorage { await saveApiKeyCredential(apiKey); return; } + case "lm-studio": { + const apiKey = await loginLmStudio(ctrl); + await saveApiKeyCredential(apiKey); + return; + } case "ollama": { const apiKey = await loginOllama(ctrl); if (!apiKey) { diff --git a/packages/ai/src/provider-models/descriptors.ts b/packages/ai/src/provider-models/descriptors.ts index 674e72dd6..d87b32e61 100644 --- a/packages/ai/src/provider-models/descriptors.ts +++ b/packages/ai/src/provider-models/descriptors.ts @@ -17,6 +17,7 @@ import { kiloModelManagerOptions, kimiCodeModelManagerOptions, litellmModelManagerOptions, + lmStudioModelManagerOptions, mistralModelManagerOptions, moonshotModelManagerOptions, nanoGptModelManagerOptions, @@ -209,6 +210,13 @@ export const PROVIDER_DESCRIPTORS: readonly ProviderDescriptor[] = [ config => litellmModelManagerOptions(config), catalog("LiteLLM", ["LITELLM_API_KEY"], { allowUnauthenticated: true }), ), + catalogDescriptor( + "lm-studio", + "llama-3-8b", + config => lmStudioModelManagerOptions(config), + catalog("LM Studio", ["LM_STUDIO_API_KEY"], { allowUnauthenticated: true }), + { allowUnauthenticated: true }, + ), catalogDescriptor( "vllm", "gpt-oss-20b", diff --git a/packages/ai/src/provider-models/openai-compat.ts b/packages/ai/src/provider-models/openai-compat.ts index 52d707f4c..98db52b5a 100644 --- a/packages/ai/src/provider-models/openai-compat.ts +++ b/packages/ai/src/provider-models/openai-compat.ts @@ -796,6 +796,37 @@ export function kimiCodeModelManagerOptions( }; } +// --------------------------------------------------------------------------- +// 12.5. LM Studio +// --------------------------------------------------------------------------- + +export interface LmStudioModelManagerConfig { + apiKey?: string; + baseUrl?: string; +} + +export function lmStudioModelManagerOptions( + config?: LmStudioModelManagerConfig, +): ModelManagerOptions<"openai-completions"> { + const apiKey = config?.apiKey; + const baseUrl = config?.baseUrl ?? Bun.env.LM_STUDIO_BASE_URL ?? "http://127.0.0.1:1234/v1"; + const references = createBundledReferenceMap<"openai-completions">("lm-studio" as any); + return { + providerId: "lm-studio", + fetchDynamicModels: () => + fetchOpenAICompatibleModels({ + api: "openai-completions", + provider: "lm-studio", + baseUrl, + apiKey, + mapModel: (entry, defaults) => { + const reference = references.get(defaults.id); + return mapWithBundledReference(entry, defaults, reference); + }, + }), + }; +} + // --------------------------------------------------------------------------- // 13. Synthetic // --------------------------------------------------------------------------- diff --git a/packages/ai/src/stream.ts b/packages/ai/src/stream.ts index a11a4f095..7aba0bd9c 100644 --- a/packages/ai/src/stream.ts +++ b/packages/ai/src/stream.ts @@ -127,6 +127,7 @@ const serviceProviderMap: Record = { moonshot: "MOONSHOT_API_KEY", nvidia: "NVIDIA_API_KEY", nanogpt: "NANO_GPT_API_KEY", + "lm-studio": "LM_STUDIO_API_KEY", ollama: "OLLAMA_API_KEY", qianfan: "QIANFAN_API_KEY", "qwen-portal": () => $pickenv("QWEN_OAUTH_TOKEN", "QWEN_PORTAL_API_KEY"), diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 155abbdd4..d66e28ef9 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -104,7 +104,8 @@ export type KnownProvider = | "together" | "venice" | "vllm" - | "xiaomi"; + | "xiaomi" + | "lm-studio"; export type Provider = KnownProvider | string; export type ThinkingLevel = "minimal" | "low" | "medium" | "high" | "xhigh"; diff --git a/packages/ai/src/utils/oauth/index.ts b/packages/ai/src/utils/oauth/index.ts index 05c33ccee..fc82f44c0 100644 --- a/packages/ai/src/utils/oauth/index.ts +++ b/packages/ai/src/utils/oauth/index.ts @@ -72,6 +72,8 @@ export { loginKilo } from "./kilo"; export { loginKimi, refreshKimiToken } from "./kimi"; // LiteLLM (API key) export { loginLiteLLM } from "./litellm"; +// LM Studio (optional API key) +export { DEFAULT_LOCAL_TOKEN, loginLmStudio } from "./lm-studio"; // MiniMax Coding Plan (API key) export { loginMiniMaxCode, loginMiniMaxCodeCn } from "./minimax-code"; // Moonshot (API key) @@ -163,6 +165,11 @@ const builtInOAuthProviders: OAuthProviderInfo[] = [ name: "LiteLLM", available: true, }, + { + id: "lm-studio", + name: "LM Studio (Local OpenAI-compatible)", + available: true, + }, { id: "ollama", name: "Ollama (Local OpenAI-compatible)", @@ -338,6 +345,7 @@ export async function refreshOAuthToken( case "synthetic": case "together": case "litellm": + case "lm-studio": case "ollama": case "xiaomi": case "zai": diff --git a/packages/ai/src/utils/oauth/lm-studio.ts b/packages/ai/src/utils/oauth/lm-studio.ts new file mode 100644 index 000000000..5a315831e --- /dev/null +++ b/packages/ai/src/utils/oauth/lm-studio.ts @@ -0,0 +1,40 @@ +/** + * LM Studio login flow. + * + * LM Studio provides an OpenAI-compatible API at a local base URL. + * It usually runs unauthenticated but can be configured to require a bearer token. + * + * This flow stores an API-key-style credential used by `/login` and auth storage. + */ + +import type { OAuthController, OAuthProvider } from "./types"; + +const PROVIDER_ID: OAuthProvider = "lm-studio"; +const _AUTH_URL = "https://lmstudio.ai/docs/api"; +const _DEFAULT_LOCAL_BASE_URL = "http://127.0.0.1:1234/v1"; +export const DEFAULT_LOCAL_TOKEN = "lm-studio-local"; + +/** + * Login to LM Studio. + * + * Opens LM Studio API docs, prompts for an optional token, + * and returns a stored key value. + */ +export async function loginLmStudio(options: OAuthController): Promise { + if (!options.onPrompt) { + throw new Error(`${PROVIDER_ID} login requires onPrompt callback`); + } + + const apiKey = await options.onPrompt({ + message: "Optional: Paste LM Studio API key (to customize endpoint URL, set LM_STUDIO_BASE_URL env var)", + placeholder: DEFAULT_LOCAL_TOKEN, + allowEmpty: true, + }); + + if (options.signal?.aborted) { + throw new Error("Login cancelled"); + } + + const trimmed = apiKey.trim(); + return trimmed || DEFAULT_LOCAL_TOKEN; +} diff --git a/packages/ai/src/utils/oauth/types.ts b/packages/ai/src/utils/oauth/types.ts index 80f0dcb75..eff139194 100644 --- a/packages/ai/src/utils/oauth/types.ts +++ b/packages/ai/src/utils/oauth/types.ts @@ -21,6 +21,7 @@ export type OAuthProvider = | "kimi-code" | "kilo" | "litellm" + | "lm-studio" | "minimax-code" | "minimax-code-cn" | "moonshot" diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 24d0f5ccf..09d583c5f 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -6,6 +6,11 @@ - Updated Anthropic Foundry environment variable documentation and CLI help text to the canonical names: `CLAUDE_CODE_USE_FOUNDRY`, `CLAUDE_CODE_CLIENT_CERT`, and `CLAUDE_CODE_CLIENT_KEY` - Documented Foundry-specific Anthropic runtime configuration (`FOUNDRY_BASE_URL`, `ANTHROPIC_FOUNDRY_API_KEY`, `ANTHROPIC_CUSTOM_HEADERS`, `NODE_EXTRA_CA_CERTS`) in environment variable reference docs +### Added + +- Added LM Studio integration to the model registry and discovery flow. +- Added support for authenticating with LM Studio using the `/login lm-studio` command. + ## [13.5.6] - 2026-03-01 ### Changed diff --git a/packages/coding-agent/src/config/model-registry.ts b/packages/coding-agent/src/config/model-registry.ts index bbcb3677d..ac84e28c6 100644 --- a/packages/coding-agent/src/config/model-registry.ts +++ b/packages/coding-agent/src/config/model-registry.ts @@ -3,6 +3,7 @@ import { type AssistantMessageEventStream, type Context, createModelManager, + DEFAULT_LOCAL_TOKEN, getBundledModels, getBundledProviders, getGitHubCopilotBaseUrl, @@ -131,7 +132,7 @@ const ModelOverrideSchema = Type.Object({ type ModelOverride = Static; const ProviderDiscoverySchema = Type.Object({ - type: Type.Union([Type.Literal("ollama")]), + type: Type.Union([Type.Literal("ollama"), Type.Literal("lm-studio")]), }); const ProviderAuthSchema = Type.Union([Type.Literal("apiKey"), Type.Literal("none")]); @@ -594,14 +595,24 @@ export class ModelRegistry { } #addImplicitDiscoverableProviders(configuredProviders: Set): void { - if (configuredProviders.has("ollama")) return; - this.#discoverableProviders.push({ - provider: "ollama", - api: "openai-completions", - baseUrl: Bun.env.OLLAMA_BASE_URL || "http://127.0.0.1:11434", - discovery: { type: "ollama" }, - }); - this.#keylessProviders.add("ollama"); + if (!configuredProviders.has("ollama")) { + this.#discoverableProviders.push({ + provider: "ollama", + api: "openai-completions", + baseUrl: Bun.env.OLLAMA_BASE_URL || "http://127.0.0.1:11434", + discovery: { type: "ollama" }, + }); + this.#keylessProviders.add("ollama"); + } + if (!configuredProviders.has("lm-studio")) { + this.#discoverableProviders.push({ + provider: "lm-studio", + api: "openai-completions", + baseUrl: Bun.env.LM_STUDIO_BASE_URL || "http://127.0.0.1:1234/v1", + discovery: { type: "lm-studio" }, + }); + this.#keylessProviders.add("lm-studio"); + } } #loadCustomModels(): CustomModelsResult { @@ -724,6 +735,8 @@ export class ModelRegistry { switch (providerConfig.discovery.type) { case "ollama": return this.#discoverOllamaModels(providerConfig); + case "lm-studio": + return this.#discoverLmStudioModels(providerConfig); } } @@ -877,6 +890,71 @@ export class ModelRegistry { } } + async #discoverLmStudioModels(providerConfig: DiscoveryProviderConfig): Promise[]> { + const endpoint = providerConfig.baseUrl || "http://127.0.0.1:1234/v1"; + let baseUrl = endpoint; + if (!baseUrl.endsWith("/v1")) { + baseUrl = baseUrl.endsWith("/") ? `${baseUrl}v1` : `${baseUrl}/v1`; + } else if (baseUrl.endsWith("/v1/")) { + baseUrl = baseUrl.slice(0, -1); + } + const modelsUrl = `${baseUrl}/models`; + + const headers: Record = { ...(providerConfig.headers ?? {}) }; + const apiKey = await this.authStorage.getApiKey("lm-studio"); + if (apiKey && apiKey !== DEFAULT_LOCAL_TOKEN && apiKey !== kNoAuth) { + headers.Authorization = `Bearer ${apiKey}`; + } + + try { + const response = await fetch(modelsUrl, { + headers, + signal: AbortSignal.timeout(3000), + }); + if (!response.ok) { + logger.warn("model discovery failed for provider", { + provider: providerConfig.provider, + status: response.status, + url: modelsUrl, + }); + return []; + } + const payload = (await response.json()) as { data?: Array<{ id: string }> }; + const models = payload.data ?? []; + const discovered: Model[] = []; + for (const item of models) { + const id = item.id; + if (!id) continue; + discovered.push({ + id, + name: id, + api: providerConfig.api, + provider: providerConfig.provider, + baseUrl, + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 128000, + maxTokens: 8192, + headers, + compat: { + supportsStore: false, + supportsDeveloperRole: false, + supportsReasoningEffort: false, + }, + }); + } + return this.#applyProviderModelOverrides(providerConfig.provider, discovered); + } catch (error) { + logger.warn("model discovery failed for provider", { + provider: providerConfig.provider, + url: modelsUrl, + error: error instanceof Error ? error.message : String(error), + }); + return []; + } + } + #normalizeOllamaBaseUrl(baseUrl?: string): string { const raw = baseUrl || "http://127.0.0.1:11434"; try { diff --git a/packages/coding-agent/test/lm-studio-fix.test.ts b/packages/coding-agent/test/lm-studio-fix.test.ts new file mode 100644 index 000000000..38eafb743 --- /dev/null +++ b/packages/coding-agent/test/lm-studio-fix.test.ts @@ -0,0 +1,129 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import * as fs from "node:fs"; +import * as os from "node:os"; +import * as path from "node:path"; +import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; +import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; +import { Snowflake } from "@oh-my-pi/pi-utils"; + +describe("ModelRegistry LM Studio Fixes", () => { + let tempDir: string; + let modelsJsonPath: string; + let authStorage: AuthStorage; + + beforeEach(async () => { + tempDir = path.join(os.tmpdir(), `pi-test-lm-studio-fixes-${Snowflake.next()}`); + fs.mkdirSync(tempDir, { recursive: true }); + modelsJsonPath = path.join(tempDir, "models.json"); + authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); + }); + + afterEach(() => { + if (tempDir && fs.existsSync(tempDir)) { + fs.rmSync(tempDir, { recursive: true }); + } + }); + + test("auto-discovers both ollama and lm-studio models independently", async () => { + const originalFetch = globalThis.fetch; + globalThis.fetch = (async (input: string | URL | Request) => { + const url = String(input); + if (url.includes(":11434/api/tags")) { + return new Response(JSON.stringify({ models: [{ name: "ollama-model" }] }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + if (url.includes(":1234/v1/models")) { + return new Response(JSON.stringify({ data: [{ id: "lm-studio-model" }] }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + return new Response(null, { status: 404 }); + }) as unknown as typeof fetch; + + try { + // Mock environment variables + const registry = new ModelRegistry(authStorage, modelsJsonPath); + await registry.refresh(); + + const allModels = registry.getAll(); + expect(allModels.some((m: any) => m.provider === "ollama" && m.id === "ollama-model")).toBe(true); + expect(allModels.some((m: any) => m.provider === "lm-studio" && m.id === "lm-studio-model")).toBe(true); + + const available = registry.getAvailable(); + expect(available.some((m: any) => m.provider === "ollama")).toBe(true); + expect(available.some((m: any) => m.provider === "lm-studio")).toBe(true); + } finally { + globalThis.fetch = originalFetch; + } + }); + + test("lm-studio discovery handles trailing slashes in baseUrl correctly", async () => { + const originalFetch = globalThis.fetch; + let requestedUrl = ""; + globalThis.fetch = (async (input: string | URL | Request) => { + requestedUrl = String(input); + return new Response(JSON.stringify({ data: [{ id: "model-1" }] }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + }) as unknown as typeof fetch; + + try { + // Scenario 1: No trailing slash + fs.writeFileSync( + modelsJsonPath, + JSON.stringify({ + providers: { + "lm-studio": { + baseUrl: "http://127.0.0.1:1234/v1", + api: "openai-completions", + discovery: { type: "lm-studio" }, + }, + }, + }), + ); + let registry = new ModelRegistry(authStorage, modelsJsonPath); + await registry.refresh(); + expect(requestedUrl).toBe("http://127.0.0.1:1234/v1/models"); + + // Scenario 2: With trailing slash + fs.writeFileSync( + modelsJsonPath, + JSON.stringify({ + providers: { + "lm-studio": { + baseUrl: "http://127.0.0.1:1234/v1/", + api: "openai-completions", + discovery: { type: "lm-studio" }, + }, + }, + }), + ); + registry = new ModelRegistry(authStorage, modelsJsonPath); + await registry.refresh(); + expect(requestedUrl).toBe("http://127.0.0.1:1234/v1/models"); + + // Scenario 3: Custom port without /v1 + fs.writeFileSync( + modelsJsonPath, + JSON.stringify({ + providers: { + "lm-studio": { + baseUrl: "http://127.0.0.1:9999", + api: "openai-completions", + discovery: { type: "lm-studio" }, + }, + }, + }), + ); + registry = new ModelRegistry(authStorage, modelsJsonPath); + await registry.refresh(); + expect(requestedUrl).toBe("http://127.0.0.1:9999/v1/models"); + } finally { + globalThis.fetch = originalFetch; + } + }); +});