From 6c4623f2ff49c1973b441bca3b5d0552b8cc7547 Mon Sep 17 00:00:00 2001 From: Andrey Avtomonov Date: Thu, 14 May 2026 14:35:58 +0200 Subject: [PATCH] feat(cli): enforce required database selection and improve tree-picker UX (#86) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(cli): enforce required database selection and improve tree-picker UX - Require at least one database driver via prompt `required: true` instead of looping on empty selection; remove the now-dead retry/back-on-empty branch. - Surface the recommended option with a "(recommended)" hint in the depth and query-history prompts. - Tree picker: add `◧` partial glyph for parents whose descendants are checked, and make `a` toggle select-all-visible / select-none. * fix(cli): drop unused export from tree-picker toggleSelectAllVisible Knip flagged the export as unused; the function is only consumed by the internal reducer via the 'toggle-select-all-visible' command, so demote it to a module-local helper to keep CI's dead-code check green. * test(cli): drop empty-selection warning assertion from setup test The empty-selection retry/warning loop in `chooseDrivers` was removed in favor of `multiselect`'s `required: true`, so the legacy warning string is unreachable. Update the test to assert the simpler back-from-selection return-to-embeddings flow. --- .../cli/src/setup-database-context-depth.ts | 8 ++- packages/cli/src/setup-databases.test.ts | 33 ++-------- packages/cli/src/setup-databases.ts | 32 ++++----- packages/cli/src/setup.test.ts | 8 +-- packages/cli/src/tree-picker-state.test.ts | 66 +++++++++++++++++++ packages/cli/src/tree-picker-state.ts | 24 +++++++ packages/cli/src/tree-picker-tui.test.tsx | 27 +++++++- packages/cli/src/tree-picker-tui.tsx | 8 ++- 8 files changed, 146 insertions(+), 60 deletions(-) diff --git a/packages/cli/src/setup-database-context-depth.ts b/packages/cli/src/setup-database-context-depth.ts index 27683b61..39e75f71 100644 --- a/packages/cli/src/setup-database-context-depth.ts +++ b/packages/cli/src/setup-database-context-depth.ts @@ -46,12 +46,16 @@ async function chooseSetupDatabaseContextDepth(input: { const options = recommended === 'deep' ? [ - { value: 'deep', label: 'Deep: AI descriptions, embeddings, relationships, slower' }, + { + value: 'deep', + label: 'Deep: AI descriptions, embeddings, relationships, slower', + hint: 'recommended', + }, { value: 'fast', label: 'Fast: schema only, no AI, quickest' }, { value: 'back', label: 'Back' }, ] : [ - { value: 'fast', label: 'Fast: schema only, no AI, quickest' }, + { value: 'fast', label: 'Fast: schema only, no AI, quickest', hint: 'recommended' }, { value: 'deep', label: 'Deep: AI descriptions, embeddings, relationships, slower' }, { value: 'back', label: 'Back' }, ]; diff --git a/packages/cli/src/setup-databases.test.ts b/packages/cli/src/setup-databases.test.ts index 276d95f1..6672633d 100644 --- a/packages/cli/src/setup-databases.test.ts +++ b/packages/cli/src/setup-databases.test.ts @@ -149,31 +149,10 @@ describe('setup databases step', () => { { value: 'bigquery', label: 'BigQuery' }, { value: 'snowflake', label: 'Snowflake' }, ], - required: false, + required: true, }); }); - it('requires choosing a database after an empty interactive selection', async () => { - const io = makeIo(); - const prompts = makePromptAdapter({ - multiselectValues: [[], ['back']], - selectValues: ['choose'], - }); - - const result = await runKtxSetupDatabasesStep( - { projectDir: tempDir, inputMode: 'auto', skipDatabases: false, databaseSchemas: [] }, - io.io, - { prompts }, - ); - - expect(result.status).toBe('back'); - expect(prompts.select).not.toHaveBeenCalled(); - expect(io.stdout()).toContain( - 'KTX cannot work without at least one database. Select a database or press Escape to go back.', - ); - expect(prompts.multiselect).toHaveBeenCalledTimes(2); - }); - it('lets Back from connection method selection return to database selection when adding a new driver', async () => { const prompts = makePromptAdapter({ multiselectValues: [['postgres'], ['back']], @@ -744,10 +723,10 @@ describe('setup databases step', () => { expect(config.setup?.database_connection_ids).toEqual(['postgres-warehouse', 'mysql-warehouse']); }); - it('returns to configured primary menu when submitting empty driver selection after adding a source', async () => { + it('returns to configured primary menu when pressing back on driver selection after adding a source', async () => { const io = makeIo(); const prompts = makePromptAdapter({ - multiselectValues: [['postgres'], []], + multiselectValues: [['postgres'], ['back']], selectValues: ['url', 'add', 'continue'], textValues: ['', 'env:DATABASE_URL'], }); @@ -787,7 +766,7 @@ describe('setup databases step', () => { }); }); - it('returns to configured primary menu when submitting empty driver selection with pre-existing source', async () => { + it('returns to configured primary menu when pressing back on driver selection with pre-existing source', async () => { await writeFile( join(tempDir, 'ktx.yaml'), [ @@ -806,7 +785,7 @@ describe('setup databases step', () => { await writeKtxSetupState(tempDir, { completed_steps: ['databases'] }); const io = makeIo(); const prompts = makePromptAdapter({ - multiselectValues: [[]], + multiselectValues: [['back']], selectValues: ['add', 'continue'], }); @@ -2007,7 +1986,7 @@ describe('setup databases step', () => { expect(prompts.select).toHaveBeenCalledWith({ message: 'Enable query-history ingest for this PostgreSQL connection?', options: [ - { value: 'yes', label: 'Enable query history' }, + { value: 'yes', label: 'Enable query history (recommended)' }, { value: 'no', label: 'Do not enable query history' }, { value: 'back', label: 'Back' }, ], diff --git a/packages/cli/src/setup-databases.ts b/packages/cli/src/setup-databases.ts index add31a4d..1d41b4da 100644 --- a/packages/cli/src/setup-databases.ts +++ b/packages/cli/src/setup-databases.ts @@ -932,7 +932,7 @@ async function maybeApplyHistoricSqlConfig(input: { const choice = await input.prompts.select({ message: `Enable query-history ingest for this ${driverLabel(input.driver)} connection?`, options: [ - { value: 'yes', label: 'Enable query history' }, + { value: 'yes', label: 'Enable query history (recommended)' }, { value: 'no', label: 'Do not enable query history' }, { value: 'back', label: 'Back' }, ], @@ -1756,27 +1756,17 @@ async function chooseDrivers( ); return 'missing-input'; } - while (true) { - const initialValues = unique(options?.initialDrivers ?? []); - const choices = await prompts.multiselect({ - message: withMultiselectNavigation('Which databases should KTX connect to?'), - options: [...DRIVER_OPTIONS], - ...(initialValues.length > 0 ? { initialValues } : {}), - required: options?.hasPrimarySources === true, - }); - if (choices.includes('back')) { - return 'back'; - } - if (choices.length > 0) { - return choices as KtxSetupDatabaseDriver[]; - } - - if (options?.hasPrimarySources) { - return 'back'; - } - - io.stdout.write('│ KTX cannot work without at least one database. Select a database or press Escape to go back.\n'); + const initialValues = unique(options?.initialDrivers ?? []); + const choices = await prompts.multiselect({ + message: withMultiselectNavigation('Which databases should KTX connect to?'), + options: [...DRIVER_OPTIONS], + ...(initialValues.length > 0 ? { initialValues } : {}), + required: true, + }); + if (choices.includes('back')) { + return 'back'; } + return choices as KtxSetupDatabaseDriver[]; } async function chooseConnectionIdForDriver(input: { diff --git a/packages/cli/src/setup.test.ts b/packages/cli/src/setup.test.ts index b0c28d7c..7e479c0e 100644 --- a/packages/cli/src/setup.test.ts +++ b/packages/cli/src/setup.test.ts @@ -1094,7 +1094,7 @@ describe('setup status', () => { expect(embeddings).toHaveBeenCalledTimes(1); }); - it('lets Back from database selection return to embedding setup after an empty selection warning', async () => { + it('lets Back from database selection return to embedding setup', async () => { const testIo = makeIo(); const modelResults = [ { status: 'ready' as const, projectDir: tempDir }, @@ -1106,9 +1106,8 @@ describe('setup status', () => { { status: 'back' as const, projectDir: tempDir }, ]; const embeddings = vi.fn(async () => embeddingResults.shift() ?? { status: 'back' as const, projectDir: tempDir }); - const databaseMultiselectValues = [[], ['back']]; const databasePrompts = { - multiselect: vi.fn(async () => databaseMultiselectValues.shift() ?? ['back']), + multiselect: vi.fn(async () => ['back']), select: vi.fn(async () => 'back'), text: vi.fn(), password: vi.fn(), @@ -1142,9 +1141,6 @@ describe('setup status', () => { ).resolves.toBe(0); expect(databasePrompts.select).not.toHaveBeenCalled(); - expect(testIo.stdout()).toContain( - 'KTX cannot work without at least one database. Select a database or press Escape to go back.', - ); expect(embeddings).toHaveBeenCalledTimes(2); expect(embeddings).toHaveBeenNthCalledWith(2, expect.objectContaining({ forcePrompt: true }), testIo.io); expect(testIo.stderr()).not.toContain('No databases selected.'); diff --git a/packages/cli/src/tree-picker-state.test.ts b/packages/cli/src/tree-picker-state.test.ts index 52e63f3d..20c2001e 100644 --- a/packages/cli/src/tree-picker-state.test.ts +++ b/packages/cli/src/tree-picker-state.test.ts @@ -6,6 +6,7 @@ import { clearExpiredTransientHint, filterTree, flattenSelection, + hasPartialChildren, moveCursor, reducer, selectAllVisible, @@ -134,6 +135,54 @@ describe('selection invariants', () => { }); }); +describe('hasPartialChildren', () => { + const ROOT = 'r0000000-0000-0000-0000-000000000000'; + const CHILD_A = 'a0000000-0000-0000-0000-000000000000'; + const CHILD_B = 'b0000000-0000-0000-0000-000000000000'; + const GRANDCHILD = 'g0000000-0000-0000-0000-000000000000'; + const LEAF = 'l0000000-0000-0000-0000-000000000000'; + + function tree() { + return buildPickerTree([ + { id: ROOT, title: 'Root', parentId: null }, + { id: CHILD_A, title: 'Child A', parentId: ROOT }, + { id: CHILD_B, title: 'Child B', parentId: ROOT }, + { id: GRANDCHILD, title: 'Grandchild', parentId: CHILD_A }, + { id: LEAF, title: 'Leaf', parentId: null }, + ]); + } + + function byId() { + return new Map(tree().map((node) => [node.id, node])); + } + + it('returns false for a leaf node with no children', () => { + expect(hasPartialChildren(LEAF, new Set(), byId())).toBe(false); + expect(hasPartialChildren(LEAF, new Set([LEAF]), byId())).toBe(false); + }); + + it('returns false when the node itself is checked', () => { + expect(hasPartialChildren(ROOT, new Set([ROOT]), byId())).toBe(false); + }); + + it('returns false when an ancestor is checked (node is locked)', () => { + expect(hasPartialChildren(CHILD_A, new Set([ROOT]), byId())).toBe(false); + }); + + it('returns true when a direct child is checked', () => { + expect(hasPartialChildren(ROOT, new Set([CHILD_A]), byId())).toBe(true); + }); + + it('returns true when a deep descendant is checked', () => { + expect(hasPartialChildren(ROOT, new Set([GRANDCHILD]), byId())).toBe(true); + }); + + it('returns false when no descendant is checked', () => { + expect(hasPartialChildren(ROOT, new Set([LEAF]), byId())).toBe(false); + expect(hasPartialChildren(ROOT, new Set(), byId())).toBe(false); + }); +}); + describe('search and cursor movement', () => { it('filters by title and path while deriving auto-expanded ancestors', () => { const state = buildInitialState({ @@ -188,6 +237,23 @@ describe('bulk actions and reducer effects', () => { expect([...selectNone(selected).checked]).toEqual([]); }); + it('toggle-select-all-visible selects all visible roots, then clears on repeat', () => { + const state = buildInitialState({ + tree: buildPickerTree(pages()), + existingSelectedIds: [], + }); + + const firstPress = reducer(state, 'toggle-select-all-visible').next; + expect(firstPress.checked.size).toBeGreaterThan(0); + const allSelectedIds = [...firstPress.checked].sort(); + + const secondPress = reducer(firstPress, 'toggle-select-all-visible').next; + expect([...secondPress.checked]).toEqual([]); + + const thirdPress = reducer(secondPress, 'toggle-select-all-visible').next; + expect([...thirdPress.checked].sort()).toEqual(allSelectedIds); + }); + it('saves immediately when confirm is not required and prompts confirmation when requireConfirmOnSave is true', () => { const noConfirm = toggleChecked( buildInitialState({ diff --git a/packages/cli/src/tree-picker-state.ts b/packages/cli/src/tree-picker-state.ts index 9d9b3c68..3f9d13bf 100644 --- a/packages/cli/src/tree-picker-state.ts +++ b/packages/cli/src/tree-picker-state.ts @@ -44,6 +44,7 @@ export type PickerCommand = | 'collapse-all' | 'toggle-check' | 'select-all-visible' + | 'toggle-select-all-visible' | 'select-none' | 'clear-transient-hint' | 'search-start' @@ -228,6 +229,17 @@ export function isAncestorChecked(nodeId: string, checked: Set, byId: Ma return ancestorsOf(nodeId, byId).some((ancestorId) => checked.has(ancestorId)); } +export function hasPartialChildren( + nodeId: string, + checked: Set, + byId: Map, +): boolean { + if (checked.has(nodeId) || isAncestorChecked(nodeId, checked, byId)) { + return false; + } + return descendantsOf(nodeId, byId).some((descendantId) => checked.has(descendantId)); +} + function checkedAncestor(nodeId: string, state: PickerState): TreePickerNode | null { for (const ancestorId of ancestorsOf(nodeId, state.byId)) { if (state.checked.has(ancestorId)) { @@ -350,6 +362,16 @@ export function selectNone(state: PickerState): PickerState { return cloneState(state, { checked: new Set(), transientHint: null }); } +function toggleSelectAllVisible(state: PickerState): PickerState { + const next = selectAllVisible(state); + const unchanged = + next.checked.size === state.checked.size && [...next.checked].every((id) => state.checked.has(id)); + if (unchanged && state.checked.size > 0) { + return selectNone(state); + } + return next; +} + function setExpanded(state: PickerState, nodeId: string, value: boolean | 'toggle'): PickerState { const expanded = new Set(state.expanded); const nextValue = value === 'toggle' ? !expanded.has(nodeId) : value; @@ -487,6 +509,8 @@ export function reducer(state: PickerState, cmd: PickerCommand, now = Date.now() return { next: toggleChecked(state, state.cursorId, now), effect: null }; case 'select-all-visible': return { next: selectAllVisible(state), effect: null }; + case 'toggle-select-all-visible': + return { next: toggleSelectAllVisible(state), effect: null }; case 'select-none': return { next: selectNone(state), effect: null }; case 'clear-transient-hint': diff --git a/packages/cli/src/tree-picker-tui.test.tsx b/packages/cli/src/tree-picker-tui.test.tsx index 8c4f8d1e..3a8dcc7b 100644 --- a/packages/cli/src/tree-picker-tui.test.tsx +++ b/packages/cli/src/tree-picker-tui.test.tsx @@ -90,7 +90,7 @@ describe('treePickerCommandForInkInput', () => { expect(treePickerCommandForInkInput('', { leftArrow: true }, state().search, null)).toBe('cursor-left'); expect(treePickerCommandForInkInput(' ', {}, state().search, null)).toBe('toggle-check'); expect(treePickerCommandForInkInput('/', {}, state().search, null)).toBe('search-start'); - expect(treePickerCommandForInkInput('a', {}, state().search, null)).toBe('select-all-visible'); + expect(treePickerCommandForInkInput('a', {}, state().search, null)).toBe('toggle-select-all-visible'); expect(treePickerCommandForInkInput('n', {}, state().search, null)).toBe('select-none'); expect(treePickerCommandForInkInput('', { return: true }, state().search, null)).toBe('save-request'); expect(treePickerCommandForInkInput('', { escape: true }, state().search, null)).toBe('quit'); @@ -198,6 +198,31 @@ describe('TreePickerApp', () => { expect(frame).toContain(' ◼ Architecture'); }); + it('renders the partial glyph on a parent whose descendant is checked', () => { + const partialPages: TreePickerNodeInput[] = [ + { id: IDS.engineering, title: 'Engineering Docs', archived: false, parentId: null }, + { id: IDS.architecture, title: 'Architecture', archived: false, parentId: IDS.engineering }, + ]; + const initialState = buildInitialState({ + tree: buildPickerTree(partialPages), + existingSelectedIds: [IDS.architecture], + }); + const { lastFrame } = renderInkTest( + , + ); + + const frame = lastFrame() ?? ''; + expect(frame).toContain('◧ Engineering Docs ▾'); + expect(frame).toContain(' ◼ Architecture'); + expect(frame).not.toContain('◻ Engineering Docs'); + }); + it('supports keyboard selection, confirm-on-save, and save callback', async () => { const onExit = vi.fn(); const { stdin, lastFrame } = renderInkTest( diff --git a/packages/cli/src/tree-picker-tui.tsx b/packages/cli/src/tree-picker-tui.tsx index 9cdbef8d..a2c4a9ab 100644 --- a/packages/cli/src/tree-picker-tui.tsx +++ b/packages/cli/src/tree-picker-tui.tsx @@ -4,6 +4,7 @@ import { type ReactNode, useEffect, useMemo, useRef, useState } from 'react'; import { filterTree, flattenSelection, + hasPartialChildren, isAncestorChecked, reducer, visibleNodeIds, @@ -165,7 +166,7 @@ export function treePickerCommandForInkInput( if (key.return) return 'save-request'; if (input === ' ') return 'toggle-check'; if (input === '/') return 'search-start'; - if (input === 'a') return 'select-all-visible'; + if (input === 'a') return 'toggle-select-all-visible'; if (input === 'n') return 'select-none'; if (key.escape) return 'quit'; return null; @@ -178,8 +179,9 @@ function PickerRow(props: { state: PickerState; nodeId: string; width: number; t const locked = isAncestorChecked(node.id, props.state.checked, props.state.byId); const checked = props.state.checked.has(node.id); const isSelected = checked || locked; - const glyph = isSelected ? '◼' : '◻'; - const glyphColor = checked || locked ? props.theme.selected : props.theme.muted; + const partial = !isSelected && hasPartialChildren(node.id, props.state.checked, props.state.byId); + const glyph = isSelected ? '◼' : partial ? '◧' : '◻'; + const glyphColor = isSelected || partial ? props.theme.selected : props.theme.muted; const childAffordance = node.childIds.length > 0 ? (props.state.expanded.has(node.id) ? ' ▾' : ` ▸ (${node.childIds.length})`) : ''; const indent = ' '.repeat(node.depth * 2);