Wait for idle instead of polling traced queries in LangSmith tests

tests/contrib/langsmith/test_integration.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import asyncio
 import uuid
 from collections.abc import Callable
 from datetime import timedelta
@@ -18,19 +17,17 @@
     Client,
     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,
     build_trace_trees,
     find_trace_trees,
     make_mock_ls_client,
 )
-from tests.helpers import new_worker
+from tests.helpers import new_worker, wait_for_workflow_idle
 from tests.helpers.nexus import make_nexus_endpoint_name
 from tests.helpers.trace import assert_trace_hierarchy
 
@@ -78,7 +75,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),
         )
 
 
@@ -112,7 +110,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
 
@@ -127,7 +126,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),
     )
 
 
@@ -165,12 +165,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(
@@ -193,7 +195,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"
@@ -261,7 +264,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),
         )
 
@@ -272,7 +275,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),
         )
 
@@ -312,22 +315,17 @@ 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."""
-    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)
+# CI has stalled workflow tasks for tens of seconds; give the untraced idle probe
+# most of the 60s per-test budget instead of the helper's 10s default.
+_IDLE_TIMEOUT = timedelta(seconds=30)
+
+
+@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 +561,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 +573,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 +588,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 +602,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, timeout=_IDLE_TIMEOUT)
+                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 +627,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, timeout=_IDLE_TIMEOUT)
                 await handle_2.execute_update(
                     ComprehensiveWorkflow.my_unvalidated_update, "test"
                 )
@@ -704,18 +705,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 +775,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 +787,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 +801,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 +815,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, timeout=_IDLE_TIMEOUT)
                 # 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 +839,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, timeout=_IDLE_TIMEOUT)
                 await handle_2.execute_update(
                     ComprehensiveWorkflow.my_unvalidated_update, "test"
                 )
@@ -883,10 +884,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"]
+        )
 
 
 # ---------------------------------------------------------------------------
@@ -930,7 +931,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}"
 
@@ -1252,14 +1254,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, timeout=_IDLE_TIMEOUT)
 
             # Built-in queries — should NOT be traced
             await handle.query("__temporal_workflow_metadata")
@@ -1270,10 +1265,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",
             ],