Validators steer, evals grade: picking a model on cost per pass

1
"""Tool surface, re-exported so callers import from ``agent.tools``."""
2

3
from agent.tools.filesystem import (
4
    WorkspaceDeps,
5
    delete_file,
6
    edit_file,
7
    list_files,
8
    read_file,
9
    write_file,
10
)
11
from agent.tools.validators import (
12
    checkov,
13
    terraform_init,
14
    terraform_validate,
15
    tflint,
16
    validate_workspace,
17
)
18

19
__all__ = [
20
    "WorkspaceDeps",
21
    "checkov",
22
    "delete_file",
23
    "edit_file",
24
    "list_files",
25
    "read_file",
26
    "terraform_init",
27
    "terraform_validate",
28
    "tflint",
29
    "validate_workspace",
30
    "write_file",
31
]

1
"""Filesystem tools: list, read, write, edit, delete inside the workspace.
2

3
All tools take ``RunContext[WorkspaceDeps]`` so the workspace root is
4
threaded through ``ctx.deps.root`` instead of read from a module global.
5
That keeps tests, evals, and multi-tenant runs from sharing state.
6
"""
7

8
from pathlib import Path
9

10
from pydantic import BaseModel, ConfigDict, Field
11
from pydantic_ai import ModelRetry, RunContext
12

13

14
class WorkspaceDeps(BaseModel):
15
    """Per-run dependencies threaded through ``RunContext``.
16

17
    ``root`` is the directory the agent is allowed to read, write, and
18
    validate inside. Tools resolve every path relative to it and reject
19
    anything that escapes the root, so the agent cannot reach outside
20
    the workspace via ``..`` or absolute paths.
21
    """
22

23
    model_config = ConfigDict(arbitrary_types_allowed=True)
24

25
    root: Path
26
    files_read: set[Path] = Field(default_factory=set)
27
    # Mutations are tracked here rather than diffed from the filesystem, so the
28
    # no-op check in core.py keeps working once runs start from a seeded
29
    # project instead of an empty workspace.
30
    files_changed: set[Path] = Field(default_factory=set)
31

32

33

34

35
def list_files(ctx: RunContext[WorkspaceDeps], path: str = ".") -> list[str]:
36
    """List files under ``path`` relative to the workspace root."""
37
    return [
38
        str(p.relative_to(ctx.deps.root)) for p in _resolve_absolute_folder(ctx, path).iterdir()
39
    ]
40

41

42
def read_file(ctx: RunContext[WorkspaceDeps], path: str) -> str:
43
    """Read the file at ``path`` and return its contents."""
44
    file = _resolve_absolute_file(ctx, path)
45
    with file.open() as f:
46
        ctx.deps.files_read.add(file)
47
        return f.read()
48

49

50
def write_file(ctx: RunContext[WorkspaceDeps], path: str, contents: str) -> None:
51
    """Create or overwrite the file at ``path`` with ``contents``."""
52
    file = _resolve_absolute_path(ctx, path)
53
    if file.exists() and file not in ctx.deps.files_read:
54
        raise ModelRetry(
55
            f"File {path} already exists. If you want to overwrite it, then delete it first."
56
        )
57
    with file.open("w") as f:
58
        # agent wrote it and knows the content
59
        ctx.deps.files_read.add(file)
60
        ctx.deps.files_changed.add(file)
61
        f.write(contents)
62

63

64
def edit_file(
65
    ctx: RunContext[WorkspaceDeps],
66
    path: str,
67
    old_string: str,
68
    new_string: str,
69
) -> None:
70
    """Replace ``old_string`` with ``new_string`` in the file at ``path``."""
71
    file = _resolve_absolute_file(ctx, path)
72
    if file not in ctx.deps.files_read:
73
        raise ModelRetry(f"File {path} was not read. If you want to edit it, then read it first.")
74
    with file.open("r+") as f:
75
        contents = f.read()
76
        if old_string not in contents:
77
            raise ModelRetry(f"String {old_string} not found in file {path}.")
78
        if contents.count(old_string) > 1:
79
            raise ModelRetry(f"String {old_string} found more than once in file {path}.")
80
        f.seek(0)
81
        f.write(contents.replace(old_string, new_string))
82
        f.truncate()
83
        ctx.deps.files_changed.add(file)
84

85

86

87

88
def delete_file(ctx: RunContext[WorkspaceDeps], path: str) -> None:
89
    """Delete the file at ``path``."""
90
    file = _resolve_absolute_file(ctx, path)
91
    if file not in ctx.deps.files_read:
92
        raise ModelRetry(f"File {path} was not read. If you want to delete it, then read it first.")
93
    file.unlink()
94
    ctx.deps.files_changed.add(file)
95

96

97
def _resolve_absolute_folder(ctx: RunContext[WorkspaceDeps], path: str):
98
    if not (absolute_path := _resolve_absolute_path(ctx, path)).is_dir():
99
        raise ModelRetry(f"Path {path} must be a directory.")
100
    return absolute_path
101

102

103
def _resolve_absolute_file(ctx: RunContext[WorkspaceDeps], path: str):
104
    absolute_path = _resolve_absolute_path(ctx, path)
105
    if not absolute_path.exists():
106
        raise ModelRetry(f"File {path} does not exist.")
107
    if not absolute_path.is_file():
108
        raise ModelRetry(f"Path {path} must be a file.")
109
    return absolute_path
110

111

112
def _resolve_absolute_path(ctx: RunContext[WorkspaceDeps], path: str):
113
    root = ctx.deps.root.resolve()
114
    absolute_path = (root / Path(path)).resolve()
115
    if not absolute_path.is_relative_to(root):
116
        raise ModelRetry(
117
            f"Path {path} must be relative to the workspace root, "
118
            f"it cannot be absolute or walk up the directory tree."
119
        )
120
    return absolute_path

