Merge pull request #110 from hashicorp/bug/agentpoolworkspaces

KshitijaChoudhari · web-flow · commit d3f5401c0ed1 · 2026-03-27T14:43:28.000+05:30
bug/agentpoolworkspaces - Agent Pools using non-existent endpoints and failing to add allowed workspaces. (99)
diff --git a/examples/agent_pool.py b/examples/agent_pool.py
@@ -3,13 +3,14 @@
 This example demonstrates:
 1. Agent Pool CRUD operations (Create, Read, Update, Delete)
 2. Agent token creation and management
-3. Using the organization SDK client
+3. Workspace assignment using assign_to_workspaces and remove_from_workspaces
 4. Proper error handling
 
 Make sure to set the following environment variables:
 - TFE_TOKEN: Your Terraform Cloud/Enterprise API token
 - TFE_ADDRESS: Your Terraform Cloud/Enterprise URL (optional, defaults to https://app.terraform.io)
 - TFE_ORG: Your organization name
+- TFE_WORKSPACE_ID: A workspace ID for testing workspace assignment (optional)
 
 Usage:
     export TFE_TOKEN="your-token-here"
@@ -24,8 +25,10 @@
 from pytfe.errors import NotFound
 from pytfe.models import (
     AgentPoolAllowedWorkspacePolicy,
+    AgentPoolAssignToWorkspacesOptions,
     AgentPoolCreateOptions,
     AgentPoolListOptions,
+    AgentPoolRemoveFromWorkspacesOptions,
     AgentPoolUpdateOptions,
     AgentTokenCreateOptions,
 )
@@ -37,6 +40,9 @@ def main():
     token = os.environ.get("TFE_TOKEN")
     org = os.environ.get("TFE_ORG")
     address = os.environ.get("TFE_ADDRESS", "https://app.terraform.io")
+    workspace_id = os.environ.get(
+        "TFE_WORKSPACE_ID"
+    )  # optional, for workspace assignment
 
     if not token:
         print("TFE_TOKEN environment variable is required")
@@ -96,7 +102,27 @@ def main():
         updated_pool = client.agent_pools.update(new_pool.id, update_options)
         print(f"Updated agent pool name to: {updated_pool.name}")
 
-        # Example 5: Create an agent token
+        # Example 5: Workspace assignment
+        # assign_to_workspaces sends PATCH /agent-pools/:id with relationships.allowed-workspaces
+        # remove_from_workspaces sends PATCH /agent-pools/:id with relationships.excluded-workspaces
+        if workspace_id:
+            print("\n Assigning workspace to agent pool...")
+            updated_pool = client.agent_pools.assign_to_workspaces(
+                new_pool.id,
+                AgentPoolAssignToWorkspacesOptions(workspace_ids=[workspace_id]),
+            )
+            print(f"  Assigned workspace {workspace_id} to pool {updated_pool.name}")
+
+            print("\n Removing workspace from agent pool...")
+            updated_pool = client.agent_pools.remove_from_workspaces(
+                new_pool.id,
+                AgentPoolRemoveFromWorkspacesOptions(workspace_ids=[workspace_id]),
+            )
+            print(f"  Removed workspace {workspace_id} from pool {updated_pool.name}")
+        else:
+            print("\n Skipping workspace assignment (set TFE_WORKSPACE_ID to test)")
+
+        # Example 6: Create an agent token
         print("\n Creating agent token...")
         token_options = AgentTokenCreateOptions(
             description="SDK example token"  # Optional description
@@ -107,7 +133,7 @@ def main():
         if agent_token.token:
             print(f"  Token (first 10 chars): {agent_token.token[:10]}...")
 
-        # Example 6: List agent tokens
+        # Example 7: List agent tokens
         print("\n Listing agent tokens...")
         tokens = client.agent_tokens.list(new_pool.id)
 
@@ -117,7 +143,7 @@ def main():
         for token in token_list:
             print(f"  - {token.description or 'No description'} (ID: {token.id})")
 
-        # Example 7: Clean up - delete the token and pool
+        # Example 8: Clean up - delete the token and pool
         print("\n Cleaning up...")
         client.agent_tokens.delete(agent_token.id)
         print("Deleted agent token")
diff --git a/src/pytfe/models/agent.py b/src/pytfe/models/agent.py
@@ -82,6 +82,10 @@ class AgentPoolCreateOptions(BaseModel):
     organization_scoped: bool | None = None
     # Optional: Allowed workspace policy
     allowed_workspace_policy: AgentPoolAllowedWorkspacePolicy | None = None
+    # Optional: IDs of workspaces allowed to use this pool (sent as relationships.allowed-workspaces)
+    allowed_workspace_ids: list[str] = Field(default_factory=list)
+    # Optional: IDs of workspaces excluded from this pool (sent as relationships.excluded-workspaces)
+    excluded_workspace_ids: list[str] = Field(default_factory=list)
 
 
 class AgentPoolUpdateOptions(BaseModel):
@@ -93,6 +97,10 @@ class AgentPoolUpdateOptions(BaseModel):
     organization_scoped: bool | None = None
     # Optional: Allowed workspace policy
     allowed_workspace_policy: AgentPoolAllowedWorkspacePolicy | None = None
+    # Optional: Full replacement list of workspace IDs allowed to use this pool
+    allowed_workspace_ids: list[str] = Field(default_factory=list)
+    # Optional: Full replacement list of workspace IDs excluded from this pool
+    excluded_workspace_ids: list[str] = Field(default_factory=list)
 
 
 class AgentPoolReadOptions(BaseModel):
diff --git a/src/pytfe/resources/agent_pools.py b/src/pytfe/resources/agent_pools.py
@@ -203,7 +203,27 @@ def create(self, organization: str, options: AgentPoolCreateOptions) -> AgentPoo
                 options.allowed_workspace_policy.value
             )
 
-        payload = {"data": {"type": "agent-pools", "attributes": attributes}}
+        relationships: dict[str, Any] = {}
+        if options.allowed_workspace_ids:
+            relationships["allowed-workspaces"] = {
+                "data": [
+                    {"type": "workspaces", "id": ws_id}
+                    for ws_id in options.allowed_workspace_ids
+                ]
+            }
+        if options.excluded_workspace_ids:
+            relationships["excluded-workspaces"] = {
+                "data": [
+                    {"type": "workspaces", "id": ws_id}
+                    for ws_id in options.excluded_workspace_ids
+                ]
+            }
+
+        payload: dict[str, Any] = {
+            "data": {"type": "agent-pools", "attributes": attributes}
+        }
+        if relationships:
+            payload["data"]["relationships"] = relationships
 
         response = self.t.request("POST", path, json_body=payload)
         data = response.json()["data"]
@@ -320,13 +340,31 @@ def update(self, agent_pool_id: str, options: AgentPoolUpdateOptions) -> AgentPo
                 options.allowed_workspace_policy.value
             )
 
