Merge pull request #3 from dreadnode/fix/eng-92-cli-agent-statelink-m…

…echanics

Fix/eng 92 cli agent statelink mechanics

dreadnode_cli/agent/cli.py

@@ -97,7 +97,19 @@ def push(
     env = {env_var.split("=")[0]: env_var.split("=")[1] for env_var in env_vars or []}
 
     agent_config = AgentConfig.read(directory)
-    server_config = UserConfig.read().get_server_config()
+    user_config = UserConfig.read()
+
+    if not user_config.active_profile_name:
+        raise Exception("No server profile is set, use [bold]dreadnode login[/] to authenticate")
+
+    if not agent_config.is_linked_to_profile(user_config.active_profile_name):
+        linked_profiles = ", ".join(agent_config.linked_profiles)
+        plural = "s" if len(agent_config.linked_profiles) > 1 else ""
+        raise Exception(
+            f"This agent is linked to the [magenta]{linked_profiles}[/] server profile{plural}, but the current server profile is [yellow]{user_config.active_profile_name}[/], create the agent again."
+        )
+
+    server_config = user_config.get_server_config()
 
     registry = get_registry(server_config)
 
@@ -124,7 +136,7 @@ def push(
         notes = notes or Prompt.ask("Notes?")
 
         agent = client.create_strike_agent(container, name, strike=agent_config.strike, notes=notes)
-        agent_config.add_link(agent.key, agent.id).write(directory)
+        agent_config.add_link(agent.key, agent.id, user_config.active_profile_name).write(directory)
     else:
         active_agent_id = agent_config.active
         if active_agent_id is None:
@@ -134,7 +146,16 @@ def push(
         print(":robot: Creating a new version ...")
         notes = notes or Prompt.ask("Notes?")
 
-        agent = client.create_strike_agent_version(str(active_agent_id), container, notes)
+        try:
+            agent = client.create_strike_agent_version(str(active_agent_id), container, notes)
+        except Exception as e:
+            # 404 is expected if the agent was created on a different server profile
+            if str(e).startswith("404"):
+                raise Exception(
+                    f"Agent '{active_agent_id}' not found for the current server profile, create the agent again."
+                ) from e
+            else:
+                raise e
 
     print(format_agent(agent))
 

dreadnode_cli/agent/config.py

@@ -8,6 +8,7 @@
 
 
 class AgentLink(pydantic.BaseModel):
+    profile: str
     id: UUID
     runs: list[UUID] = []
 
@@ -42,12 +43,19 @@ def write(self, directory: pathlib.Path = pathlib.Path(".")) -> None:
         with (directory / AGENT_CONFIG_FILENAME).open("w") as f:
             YAML().dump(self.model_dump(mode="json"), f)
 
-    def add_link(self, key: str, id: UUID) -> "AgentConfig":
+    def add_link(self, key: str, id: UUID, profile: str) -> "AgentConfig":
         if key not in self.links:
-            self.links[key] = AgentLink(id=id)
+            self.links[key] = AgentLink(id=id, profile=profile)
         self.active = key
         return self
 
+    @property
+    def linked_profiles(self) -> list[str]:
+        return list({link.profile for link in self.links.values()})
+
+    def is_linked_to_profile(self, profile: str) -> bool:
+        return any(link.profile == profile for link in self.links.values())
+
     def add_run(self, id: UUID) -> "AgentConfig":
         self.active_link.runs.append(id)
         return self

dreadnode_cli/agent/tests/test_config.py

@@ -21,19 +21,23 @@ def test_agent_config_add_link() -> None:
     config = AgentConfig(project_name="test")
     id = UUID("00000000-0000-0000-0000-000000000000")
 
-    config.add_link("test", id)
+    config.add_link("test", id, "test")
 
     assert config.active == "test"
     assert config.links["test"].id == id
     assert config.links["test"].runs == []
+    assert config.links["test"].profile == "test"
+    assert config.linked_profiles == ["test"]
+    assert config.is_linked_to_profile("test")
+    assert not config.is_linked_to_profile("other")
 
 
 def test_agent_config_add_run() -> None:
     config = AgentConfig(project_name="test")
     agent_id = UUID("00000000-0000-0000-0000-000000000000")
     run_id = UUID("11111111-1111-1111-1111-111111111111")
 
-    config.add_link("test", agent_id)
+    config.add_link("test", agent_id, "test")
     config.add_run(run_id)
 
     assert config.links["test"].runs == [run_id]
@@ -44,7 +48,7 @@ def test_agent_config_write_read(tmp_path: Path) -> None:
     agent_id = UUID("00000000-0000-0000-0000-000000000000")
     run_id = UUID("11111111-1111-1111-1111-111111111111")
 
-    config.add_link("test", agent_id)
+    config.add_link("test", agent_id, "test")
     config.add_run(run_id)
     config.write(tmp_path)
 
@@ -53,6 +57,7 @@ def test_agent_config_write_read(tmp_path: Path) -> None:
     assert loaded.active == "test"
     assert loaded.links["test"].id == agent_id
     assert loaded.links["test"].runs == [run_id]
+    assert loaded.links["test"].profile == "test"
 
 
 def test_agent_config_update_active() -> None:
@@ -61,11 +66,11 @@ def test_agent_config_update_active() -> None:
     id2 = UUID("11111111-1111-1111-1111-111111111111")
 
     # Add first link
-    config.add_link("test1", id1)
+    config.add_link("test1", id1, "test1")
     assert config.active == "test1"
 
     # Add second link
-    config.add_link("test2", id2)
+    config.add_link("test2", id2, "test2")
     assert config.active == "test2"
 
     # Remove active link

dreadnode_cli/config.py

@@ -50,6 +50,12 @@ def write(self) -> None:
         with USER_CONFIG_PATH.open("w") as f:
             YAML().dump(self.model_dump(mode="json"), f)
 
+    @property
+    def active_profile_name(self) -> str | None:
+        """Get the name of the active profile."""
+        self._update_active()
+        return self.active
+
     def get_server_config(self, profile: str | None = None) -> ServerConfig:
         """Get the server configuration for the given profile or None if not set."""