Skip to content

Add original execution IDs to Temporal workflow spans #958

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Ehsan-hash wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Add original execution IDs to Temporal workflow spans #958

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
921d9b4
test: cover Temporal original run metadata across resets
Ehsan-hash Oct 16, 2026
5088168
Add original execution IDs to Temporal workflow spans
Ehsan-hash Oct 16, 2026
461ef50
test: consolidate Temporal environment fixtures
Ehsan-hash Oct 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 10 additions & 10 deletions py/src/braintrust/integrations/temporal/plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand All @@ -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)
Expand All @@ -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()
Expand Down
89 changes: 86 additions & 3 deletions py/src/braintrust/integrations/temporal/test_temporal.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -273,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
Expand All @@ -283,13 +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:
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


Expand Down Expand Up @@ -367,6 +375,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
)
Expand All @@ -377,6 +392,74 @@ 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.parametrize("temporal_env", ["local"], indirect=True)
@pytest.mark.asyncio
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_env.client,
task_queue=task_queue,
workflows=[ReplayAfterSignalWorkflow],
activities=[simple_activity],
max_cached_workflows=0,
plugins=[BraintrustPlugin(logger=memory_logger)],
):
handle = await temporal_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_env.client.service_client.workflow_service.reset_workflow_execution(
temporalio.api.workflowservice.v1.ResetWorkflowExecutionRequest(
namespace=temporal_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_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",
[
Expand Down
Loading