Audit trail via OTLP: every agent run as a trace

1
"""AWS Lambda handler for the terraform-pr-agent.
2

3
First invocation:
4
- Reads the Logfire token from SSM via the Parameters and Secrets Lambda
5
  Extension at http://localhost:2773 when LOGFIRE_TOKEN_PARAMETER is set
6
  in the env; sets LOGFIRE_TOKEN so the logfire SDK picks it up.
7
- Configures logfire with send_to_logfire="if-token-present" and registers
8
  PerTraceAuditProcessor, a custom OTel SpanProcessor that buffers spans
9
  by trace_id and ships one OTLP-JSON Firehose record when the trace's
10
  root span ends.
11
- Instruments pydantic-ai with version=5 for the spec-compliant span
12
  names (invoke_agent, chat, execute_tool) and current GenAI semantic
13
  conventions. Pinned explicitly so the schema readers see in the
14
  audit copy stays stable across pydantic-ai releases.
15

16
The extension's HTTP server rejects requests during INIT with a
17
"not ready to serve traffic" 400, so this work runs on the first
18
INVOKE and is memoised with @cache for subsequent warm invocations.
19

20
Handler:
21
- Reads `prompt` from the invocation event.
22
- Runs the agent synchronously and returns the structured output.
23

24
The audit copy lands in Firehose from inside the processor's on_end
25
when the agent root span closes, so the handler has no flush logic.
26
A Firehose-side failure raises on the same thread as agent.run_sync
27
and propagates as a Lambda 5xx; the system-of-record copy is never
28
silently dropped.
29
"""
30

31
import json
32
import os
33
import threading
34
import urllib.parse
35
import urllib.request
36
from collections.abc import Callable, Sequence
37
from functools import cache
38
from typing import NotRequired, TypedDict
39

40
import boto3
41
import logfire
42
from google.protobuf import json_format
43
from logfire.sampling import SamplingOptions
44
from opentelemetry.context import (
45
    _SUPPRESS_INSTRUMENTATION_KEY,
46
    attach,
47
    detach,
48
    set_value,
49
)
50
from opentelemetry.exporter.otlp.proto.common._internal.trace_encoder import (
51
    encode_spans,
52
)
53
from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor
54
from pydantic_ai import Agent
55
from pydantic_ai.models.bedrock import BedrockConverseModel
56

57

58
def _fetch_logfire_token() -> str | None:
59
    """Read the Logfire token from the Parameters and Secrets extension.
60

61
    Returns None when no token parameter is configured, so the function
62
    runs fine without the Logfire integration. The extension caches the
63
    value across invocations, so this is cheap on warm starts.
64
    """
65
    parameter_name = os.environ.get("LOGFIRE_TOKEN_PARAMETER")
66
    if not parameter_name:
67
        return None
68

69
    session_token = os.environ["AWS_SESSION_TOKEN"]
70

71
    # safe="" forces urllib to percent-encode the leading and embedded
72
    # slashes in a hierarchical parameter name (e.g. /a/b/c -> %2Fa%2Fb%2Fc).
73
    # The Parameters and Secrets extension rejects unencoded slashes with
74
    # HTTP 400; AWS' own Python sample shows the same %2F-encoded form.
75
    url = (
76
        "http://localhost:2773/systemsmanager/parameters/get"
77
        f"?name={urllib.parse.quote(parameter_name, safe='')}&withDecryption=true"
78
    )
79
    req = urllib.request.Request(
80
        url,
81
        headers={"X-Aws-Parameters-Secrets-Token": session_token},
82
    )
83
    with urllib.request.urlopen(req, timeout=2) as resp:
84
        payload = json.load(resp)
85
    return payload["Parameter"]["Value"]
86

87

88

89

90
class PerTraceAuditProcessor(SpanProcessor):
91
    """Buffer spans by trace_id, ship as one batch when the root ends.
92

93
    The OTel SDK has no `OnTraceComplete` hook, so this implements it
94
    against the only signal available: `on_end` fires synchronously and
95
    `span.parent is None` on a root. Late children (spans ended on a
96
    transport thread after the root has already shipped) are dropped,
97
    mirroring logfire's tail sampler. See pydantic/logfire#1034.
98
    """
99

100
    def __init__(
101
        self,
102
        on_trace_complete: Callable[[Sequence[ReadableSpan]], None],
103
    ) -> None:
104
        self._on_trace_complete = on_trace_complete
105
        self._buffers: dict[int, list[ReadableSpan]] = {}
106
        self._shipped: set[int] = set()
107
        self._lock = threading.Lock()
108

109
    def on_end(self, span: ReadableSpan) -> None:
110
        if not (span.context and span.context.trace_flags.sampled):
111
            return
112
        trace_id = span.context.trace_id
113
        with self._lock:
114
            if trace_id in self._shipped:
115
                return
116
            self._buffers.setdefault(trace_id, []).append(span)
117
            if span.parent is not None:
118
                return
119
            spans = self._buffers.pop(trace_id)
120
            self._shipped.add(trace_id)
121
        self._ship(spans)
122

123
    def force_flush(self, timeout_millis: int = 30000) -> bool:
124
        with self._lock:
125
            pending = list(self._buffers.values())
126
            self._shipped.update(self._buffers)
127
            self._buffers.clear()
128
        for spans in pending:
129
            self._ship(spans)
130
        return True
131

132
    def shutdown(self) -> None:
133
        self.force_flush()
134

135
    def _ship(self, spans: Sequence[ReadableSpan]) -> None:
136
        # Suppress instrumentation around the callback so an instrumented
137
        # boto3/requests client inside it does not emit a span that
138
        # re-enters on_end for a sibling trace.
139
        token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True))
140
        try:
141
            self._on_trace_complete(spans)
142
        finally:
143
            detach(token)
144

145

146

147

148
_firehose = boto3.client("firehose")
149
_DELIVERY_STREAM = os.environ["FIREHOSE_DELIVERY_STREAM"]
150

151

152
def _ship_trace(spans: Sequence[ReadableSpan]) -> None:
153
    """Serialise one trace as OTLP-JSON and ship it as a single Firehose record."""
154
    payload = json_format.MessageToJson(encode_spans(spans), indent=None) + "\n"
155
    _firehose.put_record(
156
        DeliveryStreamName=_DELIVERY_STREAM,
157
        Record={"Data": payload.encode("utf-8")},
158
    )
159

160

161

162

163
@cache
164
def _init_logfire() -> None:
165
    """Wire Logfire and the audit span processor once per warm
166
    container, on the first INVOKE.
167

168
    The Parameters and Secrets extension is not ready to serve traffic
169
    during the Lambda INIT phase, so the token fetch (and the matching
170
    logfire setup) cannot run at module import time. @cache memoises
171
    on the empty argument tuple, so this runs exactly once per
172
    container and is a no-op on every subsequent invocation.
173
    """
174
    token = _fetch_logfire_token()
175
    if token:
176
        os.environ["LOGFIRE_TOKEN"] = token
177
    # head=1.0 and tail=None are today's Logfire defaults; pinned here
178
    # because this is an audit pipeline, so every trace must reach S3.
179
    # Volume is low (one trace per Lambda invocation) and the audit
180
    # requirement outweighs Logfire ingest cost. Splitting the rates
181
    # (e.g. 1% to Logfire, 100% to S3) is possible with a small extra
182
    # sampler; see the post.
183
    logfire.configure(
184
        send_to_logfire="if-token-present",
185
        sampling=SamplingOptions(head=1.0, tail=None),
186
        additional_span_processors=[PerTraceAuditProcessor(_ship_trace)],
187
    )
188
    # include_content=True is the pydantic-ai default; pinned because the
189
    # audit copy needs the actual prompts, tool args, and responses to be
190
    # useful for after-the-fact forensics. If that ever becomes a
191
    # compliance problem (PII, secrets in prompts), flip to False and
192
    # accept a metadata-only audit trail.
193
    logfire.instrument_pydantic_ai(version=5, include_content=True)
194

195

196

197

198
_INFERENCE_PROFILE_ARN = os.environ["BEDROCK_INFERENCE_PROFILE_ARN"]
199
_MODEL_ID = os.environ["BEDROCK_MODEL_ID"]
200

201

202
def _build_model() -> BedrockConverseModel:
203
    return BedrockConverseModel(
204
        _MODEL_ID,
205
        settings={"bedrock_inference_profile": _INFERENCE_PROFILE_ARN},
206
    )
207

208

209

210

211
SYSTEM_PROMPT = (
212
    "You are the terraform-pr-agent. For now you are a placeholder; "
213
    "respond briefly to whatever prompt you are given."
214
)
215

