pi-free 1.0.0

package/tests/fireworks.test.ts

@@ -0,0 +1,148 @@
+/**
+ * Fireworks Provider Tests
+ */
+
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+// Mock dependencies
+vi.mock("../config.ts", () => ({
+	FIREWORKS_API_KEY: "test-fireworks-key",
+	FIREWORKS_SHOW_PAID: true,
+	PROVIDER_FIREWORKS: "fireworks",
+	applyHidden: (models: unknown[]) => models,
+}));
+
+vi.mock("../constants.ts", () => ({
+	BASE_URL_FIREWORKS: "https://api.fireworks.ai/inference/v1",
+}));
+
+vi.mock("../provider-helper.ts", () => ({
+	createReRegister: vi.fn(() => vi.fn()),
+	setupProvider: vi.fn(),
+}));
+
+vi.mock("../lib/logger.ts", () => ({
+	createLogger: () => ({
+		warn: vi.fn(),
+		info: vi.fn(),
+		debug: vi.fn(),
+		error: vi.fn(),
+	}),
+}));
+
+import { setupProvider } from "../provider-helper.ts";
+
+describe("Fireworks Provider", () => {
+	let mockPi: ExtensionAPI;
+	let mockRegisterProvider: ReturnType<typeof vi.fn>;
+
+	beforeEach(() => {
+		vi.clearAllMocks();
+		mockRegisterProvider = vi.fn();
+
+		mockPi = {
+			registerProvider: mockRegisterProvider,
+			registerCommand: vi.fn(),
+		} as unknown as ExtensionAPI;
+	});
+
+	describe("initialization", () => {
+		it("should register provider with hardcoded models", async () => {
+			const { default: fireworksProvider } = await import(
+				"../providers/fireworks.ts"
+			);
+			await fireworksProvider(mockPi);
+
+			expect(mockRegisterProvider).toHaveBeenCalledWith(
+				"fireworks",
+				expect.objectContaining({
+					baseUrl: "https://api.fireworks.ai/inference/v1",
+					apiKey: "FIREWORKS_API_KEY",
+					api: "openai-completions",
+					models: expect.any(Array),
+				}),
+			);
+		});
+
+		it("should set API key in environment", async () => {
+			delete process.env.FIREWORKS_API_KEY;
+
+			const { default: fireworksProvider } = await import(
+				"../providers/fireworks.ts"
+			);
+			await fireworksProvider(mockPi);
+
+			expect(process.env.FIREWORKS_API_KEY).toBe("test-fireworks-key");
+		});
+
+		it("should skip registration without API key", async () => {
+			// Mock no API key by temporarily clearing the module
+			const apiKeySpy = vi
+				.spyOn(await import("../config.ts"), "FIREWORKS_API_KEY", "get")
+				.mockReturnValue(undefined as any);
+
+			const { default: fireworksProvider } = await import(
+				"../providers/fireworks.ts"
+			);
+			await fireworksProvider(mockPi);
+
+			// Should not register provider
+			expect(mockRegisterProvider).not.toHaveBeenCalled();
+
+			// Restore the mock so subsequent tests have the API key
+			apiKeySpy.mockRestore();
+		});
+	});
+
+	describe("model configuration", () => {
+		it("should have hardcoded models with correct structure", async () => {
+			const { default: fireworksProvider } = await import(
+				"../providers/fireworks.ts"
+			);
+			await fireworksProvider(mockPi);
+
+			expect(mockRegisterProvider).toHaveBeenCalled();
+			const registerCall = mockRegisterProvider.mock.calls[0];
+			expect(registerCall).toBeDefined();
+			const models = registerCall?.[1]?.models;
+
+			expect(models).toBeInstanceOf(Array);
+			expect(models.length).toBeGreaterThan(0);
+
+			// Check first model has required properties
+			const firstModel = models[0];
+			expect(firstModel).toHaveProperty("id");
+			expect(firstModel).toHaveProperty("name");
+			expect(firstModel).toHaveProperty("reasoning");
+			expect(firstModel).toHaveProperty("input");
+			expect(firstModel).toHaveProperty("cost");
+			expect(firstModel).toHaveProperty("contextWindow");
+			expect(firstModel).toHaveProperty("maxTokens");
+
+			// Verify non-zero costs (paid model, not free)
+			expect(firstModel.cost.input).toBeGreaterThan(0);
+			expect(firstModel.cost.output).toBeGreaterThan(0);
+		});
+	});
+
+	describe("setupProvider integration", () => {
+		it("should call setupProvider with stored models", async () => {
+			const { default: fireworksProvider } = await import(
+				"../providers/fireworks.ts"
+			);
+			await fireworksProvider(mockPi);
+
+			expect(setupProvider).toHaveBeenCalledWith(
+				mockPi,
+				expect.objectContaining({
+					providerId: "fireworks",
+				}),
+				expect.objectContaining({
+					free: expect.any(Array),
+					all: expect.any(Array),
+				}),
+			);
+		});
+	});
+});

