pi-llama-cpp 0.10.0 → 0.11.0

package/tests/events.test.ts

@@ -1,59 +1,29 @@
-import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 import { THINKING_BUDGETS } from "../src/constants";
 import { Status } from "../src/enums/status";
-import { createMockModel, createMockServer } from "./mocks";
-
-// Mock settings — reactToModelSelect and autoloadOnMessage
-const mockSettings = {
-  resolveReactToModelSelect: vi.fn(() => true),
-  resolveAutoloadOnMessage: vi.fn(() => false),
-  resolveThinkingLevel: vi.fn(() => "medium"),
-  resolveThinkingBudgets: vi.fn(() => ({ ...THINKING_BUDGETS })),
-};
-
-// Wire resolveThinkingBudgets to use the SettingsManager mock when set
-mockSettings.resolveThinkingBudgets.mockImplementation(() => {
-  const userBudgets = mockSettingsManager.getThinkingBudgets();
-  if (userBudgets) {
-    return { ...THINKING_BUDGETS, ...userBudgets };
-  }
-  return { ...THINKING_BUDGETS };
-});
-
-// 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,
-    },
-  };
-});
-
-vi.mock("../src/managers/settings", () => ({
-  settings: mockSettings,
-}));
-
-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;
-});
+import { EventManager } from "../src/managers/events";
+import { ServerManager } from "../src/managers/server";
+import type { Server } from "../src/server";
+import { createMockModel, createMockServer, makeSettingsStub } from "./mocks";
+
+/**
+ * Injected settings stub (EventManager and — in the live-list test — the
+ * real ServerManager). Recreated in `beforeEach` so per-case overrides
+ * (budgets, autoload, …) cannot leak between tests.
+ */
+let settingsStub = makeSettingsStub();
+
+/**
+ * Builds an EventManager on a ServerManager stub exposing `servers`.
+ * (The real ServerManager is exercised in the live-list test below.)
+ */
+const createEventManager = (...servers: Server[]) =>
+  new EventManager({ servers } as unknown as ServerManager, settingsStub);
 
 beforeEach(() => {
   vi.restoreAllMocks();
   EventManager.resetInflightModel();
-  mockSettingsManager.getDefaultThinkingLevel.mockReturnValue("medium");
-  mockSettingsManager.getThinkingBudgets.mockReturnValue(undefined);
+  settingsStub = makeSettingsStub();
 });
 
 const createPayload = (modelId: string) => ({
@@ -97,12 +67,14 @@ describe("EventManager.onBeforeProviderRequest", () => {
     ])(
       'level "$level" should return $expected',
       async ({ level, expected }) => {
-        mockSettingsManager.getDefaultThinkingLevel.mockReturnValue(level);
+        vi.mocked(settingsStub.resolveThinkingLevel).mockReturnValue(
+          level as any,
+        );
 
         const server = createMockServer({
           models: ["model-a"].map((id) => createMockModel(id)),
         });
-        const eventManager = new EventManager([server]);
+        const eventManager = createEventManager(server);
         const event = { payload: createPayload("model-a") };
 
         const ctx = createMockCtx(level as any);
@@ -117,12 +89,12 @@ describe("EventManager.onBeforeProviderRequest", () => {
     );
 
     it("should preserve original payload fields alongside new ones", async () => {
-      mockSettingsManager.getDefaultThinkingLevel.mockReturnValue("low");
+      vi.mocked(settingsStub.resolveThinkingLevel).mockReturnValue("low");
 
       const server = createMockServer({
         models: ["model-b"].map((id) => createMockModel(id)),
       });
-      const eventManager = new EventManager([server]);
+      const eventManager = createEventManager(server);
       const event = {
         payload: {
           model: "model-b",
@@ -148,7 +120,7 @@ describe("EventManager.onBeforeProviderRequest", () => {
       const server = createMockServer({
         models: ["model-a"].map((id) => createMockModel(id)),
       });
-      const eventManager = new EventManager([server]);
+      const eventManager = createEventManager(server);
       const event = { payload: createNonLlamaPayload() };
 
       const ctx = createMockCtx();
@@ -166,7 +138,7 @@ describe("EventManager.onBeforeProviderRequest", () => {
       const server = createMockServer({
         models: ["model-a"].map((id) => createMockModel(id)),
       });
-      const eventManager = new EventManager([server]);
+      const eventManager = createEventManager(server);
       const event = { payload: { messages: [] } };
 
       const ctx = createMockCtx();
@@ -181,13 +153,16 @@ describe("EventManager.onBeforeProviderRequest", () => {
 
   describe("user-defined budget overrides", () => {
     it("should use user-defined budgets instead of defaults", async () => {
-      mockSettingsManager.getDefaultThinkingLevel.mockReturnValue("low");
-      mockSettingsManager.getThinkingBudgets.mockReturnValue({ low: 4096 });
+      vi.mocked(settingsStub.resolveThinkingLevel).mockReturnValue("low");
+      vi.mocked(settingsStub.resolveThinkingBudgets).mockReturnValue({
+        ...THINKING_BUDGETS,
+        low: 4096,
+      });
 
       const server = createMockServer({
         models: ["model-a"].map((id) => createMockModel(id)),
       });
-      const eventManager = new EventManager([server]);
+      const eventManager = createEventManager(server);
       const event = { payload: createPayload("model-a") };
 
       const ctx = createMockCtx("low");
@@ -200,13 +175,16 @@ describe("EventManager.onBeforeProviderRequest", () => {
     });
 
     it("should merge user budgets with defaults (partial override)", async () => {
-      mockSettingsManager.getDefaultThinkingLevel.mockReturnValue("medium");
-      mockSettingsManager.getThinkingBudgets.mockReturnValue({ low: 4096 });
+      vi.mocked(settingsStub.resolveThinkingLevel).mockReturnValue("medium");
+      vi.mocked(settingsStub.resolveThinkingBudgets).mockReturnValue({
+        ...THINKING_BUDGETS,
+        low: 4096,
+      });
 
       const server = createMockServer({
         models: ["model-a"].map((id) => createMockModel(id)),
       });
-      const eventManager = new EventManager([server]);
+      const eventManager = createEventManager(server);
       const event = { payload: createPayload("model-a") };
 
       const ctx = createMockCtx("medium");
@@ -224,8 +202,9 @@ describe("EventManager.onBeforeProviderRequest", () => {
 
   describe("edge cases", () => {
     it("should ignore invalid keys in user budgets (they are silently dropped)", async () => {
-      mockSettingsManager.getDefaultThinkingLevel.mockReturnValue("medium");
-      mockSettingsManager.getThinkingBudgets.mockReturnValue({
+      vi.mocked(settingsStub.resolveThinkingLevel).mockReturnValue("medium");
+      vi.mocked(settingsStub.resolveThinkingBudgets).mockReturnValue({
+        ...THINKING_BUDGETS,
         foo: 999,
         bar: 123,
       } as any);
@@ -233,7 +212,7 @@ describe("EventManager.onBeforeProviderRequest", () => {
       const server = createMockServer({
         models: ["model-a"].map((id) => createMockModel(id)),
       });
-      const eventManager = new EventManager([server]);
+      const eventManager = createEventManager(server);
       const event = { payload: createPayload("model-a") };
 
       const ctx = createMockCtx("medium");
@@ -247,15 +226,16 @@ describe("EventManager.onBeforeProviderRequest", () => {
     });
 
     it("should not allow overriding 'off' — thinking stays disabled", async () => {
-      mockSettingsManager.getDefaultThinkingLevel.mockReturnValue("off");
-      mockSettingsManager.getThinkingBudgets.mockReturnValue({
+      vi.mocked(settingsStub.resolveThinkingLevel).mockReturnValue("off");
+      vi.mocked(settingsStub.resolveThinkingBudgets).mockReturnValue({
+        ...THINKING_BUDGETS,
         off: 99999,
-      } as any);
+      });
 
       const server = createMockServer({
         models: ["model-a"].map((id) => createMockModel(id)),
       });
-      const eventManager = new EventManager([server]);
+      const eventManager = createEventManager(server);
       const event = { payload: createPayload("model-a") };
 
       const ctx = createMockCtx("off");
@@ -271,15 +251,16 @@ describe("EventManager.onBeforeProviderRequest", () => {
     });
 
     it("should not inject budget for 'max' — unlimited reasoning", async () => {
-      mockSettingsManager.getDefaultThinkingLevel.mockReturnValue("max");
-      mockSettingsManager.getThinkingBudgets.mockReturnValue({
+      vi.mocked(settingsStub.resolveThinkingLevel).mockReturnValue("max");
+      vi.mocked(settingsStub.resolveThinkingBudgets).mockReturnValue({
+        ...THINKING_BUDGETS,
         max: 1,
-      } as any);
+      });
 
       const server = createMockServer({
         models: ["model-a"].map((id) => createMockModel(id)),
       });
-      const eventManager = new EventManager([server]);
+      const eventManager = createEventManager(server);
       const event = { payload: createPayload("model-a") };
 
       const ctx = createMockCtx("max");
@@ -293,13 +274,16 @@ describe("EventManager.onBeforeProviderRequest", () => {
     });
 
     it("should handle empty user budgets gracefully", async () => {
-      mockSettingsManager.getDefaultThinkingLevel.mockReturnValue("high");
-      mockSettingsManager.getThinkingBudgets.mockReturnValue({});
+      vi.mocked(settingsStub.resolveThinkingLevel).mockReturnValue("high");
+      // No user overrides — the resolver returns the default budget table
+      vi.mocked(settingsStub.resolveThinkingBudgets).mockReturnValue({
+        ...THINKING_BUDGETS,
+      });
 
       const server = createMockServer({
         models: ["model-a"].map((id) => createMockModel(id)),
       });
-      const eventManager = new EventManager([server]);
+      const eventManager = createEventManager(server);
       const event = { payload: createPayload("model-a") };
 
       const ctx = createMockCtx("high");
@@ -315,14 +299,14 @@ describe("EventManager.onBeforeProviderRequest", () => {
 
 describe("EventManager.onModelSelect", () => {
   beforeEach(() => {
-    mockSettings.resolveReactToModelSelect.mockReturnValue(true);
+    vi.mocked(settingsStub.resolveReactToModelSelect).mockReturnValue(true);
   });
 
   it("should load the model when reactToModelSelect is true", async () => {
     const server = createMockServer({
       models: ["model-a"].map((id) => createMockModel(id)),
     });
-    const eventManager = new EventManager([server]);
+    const eventManager = createEventManager(server);
     const ctx = createMockCtx();
 
     const event = {
@@ -335,12 +319,12 @@ describe("EventManager.onModelSelect", () => {
   });
 
   it("should return early when reactToModelSelect is false", async () => {
-    mockSettings.resolveReactToModelSelect.mockReturnValue(false);
+    vi.mocked(settingsStub.resolveReactToModelSelect).mockReturnValue(false);
 
     const server = createMockServer({
       models: ["model-a"].map((id) => createMockModel(id)),
     });
-    const eventManager = new EventManager([server]);
+    const eventManager = createEventManager(server);
     const ctx = createMockCtx();
 
     const event = {
@@ -355,7 +339,7 @@ describe("EventManager.onModelSelect", () => {
 
 describe("EventManager.autoLoadIfNeeded", () => {
   it("should load the model when autoloadOnMessage is true and model is UNLOADED", async () => {
-    mockSettings.resolveAutoloadOnMessage.mockReturnValue(true);
+    vi.mocked(settingsStub.resolveAutoloadOnMessage).mockReturnValue(true);
 
     const server = createMockServer({
       models: [
@@ -365,7 +349,7 @@ describe("EventManager.autoLoadIfNeeded", () => {
         }),
       ],
     });
-    const eventManager = new EventManager([server]);
+    const eventManager = createEventManager(server);
     const model = server.models[0];
 
     await (eventManager as any).autoLoadIfNeeded(model);
@@ -374,7 +358,7 @@ describe("EventManager.autoLoadIfNeeded", () => {
   });
 
   it("should not load the model when autoloadOnMessage is false", async () => {
-    mockSettings.resolveAutoloadOnMessage.mockReturnValue(false);
+    vi.mocked(settingsStub.resolveAutoloadOnMessage).mockReturnValue(false);
 
     const server = createMockServer({
       models: [
@@ -384,7 +368,7 @@ describe("EventManager.autoLoadIfNeeded", () => {
         }),
       ],
     });
-    const eventManager = new EventManager([server]);
+    const eventManager = createEventManager(server);
     const model = server.models[0];
 
     await (eventManager as any).autoLoadIfNeeded(model);
@@ -393,7 +377,7 @@ describe("EventManager.autoLoadIfNeeded", () => {
   });
 
   it("should not load the model when model is already LOADED", async () => {
-    mockSettings.resolveAutoloadOnMessage.mockReturnValue(true);
+    vi.mocked(settingsStub.resolveAutoloadOnMessage).mockReturnValue(true);
 
     const server = createMockServer({
       models: [
@@ -403,7 +387,7 @@ describe("EventManager.autoLoadIfNeeded", () => {
         }),
       ],
     });
-    const eventManager = new EventManager([server]);
+    const eventManager = createEventManager(server);
     const model = server.models[0];
 
     await (eventManager as any).autoLoadIfNeeded(model);
@@ -412,7 +396,7 @@ describe("EventManager.autoLoadIfNeeded", () => {
   });
 
   it("should not load the model when model is SLEEPING", async () => {
-    mockSettings.resolveAutoloadOnMessage.mockReturnValue(true);
+    vi.mocked(settingsStub.resolveAutoloadOnMessage).mockReturnValue(true);
 
     const server = createMockServer({
       models: [
@@ -422,7 +406,7 @@ describe("EventManager.autoLoadIfNeeded", () => {
         }),
       ],
     });
-    const eventManager = new EventManager([server]);
+    const eventManager = createEventManager(server);
     const model = server.models[0];
 
     await (eventManager as any).autoLoadIfNeeded(model);
@@ -430,3 +414,48 @@ describe("EventManager.autoLoadIfNeeded", () => {
     expect(model.load).not.toHaveBeenCalled();
   });
 });
+
+describe("EventManager with a live ServerManager", () => {
+  it("should observe servers added after construction", async () => {
+    const serverA = createMockServer({
+      models: [createMockModel("model-a")],
+    });
+    const serverB = createMockServer({
+      baseUrl: "http://127.0.0.1:8081",
+      models: [
+        createMockModel("model-b", { serverUrl: "http://127.0.0.1:8081" }),
+      ],
+    });
+
+    vi.mocked(settingsStub.resolveServers).mockReturnValue([serverA]);
+    const serverManager = new ServerManager(settingsStub);
+    const mockPi = { registerProvider: vi.fn(), unregisterProvider: vi.fn() };
+    await serverManager.update(mockPi as any);
+
+    const eventManager = new EventManager(serverManager, settingsStub);
+
+    // Second scan adds serverB — no manager re-construction
+    vi.mocked(settingsStub.resolveServers).mockReturnValue([serverA, serverB]);
+    await serverManager.update(mockPi as any);
+
+    // onBeforeProviderRequest sees the new server's models
+    const ctx = createMockCtx("medium");
+    const result = (await eventManager.onBeforeProviderRequest(
+      { payload: createPayload("model-b") } as any,
+      ctx,
+    )) as Record<string, unknown>;
+    expect(result.thinking_budget_tokens).toBe(THINKING_BUDGETS.medium);
+
+    // onModelSelect sees the new server's models too
+    vi.mocked(settingsStub.resolveReactToModelSelect).mockReturnValue(true);
+    const selectCtx = createMockCtx();
+    await eventManager.onModelSelect(
+      { model: { provider: serverB.providerId, id: "model-b" } } as any,
+      selectCtx,
+    );
+    expect(selectCtx.ui.notify).toHaveBeenCalledWith(
+      "Loading model-b...",
+      "info",
+    );
+  });
+});

package/tests/mocks.ts

@@ -1,50 +1,160 @@
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
 import { vi } from "vitest";
+import type { ApiClient } from "../src/api/client";
+import {
+  API_KEY_PLACEHOLDER,
+  AUTOLOAD_ON_MESSAGE,
+  POLLING_TIMEOUT,
+  REACT_TO_MODEL_SELECT,
+  SERVER_TIMEOUT,
+  SORT_BY,
+  THINKING_BUDGETS,
+} from "../src/constants";
 import { Mode } from "../src/enums/mode";
-import { ServerStatus } from "../src/enums/serverStatus";
 import { Status } from "../src/enums/status";
+import type { LlamaServer } from "../src/interfaces/settings";
+import type { LlamaSettingsManager } from "../src/managers/settings";
 import { BaseModel } from "../src/models/baseModel";
 import { Server } from "../src/server";
+import type { SSEManager } from "../src/sse/manager";
 
 /** Shared mock RPC — each test configures it */
 export const mockRpc = vi.fn();
 
-/** Default mock server that assumes everything works */
+/**
+ * Hand-made stub of `LlamaSettingsManager` for constructor injection —
+ * replaces whole-module `vi.mock` blocks and per-case `vi.spyOn` on the
+ * real singleton. Defaults give every consumer-used member a benign value;
+ * a case that cares passes just that member as an override. The single
+ * `as unknown as` cast lives here so test bodies stay cast-free.
+ */
+export const makeSettingsStub = (
+  overrides: Partial<LlamaSettingsManager> = {},
+): LlamaSettingsManager =>
+  ({
+    resolveTimeouts: vi.fn(() => ({
+      pollingTimeout: POLLING_TIMEOUT,
+      serverTimeout: SERVER_TIMEOUT,
+    })),
+    resolveServers: vi.fn((): Server[] => []),
+    resolveSortBy: vi.fn(() => SORT_BY),
+    resolveApiKey: vi.fn(() => API_KEY_PLACEHOLDER),
+    resolveReactToModelSelect: vi.fn(() => REACT_TO_MODEL_SELECT),
+    resolveAutoloadOnMessage: vi.fn(() => AUTOLOAD_ON_MESSAGE),
+    resolveThinkingLevel: vi.fn(() => undefined),
+    resolveThinkingBudgets: vi.fn(() => ({ ...THINKING_BUDGETS })),
+    llamaServers: [] as LlamaServer[],
+    takeWarnings: vi.fn((): string[] => []),
+    setLlamaSetting: vi.fn(() => Promise.resolve()),
+    ...overrides,
+  }) as unknown as LlamaSettingsManager;
+
+/**
+ * Fake `ApiClient`/`SSEManager` pair for injection via `ServerDeps`. The
+ * ApiClient delegates every request to the shared `mockRpc`, keyed by the
+ * endpoint path (and body for POSTs) — the same contract the old hand-rolled
+ * mock server used, so tests keep configuring responses on `mockRpc`. The
+ * SSEManager is inert: `disconnect` is a spy and `probeSSE` resolves false,
+ * steering load flows onto the HTTP polling path.
+ */
+export const createFakeClients = (): {
+  apiClient: ApiClient;
+  sseManager: SSEManager;
+} => {
+  const apiClient = {
+    get: (endpoint: string) => mockRpc(endpoint),
+    post: (endpoint: string, body?: Record<string, unknown>) =>
+      mockRpc(endpoint, body),
+    clearCache: vi.fn(),
+  } as unknown as ApiClient;
+  const sseManager = {
+    disconnect: vi.fn(),
+    probeSSE: vi.fn(async () => false),
+    subscribeToStatus: vi.fn(),
+    subscribeToProgress: vi.fn(() => () => {}),
+  } as unknown as SSEManager;
+  return { apiClient, sseManager };
+};
+
+/**
+ * Overrides for {@link createMockServer}. Identity (`baseUrl`/`customId`/
+ * `customName`), `apiKey` and the timeout getters are constructor wiring
+ * (ServerOptions + a synthesized settings stub) — a real Server has no
+ * assignable properties by those names. Everything else is shadowed onto the
+ * instance as-is (e.g. a custom `initialize`).
+ */
+export type MockServerOverrides = Partial<
+  Omit<
+    Server,
+    | "baseUrl"
+    | "providerId"
+    | "providerName"
+    | "pollingTimeout"
+    | "serverTimeout"
+    | "sseManager"
+    | "models"
+  >
+> & {
+  baseUrl?: string;
+  customId?: string;
+  customName?: string;
+  apiKey?: string;
+  models?: BaseModel[];
+  pollingTimeout?: number;
+  serverTimeout?: number;
+};
+
+/**
+ * Builds a real `Server` wired with fake collaborators (`createFakeClients`),
+ * replacing the old hand-rolled `Partial<Server>` mock whose `isReady` /
+ * `initialize` re-implementations drifted from the real ones.
+ *
+ * Seeded `models` land directly in `server.models`; `initialize` is stubbed
+ * to a no-op so a `ServerManager.update()` scan cannot wipe the seeds with
+ * fetched data — pass an `initialize` override to restore the real flow.
+ */
 export const createMockServer = (
-  overrides: Partial<Server & { apiKey?: string }> = {},
+  overrides: MockServerOverrides = {},
 ): Server => {
-  const models: BaseModel[] = [];
-  const server: Partial<Server> = {
-    baseUrl: "http://127.0.0.1:8080",
+  const {
+    baseUrl,
+    customId,
+    customName,
+    apiKey,
     models,
-    getApiKey: () => overrides.apiKey ?? "",
-    fetchModels: () => mockRpc("/v1/models"),
-    fetchModelProps: (modelId: string) =>
-      mockRpc(`/props?model=${modelId}&autoload=false`),
-    fetchServerHealth: () => mockRpc("/health"),
-    fetchServerProps: () => mockRpc("/props?autoload=false"),
-    postRequest: (resource: "load" | "unload", model: string) =>
-      mockRpc(`/models/${resource}`, { model }),
-    isReady: async (timeout: number) => {
-      try {
-        const r = await mockRpc("/health");
-        return r.status === "ok"
-          ? ServerStatus.READY
-          : ServerStatus.UNREACHABLE;
-      } catch {
-        return ServerStatus.UNREACHABLE;
-      }
+    pollingTimeout,
+    serverTimeout,
+    initialize,
+    ...members
+  } = overrides;
+
+  const settings = makeSettingsStub({
+    ...(apiKey !== undefined && { resolveApiKey: vi.fn(() => apiKey) }),
+    ...((pollingTimeout !== undefined || serverTimeout !== undefined) && {
+      resolveTimeouts: vi.fn(() => ({
+        pollingTimeout: pollingTimeout ?? POLLING_TIMEOUT,
+        serverTimeout: serverTimeout ?? SERVER_TIMEOUT,
+      })),
+    }),
+  });
+
+  const { apiClient, sseManager } = createFakeClients();
+  const server = new Server(
+    settings,
+    {
+      baseUrl: baseUrl ?? "http://127.0.0.1:8080",
+      customId,
+      customName,
     },
-    initialize: async () => {
-      const { data } = (await mockRpc("/v1/models")) as {
-        data: BaseModel[];
-      };
-      models.length = 0;
-      models.push(...(data ?? []));
+    {
+      createApiClient: () => apiClient,
+      createSSEManager: () => sseManager,
     },
-    ...overrides,
-  };
-  return server as Server;
+  );
+  if (models) server.models.push(...models);
+  return Object.assign(server, members, {
+    initialize: initialize ?? (async () => {}),
+  });
 };
 
 /** Helper to create a mock BaseModel */
@@ -72,10 +182,12 @@ export const createMockModel = (
 export const createMockCtx = (
   selectFn: (prompt: string, options: string[]) => string | null,
 ) => ({
+  mode: "tui",
   cwd: "/tmp/test",
   ui: {
     select: vi.fn(selectFn),
     notify: vi.fn(),
+    custom: vi.fn(),
     theme: {
       fg: (color: string, text: string) => text,
     },
@@ -89,6 +201,7 @@ export const createMockCtx = (
 export const createMockPi = () => ({
   setModel: vi.fn(),
   registerProvider: vi.fn(),
+  unregisterProvider: vi.fn(),
 });
 
 /** Create a mock Pi context for EventManager */