216

217
agent = Agent(_build_model(), system_prompt=SYSTEM_PROMPT)
218

219

220
class HandlerEvent(TypedDict):
221
    prompt: NotRequired[str]
222

223

224
class HandlerResponse(TypedDict):
225
    status: str
226
    output: str
227

228

229
def handler(event: HandlerEvent, context: object) -> HandlerResponse:
230
    """Lambda entry point.
231

232
    Falls back to a default prompt so the function can be smoke-tested
233
    with an empty payload.
234

235
    The audit copy ships from inside PerTraceAuditProcessor.on_end when
236
    the agent's root span closes, so the handler does not need a finally
237
    block: a Firehose failure raises on the same thread as agent.run_sync
238
    and propagates as a Lambda 5xx. A failed agent run also closes its
239
    root span (with status=ERROR) before the exception unwinds, so the
240
    partial trace still ships.
241
    """
242
    _init_logfire()
243
    prompt = event.get("prompt", "Say hello.")
244
    result = agent.run_sync(prompt)
245
    return {
246
        "status": "ok",
247
        "output": str(result.output),
248
    }

1
# alias/aws/sns is the AWS-managed SNS key; a CMK is unnecessary for an
2
# operational alarm topic that publishes alarm-fired events, not audit data.
3
#trivy:ignore:avd-aws-0136
4
resource "aws_sns_topic" "agent_alerts" {
5
  name              = "terraform-pr-agent-alerts"
6
  kms_master_key_id = "alias/aws/sns"
7
}
8

9
# The subscription stays PENDING until the recipient clicks the AWS
10
# confirmation email. Alarms fire either way, but no email goes out until the
11
# subscription is confirmed.
12
resource "aws_sns_topic_subscription" "agent_alerts_email" {
13
  topic_arn = aws_sns_topic.agent_alerts.arn
14
  protocol  = "email"
15
  endpoint  = var.alert_email
16
}
17

18
resource "aws_cloudwatch_metric_alarm" "daily_tokens" {
19
  alarm_name          = "${aws_bedrock_inference_profile.agent.name}-daily-tokens"
20
  alarm_description   = "Daily input + output token usage for the agent inference profile exceeded the configured threshold."
21
  comparison_operator = "GreaterThanThreshold"
22
  evaluation_periods  = 1
23
  threshold           = var.daily_token_alarm_threshold
24
  treat_missing_data  = "notBreaching"
25

26
  metric_query {
27
    id          = "total"
28
    expression  = "input + output"
29
    label       = "Total tokens (input + output)"
30
    return_data = true
31
  }
32

33
  metric_query {
34
    id = "input"
35
    metric {
36
      namespace   = "AWS/Bedrock"
37
      metric_name = "InputTokenCount"
38
      dimensions  = { ModelId = local.profile_id }
39
      stat        = "Sum"
40
      period      = 86400
41
    }
42
  }
43

44
  metric_query {
45
    id = "output"
46
    metric {
47
      namespace   = "AWS/Bedrock"
48
      metric_name = "OutputTokenCount"
49
      dimensions  = { ModelId = local.profile_id }
50
      stat        = "Sum"
51
      period      = 86400
52
    }
53
  }
54

55
  alarm_actions = [aws_sns_topic.agent_alerts.arn]
56
}

1
locals {
2
  audit_bucket_name = "terraform-pr-agent-audit-${data.aws_caller_identity.current.account_id}-${data.aws_region.current.region}"
3
}
4

5
# Access logging would require a second bucket and is out of scope: the audit
6
# copy here is the system of record for what the agent did, not for who read it.
7
#trivy:ignore:avd-aws-0089
8
resource "aws_s3_bucket" "audit" {
9
  bucket              = local.audit_bucket_name
10
  object_lock_enabled = true
11
}
12

13
resource "aws_s3_bucket_versioning" "audit" {
14
  bucket = aws_s3_bucket.audit.id
15

16
  versioning_configuration {
17
    status = "Enabled"
18
  }
19
}
20

21
resource "aws_s3_bucket_public_access_block" "audit" {
22
  bucket                  = aws_s3_bucket.audit.id
23
  block_public_acls       = true
24
  block_public_policy     = true
25
  ignore_public_acls      = true
26
  restrict_public_buckets = true
27
}
28

