From b9aa9e8af353766ce61922fa10a26b5d50cda42d Mon Sep 17 00:00:00 2001 From: Pablo Pardo Garcia Date: Fri, 4 Sep 2026 15:55:21 +0200 Subject: [PATCH] feat: context-scoped multi-workspace routing --- src/rius/__init__.py | 5 +- src/rius/client.py | 74 +++++++++++- src/rius/semconv.py | 9 ++ src/rius/workspace.py | 212 ++++++++++++++++++++++++++++++++++ tests/test_workspace.py | 247 ++++++++++++++++++++++++++++++++++++++++ 5 files changed, 545 insertions(+), 2 deletions(-) create mode 100644 src/rius/workspace.py create mode 100644 tests/test_workspace.py diff --git a/src/rius/__init__.py b/src/rius/__init__.py index ca7fa66..cbf7f9e 100644 --- a/src/rius/__init__.py +++ b/src/rius/__init__.py @@ -2,7 +2,7 @@ __version__ = "0.12.0" # x-release-please-version -from .client import GlassflowClient, build_span_exporter, get_tracer, init +from .client import GlassflowClient, build_span_exporter, get_tracer, init, register_workspace from .config import GlassflowConfig, resolve_config from .export_health import ProbeTransport from .generation import Generation, start_as_current_generation, start_generation @@ -10,6 +10,7 @@ from .semconv import SpanKind from .session import session from .spans import Observation, start_as_current_span, start_span +from .workspace import workspace __all__ = [ "Generation", @@ -23,10 +24,12 @@ "get_tracer", "init", "observe", + "register_workspace", "resolve_config", "session", "start_as_current_generation", "start_as_current_span", "start_generation", "start_span", + "workspace", ] diff --git a/src/rius/client.py b/src/rius/client.py index f69b22c..38a851f 100644 --- a/src/rius/client.py +++ b/src/rius/client.py @@ -29,6 +29,7 @@ from .pending import PendingSpanProcessor from .semconv import SERVICE_INSTANCE_ID, TRACER_NAME from .session import SessionSpanProcessor +from .workspace import ExporterFactory, RoutingSpanExporter, WorkspaceSpanProcessor logger = logging.getLogger(__name__) @@ -78,14 +79,30 @@ def __init__( heartbeat: HeartbeatSender | None = None, export_health: ExportOutcomeExporter | None = None, connectivity_thread: threading.Thread | None = None, + routing: RoutingSpanExporter | None = None, ) -> None: self._provider = provider self.config = config self._heartbeat = heartbeat self._export_health = export_health self._connectivity_thread = connectivity_thread + self._routing = routing self._is_shutdown = False + def register_workspace(self, alias: str, api_key: str) -> None: + """Add (or rotate the key of) a workspace destination at runtime. + + Requires routing to be enabled at ``init`` time via ``workspaces=`` + (an empty dict opts in with no static routes). Spans started inside + ``rius.workspace(alias)`` are then exported with ``api_key``. + """ + if self._routing is None: + raise RuntimeError( + "workspace routing is not enabled: pass workspaces={...} to init() " + "(an empty dict is fine) to opt in before registering destinations" + ) + self._routing.register(alias, api_key) + def get_tracer(self, name: str = TRACER_NAME) -> trace.Tracer: """Return a tracer bound to this client's provider. @@ -145,6 +162,8 @@ def init( partial_spans: bool | None = None, partial_spans_delay: float | None = None, session_id: str | None = None, + workspaces: dict[str, str] | None = None, + workspace_exporter_factory: ExporterFactory | None = None, set_global: bool = True, ) -> GlassflowClient: """Initialize the SDK: build a tracer provider that exports OTLP traces. @@ -197,6 +216,17 @@ def init( process's traces into one session. For one-run-per-process agents; a server handling many sessions scopes each one with ``rius.session()`` instead, which overrides this default. + workspaces: Enable multi-workspace routing: a mapping of alias to + API key. Spans started inside ``rius.workspace(alias)`` are + exported with that workspace's key; spans outside any scope use + the default ``api_key``. Pass ``{}`` to opt in with no static + routes and register destinations later via + ``register_workspace()``. One trace must stay inside one + workspace; see ``rius.workspace``. + workspace_exporter_factory: Override how per-workspace exporters are + built from an API key (useful for testing, like + ``span_exporter``). Defaults to the standard OTLP exporter + against the configured endpoint. set_global: Register the provider as the global OpenTelemetry provider. """ global _current_client @@ -227,6 +257,8 @@ def init( partial_spans=partial_spans, partial_spans_delay=partial_spans_delay, session_id=session_id, + workspaces=workspaces, + workspace_exporter_factory=workspace_exporter_factory, set_global=set_global, ) @@ -251,6 +283,8 @@ def _do_init( partial_spans: bool | None, partial_spans_delay: float | None, session_id: str | None, + workspaces: dict[str, str] | None, + workspace_exporter_factory: ExporterFactory | None, set_global: bool, ) -> GlassflowClient: global _current_client @@ -292,6 +326,7 @@ def _do_init( export_health: ExportOutcomeExporter | None = None connectivity_thread: threading.Thread | None = None + routing: RoutingSpanExporter | None = None if not config.disabled: if span_exporter is None: if _missing_managed_credentials(config): @@ -316,6 +351,12 @@ def _do_init( ) connectivity_thread.start() exporter = span_exporter if span_exporter is not None else build_span_exporter(config) + if workspaces is not None: + # Innermost in the chain, so masking and export-health wrap the + # whole fan-out and apply to every destination alike. + factory = workspace_exporter_factory or _workspace_exporter_factory(config) + routing = RoutingSpanExporter(exporter, factory, routes=workspaces) + exporter = routing if not config.capture_content or mask is not None: exporter = MaskingSpanExporter( exporter, capture_content=config.capture_content, mask=mask @@ -327,8 +368,12 @@ def _do_init( batch_processor = BatchSpanProcessor(export_health) # Registered BEFORE the pending processor: both act at on_start, and # the pending snapshot is built from the attributes already on the - # span, so the session id must be stamped first to ride it. + # span, so the session id (and the workspace route, which decides + # which destination the snapshot itself goes to) must be stamped + # first to ride it. provider.add_span_processor(SessionSpanProcessor(config.session_id)) + if routing is not None: + provider.add_span_processor(WorkspaceSpanProcessor()) if config.partial_spans: # Pending snapshots ride the SAME batch pipeline as final spans # (exporter, retries, masking); see pending.py for the contract. @@ -376,6 +421,7 @@ def _do_init( heartbeat=sender, export_health=export_health, connectivity_thread=connectivity_thread, + routing=routing, ) if set_global: _current_client = client @@ -385,3 +431,29 @@ def _do_init( def get_tracer(name: str = TRACER_NAME) -> trace.Tracer: """Return a tracer from the globally configured provider.""" return trace.get_tracer(name, __version__) + + +def _workspace_exporter_factory(config: GlassflowConfig) -> Callable[[str], SpanExporter]: + """Per-workspace OTLP exporters: same endpoint, that workspace's key.""" + + def build(api_key: str) -> SpanExporter: + headers = { + **{k: v for k, v in (config.headers or {}).items() if k.lower() != "authorization"}, + "Authorization": f"Bearer {api_key}", + } + return OTLPSpanExporter(endpoint=config.traces_endpoint, headers=headers) + + return build + + +def register_workspace(alias: str, api_key: str) -> None: + """Add (or rotate the key of) a workspace destination on the global client. + + The module-level twin of ``client.register_workspace()``. Requires a + global ``init(workspaces=...)`` to have opted into routing. + """ + with _lock: + client = _current_client + if client is None: + raise RuntimeError("rius.init() has not been called (no global client)") + client.register_workspace(alias, api_key) diff --git a/src/rius/semconv.py b/src/rius/semconv.py index 112c683..bf75374 100644 --- a/src/rius/semconv.py +++ b/src/rius/semconv.py @@ -33,6 +33,12 @@ # alongside it, one name for one fact. SESSION_ID = "session.id" +# Process-local routing marker for multi-workspace export (see workspace.py). +# Stamped at span start so pending snapshots route too, and ALWAYS stripped by +# the routing exporter before spans leave the process: the destination's API +# key is what tells the backend which workspace a span belongs to. +WORKSPACE_ROUTE = "rius.workspace" + # OTel GenAI (subset we emit) GEN_AI_OPERATION_NAME = "gen_ai.operation.name" GEN_AI_PROVIDER_NAME = "gen_ai.provider.name" @@ -81,6 +87,9 @@ # Identity, not content: a pending span must be groupable into its # session while still running, that is the live view's whole point. SESSION_ID, + # Routing, not content: a crashed run's snapshot must land in the same + # workspace its final span would have. Stripped at export either way. + WORKSPACE_ROUTE, } ) # gen_ai.request.* (model, temperature, ...) is identity, not content. diff --git a/src/rius/workspace.py b/src/rius/workspace.py new file mode 100644 index 0000000..e0e3315 --- /dev/null +++ b/src/rius/workspace.py @@ -0,0 +1,212 @@ +"""Workspaces: route spans from one process to per-customer destinations. + +One client, one provider, one batch pipeline; the *destination* is a +context-scoped property. ``workspace(alias)`` sets an OTel context key +(exactly like ``session()``), ``WorkspaceSpanProcessor`` stamps it as a +transient attribute at span start, and ``RoutingSpanExporter`` partitions +each export batch by that attribute, strips it, and forwards every partition +to the exporter registered for its alias. Spans started outside any scope go +to the default destination. + +Because the alias rides OTel context, everything started in scope routes +together: ``observe`` wrappers, generations, sessions, and spans created by +auto-instrumentation, which is the property a second client can never give +(instrumentors are process-global and bound to one provider). + +Two rules the design enforces or warns about: + +- The routing attribute never reaches the wire. The destination's API key is + what tells the backend which workspace a span belongs to; the alias is + process-local configuration, so the exporter strips it from copies (the + same copy-on-write discipline as ``MaskingSpanExporter``). +- One trace, one workspace. The backend derives the workspace from the API + key per request, so a trace split across scopes would come apart. Starting + a span under a different alias than its parent's logs a warning; switch + workspaces at request boundaries, not inside a trace. +""" + +from __future__ import annotations + +import copy +import logging +import threading +from collections.abc import Callable, Iterator, Sequence +from contextlib import contextmanager + +from opentelemetry import context as otel_context +from opentelemetry import trace +from opentelemetry.context import Context +from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor +from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult + +from .semconv import WORKSPACE_ROUTE + +logger = logging.getLogger(__name__) + +_WORKSPACE_KEY = otel_context.create_key("rius-workspace-alias") + +ExporterFactory = Callable[[str], SpanExporter] + + +@contextmanager +def workspace(alias: str) -> Iterator[str]: + """Scope every span started in the block to one workspace destination. + + ``alias`` names a workspace registered via ``init(workspaces={...})`` or + ``register_workspace()``; the block's spans are exported with that + workspace's API key. Scopes nest and unwind with the block, and follow + async tasks the way all OTel context does, but a trace must stay inside + one workspace: enter the scope at a request boundary, before the root + span starts. + + Example: + + ```python + with rius.workspace("acme"): + handle(request) # every span of the request lands in acme's workspace + ``` + """ + if not alias: + raise ValueError("workspace alias must be a non-empty string") + token = otel_context.attach(otel_context.set_value(_WORKSPACE_KEY, alias)) + try: + yield alias + finally: + otel_context.detach(token) + + +class WorkspaceSpanProcessor(SpanProcessor): + """Stamps the active workspace alias on every span at start. + + Stamping happens in ``on_start`` for the same reason sessions stamp + there: pending snapshots are built from start-time attributes, and a + snapshot must route to the same workspace its final span will. + """ + + def on_start(self, span: Span, parent_context: Context | None = None) -> None: + alias = otel_context.get_value(_WORKSPACE_KEY, context=parent_context) + if alias is None: + return + parent = trace.get_current_span(parent_context) + parent_alias = None + if isinstance(parent, ReadableSpan) and parent.attributes: + parent_alias = parent.attributes.get(WORKSPACE_ROUTE) + if parent_alias is not None and parent_alias != alias: + logger.warning( + "span %r starts under workspace %r but its parent is stamped %r; " + "a trace cannot straddle two workspaces (the backend derives the " + "workspace from the API key). Switch workspaces at request " + "boundaries, before the root span starts.", + span.name, + alias, + parent_alias, + ) + span.set_attribute(WORKSPACE_ROUTE, str(alias)) + + def on_end(self, span: ReadableSpan) -> None: # pragma: no cover - no-op + pass + + def shutdown(self) -> None: # pragma: no cover - no-op + pass + + def force_flush(self, timeout_millis: int = 30_000) -> bool: # pragma: no cover + return True + + +class RoutingSpanExporter(SpanExporter): + """Partition each batch by the routing attribute and fan out. + + Wraps the default exporter plus one lazily-created exporter per + registered workspace key. Sits innermost in the export chain, so masking + and export-health wrap the whole fan-out. Spans with no stamp go to the + default destination; a stamp whose alias is not registered also goes to + the default (the vendor's own workspace) with a warning, which keeps the + data rather than dropping it and cannot leak one customer's spans to + another. + """ + + def __init__( + self, + default_exporter: SpanExporter, + exporter_factory: ExporterFactory, + routes: dict[str, str] | None = None, + ) -> None: + self._default = default_exporter + self._factory = exporter_factory + self._routes: dict[str, str] = dict(routes or {}) + self._exporters: dict[str, SpanExporter] = {} + self._warned_aliases: set[str] = set() + self._lock = threading.Lock() + + def register(self, alias: str, api_key: str) -> None: + """Add or replace a route. Replacing supports key rotation.""" + if not alias or not api_key: + raise ValueError("workspace alias and api_key must be non-empty strings") + with self._lock: + self._routes[alias] = api_key + self._warned_aliases.discard(alias) + + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: + groups: dict[str | None, list[ReadableSpan]] = {} + for span in spans: + alias = None + if span.attributes and WORKSPACE_ROUTE in span.attributes: + raw = span.attributes[WORKSPACE_ROUTE] + alias = raw if isinstance(raw, str) else str(raw) + span = self._stripped(span) + groups.setdefault(alias, []).append(span) + + result = SpanExportResult.SUCCESS + for alias, group in groups.items(): + exporter = self._resolve(alias) + if exporter.export(group) is not SpanExportResult.SUCCESS: + result = SpanExportResult.FAILURE + return result + + def _resolve(self, alias: str | None) -> SpanExporter: + if alias is None: + return self._default + with self._lock: + api_key = self._routes.get(alias) + if api_key is None: + if alias not in self._warned_aliases: + self._warned_aliases.add(alias) + logger.warning( + "no workspace registered for alias %r; its spans go to the " + "default destination. Register it with " + "rius.register_workspace(%r, api_key) or in " + "init(workspaces={...}).", + alias, + alias, + ) + return self._default + exporter = self._exporters.get(api_key) + if exporter is None: + exporter = self._factory(api_key) + self._exporters[api_key] = exporter + return exporter + + @staticmethod + def _stripped(span: ReadableSpan) -> ReadableSpan: + # Copy-on-write, same discipline as MaskingSpanExporter: the attribute + # dict is shared with every other processor on the provider. + new_attributes = dict(span.attributes or {}) + del new_attributes[WORKSPACE_ROUTE] + stripped = copy.copy(span) + stripped._attributes = new_attributes + return stripped + + def force_flush(self, timeout_millis: int = 30_000) -> bool: + ok = self._default.force_flush(timeout_millis) + with self._lock: + exporters = list(self._exporters.values()) + for exporter in exporters: + ok = exporter.force_flush(timeout_millis) and ok + return ok + + def shutdown(self) -> None: + self._default.shutdown() + with self._lock: + exporters = list(self._exporters.values()) + for exporter in exporters: + exporter.shutdown() diff --git a/tests/test_workspace.py b/tests/test_workspace.py new file mode 100644 index 0000000..f7b84ab --- /dev/null +++ b/tests/test_workspace.py @@ -0,0 +1,247 @@ +"""Workspaces: context-scoped routing of spans to per-customer destinations. + +Wire contract under test: spans started inside ``workspace(alias)`` are +delivered by the exporter registered for that alias, spans outside any scope +go to the default exporter, and the transient routing attribute is stripped +before spans leave the process (the destination's API key already says which +workspace a span belongs to; the alias is process-local configuration). +""" + +from __future__ import annotations + +import logging + +import pytest +from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from rius import init, workspace +from rius.semconv import GLASSFLOW_SPAN_PENDING, PENDING_IDENTITY_ATTRIBUTES, WORKSPACE_ROUTE + + +def _routed_client(**kwargs: object): + """A scoped client with one InMemory exporter per workspace key.""" + default = InMemorySpanExporter() + per_key: dict[str, InMemorySpanExporter] = {} + + def factory(api_key: str) -> SpanExporter: + per_key[api_key] = InMemorySpanExporter() + return per_key[api_key] + + client = init( + span_exporter=default, + workspaces={"acme": "key-acme", "globex": "key-globex"}, + workspace_exporter_factory=factory, + set_global=False, + service_name="test-svc", + instruments=[], + **kwargs, # type: ignore[arg-type] + ) + return client, default, per_key + + +# --- the scoped API --- + + +def test_scope_routes_every_span_to_the_alias_exporter() -> None: + client, default, per_key = _routed_client() + tracer = client.get_tracer() + with ( + workspace("acme"), + tracer.start_as_current_span("root"), + tracer.start_as_current_span("child"), + ): + pass + client.flush() + assert default.get_finished_spans() == () + spans = per_key["key-acme"].get_finished_spans() + assert [s.name for s in spans] == ["child", "root"] + + +def test_no_scope_routes_to_the_default_exporter() -> None: + client, default, per_key = _routed_client() + with client.get_tracer().start_as_current_span("bare"): + pass + client.flush() + assert len(default.get_finished_spans()) == 1 + assert "key-acme" not in per_key or per_key["key-acme"].get_finished_spans() == () + + +def test_two_scopes_partition_one_batch() -> None: + client, default, per_key = _routed_client() + tracer = client.get_tracer() + with workspace("acme"), tracer.start_as_current_span("for-acme"): + pass + with workspace("globex"), tracer.start_as_current_span("for-globex"): + pass + with tracer.start_as_current_span("for-default"): + pass + client.flush() + assert [s.name for s in per_key["key-acme"].get_finished_spans()] == ["for-acme"] + assert [s.name for s in per_key["key-globex"].get_finished_spans()] == ["for-globex"] + assert [s.name for s in default.get_finished_spans()] == ["for-default"] + + +def test_routing_attribute_never_reaches_the_wire() -> None: + client, default, per_key = _routed_client() + with workspace("acme"), client.get_tracer().start_as_current_span("root"): + pass + with client.get_tracer().start_as_current_span("bare"): + pass + client.flush() + for exporter in (default, per_key["key-acme"]): + for span in exporter.get_finished_spans(): + assert WORKSPACE_ROUTE not in (span.attributes or {}) + + +def test_scope_ends_at_the_block() -> None: + client, default, per_key = _routed_client() + tracer = client.get_tracer() + with workspace("acme"), tracer.start_as_current_span("inside"): + pass + with tracer.start_as_current_span("after"): + pass + client.flush() + assert [s.name for s in per_key["key-acme"].get_finished_spans()] == ["inside"] + assert [s.name for s in default.get_finished_spans()] == ["after"] + + +def test_unknown_alias_falls_back_to_default_with_a_warning( + caplog: pytest.LogCaptureFixture, +) -> None: + client, default, per_key = _routed_client() + with caplog.at_level(logging.WARNING, logger="rius.workspace"): + with workspace("no-such-customer"), client.get_tracer().start_as_current_span("lost"): + pass + client.flush() + assert [s.name for s in default.get_finished_spans()] == ["lost"] + assert any("no-such-customer" in r.message for r in caplog.records) + + +def test_switching_workspace_inside_a_trace_warns(caplog: pytest.LogCaptureFixture) -> None: + client, default, per_key = _routed_client() + tracer = client.get_tracer() + with ( + caplog.at_level(logging.WARNING, logger="rius.workspace"), + workspace("acme"), + tracer.start_as_current_span("root"), + workspace("globex"), + tracer.start_as_current_span("child"), + ): + pass + client.flush() + assert any("workspace" in r.message and "trace" in r.message for r in caplog.records) + + +# --- dynamic registration --- + + +def test_register_workspace_after_init() -> None: + client, default, per_key = _routed_client() + client.register_workspace("initech", "key-initech") + with workspace("initech"), client.get_tracer().start_as_current_span("late"): + pass + client.flush() + assert [s.name for s in per_key["key-initech"].get_finished_spans()] == ["late"] + + +def test_register_workspace_without_routing_raises() -> None: + exporter = InMemorySpanExporter() + client = init( + span_exporter=exporter, + set_global=False, + service_name="test-svc", + instruments=[], + ) + with pytest.raises(RuntimeError, match="workspaces"): + client.register_workspace("acme", "key-acme") + + +def test_empty_workspaces_dict_enables_pure_dynamic_routing() -> None: + default = InMemorySpanExporter() + per_key: dict[str, InMemorySpanExporter] = {} + + def factory(api_key: str) -> SpanExporter: + per_key[api_key] = InMemorySpanExporter() + return per_key[api_key] + + client = init( + span_exporter=default, + workspaces={}, + workspace_exporter_factory=factory, + set_global=False, + service_name="test-svc", + instruments=[], + ) + client.register_workspace("acme", "key-acme") + with workspace("acme"), client.get_tracer().start_as_current_span("routed"): + pass + client.flush() + assert [s.name for s in per_key["key-acme"].get_finished_spans()] == ["routed"] + + +# --- lifecycle fan-out --- + + +class _RecordingExporter(SpanExporter): + def __init__(self) -> None: + self.exported: list[str] = [] + self.flushed = False + self.shut_down = False + + def export(self, spans) -> SpanExportResult: # type: ignore[no-untyped-def] + self.exported.extend(s.name for s in spans) + return SpanExportResult.SUCCESS + + def force_flush(self, timeout_millis: int = 30_000) -> bool: + self.flushed = True + return True + + def shutdown(self) -> None: + self.shut_down = True + + +def test_shutdown_fans_out_and_flush_delivers_to_every_destination() -> None: + # Note: OTel's BatchSpanProcessor.force_flush drains its queue through + # export() and never calls exporter.force_flush(), so the delivery + # contract under test is "flushed spans reach the destination" plus + # "shutdown reaches every destination exporter". + default = _RecordingExporter() + created: dict[str, _RecordingExporter] = {} + + def factory(api_key: str) -> SpanExporter: + created[api_key] = _RecordingExporter() + return created[api_key] + + client = init( + span_exporter=default, + workspaces={"acme": "key-acme"}, + workspace_exporter_factory=factory, + set_global=False, + service_name="test-svc", + instruments=[], + ) + with workspace("acme"), client.get_tracer().start_as_current_span("s"): + pass + client.flush() + assert created["key-acme"].exported == ["s"] + client.shutdown() + assert default.shut_down + assert created["key-acme"].shut_down + + +# --- pending snapshots route too --- + + +def test_pending_snapshot_routes_with_the_scope() -> None: + assert WORKSPACE_ROUTE in PENDING_IDENTITY_ATTRIBUTES + client, default, per_key = _routed_client(partial_spans=True) + tracer = client.get_tracer() + with workspace("acme"), tracer.start_as_current_span("long-run"): + client.flush() + pending = per_key["key-acme"].get_finished_spans() + assert len(pending) == 1 + assert pending[0].attributes is not None + assert pending[0].attributes.get(GLASSFLOW_SPAN_PENDING) is True + assert WORKSPACE_ROUTE not in pending[0].attributes + client.flush()