fix(workbench): require models when adding providers (#458)

Amp-Thread-ID: https://ampcode.com/threads/T-01a013f3-4dbc-7295-8590-ef97be89b856

Author: lucas77778

packages/client/workbench/src/settings/providers/__tests__/add-flow.test.tsx

@@ -5,6 +5,7 @@ import { act, cleanup, fireEvent, render, screen, waitFor } from '@testing-libra
 import { afterEach, describe, expect, it, vi } from 'vitest';
 import type { AgentRuntimeOnboarding } from '../../../agent-runtime/onboarding';
 import { AddAccountForm, ServiceCatalogView } from '../add-flow';
+import type { ModelSources } from '../model-selection';
 
 function translateKey(key: string): string {
   return key;
@@ -39,6 +40,13 @@ function signedOutRuntimes(): AgentRuntimes {
   };
 }
 
+function addModel(id = 'test-model'): void {
+  fireEvent.change(screen.getByPlaceholderText('models.addPlaceholder'), {
+    target: { value: id },
+  });
+  fireEvent.click(screen.getByRole('button', { name: 'models.add' }));
+}
+
 describe('subscription account creation', () => {
   it('starts Claude login and creates the account only from the success callback', () => {
     const login = vi.fn();
@@ -57,6 +65,8 @@ describe('subscription account creation', () => {
       />,
     );
 
+    expect(screen.queryByRole('button', { name: 'login' })).toBeNull();
+    addModel('claude-sonnet-5');
     fireEvent.click(screen.getByRole('button', { name: 'login' }));
     expect(login).toHaveBeenCalledWith('claude-code', expect.any(Function));
     expect(onSubmit).not.toHaveBeenCalled();
@@ -67,6 +77,7 @@ describe('subscription account creation', () => {
       expect.objectContaining({
         service: 'claude-sub',
         credential: { type: 'oauth', agent: 'claude-code' },
+        models: [{ id: 'claude-sonnet-5' }],
       }),
     );
   });
@@ -126,12 +137,17 @@ describe('subscription account creation', () => {
       />,
     );
 
-    fireEvent.click(screen.getByRole('button', { name: 'form.submit' }));
+    const submit = screen.getByRole('button', { name: 'form.submit' });
+    expect(submit).toHaveProperty('disabled', true);
+    expect(screen.getByText('models.required')).toBeTruthy();
+    addModel('gpt-5.6-sol');
+    fireEvent.click(submit);
     expect(login).not.toHaveBeenCalled();
     expect(onSubmit).toHaveBeenCalledWith(
       expect.objectContaining({
         service: 'chatgpt-sub',
         credential: { type: 'oauth', agent: 'codex' },
+        models: [{ id: 'gpt-5.6-sol' }],
       }),
     );
   });
@@ -154,12 +170,16 @@ describe('non-subscription account creation', () => {
 
     fireEvent.change(screen.getByPlaceholderText('sk-ant-…'), { target: { value: 'sk-ant-test' } });
     fireEvent.click(screen.getByRole('button', { name: 'form.submit' }));
+    await waitFor(() => expect(onSubmit).not.toHaveBeenCalled());
+    addModel('claude-sonnet-5');
+    fireEvent.click(screen.getByRole('button', { name: 'form.submit' }));
     expect(login).not.toHaveBeenCalled();
     await waitFor(() =>
       expect(onSubmit).toHaveBeenCalledWith(
         expect.objectContaining({
           service: 'anthropic-api',
           credential: { type: 'api-key', key: 'sk-ant-test' },
+          models: [{ id: 'claude-sonnet-5' }],
         }),
       ),
     );
@@ -187,13 +207,15 @@ describe('non-subscription account creation', () => {
     const secret = document.querySelector<HTMLInputElement>('input[type="password"]');
     if (!secret) throw new Error('credential input missing');
     fireEvent.change(secret, { target: { value: 'stepfun-test-key' } });
+    addModel('step-3.5-flash');
     fireEvent.click(screen.getByRole('button', { name: 'form.submit' }));
 
     await waitFor(() =>
       expect(onSubmit).toHaveBeenCalledWith(
         expect.objectContaining({
           service: 'stepfun',
           credential: { type: 'api-key', key: 'stepfun-test-key' },
+          models: [{ id: 'step-3.5-flash' }],
         }),
       ),
     );
@@ -217,10 +239,17 @@ describe('non-subscription account creation', () => {
 
   it('adds LinkCode Gateway only after the explicit user action', async () => {
     const createKey = vi.fn().mockResolvedValue('lc-gateway-key');
+    const probeInline = vi.fn().mockResolvedValue([{ id: 'anthropic/claude-sonnet-5' }]);
+    const sources: ModelSources = {
+      probeInline,
+      probeAccount: vi.fn(),
+      oauth: vi.fn(),
+    };
     const onSubmit = vi.fn();
     render(
       <AddAccountForm
         serviceId="linkcode-gateway"
+        sources={sources}
         runtimes={undefined}
         onboarding={onboarding()}
         busy={false}
@@ -241,10 +270,15 @@ describe('non-subscription account creation', () => {
 
     await waitFor(() => expect(onSubmit).toHaveBeenCalledTimes(1));
     expect(createKey).toHaveBeenCalledWith('serviceName.linkcode-gateway');
+    expect(probeInline).toHaveBeenCalledWith('linkcode-gateway', {
+      type: 'auth-token',
+      token: 'lc-gateway-key',
+    });
     expect(onSubmit).toHaveBeenCalledWith(
       expect.objectContaining({
         service: 'linkcode-gateway',
         credential: { type: 'auth-token', token: 'lc-gateway-key' },
+        models: [{ id: 'anthropic/claude-sonnet-5' }],
       }),
     );
     expect(onSubmit.mock.calls[0]?.[0]).not.toHaveProperty('endpoint');
@@ -282,6 +316,7 @@ describe('non-subscription account creation', () => {
     const secret = container.querySelector('input[type="password"]');
     if (!secret) throw new Error('credential input missing');
     fireEvent.change(secret, { target: { value: 'cf-token' } });
+    addModel('gateway-model');
     fireEvent.click(screen.getByRole('button', { name: 'form.submit' }));
 
     await waitFor(() => expect(onSubmit).toHaveBeenCalled());
@@ -290,8 +325,40 @@ describe('non-subscription account creation', () => {
       service: 'cloudflare-gateway',
       credential: { type: 'auth-token', token: 'cf-token' },
       endpointParams: { account_id: '8f3a', gateway_id: 'prod' },
+      models: [{ id: 'gateway-model' }],
     });
     // One key can resolve to a different endpoint per agent, so none is pinned here.
     expect(account).not.toHaveProperty('endpoint');
   });
+
+  it('requires a model for a custom endpoint', async () => {
+    const onSubmit = vi.fn();
+    const { container } = render(
+      <AddAccountForm
+        serviceId="custom"
+        runtimes={undefined}
+        onboarding={onboarding()}
+        busy={false}
+        onBack={vi.fn()}
+        onSubmit={onSubmit}
+      />,
+    );
+
+    fireEvent.change(screen.getByRole('textbox', { name: 'form.label' }), {
+      target: { value: 'Private endpoint' },
+    });
+    const secret = container.querySelector('input[type="password"]');
+    if (!secret) throw new Error('credential input missing');
+    fireEvent.change(secret, { target: { value: 'private-key' } });
+    fireEvent.click(screen.getByRole('button', { name: 'form.submit' }));
+    await waitFor(() => expect(onSubmit).not.toHaveBeenCalled());
+
+    addModel('private-model');
+    fireEvent.click(screen.getByRole('button', { name: 'form.submit' }));
+    await waitFor(() =>
+      expect(onSubmit).toHaveBeenCalledWith(
+        expect.objectContaining({ models: [{ id: 'private-model' }] }),
+      ),
+    );
+  });
 });

packages/client/workbench/src/settings/providers/add-flow.tsx

@@ -9,7 +9,13 @@ import {
   serviceProtocols,
   templatePlaceholders,
 } from '@linkcode/providers';
-import type { Account, AccountModel, AccountProtocol, AgentRuntimes } from '@linkcode/schema';
+import type {
+  Account,
+  AccountModel,
+  AccountProtocol,
+  AccountSecret,
+  AgentRuntimes,
+} from '@linkcode/schema';
 import { AccountModelSchema } from '@linkcode/schema';
 import { AgentOnboardingCard, ServiceIcon } from '@linkcode/ui';
 import { Button } from 'coss-ui/components/button';
@@ -49,13 +55,13 @@ function newAccountBase(label: string): Pick<Account, 'id' | 'label' | 'createdA
 function oauthAccount(
   service: Extract<ServiceDescriptor, { kind: 'oauth' }>,
   label: string,
-  models: AccountModel[] = [],
+  models: AccountModel[],
 ): Account {
   return {
     ...newAccountBase(label),
     service: service.id,
     credential: { type: 'oauth', agent: service.agent },
-    ...(models.length > 0 && { models }),
+    models,
   };
 }
 
@@ -75,7 +81,7 @@ function catalogAccount(service: EndpointService, draft: CatalogDraft): Account
         ? { type: 'auth-token', token: draft.secret }
         : { type: 'api-key', key: draft.secret },
     ...(!isObjectEmpty(trimmed) && { endpointParams: trimmed }),
-    ...(draft.models.length > 0 && { models: draft.models }),
+    models: draft.models,
   };
 }
 
@@ -214,6 +220,7 @@ export function AddAccountForm({
         <LinkCodeGatewayForm
           service={service}
           access={linkCodeGateway}
+          sources={sources}
           busy={busy}
           onSubmit={onSubmit}
         />
@@ -239,11 +246,13 @@ type LinkCodeGatewayDraft = z.infer<typeof LinkCodeGatewayDraftSchema>;
 function LinkCodeGatewayForm({
   service,
   access,
+  sources,
   busy,
   onSubmit,
 }: {
   service: Extract<ServiceDescriptor, { kind: 'endpoint' }>;
   access: LinkCodeGatewayAccess | undefined;
+  sources: ModelSources | undefined;
   busy: boolean;
   onSubmit: (account: Account) => void;
 }): React.ReactNode {
@@ -257,6 +266,7 @@ function LinkCodeGatewayForm({
     resolver: zodResolver(LinkCodeGatewayDraftSchema),
     defaultValues: { label: t(`serviceName.${service.id}`) },
   });
+  const [createdKey, setCreatedKey] = useState<string | undefined>(undefined);
 
   if (!access?.signedIn) {
     return (
@@ -281,11 +291,17 @@ function LinkCodeGatewayForm({
       className="flex flex-col gap-3"
       onSubmit={handleSubmit(async ({ label }) => {
         try {
-          const key = await access.createKey(label);
+          if (!sources) throw new Error(t('models.fetchFailed'));
+          const key = createdKey ?? (await access.createKey(label));
+          setCreatedKey(key);
+          const credential: AccountSecret = { type: 'auth-token', token: key };
+          const models = await sources.probeInline(service.id, credential);
+          if (models.length === 0) throw new Error(t('models.required'));
           onSubmit({
             ...newAccountBase(label),
             service: service.id,
-            credential: { type: 'auth-token', token: key },
+            credential,
+            models,
           });
         } catch (error) {
           setError('root', {
@@ -305,7 +321,7 @@ function LinkCodeGatewayForm({
         </p>
       ) : null}
       <div className="flex justify-end pt-1">
-        <Button type="submit" size="sm" disabled={busy || isSubmitting}>
+        <Button type="submit" size="sm" disabled={busy || isSubmitting || !sources}>
           {t('linkCodeUseGateway')}
         </Button>
       </div>
@@ -444,6 +460,7 @@ function OauthCreateForm({
   const cue = onboarding.cues[service.agent] ?? { state: 'needs-login', phase: 'idle' as const };
   const loginInProgress =
     cue.state === 'needs-login' && (cue.phase === 'opening' || cue.phase === 'awaiting-code');
+  const hasModels = models.length > 0;
 
   return (
     <div className="flex flex-col gap-3">
@@ -461,6 +478,7 @@ function OauthCreateForm({
         disabled={busy || loginInProgress}
         onChange={setModels}
         onFetch={fetchModels}
+        required
         selected={models}
       />
       {loggedIn ? (
@@ -474,7 +492,7 @@ function OauthCreateForm({
             <Button
               type="button"
               size="sm"
-              disabled={busy || label.trim() === ''}
+              disabled={busy || label.trim() === '' || !hasModels}
               onClick={() => onSubmit(oauthAccount(service, label, models))}
             >
               {t('form.submit')}
@@ -488,7 +506,7 @@ function OauthCreateForm({
           onDownload={onboarding.download}
           onContinueUnverified={onboarding.acknowledgeUnverified}
           onLogin={
-            busy || label.trim() === ''
+            !hasModels || busy || label.trim() === ''
               ? undefined
               : (kind) => {
                   onboarding.login(kind, () => onSubmit(oauthAccount(service, label, models)));
@@ -506,7 +524,7 @@ const CatalogDraftSchema = z.object({
   label: z.string().min(1),
   secret: z.string().min(1),
   placeholders: z.record(z.string(), z.string()),
-  models: z.array(AccountModelSchema),
+  models: z.array(AccountModelSchema).min(1),
 });
 type CatalogDraft = z.infer<typeof CatalogDraftSchema>;
 
@@ -606,6 +624,7 @@ function CatalogAccountForm({
             disabled={busy}
             onChange={field.onChange}
             onFetch={fetchModels}
+            required
             selected={field.value}
           />
         )}
@@ -630,6 +649,9 @@ const CustomDraftSchema = z.object({
   protocol: z.string(),
   models: z.array(AccountModelSchema),
 });
+const CustomCreateDraftSchema = CustomDraftSchema.extend({
+  models: z.array(AccountModelSchema).min(1),
+});
 type CustomDraft = z.infer<typeof CustomDraftSchema>;
 
 /** The full free-form account form (any endpoint, any protocol) — no catalog seeding. */
@@ -651,7 +673,7 @@ function CustomAccountForm({
     handleSubmit,
     formState: { isSubmitting },
   } = useForm<CustomDraft>({
-    resolver: zodResolver(CustomDraftSchema),
+    resolver: zodResolver(account === undefined ? CustomCreateDraftSchema : CustomDraftSchema),
     defaultValues: {
       label: account?.label ?? '',
       type:
@@ -747,6 +769,7 @@ function CustomAccountForm({
             disabled={busy}
             onChange={field.onChange}
             onFetch={fetchModels}
+            required={account === undefined}
             selected={field.value}
           />
         )}

packages/client/workbench/src/settings/providers/model-selection.tsx

@@ -21,6 +21,7 @@ export interface ModelSelectionProps {
   selected: AccountModel[];
   onChange: (models: AccountModel[]) => void;
   disabled?: boolean;
+  required?: boolean;
 }
 
 /**
@@ -72,6 +73,7 @@ export function ModelSelection({
   selected,
   onChange,
   disabled = false,
+  required = false,
 }: ModelSelectionProps): React.ReactNode {
   const t = useTranslations('settings.providers');
   const [fetched, setFetched] = useState<AccountModel[]>([]);
@@ -112,7 +114,14 @@ export function ModelSelection({
   return (
     <div className="flex flex-col gap-2">
       <div className="flex items-center justify-between gap-2">
-        <span className="font-medium text-sm">{t('models.label')}</span>
+        <span className="font-medium text-sm">
+          {t('models.label')}
+          {required ? (
+            <span aria-hidden="true" className="text-destructive">
+              {' *'}
+            </span>
+          ) : null}
+        </span>
         {onFetch ? (
           <Button
             type="button"
@@ -131,6 +140,11 @@ export function ModelSelection({
       <p className="text-muted-foreground text-xs">
         {onFetch ? t('models.hint') : t('models.hintUnlistable')}
       </p>
+      {required && selected.length === 0 ? (
+        <p aria-live="polite" className="text-destructive text-xs">
+          {t('models.required')}
+        </p>
+      ) : null}
       {error !== undefined ? <p className="text-destructive text-xs">{error}</p> : null}
       {listed.length > 0 ? (
         <div className="flex max-h-56 flex-col gap-1 overflow-y-auto rounded-lg border border-border p-2">