pi-llama-cpp 0.4.0 → 0.5.1

package/README.md

@@ -111,6 +111,7 @@ The extension determines the context size as follows:
 | ---------------- | ------------------------------------------------------------------------------------------ |
 | `/models`        | Browse your models with live status. Select a model to load, switch, or unload it.         |
 | `/models info`   | Show detailed information for all available models at once.                                |
+| `/models unload` | Unload all loaded models at once (Note: this only makes sense in router mode).             |
 
 > **Note:** When the llama.cpp server is unreachable, `/models` displays an error notification with the configured server URL.
 
@@ -142,9 +143,9 @@ When you trigger a load, switch, or retry action, the extension polls the server
 
 Each model exposed to Pi includes the following defaults:
 
-- **`maxTokens`** — `32000` (maximum possible tokens per response according to Pi's source code)
+- **`maxTokens`** — dynamically set to the model's context window (detected from llama-server)
 - **`reasoning`** — `true` (assumed, as llama.cpp's `/models` endpoint does not expose it)
-- **`cost`** — all zero (local model)
+- **`cost`** — all zero (local models)
 
 ## Dependencies
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-llama-cpp",
-  "version": "0.4.0",
+  "version": "0.5.1",
   "description": "Pi extension for llama.cpp integration. Supports both router and single modes.",
   "keywords": [
     "pi",
@@ -24,8 +24,7 @@
     ]
   },
   "scripts": {
-    "test": "vitest",
-    "test:run": "vitest run"
+    "test": "vitest run"
   },
   "prettier": {
     "plugins": [
@@ -36,8 +35,8 @@
     "@earendil-works/pi-coding-agent": "*"
   },
   "devDependencies": {
-    "@types/node": "^25.6.0",
+    "@types/node": "^25.9.1",
     "prettier-plugin-organize-imports": "^4.3.0",
-    "vitest": "^4.1.5"
+    "vitest": "^4.1.7"
   }
 }

package/src/commands/models.ts

@@ -1,14 +1,43 @@
 import type {
   ExtensionAPI,
   ExtensionCommandContext,
+  ExtensionContext,
+  SessionBeforeSwitchEvent,
 } from "@earendil-works/pi-coding-agent";
-import { PROVIDER_ID, PROVIDER_NAME } from "../constants";
+import { PROVIDER_ID, PROVIDER_NAME, READABLE_TIMEOUT } from "../constants";
 import { Action } from "../enums/action";
 import { Mode } from "../enums/mode";
 import { Status } from "../enums/status";
 import { BaseModel } from "../models/baseModel";
 import { resolveUrl } from "../tools/resolver";
 
+// In-flight model reference — handler gates on this.
+let inflightModel: BaseModel | null = null;
+
+export const resetInflightModel = () => (inflightModel = null);
+
+/**
+ * Session-switch handler. Registered once at extension init.
+ * Only notifies if a model load is actually in-flight.
+ */
+export const onSessionBeforeSwitch = async (
+  _event: SessionBeforeSwitchEvent,
+  ctx: ExtensionContext,
+) => {
+  if (!inflightModel) return;
+
+  const messages = [
+    `Session change detected while model '${inflightModel.name}' was still loading.`,
+    "Model load will continue in the background, but UI might not update.",
+    "",
+    "Verify that your new model is loaded, or use /models to re-select it afterwards.",
+  ];
+  ctx.ui.notify(messages.join("\n"), "warning");
+
+  // Show the notification for a reasonable amount of time
+  await new Promise((r) => setTimeout(r, READABLE_TIMEOUT));
+};
+
 /**
  * Select a model from the list. Returns null if user cancels.
  *
@@ -167,6 +196,7 @@ export const modelsCommand = async (
   const loadActions = [Action.LOAD, Action.SWITCH, Action.RETRY];
   if (loadActions.includes(action)) {
     ctx.ui.notify(`Loading ${model.name}...`, "info");
+    inflightModel = model;
 
     const onSuccess = async () => {
       const piModel = ctx.modelRegistry.find(PROVIDER_ID, model.id);
@@ -175,7 +205,7 @@ export const modelsCommand = async (
       }
 
       if ((await model.getStatus()) === Status.FAILED) {
-        throw new Error(`Failed to load model  ${model.name}`);
+        throw new Error(`Failed to load model ${model.name}`);
       }
 
       await pi.setModel(piModel);
@@ -184,10 +214,15 @@ export const modelsCommand = async (
 
     const onFailure = (err: any) => {
       const message = err instanceof Error ? err.message : String(err);
-      ctx.ui.notify(message, "error");
+
+      try {
+        ctx.ui.notify(message, "error");
+      } catch {
+        // ctx went stale between error and notification
+      }
     };
 
     // Load the model without blocking the UI
-    model.load().then(onSuccess).catch(onFailure);
+    model.load().then(onSuccess).catch(onFailure).finally(resetInflightModel);
   }
 };

package/src/constants.ts

@@ -23,11 +23,6 @@ export const API_KEY_PLACEHOLDER = "sk-placeholder";
  */
 export const DEFAULT_CTX = 128000;
 
-/**
- * Maximum number of tokens a model can generate in a single response
- */
-export const MAX_TOKENS = 32000;
-
 /**
  * Polling interval (ms) for checking model load status
  */
@@ -37,3 +32,8 @@ export const POLLING_INTERVAL = 500;
  * Maximum time (ms) to wait for model loading before giving up
  */
 export const POLLING_TIMEOUT = 60000;
+
+/**
+ * Reasonable time to read notifications if context goes stale
+ */
+export const READABLE_TIMEOUT = 15000;

package/src/index.ts

@@ -2,6 +2,8 @@ import type {
   ExtensionAPI,
   ExtensionCommandContext,
 } from "@earendil-works/pi-coding-agent";
+import type { AutocompleteItem } from "@earendil-works/pi-tui";
+import { onSessionBeforeSwitch } from "./commands/models";
 import { PROVIDER_NAME } from "./constants";
 import { onModelSelect } from "./events";
 import { CommandManager } from "./manager";
@@ -13,10 +15,28 @@ export default async function (pi: ExtensionAPI) {
   // Command: /models
   pi.registerCommand("models", {
     description: `Browse ${PROVIDER_NAME} models`,
+    getArgumentCompletions: (prefix: string): AutocompleteItem[] | null => {
+      const available = [
+        {
+          value: "info",
+          label: "info",
+          description: "Show information of all models",
+        },
+        {
+          value: "unload",
+          label: "unload",
+          description: "Unload all models",
+        },
+      ];
+
+      const filtered = available.filter((a) => a.value.startsWith(prefix));
+      return filtered.length > 0 ? filtered : null;
+    },
     handler: async (args: string, ctx: ExtensionCommandContext) =>
       await manager.run(args, ctx),
   });
 
   // Events registration
   pi.on("model_select", onModelSelect);
+  pi.on("session_before_switch", onSessionBeforeSwitch);
 }

package/src/manager.ts

@@ -66,13 +66,15 @@ export class CommandManager {
    * @param args Arguments passed to the command
    * @param ctx The context used by Pi
    * @param pi The Pi extension
-   * @returns A command handler
    */
   async run(args: string, ctx: ExtensionCommandContext) {
     if (!(await isServerReady())) {
       return await notFoundCommand(ctx);
     }
 
+    // Refresh the model list from the server
+    await this.update();
+
     // Command: `/models info`
     if (args === "info") {
       const info = await Promise.all(this.serverModels.map((m) => m.getInfo()));
@@ -81,7 +83,14 @@ export class CommandManager {
       return;
     }
 
-    // Command: `/models`
+    // Command: `/models unload`
+    if (args === "unload") {
+      await Promise.all(this.serverModels.map((m) => m.unload()));
+      ctx.ui.notify(`Unloaded all ${PROVIDER_NAME} models`, "info");
+      return;
+    }
+
+    // Command: `/models` (interactive menu)
     return await modelsCommand(ctx, this.pi, this.serverModels);
   }
 }

package/src/models/baseModel.ts

@@ -1,5 +1,5 @@
 import type { ProviderModelConfig } from "@earendil-works/pi-coding-agent";
-import { MAX_TOKENS, POLLING_INTERVAL, POLLING_TIMEOUT } from "../constants";
+import { POLLING_INTERVAL, POLLING_TIMEOUT } from "../constants";
 import { Mode } from "../enums/mode";
 import { Status } from "../enums/status";
 import { DataProperty, ModelsEndpoint } from "../interfaces/endpoints/models";
@@ -127,7 +127,7 @@ export abstract class BaseModel {
       input: await this.getCapabilities(),
       contextWindow: await this.getContextSize(),
       cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-      maxTokens: MAX_TOKENS,
+      maxTokens: await this.getContextSize(),
     };
 
     return response;

package/src/tools/resolver.ts

@@ -1,3 +1,4 @@
+import { getAgentDir } from "@earendil-works/pi-coding-agent";
 import { access, constants, readFile } from "node:fs/promises";
 import { join } from "node:path";
 import {
@@ -59,7 +60,7 @@ const readConfigValue = async <T>(
  * @returns The API key, as defined by the auth.json file
  */
 export const resolveApiKey = async (): Promise<string> => {
-  const authPath = join(process.env.HOME || ".", ".pi", "agent", "auth.json");
+  const authPath = join(getAgentDir(), "settings.json");
   if (!(await fileExists(authPath))) return API_KEY_PLACEHOLDER;
 
   const cfg = await readConfigValue<AuthFile>(authPath, PROVIDER_ID);
@@ -71,13 +72,7 @@ export const resolveApiKey = async (): Promise<string> => {
  * @returns The URL, if found.
  */
 const resolveGlobalUrl = async (): Promise<string | null> => {
-  const globalPath = join(
-    process.env.HOME || ".",
-    ".pi",
-    "agent",
-    "settings.json",
-  );
-
+  const globalPath = join(getAgentDir(), "settings.json");
   if (!(await fileExists(globalPath))) return null;
 
   return readConfigValue<Record<string, string>>(globalPath, "llamaServerUrl");

package/tests/commandManager.test.ts

@@ -1,6 +1,6 @@
-import { describe, expect, it, vi, beforeEach } from "vitest";
-import { CommandManager } from "../src/manager";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 import { PROVIDER_ID, PROVIDER_NAME } from "../src/constants";
+import { CommandManager } from "../src/manager";
 
 // Mock modules at top level (vi.mock is hoisted)
 vi.mock("../src/tools/retriever", () => ({
@@ -14,8 +14,8 @@ vi.mock("../src/tools/resolver", () => ({
 }));
 
 // Import mocked functions after vi.mock
+import { resolveApiKey, resolveUrl } from "../src/tools/resolver";
 import { isServerReady, listModels } from "../src/tools/retriever";
-import { resolveUrl, resolveApiKey } from "../src/tools/resolver";
 
 const mockPi = {
   registerProvider: vi.fn(),
@@ -47,7 +47,9 @@ describe("CommandManager", () => {
     const mockModel = {
       name: "test-model",
       id: "test-model",
-      toProviderConfig: vi.fn().mockResolvedValue({ id: "test-model", maxTokens: 32000 }),
+      toProviderConfig: vi
+        .fn()
+        .mockResolvedValue({ id: "test-model", maxTokens: 32000 }),
     };
     (isServerReady as any).mockResolvedValue(true);
     (listModels as any).mockResolvedValue([mockModel]);
@@ -92,11 +94,40 @@ describe("CommandManager", () => {
       ui: { notify: notifyFn, theme: { fg: (_c: string, t: string) => t } },
     } as any);
 
+    expect(notifyFn).toHaveBeenCalledWith("Model info for test-model", "info");
+    // Called once in initialize() and once in run() to refresh the model list
+    expect(listModels).toHaveBeenCalledTimes(2);
+  });
+
+  it("should unload all models when args is 'unload'", async () => {
+    const mockModel1 = {
+      name: "model-1",
+      id: "model-1",
+      unload: vi.fn().mockResolvedValue(undefined),
+      toProviderConfig: vi.fn().mockResolvedValue({ id: "model-1" }),
+    };
+    const mockModel2 = {
+      name: "model-2",
+      id: "model-2",
+      unload: vi.fn().mockResolvedValue(undefined),
+      toProviderConfig: vi.fn().mockResolvedValue({ id: "model-2" }),
+    };
+    (isServerReady as any).mockResolvedValue(true);
+    (listModels as any).mockResolvedValue([mockModel1, mockModel2]);
+
+    const notifyFn = vi.fn();
+    const manager = new CommandManager(mockPi as any);
+    await manager.initialize();
+    await manager.run("unload", {
+      ui: { notify: notifyFn },
+    } as any);
+
+    expect(mockModel1.unload).toHaveBeenCalled();
+    expect(mockModel2.unload).toHaveBeenCalled();
     expect(notifyFn).toHaveBeenCalledWith(
-      "Model info for test-model",
+      "Unloaded all Llama.cpp models",
       "info",
     );
-    expect(listModels).toHaveBeenCalledOnce();
   });
 
   it("should dispatch modelsCommand when args is empty", async () => {

package/tests/modelsCommand.test.ts

@@ -1,5 +1,14 @@
-import { describe, expect, it, vi } from "vitest";
-import { modelsCommand } from "../src/commands/models";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+// Set up fake timers before any imports so setTimeout is mocked globally
+vi.useFakeTimers();
+
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import {
+  modelsCommand,
+  onSessionBeforeSwitch,
+  resetInflightModel,
+} from "../src/commands/models";
 import { Action } from "../src/enums/action";
 import { Mode } from "../src/enums/mode";
 import { Status } from "../src/enums/status";
@@ -48,11 +57,27 @@ const createMockCtx = (
   },
 });
 
+const createMockPiContext = (notifyFn: ReturnType<typeof vi.fn>) =>
+  ({
+    ui: {
+      notify: notifyFn,
+    },
+  }) as any as ExtensionContext;
+
 const createMockPi = () => ({
   setModel: vi.fn(),
   registerProvider: vi.fn(),
 });
 
+beforeEach(() => {
+  vi.clearAllTimers();
+  resetInflightModel();
+});
+
+afterEach(() => {
+  vi.clearAllTimers();
+});
+
 describe("modelsCommand", () => {
   it("should return early on cancel (null model selection)", async () => {
     const models = [createMockModel("model-a")];
@@ -113,6 +138,117 @@ describe("modelsCommand", () => {
     await vi.waitFor(() => expect(pi.setModel).toHaveBeenCalled());
   });
 
+  it("should show warning when session changes during model load", async () => {
+    // Create a deferred promise so we can control when the load completes
+    let resolveLoad: () => void;
+    const loadPromise = new Promise<void>((resolve) => {
+      resolveLoad = resolve;
+    });
+    const model = createMockModel("model-a", {
+      load: () => loadPromise,
+      getStatus: vi.fn().mockResolvedValue(Status.UNLOADED),
+    });
+    const models = [model];
+
+    let selectCallCount = 0;
+    const ctx = createMockCtx(() => {
+      selectCallCount++;
+      // 1st: select model, 2nd: select LOAD
+      if (selectCallCount === 1) return "model-a";
+      if (selectCallCount === 2) return Action.LOAD;
+      return null;
+    });
+    const pi = createMockPi();
+
+    // Start the load (non-blocking)
+    const modelsPromise = modelsCommand(ctx as any, pi as any, models);
+
+    // Advance past the microtask that sets inflightModel
+    await vi.advanceTimersByTimeAsync(0);
+
+    // Simulate session switch while model is still loading
+    // onSessionBeforeSwitch awaits READABLE_TIMEOUT (15s) for the notification
+    const switchPromise = onSessionBeforeSwitch(
+      {} as any,
+      createMockPiContext(ctx.ui.notify as any),
+    );
+    await vi.advanceTimersByTimeAsync(15000);
+    await switchPromise;
+
+    // Should have shown a warning notification
+    expect(ctx.ui.notify).toHaveBeenCalledWith(
+      expect.stringContaining("Session change detected"),
+      "warning",
+    );
+    expect(ctx.ui.notify).toHaveBeenCalledWith(
+      expect.stringContaining("model-a"),
+      "warning",
+    );
+
+    // Complete the load so inflightModel is cleared
+    resolveLoad!();
+    await modelsPromise;
+  });
+
+  it("should not warn when no model is loading", async () => {
+    const notifyFn = vi.fn();
+    const ctx = createMockPiContext(notifyFn);
+
+    await onSessionBeforeSwitch({} as any, ctx);
+
+    expect(notifyFn).not.toHaveBeenCalled();
+    // No timers should be scheduled
+    expect(vi.getTimerCount()).toBe(0);
+  });
+
+  it("should clear inflightModel after load completes successfully", async () => {
+    const loadFn = vi.fn().mockResolvedValue(undefined);
+    const model = createMockModel("model-a", {
+      load: loadFn,
+      getStatus: vi.fn().mockResolvedValue(Status.UNLOADED),
+    });
+    const models = [model];
+    const ctx = createMockCtx((prompt) => {
+      if (prompt.includes("models")) return "model-a";
+      return Action.LOAD;
+    });
+    const pi = createMockPi();
+
+    await modelsCommand(ctx as any, pi as any, models);
+    await vi.waitFor(() => expect(loadFn).toHaveBeenCalled());
+    await vi.waitFor(() => expect(pi.setModel).toHaveBeenCalled());
+
+    // inflightModel should be cleared after completion
+    // (verified indirectly: calling onSessionBeforeSwitch should not warn)
+    await vi.advanceTimersByTimeAsync(0);
+    const notifyFn = vi.fn();
+    await onSessionBeforeSwitch({} as any, createMockPiContext(notifyFn));
+    expect(notifyFn).not.toHaveBeenCalled();
+  });
+
+  it("should clear inflightModel after load fails", async () => {
+    const loadFn = vi.fn().mockRejectedValue(new Error("Load failed"));
+    const model = createMockModel("model-a", {
+      load: loadFn,
+      getStatus: vi.fn().mockResolvedValue(Status.FAILED),
+    });
+    const models = [model];
+    const ctx = createMockCtx((prompt) => {
+      if (prompt.includes("models")) return "model-a";
+      return Action.RETRY;
+    });
+    const pi = createMockPi();
+
+    await modelsCommand(ctx as any, pi as any, models);
+    await vi.waitFor(() => expect(loadFn).toHaveBeenCalled());
+
+    // inflightModel should be cleared after failure
+    await vi.advanceTimersByTimeAsync(0);
+    const notifyFn = vi.fn();
+    await onSessionBeforeSwitch({} as any, createMockPiContext(notifyFn));
+    expect(notifyFn).not.toHaveBeenCalled();
+  });
+
   it("should loop back to model selection when action is cancelled", async () => {
     const model = createMockModel("model-a");
     const models = [model];

package/tests/singleModel.test.ts

@@ -65,7 +65,9 @@ describe("SingleModel getStatus", () => {
     const status = await model.getStatus();
 
     expect(status).toBe(Status.LOADED);
-    expect(mockRpc).toHaveBeenCalledWith(`/props?model=${model.id}&autoload=false`);
+    expect(mockRpc).toHaveBeenCalledWith(
+      `/props?model=${model.id}&autoload=false`,
+    );
   });
 
   it("should return SLEEPING when is_sleeping is true", async () => {