fix: preserve trace context across local ingest handoffs

deploy/helm/sie-cluster/templates/_otel-collector-config.tpl

@@ -410,7 +410,7 @@ processors:
           - set(status.message, "")
           - set(trace_state, "")
           - set(links, [])
-          - 'set(name, "other") where name != "gateway.request" and name != "gateway.publish" and name != "gateway.proxy" and name != "gateway.proxy_chat" and name != "gateway.proxy_request" and name != "gateway.proxy_generate" and name != "sidecar.dispatch" and name != "worker.run_batch" and name != "worker.run_batch.request" and name != "sidecar.dispatch.request" and name != "worker.streaming_processor" and name != "encode" and name != "score" and name != "extract" and name != "generate" and name != "openai_embeddings" and name != "chat_completions" and name != "rerank" and name != "other"'
+          - 'set(name, "other") where name != "gateway.request" and name != "gateway.dispatch" and name != "gateway.dispatch.ipc" and name != "gateway.dispatch.i6pn" and name != "dispatcher.dispatch" and name != "dispatcher.modal.remote" and name != "dispatcher.modal.stream" and name != "dispatcher.modal.spawn" and name != "dispatcher.websocket" and name != "dispatcher.i6pn" and name != "dispatcher.fallback" and name != "sidecar.local_ingest" and name != "worker.local_ingest" and name != "gateway.publish" and name != "gateway.proxy" and name != "gateway.proxy_chat" and name != "gateway.proxy_request" and name != "gateway.proxy_generate" and name != "sidecar.dispatch" and name != "worker.run_batch" and name != "worker.run_batch.request" and name != "sidecar.dispatch.request" and name != "worker.streaming_processor" and name != "encode" and name != "score" and name != "extract" and name != "generate" and name != "openai_embeddings" and name != "chat_completions" and name != "rerank" and name != "other"'
 {{- end }}
 {{- if $logsEnabled }}
   # Logs are allowlisted just like metrics are declared: only the fixed,

packages/sie_gateway/src/observability/tracing.rs

