From e5df89be095378a4cf5991925eda09249b22cee8 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 3 Nov 2026 11:29:35 -0500 Subject: [PATCH] test(temporal): cover chat spans across cache eviction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > 🤖 **This text was AI-generated.** Cache eviction could lose the original Braintrust span for a long-running Pydantic AI chat. Before the fix in PR #1054, replay created multiple workflow spans for the same chat: ```text first message -> cache eviction -> replay -> duplicate workflow spans ``` Add a two-turn chat regression for both Temporal durability APIs. Force cache eviction, stop the worker while the chat waits, and resume it on a fresh worker: ```text workflow span └─ chat_session (same IDs and start time after restart) ├─ first agent run └─ second agent run ``` Check that open spans remain unfinished during eviction and retain their input, metadata, and parent relationships when the chat completes. Share the durability API cases with the existing subagent regression. Refs: https://linear.app/braintrustdata/issue/SDK-781 Refs: https://github.com/braintrustdata/braintrust-sdk-python/pull/1054 --- .../temporal/test_temporal_pydantic_ai.py | 142 ++++++++++++++++-- 1 file changed, 129 insertions(+), 13 deletions(-) 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)