29
resource "aws_s3_bucket_server_side_encryption_configuration" "audit" {
30
  bucket = aws_s3_bucket.audit.id
31

32
  rule {
33
    apply_server_side_encryption_by_default {
34
      sse_algorithm     = "aws:kms"
35
      kms_master_key_id = aws_kms_key.audit.arn
36
    }
37

38
    bucket_key_enabled = true
39
  }
40
}
41

42
data "aws_iam_policy_document" "audit_bucket_resource_policy" {
43
  statement {
44
    sid     = "DenyInsecureTransport"
45
    effect  = "Deny"
46
    actions = ["s3:*"]
47
    resources = [
48
      aws_s3_bucket.audit.arn,
49
      "${aws_s3_bucket.audit.arn}/*",
50
    ]
51
    principals {
52
      type        = "*"
53
      identifiers = ["*"]
54
    }
55
    condition {
56
      test     = "Bool"
57
      variable = "aws:SecureTransport"
58
      values   = ["false"]
59
    }
60
  }
61
}
62

63
resource "aws_s3_bucket_policy" "audit" {
64
  bucket = aws_s3_bucket.audit.id
65
  policy = data.aws_iam_policy_document.audit_bucket_resource_policy.json
66
}
67

68
# GOVERNANCE leaves the bucket deletable for the tutorial; switch to "COMPLIANCE" in production.
69
resource "aws_s3_bucket_object_lock_configuration" "audit" {
70
  bucket = aws_s3_bucket.audit.id
71

72
  rule {
73
    default_retention {
74
      mode = "GOVERNANCE"
75
      days = var.audit_retention_days
76
    }
77
  }
78
}
79

80
resource "aws_s3_bucket_lifecycle_configuration" "audit" {
81
  bucket = aws_s3_bucket.audit.id
82

83
  rule {
84
    id     = "transition-to-glacier-ir"
85
    status = "Enabled"
86

87
    filter {}
88

89
    transition {
90
      days          = 90
91
      storage_class = "GLACIER_IR"
92
    }
93

94
    noncurrent_version_transition {
95
      noncurrent_days = 90
96
      storage_class   = "GLACIER_IR"
97
    }
98
  }
99

100
  rule {
101
    id     = "expire-after-retention"
102
    status = "Enabled"
103

104
    filter {}
105

106
    expiration {
107
      days = var.audit_retention_days
108
    }
109

110
    noncurrent_version_expiration {
111
      noncurrent_days = 1
112
    }
113
  }
114
}
115

116
output "audit_bucket_name" {
117
  value = aws_s3_bucket.audit.id
118
}
119

120
output "audit_bucket_arn" {
121
  value = aws_s3_bucket.audit.arn
122
}

1
locals {
2
  bedrock_model_id            = "anthropic.claude-haiku-4-5-20251001-v1:0"
3
  bedrock_cross_region_prefix = "eu"
4

5
  system_inference_profile_arn = format(
6
    "arn:aws:bedrock:%s:%s:inference-profile/%s.%s",
7
    data.aws_region.current.region,
8
    data.aws_caller_identity.current.account_id,
9
    local.bedrock_cross_region_prefix,
10
    local.bedrock_model_id,
11
  )
12
}
13

14
resource "aws_bedrock_inference_profile" "agent" {
15
  name        = "terraform-pr-agent"
16
  description = "Application inference profile for the Terraform PR Agent series."
17

18
  model_source {
19
    copy_from = local.system_inference_profile_arn
20
  }
21
}
22

23
output "inference_profile_arn" {
24
  value = aws_bedrock_inference_profile.agent.arn
25
}

1
locals {
2
  cloudwatch_region = data.aws_region.current.region
3

4
  dashboard_label = aws_bedrock_inference_profile.agent.name
5
  profile_id      = aws_bedrock_inference_profile.agent.id
6
}
7

