fixup

Author: realark

py/src/braintrust/logger.py

@@ -77,9 +77,9 @@
     TRACEPARENT_HEADER,
     TRACESTATE_HEADER,
     PropagatedState,
-    format_baggage,
     format_traceparent,
     get_header,
+    merge_baggage,
     parse_baggage,
     parse_traceparent,
 )
@@ -2543,12 +2543,11 @@ def _inject_into_carrier(
     if tracestate:
         _set_header(carrier, TRACESTATE_HEADER, tracestate)
 
-    # Merge braintrust.parent into any existing baggage, preserving other keys.
+    # Merge braintrust.parent into any existing baggage. Other vendors' members
+    # are forwarded byte-for-byte (see merge_baggage) so we never rewrite their
+    # percent-encoding.
     existing = get_header(carrier, BAGGAGE_HEADER)
-    entries = parse_baggage(existing) if existing else {}
-    if braintrust_parent:
-        entries[BRAINTRUST_PARENT_KEY] = braintrust_parent
-    baggage_value = format_baggage(entries)
+    baggage_value = merge_baggage(existing, braintrust_parent)
     if baggage_value is not None:
         _set_header(carrier, BAGGAGE_HEADER, baggage_value)
 

py/src/braintrust/propagation.py

@@ -13,6 +13,7 @@
 import logging
 import re
 from typing import NamedTuple
+from urllib.parse import unquote
 
 
 __all__ = [
@@ -27,6 +28,7 @@
     "format_traceparent",
     "parse_baggage",
     "format_baggage",
+    "merge_baggage",
     "get_header",
 ]
 
@@ -181,6 +183,48 @@ def format_baggage(entries):
     return ",".join(parts)
 
 
+def merge_baggage(existing, braintrust_parent):
+    """Merge a ``braintrust.parent`` value into an existing ``baggage`` header.
+
+    Unlike ``format_baggage(parse_baggage(...))``, this preserves every other
+    vendor's baggage member byte-for-byte: their raw ``key=value`` substrings
+    are forwarded exactly as received rather than decoded and re-encoded. This
+    keeps Braintrust a transparent relay so we never silently rewrite another
+    vendor's percent-encoding (e.g. ``path=a%2Fb`` must not become
+    ``path=a/b``).
+
+    Only the ``braintrust.parent`` member is (re)serialized, by us, from the
+    ``braintrust_parent`` argument. Any pre-existing ``braintrust.parent``
+    member in ``existing`` is dropped in favor of the supplied value.
+
+    Returns the merged header value, or None if there is nothing to emit (so
+    callers omit the header rather than emit an empty one).
+    """
+    members = []
+    if existing and isinstance(existing, str):
+        capped = existing[:_MAX_BAGGAGE_LENGTH] if len(existing) > _MAX_BAGGAGE_LENGTH else existing
+        for raw_member in capped.split(","):
+            member = raw_member.strip()
+            if not member or "=" not in member:
+                continue
+            # Identify the key (ignoring ';'-delimited properties) only to skip
+            # any inbound braintrust.parent; everything else is forwarded raw.
+            key_part = member.split(";", 1)[0].partition("=")[0]
+            key = _percent_decode(key_part.strip())
+            if key == BRAINTRUST_PARENT_KEY:
+                continue
+            members.append(member)
+
+    if braintrust_parent:
+        encoded_key = _percent_encode(BRAINTRUST_PARENT_KEY)
+        encoded_val = _percent_encode(str(braintrust_parent))
+        members.append(f"{encoded_key}={encoded_val}")
+
+    if not members:
+        return None
+    return ",".join(members)
+
+
 def _is_hex(value, length):
     if not isinstance(value, str) or len(value) != length:
         return False
@@ -191,17 +235,17 @@ def _is_hex(value, length):
     return all(c in "0123456789abcdef" for c in value)
 
 
-# Baggage values percent-encode a small set of characters. We keep this minimal
-# and robust: encode only the characters that would break the baggage grammar
-# (the member/list delimiters and property separator), plus ``%`` so the
-# encoding is self-describing.
+# Per W3C Baggage, member values are percent-encoded. We decode inbound values
+# with a full RFC3986 ``unquote`` so we interoperate with standard encoders
+# (e.g. OpenTelemetry's propagator percent-encodes ``:`` as ``%3A``). On emit we
+# percent-encode our own values (the ``braintrust.parent`` entry and any key we
+# re-serialize). Byte-for-byte pass-through of *other* vendors' baggage is
+# handled separately by :func:`merge_baggage`, which forwards their raw member
+# strings unchanged rather than round-tripping them through this codec.
 #
-# Decoding is the exact inverse of this set rather than a full RFC3986
-# ``unquote``. This keeps the codec a transparent relay: a percent-escape we did
-# not emit (e.g. an upstream vendor's ``%2F``) is left byte-for-byte intact
-# through a parse -> format round trip, instead of being silently rewritten
-# (``path=a%2Fb`` -> ``path=a/b``). Because ``%`` itself is always encoded as
-# ``%25``, our own values still round-trip exactly.
+# Encoded set is kept minimal: the characters that would otherwise break the
+# baggage grammar (member/list delimiters and the property separator) plus
+# ``%`` so the encoding is self-describing.
 _PERCENT_ENCODE = {
     ",": "%2C",
     ";": "%3B",
@@ -210,39 +254,18 @@ def _is_hex(value, length):
     "%": "%25",
 }
 
-# Inverse of _PERCENT_ENCODE. Both hex-case variants are accepted on decode so a
-# value we encoded (always uppercase) survives a round trip, while still being
-# lenient about an upstream that happens to lowercase these specific escapes.
-_PERCENT_DECODE = {}
-for _ch, _esc in _PERCENT_ENCODE.items():
-    _PERCENT_DECODE[_esc.lower()] = _ch
-    _PERCENT_DECODE[_esc.upper()] = _ch
-
-_PERCENT_ESCAPE_RE = re.compile(r"%[0-9a-fA-F]{2}")
-
 
 def _percent_encode(value):
     out = []
-    i = 0
-    n = len(value)
-    while i < n:
-        ch = value[i]
-        if ch == "%" and _PERCENT_ESCAPE_RE.match(value, i):
-            # Already a well-formed percent-escape (ours or an upstream
-            # vendor's). Forward it verbatim so the value is byte-preserved and
-            # we never double-encode (``%2F`` must not become ``%252F``).
-            out.append(value[i : i + 3])
-            i += 3
-            continue
+    for ch in value:
         out.append(_PERCENT_ENCODE.get(ch, ch))
-        i += 1
     return "".join(out)
 
 
 def _percent_decode(value):
     if "%" not in value:
         return value
-    # Only decode the escapes we ourselves emit; leave any other percent-escape
-    # (an upstream vendor's encoding) untouched so the value is forwarded
-    # byte-for-byte.
-    return _PERCENT_ESCAPE_RE.sub(lambda m: _PERCENT_DECODE.get(m.group(0), m.group(0)), value)
+    try:
+        return unquote(value)
+    except Exception:
+        return value

py/src/braintrust/test_propagation.py

@@ -21,6 +21,7 @@
     format_baggage,
     format_traceparent,
     get_header,
+    merge_baggage,
     parse_baggage,
     parse_traceparent,
 )
