Skip to content

Wait for idle instead of polling traced queries in LangSmith tests #1830

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
DABH wants to merge 5 commits into
base: main
Choose a base branch
from
+95 −86
Open

Wait for idle instead of polling traced queries in LangSmith tests #1830

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
777bb7f
Wait for idle instead of polling traced queries in LangSmith tests
DABH Sep 10, 2026
97deb60
Expect the RunWorkflow root in the built-in query filtering test
DABH Sep 10, 2026
a7a41ba
Stop activity retries from adding runs to the LangSmith trace assertions
DABH Sep 10, 2026
6d1d137
Merge remote-tracking branch 'origin/main' into flake/langsmith-traci…
DABH Sep 10, 2026
0b34e2e
Merge branch 'main' into flake/langsmith-tracing-query-polls
DABH Sep 11, 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
145 changes: 77 additions & 68 deletions tests/contrib/langsmith/test_integration.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,13 +16,12 @@
from temporalio import activity, common, nexus, workflow
from temporalio.client import (
Client,
WorkflowExecutionStatus,
WorkflowFailureError,
WorkflowHandle,
WorkflowQueryFailedError,
)
from temporalio.contrib.langsmith import LangSmithPlugin
from temporalio.exceptions import ApplicationError
from temporalio.service import RPCError
from temporalio.testing import WorkflowEnvironment
from tests.contrib.langsmith.conftest import (
InMemoryRunCollector,
Expand Down Expand Up @@ -78,7 +77,8 @@ class TraceableActivityWorkflow:
async def run(self, _input: str = "") -> str:
return await workflow.execute_activity(
traceable_activity,
start_to_close_timeout=timedelta(seconds=10),
start_to_close_timeout=timedelta(minutes=1),
retry_policy=common.RetryPolicy(maximum_attempts=1),
)


Expand Down Expand Up @@ -112,7 +112,8 @@ class SimpleWorkflow:
async def run(self) -> str:
result = await workflow.execute_activity(
simple_activity,
start_to_close_timeout=timedelta(seconds=10),
start_to_close_timeout=timedelta(minutes=1),
retry_policy=common.RetryPolicy(maximum_attempts=1),
)
return result

Expand All @@ -127,7 +128,8 @@ async def _step_with_activity() -> str:
"""A @traceable step that wraps an activity call."""
return await workflow.execute_activity(
nested_traceable_activity,
start_to_close_timeout=timedelta(seconds=10),
start_to_close_timeout=timedelta(minutes=1),
retry_policy=common.RetryPolicy(maximum_attempts=1),
)


Expand Down Expand Up @@ -165,12 +167,14 @@ def __init__(self) -> None:
async def run(self) -> str:
await workflow.execute_activity(
nested_traceable_activity,
start_to_close_timeout=timedelta(seconds=10),
start_to_close_timeout=timedelta(minutes=1),
retry_policy=common.RetryPolicy(maximum_attempts=1),
)
await _step_with_activity()
await workflow.execute_local_activity(
nested_traceable_activity,
start_to_close_timeout=timedelta(seconds=10),
start_to_close_timeout=timedelta(minutes=1),
retry_policy=common.RetryPolicy(maximum_attempts=1),
)
await _outer_chain("from-workflow")
await workflow.execute_child_workflow(
Expand All @@ -193,7 +197,8 @@ async def run(self) -> str:
await workflow.wait_condition(lambda: self._signal_received)
await workflow.execute_activity(
nested_traceable_activity,
start_to_close_timeout=timedelta(seconds=10),
start_to_close_timeout=timedelta(minutes=1),
retry_policy=common.RetryPolicy(maximum_attempts=1),
)
await workflow.wait_condition(lambda: self._complete)
return "comprehensive-done"
Expand Down Expand Up @@ -261,7 +266,7 @@ class ActivityFailureWorkflow:
async def run(self) -> str:
return await workflow.execute_activity(
failing_activity,
start_to_close_timeout=timedelta(seconds=10),
start_to_close_timeout=timedelta(minutes=1),
retry_policy=common.RetryPolicy(maximum_attempts=1),
)

Expand All @@ -272,7 +277,7 @@ class BenignErrorWorkflow:
async def run(self) -> str:
return await workflow.execute_activity(
benign_failing_activity,
start_to_close_timeout=timedelta(seconds=10),
start_to_close_timeout=timedelta(minutes=1),
retry_policy=common.RetryPolicy(maximum_attempts=1),
)

Expand Down Expand Up @@ -312,22 +317,28 @@ def _make_temporal_client(
return Client(**config)


@traceable(name="poll_query")
async def _poll_query(
handle: WorkflowHandle[Any, Any],
query: Callable[..., Any],
*,
expected: Any = True,
) -> bool:
"""Poll a workflow query until it returns the expected value."""
async def _wait_for_workflow_idle(handle: WorkflowHandle[Any, Any]) -> None:
"""Wait until the running workflow has no pending task, activity, child or Nexus op."""
while True:
try:
result = await handle.query(query)
if result == expected:
return True
except (WorkflowQueryFailedError, RPCError):
pass # Query not yet available (workflow hasn't started)
await asyncio.sleep(1)
desc = await handle.describe()
assert desc.status == WorkflowExecutionStatus.RUNNING, desc.status
raw = desc.raw_description
if not (
raw.HasField("pending_workflow_task")
or raw.pending_activities
or raw.pending_children
or raw.pending_nexus_operations
):
return
await asyncio.sleep(0.2)


@traceable(name="query_pipeline")
async def _query_pipeline(
handle: WorkflowHandle[Any, Any], query: Callable[..., Any]
) -> Any:
"""Run a single workflow query under a @traceable root."""
return await handle.query(query)


# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -563,7 +574,7 @@ async def test_comprehensive_with_temporal_runs(
"""Full trace hierarchy with worker restart mid-workflow.

user_pipeline only wraps start_workflow (completing before the worker
starts), so poll/signal/query traces are naturally separate root traces.
starts), so query/signal/update traces are naturally separate root traces.
"""
if env.supports_time_skipping:
pytest.skip("Time-skipping server doesn't persist headers.")
Expand All @@ -575,6 +586,8 @@ async def test_comprehensive_with_temporal_runs(
temporal_client_1 = _make_temporal_client(
client, mock_ls, add_temporal_runs=True
)
# Raw-client handle (no LangSmith interceptor) for untraced readiness checks
raw_handle = client.get_workflow_handle(workflow_id)

@traceable(name="user_pipeline")
async def user_pipeline() -> WorkflowHandle[Any, Any]:
Expand All @@ -588,7 +601,7 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]:
# Start workflow — no worker yet, just a server RPC
handle = await user_pipeline()

# Phase 1: worker picks up workflow, poll until signal wait
# Phase 1: worker runs the workflow up to the signal wait
async with new_worker(
temporal_client_1,
ComprehensiveWorkflow,
Expand All @@ -602,14 +615,14 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]:
make_nexus_endpoint_name(worker.task_queue),
worker.task_queue,
)
assert await _poll_query(
handle,
ComprehensiveWorkflow.is_waiting_for_signal,
expected=True,
await _wait_for_workflow_idle(raw_handle)
assert await _query_pipeline(
handle, ComprehensiveWorkflow.is_waiting_for_signal
), "Workflow never reached signal wait point"
# Raw-client query (no LangSmith interceptor) — root-level trace
raw_handle = client.get_workflow_handle(workflow_id)
await raw_handle.query(ComprehensiveWorkflow.is_waiting_for_signal)
# Raw-client query — root-level trace
assert await raw_handle.query(
ComprehensiveWorkflow.is_waiting_for_signal
)

# Phase 2: fresh worker, signal to resume, complete
temporal_client_2 = _make_temporal_client(
Expand All @@ -627,6 +640,7 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]:
handle_2 = temporal_client_2.get_workflow_handle(workflow_id)
await handle_2.query(ComprehensiveWorkflow.my_query)
await handle_2.signal(ComprehensiveWorkflow.my_signal, "hello")
await _wait_for_workflow_idle(raw_handle)
await handle_2.execute_update(
ComprehensiveWorkflow.my_unvalidated_update, "test"
)
Expand Down Expand Up @@ -704,18 +718,15 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]:
]
assert_trace_hierarchy(workflow_trace_trees, expected_workflow)

# poll_query trace (separate root, variable number of iterations)
poll_trace_trees = find_trace_trees(trace_trees, "poll_query")
assert len(poll_trace_trees) == 1
poll = poll_trace_trees[0]
assert poll.name == "poll_query"
poll_children = poll.children
for poll_child in poll_children:
assert poll_child.name == "QueryWorkflow:is_waiting_for_signal"
assert [child.name for child in poll_child.children] == [
"HandleQuery:is_waiting_for_signal"
]
assert not poll_child.children[0].children
# query_pipeline trace: the worker-side handler nests under the client query
assert_trace_hierarchy(
find_trace_trees(trace_trees, "query_pipeline"),
[
"query_pipeline",
" QueryWorkflow:is_waiting_for_signal",
" HandleQuery:is_waiting_for_signal",
],
)

# Raw-client query — no parent context, appears as root
raw_query_trace_trees = [
Expand Down Expand Up @@ -777,7 +788,7 @@ async def test_comprehensive_without_temporal_runs(
"""Same workflow with add_temporal_runs=False and worker restart.

Only @traceable runs appear. Context propagation via headers still works.
user_pipeline only wraps start_workflow, so poll traces are separate roots.
user_pipeline only wraps start_workflow, so the query trace is a separate root.
"""
if env.supports_time_skipping:
pytest.skip("Time-skipping server doesn't persist headers.")
Expand All @@ -789,6 +800,8 @@ async def test_comprehensive_without_temporal_runs(
temporal_client_1 = _make_temporal_client(
client, mock_ls, add_temporal_runs=False
)
# Raw-client handle (no LangSmith interceptor) for untraced readiness checks
raw_handle = client.get_workflow_handle(workflow_id)

@traceable(name="user_pipeline")
async def user_pipeline() -> WorkflowHandle[Any, Any]:
Expand All @@ -801,7 +814,7 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]:
with tracing_context(client=mock_ls, enabled=True):
handle = await user_pipeline()

# Phase 1: worker picks up workflow, poll until signal wait
# Phase 1: worker runs the workflow up to the signal wait
async with new_worker(
temporal_client_1,
ComprehensiveWorkflow,
Expand All @@ -815,13 +828,13 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]:
make_nexus_endpoint_name(worker.task_queue),
worker.task_queue,
)
await _wait_for_workflow_idle(raw_handle)
# Raw-client query — no interceptor, produces nothing
raw_handle = client.get_workflow_handle(workflow_id)
await raw_handle.query(ComprehensiveWorkflow.is_waiting_for_signal)
assert await _poll_query(
handle,
ComprehensiveWorkflow.is_waiting_for_signal,
expected=True,
assert await raw_handle.query(
ComprehensiveWorkflow.is_waiting_for_signal
)
assert await _query_pipeline(
handle, ComprehensiveWorkflow.is_waiting_for_signal
), "Workflow never reached signal wait point"

# Phase 2: fresh worker, signal to resume, complete
Expand All @@ -839,6 +852,7 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]:
):
handle_2 = temporal_client_2.get_workflow_handle(workflow_id)
await handle_2.signal(ComprehensiveWorkflow.my_signal, "hello")
await _wait_for_workflow_idle(raw_handle)
await handle_2.execute_update(
ComprehensiveWorkflow.my_unvalidated_update, "test"
)
Expand Down Expand Up @@ -883,10 +897,10 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]:
]
assert_trace_hierarchy(workflow_trace_trees, expected_workflow)

