feat: update provider row after saving credentials (#8914)

Author: kalvinnchau

ui/goose2/src/features/settings/ui/ModelProviderRow.tsx

@@ -78,7 +78,6 @@ export function ModelProviderRow({
   const [setupOutput, setSetupOutput] = useState<SetupOutputLine[]>([]);
   const [setupError, setSetupError] = useState("");
   const [showSavedState, setShowSavedState] = useState(false);
-  const [preserveSetupLayout, setPreserveSetupLayout] = useState(false);
   const setupLineCounter = useRef(0);
   const hasLoadedConfig = useRef(false);
   const shouldRestorePanelFocus = useRef(false);
@@ -170,7 +169,6 @@ export function ModelProviderRow({
     setEditingKey(null);
     setError("");
     setShowSavedState(false);
-    setPreserveSetupLayout(false);
 
     const unlisten = await onModelSetupOutput(provider.id, appendSetupOutput);
 
@@ -195,7 +193,6 @@ export function ModelProviderRow({
     setExpanded((current) => {
       if (current) {
         setShowSavedState(false);
-        setPreserveSetupLayout(false);
       }
       return !current;
     });
@@ -295,8 +292,7 @@ export function ModelProviderRow({
         })),
       );
       await loadConfig();
-      setShowSavedState(true);
-      setPreserveSetupLayout(true);
+      setShowSavedState(false);
     } catch (nextError) {
       setError(
         nextError instanceof Error ? nextError.message : "Failed to save",
@@ -312,7 +308,6 @@ export function ModelProviderRow({
       setEditingKey(null);
       setError("");
       setShowSavedState(false);
-      setPreserveSetupLayout(false);
     } catch (nextError) {
       setError(
         nextError instanceof Error ? nextError.message : "Failed to remove",
@@ -403,7 +398,7 @@ export function ModelProviderRow({
       );
     }
 
-    if (hasFields && isConnected && !preserveSetupLayout) {
+    if (hasFields && isConnected) {
       return (
         <ConnectedFieldsPanel
           panelRef={panelRef}

ui/goose2/src/features/settings/ui/__tests__/ModelProviderRow.test.tsx

@@ -1,6 +1,6 @@
 import { render, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
-import type { ComponentType } from "react";
+import { useState, type ComponentType } from "react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { getModelProviders } from "@/features/providers/providerCatalog";
 import { ModelProviderRow } from "../ModelProviderRow";
@@ -110,4 +110,59 @@ describe("ModelProviderRow", () => {
 
     expect(screen.getByText(/model refresh failed/i)).toBeInTheDocument();
   });
+
+  it("switches from setup save to connected controls after first configuration", async () => {
+    const user = userEvent.setup();
+    let saved = false;
+
+    function SetupSaveRow() {
+      const [status, setStatus] = useState<"connected" | "not_configured">(
+        "not_configured",
+      );
+
+      return (
+        <ModelProviderRow
+          provider={modelProvider("google", status)}
+          onGetConfig={async () =>
+            saved
+              ? [
+                  {
+                    key: "GOOGLE_API_KEY",
+                    value: null,
+                    isSet: true,
+                    isSecret: true,
+                    required: true,
+                  },
+                ]
+              : []
+          }
+          onSaveFields={async () => {
+            saved = true;
+            setStatus("connected");
+          }}
+          onRemoveConfig={onRemoveConfig}
+          onCompleteNativeSetup={onCompleteNativeSetup}
+        />
+      );
+    }
+
+    render(<SetupSaveRow />);
+
+    await user.click(screen.getByRole("button", { name: /google gemini/i }));
+    await user.type(
+      await screen.findByPlaceholderText(/paste your api key/i),
+      "google-token",
+    );
+    await user.click(screen.getByRole("button", { name: /^save$/i }));
+
+    await waitFor(() =>
+      expect(
+        screen.getByRole("button", { name: /disconnect/i }),
+      ).toBeInTheDocument(),
+    );
+    expect(screen.getByRole("button", { name: /edit/i })).toBeInTheDocument();
+    expect(
+      screen.queryByRole("button", { name: /saved/i }),
+    ).not.toBeInTheDocument();
+  });
 });