pi-llama-cpp 0.5.1 → 0.7.0

package/tests/commandManager.test.ts

@@ -1,153 +1,202 @@
 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(),
-};
+import { Action } from "../src/enums/action";
+import { CommandManager } from "../src/managers/command";
+import { ServerManager } from "../src/managers/server";
+import {
+  createMockCtx,
+  createMockModel,
+  createMockPi,
+  createMockServer,
+  mockRpc,
+} from "./mocks";
 
 beforeEach(() => {
   vi.clearAllMocks();
-  (resolveUrl as any).mockResolvedValue("http://127.0.0.1:8080");
-  (resolveApiKey as any).mockResolvedValue("test-key");
+  mockRpc.mockResolvedValue({ data: [] });
 });
 
 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: [],
-    });
+  let serverManager: ServerManager;
+  let commandManager: CommandManager;
+  let mockPi: ReturnType<typeof createMockPi>;
+
+  beforeEach(() => {
+    mockPi = createMockPi();
+    serverManager = new ServerManager([]);
+    commandManager = new CommandManager(serverManager);
   });
 
-  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 }],
+  describe("getArgumentCompletions", () => {
+    it("should provide completions for /models", () => {
+      const completions = commandManager.getArgumentCompletions("");
+      expect(completions).toHaveLength(2);
+      expect(completions?.map((c) => c.value)).toEqual(["info", "unload"]);
     });
-  });
 
-  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);
+    it("should filter completions by prefix", () => {
+      const completions = commandManager.getArgumentCompletions("u");
+      expect(completions).toHaveLength(1);
+      expect(completions?.[0].value).toBe("unload");
+    });
 
-    expect(mockPi.registerProvider).not.toHaveBeenCalled();
+    it("should return null when no completions match", () => {
+      const completions = commandManager.getArgumentCompletions("zzz");
+      expect(completions).toBeNull();
+    });
   });
 
-  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");
-    // Called once in initialize() and once in run() to refresh the model list
-    expect(listModels).toHaveBeenCalledTimes(2);
-  });
+  describe("handleCommand", () => {
+    it("should unload all models when args is 'unload'", async () => {
+      const model1 = createMockModel("model-1");
+      const model2 = createMockModel("model-2");
+      const server = createMockServer({
+        baseUrl: "http://127.0.0.1:8080",
+        models: [model1, model2],
+      });
+      serverManager = new ServerManager([server] as any);
+      commandManager = new CommandManager(serverManager);
+
+      const ctx = {
+        ui: {
+          notify: vi.fn(),
+          theme: { fg: (_: string, text: string) => text },
+        },
+      } as any;
+
+      await commandManager.handleCommand("unload", ctx, mockPi as any);
+
+      expect(model1.unload).toHaveBeenCalled();
+      expect(model2.unload).toHaveBeenCalled();
+      expect(ctx.ui.notify).toHaveBeenCalledWith(
+        "Unloaded all Llama.cpp models",
+        "info",
+      );
+    });
 
-  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 show model info when args is 'info'", async () => {
+      const model1 = createMockModel("model-1");
+      const model2 = createMockModel("model-2");
+      const server = createMockServer({
+        baseUrl: "http://127.0.0.1:8080",
+        models: [model1, model2],
+      });
+      serverManager = new ServerManager([server] as any);
+      commandManager = new CommandManager(serverManager);
+
+      const ctx = {
+        ui: {
+          notify: vi.fn(),
+          theme: { fg: (_: string, text: string) => text },
+        },
+      } as any;
+
+      await commandManager.handleCommand("info", ctx, mockPi as any);
+
+      expect(model1.getInfo).toHaveBeenCalled();
+      expect(model2.getInfo).toHaveBeenCalled();
+    });
   });
 
