File size: 3,330 Bytes
c745a99
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
"""Environment designer — provisions initial AWS state for each task.

Currently supports raw AWS CLI setup commands. Designed to be extended
with CloudFormation YAML template support so that each difficulty level
can declaratively define its starting infrastructure.
"""

from __future__ import annotations

import logging
from enum import Enum

from pydantic import BaseModel, Field

from models import SetupCommand, Task
from server.services.environment_strategy import EnvironmentStrategy
from server.services.drift_engine import DriftEngine

logger = logging.getLogger(__name__)


class ProvisionMethod(str, Enum):
    """How the initial environment state is provisioned."""

    CLI_COMMANDS = "cli_commands"
    CLOUDFORMATION = "cloudformation"


class ProvisionResult(BaseModel):
    """Outcome of provisioning the environment for a task."""

    success: bool = True
    method: ProvisionMethod = ProvisionMethod.CLI_COMMANDS
    resources_created: int = 0
    errors: list[str] = Field(default_factory=list)


class EnvironmentDesigner:
    """Provisions the initial AWS state required by a task before the agent acts.

    Usage::

        designer = EnvironmentDesigner(backend)
        result = designer.apply(task)
        if not result.success:
            logger.error("Failed to set up environment: %s", result.errors)
    """

    def __init__(self, backend: EnvironmentStrategy) -> None:
        self._backend = backend
        self._drift_engine = DriftEngine(backend)

    def apply(self, task: Task) -> ProvisionResult:
        """Apply the task's environment setup to MiniStack.

        Dispatches to the appropriate provisioning method based on what the
        task defines. Currently supports ``setup_commands``; CloudFormation
        support can be added by extending this method.

        Returns:
            A ``ProvisionResult`` summarising what happened.
        """
        if not task.setup_commands:
            return ProvisionResult(resources_created=0)

        result = self._apply_cli_commands(task.setup_commands)

        # Apply random configuration drifts after provisioning correct state
        if task.possible_drifts:
            applied = self._drift_engine.apply_drift(task)
            logger.info("Applied %d configuration drifts", len(applied))

        return result

    # -- Provisioning strategies ----------------------------------------------

    def _apply_cli_commands(self, commands: list[SetupCommand]) -> ProvisionResult:
        """Execute a list of setup commands against MiniStack."""
        errors: list[str] = []
        resources_created = 0

        for setup_cmd in commands:
            success, _stdout, stderr = self._backend.execute_command(setup_cmd.command)
            if success:
                resources_created += 1
            else:
                msg = f"Setup command failed: {setup_cmd.command}{stderr}"
                if setup_cmd.ignore_failure:
                    logger.info("Ignoring failed setup command: %s", msg)
                else:
                    logger.warning(msg)
                    errors.append(msg)

        return ProvisionResult(
            success=len(errors) == 0,
            method=ProvisionMethod.CLI_COMMANDS,
            resources_created=resources_created,
            errors=errors,
        )