diff --git a/py/src/braintrust/integrations/temporal/test_temporal_pydantic_ai.py b/py/src/braintrust/integrations/temporal/test_temporal_pydantic_ai.py index 66064342..7c5240a3 100644 --- a/py/src/braintrust/integrations/temporal/test_temporal_pydantic_ai.py +++ b/py/src/braintrust/integrations/temporal/test_temporal_pydantic_ai.py @@ -1,7 +1,9 @@ """Tests for BraintrustPlugin combined with Pydantic AI's Temporal durability integration.""" +import asyncio import uuid import warnings +from copy import deepcopy import pytest @@ -14,7 +16,13 @@ import temporalio.workflow from braintrust.integrations.pydantic_ai import PydanticAIIntegration from braintrust.integrations.temporal import BraintrustPlugin -from braintrust.integrations.temporal.test_temporal import _spans_named, memory_logger, temporal_env # noqa: F401 +from braintrust.integrations.temporal.test_temporal import ( # noqa: F401 + _spans_named, + _wait_for_workflow_state, + memory_logger, + temporal_env, +) +from braintrust.merge_row_batch import merge_row_batch from pydantic_ai import Agent from pydantic_ai.durable_exec.temporal import AgentPlugin, PydanticAIPlugin from pydantic_ai.models.test import TestModel @@ -23,6 +31,15 @@ _HAS_DURABILITY_CAPABILITY = hasattr(pydantic_ai_temporal, "TemporalDurability") +_DURABILITY_APIS = [ + "wrapper", + pytest.param( + "capability", + marks=pytest.mark.skipif( + not _HAS_DURABILITY_CAPABILITY, reason="TemporalDurability needs a newer pydantic-ai" + ), + ), +] def _api_for_current_workflow() -> str: @@ -51,6 +68,9 @@ async def _delegate(task: str) -> str: _wrapper_parent = pydantic_ai_temporal.TemporalAgent( _wrapper_parent_agent, tool_activity_config={"": {"delegate_task": False}} ) + _wrapper_chat = pydantic_ai_temporal.TemporalAgent( + Agent(TestModel(custom_output_text="chat done"), name="wrapper_chat_agent") + ) if _HAS_DURABILITY_CAPABILITY: _capability_sub = Agent( @@ -65,11 +85,20 @@ async def _delegate(task: str) -> str: ) # `metadata={"temporal": False}` replaces `tool_activity_config` on the capability path. _capability_parent.tool_plain(name="delegate_task", metadata={"temporal": False})(_delegate) + _capability_chat = Agent( + TestModel(custom_output_text="chat done"), + name="capability_chat_agent", + capabilities=[pydantic_ai_temporal.TemporalDurability()], + ) _AGENTS = { "wrapper": (_wrapper_parent, _wrapper_sub), **({"capability": (_capability_parent, _capability_sub)} if _HAS_DURABILITY_CAPABILITY else {}), } +_CHAT_AGENTS = { + "wrapper": _wrapper_chat, + **({"capability": _capability_chat} if _HAS_DURABILITY_CAPABILITY else {}), +} @temporalio.workflow.defn @@ -88,21 +117,40 @@ async def run(self, prompt: str) -> str: return (await parent_agent.run(prompt)).output +@temporalio.workflow.defn +class ChatWorkflow: + """Keep a chat span open between agent runs while waiting for a later message.""" + + def __init__(self) -> None: + self._message: str | None = None + self._state = "starting" + + @temporalio.workflow.run + async def run(self, prompt: str) -> list[str]: + agent = _CHAT_AGENTS[_api_for_current_workflow()] + with braintrust.start_span(name="chat_session", type="task") as span: + first = await agent.run(prompt) + span.log(input=prompt, metadata={"first_turn_complete": True}) + self._state = "waiting" + await temporalio.workflow.wait_condition(lambda: self._message is not None) + second = await agent.run(self._message, message_history=first.all_messages()) + output = [first.output, second.output] + span.log(output=output) + return output + + @temporalio.workflow.signal + def send_message(self, message: str) -> None: + self._message = message + + @temporalio.workflow.query + def state(self) -> str: + return self._state + + # With no workflow cache, every workflow task replays the run from the start, as after a worker # restart or cache eviction. Replays must not log the in-workflow Pydantic AI spans again. @pytest.mark.parametrize("max_cached_workflows", [1000, 0], ids=["cached", "replay"]) -@pytest.mark.parametrize( - "api", - [ - "wrapper", - pytest.param( - "capability", - marks=pytest.mark.skipif( - not _HAS_DURABILITY_CAPABILITY, reason="TemporalDurability needs a newer pydantic-ai" - ), - ), - ], -) +@pytest.mark.parametrize("api", _DURABILITY_APIS) @pytest.mark.asyncio async def test_subagent_child_workflow_nests_under_delegate_task_span( temporal_env, # noqa: F811 @@ -158,3 +206,71 @@ async def test_subagent_child_workflow_nests_under_delegate_task_span( if span["span_parents"][0] in {other["span_id"] for other in chat_spans + model_activity_spans} } assert len(pairs) == 3 + + +@pytest.mark.parametrize("api", _DURABILITY_APIS) +@pytest.mark.asyncio +async def test_chat_spans_survive_worker_restart(temporal_env, memory_logger, api): # noqa: F811 + assert PydanticAIIntegration.setup() + agent = _CHAT_AGENTS[api] + client = await Client.connect( + temporal_env.client.service_client.config.target_host, + namespace=temporal_env.client.namespace, + plugins=[PydanticAIPlugin()], + ) + task_queue = f"test-chat-restart-{uuid.uuid4()}" + + def worker(): + return Worker( + client, + task_queue=task_queue, + workflows=[ChatWorkflow], + max_cached_workflows=0, + plugins=[AgentPlugin(agent), BraintrustPlugin(logger=memory_logger)], + ) + + async with worker(): + handle = await client.start_workflow( + ChatWorkflow.run, + "first message", + id=f"{api}:{uuid.uuid4()}", + task_queue=task_queue, + ) + await _wait_for_workflow_state(handle, "waiting") + + # Shutdown evicts any cached workflow. Capture what was actually exported before replay: + # a still-running chat must remain open, including after its coroutine has been unwound. + braintrust.flush() + before_restart = memory_logger.pop() + (original_workflow,) = _spans_named(before_restart, "temporal.workflow.ChatWorkflow") + (original_chat,) = _spans_named(before_restart, "chat_session") + for span in (original_workflow, original_chat): + assert "end" not in span["metrics"] + assert "error" not in span + + # A new worker must rebuild the context from history, then append the next turn to the + # original trace. No in-memory span or interceptor state is carried across workers. + async with worker(): + await _wait_for_workflow_state(handle, "waiting") + await handle.signal(ChatWorkflow.send_message, "second message") + assert await asyncio.wait_for(handle.result(), timeout=15) == ["chat done", "chat done"] + + braintrust.flush() + spans = merge_row_batch(deepcopy(before_restart) + memory_logger.pop()) + (workflow_span,) = _spans_named(spans, "temporal.workflow.ChatWorkflow") + (chat_span,) = _spans_named(spans, "chat_session") + for original, resumed in ((original_workflow, workflow_span), (original_chat, chat_span)): + for field in ("id", "span_id", "root_span_id"): + assert resumed[field] == original[field] + assert resumed["metrics"]["start"] == original["metrics"]["start"] + assert resumed["metrics"]["end"] > original["metrics"]["start"] + assert "error" not in resumed + assert chat_span["input"] == "first message" + assert chat_span["metadata"]["first_turn_complete"] is True + assert chat_span["output"] == ["chat done", "chat done"] + assert chat_span["span_parents"] == [workflow_span["span_id"]] + agent_spans = _spans_named(spans, f"agent_run [{agent.name}]") + assert len(agent_spans) == 2 + assert {span["root_span_id"] for span in spans} == {workflow_span["root_span_id"]} + assert all(span["span_parents"] == [chat_span["span_id"]] for span in agent_spans) + assert all(chat_span["metrics"]["end"] >= span["metrics"]["end"] for span in agent_spans)