feat(fleet-mcp): Add workspace update & restart functionality

[mvgijssel]

Overview

This PR implements programmatic workspace update and restart operations for Coder workspaces, providing a robust solution for the two-step workflow required by Coder's REST API.

Context

Coder does not provide a single "update & restart" endpoint. Instead, updates require:

1. Stop the workspace and wait for completion
2. Start with the new template version

This implementation provides a clean abstraction over this workflow with proper error handling, polling, and timeout management.

Related Issue: coder/coder#19331

Implementation

Core Components

1. CoderClient Layer (coder_client.py)

• update_workspace() - Orchestrates the two-step update workflow
  • Stops workspace with configurable timeout (default: 60 attempts @ 1s)
  • Starts with new template version with configurable timeout (default: 120 attempts @ 1s)
  • Returns updated workspace data
• get_template_version() - Retrieves template version details by ID
• Enhanced _wait_for_build_completion() - Smart polling for build operations
  • Handles completion states: "stopped", "running", "failed", "canceled", "deleted"
  • Graceful handling of 404 responses (build cleanup)

2. AgentRepository Layer (agent_repository.py)

• update() - Updates agent workspace to new template version
  • Supports explicit template version ID
  • Auto-detects active template version when ID is None
  • Validates template has active version available

3. AgentService Layer (agent_service.py)

• update_agent() - Business logic with validation
  • Case-insensitive agent name handling
  • Template version ID validation
  • Error handling for missing agents/templates

4. MCP Tool (update_agent.py)

• Exposed as update_agent tool via FastMCP
• Integrated into __main__.py with proper routing
• Returns UpdateAgentResponse with agent state and success message

5. Standalone CLI (scripts/update_workspace.py)

• Command-line interface for direct usage
• Features:
  • Lookup by workspace ID or name (case-insensitive)
  • Configurable stop/start timeouts
  • Environment variable configuration (CODER_URL, CODER_SESSION_TOKEN)
  • Comprehensive logging with levels (DEBUG, INFO, WARNING, ERROR)
  • Graceful error handling and keyboard interrupt support

New Models

• UpdateAgentResponse - Response model for update operations

Testing

Test Coverage

1. test_coder_client_update.py (5 tests)

   • ✅ Successful workspace update with polling
   • ✅ Workspace not found error handling
   • ✅ Stop build timeout handling
   • ✅ Template version retrieval
   • ✅ Template version not found error

2. test_agent_repository_update.py (4 tests)

   • ✅ Update with explicit template version ID
   • ✅ Update with automatic active version detection
   • ✅ Agent not found error handling
   • ✅ Template with no active version error

Test Results:

test_coder_client_update.py ............ 5 passed
test_agent_repository_update.py ........ 4 passed
All tests passing ✅

Usage Examples

Via MCP Tool

# Update to specific version
await update_agent(
    agent_service,
    agent_name="my-agent",
    template_version_id="abc-123-def-456"
)

# Update to latest active version
await update_agent(agent_service, agent_name="my-agent")

Via Standalone CLI

# Update workspace by name to latest active version
export CODER_URL="https://coder.example.com"
export CODER_SESSION_TOKEN="your-token"
./scripts/update_workspace.py --workspace-name my-agent

# Update by ID to specific version
./scripts/update_workspace.py \
    --workspace-id ws-abc-123 \
    --template-version-id ver-def-456

# With custom timeouts
./scripts/update_workspace.py \
    --workspace-name my-agent \
    --stop-timeout 120 \
    --start-timeout 240 \
    --log-level DEBUG

Code Quality

• ✅ All files formatted with trunk fmt
• ✅ All linting checks passed with trunk check
• ✅ Follows existing codebase patterns and architecture
• ✅ Comprehensive docstrings with type hints
• ✅ Error handling and logging throughout

Files Changed

Modified

• src/fleet_mcp/__main__.py - Register update_agent MCP tool
• src/fleet_mcp/clients/coder_client.py - Add update methods
• src/fleet_mcp/models/__init__.py - Export UpdateAgentResponse
• src/fleet_mcp/models/responses.py - Add UpdateAgentResponse
• src/fleet_mcp/repositories/agent_repository.py - Add update method
• src/fleet_mcp/services/agent_service.py - Add update_agent method

Created

• src/fleet_mcp/tools/update_agent.py - MCP tool implementation
• scripts/update_workspace.py - Standalone CLI script (executable)
• tests/clients/test_coder_client_update.py - Client layer tests
• tests/repositories/test_agent_repository_update.py - Repository layer tests

Impact

Benefits

• ✅ Enables programmatic workspace template updates
• ✅ Provides both MCP tool and CLI interfaces
• ✅ Robust error handling and timeout management
• ✅ Automatic active version detection
• ✅ Comprehensive test coverage

Breaking Changes

• None - This is purely additive functionality

Checklist

• Implementation follows existing architecture patterns
• Comprehensive test coverage added
• All tests passing
• Code formatted and linted
• Documentation in docstrings
• CLI tool with --help documentation
• Error handling implemented
• No breaking changes

Test Plan

1. Unit Tests: Run nx test fleet-mcp
2. Manual Testing:

   # Test MCP tool via fleet-mcp server
   # Test CLI script with real workspace
   ./scripts/update_workspace.py --workspace-name test-ws

🤖 Generated with Claude Code