Skip to content

Commit 6726308

Browse files
elkaixelkaix
committed
fix(tui): speed up model picker and dedupe aliases
1 parent 8717f4e commit 6726308

6 files changed

Lines changed: 399 additions & 84 deletions

File tree

‎apps/pythinker-code/src/tui/commands/config.ts‎

Lines changed: 106 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,12 @@ import {
2121
ExperimentsSelectorComponent,
2222
type ExperimentalFeatureDraftChange,
2323
} from '../components/dialogs/experiments-selector';
24-
import { modelDisplayName } from '../components/dialogs/model-selector';
24+
import {
25+
modelDisplayName,
26+
modelIdentity,
27+
normalizeModelChoices,
28+
resolveNormalizedModelAlias,
29+
} from '../components/dialogs/model-selector';
2530
import { TabbedModelSelectorComponent } from '../components/dialogs/tabbed-model-selector';
2631
import { PermissionSelectorComponent } from '../components/dialogs/permission-selector';
2732
import { SettingsSelectorComponent, type SettingsSelection } from '../components/dialogs/settings-selector';
@@ -47,7 +52,6 @@ import type { SlashCommandHost } from './dispatch';
4752
// Plan / Config commands
4853
// ---------------------------------------------------------------------------
4954

50-
const MODEL_PICKER_REFRESH_TIMEOUT_MS = 2_000;
5155

5256
export async function handlePlanCommand(host: SlashCommandHost, args: string): Promise<void> {
5357
const session = host.session;
@@ -475,17 +479,25 @@ function resolveWorkspaceConfigPath(input: string, workDir: string): string {
475479
}
476480

477481
export async function handleModelCommand(host: SlashCommandHost, args: string): Promise<void> {
478-
const alias = args.trim();
479-
await refreshModelsForPicker(host);
480-
if (alias.length === 0) {
481-
showModelPicker(host);
482+
const requestedAlias = args.trim();
483+
const normalized = normalizeModelChoices(host.state.appState.availableModels);
484+
const selectedValue =
485+
requestedAlias.length === 0
486+
? undefined
487+
: resolveNormalizedModelAlias(
488+
normalized,
489+
requestedAlias,
490+
host.state.appState.availableModels[requestedAlias],
491+
);
492+
if (requestedAlias.length > 0 && selectedValue === undefined) {
493+
host.showError(`Unknown model alias: ${requestedAlias}`);
482494
return;
483495
}
484-
if (host.state.appState.availableModels[alias] === undefined) {
485-
host.showError(`Unknown model alias: ${alias}`);
486-
return;
496+
497+
const picker = showModelPicker(host, selectedValue);
498+
if (picker !== undefined) {
499+
void refreshModelsForOpenPicker(host, picker, selectedValue);
487500
}
488-
showModelPicker(host, alias);
489501
}
490502

491503
// ---------------------------------------------------------------------------
@@ -508,35 +520,61 @@ function showEditorPicker(host: SlashCommandHost): void {
508520
);
509521
}
510522

511-
async function refreshModelsForPicker(host: SlashCommandHost): Promise<void> {
523+
async function refreshModelsForOpenPicker(
524+
host: SlashCommandHost,
525+
picker: TabbedModelSelectorComponent,
526+
selectedValue: string | undefined,
527+
): Promise<void> {
528+
const availableModels = host.state.appState.availableModels;
529+
const normalized = normalizeModelChoices(availableModels);
530+
const currentModel = availableModels[host.state.appState.model];
531+
512532
try {
513-
const result = await withTimeout(
514-
host.authFlow.refreshOAuthProviderModels(),
515-
MODEL_PICKER_REFRESH_TIMEOUT_MS,
516-
);
517-
if (result === undefined) return;
533+
const result = await host.authFlow.refreshOAuthProviderModels();
518534
for (const f of result.failed) {
519535
host.showStatus(`Skipped refreshing ${f.provider}: ${f.reason}`, 'warning');
520536
}
521537
} catch (error) {
522538
host.showStatus(`Skipped refreshing models: ${formatErrorMessage(error)}`, 'warning');
539+
return;
523540
}
524-
}
525541

526-
async function withTimeout<T>(promise: Promise<T>, timeoutMs: number): Promise<T | undefined> {
527-
let timeout: ReturnType<typeof setTimeout> | undefined;
528-
try {
529-
return await Promise.race([
530-
promise,
531-
new Promise<undefined>((resolve) => {
532-
timeout = setTimeout(() => {
533-
resolve(undefined);
534-
}, timeoutMs);
535-
}),
536-
]);
537-
} finally {
538-
if (timeout !== undefined) clearTimeout(timeout);
542+
if (host.state.editorContainer.children[0] !== picker) return;
543+
544+
const liveSelectedAlias = picker.selectedAlias() ?? selectedValue;
545+
const selectedModel =
546+
liveSelectedAlias === undefined
547+
? undefined
548+
: normalized.models[liveSelectedAlias] ?? availableModels[liveSelectedAlias];
549+
const activeTabId = picker.activeTabId();
550+
551+
const refreshed = normalizeModelChoices(host.state.appState.availableModels);
552+
if (currentModel !== undefined) {
553+
const refreshedCurrent = resolveNormalizedModelAlias(
554+
refreshed,
555+
host.state.appState.model,
556+
currentModel,
557+
);
558+
if (refreshedCurrent === undefined) return;
559+
if (modelIdentity(refreshed.models[refreshedCurrent]) !== modelIdentity(currentModel)) {
560+
return;
561+
}
562+
}
563+
564+
let refreshedSelected = liveSelectedAlias;
565+
if (selectedModel !== undefined) {
566+
refreshedSelected = resolveNormalizedModelAlias(
567+
refreshed,
568+
liveSelectedAlias ?? '',
569+
selectedModel,
570+
);
571+
if (refreshedSelected === undefined) return;
572+
if (modelIdentity(refreshed.models[refreshedSelected]) !== modelIdentity(selectedModel)) {
573+
return;
574+
}
539575
}
576+
577+
showModelPicker(host, refreshedSelected, activeTabId);
540578
}
541579

542580
async function applyEditorChoice(host: SlashCommandHost, value: string): Promise<void> {
@@ -572,30 +610,49 @@ async function applyEditorChoice(host: SlashCommandHost, value: string): Promise
572610
);
573611
}
574612

575-
export function showModelPicker(host: SlashCommandHost, selectedValue: string = host.state.appState.model): void {
576-
const entries = Object.entries(host.state.appState.availableModels);
613+
export function showModelPicker(
614+
host: SlashCommandHost,
615+
selectedValue?: string,
616+
initialTabId?: string,
617+
): TabbedModelSelectorComponent | undefined {
618+
const normalized = normalizeModelChoices(host.state.appState.availableModels);
619+
const entries = Object.entries(normalized.models);
577620
if (entries.length === 0) {
578621
host.showNotice(
579622
'No models configured',
580623
'Run /login to sign in to Pythinker, or /provider to add another provider from a model catalog.',
581624
);
582-
return;
583-
}
584-
host.mountEditorReplacement(
585-
new TabbedModelSelectorComponent({
586-
models: host.state.appState.availableModels,
587-
currentValue: host.state.appState.model,
588-
selectedValue,
589-
currentEffort: host.state.appState.thinkingLevel,
590-
onSelect: ({ alias, effort }) => {
591-
host.restoreEditor();
592-
void performModelSwitch(host, alias, effort);
593-
},
594-
onCancel: () => {
595-
host.restoreEditor();
596-
},
597-
}),
598-
);
625+
return undefined;
626+
}
627+
const currentValue =
628+
resolveNormalizedModelAlias(
629+
normalized,
630+
host.state.appState.model,
631+
host.state.appState.availableModels[host.state.appState.model],
632+
) ?? host.state.appState.model;
633+
const selectedCandidate = selectedValue ?? host.state.appState.model;
634+
const resolvedSelectedValue =
635+
resolveNormalizedModelAlias(
636+
normalized,
637+
selectedCandidate,
638+
host.state.appState.availableModels[selectedCandidate],
639+
) ?? currentValue;
640+
const picker = new TabbedModelSelectorComponent({
641+
models: normalized.models,
642+
currentValue,
643+
selectedValue: resolvedSelectedValue,
644+
currentEffort: host.state.appState.thinkingLevel,
645+
initialTabId,
646+
onSelect: ({ alias, effort }) => {
647+
host.restoreEditor();
648+
void performModelSwitch(host, alias, effort);
649+
},
650+
onCancel: () => {
651+
host.restoreEditor();
652+
},
653+
});
654+
host.mountEditorReplacement(picker);
655+
return picker;
599656
}
600657

601658
async function performModelSwitch(host: SlashCommandHost, alias: string, effort: string): Promise<void> {

‎apps/pythinker-code/src/tui/components/dialogs/model-selector.ts‎

Lines changed: 97 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,12 @@ export interface ModelSelection {
4747
readonly effort: string;
4848
}
4949

50+
export interface NormalizedModelChoices {
51+
readonly models: Record<string, ModelAlias>;
52+
readonly aliasMap: Record<string, string>;
53+
readonly identityAliases: Record<string, string>;
54+
}
55+
5056
export function modelDisplayName(alias: string, model: ModelAlias | undefined): string {
5157
return model?.displayName ?? model?.model ?? alias;
5258
}
@@ -57,10 +63,75 @@ export function providerDisplayName(provider: string): string {
5763
return provider;
5864
}
5965

66+
export function canonicalModelAlias(model: Pick<ModelAlias, 'provider' | 'model'>): string {
67+
return `${model.provider}/${model.model}`;
68+
}
69+
70+
export function modelIdentity(
71+
model: Pick<ModelAlias, 'provider' | 'model'> | undefined,
72+
): string | undefined {
73+
return model === undefined ? undefined : `${model.provider}\u0000${model.model}`;
74+
}
75+
76+
export function normalizeModelChoices(
77+
models: Record<string, ModelAlias>,
78+
): NormalizedModelChoices {
79+
const aliasMap: Record<string, string> = {};
80+
const identityAliases: Record<string, string> = {};
81+
const sourceAliasesByIdentity = new Map<string, string[]>();
82+
const representativeByIdentity = new Map<string, { alias: string; model: ModelAlias }>();
83+
const identityOrder: string[] = [];
84+
85+
for (const [alias, cfg] of Object.entries(models)) {
86+
const identity = modelIdentity(cfg);
87+
if (identity === undefined) continue;
88+
const sourceAliases = sourceAliasesByIdentity.get(identity);
89+
if (sourceAliases === undefined) {
90+
sourceAliasesByIdentity.set(identity, [alias]);
91+
representativeByIdentity.set(identity, { alias, model: cfg });
92+
identityOrder.push(identity);
93+
continue;
94+
}
95+
96+
sourceAliases.push(alias);
97+
const representative = representativeByIdentity.get(identity);
98+
if (representative !== undefined && alias === canonicalModelAlias(cfg)) {
99+
representative.alias = alias;
100+
representative.model = cfg;
101+
}
102+
}
103+
104+
const normalized: Record<string, ModelAlias> = {};
105+
for (const identity of identityOrder) {
106+
const representative = representativeByIdentity.get(identity);
107+
if (representative === undefined) continue;
108+
normalized[representative.alias] = representative.model;
109+
identityAliases[identity] = representative.alias;
110+
for (const sourceAlias of sourceAliasesByIdentity.get(identity) ?? []) {
111+
aliasMap[sourceAlias] = representative.alias;
112+
}
113+
aliasMap[representative.alias] = representative.alias;
114+
}
115+
116+
return { models: normalized, aliasMap, identityAliases };
117+
}
118+
119+
export function resolveNormalizedModelAlias(
120+
normalized: NormalizedModelChoices,
121+
alias: string,
122+
fallbackModel?: Pick<ModelAlias, 'provider' | 'model'>,
123+
): string | undefined {
124+
const mapped = normalized.aliasMap[alias];
125+
if (mapped !== undefined) return mapped;
126+
const identity = modelIdentity(fallbackModel);
127+
return identity === undefined ? undefined : normalized.identityAliases[identity];
128+
}
129+
60130
export function createModelChoiceOptions(
61131
models: Record<string, ModelAlias>,
62132
): readonly ChoiceOption[] {
63-
return Object.entries(models).map(([alias, cfg]) => ({
133+
const normalized = normalizeModelChoices(models);
134+
return Object.entries(normalized.models).map(([alias, cfg]) => ({
64135
value: alias,
65136
label: `${modelDisplayName(alias, cfg)} (${providerDisplayName(cfg.provider)})`,
66137
}));
@@ -102,6 +173,8 @@ function createModelChoices(models: Record<string, ModelAlias>): readonly ModelC
102173
export class ModelSelectorComponent extends Container implements Focusable {
103174
focused = false;
104175
private readonly opts: ModelSelectorOptions;
176+
private readonly models: Record<string, ModelAlias>;
177+
private readonly currentValue: string;
105178
private readonly list: SearchableList<ModelChoice>;
106179
/** Per-model effort override set by ←/→; absent → the default draft. */
107180
private readonly effortOverrides = new Map<string, string>();
@@ -111,8 +184,22 @@ export class ModelSelectorComponent extends Container implements Focusable {
111184
constructor(opts: ModelSelectorOptions) {
112185
super();
113186
this.opts = opts;
114-
const choices = createModelChoices(opts.models);
115-
const selectedValue = opts.selectedValue ?? opts.currentValue;
187+
const normalized = normalizeModelChoices(opts.models);
188+
this.models = normalized.models;
189+
this.currentValue =
190+
resolveNormalizedModelAlias(
191+
normalized,
192+
opts.currentValue,
193+
opts.models[opts.currentValue],
194+
) ?? opts.currentValue;
195+
const choices = createModelChoices(this.models);
196+
const selectedCandidate = opts.selectedValue ?? opts.currentValue;
197+
const selectedValue =
198+
resolveNormalizedModelAlias(
199+
normalized,
200+
selectedCandidate,
201+
opts.models[selectedCandidate],
202+
) ?? this.currentValue;
116203
const selectedIdx = choices.findIndex((choice) => choice.alias === selectedValue);
117204
this.list = new SearchableList({
118205
items: choices,
@@ -136,7 +223,7 @@ export class ModelSelectorComponent extends Container implements Focusable {
136223
private draftFor(choice: ModelChoice): string {
137224
const override = this.effortOverrides.get(choice.alias);
138225
if (override !== undefined) return override;
139-
if (choice.alias === this.opts.currentValue) {
226+
if (choice.alias === this.currentValue) {
140227
return coerceEffortForModel(choice.model, this.opts.currentEffort);
141228
}
142229
return effortLevelsForModel(choice.model).find((level) => level !== 'off') ?? 'off';
@@ -174,7 +261,7 @@ export class ModelSelectorComponent extends Container implements Focusable {
174261
override render(width: number): string[] {
175262
const searchable = this.opts.searchable === true;
176263
const view = this.list.view();
177-
const totalCount = Object.keys(this.opts.models).length;
264+
const totalCount = Object.keys(this.models).length;
178265

179266
const titleSuffix =
180267
searchable && view.query.length === 0
@@ -232,7 +319,7 @@ export class ModelSelectorComponent extends Container implements Focusable {
232319
const choice = view.items[i];
233320
if (choice === undefined) continue;
234321
const isSelected = i === view.selectedIndex;
235-
const isCurrent = choice.alias === this.opts.currentValue;
322+
const isCurrent = choice.alias === this.currentValue;
236323
const pointer = isSelected ? SELECT_POINTER : ' ';
237324
const truncatedName = truncateToWidth(choice.name, nameWidth, '…');
238325
const namePad = ' '.repeat(Math.max(0, nameWidth - visibleWidth(truncatedName)));
@@ -281,6 +368,10 @@ export class ModelSelectorComponent extends Container implements Focusable {
281368
return lines.map((line) => truncateToWidth(line, width));
282369
}
283370

371+
selectedAlias(): string | undefined {
372+
return this.selectedChoice()?.alias;
373+
}
374+
284375
private selectedChoice(): ModelChoice | undefined {
285376
return this.list.selected();
286377
}

0 commit comments

Comments
 (0)