File tools and validate: the agent's feedback loop

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
from __future__ import annotations
12

13
import uuid
14
from pathlib import Path
15
from tempfile import TemporaryDirectory
16

17
from pydantic import BaseModel, ConfigDict
18
from pydantic_ai import Agent
19

20
from agent.env import require_env
21
from agent.models import _build_model
22
from agent.runs import _persist_run, _relock_providers
23
from agent.tools import (
24
    WorkspaceDeps,
25
    _validate,
26
    delete_file,
27
    edit_file,
28
    list_files,
29
    read_file,
30
    terraform_init,
31
    terraform_validate,
32
    write_file,
33
)
34

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

57
SYSTEM_PROMPT = (
58
    "You are the terraform-pr-agent. You operate on a Terraform workspace "
59
    "through file tools (list_files, read_file, write_file, edit_file, "
60
    "delete_file) and two terraform tools. Use the file tools to explore, "
61
    "write, and edit. Run terraform_init before your first validate and "
62
    "again whenever you add or change provider or module requirements. "
63
    "Call terraform_validate after you write or change files to confirm "
64
    "the workspace still parses; treat its output as feedback and edit "
65
    "until it is clean.\n\n" + _PROVIDER_PIN_RULE
66
)
67

68
# One Agent instance is reused across invocations: the tools reach the workspace
69
# through RunContext.deps, so each run_sync scopes them to a fresh WorkspaceDeps.
70
# The model carries no default; it is built from the registry at INVOKE and
71
# passed per run, so switching DEFAULT_MODEL needs no code change.
72
agent = Agent(
73
    deps_type=WorkspaceDeps,
74
    system_prompt=SYSTEM_PROMPT,
75
    tools=[
76
        list_files,
77
        read_file,
78
        write_file,
79
        edit_file,
80
        delete_file,
81
        terraform_init,
82
        terraform_validate,
83
    ],
84
    # Tools raise ModelRetry on failure; pydantic-ai ends the run once one tool
85
    # fails more than `retries` times in a row (a success resets the count). The
86
    # default of 1 would end the run on the second straight failing validate,
87
    # which is a normal part of the write-validate-edit loop, so the budget is
88
    # raised well past anything a converging run produces. The per-run turn cap
89
    # stays the runaway guard.
90
    retries=10,
91
)
92

93

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

99
_RETRY_PROMPT = (
100
    "terraform validate still reports errors after you finished. "
101
    "Fix them and validate again.\n\n{output}"
102
)
103

104

105
class ValidateDidNotConverge(RuntimeError):
106
    """terraform validate still failed after the caller-side retry budget."""
107

108

109
class RunResult(BaseModel):
110
    model_config = ConfigDict(frozen=True)
111

112
    run_id: str
113
    model: str
114
    output: str
115

116

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

120
    The run id is the trace's gen_ai.conversation.id and the runs/<run_id>/
121
    prefix, so one identifier joins trace, artifacts, and result. After the run
122
    the caller re-validates the workspace and feeds any failure back as a
123
    follow-up turn; a run that still fails after the retry budget raises, so the
124
    workspace ships under status error for debugging. ``model`` overrides
125
    DEFAULT_MODEL when given.
