From 4451e5028185a7d46d1c7bc596d3b0b07d4821f0 Mon Sep 17 00:00:00 2001 From: Sarah Wooders Date: Tue, 10 Mar 2026 14:17:51 -0700 Subject: [PATCH] fix: show custom BYOK models in selector (#1321) Co-authored-by: Letta Code Co-authored-by: cpacker --- src/cli/components/ModelSelector.tsx | 114 ++++++++++++++---- ...odel-selector-byok-custom-provider.test.ts | 44 +++++++ 2 files changed, 132 insertions(+), 26 deletions(-) create mode 100644 src/tests/cli/model-selector-byok-custom-provider.test.ts diff --git a/src/cli/components/ModelSelector.tsx b/src/cli/components/ModelSelector.tsx index bcafc1b..e696784 100644 --- a/src/cli/components/ModelSelector.tsx +++ b/src/cli/components/ModelSelector.tsx @@ -8,6 +8,10 @@ import { getCachedModelHandles, } from "../../agent/available-models"; import { models } from "../../agent/model"; +import { + listProviders, + type ProviderResponse, +} from "../../providers/byok-providers"; import { useTerminalWidth } from "../hooks/useTerminalWidth"; import { colors } from "./colors"; import { Text } from "./Text"; @@ -26,7 +30,57 @@ type ModelCategory = | "server-all"; // BYOK provider prefixes (ChatGPT OAuth + lc-* providers from /connect) -const BYOK_PROVIDER_PREFIXES = ["chatgpt-plus-pro/", "lc-"]; +const STATIC_BYOK_PROVIDER_PREFIXES = ["chatgpt-plus-pro/", "lc-"]; + +const PROVIDER_TYPE_TO_BASE_PROVIDER: Record = { + chatgpt_oauth: "chatgpt-plus-pro", + anthropic: "anthropic", + openai: "openai", + zai: "zai", + google_ai: "google_ai", + google_vertex: "google_vertex", + minimax: "minimax", + openrouter: "openrouter", + bedrock: "bedrock", +}; + +export function buildByokProviderAliases( + providers: Array>, +): Record { + const aliases: Record = { + "lc-anthropic": "anthropic", + "lc-openai": "openai", + "lc-zai": "zai", + "lc-gemini": "google_ai", + "chatgpt-plus-pro": "chatgpt-plus-pro", + }; + + for (const provider of providers) { + const baseProvider = PROVIDER_TYPE_TO_BASE_PROVIDER[provider.provider_type]; + if (baseProvider) { + aliases[provider.name] = baseProvider; + } + } + + return aliases; +} + +export function isByokHandleForSelector( + handle: string, + byokProviderAliases: Record, +): boolean { + if ( + STATIC_BYOK_PROVIDER_PREFIXES.some((prefix) => handle.startsWith(prefix)) + ) { + return true; + } + + const slashIndex = handle.indexOf("/"); + if (slashIndex === -1) return false; + + const provider = handle.slice(0, slashIndex); + return provider in byokProviderAliases; +} // Get tab order for model categories. // For self-hosted servers, only show server-specific tabs. @@ -125,6 +179,9 @@ export function ModelSelector({ const [isCached, setIsCached] = useState(cachedHandlesAtMount !== null); const [refreshing, setRefreshing] = useState(false); const [searchQuery, setSearchQuery] = useState(""); + const [byokProviderAliases, setByokProviderAliases] = useState< + Record + >(() => buildByokProviderAliases([])); const mountedRef = useRef(true); useEffect(() => { @@ -170,6 +227,19 @@ export function ModelSelector({ loadModels.current(forceRefreshOnMount ?? false); }, [forceRefreshOnMount]); + useEffect(() => { + (async () => { + try { + const providers = await listProviders(); + if (!mountedRef.current) return; + setByokProviderAliases(buildByokProviderAliases(providers)); + } catch { + if (!mountedRef.current) return; + setByokProviderAliases(buildByokProviderAliases([])); + } + })(); + }, []); + const pickPreferredStaticModel = useCallback( (handle: string): UiModel | undefined => { const staticCandidates = typedModels.filter((m) => m.handle === handle); @@ -242,11 +312,10 @@ export function ModelSelector({ pickPreferredStaticModel, ]); - // BYOK models: models from chatgpt-plus-pro or lc-* providers + // BYOK models: models from ChatGPT OAuth, standard lc-* providers, or any connected custom BYOK provider const isByokHandle = useCallback( - (handle: string) => - BYOK_PROVIDER_PREFIXES.some((prefix) => handle.startsWith(prefix)), - [], + (handle: string) => isByokHandleForSelector(handle, byokProviderAliases), + [byokProviderAliases], ); // Letta API (all): all non-BYOK handles from API, including recommended models. @@ -291,32 +360,25 @@ export function ModelSelector({ searchQuery, ]); - // Provider name mappings for BYOK -> models.json lookup - // Maps BYOK provider prefix to models.json provider prefix - const BYOK_PROVIDER_ALIASES: Record = { - "lc-anthropic": "anthropic", - "lc-openai": "openai", - "lc-zai": "zai", - "lc-gemini": "google_ai", - "chatgpt-plus-pro": "chatgpt-plus-pro", // No change needed - }; - // Convert BYOK handle to base provider handle for models.json lookup // e.g., "lc-anthropic/claude-3-5-haiku" -> "anthropic/claude-3-5-haiku" // e.g., "lc-gemini/gemini-2.0-flash" -> "google_ai/gemini-2.0-flash" - const toBaseHandle = useCallback((handle: string): string => { - const slashIndex = handle.indexOf("/"); - if (slashIndex === -1) return handle; + const toBaseHandle = useCallback( + (handle: string): string => { + const slashIndex = handle.indexOf("/"); + if (slashIndex === -1) return handle; - const provider = handle.slice(0, slashIndex); - const model = handle.slice(slashIndex + 1); - const baseProvider = BYOK_PROVIDER_ALIASES[provider]; + const provider = handle.slice(0, slashIndex); + const model = handle.slice(slashIndex + 1); + const baseProvider = byokProviderAliases[provider]; - if (baseProvider) { - return `${baseProvider}/${model}`; - } - return handle; - }, []); + if (baseProvider) { + return `${baseProvider}/${model}`; + } + return handle; + }, + [byokProviderAliases], + ); // BYOK (recommended): BYOK API handles that have matching entries in models.json const byokModels = useMemo(() => { diff --git a/src/tests/cli/model-selector-byok-custom-provider.test.ts b/src/tests/cli/model-selector-byok-custom-provider.test.ts new file mode 100644 index 0000000..110a39b --- /dev/null +++ b/src/tests/cli/model-selector-byok-custom-provider.test.ts @@ -0,0 +1,44 @@ +import { describe, expect, test } from "bun:test"; +import { + buildByokProviderAliases, + isByokHandleForSelector, +} from "../../cli/components/ModelSelector"; + +describe("ModelSelector custom BYOK provider detection", () => { + test("treats connected custom OpenAI providers as BYOK", () => { + const aliases = buildByokProviderAliases([ + { + name: "openai-sarah", + provider_type: "openai", + }, + ]); + + expect(aliases["openai-sarah"]).toBe("openai"); + expect(isByokHandleForSelector("openai-sarah/gpt-5.4-fast", aliases)).toBe( + true, + ); + }); + + test("maps custom OpenAI provider handles back to base openai handles", () => { + const aliases = buildByokProviderAliases([ + { + name: "openai-sarah", + provider_type: "openai", + }, + ]); + + const provider = "openai-sarah"; + const model = "gpt-5.4-fast"; + const baseProvider = aliases[provider]; + + expect(`${baseProvider}/${model}`).toBe("openai/gpt-5.4-fast"); + }); + + test("preserves existing lc-* aliases", () => { + const aliases = buildByokProviderAliases([]); + + expect(isByokHandleForSelector("lc-openai/gpt-5.4-fast", aliases)).toBe( + true, + ); + }); +});