1
"""Validation tools: the checks the agent runs on the workspace."""
2

3
import subprocess
4
import threading
5
from pathlib import Path
6

7
from agent.memory import track_memory
8
from agent.tools.filesystem import WorkspaceDeps
9
from pydantic import BaseModel
10
from pydantic_ai import ModelRetry, RunContext
11

12
# Every failure report becomes a tool result the model carries for the rest of
13
# the run, so an uncapped one is paid for on every later turn and lands in the
14
# audit record. checkov's report is the offender; its head carries the
15
# failures that matter, and the agent re-runs the check anyway.
16
MAX_REPORT_CHARS = 8_000
17

18

19
def _capped(text: str) -> str:
20
    if len(text) <= MAX_REPORT_CHARS:
21
        return text
22
    return f"{text[:MAX_REPORT_CHARS]}\n[{len(text) - MAX_REPORT_CHARS} more characters cut]"
23

24

25

26

27
# terraform's provider cache is not concurrency safe: with several inits
28
# installing at once the installer's behaviour is undefined, and a workspace
29
# that init'd against a half-written cache fails its next command with
30
# "Failed to obtain provider schema". One run is alone in its Lambda container,
31
# so this only ever contends in an eval sweep, where cases run in parallel.
32
# https://developer.hashicorp.com/terraform/cli/config/config-file
33
_INSTALL_LOCK = threading.Lock()
34

35

36

37

38
class CommandResult(BaseModel):
39
    success: bool
40
    stdout: str
41
    stderr: str
42

43
    def format_error_for_agent(self) -> str:
44
        """The failure as the model sees it, capped. Callers keep the full output."""
45
        if not self.success:
46
            return f"failed:\n{_capped(self.stdout)}\n{_capped(self.stderr)}"
47
        raise RuntimeError("CommandResult.format_error_for_agent() called on success")
48

49

50
class Command(BaseModel):
51
    name: str
52
    commands: list[str]
53

54
    @property
55
    def installs_providers(self) -> bool:
56
        return "init" in self.commands
57

58
    def run(self, path: Path) -> CommandResult:
59
        if self.installs_providers:
60
            with _INSTALL_LOCK:
61
                return self._run(path)
62
        return self._run(path)
63

64
    def _run(self, path: Path) -> CommandResult:
65
        # Inside the lock, so the span measures the command and not the wait.
66
        with track_memory(self.name):
67
            result = subprocess.run(self.commands, cwd=path, capture_output=True, text=True)
68
            return CommandResult(
69
                success=result.returncode == 0, stdout=result.stdout, stderr=result.stderr
70
            )
71

72
    def run_root(self, ctx: RunContext[WorkspaceDeps]) -> CommandResult:
73
        return self.run(ctx.deps.root)
74

75
    def run_in_tool(self, path: Path) -> str:
76
        result = self.run(path)
77
        if not result.success:
78
            raise ModelRetry(f"{self.name} {result.format_error_for_agent()}")
79
        return f"OK: {self.name} passed."
80

81

82

83
TERRAFORM_INIT = Command(
84
    name="terraform_init",
85
    commands=["terraform", "init", "-backend=false", "-input=false", "-no-color"],
86
)
87

88
TERRAFORM_VALIDATE = Command(
89
    name="terraform_validate",
90
    commands=["terraform", "validate", "-no-color"],
91
)
92

93
TERRAFORM_FMT = Command(
94
    name="terraform_fmt",
95
    commands=["terraform", "fmt", "-recursive"],
96
)
97

98
TFLINT = Command(
99
    name="tflint",
100
    commands=["tflint", "--format", "compact"],
101
)
102

103
# Checks skipped by default: cost or architecture posture decisions, not
104
# security baselines, and routinely disabled in real projects. Genuine
105
# baselines (public access blocks, encryption at rest, IAM wildcards) stay on.
106
_CHECKOV_DEFAULT_SKIPS = [
107
    "CKV_AWS_18",  # S3 access logging on every bucket
108
    "CKV_AWS_144",  # S3 cross-region replication
109
    "CKV_AWS_145",  # S3 must use KMS; SSE-S3 (CKV_AWS_19) still enforced
110
    "CKV2_AWS_61",  # S3 lifecycle configuration on every bucket
111
    "CKV2_AWS_62",  # S3 event notifications on every bucket
112
    "CKV_AWS_50",  # Lambda X-Ray tracing
113
    "CKV_AWS_115",  # Lambda reserved concurrency
114
    "CKV_AWS_116",  # Lambda dead-letter queue
115
    "CKV_AWS_117",  # Lambda attached to a VPC
116
    "CKV_AWS_272",  # Lambda code signing
117
    "CKV_AWS_338",  # CloudWatch log retention of at least a year
118
]
119

120

121
def _checkov_command(root: Path) -> Command:
122
    """Default skips apply only when the workspace brings no config of its own:
123
    checkov auto-discovers .checkov.yaml in the scanned directory, and a project
124
    that states its policy wins over ours.
125
    """
126
    commands = ["checkov", "-d", ".", "--quiet", "--compact"]
127
    if not any((root / name).exists() for name in (".checkov.yaml", ".checkov.yml")):
128
        commands += ["--skip-check", ",".join(_CHECKOV_DEFAULT_SKIPS)]
129
    return Command(name="checkov", commands=commands)
130

131

132

133

134
def terraform_init(ctx: RunContext[WorkspaceDeps]) -> str:
135
    """Run ``terraform init`` in the workspace.
136

137
    Required once before the first ``terraform_validate`` and again after
138
    provider or module requirements change.
139
    """
140
    return TERRAFORM_INIT.run_in_tool(ctx.deps.root)
141

142

143
def terraform_validate(ctx: RunContext[WorkspaceDeps]) -> str:
144
    """Run ``terraform validate`` in the workspace and return its output."""
