import { mkdir, mkdtemp, readFile, rm, writeFile } from 'node:fs/promises'; import { tmpdir } from 'node:os'; import { join } from 'node:path'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { contextBuildCommands, readKloSetupContextState, runKloSetupContextCommand, runKloSetupContextStep, writeKloSetupContextState, } from './setup-context.js'; function makeIo() { let stdout = ''; let stderr = ''; return { io: { stdout: { write: (chunk: string) => { stdout += chunk; }, }, stderr: { write: (chunk: string) => { stderr += chunk; }, }, }, stdout: () => stdout, stderr: () => stderr, }; } async function writeReadyProject(projectDir: string) { await writeFile( join(projectDir, 'klo.yaml'), [ 'project: revenue', 'setup:', ' database_connection_ids:', ' - warehouse', ' completed_steps:', ' - project', ' - llm', ' - embeddings', ' - databases', ' - sources', 'connections:', ' warehouse:', ' driver: postgres', ' url: env:DATABASE_URL', ' docs:', ' driver: notion', ' auth_token_ref: env:NOTION_TOKEN', ' crawl_mode: all_accessible', 'llm:', ' provider:', ' backend: anthropic', ' models:', ' default: claude-sonnet-4-6', 'ingest:', ' embeddings:', ' backend: openai', ' model: text-embedding-3-small', ' dimensions: 1536', 'scan:', ' enrichment:', ' mode: llm', '', ].join('\n'), 'utf-8', ); } async function writeScanReport( projectDir: string, syncId: string, report: { mode: string; tableDescriptions: string; columnDescriptions: string; embeddings: string }, ) { const reportDir = join(projectDir, 'raw-sources', 'warehouse', 'live-database', syncId); await mkdir(reportDir, { recursive: true }); await writeFile( join(reportDir, 'scan-report.json'), `${JSON.stringify( { connectionId: 'warehouse', mode: report.mode, dryRun: false, artifactPaths: { manifestShards: ['semantic-layer/warehouse/_schema/public.yaml'], enrichmentArtifacts: report.mode === 'enriched' ? [`raw-sources/warehouse/live-database/${syncId}/enrichment/descriptions.json`] : [], }, enrichment: { tableDescriptions: report.tableDescriptions, columnDescriptions: report.columnDescriptions, embeddings: report.embeddings, }, enrichmentState: { completedStages: report.tableDescriptions === 'completed' ? ['descriptions', 'embeddings'] : [], failedStages: report.tableDescriptions === 'failed' ? ['descriptions'] : [], }, createdAt: syncId, }, null, 2, )}\n`, ); } async function writeReadyEnrichedScanReport(projectDir: string, syncId = '2026-05-09T10:00:00.000Z') { await writeScanReport(projectDir, syncId, { mode: 'enriched', tableDescriptions: 'completed', columnDescriptions: 'completed', embeddings: 'completed', }); } describe('setup context build state', () => { let tempDir: string; beforeEach(async () => { tempDir = await mkdtemp(join(tmpdir(), 'klo-setup-context-')); }); afterEach(async () => { await rm(tempDir, { recursive: true, force: true }); }); it('reads missing state as not started and writes durable command metadata without secrets', async () => { await expect(readKloSetupContextState(tempDir)).resolves.toMatchObject({ status: 'not_started' }); await writeKloSetupContextState(tempDir, { runId: 'setup-context-local-abc123', status: 'running', startedAt: '2026-05-09T10:00:00.000Z', updatedAt: '2026-05-09T10:00:00.000Z', primarySourceConnectionIds: ['warehouse'], contextSourceConnectionIds: ['docs'], reportIds: [], artifactPaths: [], retryableFailedTargets: [], commands: contextBuildCommands(tempDir, 'setup-context-local-abc123'), }); const state = await readKloSetupContextState(tempDir); expect(state).toMatchObject({ runId: 'setup-context-local-abc123', status: 'running', primarySourceConnectionIds: ['warehouse'], contextSourceConnectionIds: ['docs'], commands: { watch: `klo setup context watch setup-context-local-abc123 --project-dir ${tempDir}`, status: `klo setup context status setup-context-local-abc123 --project-dir ${tempDir}`, resume: `klo setup --project-dir ${tempDir}`, }, }); expect(JSON.stringify(state)).not.toContain('DATABASE_URL'); expect(JSON.stringify(state)).not.toContain('NOTION_TOKEN'); }); it('runs setup context build, verifies readiness, and marks context complete', async () => { await writeReadyProject(tempDir); const io = makeIo(); const runContextBuildMock = vi.fn(async () => ({ exitCode: 0, detached: false })); const verifyContextReady = vi.fn(async () => ({ ready: true, agentContextReady: true, semanticSearchReady: true, details: ['warehouse: enriched scan complete', 'docs: memory update complete'], })); await expect( runKloSetupContextStep( { projectDir: tempDir, inputMode: 'disabled' }, io.io, { runIdFactory: () => 'setup-context-local-abc123', now: () => new Date('2026-05-09T10:00:00.000Z'), runContextBuild: runContextBuildMock, verifyContextReady, }, ), ).resolves.toEqual({ status: 'ready', projectDir: tempDir, runId: 'setup-context-local-abc123' }); expect(runContextBuildMock).toHaveBeenCalledWith( expect.objectContaining({ projectDir: tempDir }), expect.objectContaining({ projectDir: tempDir, inputMode: 'disabled', scanMode: 'enriched', detectRelationships: true, }), io.io, expect.objectContaining({ onDetach: expect.any(Function) }), ); expect(verifyContextReady).toHaveBeenCalledWith(tempDir); expect(await readFile(join(tempDir, 'klo.yaml'), 'utf-8')).toContain(' - context'); await expect(readKloSetupContextState(tempDir)).resolves.toMatchObject({ runId: 'setup-context-local-abc123', status: 'completed', completedAt: '2026-05-09T10:00:00.000Z', }); expect(io.stdout()).toContain('KLO context is ready for agents.'); }); it('marks context complete without prompting when initial source ingest already made agent context', async () => { await writeReadyProject(tempDir); await mkdir(join(tempDir, 'semantic-layer', 'dbt-main'), { recursive: true }); await mkdir(join(tempDir, 'knowledge', 'global'), { recursive: true }); await writeFile(join(tempDir, 'semantic-layer', 'dbt-main', 'mart_revenue_daily.yaml'), 'name: mart_revenue_daily\n'); await writeFile(join(tempDir, 'knowledge', 'global', 'metrics.md'), '# Metrics\n'); await writeReadyEnrichedScanReport(tempDir); const io = makeIo(); const runContextBuildMock = vi.fn(async () => ({ exitCode: 0, detached: false })); await expect( runKloSetupContextStep( { projectDir: tempDir, inputMode: 'auto' }, io.io, { prompts: { select: vi.fn(async () => { throw new Error('setup should not prompt when context is already ready'); }), cancel: vi.fn(), }, runIdFactory: () => 'setup-context-local-existing', now: () => new Date('2026-05-09T10:00:00.000Z'), runContextBuild: runContextBuildMock, }, ), ).resolves.toEqual({ status: 'ready', projectDir: tempDir, runId: 'setup-context-local-existing' }); expect(runContextBuildMock).not.toHaveBeenCalled(); expect(await readFile(join(tempDir, 'klo.yaml'), 'utf-8')).toContain(' - context'); await expect(readKloSetupContextState(tempDir)).resolves.toMatchObject({ runId: 'setup-context-local-existing', status: 'completed', completedAt: '2026-05-09T10:00:00.000Z', contextSourceConnectionIds: ['docs'], }); expect(io.stdout()).toContain('KLO context is ready for agents.'); }); it('does not mark context ready until primary scans have completed description enrichment', async () => { await writeReadyProject(tempDir); await mkdir(join(tempDir, 'semantic-layer', 'dbt-main'), { recursive: true }); await writeFile(join(tempDir, 'semantic-layer', 'dbt-main', 'mart_revenue_daily.yaml'), 'name: mart_revenue_daily\n'); await writeScanReport(tempDir, '2026-05-09T09:59:00.000Z', { mode: 'structural', tableDescriptions: 'skipped', columnDescriptions: 'skipped', embeddings: 'skipped', }); const io = makeIo(); const runContextBuildMock = vi.fn(async () => { await writeReadyEnrichedScanReport(tempDir, '2026-05-09T10:00:00.000Z'); return { exitCode: 0, detached: false }; }); await expect( runKloSetupContextStep( { projectDir: tempDir, inputMode: 'disabled' }, io.io, { runIdFactory: () => 'setup-context-local-enriched-scan', now: () => new Date('2026-05-09T10:00:00.000Z'), runContextBuild: runContextBuildMock, }, ), ).resolves.toEqual({ status: 'ready', projectDir: tempDir, runId: 'setup-context-local-enriched-scan' }); expect(runContextBuildMock).toHaveBeenCalledOnce(); expect(io.stdout()).not.toContain('Existing context artifacts were found from setup ingest.'); }); it('does not treat schema-only scan shards as completed setup context', async () => { await writeReadyProject(tempDir); await mkdir(join(tempDir, 'semantic-layer', 'warehouse', '_schema'), { recursive: true }); await writeFile(join(tempDir, 'semantic-layer', 'warehouse', '_schema', 'public.yaml'), 'tables: {}\n'); const io = makeIo(); const runContextBuildMock = vi.fn(async () => { await mkdir(join(tempDir, 'knowledge', 'global'), { recursive: true }); await writeFile(join(tempDir, 'knowledge', 'global', 'metrics.md'), '# Metrics\n'); await writeReadyEnrichedScanReport(tempDir); return { exitCode: 0, detached: false }; }); await expect( runKloSetupContextStep( { projectDir: tempDir, inputMode: 'disabled' }, io.io, { runIdFactory: () => 'setup-context-local-schema-only', now: () => new Date('2026-05-09T10:00:00.000Z'), runContextBuild: runContextBuildMock, }, ), ).resolves.toEqual({ status: 'ready', projectDir: tempDir, runId: 'setup-context-local-schema-only' }); expect(runContextBuildMock).toHaveBeenCalledOnce(); expect(io.stdout()).not.toContain('Existing context artifacts were found from setup ingest.'); }); it('refuses empty setup context builds', async () => { await writeFile( join(tempDir, 'klo.yaml'), [ 'project: revenue', 'connections: {}', 'llm:', ' provider:', ' backend: anthropic', ' models:', ' default: claude-sonnet-4-6', 'ingest:', ' embeddings:', ' backend: openai', ' model: text-embedding-3-small', ' dimensions: 1536', '', ].join('\n'), 'utf-8', ); const io = makeIo(); await expect( runKloSetupContextStep( { projectDir: tempDir, inputMode: 'disabled' }, io.io, { runIdFactory: () => 'setup-context-local-empty' }, ), ).resolves.toEqual({ status: 'failed', projectDir: tempDir }); expect(io.stderr()).toContain('No primary or context sources are configured for a KLO context build.'); }); it('prints JSON setup context command status with watch and resume commands', async () => { await mkdir(join(tempDir, '.klo', 'setup'), { recursive: true }); await writeKloSetupContextState(tempDir, { runId: 'setup-context-local-abc123', status: 'detached', startedAt: '2026-05-09T10:00:00.000Z', updatedAt: '2026-05-09T10:01:00.000Z', primarySourceConnectionIds: ['warehouse'], contextSourceConnectionIds: ['docs'], reportIds: [], artifactPaths: [], retryableFailedTargets: [], commands: contextBuildCommands(tempDir, 'setup-context-local-abc123'), }); const io = makeIo(); await expect( runKloSetupContextCommand( { command: 'status', projectDir: tempDir, runId: 'setup-context-local-abc123', json: true }, io.io, ), ).resolves.toBe(0); expect(JSON.parse(io.stdout())).toMatchObject({ ready: false, status: 'detached', runId: 'setup-context-local-abc123', watchCommand: `klo setup context watch setup-context-local-abc123 --project-dir ${tempDir}`, statusCommand: `klo setup context status setup-context-local-abc123 --project-dir ${tempDir}`, }); }); it('runs direct build commands without asking for setup confirmation first', async () => { await writeReadyProject(tempDir); const io = makeIo(); const runContextBuildMock = vi.fn(async () => ({ exitCode: 0, detached: false })); await expect( runKloSetupContextCommand( { command: 'build', projectDir: tempDir, inputMode: 'auto' }, io.io, { prompts: { select: vi.fn(async () => { throw new Error('direct build should not prompt'); }), cancel: vi.fn(), }, runIdFactory: () => 'setup-context-local-direct', runContextBuild: runContextBuildMock, verifyContextReady: vi.fn(async () => ({ ready: true, agentContextReady: true, semanticSearchReady: true, details: [], })), }, ), ).resolves.toBe(0); expect(runContextBuildMock).toHaveBeenCalled(); }); });