Tag traced requests with the session id and OIDC sub

server/tests/test_otel.py

@@ -16,6 +16,7 @@ from unittest.mock import patch
 
 from tavolo.app import otel_mixin_from_settings
 from tavolo.config import Settings
+from tavolo.telemetry import session_tag_hook
 
 
 def _settings(env: dict) -> Settings:
@@ -58,8 +59,17 @@ class OTelMixinFromSettingsTests(unittest.TestCase):
             "endpoint": "http://collector:4318",
             "headers": {"Authorization": "Bearer t", "X-Tenant": "one"},
             "excluded_paths": ("/api/health",),
+            "server_response_hook": session_tag_hook,
+            "websocket_close_hook": session_tag_hook,
         }, mixin.kwargs)
 
+    def test_session_tag_hooks_are_wired(self):
+        with _stub_kaya_otel():
+            mixin = otel_mixin_from_settings(_settings({"OTEL_ENABLED": "1"}))
+        assert isinstance(mixin, _StubOTelMixin)
+        self.assertIs(session_tag_hook, mixin.kwargs["server_response_hook"])
+        self.assertIs(session_tag_hook, mixin.kwargs["websocket_close_hook"])
+
     def test_defaults_when_only_enabled(self):
         with _stub_kaya_otel():
             mixin = otel_mixin_from_settings(_settings({"OTEL_ENABLED": "on"}))

server/tests/test_telemetry.py

@@ -5,6 +5,7 @@ so the provider-dependent tests skip when they are not installed.
 """
 from __future__ import annotations
 
+import hashlib
 import logging
 import os
 import sys
@@ -19,6 +20,7 @@ from tavolo.telemetry import (
     CountingRedisSessionStore,
     TelemetryMixin,
     _NoOtelLogs,
+    session_tag_hook,
     telemetry_from_settings,
 )
 from tavolo.config import Settings
@@ -137,6 +139,119 @@ class NoOtelLogsFilterTest(unittest.TestCase):
         )
 
 
+class _RecordedSpan:
+    def __init__(self) -> None:
+        self.attributes: dict = {}
+
+    def set_attribute(self, key: str, value) -> None:
+        self.attributes[key] = value
+
+
+class SessionTagHookTest(unittest.TestCase):
+    def test_tags_authenticated_session(self) -> None:
+        span = _RecordedSpan()
+        carrier = SimpleNamespace(
+            session=Session("sid", {"oidc_user": {"sub": "alice"}})
+        )
+        session_tag_hook(span, carrier)
+        self.assertEqual(
+            hashlib.sha256(b"sid").hexdigest()[:16],
+            span.attributes["session.id"],
+        )
+        self.assertEqual("alice", span.attributes["user.id"])
+
+    def test_tags_session_without_user(self) -> None:
+        span = _RecordedSpan()
+        session_tag_hook(span, SimpleNamespace(session=Session("sid", {})))
+        self.assertIn("session.id", span.attributes)
+        self.assertNotIn("user.id", span.attributes)
+
+    def test_tags_user_when_session_id_not_yet_minted(self) -> None:
+        span = _RecordedSpan()
+        carrier = SimpleNamespace(
+            session=Session(None, {"oidc_user": {"sub": "alice"}})
+        )
+        session_tag_hook(span, carrier)
+        self.assertNotIn("session.id", span.attributes)
+        self.assertEqual("alice", span.attributes["user.id"])
+
+    def test_missing_session_is_a_noop(self) -> None:
+        span = _RecordedSpan()
+        session_tag_hook(span, SimpleNamespace())
+        self.assertEqual({}, span.attributes)
+
+    def test_malformed_user_is_ignored(self) -> None:
+        span = _RecordedSpan()
+        carrier = SimpleNamespace(session=Session("sid", {"oidc_user": "alice"}))
+        session_tag_hook(span, carrier)
+        self.assertIn("session.id", span.attributes)
+        self.assertNotIn("user.id", span.attributes)
+
+
+@unittest.skipUnless(HAVE_OTEL, "tavolo-app[otel] not installed")
+class SessionTagIntegrationTest(unittest.TestCase):
+    """The hook is wired as kaya-otel's response hook, so it sees the
+    session after the session mixin persisted and minted it."""
+
+    @async_test
+    async def test_spans_carry_hashed_session_and_oidc_sub(self) -> None:
+        from httpx import ASGITransport, AsyncClient
+        from kaya.core import HttpContext, KayaApp
+        from kaya.otel import OTelMixin
+        from kaya.session import InMemorySessionStore, SessionMixin
+
+        span_exporter = InMemorySpanExporter()
+        tracer_provider = TracerProvider(
+            resource=Resource.create({"service.name": "test"})
+        )
+        tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
+        app = KayaApp(
+            mixins=[
+                OTelMixin(
+                    service_name="test",
+                    tracer_provider=tracer_provider,
+                    meter_provider=MeterProvider(
+                        metric_readers=[InMemoryMetricReader()]
+                    ),
+                    server_response_hook=session_tag_hook,
+                    websocket_close_hook=session_tag_hook,
+                ),
+                SessionMixin(InMemorySessionStore()),
+            ]
+        )
+
+        @app.GET("/login")
+        async def login(ctx: HttpContext) -> None:
+            session = ctx.session
+            assert session is not None
+            session["oidc_user"] = {"sub": "alice"}
+            await ctx.send_str(200, "ok")
+
+        @app.GET("/me")
+        async def me(ctx: HttpContext) -> None:
+            await ctx.send_str(200, "ok")
+
+        transport = ASGITransport(app=app)  # type: ignore[arg-type]
+        async with AsyncClient(
+            transport=transport, base_url="http://127.0.0.1"
+        ) as client:
+            first = await client.get("/login")
+            second = await client.get("/me")
+            self.assertEqual(200, first.status_code)
+            self.assertEqual(200, second.status_code)
+        # The cookie value is the session id; only its digest is exported.
+        cookie = first.cookies["session_id"]
+        digest = hashlib.sha256(cookie.encode()).hexdigest()[:16]
+
+        spans = {span.name: span for span in span_exporter.get_finished_spans()}
+        self.assertEqual({"GET /login", "GET /me"}, set(spans))
+        for name, span in spans.items():
+            attributes = span.attributes
+            assert attributes is not None
+            self.assertEqual(digest, attributes["session.id"], name)
+            self.assertEqual("alice", attributes["user.id"], name)
+
+
 def _metric_points(reader) -> dict:
     """Map metric name -> list of (attributes, value) data points."""
     data = reader.get_metrics_data()