import { describe, it, expect, beforeEach, vi, afterEach } from "vitest"; import { startAgentGeneration, generateAgentSpec, getAgentGenerationSession, cleanupAgentGenerationSession, checkRateLimit, getRateLimitResetTime, parseGenerationResponse, __resetAgentGenerationState, RateLimitError, SessionNotFoundError, } from "./agent-generation.js"; // Counter for unique IPs per test let ipCounter = 0; function getUniqueIp(): string { return `127.0.0.${++ipCounter}`; } describe("agent-generation module", () => { beforeEach(() => { vi.useFakeTimers(); __resetAgentGenerationState(); }); afterEach(() => { vi.useRealTimers(); }); describe("startAgentGeneration", () => { it("creates a session with valid role description", async () => { const mockIp = getUniqueIp(); const session = await startAgentGeneration(mockIp, "Senior frontend code reviewer"); expect(session.id).toBeDefined(); expect(typeof session.id).toBe("string"); expect(session.roleDescription).toBe("Senior frontend code reviewer"); expect(session.spec).toBeUndefined(); expect(session.createdAt).toBeInstanceOf(Date); expect(session.updatedAt).toBeInstanceOf(Date); }); it("does not expose IP in the public session", async () => { const mockIp = getUniqueIp(); const session = await startAgentGeneration(mockIp, "Test role"); expect((session as Record).ip).toBeUndefined(); }); it("enforces rate limiting", async () => { const mockIp = getUniqueIp(); // Create max sessions (10 per hour) for (let i = 0; i < 10; i++) { await startAgentGeneration(mockIp, `Role ${i}`); } // 11th session should fail await expect(startAgentGeneration(mockIp, "One more")).rejects.toThrow(RateLimitError); }); it("allows new sessions after rate limit window expires", async () => { const mockIp = getUniqueIp(); for (let i = 0; i < 10; i++) { await startAgentGeneration(mockIp, `Role ${i}`); } // Advance time by 1 hour + 1 minute vi.advanceTimersByTime(61 * 60 * 1000); const session = await startAgentGeneration(mockIp, "New role after reset"); expect(session.id).toBeDefined(); }); it("generates different session IDs for each session", async () => { const ip1 = getUniqueIp(); const ip2 = getUniqueIp(); const session1 = await startAgentGeneration(ip1, "Role 1"); const session2 = await startAgentGeneration(ip2, "Role 2"); expect(session1.id).not.toBe(session2.id); }); it("rate limits independently per IP", async () => { const ip1 = getUniqueIp(); const ip2 = getUniqueIp(); for (let i = 0; i < 10; i++) { await startAgentGeneration(ip1, `Role ${i}`); } // ip2 should still work const session = await startAgentGeneration(ip2, "Role from another IP"); expect(session.id).toBeDefined(); }); }); describe("generateAgentSpec", () => { it("throws SessionNotFoundError for non-existent session", async () => { await expect( generateAgentSpec("non-existent-session-id", "/tmp") ).rejects.toThrow(SessionNotFoundError); }); }); describe("getAgentGenerationSession", () => { it("returns session after creation", async () => { const mockIp = getUniqueIp(); const created = await startAgentGeneration(mockIp, "Test role"); const retrieved = getAgentGenerationSession(created.id); expect(retrieved).toBeDefined(); expect(retrieved!.id).toBe(created.id); expect(retrieved!.roleDescription).toBe("Test role"); }); it("returns undefined for non-existent session", () => { const result = getAgentGenerationSession("non-existent"); expect(result).toBeUndefined(); }); it("returns undefined after cleanup", async () => { const mockIp = getUniqueIp(); const created = await startAgentGeneration(mockIp, "Test role"); cleanupAgentGenerationSession(created.id); const result = getAgentGenerationSession(created.id); expect(result).toBeUndefined(); }); }); describe("cleanupAgentGenerationSession", () => { it("removes session", async () => { const mockIp = getUniqueIp(); const created = await startAgentGeneration(mockIp, "Test role"); cleanupAgentGenerationSession(created.id); expect(getAgentGenerationSession(created.id)).toBeUndefined(); }); it("is idempotent for non-existent session", () => { expect(() => cleanupAgentGenerationSession("non-existent")).not.toThrow(); }); }); describe("checkRateLimit", () => { beforeEach(() => { __resetAgentGenerationState(); }); it("allows first request from new IP", () => { expect(checkRateLimit("1.2.3.4")).toBe(true); }); it("blocks after exceeding limit", () => { for (let i = 0; i < 10; i++) { checkRateLimit("1.2.3.4"); } expect(checkRateLimit("1.2.3.4")).toBe(false); }); it("allows requests after window expires", () => { for (let i = 0; i < 10; i++) { checkRateLimit("1.2.3.4"); } vi.advanceTimersByTime(61 * 60 * 1000); expect(checkRateLimit("1.2.3.4")).toBe(true); }); }); describe("getRateLimitResetTime", () => { it("returns null for unknown IP", () => { expect(getRateLimitResetTime("unknown")).toBeNull(); }); it("returns a date after first request", () => { checkRateLimit("1.2.3.4"); const resetTime = getRateLimitResetTime("1.2.3.4"); expect(resetTime).toBeInstanceOf(Date); expect(resetTime!.getTime()).toBeGreaterThan(Date.now()); }); }); describe("parseGenerationResponse", () => { it("parses valid JSON response", () => { const json = JSON.stringify({ title: "Senior Frontend Reviewer", icon: "🔍", role: "reviewer", description: "Reviews frontend code for quality", systemPrompt: "# Role\nYou are a senior frontend reviewer.", thinkingLevel: "medium", maxTurns: 25, }); const spec = parseGenerationResponse(json); expect(spec.title).toBe("Senior Frontend Reviewer"); expect(spec.icon).toBe("🔍"); expect(spec.role).toBe("reviewer"); expect(spec.description).toBe("Reviews frontend code for quality"); expect(spec.systemPrompt).toBe("# Role\nYou are a senior frontend reviewer."); expect(spec.thinkingLevel).toBe("medium"); expect(spec.maxTurns).toBe(25); }); it("parses JSON wrapped in markdown code block", () => { const inner = JSON.stringify({ title: "Test Agent", icon: "🤖", role: "custom", description: "A test agent", systemPrompt: "Test prompt", thinkingLevel: "off", maxTurns: 10, }); const wrapped = "```json\n" + inner + "\n```"; const spec = parseGenerationResponse(wrapped); expect(spec.title).toBe("Test Agent"); }); it("parses JSON with surrounding text", () => { const json = JSON.stringify({ title: "Test Agent", icon: "🤖", role: "custom", description: "A test agent", systemPrompt: "Test prompt", thinkingLevel: "off", maxTurns: 10, }); const text = `Here is the specification:\n${json}\nHope this helps!`; const spec = parseGenerationResponse(text); expect(spec.title).toBe("Test Agent"); }); it("applies defaults for missing fields", () => { const json = JSON.stringify({}); const spec = parseGenerationResponse(json); expect(spec.title).toBe("Custom Agent"); expect(spec.icon).toBe("🤖"); expect(spec.role).toBe("custom"); expect(spec.description).toBe(""); expect(spec.systemPrompt).toBe(""); expect(spec.thinkingLevel).toBe("off"); expect(spec.maxTurns).toBe(10); }); it("truncates title to 60 characters", () => { const longTitle = "A".repeat(100); const json = JSON.stringify({ title: longTitle, icon: "🤖", role: "custom", description: "test", systemPrompt: "test", thinkingLevel: "off", maxTurns: 10, }); const spec = parseGenerationResponse(json); expect(spec.title.length).toBe(60); }); it("clamps maxTurns to valid range", () => { const json = JSON.stringify({ title: "Test", maxTurns: 999, }); const spec = parseGenerationResponse(json); expect(spec.maxTurns).toBe(500); const json2 = JSON.stringify({ title: "Test", maxTurns: -5 }); const spec2 = parseGenerationResponse(json2); expect(spec2.maxTurns).toBe(1); // clamped to minimum }); it("defaults invalid thinkingLevel to off", () => { const json = JSON.stringify({ title: "Test", thinkingLevel: "ultra", }); const spec = parseGenerationResponse(json); expect(spec.thinkingLevel).toBe("off"); }); it("repairs JSON with trailing commas", () => { const broken = '{"title":"Test","icon":"X","role":"custom","description":"d","systemPrompt":"s","thinkingLevel":"off","maxTurns":10,}'; const spec = parseGenerationResponse(broken); expect(spec.title).toBe("Test"); }); it("throws for non-JSON text", () => { expect(() => parseGenerationResponse("Hello world, this is not JSON")).toThrow( "AI returned no valid JSON" ); }); it("throws for empty text", () => { expect(() => parseGenerationResponse("")).toThrow("AI returned no valid JSON"); }); }); describe("TTL cleanup", () => { it("sessions are retrievable within TTL", async () => { const mockIp = getUniqueIp(); const session = await startAgentGeneration(mockIp, "Test role"); // Session should exist within TTL expect(getAgentGenerationSession(session.id)).toBeDefined(); // Advance to just before TTL vi.advanceTimersByTime(29 * 60 * 1000); expect(getAgentGenerationSession(session.id)).toBeDefined(); }); it("session data is accessible after creation", async () => { const mockIp = getUniqueIp(); const session = await startAgentGeneration(mockIp, "Security auditor role"); const retrieved = getAgentGenerationSession(session.id); expect(retrieved).toBeDefined(); expect(retrieved!.roleDescription).toBe("Security auditor role"); }); }); });