feat: add types and helper function for context propagation

cadence/_internal/context_propagation.py

@@ -0,0 +1,117 @@
+"""Shared Cadence header codecs and context-propagation scopes."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator, Mapping, Sequence
+from contextlib import ExitStack, contextmanager
+
+from cadence.api.v1.common_pb2 import Header
+from cadence.context import ContextPropagator
+from cadence.error import ContextPropagationError
+
+
+def normalize_context_propagators(
+    propagators: Sequence[ContextPropagator] | None,
+) -> tuple[ContextPropagator, ...]:
+    """Snapshot propagators so runtime behavior cannot change after setup."""
+    return tuple(propagators or ())
+
+
+def header_to_context_fields(header: Header | None) -> dict[str, bytes]:
+    """Convert a Cadence Header to the cross-SDK ``str -> bytes`` carrier."""
+    if header is None:
+        return {}
+    return {key: bytes(payload.data) for key, payload in header.fields.items()}
+
+
+def inject_context_fields(
+    propagators: Sequence[ContextPropagator],
+) -> dict[str, bytes]:
+    """Collect and validate fields emitted by configured propagators."""
+    fields: dict[str, bytes] = {}
+    for propagator in propagators:
+        try:
+            emitted = propagator.inject()
+        except ContextPropagationError:
+            raise
+        except Exception as exc:
+            raise ContextPropagationError(
+                f"Context propagation inject failed for {type(propagator).__name__}"
+            ) from exc
+        if not isinstance(emitted, Mapping):
+            raise ContextPropagationError(
+                f"{type(propagator).__name__}.inject() must return a mapping"
+            )
+        for key, value in emitted.items():
+            if not isinstance(key, str):
+                raise ContextPropagationError(
+                    f"{type(propagator).__name__}.inject() emitted a non-string key"
+                )
+            if not isinstance(value, (bytes, bytearray, memoryview)):
+                raise ContextPropagationError(
+                    f"{type(propagator).__name__}.inject() emitted a non-bytes value "
+                    f"for header {key!r}"
+                )
+            if key in fields:
+                raise ContextPropagationError(
+                    f"Multiple context propagators emitted header {key!r}"
+                )
+            fields[key] = bytes(value)
+    return fields
+
+
+def context_header_from_propagators(
+    propagators: Sequence[ContextPropagator],
+) -> Header | None:
+    """Build an independent protobuf Header for outbound propagation."""
+    fields = inject_context_fields(propagators)
+    if not fields:
+        return None
+    header = Header()
+    for key, value in fields.items():
+        header.fields[key].data = value
+    return header
+
+
+@contextmanager
+def context_propagation_scope(
+    propagators: Sequence[ContextPropagator],
+    header_or_fields: Header | Mapping[str, bytes] | None,
+) -> Iterator[None]:
+    """Install incoming context for one logical workflow or activity invocation."""
+    fields: Mapping[str, bytes]
+    if isinstance(header_or_fields, Header):
+        fields = header_to_context_fields(header_or_fields)
+    elif header_or_fields is None:
+        fields = {}
+    else:
+        fields = header_or_fields
+
+    stack = ExitStack()
+    try:
+        for propagator in propagators:
+            try:
+                scope = propagator.extract(fields)
+                stack.enter_context(scope)
+            except ContextPropagationError:
+                raise
+            except Exception as exc:
+                raise ContextPropagationError(
+                    f"Context propagation extract failed for {type(propagator).__name__}"
+                ) from exc
+    except Exception:
+        try:
+            stack.close()
+        except Exception as cleanup_error:
+            raise ContextPropagationError(
+                "Context propagation cleanup failed after extract error"
+            ) from cleanup_error
+        raise
+
+    try:
+        yield
+    finally:
+        try:
+            stack.close()
+        except Exception as exc:
+            raise ContextPropagationError("Context propagation cleanup failed") from exc

cadence/context.py

@@ -0,0 +1,31 @@
+"""Context propagation interfaces for Cadence workflow boundaries.
+
+Propagators carry application-owned, request-scoped metadata through Cadence
+headers. Implementations commonly use :mod:`contextvars` so values are scoped
+to each client, workflow, or activity invocation.
+"""
+
+from collections.abc import Mapping
+from typing import ContextManager, Protocol
+
+from cadence.error import ContextPropagationError
+
+__all__ = ["ContextPropagationError", "ContextPropagator"]
+
+
+class ContextPropagator(Protocol):
+    """Serialize and scope application context across Cadence boundaries.
+
+    ``inject`` is called before the SDK writes a Cadence header. ``extract`` is
+    called for an incoming header and must return a context manager that restores
+    any ambient state when its scope exits.
+
+    Workflow-side ``inject`` calls are replay-sensitive: implementations must
+    return identical bytes for identical persisted workflow state.
+    """
+
+    def inject(self) -> Mapping[str, bytes]:
+        """Return the current context as opaque header fields."""
+
+    def extract(self, headers: Mapping[str, bytes]) -> ContextManager[None]:
+        """Return a scope that installs context decoded from ``headers``."""

cadence/error.py

@@ -4,6 +4,10 @@
 import grpc
 
 
+class ContextPropagationError(RuntimeError):
+    """Raised when a context propagator cannot encode or install context."""
+
+
 class ContinueAsNewError(Exception):
     def __init__(
         self,

tests/cadence/test_context_propagation_primitives.py

@@ -0,0 +1,159 @@
+from collections.abc import Iterator, Mapping
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import cast
+
+import pytest
+
+from cadence._internal.context_propagation import (
+    context_header_from_propagators,
+    context_propagation_scope,
+    header_to_context_fields,
+    normalize_context_propagators,
+)
+from cadence.api.v1.common_pb2 import Header
+from cadence.context import ContextPropagationError
+
+
+class _ContextVarPropagator:
+    def __init__(self, key: str = "x-request-id") -> None:
+        self._key = key
+        self.value: ContextVar[str | None] = ContextVar(key, default=None)
+
+    def inject(self) -> dict[str, bytes]:
+        value = self.value.get()
+        return {} if value is None else {self._key: value.encode()}
+
+    @contextmanager
+    def extract(self, headers: Mapping[str, bytes]) -> Iterator[None]:
+        raw_value = headers.get(self._key)
+        token = self.value.set(raw_value.decode() if raw_value is not None else None)
+        try:
+            yield
+        finally:
+            self.value.reset(token)
+
+
+class _StaticPropagator:
+    def __init__(self, fields: Mapping[str, bytes]) -> None:
+        self._fields = fields
+
+    def inject(self) -> Mapping[str, bytes]:
+        return self._fields
+
+    @contextmanager
+    def extract(self, _headers: Mapping[str, bytes]) -> Iterator[None]:
+        yield
+
+
+class _BrokenPropagator:
+    def inject(self) -> Mapping[str, bytes]:
+        raise ValueError("cannot inject")
+
+    @contextmanager
+    def extract(self, headers: Mapping[str, bytes]) -> Iterator[None]:
+        if headers:
+            yield
+            return
+        raise ValueError("cannot extract")
+
+
+class _ContextPropagationErrorPropagator:
+    def inject(self) -> Mapping[str, bytes]:
+        raise ContextPropagationError("propagator-specific failure")
+
+    @contextmanager
+    def extract(self, _headers: Mapping[str, bytes]) -> Iterator[None]:
+        yield
+
+
+class _RecordingPropagator(_StaticPropagator):
+    def __init__(self, events: list[str], name: str) -> None:
+        super().__init__({})
+        self._events = events
+        self._name = name
+
+    @contextmanager
+    def extract(self, _headers: Mapping[str, bytes]) -> Iterator[None]:
+        self._events.append(f"enter:{self._name}")
+        try:
+            yield
+        finally:
+            self._events.append(f"exit:{self._name}")
+
+
+class _CleanupBrokenPropagator(_StaticPropagator):
+    @contextmanager
+    def extract(self, _headers: Mapping[str, bytes]) -> Iterator[None]:
+        try:
+            yield
+        finally:
+            raise ValueError("cannot clean up")
+
+
+def test_header_codec_preserves_raw_bytes_and_contextvar_scope() -> None:
+    propagator = _ContextVarPropagator()
+    header = Header()
+    header.fields["go.trace"].data = b"\x00\xffgo"
+    header.fields["x-request-id"].data = b"request-1"
+    header.fields["java.trace"].data = b"\x80java\x00"
+
+    assert header_to_context_fields(header) == {
+        "go.trace": b"\x00\xffgo",
+        "x-request-id": b"request-1",
+        "java.trace": b"\x80java\x00",
+    }
+    with context_propagation_scope((propagator,), header):
+        assert propagator.value.get() == "request-1"
+    assert propagator.value.get() is None
+
+
+def test_injection_rejects_collisions_and_invalid_values() -> None:
+    with pytest.raises(ContextPropagationError, match="Multiple context propagators"):
+        context_header_from_propagators(
+            (
+                _StaticPropagator({"shared": b"one"}),
+                _StaticPropagator({"shared": b"two"}),
+            )
+        )
+
+    invalid = _StaticPropagator(cast(Mapping[str, bytes], {"invalid": "value"}))
+    with pytest.raises(ContextPropagationError, match="non-bytes"):
+        context_header_from_propagators((invalid,))
+
+    with pytest.raises(ContextPropagationError, match="inject failed"):
+        context_header_from_propagators((_BrokenPropagator(),))
+
+    with pytest.raises(
+        ContextPropagationError, match="propagator-specific failure"
+    ) as error:
+        context_header_from_propagators((_ContextPropagationErrorPropagator(),))
+    assert str(error.value) == "propagator-specific failure"
+
+
+def test_scope_reverses_cleanup_and_wraps_propagator_errors() -> None:
+    events: list[str] = []
+    with context_propagation_scope(
+        (_RecordingPropagator(events, "first"), _RecordingPropagator(events, "second")),
+        {},
+    ):
+        assert events == ["enter:first", "enter:second"]
+    assert events == ["enter:first", "enter:second", "exit:second", "exit:first"]
+
+    with pytest.raises(ContextPropagationError, match="extract failed"):
+        with context_propagation_scope((_BrokenPropagator(),), {}):
+            pass
+    with pytest.raises(ContextPropagationError, match="cleanup failed"):
+        with context_propagation_scope((_CleanupBrokenPropagator({}),), {}):
+            pass
+
+
+def test_normalization_snapshots_configured_propagators() -> None:
+    propagator = _StaticPropagator({})
+    configured = [propagator]
+
+    normalized = normalize_context_propagators(configured)
+    configured.clear()
+
+    assert normalized == (propagator,)
+    assert normalize_context_propagators(None) == ()