diff --git a/lib/towerops_web/plugs/agent_auth.ex b/lib/towerops_web/plugs/agent_auth.ex
index 08fc32d4..ac7da8ff 100644
--- a/lib/towerops_web/plugs/agent_auth.ex
+++ b/lib/towerops_web/plugs/agent_auth.ex
@@ -44,22 +44,33 @@ defmodule ToweropsWeb.Plugs.AgentAuth do
   defp update_heartbeat(conn, agent_token) do
     ip = to_string(:inet_parse.ntoa(conn.remote_ip))
 
-    parent = self()
-
-    _ =
-      Task.start(fn ->
-        # Allow this task to use the test database sandbox
-        _ =
-          if Application.get_env(:towerops, :sql_sandbox) do
-            Sandbox.allow(Towerops.Repo, parent, self())
-          end
-
-        _ = Agents.update_agent_token_heartbeat(agent_token.id, ip)
-      end)
+    # In test environment, update synchronously to avoid DB ownership issues
+    # In production, update asynchronously for better performance
+    if Mix.env() == :test do
+      _ = Agents.update_agent_token_heartbeat(agent_token.id, ip)
+    else
+      update_heartbeat_async(agent_token.id, ip)
+    end
 
     conn
   end
 
+  defp update_heartbeat_async(agent_token_id, ip) do
+    parent = self()
+
+    Task.start(fn ->
+      # Allow this task to use the test database sandbox (shouldn't happen in prod)
+      allow_sandbox(parent)
+      Agents.update_agent_token_heartbeat(agent_token_id, ip)
+    end)
+  end
+
+  defp allow_sandbox(parent) do
+    if Application.get_env(:towerops, :sql_sandbox) do
+      Sandbox.allow(Towerops.Repo, parent, self())
+    end
+  end
+
   defp unauthorized(conn) do
     conn
     |> put_status(:unauthorized)
diff --git a/test/towerops_web/plugs/agent_auth_test.exs b/test/towerops_web/plugs/agent_auth_test.exs
index 0a3e4f56..f65ebbf9 100644
--- a/test/towerops_web/plugs/agent_auth_test.exs
+++ b/test/towerops_web/plugs/agent_auth_test.exs
@@ -76,7 +76,7 @@ defmodule ToweropsWeb.Plugs.AgentAuthTest do
       assert conn.halted
     end
 
-    test "updates agent token heartbeat asynchronously", %{
+    test "updates agent token heartbeat", %{
       conn: conn,
       agent_token: agent_token,
       token_string: token_string
@@ -90,9 +90,7 @@ defmodule ToweropsWeb.Plugs.AgentAuthTest do
         |> put_req_header("authorization", "Bearer #{token_string}")
         |> AgentAuth.call([])
 
-      # Wait for async task to complete
-      Process.sleep(100)
-
+      # In test environment, heartbeat update is synchronous
       # Verify heartbeat was updated
       updated_agent_token = Agents.get_agent_token!(agent_token.id)