From 777bb7f85c8a0bb58128e590d1f3b5203f4b09cb Mon Sep 17 00:00:00 2001 From: DABH Date: Thu, 10 Sep 2026 00:29:53 -0500 Subject: [PATCH 1/3] Wait for idle instead of polling traced queries in LangSmith tests TestComprehensiveTracing and the plugin end-to-end test waited for the workflow to reach its signal wait by polling is_waiting_for_signal every second through the traced client, then asserted that every QueryWorkflow run had exactly one HandleQuery child. Temporal does not guarantee that: a query issued while a workflow task is pending or started is buffered by the server and attached to every workflow task start until it completes, including the retry after the task times out or its completion is rejected, so the worker evaluates it twice (two HandleQuery runs). If the client's 30s RPC timeout fires first, the server drops the query before it is dispatched at all (no HandleQuery run). CI logs of the failing runs show exactly this: a QueryWorkflow call lasting 29.998s followed by "Error reporting WFT to server" for the test workflow. The same stalls made the unguarded raw query fail with "Timeout expired" and, once a task reached attempt 3, made the server reject updates with "Workflow Task in failed state". Readiness is now observed through the raw client's describe(), which is not traced and involves no workflow task, and the traced query is issued exactly once after the workflow is idle, so it is dispatched directly as a query task and cannot be re-attached to a retried workflow task. Updates are only sent once the workflow is idle again after the signal. The query assertion becomes an exact hierarchy check instead of a loop over a variable number of polls. --- tests/contrib/langsmith/test_integration.py | 116 ++++++++++---------- tests/contrib/langsmith/test_plugin.py | 36 +++--- 2 files changed, 77 insertions(+), 75 deletions(-) diff --git a/tests/contrib/langsmith/test_integration.py b/tests/contrib/langsmith/test_integration.py index f48d9d6ac..196d247a1 100644 --- a/tests/contrib/langsmith/test_integration.py +++ b/tests/contrib/langsmith/test_integration.py @@ -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, @@ -312,22 +311,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) # --------------------------------------------------------------------------- @@ -563,7 +568,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.") @@ -575,6 +580,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]: @@ -588,7 +595,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, @@ -602,14 +609,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( @@ -627,6 +634,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" ) @@ -704,18 +712,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 = [ @@ -777,7 +782,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.") @@ -789,6 +794,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]: @@ -801,7 +808,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, @@ -815,13 +822,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 @@ -839,6 +846,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" ) @@ -883,10 +891,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"] + ) # --------------------------------------------------------------------------- @@ -1252,13 +1260,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" - + await _wait_for_workflow_idle(handle) collector.clear() # Built-in queries — should NOT be traced diff --git a/tests/contrib/langsmith/test_plugin.py b/tests/contrib/langsmith/test_plugin.py index 6e3cb2e86..b1f1e7593 100644 --- a/tests/contrib/langsmith/test_plugin.py +++ b/tests/contrib/langsmith/test_plugin.py @@ -21,7 +21,8 @@ NexusService, TraceableActivityWorkflow, _make_client_and_collector, - _poll_query, + _query_pipeline, + _wait_for_workflow_idle, nested_traceable_activity, traceable_activity, ) @@ -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: @@ -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" ) @@ -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") From 97deb60be4ab1b07c0cafe4363bb2f4cfe9fe1c2 Mon Sep 17 00:00:00 2001 From: DABH Date: Thu, 10 Sep 2026 02:19:42 -0500 Subject: [PATCH 2/3] Expect the RunWorkflow root in the built-in query filtering test The worker's RunWorkflow run reaches the collector asynchronously and can arrive after describe() already reports the first task complete, so clearing the collector after the idle wait raced it and left a stray root on fast runners. Nothing needs clearing now that readiness no longer polls a traced query, so assert the full hierarchy instead. --- tests/contrib/langsmith/test_integration.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/contrib/langsmith/test_integration.py b/tests/contrib/langsmith/test_integration.py index 196d247a1..669b9f931 100644 --- a/tests/contrib/langsmith/test_integration.py +++ b/tests/contrib/langsmith/test_integration.py @@ -1261,7 +1261,6 @@ async def test_temporal_prefixed_query_not_traced( ) await _wait_for_workflow_idle(handle) - collector.clear() # Built-in queries — should NOT be traced await handle.query("__temporal_workflow_metadata") @@ -1272,10 +1271,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", ], From a7a41ba4003855fd2c178d5ffc5cb53dc12aa57d Mon Sep 17 00:00:00 2001 From: DABH Date: Thu, 10 Sep 2026 02:59:52 -0500 Subject: [PATCH 3/3] Stop activity retries from adding runs to the LangSmith trace assertions The traced activities finish instantly; their 10s start-to-close bound only satisfies Temporal. On a stalled runner an attempt timed out, the retry ran the activity again, and the exact hierarchy assertion saw one extra run. Give the bound a minute and disable retries so a failed attempt fails loudly. --- tests/contrib/langsmith/test_integration.py | 25 +++++++++++++-------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/tests/contrib/langsmith/test_integration.py b/tests/contrib/langsmith/test_integration.py index 669b9f931..01e047e2b 100644 --- a/tests/contrib/langsmith/test_integration.py +++ b/tests/contrib/langsmith/test_integration.py @@ -77,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), ) @@ -111,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 @@ -126,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), ) @@ -164,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( @@ -192,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" @@ -260,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), ) @@ -271,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), ) @@ -938,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}"