@@ -77,7 +77,7 @@ use sie_telemetry::proxy::{
 #[cfg(test)]
 use sie_telemetry::resource::compose_service_instance_id;
 pub(crate) use sie_telemetry::resource::process_start_uuid;
-use sie_telemetry::resource::{resource_from_values, service_instance_id};
+use sie_telemetry::resource::{resource_from_env, service_instance_id};
 use sie_telemetry::transport::{
     configured_signal_endpoints, endpoint_origin_for_log, otlp_metrics_protocol,
     select_signal_protocol, trace_export_config, OtlpProtocol, SignalExportConfig,
@@ -420,7 +420,7 @@ fn otlp_resource_from_values(
     deployment_environment: &str,
     cloud_region: &str,
 ) -> Resource {
-    resource_from_values(
+    resource_from_env(
         service_name,
         instance_id,
         deployment_environment,

packages/sie_server/src/sie_server/local_ingest_client.py

@@ -4,10 +4,18 @@
 
 import asyncio
 import hashlib
-from collections.abc import AsyncIterator
+import os
+from collections.abc import AsyncGenerator
+from contextlib import nullcontext, suppress
 from typing import Any
 
 import msgpack
+from opentelemetry import trace
+from opentelemetry.context import Context
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+
+_TRACER = trace.get_tracer(__name__)
+_PROPAGATOR = TraceContextTextMapPropagator()
 
 MAX_LOCAL_INGEST_FRAME_BYTES = 64 * 1024 * 1024
 PAYLOAD_DIGEST_BYTES = 32
@@ -86,6 +94,10 @@ def build_generation_request_body(items: bytes, params: bytes, meta: dict[str, A
         "payload_digest": payload_digest,
         "timeout_ms": timeout_ms,
     }
+    for key, limit in (("traceparent", 256), ("tracestate", 512)):
+        value = meta.get(key)
+        if isinstance(value, str) and len(value) <= limit:
+            body[key] = value
     if not bound_transport:
         body["payload_digest"] = compute_payload_digest(body)
     return body
@@ -114,12 +126,73 @@ async def _read_frame(reader: asyncio.StreamReader) -> dict[str, Any]:
     return response
 
 
-async def stream_generate(
+def stream_generate(socket_path: str, items: bytes, params: bytes, meta: dict[str, Any]) -> AsyncGenerator[bytes, None]:
+    if os.environ.get("SIE_TRACING_ENABLED", "").strip().lower() not in {"1", "true", "yes", "on"}:
+        return _stream_generate(socket_path, items, params, meta)
+    return _traced_generate(socket_path, items, params, meta)
+
+
+async def _traced_generate(
+    socket_path: str, items: bytes, params: bytes, meta: dict[str, Any]
+) -> AsyncGenerator[bytes, None]:
+    carrier = {
+        key: value
+        for key, limit in (("traceparent", 256), ("tracestate", 512))
+        if isinstance(value := meta.get(key), str) and len(value) <= limit
+    }
+    span = None
+    with suppress(Exception):
+        parent = _PROPAGATOR.extract(carrier, context=Context())
+        span = _TRACER.start_span(
+            "worker.local_ingest",
+            context=parent,
+            kind=trace.SpanKind.CLIENT,
+            record_exception=False,
+            set_status_on_exception=False,
+        )
+        carrier = {}
+        _PROPAGATOR.inject(carrier, context=trace.set_span_in_context(span, parent))
+        if "traceparent" in carrier:
+            meta = {key: value for key, value in meta.items() if key not in ("traceparent", "tracestate")}
+            meta.update(carrier)
+    iterator = _stream_generate(socket_path, items, params, meta)
+    try:
+        while True:
+            try:
+                with (
+                    trace.use_span(span, record_exception=False, set_status_on_exception=False)
+                    if span is not None
+                    else nullcontext()
+                ):
+                    chunk = await anext(iterator)
+            except StopAsyncIteration:
+                break
+            yield chunk
+    except BaseException as error:
+        if span is not None and not isinstance(error, (asyncio.CancelledError, GeneratorExit)):
+            with suppress(Exception):
+                span.set_status(trace.StatusCode.ERROR)
+        raise
+    finally:
+        try:
+            with (
+                trace.use_span(span, record_exception=False, set_status_on_exception=False)
+                if span is not None
+                else nullcontext()
+            ):
+                await iterator.aclose()
+        finally:
+            if span is not None:
+                with suppress(Exception):
+                    span.end()
+
+
+async def _stream_generate(
     socket_path: str,
     items: bytes,
     params: bytes,
     meta: dict[str, Any],
-) -> AsyncIterator[bytes]:
+) -> AsyncGenerator[bytes, None]:
     """Map one caller operation onto sidecar protocol v0.2.
 
     Pulling one response before requesting the next naturally propagates
@@ -133,6 +206,11 @@ async def stream_generate(
         {"id": operation_id, "op": "publish_generate_stream", "body": body},
         use_bin_type=True,
     )
+    if len(payload) > MAX_LOCAL_INGEST_FRAME_BYTES and ("traceparent" in body or "tracestate" in body):
+        # Optional telemetry must not make a previously valid frame too large.
+        body.pop("traceparent", None)
+        body.pop("tracestate", None)
+        payload = msgpack.packb({"id": operation_id, "op": "publish_generate_stream", "body": body}, use_bin_type=True)
     if len(payload) > MAX_LOCAL_INGEST_FRAME_BYTES:
         raise ValueError(
             f"local-ingest generation frame is {len(payload)} bytes; maximum is {MAX_LOCAL_INGEST_FRAME_BYTES}"

packages/sie_server/src/sie_server/observability/worker_telemetry.py

@@ -1123,6 +1123,7 @@ def worker_resource_attributes() -> dict[str, str]:
         "service.instance.id": service_instance_id(),
         "deployment.environment": deployment_environment,
         "cloud.region": cloud_region,
+        **({"service.version": version} if (version := _clean_env("OTEL_SERVICE_VERSION")) else {}),
     }
 
 

packages/sie_server/tests/observability/test_worker_telemetry.py

@@ -634,3 +634,10 @@ def test_invalid_metrics_transport_disables_export_fail_open(
     assert "continuing without export" in caplog.text
     assert "error_type=ValueError" in caplog.text
     assert "http/json" not in caplog.text
+
+
+def test_service_version_uses_explicit_deployment_revision(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("OTEL_SERVICE_VERSION", "  " + "a" * 40 + "  ")
+    assert wt.worker_resource_attributes()["service.version"] == "a" * 40
+    monkeypatch.setenv("OTEL_SERVICE_VERSION", "  ")
+    assert "service.version" not in wt.worker_resource_attributes()

packages/sie_server/tests/test_local_ingest_client.py

@@ -4,12 +4,17 @@
 
 import asyncio
 import tempfile
-from collections.abc import Iterator
+from collections.abc import AsyncIterator, Iterator
+from contextlib import aclosing
 from pathlib import Path
 from typing import Any
 
 import msgpack
 import pytest
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 from sie_server import local_ingest_client
 
 
@@ -179,3 +184,96 @@ async def handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) ->
                 chunk
                 async for chunk in local_ingest_client.stream_generate(str(socket_path), b"items", b"params", _meta())
             ]
+
+
+def test_trace_carrier_does_not_change_bound_payload() -> None:
+    parent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
+    body = local_ingest_client.build_generation_request_body(b"items", b"params", _meta())
+    traced = local_ingest_client.build_generation_request_body(
+        b"items", b"params", {**_meta(), "traceparent": parent, "tracestate": "vendor=value"}
+    )
+    assert traced["traceparent"] == parent
+    assert traced["tracestate"] == "vendor=value"
+    assert traced["items"] == body["items"]
+    assert traced["payload_digest"] == body["payload_digest"]
+    assert local_ingest_client.compute_payload_digest(traced) == local_ingest_client.compute_payload_digest(body)
+    invalid = local_ingest_client.build_generation_request_body(
+        b"items", b"params", {**_meta(), "traceparent": 3, "tracestate": "x" * 513}
+    )
+    assert "traceparent" not in invalid
+    assert "tracestate" not in invalid
+
+
+@pytest.mark.parametrize("state", [None, "vendor=value", "invalid state"])
+async def test_traced_stream_owns_handoff_lifetime_without_leaking_context(
+    monkeypatch: pytest.MonkeyPatch, state: str | None
+) -> None:
+    exporter = InMemorySpanExporter()
+    provider = TracerProvider()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    monkeypatch.setattr(local_ingest_client, "_TRACER", provider.get_tracer("test"))
+    monkeypatch.setenv("SIE_TRACING_ENABLED", "true")
+    parent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
+    captured = []
+    closed = []
+
+    async def transport(_socket: str, items: bytes, params: bytes, meta: dict[str, Any]) -> AsyncIterator[bytes]:
+        captured.append((items, params, meta, trace.get_current_span().get_span_context()))
+        try:
+            yield b"chunk"
+            await asyncio.Future()
+        finally:
+            closed.append(trace.get_current_span().get_span_context())
+
+    monkeypatch.setattr(local_ingest_client, "_stream_generate", transport)
+    iterator = local_ingest_client.stream_generate(
+        "private/socket", b"secret items", b"secret params", {**_meta(), "traceparent": parent, "tracestate": state}
+    )
+    async with aclosing(iterator):
+        assert await anext(iterator) == b"chunk"
+        assert not trace.get_current_span().get_span_context().is_valid
+        assert not exporter.get_finished_spans()
+        await asyncio.create_task(iterator.aclose())
+    (span,) = exporter.get_finished_spans()
+    assert closed == [span.context]
+    assert captured[0][3] == span.context
+    assert span.name == "worker.local_ingest"
+    assert span.parent.span_id == int(parent.split("-")[2], 16)
+    assert captured[0][0:2] == (b"secret items", b"secret params")
+    assert captured[0][2]["traceparent"].split("-")[2] == f"{span.context.span_id:016x}"
+    assert captured[0][2].get("tracestate") == (state if state == "vendor=value" else None)
+    assert not span.attributes
+    assert not span.events
+    assert not span.status.description
+    provider.shutdown()
+
+
+async def test_optional_carrier_does_not_overflow_valid_frame(
+    unix_socket_dir: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    body = local_ingest_client.build_generation_request_body(b"items", b"params", _meta())
+    limit = len(msgpack.packb({"id": 1, "op": "publish_generate_stream", "body": body}, use_bin_type=True))
+    monkeypatch.setattr(local_ingest_client, "MAX_LOCAL_INGEST_FRAME_BYTES", limit)
+    socket_path = unix_socket_dir / "bounded.sock"
+    observed = []
+
+    async def handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
+        observed.append(await _request(reader))
+        writer.write(_frame({"id": 1, "ok": True, "body": {"final": True, "outcome": {"chunks": 0}}}))
+        await writer.drain()
+        writer.close()
+        await writer.wait_closed()
+
+    server = await asyncio.start_unix_server(handle, path=socket_path)
+    async with server:
+        chunks = [
+            chunk
+            async for chunk in local_ingest_client.stream_generate(
+                str(socket_path),
+                b"items",
+                b"params",
+                {**_meta(), "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"},
+            )
+        ]
+    assert chunks == []
+    assert observed[0]["body"] == body

packages/sie_server_rust/src/observability/resource.rs

@@ -7,7 +7,7 @@ use opentelemetry_sdk::Resource;
 
 pub(crate) use sie_telemetry::env::cleaned_env;
 use sie_telemetry::resource::{
-    instance_prefix_env, resource_from_values, service_instance_id, UNKNOWN_RESOURCE_VALUE,
+    instance_prefix_env, resource_from_env, service_instance_id, UNKNOWN_RESOURCE_VALUE,
 };
 
 pub const SERVICE_NAME: &str = "sie-worker";
@@ -24,7 +24,7 @@ pub fn telemetry_resource() -> Resource {
         .or_else(|| cleaned_env("AWS_DEFAULT_REGION"))
         .unwrap_or_else(|| UNKNOWN_RESOURCE_VALUE.to_string());
 
-    resource_from_values(
+    resource_from_env(
         SERVICE_NAME,
         &instance_id,
         &deployment_environment,