refactor(desktop): extract shared runtime/provider/model selection + env-var helpers (Phase 1B.1) (#1624)

Signed-off-by: Wes <wesbillman@users.noreply.github.com>
Co-authored-by: Brain <21994759fc7a6fa6b965551d35cfd7897d262f2495467f2d78694ddcfa6a5c7e@sprout-oss.stage.blox.sqprod.co>

desktop/src/features/agents/ui/CreateAgentDialog.tsx

@@ -46,11 +46,12 @@ import {
   requiredCredentialEnvKeys,
   runtimeSupportsLlmProviderSelection,
   shouldClearKnownModelForSelectionScope,
-  getProviderApiKeyEnvVar,
-  CUSTOM_PROVIDER_DROPDOWN_VALUE,
-  AUTO_PROVIDER_DROPDOWN_VALUE,
 } from "./personaDialogPickers";
-import { shouldClearModelForRuntimeChange } from "./personaRuntimeModel";
+import {
+  selectionOnProviderDropdownChange,
+  selectionOnRuntimeChange,
+  type RuntimeModelProviderSelection,
+} from "./runtimeModelProviderSelection";
 import {
   AgentModelField,
   AgentProviderField,
@@ -386,6 +387,22 @@ export function CreateAgentDialog({
     onOpenChange(next);
   }
 
+  const selection: RuntimeModelProviderSelection = {
+    provider,
+    model,
+    isCustomProviderEditing,
+    isCustomModelEditing,
+    envVars,
+  };
+
+  function applySelection(next: RuntimeModelProviderSelection) {
+    setProvider(next.provider);
+    setModel(next.model);
+    setIsCustomProviderEditing(next.isCustomProviderEditing);
+    setIsCustomModelEditing(next.isCustomModelEditing);
+    setEnvVars(next.envVars);
+  }
+
   function handleProviderChange(nextProviderId: string) {
     const previousRuntimeId = selectedRuntimeId;
     setSelectedRuntimeId(nextProviderId);
@@ -404,23 +421,15 @@ export function CreateAgentDialog({
       }
     }
 
-    // Clear model when switching to a runtime with a different model scope,
-    // and clear provider/model when switching away from provider-selection runtimes.
-    if (
-      shouldClearModelForRuntimeChange(previousRuntimeId, nextProviderId) ||
-      shouldClearKnownModelForSelectionScope({
-        model,
-        provider,
-        runtime: nextProviderId,
-      })
-    ) {
-      setModel("");
-      setIsCustomModelEditing(false);
-    }
-    if (!runtimeSupportsLlmProviderSelection(nextProviderId)) {
-      setProvider("");
-      setIsCustomProviderEditing(false);
-    }
+    applySelection(
+      selectionOnRuntimeChange(selection, {
+        previousRuntime: previousRuntimeId,
+        nextRuntime: nextProviderId,
+        nextRuntimeCanChooseProvider:
+          runtimeSupportsLlmProviderSelection(nextProviderId),
+        lockedRuntimeReset: "provider-only",
+      }),
+    );
   }
 
   function handleRunOnChange(value: string) {
@@ -432,48 +441,13 @@ export function CreateAgentDialog({
 
   // Provider dropdown handler for local-mode provider field — mirrors Edit.
   function handleProviderDropdownChange(nextValue: string) {
-    if (nextValue === CUSTOM_PROVIDER_DROPDOWN_VALUE) {
-      const previousApiKey = getProviderApiKeyEnvVar(provider);
-      if (previousApiKey) {
-        setEnvVars((current) => {
-          const next = { ...current };
-          delete next[previousApiKey];
-          return next;
-        });
-      }
-      setIsCustomProviderEditing(true);
-      setProvider("");
-      return;
-    }
-
-    const nextProvider =
-      nextValue === AUTO_PROVIDER_DROPDOWN_VALUE ? "" : nextValue;
-
-    // Clear the old API key when switching providers.
-    const previousApiKey = getProviderApiKeyEnvVar(provider);
-    const nextApiKey = getProviderApiKeyEnvVar(nextProvider);
-    if (previousApiKey && previousApiKey !== nextApiKey) {
-      setEnvVars((current) => {
-        const next = { ...current };
-        delete next[previousApiKey];
-        return next;
-      });
-    }
-
-    setIsCustomProviderEditing(false);
-    setProvider(nextProvider);
-
-    if (
-      !isCustomModelEditing &&
-      shouldClearKnownModelForSelectionScope({
-        model,
-        provider: nextProvider,
+    applySelection(
+      selectionOnProviderDropdownChange(selection, {
         runtime: selectedRuntimeId,
-      })
-    ) {
-      setModel("");
-      setIsCustomModelEditing(false);
-    }
+        nextValue,
+        clearModelWhenApiKeyMissing: false,
+      }),
+    );
   }
 
   // Check provider config required fields are filled.

desktop/src/features/agents/ui/EditAgentDialog.tsx

@@ -27,7 +27,6 @@ import {
   formatRuntimeOptionLabel,
   getModelSelectValue,
   getPersonaProviderOptions,
-  getProviderApiKeyEnvVar,
   hasPersonaModelOption,
   isMissingRequiredDropdownField,
   NO_RUNTIME_DROPDOWN_VALUE,
@@ -45,8 +44,13 @@ import {
   resolveAgentCommandUpdate,
   resolveInheritedRuntimeSubmission,
   resolveRuntimeProviderCapability,
-  shouldClearModelForRuntimeChange,
 } from "./personaRuntimeModel";
+import {
+  selectionOnModelDropdownChange,
+  selectionOnProviderDropdownChange,
+  selectionOnRuntimeChange,
+  type RuntimeModelProviderSelection,
+} from "./runtimeModelProviderSelection";
 import { AgentCreationPreview } from "./AgentCreationPreview";
 import type { EnvVarsValue } from "./EnvVarsEditor";
 import { useRequiredCredentialState } from "./useRequiredCredentialState";
@@ -341,14 +345,27 @@ export function EditAgentDialog({
     selectedRuntimeId,
   ]);
 
+  const selection: RuntimeModelProviderSelection = {
+    provider,
+    model,
+    isCustomProviderEditing,
+    isCustomModelEditing,
+    envVars,
+  };
+
+  function applySelection(next: RuntimeModelProviderSelection) {
+    setProvider(next.provider);
+    setModel(next.model);
+    setIsCustomProviderEditing(next.isCustomProviderEditing);
+    setIsCustomModelEditing(next.isCustomModelEditing);
+    setEnvVars(next.envVars);
+  }
+
   function handleRuntimeDropdownChange(nextValue: string) {
     const nextRuntimeId =
       nextValue === NO_RUNTIME_DROPDOWN_VALUE ? "" : nextValue;
     const previousRuntimeId = selectedRuntimeId;
     const nextRuntime = runtimes.find((r) => r.id === nextRuntimeId);
-    const nextCanChooseProvider = runtimeSupportsLlmProviderSelection(
-      nextRuntime?.id ?? nextRuntimeId,
-    );
 
     // Mark that the user has made an explicit runtime choice. The catalog-arrival
     // effect will no longer overwrite selectedRuntimeId after this point.
@@ -392,98 +409,36 @@ export function EditAgentDialog({
       setAgentArgs(newArgs);
     }
 
-    // Clear model when switching away from a runtime with a different model scope.
-    if (
-      shouldClearModelForRuntimeChange(previousRuntimeId, nextRuntimeId) ||
-      shouldClearKnownModelForSelectionScope({
-        model,
-        provider,
-        runtime: nextRuntime?.id ?? nextRuntimeId,
-      })
-    ) {
-      setModel("");
-      setIsCustomModelEditing(false);
-    }
-
-    // When switching to a provider-locked runtime, clear provider state so no
-    // conflicting provider is persisted on a runtime that doesn't support it.
-    if (!nextCanChooseProvider) {
-      const previousProviderApiKeyEnvVar = getProviderApiKeyEnvVar(provider);
-      if (previousProviderApiKeyEnvVar) {
-        setEnvVars((current) => {
-          const next = { ...current };
-          delete next[previousProviderApiKeyEnvVar];
-          return next;
-        });
-      }
-      setIsCustomModelEditing(false);
-      setIsCustomProviderEditing(false);
-      setProvider("");
-    }
+    applySelection(
+      selectionOnRuntimeChange(selection, {
+        previousRuntime: previousRuntimeId,
+        nextRuntime: nextRuntime?.id ?? nextRuntimeId,
+        nextRuntimeCanChooseProvider: runtimeSupportsLlmProviderSelection(
+          nextRuntime?.id ?? nextRuntimeId,
+        ),
+        lockedRuntimeReset: "full",
+      }),
+    );
   }
 
   function handleProviderDropdownChange(nextValue: string) {
-    if (nextValue === CUSTOM_PROVIDER_DROPDOWN_VALUE) {
-      const previousProviderApiKeyEnvVar = getProviderApiKeyEnvVar(provider);
-      if (previousProviderApiKeyEnvVar) {
-        setEnvVars((current) => {
-          const next = { ...current };
-          delete next[previousProviderApiKeyEnvVar];
-          return next;
-        });
-      }
-      setIsCustomProviderEditing(true);
-      setProvider("");
-      return;
-    }
-
-    const nextProvider =
-      nextValue === AUTO_PROVIDER_DROPDOWN_VALUE ? "" : nextValue;
-
-    // Clear the old provider API key when switching providers.
-    const previousProviderApiKeyEnvVar = getProviderApiKeyEnvVar(provider);
-    const nextProviderApiKeyEnvVar = getProviderApiKeyEnvVar(nextProvider);
-    if (
-      previousProviderApiKeyEnvVar &&
-      previousProviderApiKeyEnvVar !== nextProviderApiKeyEnvVar
-    ) {
-      setEnvVars((current) => {
-        const next = { ...current };
-        delete next[previousProviderApiKeyEnvVar];
-        return next;
-      });
-    }
-
-    setIsCustomProviderEditing(false);
-    setProvider(nextProvider);
-
-    // Clear the model when switching to a provider that requires a different
-    // explicit model selection.
-    if (
-      !isCustomModelEditing &&
-      shouldClearKnownModelForSelectionScope({
-        model,
-        provider: nextProvider,
+    applySelection(
+      selectionOnProviderDropdownChange(selection, {
         runtime: selectedRuntime?.id ?? selectedRuntimeId,
-      })
-    ) {
-      setModel("");
-      setIsCustomModelEditing(false);
-    }
+        nextValue,
+        clearModelWhenApiKeyMissing: false,
+      }),
+    );
   }
 
   function handleModelDropdownChange(nextValue: string) {
-    if (nextValue === CUSTOM_MODEL_DROPDOWN_VALUE) {
-      setIsCustomModelEditing(true);
-      return;
-    }
-    if (nextValue === AUTO_MODEL_DROPDOWN_VALUE) {
-      setIsCustomModelEditing(false);
-      setModel("");
-      return;
-    }
-    setIsCustomModelEditing(false);
-    setModel(nextValue);
+    applySelection(
+      selectionOnModelDropdownChange(selection, {
+        nextValue,
+        clearKnownModelOnCustomEntry: false,
+        isModelCustom: false,
+      }),
+    );
   }
 
   function handleOpenChange(next: boolean) {

desktop/src/features/agents/ui/PersonaDialog.tsx

@@ -64,8 +64,17 @@ import {
   shouldClearKnownModelForSelectionScope,
   sortPersonaRuntimes,
 } from "./personaDialogPickers";
-import { shouldClearModelForRuntimeChange } from "./personaRuntimeModel";
 import { RequiredFieldLabel } from "./personaProviderModelFields";
+import {
+  envVarsMergingAdvancedEdit,
+  envVarsWithProviderApiKey,
+} from "./providerEnvVarUpdates";
+import {
+  selectionOnModelDropdownChange,
+  selectionOnProviderDropdownChange,
+  selectionOnRuntimeChange,
+  type RuntimeModelProviderSelection,
+} from "./runtimeModelProviderSelection";
 import {
   MODEL_DISCOVERY_LOADING_VALUE,
   usePersonaModelDiscovery,
@@ -624,136 +633,76 @@ export function PersonaDialog({
     runtime,
   ]);
 
-  function updateProviderApiKey(envKey: string, value: string) {
-    setEnvVars((current) => {
-      if ((current[envKey] ?? "") === value) {
-        return current;
-      }
-
-      const next = { ...current };
-      if (value.length > 0) {
-        next[envKey] = value;
-      } else {
-        delete next[envKey];
-      }
-      return next;
-    });
-  }
-
-  function removeEnvVar(envKey: string) {
-    setEnvVars((current) => {
-      if (!(envKey in current)) {
-        return current;
-      }
-
-      const next = { ...current };
-      delete next[envKey];
-      return next;
-    });
-  }
-
-  function clearManagedProviderApiKeyWhenLeaving(
-    previousProvider: string,
-    nextProvider: string,
-  ) {
-    const previousEnvVar = getProviderApiKeyEnvVar(previousProvider);
-    const nextEnvVar = getProviderApiKeyEnvVar(nextProvider);
-    if (previousEnvVar && previousEnvVar !== nextEnvVar) {
-      removeEnvVar(previousEnvVar);
-    }
-  }
-
   function handleProviderApiKeyChange(value: string) {
     if (!providerApiKeyConfig) {
       return;
     }
 
-    updateProviderApiKey(providerApiKeyConfig.envVar, value);
+    setEnvVars((current) =>
+      envVarsWithProviderApiKey(current, providerApiKeyConfig.envVar, value),
+    );
   }
 
   function handleAdvancedEnvVarsChange(nextAdvancedEnvVars: EnvVarsValue) {
-    setEnvVars((current) => {
-      const managedEnvVar = providerApiKeyConfig?.envVar ?? null;
-      if (!managedEnvVar || !(managedEnvVar in current)) {
-        return nextAdvancedEnvVars;
-      }
+    setEnvVars((current) =>
+      envVarsMergingAdvancedEdit(
+        current,
+        nextAdvancedEnvVars,
+        providerApiKeyConfig?.envVar ?? null,
+      ),
+    );
+  }
 
-      return {
-        ...nextAdvancedEnvVars,
-        [managedEnvVar]: current[managedEnvVar],
-      };
-    });
+  const selection: RuntimeModelProviderSelection = {
+    provider,
+    model,
+    isCustomProviderEditing,
+    isCustomModelEditing,
+    envVars,
+  };
+
+  function applySelection(next: RuntimeModelProviderSelection) {
+    setProvider(next.provider);
+    setModel(next.model);
+    setIsCustomProviderEditing(next.isCustomProviderEditing);
+    setIsCustomModelEditing(next.isCustomModelEditing);
+    setEnvVars(next.envVars);
   }
 
   function handleRuntimeDropdownChange(nextValue: string) {
     const nextRuntime =
       nextValue === NO_RUNTIME_DROPDOWN_VALUE ? "" : nextValue;
-    const previousRuntime = runtime;
-    const nextRuntimeCanChooseLlmProvider =
-      nextRuntime.trim().length > 0 &&
-      runtimeSupportsLlmProviderSelection(nextRuntime);
     setRuntime(nextRuntime);
-    if (
-      shouldClearModelForRuntimeChange(previousRuntime, nextRuntime) ||
-      shouldClearKnownModelForSelectionScope({
-        model,
-        provider,
-        runtime: nextRuntime,
-      })
-    ) {
-      setModel("");
-      setIsCustomModelEditing(false);
-    }
-    if (!nextRuntimeCanChooseLlmProvider) {
-      clearManagedProviderApiKeyWhenLeaving(provider, "");
-      setIsCustomModelEditing(false);
-      setIsCustomProviderEditing(false);
-      setProvider("");
-    }
+    applySelection(
+      selectionOnRuntimeChange(selection, {
+        previousRuntime: runtime,
+        nextRuntime,
+        nextRuntimeCanChooseProvider:
+          nextRuntime.trim().length > 0 &&
+          runtimeSupportsLlmProviderSelection(nextRuntime),
+        lockedRuntimeReset: "full",
+      }),
+    );
   }
 
   function handleProviderDropdownChange(nextValue: string) {
-    if (nextValue === CUSTOM_PROVIDER_DROPDOWN_VALUE) {
-      clearManagedProviderApiKeyWhenLeaving(provider, "");
-      setIsCustomProviderEditing(true);
-      setProvider("");
-      return;
-    }
-
-    const nextProvider =
-      nextValue === AUTO_PROVIDER_DROPDOWN_VALUE ? "" : nextValue;
-    clearManagedProviderApiKeyWhenLeaving(provider, nextProvider);
-    setIsCustomProviderEditing(false);
-    setProvider(nextProvider);
-    const requiredEnvVar = getProviderApiKeyEnvVar(nextProvider);
-    if (requiredEnvVar && !envVars[requiredEnvVar]?.trim()) {
-      setModel("");
-      setIsCustomModelEditing(false);
-    }
-    if (
-      !isCustomModelEditing &&
-      shouldClearKnownModelForSelectionScope({
-        model,
-        provider: nextProvider,
+    applySelection(
+      selectionOnProviderDropdownChange(selection, {
         runtime,
-      })
-    ) {
-      setModel("");
-      setIsCustomModelEditing(false);
-    }
+        nextValue,
+        clearModelWhenApiKeyMissing: true,
+      }),
+    );
   }
 
   function handleModelDropdownChange(nextValue: string) {
-    if (nextValue === CUSTOM_MODEL_DROPDOWN_VALUE) {
-      setIsCustomModelEditing(true);
-      if (!isModelCustom) {
-        setModel("");
-      }
-      return;
-    }
-
-    setIsCustomModelEditing(false);
-    setModel(nextValue === AUTO_MODEL_DROPDOWN_VALUE ? "" : nextValue);
+    applySelection(
+      selectionOnModelDropdownChange(selection, {
+        nextValue,
+        clearKnownModelOnCustomEntry: true,
+        isModelCustom,
+      }),
+    );
   }
 
   return (

desktop/src/features/agents/ui/providerEnvVarUpdates.test.mjs

@@ -0,0 +1,101 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+
+import {
+  envVarsClearingManagedApiKey,
+  envVarsMergingAdvancedEdit,
+  envVarsWithProviderApiKey,
+  envVarsWithoutKey,
+} from "./providerEnvVarUpdates.ts";
+
+test("envVarsWithProviderApiKey sets a non-empty value", () => {
+  const current = { OTHER: "x" };
+  const next = envVarsWithProviderApiKey(current, "ANTHROPIC_API_KEY", "sk-1");
+  assert.deepEqual(next, { OTHER: "x", ANTHROPIC_API_KEY: "sk-1" });
+  assert.notEqual(next, current);
+});
+
+test("envVarsWithProviderApiKey removes the key on empty value", () => {
+  const next = envVarsWithProviderApiKey(
+    { ANTHROPIC_API_KEY: "sk-1", OTHER: "x" },
+    "ANTHROPIC_API_KEY",
+    "",
+  );
+  assert.deepEqual(next, { OTHER: "x" });
+});
+
+test("envVarsWithProviderApiKey returns the same reference on no-op", () => {
+  const current = { ANTHROPIC_API_KEY: "sk-1" };
+  assert.equal(
+    envVarsWithProviderApiKey(current, "ANTHROPIC_API_KEY", "sk-1"),
+    current,
+  );
+  const empty = {};
+  assert.equal(
+    envVarsWithProviderApiKey(empty, "ANTHROPIC_API_KEY", ""),
+    empty,
+  );
+});
+
+test("envVarsWithoutKey removes a present key", () => {
+  assert.deepEqual(envVarsWithoutKey({ A: "1", B: "2" }, "A"), { B: "2" });
+});
+
+test("envVarsWithoutKey returns the same reference when the key is absent", () => {
+  const current = { A: "1" };
+  assert.equal(envVarsWithoutKey(current, "B"), current);
+});
+
+test("envVarsClearingManagedApiKey clears the previous provider's key on switch", () => {
+  const next = envVarsClearingManagedApiKey(
+    { ANTHROPIC_API_KEY: "sk-1", KEEP: "x" },
+    "anthropic",
+    "openai",
+  );
+  assert.deepEqual(next, { KEEP: "x" });
+});
+
+test("envVarsClearingManagedApiKey clears when leaving to a custom/empty provider", () => {
+  // The dialogs' CUSTOM-provider paths delete unconditionally; empty next
+  // provider has no managed key, so the inequality always holds — same rule.
+  const next = envVarsClearingManagedApiKey(
+    { ANTHROPIC_API_KEY: "sk-1" },
+    "anthropic",
+    "",
+  );
+  assert.deepEqual(next, {});
+});
+
+test("envVarsClearingManagedApiKey is a no-op when the managed key is shared or absent", () => {
+  const current = { ANTHROPIC_API_KEY: "sk-1" };
+  assert.equal(
+    envVarsClearingManagedApiKey(current, "anthropic", "anthropic"),
+    current,
+  );
+  const noManaged = { X: "1" };
+  assert.equal(
+    envVarsClearingManagedApiKey(noManaged, "", "openai"),
+    noManaged,
+  );
+});
+
+test("envVarsMergingAdvancedEdit preserves the managed key over the advanced edit", () => {
+  const next = envVarsMergingAdvancedEdit(
+    { ANTHROPIC_API_KEY: "sk-1", OLD: "x" },
+    { NEW: "y" },
+    "ANTHROPIC_API_KEY",
+  );
+  assert.deepEqual(next, { NEW: "y", ANTHROPIC_API_KEY: "sk-1" });
+});
+
+test("envVarsMergingAdvancedEdit passes the edit through when no managed key is set", () => {
+  const advanced = { NEW: "y" };
+  assert.equal(
+    envVarsMergingAdvancedEdit({ OLD: "x" }, advanced, null),
+    advanced,
+  );
+  assert.equal(
+    envVarsMergingAdvancedEdit({ OLD: "x" }, advanced, "ANTHROPIC_API_KEY"),
+    advanced,
+  );
+});

desktop/src/features/agents/ui/providerEnvVarUpdates.ts

@@ -0,0 +1,78 @@
+import type { EnvVarsValue } from "./EnvVarsEditor";
+import { getProviderApiKeyEnvVar } from "./personaDialogPickers";
+
+/**
+ * Pure env-var update helpers shared by the persona / create-agent /
+ * edit-agent dialogs. Every function returns the SAME reference when nothing
+ * changes, so `setEnvVars(fn(current))` skips a no-op re-render.
+ */
+
+/** Set `envKey` to `value`, or remove it when `value` is empty. */
+export function envVarsWithProviderApiKey(
+  current: EnvVarsValue,
+  envKey: string,
+  value: string,
+): EnvVarsValue {
+  if ((current[envKey] ?? "") === value) {
+    return current;
+  }
+
+  const next = { ...current };
+  if (value.length > 0) {
+    next[envKey] = value;
+  } else {
+    delete next[envKey];
+  }
+  return next;
+}
+
+/** Remove `envKey` when present. */
+export function envVarsWithoutKey(
+  current: EnvVarsValue,
+  envKey: string,
+): EnvVarsValue {
+  if (!(envKey in current)) {
+    return current;
+  }
+
+  const next = { ...current };
+  delete next[envKey];
+  return next;
+}
+
+/**
+ * Clear the previous provider's managed API key when switching providers.
+ * No-op when the previous provider has no managed key or the next provider
+ * uses the same one.
+ */
+export function envVarsClearingManagedApiKey(
+  current: EnvVarsValue,
+  previousProvider: string,
+  nextProvider: string,
+): EnvVarsValue {
+  const previousEnvVar = getProviderApiKeyEnvVar(previousProvider);
+  const nextEnvVar = getProviderApiKeyEnvVar(nextProvider);
+  if (previousEnvVar && previousEnvVar !== nextEnvVar) {
+    return envVarsWithoutKey(current, previousEnvVar);
+  }
+  return current;
+}
+
+/**
+ * Apply an Advanced-section env-vars edit while preserving the managed
+ * provider API key (which is edited via its own field, not the editor).
+ */
+export function envVarsMergingAdvancedEdit(
+  current: EnvVarsValue,
+  nextAdvancedEnvVars: EnvVarsValue,
+  managedEnvKey: string | null,
+): EnvVarsValue {
+  if (!managedEnvKey || !(managedEnvKey in current)) {
+    return nextAdvancedEnvVars;
+  }
+
+  return {
+    ...nextAdvancedEnvVars,
+    [managedEnvKey]: current[managedEnvKey],
+  };
+}

desktop/src/features/agents/ui/runtimeModelProviderSelection.test.mjs

@@ -0,0 +1,230 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+
+import {
+  selectionOnModelDropdownChange,
+  selectionOnProviderDropdownChange,
+  selectionOnRuntimeChange,
+} from "./runtimeModelProviderSelection.ts";
+
+const base = {
+  provider: "",
+  model: "",
+  isCustomProviderEditing: false,
+  isCustomModelEditing: false,
+  envVars: {},
+};
+
+// --- selectionOnRuntimeChange ---
+
+test("runtime change to a provider-locked runtime, full reset (Persona/Edit): clears provider, custom flags, and managed API key", () => {
+  const next = selectionOnRuntimeChange(
+    {
+      ...base,
+      provider: "anthropic",
+      model: "claude-4",
+      isCustomProviderEditing: true,
+      isCustomModelEditing: true,
+      envVars: { ANTHROPIC_API_KEY: "sk-1", KEEP: "x" },
+    },
+    {
+      previousRuntime: "buzz-agent",
+      nextRuntime: "claude",
+      nextRuntimeCanChooseProvider: false,
+      lockedRuntimeReset: "full",
+    },
+  );
+  assert.equal(next.provider, "");
+  assert.equal(next.isCustomProviderEditing, false);
+  assert.equal(next.isCustomModelEditing, false);
+  assert.deepEqual(next.envVars, { KEEP: "x" });
+});
+
+test("runtime change to a provider-locked runtime, provider-only reset (Create): keeps env vars and custom-model flag", () => {
+  const next = selectionOnRuntimeChange(
+    {
+      ...base,
+      provider: "anthropic",
+      envVars: { ANTHROPIC_API_KEY: "sk-1" },
+      isCustomModelEditing: true,
+      model: "my-custom",
+    },
+    {
+      previousRuntime: "buzz-agent",
+      nextRuntime: "claude",
+      nextRuntimeCanChooseProvider: false,
+      lockedRuntimeReset: "provider-only",
+    },
+  );
+  assert.equal(next.provider, "");
+  assert.equal(next.isCustomProviderEditing, false);
+  assert.deepEqual(next.envVars, { ANTHROPIC_API_KEY: "sk-1" });
+});
+
+test("runtime change between provider-selection runtimes keeps provider state", () => {
+  const current = {
+    ...base,
+    provider: "anthropic",
+    envVars: { ANTHROPIC_API_KEY: "sk-1" },
+  };
+  const next = selectionOnRuntimeChange(current, {
+    previousRuntime: "goose",
+    nextRuntime: "buzz-agent",
+    nextRuntimeCanChooseProvider: true,
+    lockedRuntimeReset: "full",
+  });
+  assert.equal(next.provider, "anthropic");
+  assert.deepEqual(next.envVars, { ANTHROPIC_API_KEY: "sk-1" });
+});
+
+// --- selectionOnProviderDropdownChange ---
+
+test("provider switch clears the previous managed API key and sets the provider", () => {
+  const next = selectionOnProviderDropdownChange(
+    {
+      ...base,
+      provider: "anthropic",
+      envVars: { ANTHROPIC_API_KEY: "sk-1", KEEP: "x" },
+    },
+    {
+      runtime: "buzz-agent",
+      nextValue: "openai",
+      clearModelWhenApiKeyMissing: false,
+    },
+  );
+  assert.equal(next.provider, "openai");
+  assert.equal(next.isCustomProviderEditing, false);
+  assert.deepEqual(next.envVars, { KEEP: "x" });
+});
+
+test("custom-provider entry clears the managed key and enters custom editing", () => {
+  const next = selectionOnProviderDropdownChange(
+    {
+      ...base,
+      provider: "anthropic",
+      envVars: { ANTHROPIC_API_KEY: "sk-1" },
+    },
+    {
+      runtime: "buzz-agent",
+      nextValue: "__custom_provider__",
+      clearModelWhenApiKeyMissing: false,
+    },
+  );
+  assert.equal(next.isCustomProviderEditing, true);
+  assert.equal(next.provider, "");
+  assert.deepEqual(next.envVars, {});
+});
+
+test("auto-provider selection maps to empty provider", () => {
+  const next = selectionOnProviderDropdownChange(
+    { ...base, provider: "anthropic", envVars: { ANTHROPIC_API_KEY: "sk-1" } },
+    {
+      runtime: "buzz-agent",
+      nextValue: "__auto_provider__",
+      clearModelWhenApiKeyMissing: false,
+    },
+  );
+  assert.equal(next.provider, "");
+  assert.deepEqual(next.envVars, {});
+});
+
+test("Persona mode clears the model when the new provider's API key is missing", () => {
+  const next = selectionOnProviderDropdownChange(
+    { ...base, model: "claude-4", provider: "" },
+    {
+      runtime: "buzz-agent",
+      nextValue: "anthropic",
+      clearModelWhenApiKeyMissing: true,
+    },
+  );
+  assert.equal(next.model, "");
+});
+
+test("Create/Edit mode keeps the model when the new provider's API key is missing", () => {
+  // claude-4 is scope-agnostic here: shouldClearKnownModelForSelectionScope
+  // only clears known models for the selection scope, and a custom string
+  // stays put — mirroring the dialogs' behavior without the persona flag.
+  const next = selectionOnProviderDropdownChange(
+    { ...base, model: "my-custom-model", provider: "" },
+    {
+      runtime: "buzz-agent",
+      nextValue: "anthropic",
+      clearModelWhenApiKeyMissing: false,
+    },
+  );
+  assert.equal(next.model, "my-custom-model");
+});
+
+test("custom-model editing suppresses the model-scope clear on provider switch", () => {
+  const next = selectionOnProviderDropdownChange(
+    { ...base, model: "anything", isCustomModelEditing: true },
+    {
+      runtime: "buzz-agent",
+      nextValue: "openai",
+      clearModelWhenApiKeyMissing: false,
+    },
+  );
+  assert.equal(next.model, "anything");
+  assert.equal(next.isCustomModelEditing, true);
+});
+
+// --- selectionOnModelDropdownChange ---
+
+test("custom-model entry with clear (Persona) drops a known model", () => {
+  const next = selectionOnModelDropdownChange(
+    { ...base, model: "known-model" },
+    {
+      nextValue: "__custom_model__",
+      clearKnownModelOnCustomEntry: true,
+      isModelCustom: false,
+    },
+  );
+  assert.equal(next.isCustomModelEditing, true);
+  assert.equal(next.model, "");
+});
+
+test("custom-model entry keeps an already-custom model (Persona) and any model (Edit)", () => {
+  const personaCustom = selectionOnModelDropdownChange(
+    { ...base, model: "already-custom" },
+    {
+      nextValue: "__custom_model__",
+      clearKnownModelOnCustomEntry: true,
+      isModelCustom: true,
+    },
+  );
+  assert.equal(personaCustom.model, "already-custom");
+
+  const edit = selectionOnModelDropdownChange(
+    { ...base, model: "known-model" },
+    {
+      nextValue: "__custom_model__",
+      clearKnownModelOnCustomEntry: false,
+      isModelCustom: false,
+    },
+  );
+  assert.equal(edit.model, "known-model");
+  assert.equal(edit.isCustomModelEditing, true);
+});
+
+test("auto-model selection clears the model; concrete selection sets it", () => {
+  const auto = selectionOnModelDropdownChange(
+    { ...base, model: "old", isCustomModelEditing: true },
+    {
+      nextValue: "__auto_model__",
+      clearKnownModelOnCustomEntry: false,
+      isModelCustom: false,
+    },
+  );
+  assert.equal(auto.model, "");
+  assert.equal(auto.isCustomModelEditing, false);
+
+  const concrete = selectionOnModelDropdownChange(
+    { ...base, model: "" },
+    {
+      nextValue: "gpt-5",
+      clearKnownModelOnCustomEntry: false,
+      isModelCustom: false,
+    },
+  );
+  assert.equal(concrete.model, "gpt-5");
+});

desktop/src/features/agents/ui/runtimeModelProviderSelection.ts

@@ -0,0 +1,166 @@
+import type { EnvVarsValue } from "./EnvVarsEditor";
+import {
+  AUTO_MODEL_DROPDOWN_VALUE,
+  AUTO_PROVIDER_DROPDOWN_VALUE,
+  CUSTOM_MODEL_DROPDOWN_VALUE,
+  CUSTOM_PROVIDER_DROPDOWN_VALUE,
+  getProviderApiKeyEnvVar,
+  shouldClearKnownModelForSelectionScope,
+} from "./personaDialogPickers";
+import { shouldClearModelForRuntimeChange } from "./personaRuntimeModel";
+import {
+  envVarsClearingManagedApiKey,
+  envVarsWithoutKey,
+} from "./providerEnvVarUpdates";
+
+/**
+ * Pure transition functions for the runtime -> LLM provider -> model dropdown
+ * state machine shared by the persona / create-agent / edit-agent dialogs.
+ * Each dialog applies the returned state to its own setters and layers its
+ * dialog-specific side effects (inherit pins, command sync, catalog memory)
+ * at the call site. Divergent behaviors are parameterized, never merged.
+ */
+export type RuntimeModelProviderSelection = {
+  provider: string;
+  model: string;
+  isCustomProviderEditing: boolean;
+  isCustomModelEditing: boolean;
+  envVars: EnvVarsValue;
+};
+
+export function selectionOnRuntimeChange(
+  current: RuntimeModelProviderSelection,
+  params: {
+    previousRuntime: string;
+    nextRuntime: string;
+    /** Caller-computed: whether the next runtime supports provider selection. */
+    nextRuntimeCanChooseProvider: boolean;
+    /**
+     * Persona/Edit clear the managed API key and custom-model editing flag
+     * when switching to a provider-locked runtime ("full"); Create clears
+     * only the provider selection ("provider-only").
+     */
+    lockedRuntimeReset: "full" | "provider-only";
+  },
+): RuntimeModelProviderSelection {
+  const next = { ...current };
+
+  if (
+    shouldClearModelForRuntimeChange(
+      params.previousRuntime,
+      params.nextRuntime,
+    ) ||
+    shouldClearKnownModelForSelectionScope({
+      model: current.model,
+      provider: current.provider,
+      runtime: params.nextRuntime,
+    })
+  ) {
+    next.model = "";
+    next.isCustomModelEditing = false;
+  }
+
+  if (!params.nextRuntimeCanChooseProvider) {
+    if (params.lockedRuntimeReset === "full") {
+      next.envVars = envVarsClearingManagedApiKey(
+        next.envVars,
+        current.provider,
+        "",
+      );
+      next.isCustomModelEditing = false;
+    }
+    next.isCustomProviderEditing = false;
+    next.provider = "";
+  }
+
+  return next;
+}
+
+export function selectionOnProviderDropdownChange(
+  current: RuntimeModelProviderSelection,
+  params: {
+    /** Runtime id used for the model-scope clearing rule. */
+    runtime: string;
+    nextValue: string;
+    /**
+     * Persona-only: clear the model when the newly selected provider's API
+     * key is not yet filled (model discovery cannot run without it).
+     */
+    clearModelWhenApiKeyMissing: boolean;
+  },
+): RuntimeModelProviderSelection {
+  const next = { ...current };
+
+  if (params.nextValue === CUSTOM_PROVIDER_DROPDOWN_VALUE) {
+    const previousEnvVar = getProviderApiKeyEnvVar(current.provider);
+    if (previousEnvVar) {
+      next.envVars = envVarsWithoutKey(next.envVars, previousEnvVar);
+    }
+    next.isCustomProviderEditing = true;
+    next.provider = "";
+    return next;
+  }
+
+  const nextProvider =
+    params.nextValue === AUTO_PROVIDER_DROPDOWN_VALUE ? "" : params.nextValue;
+  next.envVars = envVarsClearingManagedApiKey(
+    next.envVars,
+    current.provider,
+    nextProvider,
+  );
+  next.isCustomProviderEditing = false;
+  next.provider = nextProvider;
+
+  if (params.clearModelWhenApiKeyMissing) {
+    const requiredEnvVar = getProviderApiKeyEnvVar(nextProvider);
+    if (requiredEnvVar && !next.envVars[requiredEnvVar]?.trim()) {
+      next.model = "";
+      next.isCustomModelEditing = false;
+    }
+  }
+
+  // Guard on the PRE-transition editing flag, matching all three dialogs
+  // (their handlers read the render-scope value).
+  if (
+    !current.isCustomModelEditing &&
+    shouldClearKnownModelForSelectionScope({
+      model: current.model,
+      provider: nextProvider,
+      runtime: params.runtime,
+    })
+  ) {
+    next.model = "";
+    next.isCustomModelEditing = false;
+  }
+
+  return next;
+}
+
+export function selectionOnModelDropdownChange(
+  current: RuntimeModelProviderSelection,
+  params: {
+    nextValue: string;
+    /**
+     * Persona clears a known (non-custom) model when entering custom mode;
+     * Create/Edit keep it as the editable starting value.
+     */
+    clearKnownModelOnCustomEntry: boolean;
+    /** Caller-computed: whether the current model is outside the known options. */
+    isModelCustom: boolean;
+  },
+): RuntimeModelProviderSelection {
+  const next = { ...current };
+
+  if (params.nextValue === CUSTOM_MODEL_DROPDOWN_VALUE) {
+    next.isCustomModelEditing = true;
+    if (params.clearKnownModelOnCustomEntry && !params.isModelCustom) {
+      next.model = "";
+    }
+    return next;
+  }
+
+  next.isCustomModelEditing = false;
+  next.model =
+    params.nextValue === AUTO_MODEL_DROPDOWN_VALUE ? "" : params.nextValue;
+  return next;
+}