import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; import express from "express"; import request from "supertest"; import { puterProxyService } from "../services/puter-proxy.js"; import { puterProxyRoutes } from "../routes/puter-proxy.js"; // ─── Helpers ────────────────────────────────────────────────────────────────── function makeTextEncoder() { return new TextEncoder(); } /** * Build a minimal ReadableStream that emits pre-formatted SSE lines. */ function buildSseStream(chunks: string[]): ReadableStream { const enc = makeTextEncoder(); let idx = 0; return new ReadableStream({ pull(controller) { if (idx < chunks.length) { controller.enqueue(enc.encode(chunks[idx++])); } else { controller.close(); } }, }); } function sseChunk(data: object): string { return `data: ${JSON.stringify(data)}\n\n`; } // ─── Shared mock db factory ─────────────────────────────────────────────────── function makeMockDb(overrides: { existingSecret?: object | null; resolvedToken?: string; createSecret?: ReturnType; rotateSecret?: ReturnType; createCostEvent?: ReturnType; } = {}) { const { existingSecret = null, resolvedToken = "test-bearer-token", createSecret = vi.fn().mockResolvedValue({ id: "secret-1", name: "puter_auth_token" }), rotateSecret = vi.fn().mockResolvedValue({ id: "secret-1", name: "puter_auth_token" }), createCostEvent = vi.fn().mockResolvedValue({ id: "event-1" }), } = overrides; // We return a db object that puterProxyService will pass to secretService and costService. // Since we mock at the module level via the db interactions, we need the service // to call these. We'll mock the actual db calls. return { _mocks: { createSecret, rotateSecret, createCostEvent }, // The secretService and costService are created inside puterProxyService using the db. // We inject mock data via the module mocking strategy below. existingSecret, resolvedToken, }; } // ─── Mock secretService and costService at module level ────────────────────── const mockGetByName = vi.fn(); const mockCreate = vi.fn(); const mockRotate = vi.fn(); const mockResolveSecretValue = vi.fn(); const mockCreateEvent = vi.fn(); vi.mock("../services/secrets.js", () => ({ secretService: vi.fn(() => ({ getByName: mockGetByName, create: mockCreate, rotate: mockRotate, resolveSecretValue: mockResolveSecretValue, })), })); vi.mock("../services/costs.js", () => ({ costService: vi.fn(() => ({ createEvent: mockCreateEvent, })), })); // ─── Test suite ─────────────────────────────────────────────────────────────── describe("puterProxyService", () => { const companyId = "company-123"; const db = {} as any; beforeEach(() => { vi.clearAllMocks(); }); afterEach(() => { vi.unstubAllGlobals(); }); // Test 1: storeToken creates a new secret when none exists it("storeToken creates a new secret via secretService when none exists", async () => { mockGetByName.mockResolvedValue(null); mockCreate.mockResolvedValue({ id: "secret-1", name: "puter_auth_token" }); const svc = puterProxyService(db); await svc.storeToken(companyId, "my-puter-token"); expect(mockGetByName).toHaveBeenCalledWith(companyId, "puter_auth_token"); expect(mockCreate).toHaveBeenCalledWith( companyId, expect.objectContaining({ name: "puter_auth_token", provider: "local_encrypted", value: "my-puter-token", }), ); expect(mockRotate).not.toHaveBeenCalled(); }); // Test 2: storeToken rotates existing secret it("storeToken rotates an existing secret when puter_auth_token already exists", async () => { const existingSecret = { id: "existing-secret-id", name: "puter_auth_token" }; mockGetByName.mockResolvedValue(existingSecret); mockRotate.mockResolvedValue({ id: "existing-secret-id", name: "puter_auth_token" }); const svc = puterProxyService(db); await svc.storeToken(companyId, "new-puter-token"); expect(mockGetByName).toHaveBeenCalledWith(companyId, "puter_auth_token"); expect(mockRotate).toHaveBeenCalledWith( "existing-secret-id", expect.objectContaining({ value: "new-puter-token" }), ); expect(mockCreate).not.toHaveBeenCalled(); }); // Test 3: resolveToken retrieves stored token value it("resolveToken retrieves the stored token value via secretService.resolveSecretValue", async () => { const existingSecret = { id: "secret-id", name: "puter_auth_token" }; mockGetByName.mockResolvedValue(existingSecret); mockResolveSecretValue.mockResolvedValue("resolved-bearer-token"); const svc = puterProxyService(db); const token = await svc.resolveToken(companyId); expect(mockGetByName).toHaveBeenCalledWith(companyId, "puter_auth_token"); expect(mockResolveSecretValue).toHaveBeenCalledWith(companyId, "secret-id", "latest"); expect(token).toBe("resolved-bearer-token"); }); // Test 4: chatStream sends POST to Puter with correct headers and body it("chatStream sends POST to Puter OpenAI-compat endpoint with Authorization Bearer header, stream: true, stream_options", async () => { mockGetByName.mockResolvedValue({ id: "sec-1", name: "puter_auth_token" }); mockResolveSecretValue.mockResolvedValue("my-bearer-token"); mockCreateEvent.mockResolvedValue({ id: "ev-1" }); const enc = new TextEncoder(); const chunks = [ sseChunk({ choices: [{ delta: { content: "Hello" } }] }), sseChunk({ choices: [{ delta: { content: " world" } }] }), sseChunk({ choices: [{ delta: {} }], usage: { prompt_tokens: 10, completion_tokens: 5 } }), "data: [DONE]\n\n", ]; const mockStream = buildSseStream(chunks); const mockResponse = { ok: true, body: mockStream, }; const fetchSpy = vi.fn().mockResolvedValue(mockResponse); vi.stubGlobal("fetch", fetchSpy); const svc = puterProxyService(db); const messages = [{ role: "user", content: "Hi" }]; const result: string[] = []; for await (const chunk of svc.chatStream(companyId, "agent-1", messages, "claude-3-5-haiku-20241022", undefined)) { result.push(chunk); } expect(fetchSpy).toHaveBeenCalledOnce(); const [url, opts] = fetchSpy.mock.calls[0]; expect(url).toContain("api.puter.com/puterai/openai/v1/chat/completions"); expect(opts.method).toBe("POST"); expect(opts.headers["Authorization"]).toBe("Bearer my-bearer-token"); const body = JSON.parse(opts.body); expect(body.stream).toBe(true); expect(body.stream_options).toEqual({ include_usage: true }); }); // Test 5: chatStream yields content strings from SSE data chunks it("chatStream yields content strings from SSE data chunks", async () => { mockGetByName.mockResolvedValue({ id: "sec-1", name: "puter_auth_token" }); mockResolveSecretValue.mockResolvedValue("my-bearer-token"); mockCreateEvent.mockResolvedValue({ id: "ev-1" }); const chunks = [ sseChunk({ choices: [{ delta: { content: "Hello" } }] }), sseChunk({ choices: [{ delta: { content: " world" } }] }), sseChunk({ choices: [{ delta: {} }], usage: { prompt_tokens: 5, completion_tokens: 3 } }), "data: [DONE]\n\n", ]; vi.stubGlobal("fetch", vi.fn().mockResolvedValue({ ok: true, body: buildSseStream(chunks), })); const svc = puterProxyService(db); const result: string[] = []; for await (const chunk of svc.chatStream(companyId, "agent-1", [{ role: "user", content: "hi" }], undefined, undefined)) { if (chunk) result.push(chunk); } expect(result).toContain("Hello"); expect(result).toContain(" world"); }); // Test 6: chatStream records cost event when agentId is provided it("chatStream records a cost event with provider=puter, billingType=subscription_included, costCents=0 when agentId is provided", async () => { mockGetByName.mockResolvedValue({ id: "sec-1", name: "puter_auth_token" }); mockResolveSecretValue.mockResolvedValue("my-bearer-token"); mockCreateEvent.mockResolvedValue({ id: "ev-1" }); const chunks = [ sseChunk({ choices: [{ delta: { content: "Hi" } }], usage: { prompt_tokens: 10, completion_tokens: 5 } }), "data: [DONE]\n\n", ]; vi.stubGlobal("fetch", vi.fn().mockResolvedValue({ ok: true, body: buildSseStream(chunks), })); const svc = puterProxyService(db); for await (const _ of svc.chatStream(companyId, "agent-xyz", [{ role: "user", content: "hi" }], "claude-3-5-haiku-20241022", undefined)) { // consume stream } // Allow microtask queue to flush (non-blocking cost recording) await new Promise((resolve) => setTimeout(resolve, 50)); expect(mockCreateEvent).toHaveBeenCalledOnce(); expect(mockCreateEvent).toHaveBeenCalledWith( companyId, expect.objectContaining({ agentId: "agent-xyz", provider: "puter", billingType: "subscription_included", costCents: 0, }), ); }); // Test 7: chatStream skips cost recording when agentId is null/undefined it("chatStream skips cost recording when agentId is null/undefined", async () => { mockGetByName.mockResolvedValue({ id: "sec-1", name: "puter_auth_token" }); mockResolveSecretValue.mockResolvedValue("my-bearer-token"); const chunks = [ sseChunk({ choices: [{ delta: { content: "No cost" } }] }), "data: [DONE]\n\n", ]; vi.stubGlobal("fetch", vi.fn().mockResolvedValue({ ok: true, body: buildSseStream(chunks), })); const svc = puterProxyService(db); for await (const _ of svc.chatStream(companyId, null, [{ role: "user", content: "hi" }], undefined, undefined)) { // consume } await new Promise((resolve) => setTimeout(resolve, 50)); expect(mockCreateEvent).not.toHaveBeenCalled(); }); }); // ─── Route tests ───────────────────────────────────────────────────────────── function buildTestApp(db: any) { const app = express(); app.use(express.json()); // Inject a board actor for all requests app.use((req: any, _res: any, next: any) => { req.actor = { type: "board", source: "local_implicit", companyIds: ["company-123"], isInstanceAdmin: true, }; next(); }); app.use(puterProxyRoutes(db)); return app; } describe("puterProxyRoutes", () => { const db = {} as any; beforeEach(() => { vi.clearAllMocks(); }); afterEach(() => { vi.unstubAllGlobals(); }); // Test 8: POST /api/puter-proxy/token stores token and returns 200 it("POST /puter-proxy/token stores token and returns 200", async () => { mockGetByName.mockResolvedValue(null); mockCreate.mockResolvedValue({ id: "sec-1", name: "puter_auth_token" }); const app = buildTestApp(db); const res = await request(app) .post("/puter-proxy/token") .send({ companyId: "company-123", token: "puter-abc123" }); expect(res.status).toBe(200); expect(res.body).toEqual({ ok: true }); expect(mockCreate).toHaveBeenCalledWith( "company-123", expect.objectContaining({ name: "puter_auth_token", value: "puter-abc123" }), ); }); // Test 9: POST /api/puter-proxy/chat sets SSE headers and streams response chunks it("POST /puter-proxy/chat sets SSE headers and streams response chunks", async () => { mockGetByName.mockResolvedValue({ id: "sec-1", name: "puter_auth_token" }); mockResolveSecretValue.mockResolvedValue("bearer-xyz"); mockCreateEvent.mockResolvedValue({ id: "ev-1" }); const chunks = [ sseChunk({ choices: [{ delta: { content: "Stream token" } }] }), "data: [DONE]\n\n", ]; vi.stubGlobal("fetch", vi.fn().mockResolvedValue({ ok: true, body: buildSseStream(chunks), })); const app = buildTestApp(db); const res = await request(app) .post("/puter-proxy/chat") .send({ companyId: "company-123", agentId: "agent-1", messages: [{ role: "user", content: "Hello" }], }); expect(res.status).toBe(200); expect(res.headers["content-type"]).toMatch(/text\/event-stream/); expect(res.text).toContain("Stream token"); }); // Test 10: POST /api/puter-proxy/chat works without agentId in request body it("POST /puter-proxy/chat works without agentId in request body (agentId is optional)", async () => { mockGetByName.mockResolvedValue({ id: "sec-1", name: "puter_auth_token" }); mockResolveSecretValue.mockResolvedValue("bearer-xyz"); const chunks = [ sseChunk({ choices: [{ delta: { content: "No agent" } }] }), "data: [DONE]\n\n", ]; vi.stubGlobal("fetch", vi.fn().mockResolvedValue({ ok: true, body: buildSseStream(chunks), })); const app = buildTestApp(db); const res = await request(app) .post("/puter-proxy/chat") .send({ companyId: "company-123", // No agentId messages: [{ role: "user", content: "Hello without agent" }], }); expect(res.status).toBe(200); expect(res.headers["content-type"]).toMatch(/text\/event-stream/); expect(res.text).toContain("No agent"); await new Promise((resolve) => setTimeout(resolve, 50)); // No cost event since no agentId expect(mockCreateEvent).not.toHaveBeenCalled(); }); });