-  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" }),
+  describe("/models interactive menu", () => {
+    const CHOICE = "model-a   [Server: http://127.0.0.1:8080]";
+
+    /**
+     * Helper to create a CommandManager with mock servers and models.
+     */
+    const createCommandManager = (
+      models: ReturnType<typeof createMockModel>[],
+    ) => {
+      const mockPi = createMockPi();
+      const servers = models.map((model) =>
+        createMockServer({
+          baseUrl: model.serverUrl,
+          models: [model],
+        }),
+      );
+      const serverManager = new ServerManager(servers as any);
+      return {
+        commandManager: new CommandManager(serverManager),
+        serverManager,
+        mockPi,
+      };
     };
-    (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();
+
+    it("should return early on cancel (null model selection)", async () => {
+      const models = [createMockModel("model-a")];
+      const { commandManager, mockPi } = createCommandManager(models);
+      const ctx = createMockCtx(() => null);
+
+      await commandManager.handleCommand("", ctx as any, mockPi as any);
+
+      expect(ctx.ui.notify).not.toHaveBeenCalled();
+    });
+
+    it("should show info when INFO action is selected", async () => {
+      const model = createMockModel("model-a");
+      const { commandManager, mockPi } = createCommandManager([model]);
+      let selectCallCount = 0;
+      const ctx = createMockCtx(() => {
+        selectCallCount++;
+        if (selectCallCount === 1) return CHOICE;
+        return Action.INFO;
+      });
+
+      await commandManager.handleCommand("", ctx as any, mockPi as any);
+
+      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 { commandManager, mockPi } = createCommandManager([model]);
+      let selectCallCount = 0;
+      const ctx = createMockCtx(() => {
+        selectCallCount++;
+        if (selectCallCount === 1) return CHOICE;
+        return Action.UNLOAD;
+      });
+
+      await commandManager.handleCommand("", ctx as any, mockPi as any);
+
+      expect(model.unload).toHaveBeenCalled();
+      expect(ctx.ui.notify).toHaveBeenCalledWith("Unloaded model-a", "info");
+    });
+
+    it("should switch model when SWITCH action is selected", async () => {
+      const model = createMockModel("model-a");
+      const { commandManager, mockPi } = createCommandManager([model]);
+      let selectCallCount = 0;
+      const ctx = createMockCtx(() => {
+        selectCallCount++;
+        if (selectCallCount === 1) return CHOICE;
+        return Action.SWITCH;
+      });
+
+      await commandManager.handleCommand("", ctx as any, mockPi as any);
+
+      expect(mockPi.setModel).toHaveBeenCalled();
+      expect(ctx.ui.notify).toHaveBeenCalledWith("Model model-a ready", "info");
+    });
+
+    it("should loop back to model selection when action is cancelled", async () => {
+      const model = createMockModel("model-a");
+      const { commandManager, mockPi } = createCommandManager([model]);
+
+      let selectCallCount = 0;
+      const ctx = createMockCtx(() => {
+        selectCallCount++;
+        // 1st: select model-a, 2nd: cancel action, 3rd: cancel model => exit
+        if (selectCallCount === 1) return CHOICE;
+        return null;
+      });
+
+      await commandManager.handleCommand("", ctx as any, mockPi as any);
+
+      expect(ctx.ui.select).toHaveBeenCalledTimes(3);
+      expect(ctx.ui.notify).not.toHaveBeenCalled();
+    });
   });
 });

package/tests/events.test.ts