@@ -139,50 +140,73 @@ def test_format_encodes_keys_with_reserved_chars(self):
         # The serialized header must round-trip back to the original key/value.
         assert parse_baggage(formatted) == entries
 
-    def test_unrelated_baggage_value_is_byte_preserved(self):
+    def test_parse_decodes_standard_encoder_values(self):
+        # Per the W3C Baggage spec, values are percent-encoded, and standard
+        # encoders (e.g. OpenTelemetry's propagator) percent-encode `:` as
+        # `%3A`. We must fully decode inbound values to interoperate, otherwise
+        # `braintrust.parent=project_id%3Aabc` is not recognized downstream.
+        assert parse_baggage(f"{BRAINTRUST_PARENT_KEY}=project_id%3Aabc123") == {
+            BRAINTRUST_PARENT_KEY: "project_id:abc123"
+        }
+
+
+class TestMergeBaggage:
+    def test_adds_braintrust_parent_when_no_existing(self):
+        merged = merge_baggage(None, "project_id:abc")
+        assert parse_baggage(merged) == {BRAINTRUST_PARENT_KEY: "project_id:abc"}
+
+    def test_none_parent_and_no_existing_returns_none(self):
+        assert merge_baggage(None, None) is None
+        assert merge_baggage("", None) is None
+
+    def test_preserves_unrelated_baggage_byte_for_byte(self):
         # An upstream vendor may percent-encode octets outside the small set
         # Braintrust itself encodes (e.g. `/` as `%2F`). Forwarding such baggage