126
    """
127
    model_name = model or require_env("DEFAULT_MODEL")
128
    run_id = str(uuid.uuid4())
129
    with TemporaryDirectory(dir="/tmp") as workspace:
130
        root = Path(workspace)
131
        deps = WorkspaceDeps(root=root)
132
        try:
133
            built = _build_model(model_name)
134
            result = agent.run_sync(
135
                prompt,
136
                deps=deps,
137
                conversation_id=run_id,
138
                model=built,
139
                metadata={"model": model_name},
140
            )
141
            # The agent can report done while terraform validate still fails.
142
            # Re-validate ourselves and feed any error back as a follow-up turn,
143
            # reusing the run id and message history so each retry is an
144
            # invoke_agent span under the one invocation trace. Give up after the
145
            # budget and raise so the failure is honest rather than a clean run
146
            # over a broken workspace.
147
            ok, output = _validate(root)
148
            attempts = 0
149
            while not ok and attempts < _MAX_VALIDATE_RETRIES:
150
                attempts += 1
151
                result = agent.run_sync(
152
                    _RETRY_PROMPT.format(output=output),
153
                    deps=deps,
154
                    conversation_id=run_id,
155
                    model=built,
156
                    message_history=result.all_messages(),
157
                    metadata={"model": model_name},
158
                )
159
                ok, output = _validate(root)
160
            if not ok:
161
                raise ValidateDidNotConverge(output)
162
        except Exception as error:
163
            _persist_run(run_id, root, status="error", error=repr(error))
164
            raise
165
        _relock_providers(root)
166
        _persist_run(run_id, root, status="ok")
167
    return RunResult(run_id=run_id, model=model_name, output=str(result.output))

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

3
from __future__ import annotations
4

5
import os
6

7

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

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

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 __future__ import annotations
11

12
from typing import NotRequired, TypedDict
13

14
import logfire
15

16
from agent import observability
17
from agent.core import execute
18

19

20
class HandlerEvent(TypedDict):
21
    prompt: str
22
    model: NotRequired[str]
23

24

25
class HandlerResponse(TypedDict):
26
    status: str
27
    run_id: str
28
    model: str
29
    output: str
30

31

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

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

51

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

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

62

63
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
from __future__ import annotations
16

17
import resource
18
from collections.abc import Iterator
19
from contextlib import contextmanager
20
from pathlib import Path
21

22
import logfire
23

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

30

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

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

50

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

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

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

3
from __future__ import annotations
4

5
import json
6
import os
7
from functools import cache
8

9
from httpx import AsyncClient, HTTPStatusError, Response
10
from pydantic_ai.models import Model
11
from pydantic_ai.models.bedrock import BedrockConverseModel
12
from pydantic_ai.models.mistral import MistralModel
13
from pydantic_ai.providers.mistral import MistralProvider
14
from pydantic_ai.retries import AsyncTenacityTransport, RetryConfig, wait_retry_after
15
from tenacity import retry_if_exception_type, stop_after_attempt, wait_exponential
16

17
from agent.env import require_env
18
from agent.ssm import fetch_parameter
19

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

24

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

29

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

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

54

55
@cache
56
def _build_model(name: str) -> Model:
57
    """Build the pydantic-ai model registered under ``name``.
58

59
    The registry is an SSM String parameter (MODELS_PARAMETER): each entry names
60
    a provider and model id, and Bedrock entries carry the inference-profile ARN.
61
    Bedrock authenticates via the Lambda role; Mistral reads an API key from a
62
    SecureString. Memoised per name, so the lookup is one GetParameter per
63
    container.
64
    """
65
    registry = json.loads(fetch_parameter(require_env("MODELS_PARAMETER")))
66
    config = registry[name]
67
    provider = config["provider"]
68
    if provider == "bedrock":
69
        return BedrockConverseModel(
70
            config["model_id"],
71
            settings={"bedrock_inference_profile": config["inference_profile_arn"]},
72
        )
73
    if provider == "mistral":
74
        key_param = os.environ.get("MISTRAL_API_KEY_PARAMETER")
75
        if not key_param:
76
            raise RuntimeError(
77
                f"model {name!r} uses the Mistral API, but MISTRAL_API_KEY_PARAMETER "
78
                "is not set. Set MISTRAL_API_KEY and re-apply so the key is wired, or "
79
                "select a Bedrock model via DEFAULT_MODEL or the event's model field."
80
            )
81
        return MistralModel(
82
            config["model_id"],
83
            provider=MistralProvider(
84
                api_key=fetch_parameter(key_param),
85
                http_client=_retrying_http_client(),
86
            ),
87
        )
88
    raise ValueError(f"unknown provider {provider!r} for model {name!r}")

1
"""Telemetry pipeline: structured logs, the per-trace audit copy, and EMF metrics.
2