@@ -0,0 +1,256 @@
+import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
+import { DEFAULT_THINKING_BUDGETS } from "../src/constants";
+import { createMockModel, createMockServer } from "./mocks";
+
+// Create a mutable mock object shared across tests
+const mockSettingsManager = {
+  getDefaultThinkingLevel: vi.fn(() => "medium"),
+  getThinkingBudgets: vi.fn<() => Record<string, number> | undefined>(),
+};
+
+vi.mock("@earendil-works/pi-coding-agent", async (importOriginal) => {
+  const actual =
+    await importOriginal<typeof import("@earendil-works/pi-coding-agent")>();
+  return {
+    ...actual,
+    SettingsManager: {
+      create: () => mockSettingsManager,
+    },
+  };
+});
+
+let EventManager: typeof import("../src/managers/events").EventManager;
+
+beforeAll(async () => {
+  const mod = await vi.importActual("../src/managers/events");
+  EventManager =
+    mod.EventManager as typeof import("../src/managers/events").EventManager;
+});
+
+beforeEach(() => {
+  vi.restoreAllMocks();
+  EventManager.resetInflightModel();
+  mockSettingsManager.getDefaultThinkingLevel.mockReturnValue("medium");
+  mockSettingsManager.getThinkingBudgets.mockReturnValue(undefined);
+});
+
+const createPayload = (modelId: string) => ({
+  model: modelId,
+  messages: [{ role: "user", content: "hello" }],
+});
+
+const createNonLlamaPayload = () => ({
+  model: "gpt-4",
+  messages: [{ role: "user", content: "hello" }],
+});
+
+describe("EventManager.onBeforeProviderRequest", () => {
+  describe("normal usage — each thinking level", () => {
+    it.each([
+      {
+        level: "off",
+        expected: { chat_template_kwargs: { enable_thinking: false } },
+      },
+      { level: "minimal", expected: { thinking_budget_tokens: 1024 } },
+      { level: "low", expected: { thinking_budget_tokens: 2048 } },
+      { level: "medium", expected: { thinking_budget_tokens: 8192 } },
+      { level: "high", expected: { thinking_budget_tokens: 16384 } },
+      { level: "xhigh", expected: {} },
+    ])(
+      'level "$level" should return $expected',
+      async ({ level, expected }) => {
+        mockSettingsManager.getDefaultThinkingLevel.mockReturnValue(level);
+
+        const server = createMockServer({
+          models: ["model-a"].map((id) => createMockModel(id)),
+        });
+        const eventManager = new EventManager([server]);
+        const event = { payload: createPayload("model-a") };
+
+        const result = (await eventManager.onBeforeProviderRequest(
+          event as any,
+        )) as Record<string, unknown>;
+
+        expect(result.model).toBe("model-a");
+        expect(result).toMatchObject(expected);
+      },
+    );
+
+    it("should preserve original payload fields alongside new ones", async () => {
+      mockSettingsManager.getDefaultThinkingLevel.mockReturnValue("low");
+
+      const server = createMockServer({
+        models: ["model-b"].map((id) => createMockModel(id)),
+      });
+      const eventManager = new EventManager([server]);
+      const event = {
+        payload: {
+          model: "model-b",
+          messages: [{ role: "user", content: "test" }],
+          temperature: 0.7,
+        },
+      };
+
+      const result = (await eventManager.onBeforeProviderRequest(
+        event as any,
+      )) as Record<string, unknown>;
+
+      expect(result.messages).toEqual([{ role: "user", content: "test" }]);
+      expect(result.temperature).toBe(0.7);
+      expect(result.thinking_budget_tokens).toBe(DEFAULT_THINKING_BUDGETS.low);
+    });
+  });
+
+  describe("non-llama.cpp models", () => {
+    it("should return the payload unchanged for unknown models", async () => {
+      const server = createMockServer({
+        models: ["model-a"].map((id) => createMockModel(id)),
+      });
+      const eventManager = new EventManager([server]);
+      const event = { payload: createNonLlamaPayload() };
+
+      const result = await eventManager.onBeforeProviderRequest(event as any);
+
+      expect(result).toEqual(createNonLlamaPayload());
+    });
+  });
+
+  describe("missing model in payload", () => {
+    it("should return the payload unchanged when model is absent", async () => {
+      const server = createMockServer({
+        models: ["model-a"].map((id) => createMockModel(id)),
+      });
+      const eventManager = new EventManager([server]);
+      const event = { payload: { messages: [] } };
+
+      const result = await eventManager.onBeforeProviderRequest(event as any);
+
+      expect(result).toEqual({ messages: [] });
+    });
+  });
+
+  describe("user-defined budget overrides", () => {
+    it("should use user-defined budgets instead of defaults", async () => {
+      mockSettingsManager.getDefaultThinkingLevel.mockReturnValue("low");
+      mockSettingsManager.getThinkingBudgets.mockReturnValue({ low: 4096 });
+
+      const server = createMockServer({
+        models: ["model-a"].map((id) => createMockModel(id)),
+      });
+      const eventManager = new EventManager([server]);
+      const event = { payload: createPayload("model-a") };
+
+      const result = (await eventManager.onBeforeProviderRequest(
+        event as any,
+      )) as Record<string, unknown>;
+
+      expect(result.thinking_budget_tokens).toBe(4096);
+    });
+
+    it("should merge user budgets with defaults (partial override)", async () => {
+      mockSettingsManager.getDefaultThinkingLevel.mockReturnValue("medium");
+      mockSettingsManager.getThinkingBudgets.mockReturnValue({ low: 4096 });
+
+      const server = createMockServer({
+        models: ["model-a"].map((id) => createMockModel(id)),
+      });
+      const eventManager = new EventManager([server]);
+      const event = { payload: createPayload("model-a") };
+
+      const result = (await eventManager.onBeforeProviderRequest(
+        event as any,
+      )) as Record<string, unknown>;
+
+      // medium uses default since user only overrode low
+      expect(result.thinking_budget_tokens).toBe(
+        DEFAULT_THINKING_BUDGETS.medium,
+      );
+    });
+  });
+
+  // ─── Edge cases ─────────────────────────────────────────────────────
+
+  describe("edge cases", () => {
+    it("should ignore invalid keys in user budgets (they are silently dropped)", async () => {
+      mockSettingsManager.getDefaultThinkingLevel.mockReturnValue("medium");
+      mockSettingsManager.getThinkingBudgets.mockReturnValue({
+        foo: 999,
+        bar: 123,
+      } as any);
+
+      const server = createMockServer({
+        models: ["model-a"].map((id) => createMockModel(id)),
+      });
+      const eventManager = new EventManager([server]);
+      const event = { payload: createPayload("model-a") };
+
+      const result = (await eventManager.onBeforeProviderRequest(
+        event as any,
+      )) as Record<string, unknown>;
+
+      // Should fall back to default since "medium" is not in user budgets
+      expect(result.thinking_budget_tokens).toBe(
+        DEFAULT_THINKING_BUDGETS.medium,
+      );
+    });
+
+    it("should not allow overriding 'off' — thinking stays disabled", async () => {
+      mockSettingsManager.getDefaultThinkingLevel.mockReturnValue("off");
+      mockSettingsManager.getThinkingBudgets.mockReturnValue({
+        off: 99999,
+      } as any);
+
+      const server = createMockServer({
+        models: ["model-a"].map((id) => createMockModel(id)),
+      });
+      const eventManager = new EventManager([server]);
+      const event = { payload: createPayload("model-a") };
+
+      const result = (await eventManager.onBeforeProviderRequest(
+        event as any,
+      )) as Record<string, unknown>;
+
+      expect(result).toMatchObject({
+        chat_template_kwargs: { enable_thinking: false },
+      });
+      expect(result).not.toHaveProperty("thinking_budget_tokens");
+    });
+
+    it("should not allow overriding 'xhigh' — no budget is injected", async () => {
+      mockSettingsManager.getDefaultThinkingLevel.mockReturnValue("xhigh");
+      mockSettingsManager.getThinkingBudgets.mockReturnValue({
+        xhigh: 1,
+      } as any);
+
+      const server = createMockServer({
+        models: ["model-a"].map((id) => createMockModel(id)),
+      });
+      const eventManager = new EventManager([server]);
+      const event = { payload: createPayload("model-a") };
+
+      const result = (await eventManager.onBeforeProviderRequest(
+        event as any,
+      )) as Record<string, unknown>;
+
+      expect(result).toEqual(createPayload("model-a"));
+      expect(result).not.toHaveProperty("thinking_budget_tokens");
+    });
+
+    it("should handle empty user budgets gracefully", async () => {
+      mockSettingsManager.getDefaultThinkingLevel.mockReturnValue("high");
+      mockSettingsManager.getThinkingBudgets.mockReturnValue({});
+
+      const server = createMockServer({
+        models: ["model-a"].map((id) => createMockModel(id)),
+      });
+      const eventManager = new EventManager([server]);
+      const event = { payload: createPayload("model-a") };
+
+      const result = (await eventManager.onBeforeProviderRequest(
+        event as any,
+      )) as Record<string, unknown>;
+
+      expect(result.thinking_budget_tokens).toBe(DEFAULT_THINKING_BUDGETS.high);
+    });
+  });
+});