diff --git a/AGENTS.md b/AGENTS.md index 0f6f2c0f..7e3bbfe1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -50,6 +50,7 @@ Handler entry tests: `cdk/test/handlers/orchestrate-task.test.ts`, `create-task. - **`prek install`** fails if Git **`core.hooksPath`** is set — another hook manager owns hooks; see [CONTRIBUTING.md](./CONTRIBUTING.md). - **Editing on `main` directly** — ALWAYS create a worktree with a feature branch for changes, even trivial ones. Main should stay clean; all work flows through worktree → branch → PR → merge. - **Git worktrees** — Always **`git fetch origin main`** before creating a new worktree to ensure you branch from the latest remote state. `node_modules/` and `agent/.venv/` are per-tree (not shared). Run **`mise run install`** in each new worktree before building. All CDK path references (`__dirname`-relative) and mise `config_roots` resolve correctly without extra setup. +- **Instantiating AWS SDK clients without the ABCA solution User-Agent (#319)** — every outbound AWS API call must carry the solution-tracking UA segments. Never write naked `boto3.client(...)` / `new XClient({})`: in `agent/src/` use `aws_session.tenant_client`/`tenant_resource` (tenant data) or `aws_session.platform_client` (ambient-chain calls); in `cdk/src/handlers/` pass `abcaUserAgent()` and wrap with `withAbcaTrace()` from `shared/ua.ts`; in `cli/src/` same via `cli/src/ua.ts`. The three `ua` modules (`agent/src/ua.py`, `cdk/src/handlers/shared/ua.ts`, `cli/src/ua.ts`) must stay identical in solution id, wire format, and sanitization. - **Bumping Cedar engines in isolation** — `cedarpy` (Python, `agent/pyproject.toml`) and `@cedar-policy/cedar-wasm` (TypeScript, `cdk/package.json`) are two language bindings over the same Cedar Rust core. They MUST move together; even patch-version drift between bindings can yield divergent `(decision, matching_rule_ids)` on the same `(policy, input)` — invisible to per-side unit tests, caught (only) by `contracts/cedar-parity/` golden fixtures in CI. If you bump one engine you MUST bump the other to a tested-compatible version AND refresh the parity fixtures in the same commit. Both pins are EXACT (no `^`/`~`). See `docs/design/CEDAR_HITL_GATES.md` §15.6 (decision #23) and the parity-contract banner in `mise.toml`. **DO NOT** accept upstream's "Update branch" or auto-merge suggestions on cedarpy without verifying parity with cedar-wasm. ### Tech stack diff --git a/agent/src/aws_session.py b/agent/src/aws_session.py index 2c6a906c..97e0ba41 100644 --- a/agent/src/aws_session.py +++ b/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} diff --git a/agent/src/config.py b/agent/src/config.py index c33dd6cb..f1249819 100644 --- a/agent/src/config.py +++ b/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,14 +101,16 @@ 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") # nosemgrep: py-silent-success-masking -- optional Linear MCP; boto3 unavailable 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. diff --git a/agent/src/memory.py b/agent/src/memory.py index 9d2654b2..aa89d1e0 100644 --- a/agent/src/memory.py +++ b/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 diff --git a/agent/src/server.py b/agent/src/server.py index d9ae1d7c..47f72fde 100644 --- a/agent/src/server.py +++ b/agent/src/server.py @@ -166,10 +166,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): @@ -193,10 +193,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): diff --git a/agent/src/shell.py b/agent/src/shell.py index d6dea355..bf0d7bdc 100644 --- a/agent/src/shell.py +++ b/agent/src/shell.py @@ -75,10 +75,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) diff --git a/agent/src/telemetry.py b/agent/src/telemetry.py index b91f2b4e..560daa7b 100644 --- a/agent/src/telemetry.py +++ b/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): diff --git a/agent/src/ua.py b/agent/src/ua.py new file mode 100644 index 00000000..88ea61b3 --- /dev/null +++ b/agent/src/ua.py @@ -0,0 +1,149 @@ +"""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 + +# 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 ``-``.""" + 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``, + 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: + 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. + was_bytes = isinstance(current, bytes) + if was_bytes: + current = current.decode("ascii", errors="replace") + 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") + + +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_aws_session.py b/agent/tests/test_aws_session.py index fd4e5562..bcda599c 100644 --- a/agent/tests/test_aws_session.py +++ b/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"' + b"arn:aws:iam::123456789012:user/t" + b"AIDA123456789012" + b"" + ) + + 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 diff --git a/agent/tests/test_ua.py b/agent/tests/test_ua.py new file mode 100644 index 00000000..a9703e2b --- /dev/null +++ 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"' + b"arn:aws:iam::123456789012:user/t" + b"AIDA123456789012" + b"" + ) + 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() diff --git a/cdk/src/constructs/approval-metrics-publisher-consumer.ts b/cdk/src/constructs/approval-metrics-publisher-consumer.ts index 01f75452..1596bd85 100644 --- a/cdk/src/constructs/approval-metrics-publisher-consumer.ts +++ b/cdk/src/constructs/approval-metrics-publisher-consumer.ts @@ -18,7 +18,7 @@ */ import * as path from 'path'; -import { Duration, RemovalPolicy } from 'aws-cdk-lib'; +import { Duration, RemovalPolicy, Stack } from 'aws-cdk-lib'; import * as dynamodb from 'aws-cdk-lib/aws-dynamodb'; import { FilterCriteria, FilterRule, StartingPosition, Architecture, Runtime } from 'aws-cdk-lib/aws-lambda'; import { DynamoEventSource, SqsDlq } from 'aws-cdk-lib/aws-lambda-event-sources'; @@ -123,6 +123,11 @@ export class ApprovalMetricsPublisherConsumer extends Construct { timeout: Duration.minutes(1), memorySize: 256, logGroup, + // Outbound SDK User-Agent solution tracking (#319). + environment: { + ABCA_STACK_NAME: Stack.of(this).stackName, + ABCA_COMPONENT: 'orchestr', + }, bundling: { externalModules: ['@aws-sdk/*'], }, diff --git a/cdk/src/constructs/concurrency-reconciler.ts b/cdk/src/constructs/concurrency-reconciler.ts index 703c7faf..15aa9089 100644 --- a/cdk/src/constructs/concurrency-reconciler.ts +++ b/cdk/src/constructs/concurrency-reconciler.ts @@ -18,7 +18,7 @@ */ import * as path from 'path'; -import { Duration } from 'aws-cdk-lib'; +import { Duration, Stack } from 'aws-cdk-lib'; import * as dynamodb from 'aws-cdk-lib/aws-dynamodb'; import * as events from 'aws-cdk-lib/aws-events'; import * as targets from 'aws-cdk-lib/aws-events-targets'; @@ -71,6 +71,8 @@ export class ConcurrencyReconciler extends Construct { environment: { TASK_TABLE_NAME: props.taskTable.tableName, USER_CONCURRENCY_TABLE_NAME: props.userConcurrencyTable.tableName, + ABCA_STACK_NAME: Stack.of(this).stackName, + ABCA_COMPONENT: 'orchestr', }, bundling: { externalModules: ['@aws-sdk/*'], diff --git a/cdk/src/constructs/ecs-agent-cluster.ts b/cdk/src/constructs/ecs-agent-cluster.ts index acb89518..68387561 100644 --- a/cdk/src/constructs/ecs-agent-cluster.ts +++ b/cdk/src/constructs/ecs-agent-cluster.ts @@ -123,6 +123,9 @@ export class EcsAgentCluster extends Construct { }), environment: { CLAUDE_CODE_USE_BEDROCK: '1', + // Outbound SDK User-Agent solution tracking (#319); component label + // ('agent') is hardcoded in agent/src/ua.py. + ABCA_STACK_NAME: Stack.of(this).stackName, TASK_TABLE_NAME: props.taskTable.tableName, TASK_EVENTS_TABLE_NAME: props.taskEventsTable.tableName, USER_CONCURRENCY_TABLE_NAME: props.userConcurrencyTable.tableName, diff --git a/cdk/src/constructs/fanout-consumer.ts b/cdk/src/constructs/fanout-consumer.ts index fce69a11..e24d1916 100644 --- a/cdk/src/constructs/fanout-consumer.ts +++ b/cdk/src/constructs/fanout-consumer.ts @@ -18,7 +18,7 @@ */ import * as path from 'path'; -import { Duration, RemovalPolicy } from 'aws-cdk-lib'; +import { Duration, RemovalPolicy, Stack } from 'aws-cdk-lib'; import * as cloudwatch from 'aws-cdk-lib/aws-cloudwatch'; import * as dynamodb from 'aws-cdk-lib/aws-dynamodb'; import * as iam from 'aws-cdk-lib/aws-iam'; @@ -164,6 +164,10 @@ export class FanOutConsumer extends Construct { }, }); + // Outbound SDK User-Agent solution tracking (#319). + this.fn.addEnvironment('ABCA_STACK_NAME', Stack.of(this).stackName); + this.fn.addEnvironment('ABCA_COMPONENT', 'orchestr'); + // GitHub dispatcher plumbing. Each grant/env var is guarded so the // fan-out plane still deploys cleanly in a dev environment that // hasn't onboarded the RepoTable or a platform GitHub token yet — diff --git a/cdk/src/constructs/github-screenshot-integration.ts b/cdk/src/constructs/github-screenshot-integration.ts index 4704f849..88042120 100644 --- a/cdk/src/constructs/github-screenshot-integration.ts +++ b/cdk/src/constructs/github-screenshot-integration.ts @@ -121,6 +121,13 @@ export class GitHubScreenshotIntegration extends Construct { const removalPolicy = props.removalPolicy ?? RemovalPolicy.DESTROY; + // Outbound SDK User-Agent solution tracking (#319), spread into every + // handler environment in this construct. + const abcaEnv: Record = { + ABCA_STACK_NAME: Stack.of(this).stackName, + ABCA_COMPONENT: 'webhook', + }; + // --- Screenshot bucket (private; served via CloudFront with OAC) --- this.screenshotBucket = new ScreenshotBucket(this, 'ScreenshotBucket', { removalPolicy, @@ -174,6 +181,7 @@ export class GitHubScreenshotIntegration extends Construct { memorySize: 512, deadLetterQueue: processorDlq, environment: { + ...abcaEnv, SCREENSHOT_BUCKET_NAME: this.screenshotBucket.bucket.bucketName, SCREENSHOT_PUBLIC_HOST: this.screenshotBucket.distribution.domainName, GITHUB_TOKEN_SECRET_ARN: props.githubTokenSecret.secretArn, @@ -266,6 +274,7 @@ export class GitHubScreenshotIntegration extends Construct { architecture: Architecture.ARM_64, timeout: Duration.seconds(10), environment: { + ...abcaEnv, GITHUB_WEBHOOK_SECRET_ARN: this.webhookSecret.secretArn, GITHUB_WEBHOOK_DEDUP_TABLE_NAME: this.webhookDedupTable.tableName, GITHUB_WEBHOOK_PROCESSOR_FUNCTION_NAME: this.webhookProcessorFn.functionName, diff --git a/cdk/src/constructs/linear-integration.ts b/cdk/src/constructs/linear-integration.ts index 3fa701f7..62968489 100644 --- a/cdk/src/constructs/linear-integration.ts +++ b/cdk/src/constructs/linear-integration.ts @@ -144,6 +144,13 @@ export class LinearIntegration extends Construct { }; // --- Task creation environment (matches TaskApi / SlackIntegration pattern) --- + // Outbound SDK User-Agent solution tracking (#319), spread into every + // handler environment in this construct. + const abcaEnv: Record = { + ABCA_STACK_NAME: Stack.of(this).stackName, + ABCA_COMPONENT: 'webhook', + }; + const createTaskEnv: Record = { TASK_TABLE_NAME: props.taskTable.tableName, TASK_EVENTS_TABLE_NAME: props.taskEventsTable.tableName, @@ -190,6 +197,7 @@ export class LinearIntegration extends Construct { // 30s deadline on cold starts. memorySize: 512, environment: { + ...abcaEnv, ...createTaskEnv, LINEAR_PROJECT_MAPPING_TABLE_NAME: this.projectMappingTable.tableName, LINEAR_USER_MAPPING_TABLE_NAME: this.userMappingTable.tableName, @@ -248,6 +256,7 @@ export class LinearIntegration extends Construct { architecture: Architecture.ARM_64, timeout: Duration.seconds(10), environment: { + ...abcaEnv, LINEAR_WEBHOOK_SECRET_ARN: this.webhookSecret.secretArn, LINEAR_WEBHOOK_DEDUP_TABLE_NAME: this.webhookDedupTable.tableName, LINEAR_WEBHOOK_PROCESSOR_FUNCTION_NAME: webhookProcessorFn.functionName, @@ -287,6 +296,7 @@ export class LinearIntegration extends Construct { architecture: Architecture.ARM_64, timeout: Duration.seconds(10), environment: { + ...abcaEnv, LINEAR_USER_MAPPING_TABLE_NAME: this.userMappingTable.tableName, }, bundling: commonBundling, diff --git a/cdk/src/constructs/pending-upload-cleanup.ts b/cdk/src/constructs/pending-upload-cleanup.ts index 2743b1de..ee245693 100644 --- a/cdk/src/constructs/pending-upload-cleanup.ts +++ b/cdk/src/constructs/pending-upload-cleanup.ts @@ -18,7 +18,7 @@ */ import * as path from 'path'; -import { Duration } from 'aws-cdk-lib'; +import { Duration, Stack } from 'aws-cdk-lib'; import * as dynamodb from 'aws-cdk-lib/aws-dynamodb'; import * as events from 'aws-cdk-lib/aws-events'; import * as targets from 'aws-cdk-lib/aws-events-targets'; @@ -93,6 +93,8 @@ export class PendingUploadCleanup extends Construct { ATTACHMENTS_BUCKET_NAME: props.attachmentsBucket.bucketName, PENDING_UPLOAD_TIMEOUT_SECONDS: String(timeoutSeconds), TASK_RETENTION_DAYS: String(retentionDays), + ABCA_STACK_NAME: Stack.of(this).stackName, + ABCA_COMPONENT: 'orchestr', }, bundling: { externalModules: ['@aws-sdk/*'], diff --git a/cdk/src/constructs/slack-integration.ts b/cdk/src/constructs/slack-integration.ts index 9b6d2e99..7b4abf91 100644 --- a/cdk/src/constructs/slack-integration.ts +++ b/cdk/src/constructs/slack-integration.ts @@ -157,6 +157,13 @@ export class SlackIntegration extends Construct { authorizationType: apigw.AuthorizationType.NONE, }; + // Outbound SDK User-Agent solution tracking (#319), spread into every + // handler environment in this construct. + const abcaEnv: Record = { + ABCA_STACK_NAME: Stack.of(this).stackName, + ABCA_COMPONENT: 'webhook', + }; + // --- Task creation environment (matches TaskApi createTaskEnv pattern) --- const createTaskEnv: Record = { TASK_TABLE_NAME: props.taskTable.tableName, @@ -186,6 +193,7 @@ export class SlackIntegration extends Construct { architecture: Architecture.ARM_64, timeout: Duration.seconds(15), environment: { + ...abcaEnv, SLACK_INSTALLATION_TABLE_NAME: this.installationTable.tableName, SLACK_CLIENT_ID_SECRET_ARN: this.clientIdSecret.secretArn, SLACK_CLIENT_SECRET_ARN: this.clientSecret.secretArn, @@ -218,6 +226,7 @@ export class SlackIntegration extends Construct { architecture: Architecture.ARM_64, timeout: Duration.seconds(10), environment: { + ...abcaEnv, SLACK_INSTALLATION_TABLE_NAME: this.installationTable.tableName, SLACK_SIGNING_SECRET_ARN: this.signingSecret.secretArn, }, @@ -250,6 +259,7 @@ export class SlackIntegration extends Construct { timeout: Duration.seconds(30), memorySize: 512, environment: { + ...abcaEnv, ...createTaskEnv, SLACK_USER_MAPPING_TABLE_NAME: this.userMappingTable.tableName, SLACK_INSTALLATION_TABLE_NAME: this.installationTable.tableName, @@ -295,6 +305,7 @@ export class SlackIntegration extends Construct { architecture: Architecture.ARM_64, timeout: Duration.seconds(10), environment: { + ...abcaEnv, SLACK_SIGNING_SECRET_ARN: this.signingSecret.secretArn, TASK_TABLE_NAME: props.taskTable.tableName, SLACK_USER_MAPPING_TABLE_NAME: this.userMappingTable.tableName, @@ -314,6 +325,7 @@ export class SlackIntegration extends Construct { architecture: Architecture.ARM_64, timeout: Duration.seconds(3), environment: { + ...abcaEnv, SLACK_SIGNING_SECRET_ARN: this.signingSecret.secretArn, SLACK_COMMAND_PROCESSOR_FUNCTION_NAME: commandProcessorFn.functionName, }, @@ -333,6 +345,7 @@ export class SlackIntegration extends Construct { architecture: Architecture.ARM_64, timeout: Duration.seconds(10), environment: { + ...abcaEnv, SLACK_USER_MAPPING_TABLE_NAME: this.userMappingTable.tableName, }, bundling: commonBundling, diff --git a/cdk/src/constructs/stranded-task-reconciler.ts b/cdk/src/constructs/stranded-task-reconciler.ts index 5e9d9be5..e5c89b32 100644 --- a/cdk/src/constructs/stranded-task-reconciler.ts +++ b/cdk/src/constructs/stranded-task-reconciler.ts @@ -18,7 +18,7 @@ */ import * as path from 'path'; -import { Duration } from 'aws-cdk-lib'; +import { Duration, Stack } from 'aws-cdk-lib'; import * as dynamodb from 'aws-cdk-lib/aws-dynamodb'; import * as events from 'aws-cdk-lib/aws-events'; import * as targets from 'aws-cdk-lib/aws-events-targets'; @@ -111,6 +111,8 @@ export class StrandedTaskReconciler extends Construct { STRANDED_TIMEOUT_SECONDS: String(strandedTimeout), APPROVAL_STRANDED_TIMEOUT_SECONDS: String(approvalStrandedTimeout), TASK_RETENTION_DAYS: String(retentionDays), + ABCA_STACK_NAME: Stack.of(this).stackName, + ABCA_COMPONENT: 'orchestr', }, bundling: { externalModules: ['@aws-sdk/*'], diff --git a/cdk/src/constructs/task-api.ts b/cdk/src/constructs/task-api.ts index b350a084..7df87d77 100644 --- a/cdk/src/constructs/task-api.ts +++ b/cdk/src/constructs/task-api.ts @@ -443,6 +443,9 @@ export class TaskApi extends Construct { TASK_TABLE_NAME: props.taskTable.tableName, TASK_EVENTS_TABLE_NAME: props.taskEventsTable.tableName, TASK_RETENTION_DAYS: String(props.taskRetentionDays ?? 90), + // Outbound SDK User-Agent solution tracking (#319) — see handlers/shared/ua.ts + ABCA_STACK_NAME: Stack.of(this).stackName, + ABCA_COMPONENT: 'api', }; // The Node.js Lambda runtime ships an AWS SDK, but its pinned version // lags current. `@aws-sdk/client-bedrock-agentcore` in particular has @@ -931,6 +934,8 @@ export class TaskApi extends Construct { const webhookEnv: Record = { WEBHOOK_TABLE_NAME: props.webhookTable.tableName, WEBHOOK_RETENTION_DAYS: String(props.webhookRetentionDays ?? 30), + ABCA_STACK_NAME: Stack.of(this).stackName, + ABCA_COMPONENT: 'webhook', }; // --- Webhook management Lambdas (Cognito-authenticated) --- diff --git a/cdk/src/constructs/task-orchestrator.ts b/cdk/src/constructs/task-orchestrator.ts index a5cf8be7..42ac1731 100644 --- a/cdk/src/constructs/task-orchestrator.ts +++ b/cdk/src/constructs/task-orchestrator.ts @@ -223,6 +223,9 @@ export class TaskOrchestrator extends Construct { TASK_EVENTS_TABLE_NAME: props.taskEventsTable.tableName, USER_CONCURRENCY_TABLE_NAME: props.userConcurrencyTable.tableName, RUNTIME_ARN: props.runtimeArn, + // Outbound SDK User-Agent solution tracking (#319) + ABCA_STACK_NAME: Stack.of(this).stackName, + ABCA_COMPONENT: 'orchestr', MAX_CONCURRENT_TASKS_PER_USER: String(maxConcurrent), TASK_RETENTION_DAYS: String(props.taskRetentionDays ?? 90), ...(props.repoTable && { REPO_TABLE_NAME: props.repoTable.tableName }), diff --git a/cdk/src/handlers/approve-task.ts b/cdk/src/handlers/approve-task.ts index 393ade46..e36c65d7 100644 --- a/cdk/src/handlers/approve-task.ts +++ b/cdk/src/handlers/approve-task.ts @@ -27,8 +27,9 @@ import { logger } from './shared/logger'; import { formatMinuteBucket } from './shared/rate-limit'; import { ErrorCode, errorResponse, successResponse } from './shared/response'; import type { ApprovalRequest, ApprovalResponse, ApprovalScope } from './shared/types'; +import { abcaUserAgent, setAbcaTrace, withAbcaTrace } from './shared/ua'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); const TASK_TABLE_NAME = process.env.TASK_TABLE_NAME; const TASK_APPROVALS_TABLE_NAME = process.env.TASK_APPROVALS_TABLE_NAME; const EVENTS_TABLE_NAME = process.env.TASK_EVENTS_TABLE_NAME; @@ -65,6 +66,7 @@ const AUDIT_EVENT_RETENTION_DAYS = Number(process.env.TASK_RETENTION_DAYS ?? '90 */ export async function handler(event: APIGatewayProxyEvent): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { // 1. Auth diff --git a/cdk/src/handlers/cancel-task.ts b/cdk/src/handlers/cancel-task.ts index 72b3864f..e41213b3 100644 --- a/cdk/src/handlers/cancel-task.ts +++ b/cdk/src/handlers/cancel-task.ts @@ -28,11 +28,12 @@ import { extractUserId } from './shared/gateway'; import { logger } from './shared/logger'; import { ErrorCode, errorResponse, successResponse } from './shared/response'; import type { TaskRecord } from './shared/types'; +import { abcaUserAgent, setAbcaTrace, withAbcaTrace } from './shared/ua'; import { computeTtlEpoch } from './shared/validation'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); -const agentCoreClient = new BedrockAgentCoreClient({}); -const ecsClient = new ECSClient({}); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); +const agentCoreClient = withAbcaTrace(new BedrockAgentCoreClient(abcaUserAgent())); +const ecsClient = withAbcaTrace(new ECSClient(abcaUserAgent())); const TABLE_NAME = process.env.TASK_TABLE_NAME!; const EVENTS_TABLE_NAME = process.env.TASK_EVENTS_TABLE_NAME!; const TASK_RETENTION_DAYS = Number(process.env.TASK_RETENTION_DAYS ?? '90'); @@ -44,6 +45,7 @@ const ECS_CLUSTER_ARN = process.env.ECS_CLUSTER_ARN; */ export async function handler(event: APIGatewayProxyEvent): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { // 1. Extract authenticated user diff --git a/cdk/src/handlers/cleanup-pending-uploads.ts b/cdk/src/handlers/cleanup-pending-uploads.ts index d5a4c055..81668984 100644 --- a/cdk/src/handlers/cleanup-pending-uploads.ts +++ b/cdk/src/handlers/cleanup-pending-uploads.ts @@ -44,9 +44,10 @@ import { DeleteObjectsCommand, ListObjectVersionsCommand, S3Client } from '@aws- import { ulid } from 'ulid'; import { ATTACHMENT_OBJECT_KEY_PREFIX } from '../constructs/attachments-bucket'; import { logger } from './shared/logger'; +import { abcaUserAgent, withAbcaTrace } from './shared/ua'; -const ddb = new DynamoDBClient({}); -const s3 = new S3Client({}); +const ddb = withAbcaTrace(new DynamoDBClient(abcaUserAgent())); +const s3 = withAbcaTrace(new S3Client(abcaUserAgent())); const TASK_TABLE = process.env.TASK_TABLE_NAME!; const EVENTS_TABLE = process.env.TASK_EVENTS_TABLE_NAME!; diff --git a/cdk/src/handlers/confirm-uploads.ts b/cdk/src/handlers/confirm-uploads.ts index 2457a4d3..6592bede 100644 --- a/cdk/src/handlers/confirm-uploads.ts +++ b/cdk/src/handlers/confirm-uploads.ts @@ -35,11 +35,12 @@ import { estimateImageTokensFromBuffer } from './shared/image-tokens'; import { logger } from './shared/logger'; import { ErrorCode, errorResponse, successResponse } from './shared/response'; import { type AttachmentRecord, createAttachmentRecord, type TaskRecord, toTaskDetail } from './shared/types'; +import { abcaUserAgent, setAbcaTrace, withAbcaTrace } from './shared/ua'; import { computeTtlEpoch } from './shared/validation'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); -const s3Client = new S3Client({}); -const lambdaClient = process.env.ORCHESTRATOR_FUNCTION_ARN ? new LambdaClient({}) : undefined; +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); +const s3Client = withAbcaTrace(new S3Client(abcaUserAgent())); +const lambdaClient = process.env.ORCHESTRATOR_FUNCTION_ARN ? withAbcaTrace(new LambdaClient(abcaUserAgent())) : undefined; const TABLE_NAME = process.env.TASK_TABLE_NAME!; const EVENTS_TABLE_NAME = process.env.TASK_EVENTS_TABLE_NAME!; @@ -80,6 +81,7 @@ interface S3ObjectMeta { */ export async function handler(event: APIGatewayProxyEvent, context: Context): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { // 1. Auth @@ -697,7 +699,7 @@ async function buildScreeningConfig(): Promise { if (!process.env.GUARDRAIL_ID || !process.env.GUARDRAIL_VERSION) return undefined; if (!_bedrockClient) { const { BedrockRuntimeClient } = await import('@aws-sdk/client-bedrock-runtime'); - _bedrockClient = new BedrockRuntimeClient({}); + _bedrockClient = withAbcaTrace(new BedrockRuntimeClient(abcaUserAgent())); } return { guardrailId: process.env.GUARDRAIL_ID, diff --git a/cdk/src/handlers/create-task.ts b/cdk/src/handlers/create-task.ts index ac909f23..1ff7aa4d 100644 --- a/cdk/src/handlers/create-task.ts +++ b/cdk/src/handlers/create-task.ts @@ -24,6 +24,7 @@ import { buildChannelMetadata, extractUserId } from './shared/gateway'; import { logger } from './shared/logger'; import { ErrorCode, errorResponse } from './shared/response'; import type { CreateTaskRequest } from './shared/types'; +import { setAbcaTrace } from './shared/ua'; import { parseBody } from './shared/validation'; /** @@ -31,6 +32,7 @@ import { parseBody } from './shared/validation'; */ export async function handler(event: APIGatewayProxyEvent): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { // 1. Extract authenticated user diff --git a/cdk/src/handlers/create-webhook.ts b/cdk/src/handlers/create-webhook.ts index 614290e4..2c1f57cf 100644 --- a/cdk/src/handlers/create-webhook.ts +++ b/cdk/src/handlers/create-webhook.ts @@ -27,10 +27,11 @@ import { extractUserId } from './shared/gateway'; import { logger } from './shared/logger'; import { ErrorCode, errorResponse, successResponse } from './shared/response'; import type { CreateWebhookRequest, CreateWebhookResponse, WebhookRecord } from './shared/types'; +import { abcaUserAgent, setAbcaTrace, withAbcaTrace } from './shared/ua'; import { isValidWebhookName, parseBody } from './shared/validation'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); -const sm = new SecretsManagerClient({}); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); +const sm = withAbcaTrace(new SecretsManagerClient(abcaUserAgent())); const TABLE_NAME = process.env.WEBHOOK_TABLE_NAME!; const SECRET_PREFIX = 'bgagent/webhook/'; @@ -39,6 +40,7 @@ const SECRET_PREFIX = 'bgagent/webhook/'; */ export async function handler(event: APIGatewayProxyEvent): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { const userId = extractUserId(event); diff --git a/cdk/src/handlers/delete-webhook.ts b/cdk/src/handlers/delete-webhook.ts index 7ac52ee1..6100fd1b 100644 --- a/cdk/src/handlers/delete-webhook.ts +++ b/cdk/src/handlers/delete-webhook.ts @@ -26,10 +26,11 @@ import { extractUserId } from './shared/gateway'; import { logger } from './shared/logger'; import { ErrorCode, errorResponse, successResponse } from './shared/response'; import { type WebhookRecord, toWebhookDetail } from './shared/types'; +import { abcaUserAgent, setAbcaTrace, withAbcaTrace } from './shared/ua'; import { computeTtlEpoch } from './shared/validation'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); -const sm = new SecretsManagerClient({}); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); +const sm = withAbcaTrace(new SecretsManagerClient(abcaUserAgent())); const TABLE_NAME = process.env.WEBHOOK_TABLE_NAME!; const SECRET_PREFIX = 'bgagent/webhook/'; const WEBHOOK_RETENTION_DAYS = Number(process.env.WEBHOOK_RETENTION_DAYS ?? '30'); @@ -39,6 +40,7 @@ const WEBHOOK_RETENTION_DAYS = Number(process.env.WEBHOOK_RETENTION_DAYS ?? '30' */ export async function handler(event: APIGatewayProxyEvent): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { const userId = extractUserId(event); diff --git a/cdk/src/handlers/deny-task.ts b/cdk/src/handlers/deny-task.ts index eb043c99..8fa280fb 100644 --- a/cdk/src/handlers/deny-task.ts +++ b/cdk/src/handlers/deny-task.ts @@ -27,8 +27,9 @@ import { logger } from './shared/logger'; import { formatMinuteBucket } from './shared/rate-limit'; import { ErrorCode, errorResponse, successResponse } from './shared/response'; import { DENY_REASON_MAX_LENGTH, type DenyRequest, type DenyResponse } from './shared/types'; +import { abcaUserAgent, setAbcaTrace, withAbcaTrace } from './shared/ua'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); const TASK_TABLE_NAME = process.env.TASK_TABLE_NAME; const TASK_APPROVALS_TABLE_NAME = process.env.TASK_APPROVALS_TABLE_NAME; const EVENTS_TABLE_NAME = process.env.TASK_EVENTS_TABLE_NAME; @@ -63,6 +64,7 @@ const AUDIT_EVENT_RETENTION_DAYS = Number(process.env.TASK_RETENTION_DAYS ?? '90 */ export async function handler(event: APIGatewayProxyEvent): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { // 1. Auth diff --git a/cdk/src/handlers/fanout-task-events.ts b/cdk/src/handlers/fanout-task-events.ts index 3eabd8c4..cf2d883d 100644 --- a/cdk/src/handlers/fanout-task-events.ts +++ b/cdk/src/handlers/fanout-task-events.ts @@ -53,6 +53,7 @@ import { logger } from './shared/logger'; import { coerceNumericOrNull } from './shared/numeric'; import { loadRepoConfig } from './shared/repo-config'; import type { ChannelConfig, TaskNotificationsConfig, TaskRecord } from './shared/types'; +import { abcaUserAgent, withAbcaTrace } from './shared/ua'; import { dispatchSlackEvent, SlackApiError } from './slack-notify'; // Re-export the shared types so existing test imports (and any future @@ -362,7 +363,7 @@ export function shouldFanOut(event: FanOutEvent, overrides?: TaskNotificationsCo * internally (the Slack API rejecting a message — e.g. * ``channel_not_found`` — is not recoverable by a Lambda retry). */ -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); /** * Slack dispatcher — hands the event to the in-module diff --git a/cdk/src/handlers/get-pending.ts b/cdk/src/handlers/get-pending.ts index 7abb4505..a2b8b969 100644 --- a/cdk/src/handlers/get-pending.ts +++ b/cdk/src/handlers/get-pending.ts @@ -26,8 +26,9 @@ import { logger } from './shared/logger'; import { formatMinuteBucket } from './shared/rate-limit'; import { ErrorCode, errorResponse, successResponse } from './shared/response'; import type { GetPendingResponse, PendingApprovalSummary, Severity } from './shared/types'; +import { abcaUserAgent, setAbcaTrace, withAbcaTrace } from './shared/ua'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); const TASK_APPROVALS_TABLE_NAME = process.env.TASK_APPROVALS_TABLE_NAME; if (!TASK_APPROVALS_TABLE_NAME) { throw new Error('get-pending handler requires TASK_APPROVALS_TABLE_NAME env var'); @@ -54,6 +55,7 @@ const PENDING_LIST_LIMIT = 100; */ export async function handler(event: APIGatewayProxyEvent): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { const userId = extractUserId(event); diff --git a/cdk/src/handlers/get-policies.ts b/cdk/src/handlers/get-policies.ts index eb8621ed..6aa35877 100644 --- a/cdk/src/handlers/get-policies.ts +++ b/cdk/src/handlers/get-policies.ts @@ -32,8 +32,9 @@ import { formatMinuteBucket } from './shared/rate-limit'; import { checkRepoOnboarded, loadRepoConfig } from './shared/repo-config'; import { ErrorCode, errorResponse, successResponse } from './shared/response'; import type { GetPoliciesResponse, PolicyRuleSummary } from './shared/types'; +import { abcaUserAgent, setAbcaTrace, withAbcaTrace } from './shared/ua'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); const TASK_APPROVALS_TABLE_NAME = process.env.TASK_APPROVALS_TABLE_NAME; const POLICIES_RATE_LIMIT_PER_MINUTE = Number(process.env.POLICIES_RATE_LIMIT_PER_MINUTE ?? '30'); @@ -64,6 +65,7 @@ const cache = new Map(); */ export async function handler(event: APIGatewayProxyEvent): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { const userId = extractUserId(event); diff --git a/cdk/src/handlers/get-task-events.ts b/cdk/src/handlers/get-task-events.ts index f631a182..4095df50 100644 --- a/cdk/src/handlers/get-task-events.ts +++ b/cdk/src/handlers/get-task-events.ts @@ -25,6 +25,7 @@ import { extractUserId } from './shared/gateway'; import { logger } from './shared/logger'; import { ErrorCode, errorResponse, paginatedResponse } from './shared/response'; import type { EventRecord, TaskRecord } from './shared/types'; +import { abcaUserAgent, setAbcaTrace, withAbcaTrace } from './shared/ua'; import { decodePaginationToken, encodePaginationToken, @@ -32,7 +33,7 @@ import { parseLimit, } from './shared/validation'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); const TABLE_NAME = process.env.TASK_TABLE_NAME!; const EVENTS_TABLE_NAME = process.env.TASK_EVENTS_TABLE_NAME!; const LOG_LEVEL = (process.env.LOG_LEVEL ?? 'INFO').toUpperCase(); @@ -69,6 +70,7 @@ type QueryMode = 'from_beginning' | 'next_token' | 'after' | 'desc'; */ export async function handler(event: APIGatewayProxyEvent): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { // 1. Extract authenticated user diff --git a/cdk/src/handlers/get-task.ts b/cdk/src/handlers/get-task.ts index 25c51ac5..08efe3e1 100644 --- a/cdk/src/handlers/get-task.ts +++ b/cdk/src/handlers/get-task.ts @@ -25,8 +25,9 @@ import { extractUserId } from './shared/gateway'; import { logger } from './shared/logger'; import { ErrorCode, errorResponse, successResponse } from './shared/response'; import { type TaskRecord, toTaskDetail } from './shared/types'; +import { abcaUserAgent, setAbcaTrace, withAbcaTrace } from './shared/ua'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); const TABLE_NAME = process.env.TASK_TABLE_NAME!; /** @@ -34,6 +35,7 @@ const TABLE_NAME = process.env.TASK_TABLE_NAME!; */ export async function handler(event: APIGatewayProxyEvent): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { // 1. Extract authenticated user diff --git a/cdk/src/handlers/get-trace-url.ts b/cdk/src/handlers/get-trace-url.ts index bdf1a553..1dbef1aa 100644 --- a/cdk/src/handlers/get-trace-url.ts +++ b/cdk/src/handlers/get-trace-url.ts @@ -28,9 +28,10 @@ import { extractUserId } from './shared/gateway'; import { logger } from './shared/logger'; import { ErrorCode, errorResponse, successResponse } from './shared/response'; import type { TaskRecord } from './shared/types'; +import { abcaUserAgent, setAbcaTrace, withAbcaTrace } from './shared/ua'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); -const s3 = new S3Client({}); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); +const s3 = withAbcaTrace(new S3Client(abcaUserAgent())); const TABLE_NAME = process.env.TASK_TABLE_NAME!; const TRACE_BUCKET_NAME = process.env.TRACE_ARTIFACTS_BUCKET_NAME!; @@ -63,6 +64,7 @@ export const TRACE_URL_TTL_SECONDS = 900; */ export async function handler(event: APIGatewayProxyEvent): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { const userId = extractUserId(event); diff --git a/cdk/src/handlers/github-webhook-processor.ts b/cdk/src/handlers/github-webhook-processor.ts index b683f565..cb56f0d9 100644 --- a/cdk/src/handlers/github-webhook-processor.ts +++ b/cdk/src/handlers/github-webhook-processor.ts @@ -29,8 +29,9 @@ import { postIssueComment } from './shared/linear-feedback'; import { extractLinearIdentifier, findLinearIssueByIdentifier } from './shared/linear-issue-lookup'; import { logger } from './shared/logger'; import { buildScreenshotKey, encodeMarkdownUrl, isAllowedScreenshotUrl } from './shared/screenshot-url'; +import { abcaUserAgent, withAbcaTrace } from './shared/ua'; -const s3 = new S3Client({}); +const s3 = withAbcaTrace(new S3Client(abcaUserAgent())); const SCREENSHOT_BUCKET = process.env.SCREENSHOT_BUCKET_NAME!; // CloudFront distribution domain — `.cloudfront.net`. Used as diff --git a/cdk/src/handlers/github-webhook.ts b/cdk/src/handlers/github-webhook.ts index 82533863..645bbb7d 100644 --- a/cdk/src/handlers/github-webhook.ts +++ b/cdk/src/handlers/github-webhook.ts @@ -27,9 +27,10 @@ import { } from './shared/github-deployment-status'; import { verifyGitHubRequest } from './shared/github-webhook-verify'; import { logger } from './shared/logger'; +import { abcaUserAgent, withAbcaTrace } from './shared/ua'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); -const lambdaClient = new LambdaClient({}); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); +const lambdaClient = withAbcaTrace(new LambdaClient(abcaUserAgent())); const WEBHOOK_SECRET_ARN = process.env.GITHUB_WEBHOOK_SECRET_ARN!; const DEDUP_TABLE_NAME = process.env.GITHUB_WEBHOOK_DEDUP_TABLE_NAME!; diff --git a/cdk/src/handlers/linear-link.ts b/cdk/src/handlers/linear-link.ts index 41a480d3..74fb9440 100644 --- a/cdk/src/handlers/linear-link.ts +++ b/cdk/src/handlers/linear-link.ts @@ -24,9 +24,10 @@ import { ulid } from 'ulid'; import { extractUserId } from './shared/gateway'; import { logger } from './shared/logger'; import { ErrorCode, errorResponse, successResponse } from './shared/response'; +import { abcaUserAgent, setAbcaTrace, withAbcaTrace } from './shared/ua'; import { parseBody } from './shared/validation'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); const USER_MAPPING_TABLE = process.env.LINEAR_USER_MAPPING_TABLE_NAME!; @@ -48,6 +49,7 @@ interface LinkRequest { */ export async function handler(event: APIGatewayProxyEvent): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { const userId = extractUserId(event); diff --git a/cdk/src/handlers/linear-webhook-processor.ts b/cdk/src/handlers/linear-webhook-processor.ts index 15d2b4b3..696189e3 100644 --- a/cdk/src/handlers/linear-webhook-processor.ts +++ b/cdk/src/handlers/linear-webhook-processor.ts @@ -25,8 +25,9 @@ import { reportIssueFailure } from './shared/linear-feedback'; import { resolveLinearOauthToken } from './shared/linear-oauth-resolver'; import { logger } from './shared/logger'; import type { Attachment } from './shared/types'; +import { abcaUserAgent, withAbcaTrace } from './shared/ua'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); const PROJECT_MAPPING_TABLE = process.env.LINEAR_PROJECT_MAPPING_TABLE_NAME!; const USER_MAPPING_TABLE = process.env.LINEAR_USER_MAPPING_TABLE_NAME!; diff --git a/cdk/src/handlers/linear-webhook.ts b/cdk/src/handlers/linear-webhook.ts index 65fdc5d6..7331cd7d 100644 --- a/cdk/src/handlers/linear-webhook.ts +++ b/cdk/src/handlers/linear-webhook.ts @@ -27,9 +27,10 @@ import { verifyLinearRequestForWorkspace, } from './shared/linear-verify'; import { logger } from './shared/logger'; +import { abcaUserAgent, withAbcaTrace } from './shared/ua'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); -const lambdaClient = new LambdaClient({}); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); +const lambdaClient = withAbcaTrace(new LambdaClient(abcaUserAgent())); const WEBHOOK_SECRET_ARN = process.env.LINEAR_WEBHOOK_SECRET_ARN!; const DEDUP_TABLE_NAME = process.env.LINEAR_WEBHOOK_DEDUP_TABLE_NAME!; diff --git a/cdk/src/handlers/list-tasks.ts b/cdk/src/handlers/list-tasks.ts index eff4c859..245561dc 100644 --- a/cdk/src/handlers/list-tasks.ts +++ b/cdk/src/handlers/list-tasks.ts @@ -25,9 +25,10 @@ import { extractUserId } from './shared/gateway'; import { logger } from './shared/logger'; import { ErrorCode, errorResponse, paginatedResponse } from './shared/response'; import { type TaskRecord, toTaskSummary } from './shared/types'; +import { abcaUserAgent, setAbcaTrace, withAbcaTrace } from './shared/ua'; import { decodePaginationToken, encodePaginationToken, parseLimit, parseStatusFilter } from './shared/validation'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); const TABLE_NAME = process.env.TASK_TABLE_NAME!; /** @@ -35,6 +36,7 @@ const TABLE_NAME = process.env.TASK_TABLE_NAME!; */ export async function handler(event: APIGatewayProxyEvent): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { // 1. Extract authenticated user diff --git a/cdk/src/handlers/list-webhooks.ts b/cdk/src/handlers/list-webhooks.ts index 94c4f19d..2ca75f31 100644 --- a/cdk/src/handlers/list-webhooks.ts +++ b/cdk/src/handlers/list-webhooks.ts @@ -25,9 +25,10 @@ import { extractUserId } from './shared/gateway'; import { logger } from './shared/logger'; import { ErrorCode, errorResponse, paginatedResponse } from './shared/response'; import { type WebhookRecord, toWebhookDetail } from './shared/types'; +import { abcaUserAgent, setAbcaTrace, withAbcaTrace } from './shared/ua'; import { decodePaginationToken, encodePaginationToken, parseLimit } from './shared/validation'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); const TABLE_NAME = process.env.WEBHOOK_TABLE_NAME!; /** @@ -35,6 +36,7 @@ const TABLE_NAME = process.env.WEBHOOK_TABLE_NAME!; */ export async function handler(event: APIGatewayProxyEvent): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { const userId = extractUserId(event); diff --git a/cdk/src/handlers/nudge-task.ts b/cdk/src/handlers/nudge-task.ts index 5fbdcb60..ff6dc0dd 100644 --- a/cdk/src/handlers/nudge-task.ts +++ b/cdk/src/handlers/nudge-task.ts @@ -28,8 +28,9 @@ import { logger } from './shared/logger'; import { formatMinuteBucket } from './shared/rate-limit'; import { ErrorCode, errorResponse, successResponse } from './shared/response'; import { NUDGE_MAX_MESSAGE_LENGTH, type NudgeRecord, type NudgeRequest, type TaskRecord } from './shared/types'; +import { abcaUserAgent, setAbcaTrace, withAbcaTrace } from './shared/ua'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); const TASK_TABLE_NAME = process.env.TASK_TABLE_NAME; const NUDGES_TABLE_NAME = process.env.NUDGES_TABLE_NAME; if (!TASK_TABLE_NAME || !NUDGES_TABLE_NAME) { @@ -62,6 +63,7 @@ const RATE_LIMIT_ROW_TTL_SECONDS = 120; */ export async function handler(event: APIGatewayProxyEvent): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { // 1. Auth diff --git a/cdk/src/handlers/orchestrate-task.ts b/cdk/src/handlers/orchestrate-task.ts index 508ce152..6d406096 100644 --- a/cdk/src/handlers/orchestrate-task.ts +++ b/cdk/src/handlers/orchestrate-task.ts @@ -36,6 +36,7 @@ import { } from './shared/orchestrator'; import { runPreflightChecks } from './shared/preflight'; import type { TaskRecord } from './shared/types'; +import { setAbcaTrace } from './shared/ua'; import { workflowIsReadOnly, workflowRequiresRepo } from './shared/workflows'; interface OrchestrateTaskEvent { @@ -49,6 +50,7 @@ const MAX_CONSECUTIVE_ECS_COMPLETED_POLLS = 5; const durableHandler: DurableExecutionHandler = async (event, context) => { const { task_id: taskId } = event; + setAbcaTrace(taskId); // Step 1: Load task record const task = await context.step('load-task', async () => { diff --git a/cdk/src/handlers/reconcile-concurrency.ts b/cdk/src/handlers/reconcile-concurrency.ts index 2f716514..415d5a56 100644 --- a/cdk/src/handlers/reconcile-concurrency.ts +++ b/cdk/src/handlers/reconcile-concurrency.ts @@ -19,8 +19,9 @@ import { DynamoDBClient, ScanCommand, QueryCommand, UpdateItemCommand } from '@aws-sdk/client-dynamodb'; import { logger } from './shared/logger'; +import { abcaUserAgent, withAbcaTrace } from './shared/ua'; -const ddb = new DynamoDBClient({}); +const ddb = withAbcaTrace(new DynamoDBClient(abcaUserAgent())); const TASK_TABLE = process.env.TASK_TABLE_NAME!; const CONCURRENCY_TABLE = process.env.USER_CONCURRENCY_TABLE_NAME!; diff --git a/cdk/src/handlers/reconcile-stranded-tasks.ts b/cdk/src/handlers/reconcile-stranded-tasks.ts index 8688601f..7dc6bc4e 100644 --- a/cdk/src/handlers/reconcile-stranded-tasks.ts +++ b/cdk/src/handlers/reconcile-stranded-tasks.ts @@ -48,8 +48,9 @@ import { } from '@aws-sdk/client-dynamodb'; import { ulid } from 'ulid'; import { logger } from './shared/logger'; +import { abcaUserAgent, withAbcaTrace } from './shared/ua'; -const ddb = new DynamoDBClient({}); +const ddb = withAbcaTrace(new DynamoDBClient(abcaUserAgent())); const TASK_TABLE = process.env.TASK_TABLE_NAME!; const EVENTS_TABLE = process.env.TASK_EVENTS_TABLE_NAME!; const CONCURRENCY_TABLE = process.env.USER_CONCURRENCY_TABLE_NAME!; diff --git a/cdk/src/handlers/shared/agentcore-browser.ts b/cdk/src/handlers/shared/agentcore-browser.ts index deacd937..f55a1d0f 100644 --- a/cdk/src/handlers/shared/agentcore-browser.ts +++ b/cdk/src/handlers/shared/agentcore-browser.ts @@ -28,6 +28,7 @@ import { HttpRequest } from '@smithy/protocol-http'; import { SignatureV4 } from '@smithy/signature-v4'; import WebSocket, { type RawData } from 'ws'; import { logger } from './logger'; +import { abcaUserAgent, withAbcaTrace } from './ua'; const REGION = process.env.AWS_REGION ?? process.env.AWS_DEFAULT_REGION ?? 'us-east-1'; @@ -88,7 +89,7 @@ interface CdpMessage { */ export async function captureScreenshot(url: string, opts: { timeoutMs?: number } = {}): Promise { const timeoutMs = opts.timeoutMs ?? DEFAULT_TIMEOUT_MS; - const client = new BedrockAgentCoreClient({ region: REGION }); + const client = withAbcaTrace(new BedrockAgentCoreClient({ region: REGION, ...abcaUserAgent() })); const startResp = await client.send(new StartBrowserSessionCommand({ browserIdentifier: AWS_BROWSER_IDENTIFIER, diff --git a/cdk/src/handlers/shared/context-hydration.ts b/cdk/src/handlers/shared/context-hydration.ts index e2ef1f20..f7565d81 100644 --- a/cdk/src/handlers/shared/context-hydration.ts +++ b/cdk/src/handlers/shared/context-hydration.ts @@ -23,6 +23,7 @@ import { logger } from './logger'; import { loadMemoryContext, type MemoryContext } from './memory'; import { sanitizeExternalContent } from './sanitization'; import { type TaskRecord } from './types'; +import { abcaUserAgent, withAbcaTrace } from './ua'; import { workflowIsReadOnly, workflowUsesPr } from './workflows'; // --------------------------------------------------------------------------- @@ -130,7 +131,7 @@ const USER_PROMPT_TOKEN_BUDGET = Number(process.env.USER_PROMPT_TOKEN_BUDGET ?? const GITHUB_API_TIMEOUT_MS = 30_000; const GUARDRAIL_ID = process.env.GUARDRAIL_ID; const GUARDRAIL_VERSION = process.env.GUARDRAIL_VERSION; -const bedrockClient = (GUARDRAIL_ID && GUARDRAIL_VERSION) ? new BedrockRuntimeClient({}) : undefined; +const bedrockClient = (GUARDRAIL_ID && GUARDRAIL_VERSION) ? withAbcaTrace(new BedrockRuntimeClient(abcaUserAgent())) : undefined; if (GUARDRAIL_ID && !GUARDRAIL_VERSION) { logger.error('GUARDRAIL_ID is set but GUARDRAIL_VERSION is missing — guardrail screening disabled', { metric_type: 'guardrail_misconfiguration', @@ -313,7 +314,7 @@ export async function screenWithGuardrail(text: string, taskId: string): Promise const tokenCache = new Map(); const CACHE_TTL_MS = 5 * 60 * 1000; // 5 minutes -const smClient = new SecretsManagerClient({}); +const smClient = withAbcaTrace(new SecretsManagerClient(abcaUserAgent())); /** * Resolve the GitHub token from Secrets Manager with per-ARN caching. diff --git a/cdk/src/handlers/shared/create-task-core.ts b/cdk/src/handlers/shared/create-task-core.ts index 0315b325..bf82ff53 100644 --- a/cdk/src/handlers/shared/create-task-core.ts +++ b/cdk/src/handlers/shared/create-task-core.ts @@ -53,6 +53,7 @@ import { type TaskRecord, toTaskDetail, } from './types'; +import { abcaUserAgent, withAbcaTrace } from './ua'; import { computeTtlEpoch, DEFAULT_MAX_TURNS, hasTaskSpec, isValidIdempotencyKey, isValidRepo, isValidTaskDescriptionLength, MAX_ATTACHMENT_SIZE_BYTES, MAX_TASK_DESCRIPTION_LENGTH, validateAttachments, validateMaxBudgetUsd, validateMaxTurns, validatePrNumber } from './validation'; import { disallowedWorkflowModel, getWorkflowDescriptor, isValidWorkflowRef, resolveWorkflowRef, resolveWorkflowRefError } from './workflows'; import { ATTACHMENT_OBJECT_KEY_PREFIX } from '../../constructs/attachments-bucket'; @@ -68,10 +69,10 @@ export interface TaskCreationContext { readonly idempotencyKey?: string; } -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); -const lambdaClient = process.env.ORCHESTRATOR_FUNCTION_ARN ? new LambdaClient({}) : undefined; +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); +const lambdaClient = process.env.ORCHESTRATOR_FUNCTION_ARN ? withAbcaTrace(new LambdaClient(abcaUserAgent())) : undefined; const bedrockClient = (process.env.GUARDRAIL_ID && process.env.GUARDRAIL_VERSION) - ? new BedrockRuntimeClient({}) : undefined; + ? withAbcaTrace(new BedrockRuntimeClient(abcaUserAgent())) : undefined; if (process.env.GUARDRAIL_ID && !process.env.GUARDRAIL_VERSION) { logger.error('GUARDRAIL_ID is set but GUARDRAIL_VERSION is missing — guardrail screening disabled', { metric_type: 'guardrail_misconfiguration', @@ -81,7 +82,7 @@ const TABLE_NAME = process.env.TASK_TABLE_NAME!; const EVENTS_TABLE_NAME = process.env.TASK_EVENTS_TABLE_NAME!; const TASK_RETENTION_DAYS = Number(process.env.TASK_RETENTION_DAYS ?? '90'); const ATTACHMENTS_BUCKET = process.env.ATTACHMENTS_BUCKET_NAME; -const s3Client = ATTACHMENTS_BUCKET ? new S3Client({}) : undefined; +const s3Client = ATTACHMENTS_BUCKET ? withAbcaTrace(new S3Client(abcaUserAgent())) : undefined; /** Human-readable description of a workflow's required-input contract (for 400s). */ function describeRequiredInputs(requiredInputs: { allOf?: readonly string[]; oneOf?: readonly string[] }): string { diff --git a/cdk/src/handlers/shared/github-webhook-verify.ts b/cdk/src/handlers/shared/github-webhook-verify.ts index 5085efa8..a3e1d967 100644 --- a/cdk/src/handlers/shared/github-webhook-verify.ts +++ b/cdk/src/handlers/shared/github-webhook-verify.ts @@ -21,8 +21,9 @@ import * as crypto from 'crypto'; import { GetSecretValueCommand, SecretsManagerClient } from '@aws-sdk/client-secrets-manager'; import { isUsableHmacSecret } from './hmac-secret'; import { logger } from './logger'; +import { abcaUserAgent, withAbcaTrace } from './ua'; -const sm = new SecretsManagerClient({}); +const sm = withAbcaTrace(new SecretsManagerClient(abcaUserAgent())); /** * In-memory secret cache (5-minute TTL). Same pattern as `linear-verify.ts` diff --git a/cdk/src/handlers/shared/linear-issue-lookup.ts b/cdk/src/handlers/shared/linear-issue-lookup.ts index b2373887..2db05862 100644 --- a/cdk/src/handlers/shared/linear-issue-lookup.ts +++ b/cdk/src/handlers/shared/linear-issue-lookup.ts @@ -21,8 +21,9 @@ import { DynamoDBClient } from '@aws-sdk/client-dynamodb'; import { DynamoDBDocumentClient, ScanCommand } from '@aws-sdk/lib-dynamodb'; import { resolveLinearOauthToken } from './linear-oauth-resolver'; import { logger } from './logger'; +import { abcaUserAgent, withAbcaTrace } from './ua'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); /** * Linear issue identifier shape, e.g. `ABCA-42`. Linear identifiers are diff --git a/cdk/src/handlers/shared/linear-oauth-resolver.ts b/cdk/src/handlers/shared/linear-oauth-resolver.ts index 48cc6f89..3a108e9f 100644 --- a/cdk/src/handlers/shared/linear-oauth-resolver.ts +++ b/cdk/src/handlers/shared/linear-oauth-resolver.ts @@ -25,6 +25,7 @@ import { } from '@aws-sdk/client-secrets-manager'; import { DynamoDBDocumentClient, GetCommand } from '@aws-sdk/lib-dynamodb'; import { logger } from './logger'; +import { abcaUserAgent, withAbcaTrace } from './ua'; /** * Lambda-side resolver for the per-workspace Linear OAuth token written @@ -152,8 +153,8 @@ export async function resolveLinearOauthToken( options: ResolverOptions = {}, ): Promise { const region = options.region ?? process.env.AWS_REGION ?? 'us-east-1'; - const ddb = options.dynamoDbClient ?? DynamoDBDocumentClient.from(new DynamoDBClient({ region })); - const sm = options.secretsManagerClient ?? new SecretsManagerClient({ region }); + const ddb = options.dynamoDbClient ?? DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient({ region, ...abcaUserAgent() }))); + const sm = options.secretsManagerClient ?? withAbcaTrace(new SecretsManagerClient({ region, ...abcaUserAgent() })); // ─── Step 1: Registry row ──────────────────────────────────────── const row = await getRegistryRow(ddb, registryTableName, linearWorkspaceId); diff --git a/cdk/src/handlers/shared/linear-verify.ts b/cdk/src/handlers/shared/linear-verify.ts index bb4799d2..4745d344 100644 --- a/cdk/src/handlers/shared/linear-verify.ts +++ b/cdk/src/handlers/shared/linear-verify.ts @@ -24,9 +24,10 @@ import { DynamoDBDocumentClient } from '@aws-sdk/lib-dynamodb'; import { isUsableHmacSecret } from './hmac-secret'; import { getOauthSecretStrict, getRegistryRowStrict } from './linear-oauth-resolver'; import { logger } from './logger'; +import { abcaUserAgent, withAbcaTrace } from './ua'; -const sm = new SecretsManagerClient({}); -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const sm = withAbcaTrace(new SecretsManagerClient(abcaUserAgent())); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); /** Prefix for Linear-related secrets in Secrets Manager. */ export const LINEAR_SECRET_PREFIX = 'bgagent/linear/'; diff --git a/cdk/src/handlers/shared/memory.ts b/cdk/src/handlers/shared/memory.ts index bd8d5025..0281255d 100644 --- a/cdk/src/handlers/shared/memory.ts +++ b/cdk/src/handlers/shared/memory.ts @@ -25,6 +25,7 @@ import { } from '@aws-sdk/client-bedrock-agentcore'; import { logger } from './logger'; import { sanitizeExternalContent } from './sanitization'; +import { abcaUserAgent, withAbcaTrace } from './ua'; import type { TaskStatusType } from '../../constructs/task-status'; // --------------------------------------------------------------------------- @@ -146,7 +147,7 @@ function processMemoryRecords( let agentCoreClient: BedrockAgentCoreClient | undefined; function getClient(): BedrockAgentCoreClient { if (!agentCoreClient) { - agentCoreClient = new BedrockAgentCoreClient({}); + agentCoreClient = withAbcaTrace(new BedrockAgentCoreClient(abcaUserAgent())); } return agentCoreClient; } diff --git a/cdk/src/handlers/shared/orchestrator.ts b/cdk/src/handlers/shared/orchestrator.ts index 27a8a0a6..b47d72e7 100644 --- a/cdk/src/handlers/shared/orchestrator.ts +++ b/cdk/src/handlers/shared/orchestrator.ts @@ -29,10 +29,11 @@ import { computePromptVersion } from './prompt-version'; import { loadRepoConfig, type BlueprintConfig, type ComputeType } from './repo-config'; import { resolveUrlAttachments } from './resolve-url-attachments'; import { APPROVAL_GATE_CAP_MAX, APPROVAL_GATE_CAP_MIN, type AgentAttachmentPayload, type AttachmentRecord, type TaskRecord } from './types'; +import { abcaUserAgent, withAbcaTrace } from './ua'; import { computeTtlEpoch, DEFAULT_MAX_TURNS } from './validation'; import { TaskStatus, TERMINAL_STATUSES, VALID_TRANSITIONS, type TaskStatusType } from '../../constructs/task-status'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); const TABLE_NAME = process.env.TASK_TABLE_NAME!; const EVENTS_TABLE_NAME = process.env.TASK_EVENTS_TABLE_NAME!; @@ -431,7 +432,7 @@ export async function hydrateAndTransition(task: TaskRecord, blueprintConfig?: B ? { guardrailId: process.env.GUARDRAIL_ID, guardrailVersion: process.env.GUARDRAIL_VERSION, - bedrockClient: new BedrockRuntimeClient({}), + bedrockClient: withAbcaTrace(new BedrockRuntimeClient(abcaUserAgent())), } : undefined; @@ -481,7 +482,7 @@ export async function hydrateAndTransition(task: TaskRecord, blueprintConfig?: B task.task_id, task.user_id, { - s3Client: new S3Client({}), + s3Client: withAbcaTrace(new S3Client(abcaUserAgent())), bucketName: ATTACHMENTS_BUCKET_NAME, screeningConfig, githubToken, diff --git a/cdk/src/handlers/shared/repo-config.ts b/cdk/src/handlers/shared/repo-config.ts index 1753d568..90f79732 100644 --- a/cdk/src/handlers/shared/repo-config.ts +++ b/cdk/src/handlers/shared/repo-config.ts @@ -20,6 +20,7 @@ import { DynamoDBClient } from '@aws-sdk/client-dynamodb'; import { DynamoDBDocumentClient, GetCommand } from '@aws-sdk/lib-dynamodb'; import { logger } from './logger'; +import { abcaUserAgent, withAbcaTrace } from './ua'; /** * Per-repository configuration written by the Blueprint CDK construct @@ -79,7 +80,7 @@ export interface BlueprintConfig { readonly approval_gate_cap?: number; } -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); /** * Combined result of a single RepoTable GetItem used by the submit diff --git a/cdk/src/handlers/shared/slack-verify.ts b/cdk/src/handlers/shared/slack-verify.ts index c7ff6e7e..e6d50235 100644 --- a/cdk/src/handlers/shared/slack-verify.ts +++ b/cdk/src/handlers/shared/slack-verify.ts @@ -21,8 +21,9 @@ import * as crypto from 'crypto'; import { GetSecretValueCommand, SecretsManagerClient } from '@aws-sdk/client-secrets-manager'; import { isUsableHmacSecret } from './hmac-secret'; import { logger } from './logger'; +import { abcaUserAgent, withAbcaTrace } from './ua'; -const sm = new SecretsManagerClient({}); +const sm = withAbcaTrace(new SecretsManagerClient(abcaUserAgent())); /** Prefix for Slack-related secrets in Secrets Manager. */ export const SLACK_SECRET_PREFIX = 'bgagent/slack/'; diff --git a/cdk/src/handlers/shared/strategies/agentcore-strategy.ts b/cdk/src/handlers/shared/strategies/agentcore-strategy.ts index d10e9bde..ab70344c 100644 --- a/cdk/src/handlers/shared/strategies/agentcore-strategy.ts +++ b/cdk/src/handlers/shared/strategies/agentcore-strategy.ts @@ -22,11 +22,12 @@ import { BedrockAgentCoreClient, InvokeAgentRuntimeCommand, StopRuntimeSessionCo import type { ComputeStrategy, SessionHandle, SessionStatus } from '../compute-strategy'; import { logger } from '../logger'; import type { BlueprintConfig } from '../repo-config'; +import { abcaUserAgent, withAbcaTrace } from '../ua'; let sharedClient: BedrockAgentCoreClient | undefined; function getClient(): BedrockAgentCoreClient { if (!sharedClient) { - sharedClient = new BedrockAgentCoreClient({}); + sharedClient = withAbcaTrace(new BedrockAgentCoreClient(abcaUserAgent())); } return sharedClient; } diff --git a/cdk/src/handlers/shared/strategies/ecs-strategy.ts b/cdk/src/handlers/shared/strategies/ecs-strategy.ts index e032c6ce..0d4052cc 100644 --- a/cdk/src/handlers/shared/strategies/ecs-strategy.ts +++ b/cdk/src/handlers/shared/strategies/ecs-strategy.ts @@ -21,11 +21,12 @@ import { ECSClient, RunTaskCommand, DescribeTasksCommand, StopTaskCommand } from import type { ComputeStrategy, SessionHandle, SessionStatus } from '../compute-strategy'; import { logger } from '../logger'; import type { BlueprintConfig } from '../repo-config'; +import { abcaUserAgent, withAbcaTrace } from '../ua'; let sharedClient: ECSClient | undefined; function getClient(): ECSClient { if (!sharedClient) { - sharedClient = new ECSClient({}); + sharedClient = withAbcaTrace(new ECSClient(abcaUserAgent())); } return sharedClient; } diff --git a/cdk/src/handlers/shared/ua.ts b/cdk/src/handlers/shared/ua.ts new file mode 100644 index 00000000..9eb352e2 --- /dev/null +++ b/cdk/src/handlers/shared/ua.ts @@ -0,0 +1,196 @@ +/** + * MIT No Attribution + * + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * the Software without restriction, including without limitation the rights to + * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of + * the Software, and to permit persons to whom the Software is furnished to do so. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ + +/** + * Outbound AWS SDK User-Agent solution tracking (#319). + * + * Every AWS API call made by the Lambda handlers carries two ABCA + * solution-tracking segments in the `User-Agent` header: + * + * app/uksb-wt64nei4u6/{STACKNAME} (only when ABCA_STACK_NAME set) + * md/uksb-wt64nei4u6#{COMPONENT}[#{TRACE}] + * + * Both ride the verbatim `customUserAgent` path — NOT the sanitizing + * `userAgentAppId` config field, whose allowed charset excludes `/` and would + * mangle the `uksb-wt64nei4u6/` separator into `-`. Because the raw path + * applies only pass-through escaping to our token-safe characters, this + * module sanitizes `{STACKNAME}` and `{TRACE}` itself. + * + * The static part is baked once at client construction via + * {@link abcaUserAgent}. The optional `#{TRACE}` suffix (the handler's ulid + * request id, or a task id) is appended **per request** by the middleware + * {@link withAbcaTrace} adds — never via client config — so module-level + * cached clients keep their connection pools across trace changes. + * + * Trace state is a module-level variable: a Lambda execution environment + * processes one invocation at a time, so ambient module state is safe (and + * survives across the SDK's async internals where per-call threading can't). + * + * Counterparts: `agent/src/ua.py` (Python agent runtime) and `cli/src/ua.ts` + * (bgagent CLI). Solution id, wire format, and sanitization rules must stay + * identical across all three. + */ + +/** + * AWS solution-tracking id for ABCA. Deploy-time counterpart (#292) lives in + * the CloudFormation stack description in `cdk/src/main.ts`. Per-surface + * literal by design — see PR #338. + */ +export const SOLUTION_ID = 'uksb-wt64nei4u6'; + +/** + * Env var carrying the stable per-component label (`api`, `webhook`, + * `orchestr`) — set per-Lambda by the CDK constructs. Shared handler modules + * are bundled into multiple Lambdas, so identity must come from the + * environment, not from code. + */ +export const COMPONENT_ENV = 'ABCA_COMPONENT'; + +/** Env var carrying the deployed CloudFormation stack name (set by CDK). */ +export const STACK_NAME_ENV = 'ABCA_STACK_NAME'; + +/** Default component label when ABCA_COMPONENT is absent (REST API surface). */ +const DEFAULT_COMPONENT = 'api'; + +/** + * App-id budget: the documented 50-char cap on the value, minus + * `uksb-wt64nei4u6/` (16 chars), leaves 34 for the stack name. + */ +const STACK_NAME_MAX = 34; + +/** + * RFC 7230 token charset (the UA product-token alphabet). `/` and `#` are + * deliberately excluded — they are the structural separators of the scheme. + * Mirrors `_ALLOWED` in `agent/src/ua.py`. + */ +const UA_TOKEN_SAFE = /[^A-Za-z0-9!$%&'*+\-.^_`|~]/g; + +let currentTrace: string | undefined; + +/** Replace every non-UA-token char (incl. non-ASCII) with `-`. */ +export function sanitizeUaValue(raw: string): string { + return raw.replace(UA_TOKEN_SAFE, '-'); +} + +/** The component label for this Lambda (from env, sanitized). */ +function componentLabel(): string { + return sanitizeUaValue(process.env[COMPONENT_ENV]?.trim() || DEFAULT_COMPONENT); +} + +/** The static `md/` segment as it renders on the wire. */ +function mdSegment(): string { + return `md/${SOLUTION_ID}#${componentLabel()}`; +} + +/** + * Client config fragment carrying the static ABCA UA segments. + * + * Spread into any SDK v3 client constructor: + * `new DynamoDBClient({ ...abcaUserAgent() })`. Each entry is a + * `[name, value?]` user-agent pair. The SDK's escaper treats the two + * positions differently: the *name* is split on `/`, each part escaped + * (where `#` is NOT allowed and becomes `-`), and rejoined with `/`; the + * *value* allows `#` and is joined to the name with `#`. So: + * + * - the `app/` segment is a single-element pair — its only separators are + * slashes, which survive the name path (this is what keeps the literal + * `/` that the sanitizing app-id config field would destroy); + * - the `md/` segment is a two-element pair `['md/{id}', component]`, + * rendering `md/{id}#component` — the `#` comes from the SDK's own + * name#value join, not from our string (a `#` inside the name would be + * escaped to `-`). + */ +export function abcaUserAgent(): { customUserAgent: ([string] | [string, string])[] } { + const pairs: ([string] | [string, string])[] = []; + const stackName = process.env[STACK_NAME_ENV]?.trim(); + if (stackName) { + // Sanitize FIRST, then clip, so a replaced char can't be re-split. + const clipped = sanitizeUaValue(stackName).slice(0, STACK_NAME_MAX); + pairs.push([`app/${SOLUTION_ID}/${clipped}`]); + } + pairs.push([`md/${SOLUTION_ID}`, componentLabel()]); + return { customUserAgent: pairs }; +} + +/** + * Set (or clear, by omitting the argument) the ambient trace handle. + * Handlers call this with their per-invocation request id right after + * minting it; the orchestrator uses the task id. + */ +export function setAbcaTrace(handle?: string): void { + currentTrace = handle || undefined; +} + +/** Current trace handle, sanitized to UA-token-safe ASCII, or undefined. */ +export function getAbcaTrace(): string | undefined { + return currentTrace ? sanitizeUaValue(currentTrace) : undefined; +} + +/** + * Minimal structural view of an SDK v3 client middleware stack — enough to + * add the trace middleware without importing @smithy/types (which is not a + * declared dependency of the handlers). + */ +interface MiddlewareStackLike { + addRelativeTo(middleware: unknown, options: Record): void; +} + +/** + * Append `#{TRACE}` to the outgoing User-Agent headers on every request. + * + * Adds a middleware right after the SDK's own `getUserAgentMiddleware` + * (step `build`) that splices the current trace onto the static `md/` + * segment in both `user-agent` and `x-amz-user-agent`. Only the header + * strings change — the client, its config, and its connection pool are + * untouched, so cached/module-level clients are reused freely across traces. + * + * No-ops when the client has no middleware stack: ~40 existing test suites + * mock client constructors as `jest.fn(() => ({}))`, and module-level + * instrumentation must not crash under those mocks. Real SDK clients always + * have a stack, so the guard is test-environment-only. + */ +export function withAbcaTrace(client: T): T { + const stack = (client as { middlewareStack?: MiddlewareStackLike }).middlewareStack; + if (!stack || typeof stack.addRelativeTo !== 'function') { + return client; + } + const md = mdSegment(); + stack.addRelativeTo( + (next: (args: unknown) => Promise) => async (args: unknown) => { + const trace = getAbcaTrace(); + const request = (args as { request?: { headers?: Record } }).request; + if (trace && request?.headers) { + for (const header of ['user-agent', 'x-amz-user-agent']) { + const value = request.headers[header]; + if (value && value.includes(md)) { + request.headers[header] = value.replace(md, `${md}#${trace}`); + } + } + } + return next(args); + }, + { + name: 'abcaUaTraceMiddleware', + relation: 'after', + toMiddleware: 'getUserAgentMiddleware', + override: true, + }, + ); + return client; +} diff --git a/cdk/src/handlers/slack-command-processor.ts b/cdk/src/handlers/slack-command-processor.ts index 3832ffa0..f7cb460b 100644 --- a/cdk/src/handlers/slack-command-processor.ts +++ b/cdk/src/handlers/slack-command-processor.ts @@ -25,6 +25,7 @@ import { logger } from './shared/logger'; import { slackFetch } from './shared/slack-api'; import { getSlackSecret, SLACK_SECRET_PREFIX } from './shared/slack-verify'; import type { Attachment } from './shared/types'; +import { abcaUserAgent, withAbcaTrace } from './shared/ua'; import type { SlackCommandPayload } from './slack-commands'; /** @@ -76,7 +77,7 @@ function normalizeEvent(event: RawEvent): CommandProcessorEvent { return { ...event, source: 'slash' }; } -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); const USER_MAPPING_TABLE = process.env.SLACK_USER_MAPPING_TABLE_NAME!; const INSTALLATION_TABLE = process.env.SLACK_INSTALLATION_TABLE_NAME!; diff --git a/cdk/src/handlers/slack-commands.ts b/cdk/src/handlers/slack-commands.ts index f89b47c8..591b0886 100644 --- a/cdk/src/handlers/slack-commands.ts +++ b/cdk/src/handlers/slack-commands.ts @@ -21,8 +21,9 @@ import { InvokeCommand, LambdaClient } from '@aws-sdk/client-lambda'; import type { APIGatewayProxyEvent, APIGatewayProxyResult } from 'aws-lambda'; import { logger } from './shared/logger'; import { getSlackSecret, verifySlackRequest } from './shared/slack-verify'; +import { abcaUserAgent, withAbcaTrace } from './shared/ua'; -const lambdaClient = new LambdaClient({}); +const lambdaClient = withAbcaTrace(new LambdaClient(abcaUserAgent())); const SIGNING_SECRET_ARN = process.env.SLACK_SIGNING_SECRET_ARN!; const PROCESSOR_FUNCTION_NAME = process.env.SLACK_COMMAND_PROCESSOR_FUNCTION_NAME!; diff --git a/cdk/src/handlers/slack-events.ts b/cdk/src/handlers/slack-events.ts index 954f53e7..74d26c17 100644 --- a/cdk/src/handlers/slack-events.ts +++ b/cdk/src/handlers/slack-events.ts @@ -25,11 +25,12 @@ import type { APIGatewayProxyEvent, APIGatewayProxyResult } from 'aws-lambda'; import { logger } from './shared/logger'; import { slackFetch } from './shared/slack-api'; import { getSlackSecret, SLACK_SECRET_PREFIX, verifySlackRequest } from './shared/slack-verify'; +import { abcaUserAgent, withAbcaTrace } from './shared/ua'; import type { MentionEvent, SlackFileRef } from './slack-command-processor'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); -const sm = new SecretsManagerClient({}); -const lambdaClient = new LambdaClient({}); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); +const sm = withAbcaTrace(new SecretsManagerClient(abcaUserAgent())); +const lambdaClient = withAbcaTrace(new LambdaClient(abcaUserAgent())); const TABLE_NAME = process.env.SLACK_INSTALLATION_TABLE_NAME!; const SIGNING_SECRET_ARN = process.env.SLACK_SIGNING_SECRET_ARN!; diff --git a/cdk/src/handlers/slack-interactions.ts b/cdk/src/handlers/slack-interactions.ts index ac25ccff..199debf2 100644 --- a/cdk/src/handlers/slack-interactions.ts +++ b/cdk/src/handlers/slack-interactions.ts @@ -22,8 +22,9 @@ import { DynamoDBDocumentClient, GetCommand, UpdateCommand } from '@aws-sdk/lib- import type { APIGatewayProxyEvent, APIGatewayProxyResult } from 'aws-lambda'; import { logger } from './shared/logger'; import { getSlackSecret, SLACK_SECRET_PREFIX, verifySlackRequest } from './shared/slack-verify'; +import { abcaUserAgent, withAbcaTrace } from './shared/ua'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); const SIGNING_SECRET_ARN = process.env.SLACK_SIGNING_SECRET_ARN!; const TASK_TABLE = process.env.TASK_TABLE_NAME!; diff --git a/cdk/src/handlers/slack-link.ts b/cdk/src/handlers/slack-link.ts index 60ba20dd..9b9e9b7b 100644 --- a/cdk/src/handlers/slack-link.ts +++ b/cdk/src/handlers/slack-link.ts @@ -24,9 +24,10 @@ import { ulid } from 'ulid'; import { extractUserId } from './shared/gateway'; import { logger } from './shared/logger'; import { ErrorCode, errorResponse, successResponse } from './shared/response'; +import { abcaUserAgent, setAbcaTrace, withAbcaTrace } from './shared/ua'; import { parseBody } from './shared/validation'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); const USER_MAPPING_TABLE = process.env.SLACK_USER_MAPPING_TABLE_NAME!; @@ -43,6 +44,7 @@ interface LinkRequest { */ export async function handler(event: APIGatewayProxyEvent): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { const userId = extractUserId(event); diff --git a/cdk/src/handlers/slack-oauth-callback.ts b/cdk/src/handlers/slack-oauth-callback.ts index 872d5581..47cb17f6 100644 --- a/cdk/src/handlers/slack-oauth-callback.ts +++ b/cdk/src/handlers/slack-oauth-callback.ts @@ -23,9 +23,10 @@ import { DynamoDBDocumentClient, PutCommand } from '@aws-sdk/lib-dynamodb'; import type { APIGatewayProxyEvent, APIGatewayProxyResult } from 'aws-lambda'; import { logger } from './shared/logger'; import { getSlackSecret, SLACK_SECRET_PREFIX } from './shared/slack-verify'; +import { abcaUserAgent, withAbcaTrace } from './shared/ua'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); -const sm = new SecretsManagerClient({}); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); +const sm = withAbcaTrace(new SecretsManagerClient(abcaUserAgent())); const TABLE_NAME = process.env.SLACK_INSTALLATION_TABLE_NAME!; const CLIENT_ID_SECRET_ARN = process.env.SLACK_CLIENT_ID_SECRET_ARN!; diff --git a/cdk/src/handlers/webhook-authorizer.ts b/cdk/src/handlers/webhook-authorizer.ts index 91aeb85d..c98fbfbc 100644 --- a/cdk/src/handlers/webhook-authorizer.ts +++ b/cdk/src/handlers/webhook-authorizer.ts @@ -22,8 +22,9 @@ import { DynamoDBDocumentClient, GetCommand } from '@aws-sdk/lib-dynamodb'; import type { APIGatewayRequestAuthorizerEvent, APIGatewayAuthorizerResult } from 'aws-lambda'; import { logger } from './shared/logger'; import type { WebhookRecord } from './shared/types'; +import { abcaUserAgent, withAbcaTrace } from './shared/ua'; -const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); +const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient(abcaUserAgent()))); const TABLE_NAME = process.env.WEBHOOK_TABLE_NAME!; function generatePolicy( diff --git a/cdk/src/handlers/webhook-create-task.ts b/cdk/src/handlers/webhook-create-task.ts index 4b7a86db..a21ed711 100644 --- a/cdk/src/handlers/webhook-create-task.ts +++ b/cdk/src/handlers/webhook-create-task.ts @@ -27,9 +27,10 @@ import { isUsableHmacSecret } from './shared/hmac-secret'; import { logger } from './shared/logger'; import { ErrorCode, errorResponse } from './shared/response'; import type { CreateTaskRequest } from './shared/types'; +import { abcaUserAgent, setAbcaTrace, withAbcaTrace } from './shared/ua'; import { parseBody } from './shared/validation'; -const sm = new SecretsManagerClient({}); +const sm = withAbcaTrace(new SecretsManagerClient(abcaUserAgent())); const SECRET_PREFIX = 'bgagent/webhook/'; // In-memory secret cache with 5-minute TTL @@ -104,6 +105,7 @@ function verifySignature(body: string, secret: string, signature: string): boole */ export async function handler(event: APIGatewayProxyEvent): Promise { const requestId = ulid(); + setAbcaTrace(requestId); try { // 1. Extract webhook auth context (injected by REQUEST authorizer) diff --git a/cdk/src/stacks/agent.ts b/cdk/src/stacks/agent.ts index 0ee8e5cd..f19d79d3 100644 --- a/cdk/src/stacks/agent.ts +++ b/cdk/src/stacks/agent.ts @@ -295,6 +295,9 @@ export class AgentStack extends Stack { // `docs/design/INTERACTIVE_AGENTS.md` §3.1 and AD-1. const runtimeEnvironmentVariables = { GITHUB_TOKEN_SECRET_ARN: githubTokenSecret.secretArn, + // Outbound SDK User-Agent solution tracking (#319); the agent + // hardcodes its component label ('agent') in agent/src/ua.py. + ABCA_STACK_NAME: this.stackName, AWS_REGION: process.env.AWS_REGION ?? 'us-east-1', CLAUDE_CODE_USE_BEDROCK: '1', ANTHROPIC_LOG: 'debug', diff --git a/cdk/test/constructs/task-api.test.ts b/cdk/test/constructs/task-api.test.ts index 32192ff7..5709d415 100644 --- a/cdk/test/constructs/task-api.test.ts +++ b/cdk/test/constructs/task-api.test.ts @@ -109,6 +109,26 @@ describe('TaskApi construct', () => { }); }); + test('Lambdas carry ABCA UA solution-tracking env vars (#319)', () => { + // Every API-surface Lambda gets the stack name + 'api' component label; + // webhook-surface Lambdas get 'webhook'. + baseTemplate.hasResourceProperties('AWS::Lambda::Function', { + Environment: { + Variables: Match.objectLike({ + ABCA_STACK_NAME: 'TestStack', + ABCA_COMPONENT: 'api', + }), + }, + }); + webhookTemplate.hasResourceProperties('AWS::Lambda::Function', { + Environment: { + Variables: Match.objectLike({ + ABCA_COMPONENT: 'webhook', + }), + }, + }); + }); + test('creates a REST API with correct stage name', () => { baseTemplate.resourceCountIs('AWS::ApiGateway::RestApi', 1); baseTemplate.hasResourceProperties('AWS::ApiGateway::RestApi', { diff --git a/cdk/test/constructs/task-orchestrator.test.ts b/cdk/test/constructs/task-orchestrator.test.ts index ebcba02e..0841113c 100644 --- a/cdk/test/constructs/task-orchestrator.test.ts +++ b/cdk/test/constructs/task-orchestrator.test.ts @@ -130,6 +130,17 @@ describe('TaskOrchestrator construct', () => { }); }); + test('orchestrator Lambda carries ABCA UA solution-tracking env vars (#319)', () => { + baseTemplate.hasResourceProperties('AWS::Lambda::Function', { + Environment: { + Variables: Match.objectLike({ + ABCA_STACK_NAME: 'TestStack', + ABCA_COMPONENT: 'orchestr', + }), + }, + }); + }); + test('Lambda function has correct environment variables', () => { baseTemplate.hasResourceProperties('AWS::Lambda::Function', { Environment: { diff --git a/cdk/test/handlers/shared/ua.test.ts b/cdk/test/handlers/shared/ua.test.ts new file mode 100644 index 00000000..39df847c --- /dev/null +++ b/cdk/test/handlers/shared/ua.test.ts @@ -0,0 +1,231 @@ +/** + * MIT No Attribution + * + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * the Software without restriction, including without limitation the rights to + * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of + * the Software, and to permit persons to whom the Software is furnished to do so. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ + +import { DynamoDBClient, ListTablesCommand } from '@aws-sdk/client-dynamodb'; +import { + SOLUTION_ID, + COMPONENT_ENV, + STACK_NAME_ENV, + sanitizeUaValue, + abcaUserAgent, + setAbcaTrace, + getAbcaTrace, + withAbcaTrace, +} from '../../../src/handlers/shared/ua'; + +const ORIGINAL_ENV = process.env; + +beforeEach(() => { + process.env = { ...ORIGINAL_ENV }; + delete process.env[COMPONENT_ENV]; + delete process.env[STACK_NAME_ENV]; + setAbcaTrace(undefined); +}); + +afterAll(() => { + process.env = ORIGINAL_ENV; +}); + +describe('sanitizeUaValue', () => { + // Vectors mirror agent/tests/test_ua.py — the two implementations must + // agree character-for-character. + test('passes through UA-token-safe characters', () => { + expect(sanitizeUaValue('backgroundagent-dev')).toBe('backgroundagent-dev'); + expect(sanitizeUaValue("A1!$%&'*+-.^_`|~z")).toBe("A1!$%&'*+-.^_`|~z"); + }); + + test('replaces structural separators / and #', () => { + expect(sanitizeUaValue('a/b#c')).toBe('a-b-c'); + }); + + test('replaces non-ASCII characters', () => { + expect(sanitizeUaValue('stäck')).toBe('st-ck'); + expect(sanitizeUaValue('名前')).toBe('--'); + }); + + test('replaces whitespace and controls', () => { + expect(sanitizeUaValue('a b\tc\nd')).toBe('a-b-c-d'); + }); + + test('empty string', () => { + expect(sanitizeUaValue('')).toBe(''); + }); +}); + +describe('abcaUserAgent', () => { + test('without stack name emits only the md pair with default component', () => { + expect(abcaUserAgent()).toEqual({ + customUserAgent: [[`md/${SOLUTION_ID}`, 'api']], + }); + }); + + test('component label from ABCA_COMPONENT', () => { + process.env[COMPONENT_ENV] = 'orchestr'; + expect(abcaUserAgent()).toEqual({ + customUserAgent: [[`md/${SOLUTION_ID}`, 'orchestr']], + }); + }); + + test('with stack name emits the app segment first', () => { + process.env[STACK_NAME_ENV] = 'backgroundagent-dev'; + expect(abcaUserAgent()).toEqual({ + customUserAgent: [ + [`app/${SOLUTION_ID}/backgroundagent-dev`], + [`md/${SOLUTION_ID}`, 'api'], + ], + }); + }); + + test('stack name sanitized first, then clipped to 34 (value <= 50)', () => { + process.env[STACK_NAME_ENV] = 'my/stack#nämé' + 'x'.repeat(40); + const [appPair] = abcaUserAgent().customUserAgent; + const appValue = appPair[0].replace(/^app\//, ''); + expect(appValue.startsWith(`${SOLUTION_ID}/my-stack-n-m-`)).toBe(true); + expect(appValue.length).toBeLessThanOrEqual(50); + const stackPart = appValue.slice(`${SOLUTION_ID}/`.length); + expect(stackPart).toHaveLength(34); + expect(stackPart).not.toMatch(/[/#]/); + }); + + test('longest realistic stack name stays exactly at the 50 budget', () => { + process.env[STACK_NAME_ENV] = 'a'.repeat(128); + const [appPair] = abcaUserAgent().customUserAgent; + expect(appPair[0].replace(/^app\//, '')).toHaveLength(50); + }); + + test('blank stack name omits the app segment', () => { + process.env[STACK_NAME_ENV] = ' '; + expect(abcaUserAgent().customUserAgent).toHaveLength(1); + }); + + test('hostile component label is sanitized', () => { + process.env[COMPONENT_ENV] = 'or/ch#str'; + expect(abcaUserAgent().customUserAgent[0]).toEqual([`md/${SOLUTION_ID}`, 'or-ch-str']); + }); +}); + +describe('trace state', () => { + test('defaults to undefined', () => { + expect(getAbcaTrace()).toBeUndefined(); + }); + + test('set and get', () => { + setAbcaTrace('01KTVYABCDEF'); + expect(getAbcaTrace()).toBe('01KTVYABCDEF'); + }); + + test('sanitized on read', () => { + setAbcaTrace('trace/with#bad chars'); + expect(getAbcaTrace()).toBe('trace-with-bad-chars'); + }); + + test('empty string clears', () => { + setAbcaTrace('x'); + setAbcaTrace(''); + expect(getAbcaTrace()).toBeUndefined(); + }); +}); + +describe('withAbcaTrace on mock clients', () => { + test('no-ops on a bare object (jest constructor-mock shape)', () => { + const fake = {}; + expect(withAbcaTrace(fake)).toBe(fake); + }); +}); + +describe('wire capture', () => { + // A real DynamoDBClient with a stub requestHandler: the full middleware + // stack (including the SDK's user-agent middleware and ours) runs, the + // handler records the final headers and returns a canned 200 — no network. + + interface CapturedRequest { + headers: Record; + } + + function makeClient(captured: CapturedRequest[]) { + const client = new DynamoDBClient({ + region: 'us-east-1', + credentials: { accessKeyId: 'testing', secretAccessKey: 'testing' }, + ...abcaUserAgent(), + requestHandler: { + handle: async (request: CapturedRequest) => { + captured.push(request); + return { + response: { + statusCode: 200, + headers: { 'content-type': 'application/x-amz-json-1.0' }, + body: Uint8Array.from(Buffer.from(JSON.stringify({ TableNames: [] }))), + }, + }; + }, + }, + }); + return withAbcaTrace(client); + } + + test('both segments intact in the emitted header; literal slash survives', async () => { + process.env[STACK_NAME_ENV] = 'backgroundagent-dev'; + process.env[COMPONENT_ENV] = 'api'; + const captured: CapturedRequest[] = []; + const client = makeClient(captured); + + await client.send(new ListTablesCommand({})); + + const ua = captured[0].headers['user-agent']; + // The '/' separator survived: the segment rode the raw customUserAgent + // path, NOT the sanitizing app-id config field (which would emit '-'). + expect(ua).toContain(`app/${SOLUTION_ID}/backgroundagent-dev`); + expect(ua).toContain(`md/${SOLUTION_ID}#api`); + // Trace-absent: no trailing '#' after the component label. + expect(ua).not.toContain(`md/${SOLUTION_ID}#api#`); + // customUserAgent also lands in x-amz-user-agent on node. + expect(captured[0].headers['x-amz-user-agent']).toContain(`md/${SOLUTION_ID}#api`); + }); + + test('same cached client emits different traces per request', async () => { + process.env[COMPONENT_ENV] = 'api'; + const captured: CapturedRequest[] = []; + const client = makeClient(captured); + + setAbcaTrace('01KTVYTRACE1'); + await client.send(new ListTablesCommand({})); + setAbcaTrace('01KTVYTRACE2'); + await client.send(new ListTablesCommand({})); + setAbcaTrace(undefined); + await client.send(new ListTablesCommand({})); + + expect(captured[0].headers['user-agent']).toContain(`md/${SOLUTION_ID}#api#01KTVYTRACE1`); + expect(captured[1].headers['user-agent']).toContain(`md/${SOLUTION_ID}#api#01KTVYTRACE2`); + expect(captured[2].headers['user-agent']).toContain(`md/${SOLUTION_ID}#api`); + expect(captured[2].headers['user-agent']).not.toContain(`md/${SOLUTION_ID}#api#`); + }); + + test('trace is sanitized at the wire', async () => { + process.env[COMPONENT_ENV] = 'api'; + const captured: CapturedRequest[] = []; + const client = makeClient(captured); + + setAbcaTrace('evil/trace#☃ value'); + await client.send(new ListTablesCommand({})); + + expect(captured[0].headers['user-agent']).toContain( + `md/${SOLUTION_ID}#api#evil-trace---value`, + ); + }); +}); diff --git a/cli/src/auth.ts b/cli/src/auth.ts index 665a7ce6..9c14cfba 100644 --- a/cli/src/auth.ts +++ b/cli/src/auth.ts @@ -26,6 +26,7 @@ import { loadConfig, loadCredentials, saveCredentials } from './config'; import { debug } from './debug'; import { CliError } from './errors'; import { Credentials } from './types'; +import { abcaUserAgent, withAbcaTrace } from './ua'; const TOKEN_REFRESH_BUFFER_MINUTES = 5; const TOKEN_REFRESH_BUFFER_MS = TOKEN_REFRESH_BUFFER_MINUTES * 60 * 1000; @@ -45,7 +46,7 @@ let inFlightRefresh: Promise | null = null; export async function login(username: string, password: string): Promise { const config = loadConfig(); debug(`Cognito region: ${config.region}, client_id: ${config.client_id}, user_pool_id: ${config.user_pool_id}`); - const client = new CognitoIdentityProviderClient({ region: config.region }); + const client = withAbcaTrace(new CognitoIdentityProviderClient({ region: config.region, ...abcaUserAgent() })); const result = await client.send(new InitiateAuthCommand({ AuthFlow: AuthFlowType.USER_PASSWORD_AUTH, @@ -126,7 +127,7 @@ function isExpired(creds: Credentials): boolean { async function refreshToken(creds: Credentials): Promise { const config = loadConfig(); - const client = new CognitoIdentityProviderClient({ region: config.region }); + const client = withAbcaTrace(new CognitoIdentityProviderClient({ region: config.region, ...abcaUserAgent() })); try { const result = await client.send(new InitiateAuthCommand({ diff --git a/cli/src/bin/bgagent.ts b/cli/src/bin/bgagent.ts index bc0ed285..3e3db096 100644 --- a/cli/src/bin/bgagent.ts +++ b/cli/src/bin/bgagent.ts @@ -41,6 +41,11 @@ import { makeWatchCommand } from '../commands/watch'; import { makeWebhookCommand } from '../commands/webhook'; import { setVerbose } from '../debug'; import { ApiError, CliError } from '../errors'; +import { setAbcaTrace } from '../ua'; + +// User-Agent solution tracking (#319): the pid correlates all AWS calls +// made by this CLI invocation. +setAbcaTrace(String(process.pid)); const program = new Command(); diff --git a/cli/src/commands/admin.ts b/cli/src/commands/admin.ts index dad20771..f0e44554 100644 --- a/cli/src/commands/admin.ts +++ b/cli/src/commands/admin.ts @@ -29,6 +29,7 @@ import { Command } from 'commander'; import { getConfigDir, loadConfig, SECRET_FILE_MODE } from '../config'; import { CliError } from '../errors'; import { CliConfig } from '../types'; +import { abcaUserAgent, withAbcaTrace } from '../ua'; /** * Generate a strong temporary password meeting Cognito's default policy: @@ -151,7 +152,7 @@ export function makeAdminCommand(): Command { const tempPassword = opts.tempPassword ?? generateTempPassword(); - const cognito = new CognitoIdentityProviderClient({ region }); + const cognito = withAbcaTrace(new CognitoIdentityProviderClient({ region, ...abcaUserAgent() })); try { await cognito.send(new AdminCreateUserCommand({ UserPoolId: config.user_pool_id, diff --git a/cli/src/commands/configure.ts b/cli/src/commands/configure.ts index a5d2eb15..09856d73 100644 --- a/cli/src/commands/configure.ts +++ b/cli/src/commands/configure.ts @@ -37,6 +37,7 @@ export function makeConfigureCommand(): Command { .option('--region ', 'AWS region') .option('--user-pool-id ', 'Cognito User Pool ID') .option('--client-id ', 'Cognito App Client ID') + .option('--stack-name ', 'Deployed stack name (User-Agent solution tracking)') .option('--from-bundle ', 'Base64 config bundle from `bgagent admin invite-user`') .action((opts) => { // --from-bundle is mutually exclusive with the individual flags. Mixing @@ -58,6 +59,7 @@ export function makeConfigureCommand(): Command { ...(opts.region !== undefined ? { region: opts.region } : {}), ...(opts.userPoolId !== undefined ? { user_pool_id: opts.userPoolId } : {}), ...(opts.clientId !== undefined ? { client_id: opts.clientId } : {}), + ...(opts.stackName !== undefined ? { stack_name: opts.stackName } : {}), }; } const merged: Partial = { diff --git a/cli/src/commands/github.ts b/cli/src/commands/github.ts index 53d811aa..3d3ea3a3 100644 --- a/cli/src/commands/github.ts +++ b/cli/src/commands/github.ts @@ -26,6 +26,7 @@ import { import { Command } from 'commander'; import { loadConfig } from '../config'; import { CliError } from '../errors'; +import { abcaUserAgent, withAbcaTrace } from '../ua'; /** Width of the `═` banner rules printed around webhook-info output. */ const BANNER_WIDTH = 72; @@ -116,7 +117,7 @@ export function makeGithubCommand(): Command { ); } - const sm = new SecretsManagerClient({ region }); + const sm = withAbcaTrace(new SecretsManagerClient({ region, ...abcaUserAgent() })); // Show whether a secret is already configured so the operator // doesn't accidentally rotate it without realising. Linear's @@ -166,7 +167,7 @@ export function makeGithubCommand(): Command { // ─── Stack-output helper ───────────────────────────────────────────────────── async function getStackOutput(region: string, stackName: string, outputKey: string): Promise { - const cf = new CloudFormationClient({ region }); + const cf = withAbcaTrace(new CloudFormationClient({ region, ...abcaUserAgent() })); try { const result = await cf.send(new DescribeStacksCommand({ StackName: stackName })); const stack = result.Stacks?.[0]; diff --git a/cli/src/commands/linear.ts b/cli/src/commands/linear.ts index 475dbebe..b5a82fa4 100644 --- a/cli/src/commands/linear.ts +++ b/cli/src/commands/linear.ts @@ -45,6 +45,7 @@ import { StoredLinearOauthToken, } from '../linear-oauth'; import { awaitOauthCallback, CALLBACK_URL } from '../oauth-callback-server'; +import { abcaUserAgent, withAbcaTrace } from '../ua'; /** Default label that triggers an ABCA task when applied to a Linear issue. */ const DEFAULT_LABEL_FILTER = 'bgagent'; @@ -597,7 +598,7 @@ export function makeLinearCommand(): Command { // ─── Step 4: Persist token to per-workspace Secrets Manager ─── process.stdout.write(' → Storing OAuth token...'); - const sm = new SecretsManagerClient({ region }); + const sm = withAbcaTrace(new SecretsManagerClient({ region, ...abcaUserAgent() })); const now = new Date().toISOString(); const stored: StoredLinearOauthToken = { access_token: tokenResponse.access_token, @@ -625,7 +626,7 @@ export function makeLinearCommand(): Command { console.log(` ✓ (${secretName})`); // ─── Step 5: Persist registry + user-mapping rows ───────────── - const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({ region })); + const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient({ region, ...abcaUserAgent() }))); // Best-effort: fetch team keys so the screenshot processor can // prefix-route Linear issue lookups (e.g. ABCA-42 → workspace @@ -829,8 +830,8 @@ export function makeLinearCommand(): Command { ); } - const sm = new SecretsManagerClient({ region }); - const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({ region })); + const sm = withAbcaTrace(new SecretsManagerClient({ region, ...abcaUserAgent() })); + const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient({ region, ...abcaUserAgent() }))); // ─── Linear OAuth app credentials ────────────────────────────── // Always prompt — never accept secrets via flags (shell history @@ -1099,7 +1100,7 @@ export function makeLinearCommand(): Command { const config = loadConfig(); const region = opts.region || config.region; - const sm = new SecretsManagerClient({ region }); + const sm = withAbcaTrace(new SecretsManagerClient({ region, ...abcaUserAgent() })); const secretName = linearOauthSecretName(slug); // ─── Read existing bundle ─────────────────────────────────── @@ -1218,8 +1219,8 @@ export function makeLinearCommand(): Command { const callerCognitoSub = extractCognitoSub(); // ─── Resolve workspace + OAuth secret arn ────────────────────── - const sm = new SecretsManagerClient({ region }); - const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({ region })); + const sm = withAbcaTrace(new SecretsManagerClient({ region, ...abcaUserAgent() })); + const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient({ region, ...abcaUserAgent() }))); const registryScan = await ddb.send(new ScanCommand({ TableName: workspaceRegistryTable!, FilterExpression: 'workspace_slug = :slug AND #status = :active', @@ -1342,7 +1343,7 @@ export function makeLinearCommand(): Command { } const now = new Date().toISOString(); - const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({ region })); + const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient({ region, ...abcaUserAgent() }))); await ddb.send(new PutCommand({ TableName: tableName, Item: { @@ -1373,7 +1374,7 @@ export function makeLinearCommand(): Command { .action(async (opts) => { const config = loadConfig(); const region = opts.region || config.region; - const sm = new SecretsManagerClient({ region }); + const sm = withAbcaTrace(new SecretsManagerClient({ region, ...abcaUserAgent() })); // Resolve the set of workspace slugs to query. Either an // explicit `--slug` (one workspace) or every Linear workspace @@ -1908,7 +1909,7 @@ export async function autoLinkTokenOwner(args: { return; } - const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({ region: args.region })); + const ddb = DynamoDBDocumentClient.from(withAbcaTrace(new DynamoDBClient({ region: args.region, ...abcaUserAgent() }))); await ddb.send(new PutCommand({ TableName: args.userMappingTable, Item: { @@ -1947,7 +1948,7 @@ function extractCognitoSub(): string { async function getStackOutput(region: string, stackName: string, outputKey: string): Promise { try { - const cfn = new CloudFormationClient({ region }); + const cfn = withAbcaTrace(new CloudFormationClient({ region, ...abcaUserAgent() })); const result = await cfn.send(new DescribeStacksCommand({ StackName: stackName })); const outputs = result.Stacks?.[0]?.Outputs ?? []; const output = outputs.find((o) => o.OutputKey === outputKey); diff --git a/cli/src/commands/slack.ts b/cli/src/commands/slack.ts index bc270fa2..e6bdfc60 100644 --- a/cli/src/commands/slack.ts +++ b/cli/src/commands/slack.ts @@ -27,6 +27,7 @@ import { Command } from 'commander'; import { ApiClient } from '../api-client'; import { loadConfig } from '../config'; import { formatJson } from '../format'; +import { abcaUserAgent, withAbcaTrace } from '../ua'; export function makeSlackCommand(): Command { const slack = new Command('slack') @@ -208,7 +209,7 @@ async function promptAndStoreCredentials(region: string, arns: SecretArns): Prom // Store in Secrets Manager. console.log(''); - const sm = new SecretsManagerClient({ region }); + const sm = withAbcaTrace(new SecretsManagerClient({ region, ...abcaUserAgent() })); const secrets = [ { id: arns.signingSecretArn, value: signingSecret, label: 'signing secret' }, @@ -345,7 +346,7 @@ function findRepoRoot(): string { async function getStackOutput(region: string, stackName: string, outputKey: string): Promise { try { - const cfn = new CloudFormationClient({ region }); + const cfn = withAbcaTrace(new CloudFormationClient({ region, ...abcaUserAgent() })); const result = await cfn.send(new DescribeStacksCommand({ StackName: stackName })); const outputs = result.Stacks?.[0]?.Outputs ?? []; const output = outputs.find((o) => o.OutputKey === outputKey); diff --git a/cli/src/types.ts b/cli/src/types.ts index c27c23ae..57f7a176 100644 --- a/cli/src/types.ts +++ b/cli/src/types.ts @@ -372,6 +372,8 @@ export interface CliConfig { readonly region: string; readonly user_pool_id: string; readonly client_id: string; + /** Deployed stack name for User-Agent solution tracking (#319); optional. */ + readonly stack_name?: string; } /** Cached credentials stored in ~/.bgagent/credentials.json. diff --git a/cli/src/ua.ts b/cli/src/ua.ts new file mode 100644 index 00000000..c9c7c0a8 --- /dev/null +++ b/cli/src/ua.ts @@ -0,0 +1,137 @@ +/** + * MIT No Attribution + * + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * the Software without restriction, including without limitation the rights to + * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of + * the Software, and to permit persons to whom the Software is furnished to do so. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ + +/** + * Outbound AWS SDK User-Agent solution tracking (#319) — CLI surface. + * + * Every AWS API call made by `bgagent` carries: + * + * app/uksb-wt64nei4u6/{STACKNAME} (only when config has stack_name) + * md/uksb-wt64nei4u6#cli[#{PID}] + * + * CLI-local mirror of `cdk/src/handlers/shared/ua.ts` (the CLI package + * cannot import from the CDK package — same mirroring convention as + * `cli/src/types.ts`). Solution id, wire format, and sanitization rules + * must stay identical; `agent/src/ua.py` is the Python counterpart. + * + * The component label is hardcoded (`cli`); the stack name comes from the + * optional `stack_name` field in `~/.bgagent/config.json`. The trace handle + * is the CLI process pid, set once at startup in `bin/bgagent.ts` and + * appended per-request by the {@link withAbcaTrace} middleware. + */ + +import { tryLoadConfig } from './config'; + +/** + * AWS solution-tracking id for ABCA. Deploy-time counterpart (#292) lives in + * the CloudFormation stack description in `cdk/src/main.ts`. + */ +export const SOLUTION_ID = 'uksb-wt64nei4u6'; + +/** Stable per-component label: this surface IS the bgagent CLI. */ +const COMPONENT = 'cli'; + +/** App-id budget: 50-char value cap minus `uksb-wt64nei4u6/` (16) = 34. */ +const STACK_NAME_MAX = 34; + +/** + * RFC 7230 token charset; `/` and `#` deliberately excluded (structural + * separators of the scheme). Mirrors the CDK and Python implementations. + */ +const UA_TOKEN_SAFE = /[^A-Za-z0-9!$%&'*+\-.^_`|~]/g; + +let currentTrace: string | undefined; + +/** Replace every non-UA-token char (incl. non-ASCII) with `-`. */ +export function sanitizeUaValue(raw: string): string { + return raw.replace(UA_TOKEN_SAFE, '-'); +} + +/** + * Client config fragment carrying the static ABCA UA segments. Spread into + * every SDK client constructor: `new SecretsManagerClient({ region, ...abcaUserAgent() })`. + * + * Pair semantics (mirrors the CDK module): the `app/` segment is a + * single-element pair so its literal `/` separators survive the SDK's + * name-position escaping; the `md/` pair lets the SDK's own `name#value` + * join produce the `#`. + */ +export function abcaUserAgent(): { customUserAgent: ([string] | [string, string])[] } { + const pairs: ([string] | [string, string])[] = []; + const stackName = tryLoadConfig()?.stack_name?.trim(); + if (stackName) { + // Sanitize FIRST, then clip, so a replaced char can't be re-split. + const clipped = sanitizeUaValue(stackName).slice(0, STACK_NAME_MAX); + pairs.push([`app/${SOLUTION_ID}/${clipped}`]); + } + pairs.push([`md/${SOLUTION_ID}`, COMPONENT]); + return { customUserAgent: pairs }; +} + +/** Set (or clear) the ambient trace handle (the CLI pid, set at startup). */ +export function setAbcaTrace(handle?: string): void { + currentTrace = handle || undefined; +} + +/** Current trace handle, sanitized to UA-token-safe ASCII, or undefined. */ +export function getAbcaTrace(): string | undefined { + return currentTrace ? sanitizeUaValue(currentTrace) : undefined; +} + +/** Structural view of a client middleware stack (avoids @smithy/types dep). */ +interface MiddlewareStackLike { + addRelativeTo(middleware: unknown, options: Record): void; +} + +/** + * Append `#{TRACE}` to the outgoing User-Agent headers on every request by + * splicing onto the static `md/` segment, after the SDK's own + * `getUserAgentMiddleware` has rendered the headers. Mutates only the + * header strings; the client and its connection pool are untouched. + * No-ops on clients without a middleware stack (jest constructor mocks). + */ +export function withAbcaTrace(client: T): T { + const stack = (client as { middlewareStack?: MiddlewareStackLike }).middlewareStack; + if (!stack || typeof stack.addRelativeTo !== 'function') { + return client; + } + const md = `md/${SOLUTION_ID}#${COMPONENT}`; + stack.addRelativeTo( + (next: (args: unknown) => Promise) => async (args: unknown) => { + const trace = getAbcaTrace(); + const request = (args as { request?: { headers?: Record } }).request; + if (trace && request?.headers) { + for (const header of ['user-agent', 'x-amz-user-agent']) { + const value = request.headers[header]; + if (value && value.includes(md)) { + request.headers[header] = value.replace(md, `${md}#${trace}`); + } + } + } + return next(args); + }, + { + name: 'abcaUaTraceMiddleware', + relation: 'after', + toMiddleware: 'getUserAgentMiddleware', + override: true, + }, + ); + return client; +} diff --git a/cli/test/auth.test.ts b/cli/test/auth.test.ts index 24ab4fec..ae271a69 100644 --- a/cli/test/auth.test.ts +++ b/cli/test/auth.test.ts @@ -52,6 +52,19 @@ describe('auth', () => { }); describe('login', () => { + test('Cognito client carries the ABCA solution customUserAgent (#319)', async () => { + const { CognitoIdentityProviderClient } = jest.requireMock( + '@aws-sdk/client-cognito-identity-provider', + ); + mockSend.mockResolvedValue({ + AuthenticationResult: { IdToken: 'i', AccessToken: 'a', RefreshToken: 'r', ExpiresIn: 3600 }, + }); + await login('user', 'pass'); + const calls = (CognitoIdentityProviderClient as jest.Mock).mock.calls; + const config = calls[calls.length - 1]?.[0]; + expect(config.customUserAgent).toEqual([['md/uksb-wt64nei4u6', 'cli']]); + }); + test('saves credentials on successful login', async () => { mockSend.mockResolvedValue({ AuthenticationResult: { diff --git a/cli/test/ua.test.ts b/cli/test/ua.test.ts new file mode 100644 index 00000000..689505ca --- /dev/null +++ b/cli/test/ua.test.ts @@ -0,0 +1,177 @@ +/** + * MIT No Attribution + * + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * the Software without restriction, including without limitation the rights to + * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of + * the Software, and to permit persons to whom the Software is furnished to do so. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ + +import * as fs from 'fs'; +import * as os from 'os'; +import * as path from 'path'; +import { saveConfig } from '../src/config'; +import { + SOLUTION_ID, + sanitizeUaValue, + abcaUserAgent, + setAbcaTrace, + getAbcaTrace, + withAbcaTrace, +} from '../src/ua'; + +let tmpDir: string; + +function writeConfig(stackName?: string): void { + saveConfig({ + api_url: 'https://api.example.com', + region: 'us-east-1', + user_pool_id: 'us-east-1_abc', + client_id: 'client123', + ...(stackName !== undefined ? { stack_name: stackName } : {}), + }); +} + +beforeEach(() => { + tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'bgagent-ua-test-')); + process.env.BGAGENT_CONFIG_DIR = tmpDir; + setAbcaTrace(undefined); +}); + +afterEach(() => { + delete process.env.BGAGENT_CONFIG_DIR; + fs.rmSync(tmpDir, { recursive: true, force: true }); +}); + +describe('sanitizeUaValue', () => { + // Vectors mirror cdk/test/handlers/shared/ua.test.ts and + // agent/tests/test_ua.py — all three implementations must agree. + test('passes through UA-token-safe characters', () => { + expect(sanitizeUaValue('backgroundagent-dev')).toBe('backgroundagent-dev'); + expect(sanitizeUaValue("A1!$%&'*+-.^_`|~z")).toBe("A1!$%&'*+-.^_`|~z"); + }); + + test('replaces structural separators, non-ASCII, whitespace', () => { + expect(sanitizeUaValue('a/b#c')).toBe('a-b-c'); + expect(sanitizeUaValue('stäck')).toBe('st-ck'); + expect(sanitizeUaValue('a b\tc')).toBe('a-b-c'); + }); +}); + +describe('abcaUserAgent', () => { + test('without configured stack name emits only the md pair', () => { + writeConfig(); + expect(abcaUserAgent()).toEqual({ + customUserAgent: [[`md/${SOLUTION_ID}`, 'cli']], + }); + }); + + test('with no config file at all still emits the md pair', () => { + expect(abcaUserAgent()).toEqual({ + customUserAgent: [[`md/${SOLUTION_ID}`, 'cli']], + }); + }); + + test('with stack name emits the app segment first', () => { + writeConfig('backgroundagent-dev'); + expect(abcaUserAgent()).toEqual({ + customUserAgent: [ + [`app/${SOLUTION_ID}/backgroundagent-dev`], + [`md/${SOLUTION_ID}`, 'cli'], + ], + }); + }); + + test('stack name sanitized first, then clipped to 34 (value <= 50)', () => { + writeConfig('my/stack#nämé' + 'x'.repeat(40)); + const [appPair] = abcaUserAgent().customUserAgent; + const appValue = appPair[0].replace(/^app\//, ''); + expect(appValue.startsWith(`${SOLUTION_ID}/my-stack-n-m-`)).toBe(true); + expect(appValue.length).toBeLessThanOrEqual(50); + expect(appValue.slice(`${SOLUTION_ID}/`.length)).toHaveLength(34); + }); +}); + +describe('trace state', () => { + test('set, get, sanitize, clear', () => { + expect(getAbcaTrace()).toBeUndefined(); + setAbcaTrace('12345'); + expect(getAbcaTrace()).toBe('12345'); + setAbcaTrace('bad/pid#1'); + expect(getAbcaTrace()).toBe('bad-pid-1'); + setAbcaTrace(undefined); + expect(getAbcaTrace()).toBeUndefined(); + }); +}); + +describe('withAbcaTrace', () => { + test('no-ops on a bare object (jest constructor-mock shape)', () => { + const fake = {}; + expect(withAbcaTrace(fake)).toBe(fake); + }); +}); + +describe('wire capture', () => { + // Real Cognito client + stub requestHandler: the full middleware stack + // runs, the handler records the final headers — no network. + interface CapturedRequest { + headers: Record; + } + + test('both segments intact; same client emits per-request traces', async () => { + writeConfig('backgroundagent-dev'); + const { + CognitoIdentityProviderClient, + GetUserCommand, + } = jest.requireActual('@aws-sdk/client-cognito-identity-provider'); + + const captured: CapturedRequest[] = []; + const client = withAbcaTrace( + new CognitoIdentityProviderClient({ + region: 'us-east-1', + credentials: { accessKeyId: 'testing', secretAccessKey: 'testing' }, + ...abcaUserAgent(), + requestHandler: { + handle: async (request: CapturedRequest) => { + captured.push(request); + return { + response: { + statusCode: 200, + headers: { 'content-type': 'application/x-amz-json-1.1' }, + body: Uint8Array.from( + Buffer.from(JSON.stringify({ Username: 'u', UserAttributes: [] })), + ), + }, + }; + }, + }, + }), + ); + + setAbcaTrace('4242'); + await client.send(new GetUserCommand({ AccessToken: 'token' })); + setAbcaTrace('9999'); + await client.send(new GetUserCommand({ AccessToken: 'token' })); + setAbcaTrace(undefined); + await client.send(new GetUserCommand({ AccessToken: 'token' })); + + const ua0 = captured[0].headers['user-agent']; + // Literal '/' survived — raw customUserAgent path, not the app-id field. + expect(ua0).toContain(`app/${SOLUTION_ID}/backgroundagent-dev`); + expect(ua0).toContain(`md/${SOLUTION_ID}#cli#4242`); + expect(captured[1].headers['user-agent']).toContain(`md/${SOLUTION_ID}#cli#9999`); + // Trace-absent: segment ends at the component label, no trailing '#'. + expect(captured[2].headers['user-agent']).toContain(`md/${SOLUTION_ID}#cli`); + expect(captured[2].headers['user-agent']).not.toContain(`md/${SOLUTION_ID}#cli#`); + }); +});