/** * @vitest-environment node */ import { beforeEach, describe, expect, it, vi } from 'vitest' const { e2bFlag, betaFlag, mockLoadCompiledDoc, mockRunSandboxTask } = vi.hoisted(() => ({ e2bFlag: { value: true }, betaFlag: { value: false }, mockLoadCompiledDoc: vi.fn(), mockRunSandboxTask: vi.fn(), })) vi.mock('@/lib/execution/e2b', () => ({ executeInE2B: vi.fn(), executeShellInE2B: vi.fn(), })) vi.mock('@/lib/execution/languages', () => ({ CodeLanguage: { javascript: 'javascript', python: 'python' }, })) vi.mock('@/lib/execution/sandbox/run-task', () => ({ runSandboxTask: mockRunSandboxTask, })) vi.mock('@/lib/uploads/contexts/workspace/workspace-file-manager', () => ({ getWorkspaceFile: vi.fn(), fetchWorkspaceFileBuffer: vi.fn(), })) vi.mock('./doc-compiled-store', () => ({ loadCompiledDoc: mockLoadCompiledDoc, storeCompiledDoc: vi.fn(), })) vi.mock('@/lib/core/config/feature-flags', () => ({ isFeatureEnabled: vi.fn(async () => betaFlag.value), })) vi.mock('@/lib/core/config/env-flags', () => ({ get isE2BDocEnabled() { return e2bFlag.value }, })) vi.mock('@/app/api/files/utils', () => ({ getContentType: (name: string) => name.endsWith('.pdf') ? 'application/pdf' : name.endsWith('.txt') ? 'text/plain' : 'application/octet-stream', })) import { DocCompileUserError, resolveServableDocBytes } from './doc-compile' const WORKSPACE_ID = '550e8400-e29b-41d4-a716-446655440000' const PDF_MAGIC = Buffer.from('%PDF-1.7\n...binary...') const PDF_SOURCE = Buffer.from('from reportlab.pdfgen import canvas\n# generates a PDF', 'utf-8') const ZIP_MAGIC = Buffer.from([0x50, 0x4b, 0x03, 0x04, 0x00, 0x01]) const XLSX_SOURCE = Buffer.from('from openpyxl import Workbook\n# generates an xlsx', 'utf-8') describe('resolveServableDocBytes', () => { beforeEach(() => { vi.clearAllMocks() e2bFlag.value = true betaFlag.value = false }) it('swaps generated-doc source for the compiled artifact + binary content type', async () => { const artifact = Buffer.from('%PDF-compiled-binary') mockLoadCompiledDoc.mockResolvedValue(artifact) const result = await resolveServableDocBytes({ rawBuffer: PDF_SOURCE, fileName: 'report.pdf', workspaceId: WORKSPACE_ID, }) expect(result.buffer).toBe(artifact) expect(result.contentType).toBe('application/pdf') expect(mockLoadCompiledDoc).toHaveBeenCalledWith( WORKSPACE_ID, PDF_SOURCE.toString('utf-8'), 'pdf' ) }) it('passes through a real binary PDF (carries the %PDF magic) without an artifact lookup', async () => { const result = await resolveServableDocBytes({ rawBuffer: PDF_MAGIC, fileName: 'uploaded.pdf', workspaceId: WORKSPACE_ID, }) expect(result.buffer).toBe(PDF_MAGIC) expect(result.contentType).toBe('application/pdf') expect(mockLoadCompiledDoc).not.toHaveBeenCalled() }) it('throws DocCompileUserError when a generated doc artifact is not ready (E2B regime)', async () => { mockLoadCompiledDoc.mockResolvedValue(null) e2bFlag.value = true await expect( resolveServableDocBytes({ rawBuffer: PDF_SOURCE, fileName: 'report.pdf', workspaceId: WORKSPACE_ID, }) ).rejects.toBeInstanceOf(DocCompileUserError) expect(mockRunSandboxTask).not.toHaveBeenCalled() }) it('compiles via the sandbox when E2B is disabled and no artifact is stored', async () => { mockLoadCompiledDoc.mockResolvedValue(null) e2bFlag.value = false const compiled = Buffer.from('%PDF-isolated-vm-binary') mockRunSandboxTask.mockResolvedValue(compiled) const result = await resolveServableDocBytes({ rawBuffer: PDF_SOURCE, fileName: 'report.pdf', workspaceId: WORKSPACE_ID, }) expect(result.buffer).toBe(compiled) expect(result.contentType).toBe('application/pdf') expect(mockRunSandboxTask).toHaveBeenCalledWith( 'pdf-generate', { code: PDF_SOURCE.toString('utf-8'), workspaceId: WORKSPACE_ID }, expect.objectContaining({}) ) }) it('passes non-doc files through untouched with their extension content type', async () => { const text = Buffer.from('hello world', 'utf-8') const result = await resolveServableDocBytes({ rawBuffer: text, fileName: 'notes.txt', workspaceId: WORKSPACE_ID, }) expect(result.buffer).toBe(text) expect(result.contentType).toBe('text/plain') expect(mockLoadCompiledDoc).not.toHaveBeenCalled() }) it('passes through a real binary XLSX (ZIP magic) without an artifact lookup', async () => { const result = await resolveServableDocBytes({ rawBuffer: ZIP_MAGIC, fileName: 'sheet.xlsx', workspaceId: WORKSPACE_ID, }) expect(result.buffer).toBe(ZIP_MAGIC) expect(mockLoadCompiledDoc).not.toHaveBeenCalled() }) it('throws when a generated XLSX artifact is not ready (E2B + mothership-beta enabled)', async () => { mockLoadCompiledDoc.mockResolvedValue(null) e2bFlag.value = true betaFlag.value = true await expect( resolveServableDocBytes({ rawBuffer: XLSX_SOURCE, fileName: 'sheet.xlsx', workspaceId: WORKSPACE_ID, }) ).rejects.toBeInstanceOf(DocCompileUserError) expect(mockRunSandboxTask).not.toHaveBeenCalled() }) it('returns raw XLSX source when there is no workspaceId (xlsx has no isolated-vm path)', async () => { betaFlag.value = true const result = await resolveServableDocBytes({ rawBuffer: XLSX_SOURCE, fileName: 'sheet.xlsx', workspaceId: undefined, }) expect(result.buffer).toBe(XLSX_SOURCE) expect(mockLoadCompiledDoc).not.toHaveBeenCalled() expect(mockRunSandboxTask).not.toHaveBeenCalled() }) })