diff --git a/packages/coding-agent/src/extensions/llama/provider.ts b/packages/coding-agent/src/extensions/llama/provider.ts index f7801d991..1692a7d45 100644 --- a/packages/coding-agent/src/extensions/llama/provider.ts +++ b/packages/coding-agent/src/extensions/llama/provider.ts @@ -111,11 +111,20 @@ export function createLlamaProvider(): LlamaProviderController { }, getModels: () => models, refreshModels: async (context: RefreshModelsContext): Promise => { + const stored = await context.store.read(); + if (stored) { + models = stored.models.filter( + (model): model is Model<"openai-completions"> => + model.provider === LLAMA_PROVIDER_ID && model.api === "openai-completions", + ); + } + if (!context.allowNetwork || context.signal?.aborted || context.credential?.type !== "api_key") return; const serverUrl = credentialServerUrl(context.credential); if (!serverUrl) return; const catalog = await new LlamaClient(serverUrl, context.credential.key).list({ signal: context.signal }); setCatalog(catalog, serverUrl); + if (!context.signal?.aborted) await context.store.write({ models, checkedAt: Date.now() }); }, stream: (model, context, options) => stream(model, context, options as ProviderStreamOptions | undefined), streamSimple: (model, context, options) => streamSimple(model, context, options), diff --git a/packages/coding-agent/test/llama-extension.test.ts b/packages/coding-agent/test/llama-extension.test.ts index cb286cba0..3af25e563 100644 --- a/packages/coding-agent/test/llama-extension.test.ts +++ b/packages/coding-agent/test/llama-extension.test.ts @@ -1,7 +1,7 @@ import { once } from "node:events"; import { createServer, type RequestListener, type Server, type ServerResponse } from "node:http"; import type { AddressInfo } from "node:net"; -import type { AuthContext, AuthPrompt } from "@earendil-works/pi-ai"; +import type { AuthContext, AuthPrompt, ModelsStoreEntry } from "@earendil-works/pi-ai"; import { afterEach, describe, expect, it } from "vitest"; import { createEventBus } from "../src/core/event-bus.ts"; import { createExtensionRuntime, loadExtensionFromFactory } from "../src/core/extensions/loader.ts"; @@ -86,6 +86,50 @@ describe("llama.cpp extension", () => { ]); }); + it("persists and restores loaded models for cache-only startup refreshes", async () => { + let cachedEntry: ModelsStoreEntry | undefined; + const store = { + read: async () => cachedEntry, + write: async (entry: ModelsStoreEntry) => { + cachedEntry = structuredClone(entry); + }, + delete: async () => { + cachedEntry = undefined; + }, + }; + const { url } = await listen((request, response) => { + if (request.url === "/models") { + json(response, { + data: [ + { id: "loaded", status: { value: "loaded" }, meta: { n_ctx: 32768 } }, + { id: "unloaded", status: { value: "unloaded" } }, + ], + }); + return; + } + response.writeHead(404).end(); + }); + + const first = createLlamaProvider(); + await first.provider.refreshModels?.({ + credential: { type: "api_key", key: "local", env: { LLAMA_BASE_URL: url } }, + store, + allowNetwork: true, + }); + expect(first.provider.getModels().map((model) => model.id)).toEqual(["loaded"]); + expect(cachedEntry?.models.map((model) => model.id)).toEqual(["loaded"]); + + const second = createLlamaProvider(); + await second.provider.refreshModels?.({ + credential: { type: "api_key", key: "local", env: { LLAMA_BASE_URL: url } }, + store, + allowNetwork: false, + }); + expect(second.provider.getModels()).toEqual([ + expect.objectContaining({ id: "loaded", baseUrl: `${url}/v1`, contextWindow: 32768 }), + ]); + }); + it("stays dormant until configured and stores URL plus optional key", async () => { const { provider } = createLlamaProvider(); const auth = provider.auth.apiKey!;