diff --git a/.nx/version-plans/version-plan-1763465747542.md b/.nx/version-plans/version-plan-1763465747542.md new file mode 100644 index 000000000..aa196e181 --- /dev/null +++ b/.nx/version-plans/version-plan-1763465747542.md @@ -0,0 +1,16 @@ +--- +fleet-mcp: minor +--- + +feat: Add workspace update & restart functionality + +Implements programmatic workspace update and restart operations for Coder workspaces using the two-step workflow required by Coder's REST API. + +**New Features:** +- Add `update_workspace()` method to CoderClient for two-step update workflow +- Add `update_agent` MCP tool for fleet management +- Add `scripts/update_workspace.py` standalone CLI tool +- Support both explicit and automatic template version selection +- Comprehensive test coverage with 9 new tests + +**Related Issue:** https://github.com/coder/coder/issues/19331 diff --git a/libs/fleet-mcp/scripts/update_workspace.py b/libs/fleet-mcp/scripts/update_workspace.py new file mode 100755 index 000000000..873ea7858 --- /dev/null +++ b/libs/fleet-mcp/scripts/update_workspace.py @@ -0,0 +1,301 @@ +#!/usr/bin/env python3 +r"""Standalone CLI script for updating Coder workspaces to new template versions. + +This script provides a command-line interface for the two-step workspace update workflow +required by Coder's REST API. It can be used independently of the fleet-mcp server. + +Usage: + # Update workspace to specific template version + python update_workspace.py --workspace-id abc-123 --template-version-id def-456 + + # Update workspace by name to latest active version + python update_workspace.py --workspace-name my-workspace + + # Update with custom timeout settings + python update_workspace.py --workspace-id abc-123 --template-version-id def-456 \\ + --stop-timeout 120 --start-timeout 240 + +Environment Variables: + CODER_URL: Base URL for Coder API (required) + CODER_SESSION_TOKEN: Authentication token (required) + +Example: + export CODER_URL="https://coder.example.com" + export CODER_SESSION_TOKEN="your-token-here" + python update_workspace.py --workspace-name my-agent +""" + +import argparse +import asyncio +import logging +import os +import sys +from typing import Optional + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) +logger = logging.getLogger(__name__) + + +async def update_workspace( + workspace_id: str, + template_version_id: Optional[str] = None, + stop_timeout: int = 60, + start_timeout: int = 120, + coder_url: Optional[str] = None, + token: Optional[str] = None, +) -> dict: + """Update a workspace to a new template version. + + Args: + workspace_id: Workspace UUID + template_version_id: Template version UUID (if None, uses active version) + stop_timeout: Maximum seconds to wait for stop completion + start_timeout: Maximum seconds to wait for start completion + coder_url: Coder API base URL (defaults to CODER_URL env var) + token: Coder session token (defaults to CODER_SESSION_TOKEN env var) + + Returns: + Updated workspace data + + Raises: + ValueError: If required credentials are missing + Exception: If workspace update fails + """ + # Import CoderClient here to avoid circular imports + from fleet_mcp.clients import CoderClient + + # Validate credentials + coder_url = coder_url or os.getenv("CODER_URL") + token = token or os.getenv("CODER_SESSION_TOKEN") + + if not coder_url or not token: + raise ValueError( + "CODER_URL and CODER_SESSION_TOKEN must be set via environment variables or parameters" + ) + + logger.info(f"Connecting to Coder API at {coder_url}") + + async with CoderClient(base_url=coder_url, token=token) as client: + # If no template_version_id provided, get the active version + if template_version_id is None: + logger.info( + f"No template version specified, fetching active version for workspace {workspace_id}" + ) + workspace = await client.get_workspace(workspace_id) + template_id = workspace.get("template_id") + + if not template_id: + raise ValueError( + f"Workspace {workspace_id} has no template_id in workspace data" + ) + + template = await client.get_template(template_id) + template_version_id = template.get("active_version_id") + + if not template_version_id: + raise ValueError( + f"Template {template_id} has no active version available" + ) + + logger.info(f"Using active template version: {template_version_id}") + + # Perform the update + logger.info( + f"Updating workspace {workspace_id} to template version {template_version_id}" + ) + logger.info( + f"Timeouts: stop={stop_timeout}s, start={start_timeout}s (1s poll interval)" + ) + + workspace = await client.update_workspace( + workspace_id=workspace_id, + template_version_id=template_version_id, + max_stop_attempts=stop_timeout, + max_start_attempts=start_timeout, + ) + + logger.info(f"Successfully updated workspace {workspace_id}") + return workspace + + +async def update_workspace_by_name( + workspace_name: str, + template_version_id: Optional[str] = None, + stop_timeout: int = 60, + start_timeout: int = 120, + coder_url: Optional[str] = None, + token: Optional[str] = None, +) -> dict: + """Update a workspace by name to a new template version. + + Args: + workspace_name: Workspace name + template_version_id: Template version UUID (if None, uses active version) + stop_timeout: Maximum seconds to wait for stop completion + start_timeout: Maximum seconds to wait for start completion + coder_url: Coder API base URL (defaults to CODER_URL env var) + token: Coder session token (defaults to CODER_SESSION_TOKEN env var) + + Returns: + Updated workspace data + + Raises: + ValueError: If required credentials are missing or workspace not found + Exception: If workspace update fails + """ + # Import CoderClient here to avoid circular imports + from fleet_mcp.clients import CoderClient + + # Validate credentials + coder_url = coder_url or os.getenv("CODER_URL") + token = token or os.getenv("CODER_SESSION_TOKEN") + + if not coder_url or not token: + raise ValueError( + "CODER_URL and CODER_SESSION_TOKEN must be set via environment variables or parameters" + ) + + logger.info(f"Looking up workspace by name: {workspace_name}") + + async with CoderClient(base_url=coder_url, token=token) as client: + # Find workspace by name + workspaces = await client.list_workspaces(owner="me") + workspace = None + for ws in workspaces: + if ws.get("name", "").lower() == workspace_name.lower(): + workspace = ws + break + + if not workspace: + raise ValueError(f"Workspace '{workspace_name}' not found") + + workspace_id = workspace.get("id") + logger.info(f"Found workspace {workspace_name} with ID: {workspace_id}") + + # Delegate to update_workspace + return await update_workspace( + workspace_id=workspace_id, + template_version_id=template_version_id, + stop_timeout=stop_timeout, + start_timeout=start_timeout, + coder_url=coder_url, + token=token, + ) + + +def main(): + """Main CLI entry point.""" + parser = argparse.ArgumentParser( + description="Update a Coder workspace to a new template version", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + + # Workspace identification (mutually exclusive) + id_group = parser.add_mutually_exclusive_group(required=True) + id_group.add_argument( + "--workspace-id", + help="Workspace UUID to update", + ) + id_group.add_argument( + "--workspace-name", + help="Workspace name to update (case-insensitive)", + ) + + # Template version + parser.add_argument( + "--template-version-id", + help="Template version UUID to update to (if not specified, uses active version)", + ) + + # Timeout settings + parser.add_argument( + "--stop-timeout", + type=int, + default=60, + help="Maximum seconds to wait for stop completion (default: 60)", + ) + parser.add_argument( + "--start-timeout", + type=int, + default=120, + help="Maximum seconds to wait for start completion (default: 120)", + ) + + # Credentials (optional, can use env vars) + parser.add_argument( + "--coder-url", + help="Coder API base URL (default: CODER_URL env var)", + ) + parser.add_argument( + "--token", + help="Coder session token (default: CODER_SESSION_TOKEN env var)", + ) + + # Logging level + parser.add_argument( + "--log-level", + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + default="INFO", + help="Logging level (default: INFO)", + ) + + args = parser.parse_args() + + # Set logging level + logging.getLogger().setLevel(getattr(logging, args.log_level)) + + try: + # Run the update + if args.workspace_id: + result = asyncio.run( + update_workspace( + workspace_id=args.workspace_id, + template_version_id=args.template_version_id, + stop_timeout=args.stop_timeout, + start_timeout=args.start_timeout, + coder_url=args.coder_url, + token=args.token, + ) + ) + else: + result = asyncio.run( + update_workspace_by_name( + workspace_name=args.workspace_name, + template_version_id=args.template_version_id, + stop_timeout=args.stop_timeout, + start_timeout=args.start_timeout, + coder_url=args.coder_url, + token=args.token, + ) + ) + + logger.info("Update completed successfully") + logger.info(f"Workspace name: {result.get('name')}") + logger.info(f"Workspace ID: {result.get('id')}") + logger.info(f"Template: {result.get('template_display_name')}") + logger.info( + f"Latest build status: {result.get('latest_build', {}).get('status')}" + ) + + return 0 + + except KeyboardInterrupt: + logger.error("Operation cancelled by user") + return 130 + except Exception as e: + logger.error(f"Failed to update workspace: {e}") + if args.log_level == "DEBUG": + import traceback + + traceback.print_exc() + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/libs/fleet-mcp/src/fleet_mcp/__main__.py b/libs/fleet-mcp/src/fleet_mcp/__main__.py index 19f75aa8e..fd9145add 100644 --- a/libs/fleet-mcp/src/fleet_mcp/__main__.py +++ b/libs/fleet-mcp/src/fleet_mcp/__main__.py @@ -258,6 +258,31 @@ async def restart_agent( return result.model_dump() +@mcp.tool() +async def update_agent( + agent_name: Annotated[ + str, + Field(min_length=1, max_length=32, description="Name of the agent to update"), + ], + template_version_id: Annotated[ + str | None, + Field( + None, + description="Template version UUID to update to. If not provided, uses the active version of the agent's template.", + ), + ] = None, +) -> dict: + """Update an agent's workspace to a new template version.""" + from .tools.update_agent import update_agent as update_agent_impl + + result = await update_agent_impl( + get_agent_service(), + agent_name=agent_name, + template_version_id=template_version_id, + ) + return result.model_dump() + + # ======================================================================== # User Story 3: Task Assignment and Cancellation Tools # ======================================================================== diff --git a/libs/fleet-mcp/src/fleet_mcp/clients/coder_client.py b/libs/fleet-mcp/src/fleet_mcp/clients/coder_client.py index 863e9ee7c..06001e7bf 100644 --- a/libs/fleet-mcp/src/fleet_mcp/clients/coder_client.py +++ b/libs/fleet-mcp/src/fleet_mcp/clients/coder_client.py @@ -232,13 +232,86 @@ async def restart_workspace(self, workspace_id: str) -> dict[str, Any]: except httpx.RequestError as e: raise HTTPError(f"Failed to connect to Coder API: {e}") from e + async def update_workspace( + self, + workspace_id: str, + template_version_id: str, + max_stop_attempts: int = 60, + max_start_attempts: int = 120, + ) -> dict[str, Any]: + """Update a workspace to a new template version by stopping and restarting. + + This implements the two-step workflow required by Coder API: + 1. Stop the workspace and wait for completion + 2. Start with the new template version + + Note: Coder does not provide a single "update & restart" endpoint. + See: https://github.com/coder/coder/issues/19331 + + Args: + workspace_id: Workspace UUID + template_version_id: New template version UUID to apply + max_stop_attempts: Maximum polling attempts for stop completion (default: 60) + max_start_attempts: Maximum polling attempts for start completion (default: 120) + + Returns: + Started workspace data from Coder API with updated template version + + Raises: + NotFoundError: If workspace doesn't exist + HTTPError: If API request fails or build doesn't complete in time + """ + try: + # Step 1: Stop the workspace + stop_response = await self.client.post( + f"{self.base_url}/api/v2/workspaces/{workspace_id}/builds", + json={"transition": "stop"}, + ) + stop_response.raise_for_status() + stop_build_data = stop_response.json() + stop_build_id = stop_build_data.get("id") + + # Step 2: Wait for the stop build to complete + if stop_build_id: + await self._wait_for_build_completion( + stop_build_id, max_attempts=max_stop_attempts + ) + + # Step 3: Start with new template version + start_response = await self.client.post( + f"{self.base_url}/api/v2/workspaces/{workspace_id}/builds", + json={ + "transition": "start", + "template_version_id": template_version_id, + }, + ) + start_response.raise_for_status() + start_build_data = start_response.json() + start_build_id = start_build_data.get("id") + + # Step 4: Wait for the start build to complete + if start_build_id: + await self._wait_for_build_completion( + start_build_id, max_attempts=max_start_attempts + ) + + # Step 5: Return the updated workspace data + return await self.get_workspace(workspace_id) + except httpx.HTTPStatusError as e: + if e.response.status_code == 404: + raise NotFoundError(f"Workspace {workspace_id} not found") from e + self._handle_http_error(e) + except httpx.RequestError as e: + raise HTTPError(f"Failed to connect to Coder API: {e}") from e + async def _wait_for_build_completion( self, build_id: str, max_attempts: int = 60, delay_seconds: float = 1.0 ) -> None: """Wait for a workspace build to complete. - Polls the build status until it reaches a terminal state (stopped, failed, canceled) - or the maximum number of attempts is reached. + Polls the build status until it reaches a completion state: + - For stop/delete builds: "stopped", "failed", "canceled", "deleted" + - For start builds: "running", "failed", "canceled" Args: build_id: Workspace build UUID @@ -250,7 +323,8 @@ async def _wait_for_build_completion( """ import asyncio - terminal_states = {"stopped", "failed", "canceled"} + # States that indicate build completion (any transition type) + completion_states = {"stopped", "running", "failed", "canceled", "deleted"} for _attempt in range(max_attempts): try: @@ -261,7 +335,7 @@ async def _wait_for_build_completion( build_data = response.json() status = build_data.get("status", "").lower() - if status in terminal_states: + if status in completion_states: return # Build completed # Wait before next poll @@ -446,6 +520,32 @@ async def list_workspace_presets(self, template_id: str) -> list[dict[str, Any]] except httpx.RequestError as e: raise HTTPError(f"Failed to connect to Coder API: {e}") from e + async def get_template_version(self, version_id: str) -> dict[str, Any]: + """Get template version details by ID. + + Args: + version_id: Template version UUID + + Returns: + Template version dictionary from Coder API + + Raises: + NotFoundError: If template version doesn't exist + HTTPError: If API request fails + """ + try: + response = await self.client.get( + f"{self.base_url}/api/v2/templateversions/{version_id}" + ) + response.raise_for_status() + return response.json() + except httpx.HTTPStatusError as e: + if e.response.status_code == 404: + raise NotFoundError(f"Template version {version_id} not found") from e + self._handle_http_error(e) + except httpx.RequestError as e: + raise HTTPError(f"Failed to connect to Coder API: {e}") from e + # ======================================================================== # Task Operations (via AgentAPI and Experimental Task API) # ======================================================================== diff --git a/libs/fleet-mcp/src/fleet_mcp/models/__init__.py b/libs/fleet-mcp/src/fleet_mcp/models/__init__.py index d440c07b4..eb8c1a418 100644 --- a/libs/fleet-mcp/src/fleet_mcp/models/__init__.py +++ b/libs/fleet-mcp/src/fleet_mcp/models/__init__.py @@ -30,6 +30,7 @@ ShowLogsResponse, ShowTaskHistoryResponse, StartTaskResponse, + UpdateAgentResponse, ) from .task import ConversationLog, LogEntry, Task, TaskHistory @@ -56,6 +57,7 @@ "CreateAgentResponse", "DeleteAgentResponse", "RestartAgentResponse", + "UpdateAgentResponse", "StartTaskResponse", "CancelTaskResponse", "ShowTaskHistoryResponse", diff --git a/libs/fleet-mcp/src/fleet_mcp/models/responses.py b/libs/fleet-mcp/src/fleet_mcp/models/responses.py index f7de47b5f..3a4ad8e75 100644 --- a/libs/fleet-mcp/src/fleet_mcp/models/responses.py +++ b/libs/fleet-mcp/src/fleet_mcp/models/responses.py @@ -110,6 +110,13 @@ class RestartAgentResponse(BaseModel): message: str +class UpdateAgentResponse(BaseModel): + """Response for update_agent tool.""" + + agent: Agent + message: str + + class StartTaskResponse(BaseModel): """Response for start_agent_task tool.""" diff --git a/libs/fleet-mcp/src/fleet_mcp/repositories/agent_repository.py b/libs/fleet-mcp/src/fleet_mcp/repositories/agent_repository.py index 88aa4922b..b365eb1fd 100644 --- a/libs/fleet-mcp/src/fleet_mcp/repositories/agent_repository.py +++ b/libs/fleet-mcp/src/fleet_mcp/repositories/agent_repository.py @@ -168,6 +168,62 @@ async def restart(self, agent_name: str) -> Agent: except Exception as e: raise CoderAPIError(f"Failed to restart agent '{agent_name}': {e}") from e + async def update( + self, agent_name: str, template_version_id: str | None = None + ) -> Agent: + """Update an agent to a new template version and restart. + + This method implements the two-step workflow required by Coder API: + 1. Stop the workspace and wait for completion + 2. Start with the new template version + + Args: + agent_name: Agent name (workspace name) + template_version_id: Template version UUID to update to. + If None, uses the active version of the agent's template. + + Returns: + Updated Agent domain model after update and restart + + Raises: + AgentNotFoundError: If agent doesn't exist + CoderAPIError: If workspace update fails + """ + try: + # Find workspace by name + agent = await self.get_by_name(agent_name) + + # If no template_version_id provided, get the active version + if template_version_id is None: + # Get the workspace to find its template_id + workspace = await self.client.get_workspace(agent.workspace_id) + template_id = workspace.get("template_id") + + if not template_id: + raise CoderAPIError( + f"Agent '{agent_name}' has no template_id in workspace data" + ) + + # Get the template's active version + template = await self.client.get_template(template_id) + template_version_id = template.get("active_version_id") + + if not template_version_id: + raise CoderAPIError( + f"Template {template_id} has no active version available" + ) + + # Update the workspace with the new template version + workspace = await self.client.update_workspace( + agent.workspace_id, template_version_id + ) + + return await self._workspace_to_agent(workspace) + except AgentNotFoundError: + raise + except Exception as e: + raise CoderAPIError(f"Failed to update agent '{agent_name}': {e}") from e + def _are_all_workspace_agents_healthy(self, workspace: dict[str, Any]) -> bool: """Check if all agents within the workspace are healthy and ready. diff --git a/libs/fleet-mcp/src/fleet_mcp/services/agent_service.py b/libs/fleet-mcp/src/fleet_mcp/services/agent_service.py index 283b12596..98d50aaee 100644 --- a/libs/fleet-mcp/src/fleet_mcp/services/agent_service.py +++ b/libs/fleet-mcp/src/fleet_mcp/services/agent_service.py @@ -244,6 +244,42 @@ async def restart_agent(self, name: str) -> Agent: normalized_name = name.lower() return await self.agent_repo.restart(normalized_name) + async def update_agent( + self, name: str, template_version_id: str | None = None + ) -> Agent: + """Update an agent to a new template version (case-insensitive). + + This updates the agent's workspace to use a new template version by: + 1. Stopping the workspace + 2. Starting it with the new template version + + Args: + name: Agent name (case-insensitive) + template_version_id: Template version UUID to update to. + If None, uses the active version of the agent's template. + + Returns: + Updated Agent domain model + + Raises: + AgentNotFoundError: If agent doesn't exist + ValidationError: If name or template_version_id is invalid + CoderAPIError: If update operation fails + """ + # Validate agent name format + self._validate_agent_name(name) + + # Validate template_version_id if provided + if template_version_id is not None: + if not template_version_id or not template_version_id.strip(): + raise ValidationError( + "template_version_id", "Template version ID cannot be empty" + ) + + # Normalize to lowercase for case-insensitive lookup + normalized_name = name.lower() + return await self.agent_repo.update(normalized_name, template_version_id) + def _validate_agent_name(self, name: str) -> None: """Validate agent name format according to Coder workspace constraints. diff --git a/libs/fleet-mcp/src/fleet_mcp/tools/update_agent.py b/libs/fleet-mcp/src/fleet_mcp/tools/update_agent.py new file mode 100644 index 000000000..667f1ec3a --- /dev/null +++ b/libs/fleet-mcp/src/fleet_mcp/tools/update_agent.py @@ -0,0 +1,75 @@ +"""MCP tool for updating an agent to a new template version.""" + +from typing import Annotated + +from pydantic import Field + +from ..models import UpdateAgentResponse +from ..services import AgentService + + +async def update_agent( + agent_service: AgentService, + agent_name: Annotated[ + str, + Field( + min_length=1, + max_length=32, + description="Name of the agent to update", + ), + ], + template_version_id: Annotated[ + str | None, + Field( + default=None, + description="Template version UUID to update to. If not provided, uses the active version of the agent's template.", + ), + ] = None, +) -> UpdateAgentResponse: + """Update an agent's workspace to a new template version. + + Tool: update_agent + User Story: US2 (Agent Lifecycle Management) + Architecture: Layer 1 (Tool Layer) + + This updates the agent's workspace by stopping it and restarting with + a new template version. This is useful for applying template updates, + configuration changes, or environment updates. + + The update process follows Coder's two-step workflow: + 1. Stop the workspace and wait for completion + 2. Start with the new template version + + Args: + agent_service: Service instance for agent business logic + agent_name: Name of the agent to update + template_version_id: Template version UUID to update to. + If None, uses the active version of the agent's template. + + Returns: + UpdateAgentResponse with updated agent and success message + + Raises: + AgentNotFoundError: If agent doesn't exist + ValidationError: If agent_name or template_version_id is invalid + CoderAPIError: If update fails + + Example: + # Update to specific version + result = await update_agent( + agent_service, + agent_name="my-agent", + template_version_id="abc-123-def-456" + ) + + # Update to latest active version + result = await update_agent(agent_service, agent_name="my-agent") + """ + agent = await agent_service.update_agent(agent_name, template_version_id) + + message = ( + f"Agent '{agent_name}' updated to template version " + f"{template_version_id if template_version_id else 'latest active version'}" + ) + + return UpdateAgentResponse(agent=agent, message=message) diff --git a/libs/fleet-mcp/tests/clients/test_coder_client_update.py b/libs/fleet-mcp/tests/clients/test_coder_client_update.py new file mode 100644 index 000000000..1f3bf9a15 --- /dev/null +++ b/libs/fleet-mcp/tests/clients/test_coder_client_update.py @@ -0,0 +1,216 @@ +"""Tests for CoderClient update_workspace method.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fleet_mcp.clients import CoderClient +from fleet_mcp.clients.exceptions import HTTPError, NotFoundError + + +@pytest.fixture +def mock_httpx_client(): + """Create a mock httpx.AsyncClient.""" + client = MagicMock() + client.aclose = AsyncMock() + return client + + +@pytest.fixture +def coder_client(mock_httpx_client): + """Create a CoderClient with mocked httpx client.""" + with patch("fleet_mcp.clients.coder_client.httpx.AsyncClient") as mock_client_class: + mock_client_class.return_value = mock_httpx_client + client = CoderClient(base_url="https://coder.example.com", token="test-token") + yield client + + +@pytest.mark.asyncio +async def test_update_workspace_success(coder_client, mock_httpx_client): + """Test successful workspace update to new template version.""" + workspace_id = "ws-123" + template_version_id = "ver-456" + + # Mock stop response + stop_response = MagicMock() + stop_response.raise_for_status = MagicMock() + stop_response.json.return_value = {"id": "build-stop-123", "status": "stopping"} + + # Mock stop build status (completed) + stop_status_response = MagicMock() + stop_status_response.raise_for_status = MagicMock() + stop_status_response.json.return_value = { + "id": "build-stop-123", + "status": "stopped", + } + + # Mock start response + start_response = MagicMock() + start_response.raise_for_status = MagicMock() + start_response.json.return_value = {"id": "build-start-456", "status": "starting"} + + # Mock start build status (completed) + start_status_response = MagicMock() + start_status_response.raise_for_status = MagicMock() + start_status_response.json.return_value = { + "id": "build-start-456", + "status": "running", + } + + # Mock final workspace data + final_workspace_response = MagicMock() + final_workspace_response.raise_for_status = MagicMock() + final_workspace_response.json.return_value = { + "id": workspace_id, + "name": "test-workspace", + "template_id": "tpl-123", + "latest_build": { + "id": "build-start-456", + "status": "running", + "template_version_id": template_version_id, + }, + } + + # Set up mock responses in order + mock_httpx_client.post = AsyncMock(side_effect=[stop_response, start_response]) + mock_httpx_client.get = AsyncMock( + side_effect=[ + stop_status_response, # Check stop build status + start_status_response, # Check start build status (running = completed) + final_workspace_response, # Get final workspace data + ] + ) + + # Execute update + result = await coder_client.update_workspace(workspace_id, template_version_id) + + # Verify the result + assert result["id"] == workspace_id + assert result["name"] == "test-workspace" + assert result["latest_build"]["template_version_id"] == template_version_id + + # Verify API calls + assert mock_httpx_client.post.call_count == 2 + assert mock_httpx_client.get.call_count == 3 + + # Verify stop build call + stop_call = mock_httpx_client.post.call_args_list[0] + assert f"/workspaces/{workspace_id}/builds" in stop_call[0][0] + assert stop_call[1]["json"] == {"transition": "stop"} + + # Verify start build call with template version + start_call = mock_httpx_client.post.call_args_list[1] + assert f"/workspaces/{workspace_id}/builds" in start_call[0][0] + assert start_call[1]["json"] == { + "transition": "start", + "template_version_id": template_version_id, + } + + +@pytest.mark.asyncio +async def test_update_workspace_not_found(coder_client, mock_httpx_client): + """Test update_workspace with non-existent workspace.""" + import httpx + + workspace_id = "ws-nonexistent" + template_version_id = "ver-456" + + # Mock 404 response + error_response = MagicMock() + error_response.status_code = 404 + error_response.json.return_value = {"message": "Workspace not found"} + + stop_response = MagicMock() + stop_response.raise_for_status.side_effect = httpx.HTTPStatusError( + "404", request=MagicMock(), response=error_response + ) + stop_response.status_code = 404 + + mock_httpx_client.post = AsyncMock(return_value=stop_response) + + # Execute and expect NotFoundError + with pytest.raises(NotFoundError, match="not found"): + await coder_client.update_workspace(workspace_id, template_version_id) + + +@pytest.mark.asyncio +async def test_update_workspace_stop_timeout(coder_client, mock_httpx_client): + """Test update_workspace with stop build timeout.""" + workspace_id = "ws-123" + template_version_id = "ver-456" + + # Mock stop response + stop_response = MagicMock() + stop_response.raise_for_status = MagicMock() + stop_response.json.return_value = {"id": "build-stop-123", "status": "stopping"} + + # Mock stop build status (never completes) + stop_status_response = MagicMock() + stop_status_response.raise_for_status = MagicMock() + stop_status_response.json.return_value = { + "id": "build-stop-123", + "status": "stopping", + } + + mock_httpx_client.post = AsyncMock(return_value=stop_response) + mock_httpx_client.get = AsyncMock(return_value=stop_status_response) + + # Execute with very short timeout and expect HTTPError + with pytest.raises(HTTPError, match="did not complete"): + await coder_client.update_workspace( + workspace_id, template_version_id, max_stop_attempts=2 + ) + + +@pytest.mark.asyncio +async def test_get_template_version(coder_client, mock_httpx_client): + """Test get_template_version method.""" + version_id = "ver-123" + + # Mock response + response = MagicMock() + response.raise_for_status = MagicMock() + response.json.return_value = { + "id": version_id, + "name": "v1.2.3", + "template_id": "tpl-456", + "created_at": "2025-01-01T00:00:00Z", + } + + mock_httpx_client.get = AsyncMock(return_value=response) + + # Execute + result = await coder_client.get_template_version(version_id) + + # Verify + assert result["id"] == version_id + assert result["name"] == "v1.2.3" + assert result["template_id"] == "tpl-456" + + # Verify API call + mock_httpx_client.get.assert_called_once() + call_args = mock_httpx_client.get.call_args + assert f"/templateversions/{version_id}" in call_args[0][0] + + +@pytest.mark.asyncio +async def test_get_template_version_not_found(coder_client, mock_httpx_client): + """Test get_template_version with non-existent version.""" + import httpx + + version_id = "ver-nonexistent" + + # Mock 404 response + error_response = MagicMock() + error_response.status_code = 404 + + response = MagicMock() + response.raise_for_status.side_effect = httpx.HTTPStatusError( + "404", request=MagicMock(), response=error_response + ) + response.status_code = 404 + + mock_httpx_client.get = AsyncMock(return_value=response) + + # Execute and expect NotFoundError + with pytest.raises(NotFoundError, match="not found"): + await coder_client.get_template_version(version_id) diff --git a/libs/fleet-mcp/tests/repositories/test_agent_repository_update.py b/libs/fleet-mcp/tests/repositories/test_agent_repository_update.py new file mode 100644 index 000000000..0ee46745e --- /dev/null +++ b/libs/fleet-mcp/tests/repositories/test_agent_repository_update.py @@ -0,0 +1,272 @@ +"""Tests for AgentRepository update method.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fleet_mcp.models import Agent, AgentStatus +from fleet_mcp.models.errors import AgentNotFoundError, CoderAPIError +from fleet_mcp.repositories import AgentRepository + + +@pytest.fixture +def mock_coder_client(): + """Create a mock CoderClient.""" + client = MagicMock() + client.list_workspaces = AsyncMock() + client.get_workspace = AsyncMock() + client.get_template = AsyncMock() + client.update_workspace = AsyncMock() + client.list_workspace_presets = AsyncMock() + return client + + +@pytest.fixture +def agent_repository(mock_coder_client): + """Create an AgentRepository with mocked client.""" + return AgentRepository(mock_coder_client) + + +@pytest.mark.asyncio +async def test_update_agent_with_version_id(agent_repository, mock_coder_client): + """Test updating agent with explicit template version ID.""" + agent_name = "test-agent" + workspace_id = "ws-123" + template_version_id = "ver-456" + + # Mock workspace lookup + mock_coder_client.list_workspaces.return_value = [ + { + "id": workspace_id, + "name": agent_name, + "template_id": "tpl-789", + "template_display_name": "Setup", + "created_at": "2025-01-01T00:00:00Z", + "updated_at": "2025-01-01T00:00:00Z", + } + ] + + # Mock full workspace details + mock_coder_client.get_workspace.return_value = { + "id": workspace_id, + "name": agent_name, + "template_id": "tpl-789", + "template_display_name": "Setup", + "created_at": "2025-01-01T00:00:00Z", + "updated_at": "2025-01-01T00:00:00Z", + "latest_build": { + "id": "build-123", + "status": "running", + "template_version_id": template_version_id, + "template_version_preset_id": "preset-111", + "resources": [ + { + "agents": [ + { + "status": "connected", + "lifecycle_state": "ready", + "apps": [], + } + ] + } + ], + }, + } + + # Mock workspace update + mock_coder_client.update_workspace.return_value = { + "id": workspace_id, + "name": agent_name, + "template_id": "tpl-789", + "template_display_name": "Setup", + "created_at": "2025-01-01T00:00:00Z", + "updated_at": "2025-01-01T01:00:00Z", + "latest_build": { + "id": "build-456", + "status": "running", + "template_version_id": template_version_id, + "template_version_preset_id": "preset-111", + "resources": [ + { + "agents": [ + { + "status": "connected", + "lifecycle_state": "ready", + "apps": [], + } + ] + } + ], + }, + } + + # Mock presets + mock_coder_client.list_workspace_presets.return_value = [ + {"ID": "preset-111", "Name": "coder"} + ] + + # Execute update + result = await agent_repository.update(agent_name, template_version_id) + + # Verify result + assert isinstance(result, Agent) + assert result.name == agent_name + assert result.workspace_id == workspace_id + assert result.status == AgentStatus.IDLE + + # Verify API calls + mock_coder_client.list_workspaces.assert_called_once() + mock_coder_client.update_workspace.assert_called_once_with( + workspace_id, template_version_id + ) + + +@pytest.mark.asyncio +async def test_update_agent_without_version_id(agent_repository, mock_coder_client): + """Test updating agent to active template version.""" + agent_name = "test-agent" + workspace_id = "ws-123" + template_id = "tpl-789" + active_version_id = "ver-active-999" + + # Mock workspace lookup + mock_coder_client.list_workspaces.return_value = [ + { + "id": workspace_id, + "name": agent_name, + "template_id": template_id, + "template_display_name": "Setup", + "created_at": "2025-01-01T00:00:00Z", + "updated_at": "2025-01-01T00:00:00Z", + } + ] + + # Mock workspace details (for getting template_id) + mock_coder_client.get_workspace.return_value = { + "id": workspace_id, + "name": agent_name, + "template_id": template_id, + "template_display_name": "Setup", + "created_at": "2025-01-01T00:00:00Z", + "updated_at": "2025-01-01T00:00:00Z", + "latest_build": { + "id": "build-123", + "status": "running", + "template_version_id": "ver-old-888", + "template_version_preset_id": "preset-111", + "resources": [ + { + "agents": [ + { + "status": "connected", + "lifecycle_state": "ready", + "apps": [], + } + ] + } + ], + }, + } + + # Mock template with active version + mock_coder_client.get_template.return_value = { + "id": template_id, + "name": "Setup", + "active_version_id": active_version_id, + } + + # Mock workspace update result + mock_coder_client.update_workspace.return_value = { + "id": workspace_id, + "name": agent_name, + "template_id": template_id, + "template_display_name": "Setup", + "created_at": "2025-01-01T00:00:00Z", + "updated_at": "2025-01-01T02:00:00Z", + "latest_build": { + "id": "build-789", + "status": "running", + "template_version_id": active_version_id, + "template_version_preset_id": "preset-111", + "resources": [ + { + "agents": [ + { + "status": "connected", + "lifecycle_state": "ready", + "apps": [], + } + ] + } + ], + }, + } + + # Mock presets + mock_coder_client.list_workspace_presets.return_value = [ + {"ID": "preset-111", "Name": "coder"} + ] + + # Execute update without version ID + result = await agent_repository.update(agent_name, None) + + # Verify result + assert isinstance(result, Agent) + assert result.name == agent_name + assert result.workspace_id == workspace_id + + # Verify API calls + mock_coder_client.get_workspace.assert_called() + mock_coder_client.get_template.assert_called_once_with(template_id) + mock_coder_client.update_workspace.assert_called_once_with( + workspace_id, active_version_id + ) + + +@pytest.mark.asyncio +async def test_update_agent_not_found(agent_repository, mock_coder_client): + """Test updating non-existent agent.""" + agent_name = "nonexistent-agent" + + # Mock empty workspace list + mock_coder_client.list_workspaces.return_value = [] + + # Execute and expect AgentNotFoundError + with pytest.raises(AgentNotFoundError, match="not found"): + await agent_repository.update(agent_name, "ver-123") + + +@pytest.mark.asyncio +async def test_update_agent_no_active_version(agent_repository, mock_coder_client): + """Test updating agent when template has no active version.""" + agent_name = "test-agent" + workspace_id = "ws-123" + template_id = "tpl-789" + + # Mock workspace lookup + mock_coder_client.list_workspaces.return_value = [ + { + "id": workspace_id, + "name": agent_name, + "template_id": template_id, + "template_display_name": "Setup", + } + ] + + # Mock workspace details + mock_coder_client.get_workspace.return_value = { + "id": workspace_id, + "name": agent_name, + "template_id": template_id, + "latest_build": {"status": "running"}, + } + + # Mock template with no active version + mock_coder_client.get_template.return_value = { + "id": template_id, + "name": "Setup", + "active_version_id": None, # No active version + } + + # Execute and expect CoderAPIError + with pytest.raises(CoderAPIError, match="no active version"): + await agent_repository.update(agent_name, None)