diff --git a/src/features/onboarding/hooks/use-agent-status.ts b/src/features/onboarding/hooks/use-agent-status.ts
new file mode 100644
index 00000000..4ddccef6
--- /dev/null
+++ b/src/features/onboarding/hooks/use-agent-status.ts
@@ -0,0 +1,23 @@
+"use client";
+
+import { useQuery } from "@tanstack/react-query";
+import { useOnboarding } from "@onboardjs/react";
+import { getAgentStatusAction } from "@/features/onboarding/actions/get-agent-status.action";
+import type { OnboardingAgent } from "@/features/onboarding/types";
+
+export const useAgentStatus = () => {
+  const { state } = useOnboarding();
+  const agents = (state?.context.flowData.agents ?? []) as OnboardingAgent[];
+  const firstAgentId = agents[0]?.id;
+
+  return useQuery({
+    queryKey: ["onboarding-agent-status", firstAgentId],
+    queryFn: async () => {
+      if (!firstAgentId) return { connected: false };
+      const result = await getAgentStatusAction({ agentId: firstAgentId });
+      return result?.data ?? { connected: false };
+    },
+    refetchInterval: 3_000,
+    enabled: !!firstAgentId,
+  });
+};
diff --git a/src/features/onboarding/hooks/use-generate-edge-key.ts b/src/features/onboarding/hooks/use-generate-edge-key.ts
new file mode 100644
index 00000000..c14f9ee4
--- /dev/null
+++ b/src/features/onboarding/hooks/use-generate-edge-key.ts
@@ -0,0 +1,17 @@
+"use client";
+
+import { useQuery } from "@tanstack/react-query";
+import { generateEdgeKeyAction } from "@/features/onboarding/actions/generate-edge-key.action";
+
+export const useGenerateEdgeKey = (agentId: string) => {
+  return useQuery({
+    queryKey: ["onboarding-edge-key", agentId],
+    queryFn: async () => {
+      const result = await generateEdgeKeyAction({ agentId });
+      if (!result?.data?.key) throw new Error("Failed to generate key");
+      return result.data.key;
+    },
+    staleTime: Infinity,
+    enabled: !!agentId,
+  });
+};
diff --git a/src/features/onboarding/steps/step-agent-key.tsx b/src/features/onboarding/steps/step-agent-key.tsx
index 459eb864..0436a6c9 100644
--- a/src/features/onboarding/steps/step-agent-key.tsx
+++ b/src/features/onboarding/steps/step-agent-key.tsx
@@ -1,62 +1,53 @@
 "use client";
 
 import { useOnboarding } from "@onboardjs/react";
-import { useQuery } from "@tanstack/react-query";
-import { Button } from "@/components/ui/button";
 import { Loader2 } from "lucide-react";
+import { Button } from "@/components/ui/button";
 import { CodeSnippet } from "@/components/common/code-snippet";