8
resource "aws_cloudwatch_dashboard" "agent" {
9
  dashboard_name = aws_bedrock_inference_profile.agent.name
10
  dashboard_body = jsonencode({
11
    widgets = [
12
      {
13
        type   = "metric"
14
        x      = 0
15
        y      = 0
16
        width  = 12
17
        height = 6
18
        properties = {
19
          title  = "Tokens"
20
          region = local.cloudwatch_region
21
          view   = "timeSeries"
22
          stat   = "Sum"
23
          period = 60
24
          metrics = [
25
            # here we plug in the labels so this looks nicer
26
            ["AWS/Bedrock", "InputTokenCount", "ModelId", local.profile_id, { label = "${local.dashboard_label} / input" }],
27
            # the dot syntax allows not to repeat the same as above
28
            [".", "OutputTokenCount", ".", ".", { label = "${local.dashboard_label} / output" }],
29
          ]
30
        }
31
      },
32
      {
33
        type   = "metric"
34
        x      = 12
35
        y      = 0
36
        width  = 12
37
        height = 6
38
        properties = {
39
          title  = "Cache tokens"
40
          region = local.cloudwatch_region
41
          view   = "timeSeries"
42
          stat   = "Sum"
43
          period = 60
44
          metrics = [
45
            ["AWS/Bedrock", "CacheReadInputTokenCount", "ModelId", local.profile_id, { label = "${local.dashboard_label} / cache read" }],
46
            [".", "CacheWriteInputTokenCount", ".", ".", { label = "${local.dashboard_label} / cache write" }],
47
          ]
48
        }
49
      },
50
      {
51
        type   = "metric"
52
        x      = 0
53
        y      = 6
54
        width  = 12
55
        height = 6
56
        properties = {
57
          title  = "Invocations and errors"
58
          region = local.cloudwatch_region
59
          view   = "timeSeries"
60
          stat   = "Sum"
61
          period = 60
62
          metrics = [
63
            ["AWS/Bedrock", "Invocations", "ModelId", local.profile_id, { label = "${local.dashboard_label} / invocations" }],
64
            [".", "InvocationClientErrors", ".", ".", { label = "${local.dashboard_label} / 4xx" }],
65
            [".", "InvocationServerErrors", ".", ".", { label = "${local.dashboard_label} / 5xx" }],
66
            [".", "InvocationThrottles", ".", ".", { label = "${local.dashboard_label} / throttles" }],
67
          ]
68
        }
69
      },
70
      {
71
        type   = "metric"
72
        x      = 12
73
        y      = 6
74
        width  = 12
75
        height = 6
76
        properties = {
77
          title  = "Latency (ms)"
78
          region = local.cloudwatch_region
79
          view   = "timeSeries"
80
          period = 60
81
          metrics = [
82
            ["AWS/Bedrock", "InvocationLatency", "ModelId", local.profile_id, { label = "${local.dashboard_label} / avg", stat = "Average" }],
83
            [".", ".", ".", ".", { label = "${local.dashboard_label} / p99", stat = "p99" }],
84
          ]
85
        }
86
      },
87
    ]
88
  })
89
}
90

91
output "cloudwatch_dashboard_url" {
92
  value = format(
93
    "https://%s.console.aws.amazon.com/cloudwatch/home?region=%s#dashboards/dashboard/%s",
94
    local.cloudwatch_region,
95
    local.cloudwatch_region,
96
    aws_cloudwatch_dashboard.agent.dashboard_name,
97
  )
98
}

1
data "aws_iam_policy_document" "firehose_assume" {
2
  statement {
3
    actions = ["sts:AssumeRole"]
4
    principals {
5
      type        = "Service"
6
      identifiers = ["firehose.amazonaws.com"]
7
    }
8
  }
9
}
10

11
resource "aws_iam_role" "firehose" {
12
  name               = "terraform-pr-agent-firehose"
13
  assume_role_policy = data.aws_iam_policy_document.firehose_assume.json
14
}
15

16
data "aws_iam_policy_document" "firehose_permissions" {
17
  statement {
18
    actions = [
19
      "s3:PutObject",
20
      "s3:PutObjectRetention",
21
      "s3:GetBucketLocation",
22
      "s3:ListBucket",
23
    ]
24
    # /* is how S3 IAM grants object-scoped actions; bounded to the audit bucket.
25
    #trivy:ignore:avd-aws-0057
26
    resources = [
27
      aws_s3_bucket.audit.arn,
28
      "${aws_s3_bucket.audit.arn}/*",
29
    ]
30
  }
31

32
  statement {
33
    actions = [
34
      "kms:GenerateDataKey",
35
      "kms:Decrypt",
36
    ]
37
    resources = [aws_kms_key.audit.arn]
38
  }
39

40
  statement {
41
    actions = ["logs:PutLogEvents"]
42
    # :* covers log streams inside the named Firehose log group only.
43
    #trivy:ignore:avd-aws-0057
44
    resources = ["${aws_cloudwatch_log_group.firehose.arn}:*"]
45
  }
46
}
47

