diff --git a/py/noxfile.py b/py/noxfile.py index 5a300353..5ecb41e7 100644 --- a/py/noxfile.py +++ b/py/noxfile.py @@ -827,6 +827,11 @@ def test_transformers(session, version): def test_temporal(session, version): _install_test_deps(session) _install_matrix_dep(session, "temporalio", version) + # Pydantic AI's Temporal durability extra needs temporalio>=1.27 and Python 3.11+, + # so the Temporal + Pydantic AI tests only run against the latest temporalio pin. + if version == LATEST and sys.version_info >= (3, 11): + pydantic_ai_version = _MATRIX["pydantic-ai-integration"]["latest"].split("==", 1)[1] + _session_install(session, f"pydantic-ai-slim[temporal]=={pydantic_ai_version}", silent=SILENT_INSTALLS) _run_tests(session, f"{INTEGRATION_DIR}/temporal", version=version) @@ -1036,11 +1041,16 @@ def pylint(session): # The lint group skips crewai to avoid vulnerable transitive chromadb # versions, so skip the matching example too. files = [f for f in files if not f.startswith("../examples/crewai/")] + # Tests outside an integration's own directory that import it. + cross_integration_files = { + "pydantic_ai": {f"src/{INTEGRATION_DIR}/temporal/test_temporal_pydantic_ai.py"}, + } for integration in sorted(unlintable_integrations): files = [ f for f in files if not f.startswith(f"src/{INTEGRATION_DIR}/{integration}/") + and f not in cross_integration_files.get(integration, set()) and f != f"src/{INTEGRATION_DIR}/auto_test_scripts/test_auto_{integration}.py" and not f.startswith(f"src/{INTEGRATION_DIR}/auto_test_scripts/test_manual_wrap_{integration}_") and not f.startswith(f"../examples/{integration}/") diff --git a/py/src/braintrust/integrations/temporal/plugin.py b/py/src/braintrust/integrations/temporal/plugin.py index 60881c9f..3841ad09 100644 --- a/py/src/braintrust/integrations/temporal/plugin.py +++ b/py/src/braintrust/integrations/temporal/plugin.py @@ -82,7 +82,7 @@ import dataclasses import hashlib import uuid -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import Any import braintrust @@ -256,6 +256,8 @@ def __init__(self, next: temporalio.worker.WorkflowInboundInterceptor) -> None: self.payload_converter = temporalio.converter.PayloadConverter.default self._parent_span_context: Any | None = None self._workflow_span: Any | None = None + self._span_hooks: _WorkflowSpanHooks | None = None + self._span_hooks_token: Any | None = None def init(self, outbound: temporalio.worker.WorkflowOutboundInterceptor) -> None: super().init(_BraintrustWorkflowOutboundInterceptor(outbound, self)) @@ -282,25 +284,33 @@ async def execute_workflow(self, input: temporalio.worker.ExecuteWorkflowInput) parent = parent_span_context or logger.export() if temporalio.workflow.patched(_SPAN_CONTEXT_PATCH_ID): ids = _workflow_span_ids(info, parent) - if temporalio.workflow.unsafe.is_replaying(): - self._parent_span_context = _workflow_span_context(parent, ids) - else: - span = logger.start_span( - internal={"instrumentation": "temporal-auto"}, - name=f"temporal.workflow.{info.workflow_type}", - type="task", - parent=parent, - id=ids["row_id"], - span_id=ids["span_id"], - root_span_id=ids["root_span_id"], - metadata={ - "temporal.workflow_type": info.workflow_type, - "temporal.workflow_id": info.workflow_id, - "temporal.run_id": info.run_id, - }, - ) - span.set_current() - self._parent_span_context = _workflow_span_context(parent, ids) + # Spans opened in workflow code (the workflow span, user spans, Pydantic AI + # agent/tool spans) get deterministic IDs and are not written while replaying, + # so a replay rebuilds the same span tree instead of logging a second copy. + self._span_hooks = _WorkflowSpanHooks( + _workflow_span_id_base(info), + legacy_ids=SpanComponentsV4.get_version(parent) < 4, + bind_logger=get_logger, + ) + self._span_hooks_token = braintrust.logger._internal_set_span_hooks(self._span_hooks) + # Always create the workflow span, even on replay, so it is the current span for + # in-workflow spans and its end is logged if the workflow finishes after a replay. + span = logger.start_span( + internal={"instrumentation": "temporal-auto"}, + name=f"temporal.workflow.{info.workflow_type}", + type="task", + parent=parent, + id=ids["row_id"], + span_id=ids["span_id"], + root_span_id=ids["root_span_id"], + metadata={ + "temporal.workflow_type": info.workflow_type, + "temporal.workflow_id": info.workflow_id, + "temporal.run_id": info.run_id, + }, + ) + span.set_current() + self._parent_span_context = _workflow_span_context(parent, ids) elif not temporalio.workflow.unsafe.is_replaying(): span = logger.start_span( internal={"instrumentation": "temporal-auto"}, @@ -326,10 +336,13 @@ async def execute_workflow(self, input: temporalio.worker.ExecuteWorkflowInput) self._workflow_span.log(error=e) raise finally: - if self._workflow_span: - with temporalio.workflow.unsafe.sandbox_unrestricted(): + with temporalio.workflow.unsafe.sandbox_unrestricted(): + if self._workflow_span: self._workflow_span.unset_current() self._workflow_span.end() + if self._span_hooks_token is not None: + braintrust.logger._internal_reset_span_hooks(self._span_hooks_token) + self._span_hooks_token = None class _BraintrustWorkflowOutboundInterceptor(temporalio.worker.WorkflowOutboundInterceptor): @@ -348,12 +361,25 @@ def _add_span_context_to_headers( ) -> Mapping[str, temporalio.api.common.v1.Payload]: # Always propagate context — even during replay — so activities can find their parent. # Replay safety is enforced in the activity interceptor, which only spans on real execution. - if self.root._parent_span_context: - payloads = self.root.payload_converter.to_payloads([self.root._parent_span_context]) + parent = self._current_parent_span_context() + if parent: + payloads = self.root.payload_converter.to_payloads([parent]) if payloads: return {**headers, _HEADER_KEY: payloads[0]} return headers + def _current_parent_span_context(self) -> Any | None: + # Prefer a span opened inside workflow code (e.g. a Pydantic AI tool span) so activities + # and child workflows nest under it. Without span hooks (workflows started before the + # span-context patch), spans created during replay have fresh IDs, so only trust the + # current span on real execution there. + if self.root._span_hooks is not None or not temporalio.workflow.unsafe.is_replaying(): + with temporalio.workflow.unsafe.sandbox_unrestricted(): + span = _current_span() + if span is not braintrust.logger.NOOP_SPAN and span is not self.root._workflow_span: + return span.export() + return self.root._parent_span_context + def start_activity(self, input: temporalio.worker.StartActivityInput) -> temporalio.workflow.ActivityHandle: input.headers = self._add_span_context_to_headers(input.headers) return super().start_activity(input) @@ -378,10 +404,48 @@ def _modify_workflow_runner(existing: WorkflowRunner | None) -> WorkflowRunner | return existing -def _workflow_span_ids(info: temporalio.workflow.Info, parent: str) -> dict[str, str]: +class _WorkflowSpanHooks: + """Span hooks for one workflow run: deterministic IDs and no writes while replaying. + + Workflow code is deterministic, so the Nth span opened in a run is the same span on every + replay. Numbering spans per run therefore gives a replayed span the ID it was first logged with. + """ + + def __init__(self, base: str, legacy_ids: bool, bind_logger: Callable[[], Any]) -> None: + self._base = base + self._legacy_ids = legacy_ids + # Points this thread at the plugin's logger. The logger override is thread-local, and + # Temporal can run each activation of a workflow on a different worker thread. + self._bind_logger = bind_logger + self._counter = 0 + + def next_ids(self) -> tuple[str, str, str]: + self._counter += 1 + ids = _deterministic_span_ids(f"{self._base}:span:{self._counter}", self._legacy_ids) + return ids["row_id"], ids["span_id"], ids["root_span_id"] + + def should_write(self) -> bool: + try: + if temporalio.workflow.unsafe.is_replaying(): + return False + except Exception: + # Not in workflow code (e.g. the span is ended from outside the workflow event loop). + return True + self._bind_logger() + return True + + +def _workflow_span_id_base(info: temporalio.workflow.Info) -> str: namespace = getattr(info, "namespace", "") - base = f"{_SPAN_ID_NAMESPACE}:{namespace}:{info.workflow_type}:{info.workflow_id}:{info.run_id}" - if SpanComponentsV4.get_version(parent) >= 4: + return f"{_SPAN_ID_NAMESPACE}:{namespace}:{info.workflow_type}:{info.workflow_id}:{info.run_id}" + + +def _workflow_span_ids(info: temporalio.workflow.Info, parent: str) -> dict[str, str]: + return _deterministic_span_ids(_workflow_span_id_base(info), SpanComponentsV4.get_version(parent) < 4) + + +def _deterministic_span_ids(base: str, legacy_ids: bool) -> dict[str, str]: + if not legacy_ids: digest = hashlib.sha256(base.encode()).hexdigest() return { "row_id": str(uuid.uuid5(uuid.NAMESPACE_URL, base)), diff --git a/py/src/braintrust/integrations/temporal/test_temporal.py b/py/src/braintrust/integrations/temporal/test_temporal.py index 2260fdc4..2d1e1554 100644 --- a/py/src/braintrust/integrations/temporal/test_temporal.py +++ b/py/src/braintrust/integrations/temporal/test_temporal.py @@ -1,7 +1,9 @@ """Unit tests for Braintrust Temporal interceptor.""" import asyncio +import concurrent.futures import os +import threading import uuid from collections import Counter from dataclasses import dataclass @@ -219,6 +221,29 @@ async def run(self, input: TaskInput) -> int: return child_result +@temporalio.workflow.defn +class DelegatingParentWorkflow: + """Parent workflow that starts a child workflow and an activity inside an in-workflow span. + + Mirrors Pydantic AI's TemporalAgent, where a tool such as ``delegate_task`` runs in + workflow code and starts a subagent child workflow. + """ + + @temporalio.workflow.run + async def run(self, input: TaskInput) -> int: + with braintrust.start_span(name="delegate_task", type="tool"): + child_result = await temporalio.workflow.execute_child_workflow( + ChildWorkflow.run, + input, + id=f"child-{temporalio.workflow.info().workflow_id}", + ) + return await temporalio.workflow.execute_activity( + simple_activity, + TaskInput(value=child_result), + start_to_close_timeout=timedelta(seconds=10), + ) + + @temporalio.workflow.defn class ReplayAfterSignalWorkflow: """Workflow that schedules an activity after a later workflow task.""" @@ -252,6 +277,37 @@ def state(self) -> str: return self._state +@temporalio.workflow.defn +class InWorkflowSpanReplayWorkflow: + """Opens in-workflow spans across workflow tasks, like a Pydantic AI agent_run span.""" + + def __init__(self) -> None: + self._continue = False + self._state = "starting" + + @temporalio.workflow.run + async def run(self, input: TaskInput) -> int: + with braintrust.start_span(name="agent_run", type="task"): + with braintrust.start_span(name="first_tool", type="tool"): + first = await temporalio.workflow.execute_activity( + simple_activity, input, start_to_close_timeout=timedelta(seconds=10) + ) + self._state = "waiting" + await temporalio.workflow.wait_condition(lambda: self._continue) + with braintrust.start_span(name="second_tool", type="tool"): + return await temporalio.workflow.execute_activity( + simple_activity, TaskInput(value=first), start_to_close_timeout=timedelta(seconds=10) + ) + + @temporalio.workflow.signal + def continue_workflow(self) -> None: + self._continue = True + + @temporalio.workflow.query + def state(self) -> str: + return self._state + + class TestAutoInstrumentation: """Tests for Temporal auto-instrumentation helpers.""" @@ -272,6 +328,28 @@ def test_contrib_temporal_compat_import_deprecated(self): # Integration Tests +class _NewThreadPerTaskExecutor(concurrent.futures.ThreadPoolExecutor): + """Runs every workflow activation on a new thread. + + A shared worker pool can hand each activation of a workflow to a different thread; this makes + that happen every time, so thread-local state cannot be relied on between activations. + """ + + def submit(self, fn, /, *args, **kwargs): + future: concurrent.futures.Future = concurrent.futures.Future() + + def run() -> None: + if not future.set_running_or_notify_cancel(): + return + try: + future.set_result(fn(*args, **kwargs)) + except BaseException as e: + future.set_exception(e) + + threading.Thread(target=run, daemon=True).start() + return future + + @pytest_asyncio.fixture(scope="function") async def temporal_env(): """Create a Temporal test environment. @@ -307,7 +385,7 @@ def _spans_named(spans: list[dict[str, Any]], name: str) -> list[dict[str, Any]] async def _wait_for_workflow_state(handle: Any, expected: str) -> None: for _ in range(50): - if await handle.query(ReplayAfterSignalWorkflow.state) == expected: + if await handle.query("state") == expected: return await asyncio.sleep(0.1) raise AssertionError(f"Workflow did not reach state {expected!r}") @@ -366,6 +444,9 @@ async def test_plugin_activity_after_replay_stays_under_workflow_span( assert len(workflow_spans) == 1 assert len(activity_spans) == 2 workflow_span = workflow_spans[0] + # Evicting the workflow between tasks unwinds its coroutine; the span must still end when the + # workflow really finishes, not when it was first evicted. + assert workflow_span["metrics"]["end"] >= max(span["metrics"]["start"] for span in activity_spans) assert workflow_span["context"]["span_origin"]["instrumentation"]["name"] == "temporal-auto" assert all( span["context"]["span_origin"]["instrumentation"]["name"] == "temporal-auto" for span in activity_spans @@ -449,3 +530,82 @@ async def test_plugin_workflow_tracing( assert [bool(span.get("error")) for span in activity_spans] == expected_activity_errors if any(expected_activity_errors): assert "Simulated failure on first attempt" in activity_spans[0]["error"] + + @pytest.mark.asyncio + @pytest.mark.parametrize("new_thread_per_task", [False, True], ids=["shared_pool", "new_thread_per_task"]) + async def test_child_workflow_and_activity_nest_under_in_workflow_span( + self, temporal_env, memory_logger, new_thread_per_task + ): + task_queue = f"test-queue-delegate-{uuid.uuid4()}" + executor = _NewThreadPerTaskExecutor() if new_thread_per_task else None + async with Worker( + temporal_env.client, + task_queue=task_queue, + workflows=[DelegatingParentWorkflow, ChildWorkflow], + activities=[simple_activity], + workflow_task_executor=executor, + plugins=[BraintrustPlugin(logger=memory_logger)], + ): + await temporal_env.client.execute_workflow( + DelegatingParentWorkflow.run, + TaskInput(value=1), + id=f"test-workflow-{uuid.uuid4()}", + task_queue=task_queue, + ) + + braintrust.flush() + spans = memory_logger.pop() + # Every row of every span reaches the plugin's logger, whichever thread ran the activation. + assert all("span_attributes" in span for span in spans) + assert all(span["metrics"].get("end") is not None for span in spans) + (parent_wf,) = _spans_named(spans, "temporal.workflow.DelegatingParentWorkflow") + (tool,) = _spans_named(spans, "delegate_task") + (child_wf,) = _spans_named(spans, "temporal.workflow.ChildWorkflow") + activities = _spans_named(spans, "temporal.activity.simple_activity") + + assert tool["span_parents"] == [parent_wf["span_id"]] + assert child_wf["span_parents"] == [tool["span_id"]] + parent_activity = next( + a + for a in activities + if a["metadata"]["temporal.workflow_id"] != child_wf["metadata"]["temporal.workflow_id"] + ) + assert parent_activity["span_parents"] == [tool["span_id"]] + + @pytest.mark.asyncio + async def test_in_workflow_spans_are_not_duplicated_on_replay(self, temporal_env, memory_logger): + task_queue = f"test-queue-span-replay-{uuid.uuid4()}" + async with Worker( + temporal_env.client, + task_queue=task_queue, + workflows=[InWorkflowSpanReplayWorkflow], + activities=[simple_activity], + max_cached_workflows=0, + plugins=[BraintrustPlugin(logger=memory_logger)], + ): + handle = await temporal_env.client.start_workflow( + InWorkflowSpanReplayWorkflow.run, + TaskInput(value=10), + id=f"test-workflow-{uuid.uuid4()}", + task_queue=task_queue, + ) + await _wait_for_workflow_state(handle, "waiting") + await handle.signal(InWorkflowSpanReplayWorkflow.continue_workflow) + assert await handle.result() == 30 + + braintrust.flush() + spans = memory_logger.pop() + (workflow_span,) = _spans_named(spans, "temporal.workflow.InWorkflowSpanReplayWorkflow") + (agent_run,) = _spans_named(spans, "agent_run") + (first_tool,) = _spans_named(spans, "first_tool") + (second_tool,) = _spans_named(spans, "second_tool") + # Spans still open across a replay must end when the workflow really finishes. + assert workflow_span["metrics"]["end"] >= second_tool["metrics"]["start"] + assert agent_run["metrics"]["end"] >= second_tool["metrics"]["start"] + assert agent_run["span_parents"] == [workflow_span["span_id"]] + assert first_tool["span_parents"] == [agent_run["span_id"]] + assert second_tool["span_parents"] == [agent_run["span_id"]] + activities = _spans_named(spans, "temporal.activity.simple_activity") + assert sorted(a["span_parents"][0] for a in activities) == sorted( + [first_tool["span_id"], second_tool["span_id"]] + ) diff --git a/py/src/braintrust/integrations/temporal/test_temporal_pydantic_ai.py b/py/src/braintrust/integrations/temporal/test_temporal_pydantic_ai.py new file mode 100644 index 00000000..66064342 --- /dev/null +++ b/py/src/braintrust/integrations/temporal/test_temporal_pydantic_ai.py @@ -0,0 +1,160 @@ +"""Tests for BraintrustPlugin combined with Pydantic AI's Temporal durability integration.""" + +import uuid +import warnings + +import pytest + + +pytest.importorskip("temporalio") +pytest.importorskip("pydantic_ai.durable_exec.temporal") + +import braintrust +import pydantic_ai.durable_exec.temporal as pydantic_ai_temporal +import temporalio.workflow +from braintrust.integrations.pydantic_ai import PydanticAIIntegration +from braintrust.integrations.temporal import BraintrustPlugin +from braintrust.integrations.temporal.test_temporal import _spans_named, memory_logger, temporal_env # noqa: F401 +from pydantic_ai import Agent +from pydantic_ai.durable_exec.temporal import AgentPlugin, PydanticAIPlugin +from pydantic_ai.models.test import TestModel +from temporalio.client import Client +from temporalio.worker import Worker + + +_HAS_DURABILITY_CAPABILITY = hasattr(pydantic_ai_temporal, "TemporalDurability") + + +def _api_for_current_workflow() -> str: + # Workflow IDs start with the API under test; see the test below. + return temporalio.workflow.info().workflow_id.split(":", 1)[0] + + +async def _delegate(task: str) -> str: + # Runs in workflow code and starts the sub-agent as a child workflow, like a `delegate_task` tool. + return await temporalio.workflow.execute_child_workflow( + SubAgentWorkflow.run, + task, + id=f"{temporalio.workflow.info().workflow_id}:sub", + ) + + +# Deprecated `TemporalAgent` wrapper: the only API on older pydantic-ai releases. +with warnings.catch_warnings(): + warnings.simplefilter("ignore") + _wrapper_sub = pydantic_ai_temporal.TemporalAgent( + Agent(TestModel(custom_output_text="sub done"), name="wrapper_sub_agent") + ) + _wrapper_parent_agent = Agent(TestModel(call_tools=["delegate_task"]), name="wrapper_parent_agent") + _wrapper_parent_agent.tool_plain(name="delegate_task")(_delegate) + # delegate_task starts a child workflow, so it has to run in workflow code, not an activity. + _wrapper_parent = pydantic_ai_temporal.TemporalAgent( + _wrapper_parent_agent, tool_activity_config={"": {"delegate_task": False}} + ) + +if _HAS_DURABILITY_CAPABILITY: + _capability_sub = Agent( + TestModel(custom_output_text="sub done"), + name="capability_sub_agent", + capabilities=[pydantic_ai_temporal.TemporalDurability()], + ) + _capability_parent = Agent( + TestModel(call_tools=["delegate_task"]), + name="capability_parent_agent", + capabilities=[pydantic_ai_temporal.TemporalDurability()], + ) + # `metadata={"temporal": False}` replaces `tool_activity_config` on the capability path. + _capability_parent.tool_plain(name="delegate_task", metadata={"temporal": False})(_delegate) + +_AGENTS = { + "wrapper": (_wrapper_parent, _wrapper_sub), + **({"capability": (_capability_parent, _capability_sub)} if _HAS_DURABILITY_CAPABILITY else {}), +} + + +@temporalio.workflow.defn +class SubAgentWorkflow: + @temporalio.workflow.run + async def run(self, prompt: str) -> str: + _, sub_agent = _AGENTS[_api_for_current_workflow()] + return (await sub_agent.run(prompt)).output + + +@temporalio.workflow.defn +class ParentAgentWorkflow: + @temporalio.workflow.run + async def run(self, prompt: str) -> str: + parent_agent, _ = _AGENTS[_api_for_current_workflow()] + return (await parent_agent.run(prompt)).output + + +# With no workflow cache, every workflow task replays the run from the start, as after a worker +# restart or cache eviction. Replays must not log the in-workflow Pydantic AI spans again. +@pytest.mark.parametrize("max_cached_workflows", [1000, 0], ids=["cached", "replay"]) +@pytest.mark.parametrize( + "api", + [ + "wrapper", + pytest.param( + "capability", + marks=pytest.mark.skipif( + not _HAS_DURABILITY_CAPABILITY, reason="TemporalDurability needs a newer pydantic-ai" + ), + ), + ], +) +@pytest.mark.asyncio +async def test_subagent_child_workflow_nests_under_delegate_task_span( + temporal_env, # noqa: F811 + memory_logger, # noqa: F811 + api, + max_cached_workflows, +): + assert PydanticAIIntegration.setup() + parent_agent, sub_agent = _AGENTS[api] + client = await Client.connect( + temporal_env.client.service_client.config.target_host, + namespace=temporal_env.client.namespace, + plugins=[PydanticAIPlugin()], + ) + task_queue = f"test-queue-pydantic-ai-{api}-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[ParentAgentWorkflow, SubAgentWorkflow], + max_cached_workflows=max_cached_workflows, + plugins=[AgentPlugin(parent_agent), AgentPlugin(sub_agent), BraintrustPlugin(logger=memory_logger)], + ): + await client.execute_workflow( + ParentAgentWorkflow.run, + "delegate something", + id=f"{api}:{uuid.uuid4()}", + task_queue=task_queue, + ) + + braintrust.flush() + spans = memory_logger.pop() + (parent_workflow_span,) = _spans_named(spans, "temporal.workflow.ParentAgentWorkflow") + (parent_agent_span,) = _spans_named(spans, f"agent_run [{parent_agent.name}]") + (tool_span,) = _spans_named(spans, "delegate_task") + (child_workflow_span,) = _spans_named(spans, "temporal.workflow.SubAgentWorkflow") + (sub_agent_span,) = _spans_named(spans, f"agent_run [{sub_agent.name}]") + + assert parent_agent_span["span_parents"] == [parent_workflow_span["span_id"]] + assert tool_span["span_parents"] == [parent_agent_span["span_id"]] + assert child_workflow_span["span_parents"] == [tool_span["span_id"]] + assert sub_agent_span["span_parents"] == [child_workflow_span["span_id"]] + assert {span["root_span_id"] for span in spans} == {parent_workflow_span["root_span_id"]} + + # Each model request is traced once: one `chat` span paired with its model activity span. The + # wrapper API traces the request in workflow code (activity under `chat`); the capability + # API hands it to the activity first (`chat` under activity). + chat_spans = [span for span in spans if span["span_attributes"]["name"].startswith("chat ")] + model_activity_spans = [span for span in spans if span["span_attributes"]["name"].endswith("__model_request")] + assert len(chat_spans) == len(model_activity_spans) == 3 + pairs = { + frozenset((span["span_id"], span["span_parents"][0])) + for span in chat_spans + model_activity_spans + if span["span_parents"][0] in {other["span_id"] for other in chat_spans + model_activity_spans} + } + assert len(pairs) == 3 diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 41dbb027..f89c4852 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -35,6 +35,7 @@ Any, Generic, Literal, + Protocol, TypedDict, TypeVar, cast, @@ -4987,6 +4988,40 @@ def as_dataset(self, batch_size: int | None = None) -> Iterator[ExperimentDatase _EXEC_COUNTER = 0 +class _SpanHooks(Protocol): + """Internal hooks that let a runtime control span IDs and writes in its own context. + + Used by the Temporal integration: spans opened in workflow code need deterministic IDs and must + not be written while the workflow replays its history. Not a public API. + """ + + def next_ids(self) -> tuple[str, str, str]: + """Return ``(row_id, span_id, trace_id)`` for a new span that has no explicit IDs.""" + ... + + def should_write(self) -> bool: + """Return False to drop a span write.""" + ... + + +_span_hooks: contextvars.ContextVar["_SpanHooks | None"] = contextvars.ContextVar( + "braintrust_span_hooks", default=None +) +# Flipped the first time hooks are installed, so spans skip the context lookup until a runtime uses them. +_span_hooks_enabled = False + + +def _internal_set_span_hooks(hooks: "_SpanHooks | None") -> contextvars.Token: + """Install span hooks for the current context. Returns a token for `_internal_reset_span_hooks`.""" + global _span_hooks_enabled + _span_hooks_enabled = True + return _span_hooks.set(hooks) + + +def _internal_reset_span_hooks(token: contextvars.Token) -> None: + _span_hooks.reset(token) + + class SpanImpl(Span): """Primary implementation of the `Span` interface. See the `Span` interface for full details on each method. @@ -5091,14 +5126,26 @@ def __init__( # TODO: can be simplified after `event` is typed. id = event.pop("id", None) - if id is None or not isinstance(id, str): + if not isinstance(id, str): + id = None + span_id = internal.get("span_id", span_id) + root_span_id = internal.get("root_span_id", root_span_id) + + self._hooks = _span_hooks.get() if _span_hooks_enabled else None + if self._hooks is not None and (id is None or span_id is None): + hook_row_id, hook_span_id, hook_trace_id = self._hooks.next_ids() + id = id or hook_row_id + span_id = span_id or hook_span_id + root_span_id = root_span_id or hook_trace_id + + if id is None: id = str(uuid.uuid4()) self._id = id # Resolve all span IDs (span_id, root_span_id, span_parents) span_ids = _resolve_span_ids( - span_id=internal.get("span_id", span_id), - root_span_id=internal.get("root_span_id", root_span_id), + span_id=span_id, + root_span_id=root_span_id, parent_span_ids=parent_span_ids, lookup_span_parent=lookup_span_parent, id_generator=self.state.id_generator, @@ -5148,6 +5195,8 @@ def log(self, **event: Any) -> None: return self.log_internal(event=event, internal_data=None) def log_internal(self, event: dict[str, Any] | None = None, internal_data: dict[str, Any] | None = None) -> None: + if self._hooks is not None and not self._hooks.should_write(): + return serializable_partial_record, lazy_partial_record = split_logging_data(event, internal_data) # We both check for serializability and round-trip `partial_record` diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index 9bb585fa..514e4918 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -3829,3 +3829,55 @@ def test_proxy_conn_base_url(proxy_url, expected_base_url): state = BraintrustState() state.proxy_url = proxy_url assert state.proxy_conn().base_url == expected_base_url + + +class _RecordingSpanHooks: + def __init__(self): + self.counter = 0 + self.write = True + + def next_ids(self): + self.counter += 1 + return f"row-{self.counter}", f"span-{self.counter}", f"trace-{self.counter}" + + def should_write(self): + return self.write + + +def test_span_hooks_supply_default_ids_and_gate_writes(with_memory_logger): + hooks = _RecordingSpanHooks() + test_logger = init_test_logger(__name__) + token = braintrust.logger._internal_set_span_hooks(hooks) + try: + with test_logger.start_span(name="root") as root: + with root.start_span(name="child"): + pass + # Explicit IDs win over the hooks and do not consume a hook ID. + with test_logger.start_span(name="explicit", id="explicit-row", span_id="explicit-span"): + pass + hooks.write = False + with root.start_span(name="dropped"): + pass + hooks.write = True + finally: + braintrust.logger._internal_reset_span_hooks(token) + + with test_logger.start_span(name="after_reset") as after_reset: + pass + + rows = {row["span_attributes"]["name"]: row for row in with_memory_logger.pop()} + assert set(rows) == {"root", "child", "explicit", "after_reset"} + assert (rows["root"]["id"], rows["root"]["span_id"], rows["root"]["root_span_id"]) == ( + "row-1", + "span-1", + "trace-1", + ) + assert (rows["child"]["id"], rows["child"]["span_id"], rows["child"]["root_span_id"]) == ( + "row-2", + "span-2", + "trace-1", + ) + assert rows["child"]["span_parents"] == ["span-1"] + assert (rows["explicit"]["id"], rows["explicit"]["span_id"]) == ("explicit-row", "explicit-span") + assert hooks.counter == 3 + assert not after_reset.span_id.startswith("span-")