3
``configure()`` runs at INIT (see handler.py) so the tracer provider and the
4
audit processor exist before instrument_aws_lambda opens the first invocation
5
span. The audit copy ships from inside the processor when the trace's root span
6
ends, so the handler needs no flush logic.
7
"""
8

9
from __future__ import annotations
10

11
import json
12
import os
13
import threading
14
from collections.abc import Callable, Sequence
15

16
import boto3
17
import logfire
18
import structlog
19
from google.protobuf import json_format
20
from logfire.sampling import SamplingOptions
21
from opentelemetry.context import (
22
    _SUPPRESS_INSTRUMENTATION_KEY,
23
    attach,
24
    detach,
25
    set_value,
26
)
27
from opentelemetry.exporter.otlp.proto.common._internal.trace_encoder import (
28
    encode_spans,
29
)
30
from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor
31
from opentelemetry.trace.status import StatusCode
32

33
from agent.env import require_env
34
from agent.ssm import fetch_parameter
35

36
# JSON logs to stdout, which CloudWatch Logs ingests as-is. The same stream
37
# carries the EMF metric envelope (see _emit_emf), so one structured sink covers
38
# both logs and metrics.
39
structlog.configure(
40
    processors=[
41
        structlog.processors.add_log_level,
42
        structlog.processors.TimeStamper(fmt="iso"),
43
        structlog.processors.EventRenamer("message"),
44
        structlog.processors.JSONRenderer(),
45
    ],
46
    logger_factory=structlog.PrintLoggerFactory(),
47
    cache_logger_on_first_use=True,
48
)
49
log = structlog.get_logger()
50

51

52
class PerTraceAuditProcessor(SpanProcessor):
53
    """Buffer spans by trace_id, ship as one batch when the local root ends.
54

55
    The OTel SDK has no `OnTraceComplete` hook, so this implements it against the
56
    only signal available: `on_end` fires synchronously, and a span is this
57
    process's local root when it has no parent, or a remote parent (context
58
    propagated in, e.g. from API Gateway), which `SpanContext.is_remote` marks.
59
    Late children (ended on a transport thread after the root shipped) are
60
    dropped, mirroring logfire's tail sampler. See pydantic/logfire#1034.
61
    """
62

63
    def __init__(
64
        self,
65
        on_trace_complete: Callable[[Sequence[ReadableSpan]], None],
66
    ) -> None:
67
        self._on_trace_complete = on_trace_complete
68
        self._buffers: dict[int, list[ReadableSpan]] = {}
69
        self._shipped: set[int] = set()
70
        self._lock = threading.Lock()
71

72
    def on_end(self, span: ReadableSpan) -> None:
73
        if not (span.context and span.context.trace_flags.sampled):
74
            return
75
        trace_id = span.context.trace_id
76
        with self._lock:
77
            if trace_id in self._shipped:
78
                return
79
            self._buffers.setdefault(trace_id, []).append(span)
80
            if span.parent is not None and not span.parent.is_remote:
81
                return
82
            spans = self._buffers.pop(trace_id)
83
            self._shipped.add(trace_id)
84
        self._ship(spans)
85

86
    def force_flush(self, timeout_millis: int = 30000) -> bool:
87
        with self._lock:
88
            pending = list(self._buffers.values())
89
            self._shipped.update(self._buffers)
90
            self._buffers.clear()
91
        for spans in pending:
92
            self._ship(spans)
93
        return True
94

95
    def shutdown(self) -> None:
96
        self.force_flush()
97

98
    def _ship(self, spans: Sequence[ReadableSpan]) -> None:
99
        # Suppress instrumentation around the callback so an instrumented
100
        # boto3 client inside it does not emit a span that re-enters on_end.
101
        token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True))
102
        try:
103
            self._on_trace_complete(spans)
104
        finally:
105
            detach(token)
106

107

108

109

110
_firehose = boto3.client("firehose")
111
_DELIVERY_STREAM = require_env("FIREHOSE_DELIVERY_STREAM")
112

113

114
def _ship_trace(spans: Sequence[ReadableSpan]) -> None:
115
    """Serialise one trace as OTLP-JSON and ship it as a single Firehose record."""
116
    payload = json_format.MessageToJson(encode_spans(spans), indent=None) + "\n"
117
    _firehose.put_record(
118
        DeliveryStreamName=_DELIVERY_STREAM,
119
        Record={"Data": payload.encode("utf-8")},
120
    )
121

122

123

124

125
def _emf_record(span: ReadableSpan) -> dict:
126
    """Build the EMF envelope for one agent-run span.
127

128
    pydantic-ai records gen_ai.usage.* on the agent-run span as the run total,
129
    so a single read is the correct total. The model dimension is the registry
130
    key the handler passed as run metadata, read back so a Bedrock run and a
131
    Mistral run land on one set of widgets.
132
    """
133
    attributes = span.attributes or {}
134
    model = json.loads(attributes["metadata"]).get("model", "unknown")
135
    errored = span.status.status_code is StatusCode.ERROR
136
    return {
137
        "_aws": {
138
            "Timestamp": span.end_time // 1_000_000,
139
            "CloudWatchMetrics": [
140
                {
141
                    "Namespace": require_env("METRICS_NAMESPACE"),
142
                    "Dimensions": [["Model"]],
143
                    "Metrics": [
144
                        {"Name": "InputTokens", "Unit": "Count"},
145
                        {"Name": "OutputTokens", "Unit": "Count"},
146
                        {"Name": "CacheReadTokens", "Unit": "Count"},
147
                        {"Name": "CacheWriteTokens", "Unit": "Count"},
148
                        {"Name": "Latency", "Unit": "Milliseconds"},
149
                        {"Name": "Invocations", "Unit": "Count"},
150
                        {"Name": "Errors", "Unit": "Count"},
151
                    ],
152
                }
153
            ],
154
        },
155
        "Model": model,
156
        "InputTokens": attributes.get("gen_ai.usage.input_tokens", 0),
157
        "OutputTokens": attributes.get("gen_ai.usage.output_tokens", 0),
158
        # pydantic-ai sets these only when non-zero; providers without prompt
159
        # caching (the Mistral API) never report them, so default to 0.
160
        "CacheReadTokens": attributes.get("gen_ai.usage.cache_read.input_tokens", 0),
161
        "CacheWriteTokens": attributes.get("gen_ai.usage.cache_creation.input_tokens", 0),
162
        "Latency": (span.end_time - span.start_time) / 1_000_000,
163
        "Invocations": 1,
164
        "Errors": 1 if errored else 0,
165
    }
166

167

168
def _emit_emf(spans: Sequence[ReadableSpan]) -> None:
169
    """Emit one EMF metric line per agent run in the trace.
170

171
    Off Bedrock there are no AWS/Bedrock metrics, so the dashboard reads these.
172
    The trace root is the Lambda invocation span, so metrics come off the
173
    agent-run spans nested under it, found by the `metadata` attribute
174
    pydantic-ai stamps on every run (a caller-side retry holds several, so one
175
    line each). CloudWatch Logs extracts the metrics from the structured line.
176
    """
177
    for span in spans:
178
        if span.attributes and "metadata" in span.attributes:
179
            log.info("trace_metrics", **_emf_record(span))
180

181

182
def _on_trace_complete(spans: Sequence[ReadableSpan]) -> None:
183
    """Ship the audit copy, then emit metrics: one hook, two sinks."""
184
    _ship_trace(spans)
185
    _emit_emf(spans)
186

187

188

189

190
def _logfire_token() -> str | None:
191
    """Logfire token from SSM, or None when the integration is not wired."""
192
    name = os.environ.get("LOGFIRE_TOKEN_PARAMETER")
193
    return fetch_parameter(name) if name else None
194

195

196
def configure() -> None:
197
    """Wire logfire at INIT: register the audit processor and instrument pydantic-ai.
198

199
    head=1.0 / tail=None because this is an audit pipeline: every trace must
200
    reach S3, and volume is low (one trace per invocation). Splitting the rates
201
    (say 1% to Logfire, 100% to S3) is possible with an extra sampler.
202
    include_content=True keeps the audit copy useful for forensics; flip it to
203
    False if prompts ever carry PII or secrets. version=5 pins the GenAI span
204
    schema so the audit copy stays stable across pydantic-ai releases.
205
    """
206
    if token := _logfire_token():
207
        os.environ["LOGFIRE_TOKEN"] = token
208
    logfire.configure(
209
        send_to_logfire="if-token-present",
210
        sampling=SamplingOptions(head=1.0, tail=None),
211
        additional_span_processors=[PerTraceAuditProcessor(_on_trace_complete)],
212
    )
213
    logfire.instrument_pydantic_ai(version=5, include_content=True)

1
"""Park a run's workspace in S3, and re-lock providers before it ships."""
2

3
from __future__ import annotations
4

5
import json
6
import subprocess
7
from collections.abc import Iterator
8
from pathlib import Path
9

10
import boto3
11

12
from agent.env import require_env
13
from agent.memory import track_memory
14
from agent.observability import log
15

16

17
def _persist_run(run_id: str, workspace: Path, status: str, error: str | None = None) -> None:
18
    """Park the workspace and a minimal result marker under runs/<run_id>/.
19

20
    result.json carries only what the audit trace does not: prompt, output, and
21
    messages already live in the trace under the same conversation id, so copying
22
    them here would create a second source of truth.
23
    """
24
    bucket = require_env("RUNS_BUCKET")
25
    s3 = boto3.client("s3")
26
    for file in _workspace_files(workspace):
27
        key = f"runs/{run_id}/workspace/{file.relative_to(workspace)}"
28
        s3.put_object(Bucket=bucket, Key=key, Body=file.read_bytes())
29
    result = {"status": status} | ({"error": error} if error else {})
30
    s3.put_object(
31
        Bucket=bucket,
32
        Key=f"runs/{run_id}/result.json",
33
        Body=json.dumps(result).encode(),
34
    )
35

36

37
def _workspace_files(workspace: Path) -> Iterator[Path]:
38
    """Every file except .terraform/, which is init scratch plus the provider
39
    downloaded into /tmp, gigabytes of noise per run. The top-level
40
    .terraform.lock.hcl is the reproducibility record and stays.
41
    """
42
    for path in sorted(workspace.rglob("*")):
43
        if ".terraform" in path.relative_to(workspace).parts:
44
            continue
45
        if path.is_file():
46
            yield path
47

48

49

50

51
# init inside the arm64 Lambda locks only linux_arm64, so a reviewer or CI on
52
# another platform hits a checksum error. Re-lock the platforms they are likely
53
# to run before the workspace ships. One call per platform rather than one with
54
# three -platform flags: `providers lock` is additive, so each call merges its
55
# platform and leaves the others, and a fresh process per platform keeps the
56
# peak to one platform's footprint. Best-effort and independent: a failure on
57
# one is logged and the rest still run.
58
_LOCK_PLATFORMS = ("linux_amd64", "linux_arm64", "darwin_arm64")
59

60

61
def _relock_providers(workspace: Path) -> None:
62
    if not (workspace / ".terraform.lock.hcl").exists():
63
        return
64
    for platform in _LOCK_PLATFORMS:
65
        with track_memory(f"relock.{platform}"):
66
            result = subprocess.run(
67
                ["terraform", "providers", "lock", "-no-color", f"-platform={platform}"],
68
                cwd=workspace,
69
                capture_output=True,
70
                text=True,
71
            )
72
        if result.returncode != 0:
73
            log.warning(
74
                "provider re-lock failed",
75
                platform=platform,
76
                stdout=result.stdout,
77
                stderr=result.stderr,
78
            )

1
"""SSM Parameter Store access, shared by the model registry and the Logfire token fetch."""
2

3
from __future__ import annotations
4

5
import boto3
6

7

8
def fetch_parameter(name: str) -> str:
9
    """Read one parameter from SSM Parameter Store.
10

11
    WithDecryption is a no-op on a plain String, so this covers String and
12
    SecureString alike. A direct GetParameter call, because container-image
13
    Lambdas cannot attach the Parameters and Secrets extension layer the zip
14
    package used. The client is built per call so moto can intercept it in tests.
15
    """
16
    response = boto3.client("ssm").get_parameter(Name=name, WithDecryption=True)
17
    return response["Parameter"]["Value"]

1
"""Tool surface for the terraform-pr-agent.
2

3
Each function below is registered on the Agent in handler.py and exposed
4
to the model as a callable tool. The bodies are intentionally empty in
5
this scaffold: the post fills them in section by section, so the reader
6
can see the agent's behaviour change as each tool comes online.
7

8
All tools take ``RunContext[WorkspaceDeps]`` so the workspace root is
9
threaded through ``ctx.deps.root`` instead of read from a module global.
10
That keeps tests, evals, and multi-tenant runs from sharing state.
11
"""
12

13
from __future__ import annotations
14

15
import subprocess
16
from pathlib import Path
17

18
from pydantic import BaseModel, ConfigDict, Field
19
from pydantic_ai import ModelRetry, RunContext
20

21
from agent.memory import track_memory
22

23

24
class WorkspaceDeps(BaseModel):
25
    """Per-run dependencies threaded through ``RunContext``.
26

27
    ``root`` is the directory the agent is allowed to read, write, and
28
    validate inside. Tools resolve every path relative to it and reject
29
    anything that escapes the root, so the agent cannot reach outside
30
    the workspace via ``..`` or absolute paths.
31
    """
32

33
    model_config = ConfigDict(arbitrary_types_allowed=True)
34

35
    root: Path
36
    files_read: set[Path] = Field(default_factory=set)
37

38

39

40

41
def list_files(ctx: RunContext[WorkspaceDeps], path: str = ".") -> list[str]:
42
    """List files under ``path`` relative to the workspace root."""
43
    return [
44
        str(p.relative_to(ctx.deps.root)) for p in _resolve_absolute_folder(ctx, path).iterdir()
45
    ]
46

47

48
def read_file(ctx: RunContext[WorkspaceDeps], path: str) -> str:
49
    """Read the file at ``path`` and return its contents."""
50
    file = _resolve_absolute_file(ctx, path)
51
    if not file.exists():
52
        raise ModelRetry(f"File {path} does not exist.")
53
    with file.open() as f:
54
        ctx.deps.files_read.add(file)
55
        return f.read()
56

57

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

70

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

93

94

95

96
def delete_file(ctx: RunContext[WorkspaceDeps], path: str) -> None:
97
    """Delete the file at ``path``."""
98
    file = _resolve_absolute_file(ctx, path)
99
    if file not in ctx.deps.files_read:
100
        raise ModelRetry(f"File {path} was not read. If you want to delete it, then read it first.")
101
    file.unlink()
102

103

104
def terraform_init(ctx: RunContext[WorkspaceDeps]) -> str:
105
    """Run ``terraform init`` in the workspace.
106

107
    Required once before the first ``terraform_validate`` and again after
108
    provider or module requirements change.
109
    """
110
    with track_memory("terraform_init"):
111
        init = subprocess.run(
112
            ["terraform", "init", "-backend=false", "-input=false", "-no-color"],
113
            cwd=ctx.deps.root,
114
            capture_output=True,
115
            text=True,
116
        )
117
    if init.returncode != 0:
118
        raise ModelRetry(f"terraform init failed:\n{init.stdout}{init.stderr}")
119
    return "OK: terraform init completed."
120

121

122
def _validate(root: Path) -> tuple[bool, str]:
123
    """Run ``terraform validate`` in ``root``; return (passed, combined output).
124

125
    The agent reaches this through the tool below; the caller-side retry in
126
    handler.py calls it directly to re-check the workspace after the run.
127
    """
128
    with track_memory("terraform_validate"):
129
        result = subprocess.run(
130
            ["terraform", "validate", "-no-color"],
131
            cwd=root,
132
            capture_output=True,
133
            text=True,
134
        )
135
    return result.returncode == 0, f"{result.stdout}{result.stderr}"
136

137

138
def terraform_validate(ctx: RunContext[WorkspaceDeps]) -> str:
139
    """Run ``terraform validate`` in the workspace and return its output."""
140
    ok, output = _validate(ctx.deps.root)
141
    if ok:
142
        return "OK: terraform validate passed."
143
    raise ModelRetry(f"terraform validate failed:\n{output}")
144

145

146

147

148
def _resolve_absolute_folder(ctx: RunContext[WorkspaceDeps], path: str):
149
    if not (absolute_path := _resolve_absolute_path(ctx, path)).is_dir():
150
        raise ModelRetry(f"Path {path} must be a directory.")
151
    return absolute_path
152

153

154
def _resolve_absolute_file(ctx: RunContext[WorkspaceDeps], path: str):
155
    if not (absolute_path := _resolve_absolute_path(ctx, path)).is_file():
156
        raise ModelRetry(f"Path {path} must be a file.")
157
    return absolute_path
158

159

160
def _resolve_absolute_path(ctx: RunContext[WorkspaceDeps], path: str):
161
    root = ctx.deps.root.resolve()
162
    absolute_path = (root / Path(path)).resolve()
163
    if not absolute_path.is_relative_to(root):
164
        raise ModelRetry(
165
            f"Path {path} must be relative to the workspace root, "
166
            f"it cannot be absolute or walk up the directory tree."
167
        )
168
    return absolute_path

1
FROM public.ecr.aws/lambda/python:3.13
2