145
    return TERRAFORM_VALIDATE.run_in_tool(ctx.deps.root)
146

147

148
def tflint(ctx: RunContext[WorkspaceDeps]) -> str:
149
    return TFLINT.run_in_tool(ctx.deps.root)
150

151

152
def checkov(ctx: RunContext[WorkspaceDeps]) -> str:
153
    return _checkov_command(ctx.deps.root).run_in_tool(ctx.deps.root)
154

155

156

157

158
def validate_workspace(path: Path) -> CommandResult:
159
    # Normalize formatting before gating; a fmt failure (unparsable HCL) is
160
    # ignored here because terraform validate reports it better one step later.
161
    TERRAFORM_FMT.run(path)
162
    for command in [TERRAFORM_VALIDATE, TFLINT, _checkov_command(path)]:
163
        result = command.run(path)
164
        if not result.success:
165
            break
166
    return result

"""Define the agent and execute one run.

Runtime-agnostic: a fresh /tmp workspace threaded through the file and validate
8 collapsed lines
tools, one run id that joins the audit trace, the parked artifacts, and the
result. The telemetry pipeline lives in observability.py, model construction in
models.py, and run persistence plus the provider re-lock in runs.py; the Lambda
envelope and INIT wiring live in lambda_entry.py. This module knows nothing
about Lambda events.
"""

import uuid
from pathlib import Path
from tempfile import TemporaryDirectory

from pydantic import BaseModel, ConfigDict
from pydantic_ai import Agent

from agent.env import require_env
from agent.models import _build_model
from agent.runs import _persist_run, _relock_providers
from agent.tools import (
    WorkspaceDeps,
    checkov,
    delete_file,
    edit_file,
    list_files,
    read_file,
    terraform_init,
    terraform_validate,
    tflint,
    validate_workspace,
    write_file,
)
from pydantic import BaseModel, ConfigDict
from pydantic_ai import Agent, ModelRetry, RunContext
from pydantic_ai.settings import ModelSettings
from pydantic_ai.usage import UsageLimits

# A soft reproducibility nudge following the standard Terraform pattern: a
# version constraint in the config, the exact version and checksums in the lock
20 collapsed lines
# file. The tool-call spans in the trace are the ground truth for what the agent
# actually wrote.
_PROVIDER_PIN_RULE = (
    "When you add the AWS provider, give it a version constraint such as "
    '"~> 6.0" rather than leaving it unconstrained. terraform init records '
    "the exact resolved version and checksums in .terraform.lock.hcl, "
    "which travels with the workspace and is the reproducibility record, "
    "so the constraint does not need to be an exact pin. Pin an exact "
    "version only when the user asks for one. Example:\n"
    "\n"
    "terraform {\n"
    "  required_providers {\n"
    "    aws = {\n"
    '      source  = "hashicorp/aws"\n'
    '      version = "~> 6.0"\n'
    "    }\n"
    "  }\n"
    "}"
)

SYSTEM_PROMPT = (
    "You are the terraform-pr-agent. You operate on a Terraform workspace "
    "through file tools (list_files, read_file, write_file, edit_file, "
    "delete_file) and two terraform tools. Use the file tools to explore, "
    "write, and edit. Run terraform_init before your first validate and "
    "again whenever you add or change provider or module requirements. "
    "Call terraform_validate after you write or change files to confirm "
    "the workspace still parses; treat its output as feedback and edit "
    "until it is clean.\n\n" + _PROVIDER_PIN_RULE
    "delete_file) and a set of terraform and validation tools. Use the file "
    "tools to explore, write, and edit. Run terraform_init before your first "
    "validate and again whenever you add or change provider or module "
    "requirements. Call terraform_validate after you write or change files to "
    "confirm the workspace still parses; treat its output as feedback and edit "
    "until it is clean. When the workspace is done, report your work through "
    "the final result.\n\n"
    "Give every variable you declare a default value, so the workspace plans "
    "without prompting for input.\n\n"
    "Run tflint and checkov and clear all findings before reporting done.\n\n" + _PROVIDER_PIN_RULE
)


class TaskResult(BaseModel):
    """The agent's structured self-report on a finished run.

    Each required field pushes the agent to consider that dimension of its
    work; the tool-call spans in the trace remain the ground truth that
    exposes any embellishment. Because this is the agent's output type,
    ending a run now requires calling the final_result output tool, so a
    text-only reply can no longer end a run silently.
    """

    model_config = ConfigDict(frozen=True)

    summary: str
    """One line on what was done."""
    solution_description: str
    """How the problem was solved, including architectural choices."""
    validations_run: list[str]
    """Which validation tools were invoked during the run."""
    issues_addressed: list[str]
    """Security or correctness problems identified and fixed."""
    known_limitations: list[str]
    """What was not handled; surfaces for human review."""
    ready_for_review: bool
    """The agent's self-assessment that the workspace is PR-ready."""



# One Agent instance is reused across invocations: the tools reach the workspace
# through RunContext.deps, so each run_sync scopes them to a fresh WorkspaceDeps.
# The model carries no default; it is built from the registry at INVOKE and
# passed per run, so switching DEFAULT_MODEL needs no code change.
#
# With TaskResult as the output type, plain text is not an allowed way to end
# the run, so pydantic-ai asks providers to force tool use on every turn
# (Bedrock toolChoice "any"). Note Bedrock rejects that combination once
# extended thinking is enabled on a model; nothing in the registry enables it.
agent = Agent(
    deps_type=WorkspaceDeps,
    output_type=TaskResult,
    system_prompt=SYSTEM_PROMPT,
    tools=[
        list_files,
3 collapsed lines
        read_file,
        write_file,
        edit_file,
        delete_file,
        terraform_init,
        terraform_validate,
        tflint,
        checkov,
    ],
    # Tools raise ModelRetry on failure; pydantic-ai ends the run once one tool
    # fails more than `retries` times in a row (a success resets the count). The
    # default of 1 would end the run on the second straight failing validate,
    # which is a normal part of the write-validate-edit loop, so the budget is
    # raised well past anything a converging run produces. The per-run turn cap
    # stays the runaway guard.
    retries=10,
    retries={"tools": 3, "output": 3},
)

# The two ceilings on how long one run's message history can get, and with it
# the biggest audit record the trace produces (see observability._span_record).
# request_limit is pydantic-ai's own default, set here so the bound is ours;
# max_tokens caps what each turn adds. The third term, tool output, is capped
# in tools/validators.py.
RUN_LIMITS = UsageLimits(request_limit=50)
RUN_SETTINGS = ModelSettings(max_tokens=4096)

# Caller-side backstop budget: how many follow-up runs to spend trying to get a
# clean validate after the agent reports done. The per-tool `retries` above
# guards the loop inside one run; this guards the run as a whole.
_MAX_VALIDATE_RETRIES = 3

_RETRY_PROMPT = (
    "terraform validate still reports errors after you finished. "
    "terraform validate still reports errors after you reported done. "
    "Fix them and validate again.\n\n{output}"
)


class ValidateDidNotConverge(RuntimeError):
    """terraform validate still failed after the caller-side retry budget."""
@agent.output_validator
def _validate_final_workspace(ctx: RunContext[WorkspaceDeps], output: TaskResult) -> TaskResult:
    # The deliverable is the workspace, not this object: a run that changed
    # nothing produced nothing, however plausible its self-report reads.
    if not ctx.deps.files_changed:
        raise ModelRetry(
            "You reported done but made no changes to the workspace. "
            "Use the file tools to implement the request, validate, then report done again."
        )
    validation_result = validate_workspace(ctx.deps.root)
    if not validation_result.success:
        raise ModelRetry(_RETRY_PROMPT.format(output=validation_result.format_error_for_agent()))
    return output




class RunResult(BaseModel):
    model_config = ConfigDict(frozen=True)

    run_id: str
    model: str
    output: str
    output: TaskResult
    input_tokens: int
    output_tokens: int


def execute(prompt: str, model: str | None = None) -> RunResult:
def execute(prompt: str, model: str | None = None, workspace: Path | None = None) -> RunResult:
    """Run the agent once and return the result.

    The run id is the trace's gen_ai.conversation.id and the runs/<run_id>/
    prefix, so one identifier joins trace, artifacts, and result. After the run
    the caller re-validates the workspace and feeds any failure back as a
    follow-up turn; a run that still fails after the retry budget raises, so the
    workspace ships under status error for debugging. ``model`` overrides
    DEFAULT_MODEL when given.
    prefix, so one identifier joins trace, artifacts, and result. The output
    validator re-checks the workspace every time the agent reports done, so a
    returned run is a validated run; one that exhausts the output budget raises
    UnexpectedModelBehavior and the workspace ships under status error for
    debugging. ``model`` overrides DEFAULT_MODEL when given.

    ``workspace`` runs the agent in that directory instead of a throwaway
    tempdir, and hands its lifetime to the caller: a run that owns its own
    directory parks the finished tree in S3 before deleting it, one that does
    not leaves the files where the caller can still read them. The eval
    harness passes a workspace for exactly that reason, so a sweep needs
    neither the runs bucket nor the provider re-lock the review flow wants.
    """
    model_name = model or require_env("DEFAULT_MODEL")
    run_id = str(uuid.uuid4())
    with TemporaryDirectory(dir="/tmp") as workspace:
        root = Path(workspace)
        deps = WorkspaceDeps(root=root)
        try:
            built = _build_model(model_name)
            result = agent.run_sync(
                prompt,
                deps=deps,
                conversation_id=run_id,
                model=built,
                metadata={"model": model_name},
            )
            # The agent can report done while terraform validate still fails.
            # Re-validate ourselves and feed any error back as a follow-up turn,
            # reusing the run id and message history so each retry is an
            # invoke_agent span under the one invocation trace. Give up after the
            # budget and raise so the failure is honest rather than a clean run
            # over a broken workspace.
            ok, output = validate_workspace(root)
            attempts = 0
            while not ok and attempts < _MAX_VALIDATE_RETRIES:
                attempts += 1
                result = agent.run_sync(
                    _RETRY_PROMPT.format(output=output),
                    deps=deps,
                    conversation_id=run_id,
                    model=built,
                    message_history=result.all_messages(),
                    metadata={"model": model_name},
                )
                ok, output = validate_workspace(root)
            if not ok:
                raise ValidateDidNotConverge(output)
        except Exception as error:
    if workspace is not None:
        return _execute_in(prompt, model_name, run_id, workspace, persist=False)
    with TemporaryDirectory(dir="/tmp") as tmp:
        return _execute_in(prompt, model_name, run_id, Path(tmp), persist=True)


def _execute_in(prompt: str, model_name: str, run_id: str, root: Path, persist: bool) -> RunResult:
    deps = WorkspaceDeps(root=root)
    try:
        built = _build_model(model_name)
        result = agent.run_sync(
            prompt,
            deps=deps,
            conversation_id=run_id,
            model=built,
            metadata={"model": model_name},
            usage_limits=RUN_LIMITS,
            model_settings=RUN_SETTINGS,
        )
    except Exception as error:
        if persist:
            _persist_run(run_id, root, status="error", error=repr(error))
            raise
        raise
    if persist:
        _relock_providers(root)
        _persist_run(run_id, root, status="ok")
    return RunResult(run_id=run_id, model=model_name, output=str(result.output))
    usage = result.usage
    return RunResult(
        run_id=run_id,
        model=model_name,
        output=result.output,
        input_tokens=usage.input_tokens,
        output_tokens=usage.output_tokens,
    )

1
"""Define the agent and execute one run.
2

3
Runtime-agnostic: a fresh /tmp workspace threaded through the file and validate
4
tools, one run id that joins the audit trace, the parked artifacts, and the
5
result. The telemetry pipeline lives in observability.py, model construction in
6
models.py, and run persistence plus the provider re-lock in runs.py; the Lambda
7
envelope and INIT wiring live in lambda_entry.py. This module knows nothing
8
about Lambda events.
9
"""
10

11
import uuid
12
from pathlib import Path
13
from tempfile import TemporaryDirectory
14

15
from agent.env import require_env
16
from agent.models import _build_model
17
from agent.runs import _persist_run, _relock_providers
18
from agent.tools import (
19
    WorkspaceDeps,
20
    checkov,
21
    delete_file,
22
    edit_file,
23
    list_files,
24
    read_file,
25
    terraform_init,
26
    terraform_validate,
27
    tflint,
28
    validate_workspace,
29
    write_file,
30
)
31
from pydantic import BaseModel, ConfigDict
32
from pydantic_ai import Agent, ModelRetry, RunContext
33
from pydantic_ai.settings import ModelSettings
34
from pydantic_ai.usage import UsageLimits
35

36
# A soft reproducibility nudge following the standard Terraform pattern: a
37
# version constraint in the config, the exact version and checksums in the lock
38
# file. The tool-call spans in the trace are the ground truth for what the agent
39
# actually wrote.
40
_PROVIDER_PIN_RULE = (
41
    "When you add the AWS provider, give it a version constraint such as "
42
    '"~> 6.0" rather than leaving it unconstrained. terraform init records '
43
    "the exact resolved version and checksums in .terraform.lock.hcl, "
44
    "which travels with the workspace and is the reproducibility record, "
45
    "so the constraint does not need to be an exact pin. Pin an exact "
46
    "version only when the user asks for one. Example:\n"
47
    "\n"
48
    "terraform {\n"
49
    "  required_providers {\n"
50
    "    aws = {\n"
51
    '      source  = "hashicorp/aws"\n'
52
    '      version = "~> 6.0"\n'
53
    "    }\n"
54
    "  }\n"
55
    "}"
56
)
57

58
SYSTEM_PROMPT = (
59
    "You are the terraform-pr-agent. You operate on a Terraform workspace "
60
    "through file tools (list_files, read_file, write_file, edit_file, "
61
    "delete_file) and a set of terraform and validation tools. Use the file "
62
    "tools to explore, write, and edit. Run terraform_init before your first "
63
    "validate and again whenever you add or change provider or module "
64
    "requirements. Call terraform_validate after you write or change files to "
65
    "confirm the workspace still parses; treat its output as feedback and edit "
66
    "until it is clean. When the workspace is done, report your work through "
67
    "the final result.\n\n"
68
    "Give every variable you declare a default value, so the workspace plans "
69
    "without prompting for input.\n\n"
70
    "Run tflint and checkov and clear all findings before reporting done.\n\n" + _PROVIDER_PIN_RULE
71
)
72

73

74
class TaskResult(BaseModel):
75
    """The agent's structured self-report on a finished run.
76

77
    Each required field pushes the agent to consider that dimension of its
78
    work; the tool-call spans in the trace remain the ground truth that
79
    exposes any embellishment. Because this is the agent's output type,
80
    ending a run now requires calling the final_result output tool, so a
81
    text-only reply can no longer end a run silently.
82
    """
83

84
    model_config = ConfigDict(frozen=True)
85

86
    summary: str
87
    """One line on what was done."""
88
    solution_description: str
89
    """How the problem was solved, including architectural choices."""
90
    validations_run: list[str]
91
    """Which validation tools were invoked during the run."""
92
    issues_addressed: list[str]
93
    """Security or correctness problems identified and fixed."""
94
    known_limitations: list[str]
95
    """What was not handled; surfaces for human review."""
96
    ready_for_review: bool
97
    """The agent's self-assessment that the workspace is PR-ready."""
98

99

100

101
# One Agent instance is reused across invocations: the tools reach the workspace
102
# through RunContext.deps, so each run_sync scopes them to a fresh WorkspaceDeps.
103
# The model carries no default; it is built from the registry at INVOKE and
104
# passed per run, so switching DEFAULT_MODEL needs no code change.
105
#
106
# With TaskResult as the output type, plain text is not an allowed way to end
107
# the run, so pydantic-ai asks providers to force tool use on every turn
108
# (Bedrock toolChoice "any"). Note Bedrock rejects that combination once
109
# extended thinking is enabled on a model; nothing in the registry enables it.
110
agent = Agent(
111
    deps_type=WorkspaceDeps,
112
    output_type=TaskResult,
113
    system_prompt=SYSTEM_PROMPT,
114
    tools=[
115
        list_files,
116
        read_file,
117
        write_file,
118
        edit_file,
119
        delete_file,
120
        terraform_init,
121
        terraform_validate,
122
        tflint,
123
        checkov,
124
    ],
125
    retries={"tools": 3, "output": 3},
126
)
127

128
# The two ceilings on how long one run's message history can get, and with it
129
# the biggest audit record the trace produces (see observability._span_record).
130
# request_limit is pydantic-ai's own default, set here so the bound is ours;
131
# max_tokens caps what each turn adds. The third term, tool output, is capped
132
# in tools/validators.py.
133
RUN_LIMITS = UsageLimits(request_limit=50)
134
RUN_SETTINGS = ModelSettings(max_tokens=4096)
135

136

137
_RETRY_PROMPT = (
138
    "terraform validate still reports errors after you reported done. "
139
    "Fix them and validate again.\n\n{output}"
140
)
141

142

143
@agent.output_validator
144
def _validate_final_workspace(ctx: RunContext[WorkspaceDeps], output: TaskResult) -> TaskResult:
145
    # The deliverable is the workspace, not this object: a run that changed
146
    # nothing produced nothing, however plausible its self-report reads.
147
    if not ctx.deps.files_changed:
148
        raise ModelRetry(
149
            "You reported done but made no changes to the workspace. "
150
            "Use the file tools to implement the request, validate, then report done again."
151
        )
152
    validation_result = validate_workspace(ctx.deps.root)
153
    if not validation_result.success:
154
        raise ModelRetry(_RETRY_PROMPT.format(output=validation_result.format_error_for_agent()))
155
    return output
156

157

158

159

160
class RunResult(BaseModel):
161
    model_config = ConfigDict(frozen=True)
162

163
    run_id: str
164
    model: str
165
    output: TaskResult
166
    input_tokens: int
167
    output_tokens: int
168

169

170
def execute(prompt: str, model: str | None = None, workspace: Path | None = None) -> RunResult:
171
    """Run the agent once and return the result.
172

173
    The run id is the trace's gen_ai.conversation.id and the runs/<run_id>/
174
    prefix, so one identifier joins trace, artifacts, and result. The output
175
    validator re-checks the workspace every time the agent reports done, so a
176
    returned run is a validated run; one that exhausts the output budget raises
177
    UnexpectedModelBehavior and the workspace ships under status error for
178
    debugging. ``model`` overrides DEFAULT_MODEL when given.
179

180
    ``workspace`` runs the agent in that directory instead of a throwaway
181
    tempdir, and hands its lifetime to the caller: a run that owns its own
182
    directory parks the finished tree in S3 before deleting it, one that does
183
    not leaves the files where the caller can still read them. The eval
184
    harness passes a workspace for exactly that reason, so a sweep needs
185
    neither the runs bucket nor the provider re-lock the review flow wants.
186
    """
187
    model_name = model or require_env("DEFAULT_MODEL")
188
    run_id = str(uuid.uuid4())
189
    if workspace is not None:
190
        return _execute_in(prompt, model_name, run_id, workspace, persist=False)
191
    with TemporaryDirectory(dir="/tmp") as tmp:
192
        return _execute_in(prompt, model_name, run_id, Path(tmp), persist=True)
193

194

195
def _execute_in(prompt: str, model_name: str, run_id: str, root: Path, persist: bool) -> RunResult:
196
    deps = WorkspaceDeps(root=root)
197
    try:
198
        built = _build_model(model_name)
199
        result = agent.run_sync(
200
            prompt,
201
            deps=deps,
202
            conversation_id=run_id,
203
            model=built,
204
            metadata={"model": model_name},
205
            usage_limits=RUN_LIMITS,
206
            model_settings=RUN_SETTINGS,
207
        )
208
    except Exception as error:
209
        if persist:
210
            _persist_run(run_id, root, status="error", error=repr(error))
211
        raise
212
    if persist:
213
        _relock_providers(root)
214
        _persist_run(run_id, root, status="ok")
215
    usage = result.usage
216
    return RunResult(
217
        run_id=run_id,
218
        model=model_name,
219
        output=result.output,
220
        input_tokens=usage.input_tokens,
221
        output_tokens=usage.output_tokens,
222
    )

1
"""Read a required environment variable, failing loudly when it is missing."""
2

3
import os
4

5

6
def require_env(name: str) -> str:
7
    """Return the value of `name`, or raise if it is unset or empty.
8

9
    A missing required variable is a deployment fault, not a runtime branch, so
10
    we surface it the same way everywhere instead of degrading into a silent
11
    no-op. Empty is treated as unset: a blank value is never a real config.
12
    """
13
    value = os.environ.get(name)
14
    if not value:
15
        raise RuntimeError(f"required environment variable {name} is unset")
16
    return value

"""The Lambda boundary: parse the event, run the agent, shape the response.

This module also owns the INIT wiring. The container CMD targets
3 collapsed lines
agent.lambda_entry.handler, so a unit test that imports agent.core never
configures logfire or registers the Firehose audit processor. Everything
Lambda-specific lives here, off the core, which is why no runtime-detection
check is needed to keep it out of tests.
"""

from typing import NotRequired, TypedDict
from typing import Any, NotRequired, TypedDict

import logfire

from agent import observability
from agent.core import execute

7 collapsed lines

class HandlerEvent(TypedDict):
    prompt: str
    model: NotRequired[str]


class HandlerResponse(TypedDict):
    status: str
    run_id: str
    model: str
    output: str
    output: dict[str, Any]


def handler(event: HandlerEvent, context: object) -> HandlerResponse:
2 collapsed lines
    """Lambda entry point: require a prompt, run the agent, wrap the result.

    ``prompt`` is required; an event without one is a caller error and fails
    fast rather than running a default. ``model`` overrides DEFAULT_MODEL when
    given. A run that does not converge raises, so the Lambda reports 5xx and
    the workspace ships under status error for debugging.
    the workspace ships under status error for debugging. ``output`` is the
    agent's TaskResult self-report, serialized for the JSON response.
    """
    prompt = event.get("prompt")
    if not prompt:
3 collapsed lines
        raise ValueError("event missing required 'prompt'")
    result = execute(prompt, event.get("model"))
    return {
        "status": "ok",
        "run_id": result.run_id,
        "model": result.model,
        "output": result.output,
        "output": result.output.model_dump(),
    }


9 collapsed lines
def bootstrap() -> None:
    """Stand up telemetry, then attach the Lambda runtime adapter.

    configure() first so the tracer provider exists when the handler is wrapped.
    instrument_aws_lambda wraps the target named by _HANDLER
    (agent.lambda_entry.handler) in place, so each invocation becomes one trace.
    """
    observability.configure()
    logfire.instrument_aws_lambda(handler)


bootstrap()

1
"""The Lambda boundary: parse the event, run the agent, shape the response.
2

3
This module also owns the INIT wiring. The container CMD targets
4
agent.lambda_entry.handler, so a unit test that imports agent.core never
5
configures logfire or registers the Firehose audit processor. Everything
6
Lambda-specific lives here, off the core, which is why no runtime-detection
7
check is needed to keep it out of tests.
8
"""
9

10
from typing import Any, NotRequired, TypedDict
11

12
import logfire
13
from agent import observability
14
from agent.core import execute
15

16

17
class HandlerEvent(TypedDict):
18
    prompt: str
19
    model: NotRequired[str]
20

21

22
class HandlerResponse(TypedDict):
23
    status: str
24
    run_id: str
25
    model: str
26
    output: dict[str, Any]
27

28

29
def handler(event: HandlerEvent, context: object) -> HandlerResponse:
30
    """Lambda entry point: require a prompt, run the agent, wrap the result.
31

32
    ``prompt`` is required; an event without one is a caller error and fails
33
    fast rather than running a default. ``model`` overrides DEFAULT_MODEL when
34
    given. A run that does not converge raises, so the Lambda reports 5xx and
35
    the workspace ships under status error for debugging. ``output`` is the
36
    agent's TaskResult self-report, serialized for the JSON response.
37
    """
38
    prompt = event.get("prompt")
39
    if not prompt:
40
        raise ValueError("event missing required 'prompt'")
41
    result = execute(prompt, event.get("model"))
42
    return {
43
        "status": "ok",
44
        "run_id": result.run_id,
45
        "model": result.model,
46
        "output": result.output.model_dump(),
47
    }
48

49

50
def bootstrap() -> None:
51
    """Stand up telemetry, then attach the Lambda runtime adapter.
52

53
    configure() first so the tracer provider exists when the handler is wrapped.
54
    instrument_aws_lambda wraps the target named by _HANDLER
55
    (agent.lambda_entry.handler) in place, so each invocation becomes one trace.
56
    """
57
    observability.configure()
58
    logfire.instrument_aws_lambda(handler)
59

60

61
bootstrap()

1
"""Memory probes for the terraform steps.
2

3
Lambda's Max Memory Used is misleading for this function. The terraform steps
4
download and unpack the AWS provider (~800 MB) into /tmp, and that file IO
5
fills the kernel page cache, which the cgroup-based billed figure counts but
6
which the kernel reclaims under pressure, so it is not OOM risk. ``track_memory``
7
tags a logfire span with three numbers so a run's trace separates real demand
8
from cache: the sandbox's used memory (MemTotal - MemAvailable, the
9
non-reclaimable memory that actually risks OOM), the reclaimable page cache
10
(Cached + Buffers, the bulk of the billed peak), and the peak resident size of
11
the largest terraform subprocess. Read together they show real demand stays
12
under ~1 GB while the billed peak runs to ~2 GB of reclaimable cache.
13
"""
14

15
import resource
16
from collections.abc import Iterator
17
from contextlib import contextmanager
18
from pathlib import Path
19

20
import logfire
21

22
# /proc/meminfo is always present in the Lambda sandbox; the cgroup memory
23
# files are not (neither the v2 /sys/fs/cgroup/memory.current nor the v1
24
# /sys/fs/cgroup/memory/memory.usage_in_bytes is readable there, so reading
25
# them silently returned None). Values are in kibibytes.
26
_MEMINFO = Path("/proc/meminfo")
27

28

29
def _memory_snapshot() -> tuple[int, int] | None:
30
    """(used, cache) bytes for the whole sandbox, or None when /proc/meminfo is absent.
31

32
    ``used`` is MemTotal - MemAvailable, the non-reclaimable memory in use (the
33
    Python runtime plus any live subprocess), which is what actually risks OOM.
34
    ``cache`` is Cached + Buffers, the reclaimable page cache that file IO on
35
    /tmp fills; Lambda's Max Memory Used counts it but real demand does not, so
36
    recording both shows why the billed peak overstates what the function needs.
37
    Absent locally (macOS) and in tests, where the caller leaves the attributes
38
    unset.
39
    """
40
    try:
41
        fields = dict(line.split(":", 1) for line in _MEMINFO.read_text().splitlines())
42
        used_kib = int(fields["MemTotal"].split()[0]) - int(fields["MemAvailable"].split()[0])
43
        cache_kib = int(fields["Cached"].split()[0]) + int(fields["Buffers"].split()[0])
44
    except (OSError, KeyError, ValueError, IndexError):
45
        return None
46
    return used_kib * 1024, cache_kib * 1024
47

48

49
@contextmanager
50
def track_memory(step: str) -> Iterator[None]:
51
    """Run a block in a logfire span tagged with its memory cost.
52

53
    On exit, records the sandbox's used memory and reclaimable page cache, and
54
    the peak resident size of the largest subprocess waited for so far.
55
    ru_maxrss is reported in kilobytes on Linux, so it is scaled to bytes; it is
56
    a monotonic high-water mark across all children, not a per-call delta, so
57
    compare it between steps to see which one grew the subprocess most. Built
58
    with contextmanager, so it doubles as a decorator: ``@track_memory("relock")``.
59
    """
60
    # _span_name forces the OTel span name to the interpolated value, so the
61
    # trace reads "memory.terraform_validate". Without it logfire keeps the
62
    # low-cardinality template "memory.{step}" as the name (its f-string magic
63
    # reconstructs the template), leaving the per-step value only in the attribute.
64
    with logfire.span("memory.{step}", _span_name=f"memory.{step}", step=step) as span:
65
        try:
66
            yield
67
        finally:
68
            if (snapshot := _memory_snapshot()) is not None:
69
                used, cache = snapshot
70
                span.set_attribute("mem.used_bytes", used)
71
                span.set_attribute("mem.cache_bytes", cache)