48
resource "aws_iam_role_policy" "firehose_permissions" {
49
  name   = "terraform-pr-agent-firehose-permissions"
50
  role   = aws_iam_role.firehose.id
51
  policy = data.aws_iam_policy_document.firehose_permissions.json
52
}
53

54
# Operational delivery log for Firehose; carries error metadata, not audit
55
# payload. The audit copy in S3 is the protected artefact and already uses a CMK.
56
#trivy:ignore:avd-aws-0017
57
resource "aws_cloudwatch_log_group" "firehose" {
58
  name              = "/aws/kinesisfirehose/terraform-pr-agent-audit"
59
  retention_in_days = 7
60
}
61

62
resource "aws_cloudwatch_log_stream" "firehose_s3" {
63
  name           = "S3Delivery"
64
  log_group_name = aws_cloudwatch_log_group.firehose.name
65
}
66

67
resource "aws_kinesis_firehose_delivery_stream" "audit" {
68
  name        = "terraform-pr-agent-audit"
69
  destination = "extended_s3"
70

71
  extended_s3_configuration {
72
    role_arn            = aws_iam_role.firehose.arn
73
    bucket_arn          = aws_s3_bucket.audit.arn
74
    prefix              = "traces/year=!{timestamp:yyyy}/month=!{timestamp:MM}/day=!{timestamp:dd}/hour=!{timestamp:HH}/"
75
    error_output_prefix = "errors/!{firehose:error-output-type}/year=!{timestamp:yyyy}/month=!{timestamp:MM}/day=!{timestamp:dd}/"
76
    buffering_size      = 5
77
    buffering_interval  = 60
78
    compression_format  = "GZIP"
79

80
    cloudwatch_logging_options {
81
      enabled         = true
82
      log_group_name  = aws_cloudwatch_log_group.firehose.name
83
      log_stream_name = aws_cloudwatch_log_stream.firehose_s3.name
84
    }
85
  }
86
}
87

88
output "firehose_stream_name" {
89
  value = aws_kinesis_firehose_delivery_stream.audit.name
90
}
91

92
output "firehose_stream_arn" {
93
  value = aws_kinesis_firehose_delivery_stream.audit.arn
94
}

1
data "aws_iam_policy_document" "bedrock_invoke" {
2
  statement {
3
    actions = [
4
      "bedrock:Converse",
5
      "bedrock:ConverseStream",
6
      "bedrock:InvokeModel",
7
      "bedrock:InvokeModelWithResponseStream",
8
    ]
9
    # Bedrock foundation-model ARNs do not pin to the caller region; the
10
    # inference profile fans out cross-region, so the * region segment is required.
11
    #trivy:ignore:avd-aws-0057
12
    resources = [
13
      aws_bedrock_inference_profile.agent.arn,
14
      local.system_inference_profile_arn,
15
      "arn:aws:bedrock:*::foundation-model/${local.bedrock_model_id}",
16
    ]
17
  }
18

19
  # Anthropic models on Bedrock are distributed via AWS Marketplace. On the
20
  # first invocation from a new account, Bedrock auto-subscribes the account
21
  # to the model product, which requires the invoking principal to hold these
22
  # Marketplace actions. Once the subscription is active these calls become
23
  # no-ops, but the principal still needs ViewSubscriptions on every call so
24
  # Bedrock can confirm the subscription is in place.
25
  statement {
26
    actions = [
27
      "aws-marketplace:Subscribe",
28
      "aws-marketplace:Unsubscribe",
29
      "aws-marketplace:ViewSubscriptions",
30
    ]
31
    # Marketplace subscription actions are global by design.
32
    #trivy:ignore:avd-aws-0057
33
    resources = ["*"]
34
  }
35
}
36

37
resource "aws_iam_policy" "bedrock_invoke" {
38
  name        = "terraform-pr-agent-bedrock-invoke"
39
  description = "Invoke Claude via the terraform-pr-agent inference profile."
40
  policy      = data.aws_iam_policy_document.bedrock_invoke.json
41
}
42

43
data "aws_iam_policy_document" "agent_assume" {
44
  statement {
45
    actions = ["sts:AssumeRole"]
46
    principals {
47
      type        = "AWS"
48
      identifiers = ["arn:aws:iam::${data.aws_caller_identity.current.account_id}:root"]
49
    }
50
  }
51
}
52