3
COPY handler.py ${LAMBDA_TASK_ROOT}/
4
CMD ["handler.handler"]

1
def handler(event, context):
2
    return {
3
        "status": "placeholder",
4
        "note": "run scripts/build-lambda.sh to build and deploy the real image",
5
    }

1
locals {
2
  cloudwatch_region = data.aws_region.current.region
3
  dashboard_name    = "terraform-pr-agent"
4
  lambda_name       = aws_lambda_function.agent.function_name
5

6
  # The model widgets read the EMF metrics the handler emits (namespace
7
  # local.metrics_namespace, dimensioned by Model), not AWS/Bedrock, so a
8
  # Bedrock model and a Mistral-API model land in the same widgets. One line
9
  # per registry model, built by iterating keys(local.models).
10
  model_keys = keys(local.models)
11
}
12

13
resource "aws_cloudwatch_dashboard" "agent" {
14
  dashboard_name = local.dashboard_name
15
  dashboard_body = jsonencode({
16
    widgets = [
17
      {
18
        type   = "text"
19
        x      = 0
20
        y      = 0
21
        width  = 24
22
        height = 2
23
        properties = {
24
          markdown = "## Lambda\nContainer-image function health: invocations and errors, end to end duration, cold start init duration (Lambda Insights emits `init_duration` only on a cold start), and the memory and `/tmp` footprint behind the `memory_size` and `ephemeral_storage` sizing."
25
        }
26
      },
27
      {
28
        type   = "metric"
29
        x      = 0
30
        y      = 2
31
        width  = 12
32
        height = 6
33
        properties = {
34
          title  = "Lambda invocations and errors"
35
          region = local.cloudwatch_region
36
          view   = "timeSeries"
37
          stat   = "Sum"
38
          period = 60
39
          metrics = [
40
            ["AWS/Lambda", "Invocations", "FunctionName", local.lambda_name, { label = "${local.lambda_name} / invocations" }],
41
            [".", "Errors", ".", ".", { label = "${local.lambda_name} / errors" }],
42
            [".", "Throttles", ".", ".", { label = "${local.lambda_name} / throttles" }],
43
          ]
44
        }
45
      },
46
      {
47
        type   = "metric"
48
        x      = 12
49
        y      = 2
50
        width  = 12
51
        height = 6
52
        properties = {
53
          title  = "Lambda duration (ms)"
54
          region = local.cloudwatch_region
55
          view   = "timeSeries"
56
          period = 60
57
          metrics = [
58
            ["AWS/Lambda", "Duration", "FunctionName", local.lambda_name, { label = "${local.lambda_name} / avg", stat = "Average" }],
59
            [".", ".", ".", ".", { label = "${local.lambda_name} / p99", stat = "p99" }],
60
          ]
61
        }
62
      },
63
      {
64
        type   = "metric"
65
        x      = 0
66
        y      = 8
67
        width  = 12
68
        height = 6
69
        properties = {
70
          title  = "Cold start init duration (ms)"
71
          region = local.cloudwatch_region
72
          view   = "timeSeries"
73
          period = 60
74
          metrics = [
75
            # Insights reports init_duration only when an init phase happened,
76
            # so points appear only on cold starts.
77
            ["LambdaInsights", "init_duration", "function_name", local.lambda_name, { label = "${local.lambda_name} / init avg (ms)", stat = "Average" }],
78
            [".", ".", ".", ".", { label = "${local.lambda_name} / init max (ms)", stat = "Maximum" }],
79
          ]
80
        }
81
      },
82
      {
83
        type   = "metric"
84
        x      = 12
85
        y      = 8
86
        width  = 12
87
        height = 6
88
        properties = {
89
          title  = "Memory used (MB)"
90
          region = local.cloudwatch_region
91
          view   = "timeSeries"
92
          period = 60
93
          metrics = [
94
            # used_memory_max is the cgroup figure (Max Memory Used), which
95
            # counts the reclaimable /tmp page cache and so reads ~2 GB while
96
            # real demand is under 1 GB. Backs the memory_size comment in
97
            # lambda.tf.
98
            ["LambdaInsights", "used_memory_max", "function_name", local.lambda_name, { label = "${local.lambda_name} / memory max (MB)", stat = "Maximum" }],
99
          ]
100
        }
101
      },
102
      {
103
        type   = "metric"
104
        x      = 0
105
        y      = 14
106
        width  = 12
107
        height = 6
108
        properties = {
109
          title  = "/tmp used (bytes)"
110
          region = local.cloudwatch_region
111
          view   = "timeSeries"
112
          period = 60
113
          metrics = [
114
            # tmp_used tracks the ~800 MB provider download into /tmp, behind
115
            # the 4 GB ephemeral_storage sizing in lambda.tf.
116
            ["LambdaInsights", "tmp_used", "function_name", local.lambda_name, { label = "${local.lambda_name} / tmp used (B)", stat = "Maximum" }],
117
          ]
118
        }
119
      },
120
      {
121
        type   = "text"
122
        x      = 0
123
        y      = 20
124
        width  = 24
125
        height = 2
126
        properties = {
127
          markdown = "## Model\nPer-model usage from the handler's EMF metrics (namespace `local.metrics_namespace`, dimensioned by `Model`), so Bedrock and Mistral-API models share one set of widgets: tokens, invocations and errors, latency, and cache reads and writes."
128
        }
129
      },
130
      {
131
        type   = "metric"
132
        x      = 0
133
        y      = 22
134
        width  = 12
135
        height = 6
136
        properties = {
137
          title  = "Tokens"
138
          region = local.cloudwatch_region
139
          view   = "timeSeries"
140
          stat   = "Sum"
141
          period = 60
142
          metrics = [
143
            for m in flatten([
144
              for key in local.model_keys : [
145
                { metric = "InputTokens", label = "${key} / input", key = key },
146
                { metric = "OutputTokens", label = "${key} / output", key = key },
147
              ]
148
            ]) : [local.metrics_namespace, m.metric, "Model", m.key, { label = m.label }]
149
          ]
150
        }
151
      },
152
      {
153
        type   = "metric"
154
        x      = 12
155
        y      = 22
156
        width  = 12
157
        height = 6
158
        properties = {
159
          title  = "Invocations and errors"
160
          region = local.cloudwatch_region
161
          view   = "timeSeries"
162
          stat   = "Sum"
163
          period = 60
164
          metrics = [
165
            for m in flatten([
166
              for key in local.model_keys : [
167
                { metric = "Invocations", label = "${key} / invocations", key = key },
168
                { metric = "Errors", label = "${key} / errors", key = key },
169
              ]
170
            ]) : [local.metrics_namespace, m.metric, "Model", m.key, { label = m.label }]
171
          ]
172
        }
173
      },
174
      {
175
        type   = "metric"
176
        x      = 0
177
        y      = 28
178
        width  = 12
179
        height = 6
180
        properties = {
181
          title  = "Latency (ms)"
182
          region = local.cloudwatch_region
183
          view   = "timeSeries"
184
          period = 60
185
          metrics = [
186
            for m in flatten([
187
              for key in local.model_keys : [
188
                { label = "${key} / avg", key = key, stat = "Average" },
189
                { label = "${key} / p99", key = key, stat = "p99" },
190
              ]
191
            ]) : [local.metrics_namespace, "Latency", "Model", m.key, { label = m.label, stat = m.stat }]
192
          ]
193
        }
194
      },
195
      {
196
        type   = "metric"
197
        x      = 12
198
        y      = 28
199
        width  = 12
200
        height = 6
201
        properties = {
202
          title  = "Cache tokens"
203
          region = local.cloudwatch_region
204
          view   = "timeSeries"
205
          stat   = "Sum"
206
          period = 60
207
          metrics = [
208
            for m in flatten([
209
              for key in local.model_keys : [
210
                { metric = "CacheReadTokens", label = "${key} / cache read", key = key },
211
                { metric = "CacheWriteTokens", label = "${key} / cache write", key = key },
212
              ]
213
            ]) : [local.metrics_namespace, m.metric, "Model", m.key, { label = m.label }]
214
          ]
215
        }
216
      },
217
    ]
218
  })
219
}
220

221
output "cloudwatch_dashboard_url" {
222
  value = format(
223
    "https://%s.console.aws.amazon.com/cloudwatch/home?region=%s#dashboards/dashboard/%s",
224
    local.cloudwatch_region,
225
    local.cloudwatch_region,
226
    aws_cloudwatch_dashboard.agent.dashboard_name,
227
  )
228
}