Replace traced readiness query polls with untraced waits

tests/contrib/langgraph/test_summary_fn.py

@@ -19,7 +19,7 @@
 from temporalio.contrib.langgraph import LangGraphPlugin, graph
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import Replayer, Worker
-from tests.helpers import assert_eq_eventually
+from tests.helpers import wait_for_workflow_idle
 
 SummaryFn = Callable[[tuple[Any, ...], dict[str, Any]], "str | None"]
 
@@ -235,27 +235,17 @@ class WorkflowNodeSummaryWorkflow:
     def __init__(self) -> None:
         self.app = graph("wf-node-graph").compile()
         self._done = False
-        self._invoked = False
 
     @workflow.run
     async def run(self, input: str) -> Any:
         result = await self.app.ainvoke({"value": input})
-        self._invoked = True
         await workflow.wait_condition(lambda: self._done)
         return result
 
     @workflow.signal
     def finish(self) -> None:
         self._done = True
 
-    @workflow.query
-    def ran(self) -> bool:
-        return workflow.get_current_details() != ""
-
-    @workflow.query
-    def invoked(self) -> bool:
-        return self._invoked
-
 
 async def test_workflow_node_sets_current_details(
     client: Client, env: WorkflowEnvironment
@@ -285,9 +275,7 @@ async def test_workflow_node_sets_current_details(
             id=f"wf-node-{uuid.uuid4()}",
             task_queue=task_queue,
         )
-        await assert_eq_eventually(
-            True, lambda: handle.query(WorkflowNodeSummaryWorkflow.ran)
-        )
+        await wait_for_workflow_idle(handle)
         md: temporalio.api.sdk.v1.WorkflowMetadata = await handle.query(
             "__temporal_workflow_metadata",
             result_type=temporalio.api.sdk.v1.WorkflowMetadata,
@@ -328,9 +316,7 @@ async def test_workflow_node_clears_current_details_on_empty(
             id=f"wf-node-clear-{uuid.uuid4()}",
             task_queue=task_queue,
         )
-        await assert_eq_eventually(
-            True, lambda: handle.query(WorkflowNodeSummaryWorkflow.invoked)
-        )
+        await wait_for_workflow_idle(handle)
         md: temporalio.api.sdk.v1.WorkflowMetadata = await handle.query(
             "__temporal_workflow_metadata",
             result_type=temporalio.api.sdk.v1.WorkflowMetadata,

tests/contrib/openai_agents/test_openai_tracing.py

@@ -10,6 +10,7 @@
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 
 from temporalio import activity, workflow
+from temporalio.api.enums.v1 import EventType
 from temporalio.client import Client
 from temporalio.contrib.openai_agents import _temporal_openai_agents
 from temporalio.contrib.openai_agents.testing import (
@@ -24,7 +25,7 @@
     ResearchWorkflow,
     research_mock_model,
 )
-from tests.helpers import assert_eq_eventually, new_worker
+from tests.helpers import assert_event_subsequence, new_worker
 
 
 class MemoryTracingProcessor(TracingProcessor):
@@ -242,11 +243,26 @@ async def simple_no_context_activity() -> str:
     return "success"
 
 
+async def wait_for_activity_processed(client: Client, workflow_id: str) -> None:
+    """Wait, via an untraced history poll, until the workflow task that handled the activity result completed.
+
+    Assumes the workflow's first completed activity is the one right before its park point, which
+    holds for every workflow in this module.
+    """
+    await assert_event_subsequence(
+        client.get_workflow_handle(workflow_id),
+        [
+            EventType.EVENT_TYPE_ACTIVITY_TASK_COMPLETED,
+            EventType.EVENT_TYPE_WORKFLOW_TASK_COMPLETED,
+        ],
+        timeout=timedelta(seconds=10),
+    )
+
+
 @workflow.defn
 class TraceWorkflow:
     def __init__(self) -> None:
         self._proceed = False
-        self._ready = False
 
     @workflow.run
     async def run(self):
@@ -256,14 +272,9 @@ async def run(self):
                 simple_no_context_activity,
                 start_to_close_timeout=timedelta(seconds=10),
             )
-            self._ready = True
             await workflow.wait_condition(lambda: self._proceed)
         return "done"
 
-    @workflow.query
-    def ready(self) -> bool:
-        return self._ready
-
     @workflow.signal
     def proceed(self) -> None:
         self._proceed = True
@@ -273,7 +284,6 @@ def proceed(self) -> None:
 class SelfTracingWorkflow:
     def __init__(self) -> None:
         self._proceed = False
-        self._ready = False
 
     @workflow.run
     async def run(self):
@@ -284,14 +294,9 @@ async def run(self):
                     simple_no_context_activity,
                     start_to_close_timeout=timedelta(seconds=10),
                 )
-                self._ready = True
                 await workflow.wait_condition(lambda: self._proceed)
         return "done"
 
-    @workflow.query
-    def ready(self) -> bool:
-        return self._ready
-
     @workflow.signal
     def proceed(self) -> None:
         self._proceed = True
@@ -366,11 +371,7 @@ async def test_external_trace_to_workflow_spans(
             max_cached_workflows=0,
             task_queue=task_queue,
         ):
-            # Wait for workflow to be ready
-            async def ready() -> bool:
-                return await workflow_handle.query(TraceWorkflow.ready)
-
-            await assert_eq_eventually(True, ready)
+            await wait_for_activity_processed(client, workflow_handle.id)
 
     # Second worker: Complete the workflow with fresh objects (new instrumentation)
     async with AgentEnvironment(
@@ -458,11 +459,7 @@ async def test_external_trace_and_span_to_workflow_spans(
             max_cached_workflows=0,
             task_queue=task_queue,
         ):
-            # Wait for workflow to be ready
-            async def ready() -> bool:
-                return await workflow_handle.query(TraceWorkflow.ready)
-
-            await assert_eq_eventually(True, ready)
+            await wait_for_activity_processed(client, workflow_handle.id)
 
     # Second worker: Complete the workflow with fresh objects (new instrumentation)
     async with AgentEnvironment(
@@ -554,11 +551,7 @@ async def test_workflow_only_trace_to_spans(
             )
             workflow_id = workflow_handle.id
 
-            # Wait for workflow to be ready
-            async def ready() -> bool:
-                return await workflow_handle.query(SelfTracingWorkflow.ready)
-
-            await assert_eq_eventually(True, ready)
+            await wait_for_activity_processed(client, workflow_handle.id)
 
     # Second worker: Complete the workflow with fresh objects (new instrumentation)
     async with AgentEnvironment(
@@ -805,7 +798,6 @@ def is_descendant_of(child: ReadableSpan, ancestor_span_id: int) -> bool:
 class OtelSpanWorkflow:
     def __init__(self) -> None:
         self._proceed = False
-        self._ready = False
 
     @workflow.run
     async def run(self):
@@ -818,14 +810,9 @@ async def run(self):
                     simple_no_context_activity,
                     start_to_close_timeout=timedelta(seconds=10),
                 )
-                self._ready = True
                 await workflow.wait_condition(lambda: self._proceed)
         return "done"
 
-    @workflow.query
-    def ready(self) -> bool:
-        return self._ready
-
     @workflow.signal
     def proceed(self) -> None:
         self._proceed = True
@@ -868,11 +855,7 @@ async def test_sdk_trace_to_otel_span_parenting(
                 )
                 workflow_id = workflow_handle.id
 
-                # Wait for workflow to be ready
-                async def ready() -> bool:
-                    return await workflow_handle.query(OtelSpanWorkflow.ready)
-
-                await assert_eq_eventually(True, ready)
+                await wait_for_activity_processed(client, workflow_handle.id)
 
     # Second worker: Complete the workflow with fresh objects (new instrumentation)
     async with AgentEnvironment(

tests/helpers/__init__.py

@@ -16,6 +16,8 @@
     cast,
 )
 
+import pytest
+
 from temporalio.api.common.v1 import WorkflowExecution
 from temporalio.api.enums.v1 import EventType as EventType
 from temporalio.api.enums.v1 import IndexedValueType
@@ -31,7 +33,12 @@
     PollWorkflowExecutionUpdateRequest,
     UnpauseActivityRequest,
 )
-from temporalio.client import BuildIdOpAddNewDefault, Client, WorkflowHandle
+from temporalio.client import (
+    BuildIdOpAddNewDefault,
+    Client,
+    WorkflowExecutionStatus,
+    WorkflowHandle,
+)
 from temporalio.common import SearchAttributeKey
 from temporalio.converter import DataConverter
 from temporalio.service import RPCError, RPCStatusCode
@@ -83,9 +90,9 @@ async def assert_eventually(
             if timedelta(seconds=time.monotonic() - start_sec) >= timeout:
                 raise
         except RPCError as e:
-            if retry_on_rpc_cancelled and e.status == RPCStatusCode.CANCELLED:
-                continue
-            else:
+            if not (retry_on_rpc_cancelled and e.status == RPCStatusCode.CANCELLED):
+                raise
+            if timedelta(seconds=time.monotonic() - start_sec) >= timeout:
                 raise
         await asyncio.sleep(interval.total_seconds())
 
@@ -103,6 +110,44 @@ async def check() -> None:
     await assert_eventually(check, timeout=timeout, interval=interval)
 
 
+async def wait_for_workflow_idle(
+    handle: WorkflowHandle[Any, Any],
+    *,
+    timeout: timedelta = timedelta(seconds=10),
+    interval: timedelta = timedelta(milliseconds=200),
+) -> None:
+    """Wait until the running workflow has no pending workflow task, activity, child or Nexus op.
+
+    Use as an untraced readiness probe before a traced query, signal or update: pass a handle
+    from a client without tracing interceptors. Fails fast if the workflow is no longer running.
+
+    Limits: a workflow blocked only on a timer also counts as idle, and the time-skipping test
+    server never reports ``pending_workflow_task``, so this returns immediately there.
+    """
+
+    async def check() -> None:
+        desc = await handle.describe()
+        if desc.status != WorkflowExecutionStatus.RUNNING:
+            status = desc.status.name if desc.status is not None else "UNKNOWN"
+            pytest.fail(f"Workflow {handle.id} is {status}, not RUNNING")
+        raw = desc.raw_description
+        pending = [
+            name
+            for name, present in (
+                ("workflow task", raw.HasField("pending_workflow_task")),
+                ("activities", bool(raw.pending_activities)),
+                ("child workflows", bool(raw.pending_children)),
+                ("Nexus operations", bool(raw.pending_nexus_operations)),
+            )
+            if present
+        ]
+        assert not pending, (
+            f"Workflow {handle.id} still has pending {', '.join(pending)} after {timeout}"
+        )
+
+    await assert_eventually(check, timeout=timeout, interval=interval)
+
+
 async def assert_task_fail_eventually(
     handle: WorkflowHandle, *, message_contains: str | None = None
 ) -> None: