aws-samples
diff --git a/‎agent/src/observability.py‎
Lines changed: 41 additions & 11 deletions b/‎agent/src/observability.py‎
Lines changed: 41 additions & 11 deletions
diff --git a/‎agent/src/pipeline.py‎
Lines changed: 6 additions & 1 deletion b/‎agent/src/pipeline.py‎
Lines changed: 6 additions & 1 deletion
diff --git a/‎agent/src/progress_writer.py‎
Lines changed: 25 additions & 1 deletion b/‎agent/src/progress_writer.py‎
Lines changed: 25 additions & 1 deletion
diff --git a/‎agent/src/runner.py‎
Lines changed: 6 additions & 1 deletion b/‎agent/src/runner.py‎
Lines changed: 6 additions & 1 deletion
diff --git a/‎agent/src/server.py‎
Lines changed: 7 additions & 5 deletions b/‎agent/src/server.py‎
Lines changed: 7 additions & 5 deletions
diff --git a/‎agent/tests/test_observability.py‎
Lines changed: 43 additions & 0 deletions b/‎agent/tests/test_observability.py‎
Lines changed: 43 additions & 0 deletions
diff --git a/‎agent/tests/test_progress_writer.py‎
Lines changed: 53 additions & 0 deletions b/‎agent/tests/test_progress_writer.py‎
Lines changed: 53 additions & 0 deletions
diff --git a/‎agent/tests/test_server.py‎
Lines changed: 35 additions & 0 deletions b/‎agent/tests/test_server.py‎
Lines changed: 35 additions & 0 deletions
diff --git a/‎cdk/src/handlers/get-task-events.ts‎
Lines changed: 7 additions & 1 deletion b/‎cdk/src/handlers/get-task-events.ts‎
Lines changed: 7 additions & 1 deletion
diff --git a/‎cdk/src/handlers/get-task-replay.ts‎
Lines changed: 4 additions & 0 deletions b/‎cdk/src/handlers/get-task-replay.ts‎
Lines changed: 4 additions & 0 deletions
@@ -64,24 +64,54 @@ def current_otel_trace_id() -> str | None:
     bundle) so operators can join the task to its CloudWatch/X-Ray trace. Returns
     ``None`` when there is no recording span (e.g. tracing disabled locally) or
     the context is invalid, so callers can treat it as a graceful-missing field.
+
+    Never raises: a broken/misconfigured tracer degrades to ``None`` rather than
+    propagating. Callers read this inside DDB-write try-blocks (progress_writer),
+    where a raised trace error would otherwise be misclassified as a DDB failure
+    and trip the shared progress circuit breaker.
     """
-    span = trace.get_current_span()
-    ctx = span.get_span_context()
-    if not ctx.is_valid:
+    try:
+        span = trace.get_current_span()
+        ctx = span.get_span_context()
+        if not ctx.is_valid:
+            return None
+        # format_trace_id renders the 128-bit id as zero-padded 32-char hex — the
+        # OTEL format, so it joins directly in CloudWatch Transaction Search. Note
+        # the X-Ray console renders trace ids as ``1-{8hex}-{24hex}``; to look this
+        # up there, transform to that form (the timestamp is the first 8 hex chars).
+        return trace.format_trace_id(ctx.trace_id)
+    except Exception:
+        # nosemgrep: py-silent-success-masking -- trace id is a graceful-missing
+        # correlation field; a tracer fault must not fail the caller's write path.
         return None
-    # format_trace_id renders the 128-bit id as zero-padded 32-char hex — the
-    # OTEL format, so it joins directly in CloudWatch Transaction Search. Note
-    # the X-Ray console renders trace ids as ``1-{8hex}-{24hex}``; to look this
-    # up there, transform to that form (the timestamp is the first 8 hex chars).
-    return trace.format_trace_id(ctx.trace_id)
 
 
-def set_session_id(session_id: str) -> None:
-    """Propagate *session_id* via OTEL baggage for AgentCore session correlation.
+def propagate_correlation_context(
+    session_id: str,
+    user_id: str = "",
+    # ``user_id`` uses ""-means-absent (Cognito sub, mirrors AgentConfig.user_id
+    # which is never None); ``repo`` uses None-means-absent (mirrors the optional
+    # TaskRecord.repo). Both conventions flow from upstream config types; the
+    # ``if x:`` guards below flatten either to "don't set the baggage key".
+    repo: str | None = None,
+) -> None:
+    """Propagate the correlation envelope via OTEL baggage.
+
+    *session_id* correlates custom spans to the AgentCore session; *user_id*
+    and *repo* (#245) carry the platform identity and target repo so baggage
+    survives across pipeline phases on the task thread. Empty/None fields are
+    not set — so this runs (and is useful) even when *session_id* is empty but
+    the identity is known. *repo* is None for repo-less workflows (#248 Phase 3).
 
     The attached context is intentionally not detached: the background thread
     runs a single task then exits, so the context is garbage-collected with the
     thread.
     """