-        payload = {
+        relationships: dict[str, Any] = {}
+        if options.allowed_workspace_ids:
+            relationships["allowed-workspaces"] = {
+                "data": [
+                    {"type": "workspaces", "id": ws_id}
+                    for ws_id in options.allowed_workspace_ids
+                ]
+            }
+        if options.excluded_workspace_ids:
+            relationships["excluded-workspaces"] = {
+                "data": [
+                    {"type": "workspaces", "id": ws_id}
+                    for ws_id in options.excluded_workspace_ids
+                ]
+            }
+
+        payload: dict[str, Any] = {
             "data": {
                 "type": "agent-pools",
                 "id": agent_pool_id,
                 "attributes": attributes,
             }
         }
+        if relationships:
+            payload["data"]["relationships"] = relationships
 
         response = self.t.request("PATCH", path, json_body=payload)
         data = response.json()["data"]
@@ -371,13 +409,20 @@ def delete(self, agent_pool_id: str) -> None:
 
     def assign_to_workspaces(
         self, agent_pool_id: str, options: AgentPoolAssignToWorkspacesOptions
-    ) -> None:
-        """Assign an agent pool to workspaces.
+    ) -> AgentPool:
+        """Assign an agent pool to workspaces by updating the allowed-workspaces
+        relationship via PATCH /agent-pools/:id.
+
+        The provided workspace IDs become the new complete list of allowed
+        workspaces for this pool (full replacement, not append).
 
         Args:
             agent_pool_id: Agent pool ID
             options: Assignment options containing workspace IDs
 
+        Returns:
+            Updated AgentPool object
+
         Raises:
             ValueError: If parameters are invalid
             TFEError: If API request fails
@@ -388,26 +433,67 @@ def assign_to_workspaces(
         if not options.workspace_ids:
             raise ValueError("At least one workspace ID is required")
 
-        path = f"/api/v2/agent-pools/{agent_pool_id}/relationships/workspaces"
-
-        # Create data payload with workspace references
-        workspace_data = []
         for workspace_id in options.workspace_ids:
             if not valid_string_id(workspace_id):
                 raise ValueError(f"Invalid workspace ID: {workspace_id}")
-            workspace_data.append({"type": "workspaces", "id": workspace_id})
 
-        payload = {"data": workspace_data}
-        self.t.request("POST", path, json_body=payload)
+        path = f"/api/v2/agent-pools/{agent_pool_id}"
+        payload: dict[str, Any] = {
+            "data": {
+                "type": "agent-pools",
+                "id": agent_pool_id,
+                "attributes": {},
+                "relationships": {
+                    "allowed-workspaces": {
+                        "data": [
+                            {"type": "workspaces", "id": ws_id}
+                            for ws_id in options.workspace_ids
+                        ]
+                    }
+                },
+            }
+        }
+        response = self.t.request("PATCH", path, json_body=payload)
+        data = response.json()["data"]
+
+        # Extract agent pool data from response
+        attr = data.get("attributes", {}) or {}
+        agent_pool_data = {
+            "id": _safe_str(data.get("id")),
+            "name": _safe_str(attr.get("name")),
+            "created_at": attr.get("created-at"),
+            "organization_scoped": attr.get("organization-scoped"),
+            "allowed_workspace_policy": attr.get("allowed-workspace-policy"),
+            "agent_count": attr.get("agent-count", 0),
+        }
+
+        return AgentPool(
+            id=_safe_str(agent_pool_data["id"]) or "",
+            name=_safe_str(agent_pool_data["name"]),
+            created_at=cast(Any, agent_pool_data["created_at"]),
+            organization_scoped=_safe_bool(agent_pool_data["organization_scoped"]),
+            allowed_workspace_policy=_safe_workspace_policy(
+                agent_pool_data["allowed_workspace_policy"]
+            ),
+            agent_count=_safe_int(agent_pool_data["agent_count"]),
+        )
 
     def remove_from_workspaces(
         self, agent_pool_id: str, options: AgentPoolRemoveFromWorkspacesOptions
-    ) -> None:
-        """Remove an agent pool from workspaces.
+    ) -> AgentPool:
+        """Exclude workspaces from an agent pool by updating the excluded-workspaces
+        relationship via PATCH /agent-pools/:id.
+
+        Use this for organization-scoped pools where most workspaces are allowed
+        but you want to block specific ones.  The provided list becomes the new
+        complete excluded-workspaces list (full replacement, not append).
 
         Args:
             agent_pool_id: Agent pool ID
-            options: Removal options containing workspace IDs
+            options: Removal options containing workspace IDs to exclude
+
+        Returns:
+            Updated AgentPool object
 
         Raises:
             ValueError: If parameters are invalid
@@ -419,14 +505,47 @@ def remove_from_workspaces(
         if not options.workspace_ids:
             raise ValueError("At least one workspace ID is required")
 
-        path = f"/api/v2/agent-pools/{agent_pool_id}/relationships/workspaces"
-
-        # Create data payload with workspace references
-        workspace_data = []
         for workspace_id in options.workspace_ids:
             if not valid_string_id(workspace_id):
                 raise ValueError(f"Invalid workspace ID: {workspace_id}")
-            workspace_data.append({"type": "workspaces", "id": workspace_id})
 
-        payload = {"data": workspace_data}
-        self.t.request("DELETE", path, json_body=payload)
+        path = f"/api/v2/agent-pools/{agent_pool_id}"
+        payload: dict[str, Any] = {
+            "data": {
+                "type": "agent-pools",
+                "id": agent_pool_id,
+                "attributes": {},
+                "relationships": {
+                    "excluded-workspaces": {
+                        "data": [
+                            {"type": "workspaces", "id": ws_id}
+                            for ws_id in options.workspace_ids
+                        ]
+                    }
+                },
+            }
+        }
+        response = self.t.request("PATCH", path, json_body=payload)
+        data = response.json()["data"]
+
+        # Extract agent pool data from response
+        attr = data.get("attributes", {}) or {}
+        agent_pool_data = {
+            "id": _safe_str(data.get("id")),
+            "name": _safe_str(attr.get("name")),
+            "created_at": attr.get("created-at"),
+            "organization_scoped": attr.get("organization-scoped"),
+            "allowed_workspace_policy": attr.get("allowed-workspace-policy"),
+            "agent_count": attr.get("agent-count", 0),
+        }
+
+        return AgentPool(
+            id=_safe_str(agent_pool_data["id"]) or "",
+            name=_safe_str(agent_pool_data["name"]),
+            created_at=cast(Any, agent_pool_data["created_at"]),
+            organization_scoped=_safe_bool(agent_pool_data["organization_scoped"]),
+            allowed_workspace_policy=_safe_workspace_policy(
+                agent_pool_data["allowed_workspace_policy"]
+            ),
+            agent_count=_safe_int(agent_pool_data["agent_count"]),
+        )
diff --git a/tests/units/test_agent_pools.py b/tests/units/test_agent_pools.py
