feat: add OpenTelemetry session correlation to main

.github/workflows/tests.yml

@@ -29,10 +29,11 @@ jobs:
           python-version: "3.12"
 
       - name: Install test dependencies
-        # Hash-pinned per OpenSSF Scorecard Pinned-Dependencies. Reuses the
-        # property-test lockfile (pytest, hypothesis, and the package runtime
-        # dependencies), which covers the full suite without unpinned installs.
-        run: python -m pip install --require-hashes -r fuzz/requirements-property.txt
+        # Hash-pinned per OpenSSF Scorecard Pinned-Dependencies. Install both
+        # runtime integrations and the property-test tools exercised below.
+        run: |
+          python -m pip install --require-hashes -r requirements.lock
+          python -m pip install --require-hashes -r fuzz/requirements-property.txt
 
       - name: Run full test suite
         run: python -m pytest -q

CHANGELOG.d/otel-client-disconnect.md

@@ -0,0 +1 @@
+Fixed gateway diagnostics so caller disconnects do not produce a second error response or expose raw request paths in logs, and caller-controlled trace baggage does not cross the provider boundary.

README.md

@@ -237,8 +237,9 @@ is read from a **KV config store**, never `os.getenv`.
   backend (local in-process backend standalone), and records one usage-ledger row
   per original vector with the full attribution dimensions (service, team,
   group, company, provider) carried in `metadata`.
-- **Health.** `GET /healthz` is an unauthenticated liveness probe; it never
-  claims that an upstream chat worker is serving. Admins can use
+- **Health.** `GET /healthz` is an unauthenticated liveness probe that returns
+  only service identity and process status; it never discloses worker topology,
+  backend names, usage volume, or upstream readiness. Admins can use
   `GET /api/v1/provider_readiness/latest?refresh=true` for one bounded,
   non-retrying chat probe per enabled worker.
 - **Standalone + optional pg-llm-batch integration.** The hub runs standalone

contextual_orchestrator/__main__.py

@@ -9,6 +9,7 @@
 from dataclasses import replace
 
 from .cost_ledger import PriceBook
+from .cost_router import CostRoutingCoordinator
 from .credentials import get_credential, register_credential
 from .kv_config import InMemoryConfigStore
 from .model_discovery import (
@@ -33,6 +34,20 @@
 DEFAULT_INFERENCE_CREDENTIAL_NAME = "CONTEXTUAL_ORCHESTRATOR_INFERENCE_TOKEN"
 
 
+def _bootstrap_telemetry_config() -> InMemoryConfigStore:
+    """Load non-secret OTEL deployment settings into the process KV at startup."""
+    config = InMemoryConfigStore()
+    for environment_name, key in (
+        ("OTEL_EXPORTER_OTLP_ENDPOINT", "exporter_otlp_endpoint"),
+        ("OTEL_SERVICE_NAME", "service_name"),
+        ("OTEL_SDK_DISABLED", "sdk_disabled"),
+    ):
+        value = os.environ.get(environment_name, "").strip()
+        if value:
+            config.set("telemetry", key, value)
+    return config
+
+
 def _positive_int(value: str) -> int:
     """Parse a strictly positive integer for an argparse option."""
     try:
@@ -429,6 +444,10 @@ def main(argv: list[str] | None = None) -> None:
                 admin_session_secure_cookie=not args.insecure_admin_session_cookie,
             ),
             clearfolio_url=args.clearfolio_url,
+            coordinator=CostRoutingCoordinator(
+                orchestrator,
+                config_store=_bootstrap_telemetry_config(),
+            ),
         )
         return
 

contextual_orchestrator/batch_routing.py

@@ -27,6 +27,7 @@
 import time
 import uuid
 from concurrent.futures import ThreadPoolExecutor
+from contextvars import copy_context
 from dataclasses import dataclass, field
 from typing import Any, Callable, Dict, List, Optional, Protocol
 
@@ -261,8 +262,13 @@ def run(request: BatchRequest) -> BatchResultItem:
         if self.max_concurrency == 1 or len(requests) <= 1:
             items = [run(request) for request in requests]
         else:
+            def run_with_context(item: tuple[Any, BatchRequest]) -> BatchResultItem:
+                context, request = item
+                return context.run(run, request)
+
+            contexts_and_requests = [(copy_context(), request) for request in requests]
             with ThreadPoolExecutor(max_workers=min(self.max_concurrency, len(requests))) as pool:
-                items = list(pool.map(run, requests))
+                items = list(pool.map(run_with_context, contexts_and_requests))
         self._results[job_id] = items
         return BatchJob(job_id=job_id, backend=self.name, status="completed", request_count=len(requests))
 

contextual_orchestrator/orchestrator.py

@@ -5,7 +5,7 @@
 from collections import Counter, deque, OrderedDict
 from collections.abc import Iterable, Mapping
 from contextlib import contextmanager
-from contextvars import ContextVar
+from contextvars import ContextVar, copy_context
 from concurrent.futures import ThreadPoolExecutor
 import copy
 from dataclasses import dataclass, replace
@@ -36,6 +36,7 @@
 )
 from .conventions import require_object_name
 from .credentials import NotConfigured, get_credential
+from .telemetry import inject_trace_context, traced
 from .pii_protection import (
     DEFAULT_PII_KEY_NAME,
     ENCRYPTED_FIELDS_KEY,
@@ -928,7 +929,18 @@ def chat(
         if _is_direct_mlx_provider_url(agent.base_url) and self.chat_template_args:
             payload["chat_template_kwargs"] = self.chat_template_args
         payload = self.apply_effort_profile(agent, payload, effort_profile)
-        with _local_provider_slot(agent, self.local_concurrency, self.timeout):
+        parsed_provider = urlparse(agent.base_url)
+        with traced(
+            f"chat {agent.model}",
+            {
+                "gen_ai.operation.name": "chat",
+                "gen_ai.provider.name": agent.provider_name or parsed_provider.hostname or agent.id,
+                "gen_ai.request.model": agent.model,
+                "contextual_orchestrator.agent_id": agent.id,
+                "server.address": parsed_provider.hostname or "",
+                "server.port": parsed_provider.port or (443 if parsed_provider.scheme == "https" else 80),
+            },
+        ), _local_provider_slot(agent, self.local_concurrency, self.timeout):
             return self._send_with_retry(agent, payload, destination)
 
     def apply_effort_profile(
@@ -1080,6 +1092,7 @@ def _send(
         headers = {"content-type": "application/json"}
         if api_key:
             headers["authorization"] = f"{agent.auth_scheme} {api_key}"
+        inject_trace_context(headers)
         request = urllib.request.Request(
             self._provider_url(agent, "/chat/completions"),
             data=json.dumps(payload).encode("utf-8"),
@@ -1244,8 +1257,19 @@ def stream_chat(
         if _is_direct_mlx_provider_url(agent.base_url) and self.chat_template_args:
             payload["chat_template_kwargs"] = self.chat_template_args
         payload = self.apply_effort_profile(agent, payload, effort_profile)
-        with _local_provider_slot(agent, self.local_concurrency, self.timeout):  # pragma: no cover
-            yield from self._stream_send(agent, payload, destination)  # pragma: no cover
+        parsed_provider = urlparse(agent.base_url)
+        with traced(
+            f"chat {agent.model}",
+            {
+                "gen_ai.operation.name": "chat",
+                "gen_ai.provider.name": agent.provider_name or parsed_provider.hostname or agent.id,
+                "gen_ai.request.model": agent.model,
+                "contextual_orchestrator.agent_id": agent.id,
+                "server.address": parsed_provider.hostname or "",
+                "server.port": parsed_provider.port or (443 if parsed_provider.scheme == "https" else 80),
+            },
+        ), _local_provider_slot(agent, self.local_concurrency, self.timeout):  # pragma: no cover
+            yield from self._stream_send(agent, payload, destination)
 
     def _stream_send(
         self, agent: ModelAgent, payload: dict[str, Any], destination: ProviderDestination | None = None
@@ -1255,6 +1279,7 @@ def _stream_send(
         headers = {"content-type": "application/json", "accept": "text/event-stream"}
         if api_key:
             headers["authorization"] = f"{agent.auth_scheme} {api_key}"
+        inject_trace_context(headers)
         request = urllib.request.Request(
             self._provider_url(agent, "/chat/completions"),
             data=json.dumps(payload).encode("utf-8"),
@@ -1318,18 +1343,35 @@ def proxy_send(
         if agent.base_url.startswith("mock://"):
             return self._mock_raw(agent, normalized_endpoint, payload)
         destination = self._validate_provider(agent)  # pragma: no cover
-        if normalized_endpoint == "responses" and _is_local_provider_url(agent.base_url):
-            chat_payload = _responses_to_chat_payload(payload)
-            chat_payload.setdefault("max_tokens", self.request_settings_snapshot()["max_output_tokens"])
-            if _is_direct_mlx_provider_url(agent.base_url) and self.chat_template_args:
-                chat_payload["chat_template_kwargs"] = self.chat_template_args
-            with _local_provider_slot(agent, self.local_concurrency, self.timeout):
-                chat_response = self._send_raw_with_retry(
-                    agent, "chat/completions", chat_payload, destination
-                )
-            return _chat_to_responses_payload(chat_response, payload)
-        with _local_provider_slot(agent, self.local_concurrency, self.timeout):  # pragma: no cover
-            return self._send_raw_with_retry(agent, normalized_endpoint, payload, destination)
+        parsed_provider = urlparse(agent.base_url)
+        operation_name = {
+            "chat/completions": "chat",
+            "completions": "text_completion",
+            "responses": "generate_content",
+        }.get(normalized_endpoint, "generate_content")
+        with traced(
+            f"{operation_name} {agent.model}",
+            {
+                "gen_ai.operation.name": operation_name,
+                "gen_ai.provider.name": agent.provider_name or parsed_provider.hostname or agent.id,
+                "gen_ai.request.model": agent.model,
+                "contextual_orchestrator.agent_id": agent.id,
+                "server.address": parsed_provider.hostname or "",
+                "server.port": parsed_provider.port or (443 if parsed_provider.scheme == "https" else 80),
+            },
+        ):
+            if normalized_endpoint == "responses" and _is_local_provider_url(agent.base_url):
+                chat_payload = _responses_to_chat_payload(payload)
+                chat_payload.setdefault("max_tokens", self.request_settings_snapshot()["max_output_tokens"])
+                if _is_direct_mlx_provider_url(agent.base_url) and self.chat_template_args:
+                    chat_payload["chat_template_kwargs"] = self.chat_template_args
+                with _local_provider_slot(agent, self.local_concurrency, self.timeout):
+                    chat_response = self._send_raw_with_retry(
+                        agent, "chat/completions", chat_payload, destination
+                    )
+                return _chat_to_responses_payload(chat_response, payload)
+            with _local_provider_slot(agent, self.local_concurrency, self.timeout):  # pragma: no cover
+                return self._send_raw_with_retry(agent, normalized_endpoint, payload, destination)
 
     def _send_raw_with_retry(
         self,
@@ -1365,6 +1407,7 @@ def _send_raw(
         headers = {"content-type": "application/json"}
         if api_key:
             headers["authorization"] = f"{agent.auth_scheme} {api_key}"
+        inject_trace_context(headers)
         request = urllib.request.Request(
             self._provider_url(agent, f"/{endpoint.lstrip('/')}"),
             data=json.dumps(payload).encode("utf-8"),
@@ -1537,7 +1580,7 @@ def _local_batch_chat(
         agent: ModelAgent,
         requests: dict[str, list[ChatMessage]],
         temperature: float | None,
-        effort_profile: ReasoningEffortProfile | None,
+        effort_profile: ReasoningEffortProfile | None = None,
     ) -> dict[str, dict[str, Any]]:
         """Run local OpenAI-compatible requests concurrently through mlx-lm."""
         request_settings = self.request_settings_snapshot()
@@ -1558,7 +1601,10 @@ def complete(custom_id: str, messages: list[ChatMessage]) -> tuple[str, dict[str
         if self.local_concurrency == 1 or len(requests) <= 1:
             return dict(complete(custom_id, messages) for custom_id, messages in requests.items())
         with ThreadPoolExecutor(max_workers=min(self.local_concurrency, len(requests))) as pool:
-            futures = [pool.submit(complete, custom_id, messages) for custom_id, messages in requests.items()]
+            futures = [
+                pool.submit(copy_context().run, complete, custom_id, messages)
+                for custom_id, messages in requests.items()
+            ]
             return dict(future.result() for future in futures)
 
     def _batch_run(

contextual_orchestrator/server.py

@@ -8,6 +8,7 @@
 import base64
 import hashlib
 import json
+import logging
 import secrets
 import struct
 import threading
@@ -34,6 +35,18 @@
 )
 from .pii_protection import DEFAULT_PURPOSE_BY_SCOPE, PURPOSES_BY_SCOPE
 from .tool_fallback import ToolFallbackStoppedError
+from .telemetry import (
+    attach_trace_context,
+    configure_telemetry,
+    current_session_id,
+    detach_trace_context,
+    reset_session_id,
+    session_id_from_headers,
+    session_id_from_request,
+    set_session_id,
+)
+
+_LOGGER = logging.getLogger(__name__)
 
 # OpenAI request params forwarded verbatim to the provider on passthrough.
 OPENAI_PASSTHROUGH_PARAM_KEYS = {
@@ -4787,6 +4800,7 @@ def build_server(
     security = security or SecurityConfig()
     security.check_bind(host)
     coordinator = coordinator or CostRoutingCoordinator(orchestrator)
+    configure_telemetry(config=coordinator.config)
     if clearfolio_url is not None:
         parsed_viewer = urllib.parse.urlparse(clearfolio_url)
         if parsed_viewer.scheme not in {"http", "https"} or not parsed_viewer.netloc:
@@ -4795,6 +4809,45 @@ def build_server(
 
     class Handler(BaseHTTPRequestHandler):
         """Handle authenticated orchestration, administration, and health routes."""
+        _session_token = None
+        _trace_token = None
+
+        def _bind_session(self, session_id: str | None) -> None:
+            """Bind validated request correlation to this handler context."""
+            if session_id is None:
+                return
+            if self._session_token is not None:
+                reset_session_id(self._session_token)
+            self._session_token = set_session_id(session_id)
+
+        def _bind_trace(self) -> None:
+            """Replace, rather than stack, inbound trace context on this request."""
+            if self._trace_token is not None:
+                detach_trace_context(self._trace_token)
+            self._trace_token = attach_trace_context(self.headers)
+
+        def _reset_session(self) -> None:
+            """Release request correlation state before a keep-alive request."""
+            trace_token, self._trace_token = self._trace_token, None
+            if trace_token is not None:
+                detach_trace_context(trace_token)
+            session_token, self._session_token = self._session_token, None
+            if session_token is not None:
+                reset_session_id(session_token)
+
+        def handle_one_request(self) -> None:
+            """Prevent a keep-alive connection from carrying request state."""
+            try:
+                super().handle_one_request()
+            finally:
+                self._reset_session()
+
+        def finish(self) -> None:
+            """Finish the response and release request correlation state."""
+            try:
+                super().finish()
+            finally:
+                self._reset_session()
 
         def do_GET(self) -> None:  # noqa: N802
             """Dispatch GET requests after applying the route's authorization scope."""
@@ -5241,6 +5294,14 @@ def do_POST(self) -> None:  # noqa: N802
                 )
                 self._authorize(scope, state_changing=True)
                 body = self._read_json()
+                metadata_values = [
+                    value
+                    for key in ("metadata", "client_metadata")
+                    if isinstance((value := body.get(key)), dict)
+                ]
+                request_session_id = session_id_from_request(self.headers, *metadata_values)
+                if request_session_id != current_session_id():
+                    self._bind_session(request_session_id)
                 cache_bypass = _cache_bypass_header(self.headers.get("x-cache-bypass"))
                 cache_partition = self._cache_partition()
 
@@ -6028,6 +6089,8 @@ def _authorize(
             (durable for replays), and browser-driven state-changing admin
             requests must pass the same-origin check.
             """
+            self._bind_trace()
+            self._bind_session(session_id_from_headers(self.headers))
             effective_purpose = purpose or DEFAULT_PURPOSE_BY_SCOPE.get(scope, "")
             try:
                 security.check_rate_limit(self.client_address[0])
@@ -6180,6 +6243,7 @@ def _send_error(
             message: str,
             detail: dict[str, Any] | None = None,
         ) -> None:
+            _LOGGER.warning("request_failed status=%s code=%s", status, code)
             self._send(_error_payload(code, message, {"request_id": uuid.uuid4().hex, **(detail or {})}), status)
 
         def _write_response(self, writer: Callable[[], None]) -> bool:
@@ -6200,6 +6264,7 @@ def _write_response(self, writer: Callable[[], None]) -> bool:
                 writer()
                 return True
             except (BrokenPipeError, ConnectionError, OSError):
+                _LOGGER.debug("client_disconnected")
                 return False
 
         def _send(
@@ -6333,8 +6398,16 @@ def serve(
     port: int = 8000,
     security: SecurityConfig | None = None,
     clearfolio_url: str | None = None,
+    coordinator: CostRoutingCoordinator | None = None,
 ) -> None:
     """Serve the admin console and resource-oriented orchestration API."""
-    server = build_server(orchestrator, host=host, port=port, security=security, clearfolio_url=clearfolio_url)
+    server = build_server(
+        orchestrator,
+        host=host,
+        port=port,
+        security=security,
+        clearfolio_url=clearfolio_url,
+        coordinator=coordinator,
+    )
     print(f"listening on http://{host}:{port}")
     server.serve_forever()