53
resource "aws_iam_role" "agent" {
54
  name               = "terraform-pr-agent"
55
  assume_role_policy = data.aws_iam_policy_document.agent_assume.json
56
}
57

58
resource "aws_iam_role_policy_attachment" "agent_bedrock_invoke" {
59
  role       = aws_iam_role.agent.name
60
  policy_arn = aws_iam_policy.bedrock_invoke.arn
61
}
62

63
output "agent_role_arn" {
64
  value = aws_iam_role.agent.arn
65
}

1
resource "aws_kms_key" "audit" {
2
  description             = "Encrypts the terraform-pr-agent audit bucket."
3
  enable_key_rotation     = true
4
  deletion_window_in_days = 7
5
  policy                  = data.aws_iam_policy_document.audit_kms_key_resource_policy.json
6
}
7

8
resource "aws_kms_alias" "audit" {
9
  name          = "alias/terraform-pr-agent-audit"
10
  target_key_id = aws_kms_key.audit.key_id
11
}
12

13
data "aws_iam_policy_document" "audit_kms_key_resource_policy" {
14
  statement {
15
    sid       = "EnableIAMUserPermissions"
16
    actions   = ["kms:*"]
17
    resources = ["*"]
18
    principals {
19
      type        = "AWS"
20
      identifiers = ["arn:aws:iam::${data.aws_caller_identity.current.account_id}:root"]
21
    }
22
  }
23
}
24

25
output "audit_kms_key_arn" {
26
  value = aws_kms_key.audit.arn
27
}

1
data "aws_iam_policy_document" "lambda_assume" {
2
  statement {
3
    actions = ["sts:AssumeRole"]
4
    principals {
5
      type        = "Service"
6
      identifiers = ["lambda.amazonaws.com"]
7
    }
8
  }
9
}
10

11
# trivy:ignore:avd-aws-0057
12
# Bedrock foundation-model ARNs do not pin to the caller region (the inference
13
# profile fans out cross-region), and Marketplace subscription actions are
14
# global by design.
15
data "aws_iam_policy_document" "lambda_permissions" {
16
  # Bedrock invocation. Same shape as iam.tf's bedrock_invoke; copied
17
  # here so the Lambda role is self-contained and does not require
18
  # the user-role policy to also be attached to the Lambda role.
19
  statement {
20
    actions = [
21
      "bedrock:Converse",
22
      "bedrock:ConverseStream",
23
      "bedrock:InvokeModel",
24
      "bedrock:InvokeModelWithResponseStream",
25
    ]
26
    # Bedrock foundation-model ARNs do not pin to the caller region; the
27
    # inference profile fans out cross-region, so the * region segment is required.
28
    #trivy:ignore:avd-aws-0057
29
    resources = [
30
      aws_bedrock_inference_profile.agent.arn,
31
      local.system_inference_profile_arn,
32
      "arn:aws:bedrock:*::foundation-model/${local.bedrock_model_id}",
33
    ]
34
  }
35

36
  statement {
37
    actions = [
38
      "aws-marketplace:Subscribe",
39
      "aws-marketplace:Unsubscribe",
40
      "aws-marketplace:ViewSubscriptions",
41
    ]
42
    # Marketplace subscription actions are global by design.
43
    #trivy:ignore:avd-aws-0057
44
    resources = ["*"]
45
  }
46

47
  # SSM SecureString read for the Logfire token, only when wired.
48
  dynamic "statement" {
49
    for_each = local.logfire_token_wired ? [1] : []
50
    content {
51
      actions   = ["ssm:GetParameter"]
52
      resources = [aws_ssm_parameter.logfire_token[0].arn]
53
    }
54
  }
55

56
  # KMS Decrypt on the AWS-managed SSM key, required to read a
57
  # SecureString value through the extension. Only when wired.
58
  dynamic "statement" {
59
    for_each = local.logfire_token_wired ? [1] : []
60
    content {
61
      actions = ["kms:Decrypt"]
62
      resources = [
63
        "arn:aws:kms:${data.aws_region.current.region}:${data.aws_caller_identity.current.account_id}:alias/aws/ssm",
64
      ]
65
    }
66
  }
67

68
  statement {
69
    actions = [
70
      "firehose:PutRecord",
71
      "firehose:PutRecordBatch",
72
    ]
73
    resources = [aws_kinesis_firehose_delivery_stream.audit.arn]
74
  }
75
}
76

