Repository: PrefectHQ/prefect

Issue:
All agents with default labels are not registering their labels
Currently, all agents that have default labels (e.g., `LocalAgent`) are not registering their labels when first registering with Prefect Cloud.  This is because the cloud registration step occurs in the `__init__` of the base agent class, which is being called _before_ `self.labels` is fully realized.  

To be clear, the only consequence of this is that labels are not properly being displayed in the UI - because the agent queries for work using its labels, all work is being properly directed.



Discussion:

Resolving patch:
diff --git a/src/prefect/agent/agent.py b/src/prefect/agent/agent.py
--- a/src/prefect/agent/agent.py
+++ b/src/prefect/agent/agent.py
@@ -133,9 +133,6 @@ def __init__(
         self.logger.debug(f"Prefect backend: {config.backend}")
 
         self.client = Client(api_token=token)
-        if config.backend == "cloud":
-            self._verify_token(token)
-            self.client.attach_headers({"X-PREFECT-AGENT-ID": self._register_agent()})
 
     def _verify_token(self, token: str) -> None:
         """
@@ -176,6 +173,10 @@ def start(self) -> None:
         The main entrypoint to the agent. This function loops and constantly polls for
         new flow runs to deploy
         """
+        if config.backend == "cloud":
+            self._verify_token(self.client.get_auth_token())
+            self.client.attach_headers({"X-PREFECT-AGENT-ID": self._register_agent()})
+
         try:
             self.setup()
 


Test patch:
diff --git a/tests/agent/test_agent.py b/tests/agent/test_agent.py
--- a/tests/agent/test_agent.py
+++ b/tests/agent/test_agent.py
@@ -116,7 +116,7 @@ def test_agent_log_level_debug(runner_token):
 
 def test_agent_fails_no_auth_token(cloud_api):
     with pytest.raises(AuthorizationError):
-        agent = Agent()
+        agent = Agent().start()
 
 
 def test_agent_fails_no_runner_token(monkeypatch, cloud_api):
@@ -134,7 +134,7 @@ def test_agent_fails_no_runner_token(monkeypatch, cloud_api):
     monkeypatch.setattr("requests.Session", session)
 
     with pytest.raises(AuthorizationError):
-        agent = Agent()
+        agent = Agent().start()
 
 
 def test_query_flow_runs(monkeypatch, runner_token, cloud_api):
@@ -506,7 +506,8 @@ def test_agent_registration_and_id(monkeypatch, cloud_api):
         "prefect.agent.agent.Client.register_agent", MagicMock(return_value="ID")
     )
 
-    agent = Agent()
+    agent = Agent(max_polls=1)
+    agent.start()
     assert agent._register_agent() == "ID"
     assert agent.client._attached_headers == {"X-PREFECT-AGENT-ID": "ID"}
 
diff --git a/tests/cli/test_agent.py b/tests/cli/test_agent.py
--- a/tests/cli/test_agent.py
+++ b/tests/cli/test_agent.py
@@ -24,10 +24,7 @@ def test_agent_help():
     assert "Manage Prefect agents." in result.output
 
 
-def test_agent_start_fails_no_token(monkeypatch, cloud_api):
-    start = MagicMock()
-    monkeypatch.setattr("prefect.agent.local.LocalAgent.start", start)
-
+def test_agent_start_fails_no_token(cloud_api):
     runner = CliRunner()
     result = runner.invoke(agent, ["start"])
     assert result.exit_code == 1