import { describe, it, expect, beforeEach, afterEach, vi } from "vitest"; import { refineText, validateRefineRequest, checkRateLimit, getRateLimitResetTime, __resetRefineState, ValidationError, InvalidTypeError, AiServiceError, VALID_REFINEMENT_TYPES, MIN_TEXT_LENGTH, MAX_TEXT_LENGTH, MAX_REQUESTS_PER_HOUR, RATE_LIMIT_WINDOW_MS, } from "./ai-refine.js"; // Hoisted mock factory const { mockCreateFnAgent } = vi.hoisted(() => ({ mockCreateFnAgent: vi.fn(), })); // Mock the engine module to avoid dynamic import issues in tests vi.mock("@fusion/engine", () => ({ createFnAgent: mockCreateFnAgent, })); describe("ai-refine module", () => { beforeEach(() => { vi.useFakeTimers({ shouldAdvanceTime: true }); __resetRefineState(); vi.clearAllMocks(); mockCreateFnAgent.mockResolvedValue(null); }); afterEach(() => { vi.useRealTimers(); }); describe("VALID_REFINEMENT_TYPES", () => { it("contains all four refinement types", () => { expect(VALID_REFINEMENT_TYPES).toEqual([ "clarify", "add-details", "expand", "simplify", ]); }); }); describe("validateRefineRequest", () => { it("accepts valid text and 'clarify' type", () => { const result = validateRefineRequest("Some text", "clarify"); expect(result).toEqual({ text: "Some text", type: "clarify" }); }); it("accepts valid text and 'add-details' type", () => { const result = validateRefineRequest("Some text", "add-details"); expect(result).toEqual({ text: "Some text", type: "add-details" }); }); it("accepts valid text and 'expand' type", () => { const result = validateRefineRequest("Some text", "expand"); expect(result).toEqual({ text: "Some text", type: "expand" }); }); it("accepts valid text and 'simplify' type", () => { const result = validateRefineRequest("Some text", "simplify"); expect(result).toEqual({ text: "Some text", type: "simplify" }); }); it("throws ValidationError for missing text", () => { expect(() => validateRefineRequest(undefined, "clarify")).toThrow(ValidationError); expect(() => validateRefineRequest(undefined, "clarify")).toThrow("text is required"); }); it("throws ValidationError for null text", () => { expect(() => validateRefineRequest(null, "clarify")).toThrow(ValidationError); expect(() => validateRefineRequest(null, "clarify")).toThrow("text is required"); }); it("throws ValidationError for non-string text", () => { expect(() => validateRefineRequest(123, "clarify")).toThrow(ValidationError); expect(() => validateRefineRequest(123, "clarify")).toThrow("text must be a string"); }); it("throws ValidationError for empty text", () => { expect(() => validateRefineRequest("", "clarify")).toThrow(ValidationError); expect(() => validateRefineRequest("", "clarify")).toThrow( `text must be at least ${MIN_TEXT_LENGTH} character` ); }); it("throws ValidationError for text exceeding MAX_TEXT_LENGTH", () => { const longText = "a".repeat(MAX_TEXT_LENGTH + 1); expect(() => validateRefineRequest(longText, "clarify")).toThrow(ValidationError); expect(() => validateRefineRequest(longText, "clarify")).toThrow( `text must not exceed ${MAX_TEXT_LENGTH} characters` ); }); it("accepts text at exactly MAX_TEXT_LENGTH", () => { const maxText = "a".repeat(MAX_TEXT_LENGTH); const result = validateRefineRequest(maxText, "clarify"); expect(result.text).toHaveLength(MAX_TEXT_LENGTH); }); it("accepts text at exactly MIN_TEXT_LENGTH", () => { const result = validateRefineRequest("a", "clarify"); expect(result.text).toBe("a"); }); it("throws ValidationError for missing type", () => { expect(() => validateRefineRequest("some text", undefined)).toThrow(ValidationError); expect(() => validateRefineRequest("some text", undefined)).toThrow("type is required"); }); it("throws ValidationError for null type", () => { expect(() => validateRefineRequest("some text", null)).toThrow(ValidationError); expect(() => validateRefineRequest("some text", null)).toThrow("type is required"); }); it("throws InvalidTypeError for invalid type string", () => { expect(() => validateRefineRequest("some text", "invalid")).toThrow(InvalidTypeError); expect(() => validateRefineRequest("some text", "invalid")).toThrow( "type must be one of: clarify, add-details, expand, simplify" ); }); it("throws InvalidTypeError for numeric type", () => { expect(() => validateRefineRequest("some text", 123)).toThrow(InvalidTypeError); }); }); describe("checkRateLimit", () => { it("allows first request from an IP", () => { expect(checkRateLimit("192.168.1.1")).toBe(true); }); it("allows up to MAX_REQUESTS_PER_HOUR requests", () => { const ip = "192.168.1.1"; for (let i = 0; i < MAX_REQUESTS_PER_HOUR; i++) { expect(checkRateLimit(ip)).toBe(true); } }); it("blocks request beyond MAX_REQUESTS_PER_HOUR", () => { const ip = "192.168.1.1"; // Use up the quota for (let i = 0; i < MAX_REQUESTS_PER_HOUR; i++) { checkRateLimit(ip); } // 11th request should be blocked expect(checkRateLimit(ip)).toBe(false); }); it("tracks different IPs independently", () => { const ip1 = "192.168.1.1"; const ip2 = "192.168.1.2"; // Use up quota for ip1 for (let i = 0; i < MAX_REQUESTS_PER_HOUR; i++) { checkRateLimit(ip1); } expect(checkRateLimit(ip1)).toBe(false); // ip2 should still have full quota expect(checkRateLimit(ip2)).toBe(true); }); it("resets rate limit after RATE_LIMIT_WINDOW_MS", () => { const ip = "192.168.1.1"; // Use up the quota for (let i = 0; i < MAX_REQUESTS_PER_HOUR; i++) { checkRateLimit(ip); } expect(checkRateLimit(ip)).toBe(false); // Advance time by 1 hour + 1ms vi.advanceTimersByTime(RATE_LIMIT_WINDOW_MS + 1); // Should be allowed again expect(checkRateLimit(ip)).toBe(true); }); it("resets count but tracks new window after expiry", () => { const ip = "192.168.1.1"; // Make 5 requests for (let i = 0; i < 5; i++) { checkRateLimit(ip); } // Advance time by 1 hour + 1ms vi.advanceTimersByTime(RATE_LIMIT_WINDOW_MS + 1); // First request in new window should work expect(checkRateLimit(ip)).toBe(true); // Use up remaining quota in new window for (let i = 0; i < MAX_REQUESTS_PER_HOUR - 1; i++) { expect(checkRateLimit(ip)).toBe(true); } // Next request should be blocked expect(checkRateLimit(ip)).toBe(false); }); }); describe("getRateLimitResetTime", () => { it("returns null for unknown IP", () => { expect(getRateLimitResetTime("unknown-ip")).toBeNull(); }); it("returns reset time after a request is made", () => { const ip = "192.168.1.1"; const beforeRequest = Date.now(); checkRateLimit(ip); const resetTime = getRateLimitResetTime(ip); expect(resetTime).not.toBeNull(); expect(resetTime!.getTime()).toBe(beforeRequest + RATE_LIMIT_WINDOW_MS); }); it("returns updated reset time after window resets", () => { const ip = "192.168.1.1"; checkRateLimit(ip); const firstResetTime = getRateLimitResetTime(ip); // Advance time past the window vi.advanceTimersByTime(RATE_LIMIT_WINDOW_MS + 1000); // Make another request checkRateLimit(ip); const secondResetTime = getRateLimitResetTime(ip); // Second reset time should be later than first expect(secondResetTime!.getTime()).toBeGreaterThan(firstResetTime!.getTime()); }); }); describe("error classes", () => { it("ValidationError has correct name", () => { const error = new ValidationError("test"); expect(error.name).toBe("ValidationError"); expect(error.message).toBe("test"); }); it("InvalidTypeError has correct name", () => { const error = new InvalidTypeError("test"); expect(error.name).toBe("InvalidTypeError"); expect(error.message).toBe("test"); }); it("AiServiceError has correct name", () => { const error = new AiServiceError("test"); expect(error.name).toBe("AiServiceError"); expect(error.message).toBe("test"); }); }); describe("refineText", () => { // Note: refineText requires the AI engine which is not available in tests. // These tests verify error handling when the engine is unavailable. it("throws AiServiceError when AI engine is not available", async () => { await expect(refineText("some text", "clarify", "/some/path")).rejects.toThrow( AiServiceError ); await expect(refineText("some text", "clarify", "/some/path")).rejects.toThrow( "Failed to initialize AI agent" ); }); }); describe("__resetRefineState", () => { it("clears all rate limit entries", () => { const ip = "192.168.1.1"; // Make requests to populate rate limits for (let i = 0; i < 5; i++) { checkRateLimit(ip); } expect(getRateLimitResetTime(ip)).not.toBeNull(); // Reset state __resetRefineState(); // Should be like starting fresh expect(getRateLimitResetTime(ip)).toBeNull(); expect(checkRateLimit(ip)).toBe(true); }); }); describe("prompt override regression", () => { const customPrompt = "You are a custom refine assistant."; function createRefineMockAgent(responseText: string) { return { session: { state: { messages: [ { role: "assistant", content: responseText }, ], }, prompt: vi.fn().mockResolvedValue(undefined), dispose: vi.fn(), }, }; } it("uses default prompt when no overrides provided", async () => { const mockAgent = createRefineMockAgent("Refined text here"); mockCreateFnAgent.mockResolvedValueOnce(mockAgent); await refineText("some text", "clarify", "/tmp/project"); expect(mockCreateFnAgent).toHaveBeenCalled(); const call = mockCreateFnAgent.mock.calls[0]; expect(call[0].systemPrompt).toMatch(/^You are a text refinement assistant/); }); it("uses override prompt when promptOverrides provided", async () => { const mockAgent = createRefineMockAgent("Refined text here"); mockCreateFnAgent.mockResolvedValueOnce(mockAgent); await refineText("some text", "clarify", "/tmp/project", { "ai-refine-system": customPrompt, }); expect(mockCreateFnAgent).toHaveBeenCalled(); const call = mockCreateFnAgent.mock.calls[0]; expect(call[0].systemPrompt).toBe(customPrompt); }); it("falls back to default prompt when override is empty string", async () => { const mockAgent = createRefineMockAgent("Refined text here"); mockCreateFnAgent.mockResolvedValueOnce(mockAgent); await refineText("some text", "simplify", "/tmp/project", { "ai-refine-system": "", }); expect(mockCreateFnAgent).toHaveBeenCalled(); const call = mockCreateFnAgent.mock.calls[0]; expect(call[0].systemPrompt).toMatch(/^You are a text refinement assistant/); }); it("does not introduce unexpected model/provider override fields in createFnAgent", async () => { const mockAgent = createRefineMockAgent("Refined text here"); mockCreateFnAgent.mockResolvedValueOnce(mockAgent); await refineText("some text", "expand", "/tmp/project", { "ai-refine-system": customPrompt, }); expect(mockCreateFnAgent).toHaveBeenCalled(); const call = mockCreateFnAgent.mock.calls[0]; const agentConfig = call[0]; // Verify only expected fields are present expect(agentConfig).toHaveProperty("cwd"); expect(agentConfig).toHaveProperty("systemPrompt"); expect(agentConfig).toHaveProperty("tools"); // Verify no unexpected model override fields expect(agentConfig).not.toHaveProperty("modelProvider"); expect(agentConfig).not.toHaveProperty("modelId"); expect(agentConfig).not.toHaveProperty("provider"); expect(agentConfig).not.toHaveProperty("model"); }); it("prompt overrides do not affect other prompt keys", async () => { const mockAgent = createRefineMockAgent("Refined text here"); mockCreateFnAgent.mockResolvedValueOnce(mockAgent); // Provide overrides for a different key only await refineText("some text", "clarify", "/tmp/project", { "workflow-step-refine": "Should not affect AI refine", }); expect(mockCreateFnAgent).toHaveBeenCalled(); const call = mockCreateFnAgent.mock.calls[0]; // AI refine should use its own default prompt, not affected by workflow-step-refine override expect(call[0].systemPrompt).toMatch(/^You are a text refinement assistant/); }); }); });