feat(agent): wire solution UA into aws_session and all direct boto3 sites

Session-level user_agent_extra on both the scoped (refreshable) and
plain singleton sessions covers every tenant_client/tenant_resource
caller; new platform_client() helper carries the UA + trace appender
for the ambient-chain call sites (logs x5, secretsmanager x2,
bedrock-agentcore x1) that bypass the session by design.
configure_session() doubles the task id as the UA trace handle.

The trace appender splices #{TRACE} onto the md/ segment (not the
header end) because boto3 renders 'Botocore/x.y.z' after the
session-level extra.

Task 2 of PR #338 plan. Part of #319

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: scottschreckengaust

agent/src/aws_session.py

@@ -94,15 +94,23 @@ def configure_session(user_id: str, repo: str, task_id: str) -> None:
         for key, value in (("user_id", user_id), ("repo", repo), ("task_id", task_id))
         if value
     }
+    # The task id doubles as the UA trace handle (#319): every AWS call made
+    # while this task runs carries md/...#agent#{task_id}.
+    import ua
+
+    ua.set_trace(task_id or None)
 
 
 def reset_session_cache() -> None:
-    """Drop the cached session and tags. For tests that toggle config."""
+    """Drop the cached session, tags, and UA trace. For tests that toggle config."""
     global _session, _scoped, _tags
     with _lock:
         _session = None
         _scoped = None
         _tags = {}
+    import ua
+
+    ua.set_trace(None)
 
 
 def _session_tags() -> list[dict[str, str]]:
@@ -128,6 +136,8 @@ def _build_scoped_session(role_arn: str) -> Any:
     )
     from botocore.session import get_session as get_botocore_session
 
+    import ua
+
     region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
     task_id = _tags.get("task_id", "")
     # Role session name must be <=64 chars and match [\w+=,.@-]. task_id is a
@@ -139,7 +149,8 @@ def _build_scoped_session(role_arn: str) -> Any:
     # A dedicated STS client built from the *ambient* (compute-role) chain.
     # This is the role-chaining caller; the assumed SessionRole credentials it
     # returns must NOT be used to build it, or refresh would recurse.
-    sts_client = boto3.client("sts", region_name=region)
+    sts_client = boto3.client("sts", region_name=region, config=ua.client_config())
+    ua.register_trace_appender(sts_client.meta.events)
 
     def _refresh() -> dict[str, str]:
         resp = sts_client.assume_role(
@@ -167,6 +178,12 @@ def _refresh() -> dict[str, str]:
     )
     if region:
         botocore_session.set_config_variable("region", region)
+    # Outbound UA solution tracking (#319): session-level so every client and
+    # resource derived from this singleton carries the static segments; the
+    # per-request #{TRACE} appender mutates only the header, preserving the
+    # session's connection pool across trace changes.
+    botocore_session.user_agent_extra = ua.static_user_agent_extra()
+    ua.register_trace_appender(botocore_session.get_component("event_emitter"))
     return boto3.Session(botocore_session=botocore_session)
 
 
@@ -209,10 +226,19 @@ def get_session() -> Any:
                 ) from exc
         else:
             # Scoping not requested (local/dev/tests, or pre-provisioning):
-            # plain ambient session, behaviorally identical to pre-feature code.
-            _session = boto3.Session(
-                region_name=os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
-            )
+            # plain ambient session, behaviorally identical to pre-feature code
+            # apart from the UA solution-tracking segments (#319).
+            from botocore.session import get_session as get_botocore_session
+
+            import ua
+
+            botocore_session = get_botocore_session()
+            region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
+            if region:
+                botocore_session.set_config_variable("region", region)
+            botocore_session.user_agent_extra = ua.static_user_agent_extra()
+            ua.register_trace_appender(botocore_session.get_component("event_emitter"))
+            _session = boto3.Session(botocore_session=botocore_session)
             _scoped = False
         return _session
 
@@ -235,9 +261,7 @@ def tenant_client(service_name: str, **kwargs: Any) -> Any:
     session = get_session()
     if is_scoped():
         return session.client(service_name, **kwargs)