# Poll query — separate root, just the @traceable wrapper, no Temporal children
poll_trace_trees = find_trace_trees(trace_trees, "poll_query")
assert len(poll_trace_trees) == 1
assert_trace_hierarchy(poll_trace_trees, ["poll_query"])
# Query — separate root, just the @traceable wrapper, no Temporal children
assert_trace_hierarchy(
find_trace_trees(trace_trees, "query_pipeline"), ["query_pipeline"]
)


# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -930,7 +944,8 @@ async def run(self) -> str:
# Activity with nested @traceable
await workflow.execute_activity(
nested_traceable_activity,
start_to_close_timeout=timedelta(seconds=10),
start_to_close_timeout=timedelta(minutes=1),
retry_policy=common.RetryPolicy(maximum_attempts=1),
)
return f"{r1}|{r2}|{r3}"

Expand Down Expand Up @@ -1252,14 +1267,7 @@ async def test_temporal_prefixed_query_not_traced(
task_queue=worker.task_queue,
)

# Wait for workflow to start by polling the user query
assert await _poll_query(
handle,
QueryFilteringWorkflow.my_query,
expected="query-result",
), "Workflow never started"

collector.clear()
await _wait_for_workflow_idle(handle)

# Built-in queries — should NOT be traced
await handle.query("__temporal_workflow_metadata")
Expand All @@ -1270,10 +1278,11 @@ async def test_temporal_prefixed_query_not_traced(
await handle.signal(QueryFilteringWorkflow.complete)
assert await handle.result() == "done"

# Built-in queries should be absent; only user query and signal remain.
# The built-in query leaves no run; everything else the worker did is here.
assert_trace_hierarchy(
build_trace_trees(collector),
[
"RunWorkflow:QueryFilteringWorkflow",
"HandleQuery:my_query",
"HandleSignal:complete",
],
Expand Down
36 changes: 18 additions & 18 deletions tests/contrib/langsmith/test_plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,8 @@
NexusService,
TraceableActivityWorkflow,
_make_client_and_collector,
_poll_query,
_query_pipeline,
_wait_for_workflow_idle,
nested_traceable_activity,
traceable_activity,
)
Expand Down Expand Up @@ -63,7 +64,7 @@ async def test_comprehensive_plugin_trace_hierarchy(
) -> None:
"""Plugin wired to a real Temporal worker produces the full trace hierarchy.

user_pipeline only wraps start_workflow, so poll/query/signal/update
user_pipeline only wraps start_workflow, so query/signal/update
traces are naturally separate root traces.
"""
if env.supports_time_skipping:
Expand Down Expand Up @@ -100,13 +101,15 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]:
make_nexus_endpoint_name(worker.task_queue),
worker.task_queue,
)
assert await _poll_query(
handle,
ComprehensiveWorkflow.is_waiting_for_signal,
expected=True,
# Raw-client handle (no LangSmith interceptor) for untraced readiness checks
raw_handle = client.get_workflow_handle(workflow_id)
await _wait_for_workflow_idle(raw_handle)
assert await _query_pipeline(
handle, ComprehensiveWorkflow.is_waiting_for_signal
), "Workflow never reached signal wait point"
await handle.query(ComprehensiveWorkflow.my_query)
await handle.signal(ComprehensiveWorkflow.my_signal, "hello")
await _wait_for_workflow_idle(raw_handle)
await handle.execute_update(
ComprehensiveWorkflow.my_unvalidated_update, "test"
)
Expand Down Expand Up @@ -184,18 +187,15 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]:
]
assert_trace_hierarchy(workflow_trace_trees, expected_workflow)

# poll_query trace (separate root, variable number of iterations)
poll_trace_trees = find_trace_trees(trace_trees, "poll_query")
assert len(poll_trace_trees) == 1
poll = poll_trace_trees[0]
assert poll.name == "poll_query"
poll_children = poll.children
for poll_child in poll_children:
assert poll_child.name == "QueryWorkflow:is_waiting_for_signal"
assert [child.name for child in poll_child.children] == [
"HandleQuery:is_waiting_for_signal"
]
assert not poll_child.children[0].children
# query_pipeline trace: the worker-side handler nests under the client query
assert_trace_hierarchy(
find_trace_trees(trace_trees, "query_pipeline"),
[
"query_pipeline",
" QueryWorkflow:is_waiting_for_signal",
" HandleQuery:is_waiting_for_signal",
],
)

# Each remaining operation is its own root trace
query_trace_trees = find_trace_trees(trace_trees, "QueryWorkflow:my_query")
Expand Down
Loading