pi-llama-cpp 0.3.4 → 0.5.0

package/README.md

@@ -99,17 +99,21 @@ llama-server --model path/to/model.gguf ...
 ```
 
 The extension determines the context size as follows:
-- **Router mode** — when loaded, reads `meta.n_ctx` from the `/models` endpoint; when not loaded, reads `--ctx-size` and/or `--fit-ctx` from the model's status `args` array
+- **Router mode**
+  - When loaded, reads `meta.n_ctx` from the `/models` endpoint
+  - When not loaded, reads `--ctx-size` and/or `--fit-ctx` from the server arguments, or `ctx-size` and/or `fit-ctx` keys from the **presets.ini** file.
 - **Single mode** — reads `meta.n_ctx` from the `/models` endpoint
 - Falls back to `128000` if not available
 
 ### Commands
 
-| Command   | Description                                                                                |
-| --------- | ------------------------------------------------------------------------------------------ |
-| `/models` | Browse your models with live status. Select a model to load, switch, or unload it.         |
+| Command          | Description                                                                                |
+| ---------------- | ------------------------------------------------------------------------------------------ |
+| `/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` is still available but shows the description `Llama.cpp models (offline)` and displays an error notification with the configured server URL.
+> **Note:** When the llama.cpp server is unreachable, `/models` displays an error notification with the configured server URL.
 
 ### Model Actions
 
@@ -139,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.3.4",
+  "version": "0.5.0",
   "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.
  *
@@ -130,8 +159,10 @@ export const notFoundCommand = async (
 /**
  * Handles the /models command
  *
+ * @param args Arguments passed to the command
  * @param ctx The context used by Pi
  * @param pi The Pi extension
+ * @param models List of available models
  */
 export const modelsCommand = async (
   ctx: ExtensionCommandContext,
@@ -165,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);
@@ -173,7 +205,7 @@ export const modelsCommand = async (
       }
 
       if ((await model.getStatus()) === Status.FAILED) {
-        throw new Error("Failed to load model");
+        throw new Error(`Failed to load model ${model.name}`);
       }
 
       await pi.setModel(piModel);
@@ -182,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/events.ts

@@ -18,6 +18,9 @@ export const onModelSelect = async (
   const model = models.find((m) => m.id === event.model.id);
   if (!model) return;
 
-  ctx.ui.notify(`>> Loading ${model.id}...`, "info");
-  await model.load();
+  ctx.ui.notify(`Loading ${model.name}...`, "info");
+  await model
+    .load()
+    .then(() => ctx.ui.notify(`Model ${model.name} ready`, "info"))
+    .catch(() => ctx.ui.notify(`Failed to load model ${model.name}`, "error"));
 };

package/src/index.ts

@@ -2,35 +2,41 @@ import type {
   ExtensionAPI,
   ExtensionCommandContext,
 } from "@earendil-works/pi-coding-agent";
-import { modelsCommand, notFoundCommand } from "./commands/models";
+import type { AutocompleteItem } from "@earendil-works/pi-tui";
+import { onSessionBeforeSwitch } from "./commands/models";
 import { PROVIDER_NAME } from "./constants";
 import { onModelSelect } from "./events";
-import { registerLlamaCppProvider } from "./tools/provider";
-import { isServerReady } from "./tools/retriever";
+import { CommandManager } from "./manager";
 
 export default async function (pi: ExtensionAPI) {
-  // Server verification
-  if (!(await isServerReady())) {
-    pi.registerCommand("models", {
-      description: `${PROVIDER_NAME} models (offline)`,
-      handler: async (_: string, ctx: ExtensionCommandContext) => {
-        await notFoundCommand(ctx);
-      },
-    });
-
-    return;
-  }
-
-  // Provider registration
-  const serverModels = await registerLlamaCppProvider(pi);
+  const manager = new CommandManager(pi);
+  await manager.initialize();
 
   // Command: /models
   pi.registerCommand("models", {
-    description: `Browse ${PROVIDER_NAME} models (live status)`,
-    handler: async (_: string, ctx: ExtensionCommandContext) =>
-      await modelsCommand(ctx, pi, serverModels),
+    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/interfaces/endpoints/models.ts

@@ -39,6 +39,7 @@ export interface DataProperty {
   owned_by: string;
   created: number;
   status?: StatusProperty;
+  architecture?: ArchitectureProperty;
   meta?: MetaProperty;
 }
 
@@ -50,6 +51,11 @@ interface StatusProperty {
   failed?: boolean;
 }
 
+interface ArchitectureProperty {
+  input_modalities: ("text" | "image" | "audio")[];
+  output_modalities: ["text"];
+}
+
 interface MetaProperty {
   vocab_type: number;
   n_vocab: number;

package/src/manager.ts

@@ -0,0 +1,93 @@
+import type {
+  ExtensionAPI,
+  ExtensionCommandContext,
+  ProviderModelConfig,
+} from "@earendil-works/pi-coding-agent";
+import { modelsCommand, notFoundCommand } from "./commands/models";
+import {
+  DEFAULT_LLAMA_SERVER_URL,
+  PROVIDER_ID,
+  PROVIDER_NAME,
+} from "./constants";
+import { BaseModel } from "./models/baseModel";
+import { resolveApiKey, resolveUrl } from "./tools/resolver";
+import { isServerReady, listModels } from "./tools/retriever";
+
+export class CommandManager {
+  private baseUrl: string = DEFAULT_LLAMA_SERVER_URL;
+  private serverModels: BaseModel[] = [];
+
+  constructor(private readonly pi: ExtensionAPI) {}
+
+  /**
+   * Sets up the initial state of the provider
+   */
+  async initialize() {
+    if (await isServerReady()) {
+      await this.update();
+    } else {
+      await this.register([]);
+    }
+  }
+
+  /**
+   * Ensures the models are up-to-date with the server
+   */
+  async update() {
+    this.baseUrl = `${await resolveUrl(process.cwd())}`;
+
+    this.serverModels = await listModels();
+    const modelConfigs = await Promise.all(
+      this.serverModels.map((m) => m.toProviderConfig()),
+    );
+
+    await this.register(modelConfigs);
+  }
+
+  /**
+   * Registers the provider in Pi with the given configurations
+   * Note: Registrations overload previous provider
+   *
+   * @param models Provider configurations for the models
+   */
+  async register(models: ProviderModelConfig[]) {
+    this.pi.registerProvider(PROVIDER_ID, {
+      name: PROVIDER_NAME,
+      baseUrl: this.baseUrl,
+      api: "openai-completions",
+      apiKey: await resolveApiKey(),
+      models,
+    });
+  }
+
+  /**
+   * Dispatches the /models command
+   *
+   * @param args Arguments passed to the command
+   * @param ctx The context used by Pi
+   * @param pi The Pi extension
+   */
+  async run(args: string, ctx: ExtensionCommandContext) {
+    if (!(await isServerReady())) {
+      return await notFoundCommand(ctx);
+    }
+
+    // Command: `/models info`
+    if (args === "info") {
+      const info = await Promise.all(this.serverModels.map((m) => m.getInfo()));
+      const message = ctx.ui.theme.fg("accent", info.join("\n"));
+      ctx.ui.notify(message, "info");
+      return;
+    }
+
+    // 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,13 +1,8 @@
 import type { ProviderModelConfig } from "@earendil-works/pi-coding-agent";
-import {
-  DEFAULT_CTX,
-  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 } from "../interfaces/endpoints/models";
+import { DataProperty, ModelsEndpoint } from "../interfaces/endpoints/models";
 import { PropsEndpoint } from "../interfaces/endpoints/props";
 import { rpc } from "../tools/retriever";
 
@@ -55,17 +50,7 @@ export abstract class BaseModel {
    *
    * @returns An array of capabilities, as expected by Pi
    */
-  async getCapabilities(): Promise<["text"] | ["image"]> {
-    try {
-      const { modalities } = await rpc<PropsEndpoint>(
-        `/props?model=${this.id}`,
-      );
-
-      return modalities.vision ? ["image"] : ["text"];
-    } catch {
-      return ["text"];
-    }
-  }
+  abstract getCapabilities(): Promise<("text" | "image")[]>;
 
   /**
    * Gets the load status of the model
@@ -75,7 +60,7 @@ export abstract class BaseModel {
   public async getStatus(): Promise<Status> {
     try {
       const { is_sleeping, error } = await rpc<PropsEndpoint>(
-        `/props?model=${this.id}`,
+        `/props?model=${this.id}&autoload=false`,
       );
 
       if (is_sleeping) return Status.SLEEPING;
@@ -96,15 +81,10 @@ export abstract class BaseModel {
    * @returns The detected context size
    */
   async getContextSize(): Promise<number> {
-    try {
-      const { default_generation_settings } = await rpc<PropsEndpoint>(
-        `/props?model=${this.id}`,
-      );
-      const { n_ctx } = default_generation_settings;
-      return n_ctx;
-    } catch {
-      return DEFAULT_CTX;
-    }
+    const { data } = await rpc<ModelsEndpoint>("/models");
+    const { n_ctx } = data.find((m) => m.id === this.id)?.meta!;
+
+    return n_ctx;
   }
 
   /**
@@ -147,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;
@@ -157,7 +137,8 @@ export abstract class BaseModel {
    * Loads the model in llama-server
    */
   async load(): Promise<void> {
-    if ((await this.getStatus()) === Status.LOADED) return;
+    const status = await this.getStatus();
+    if (status === Status.LOADED || status === Status.SLEEPING) return;
 
     await rpc("/models/load", { model: this.id });
     await this.pollStatus();

package/src/models/routerModel.ts

@@ -50,7 +50,7 @@ export class RouterModel extends BaseModel {
     // Grab the glitch
     while (Date.now() - startTime <= limit) {
       try {
-        await rpc<PropsEndpoint>(`/props?model=${this.id}`);
+        await rpc<PropsEndpoint>(`/props?model=${this.id}&autoload=false`);
         break;
       } catch {
         elapsed += POLLING_INTERVAL;
@@ -62,14 +62,17 @@ export class RouterModel extends BaseModel {
     return await super.pollStatus(startTime, timeout);
   }
 
-  async getCapabilities(): Promise<["text"] | ["image"]> {
-    // We can get the real capabilities if the model is already loaded
-    if ((await this.getStatus()) === Status.LOADED) {
-      return super.getCapabilities();
-    }
+  async getCapabilities(): Promise<("text" | "image")[]> {
+    const { data } = await rpc<ModelsEndpoint>(`/models`);
+    const model = data.find((d) => d.id === this.id);
+    if (!model) return ["text"];
 
-    const hasImage = this.model.status?.args?.includes("--mmproj") ?? false;
-    return hasImage ? ["image"] : ["text"];
+    const { input_modalities } = model.architecture!;
+    const response = input_modalities.filter(
+      (mod) => mod === "text" || mod === "image",
+    );
+
+    return response;
   }
 
   async getContextSize(): Promise<number> {

package/src/models/singleModel.ts

@@ -1,8 +1,18 @@
 import { Mode } from "../enums/mode";
+import { ModelsEndpoint } from "../interfaces/endpoints/models";
+import { rpc } from "../tools/retriever";
 import { BaseModel } from "./baseModel";
 
 export class SingleModel extends BaseModel {
   get mode(): Mode {
     return Mode.SINGLE;
   }
+
+  async getCapabilities(): Promise<("text" | "image")[]> {
+    const { models } = await rpc<ModelsEndpoint>(`/models`);
+    const [model] = models!;
+
+    const hasImage = model.capabilities.includes("multimodal");
+    return hasImage ? ["text", "image"] : ["text"];
+  }
 }

package/src/tools/retriever.ts

@@ -28,7 +28,7 @@ export const isServerReady = async (): Promise<boolean> => {
 export const rpc = async <T>(
   endpoint: string,
   body?: Record<string, unknown>,
-) => {
+): Promise<T> => {
   const base = await resolveUrl(process.cwd());
   const url = `${base}${endpoint}`;
 

package/tests/commandManager.test.ts

@@ -0,0 +1,152 @@
+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", () => ({
+  isServerReady: vi.fn(),
+  listModels: vi.fn(),
+}));
+
+vi.mock("../src/tools/resolver", () => ({
+  resolveUrl: vi.fn(),
+  resolveApiKey: vi.fn(),
+}));
+
+// Import mocked functions after vi.mock
+import { resolveApiKey, resolveUrl } from "../src/tools/resolver";
+import { isServerReady, listModels } from "../src/tools/retriever";
+
+const mockPi = {
+  registerProvider: vi.fn(),
+};
+
+beforeEach(() => {
+  vi.clearAllMocks();
+  (resolveUrl as any).mockResolvedValue("http://127.0.0.1:8080");
+  (resolveApiKey as any).mockResolvedValue("test-key");
+});
+
+describe("CommandManager", () => {
+  it("should register empty models when server is not ready", async () => {
+    (isServerReady as any).mockResolvedValue(false);
+
+    const manager = new CommandManager(mockPi as any);
+    await manager.initialize();
+
+    expect(mockPi.registerProvider).toHaveBeenCalledWith(PROVIDER_ID, {
+      name: PROVIDER_NAME,
+      baseUrl: "http://127.0.0.1:8080",
+      api: "openai-completions",
+      apiKey: "test-key",
+      models: [],
+    });
+  });
+
+  it("should update and register models when server is ready", async () => {
+    const mockModel = {
+      name: "test-model",
+      id: "test-model",
+      toProviderConfig: vi
+        .fn()
+        .mockResolvedValue({ id: "test-model", maxTokens: 32000 }),
+    };
+    (isServerReady as any).mockResolvedValue(true);
+    (listModels as any).mockResolvedValue([mockModel]);
+
+    const manager = new CommandManager(mockPi as any);
+    await manager.initialize();
+
+    expect(resolveUrl).toHaveBeenCalledWith(expect.any(String));
+    expect(listModels).toHaveBeenCalled();
+    expect(mockPi.registerProvider).toHaveBeenCalledWith(PROVIDER_ID, {
+      name: PROVIDER_NAME,
+      baseUrl: "http://127.0.0.1:8080",
+      api: "openai-completions",
+      apiKey: "test-key",
+      models: [{ id: "test-model", maxTokens: 32000 }],
+    });
+  });
+
+  it("should call notFoundCommand when server is not ready in run()", async () => {
+    (isServerReady as any).mockResolvedValue(false);
+
+    const manager = new CommandManager(mockPi as any);
+    await manager.run("", { ui: { notify: vi.fn() } } as any);
+
+    expect(mockPi.registerProvider).not.toHaveBeenCalled();
+  });
+
+  it("should show info for all models when args is 'info'", async () => {
+    const mockModel = {
+      name: "test-model",
+      id: "test-model",
+      getInfo: vi.fn().mockResolvedValue("Model info for test-model"),
+      toProviderConfig: vi.fn().mockResolvedValue({ id: "test-model" }),
+    };
+    (isServerReady as any).mockResolvedValue(true);
+    (listModels as any).mockResolvedValue([mockModel]);
+
+    const notifyFn = vi.fn();
+    const manager = new CommandManager(mockPi as any);
+    await manager.initialize();
+    await manager.run("info", {
+      ui: { notify: notifyFn, theme: { fg: (_c: string, t: string) => t } },
+    } as any);
+
+    expect(notifyFn).toHaveBeenCalledWith("Model info for test-model", "info");
+    expect(listModels).toHaveBeenCalledOnce();
+  });
+
+  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(
+      "Unloaded all Llama.cpp models",
+      "info",
+    );
+  });
+
+  it("should dispatch modelsCommand when args is empty", async () => {
+    const mockModel = {
+      name: "test-model",
+      id: "test-model",
+      getLabel: vi.fn().mockResolvedValue("test-model"),
+      toProviderConfig: vi.fn().mockResolvedValue({ id: "test-model" }),
+    };
+    (isServerReady as any).mockResolvedValue(true);
+    (listModels as any).mockResolvedValue([mockModel]);
+
+    const selectFn = vi.fn().mockReturnValue(null); // cancel immediately
+    const manager = new CommandManager(mockPi as any);
+    await manager.initialize();
+    await manager.run("", {
+      ui: { notify: vi.fn(), select: selectFn },
+    } as any);
+
+    // modelsCommand was called (select is invoked for model picking)
+    expect(selectFn).toHaveBeenCalled();
+  });
+});

package/tests/handlers.test.ts

@@ -62,7 +62,12 @@ const getActionsForModel = async (model: TestModel): Promise<Array<Action>> => {
     [Status.LOADED]: [Action.SWITCH, Action.UNLOAD, Action.INFO, Action.CANCEL],
     [Status.LOADING]: [Action.INFO, Action.CANCEL],
     [Status.FAILED]: [Action.RETRY, Action.CANCEL],
-    [Status.SLEEPING]: [Action.UNLOAD, Action.INFO, Action.CANCEL],
+    [Status.SLEEPING]: [
+      Action.SWITCH,
+      Action.UNLOAD,
+      Action.INFO,
+      Action.CANCEL,
+    ],
     [Status.UNLOADED]: [Action.LOAD, Action.CANCEL],
   };
 
@@ -106,7 +111,7 @@ describe("Action availability", () => {
     {
       mode: Mode.ROUTER,
       status: Status.SLEEPING,
-      expected: [Action.UNLOAD, Action.INFO, Action.CANCEL],
+      expected: [Action.SWITCH, Action.UNLOAD, Action.INFO, Action.CANCEL],
     },
     {
       mode: Mode.ROUTER,

package/tests/modelsCommand.test.ts

@@ -0,0 +1,270 @@
+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";
+import { BaseModel } from "../src/models/baseModel";
+
+// Mock the retriever module
+vi.mock("../src/tools/retriever", () => ({
+  rpc: vi.fn(),
+  isServerReady: vi.fn(),
+  listModels: vi.fn(),
+}));
+
+// Helper to create a mock BaseModel
+const createMockModel = (
+  name: string,
+  overrides: Partial<BaseModel> = {},
+): BaseModel =>
+  ({
+    name,
+    id: name,
+    mode: Mode.ROUTER,
+    capabilities: ["text"] as ["text"],
+    getStatus: vi.fn().mockResolvedValue(Status.LOADED),
+    getContextSize: vi.fn().mockResolvedValue(4096),
+    getInfo: vi.fn().mockResolvedValue(`Model: ${name}\nID: ${name}`),
+    load: vi.fn().mockResolvedValue(undefined),
+    unload: vi.fn().mockResolvedValue(undefined),
+    toProviderConfig: vi.fn().mockResolvedValue({}),
+    getLabel: vi.fn().mockResolvedValue(name),
+    ...overrides,
+  }) as unknown as BaseModel;
+
+const createMockCtx = (
+  selectFn: (prompt: string, options: string[]) => string | null,
+) => ({
+  cwd: "/tmp/test",
+  ui: {
+    select: vi.fn(selectFn),
+    notify: vi.fn(),
+    theme: {
+      fg: (color: string, text: string) => text,
+    },
+  },
+  modelRegistry: {
+    find: vi.fn().mockReturnValue({ id: "test-model-id" }),
+  },
+});
+
+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")];
+    const ctx = createMockCtx(() => null);
+    const pi = createMockPi();
+
+    await modelsCommand(ctx as any, pi as any, models);
+
+    expect(ctx.ui.notify).not.toHaveBeenCalled();
+  });
+
+  it("should show info when INFO action is selected", async () => {
+    const model = createMockModel("model-a");
+    const models = [model];
+    const ctx = createMockCtx((prompt) => {
+      if (prompt.includes("models")) return "model-a";
+      return Action.INFO;
+    });
+    const pi = createMockPi();
+
+    await modelsCommand(ctx as any, pi as any, models);
+
+    expect(ctx.ui.notify).toHaveBeenCalledWith(
+      "Model: model-a\nID: model-a",
+      "info",
+    );
+  });
+
+  it("should unload model when UNLOAD action is selected", async () => {
+    const model = createMockModel("model-a");
+    const models = [model];
+    const ctx = createMockCtx((prompt) => {
+      if (prompt.includes("models")) return "model-a";
+      return Action.UNLOAD;
+    });
+    const pi = createMockPi();
+
+    await modelsCommand(ctx as any, pi as any, models);
+
+    expect(model.unload).toHaveBeenCalled();
+    expect(ctx.ui.notify).toHaveBeenCalledWith("Unloaded model-a", "info");
+  });
+
+  it("should load model when LOAD action is selected", async () => {
+    const loadFn = vi.fn().mockResolvedValue(undefined);
+    const model = createMockModel("model-a");
+    (model.load as any) = loadFn;
+    (model.getStatus as any).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());
+  });
+
+  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];
+
+    let selectCallCount = 0;
+    const ctx = createMockCtx(() => {
+      selectCallCount++;
+      // 1st: select model-a, 2nd: cancel action, 3rd: cancel model => exit
+      if (selectCallCount === 1) return "model-a";
+      return null;
+    });
+    const pi = createMockPi();
+
+    await modelsCommand(ctx as any, pi as any, models);
+
+    expect(ctx.ui.select).toHaveBeenCalledTimes(3);
+    expect(ctx.ui.notify).not.toHaveBeenCalled();
+  });
+});

package/tests/routerModel.test.ts

@@ -130,9 +130,14 @@ describe("RouterModel context size extraction", () => {
         },
       ],
     });
-    // Second call: super.getContextSize() -> /props?model=test-model with default_generation_settings.n_ctx
+    // Second call: super.getContextSize() -> /models with meta.n_ctx
     mockRpc.mockResolvedValueOnce({
-      default_generation_settings: { n_ctx: 4096 },
+      data: [
+        {
+          id: "test-model",
+          meta: { n_ctx: 4096 },
+        },
+      ],
     });
 
     const model = new RouterModel(
@@ -149,7 +154,7 @@ describe("RouterModel context size extraction", () => {
     expect(ctxSize).toBe(4096);
   });
 
-  it("should return DEFAULT_CTX when no context size args are present and loaded", async () => {
+  it("should return n_ctx from meta when loaded without context size args", async () => {
     // First call: getStatus() -> /models
     mockRpc.mockResolvedValueOnce({
       data: [
@@ -163,17 +168,16 @@ describe("RouterModel context size extraction", () => {
         },
       ],
     });
-    // Second call: super.getContextSize() -> /models without meta.n_ctx
+    // Second call: super.getContextSize() -> /models with meta.n_ctx
     mockRpc.mockResolvedValueOnce({
       data: [
         {
           id: "test-model",
+          meta: { n_ctx: 4096 },
         },
       ],
     });
 
-    const { DEFAULT_CTX } = await import("../src/constants");
-
     const model = new RouterModel(
       createModel({
         status: {
@@ -185,13 +189,12 @@ describe("RouterModel context size extraction", () => {
     );
 
     const ctxSize = await model.getContextSize();
-    expect(ctxSize).toBe(DEFAULT_CTX);
+    expect(ctxSize).toBe(4096);
   });
 });
 
 describe("RouterModel capabilities detection", () => {
-  it("should detect image capability when modalities.vision is true", async () => {
-    // getStatus() calls /models first
+  it("should detect image capability from architecture.input_modalities", async () => {
     mockRpc.mockResolvedValueOnce({
       data: [
         {
@@ -202,21 +205,22 @@ describe("RouterModel capabilities detection", () => {
             preset: "default",
             failed: false,
           },
+          architecture: {
+            input_modalities: ["text", "image"],
+            output_modalities: ["text"],
+          },
         },
       ],
     });
-    // super.getCapabilities() calls /props?model=<id>
-    mockRpc.mockResolvedValueOnce({ modalities: { vision: true } });
 
     const model = new RouterModel(createModel());
     const capabilities = await model.getCapabilities();
 
-    expect(capabilities).toEqual(["image"]);
-    expect(mockRpc).toHaveBeenCalledWith("/props?model=test-model");
+    expect(capabilities).toEqual(["text", "image"]);
+    expect(mockRpc).toHaveBeenCalledWith("/models");
   });
 
-  it("should detect text-only capability when modalities.vision is false", async () => {
-    // getStatus() calls /models first
+  it("should detect text-only capability when only text in input_modalities", async () => {
     mockRpc.mockResolvedValueOnce({
       data: [
         {
@@ -227,11 +231,13 @@ describe("RouterModel capabilities detection", () => {
             preset: "default",
             failed: false,
           },
+          architecture: {
+            input_modalities: ["text"],
+            output_modalities: ["text"],
+          },
         },
       ],
     });
-    // super.getCapabilities() calls /props?model=<id>
-    mockRpc.mockResolvedValueOnce({ modalities: { vision: false } });
 
     const model = new RouterModel(createModel());
     const capabilities = await model.getCapabilities();
@@ -239,12 +245,11 @@ describe("RouterModel capabilities detection", () => {
     expect(capabilities).toEqual(["text"]);
   });
 
-  it("should default to text when /props call fails", async () => {
-    // getStatus() calls /models first
+  it("should return text when model not found in /models response", async () => {
     mockRpc.mockResolvedValueOnce({
       data: [
         {
-          id: "test-model",
+          id: "other-model",
           status: {
             value: "loaded",
             args: [],
@@ -254,76 +259,12 @@ describe("RouterModel capabilities detection", () => {
         },
       ],
     });
-    // super.getCapabilities() calls /props?model=<id> which fails
-    mockRpc.mockRejectedValueOnce(new Error("Connection refused"));
 
     const model = new RouterModel(createModel());
     const capabilities = await model.getCapabilities();
 
     expect(capabilities).toEqual(["text"]);
   });
-
-  it("should use status.args to detect image capability when not loaded", async () => {
-    // getStatus() calls /models first, returns unloaded
-    mockRpc.mockResolvedValueOnce({
-      data: [
-        {
-          id: "test-model",
-          status: {
-            value: "unloaded",
-            args: ["--model", "gguf", "--mmproj", "mmproj.gguf"],
-            preset: "default",
-            failed: false,
-          },
-        },
-      ],
-    });
-
-    const model = new RouterModel(
-      createModel({
-        status: {
-          value: "unloaded",
-          args: ["--model", "gguf", "--mmproj", "mmproj.gguf"],
-          preset: "default",
-          failed: false,
-        },
-      }),
-    );
-    const capabilities = await model.getCapabilities();
-
-    expect(capabilities).toEqual(["image"]);
-  });
-
-  it("should return text when not loaded and no --mmproj in args", async () => {
-    // getStatus() calls /models first, returns unloaded
-    mockRpc.mockResolvedValueOnce({
-      data: [
-        {
-          id: "test-model",
-          status: {
-            value: "unloaded",
-            args: ["--model", "gguf"],
-            preset: "default",
-            failed: false,
-          },
-        },
-      ],
-    });
-
-    const model = new RouterModel(
-      createModel({
-        status: {
-          value: "unloaded",
-          args: ["--model", "gguf"],
-          preset: "default",
-          failed: false,
-        },
-      }),
-    );
-    const capabilities = await model.getCapabilities();
-
-    expect(capabilities).toEqual(["text"]);
-  });
 });
 
 describe("RouterModel mode", () => {

package/tests/singleModel.test.ts

@@ -1,5 +1,4 @@
 import { beforeEach, describe, expect, it, vi } from "vitest";
-import { DEFAULT_CTX } from "../src/constants";
 import { Mode } from "../src/enums/mode";
 import { Status } from "../src/enums/status";
 import { ModelProperty } from "../src/interfaces/endpoints/models";
@@ -34,27 +33,22 @@ describe("SingleModel mode", () => {
 });
 
 describe("SingleModel capabilities", () => {
-  it("should detect image capability when modalities.vision is true", async () => {
-    mockRpc.mockResolvedValueOnce({ modalities: { vision: true } });
-
-    const model = createModel();
-    const capabilities = await model.getCapabilities();
-
-    expect(capabilities).toEqual(["image"]);
-    expect(mockRpc).toHaveBeenCalledWith("/props?model=test");
-  });
-
-  it("should detect text-only capability when modalities.vision is false", async () => {
-    mockRpc.mockResolvedValueOnce({ modalities: { vision: false } });
+  it("should detect image capability when multimodal is in capabilities", async () => {
+    mockRpc.mockResolvedValueOnce({
+      models: [{ id: "test", capabilities: ["multimodal"] }],
+    });
 
     const model = createModel();
     const capabilities = await model.getCapabilities();
 
-    expect(capabilities).toEqual(["text"]);
+    expect(capabilities).toEqual(["text", "image"]);
+    expect(mockRpc).toHaveBeenCalledWith("/models");
   });
 
-  it("should return text when /props call fails", async () => {
-    mockRpc.mockRejectedValueOnce(new Error("Connection refused"));
+  it("should detect text-only capability when multimodal is not in capabilities", async () => {
+    mockRpc.mockResolvedValueOnce({
+      models: [{ id: "test", capabilities: [] }],
+    });
 
     const model = createModel();
     const capabilities = await model.getCapabilities();
@@ -71,7 +65,9 @@ describe("SingleModel getStatus", () => {
     const status = await model.getStatus();
 
     expect(status).toBe(Status.LOADED);
-    expect(mockRpc).toHaveBeenCalledWith(`/props?model=${model.id}`);
+    expect(mockRpc).toHaveBeenCalledWith(
+      `/props?model=${model.id}&autoload=false`,
+    );
   });
 
   it("should return SLEEPING when is_sleeping is true", async () => {
@@ -85,24 +81,15 @@ describe("SingleModel getStatus", () => {
 });
 
 describe("SingleModel getContextSize", () => {
-  it("should return n_ctx from /props endpoint default_generation_settings", async () => {
+  it("should return n_ctx from /models endpoint meta", async () => {
     mockRpc.mockResolvedValueOnce({
-      default_generation_settings: { n_ctx: 8192 },
+      data: [{ id: "test", meta: { n_ctx: 8192 } }],
     });
 
     const model = createModel();
     const ctxSize = await model.getContextSize();
 
     expect(ctxSize).toBe(8192);
-    expect(mockRpc).toHaveBeenCalledWith("/props?model=test");
-  });
-
-  it("should return DEFAULT_CTX when /props fails", async () => {
-    mockRpc.mockRejectedValueOnce(new Error("Connection refused"));
-
-    const model = createModel();
-    const ctxSize = await model.getContextSize();
-
-    expect(ctxSize).toBe(DEFAULT_CTX);
+    expect(mockRpc).toHaveBeenCalledWith("/models");
   });
 });

package/src/tools/provider.ts

@@ -1,28 +0,0 @@
-import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
-import { PROVIDER_ID, PROVIDER_NAME } from "../constants";
-import type { BaseModel } from "../models/baseModel";
-import { resolveApiKey, resolveUrl } from "./resolver";
-import { listModels } from "./retriever";
-
-/**
- * Registers the Llama.cpp provider and returns the fetched models.
- *
- * @param pi The Pi extension API
- * @returns The list of models fetched from the server
- */
-export const registerLlamaCppProvider = async (
-  pi: ExtensionAPI,
-): Promise<BaseModel[]> => {
-  const baseUrl = `${await resolveUrl(process.cwd())}/v1`;
-  const models = await listModels();
-
-  pi.registerProvider(PROVIDER_ID, {
-    name: PROVIDER_NAME,
-    baseUrl,
-    api: "openai-completions",
-    apiKey: await resolveApiKey(),
-    models: await Promise.all(models.map((m) => m.toProviderConfig())),
-  });
-
-  return models;
-};