From 921d9b498291ef47261c84be82f0ef6d24ea90dd Mon Sep 17 00:00:00 2001 From: Ehsan-hash Date: Fri, 16 Oct 2026 10:06:05 -0700 Subject: [PATCH 1/3] test: cover Temporal original run metadata across resets --- .../integrations/temporal/test_temporal.py | 91 +++++++++++++++++++ 1 file changed, 91 insertions(+) diff --git a/py/src/braintrust/integrations/temporal/test_temporal.py b/py/src/braintrust/integrations/temporal/test_temporal.py index 2260fdc4..681c7f75 100644 --- a/py/src/braintrust/integrations/temporal/test_temporal.py +++ b/py/src/braintrust/integrations/temporal/test_temporal.py @@ -18,6 +18,8 @@ import braintrust import temporalio.activity import temporalio.api.common.v1 +import temporalio.api.enums.v1 +import temporalio.api.workflowservice.v1 import temporalio.converter import temporalio.testing import temporalio.worker @@ -293,6 +295,19 @@ async def temporal_env(): yield env +@pytest_asyncio.fixture +async def temporal_local_env(): + """Create a full local Temporal server for capabilities absent from the test server.""" + kwargs: dict[str, Any] = {} + cache_dir = os.environ.get("BRAINTRUST_TEMPORAL_TEST_SERVER_DIR") + if cache_dir: + os.makedirs(cache_dir, exist_ok=True) + kwargs["download_dest_dir"] = cache_dir + kwargs["dev_server_download_ttl"] = timedelta(days=365) + async with await temporalio.testing.WorkflowEnvironment.start_local(**kwargs) as env: + yield env + + @pytest.fixture def memory_logger(): """Set up memory logger to capture spans for testing.""" @@ -367,6 +382,13 @@ async def test_plugin_activity_after_replay_stays_under_workflow_span( assert len(activity_spans) == 2 workflow_span = workflow_spans[0] assert workflow_span["context"]["span_origin"]["instrumentation"]["name"] == "temporal-auto" + if "original_execution_run_id" in temporalio.workflow.Info.__annotations__: + assert ( + workflow_span["metadata"]["temporal.original_execution_run_id"] + == workflow_span["metadata"]["temporal.run_id"] + ) + else: + assert "temporal.original_execution_run_id" not in workflow_span["metadata"] assert all( span["context"]["span_origin"]["instrumentation"]["name"] == "temporal-auto" for span in activity_spans ) @@ -377,6 +399,75 @@ async def test_plugin_activity_after_replay_stays_under_workflow_span( assert len(client_spans) == 1 assert workflow_span["root_span_id"] == client_spans[0]["root_span_id"] + @pytest.mark.skipif( + "original_execution_run_id" not in temporalio.workflow.Info.__annotations__, + reason="Temporal version does not expose original_execution_run_id", + ) + @pytest.mark.asyncio + async def test_plugin_workflow_metadata_preserves_original_run_id_after_reset( + self, temporal_local_env, memory_logger + ): + task_queue = f"test-queue-reset-{uuid.uuid4()}" + workflow_id = f"test-workflow-reset-{uuid.uuid4()}" + async with Worker( + temporal_local_env.client, + task_queue=task_queue, + workflows=[ReplayAfterSignalWorkflow], + activities=[simple_activity], + max_cached_workflows=0, + plugins=[BraintrustPlugin(logger=memory_logger)], + ): + handle = await temporal_local_env.client.start_workflow( + ReplayAfterSignalWorkflow.run, + TaskInput(value=10), + id=workflow_id, + task_queue=task_queue, + ) + await _wait_for_workflow_state(handle, "waiting") + original_run_id = (await handle.describe()).run_id + history = await handle.fetch_history() + workflow_task_finish_event_id = next( + event.event_id + for event in history.events + if event.event_type == temporalio.api.enums.v1.EventType.EVENT_TYPE_WORKFLOW_TASK_COMPLETED + ) + await temporal_local_env.client.service_client.workflow_service.reset_workflow_execution( + temporalio.api.workflowservice.v1.ResetWorkflowExecutionRequest( + namespace=temporal_local_env.client.namespace, + workflow_execution=temporalio.api.common.v1.WorkflowExecution( + workflow_id=workflow_id, + run_id=original_run_id, + ), + reason="metadata test", + workflow_task_finish_event_id=workflow_task_finish_event_id, + request_id=str(uuid.uuid4()), + identity="braintrust-reset-test", + ) + ) + reset_handle = temporal_local_env.client.get_workflow_handle(workflow_id) + await _wait_for_workflow_state(reset_handle, "waiting") + reset_run_id = (await reset_handle.describe()).run_id + await reset_handle.signal(ReplayAfterSignalWorkflow.continue_workflow) + assert await reset_handle.result() == 30 + + braintrust.flush() + spans = memory_logger.pop() + workflow_spans = _spans_named(spans, "temporal.workflow.ReplayAfterSignalWorkflow") + activity_spans = _spans_named(spans, "temporal.activity.simple_activity") + assert original_run_id != reset_run_id + assert {span["metadata"]["temporal.run_id"] for span in workflow_spans} == { + original_run_id, + reset_run_id, + } + assert all( + span["metadata"]["temporal.original_execution_run_id"] == original_run_id for span in workflow_spans + ) + assert len({span["span_id"] for span in workflow_spans}) == 2 + assert all( + any(workflow_span["span_id"] in span.get("span_parents", []) for workflow_span in workflow_spans) + for span in activity_spans + ) + @pytest.mark.parametrize( "workflow,input_value,expected_result,expected_span_counts,expected_activity_errors", [ From 5088168339b0e79265722528f2cbff009431f3ff Mon Sep 17 00:00:00 2001 From: Ehsan-hash Date: Fri, 16 Oct 2026 10:06:35 -0700 Subject: [PATCH 2/3] Add original execution IDs to Temporal workflow spans --- .../integrations/temporal/plugin.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/py/src/braintrust/integrations/temporal/plugin.py b/py/src/braintrust/integrations/temporal/plugin.py index 60881c9f..07461a54 100644 --- a/py/src/braintrust/integrations/temporal/plugin.py +++ b/py/src/braintrust/integrations/temporal/plugin.py @@ -280,6 +280,14 @@ async def execute_workflow(self, input: temporalio.worker.ExecuteWorkflowInput) if logger: info = temporalio.workflow.info() parent = parent_span_context or logger.export() + metadata = { + "temporal.workflow_type": info.workflow_type, + "temporal.workflow_id": info.workflow_id, + "temporal.run_id": info.run_id, + } + original_execution_run_id = getattr(info, "original_execution_run_id", None) + if original_execution_run_id is not None: + metadata["temporal.original_execution_run_id"] = original_execution_run_id if temporalio.workflow.patched(_SPAN_CONTEXT_PATCH_ID): ids = _workflow_span_ids(info, parent) if temporalio.workflow.unsafe.is_replaying(): @@ -293,11 +301,7 @@ async def execute_workflow(self, input: temporalio.worker.ExecuteWorkflowInput) 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, - }, + metadata=metadata, ) span.set_current() self._parent_span_context = _workflow_span_context(parent, ids) @@ -307,11 +311,7 @@ async def execute_workflow(self, input: temporalio.worker.ExecuteWorkflowInput) name=f"temporal.workflow.{info.workflow_type}", type="task", parent=parent_span_context or None, - metadata={ - "temporal.workflow_type": info.workflow_type, - "temporal.workflow_id": info.workflow_id, - "temporal.run_id": info.run_id, - }, + metadata=metadata, ) span.set_current() self._parent_span_context = span.export() From 461ef50e1910d40abb7fc9848efd7aec5f93d18e Mon Sep 17 00:00:00 2001 From: Ehsan-hash Date: Fri, 16 Oct 2026 11:41:23 -0700 Subject: [PATCH 3/3] test: consolidate Temporal environment fixtures --- .../integrations/temporal/test_temporal.py | 40 ++++++++----------- 1 file changed, 16 insertions(+), 24 deletions(-) diff --git a/py/src/braintrust/integrations/temporal/test_temporal.py b/py/src/braintrust/integrations/temporal/test_temporal.py index 681c7f75..f03d1986 100644 --- a/py/src/braintrust/integrations/temporal/test_temporal.py +++ b/py/src/braintrust/integrations/temporal/test_temporal.py @@ -275,7 +275,7 @@ def test_contrib_temporal_compat_import_deprecated(self): @pytest_asyncio.fixture(scope="function") -async def temporal_env(): +async def temporal_env(request): """Create a Temporal test environment. If ``BRAINTRUST_TEMPORAL_TEST_SERVER_DIR`` is set, point the SDK's binary @@ -285,26 +285,19 @@ async def temporal_env(): limits CI runners). When the var is unset (the local default), the SDK falls back to its built-in temp-dir download behavior. """ + local_server = getattr(request, "param", None) == "local" kwargs: dict[str, Any] = {} cache_dir = os.environ.get("BRAINTRUST_TEMPORAL_TEST_SERVER_DIR") if cache_dir: os.makedirs(cache_dir, exist_ok=True) kwargs["download_dest_dir"] = cache_dir - kwargs["test_server_download_ttl"] = timedelta(days=365) - async with await temporalio.testing.WorkflowEnvironment.start_time_skipping(**kwargs) as env: - yield env - - -@pytest_asyncio.fixture -async def temporal_local_env(): - """Create a full local Temporal server for capabilities absent from the test server.""" - kwargs: dict[str, Any] = {} - cache_dir = os.environ.get("BRAINTRUST_TEMPORAL_TEST_SERVER_DIR") - if cache_dir: - os.makedirs(cache_dir, exist_ok=True) - kwargs["download_dest_dir"] = cache_dir - kwargs["dev_server_download_ttl"] = timedelta(days=365) - async with await temporalio.testing.WorkflowEnvironment.start_local(**kwargs) as env: + kwargs["dev_server_download_ttl" if local_server else "test_server_download_ttl"] = timedelta(days=365) + start_environment = ( + temporalio.testing.WorkflowEnvironment.start_local + if local_server + else temporalio.testing.WorkflowEnvironment.start_time_skipping + ) + async with await start_environment(**kwargs) as env: yield env @@ -403,21 +396,20 @@ async def test_plugin_activity_after_replay_stays_under_workflow_span( "original_execution_run_id" not in temporalio.workflow.Info.__annotations__, reason="Temporal version does not expose original_execution_run_id", ) + @pytest.mark.parametrize("temporal_env", ["local"], indirect=True) @pytest.mark.asyncio - async def test_plugin_workflow_metadata_preserves_original_run_id_after_reset( - self, temporal_local_env, memory_logger - ): + async def test_plugin_workflow_metadata_preserves_original_run_id_after_reset(self, temporal_env, memory_logger): task_queue = f"test-queue-reset-{uuid.uuid4()}" workflow_id = f"test-workflow-reset-{uuid.uuid4()}" async with Worker( - temporal_local_env.client, + temporal_env.client, task_queue=task_queue, workflows=[ReplayAfterSignalWorkflow], activities=[simple_activity], max_cached_workflows=0, plugins=[BraintrustPlugin(logger=memory_logger)], ): - handle = await temporal_local_env.client.start_workflow( + handle = await temporal_env.client.start_workflow( ReplayAfterSignalWorkflow.run, TaskInput(value=10), id=workflow_id, @@ -431,9 +423,9 @@ async def test_plugin_workflow_metadata_preserves_original_run_id_after_reset( for event in history.events if event.event_type == temporalio.api.enums.v1.EventType.EVENT_TYPE_WORKFLOW_TASK_COMPLETED ) - await temporal_local_env.client.service_client.workflow_service.reset_workflow_execution( + await temporal_env.client.service_client.workflow_service.reset_workflow_execution( temporalio.api.workflowservice.v1.ResetWorkflowExecutionRequest( - namespace=temporal_local_env.client.namespace, + namespace=temporal_env.client.namespace, workflow_execution=temporalio.api.common.v1.WorkflowExecution( workflow_id=workflow_id, run_id=original_run_id, @@ -444,7 +436,7 @@ async def test_plugin_workflow_metadata_preserves_original_run_id_after_reset( identity="braintrust-reset-test", ) ) - reset_handle = temporal_local_env.client.get_workflow_handle(workflow_id) + reset_handle = temporal_env.client.get_workflow_handle(workflow_id) await _wait_for_workflow_state(reset_handle, "waiting") reset_run_id = (await reset_handle.describe()).run_id await reset_handle.signal(ReplayAfterSignalWorkflow.continue_workflow)