package/tests/free-tier-limits.test.ts

@@ -0,0 +1,191 @@
+import { beforeEach, describe, expect, it } from "vitest";
+import {
+	FREE_TIER_LIMITS,
+	formatFreeTierStatus,
+	getFreeTierUsage,
+	getLimitWarning,
+	getModelUsage,
+	getProviderModelUsage,
+	getSessionUsage,
+	getTopModels,
+	incrementModelRequestCount,
+	isApproachingLimit,
+	resetUsageStats,
+} from "../usage/limits.ts";
+
+describe("Free Tier Limits", () => {
+	beforeEach(() => {
+		// Reset usage stats before each test
+		resetUsageStats();
+	});
+
+	describe("FREE_TIER_LIMITS", () => {
+		it("should have limits for kilo provider", () => {
+			expect(FREE_TIER_LIMITS.kilo).toBeDefined();
+			expect(FREE_TIER_LIMITS.kilo.requestsPerHour).toBe(200);
+		});
+
+		it("should have limits for openrouter provider", () => {
+			expect(FREE_TIER_LIMITS.openrouter).toBeDefined();
+			expect(FREE_TIER_LIMITS.openrouter.requestsPerDay).toBe(1000);
+		});
+
+		it("should have limits for nvidia provider", () => {
+			expect(FREE_TIER_LIMITS.nvidia).toBeDefined();
+			expect(FREE_TIER_LIMITS.nvidia.requestsPerMonth).toBe(1000);
+		});
+
+		it("should have limits for fireworks provider", () => {
+			expect(FREE_TIER_LIMITS.fireworks).toBeDefined();
+			expect(FREE_TIER_LIMITS.fireworks.requestsPerMonth).toBe(1000);
+		});
+	});
+
+	describe("Request Counting", () => {
+		it("should increment request count", () => {
+			// Just verify it doesn't throw
+			expect(() => {
+				incrementModelRequestCount("kilo", "gpt-4", 100, 50);
+			}).not.toThrow();
+		});
+
+		it("should track model usage", () => {
+			incrementModelRequestCount("test", "model-1", 100, 50);
+			incrementModelRequestCount("test", "model-1", 200, 100);
+
+			const usage = getModelUsage("test", "model-1");
+			expect(usage).toBeDefined();
+			expect(usage?.count).toBe(2);
+			expect(usage?.tokensIn).toBe(300);
+			expect(usage?.tokensOut).toBe(150);
+		});
+
+		it("should return undefined for unknown models", () => {
+			const usage = getModelUsage("unknown", "unknown-model");
+			expect(usage).toBeUndefined();
+		});
+
+		it("should track different models separately", () => {
+			incrementModelRequestCount("test", "model-a", 100, 50);
+			incrementModelRequestCount("test", "model-b", 200, 100);
+
+			expect(getModelUsage("test", "model-a")?.count).toBe(1);
+			expect(getModelUsage("test", "model-b")?.count).toBe(1);
+		});
+	});
+
+	describe("Provider Model Usage", () => {
+		it("should return all models for a provider", () => {
+			incrementModelRequestCount("test-provider", "model-1", 100, 50);
+			incrementModelRequestCount("test-provider", "model-2", 200, 100);
+
+			const models = getProviderModelUsage("test-provider");
+			expect(models).toHaveLength(2);
+			expect(models.map((m) => m.modelId)).toContain("model-1");
+			expect(models.map((m) => m.modelId)).toContain("model-2");
+		});
+
+		it("should return empty array for unknown provider", () => {
+			const models = getProviderModelUsage("nonexistent");
+			expect(models).toEqual([]);
+		});
+	});
+
+	describe("Top Models", () => {
+		it("should return top N models by request count", () => {
+			// Add many models with varying request counts
+			for (let i = 0; i < 5; i++) {
+				incrementModelRequestCount("test", "popular-model", 100, 50);
+			}
+			incrementModelRequestCount("test", "unpopular-model", 100, 50);
+
+			const top = getTopModels(2);
+			expect(top).toHaveLength(2);
+			expect(top[0].modelId).toBe("popular-model");
+		});
+	});
+
+	describe("Free Tier Usage", () => {
+		it("should calculate usage for kilo provider", () => {
+			// Simulate requests (out of 200/hour limit)
+			for (let i = 0; i < 100; i++) {
+				incrementModelRequestCount("kilo", "gpt-4", 10, 10);
+			}
+
+			const usage = getFreeTierUsage("kilo");
+			// requestsThisHour is capped at 50 as a rough estimate
+			expect(usage.requestsToday).toBeGreaterThanOrEqual(100);
+			expect(usage.requestsThisHour).toBeLessThanOrEqual(50);
+		});
+
+		it("should calculate usage for providers without hourly limits", () => {
+			// OpenRouter has daily limit, not hourly
+			const usage = getFreeTierUsage("openrouter");
+			expect(usage.requestsThisHour).toBe(0);
+		});
+	});
+
+	describe("Limit Warnings", () => {
+		it("should detect when approaching limit", () => {
+			// Use openrouter with daily limit (1000/day)
+			// Add 750 requests = 75% which triggers warning
+			for (let i = 0; i < 750; i++) {
+				incrementModelRequestCount("openrouter", "gpt-4", 10, 10);
+			}
+
+			expect(isApproachingLimit("openrouter")).toBe(true);
+		});
+
+		it("should not trigger warning when usage is low", () => {
+			incrementModelRequestCount("kilo", "gpt-4", 10, 10);
+			expect(isApproachingLimit("kilo")).toBe(false);
+		});
+
+		it("should return warning message when approaching limit", () => {
+			// Use openrouter with daily limit (1000/day)
+			// Add 750 requests = 75% which triggers warning
+			for (let i = 0; i < 750; i++) {
+				incrementModelRequestCount("openrouter", "gpt-4", 10, 10);
+			}
+
+			const warning = getLimitWarning("openrouter");
+			expect(warning).not.toBeNull();
+			expect(warning).toContain("%");
+		});
+
+		it("should return null when not approaching limit", () => {
+			const warning = getLimitWarning("kilo");
+			expect(warning).toBeNull();
+		});
+	});
+
+	describe("Session Usage", () => {
+		it("should generate session report", () => {
+			incrementModelRequestCount("kilo", "gpt-4", 1000, 500);
+			incrementModelRequestCount("openrouter", "mimo", 500, 250);
+
+			const report = getSessionUsage();
+			expect(report.totalRequests).toBe(2);
+			expect(report.totalTokensIn).toBe(1500);
+			expect(report.totalTokensOut).toBe(750);
+		});
+
+		it("should track providers in session report", () => {
+			incrementModelRequestCount("kilo", "gpt-4", 100, 50);
+			incrementModelRequestCount("openrouter", "mimo", 100, 50);
+
+			const report = getSessionUsage();
+			const providerNames = report.providers.map((p) => p.name);
+			expect(providerNames).toContain("kilo");
+			expect(providerNames).toContain("openrouter");
+		});
+	});
+
+	describe("Status Formatting", () => {
+		it("should format status for provider", () => {
+			const status = formatFreeTierStatus("kilo");
+			expect(typeof status).toBe("string");
+			expect(status).toContain("kilo");
+		});
+	});
+});