-    ctx = baggage.set_baggage("session.id", session_id)
+    ctx = context.get_current()
+    if session_id:
+        ctx = baggage.set_baggage("session.id", session_id, context=ctx)
+    if user_id:
+        ctx = baggage.set_baggage("user.id", user_id, context=ctx)
+    if repo:
+        ctx = baggage.set_baggage("repo.url", repo, context=ctx)
     context.attach(ctx)  # token not stored — thread-scoped lifetime
@@ -679,13 +679,18 @@ def run_task(
             "repo.url": config.repo_url,
             "issue.number": config.issue_number,
             "agent.model": config.anthropic_model,
+            # Correlation envelope (#245): user.id joins agent spans to
+            # orchestrator logs by the platform identity, not just task/repo.
+            **({"user.id": config.user_id} if config.user_id else {}),
         },
     ) as root_span:
         task_state.write_running(config.task_id)
         task_state.write_heartbeat(config.task_id)
 
         agent_result: AgentResult | None = None
-        progress = _ProgressWriter(config.task_id, trace=trace)
+        progress = _ProgressWriter(
+            config.task_id, trace=trace, user_id=config.user_id, repo=config.repo_url
+        )
         # #251: clear any blocker latched by a prior task. The agent container
         # is one-task-per-process today, but the FastAPI server thread-pool can
         # in principle dispatch a second run_task in the same process — reset
 
@@ -38,6 +38,8 @@
 from decimal import Decimal
 from typing import Literal
 
+from observability import current_otel_trace_id
+
 # Preview field cap defaults (design §10.1):
 #   - 200 chars for normal tasks — small DDB rows, cheap watch-stream bytes.
 #   - 4096 chars (4 KB) for ``--trace`` opt-in tasks — full enough to
@@ -373,8 +375,19 @@ class _ProgressWriter:
 
     _MAX_FAILURES = 3
 
-    def __init__(self, task_id: str, trace: bool = False) -> None:
+    def __init__(
+        self,
+        task_id: str,
+        trace: bool = False,
+        user_id: str = "",
+        repo: str | None = None,
+    ) -> None:
         self._task_id = task_id
+        # Correlation envelope (#245): stamped on every event so the agent-side
+        # event stream joins to orchestrator logs and the X-Ray trace. ``repo``
+        # is None for repo-less workflows (#248 Phase 3).
+        self._user_id = user_id
+        self._repo = repo
         self._table_name = os.environ.get("TASK_EVENTS_TABLE_NAME")
         self._table = None
         # Per-instance preview cap — design §10.1. ``trace=True`` raises
@@ -467,6 +480,11 @@ def _put_event(self, event_type: str, metadata: dict) -> None:
                 return
 
             now = datetime.now(UTC)
+            # Correlation envelope (#245): trace_id is read per-event from the
+            # active span (it may be None early, before the root span opens);
+            # user_id/repo are stamped when known. All are omitted when
+            # empty/None so the item shape stays stable for consumers.
+            trace_id = current_otel_trace_id()
             item = {
                 "task_id": self._task_id,
                 "event_id": _generate_ulid(),
@@ -478,6 +496,12 @@ def _put_event(self, event_type: str, metadata: dict) -> None:
                 "timestamp": now.isoformat(),
                 "ttl": int(now.timestamp()) + _TTL_SECONDS,
             }
+            if self._user_id:
+                item["user_id"] = self._user_id
+            if self._repo:
+                item["repo"] = self._repo
+            if trace_id:
+                item["trace_id"] = trace_id
             self._table.put_item(Item=item)
 
             # Success: reset the shared failure counter.  We do NOT flip
 
@@ -483,7 +483,12 @@ def _on_stderr(line: str) -> None:
     # invocations we fall back to a fresh writer with no accumulator.
     if trajectory is None:
         trajectory = _TrajectoryWriter(config.task_id or "unknown")
-    progress = _ProgressWriter(config.task_id or "unknown", trace=config.trace)
+    progress = _ProgressWriter(
+        config.task_id or "unknown",
+        trace=config.trace,
+        user_id=config.user_id,
+        repo=config.repo_url,
+    )
 
     # Map tool_use_id → tool_name so we can label ToolResultBlocks that arrive
     # in UserMessages (ToolResultBlock carries only the id, not the name).
 
@@ -26,7 +26,7 @@
 import task_state
 from config import resolve_github_token
 from models import TaskResult
