Audit trail via OTLP: every agent run as a trace

1
[project]
2
name = "terraform-pr-agent"
3
version = "0.1.0"
4
description = "The terraform-pr-agent Lambda handler."
5
requires-python = ">=3.13"
6
dependencies = [
7
    "pydantic-ai-slim[bedrock]>=1.106",
8
    "logfire>=4.35",
9
    "boto3>=1.35",
10
]

1
#!/usr/bin/env bash
2
# Build a deployable zip for the terraform-pr-agent Lambda.
3
#
4
# Runs from anywhere; cd's to the project root (the dir holding
5
# pyproject.toml). Produces build/lambda.zip ready for:
6
#
7
#   aws lambda update-function-code \
8
#       --function-name terraform-pr-agent \
9
#       --zip-file fileb://build/lambda.zip
10
#
11
# Mirrors the pattern from
12
# https://docs.astral.sh/uv/guides/integration/aws-lambda/
13
set -euo pipefail
14

15
cd "$(dirname "$0")/.."
16

17
# Make sure the lock is in sync with pyproject.toml before exporting.
18
uv sync --quiet
19

20
rm -rf build
21
mkdir -p build/packages
22

23
# Export the locked dependency set so uv pip install can consume it
24
# without re-resolving.
25
uv export --frozen --no-dev --no-editable -o build/requirements.txt
26

27
# Install deps for Lambda's runtime. --python-platform forces wheels
28
# compatible with Amazon Linux 2 on arm64 (matches the function's
29
# architectures = ["arm64"]). --no-compile-bytecode keeps the zip small
30
# and avoids spending cold-start cycles on pyc creation.
31
uv pip install \
32
    --no-installer-metadata \
33
    --no-compile-bytecode \
34
    --python-platform aarch64-manylinux2014 \
35
    --python 3.13 \
36
    --target build/packages \
37
    -r build/requirements.txt
38

39
# Drop the handler package alongside the installed deps.
40
cp -r agent build/packages/
41

42
# Zip from inside the staging dir so paths sit at the zip root.
43
( cd build/packages && zip -qr ../lambda.zip . )
44

45
echo "built: $(pwd)/build/lambda.zip ($(du -h build/lambda.zip | cut -f1))"

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

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

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

59

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

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

71
    session_token = os.environ["AWS_SESSION_TOKEN"]
72

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

89

90

91

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

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

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

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

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

134
    def shutdown(self) -> None:
135
        self.force_flush()
136

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

147

148

149

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

153

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

162

163

164

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

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

197

198

199

200
_INFERENCE_PROFILE_ARN = os.environ["BEDROCK_INFERENCE_PROFILE_ARN"]
201
_MODEL_ID = os.environ["BEDROCK_MODEL_ID"]
202

203

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

210

211

212

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

218

219
agent = Agent(_build_model(), system_prompt=SYSTEM_PROMPT)
220

221

222
class HandlerEvent(TypedDict):
223
    prompt: NotRequired[str]
224

225

226
class HandlerResponse(TypedDict):
227
    status: str
228
    output: str
229

230

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

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

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

1
CREATE PERSISTENT SECRET (
2
    TYPE s3,
3
    PROVIDER credential_chain,
4
    REFRESH auto
5
);
6
SET VARIABLE audit_bucket = getenv('AUDIT_BUCKET');
7