-        # must be a transparent relay: parse -> format must reproduce the exact
-        # inbound wire bytes for keys Braintrust does not own, otherwise we
-        # silently rewrite another vendor's value (`path=a%2Fb` -> `path=a/b`).
-        inbound = "path=a%2Fb,user=alice"
-        formatted = format_baggage(parse_baggage(inbound))
-        assert formatted == inbound
-
-    def test_round_trip_does_not_decode_unowned_percent_sequences(self):
+        # must be a transparent relay: the raw member is forwarded unchanged, so
+        # we never silently rewrite another vendor's value (`path=a%2Fb` must not
+        # become `path=a/b`).
+        merged = merge_baggage("path=a%2Fb,user=alice", "project_id:abc")
+        assert "path=a%2Fb" in merged
+        assert "user=alice" in merged
+        assert f"{BRAINTRUST_PARENT_KEY}=project_id:abc" in merged
+
+    def test_does_not_decode_unowned_percent_sequences(self):
         # `%41` is the percent-encoding of `A`. A transparent relay must not
-        # collapse `a%41b` to `aAb`; the inbound wire form must survive a
-        # parse -> format round trip unchanged.
-        inbound = "k=a%41b"
-        assert format_baggage(parse_baggage(inbound)) == inbound
-
-    def test_bare_percent_is_encoded_and_round_trips(self):
-        # A literal `%` that is not part of a valid percent-escape must still be
-        # encoded as `%25` so the value round-trips and the next hop does not
-        # mis-decode it.
-        for value in ["50% off", "100%", "a%4"]:
-            entries = {"k": value}
-            assert parse_baggage(format_baggage(entries)) == entries
-
-    def test_braintrust_owned_value_with_reserved_chars_round_trips(self):
-        # Braintrust does own its `braintrust.parent` value, which can carry an
-        # arbitrary `project_name` containing reserved characters. That value
-        # must be encoded on emit and decoded on parse.
-        entries = {BRAINTRUST_PARENT_KEY: "project_name:a,b c"}
-        assert parse_baggage(format_baggage(entries)) == entries
-
-    def test_mixed_owned_and_unowned_values_round_trip(self):
-        # A well-formed inbound header carrying both a Braintrust value with an
-        # encoded reserved char and an upstream vendor's pre-encoded value must
-        # preserve both: the owned value is decoded then re-encoded, the unowned
-        # escape is forwarded byte-for-byte.
-        inbound = f"{BRAINTRUST_PARENT_KEY}=project_name:a%2Cb,vendor=x%2Fy"
-        formatted = format_baggage(parse_baggage(inbound))
-        parsed = parse_baggage(formatted)
-        assert parsed[BRAINTRUST_PARENT_KEY] == "project_name:a,b"
-        assert parsed["vendor"] == "x%2Fy"
-        # The unowned escape survives on the wire byte-for-byte.
-        assert "vendor=x%2Fy" in formatted
+        # collapse `a%41b` to `aAb`; the inbound wire form is forwarded unchanged.
+        merged = merge_baggage("k=a%41b", None)
+        assert merged == "k=a%41b"
+
+    def test_replaces_existing_braintrust_parent(self):
+        # A stale inbound braintrust.parent must be dropped in favor of the
+        # value we supply, not duplicated.
+        merged = merge_baggage(
+            f"{BRAINTRUST_PARENT_KEY}=project_id:old,vendor=x",
+            "project_id:new",
+        )
+        parsed = parse_baggage(merged)
+        assert parsed[BRAINTRUST_PARENT_KEY] == "project_id:new"
+        assert parsed["vendor"] == "x"
+        # Only one braintrust.parent member is emitted.
+        assert merged.count(f"{BRAINTRUST_PARENT_KEY}=") == 1
+
+    def test_drops_existing_braintrust_parent_when_no_new_value(self):
+        # If we have no braintrust.parent to add, an inbound one is still
+        # consumed (it is ours to own) rather than forwarded raw.
+        merged = merge_baggage(f"{BRAINTRUST_PARENT_KEY}=project_id:old,vendor=x", None)
+        assert merged == "vendor=x"
+
+    def test_encodes_braintrust_parent_with_reserved_chars(self):
+        # Braintrust owns its braintrust.parent value, which can carry an
+        # arbitrary project_name containing reserved characters. That value must
+        # be encoded on emit and decode back cleanly.
+        merged = merge_baggage(None, "project_name:a,b c")
+        assert parse_baggage(merged) == {BRAINTRUST_PARENT_KEY: "project_name:a,b c"}
+
+    def test_skips_malformed_existing_members(self):
+        merged = merge_baggage("garbage,,k=v,no-equals", "project_id:abc")
+        parsed = parse_baggage(merged)
+        assert parsed["k"] == "v"
+        assert parsed[BRAINTRUST_PARENT_KEY] == "project_id:abc"
 
 
 def test_get_header_case_insensitive():