diff --git a/tests/contrib/langsmith/test_integration.py b/tests/contrib/langsmith/test_integration.py index f48d9d6ac..01e047e2b 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, @@ -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), ) @@ -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 @@ -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), ) @@ -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( @@ -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" @@ -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), ) @@ -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), ) @@ -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) # --------------------------------------------------------------------------- @@ -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.") @@ -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]: @@ -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, @@ -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( @@ -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" ) @@ -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 = [ @@ -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.") @@ -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]: @@ -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, @@ -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 @@ -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" ) @@ -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"] + ) # --------------------------------------------------------------------------- @@ -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}" @@ -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") @@ -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", ], 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")