package/tests/json-persistence.test.ts

@@ -0,0 +1,105 @@
+/**
+ * JSON Persistence Tests
+ */
+
+import { existsSync, rmdirSync, unlinkSync } from "node:fs";
+import { join } from "node:path";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import { createJSONLStore, createJSONStore } from "../lib/json-persistence.ts";
+
+const TEST_DIR = join(
+	process.env.HOME || process.env.USERPROFILE || "",
+	".pi-test",
+);
+
+describe("JSON Persistence", () => {
+	describe("createJSONStore", () => {
+		const testFile = join(TEST_DIR, "test-store.json");
+
+		beforeEach(() => {
+			// Clean up
+			try {
+				if (existsSync(testFile)) unlinkSync(testFile);
+			} catch {}
+		});
+
+		afterEach(() => {
+			try {
+				if (existsSync(testFile)) unlinkSync(testFile);
+				if (existsSync(TEST_DIR)) rmdirSync(TEST_DIR);
+			} catch {}
+		});
+
+		it("should return default value when file doesn't exist", () => {
+			const store = createJSONStore(testFile, { default: true });
+			const data = store.load();
+			expect(data).toEqual({ default: true });
+		});
+
+		it("should persist and load data", () => {
+			const store = createJSONStore(testFile, { count: 0 });
+			store.save({ count: 42 });
+
+			// Create new store instance to test loading from disk
+			const store2 = createJSONStore(testFile, { count: 0 });
+			const data = store2.load();
+			expect(data.count).toBe(42);
+		});
+
+		it("should cache data after first load", () => {
+			const store = createJSONStore(testFile, { value: "initial" });
+			store.save({ value: "updated" });
+
+			// First load reads from disk
+			const data1 = store.load();
+			expect(data1.value).toBe("updated");
+
+			// Second load should return cached value
+			const data2 = store.load();
+			expect(data2.value).toBe("updated");
+		});
+	});
+
+	describe("createJSONLStore", () => {
+		const testFile = join(TEST_DIR, "test-log.jsonl");
+
+		beforeEach(() => {
+			try {
+				if (existsSync(testFile)) unlinkSync(testFile);
+			} catch {}
+		});
+
+		afterEach(() => {
+			try {
+				if (existsSync(testFile)) unlinkSync(testFile);
+				if (existsSync(TEST_DIR)) rmdirSync(TEST_DIR);
+			} catch {}
+		});
+
+		it("should return empty array when file doesn't exist", () => {
+			const store = createJSONLStore<{ msg: string }>(testFile);
+			const data = store.load();
+			expect(data).toEqual([]);
+		});
+
+		it("should append entries", () => {
+			const store = createJSONLStore<{ event: string }>(testFile);
+			store.append({ event: "first" });
+			store.append({ event: "second" });
+
+			const data = store.load();
+			expect(data).toHaveLength(2);
+			expect(data[0].event).toBe("first");
+			expect(data[1].event).toBe("second");
+		});
+
+		it("should clear entries", () => {
+			const store = createJSONLStore<{ event: string }>(testFile);
+			store.append({ event: "test" });
+			store.clear();
+
+			const data = store.load();
+			expect(data).toEqual([]);
+		});
+	});
+});