8
-- hive_partitioning = true reads year=YYYY/month=MM/day=DD/ from the
9
-- object path as virtual columns, so the partition predicate in a
10
-- query below prunes objects before any file is opened.
11
CREATE OR REPLACE VIEW traces AS
12
WITH spans AS (
13
    -- Flatten the OTLP-JSON envelope into one row per span, with the
14
    -- common span fields lifted out as named columns so downstream
15
    -- CTEs and ad-hoc queries can work against `name`, `trace_id`,
16
    -- `dur_ms`, etc. without re-doing the struct navigation each time.
17
    SELECT
18
        year, month, day,
19
        span.name                                                                AS name,
20
        lower(hex(from_base64(span.traceId::VARCHAR)))                           AS trace_id,
21
        lower(hex(from_base64(span.spanId::VARCHAR)))                            AS span_id,
22
        lower(hex(from_base64(span.parentSpanId::VARCHAR)))                      AS parent_span_id,
23
        make_timestamp_ns(span.startTimeUnixNano::BIGINT)                        AS started,
24
        make_timestamp_ns(span.endTimeUnixNano::BIGINT)                          AS ended,
25
        (span.endTimeUnixNano::BIGINT - span.startTimeUnixNano::BIGINT) / 1e6    AS dur_ms,
26
        span.status.code::VARCHAR                                                AS status_code,
27
        span.attributes                                                          AS attributes,
28
        data.filename                                                            AS source_file,
29
        -- Firehose names objects <stream>-<ver>-<YYYY-MM-DD-HH-MM-SS>-<uuid>.gz;
30
        -- stripping the trailing -<uuid>.gz collapses rows from the same flush
31
        -- batch onto a stable key for grouping.
32
        regexp_replace(
33
            split_part(data.filename, '/', -1),
34
            '-[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\.gz$',
35
            ''
36
        )                                                                        AS batch_key
37
    FROM read_ndjson(
38
        's3://' || getvariable('audit_bucket') || '/traces/**/*.gz',
39
        compression = 'gzip', hive_partitioning = true, filename = true) AS data
40
       , UNNEST(data.resourceSpans) AS u1(rs)
41
       , UNNEST(rs.scopeSpans)      AS u2(ss)
42
       , UNNEST(ss.spans)           AS u3(span)
43
),
44
roots AS (
45
    -- One row per trace: the invoke_agent root span pydantic-ai emits per run.
46
    SELECT * FROM spans
47
    WHERE parent_span_id IS NULL OR parent_span_id = ''
48
),
49
chats AS (
50
    -- Per-trace summary of the LLM call: pulls the GenAI semantic
51
    -- convention attributes off the chat span and exposes each as a
52
    -- named column.
53
    SELECT
54
        trace_id,
55
        list_filter(attributes, x -> x.key = 'gen_ai.system')[1].value.stringValue                     AS gen_ai_system,
56
        list_filter(attributes, x -> x.key = 'gen_ai.operation.name')[1].value.stringValue             AS operation,
57
        list_filter(attributes, x -> x.key = 'gen_ai.request.model')[1].value.stringValue              AS request_model,
58
        list_filter(attributes, x -> x.key = 'gen_ai.response.model')[1].value.stringValue             AS response_model,
59
        list_filter(attributes, x -> x.key = 'gen_ai.usage.input_tokens')[1].value.intValue::BIGINT    AS in_tokens,
60
        list_filter(attributes, x -> x.key = 'gen_ai.usage.output_tokens')[1].value.intValue::BIGINT   AS out_tokens,
61
        list_filter(attributes, x -> x.key = 'gen_ai.response.finish_reasons')[1]
62
            .value.arrayValue.values[1].stringValue                                                    AS finish,
63
        list_filter(attributes, x -> x.key = 'gen_ai.conversation.id')[1].value.stringValue            AS conversation_id,
64
        list_filter(attributes, x -> x.key = 'gen_ai.agent.name')[1].value.stringValue                 AS agent_name,
65
        list_filter(attributes, x -> x.key = 'gen_ai.agent.call.id')[1].value.stringValue              AS agent_call_id,
66
        list_filter(attributes, x -> x.key = 'gen_ai.input.messages')[1].value.stringValue             AS input_messages,
67
        list_filter(attributes, x -> x.key = 'gen_ai.output.messages')[1].value.stringValue            AS output_messages
68
    FROM spans
69
    WHERE name LIKE 'chat %'
70
)
71
SELECT
72
    roots.started,
73
    roots.year, roots.month, roots.day,
74
    roots.trace_id,
75
    substr(roots.trace_id, 1, 8)                       AS trace,
76
    roots.batch_key,
77
    roots.source_file,
78
    roots.dur_ms,
79
    regexp_extract(chats.request_model, '[^./]+$')     AS model,
80
    chats.in_tokens,
81
    chats.out_tokens,
82
    chats.finish,
83
    chats.agent_name,
84
    chats.conversation_id,
85
    -- Convenience columns for the common single-turn shape: system at
86
    -- input[0], user at input[1], assistant at output[0]. Multi-turn
87
    -- runs invalidate the indices, so reach for input_messages and
88
    -- output_messages directly for those.
89
    json_extract_string(chats.input_messages,  '$[0].parts[0].content')  AS system_prompt,
90
    json_extract_string(chats.input_messages,  '$[1].parts[0].content')  AS user_prompt,
91
    json_extract_string(chats.output_messages, '$[0].parts[0].content')  AS assistant_response,
92
    chats.input_messages,
93
    chats.output_messages,
94
    -- Per the OTel spec, instrumentation libraries leave status unset
95
    -- on success (only application code may set it to Ok). Every OTel
96
    -- backend treats unset as "no error reported"; we render the same.
97
    CASE roots.status_code
98
        WHEN 'STATUS_CODE_ERROR' THEN 'err'
99
        ELSE 'ok'
100
    END                                                AS status
101
FROM roots
102
JOIN chats USING (trace_id);