import { mkdir, mkdtemp, readFile, rm, writeFile } from 'node:fs/promises'; import { tmpdir } from 'node:os'; import { join } from 'node:path'; import { readKtxSetupState } from '@ktx/context/project'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { contextBuildCommands, readKtxSetupContextState, runKtxSetupContextStep, writeKtxSetupContextState, } 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, 'ktx.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(), 'ktx-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(readKtxSetupContextState(tempDir)).resolves.toMatchObject({ status: 'not_started' }); await writeKtxSetupContextState(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 readKtxSetupContextState(tempDir); expect(state).toMatchObject({ runId: 'setup-context-local-abc123', status: 'running', primarySourceConnectionIds: ['warehouse'], contextSourceConnectionIds: ['docs'], commands: { watch: `ktx setup --project-dir ${tempDir}`, status: `ktx status --project-dir ${tempDir}`, resume: `ktx 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, reportIds: ['report-docs-1'], artifactPaths: ['raw-sources/warehouse/live-database/sync-1/scan-report.json'], })); const verifyContextReady = vi.fn(async () => ({ ready: true, agentContextReady: true, semanticSearchReady: true, details: ['warehouse: enriched scan complete', 'docs: memory update complete'], })); await expect( runKtxSetupContextStep( { 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, 'ktx.yaml'), 'utf-8')).not.toContain('completed_steps:'); expect((await readKtxSetupState(tempDir)).completed_steps).toContain('context'); await expect(readKtxSetupContextState(tempDir)).resolves.toMatchObject({ runId: 'setup-context-local-abc123', status: 'completed', completedAt: '2026-05-09T10:00:00.000Z', reportIds: ['report-docs-1'], artifactPaths: ['raw-sources/warehouse/live-database/sync-1/scan-report.json'], }); expect(io.stdout()).toContain('KTX context is ready for agents.'); }); it('records only failed sources as retryable when the context build fails', async () => { await writeReadyProject(tempDir); const io = makeIo(); const runContextBuildMock = vi.fn(async (_project, _args, _io, hooks) => { hooks.onSourceProgress?.([ { connectionId: 'warehouse', operation: 'scan', status: 'done', elapsedMs: 1000 }, { connectionId: 'docs', operation: 'source-ingest', status: 'failed', elapsedMs: 2000 }, ]); return { exitCode: 1, detached: false, reportIds: ['report-docs-failed'], artifactPaths: ['raw-sources/docs/notion/sync-1/ingest-report.json'], }; }); await expect( runKtxSetupContextStep( { projectDir: tempDir, inputMode: 'disabled' }, io.io, { runIdFactory: () => 'setup-context-local-failed', now: () => new Date('2026-05-09T10:00:00.000Z'), runContextBuild: runContextBuildMock, }, ), ).resolves.toEqual({ status: 'failed', projectDir: tempDir }); await expect(readKtxSetupContextState(tempDir)).resolves.toMatchObject({ runId: 'setup-context-local-failed', status: 'failed', reportIds: ['report-docs-failed'], artifactPaths: ['raw-sources/docs/notion/sync-1/ingest-report.json'], retryableFailedTargets: ['docs'], sourceProgress: [ { connectionId: 'warehouse', operation: 'scan', status: 'done', elapsedMs: 1000 }, { connectionId: 'docs', operation: 'source-ingest', status: 'failed', elapsedMs: 2000 }, ], }); }); 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( runKtxSetupContextStep( { 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, 'ktx.yaml'), 'utf-8')).not.toContain('completed_steps:'); expect((await readKtxSetupState(tempDir)).completed_steps).toContain('context'); await expect(readKtxSetupContextState(tempDir)).resolves.toMatchObject({ runId: 'setup-context-local-existing', status: 'completed', completedAt: '2026-05-09T10:00:00.000Z', contextSourceConnectionIds: ['docs'], }); expect(io.stdout()).toContain('KTX 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( runKtxSetupContextStep( { 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( runKtxSetupContextStep( { 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, 'ktx.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( runKtxSetupContextStep( { 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 KTX context build.'); }); it('watches an already-running setup context build from the resume prompt', async () => { await writeReadyProject(tempDir); await writeKtxSetupContextState(tempDir, { runId: 'setup-context-local-resume-watch', status: 'detached', 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-resume-watch'), }); const io = makeIo(); const completeRun = async () => { await writeKtxSetupContextState(tempDir, { runId: 'setup-context-local-resume-watch', status: 'completed', startedAt: '2026-05-09T10:00:00.000Z', updatedAt: '2026-05-09T10:02:00.000Z', completedAt: '2026-05-09T10:02:00.000Z', primarySourceConnectionIds: ['warehouse'], contextSourceConnectionIds: ['docs'], reportIds: [], artifactPaths: [], retryableFailedTargets: [], commands: contextBuildCommands(tempDir, 'setup-context-local-resume-watch'), }); }; const select = vi.fn(async (options: { options: Array<{ value: string; label: string }> }) => { expect(options.options.map((option) => option.label)).toContain('Watch progress'); return 'watch'; }); await expect( runKtxSetupContextStep( { projectDir: tempDir, inputMode: 'auto' }, io.io, { prompts: { select, cancel: vi.fn() }, sleep: completeRun, watchIntervalMs: 1, }, ), ).resolves.toEqual({ status: 'ready', projectDir: tempDir, runId: 'setup-context-local-resume-watch' }); expect(io.stdout()).toContain('KTX context built: detached'); expect(io.stdout()).toContain('KTX context built: yes'); }); it('auto-watches a running build without prompting when autoWatch is true', async () => { await writeReadyProject(tempDir); await writeKtxSetupContextState(tempDir, { runId: 'setup-context-local-auto-watch', status: 'detached', startedAt: '2026-05-09T10:00:00.000Z', updatedAt: '2026-05-09T10:00:00.000Z', primarySourceConnectionIds: ['warehouse'], contextSourceConnectionIds: [], reportIds: [], artifactPaths: [], retryableFailedTargets: [], commands: contextBuildCommands(tempDir, 'setup-context-local-auto-watch'), }); const io = makeIo(); const completeRun = async () => { await writeKtxSetupContextState(tempDir, { runId: 'setup-context-local-auto-watch', status: 'completed', startedAt: '2026-05-09T10:00:00.000Z', updatedAt: '2026-05-09T10:02:00.000Z', completedAt: '2026-05-09T10:02:00.000Z', primarySourceConnectionIds: ['warehouse'], contextSourceConnectionIds: [], reportIds: [], artifactPaths: [], retryableFailedTargets: [], commands: contextBuildCommands(tempDir, 'setup-context-local-auto-watch'), }); }; const select = vi.fn(async () => { throw new Error('should not prompt when autoWatch is true'); }); await expect( runKtxSetupContextStep( { projectDir: tempDir, inputMode: 'auto', autoWatch: true }, io.io, { prompts: { select, cancel: vi.fn() }, sleep: completeRun, watchIntervalMs: 1, }, ), ).resolves.toEqual({ status: 'ready', projectDir: tempDir, runId: 'setup-context-local-auto-watch' }); expect(select).not.toHaveBeenCalled(); expect(io.stdout()).toContain('KTX context built: yes'); }); it('renders the progress view when watching a build with sourceProgress', async () => { await writeReadyProject(tempDir); await writeKtxSetupContextState(tempDir, { runId: 'setup-context-local-progress', status: 'detached', 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-progress'), sourceProgress: [ { connectionId: 'warehouse', operation: 'scan' as const, status: 'done' as const, elapsedMs: 30000 }, { connectionId: 'docs', operation: 'source-ingest' as const, status: 'running' as const, startedAtMs: Date.now() - 5000 }, ], }); const io = makeIo(); const completeRun = async () => { await writeKtxSetupContextState(tempDir, { runId: 'setup-context-local-progress', status: 'completed', startedAt: '2026-05-09T10:00:00.000Z', updatedAt: '2026-05-09T10:02:00.000Z', completedAt: '2026-05-09T10:02:00.000Z', primarySourceConnectionIds: ['warehouse'], contextSourceConnectionIds: ['docs'], reportIds: [], artifactPaths: [], retryableFailedTargets: [], commands: contextBuildCommands(tempDir, 'setup-context-local-progress'), sourceProgress: [ { connectionId: 'warehouse', operation: 'scan' as const, status: 'done' as const, elapsedMs: 30000 }, { connectionId: 'docs', operation: 'source-ingest' as const, status: 'done' as const, elapsedMs: 60000 }, ], }); }; const select = vi.fn(async () => 'watch'); await expect( runKtxSetupContextStep( { projectDir: tempDir, inputMode: 'auto' }, io.io, { prompts: { select, cancel: vi.fn() }, sleep: completeRun, watchIntervalMs: 1, }, ), ).resolves.toEqual({ status: 'ready', projectDir: tempDir, runId: 'setup-context-local-progress' }); const output = io.stdout(); expect(output).toContain('Building KTX context'); expect(output).toContain('Primary sources:'); expect(output).toContain('warehouse'); expect(output).toContain('Context sources:'); expect(output).toContain('docs'); expect(output).not.toContain('KTX context built: detached'); }); it('supports d to detach from the progress watch view', async () => { await writeReadyProject(tempDir); await writeKtxSetupContextState(tempDir, { runId: 'setup-context-local-detach', status: 'running', startedAt: '2026-05-09T10:00:00.000Z', updatedAt: '2026-05-09T10:00:00.000Z', primarySourceConnectionIds: ['warehouse'], contextSourceConnectionIds: [], reportIds: [], artifactPaths: [], retryableFailedTargets: [], commands: contextBuildCommands(tempDir, 'setup-context-local-detach'), sourceProgress: [ { connectionId: 'warehouse', operation: 'scan' as const, status: 'running' as const, startedAtMs: Date.now() }, ], }); const io = makeIo(); let triggerDetach: (() => void) | null = null; await expect( runKtxSetupContextStep( { projectDir: tempDir, inputMode: 'auto', autoWatch: true }, io.io, { sleep: async () => { triggerDetach?.(); }, watchIntervalMs: 1, setupKeystroke: (onDetach) => { triggerDetach = onDetach; return () => {}; }, }, ), ).resolves.toMatchObject({ status: 'detached' }); const output = io.stdout(); expect(output).toContain('Building KTX context'); expect(output).toContain('Context build continuing in the background.'); expect(output).toContain('Resume: ktx setup --project-dir'); }); });