refactor(ui): add callback injection to workflow forms

ui/components/providers/workflow/forms/add-via-credentials-form.tsx

@@ -8,9 +8,23 @@ import { BaseCredentialsForm } from "./base-credentials-form";
 export const AddViaCredentialsForm = ({
   searchParams,
   providerUid,
+  via,
+  onSuccess,
+  onBack,
+  formId,
+  hideActions,
+  onLoadingChange,
+  onValidityChange,
 }: {
   searchParams: { type: string; id: string };
   providerUid?: string;
+  via?: string | null;
+  onSuccess?: () => void;
+  onBack?: () => void;
+  formId?: string;
+  hideActions?: boolean;
+  onLoadingChange?: (isLoading: boolean) => void;
+  onValidityChange?: (isValid: boolean) => void;
 }) => {
   const providerType = searchParams.type as ProviderType;
   const providerId = searchParams.id;
@@ -19,7 +33,7 @@ export const AddViaCredentialsForm = ({
     return await addCredentialsProvider(formData);
   };
 
-  const successNavigationUrl = `/providers/test-connection?type=${providerType}&id=${providerId}`;
+  const successNavigationUrl = "/providers";
 
   return (
     <BaseCredentialsForm
@@ -28,6 +42,13 @@ export const AddViaCredentialsForm = ({
       providerUid={providerUid}
       onSubmit={handleAddCredentials}
       successNavigationUrl={successNavigationUrl}
+      via={via}
+      onSuccess={onSuccess}
+      onBack={onBack}
+      formId={formId}
+      hideActions={hideActions}
+      onLoadingChange={onLoadingChange}
+      onValidityChange={onValidityChange}
       submitButtonText="Next"
     />
   );

ui/components/providers/workflow/forms/add-via-role-form.tsx

@@ -8,9 +8,23 @@ import { BaseCredentialsForm } from "./base-credentials-form";
 export const AddViaRoleForm = ({
   searchParams,
   providerUid,
+  via,
+  onSuccess,
+  onBack,
+  formId,
+  hideActions,
+  onLoadingChange,
+  onValidityChange,
 }: {
   searchParams: { type: string; id: string };
   providerUid?: string;
+  via?: string | null;
+  onSuccess?: () => void;
+  onBack?: () => void;
+  formId?: string;
+  hideActions?: boolean;
+  onLoadingChange?: (isLoading: boolean) => void;
+  onValidityChange?: (isValid: boolean) => void;
 }) => {
   const providerType = searchParams.type as ProviderType;
   const providerId = searchParams.id;
@@ -19,7 +33,7 @@ export const AddViaRoleForm = ({
     return await addCredentialsProvider(formData);
   };
 
-  const successNavigationUrl = `/providers/test-connection?type=${providerType}&id=${providerId}`;
+  const successNavigationUrl = "/providers";
 
   return (
     <BaseCredentialsForm
@@ -28,6 +42,13 @@ export const AddViaRoleForm = ({
       providerUid={providerUid}
       onSubmit={handleAddCredentials}
       successNavigationUrl={successNavigationUrl}
+      via={via}
+      onSuccess={onSuccess}
+      onBack={onBack}
+      formId={formId}
+      hideActions={hideActions}
+      onLoadingChange={onLoadingChange}
+      onValidityChange={onValidityChange}
       submitButtonText="Next"
     />
   );

ui/components/providers/workflow/forms/base-credentials-form.tsx

@@ -2,6 +2,7 @@
 
 import { Divider } from "@heroui/divider";
 import { ChevronLeftIcon, ChevronRightIcon, Loader2 } from "lucide-react";
+import { useEffect } from "react";
 import { Control, UseFormSetValue } from "react-hook-form";
 
 import { Button } from "@/components/shadcn";
@@ -54,6 +55,13 @@ type BaseCredentialsFormProps = {
   providerUid?: string;
   onSubmit: (formData: FormData) => Promise<ApiResponse>;
   successNavigationUrl: string;
+  via?: string | null;
+  onSuccess?: () => void;
+  onBack?: () => void;
+  formId?: string;
+  hideActions?: boolean;
+  onLoadingChange?: (isLoading: boolean) => void;
+  onValidityChange?: (isValid: boolean) => void;
   submitButtonText?: string;
   showBackButton?: boolean;
 };
@@ -64,15 +72,23 @@ export const BaseCredentialsForm = ({
   providerUid,
   onSubmit,
   successNavigationUrl,
+  via,
+  onSuccess,
+  onBack,
+  formId,
+  hideActions = false,
+  onLoadingChange,
+  onValidityChange,
   submitButtonText = "Next",
   showBackButton = true,
 }: BaseCredentialsFormProps) => {
   const {
     form,
     isLoading,
+    isValid,
     handleSubmit,
     handleBackStep,
-    searchParamsObj,
+    effectiveVia,
     externalId,
   } = useCredentialsForm({
     providerType,
@@ -80,13 +96,25 @@ export const BaseCredentialsForm = ({
     providerUid,
     onSubmit,
     successNavigationUrl,
+    via,
+    onSuccess,
+    onBack,
   });
 
+  useEffect(() => {
+    onLoadingChange?.(isLoading);
+  }, [isLoading, onLoadingChange]);
+
+  useEffect(() => {
+    onValidityChange?.(isValid);
+  }, [isValid, onValidityChange]);
+
   const templateLinks = getAWSCredentialsTemplateLinks(externalId);
 
   return (
     <Form {...form}>
       <form
+        id={formId}
         onSubmit={form.handleSubmit(handleSubmit)}
         className="flex flex-col gap-4"
       >
@@ -112,7 +140,7 @@ export const BaseCredentialsForm = ({
 
         <Divider />
 
-        {providerType === "aws" && searchParamsObj.get("via") === "role" && (
+        {providerType === "aws" && effectiveVia === "role" && (
           <AWSRoleCredentialsForm
             control={form.control as unknown as Control<AWSCredentialsRole>}
             setValue={
@@ -122,7 +150,7 @@ export const BaseCredentialsForm = ({
             templateLinks={templateLinks}
           />
         )}
-        {providerType === "aws" && searchParamsObj.get("via") !== "role" && (
+        {providerType === "aws" && effectiveVia !== "role" && (
           <AWSStaticCredentialsForm
             control={form.control as unknown as Control<AWSCredentials>}
           />
@@ -132,36 +160,30 @@ export const BaseCredentialsForm = ({
             control={form.control as unknown as Control<AzureCredentials>}
           />
         )}
-        {providerType === "m365" &&
-          searchParamsObj.get("via") === "app_client_secret" && (
-            <M365ClientSecretCredentialsForm
-              control={
-                form.control as unknown as Control<M365ClientSecretCredentials>
-              }
-            />
-          )}
-        {providerType === "m365" &&
-          searchParamsObj.get("via") === "app_certificate" && (
-            <M365CertificateCredentialsForm
-              control={
-                form.control as unknown as Control<M365CertificateCredentials>
-              }
-            />
-          )}
-        {providerType === "gcp" &&
-          searchParamsObj.get("via") === "service-account" && (
-            <GCPServiceAccountKeyForm
-              control={form.control as unknown as Control<GCPServiceAccountKey>}
-            />
-          )}
-        {providerType === "gcp" &&
-          searchParamsObj.get("via") !== "service-account" && (
-            <GCPDefaultCredentialsForm
-              control={
-                form.control as unknown as Control<GCPDefaultCredentials>
-              }
-            />
-          )}
+        {providerType === "m365" && effectiveVia === "app_client_secret" && (
+          <M365ClientSecretCredentialsForm
+            control={
+              form.control as unknown as Control<M365ClientSecretCredentials>
+            }
+          />
+        )}
+        {providerType === "m365" && effectiveVia === "app_certificate" && (
+          <M365CertificateCredentialsForm
+            control={
+              form.control as unknown as Control<M365CertificateCredentials>
+            }
+          />
+        )}
+        {providerType === "gcp" && effectiveVia === "service-account" && (
+          <GCPServiceAccountKeyForm
+            control={form.control as unknown as Control<GCPServiceAccountKey>}
+          />
+        )}
+        {providerType === "gcp" && effectiveVia !== "service-account" && (
+          <GCPDefaultCredentialsForm
+            control={form.control as unknown as Control<GCPDefaultCredentials>}
+          />
+        )}
         {providerType === "kubernetes" && (
           <KubernetesCredentialsForm
             control={form.control as unknown as Control<KubernetesCredentials>}
@@ -170,7 +192,7 @@ export const BaseCredentialsForm = ({
         {providerType === "github" && (
           <GitHubCredentialsForm
             control={form.control}
-            credentialsType={searchParamsObj.get("via") || undefined}
+            credentialsType={effectiveVia || undefined}
           />
         )}
         {providerType === "iac" && (
@@ -190,50 +212,50 @@ export const BaseCredentialsForm = ({
             }
           />
         )}
-        {providerType === "alibabacloud" &&
-          searchParamsObj.get("via") === "role" && (
-            <AlibabaCloudRoleCredentialsForm
-              control={
-                form.control as unknown as Control<AlibabaCloudCredentialsRole>
-              }
-            />
-          )}
-        {providerType === "alibabacloud" &&
-          searchParamsObj.get("via") !== "role" && (
-            <AlibabaCloudStaticCredentialsForm
-              control={
-                form.control as unknown as Control<AlibabaCloudCredentials>
-              }
-            />
-          )}
+        {providerType === "alibabacloud" && effectiveVia === "role" && (
+          <AlibabaCloudRoleCredentialsForm
+            control={
+              form.control as unknown as Control<AlibabaCloudCredentialsRole>
+            }
+          />
+        )}
+        {providerType === "alibabacloud" && effectiveVia !== "role" && (
+          <AlibabaCloudStaticCredentialsForm
+            control={
+              form.control as unknown as Control<AlibabaCloudCredentials>
+            }
+          />
+        )}
 
-        <div className="flex w-full justify-end gap-4">
-          {showBackButton && requiresBackButton(searchParamsObj.get("via")) && (
+        {!hideActions && (
+          <div className="flex w-full justify-end gap-4">
+            {showBackButton && requiresBackButton(effectiveVia) && (
+              <Button
+                type="button"
+                variant="ghost"
+                size="lg"
+                onClick={handleBackStep}
+                disabled={isLoading}
+              >
+                {!isLoading && <ChevronLeftIcon size={24} />}
+                Back
+              </Button>
+            )}
             <Button
-              type="button"
-              variant="ghost"
+              type="submit"
+              variant="default"
               size="lg"
-              onClick={handleBackStep}
               disabled={isLoading}
             >
-              {!isLoading && <ChevronLeftIcon size={24} />}
-              Back
+              {isLoading ? (
+                <Loader2 className="animate-spin" />
+              ) : (
+                <ChevronRightIcon size={24} />
+              )}
+              {isLoading ? "Loading" : submitButtonText}
             </Button>
-          )}
-          <Button
-            type="submit"
-            variant="default"
-            size="lg"
-            disabled={isLoading}
-          >
-            {isLoading ? (
-              <Loader2 className="animate-spin" />
-            ) : (
-              <ChevronRightIcon size={24} />
-            )}
-            {isLoading ? "Loading" : submitButtonText}
-          </Button>
-        </div>
+          </div>
+        )}
       </form>
     </Form>
   );

ui/components/providers/workflow/forms/connect-account-form.tsx

@@ -9,7 +9,6 @@ import { z } from "zod";
 
 import { addProvider } from "@/actions/providers/providers";
 import { AwsMethodSelector } from "@/components/providers/organizations/aws-method-selector";
-import { OrgWizardModal } from "@/components/providers/organizations/org-wizard-modal";
 import { ProviderTitleDocs } from "@/components/providers/workflow/provider-title-docs";
 import { Button } from "@/components/shadcn";
 import { useToast } from "@/components/ui";
@@ -21,6 +20,29 @@ import { RadioGroupProvider } from "../../radio-group-provider";
 
 export type FormValues = z.infer<typeof addProviderFormSchema>;
 
+export interface ConnectAccountSuccessData {
+  id: string;
+  providerType: ProviderType;
+  uid: string;
+  alias: string | null;
+}
+
+interface ConnectAccountFormProps {
+  onSuccess?: (data: ConnectAccountSuccessData) => void;
+  onSelectOrganizations?: () => void;
+  onProviderTypeChange?: (providerType: ProviderType | null) => void;
+  formId?: string;
+  hideNavigation?: boolean;
+  onUiStateChange?: (state: {
+    showBack: boolean;
+    showAction: boolean;
+    actionLabel: string;
+    actionDisabled: boolean;
+    isLoading: boolean;
+  }) => void;
+  onBackHandlerChange?: (handler: () => void) => void;
+}
+
 // Helper function for labels and placeholders
 const getProviderFieldDetails = (providerType?: ProviderType) => {
   switch (providerType) {
@@ -82,11 +104,18 @@ const getProviderFieldDetails = (providerType?: ProviderType) => {
   }
 };
 
-export const ConnectAccountForm = () => {
+export const ConnectAccountForm = ({
+  onSuccess,
+  onSelectOrganizations,
+  onProviderTypeChange,
+  formId,
+  hideNavigation = false,
+  onUiStateChange,
+  onBackHandlerChange,
+}: ConnectAccountFormProps) => {
   const { toast } = useToast();
   const [prevStep, setPrevStep] = useState(1);
   const [awsMethod, setAwsMethod] = useState<"single" | null>(null);
-  const [isOrgModalOpen, setIsOrgModalOpen] = useState(false);
   const router = useRouter();
 
   const formSchema = addProviderFormSchema;
@@ -98,9 +127,12 @@ export const ConnectAccountForm = () => {
       providerUid: "",
       providerAlias: "",
     },
+    mode: "onChange",
+    reValidateMode: "onChange",
   });
 
   const providerType = form.watch("providerType");
+  const providerUid = form.watch("providerUid");
   const providerFieldDetails = getProviderFieldDetails(providerType);
 
   const isLoading = form.formState.isSubmitting;
@@ -155,10 +187,20 @@ export const ConnectAccountForm = () => {
         // Go to the next step after successful submission
         const {
           id,
-          attributes: { provider: providerType },
+          attributes: { provider: createdProviderType, uid, alias },
         } = data.data;
 
-        router.push(`/providers/add-credentials?type=${providerType}&id=${id}`);
+        if (onSuccess) {
+          onSuccess({
+            id,
+            providerType: createdProviderType,
+            uid: uid || values.providerUid,
+            alias: alias ?? values.providerAlias ?? null,
+          });
+          return;
+        }
+
+        router.push("/providers");
       }
     } catch (error: unknown) {
       console.error("Error during submission:", error);
@@ -199,9 +241,61 @@ export const ConnectAccountForm = () => {
     }
   }, [providerType]);
 
+  useEffect(() => {
+    onProviderTypeChange?.(providerType ?? null);
+  }, [onProviderTypeChange, providerType]);
+
+  useEffect(() => {
+    onBackHandlerChange?.(() => {
+      // If in UID form after choosing single, go back to method selector
+      if (prevStep === 2 && awsMethod === "single") {
+        setAwsMethod(null);
+        form.setValue("providerUid", "");
+        form.setValue("providerAlias", "");
+        return;
+      }
+
+      setPrevStep((prev) => prev - 1);
+      // Deselect the providerType if the user is going back to the first step
+      if (prevStep === 2) {
+        form.setValue("providerType", undefined as unknown as ProviderType);
+        setAwsMethod(null);
+      }
+      // Reset the providerUid and providerAlias fields when going back
+      form.setValue("providerUid", "");
+      form.setValue("providerAlias", "");
+    });
+  }, [onBackHandlerChange, prevStep, awsMethod, providerType, isLoading, form]);
+
+  useEffect(() => {
+    const canSubmit =
+      prevStep === 2 &&
+      (providerType !== "aws" || awsMethod === "single") &&
+      providerUid.trim().length > 0 &&
+      form.formState.isValid;
+
+    onUiStateChange?.({
+      showBack: prevStep === 2,
+      showAction:
+        prevStep === 2 && (providerType !== "aws" || awsMethod === "single"),
+      actionLabel: "Next",
+      actionDisabled: !canSubmit || isLoading,
+      isLoading,
+    });
+  }, [
+    awsMethod,
+    form.formState.isValid,
+    isLoading,
+    onUiStateChange,
+    prevStep,
+    providerType,
+    providerUid,
+  ]);
+
   return (
     <Form {...form}>
       <form
+        id={formId}
         onSubmit={form.handleSubmit(onSubmitClient)}
         className="flex flex-col gap-4"
       >
@@ -219,11 +313,9 @@ export const ConnectAccountForm = () => {
             <ProviderTitleDocs providerType={providerType} />
             <AwsMethodSelector
               onSelectSingle={() => setAwsMethod("single")}
-              onSelectOrganizations={() => setIsOrgModalOpen(true)}
-            />
-            <OrgWizardModal
-              open={isOrgModalOpen}
-              onOpenChange={setIsOrgModalOpen}
+              onSelectOrganizations={() => {
+                onSelectOrganizations?.();
+              }}
             />
           </>
         )}
@@ -254,39 +346,38 @@ export const ConnectAccountForm = () => {
               />
             </>
           )}
-        {/* Navigation buttons */}
-        <div className="flex w-full justify-end gap-4">
-          {/* Show "Back" button in Step 2 (skip when on AWS method selector) */}
-          {prevStep === 2 && (
-            <Button
-              type="button"
-              variant="ghost"
-              size="lg"
-              onClick={handleBackStep}
-              disabled={isLoading}
-            >
-              {!isLoading && <ChevronLeftIcon size={24} />}
-              Back
-            </Button>
-          )}
-          {/* Show "Next" button in Step 2 (only when UID form is visible) */}
-          {prevStep === 2 &&
-            (providerType !== "aws" || awsMethod === "single") && (
+        {!hideNavigation && (
+          <div className="flex w-full justify-end gap-4">
+            {prevStep === 2 && (
               <Button
-                type="submit"
-                variant="default"
+                type="button"
+                variant="ghost"
                 size="lg"
+                onClick={handleBackStep}
                 disabled={isLoading}
               >
-                {isLoading ? (
-                  <Loader2 className="animate-spin" />
-                ) : (
-                  <ChevronRightIcon size={24} />
-                )}
-                {isLoading ? "Loading" : "Next"}
+                {!isLoading && <ChevronLeftIcon size={24} />}
+                Back
               </Button>
             )}
-        </div>
+            {prevStep === 2 &&
+              (providerType !== "aws" || awsMethod === "single") && (
+                <Button
+                  type="submit"
+                  variant="default"
+                  size="lg"
+                  disabled={isLoading}
+                >
+                  {isLoading ? (
+                    <Loader2 className="animate-spin" />
+                  ) : (
+                    <ChevronRightIcon size={24} />
+                  )}
+                  {isLoading ? "Loading" : "Next"}
+                </Button>
+              )}
+          </div>
+        )}
       </form>
     </Form>
   );

ui/components/providers/workflow/forms/select-credentials-type/alibabacloud/select-via-alibabacloud.tsx

@@ -9,10 +9,12 @@ import { RadioGroupAlibabaCloudViaCredentialsTypeForm } from "./radio-group-alib
 
 interface SelectViaAlibabaCloudProps {
   initialVia?: string;
+  onViaChange?: (via: string) => void;
 }
 
 export const SelectViaAlibabaCloud = ({
   initialVia,
+  onViaChange,
 }: SelectViaAlibabaCloudProps) => {
   const router = useRouter();
   const form = useForm({
@@ -22,6 +24,11 @@ export const SelectViaAlibabaCloud = ({
   });
 
   const handleSelectionChange = (value: string) => {
+    if (onViaChange) {
+      onViaChange(value);
+      return;
+    }
+
     const url = new URL(window.location.href);
     url.searchParams.set("via", value);
     router.push(url.toString());

ui/components/providers/workflow/forms/select-credentials-type/aws/select-via-aws.tsx

@@ -9,9 +9,13 @@ import { RadioGroupAWSViaCredentialsTypeForm } from "./radio-group-aws-via-crede
 
 interface SelectViaAWSProps {
   initialVia?: string;
+  onViaChange?: (via: string) => void;
 }
 
-export const SelectViaAWS = ({ initialVia }: SelectViaAWSProps) => {
+export const SelectViaAWS = ({
+  initialVia,
+  onViaChange,
+}: SelectViaAWSProps) => {
   const router = useRouter();
   const form = useForm({
     defaultValues: {
@@ -20,6 +24,11 @@ export const SelectViaAWS = ({ initialVia }: SelectViaAWSProps) => {
   });
 
   const handleSelectionChange = (value: string) => {
+    if (onViaChange) {
+      onViaChange(value);
+      return;
+    }
+
     const url = new URL(window.location.href);
     url.searchParams.set("via", value);
     router.push(url.toString());

ui/components/providers/workflow/forms/select-credentials-type/gcp/add-via-service-account-form.tsx

@@ -8,9 +8,23 @@ import { BaseCredentialsForm } from "../../base-credentials-form";
 export const AddViaServiceAccountForm = ({
   searchParams,
   providerUid,
+  via,
+  onSuccess,
+  onBack,
+  formId,
+  hideActions,
+  onLoadingChange,
+  onValidityChange,
 }: {
   searchParams: { type: ProviderType; id: string };
   providerUid?: string;
+  via?: string | null;
+  onSuccess?: () => void;
+  onBack?: () => void;
+  formId?: string;
+  hideActions?: boolean;
+  onLoadingChange?: (isLoading: boolean) => void;
+  onValidityChange?: (isValid: boolean) => void;
 }) => {
   const providerType = searchParams.type;
   const providerId = searchParams.id;
@@ -19,7 +33,7 @@ export const AddViaServiceAccountForm = ({
     return await addCredentialsProvider(formData);
   };
 
-  const successNavigationUrl = `/providers/test-connection?type=${providerType}&id=${providerId}`;
+  const successNavigationUrl = "/providers";
 
   return (
     <BaseCredentialsForm
@@ -28,6 +42,13 @@ export const AddViaServiceAccountForm = ({
       providerUid={providerUid}
       onSubmit={handleAddCredentials}
       successNavigationUrl={successNavigationUrl}
+      via={via}
+      onSuccess={onSuccess}
+      onBack={onBack}
+      formId={formId}
+      hideActions={hideActions}
+      onLoadingChange={onLoadingChange}
+      onValidityChange={onValidityChange}
       submitButtonText="Next"
     />
   );

ui/components/providers/workflow/forms/select-credentials-type/gcp/select-via-gcp.tsx

@@ -9,9 +9,13 @@ import { RadioGroupGCPViaCredentialsTypeForm } from "./radio-group-gcp-via-crede
 
 interface SelectViaGCPProps {
   initialVia?: string;
+  onViaChange?: (via: string) => void;
 }
 
-export const SelectViaGCP = ({ initialVia }: SelectViaGCPProps) => {
+export const SelectViaGCP = ({
+  initialVia,
+  onViaChange,
+}: SelectViaGCPProps) => {
   const router = useRouter();
   const form = useForm({
     defaultValues: {
@@ -20,6 +24,11 @@ export const SelectViaGCP = ({ initialVia }: SelectViaGCPProps) => {
   });
 
   const handleSelectionChange = (value: string) => {
+    if (onViaChange) {
+      onViaChange(value);
+      return;
+    }
+
     const url = new URL(window.location.href);
     url.searchParams.set("via", value);
     router.push(url.toString());

ui/components/providers/workflow/forms/select-credentials-type/github/select-via-github.tsx

@@ -9,9 +9,13 @@ import { RadioGroupGitHubViaCredentialsTypeForm } from "./radio-group-github-via
 
 interface SelectViaGitHubProps {
   initialVia?: string;
+  onViaChange?: (via: string) => void;
 }
 
-export const SelectViaGitHub = ({ initialVia }: SelectViaGitHubProps) => {
+export const SelectViaGitHub = ({
+  initialVia,
+  onViaChange,
+}: SelectViaGitHubProps) => {
   const router = useRouter();
   const form = useForm({
     defaultValues: {
@@ -20,6 +24,11 @@ export const SelectViaGitHub = ({ initialVia }: SelectViaGitHubProps) => {
   });
 
   const handleSelectionChange = (value: string) => {
+    if (onViaChange) {
+      onViaChange(value);
+      return;
+    }
+
     const url = new URL(window.location.href);
     url.searchParams.set("via", value);
     router.push(url.toString());

ui/components/providers/workflow/forms/select-credentials-type/m365/select-via-m365.tsx

@@ -9,9 +9,13 @@ import { RadioGroupM365ViaCredentialsTypeForm } from "./radio-group-m365-via-cre
 
 interface SelectViaM365Props {
   initialVia?: string;
+  onViaChange?: (via: string) => void;
 }
 
-export const SelectViaM365 = ({ initialVia }: SelectViaM365Props) => {
+export const SelectViaM365 = ({
+  initialVia,
+  onViaChange,
+}: SelectViaM365Props) => {
   const router = useRouter();
   const form = useForm({
     defaultValues: {
@@ -20,6 +24,11 @@ export const SelectViaM365 = ({ initialVia }: SelectViaM365Props) => {
   });
 
   const handleSelectionChange = (value: string) => {
+    if (onViaChange) {
+      onViaChange(value);
+      return;
+    }
+
     const url = new URL(window.location.href);
     url.searchParams.set("via", value);
     router.push(url.toString());

ui/components/providers/workflow/forms/test-connection-form.tsx

@@ -6,7 +6,7 @@ import { Icon } from "@iconify/react";
 import { Loader2 } from "lucide-react";
 import Link from "next/link";
 import { useRouter } from "next/navigation";
-import { useState } from "react";
+import { useEffect, useState } from "react";
 import { useForm } from "react-hook-form";
 import { z } from "zod";
 
@@ -28,40 +28,53 @@ import { ProviderInfo } from "../..";
 
 type FormValues = z.input<typeof testConnectionFormSchema>;
 
-export const TestConnectionForm = ({
-  searchParams,
-  providerData,
-}: {
-  searchParams: { type: string; id: string; updated: string };
-  providerData: {
-    data: {
-      id: string;
-      type: string;
-      attributes: {
-        uid: string;
-        connection: {
-          connected: boolean | null;
-          last_checked_at: string | null;
-        };
-        provider: ProviderType;
-        alias: string;
-        scanner_args: Record<string, unknown>;
+export interface TestConnectionProviderData {
+  data: {
+    id: string;
+    type: string;
+    attributes: {
+      uid: string;
+      connection: {
+        connected: boolean | null;
+        last_checked_at: string | null;
       };
-      relationships: {
-        secret: {
-          data: {
-            type: string;
-            id: string;
-          } | null;
-        };
+      provider: ProviderType;
+      alias: string;
+      scanner_args: Record<string, unknown>;
+    };
+    relationships: {
+      secret: {
+        data: {
+          type: string;
+          id: string;
+        } | null;
       };
     };
   };
-}) => {
+}
+
+interface TestConnectionFormProps {
+  searchParams: { type: string; id: string; updated: string };
+  providerData: TestConnectionProviderData;
+  onSuccess?: () => void;
+  onResetCredentials?: () => void;
+  formId?: string;
+  hideActions?: boolean;
+  onLoadingChange?: (isLoading: boolean) => void;
+}
+
+export const TestConnectionForm = ({
+  searchParams,
+  providerData,
+  onSuccess,
+  onResetCredentials: onResetCredentialsCallback,
+  formId,
+  hideActions = false,
+  onLoadingChange,
+}: TestConnectionFormProps) => {
   const { toast } = useToast();
   const router = useRouter();
 
-  const providerType = searchParams.type;
   const providerId = searchParams.id;
 
   const [apiErrorMessage, setApiErrorMessage] = useState<string | null>(null);
@@ -85,6 +98,10 @@ export const TestConnectionForm = ({
   const isLoading = form.formState.isSubmitting;
   const isUpdated = searchParams?.updated === "true";
 
+  useEffect(() => {
+    onLoadingChange?.(isLoading || isResettingCredentials);
+  }, [isLoading, isResettingCredentials, onLoadingChange]);
+
   const onSubmitClient = async (values: FormValues) => {
     const formData = new FormData();
     formData.append("providerId", values.providerId);
@@ -123,7 +140,13 @@ export const TestConnectionForm = ({
           error: connected ? null : error || "Unknown error",
         });
 
-        if (connected && isUpdated) return router.push("/providers");
+        if (connected && isUpdated) {
+          if (onSuccess) {
+            onSuccess();
+            return;
+          }
+          return router.push("/providers");
+        }
 
         if (connected && !isUpdated) {
           try {
@@ -150,6 +173,11 @@ export const TestConnectionForm = ({
                 description: data.error,
               });
             } else {
+              if (onSuccess) {
+                onSuccess();
+                return;
+              }
+
               setIsRedirecting(true);
               router.push("/scans");
             }
@@ -174,7 +202,7 @@ export const TestConnectionForm = ({
     }
   };
 
-  const onResetCredentials = async () => {
+  const handleResetCredentials = async () => {
     setIsResettingCredentials(true);
 
     // Check if provider the provider has no credentials
@@ -183,20 +211,23 @@ export const TestConnectionForm = ({
     const hasNoCredentials = !providerSecretId;
 
     if (hasNoCredentials) {
-      // If no credentials, redirect to add credentials page
-      router.push(
-        `/providers/add-credentials?type=${providerType}&id=${providerId}`,
-      );
+      if (onResetCredentialsCallback) {
+        onResetCredentialsCallback();
+      } else {
+        router.push("/providers");
+      }
+      setIsResettingCredentials(false);
       return;
     }
 
     // If provider has credentials, delete them first
     try {
       await deleteCredentials(providerSecretId);
-      // After successful deletion, redirect to add credentials page
-      router.push(
-        `/providers/add-credentials?type=${providerType}&id=${providerId}`,
-      );
+      if (onResetCredentialsCallback) {
+        onResetCredentialsCallback();
+      } else {
+        router.push("/providers");
+      }
     } catch (error) {
       console.error("Failed to delete credentials:", error);
     } finally {
@@ -226,6 +257,7 @@ export const TestConnectionForm = ({
   return (
     <Form {...form}>
       <form
+        id={formId}
         onSubmit={form.handleSubmit(onSubmitClient)}
         className="flex flex-col gap-4"
       >
@@ -299,52 +331,58 @@ export const TestConnectionForm = ({
 
         <input type="hidden" name="providerId" value={providerId} />
 
-        <div className="flex w-full justify-end sm:gap-6">
-          {apiErrorMessage ? (
-            <Button variant="outline" size="lg" asChild>
-              <Link href="/providers">Back to providers</Link>
-            </Button>
-          ) : connectionStatus?.error ? (
-            <Button
-              onClick={isUpdated ? () => router.back() : onResetCredentials}
-              type="button"
-              variant="secondary"
-              size="lg"
-              disabled={isResettingCredentials}
-            >
-              {isResettingCredentials ? (
-                <Loader2 className="animate-spin" />
-              ) : (
-                <CheckIcon size={24} />
-              )}
-              {isResettingCredentials
-                ? "Loading"
-                : isUpdated
-                  ? "Update credentials"
-                  : "Reset credentials"}
-            </Button>
-          ) : (
-            <Button
-              type={
-                isUpdated && connectionStatus?.connected ? "button" : "submit"
-              }
-              variant="default"
-              size="lg"
-              disabled={isLoading}
-            >
-              {isLoading ? (
-                <Loader2 className="animate-spin" />
-              ) : (
-                !isUpdated && <RocketIcon size={24} />
-              )}
-              {isLoading
-                ? "Loading"
-                : isUpdated
-                  ? "Check connection"
-                  : "Launch scan"}
-            </Button>
-          )}
-        </div>
+        {!hideActions && (
+          <div className="flex w-full justify-end sm:gap-6">
+            {apiErrorMessage ? (
+              <Button variant="outline" size="lg" asChild>
+                <Link href="/providers">Back to providers</Link>
+              </Button>
+            ) : connectionStatus?.error ? (
+              <Button
+                onClick={
+                  isUpdated
+                    ? onResetCredentialsCallback || (() => router.back())
+                    : handleResetCredentials
+                }
+                type="button"
+                variant="secondary"
+                size="lg"
+                disabled={isResettingCredentials}
+              >
+                {isResettingCredentials ? (
+                  <Loader2 className="animate-spin" />
+                ) : (
+                  <CheckIcon size={24} />
+                )}
+                {isResettingCredentials
+                  ? "Loading"
+                  : isUpdated
+                    ? "Update credentials"
+                    : "Reset credentials"}
+              </Button>
+            ) : (
+              <Button
+                type={
+                  isUpdated && connectionStatus?.connected ? "button" : "submit"
+                }
+                variant="default"
+                size="lg"
+                disabled={isLoading}
+              >
+                {isLoading ? (
+                  <Loader2 className="animate-spin" />
+                ) : (
+                  !isUpdated && <RocketIcon size={24} />
+                )}
+                {isLoading
+                  ? "Loading"
+                  : isUpdated
+                    ? "Check connection"
+                    : "Launch scan"}
+              </Button>
+            )}
+          </div>
+        )}
       </form>
     </Form>
   );

ui/components/providers/workflow/forms/update-via-credentials-form.tsx

@@ -8,9 +8,23 @@ import { BaseCredentialsForm } from "./base-credentials-form";
 export const UpdateViaCredentialsForm = ({
   searchParams,
   providerUid,
+  via,
+  onSuccess,
+  onBack,
+  formId,
+  hideActions,
+  onLoadingChange,
+  onValidityChange,
 }: {
   searchParams: { type: string; id: string; secretId?: string };
   providerUid?: string;
+  via?: string | null;
+  onSuccess?: () => void;
+  onBack?: () => void;
+  formId?: string;
+  hideActions?: boolean;
+  onLoadingChange?: (isLoading: boolean) => void;
+  onValidityChange?: (isValid: boolean) => void;
 }) => {
   const providerType = searchParams.type as ProviderType;
   const providerId = searchParams.id;
@@ -20,7 +34,7 @@ export const UpdateViaCredentialsForm = ({
     return await updateCredentialsProvider(providerSecretId, formData);
   };
 
-  const successNavigationUrl = `/providers/test-connection?type=${providerType}&id=${providerId}&updated=true`;
+  const successNavigationUrl = "/providers";
 
   return (
     <BaseCredentialsForm
@@ -29,6 +43,13 @@ export const UpdateViaCredentialsForm = ({
       providerUid={providerUid}
       onSubmit={handleUpdateCredentials}
       successNavigationUrl={successNavigationUrl}
+      via={via}
+      onSuccess={onSuccess}
+      onBack={onBack}
+      formId={formId}
+      hideActions={hideActions}
+      onLoadingChange={onLoadingChange}
+      onValidityChange={onValidityChange}
       submitButtonText="Next"
     />
   );

ui/components/providers/workflow/forms/update-via-role-form.tsx

@@ -8,9 +8,23 @@ import { BaseCredentialsForm } from "./base-credentials-form";
 export const UpdateViaRoleForm = ({
   searchParams,
   providerUid,
+  via,
+  onSuccess,
+  onBack,
+  formId,
+  hideActions,
+  onLoadingChange,
+  onValidityChange,
 }: {
   searchParams: { type: string; id: string; secretId?: string };
   providerUid?: string;
+  via?: string | null;
+  onSuccess?: () => void;
+  onBack?: () => void;
+  formId?: string;
+  hideActions?: boolean;
+  onLoadingChange?: (isLoading: boolean) => void;
+  onValidityChange?: (isValid: boolean) => void;
 }) => {
   const providerType = searchParams.type as ProviderType;
   const providerId = searchParams.id;
@@ -20,7 +34,7 @@ export const UpdateViaRoleForm = ({
     return await updateCredentialsProvider(providerSecretId, formData);
   };
 
-  const successNavigationUrl = `/providers/test-connection?type=${providerType}&id=${providerId}&updated=true`;
+  const successNavigationUrl = "/providers";
 
   return (
     <BaseCredentialsForm
@@ -29,6 +43,13 @@ export const UpdateViaRoleForm = ({
       providerUid={providerUid}
       onSubmit={handleUpdateCredentials}
       successNavigationUrl={successNavigationUrl}
+      via={via}
+      onSuccess={onSuccess}
+      onBack={onBack}
+      formId={formId}
+      hideActions={hideActions}
+      onLoadingChange={onLoadingChange}
+      onValidityChange={onValidityChange}
       submitButtonText="Next"
     />
   );

ui/components/providers/workflow/forms/update-via-service-account-key-form.tsx

@@ -8,9 +8,23 @@ import { BaseCredentialsForm } from "./base-credentials-form";
 export const UpdateViaServiceAccountForm = ({
   searchParams,
   providerUid,
+  via,
+  onSuccess,
+  onBack,
+  formId,
+  hideActions,
+  onLoadingChange,
+  onValidityChange,
 }: {
   searchParams: { type: string; id: string; secretId?: string };
   providerUid?: string;
+  via?: string | null;
+  onSuccess?: () => void;
+  onBack?: () => void;
+  formId?: string;
+  hideActions?: boolean;
+  onLoadingChange?: (isLoading: boolean) => void;
+  onValidityChange?: (isValid: boolean) => void;
 }) => {
   const providerType = searchParams.type as ProviderType;
   const providerId = searchParams.id;
@@ -20,7 +34,7 @@ export const UpdateViaServiceAccountForm = ({
     return await updateCredentialsProvider(providerSecretId, formData);
   };
 
-  const successNavigationUrl = `/providers/test-connection?type=${providerType}&id=${providerId}&updated=true`;
+  const successNavigationUrl = "/providers";
 
   return (
     <BaseCredentialsForm
@@ -29,6 +43,13 @@ export const UpdateViaServiceAccountForm = ({
       providerUid={providerUid}
       onSubmit={handleUpdateCredentials}
       successNavigationUrl={successNavigationUrl}
+      via={via}
+      onSuccess={onSuccess}
+      onBack={onBack}
+      formId={formId}
+      hideActions={hideActions}
+      onLoadingChange={onLoadingChange}
+      onValidityChange={onValidityChange}
       submitButtonText="Next"
     />
   );

ui/hooks/use-credentials-form.ts

@@ -21,6 +21,9 @@ type UseCredentialsFormProps = {
   providerUid?: string;
   onSubmit: (formData: FormData) => Promise<ApiResponse>;
   successNavigationUrl: string;
+  via?: string | null;
+  onSuccess?: () => void;
+  onBack?: () => void;
 };
 
 export const useCredentialsForm = ({
@@ -29,26 +32,29 @@ export const useCredentialsForm = ({
   providerUid,
   onSubmit,
   successNavigationUrl,
+  via: viaOverride,
+  onSuccess,
+  onBack,
 }: UseCredentialsFormProps) => {
   const router = useRouter();
   const searchParamsObj = useSearchParams();
   const { data: session } = useSession();
-  const via = searchParamsObj.get("via");
+  const effectiveVia = viaOverride ?? searchParamsObj.get("via");
 
   // Select the appropriate schema based on provider type and via parameter
   const getFormSchema = () => {
-    if (providerType === "aws" && via === "role") {
+    if (providerType === "aws" && effectiveVia === "role") {
       return addCredentialsRoleFormSchema(providerType);
     }
-    if (providerType === "alibabacloud" && via === "role") {
+    if (providerType === "alibabacloud" && effectiveVia === "role") {
       return addCredentialsRoleFormSchema(providerType);
     }
-    if (providerType === "gcp" && via === "service-account") {
+    if (providerType === "gcp" && effectiveVia === "service-account") {
       return addCredentialsServiceAccountFormSchema(providerType);
     }
     // For GitHub and M365, we need to pass the via parameter to determine which fields are required
     if (providerType === "github" || providerType === "m365") {
-      return addCredentialsFormSchema(providerType, via);
+      return addCredentialsFormSchema(providerType, effectiveVia);
     }
     return addCredentialsFormSchema(providerType);
   };
@@ -63,7 +69,7 @@ export const useCredentialsForm = ({
     };
 
     // AWS Role credentials
-    if (providerType === "aws" && via === "role") {
+    if (providerType === "aws" && effectiveVia === "role") {
       const isCloudEnv = process.env.NEXT_PUBLIC_IS_CLOUD_ENV === "true";
       const defaultCredentialsType = isCloudEnv
         ? "aws-sdk-default"
@@ -82,7 +88,7 @@ export const useCredentialsForm = ({
     }
 
     // GCP Service Account
-    if (providerType === "gcp" && via === "service-account") {
+    if (providerType === "gcp" && effectiveVia === "service-account") {
       return {
         ...baseDefaults,
         [ProviderCredentialFields.SERVICE_ACCOUNT_KEY]: "",
@@ -106,7 +112,7 @@ export const useCredentialsForm = ({
         };
       case "m365":
         // M365 credentials based on via parameter
-        if (via === "app_client_secret") {
+        if (effectiveVia === "app_client_secret") {
           return {
             ...baseDefaults,
             [ProviderCredentialFields.CLIENT_ID]: "",
@@ -114,7 +120,7 @@ export const useCredentialsForm = ({
             [ProviderCredentialFields.TENANT_ID]: "",
           };
         }
-        if (via === "app_certificate") {
+        if (effectiveVia === "app_certificate") {
           return {
             ...baseDefaults,
             [ProviderCredentialFields.CLIENT_ID]: "",
@@ -141,19 +147,19 @@ export const useCredentialsForm = ({
         };
       case "github":
         // GitHub credentials based on via parameter
-        if (via === "personal_access_token") {
+        if (effectiveVia === "personal_access_token") {
           return {
             ...baseDefaults,
             [ProviderCredentialFields.PERSONAL_ACCESS_TOKEN]: "",
           };
         }
-        if (via === "oauth_app") {
+        if (effectiveVia === "oauth_app") {
           return {
             ...baseDefaults,
             [ProviderCredentialFields.OAUTH_APP_TOKEN]: "",
           };
         }
-        if (via === "github_app") {
+        if (effectiveVia === "github_app") {
           return {
             ...baseDefaults,
             [ProviderCredentialFields.GITHUB_APP_ID]: "",
@@ -178,7 +184,7 @@ export const useCredentialsForm = ({
           [ProviderCredentialFields.ATLAS_PRIVATE_KEY]: "",
         };
       case "alibabacloud":
-        if (via === "role") {
+        if (effectiveVia === "role") {
           return {
             ...baseDefaults,
             [ProviderCredentialFields.ALIBABACLOUD_ROLE_ARN]: "",
@@ -202,7 +208,7 @@ export const useCredentialsForm = ({
   const form = useForm({
     resolver: zodResolver(formSchema),
     defaultValues: defaultValues,
-    mode: "onSubmit",
+    mode: "onChange",
     reValidateMode: "onChange",
     criteriaMode: "all", // Show all errors for each field
   });
@@ -214,6 +220,11 @@ export const useCredentialsForm = ({
 
   // Handler for back button
   const handleBackStep = () => {
+    if (onBack) {
+      onBack();
+      return;
+    }
+
     const currentParams = new URLSearchParams(window.location.search);
     currentParams.delete("via");
     router.push(`?${currentParams.toString()}`);
@@ -234,19 +245,25 @@ export const useCredentialsForm = ({
 
     const isSuccess = handleServerResponse(data);
     if (isSuccess) {
+      if (onSuccess) {
+        onSuccess();
+        return;
+      }
       router.push(successNavigationUrl);
     }
   };
 
-  const { isSubmitting, errors } = form.formState;
+  const { isSubmitting, isValid, errors } = form.formState;
 
   return {
     form,
     isLoading: isSubmitting,
+    isValid,
     errors,
     handleSubmit,
     handleBackStep,
     searchParamsObj,
+    effectiveVia,
     externalId: session?.tenantId || "",
   };
 };