feat(agent): add ua module for outbound SDK User-Agent solution tracking

scottschreckengaust · claude · scottschreckengaust · commit d126f1c13045 · 2026-06-12T23:35:37.000Z
app/uksb-wt64nei4u6/{STACKNAME} + md/uksb-wt64nei4u6#agent[#{TRACE}], emitted via botocore's verbatim user_agent_extra path (the sanitizing app-id field would mangle the '/' separator). Static part baked at construction; optional #{TRACE} appended per-request by a before-send handler so cached clients keep their connection pool. Wire-capture tests assert the emitted header via a short-circuiting before-send stub. Task 1 of PR #338 plan. Part of #319 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
diff --git a/agent/src/ua.py b/agent/src/ua.py
@@ -0,0 +1,137 @@
+"""Outbound AWS SDK User-Agent solution tracking (#319).
+
+Every AWS API call made by the agent carries two ABCA solution-tracking
+segments in the ``User-Agent`` header:
+
+    app/uksb-wt64nei4u6/{STACKNAME}            (only when ABCA_STACK_NAME set)
+    md/uksb-wt64nei4u6#agent[#{TRACE}]
+
+Both are emitted via botocore's *verbatim* ``user_agent_extra`` path — NOT
+the sanitizing ``user_agent_appid`` config field, whose allowed charset
+excludes ``/`` and would mangle the ``uksb-wt64nei4u6/`` separator into
+``-``. Because the raw path applies no sanitization, this module sanitizes
+``{STACKNAME}`` and ``{TRACE}`` itself (non-UA-token chars → ``-``).
+
+The static part is baked once at client/session construction. The optional
+``#{TRACE}`` suffix (the current task id) is appended **per request** by a
+``before-send`` event handler that mutates only the outgoing header — never
+the client config — so the singleton session in :mod:`aws_session` keeps its
+connection pool across trace changes (see issue #319 "Connection sharing").
+
+Trace state is plain lock-guarded module state rather than a ``ContextVar``:
+the task runs on a thread spawned by ``server.py``, where per-thread
+``ContextVar`` propagation is exactly the trap documented at
+``server.py`` (workload-token plumbing). One agent process works one task.
+
+The TypeScript counterparts are ``cdk/src/handlers/shared/ua.ts`` and
+``cli/src/ua.ts`` — the solution id, wire format, and sanitization rules
+must stay identical across all three.
+"""
+
+from __future__ import annotations
+
+import os
+import string
+import threading
+from typing import Any
+
+# AWS solution-tracking id for ABCA. Also appears (deploy-time counterpart,
+# #292) in the CloudFormation stack description in ``cdk/src/main.ts`` and in
+# the TS mirrors of this module. Per-surface literal by design — see PR #338.
+SOLUTION_ID = "uksb-wt64nei4u6"
+
+# Stable per-component label: this surface IS the Python agent runtime.
+COMPONENT = "agent"
+
+# Deployed CloudFormation stack name, threaded in by CDK (AgentCore runtime
+# env / ECS container env). Absent in local dev — the app/ segment is then
+# omitted entirely.
+STACK_NAME_ENV = "ABCA_STACK_NAME"
+
+# The documented app-id budget is 50 chars on the value;
+# len("uksb-wt64nei4u6/") == 16, leaving 34 for the stack name.
+_STACK_NAME_MAX = 34
+
+# RFC 7230 token charset (the UA product-token alphabet). '/' and '#' are
+# deliberately NOT here — they are the structural separators of the scheme.
+_ALLOWED = frozenset(string.ascii_letters + string.digits + "!$%&'*+-.^_`|~")
+
+_trace_lock = threading.Lock()
+_trace: str | None = None
+
+
+def sanitize_ua_value(raw: str) -> str:
+    """Replace every non-UA-token char (incl. non-ASCII) with ``-``."""
+    return "".join(c if c in _ALLOWED else "-" for c in raw)
+
+
+def static_user_agent_extra() -> str:
+    """The static UA suffix baked at client/session construction.
+
+    ``app/{SOLUTION_ID}/{stack}`` (stack sanitized FIRST, then clipped to 34
+    so a replaced multi-byte char can't be re-split) followed by
+    ``md/{SOLUTION_ID}#{COMPONENT}``. Without a stack name only the ``md/``
+    segment is emitted — never a placeholder.
+    """
+    segments = []
+    stack_name = os.environ.get(STACK_NAME_ENV, "").strip()
+    if stack_name:
+        clipped = sanitize_ua_value(stack_name)[:_STACK_NAME_MAX]
+        segments.append(f"app/{SOLUTION_ID}/{clipped}")
+    segments.append(f"md/{SOLUTION_ID}#{COMPONENT}")
+    return " ".join(segments)
+
+
+def set_trace(handle: str | None) -> None:
+    """Set (or clear, with ``None``/empty) the ambient trace handle.
+
+    Called once per task with the task id (``aws_session.configure_session``).
+    The handle is stored raw and sanitized on read.
+    """
+    global _trace
+    with _trace_lock:
+        _trace = handle or None
+
+
+def get_trace() -> str | None:
+    """Current trace handle, sanitized to UA-token-safe ASCII, or ``None``."""
+    with _trace_lock:
+        raw = _trace
+    return sanitize_ua_value(raw) if raw else None
+
+
+def register_trace_appender(events: Any) -> None:
+    """Append ``#{TRACE}`` to the outgoing User-Agent on every request.
+
+    ``events`` is a botocore event emitter — either ``client.meta.events``
+    (single client) or a botocore session's emitter (propagates to every
+    client *and resource* derived from it). Registered on ``before-send`` so
+    it runs after botocore renders the header (``user_agent_extra`` is the
+    final component, so the suffix lands exactly on our ``md/`` segment) and
+    mutates only the header — the connection pool is untouched.
+    """
+
+    def _append_trace(request: Any, **_kwargs: Any) -> None:
+        trace = get_trace()
+        if not trace:
+            return
+        current = request.headers.get("User-Agent")
+        if current is None:
+            return
+        # Headers may surface as bytes depending on the transport path.
+        if isinstance(current, bytes):
+            current = current.decode("ascii", errors="replace")
+        request.headers["User-Agent"] = f"{current}#{trace}"
+
+    events.register("before-send.*", _append_trace, unique_id="abca-ua-trace")
+
+
+def client_config() -> Any:
+    """``botocore.config.Config`` carrying the static UA suffix.
+
+    For direct ``boto3.client(...)`` call sites that don't go through a
+    shared session (see ``aws_session.platform_client``).
+    """
+    from botocore.config import Config
+
+    return Config(user_agent_extra=static_user_agent_extra())
diff --git a/agent/tests/test_ua.py b/agent/tests/test_ua.py
@@ -0,0 +1,216 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: MIT-0
+
+"""Unit tests for the outbound SDK User-Agent solution tracking (ua)."""
+
+from __future__ import annotations
+
+import threading
+
+import pytest
+
+from ua import (
+    COMPONENT,
+    SOLUTION_ID,
+    STACK_NAME_ENV,
+    client_config,
+    get_trace,
+    register_trace_appender,
+    sanitize_ua_value,
+    set_trace,
+    static_user_agent_extra,
+)
+
+
+@pytest.fixture(autouse=True)
+def _reset(monkeypatch):
+    """Clear trace state and the stack-name env var between tests."""
+    monkeypatch.delenv(STACK_NAME_ENV, raising=False)
+    set_trace(None)
+    yield
+    set_trace(None)
+
+
+class TestSanitizeUaValue:
+    def test_passthrough_for_token_safe_chars(self):
+        assert sanitize_ua_value("backgroundagent-dev") == "backgroundagent-dev"
+        assert sanitize_ua_value("A1!$%&'*+-.^_`|~z") == "A1!$%&'*+-.^_`|~z"
+
+    def test_structural_separators_replaced(self):
+        # '/' and '#' are the structural separators of the UA scheme and are
+        # NOT in the UA token charset — both must become '-'.
+        assert sanitize_ua_value("a/b#c") == "a-b-c"
+
+    def test_non_ascii_replaced(self):
+        assert sanitize_ua_value("stäck") == "st-ck"
+        assert sanitize_ua_value("名前") == "--"
+
+    def test_whitespace_and_controls_replaced(self):
+        assert sanitize_ua_value("a b\tc\nd") == "a-b-c-d"
+
+    def test_empty(self):
+        assert sanitize_ua_value("") == ""
+
+
+class TestStaticUserAgentExtra:
+    def test_without_stack_name_omits_app_segment(self):
+        extra = static_user_agent_extra()
+        assert extra == f"md/{SOLUTION_ID}#{COMPONENT}"
+        assert "app/" not in extra
+
+    def test_with_stack_name(self, monkeypatch):
+        monkeypatch.setenv(STACK_NAME_ENV, "backgroundagent-dev")
+        extra = static_user_agent_extra()
+        assert extra == (f"app/{SOLUTION_ID}/backgroundagent-dev md/{SOLUTION_ID}#{COMPONENT}")
+
+    def test_stack_name_sanitized_then_clipped(self, monkeypatch):
+        # Sanitize FIRST, then clip to 34 — a multi-byte char near the cut
+        # must already be '-' before clipping.
+        hostile = "my/stack#nämé" + "x" * 40
+        monkeypatch.setenv(STACK_NAME_ENV, hostile)
+        extra = static_user_agent_extra()
+        app_value = extra.split(" ")[0].removeprefix("app/")
+        assert app_value.startswith(f"{SOLUTION_ID}/my-stack-n-m-")
+        # uksb-wt64nei4u6/ (16) + clipped stack (<=34) <= 50.
+        assert len(app_value) <= 50
+        stack_part = app_value.removeprefix(f"{SOLUTION_ID}/")
+        assert len(stack_part) == 34
+        assert "/" not in stack_part and "#" not in stack_part
+
+    def test_longest_realistic_stack_name_within_budget(self, monkeypatch):
+        # CloudFormation stack names max out at 128 chars [A-Za-z0-9-].
+        monkeypatch.setenv(STACK_NAME_ENV, "a" * 128)
+        app_value = static_user_agent_extra().split(" ")[0].removeprefix("app/")
+        assert len(app_value) == 50
+
+    def test_blank_stack_name_omits_app_segment(self, monkeypatch):
+        monkeypatch.setenv(STACK_NAME_ENV, "   ")
+        assert static_user_agent_extra() == f"md/{SOLUTION_ID}#{COMPONENT}"
+
+
+class TestTraceState:
+    def test_default_none(self):
+        assert get_trace() is None
+
+    def test_set_and_get(self):
+        set_trace("01KTVYABCDEF")
+        assert get_trace() == "01KTVYABCDEF"
+
+    def test_sanitized_on_read(self):
+        set_trace("trace/with#bad chars")
+        assert get_trace() == "trace-with-bad-chars"
+
+    def test_none_and_empty_clear(self):
+        set_trace("x")
+        set_trace(None)
+        assert get_trace() is None
+        set_trace("y")
+        set_trace("")
+        assert get_trace() is None
+
+    def test_thread_safe_set(self):
+        # Smoke test: concurrent set_trace calls must not corrupt state.
+        def _spin(val: str):
+            for _ in range(200):
+                set_trace(val)
+
+        threads = [threading.Thread(target=_spin, args=(f"t{i}",)) for i in range(4)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+        assert get_trace() in {"t0", "t1", "t2", "t3"}
+
+
+class TestClientConfig:
+    def test_config_carries_static_extra(self, monkeypatch):
+        monkeypatch.setenv(STACK_NAME_ENV, "mystack")
+        cfg = client_config()
+        assert cfg.user_agent_extra == static_user_agent_extra()
+
+
+class TestWireCapture:
+    """Capture the actual outbound User-Agent header at the wire layer.
+
+    Uses a real botocore client with fake credentials and a registered
+    ``before-send`` stub that short-circuits the HTTP send by returning a
+    canned AWSResponse — no network, no moto.
+    """
+
+    @pytest.fixture()
+    def capture(self, monkeypatch):
+        import boto3
+        from botocore.awsrequest import AWSResponse
+
+        monkeypatch.setenv(STACK_NAME_ENV, "backgroundagent-dev")
+
+        session = boto3.Session(
+            aws_access_key_id="testing",
+            aws_secret_access_key="testing",
+            region_name="us-east-1",
+        )
+        client = session.client("sts", config=client_config())
+        register_trace_appender(client.meta.events)
+
+        captured: list[str] = []
+
+        def _short_circuit(request, **kwargs):
+            # At the before-send stage the prepared request's header values
+            # can be bytes; normalize so assertions read naturally.
+            value = request.headers["User-Agent"]
+            captured.append(value.decode("ascii") if isinstance(value, bytes) else value)
+            body = (
+                b"<GetCallerIdentityResponse "
+                b'xmlns="https://sts.amazonaws.com/doc/2011-06-15/">'
+                b"<GetCallerIdentityResult><Arn>arn:aws:iam::123456789012:user/t</Arn>"
+                b"<UserId>AIDA</UserId><Account>123456789012</Account>"
+                b"</GetCallerIdentityResult></GetCallerIdentityResponse>"
+            )
+            return AWSResponse(url=request.url, status_code=200, headers={}, raw=_FakeRaw(body))
+
+        # register_last so it runs AFTER the trace appender (register order
+        # within the same event is what guarantees we see the final header).
+        client.meta.events.register_last("before-send.sts.GetCallerIdentity", _short_circuit)
+        return client, captured
+
+    def test_both_segments_intact_no_trace(self, capture):
+        client, captured = capture
+        client.get_caller_identity()
+        ua_header = captured[0]
+        # Literal '/' in the app segment survived (raw path, NOT app-id field).
+        assert f"app/{SOLUTION_ID}/backgroundagent-dev" in ua_header
+        # Trace-absent: md segment ends exactly at the component label.
+        assert ua_header.endswith(f"md/{SOLUTION_ID}#{COMPONENT}")
+        assert not ua_header.endswith("#")
+
+    def test_trace_appended_per_request_same_client(self, capture):
+        client, captured = capture
+        set_trace("01KTVYTRACE1")
+        client.get_caller_identity()
+        set_trace("01KTVYTRACE2")
+        client.get_caller_identity()
+        set_trace(None)
+        client.get_caller_identity()
+        assert captured[0].endswith(f"md/{SOLUTION_ID}#{COMPONENT}#01KTVYTRACE1")
+        assert captured[1].endswith(f"md/{SOLUTION_ID}#{COMPONENT}#01KTVYTRACE2")
+        assert captured[2].endswith(f"md/{SOLUTION_ID}#{COMPONENT}")
+
+    def test_trace_sanitized_at_wire(self, capture):
+        client, captured = capture
+        set_trace("evil/trace#☃ value")
+        client.get_caller_identity()
+        assert captured[0].endswith(f"md/{SOLUTION_ID}#{COMPONENT}#evil-trace---value")
+
+
+class _FakeRaw:
+    """Minimal raw-body shim for AWSResponse."""
+
+    def __init__(self, data: bytes):
+        self._data = data
+
+    def read(self, *args, **kwargs):
+        data, self._data = self._data, b""
+        return data
+
+    def stream(self, *args, **kwargs):  # pragma: no cover - botocore fallback
+        yield self.read()