-import { generateEdgeKeyAction } from "@/features/onboarding/actions/generate-edge-key.action";
+import { useGenerateEdgeKey } from "@/features/onboarding/hooks/use-generate-edge-key";
 import type { OnboardingAgent } from "@/features/onboarding/types";
 
 export const StepAgentKey = () => {
-    const { next, state } = useOnboarding();
-    const agents = (state?.context.flowData.agents ?? []) as OnboardingAgent[];
+  const { next, state } = useOnboarding();
+  const agents = (state?.context.flowData.agents ?? []) as OnboardingAgent[];
 
-    return (
-        <div className="flex flex-col gap-6">
-            <div>
-                <h1 className="text-2xl font-semibold">Connect your agent</h1>
-                <p className="text-sm text-muted-foreground mt-1">
-                    Run the command below on your server. The next step waits for the agent to connect.
-                </p>
-            </div>
-            {agents.map((agent) => (
-                <AgentKeyBlock key={agent.id} agent={agent} />
-            ))}
-            <Button type="button" onClick={() => next()}>
-                I&apos;ve run the command →
-            </Button>
-        </div>
-    );
+  return (
+    <div className="flex flex-col gap-6">
+      <div>
+        <h1 className="text-2xl font-semibold">Connect your agent</h1>
+        <p className="text-sm text-muted-foreground mt-1">
+          Run the command below on your server. The next step waits for the agent to connect.
+        </p>
+      </div>
+      {agents.map((agent) => (
+        <AgentKeyBlock key={agent.id} agent={agent} />
+      ))}
+      <Button type="button" onClick={() => next()}>
+        I&apos;ve run the command →
+      </Button>
+    </div>
+  );
 };
 
 const AgentKeyBlock = ({ agent }: { agent: OnboardingAgent }) => {
-    const { data, isLoading } = useQuery({
-        queryKey: ["edge-key", agent.id],
-        queryFn: async () => {
-            const result = await generateEdgeKeyAction({ agentId: agent.id });
-            if (!result?.data?.key) throw new Error("Failed to generate key");
-            return result.data.key;
-        },
-        staleTime: Infinity,
-    });
-
-    if (isLoading) {
-        return (
-            <div className="flex items-center gap-2 text-sm text-muted-foreground">
-                <Loader2 className="size-4 animate-spin" />
-                Generating key for {agent.name}…
-            </div>
-        );
-    }
-
-    const command = `portabase agent "${agent.name}" --key ${data}`;
+  const { data, isLoading } = useGenerateEdgeKey(agent.id);
 
+  if (isLoading) {
     return (
-        <div className="flex flex-col gap-3 p-4 rounded-lg border border-border">
-            <p className="text-sm font-medium">{agent.name}</p>
-            <CodeSnippet title="Installation Command" code={command} />
-            <CodeSnippet title="Agent Key (manual)" code={data ?? ""} />
-        </div>
+      <div className="flex items-center gap-2 text-sm text-muted-foreground">
+        <Loader2 className="size-4 animate-spin" />
+        Generating key for {agent.name}…
+      </div>
     );
+  }
+
+  const command = `portabase agent "${agent.name}" --key ${data}`;
+
+  return (
+    <div className="flex flex-col gap-3 p-4 rounded-lg border border-border">
+      <p className="text-sm font-medium">{agent.name}</p>
+      <CodeSnippet title="Installation Command" code={command} />
+      <CodeSnippet title="Agent Key (manual)" code={data ?? ""} />
+    </div>
+  );
 };
diff --git a/src/features/onboarding/steps/step-agent-waiting.tsx b/src/features/onboarding/steps/step-agent-waiting.tsx
index 77f11eea..487e547e 100644
--- a/src/features/onboarding/steps/step-agent-waiting.tsx
+++ b/src/features/onboarding/steps/step-agent-waiting.tsx
@@ -2,43 +2,29 @@
 
 import { useEffect } from "react";
 import { useOnboarding } from "@onboardjs/react";
-import { useQuery } from "@tanstack/react-query";
 import { Loader2 } from "lucide-react";
-import { getAgentStatusAction } from "@/features/onboarding/actions/get-agent-status.action";
+import { useAgentStatus } from "@/features/onboarding/hooks/use-agent-status";
 import type { OnboardingAgent } from "@/features/onboarding/types";
 
 export const StepAgentWaiting = () => {
-    const { next, state } = useOnboarding();
-    const agents = (state?.context.flowData.agents ?? []) as OnboardingAgent[];
-    const firstAgentId = agents[0]?.id;
+  const { next, state } = useOnboarding();
+  const agents = (state?.context.flowData.agents ?? []) as OnboardingAgent[];
+  const { data } = useAgentStatus();
 
-    const { data } = useQuery({
-        queryKey: ["agent-status", firstAgentId],
-        queryFn: async () => {
-            if (!firstAgentId) return { connected: false };
-            const result = await getAgentStatusAction({ agentId: firstAgentId });
-            return result?.data ?? { connected: false };
-        },
-        refetchInterval: 3000,
-        enabled: !!firstAgentId,
-    });
+  useEffect(() => {
+    if (data?.connected) {
+      next();
+    }
+  }, [data?.connected, next]);
 
-    useEffect(() => {
-        if (data?.connected) {
-            next();
-        }
-    }, [data?.connected, next]);
-
-    return (
-        <div className="flex flex-col items-center justify-center gap-4 h-full text-center">
-            <Loader2 className="size-10 animate-spin text-primary" />
-            <h1 className="text-xl font-semibold">Waiting for your agent</h1>
-            <p className="text-sm text-muted-foreground">
-                Checking connectivity every 3 seconds…
-            </p>
-            {agents.length > 0 && (
-                <p className="text-xs text-muted-foreground">Agent: {agents[0].name}</p>
-            )}
-        </div>
-    );
+  return (
+    <div className="flex flex-col items-center justify-center gap-4 h-full text-center">
+      <Loader2 className="size-10 animate-spin text-primary" />
+      <h1 className="text-xl font-semibold">Waiting for your agent</h1>
+      <p className="text-sm text-muted-foreground">Checking connectivity every 3 seconds…</p>
+      {agents.length > 0 && (
+        <p className="text-xs text-muted-foreground">Agent: {agents[0].name}</p>
+      )}
+    </div>
+  );
 };