package/tests/kilo.test.ts

@@ -0,0 +1,186 @@
+/**
+ * Kilo Provider Tests
+ */
+
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+// Create mock functions first
+const mockFetchKiloModels = vi.fn();
+const mockSetupProvider = vi.fn();
+const mockLoginKilo = vi.fn();
+
+// Mock dependencies before importing the provider
+vi.mock("../kilo-auth.ts", () => ({
+	loginKilo: (...args: unknown[]) => mockLoginKilo(...args),
+	refreshKiloToken: vi.fn(),
+}));
+
+vi.mock("../kilo-models.ts", () => ({
+	fetchKiloModels: (...args: unknown[]) => mockFetchKiloModels(...args),
+	KILO_GATEWAY_BASE: "https://api.kilo.ai/api/gateway",
+}));
+
+vi.mock("../provider-helper.ts", () => ({
+	createReRegister: vi.fn(() => vi.fn()),
+	setupProvider: (...args: unknown[]) => mockSetupProvider(...args),
+	enhanceWithCI: (models: unknown[]) => models,
+}));
+
+vi.mock("../usage-widget.ts", () => ({
+	registerUsageWidget: vi.fn(),
+}));
+
+vi.mock("../util.ts", () => ({
+	logWarning: vi.fn(),
+}));
+
+import kiloProvider from "../providers/kilo.ts";
+
+describe("Kilo Provider", () => {
+	let _mockPi: ExtensionAPI;
+	let mockRegisterProvider: ReturnType<typeof vi.fn>;
+	let mockOn: ReturnType<typeof vi.fn>;
+
+	beforeEach(() => {
+		vi.clearAllMocks();
+		mockFetchKiloModels.mockReset();
+		mockSetupProvider.mockReset();
+		mockLoginKilo.mockReset();
+
+		mockRegisterProvider = vi.fn();
+		mockOn = vi.fn();
+
+		_mockPi = {
+			registerProvider: mockRegisterProvider,
+			on: mockOn,
+			registerCommand: vi.fn(),
+		} as unknown as ExtensionAPI;
+	});
+
+describe("Kilo Provider", () => {
+	let mockPi: ExtensionAPI;
+	let mockRegisterProvider: ReturnType<typeof vi.fn>;
+	let mockOn: ReturnType<typeof vi.fn>;
+
+	beforeEach(() => {
+		vi.clearAllMocks();
+		mockRegisterProvider = vi.fn();
+		mockOn = vi.fn();
+
+		mockPi = {
+			registerProvider: mockRegisterProvider,
+			on: mockOn,
+			registerCommand: vi.fn(),
+		} as unknown as ExtensionAPI;
+	});
+
+	describe("initialization", () => {
+		it("should register provider with free models on startup", async () => {
+			const mockModels = [
+				{
+					id: "gpt-4",
+					name: "GPT-4",
+					reasoning: true,
+					input: ["text"],
+					cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+					contextWindow: 128000,
+					maxTokens: 4096,
+				},
+			];
+			vi.mocked(fetchKiloModels).mockResolvedValue(mockModels);
+
+			await kiloProvider(mockPi);
+
+			expect(mockRegisterProvider).toHaveBeenCalledWith(
+				"kilo",
+				expect.objectContaining({
+					baseUrl: "https://api.kilo.ai/api/gateway",
+					apiKey: "KILO_API_KEY",
+					api: "openai-completions",
+					models: mockModels,
+					oauth: expect.any(Object),
+				}),
+			);
+		});
+
+		it("should handle model fetch failure gracefully", async () => {
+			vi.mocked(fetchKiloModels).mockRejectedValue(new Error("Network error"));
+
+			await kiloProvider(mockPi);
+
+			// Should still register with empty models
+			expect(mockRegisterProvider).toHaveBeenCalled();
+		});
+	});
+
+	describe("OAuth integration", () => {
+		it("should have oauth configuration", async () => {
+			vi.mocked(fetchKiloModels).mockResolvedValue([]);
+
+			await kiloProvider(mockPi);
+
+			const registerCall = mockRegisterProvider.mock.calls[0];
+			expect(registerCall[1]).toHaveProperty("oauth");
+			expect(registerCall[1].oauth).toHaveProperty("name", "Kilo");
+			expect(registerCall[1].oauth).toHaveProperty("login");
+			expect(registerCall[1].oauth).toHaveProperty("refreshToken");
+			expect(registerCall[1].oauth).toHaveProperty("getApiKey");
+			expect(registerCall[1].oauth).toHaveProperty("modifyModels");
+		});
+
+		it("should fetch all models after login", async () => {
+			vi.mocked(fetchKiloModels).mockResolvedValue([]);
+
+			await kiloProvider(mockPi);
+
+			const oauth = mockRegisterProvider.mock.calls[0][1].oauth;
+			const _mockCreds = {
+				access: "test-token",
+				refresh: "",
+				expires: Date.now() + 3600000,
+			};
+
+			vi.mocked(fetchKiloModels).mockResolvedValue([
+				{ id: "gpt-4", name: "GPT-4" },
+				{ id: "claude-3", name: "Claude 3" },
+			]);
+
+			await oauth.login({ onProgress: vi.fn() });
+
+			expect(fetchKiloModels).toHaveBeenCalledWith({ token: "test-token" });
+		});
+	});
+
+	describe("event handlers", () => {
+		it("should register session_start handler", async () => {
+			vi.mocked(fetchKiloModels).mockResolvedValue([]);
+
+			await kiloProvider(mockPi);
+
+			expect(mockOn).toHaveBeenCalledWith(
+				"session_start",
+				expect.any(Function),
+			);
+		});
+	});
+
+	describe("setupProvider integration", () => {
+		it("should call setupProvider with correct config", async () => {
+			vi.mocked(fetchKiloModels).mockResolvedValue([]);
+
+			await kiloProvider(mockPi);
+
+			expect(setupProvider).toHaveBeenCalledWith(
+				mockPi,
+				expect.objectContaining({
+					providerId: "kilo",
+					tosUrl: expect.stringContaining("kilo.ai"),
+					reRegister: expect.any(Function),
+				}),
+				expect.objectContaining({ free: [], all: [] }),
+			);
+		});
+	});
+});
+});