feat(context): add w3c() helper for W3C trace fields (#228)

* feat(context): add w3c() helper for W3C trace fields

* feat: use python 3.14

* feat: use main for harness tests

* fix(ci): drop unused Set up Python step so Post cleanup stops failing

The dockerized-test workflow never runs pip on the runner — Python is only
invoked inside the Dockerfile build using the base image's own Python. The
'cache: pip' directive on setup-python@v5 therefore ran its Post cleanup
against an empty /home/runner/.cache/pip and failed the job even though all
12 dockerized suite cases passed.

Also bumps actions/checkout to v5 to drop the Node 20 deprecation warning.

Author: maxday

.github/workflows/dockerized-test.yml

@@ -0,0 +1,31 @@
+name: dockerized-test
+
+permissions:
+  contents: read
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: ['*']
+  workflow_dispatch:
+
+jobs:
+  dockerized-test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v5
+
+      - name: Build the test image
+        run: |
+          docker build . \
+            -t local/test \
+            -f Dockerfile.test \
+            --build-arg BASE_IMAGE=public.ecr.aws/lambda/python:3.14
+
+      - name: Run dockerized suites
+        uses: aws/containerized-test-runner-for-aws-lambda@0863dd17b5fc19585250a2405c0f939a77b4f397 # main
+        with:
+          suiteFileArray: '["./test/dockerized/suites/*.json"]'
+          dockerImageName: 'local/test'
+          taskFolder: './test/dockerized/tasks'

Dockerfile.test

@@ -0,0 +1,12 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+ARG BASE_IMAGE=public.ecr.aws/lambda/python:3.14
+FROM $BASE_IMAGE
+
+COPY awslambdaric/*.py /tmp/awslambdaric-py/
+RUN RIC_PATH=$(python -c "import awslambdaric, os; print(os.path.dirname(awslambdaric.__file__))") \
+    && cp /tmp/awslambdaric-py/*.py "$RIC_PATH"/ \
+    && rm -rf /tmp/awslambdaric-py
+
+COPY test/dockerized/tasks /var/task/

awslambdaric/lambda_context.py

@@ -7,6 +7,11 @@
 import sys
 import time
 
+# Allowlist of W3C trace-context fields that may be surfaced through
+# ``LambdaContext.w3c()``. Any other key carried on ``clientContext.w3c`` is
+# ignored, and any allowlisted key whose value is not a string is dropped.
+W3C_ALLOWED_FIELDS = ("traceparent", "tracestate", "baggage")
+
 
 class LambdaContext(object):
     def __init__(
@@ -26,6 +31,7 @@ def __init__(
         self.function_version = os.environ.get("AWS_LAMBDA_FUNCTION_VERSION")
         self.invoked_function_arn = invoked_function_arn
         self.tenant_id = tenant_id
+        self._w3c_fields = self._extract_and_strip_w3c(client_context)
 
         self.client_context = make_obj_from_dict(ClientContext, client_context)
         if self.client_context is not None:
@@ -49,6 +55,38 @@ def get_remaining_time_in_millis(self):
         delta_ms = self._epoch_deadline_time_in_ms - epoch_now_in_ms
         return delta_ms if delta_ms > 0 else 0
 
+    def w3c(self):
+        """
+        Return the W3C trace-context at invoke time.
+        """
+        return dict(self._w3c_fields)
+
+    @staticmethod
+    def _extract_and_strip_w3c(client_context):
+        """
+        Pop ``w3c`` out of the parsed ``client_context`` dict and return a
+        normalized copy of its allowlisted string fields (see
+        ``W3C_ALLOWED_FIELDS``). Mutates ``client_context`` in place so the
+        ``w3c`` key is removed and cannot be read through
+        ``context.client_context``.
+        """
+        if not isinstance(client_context, dict):
+            return {}
+        if "w3c" not in client_context:
+            return {}
+
+        raw_w3c = client_context.pop("w3c")
+
+        if not isinstance(raw_w3c, dict):
+            return {}
+
+        fields = {}
+        for key in W3C_ALLOWED_FIELDS:
+            value = raw_w3c.get(key)
+            if isinstance(value, str):
+                fields[key] = value
+        return fields
+
     def log(self, msg):
         for handler in logging.getLogger().handlers:
             if hasattr(handler, "log_sink"):

test/dockerized/suites/ctx.json

@@ -0,0 +1,30 @@
+{
+  "tests": [
+    {
+      "name": "client_context_is_echoed_when_no_w3c_key",
+      "handler": "w3c.echo_client_context",
+      "request": {},
+      "clientContext": {
+        "custom": { "value": "hello" },
+        "env": { "stage": "beta" }
+      },
+      "assertions": [
+        {
+          "response": {
+            "custom": { "value": "hello" },
+            "env": { "stage": "beta" }
+          }
+        }
+      ]
+    },
+
+    {
+      "name": "client_context_is_null_when_header_absent",
+      "handler": "w3c.echo_client_context",
+      "request": {},
+      "assertions": [
+        { "response": null }
+      ]
+    }
+  ]
+}

test/dockerized/suites/w3c.json

@@ -0,0 +1,149 @@
+{
+  "tests": [
+    {
+      "name": "w3c_is_callable_on_context",
+      "handler": "w3c.w3c_is_callable",
+      "request": {},
+      "assertions": [
+        { "response": { "isCallable": true } }
+      ]
+    },
+
+    {
+      "name": "w3c_returns_empty_when_no_client_context_header",
+      "handler": "w3c.get_w3c",
+      "request": {},
+      "assertions": [
+        { "response": {} }
+      ]
+    },
+
+    {
+      "name": "w3c_returns_empty_when_client_context_has_no_w3c_key",
+      "handler": "w3c.get_w3c",
+      "request": {},
+      "clientContext": {
+        "custom": { "value": "test" }
+      },
+      "assertions": [
+        { "response": {} }
+      ]
+    },
+
+    {
+      "name": "w3c_returns_baggage_only",
+      "handler": "w3c.get_w3c",
+      "request": {},
+      "clientContext": {
+        "w3c": { "baggage": "userId=alice" }
+      },
+      "assertions": [
+        { "response": { "baggage": "userId=alice" } }
+      ]
+    },
+
+    {
+      "name": "w3c_returns_all_three_allowlisted_fields",
+      "handler": "w3c.get_w3c",
+      "request": {},
+      "clientContext": {
+        "w3c": {
+          "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
+          "tracestate": "rojo=00f067aa0ba902b7",
+          "baggage": "userId=alice"
+        }
+      },
+      "assertions": [
+        {
+          "response": {
+            "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
+            "tracestate": "rojo=00f067aa0ba902b7",
+            "baggage": "userId=alice"
+          }
+        }
+      ]
+    },
+
+    {
+      "name": "w3c_allowlist_drops_non_allowlisted_keys",
+      "handler": "w3c.get_w3c",
+      "request": {},
+      "clientContext": {
+        "w3c": {
+          "baggage": "keep=me",
+          "unknownField": "should-not-appear",
+          "x-custom-trace": "should-not-appear"
+        }
+      },
+      "assertions": [
+        { "response": { "baggage": "keep=me" } }
+      ]
+    },
+
+    {
+      "name": "w3c_drops_allowlisted_fields_with_non_string_values",
+      "handler": "w3c.get_w3c",
+      "request": {},
+      "clientContext": {
+        "w3c": {
+          "traceparent": 42,
+          "tracestate": null,
+          "baggage": { "nested": "no" }
+        }
+      },
+      "assertions": [
+        { "response": {} }
+      ]
+    },
+
+    {
+      "name": "w3c_treats_non_object_as_empty",
+      "handler": "w3c.get_w3c",
+      "request": {},
+      "clientContext": {
+        "w3c": "not-an-object"
+      },
+      "assertions": [
+        { "response": {} }
+      ]
+    },
+
+    {
+      "name": "w3c_treats_array_as_empty",
+      "handler": "w3c.get_w3c",
+      "request": {},
+      "clientContext": {
+        "w3c": ["baggage=abc"]
+      },
+      "assertions": [
+        { "response": {} }
+      ]
+    },
+
+    {
+      "name": "w3c_strips_source_client_context_w3c_after_construction",
+      "handler": "w3c.get_w3c_and_source",
+      "request": {},
+      "clientContext": {
+        "custom": { "value": "test" },
+        "w3c": {
+          "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
+          "baggage": "userId=alice"
+        }
+      },
+      "assertions": [
+        {
+          "response": {
+            "w3c": {
+              "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
+              "baggage": "userId=alice"
+            },
+            "clientContextIsDefined": true,
+            "clientContextHasW3c": false,
+            "clientContext": { "custom": { "value": "test" } }
+          }
+        }
+      ]
+    }
+  ]
+}

test/dockerized/tasks/w3c.py

@@ -0,0 +1,31 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+def _serialize_client_context(client_context):
+    if client_context is None:
+        return None
+    result = {}
+    for field in ("custom", "env"):
+        value = getattr(client_context, field, None)
+        if value is not None:
+            result[field] = value
+    return result
+
+def get_w3c(event, context):
+    return context.w3c()
+
+def get_w3c_and_source(event, context):
+    client_context = context.client_context
+    return {
+        "w3c": context.w3c(),
+        "clientContextIsDefined": client_context is not None,
+        "clientContextHasW3c": client_context is not None
+        and hasattr(client_context, "w3c"),
+        "clientContext": _serialize_client_context(client_context),
+    }
+
+def echo_client_context(event, context):
+    return _serialize_client_context(context.client_context)
+
+def w3c_is_callable(event, context):
+    return {"isCallable": callable(getattr(context, "w3c", None))}