77
resource "aws_iam_role" "lambda" {
78
  name               = "terraform-pr-agent-lambda"
79
  assume_role_policy = data.aws_iam_policy_document.lambda_assume.json
80
}
81

82
resource "aws_iam_role_policy_attachment" "lambda_basic_execution" {
83
  role       = aws_iam_role.lambda.name
84
  policy_arn = "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
85
}
86

87
resource "aws_iam_role_policy" "lambda_permissions" {
88
  name   = "terraform-pr-agent-lambda-permissions"
89
  role   = aws_iam_role.lambda.id
90
  policy = data.aws_iam_policy_document.lambda_permissions.json
91
}
92

93
data "archive_file" "placeholder" {
94
  type        = "zip"
95
  output_path = "${path.module}/.terraform/placeholder.zip"
96

97
  source {
98
    content  = "def handler(event, context):\n    return {\"status\": \"placeholder\", \"note\": \"run scripts/build-lambda.sh then aws lambda update-function-code\"}\n"
99
    filename = "agent/handler.py"
100
  }
101

102
  source {
103
    content  = ""
104
    filename = "agent/__init__.py"
105
  }
106
}
107

108
resource "aws_lambda_function" "agent" {
109
  function_name = "terraform-pr-agent"
110
  role          = aws_iam_role.lambda.arn
111
  runtime       = "python3.13"
112
  architectures = ["arm64"]
113
  handler       = "agent.handler.handler"
114
  timeout       = 60
115
  memory_size   = 512
116

117
  filename         = data.archive_file.placeholder.output_path
118
  source_code_hash = data.archive_file.placeholder.output_base64sha256
119

120
  layers = [var.secrets_extension_layer_arn]
121

122
  tracing_config {
123
    mode = "Active"
124
  }
125

126
  environment {
127
    variables = merge(
128
      {
129
        BEDROCK_INFERENCE_PROFILE_ARN           = aws_bedrock_inference_profile.agent.arn
130
        BEDROCK_MODEL_ID                        = local.bedrock_model_id
131
        PARAMETERS_SECRETS_EXTENSION_CACHE_SIZE = "100"
132
        PARAMETERS_SECRETS_EXTENSION_HTTP_PORT  = "2773"
133
        FIREHOSE_DELIVERY_STREAM                = aws_kinesis_firehose_delivery_stream.audit.name
134
      },
135
      local.logfire_token_wired ? {
136
        LOGFIRE_TOKEN_PARAMETER = local.logfire_token_parameter_name
137
      } : {},
138
    )
139
  }
140

141
  lifecycle {
142
    ignore_changes = [
143
      filename,
144
      source_code_hash,
145
    ]
146
  }
147

148
  depends_on = [
149
    aws_iam_role_policy_attachment.lambda_basic_execution,
150
    aws_iam_role_policy.lambda_permissions,
151
  ]
152
}
153

154
output "lambda_function_name" {
155
  value = aws_lambda_function.agent.function_name
156
}
157

158
output "lambda_function_arn" {
159
  value = aws_lambda_function.agent.arn
160
}

1
locals {
2
  logfire_token_parameter_name = "/terraform-pr-agent/logfire-token"
3
  logfire_token_wired          = var.logfire_token != ""
4
}
5

6
resource "aws_ssm_parameter" "logfire_token" {
7
  count = local.logfire_token_wired ? 1 : 0
8

9
  name        = local.logfire_token_parameter_name
10
  description = "Logfire write token. Consumed by the terraform-pr-agent Lambda."
11
  type        = "SecureString"
12
  value       = var.logfire_token
13
}
14

15
output "logfire_token_parameter_arn" {
16
  # nonsensitive: the ARN is just a resource identifier, but Terraform's
17
  # taint analysis propagates the sensitive marker from var.logfire_token
18
  # via the resource. Explicit unmark so the ARN renders in `terraform
19
  # output` (the value itself never leaks the token).
20
  value       = nonsensitive(local.logfire_token_wired ? aws_ssm_parameter.logfire_token[0].arn : "")
21
  description = "ARN of the SSM SecureString; empty when LOGFIRE_TOKEN is not set."
22
}