-    import boto3
-
-    return boto3.client(service_name, **kwargs)
+    return platform_client(service_name, **kwargs)
 
 
 def tenant_resource(service_name: str, **kwargs: Any) -> Any:
@@ -247,4 +271,43 @@ def tenant_resource(service_name: str, **kwargs: Any) -> Any:
         return session.resource(service_name, **kwargs)
     import boto3
 
-    return boto3.resource(service_name, **kwargs)
+    resource = boto3.resource(service_name, **_with_ua(kwargs))
+    # Guarded like platform_client: test doubles may lack the meta chain.
+    inner = getattr(getattr(resource, "meta", None), "client", None)
+    events = getattr(getattr(inner, "meta", None), "events", None)
+    if events is not None:
+        import ua
+
+        ua.register_trace_appender(events)
+    return resource
+
+
+def platform_client(service_name: str, **kwargs: Any) -> Any:
+    """boto3 client for platform (non-tenant) calls, with the ABCA UA (#319).
+
+    For call sites that intentionally use the ambient compute-role chain
+    (CloudWatch Logs debug writers, Secrets Manager, AgentCore memory) rather
+    than the tenant-scoped session. Same signature as ``boto3.client`` plus
+    the solution-tracking User-Agent and per-request trace appender.
+    """
+    import boto3
+
+    client = boto3.client(service_name, **_with_ua(kwargs))
+    # Real clients always expose meta.events; test doubles (MagicMock, or the
+    # bare fakes some suites install as a stub boto3 module) may not — the
+    # appender is solution telemetry, never worth failing a call site over.
+    events = getattr(getattr(client, "meta", None), "events", None)
+    if events is not None:
+        import ua
+
+        ua.register_trace_appender(events)
+    return client
+
+
+def _with_ua(kwargs: dict[str, Any]) -> dict[str, Any]:
+    """Merge the ABCA UA config into a boto3 client/resource kwargs dict."""
+    import ua
+
+    supplied = kwargs.get("config")
+    config = supplied.merge(ua.client_config()) if supplied is not None else ua.client_config()
+    return {**kwargs, "config": config}

agent/src/config.py

@@ -40,10 +40,10 @@ def resolve_github_token() -> str:
         return cached
     secret_arn = os.environ.get("GITHUB_TOKEN_SECRET_ARN")
     if secret_arn:
-        import boto3
+        from aws_session import platform_client
 
         region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
-        client = boto3.client("secretsmanager", region_name=region)
+        client = platform_client("secretsmanager", region_name=region)
         resp = client.get_secret_value(SecretId=secret_arn)
         token = resp["SecretString"]
         # Cache in env so downstream tools (git, gh CLI) work unchanged
@@ -101,13 +101,15 @@ def resolve_linear_api_token(channel_metadata: dict[str, str] | None = None) ->
         import json
         from datetime import datetime, timedelta
 
-        import boto3
+        import boto3  # noqa: F401 — availability probe; client built via platform_client
         from botocore.exceptions import BotoCoreError, ClientError
     except ImportError as e:
         log("WARN", f"resolve_linear_api_token: boto3 unavailable ({e}); skipping")
         return ""
 
-    sm = boto3.client("secretsmanager", region_name=region)
+    from aws_session import platform_client
+
+    sm = platform_client("secretsmanager", region_name=region)
 
     def _fetch_token() -> dict | None:
         """Fetch + parse the per-workspace OAuth secret.

agent/src/memory.py

@@ -35,12 +35,12 @@ def _get_client():
     global _client
     if _client is not None:
         return _client
-    import boto3
+    from aws_session import platform_client
 
     region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
     if not region:
         raise ValueError("AWS_REGION or AWS_DEFAULT_REGION must be set for memory operations")
-    _client = boto3.client("bedrock-agentcore", region_name=region)
+    _client = platform_client("bedrock-agentcore", region_name=region)
     return _client
 
 

agent/src/server.py

@@ -151,10 +151,10 @@ def _warn_cw_write_blocking(log_group: str, task_id: str | None, stamped: str) -
     covers both writers.
     """
     try:
