// Copyright (c) 2026 Doforu // Licensed under the GNU Affero General Public License v3.0 (AGPL-3.0-only). // See the LICENSE file or at https://www.gnu.org/licenses/agpl-3.0.txt /** * Tests for the AI Chat SSE endpoint. * * To run: npx tsx --test workers/ai-chat.test.ts * (requires tsx: npm install -D tsx) * * These tests verify the AI Chat SSE endpoint logic using mocked * D1 database and Workers AI bindings. */ import assert from "node:assert"; import { describe, it, mock } from "node:test"; // Imported for its side effect of loading the real module: mock.module() below // replaces the whole module, so we spread these exports and override only the // few functions under test. Without the spread, any other export that the app // pulls in fails to resolve at import time with "does not provide an export // named ...". import * as actualDb from "./db"; // ── Mock data ────────────────────────────────────────────────────── const MOCK_HISTORY = [ { id: "h1", role: "user", content: "hello", created_at: "2026-01-01T00:00:00Z" }, { id: "h2", role: "assistant", content: "hi there", created_at: "2026-01-01T00:00:01Z" }, ]; const MOCK_EMAIL = { id: "email-123", mailbox_id: "test@example.com", folder_id: "inbox", subject: "Test Subject", sender: "sender@example.com", recipient: "test@example.com", cc: null, bcc: null, date: "2026-01-01T00:00:00Z", read: false, starred: false, body: "This is a test email body for the AI to reference", in_reply_to: null, email_references: null, thread_id: null, message_id: null, raw_headers: null, attachments: [], }; const MOCK_THREAD = [ { ...MOCK_EMAIL, id: "t1", subject: "Earlier thread email", body: "Earlier body content", }, { ...MOCK_EMAIL, id: "t2", subject: "Later thread email", body: "Later body content", }, ]; // ── Mock db functions (set up BEFORE app import) ─────────────────── // Use mock.module() to intercept the db module import, since ESM // module namespace objects are non-configurable and mock.method() won't work. const mockSaveAiMessage = mock.fn< (db: any, mailboxId: string, role: string, content: string) => Promise<{ id: string }> >(async () => ({ id: "mock-msg-id" })); const mockGetAiChatHistory = mock.fn< (db: any, mailboxId: string, limit?: number) => Promise >(async () => []); const mockClearAiChatHistory = mock.fn< (db: any, mailboxId: string) => Promise >(async () => {}); const mockGetEmail = mock.fn< (db: any, mailboxId: string, id: string) => Promise >(async () => null); const mockGetThreadEmails = mock.fn< (db: any, mailboxId: string, threadId: string) => Promise >(async () => []); // ── Admin session used by every request ──────────────────────────── // // Every /api/v1/mailboxes/:mailboxId/* route sits behind requireAuth // (workers/lib/auth.ts), which authenticates a signed-in admin from the // `mailboxes_session` cookie and then looks that token up in the D1 `sessions` // table. These tests therefore present a real session token and stub only the // lookup — the middleware itself still runs untouched, so the whole auth chain // is exercised: a request without a cookie, or with a token the store does not // know, is still rejected with 401 (see the "admin session authentication" // block at the end of the file). const TEST_SESSION_TOKEN = "test-session-token"; const TEST_SESSION_CREATED_AT = new Date(Date.now() - 60 * 1000).toISOString(); type MockSession = { token: string; created_at: string; expires_at: string }; const mockGetSession = mock.fn< (db: any, token: string) => Promise >(async () => ({ token: TEST_SESSION_TOKEN, created_at: TEST_SESSION_CREATED_AT, expires_at: new Date(Date.now() + 60 * 60 * 1000).toISOString(), })); // requireAuth deletes a session row it finds expired; stub that write too so the // expiry path can be exercised without a real D1 binding. const mockDeleteSession = mock.fn<(db: any, token: string) => Promise>( async () => {}, ); const dbMocks = { // Keep every real export, then override only the stubs this file needs. ...actualDb, saveAiMessage: mockSaveAiMessage, getAiChatHistory: mockGetAiChatHistory, clearAiChatHistory: mockClearAiChatHistory, getEmail: mockGetEmail, getThreadEmails: mockGetThreadEmails, getSession: mockGetSession, deleteSession: mockDeleteSession, }; // NOTE: ./lib/auth is deliberately NOT mocked. mock.module() builds a fresh // module instance for the target, which would re-evaluate workers/lib/auth.ts // outside this mock registry and re-bind its `./db` import to the REAL module — // resurrecting the unmocked session lookup. The real requireAuth is left in // place and satisfied through the db stub instead (see the tests at the end of // the file), which is what keeps the authentication path under test. mock.module("./db", { namedExports: dbMocks }); // Now import the app (will use the mocked db module) const { app } = await import("./index"); // ── Helper: reset all mock call tracking ─────────────────────────── function resetMockCalls() { for (const m of [ mockSaveAiMessage, mockGetAiChatHistory, mockClearAiChatHistory, mockGetEmail, mockGetThreadEmails, mockGetSession, mockDeleteSession, ]) { m.mock.resetCalls(); } // Each test starts from an unexpired session unless it says otherwise. mockGetSession.mock.mockImplementation(async () => ({ token: TEST_SESSION_TOKEN, created_at: TEST_SESSION_CREATED_AT, expires_at: new Date(Date.now() + 60 * 60 * 1000).toISOString(), })); mockDeleteSession.mock.mockImplementation(async () => {}); } // ── Mock env builder ─────────────────────────────────────────────── function createMockEnv() { const mockBucket = { head: mock.fn(async (_key: string) => ({ key: "mailboxes/test@example.com.json", })), get: mock.fn(async () => null), put: mock.fn(async () => {}), delete: mock.fn(async () => {}), list: mock.fn(async () => ({ objects: [] })), } as any; const mockAi = { run: mock.fn<(model: string, options: unknown) => Promise>(), }; const env = { DB: {} as any, AI: mockAi, BUCKET: mockBucket, } as any; return { env, mockAi, mockBucket }; } // ── Request helpers ──────────────────────────────────────────────── const BASE = "http://localhost"; /** * The cookie an authenticated admin browser sends. requireAuth reads it with * hono/cookie's getCookie() and validates the token against the (stubbed) * session store before letting the request reach the route. */ const SESSION_COOKIE = `mailboxes_session=${TEST_SESSION_TOKEN}`; function postReq(path: string, body: any, extraHeaders?: Record) { return new Request(`${BASE}${path}`, { method: "POST", headers: { "Content-Type": "application/json", Cookie: SESSION_COOKIE, ...extraHeaders, }, body: JSON.stringify(body), }); } function getReq(path: string) { return new Request(`${BASE}${path}`, { method: "GET", headers: { Cookie: SESSION_COOKIE }, }); } function deleteReq(path: string) { return new Request(`${BASE}${path}`, { method: "DELETE", headers: { Cookie: SESSION_COOKIE }, }); } // ── Tests ────────────────────────────────────────────────────────── describe("AI Chat SSE endpoint", async () => { // ── Test 1: POST - validates message field ─────────────────────── it("should return 400 for empty body", async () => { resetMockCalls(); const { env } = createMockEnv(); const req = postReq("/api/v1/mailboxes/test@example.com/ai/chat", {}); const res = await app.fetch(req, env); assert.strictEqual(res.status, 400); const data = (await res.json()) as { error?: string }; assert.strictEqual(data.error, "message is required"); }); it("should return 400 for non-string message (number)", async () => { resetMockCalls(); const { env } = createMockEnv(); const req = postReq("/api/v1/mailboxes/test@example.com/ai/chat", { message: 123, }); const res = await app.fetch(req, env); assert.strictEqual(res.status, 400); const data = (await res.json()) as { error?: string }; assert.strictEqual(data.error, "message is required"); }); it("should return SSE response for valid message", async () => { resetMockCalls(); const { env, mockAi } = createMockEnv(); mockAi.run.mock.mockImplementation(async () => ({ response: "Hello! I'm your email assistant.", })); const req = postReq( "/api/v1/mailboxes/test@example.com/ai/chat", { message: "hi" }, { Accept: "text/event-stream" }, ); const res = await app.fetch(req, env); assert.strictEqual(res.status, 200); assert.strictEqual(res.headers.get("Content-Type"), "text/event-stream"); const body = await res.text(); assert.ok(body.includes("data: "), "Should have SSE data events"); assert.ok(body.includes('"token":'), "Should have token events"); assert.ok(body.includes('"done":true'), "Should have done event"); }); // ── Test 2: POST - builds AI messages with email context ───────── it("should include email and thread context in AI messages", async () => { resetMockCalls(); const { env, mockAi } = createMockEnv(); // Mock db to return email and thread mockGetEmail.mock.mockImplementation(async () => MOCK_EMAIL); mockGetThreadEmails.mock.mockImplementation(async () => MOCK_THREAD); mockAi.run.mock.mockImplementation(async () => ({ response: "Here is your email summary.", })); const req = postReq( "/api/v1/mailboxes/test@example.com/ai/chat", { message: "find email", emailContext: { emailId: "email-123", threadId: "thread-1" }, }, { Accept: "text/event-stream" }, ); const res = await app.fetch(req, env); assert.strictEqual(res.status, 200); // Verify buildAiMessages was called with email + thread context const callArgs = mockAi.run.mock.calls[0].arguments; assert.ok(callArgs, "ai.run should have been called"); // `run` receives (model, { messages }); the mock's arguments are typed // loosely, so narrow the second argument before reading it. const messages = ( callArgs[1] as { messages?: Array<{ role: string; content: string }> } )?.messages; assert.ok(messages, "messages should be passed to ai.run"); const systemMsgs = messages.filter((m: any) => m.role === "system"); // Email context should be in system messages const hasEmailContext = systemMsgs.some( (m: any) => m.content.includes("sender@example.com") && m.content.includes("Test Subject"), ); assert.ok( hasEmailContext, "Email context should be included in system messages", ); // Thread context should be in system messages const hasThreadContext = systemMsgs.some( (m: any) => m.content.includes("Earlier thread email") || m.content.includes("Earlier body content"), ); assert.ok( hasThreadContext, "Thread context should be included in system messages", ); // Verify db was queried for email and thread assert.strictEqual(mockGetEmail.mock.calls.length, 1); assert.strictEqual(mockGetThreadEmails.mock.calls.length, 1); }); // ── Test 3: POST - model fallback works ────────────────────────── it("should fallback to second model when first fails", async () => { resetMockCalls(); const { env, mockAi } = createMockEnv(); let callCount = 0; mockAi.run.mock.mockImplementation(async () => { callCount++; if (callCount === 1) throw new Error("Model A unavailable"); return { response: "Fallback model response" }; }); const req = postReq( "/api/v1/mailboxes/test@example.com/ai/chat", { message: "hello" }, { Accept: "text/event-stream" }, ); const res = await app.fetch(req, env); assert.strictEqual(res.status, 200); // Both models should have been called (primary + fallback) assert.strictEqual( mockAi.run.mock.calls.length, 2, "ai.run should be called twice (primary + fallback)", ); // Parse SSE events to reconstruct full text from token chunks const body = await res.text(); const sseEvents = body.split("\n").filter((l) => l.startsWith("data: ")); const fullText = sseEvents .map((l) => { try { const parsed = JSON.parse(l.slice(6)); return parsed.token || ""; } catch { return ""; } }) .join(""); assert.ok( fullText.includes("Fallback model response"), "Should contain fallback model output, not empty string", ); assert.ok(body.includes('"done":true'), "Should have done event"); }); // ── Test 4: POST - returns SSE error when both models fail ─────── it("should return SSE error event when both AI models fail", async () => { resetMockCalls(); const { env, mockAi } = createMockEnv(); mockAi.run.mock.mockImplementation(async () => { throw new Error("AI service unavailable"); }); const req = postReq( "/api/v1/mailboxes/test@example.com/ai/chat", { message: "hello" }, { Accept: "text/event-stream" }, ); const res = await app.fetch(req, env); assert.strictEqual(res.status, 200); assert.strictEqual(res.headers.get("Content-Type"), "text/event-stream"); const body = await res.text(); assert.ok( body.includes("AI temporarily unavailable"), "SSE error should mention AI unavailability", ); assert.ok(body.includes("error"), "SSE response should have error field"); }); // ── Test 5: GET - returns chat history ─────────────────────────── it("should return chat history from GET endpoint", async () => { resetMockCalls(); mockGetAiChatHistory.mock.mockImplementation(async () => MOCK_HISTORY); const { env } = createMockEnv(); const req = getReq("/api/v1/mailboxes/test@example.com/ai/chat"); const res = await app.fetch(req, env); assert.strictEqual(res.status, 200); const data = (await res.json()) as { messages: { role: string; content: string }[] }; assert.ok(data.messages, "Response should have messages array"); assert.strictEqual(data.messages.length, 2); assert.strictEqual(data.messages[0].role, "user"); assert.strictEqual(data.messages[0].content, "hello"); assert.strictEqual(data.messages[1].role, "assistant"); assert.strictEqual(data.messages[1].content, "hi there"); }); // ── Test 6: DELETE - clears history ────────────────────────────── it("should clear chat history from DELETE endpoint", async () => { resetMockCalls(); const { env } = createMockEnv(); const req = deleteReq("/api/v1/mailboxes/test@example.com/ai/chat"); const res = await app.fetch(req, env); assert.strictEqual(res.status, 204); // Verify clearAiChatHistory was called with correct params assert.strictEqual(mockClearAiChatHistory.mock.calls.length, 1); const args = mockClearAiChatHistory.mock.calls[0].arguments; assert.strictEqual(args[1], "test@example.com"); }); // ── Test 7: POST - validation rejects empty string ─────────────── it("should return 400 for empty string message", async () => { resetMockCalls(); const { env } = createMockEnv(); const req = postReq("/api/v1/mailboxes/test@example.com/ai/chat", { message: "", }); const res = await app.fetch(req, env); assert.strictEqual(res.status, 400); const data = (await res.json()) as { error?: string }; assert.strictEqual(data.error, "message is required"); }); // ── Test 8: POST - validation rejects overly long message ──────── it("should return 400 for overly long message (>10000 chars)", async () => { resetMockCalls(); const { env } = createMockEnv(); const longMsg = "a".repeat(10001); const req = postReq("/api/v1/mailboxes/test@example.com/ai/chat", { message: longMsg, }); const res = await app.fetch(req, env); assert.strictEqual(res.status, 400); const data = (await res.json()) as { error?: string }; assert.strictEqual( data.error, "message too long", "Should reject overly long message", ); }); }); // ── Admin session authentication ─────────────────────────────────── // // Guards the fix above: the requests in this file only reach the AI Chat routes // because they carry a valid session. These cases pin that the guard itself is // still in place, so the rest of the file cannot be made green by dropping the // cookie — or by weakening the middleware. /** * The `hint` that `requireAuth`'s 401 now carries, spelled out here as a * literal (not imported from the module under test) and asserted alongside the * unchanged `error`. This is what the browser-session 401 sends to a caller * that presented no usable session cookie. */ const UNAUTHORIZED_HINT = "Sign in at /login in your browser and send the session cookie with the request, or call POST /api/v1/auth/login to obtain one."; describe("AI Chat endpoint authentication", () => { it("rejects a request without a session cookie", async () => { resetMockCalls(); const { env } = createMockEnv(); const req = new Request(`${BASE}/api/v1/mailboxes/test@example.com/ai/chat`, { method: "GET", }); const res = await app.fetch(req, env); assert.strictEqual(res.status, 401); assert.strictEqual(sessionDbQueries(), 0, "an anonymous request must not query the session store"); assert.deepStrictEqual(await res.json(), { error: "Unauthorized", hint: UNAUTHORIZED_HINT, }); }); it("rejects a session token that is not in the store", async () => { resetMockCalls(); mockGetSession.mock.mockImplementationOnce(async () => null); const { env } = createMockEnv(); const req = new Request(`${BASE}/api/v1/mailboxes/test@example.com/ai/chat`, { method: "GET", headers: { Cookie: "mailboxes_session=not-a-real-token" }, }); const res = await app.fetch(req, env); assert.strictEqual(res.status, 401); assert.strictEqual(sessionDbQueries(), 1, "the presented token must be looked up"); assert.strictEqual( mockGetSession.mock.calls[0].arguments[1], "not-a-real-token", "the middleware must check the token the client actually sent", ); }); it("rejects an expired session", async () => { resetMockCalls(); mockGetSession.mock.mockImplementationOnce(async () => ({ token: TEST_SESSION_TOKEN, created_at: new Date(Date.now() - 8 * 24 * 60 * 60 * 1000).toISOString(), expires_at: new Date(Date.now() - 24 * 60 * 60 * 1000).toISOString(), })); const { env } = createMockEnv(); const req = getReq("/api/v1/mailboxes/test@example.com/ai/chat"); const res = await app.fetch(req, env); assert.strictEqual(res.status, 401); assert.strictEqual( mockGetSession.mock.calls[0].arguments[1], TEST_SESSION_TOKEN, "the expiry check must run against the stored session row", ); }); }); /** Number of times the (stubbed) D1 session lookup has been consulted. */ function sessionDbQueries(): number { return mockGetSession.mock.calls.length; }