pi-llama-cpp 0.5.0 → 0.6.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,270 +0,0 @@
1
- import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
2
-
3
- // Set up fake timers before any imports so setTimeout is mocked globally
4
- vi.useFakeTimers();
5
-
6
- import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
7
- import {
8
- modelsCommand,
9
- onSessionBeforeSwitch,
10
- resetInflightModel,
11
- } from "../src/commands/models";
12
- import { Action } from "../src/enums/action";
13
- import { Mode } from "../src/enums/mode";
14
- import { Status } from "../src/enums/status";
15
- import { BaseModel } from "../src/models/baseModel";
16
-
17
- // Mock the retriever module
18
- vi.mock("../src/tools/retriever", () => ({
19
- rpc: vi.fn(),
20
- isServerReady: vi.fn(),
21
- listModels: vi.fn(),
22
- }));
23
-
24
- // Helper to create a mock BaseModel
25
- const createMockModel = (
26
- name: string,
27
- overrides: Partial<BaseModel> = {},
28
- ): BaseModel =>
29
- ({
30
- name,
31
- id: name,
32
- mode: Mode.ROUTER,
33
- capabilities: ["text"] as ["text"],
34
- getStatus: vi.fn().mockResolvedValue(Status.LOADED),
35
- getContextSize: vi.fn().mockResolvedValue(4096),
36
- getInfo: vi.fn().mockResolvedValue(`Model: ${name}\nID: ${name}`),
37
- load: vi.fn().mockResolvedValue(undefined),
38
- unload: vi.fn().mockResolvedValue(undefined),
39
- toProviderConfig: vi.fn().mockResolvedValue({}),
40
- getLabel: vi.fn().mockResolvedValue(name),
41
- ...overrides,
42
- }) as unknown as BaseModel;
43
-
44
- const createMockCtx = (
45
- selectFn: (prompt: string, options: string[]) => string | null,
46
- ) => ({
47
- cwd: "/tmp/test",
48
- ui: {
49
- select: vi.fn(selectFn),
50
- notify: vi.fn(),
51
- theme: {
52
- fg: (color: string, text: string) => text,
53
- },
54
- },
55
- modelRegistry: {
56
- find: vi.fn().mockReturnValue({ id: "test-model-id" }),
57
- },
58
- });
59
-
60
- const createMockPiContext = (notifyFn: ReturnType<typeof vi.fn>) =>
61
- ({
62
- ui: {
63
- notify: notifyFn,
64
- },
65
- }) as any as ExtensionContext;
66
-
67
- const createMockPi = () => ({
68
- setModel: vi.fn(),
69
- registerProvider: vi.fn(),
70
- });
71
-
72
- beforeEach(() => {
73
- vi.clearAllTimers();
74
- resetInflightModel();
75
- });
76
-
77
- afterEach(() => {
78
- vi.clearAllTimers();
79
- });
80
-
81
- describe("modelsCommand", () => {
82
- it("should return early on cancel (null model selection)", async () => {
83
- const models = [createMockModel("model-a")];
84
- const ctx = createMockCtx(() => null);
85
- const pi = createMockPi();
86
-
87
- await modelsCommand(ctx as any, pi as any, models);
88
-
89
- expect(ctx.ui.notify).not.toHaveBeenCalled();
90
- });
91
-
92
- it("should show info when INFO action is selected", async () => {
93
- const model = createMockModel("model-a");
94
- const models = [model];
95
- const ctx = createMockCtx((prompt) => {
96
- if (prompt.includes("models")) return "model-a";
97
- return Action.INFO;
98
- });
99
- const pi = createMockPi();
100
-
101
- await modelsCommand(ctx as any, pi as any, models);
102
-
103
- expect(ctx.ui.notify).toHaveBeenCalledWith(
104
- "Model: model-a\nID: model-a",
105
- "info",
106
- );
107
- });
108
-
109
- it("should unload model when UNLOAD action is selected", async () => {
110
- const model = createMockModel("model-a");
111
- const models = [model];
112
- const ctx = createMockCtx((prompt) => {
113
- if (prompt.includes("models")) return "model-a";
114
- return Action.UNLOAD;
115
- });
116
- const pi = createMockPi();
117
-
118
- await modelsCommand(ctx as any, pi as any, models);
119
-
120
- expect(model.unload).toHaveBeenCalled();
121
- expect(ctx.ui.notify).toHaveBeenCalledWith("Unloaded model-a", "info");
122
- });
123
-
124
- it("should load model when LOAD action is selected", async () => {
125
- const loadFn = vi.fn().mockResolvedValue(undefined);
126
- const model = createMockModel("model-a");
127
- (model.load as any) = loadFn;
128
- (model.getStatus as any).mockResolvedValue(Status.UNLOADED);
129
- const models = [model];
130
- const ctx = createMockCtx((prompt) => {
131
- if (prompt.includes("models")) return "model-a";
132
- return Action.LOAD;
133
- });
134
- const pi = createMockPi();
135
-
136
- await modelsCommand(ctx as any, pi as any, models);
137
- await vi.waitFor(() => expect(loadFn).toHaveBeenCalled());
138
- await vi.waitFor(() => expect(pi.setModel).toHaveBeenCalled());
139
- });
140
-
141
- it("should show warning when session changes during model load", async () => {
142
- // Create a deferred promise so we can control when the load completes
143
- let resolveLoad: () => void;
144
- const loadPromise = new Promise<void>((resolve) => {
145
- resolveLoad = resolve;
146
- });
147
- const model = createMockModel("model-a", {
148
- load: () => loadPromise,
149
- getStatus: vi.fn().mockResolvedValue(Status.UNLOADED),
150
- });
151
- const models = [model];
152
-
153
- let selectCallCount = 0;
154
- const ctx = createMockCtx(() => {
155
- selectCallCount++;
156
- // 1st: select model, 2nd: select LOAD
157
- if (selectCallCount === 1) return "model-a";
158
- if (selectCallCount === 2) return Action.LOAD;
159
- return null;
160
- });
161
- const pi = createMockPi();
162
-
163
- // Start the load (non-blocking)
164
- const modelsPromise = modelsCommand(ctx as any, pi as any, models);
165
-
166
- // Advance past the microtask that sets inflightModel
167
- await vi.advanceTimersByTimeAsync(0);
168
-
169
- // Simulate session switch while model is still loading
170
- // onSessionBeforeSwitch awaits READABLE_TIMEOUT (15s) for the notification
171
- const switchPromise = onSessionBeforeSwitch(
172
- {} as any,
173
- createMockPiContext(ctx.ui.notify as any),
174
- );
175
- await vi.advanceTimersByTimeAsync(15000);
176
- await switchPromise;
177
-
178
- // Should have shown a warning notification
179
- expect(ctx.ui.notify).toHaveBeenCalledWith(
180
- expect.stringContaining("Session change detected"),
181
- "warning",
182
- );
183
- expect(ctx.ui.notify).toHaveBeenCalledWith(
184
- expect.stringContaining("model-a"),
185
- "warning",
186
- );
187
-
188
- // Complete the load so inflightModel is cleared
189
- resolveLoad!();
190
- await modelsPromise;
191
- });
192
-
193
- it("should not warn when no model is loading", async () => {
194
- const notifyFn = vi.fn();
195
- const ctx = createMockPiContext(notifyFn);
196
-
197
- await onSessionBeforeSwitch({} as any, ctx);
198
-
199
- expect(notifyFn).not.toHaveBeenCalled();
200
- // No timers should be scheduled
201
- expect(vi.getTimerCount()).toBe(0);
202
- });
203
-
204
- it("should clear inflightModel after load completes successfully", async () => {
205
- const loadFn = vi.fn().mockResolvedValue(undefined);
206
- const model = createMockModel("model-a", {
207
- load: loadFn,
208
- getStatus: vi.fn().mockResolvedValue(Status.UNLOADED),
209
- });
210
- const models = [model];
211
- const ctx = createMockCtx((prompt) => {
212
- if (prompt.includes("models")) return "model-a";
213
- return Action.LOAD;
214
- });
215
- const pi = createMockPi();
216
-
217
- await modelsCommand(ctx as any, pi as any, models);
218
- await vi.waitFor(() => expect(loadFn).toHaveBeenCalled());
219
- await vi.waitFor(() => expect(pi.setModel).toHaveBeenCalled());
220
-
221
- // inflightModel should be cleared after completion
222
- // (verified indirectly: calling onSessionBeforeSwitch should not warn)
223
- await vi.advanceTimersByTimeAsync(0);
224
- const notifyFn = vi.fn();
225
- await onSessionBeforeSwitch({} as any, createMockPiContext(notifyFn));
226
- expect(notifyFn).not.toHaveBeenCalled();
227
- });
228
-
229
- it("should clear inflightModel after load fails", async () => {
230
- const loadFn = vi.fn().mockRejectedValue(new Error("Load failed"));
231
- const model = createMockModel("model-a", {
232
- load: loadFn,
233
- getStatus: vi.fn().mockResolvedValue(Status.FAILED),
234
- });
235
- const models = [model];
236
- const ctx = createMockCtx((prompt) => {
237
- if (prompt.includes("models")) return "model-a";
238
- return Action.RETRY;
239
- });
240
- const pi = createMockPi();
241
-
242
- await modelsCommand(ctx as any, pi as any, models);
243
- await vi.waitFor(() => expect(loadFn).toHaveBeenCalled());
244
-
245
- // inflightModel should be cleared after failure
246
- await vi.advanceTimersByTimeAsync(0);
247
- const notifyFn = vi.fn();
248
- await onSessionBeforeSwitch({} as any, createMockPiContext(notifyFn));
249
- expect(notifyFn).not.toHaveBeenCalled();
250
- });
251
-
252
- it("should loop back to model selection when action is cancelled", async () => {
253
- const model = createMockModel("model-a");
254
- const models = [model];
255
-
256
- let selectCallCount = 0;
257
- const ctx = createMockCtx(() => {
258
- selectCallCount++;
259
- // 1st: select model-a, 2nd: cancel action, 3rd: cancel model => exit
260
- if (selectCallCount === 1) return "model-a";
261
- return null;
262
- });
263
- const pi = createMockPi();
264
-
265
- await modelsCommand(ctx as any, pi as any, models);
266
-
267
- expect(ctx.ui.select).toHaveBeenCalledTimes(3);
268
- expect(ctx.ui.notify).not.toHaveBeenCalled();
269
- });
270
- });