72
            child_peak_kb = resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss
73
            span.set_attribute("mem.child_max_rss_bytes", child_peak_kb * 1024)

"""Build a pydantic-ai model from the SSM-backed registry, per invocation."""

import json
import os
from functools import cache

from agent.ssm import fetch_parameter
from httpx import AsyncClient, HTTPStatusError, Response
from mistralai.client import Mistral
from mistralai.extra.observability.telemetry import configure_telemetry
from opentelemetry.trace import NoOpTracerProvider
from pydantic_ai.models import Model
from pydantic_ai.models.bedrock import BedrockConverseModel
from pydantic_ai.models.mistral import MistralModel
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.fireworks import FireworksProvider
from pydantic_ai.providers.mistral import MistralProvider
from pydantic_ai.retries import AsyncTenacityTransport, RetryConfig, wait_retry_after
from tenacity import retry_if_exception_type, stop_after_attempt, wait_exponential

from agent.env import require_env
from agent.ssm import fetch_parameter

# Rate limit and transient gateway errors are worth retrying; auth and bad
# request fail fast so a real problem is not retried five times.
_RETRYABLE_STATUS = frozenset({429, 502, 503, 504})
30 collapsed lines


def _raise_for_retryable(response: Response) -> None:
    if response.status_code in _RETRYABLE_STATUS:
        response.raise_for_status()


def _retrying_http_client() -> AsyncClient:
    """An httpx client that retries rate-limit and transient errors.

    pydantic-ai does not retry transport errors itself, so a rate-limited
    Mistral call (the write/init/validate loop can burst past the free-tier
    per-second cap) would fail the whole run. wait_retry_after honours the
    Retry-After header Mistral sends on a 429, falling back to exponential
    backoff. Bedrock uses boto3 with its own retry config, so this is the
    Mistral client only.
    """
    transport = AsyncTenacityTransport(
        config=RetryConfig(
            retry=retry_if_exception_type(HTTPStatusError),
            wait=wait_retry_after(
                fallback_strategy=wait_exponential(multiplier=1, max=60),
                max_wait=300,
            ),
            stop=stop_after_attempt(5),
            reraise=True,
        ),
        validate_response=_raise_for_retryable,
    )
    return AsyncClient(transport=transport)


@cache
def _registry() -> dict[str, dict[str, str]]:
    """The model registry, one GetParameter per container.

    An SSM String parameter at a terraform-owned path: each entry names a
    provider and model id, and Bedrock entries carry the inference-profile ARN.
    MODELS_PARAMETER selects a different registry, which is how the eval
    harness reaches its own copy.
    """
    parameter = os.environ.get("MODELS_PARAMETER", "/terraform-pr-agent/models")
    return json.loads(fetch_parameter(parameter))


def _build_model(name: str) -> Model:
    """Build the pydantic-ai model registered under ``name``.

    The registry is an SSM String parameter (MODELS_PARAMETER): each entry names
    a provider and model id, and Bedrock entries carry the inference-profile ARN.
    Bedrock authenticates via the Lambda role; Mistral reads an API key from a
    SecureString. Memoised per name, so the lookup is one GetParameter per
    container.
    Bedrock authenticates via the Lambda role; Mistral and Fireworks read an
    API key from a SecureString. Deliberately not memoised: the Mistral client
    owns an httpx connection pool whose primitives belong to the event loop
    that first used them, and an eval sweep runs concurrent cases on one loop
    per worker thread. Only the registry lookup, the part worth caching, is.
    """
    registry = json.loads(fetch_parameter(require_env("MODELS_PARAMETER")))
    config = registry[name]
    config = _registry()[name]
    provider = config["provider"]
    model_id = config["model_id"]
    if provider == "bedrock":
        return BedrockConverseModel(
            config["model_id"],
            model_id,
            settings={"bedrock_inference_profile": config["inference_profile_arn"]},
        )
    if provider == "mistral":
        key_param = os.environ.get("MISTRAL_API_KEY_PARAMETER")
    elif provider == "fireworks":
        key_param = os.environ.get(
            "FIREWORKS_API_KEY_PARAMETER", "/terraform-pr-agent/fireworks-api-key"
        )
        if not key_param:
            raise RuntimeError(
                f"model {name!r} uses the Fireworks API, but FIREWORKS_API_KEY_PARAMETER "
                "is not set. Set FIREWORKS_API_KEY and re-apply so the key is wired, or "
                "select a Bedrock model via DEFAULT_MODEL or the event's model field."
            )
        return OpenAIChatModel(
            model_id, provider=FireworksProvider(api_key=fetch_parameter(key_param))
        )
    elif provider == "mistral":
        key_param = os.environ.get(
            "MISTRAL_API_KEY_PARAMETER", "/terraform-pr-agent/mistral-api-key"
        )
        if not key_param:
            raise RuntimeError(
                f"model {name!r} uses the Mistral API, but MISTRAL_API_KEY_PARAMETER "
                "is not set. Set MISTRAL_API_KEY and re-apply so the key is wired, or "
                "select a Bedrock model via DEFAULT_MODEL or the event's model field."
            )
        return MistralModel(
            config["model_id"],
            provider=MistralProvider(
                api_key=fetch_parameter(key_param),
                http_client=_retrying_http_client(),
            ),
        client = Mistral(
            api_key=fetch_parameter(key_param),
            async_client=_retrying_http_client(),
        )
        # The Mistral SDK sends a second `chat` span per request when a global
        # tracer provider exists, so the audit copy and the token metrics count
        # twice. Its env knob stops only the SDK exporter, not the spans.
        configure_telemetry(client, provider=NoOpTracerProvider())
        return MistralModel(config["model_id"], provider=MistralProvider(mistral_client=client))
    raise ValueError(f"unknown provider {provider!r} for model {name!r}")