|
| 1 | +import { describe, expect, it, vi } from 'vitest'; |
| 2 | + |
| 3 | +import { handleModelCommand } from '#/tui/commands/index'; |
| 4 | +import type { SlashCommandHost } from '#/tui/commands/dispatch'; |
| 5 | + |
| 6 | +const ENTER = '\r'; |
| 7 | + |
| 8 | +interface TestPicker { |
| 9 | + handleInput(data: string): void; |
| 10 | +} |
| 11 | + |
| 12 | +function model(name: string) { |
| 13 | + return { |
| 14 | + provider: 'test', |
| 15 | + model: name, |
| 16 | + maxContextSize: 200_000, |
| 17 | + displayName: name, |
| 18 | + capabilities: [], |
| 19 | + }; |
| 20 | +} |
| 21 | + |
| 22 | +function makeHost(options: { |
| 23 | + currentModel?: string; |
| 24 | + availableModels?: Record<string, ReturnType<typeof model>>; |
| 25 | + modelRoles?: Record<string, string>; |
| 26 | + setConfig?: ReturnType<typeof vi.fn>; |
| 27 | +} = {}) { |
| 28 | + const session = { |
| 29 | + setModel: vi.fn(async () => {}), |
| 30 | + setThinking: vi.fn(async () => {}), |
| 31 | + }; |
| 32 | + const getConfig = vi.fn(async () => ({ |
| 33 | + providers: {}, |
| 34 | + modelRoles: options.modelRoles, |
| 35 | + })); |
| 36 | + const setConfig = options.setConfig ?? vi.fn(async () => {}); |
| 37 | + const host = { |
| 38 | + state: { |
| 39 | + appState: { |
| 40 | + model: options.currentModel ?? 'worker', |
| 41 | + thinkingLevel: 'off', |
| 42 | + streamingPhase: 'idle', |
| 43 | + availableModels: options.availableModels ?? { worker: model('worker') }, |
| 44 | + }, |
| 45 | + editorContainer: { children: [] }, |
| 46 | + }, |
| 47 | + session, |
| 48 | + harness: { getConfig, setConfig }, |
| 49 | + authFlow: { |
| 50 | + refreshProviderModels: vi.fn(async () => ({ failed: [] })), |
| 51 | + }, |
| 52 | + setAppState: vi.fn((patch: Record<string, unknown>) => Object.assign(host.state.appState, patch)), |
| 53 | + showError: vi.fn(), |
| 54 | + showStatus: vi.fn(), |
| 55 | + showNotice: vi.fn(), |
| 56 | + mountEditorReplacement: vi.fn(), |
| 57 | + restoreEditor: vi.fn(), |
| 58 | + track: vi.fn(), |
| 59 | + } as unknown as SlashCommandHost; |
| 60 | + return { host, session, setConfig }; |
| 61 | +} |
| 62 | + |
| 63 | +function mountedPicker(host: SlashCommandHost, index = 0): TestPicker { |
| 64 | + const mount = host.mountEditorReplacement as ReturnType<typeof vi.fn>; |
| 65 | + return mount.mock.calls[index]?.[0] as TestPicker; |
| 66 | +} |
| 67 | + |
| 68 | +describe('/model roles', () => { |
| 69 | + it('lists every built-in role as not set when no assignments exist', async () => { |
| 70 | + const { host } = makeHost(); |
| 71 | + |
| 72 | + await handleModelCommand(host, 'roles'); |
| 73 | + |
| 74 | + expect(host.showNotice).toHaveBeenCalledWith( |
| 75 | + 'Model roles', |
| 76 | + 'small: (not set)\nimplementer: (not set)\nadvisor: (not set)', |
| 77 | + ); |
| 78 | + }); |
| 79 | + |
| 80 | + it('locks a selected alias to a role without switching the session model', async () => { |
| 81 | + const { host, session, setConfig } = makeHost(); |
| 82 | + |
| 83 | + await handleModelCommand(host, 'small'); |
| 84 | + expect(host.authFlow.refreshProviderModels).toHaveBeenCalledOnce(); |
| 85 | + mountedPicker(host).handleInput(ENTER); |
| 86 | + |
| 87 | + await vi.waitFor(() => { |
| 88 | + expect(setConfig).toHaveBeenCalledWith({ modelRoles: { small: 'worker' } }); |
| 89 | + }); |
| 90 | + expect(session.setModel).not.toHaveBeenCalled(); |
| 91 | + }); |
| 92 | + |
| 93 | + it('keeps role assignment active after the picker refreshes', async () => { |
| 94 | + const { host, session, setConfig } = makeHost({ |
| 95 | + currentModel: 'parent', |
| 96 | + availableModels: { |
| 97 | + parent: model('parent'), |
| 98 | + worker: model('worker'), |
| 99 | + }, |
| 100 | + modelRoles: { small: 'worker' }, |
| 101 | + }); |
| 102 | + vi.mocked(host.mountEditorReplacement).mockImplementation((picker) => { |
| 103 | + host.state.editorContainer.children[0] = picker; |
| 104 | + }); |
| 105 | + vi.mocked(host.authFlow.refreshProviderModels).mockImplementation(async () => { |
| 106 | + host.state.appState.availableModels['reviewer'] = model('reviewer'); |
| 107 | + return { changed: [], unchanged: [], failed: [] }; |
| 108 | + }); |
| 109 | + |
| 110 | + await handleModelCommand(host, 'small'); |
| 111 | + await vi.waitFor(() => { |
| 112 | + expect(host.mountEditorReplacement).toHaveBeenCalledTimes(2); |
| 113 | + }); |
| 114 | + mountedPicker(host, 1).handleInput(ENTER); |
| 115 | + |
| 116 | + await vi.waitFor(() => { |
| 117 | + expect(setConfig).toHaveBeenCalledWith({ modelRoles: { small: 'worker' } }); |
| 118 | + }); |
| 119 | + expect(session.setModel).not.toHaveBeenCalled(); |
| 120 | + }); |
| 121 | + |
| 122 | + it('reports a role persistence failure without showing success', async () => { |
| 123 | + const setConfig = vi.fn(async () => { |
| 124 | + throw new Error('disk full'); |
| 125 | + }); |
| 126 | + const { host } = makeHost({ setConfig }); |
| 127 | + |
| 128 | + await handleModelCommand(host, 'small'); |
| 129 | + mountedPicker(host).handleInput(ENTER); |
| 130 | + |
| 131 | + await vi.waitFor(() => { |
| 132 | + expect(host.showError).toHaveBeenCalledWith(expect.stringContaining('disk full')); |
| 133 | + }); |
| 134 | + expect(host.showStatus).not.toHaveBeenCalled(); |
| 135 | + }); |
| 136 | + |
| 137 | + it('clears a role with an empty-string tombstone', async () => { |
| 138 | + const { host, setConfig } = makeHost({ modelRoles: { small: 'worker' } }); |
| 139 | + |
| 140 | + await handleModelCommand(host, 'small clear'); |
| 141 | + |
| 142 | + expect(setConfig).toHaveBeenCalledWith({ modelRoles: { small: '' } }); |
| 143 | + }); |
| 144 | + |
| 145 | + it('keeps an existing model alias on the default switch path', async () => { |
| 146 | + const { host, session } = makeHost({ |
| 147 | + currentModel: 'parent', |
| 148 | + availableModels: { |
| 149 | + parent: model('parent'), |
| 150 | + worker: model('worker'), |
| 151 | + }, |
| 152 | + }); |
| 153 | + |
| 154 | + await handleModelCommand(host, 'worker'); |
| 155 | + mountedPicker(host).handleInput(ENTER); |
| 156 | + |
| 157 | + await vi.waitFor(() => { |
| 158 | + expect(session.setModel).toHaveBeenCalledWith('worker'); |
| 159 | + }); |
| 160 | + }); |
| 161 | +}); |
0 commit comments