-        import boto3
+        from aws_session import platform_client
 
         region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
-        client = boto3.client("logs", region_name=region)
+        client = platform_client("logs", region_name=region)
 
         stream = f"server_warn/{task_id or 'server'}"
         with _ctx_for_debug.suppress(client.exceptions.ResourceAlreadyExistsException):
@@ -178,10 +178,10 @@ def _warn_cw_write_blocking(log_group: str, task_id: str | None, stamped: str) -
 def _debug_cw_write_blocking(log_group: str, task_id: str | None, stamped: str) -> None:
     """Blocking CloudWatch write — only called from a background thread."""
     try:
-        import boto3
+        from aws_session import platform_client
 
         region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
-        client = boto3.client("logs", region_name=region)
+        client = platform_client("logs", region_name=region)
 
         stream = f"server_debug/{task_id or 'server'}"
         with _ctx_for_debug.suppress(client.exceptions.ResourceAlreadyExistsException):

agent/src/shell.py

@@ -61,10 +61,10 @@ def _log_error_cw_blocking(log_group: str, task_id: str | None, stamped: str) ->
     fire on the absence of the expected stream, not on this helper).
     """
     try:
-        import boto3
+        from aws_session import platform_client
 
         region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
-        client = boto3.client("logs", region_name=region)
+        client = platform_client("logs", region_name=region)
         stream = f"agent_error/{task_id or 'unknown'}"
         with contextlib.suppress(client.exceptions.ResourceAlreadyExistsException):
             client.create_log_stream(logGroupName=log_group, logStreamName=stream)

agent/src/telemetry.py

@@ -56,10 +56,10 @@ def _emit_metrics_to_cloudwatch(json_payload: dict) -> None:
     try:
         import contextlib
 
-        import boto3
+        from aws_session import platform_client
 
         region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
-        client = boto3.client("logs", region_name=region)
+        client = platform_client("logs", region_name=region)
 
         task_id = json_payload.get("task_id", "unknown")
         log_stream = f"metrics/{task_id}"
@@ -164,10 +164,10 @@ def _ensure_client(self):
 
         import contextlib
 
-        import boto3
+        from aws_session import platform_client
 
         region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
-        self._client = boto3.client("logs", region_name=region)
+        self._client = platform_client("logs", region_name=region)
 
         log_stream = f"trajectory/{self._task_id}"
         with contextlib.suppress(self._client.exceptions.ResourceAlreadyExistsException):

agent/src/ua.py

@@ -59,6 +59,10 @@
 _trace_lock = threading.Lock()
 _trace: str | None = None
 
+# The static md/ segment the per-request appender extends. Computed once —
+# COMPONENT never varies at runtime.
+_MD_SEGMENT = f"md/{SOLUTION_ID}#{COMPONENT}"
+
 
 def sanitize_ua_value(raw: str) -> str:
     """Replace every non-UA-token char (incl. non-ASCII) with ``-``."""
@@ -105,10 +109,14 @@ def register_trace_appender(events: Any) -> None:
 
     ``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.
+    client *and resource* derived from it). Registered on ``before-send``,
+    after botocore has rendered the header; only the header string changes —
+    the connection pool is untouched.
+
+    The trace is spliced onto the ``md/`` segment rather than appended to the
+    header's end: for *resources*, boto3 sets a client-level
+    ``user_agent_extra='Resource'`` marker that renders after the
+    session-level extra, so our segment is not always last.
     """
 
     def _append_trace(request: Any, **_kwargs: Any) -> None:
@@ -119,9 +127,13 @@ def _append_trace(request: Any, **_kwargs: Any) -> None:
         if current is None:
             return
         # Headers may surface as bytes depending on the transport path.
-        if isinstance(current, bytes):
+        was_bytes = isinstance(current, bytes)
+        if was_bytes:
             current = current.decode("ascii", errors="replace")
-        request.headers["User-Agent"] = f"{current}#{trace}"
+        if _MD_SEGMENT not in current:
+            return
+        updated = current.replace(_MD_SEGMENT, f"{_MD_SEGMENT}#{trace}", 1)
+        request.headers["User-Agent"] = updated.encode("ascii") if was_bytes else updated
 
     events.register("before-send.*", _append_trace, unique_id="abca-ua-trace")
 

agent/tests/test_aws_session.py

@@ -298,3 +298,93 @@ def test_overlong_value_truncated_to_256(self, monkeypatch):
         assert len(tags["repo"]) == _MAX_TAG_VALUE_LEN == 256
         # Untruncated values are passed through unchanged.
         assert tags["user_id"] == "u-1"
+
+
+# ---------------------------------------------------------------------------
+# Outbound UA solution tracking (#319)
+# ---------------------------------------------------------------------------
+
+
+class TestUserAgentWiring:
+    def test_configure_session_sets_ua_trace(self, monkeypatch):
+        import ua
+
+        configure_session(user_id="u-1", repo="owner/repo", task_id="01KTVYTASK")
+        assert ua.get_trace() == "01KTVYTASK"
+
+    def test_reset_session_cache_clears_trace(self, monkeypatch):
+        import ua
+
+        configure_session(user_id="u-1", repo="owner/repo", task_id="t-abc")
+        reset_session_cache()
+        assert ua.get_trace() is None
+
+    def test_plain_session_emits_solution_ua(self, monkeypatch):
+        """End-to-end wire capture through the real unscoped session path:
+        the singleton session must bake the static segments and append the
+        per-request #{TRACE} without rebuilding the client."""
+        from botocore.awsrequest import AWSResponse
+
+        import ua
+
+        monkeypatch.setenv("AWS_REGION", "us-east-1")
+        monkeypatch.setenv(ua.STACK_NAME_ENV, "backgroundagent-dev")
+        monkeypatch.setenv("AWS_ACCESS_KEY_ID", "testing")
+        monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "testing")
+
+        session = get_session()
+        client = session.client("sts")
+
+        captured: list[str] = []
+
+        def _short_circuit(request, **kwargs):
+            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>"
+            )
+
+            class _Raw:
+                def __init__(self, data):
+                    self._data = data
+
+                def read(self, *a, **k):
+                    data, self._data = self._data, b""
+                    return data
+
+                def stream(self, *a, **k):
+                    yield self.read()
+
+            return AWSResponse(url=request.url, status_code=200, headers={}, raw=_Raw(body))
+
+        client.meta.events.register_last("before-send.sts.GetCallerIdentity", _short_circuit)
+
+        ua.set_trace("trace-one")
+        client.get_caller_identity()
+        ua.set_trace("trace-two")
+        client.get_caller_identity()
+
+        assert f"app/{ua.SOLUTION_ID}/backgroundagent-dev" in captured[0]
+        # boto3 appends "Botocore/x.y.z" AFTER the session-level extra, so the
+        # segment is mid-header — which is exactly why the appender splices
+        # onto the md/ segment instead of appending to the header's end.
+        assert f"md/{ua.SOLUTION_ID}#agent#trace-one " in captured[0] + " "
+        assert f"md/{ua.SOLUTION_ID}#agent#trace-two " in captured[1] + " "
+
+    def test_tenant_resource_unscoped_carries_ua(self, monkeypatch):
+        """The unscoped resource path bypasses the session; it must still
+        carry the static UA via the merged client config."""
+        import ua
+
+        monkeypatch.setenv("AWS_REGION", "us-east-1")
+        monkeypatch.setenv("AWS_ACCESS_KEY_ID", "testing")
+        monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "testing")
+        from aws_session import tenant_resource
+
+        resource = tenant_resource("dynamodb", region_name="us-east-1")
+        extra = resource.meta.client.meta.config.user_agent_extra
+        assert f"md/{ua.SOLUTION_ID}#agent" in extra