-from observability import set_session_id
+from observability import propagate_correlation_context
 from pipeline import run_task
 
 # --- _debug_cw / _warn_cw failure counter -------------------------------
@@ -452,10 +452,12 @@ def _run_task_background(
         hb_thread.start()
 
     try:
-        # Propagate session ID into this thread's OTEL context so spans
-        # are correlated with the AgentCore session in CloudWatch.
-        if session_id:
-            set_session_id(session_id)
+        # Propagate the correlation envelope into this thread's OTEL context
+        # so spans are correlated with the AgentCore session and the platform
+        # identity in CloudWatch (#245). Runs whenever any field is present —
+        # session_id may be empty while user_id/repo are known.
+        if session_id or user_id or repo_url:
+            propagate_correlation_context(session_id, user_id=user_id, repo=repo_url or None)
 
         run_task(
             repo_url=repo_url,
 
@@ -44,3 +44,46 @@ def test_returns_none_for_invalid_context(self):
         span.get_span_context.return_value = ctx
         with patch.object(observability.trace, "get_current_span", return_value=span):
             assert observability.current_otel_trace_id() is None
+
+    def test_degrades_to_none_when_tracer_raises(self):
+        # A broken/misconfigured tracer must not propagate: callers read this
+        # inside DDB-write try-blocks, where a raise would be misclassified as a
+        # DDB failure and trip the shared progress circuit breaker (#245 review).
+        with patch.object(
+            observability.trace, "get_current_span", side_effect=RuntimeError("tracer boom")
+        ):
+            assert observability.current_otel_trace_id() is None
+
+
+class TestPropagateCorrelationContext:
+    """``propagate_correlation_context`` propagates the correlation envelope
+    (#245) via OTEL baggage, setting only the fields that are present."""
+
+    def test_sets_all_envelope_fields_when_present(self):
+        with (
+            patch.object(observability, "baggage") as bag,
+            patch.object(observability, "context") as ctx,
+        ):
+            bag.set_baggage.return_value = "CTX"
+            observability.propagate_correlation_context("sess-1", user_id="user-1", repo="org/repo")
+        # Baggage-key ordering is not part of the contract — compare as a set.
+        keys = {c.args[0] for c in bag.set_baggage.call_args_list}
+        assert keys == {"session.id", "user.id", "repo.url"}
+        ctx.attach.assert_called_once()
+
+    def test_omits_absent_fields(self):
+        # Empty user_id and None repo → only session.id is set on the baggage.
+        with patch.object(observability, "baggage") as bag, patch.object(observability, "context"):
+            bag.set_baggage.return_value = "CTX"
+            observability.propagate_correlation_context("sess-1")
+        keys = {c.args[0] for c in bag.set_baggage.call_args_list}
+        assert keys == {"session.id"}
+
+    def test_empty_session_id_is_not_stamped(self):
+        # Reachable via server.py's widened trigger (user_id known, no session):
+        # an empty session_id must not write an empty-string session.id baggage.
+        with patch.object(observability, "baggage") as bag, patch.object(observability, "context"):
+            bag.set_baggage.return_value = "CTX"
+            observability.propagate_correlation_context("", user_id="user-1")
+        keys = {c.args[0] for c in bag.set_baggage.call_args_list}
+        assert keys == {"user.id"}
@@ -267,6 +267,59 @@ def test_ttl_is_90_days_from_now(self, writer, mock_table):
         assert before + ttl_90_days <= item["ttl"] <= after + ttl_90_days + 1
 
 
+# ---------------------------------------------------------------------------
+# _ProgressWriter — correlation envelope (#245)
+# ---------------------------------------------------------------------------
+
+
+class TestCorrelationEnvelope:
+    """Every emitted event stamps {user_id, repo, trace_id} when known, so the
+    agent-side event stream joins to orchestrator logs and the X-Ray trace."""
+
+    def _writer(self, monkeypatch, **kwargs):
+        monkeypatch.setenv("TASK_EVENTS_TABLE_NAME", "events-table")
+        writer = _ProgressWriter("task-42", **kwargs)
+        writer._table = MagicMock()
+        return writer
+
+    def test_stamps_user_repo_trace_when_present(self, monkeypatch):
+        writer = self._writer(monkeypatch, user_id="user-1", repo="org/repo")
+        with patch("progress_writer.current_otel_trace_id", return_value="a" * 32):
+            writer.write_agent_milestone("m", "")
+        item = writer._table.put_item.call_args[1]["Item"]
+        assert item["user_id"] == "user-1"
+        assert item["repo"] == "org/repo"
+        assert item["trace_id"] == "a" * 32
+
+    def test_repo_less_omits_repo(self, monkeypatch):
+        # repo-less workflow (#248 Phase 3): repo is None → key omitted, not null.
+        writer = self._writer(monkeypatch, user_id="user-1", repo=None)
+        with patch("progress_writer.current_otel_trace_id", return_value="b" * 32):
+            writer.write_agent_milestone("m", "")
+        item = writer._table.put_item.call_args[1]["Item"]
+        assert item["user_id"] == "user-1"
+        assert "repo" not in item
+
+    def test_missing_trace_id_omits_key(self, monkeypatch):
+        # No recording span (tracing disabled) → trace_id key omitted entirely.
+        writer = self._writer(monkeypatch, user_id="user-1", repo="org/repo")
+        with patch("progress_writer.current_otel_trace_id", return_value=None):
+            writer.write_agent_milestone("m", "")
+        item = writer._table.put_item.call_args[1]["Item"]
+        assert "trace_id" not in item
+
+    def test_defaults_omit_all_correlation_keys(self, monkeypatch):
+        # Back-compat: a writer built without the envelope (empty user_id) stamps
+        # none of the optional keys, keeping the historical item shape.
+        writer = self._writer(monkeypatch)
+        with patch("progress_writer.current_otel_trace_id", return_value=None):
+            writer.write_agent_milestone("m", "")
+        item = writer._table.put_item.call_args[1]["Item"]
+        assert "user_id" not in item
+        assert "repo" not in item
+        assert "trace_id" not in item
+
+
 # ---------------------------------------------------------------------------
 # _ProgressWriter — --trace preview cap (design §10.1)
 # ---------------------------------------------------------------------------
 
@@ -288,6 +288,41 @@ def fake_run_task(**_kwargs):
     assert heartbeat_calls[0] == "t-heartbeat"
 
 
+def test_run_task_background_propagates_correlation_envelope(monkeypatch):
+    """The background task thread propagates {session_id, user_id, repo} into
+    OTEL baggage via propagate_correlation_context (#245).
+
+    Regression guard for the widened trigger: correlation must propagate even
+    when session_id is empty but user_id/repo are known — the branch the whole
+    envelope-in-baggage feature depends on.
+    """
+    calls: list[dict] = []
+    monkeypatch.setattr(
+        server,
+        "propagate_correlation_context",
+        lambda session_id, **kw: calls.append({"session_id": session_id, **kw}),
+    )
+    monkeypatch.setattr(server, "run_task", lambda **_kwargs: None)
+    monkeypatch.setattr(server.task_state, "write_heartbeat", lambda *a, **kw: None)
+    monkeypatch.setattr(server.task_state, "write_terminal", lambda *a, **kw: None)
+
+    # No session_id, but user_id + repo_url known → propagation must still run.
+    server._run_task_background(
+        task_id="t-corr",
+        repo_url="o/r",
+        task_description="x",
+        issue_number="",
+        github_token="",
+        anthropic_model="",
+        max_turns=10,
+        max_budget_usd=None,
+        aws_region="us-east-1",
+        user_id="user-1",
+    )
+
+    assert calls == [{"session_id": "", "user_id": "user-1", "repo": "o/r"}]
+
+
 def test_validate_required_params_pr_workflows_require_pr_number():
     """PR-iteration and PR-review workflows need a pr_number regardless."""
     missing = server._validate_required_params(
 
@@ -219,12 +219,18 @@ export async function handler(event: APIGatewayProxyEvent): Promise<APIGatewayPr
       });
     }
 
-    // 6. Strip task_id from event records (redundant in response context)
+    // 6. Strip task_id from event records (redundant in response context).
+    //    Pass through the correlation envelope (#245) per-event, omitting each
+    //    field when the source didn't stamp it (e.g. task_created) so the shape
+    //    stays stable for consumers.
     const eventData = events.map(e => ({
       event_id: e.event_id,
       event_type: e.event_type,
       timestamp: e.timestamp,
       metadata: e.metadata ?? {},
+      ...(e.user_id && { user_id: e.user_id }),
+      ...(e.repo && { repo: e.repo }),
+      ...(e.trace_id && { trace_id: e.trace_id }),
     }));
 
     // For descending scans we intentionally suppress ``next_token``. DDB's
 
@@ -197,6 +197,10 @@ export function assembleBundle(
     event_type: e.event_type,
     timestamp: e.timestamp,
     metadata: e.metadata ?? {},
+    // Correlation envelope (#245): passed through per-event, omitted when absent.
+    ...(e.user_id && { user_id: e.user_id }),
+    ...(e.repo && { repo: e.repo }),
+    ...(e.trace_id && { trace_id: e.trace_id }),
   }));
 
   // Verification is non-null only when at least one gate result was persisted.