plba

src/app_runtime/tracing/store.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from contextvars import ContextVar
+from dataclasses import dataclass, replace
+
+from app_runtime.contracts.trace import TraceContextRecord
+
+
+@dataclass(frozen=True)
+class ActiveTraceContext:
+    record: TraceContextRecord
+    step: str = ""
+
+
+class TraceContextStore:
+    def __init__(self) -> None:
+        self._current: ContextVar[ActiveTraceContext | None] = ContextVar("trace_current", default=None)
+        self._stack: ContextVar[tuple[ActiveTraceContext, ...]] = ContextVar("trace_stack", default=())
+
+    def current(self) -> ActiveTraceContext | None:
+        return self._current.get()
+
+    def current_trace_id(self) -> str | None:
+        active = self.current()
+        return active.record.trace_id if active else None
+
+    def push(self, record: TraceContextRecord) -> ActiveTraceContext:
+        active = self.current()
+        stack = self._stack.get()
+        if active is not None:
+            self._stack.set(stack + (active,))
+        updated = ActiveTraceContext(record=record)
+        self._current.set(updated)
+        return updated
+
+    def pop(self) -> ActiveTraceContext | None:
+        stack = self._stack.get()
+        if not stack:
+            self._current.set(None)
+            return None
+        previous = stack[-1]
+        self._stack.set(stack[:-1])
+        self._current.set(previous)
+        return previous
+
+    def set_step(self, step: str) -> ActiveTraceContext | None:
+        active = self.current()
+        if active is None:
+            return None
+        updated = replace(active, step=step)
+        self._current.set(updated)
+        return updated