pydantic/pydantic-ai
1
0
Fork 0
mirror of https://github.com/pydantic/pydantic-ai.git synced 2026-10-11 21:59:51 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
69ea1e57201ed59e7d30fbe4a7590d193acecce0
pydantic-ai/tests/durable_exec/test_dbos.py
T

4896 lines
211 KiB
Python
Raw Blame History

from __future__ import annotations
import asyncio
import os
import re
import time
import uuid
import warnings
from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Generator, Iterator, Sequence
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass, field
from datetime import UTC, datetime
from decimal import Decimal
from pathlib import Path
from typing import Any, Literal, cast
from unittest.mock import patch
import httpx2
import pytest
from httpx import AsyncClient
from pydantic import BaseModel
from pydantic_ai import (
Agent,
AgentStreamEvent,
CancellationToken,
CapabilityEvent,
CustomEvent,
FinalResultEvent,
FunctionToolCallEvent,
FunctionToolResultEvent,
InstructionPart,
ModelMessage,
ModelRequest,
ModelResponse,
ModelSettings,
PartDeltaEvent,
PartStartEvent,
RetryPromptPart,
RunContext,
RunUsage,
TextPart,
TextPartDelta,
Tool,
ToolCallPart,
ToolReturnPart,
UserPromptPart,
)
from pydantic_ai._run_context import get_current_run_context
from pydantic_ai._warnings import PydanticAIDeprecationWarning
from pydantic_ai.capabilities import MCP, Capability, DynamicCapability, Hooks
from pydantic_ai.capabilities.abstract import AbstractCapability
from pydantic_ai.capabilities.instrumentation import Instrumentation
from pydantic_ai.direct import model_request_stream
from pydantic_ai.exceptions import (
ApprovalRequired,
CallDeferred,
ModelHTTPError,
ModelRetry,
RunCancelled,
ToolFailed,
UnexpectedModelBehavior,
UsageLimitExceeded,
UserError,
)
from pydantic_ai.models import (
ModelRequestContext,
ModelRequestParameters,
ModelResolutionContext,
StreamedResponse,
)
from pydantic_ai.models.function import AgentInfo, FunctionModel
from pydantic_ai.models.instrumented import InstrumentationSettings
from pydantic_ai.models.test import TestModel
from pydantic_ai.models.wrapper import WrapperModel
from pydantic_ai.realtime import (
RealtimeModel,
RealtimeModelProfile,
RealtimeModelSettings,
RealtimeSession,
WebRTCSession,
)
from pydantic_ai.realtime.codec import RealtimeConnection
from pydantic_ai.run import AgentRunResult
from pydantic_ai.usage import RequestUsage, UsageLimits
from pydantic_ai.workspaces import LocalWorkspaceBackend
from ..conftest import IsDatetime, IsNow, IsStr, detach_dbos_logging
try:
from dbos import DBOS, DBOSConfig, SetWorkflowID
from dbos._error import DBOSAwaitedWorkflowCancelledError
from pydantic_ai.durable_exec._toolset import unwrap_recorded_tool_call_result, wrap_tool_call_result
from pydantic_ai.durable_exec.dbos import (
DBOSAgent, # pyright: ignore[reportDeprecated]
DBOSDurability,
DBOSModel,
StepConfig,
)
from pydantic_ai.durable_exec.dbos._mcp_toolset import DBOSMCPToolset, dbosify_mcp_toolset
from pydantic_ai.toolsets.external import TOOL_SCHEMA_VALIDATOR
except ImportError: # pragma: lax no cover
pytest.skip('DBOS is not installed', allow_module_level=True)
try:
import logfire
from logfire.testing import CaptureLogfire
except ImportError: # pragma: lax no cover
pytest.skip('logfire not installed', allow_module_level=True)
try:
from fastmcp.client.transports import StdioTransport
from pydantic_ai.mcp import MCPToolset
except ImportError: # pragma: lax no cover
pytest.skip('mcp not installed', allow_module_level=True)
try:
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.openai import OpenAIProvider
except ImportError: # pragma: lax no cover
pytest.skip('openai not installed', allow_module_level=True)
from pydantic_ai import ExternalToolset, FunctionToolset
from pydantic_ai.capabilities import ProcessEventStream, ResolveModelId, SelectModel, Toolset
from pydantic_ai.tools import DeferredToolRequests, DeferredToolResults, ToolDefinition
from pydantic_ai.toolsets import AbstractToolset, CombinedToolset, ToolsetTool
from pydantic_ai.toolsets._dynamic import DynamicToolset
from .._inline_snapshot import snapshot
from ..continuation_utils import ScriptedContinuationModel, StreamSegment, scripted_response
from ..model_lifecycle_utils import LifecycleTrackingModel
from .decision_spans import ShipIt, ShipItDecisionModel, decide_span_lineage
# `DBOSAgent` is deprecated in favor of `capabilities=[DBOSDurability(...)]`.
# These tests exercise the wrapper-agent path on purpose; suppress the warning here
# rather than globally in `pyproject.toml`. The `pytestmark` entry below covers warnings
# emitted *inside* test functions; the `filterwarnings` call below covers warnings emitted
# at module import time (e.g. `simple_dbos_agent = DBOSAgent(...)`).
warnings.filterwarnings('ignore', message='`DBOSAgent` is deprecated', category=PydanticAIDeprecationWarning)
pytestmark = [
pytest.mark.vcr,
pytest.mark.xdist_group(name='dbos'),
pytest.mark.filterwarnings('ignore:`DBOSAgent` is deprecated:pydantic_ai._warnings.PydanticAIDeprecationWarning'),
]
# We need to use a custom cached HTTP client here as the default one created for OpenAIProvider will be closed automatically
# at the end of each test, but we need this one to live longer.
http_client = httpx2.AsyncClient()
@pytest.fixture(autouse=True, scope='module')
async def close_cached_httpx_client(anyio_backend: str) -> AsyncIterator[None]:
try:
yield
finally:
await http_client.aclose()
@pytest.fixture(autouse=True, scope='module')
def setup_logfire_instrumentation() -> Iterator[None]:
# Set up logfire for the tests.
logfire.configure(metrics=False)
yield
@contextmanager
def workflow_raises(exc_type: type[Exception], exc_message: str) -> Generator[None]:
"""Helper for asserting that a DBOS workflow fails with the expected error."""
with pytest.raises(Exception) as exc_info:
yield
assert isinstance(exc_info.value, Exception)
assert str(exc_info.value) == exc_message
@pytest.fixture(scope='module')
def dbos(tmp_path_factory: pytest.TempPathFactory) -> Generator[DBOS, Any, None]:
dbos_sqlite_file = tmp_path_factory.mktemp('dbos') / 'dbostest.sqlite'
dbos_config: DBOSConfig = {
'name': 'pydantic_dbos_tests',
'system_database_url': f'sqlite:///{dbos_sqlite_file}',
'run_admin_server': False,
'enable_otlp': True,
}
dbos = DBOS(config=dbos_config)
DBOS.launch()
try:
yield dbos
finally:
DBOS.destroy()
detach_dbos_logging()
model = OpenAIChatModel(
'gpt-4o',
provider=OpenAIProvider(
api_key=os.getenv('OPENAI_API_KEY', 'mock-api-key'),
http_client=http_client,
),
)
# Not necessarily need to define it outside of the function. DBOS just requires workflows to be statically defined so recovery would be able to find those workflows. It's nice to reuse it in multiple tests.
simple_agent = Agent(model, name='simple_agent')
simple_dbos_agent = DBOSAgent(simple_agent) # pyright: ignore[reportDeprecated]
def test_dbos_agent_construction_warns_deprecated() -> None:
"""The `DBOSAgent` deprecation fires at runtime; the module-level filters only suppress it."""
with pytest.warns(PydanticAIDeprecationWarning, match='`DBOSAgent` is deprecated'):
DBOSAgent(Agent(TestModel(), name='dbos_agent_deprecation_probe')) # pyright: ignore[reportDeprecated]
async def test_simple_agent_run_in_workflow(allow_model_requests: None, dbos: DBOS, openai_api_key: str) -> None:
"""Test that a simple agent can run in a DBOS workflow."""
@DBOS.workflow()
async def run_simple_agent() -> str:
result = await simple_dbos_agent.run('What is the capital of Mexico?')
return result.output
output = await run_simple_agent()
assert output == snapshot('The capital of Mexico is Mexico City.')
async def round_trip_workspace(ctx: RunContext[None]) -> str:
await ctx.workspace.write_text('note.txt', 'workspace content')
return await ctx.workspace.read_text('note.txt')
async def test_dbos_agent_rejects_a_workspace_in_its_workflow(dbos: DBOS, tmp_path: Path) -> None:
"""The deprecated wrapper has no durability capability, so a workspace's operations could not run as steps."""
agent = Agent(
TestModel(call_tools=['round_trip_workspace']),
name='dbos_agent_workspace',
deps_type=type(None),
tools=[round_trip_workspace],
)
dbos_agent = DBOSAgent(agent) # pyright: ignore[reportDeprecated]
with pytest.raises(
UserError, match='Workspaces are not supported inside a DBOS workflow through the deprecated wrapper agent'
):
await dbos_agent.run('Use the workspace.', workspace=LocalWorkspaceBackend(tmp_path))
class Deps(BaseModel):
country: str
# Wrap event_stream_handler as a DBOS step because it's non-deterministic (uses logfire)
@DBOS.step()
async def event_stream_handler(
ctx: RunContext[Deps],
stream: AsyncIterable[AgentStreamEvent],
):
logfire.info(f'{ctx.run_step=}')
async for event in stream:
logfire.info('event', event=event)
@DBOS.step()
async def runtime_event_stream_handler(
ctx: RunContext[object],
stream: AsyncIterable[AgentStreamEvent],
):
logfire.info(f'{ctx.run_step=}')
async for event in stream:
logfire.info('runtime_event', event=event)
# This doesn't need to be a step
async def get_country(ctx: RunContext[Deps]) -> str:
return ctx.deps.country
class WeatherArgs(BaseModel):
city: str
@DBOS.step()
def get_weather(args: WeatherArgs) -> str:
if args.city == 'Mexico City':
return 'sunny'
else:
return 'unknown' # pragma: no cover
@dataclass
class Answer:
label: str
answer: str
@dataclass
class Response:
answers: list[Answer]
@dataclass
class BasicSpan:
content: str
children: list[BasicSpan] = field(default_factory=list['BasicSpan'])
parent_id: int | None = field(repr=False, compare=False, default=None)
complex_agent = Agent(
model,
deps_type=Deps,
output_type=Response,
toolsets=[
FunctionToolset[Deps](tools=[get_country], id='country'),
MCPToolset(StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='mcp', init_timeout=20),
ExternalToolset(tool_defs=[ToolDefinition(name='external')], id='external'),
],
tools=[get_weather],
capabilities=[Instrumentation(settings=InstrumentationSettings())], # Enable instrumentation for testing
name='complex_agent',
)
complex_dbos_agent = DBOSAgent(complex_agent, event_stream_handler=event_stream_handler) # pyright: ignore[reportDeprecated]
seq_complex_dbos_agent = DBOSAgent( # pyright: ignore[reportDeprecated]
complex_agent,
event_stream_handler=event_stream_handler,
parallel_execution_mode='sequential',
name='seq_complex_agent',
)
async def runtime_handler_stream_function(messages: list[ModelMessage], agent_info: AgentInfo) -> AsyncIterator[str]:
del messages, agent_info
yield 'Hello'
yield ' world'
runtime_handler_stream_agent = Agent(
FunctionModel(stream_function=runtime_handler_stream_function),
name='runtime_handler_stream_agent',
)
runtime_handler_stream_dbos_agent = DBOSAgent(runtime_handler_stream_agent) # pyright: ignore[reportDeprecated]
iter_handler_stream_agent = Agent(
FunctionModel(stream_function=runtime_handler_stream_function),
name='iter_handler_stream_agent',
capabilities=[DBOSDurability(event_stream_handler=runtime_event_stream_handler)],
)
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_complex_agent_run_in_workflow(allow_model_requests: None, dbos: DBOS, capfire: CaptureLogfire) -> None:
# Set a workflow ID for testing list steps
wfid = str(uuid.uuid4())
with SetWorkflowID(wfid):
# DBOSAgent already wraps the `run` function as a DBOS workflow, so we can just call it directly.
result = await complex_dbos_agent.run(
'Tell me: the capital of the country; the weather there; the product name', deps=Deps(country='Mexico')
)
assert result.output == snapshot(
Response(
answers=[
Answer(label='Capital of the country', answer='Mexico City'),
Answer(label='Weather in the capital', answer='Sunny'),
Answer(label='Product Name', answer='Pydantic AI'),
]
)
)
# Make sure the steps are persisted correctly in the DBOS database.
steps = await dbos.list_workflow_steps_async(wfid)
assert [step['function_name'] for step in steps] == snapshot(
[
'complex_agent__mcp_server__mcp.get_tools',
'complex_agent__model.request_stream',
'event_stream_handler',
'event_stream_handler',
'complex_agent__mcp_server__mcp.call_tool',
'event_stream_handler',
'event_stream_handler',
'complex_agent__model.request_stream',
'event_stream_handler',
'get_weather',
'event_stream_handler',
'complex_agent__model.request_stream',
'event_stream_handler',
'event_stream_handler',
]
)
exporter = capfire.exporter
spans = exporter.exported_spans_as_dict()
basic_spans_by_id = {
span['context']['span_id']: BasicSpan(
parent_id=span['parent']['span_id'] if span['parent'] else None,
content=attributes.get('event') or attributes['logfire.msg'],
)
for span in spans
if (attributes := span.get('attributes'))
}
assert len(basic_spans_by_id) > 0, 'No spans were exported'
root_span = None
for basic_span in basic_spans_by_id.values():
if basic_span.parent_id is None:
root_span = basic_span
else:
parent_id = basic_span.parent_id
parent_span = basic_spans_by_id[parent_id]
parent_span.children.append(basic_span)
def _normalize_json_spans(span: BasicSpan) -> None:
"""Normalize non-deterministic tool_call_ids in JSON event spans."""
import json
for child in span.children:
if child.content.startswith('{'):
try:
data = json.loads(child.content)
_strip_volatile_fields(data)
child.content = json.dumps(data)
except json.JSONDecodeError:
pass
_normalize_json_spans(child)
def _strip_volatile_fields(obj: dict[str, Any]) -> None:
for k, v in obj.items():
if k in ('tool_call_id', 'timestamp'):
obj[k] = None
elif isinstance(v, dict):
_strip_volatile_fields(cast(dict[str, Any], v))
assert root_span is not None
_normalize_json_spans(root_span)
# Assert the root span and its structure matches expected hierarchy
assert root_span == snapshot(
BasicSpan(
content='complex_agent.run',
children=[
BasicSpan(
content='complex_agent run',
children=[
BasicSpan(
content='complex_agent__mcp_server__mcp.get_tools',
children=[BasicSpan(content='tools/list')],
),
BasicSpan(
content='chat gpt-4o',
children=[
BasicSpan(
content='complex_agent__model.request_stream',
children=[
BasicSpan(content='ctx.run_step=1'),
BasicSpan(
content='{"index": 0, "part": {"tool_name": "get_country", "args": "", "tool_call_id": null, "tool_kind": null, "id": null, "provider_name": null, "provider_details": null, "part_kind": "tool-call"}, "previous_part_kind": null, "event_kind": "part_start"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "{}", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "part": {"tool_name": "get_country", "args": "{}", "tool_call_id": null, "tool_kind": null, "id": null, "provider_name": null, "provider_details": null, "part_kind": "tool-call"}, "next_part_kind": "tool-call", "event_kind": "part_end"}'
),
BasicSpan(
content='{"index": 1, "part": {"tool_name": "get_product_name", "args": "", "tool_call_id": null, "tool_kind": null, "id": null, "provider_name": null, "provider_details": null, "part_kind": "tool-call"}, "previous_part_kind": "tool-call", "event_kind": "part_start"}'
),
BasicSpan(
content='{"index": 1, "delta": {"tool_name_delta": null, "args_delta": "{}", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 1, "part": {"tool_name": "get_product_name", "args": "{}", "tool_call_id": null, "tool_kind": null, "id": null, "provider_name": null, "provider_details": null, "part_kind": "tool-call"}, "next_part_kind": null, "event_kind": "part_end"}'
),
],
)
],
),
BasicSpan(
content='event_stream_handler',
children=[
BasicSpan(content='ctx.run_step=1'),
BasicSpan(
content='{"part": {"tool_name": "get_country", "args": "{}", "tool_call_id": null, "tool_kind": null, "id": null, "provider_name": null, "provider_details": null, "part_kind": "tool-call"}, "args_valid": true, "event_kind": "function_tool_call"}'
),
],
),
BasicSpan(
content='event_stream_handler',
children=[
BasicSpan(content='ctx.run_step=1'),
BasicSpan(
content='{"part": {"tool_name": "get_product_name", "args": "{}", "tool_call_id": null, "tool_kind": null, "id": null, "provider_name": null, "provider_details": null, "part_kind": "tool-call"}, "args_valid": true, "event_kind": "function_tool_call"}'
),
],
),
BasicSpan(content='running tool: get_country'),
BasicSpan(
content='running tool: get_product_name',
children=[
BasicSpan(
content='complex_agent__mcp_server__mcp.call_tool',
children=[BasicSpan(content='tools/call get_product_name')],
)
],
),
BasicSpan(
content='event_stream_handler',
children=[
BasicSpan(content='ctx.run_step=1'),
BasicSpan(
content='{"part": {"tool_name": "get_country", "content": "Mexico", "tool_call_id": null, "tool_kind": null, "metadata": null, "timestamp": null, "outcome": "success", "part_kind": "tool-return"}, "content": null, "event_kind": "function_tool_result"}'
),
],
),
BasicSpan(
content='event_stream_handler',
children=[
BasicSpan(content='ctx.run_step=1'),
BasicSpan(
content='{"part": {"tool_name": "get_product_name", "content": "Pydantic AI", "tool_call_id": null, "tool_kind": null, "metadata": null, "timestamp": null, "outcome": "success", "part_kind": "tool-return"}, "content": null, "event_kind": "function_tool_result"}'
),
],
),
BasicSpan(
content='chat gpt-4o',
children=[
BasicSpan(
content='complex_agent__model.request_stream',
children=[
BasicSpan(content='ctx.run_step=2'),
BasicSpan(
content='{"index": 0, "part": {"tool_name": "get_weather", "args": "", "tool_call_id": null, "tool_kind": null, "id": null, "provider_name": null, "provider_details": null, "part_kind": "tool-call"}, "previous_part_kind": null, "event_kind": "part_start"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "{\\"", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "city", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "\\":\\"", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "Mexico", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": " City", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "\\"}", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "part": {"tool_name": "get_weather", "args": "{\\"city\\":\\"Mexico City\\"}", "tool_call_id": null, "tool_kind": null, "id": null, "provider_name": null, "provider_details": null, "part_kind": "tool-call"}, "next_part_kind": null, "event_kind": "part_end"}'
),
],
)
],
),
BasicSpan(
content='event_stream_handler',
children=[
BasicSpan(content='ctx.run_step=2'),
BasicSpan(
content='{"part": {"tool_name": "get_weather", "args": "{\\"city\\":\\"Mexico City\\"}", "tool_call_id": null, "tool_kind": null, "id": null, "provider_name": null, "provider_details": null, "part_kind": "tool-call"}, "args_valid": true, "event_kind": "function_tool_call"}'
),
],
),
BasicSpan(content='running tool: get_weather', children=[BasicSpan(content='get_weather')]),
BasicSpan(
content='event_stream_handler',
children=[
BasicSpan(content='ctx.run_step=2'),
BasicSpan(
content='{"part": {"tool_name": "get_weather", "content": "sunny", "tool_call_id": null, "tool_kind": null, "metadata": null, "timestamp": null, "outcome": "success", "part_kind": "tool-return"}, "content": null, "event_kind": "function_tool_result"}'
),
],
),
BasicSpan(
content='chat gpt-4o',
children=[
BasicSpan(
content='complex_agent__model.request_stream',
children=[
BasicSpan(content='ctx.run_step=3'),
BasicSpan(
content='{"index": 0, "part": {"tool_name": "final_result", "args": "", "tool_call_id": null, "tool_kind": null, "id": null, "provider_name": null, "provider_details": null, "part_kind": "tool-call"}, "previous_part_kind": null, "event_kind": "part_start"}'
),
BasicSpan(
content='{"tool_name": "final_result", "tool_call_id": null, "event_kind": "final_result"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "{\\"", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "answers", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "\\":[", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "{\\"", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "label", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "\\":\\"", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "Capital", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": " of", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": " the", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": " country", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "\\",\\"", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "answer", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "\\":\\"", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "Mexico", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": " City", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "\\"},{\\"", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "label", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "\\":\\"", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "Weather", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": " in", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": " the", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": " capital", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "\\",\\"", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "answer", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "\\":\\"", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "Sunny", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "\\"},{\\"", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "label", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "\\":\\"", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "Product", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": " Name", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "\\",\\"", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "answer", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "\\":\\"", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "P", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "yd", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "antic", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": " AI", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "\\"}", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "delta": {"tool_name_delta": null, "args_delta": "]}", "tool_call_id": null, "provider_name": null, "provider_details": null, "part_delta_kind": "tool_call"}, "event_kind": "part_delta"}'
),
BasicSpan(
content='{"index": 0, "part": {"tool_name": "final_result", "args": "{\\"answers\\":[{\\"label\\":\\"Capital of the country\\",\\"answer\\":\\"Mexico City\\"},{\\"label\\":\\"Weather in the capital\\",\\"answer\\":\\"Sunny\\"},{\\"label\\":\\"Product Name\\",\\"answer\\":\\"Pydantic AI\\"}]}", "tool_call_id": null, "tool_kind": null, "id": null, "provider_name": null, "provider_details": null, "part_kind": "tool-call"}, "next_part_kind": null, "event_kind": "part_end"}'
),
],
)
],
),
BasicSpan(
content='event_stream_handler',
children=[
BasicSpan(content='ctx.run_step=3'),
BasicSpan(
content='{"part": {"tool_name": "final_result", "args": "{\\"answers\\":[{\\"label\\":\\"Capital of the country\\",\\"answer\\":\\"Mexico City\\"},{\\"label\\":\\"Weather in the capital\\",\\"answer\\":\\"Sunny\\"},{\\"label\\":\\"Product Name\\",\\"answer\\":\\"Pydantic AI\\"}]}", "tool_call_id": null, "tool_kind": null, "id": null, "provider_name": null, "provider_details": null, "part_kind": "tool-call"}, "args_valid": true, "event_kind": "output_tool_call"}'
),
],
),
BasicSpan(
content='event_stream_handler',
children=[
BasicSpan(content='ctx.run_step=3'),
BasicSpan(
content='{"part": {"tool_name": "final_result", "content": "Final result processed.", "tool_call_id": null, "tool_kind": null, "metadata": null, "timestamp": null, "outcome": "success", "part_kind": "tool-return"}, "event_kind": "output_tool_result"}'
),
],
),
],
)
],
)
)
async def test_dbos_agent_run_in_workflow_with_runtime_event_stream_handler(
allow_model_requests: None, dbos: DBOS, capfire: CaptureLogfire
) -> None:
wfid = str(uuid.uuid4())
with SetWorkflowID(wfid):
result = await runtime_handler_stream_dbos_agent.run(
'Say hello', event_stream_handler=runtime_event_stream_handler
)
assert result.output == snapshot('Hello world')
steps = await dbos.list_workflow_steps_async(wfid)
step_names = [step['function_name'] for step in steps]
assert step_names[0] == 'runtime_handler_stream_agent__model.request_stream'
# The per-run handler fires live, nested inside the model-request step (delivering the streamed
# events asserted below). It is no longer invoked a second time at the graph level against the
# already-consumed, empty stream, so it doesn't appear as a separate top-level workflow step.
assert 'runtime_event_stream_handler' not in step_names
exported_event_messages = [
event
for span in capfire.exporter.exported_spans_as_dict()
if (attributes := span.get('attributes'))
and attributes.get('logfire.msg') == 'runtime_event'
and isinstance((event := attributes.get('event')), str)
]
assert exported_event_messages != []
async def test_dbos_agent_iter_in_workflow_fires_event_stream_handler(
allow_model_requests: None, dbos: DBOS, capfire: CaptureLogfire
) -> None:
"""`agent.iter()` inside a DBOS workflow delivers events to the durable `event_stream_handler`.
The handler used to be skipped entirely under `iter()`, because `wrap_run_event_stream` was
applied by `run()`/`run_stream()` rather than by the node stream primitives.
"""
@DBOS.workflow()
async def run_iter_workflow() -> str | None:
async with iter_handler_stream_agent.iter('Say hello') as run:
async for _node in run:
pass
assert run.result is not None
return run.result.output
wfid = str(uuid.uuid4())
with SetWorkflowID(wfid):
output = await run_iter_workflow()
assert output == snapshot('Hello world')
exported_event_messages = [
event
for span in capfire.exporter.exported_spans_as_dict()
if (attributes := span.get('attributes'))
and attributes.get('logfire.msg') == 'runtime_event'
and isinstance((event := attributes.get('event')), str)
]
assert exported_event_messages != []
async def test_dbos_agent_event_stream_handler_property_outside_workflow(dbos: DBOS) -> None:
# Outside a DBOS workflow, the `event_stream_handler` property resolves to the effective handler
# directly, rather than the in-workflow per-event dispatcher.
agent = Agent(TestModel(), name='event_stream_handler_property_agent')
dbos_agent = DBOSAgent(agent, event_stream_handler=runtime_event_stream_handler) # pyright: ignore[reportDeprecated]
assert dbos_agent.event_stream_handler is runtime_event_stream_handler
async def test_mcp_tools_not_cached_when_disabled(allow_model_requests: None, dbos: DBOS) -> None:
"""Verify that wrapper-level caching is skipped when cache_tools=False.
With caching disabled, every model request should be preceded by a get_tools step,
rather than only the first one populating the cache.
"""
mcp_toolset = next(ts for ts in complex_dbos_agent.toolsets if isinstance(ts, DBOSMCPToolset))
wrapped = cast(MCPToolset[Deps], mcp_toolset.wrapped)
original_cache_tools = wrapped.cache_tools
wrapped.cache_tools = False
try:
wfid = str(uuid.uuid4())
with SetWorkflowID(wfid):
result = await complex_dbos_agent.run(
'Tell me: the capital of the country; the weather there; the product name', deps=Deps(country='Mexico')
)
assert result.output == snapshot(
Response(
answers=[
Answer(label='Capital of the country', answer='Mexico City'),
Answer(label='Weather in the capital', answer='Sunny'),
Answer(label='Product Name', answer='Pydantic AI'),
]
)
)
steps = await dbos.list_workflow_steps_async(wfid)
step_names = [step['function_name'] for step in steps]
# Without caching, get_tools should be called 3 times (once per model request)
assert step_names.count('complex_agent__mcp_server__mcp.get_tools') == 3
finally:
wrapped.cache_tools = original_cache_tools
# Test sequential tool call works
async def test_complex_agent_run_sequential_tool(allow_model_requests: None, dbos: DBOS) -> None:
# Set a workflow ID for testing list steps
wfid = str(uuid.uuid4())
with SetWorkflowID(wfid):
# DBOSAgent already wraps the `run` function as a DBOS workflow, so we can just call it directly.
result = await seq_complex_dbos_agent.run(
'Tell me: the capital of the country; the weather there; the product name', deps=Deps(country='Mexico')
)
assert result.output == snapshot(
Response(
answers=[
Answer(label='Capital of the country', answer='Mexico City'),
Answer(label='Weather in the capital', answer='Sunny'),
Answer(label='Product Name', answer='Pydantic AI'),
]
)
)
# Make sure the steps are persisted correctly in the DBOS database.
steps = await dbos.list_workflow_steps_async(wfid)
assert [step['function_name'] for step in steps] == snapshot(
[
'seq_complex_agent__mcp_server__mcp.get_tools',
'seq_complex_agent__model.request_stream',
'event_stream_handler',
'event_stream_handler',
'event_stream_handler',
'seq_complex_agent__mcp_server__mcp.call_tool',
'event_stream_handler',
'seq_complex_agent__model.request_stream',
'event_stream_handler',
'get_weather',
'event_stream_handler',
'seq_complex_agent__model.request_stream',
'event_stream_handler',
'event_stream_handler',
]
)
async def test_multiple_agents(allow_model_requests: None, dbos: DBOS):
"""Test that multiple agents can run in a DBOS workflow."""
# This is just a smoke test to ensure that multiple agents can run in a DBOS workflow.
# We don't need to check the output as it's already tested in the individual agent tests.
result = await simple_dbos_agent.run('What is the capital of Mexico?')
assert result.output == snapshot('The capital of Mexico is Mexico City.')
result = await complex_dbos_agent.run(
'Tell me: the capital of the country; the weather there; the product name', deps=Deps(country='Mexico')
)
assert result.output == snapshot(
Response(
answers=[
Answer(label='Capital of the Country', answer='Mexico City'),
Answer(label='Weather in Mexico City', answer='Sunny'),
Answer(label='Product Name', answer='Pydantic AI'),
]
)
)
async def test_agent_name_collision(allow_model_requests: None, dbos: DBOS):
with pytest.raises(
Exception, match="Duplicate instance registration for class 'DBOSAgent' instance 'simple_agent'"
):
DBOSAgent(simple_agent) # pyright: ignore[reportDeprecated]
async def test_agent_without_name():
with pytest.raises(
UserError,
match=re.escape(
"An agent needs to have a unique `name` in order to be used with DBOS. The name will be used to identify the agent's workflows and steps."
),
):
DBOSAgent(Agent()) # pyright: ignore[reportDeprecated]
async def test_agent_without_model():
with pytest.raises(
UserError,
match=re.escape(
'An agent needs to have a `model` in order to be used with DBOS, it cannot be set at agent run time.'
),
):
DBOSAgent(Agent(name='test_agent')) # pyright: ignore[reportDeprecated]
async def test_toolset_without_id():
# Note: this is allowed in DBOS because we don't wrap the tools automatically in a workflow. It's up to the user to define the tools as DBOS steps if they want to use them as steps in a workflow.
DBOSAgent(Agent(model=model, name='test_agent', toolsets=[FunctionToolset()])) # pyright: ignore[reportDeprecated]
async def test_mcp_toolset_without_id():
# Unlike `FunctionToolset`, an `MCPToolset` is wrapped in a `DBOSMCPToolset` whose step names and per-run
# tool-defs cache key both derive from the toolset's `id`; without one, two id-less MCP toolsets would
# collide on both. Require an `id` up front, like Temporal does for all leaf toolsets.
with pytest.raises(
UserError,
match=re.escape(
'MCP toolsets need to have a unique `id` in order to be used with DBOS. '
"The ID will be used to identify the MCP server's steps within the workflow. "
'Set it on the toolset itself with `MCPToolset(..., id=...)`, or, when the toolset '
"is contributed by a capability, set the capability's `id` "
"(for example, `MCP(url='...', id='...')`)."
),
):
DBOSAgent(Agent(model=model, name='test_agent', toolsets=[MCPToolset('https://example.com/mcp')])) # pyright: ignore[reportDeprecated]
async def test_capability_contributed_toolset_id_from_capability():
"""A capability's `id` flows to its contributed leaf toolset, so a capability combined with a
local MCP server can be used under DBOS instead of tripping the id-less-MCP guard. An `MCP` with
no explicit `id` derives one from its URL, which is what lets the contributed leaf pass the guard.
This isn't a VCR test: it inspects the constructed toolset tree and DBOS registration during local
agent construction, before any model or MCP request, so there's no network round-trip to record.
Regression for https://github.com/pydantic/pydantic-ai/issues/6334.
"""
def add(x: int) -> int:
return x + 1 # pragma: no cover
agent = Agent(
model,
name='capability_agent',
capabilities=[
Capability(id='billing', tools=[add]),
MCP(url='https://mcp.example.com/api'),
],
)
# Previously raised `UserError` from the DBOS id-less-MCP guard because the contributed MCP leaf
# had `id=None`; it now derives `mcp.example.com-api` from the URL, so construction succeeds.
dbos_agent = DBOSAgent(agent) # pyright: ignore[reportDeprecated]
leaves: list[AbstractToolset[object]] = []
for toolset in dbos_agent.toolsets:
toolset.apply(leaves.append)
# The contributed MCP leaf carries the URL-derived id; the `billing` function toolset carries the
# capability id.
assert any(isinstance(ts, MCPToolset) and ts.id == 'mcp.example.com-api' for ts in leaves)
assert any(isinstance(ts, FunctionToolset) and ts.id == 'billing' for ts in leaves)
async def test_capability_contributed_toolsets_with_colliding_derived_id():
"""Two genuinely different MCP servers whose URLs derive the same id are rejected by `Agent`.
Both would key their steps on `a.com-api` under DBOS, the second server returning the first's
cached tools. DBOS used to be what caught it, because `MCP(url=...)` left the capability
anonymous and only its toolset carried the derived id, so the agent-level capability-id
uniqueness check had nothing to compare. Now the capability takes the same derived id, and a
plain `Agent` rejects the pair before any engine is involved -- so the same mistake is reported
the same way with or without durable execution.
DBOS keeps its own leaf walk for collisions that never reach a capability id, such as two
toolsets passed directly to `toolsets=`; `test_nested_toolsets_with_colliding_id` covers that.
This isn't a VCR test: the collision is rejected during local construction, before any model or
MCP request, so there's no network round-trip to record.
"""
with pytest.raises(
match=re.escape(
"Capability id 'a.com-api' is used by multiple capabilities. Ids identify one capability within a run, so give each a distinct `id`."
),
):
DBOSAgent( # pyright: ignore[reportDeprecated]
Agent(
model,
name='colliding_capability_agent',
capabilities=[MCP(url='https://a.com/api'), MCP(url='https://a.com/v2/api')],
)
)
async def test_nested_toolsets_with_colliding_id():
"""DBOS reports the shape of the collision, not just that there is one.
Its guard walks the leaves, so it sees a duplicate however deeply the toolset tree nests it, and
names the `MCPToolset` the user has to change. The id keys the MCP server's step names across
the whole workflow, so DBOS needs it unique whether or not the server contributes instructions.
"""
with pytest.raises(
UserError,
match=re.escape(
'MCP toolsets need to have a unique `id` in order to be used with DBOS, '
"but more than one leaf toolset uses the id 'dup'. "
"The ID identifies the MCP server's steps within the workflow, so duplicates would collide. "
'Set a distinct `id` on each `MCPToolset` (or the `Capability`/`MCP` that contributes it) to disambiguate them.'
),
):
DBOSAgent( # pyright: ignore[reportDeprecated]
Agent(
model,
name='colliding_nested_agent',
toolsets=[
CombinedToolset([MCPToolset('https://a.com/api', id='dup')]),
MCPToolset('https://a.com/v2/api', id='dup'),
],
)
)
async def test_dbos_agent():
assert isinstance(complex_dbos_agent.model, DBOSModel)
assert complex_dbos_agent.model.wrapped == complex_agent.model
# DBOS only wraps the MCP server toolsets. Other toolsets are not wrapped.
toolsets = complex_dbos_agent.toolsets
assert len(toolsets) == 5
# Empty function toolset for the agent's own tools
assert isinstance(toolsets[0], FunctionToolset)
assert toolsets[0].id == '<agent>'
assert toolsets[0].tools == {}
# Function toolset for the wrapped agent's own tools
assert isinstance(toolsets[1], FunctionToolset)
assert toolsets[1].id == '<agent>'
assert toolsets[1].tools.keys() == {'get_weather'}
# Wrapped 'country' toolset
assert isinstance(toolsets[2], FunctionToolset)
assert toolsets[2].id == 'country'
assert toolsets[2].tools.keys() == {'get_country'}
# Wrapped 'mcp' MCP server
assert isinstance(toolsets[3], DBOSMCPToolset)
assert toolsets[3].id == 'mcp'
assert toolsets[3].wrapped == complex_agent.toolsets[2]
# Unwrapped 'external' toolset
assert isinstance(toolsets[4], ExternalToolset)
assert toolsets[4].id == 'external'
assert toolsets[4] == complex_agent.toolsets[3]
async def test_dbos_agent_run(allow_model_requests: None, dbos: DBOS):
# Note: this runs as a DBOS workflow because we automatically wrap the run function.
result = await simple_dbos_agent.run('What is the capital of Mexico?')
assert result.output == snapshot('The capital of Mexico is Mexico City.')
def test_dbos_agent_run_sync(allow_model_requests: None, dbos: DBOS):
# Note: this runs as a DBOS workflow because we automatically wrap the run_sync function.
# This is equivalent to test_dbos_agent_run_sync_in_workflow
result = simple_dbos_agent.run_sync('What is the capital of Mexico?')
assert result.output == snapshot('The capital of Mexico is Mexico City.')
async def test_dbos_agent_run_stream(allow_model_requests: None):
# Run stream is not a DBOS workflow, so we can use it directly.
async with simple_dbos_agent.run_stream('What is the capital of Mexico?') as result:
assert [c async for c in result.stream_text(debounce_by=None)] == snapshot(
[
'The',
'The capital',
'The capital of',
'The capital of Mexico',
'The capital of Mexico is',
'The capital of Mexico is Mexico',
'The capital of Mexico is Mexico City',
'The capital of Mexico is Mexico City.',
]
)
async def test_dbos_agent_run_stream_events(allow_model_requests: None):
# This doesn't work because `run_stream_events` calls `run` internally, which is automatically wrapped in a DBOS workflow.
with pytest.raises(
UserError,
match=re.escape(
'`agent.run_stream_events()` cannot be used with DBOS. Set an `event_stream_handler` on the agent and use `agent.run()` instead.'
),
):
async with simple_dbos_agent.run_stream_events('What is the capital of Mexico?'):
pass
async def test_dbos_agent_iter(allow_model_requests: None):
output: list[str] = []
async with simple_dbos_agent.iter('What is the capital of Mexico?') as run:
async for node in run:
if Agent.is_model_request_node(node):
async with node.stream(run.ctx) as stream:
async for chunk in stream.stream_text(debounce_by=None):
output.append(chunk)
assert output == snapshot(
[
'The',
'The capital',
'The capital of',
'The capital of Mexico',
'The capital of Mexico is',
'The capital of Mexico is Mexico',
'The capital of Mexico is Mexico City',
'The capital of Mexico is Mexico City.',
]
)
def test_dbos_agent_run_sync_in_workflow(allow_model_requests: None, dbos: DBOS):
# DBOS allows calling `run_sync` inside a workflow as a child workflow.
@DBOS.workflow()
def run_sync_workflow():
result = simple_dbos_agent.run_sync('What is the capital of Mexico?')
return result.output
output = run_sync_workflow()
assert output == snapshot('The capital of Mexico is Mexico City.')
async def test_dbos_agent_run_stream_in_workflow(allow_model_requests: None, dbos: DBOS):
@DBOS.workflow()
async def run_stream_workflow():
async with simple_dbos_agent.run_stream('What is the capital of Mexico?') as result:
pass
return await result.get_output() # pragma: no cover
with workflow_raises(
UserError,
snapshot(
'`agent.run_stream()` cannot be used inside a DBOS workflow. '
'Set an `event_stream_handler` on the agent and use `agent.run()` instead.'
),
):
await run_stream_workflow()
async def test_dbos_agent_run_stream_events_in_workflow(allow_model_requests: None, dbos: DBOS):
@DBOS.workflow()
async def run_stream_events_workflow():
async with simple_dbos_agent.run_stream_events('What is the capital of Mexico?') as event_stream:
return [event async for event in event_stream] # pragma: no cover
with workflow_raises(
UserError,
snapshot(
'`agent.run_stream_events()` cannot be used with DBOS. Set an `event_stream_handler` on the agent and use `agent.run()` instead.'
),
):
await run_stream_events_workflow()
async def test_dbos_agent_realtime_session_in_workflow():
# A realtime session opens a long-lived, non-deterministic connection, so it can't run inside a
# workflow; the guard trips before the model is ever connected.
with patch.object(DBOS, 'workflow_id', 'wf-1'):
with pytest.raises(UserError, match='cannot be used inside a DBOS workflow'):
async with simple_dbos_agent.realtime(cast('Any', object())).session():
pass # pragma: no cover
async def test_dbos_agent_realtime_signaling_in_workflow():
# Browser-call signaling issues a live provider request, so workflow code can't call it directly:
# the two helpers reach the agent through `_resolve_realtime_session`, which the wrapper guards too.
with patch.object(DBOS, 'workflow_id', 'wf-1'):
realtime = simple_dbos_agent.realtime(cast('Any', object()))
with pytest.raises(UserError, match='cannot be used directly inside a DBOS workflow'):
await realtime.answer_webrtc_offer('v=0')
with pytest.raises(UserError, match='cannot be used directly inside a DBOS workflow'):
await realtime.create_client_secret()
with pytest.raises(UserError, match='cannot be used directly inside a DBOS workflow'):
await realtime.hang_up(WebRTCSession(provider_name='openai', session_id='rtc_x'))
async def test_dbos_agent_realtime_signaling_in_step():
# Inside a step — the boundary where DBOS records non-deterministic I/O — signaling delegates to
# the wrapped agent like `run()` does. The fake model has no WebRTC support, so reaching *its*
# refusal proves the workflow guard let the call through.
with patch.object(DBOS, 'workflow_id', 'wf-1'), patch.object(DBOS, 'step_id', 7):
realtime = simple_dbos_agent.realtime(_FakeRealtimeModel())
with pytest.raises(UserError, match='does not support WebRTC'):
await realtime.create_client_secret()
with pytest.raises(UserError, match='cannot end a call from the server'):
await realtime.hang_up(WebRTCSession(provider_name='openai', session_id='rtc_x'))
class _FakeRealtimeConnection(RealtimeConnection):
async def send(self, content: Any) -> None: ... # pragma: no cover
async def __aiter__(self) -> AsyncIterator[Any]:
return
yield # pragma: no cover
class _FakeRealtimeModel(RealtimeModel):
@property
def model_name(self) -> str:
return 'fake-realtime'
@property
def system(self) -> str:
return 'fake'
@property
def profile(self) -> RealtimeModelProfile:
return RealtimeModelProfile(
supports_image_input=True,
supports_manual_turn_control=True,
supports_interruption=True,
supports_output_truncation=True,
supports_session_seeding=True,
supported_native_tools=frozenset(),
)
@asynccontextmanager
async def connect(
self,
*,
messages: Sequence[ModelMessage],
model_settings: RealtimeModelSettings | None,
model_request_parameters: ModelRequestParameters,
) -> AsyncGenerator[_FakeRealtimeConnection]:
yield _FakeRealtimeConnection()
async def test_dbos_agent_realtime_session_outside_workflow():
# Outside a workflow, the session is delegated to the wrapped agent.
async with simple_dbos_agent.realtime(_FakeRealtimeModel()).session() as session:
assert isinstance(session, RealtimeSession)
assert [event async for event in session] == []
async def test_dbos_agent_iter_in_workflow(allow_model_requests: None, dbos: DBOS):
# DBOS allows calling `iter` inside a workflow as a step.
@DBOS.workflow()
async def run_iter_workflow():
output: list[str] = []
async with simple_dbos_agent.iter('What is the capital of Mexico?') as run:
async for node in run:
if Agent.is_model_request_node(node):
async with node.stream(run.ctx) as stream:
async for chunk in stream.stream_text(debounce_by=None):
output.append(chunk)
return output
output = await run_iter_workflow()
# If called in a workflow, the output is a single concatenated string.
assert output == snapshot(
[
'The capital of Mexico is Mexico City.',
]
)
async def test_dbos_agent_run_in_workflow_with_event_stream_handler(allow_model_requests: None, dbos: DBOS) -> None:
# DBOS workflow input must be serializable, so we cannot use an inner function as an argument.
# It's fine to pass in an event_stream_handler that is defined as a top-level function.
async def simple_event_stream_handler(
ctx: RunContext,
stream: AsyncIterable[AgentStreamEvent],
):
pass
with pytest.raises(Exception) as exc_info:
await simple_dbos_agent.run('What is the capital of Mexico?', event_stream_handler=simple_event_stream_handler)
# 3.14+: _pickle.PicklingError("Can't pickle local object <function ...>")
# <=3.13: AttributeError("Can't get local object '...<locals>...'")
assert 'local object' in str(exc_info.value) and 'simple_event_stream_handler' in str(exc_info.value)
async def test_dbos_agent_run_in_workflow_with_model(allow_model_requests: None, dbos: DBOS):
# A non-DBOS model is not wrapped as steps so it's not deterministic and cannot be used in a DBOS workflow.
with workflow_raises(
UserError,
snapshot(
'Non-DBOS model cannot be set at agent run time inside a DBOS workflow, it must be set at agent creation time.'
),
):
await simple_dbos_agent.run('What is the capital of Mexico?', model=model)
async def test_dbos_cancel_suspended_response_runs_in_step(allow_model_requests: None, dbos: DBOS):
"""`DBOSModel.cancel_suspended_response` must run as a DBOS step, not inline in the workflow.
The provider teardown that cancels a server-side suspended/background job is a raw HTTP call;
wrapping it as a step makes it durable, retried, and recorded rather than running unrecorded
inside the workflow.
"""
cancelled: list[ModelResponse] = []
class RecordingModel(TestModel):
async def cancel_suspended_response(self, response: ModelResponse) -> None:
cancelled.append(response)
dbos_model = DBOSModel(
RecordingModel(),
step_name_prefix='cancel_suspended',
step_config={},
get_event_stream_handler=lambda: None,
)
response = ModelResponse(parts=[TextPart('paused')], state='suspended')
wfid = str(uuid.uuid4())
@DBOS.workflow()
async def cancel_in_workflow() -> None:
await dbos_model.cancel_suspended_response(response)
with SetWorkflowID(wfid):
await cancel_in_workflow()
steps = await dbos.list_workflow_steps_async(wfid)
assert [step['function_name'] for step in steps] == snapshot(['cancel_suspended__model.cancel_suspended_response'])
assert cancelled == [response]
async def test_dbos_agent_run_in_workflow_with_toolsets(allow_model_requests: None, dbos: DBOS):
# Since DBOS does not automatically wrap the tools in a workflow, and allows dynamic steps, we can pass in toolsets directly.
result = await simple_dbos_agent.run('What is the capital of Mexico?', toolsets=[FunctionToolset()])
assert result.output == snapshot('The capital of Mexico is Mexico City.')
async def test_dbos_agent_run_in_workflow_with_runtime_external_toolset(dbos: DBOS):
def request_external_tool(_: list[ModelMessage], __: AgentInfo) -> ModelResponse:
return ModelResponse(parts=[ToolCallPart('external', {'query': 'runtime'}, tool_call_id='call-1')])
agent = Agent(
FunctionModel(request_external_tool),
name='runtime_external_toolset_agent',
output_type=[str, DeferredToolRequests],
)
dbos_agent = DBOSAgent(agent) # pyright: ignore[reportDeprecated]
result = await dbos_agent.run(
'Call the runtime external tool.',
toolsets=[
ExternalToolset(
tool_defs=[
ToolDefinition(
name='external',
parameters_json_schema={
'type': 'object',
'properties': {'query': {'type': 'string'}},
'required': ['query'],
},
)
],
id='external',
)
],
)
assert result.output == DeferredToolRequests(
calls=[ToolCallPart('external', {'query': 'runtime'}, tool_call_id='call-1')]
)
def runtime_tool() -> str:
return 'tool-result'
async def test_dbos_agent_run_in_workflow_with_runtime_function_toolset(dbos: DBOS):
# Unlike Temporal and Prefect, DBOS runs function tools inline rather than wrapping them, so a
# `FunctionToolset` added per-run is allowed and executes like a constructor-time one. (The toolset,
# like all DBOS workflow arguments, must be serializable, so the tool is a module-level function.)
def call_then_answer(messages: list[ModelMessage], _: AgentInfo) -> ModelResponse:
if any(isinstance(part, ToolReturnPart) for message in messages for part in message.parts):
return ModelResponse(parts=[TextPart('done')])
return ModelResponse(parts=[ToolCallPart('runtime_tool', {}, tool_call_id='call-1')])
agent = Agent(FunctionModel(call_then_answer), name='runtime_function_toolset_agent')
dbos_agent = DBOSAgent(agent) # pyright: ignore[reportDeprecated]
result = await dbos_agent.run(
'Call the runtime tool.', toolsets=[FunctionToolset(tools=[runtime_tool], id='runtime_fn')]
)
assert result.output == 'done'
assert any(
isinstance(part, ToolReturnPart) and part.content == 'tool-result'
for message in result.all_messages()
for part in message.parts
)
async def test_dbos_agent_run_in_workflow_rejects_runtime_mcp_toolset(dbos: DBOS):
with workflow_raises(
UserError,
snapshot(
"MCPToolset 'runtime_mcp' cannot be added at runtime with DBOS, because toolsets that execute their own tools or resolve dynamically must be registered for durable execution when the agent is constructed. Pass them to the agent constructor instead -- not to `run(toolsets=...)` or `override(toolsets=...)`, and not via a post-construction `@agent.toolset`. Non-executing toolsets like `ExternalToolset` can be passed at runtime."
),
):
await simple_dbos_agent.run(
'Hello',
toolsets=[MCPToolset(StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='runtime_mcp')],
)
async def test_dbos_agent_run_in_workflow_rejects_runtime_dynamic_toolset(dbos: DBOS):
with workflow_raises(
UserError,
snapshot(
"DynamicToolset 'runtime_dynamic' cannot be added at runtime with DBOS, because toolsets that execute their own tools or resolve dynamically must be registered for durable execution when the agent is constructed. Pass them to the agent constructor instead -- not to `run(toolsets=...)` or `override(toolsets=...)`, and not via a post-construction `@agent.toolset`. Non-executing toolsets like `ExternalToolset` can be passed at runtime."
),
):
await simple_dbos_agent.run(
'Hello',
toolsets=[DynamicToolset(lambda _: FunctionToolset(), id='runtime_dynamic')],
)
async def test_dbos_agent_override_model_in_workflow(allow_model_requests: None, dbos: DBOS):
# We cannot override the model to a non-DBOS one in a DBOS workflow.
with workflow_raises(
UserError,
snapshot(
'Non-DBOS model cannot be contextually overridden inside a DBOS workflow, it must be set at agent creation time.'
),
):
with simple_dbos_agent.override(model=model):
pass
async def test_dbos_agent_override_toolsets_in_workflow(allow_model_requests: None, dbos: DBOS):
# Since DBOS does not automatically wrap the tools in a workflow, and allows dynamic steps, we can override toolsets directly.
@DBOS.workflow()
async def run_with_toolsets():
with simple_dbos_agent.override(toolsets=[FunctionToolset()]):
pass
await run_with_toolsets()
async def test_dbos_agent_override_tools_in_workflow(allow_model_requests: None, dbos: DBOS):
# Since DBOS does not automatically wrap the tools in a workflow, and allows dynamic steps, we can override tools directly.
@DBOS.workflow()
async def run_with_tools():
with simple_dbos_agent.override(tools=[get_weather]):
result = await simple_dbos_agent.run('What is the capital of Mexico?')
return result.output
output = await run_with_tools()
assert output == snapshot('The capital of Mexico is Mexico City.')
async def test_dbos_agent_override_deps_in_workflow(allow_model_requests: None, dbos: DBOS):
# This is allowed
@DBOS.workflow()
async def run_with_deps():
with simple_dbos_agent.override(deps=None):
result = await simple_dbos_agent.run('What is the capital of the country?')
return result.output
output = await run_with_deps()
assert output == snapshot('The capital of Mexico is Mexico City.')
async def test_dbos_model_stream_direct(allow_model_requests: None, dbos: DBOS):
@DBOS.workflow()
async def run_model_stream():
messages: list[ModelMessage] = [ModelRequest.user_text_prompt('What is the capital of Mexico?')]
async with model_request_stream(complex_dbos_agent.model, messages) as stream:
async for _ in stream:
pass
with workflow_raises(
AssertionError,
snapshot(
'A DBOS model cannot be used with `pydantic_ai.direct.model_request_stream()` as it requires a `run_context`. Set an `event_stream_handler` on the agent and use `agent.run()` instead.'
),
):
await run_model_stream()
@dataclass
class UnserializableDeps:
client: AsyncClient
unserializable_deps_agent = Agent(model, name='unserializable_deps_agent', deps_type=UnserializableDeps)
@unserializable_deps_agent.tool
async def get_model_name(ctx: RunContext[UnserializableDeps]) -> int:
return ctx.deps.client.max_redirects # pragma: lax no cover
async def test_dbos_agent_with_unserializable_deps_type(allow_model_requests: None, dbos: DBOS):
unserializable_deps_dbos_agent = DBOSAgent(unserializable_deps_agent) # pyright: ignore[reportDeprecated]
# Test this raises a serialization error because httpx.AsyncClient is not serializable.
with pytest.raises(Exception) as exc_info:
async with AsyncClient() as client:
# This will trigger the client to be unserializable
logfire.instrument_httpx(client, capture_all=True)
await unserializable_deps_dbos_agent.run('What is the model name?', deps=UnserializableDeps(client=client))
assert str(exc_info.value) == snapshot("cannot pickle '_thread.RLock' object")
# Test dynamic toolsets in an agent with DBOS
@DBOS.step()
def temperature_celsius(city: str) -> float:
return 21.0
@DBOS.step()
def temperature_fahrenheit(city: str) -> float:
return 69.8
weather_toolset = FunctionToolset(tools=[temperature_celsius, temperature_fahrenheit])
@weather_toolset.tool
@DBOS.step()
def conditions(ctx: RunContext, city: str) -> str:
if ctx.run_step % 2 == 0:
return "It's sunny" # pragma: lax no cover
else:
return "It's raining"
datetime_toolset = FunctionToolset()
@DBOS.step()
def now_func() -> datetime:
return datetime.now()
datetime_toolset.add_function(now_func, name='now')
@dataclass
class ToggleableDeps:
active: Literal['weather', 'datetime']
def toggle(self):
if self.active == 'weather':
self.active = 'datetime'
else:
self.active = 'weather'
test_model = TestModel()
dynamic_agent = Agent(name='dynamic_agent', model=test_model, deps_type=ToggleableDeps)
@dynamic_agent.toolset
def toggleable_toolset(ctx: RunContext[ToggleableDeps]) -> FunctionToolset:
if ctx.deps.active == 'weather':
return weather_toolset
else:
return datetime_toolset
@dynamic_agent.tool
def toggle(ctx: RunContext[ToggleableDeps]):
ctx.deps.toggle()
dynamic_dbos_agent = DBOSAgent(dynamic_agent) # pyright: ignore[reportDeprecated]
def test_dbos_agent_explicit_run_id(dbos: DBOS):
"""A pre-minted `run_id=` is preserved through DBOSAgent.run_sync."""
agent = Agent(TestModel(custom_output_text='ok'), name='run_id_dbos_agent')
dbos_agent = DBOSAgent(agent) # pyright: ignore[reportDeprecated]
result = dbos_agent.run_sync('Hello', run_id='run-from-dbos')
assert result.run_id == 'run-from-dbos'
assert all(m.run_id == 'run-from-dbos' for m in result.all_messages())
def test_dynamic_toolset(dbos: DBOS):
weather_deps = ToggleableDeps('weather')
result = dynamic_dbos_agent.run_sync('Toggle the toolset', deps=weather_deps)
assert result.output == snapshot(
'{"toggle":null,"temperature_celsius":21.0,"temperature_fahrenheit":69.8,"conditions":"It\'s raining"}'
)
result = dynamic_dbos_agent.run_sync('Toggle the toolset', deps=weather_deps)
assert result.output == snapshot(IsStr(regex=r'{"toggle":null,"now":".+?"}'))
# Test human-in-the-loop with DBOS agent
hitl_agent = Agent(
model,
name='hitl_agent',
output_type=[str, DeferredToolRequests],
instructions='Just call tools without asking for confirmation.',
)
@hitl_agent.tool
@DBOS.step()
def create_file(ctx: RunContext, path: str) -> None:
raise CallDeferred
@hitl_agent.tool
@DBOS.step()
def delete_file(ctx: RunContext, path: str) -> bool:
if not ctx.tool_call_approved:
raise ApprovalRequired
return True
hitl_dbos_agent = DBOSAgent(hitl_agent) # pyright: ignore[reportDeprecated]
async def test_dbos_agent_with_hitl_tool(allow_model_requests: None, dbos: DBOS):
# Main loop for the agent, keep running until we get a final string output.
@DBOS.workflow()
async def hitl_main_loop(prompt: str) -> AgentRunResult[str | DeferredToolRequests]:
messages: list[ModelMessage] = [ModelRequest.user_text_prompt(prompt)]
deferred_tool_results: DeferredToolResults | None = None
while True:
result = await hitl_dbos_agent.run(message_history=messages, deferred_tool_results=deferred_tool_results)
messages = result.all_messages()
if isinstance(result.output, DeferredToolRequests):
deferred_tool_requests = result.output
# Set deferred_tool_requests as a DBOS workflow event, so the external functions can see it.
await DBOS.set_event_async('deferred_tool_requests', deferred_tool_requests)
# Wait for the deferred tool requests to be handled externally.
deferred_tool_results = await DBOS.recv_async('deferred_tool_results', timeout_seconds=30)
else:
return result
wf_handle = await DBOS.start_workflow_async(hitl_main_loop, 'Delete the file `.env` and create `test.txt`')
while True:
await asyncio.sleep(1)
status = await wf_handle.get_status()
if status.status == 'SUCCESS':
break
assert status.status == 'PENDING'
# Wait and check if the workflow has set a deferred tool request event.
deferred_tool_requests = await DBOS.get_event_async(
wf_handle.workflow_id, 'deferred_tool_requests', timeout_seconds=1
)
if deferred_tool_requests is not None: # pragma: no branch
results = DeferredToolResults()
# Approve all calls
for tool_call in deferred_tool_requests.approvals:
results.approvals[tool_call.tool_call_id] = True
for tool_call in deferred_tool_requests.calls:
results.calls[tool_call.tool_call_id] = 'Success'
# Signal the workflow with the results.
await DBOS.send_async(wf_handle.workflow_id, results, topic='deferred_tool_results')
result = await wf_handle.get_result()
assert result.output == snapshot('The file `.env` has been deleted and `test.txt` has been created successfully.')
assert result.all_messages() == snapshot(
[
ModelRequest(
parts=[
UserPromptPart(
content='Delete the file `.env` and create `test.txt`',
timestamp=IsDatetime(),
)
],
timestamp=IsNow(tz=UTC),
instructions='Just call tools without asking for confirmation.',
run_id=IsStr(),
conversation_id=IsStr(),
),
ModelResponse(
parts=[
ToolCallPart(
tool_name='delete_file',
args='{"path": ".env"}',
tool_call_id='call_jYdIdRZHxZTn5bWCq5jlMrJi',
),
ToolCallPart(
tool_name='create_file',
args='{"path": "test.txt"}',
tool_call_id='call_TmlTVWQbzrXCZ4jNsCVNbNqu',
),
],
usage=RequestUsage(
input_tokens=71,
output_tokens=46,
output_reasoning_tokens=0,
details={
'accepted_prediction_tokens': 0,
'audio_tokens': 0,
'reasoning_tokens': 0,
'rejected_prediction_tokens': 0,
},
cost=Decimal('0.0006375'),
),
model_name=IsStr(),
timestamp=IsDatetime(),
provider_name='openai',
provider_url='https://api.openai.com/v1/',
provider_details={
'finish_reason': 'tool_calls',
'service_tier': 'default',
'timestamp': IsDatetime(),
},
provider_response_id=IsStr(),
finish_reason='tool_call',
run_id=IsStr(),
conversation_id=IsStr(),
),
ModelRequest(
parts=[
ToolReturnPart(
tool_name='delete_file',
content=True,
tool_call_id=IsStr(),
timestamp=IsDatetime(),
),
ToolReturnPart(
tool_name='create_file',
content='Success',
tool_call_id=IsStr(),
timestamp=IsDatetime(),
),
],
timestamp=IsNow(tz=UTC),
instructions='Just call tools without asking for confirmation.',
run_id=IsStr(),
conversation_id=IsStr(),
),
ModelResponse(
parts=[
TextPart(content='The file `.env` has been deleted and `test.txt` has been created successfully.')
],
usage=RequestUsage(
input_tokens=133,
output_tokens=19,
output_reasoning_tokens=0,
details={
'accepted_prediction_tokens': 0,
'audio_tokens': 0,
'reasoning_tokens': 0,
'rejected_prediction_tokens': 0,
},
cost=Decimal('0.0005225'),
),
model_name='gpt-4o-2024-08-06',
timestamp=IsDatetime(),
provider_name='openai',
provider_url='https://api.openai.com/v1/',
provider_details={
'finish_reason': 'stop',
'service_tier': 'default',
'timestamp': IsDatetime(),
},
provider_response_id=IsStr(),
finish_reason='stop',
run_id=IsStr(),
conversation_id=IsStr(),
),
]
)
def test_dbos_agent_with_hitl_tool_sync(allow_model_requests: None, dbos: DBOS):
# Main loop for the agent, keep running until we get a final string output.
@DBOS.workflow()
def hitl_main_loop_sync(prompt: str) -> AgentRunResult[str | DeferredToolRequests]:
messages: list[ModelMessage] = [ModelRequest.user_text_prompt(prompt)]
deferred_tool_results: DeferredToolResults | None = None
while True:
result = hitl_dbos_agent.run_sync(message_history=messages, deferred_tool_results=deferred_tool_results)
messages = result.all_messages()
if isinstance(result.output, DeferredToolRequests):
deferred_tool_requests = result.output
# Set deferred_tool_requests as a DBOS workflow event, so the external functions can see it.
DBOS.set_event('deferred_tool_requests', deferred_tool_requests)
# Wait for the deferred tool requests to be handled externally.
deferred_tool_results = DBOS.recv('deferred_tool_results', timeout_seconds=30)
else:
return result
wf_handle = DBOS.start_workflow(hitl_main_loop_sync, 'Delete the file `.env` and create `test.txt`')
while True:
time.sleep(1)
status = wf_handle.get_status()
if status.status == 'SUCCESS':
break
# Wait and check if the workflow has set a deferred tool request event.
deferred_tool_requests = DBOS.get_event(wf_handle.workflow_id, 'deferred_tool_requests', timeout_seconds=1)
if deferred_tool_requests is not None: # pragma: no branch
results = DeferredToolResults()
# Approve all calls
for tool_call in deferred_tool_requests.approvals:
results.approvals[tool_call.tool_call_id] = True
for tool_call in deferred_tool_requests.calls:
results.calls[tool_call.tool_call_id] = 'Success'
# Signal the workflow with the results.
DBOS.send(wf_handle.workflow_id, results, topic='deferred_tool_results')
result = wf_handle.get_result()
assert result.output == snapshot('The file `.env` has been deleted and `test.txt` has been created successfully.')
assert result.all_messages() == snapshot(
[
ModelRequest(
parts=[
UserPromptPart(
content='Delete the file `.env` and create `test.txt`',
timestamp=IsDatetime(),
)
],
timestamp=IsNow(tz=UTC),
instructions='Just call tools without asking for confirmation.',
run_id=IsStr(),
conversation_id=IsStr(),
),
ModelResponse(
parts=[
ToolCallPart(
tool_name='delete_file',
args='{"path": ".env"}',
tool_call_id='call_jYdIdRZHxZTn5bWCq5jlMrJi',
),
ToolCallPart(
tool_name='create_file',
args='{"path": "test.txt"}',
tool_call_id='call_TmlTVWQbzrXCZ4jNsCVNbNqu',
),
],
usage=RequestUsage(
input_tokens=71,
output_tokens=46,
output_reasoning_tokens=0,
details={
'accepted_prediction_tokens': 0,
'audio_tokens': 0,
'reasoning_tokens': 0,
'rejected_prediction_tokens': 0,
},
cost=Decimal('0.0006375'),
),
model_name=IsStr(),
timestamp=IsDatetime(),
provider_name='openai',
provider_url='https://api.openai.com/v1/',
provider_details={
'finish_reason': 'tool_calls',
'service_tier': 'default',
'timestamp': IsDatetime(),
},
provider_response_id=IsStr(),
finish_reason='tool_call',
run_id=IsStr(),
conversation_id=IsStr(),
),
ModelRequest(
parts=[
ToolReturnPart(
tool_name='delete_file',
content=True,
tool_call_id=IsStr(),
timestamp=IsDatetime(),
),
ToolReturnPart(
tool_name='create_file',
content='Success',
tool_call_id=IsStr(),
timestamp=IsDatetime(),
),
],
timestamp=IsNow(tz=UTC),
instructions='Just call tools without asking for confirmation.',
run_id=IsStr(),
conversation_id=IsStr(),
),
ModelResponse(
parts=[
TextPart(content='The file `.env` has been deleted and `test.txt` has been created successfully.')
],
usage=RequestUsage(
input_tokens=133,
output_tokens=19,
output_reasoning_tokens=0,
details={
'accepted_prediction_tokens': 0,
'audio_tokens': 0,
'reasoning_tokens': 0,
'rejected_prediction_tokens': 0,
},
cost=Decimal('0.0005225'),
),
model_name='gpt-4o-2024-08-06',
timestamp=IsDatetime(),
provider_name='openai',
provider_url='https://api.openai.com/v1/',
provider_details={
'finish_reason': 'stop',
'service_tier': 'default',
'timestamp': IsDatetime(),
},
provider_response_id=IsStr(),
finish_reason='stop',
run_id=IsStr(),
conversation_id=IsStr(),
),
]
)
# Test model retry
model_retry_agent = Agent(model, name='model_retry_agent')
@model_retry_agent.tool_plain
@DBOS.step()
def get_weather_in_city(city: str) -> str:
if city != 'Mexico City':
raise ModelRetry('Did you mean Mexico City?')
return 'sunny'
model_retry_dbos_agent = DBOSAgent(model_retry_agent) # pyright: ignore[reportDeprecated]
async def test_dbos_agent_with_model_retry(allow_model_requests: None, dbos: DBOS):
result = await model_retry_dbos_agent.run('What is the weather in CDMX?')
assert result.output == snapshot('The weather in Mexico City is currently sunny.')
assert result.all_messages() == snapshot(
[
ModelRequest(
parts=[
UserPromptPart(
content='What is the weather in CDMX?',
timestamp=IsDatetime(),
)
],
timestamp=IsNow(tz=UTC),
run_id=IsStr(),
conversation_id=IsStr(),
),
ModelResponse(
parts=[
ToolCallPart(
tool_name='get_weather_in_city',
args='{"city":"CDMX"}',
tool_call_id=IsStr(),
)
],
usage=RequestUsage(
input_tokens=47,
output_tokens=17,
output_reasoning_tokens=0,
details={
'accepted_prediction_tokens': 0,
'audio_tokens': 0,
'reasoning_tokens': 0,
'rejected_prediction_tokens': 0,
},
cost=Decimal('0.0002875'),
),
model_name='gpt-4o-2024-08-06',
timestamp=IsDatetime(),
provider_name='openai',
provider_url='https://api.openai.com/v1/',
provider_details={
'finish_reason': 'tool_calls',
'service_tier': 'default',
'timestamp': IsDatetime(),
},
provider_response_id=IsStr(),
finish_reason='tool_call',
run_id=IsStr(),
conversation_id=IsStr(),
),
ModelRequest(
parts=[
RetryPromptPart(
content='Did you mean Mexico City?',
tool_name='get_weather_in_city',
tool_call_id=IsStr(),
timestamp=IsDatetime(),
)
],
timestamp=IsNow(tz=UTC),
run_id=IsStr(),
conversation_id=IsStr(),
),
ModelResponse(
parts=[
ToolCallPart(
tool_name='get_weather_in_city',
args='{"city":"Mexico City"}',
tool_call_id=IsStr(),
)
],
usage=RequestUsage(
input_tokens=87,
output_tokens=17,
output_reasoning_tokens=0,
details={
'accepted_prediction_tokens': 0,
'audio_tokens': 0,
'reasoning_tokens': 0,
'rejected_prediction_tokens': 0,
},
cost=Decimal('0.0003875'),
),
model_name='gpt-4o-2024-08-06',
timestamp=IsDatetime(),
provider_name='openai',
provider_url='https://api.openai.com/v1/',
provider_details={
'finish_reason': 'tool_calls',
'service_tier': 'default',
'timestamp': IsDatetime(),
},
provider_response_id=IsStr(),
finish_reason='tool_call',
run_id=IsStr(),
conversation_id=IsStr(),
),
ModelRequest(
parts=[
ToolReturnPart(
tool_name='get_weather_in_city',
content='sunny',
tool_call_id=IsStr(),
timestamp=IsDatetime(),
)
],
timestamp=IsNow(tz=UTC),
run_id=IsStr(),
conversation_id=IsStr(),
),
ModelResponse(
parts=[TextPart(content='The weather in Mexico City is currently sunny.')],
usage=RequestUsage(
input_tokens=116,
output_tokens=10,
output_reasoning_tokens=0,
details={
'accepted_prediction_tokens': 0,
'audio_tokens': 0,
'reasoning_tokens': 0,
'rejected_prediction_tokens': 0,
},
cost=Decimal('0.00039'),
),
model_name='gpt-4o-2024-08-06',
timestamp=IsDatetime(),
provider_name='openai',
provider_url='https://api.openai.com/v1/',
provider_details={
'finish_reason': 'stop',
'service_tier': 'default',
'timestamp': IsDatetime(),
},
provider_response_id=IsStr(),
finish_reason='stop',
run_id=IsStr(),
conversation_id=IsStr(),
),
]
)
tool_failed_agent = Agent(TestModel(call_tools=['failing_tool']), name='tool_failed_agent')
@tool_failed_agent.tool_plain
@DBOS.step()
def failing_tool() -> str:
raise ToolFailed('Disk full')
tool_failed_dbos_agent = DBOSAgent(tool_failed_agent) # pyright: ignore[reportDeprecated]
async def test_dbos_agent_with_tool_failed(dbos: DBOS):
result = await tool_failed_dbos_agent.run('Call the failing tool')
tool_returns = [
(part.tool_name, part.content, part.outcome)
for message in result.all_messages()
for part in message.parts
if isinstance(part, ToolReturnPart)
]
assert tool_returns == [('failing_tool', 'Disk full', 'failed')]
class CustomModelSettings(ModelSettings, total=False):
custom_setting: str
def return_settings(messages: list[ModelMessage], agent_info: AgentInfo) -> ModelResponse:
return ModelResponse(parts=[TextPart(str(agent_info.model_settings))])
model_settings = CustomModelSettings(max_tokens=123, custom_setting='custom_value')
return_settings_model = FunctionModel(return_settings, settings=model_settings)
settings_agent = Agent(return_settings_model, name='settings_agent')
settings_dbos_agent = DBOSAgent(settings_agent) # pyright: ignore[reportDeprecated]
async def test_custom_model_settings(allow_model_requests: None, dbos: DBOS):
result = await settings_dbos_agent.run('Give me those settings')
assert result.output == snapshot("{'max_tokens': 123, 'custom_setting': 'custom_value'}")
def return_mcp_instructions(messages: list[ModelMessage], agent_info: AgentInfo) -> ModelResponse:
return ModelResponse(parts=[TextPart(agent_info.instructions or '')])
_uninit_instructions_toolset = dbosify_mcp_toolset(
MCPToolset(StdioTransport(command='python', args=['-m', 'tests.mcp_server']), include_instructions=True),
step_name_prefix='coverage_test',
step_config={},
)
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_mcp_toolset_get_instructions_falls_back_to_step(dbos: DBOS):
"""When the MCP server isn't initialized locally, DBOS wrapper fetches instructions via a step."""
run_context = RunContext(deps=0, model=TestModel(), usage=RunUsage())
instructions = await _uninit_instructions_toolset.get_instructions(run_context)
assert instructions == InstructionPart(content='Be a helpful assistant.', dynamic=False)
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_mcp_toolset_get_instructions_uses_step_when_server_warm(dbos: DBOS):
"""A warm in-process MCP server must not short-circuit the `get_instructions` step (#5884).
Whether the step runs must not depend on process warmth: skipping it when the wrapped
server already holds instructions locally makes the durable schedule diverge between the
original execution and replay/recovery in a differently-warm process.
Not a VCR test: what's asserted is *which side* serves the instructions (the DBOS step vs
the warm server's local cache) — invisible to a network cassette, since both paths issue
the same MCP traffic.
"""
inner = MCPToolset(
StdioTransport(command='python', args=['-m', 'tests.mcp_server']),
include_instructions=True,
id='warm_instructions_test',
)
wrapper = dbosify_mcp_toolset(inner, step_name_prefix='warm_instructions_test', step_config={})
run_context = RunContext(deps=0, model=TestModel(), usage=RunUsage())
step_calls: list[str] = []
original_operation = wrapper._get_instructions_operation # pyright: ignore[reportPrivateUsage]
assert original_operation is not None
async def recording_operation(
ctx: RunContext[object],
) -> str | InstructionPart | Sequence[str | InstructionPart] | None:
step_calls.append('step')
return await original_operation(ctx)
wrapper._get_instructions_operation = recording_operation # pyright: ignore[reportPrivateUsage]
# Hold the wrapped server entered so it can serve instructions locally (warm process).
async with inner:
assert await inner.get_instructions(run_context) is not None
instructions = await wrapper.get_instructions(run_context)
assert instructions == InstructionPart(content='Be a helpful assistant.', dynamic=False)
assert step_calls == ['step']
# Exercises the `DBOSMCPToolset` wrapper's `get_instructions` step path with a real `MCPToolset`.
mcptoolset_instructions_agent = Agent(
FunctionModel(return_mcp_instructions),
name='mcptoolset_instructions_agent',
toolsets=[
MCPToolset(
StdioTransport(command='python', args=['-m', 'tests.mcp_server']),
include_instructions=True,
id='mcp',
)
],
)
mcptoolset_instructions_dbos_agent = DBOSAgent(mcptoolset_instructions_agent) # pyright: ignore[reportDeprecated]
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_mcptoolset_instructions_propagate(dbos: DBOS):
"""`MCPToolset` instructions propagate through the `DBOSMCPToolset` wrapper."""
result = await mcptoolset_instructions_dbos_agent.run('Use MCP instructions')
assert result.output == snapshot('Be a helpful assistant.')
def test_dbosify_mcptoolset_dispatches_to_dbosmcptoolset():
"""`DBOSAgent` wraps `MCPToolset` in `DBOSMCPToolset`."""
from pydantic_ai.durable_exec.dbos._mcp_toolset import DBOSMCPToolset
toolset = MCPToolset('https://example.com/mcp', id='test_dispatch')
agent = Agent(model=model, name='dispatch_agent', toolsets=[toolset])
dbos_agent = DBOSAgent(agent) # pyright: ignore[reportDeprecated]
wrapped = next(ts for ts in dbos_agent._toolsets if isinstance(ts, DBOSMCPToolset)) # pyright: ignore[reportPrivateUsage]
assert wrapped.wrapped is toolset
async def test_dbos_mcptoolset_returns_cached_tool_defs(dbos: DBOS):
"""When the run's tool-defs cache is populated, `DBOSMCPToolset.get_tools` returns from it without invoking the step."""
from pydantic_ai.durable_exec.dbos._mcp_toolset import dbosify_mcp_toolset
inner = MCPToolset('https://example.com/mcp', id='cache_return_test')
wrapper = dbosify_mcp_toolset(inner, step_name_prefix='cache_return_test', step_config={})
run_context = RunContext(deps=None, model=TestModel(), usage=RunUsage(), max_retries=5)
run_context._mcp_tool_defs_cache['cache_return_test'] = { # pyright: ignore[reportPrivateUsage]
'foo': ToolDefinition(name='foo', parameters_json_schema={'type': 'object'}),
}
tools = await wrapper.get_tools(run_context)
assert list(tools.keys()) == ['foo']
# Returned ToolsetTool wraps the cached `ToolDefinition` via `tool_for_tool_def` on the wrapped MCPToolset,
# inheriting the agent-level retry count from the run context (rather than a hard-coded default).
assert tools['foo'].tool_def.name == 'foo'
assert tools['foo'].max_retries == 5
_mcp_task_dbos_agent = DBOSAgent( # pyright: ignore[reportDeprecated]
Agent(
TestModel(call_tools=['required_task_tool', 'optional_task_tool']),
name='mcp_task_dbos_agent',
toolsets=[
MCPToolset(
StdioTransport(command='python', args=['-m', 'tests.mcp_task_server']),
id='mcp_tasks',
init_timeout=20,
prefer_tasks=False,
)
],
)
)
@pytest.mark.subprocess(reason='connects to `tests.mcp_task_server` over stdio')
async def test_dbos_mcptoolset_preserves_task_routing(dbos: DBOS):
"""Effective task routing in `ToolDefinition.metadata` survives DBOS steps."""
result = await _mcp_task_dbos_agent.run('Call both tools')
assert result.output == '{"required_task_tool":"required_completed","optional_task_tool":"optional_sync"}'
def _call_mcp_then_finish(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
"""Two model steps: call an MCP tool on the first request, return text on the second.
Two model requests means `get_tools` runs twice on the MCP toolset within one run, so the
run-scoped cache (and whether it schedules a step each time) is exercised.
"""
tool_returned = any(isinstance(part, ToolReturnPart) for message in messages for part in message.parts)
if tool_returned:
return ModelResponse(parts=[TextPart('done')])
return ModelResponse(parts=[ToolCallPart('get_weather_forecast', {'location': 'Mexico City'})])
mcp_replay_agent = Agent(
FunctionModel(_call_mcp_then_finish),
name='mcp_replay_agent',
toolsets=[MCPToolset(StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='mcp', init_timeout=20)],
)
mcp_replay_dbos_agent = DBOSAgent(mcp_replay_agent) # pyright: ignore[reportDeprecated]
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_mcp_get_tools_recorded_independently_per_run(allow_model_requests: None, dbos: DBOS):
"""#5875 regression: DBOS `get_tools` step scheduling must depend only on the workflow's own history.
The run-scoped tool-defs cache collapses the per-request `get_tools` calls to a single recorded
step within a run (the #4331 win). Crucially the cache lives on the run, not on the process-shared
wrapper instance: two workflows executed back-to-back in the same worker process each record their
own `get_tools` step at the front. With the old instance-level cache, run 2 would read run 1's warm
cache and record ZERO `get_tools` steps — a different recorded step sequence than a cold run, which
is what breaks replay determinism on recovery (`DBOSUnexpectedStepError`).
"""
get_tools_step = 'mcp_replay_agent__mcp_server__mcp.get_tools'
async def run_and_list_steps() -> list[str]:
wfid = str(uuid.uuid4())
with SetWorkflowID(wfid):
result = await mcp_replay_dbos_agent.run('hello')
assert result.output == 'done'
return [step['function_name'] for step in await dbos.list_workflow_steps_async(wfid)]
run1_steps = await run_and_list_steps()
run2_steps = await run_and_list_steps()
# Within a run, the run-scoped cache collapses both requests' `get_tools` calls to one step, at the front.
assert run1_steps.count(get_tools_step) == 1
assert run1_steps[0] == get_tools_step
# Run 2 records `get_tools` independently — it does NOT inherit run 1's warm process cache (the #5875 fix).
assert run2_steps.count(get_tools_step) == 1
assert run2_steps[0] == get_tools_step
def _always_call_erroring_mcp_tool(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
"""Keep calling the MCP tool that always errors, so the agent's tool-retry budget is what stops the run."""
return ModelResponse(parts=[ToolCallPart('get_error', {})])
mcp_retry_budget_agent = Agent(
FunctionModel(_always_call_erroring_mcp_tool),
name='mcp_retry_budget_agent',
retries=3,
toolsets=[
MCPToolset(
StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='retry_budget_mcp', init_timeout=20
)
],
)
mcp_retry_budget_dbos_agent = DBOSAgent(mcp_retry_budget_agent) # pyright: ignore[reportDeprecated]
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_mcp_tool_inherits_agent_retries(allow_model_requests: None, dbos: DBOS):
"""#5180 regression: a durably-wrapped MCP tool enforces the agent's tool-retry budget, not a hard-coded 1.
The durable wrapper resolves tools inside a step and keeps only the serializable `ToolDefinition`,
rebuilding each `ToolsetTool` on the workflow side via `MCPToolset.tool_for_tool_def`. When that
rebuild ignored the run context, `Agent(retries=3)` was silently enforced as 1.
"""
with pytest.raises(UnexpectedModelBehavior, match=r"Tool 'get_error' exceeded max retries count of 3"):
await mcp_retry_budget_dbos_agent.run('hello')
async def test_dbos_mcp_toolset_get_instructions_uses_local_when_initialized(dbos: DBOS):
"""When the wrapped MCP toolset is already initialized, the DBOS wrapper short-circuits and returns the local instructions."""
run_context = RunContext(deps=0, model=TestModel(), usage=RunUsage())
# Entering the wrapped toolset populates `_instructions` so the local fast-path returns the part directly.
async with _uninit_instructions_toolset.wrapped:
instructions = await _uninit_instructions_toolset.get_instructions(run_context)
assert instructions == InstructionPart(content='Be a helpful assistant.', dynamic=False)
def test_dbos_mcp_wrapper_visit_and_replace():
"""DBOS MCP wrapper toolsets should not be replaced by visit_and_replace."""
toolsets = mcptoolset_instructions_dbos_agent._toolsets # pyright: ignore[reportPrivateUsage]
dbos_mcp_toolsets = [ts for ts in toolsets if isinstance(ts, DBOSMCPToolset)]
assert len(dbos_mcp_toolsets) >= 1
dbos_mcp_toolset = dbos_mcp_toolsets[0]
# visit_and_replace should return self for DBOS wrappers
result = dbos_mcp_toolset.visit_and_replace(lambda t: FunctionToolset(id='replaced'))
assert result is dbos_mcp_toolset
# ==========================================
# DBOSDurability capability tests
# ==========================================
def _durability_model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
"""Simple model function for durability tests."""
# The first message carries the prompt and its first part is the `UserPromptPart`, so none of these branch.
for msg in reversed(messages): # pragma: no branch
for part in msg.parts: # pragma: no branch
if isinstance(part, UserPromptPart): # pragma: no branch
return ModelResponse(parts=[TextPart(content=f'Echo: {part.content}')])
return ModelResponse(parts=[TextPart(content='no prompt')]) # pragma: no cover
_durability_fn_model = FunctionModel(_durability_model_fn)
# DBOSDurability must be created after DBOS.launch() (in fixture), but since the module-level
# agents are created at import time before DBOS is initialized, we use a fixture-based approach.
async def test_dbos_durability_simple_agent(dbos: DBOS) -> None:
"""DBOSDurability routes model requests through DBOS steps."""
agent = Agent(_durability_fn_model, name='durability_simple', capabilities=[DBOSDurability()])
@DBOS.workflow()
async def run_durable_agent() -> str:
result = await agent.run('Hello DBOS')
return result.output
output = await run_durable_agent()
assert output == 'Echo: Hello DBOS'
async def test_dbos_durability_model_error_reaches_workflow_with_its_type(dbos: DBOS) -> None:
"""A model step's `ModelHTTPError` reaches workflow code as itself, so error hooks can match on it."""
def overloaded(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
raise ModelHTTPError(503, 'overloaded-model', body={'error': 'overloaded'}, headers={'Retry-After': '7'})
async def describe(
ctx: RunContext[None], *, request_context: ModelRequestContext, error: Exception
) -> ModelResponse:
assert isinstance(error, ModelHTTPError)
return ModelResponse(parts=[TextPart(f'{type(error).__name__} {error.status_code} {error.retry_after}')])
agent = Agent(
FunctionModel(overloaded),
name='durability_model_error',
deps_type=type(None),
capabilities=[Hooks[None](model_request_error=describe), DBOSDurability()],
)
@DBOS.workflow()
async def run_durable_agent() -> str:
return (await agent.run('Hello DBOS')).output
assert await run_durable_agent() == 'ModelHTTPError 503 7.0'
async def test_dbos_durability_rejects_a_live_workspace_workflow_input(dbos: DBOS, tmp_path: Path) -> None:
"""A live backend passed into a workflow has no durable steps behind it, so the run refuses it.
The same agent still takes the backend outside a workflow, where the capability is transparent.
The supported shape inside a workflow is a workspace capability attached at construction time;
see `test_dbos_workspace.py`.
"""
agent = Agent(
TestModel(call_tools=['round_trip_workspace']),
name='dbos_durability_workspace',
deps_type=type(None),
tools=[round_trip_workspace],
capabilities=[DBOSDurability()],
)
@DBOS.workflow()
async def run_durable_agent(workspace: LocalWorkspaceBackend) -> AgentRunResult[str]:
return await agent.run('Use the workspace.', workspace=workspace)
with pytest.raises(UserError, match='no capability supplied workspaces when the agent was constructed'):
await run_durable_agent(LocalWorkspaceBackend(tmp_path))
result = await agent.run('Use the workspace.', workspace=LocalWorkspaceBackend(tmp_path))
assert result.output == '{"round_trip_workspace":"workspace content"}'
async def test_dbos_durability_rejects_cancellation_token_in_workflow(dbos: DBOS) -> None:
"""A same-process `cancellation_token` can't cross the durable boundary, so it's rejected inside
a workflow — but the same durable-capable agent still accepts one when run outside a workflow."""
agent = Agent(_durability_fn_model, name='durability_cancel_token', capabilities=[DBOSDurability()])
@DBOS.workflow()
async def run_durable_agent() -> None:
await agent.run('Hello', cancellation_token=CancellationToken())
with pytest.raises(UserError, match='`cancellation_token` cannot be used with DBOS durable execution'):
await run_durable_agent()
# Outside a workflow the capability is transparent, so the token works like a normal run.
result = await agent.run('Hello', cancellation_token=CancellationToken())
assert result.output == 'Echo: Hello'
def _cancel_own_workflow_model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
if len(messages) == 1:
return ModelResponse(parts=[ToolCallPart('cancel_own_workflow')])
return ModelResponse(parts=[TextPart('never reached')]) # pragma: no cover
async def cancel_own_workflow() -> str:
workflow_id = DBOS.workflow_id
assert workflow_id is not None
await DBOS.cancel_workflow_async(workflow_id)
return 'workflow cancelled'
async def test_dbos_durability_workflow_cancellation_carries_run_state(dbos: DBOS) -> None:
"""Cancelling a DBOS workflow aborts the run with a DBOS `BaseException` rather than a `CancelledError`
(`DBOSWorkflowCancelledError`, or `DBOSWorkflowConflictIDError` on dbos 3.2+). The run still
attaches its state to it, so `RunCancelled.from_cancellation()` recovers the history inside the workflow."""
agent = Agent(
FunctionModel(_cancel_own_workflow_model_fn),
name='durability_workflow_cancel',
tools=[cancel_own_workflow],
capabilities=[DBOSDurability()],
)
recovered: list[RunCancelled | None] = []
@DBOS.workflow()
async def run_durable_agent() -> None:
try:
await agent.run('Hello')
except BaseException as exc:
recovered.append(RunCancelled.from_cancellation(exc))
raise
with pytest.raises(DBOSAwaitedWorkflowCancelledError):
await run_durable_agent()
[cancelled] = recovered
assert cancelled is not None
tool_parts = [
(type(part).__name__, part.tool_name)
for message in cancelled.all_messages()
for part in message.parts
if isinstance(part, (ToolCallPart, ToolReturnPart))
]
assert tool_parts[0] == ('ToolCallPart', 'cancel_own_workflow')
async def test_dbos_durability_registers_legacy_workflows_opt_in(dbos: DBOS) -> None:
agent = Agent(
_durability_fn_model,
name='legacy_workflow_compat',
capabilities=[DBOSDurability(register_legacy_workflows=True)],
)
durability = DBOSDurability.from_agent(agent)
assert durability is not None
assert durability._legacy_run_workflow is not None # pyright: ignore[reportPrivateUsage]
assert durability._legacy_run_sync_workflow is not None # pyright: ignore[reportPrivateUsage]
result = await durability._legacy_run_workflow('legacy') # pyright: ignore[reportPrivateUsage]
assert result.output == 'Echo: legacy'
without_flag = Agent(_durability_fn_model, name='no_legacy_workflows', capabilities=[DBOSDurability()])
without_flag_durability = DBOSDurability.from_agent(without_flag)
assert without_flag_durability is not None
assert without_flag_durability._legacy_run_workflow is None # pyright: ignore[reportPrivateUsage]
assert without_flag_durability._legacy_run_sync_workflow is None # pyright: ignore[reportPrivateUsage]
async def test_dbos_durability_accepts_legacy_stream_step_shape(dbos: DBOS) -> None:
response = ModelResponse(parts=[TextPart(content='legacy stream')], model_name='legacy')
agent = Agent(TestModel(), name='legacy_stream_shape', capabilities=[DBOSDurability()])
durability = DBOSDurability.from_agent(agent)
assert durability is not None
result = await durability._load_streamed_activity_result( # pyright: ignore[reportPrivateUsage]
response, ModelRequestParameters()
)
assert result.response == response
assert result.events
async def test_dbos_operation_backend_registers_novel_id_generically(dbos: DBOS) -> None:
from pydantic_ai.durable_exec._operation import (
DurableOperation,
IdentityParameterTransport,
NoCacheIdentity,
TypedResultCodec,
)
from pydantic_ai.durable_exec.dbos._operation_backend import DBOSOperationBackend, DBOSOperationConfig
@dataclass(frozen=True)
class NovelOperationId:
name: str
@dataclass(frozen=True)
class Params:
value: str
async def handler(params: Params) -> str:
return f'handled:{params.value}'
operation = DurableOperation(
operation_id=cast(Any, NovelOperationId('novel')),
handler=handler,
parameter_transport=IdentityParameterTransport[Params](),
cache_identity=NoCacheIdentity[Params](),
result_codec=TypedResultCodec[str](str, mode='identity'),
config_role='capability',
)
backend = DBOSOperationBackend(agent_name='novel', config=DBOSOperationConfig(model={}, event={}, tool={}))
bound, registrations = backend.register(operation, name='novel.generic', config={})
assert await bound(Params('input')) == 'handled:input'
assert registrations == (cast(Any, bound).step,)
# Module-level like real wrapper-era handlers: `DBOSAgent.run` recorded the handler as a
# workflow input, so it had to be picklable by reference.
_legacy_handler_events: list[tuple[AgentStreamEvent, bool]] = []
async def _legacy_workflow_event_handler(ctx: RunContext[Any], stream: AsyncIterable[AgentStreamEvent]) -> None:
async for event in stream:
_legacy_handler_events.append((event, DBOS.step_id is not None))
def test_dbos_durability_legacy_run_sync_workflow(dbos: DBOS) -> None:
"""The opt-in legacy `run_sync` workflow executes the agent like `DBOSAgent.run_sync` did,
including honoring a recorded per-run `event_stream_handler` input."""
_legacy_handler_events.clear()
agent = Agent(
TestModel(custom_output_text='legacy sync'),
name='legacy_workflow_sync_compat',
capabilities=[DBOSDurability(register_legacy_workflows=True)],
)
durability = DBOSDurability.from_agent(agent)
assert durability is not None
result = durability._legacy_run_sync_workflow('legacy', event_stream_handler=_legacy_workflow_event_handler) # pyright: ignore[reportPrivateUsage]
assert result.output == 'legacy sync'
assert _legacy_handler_events
async def test_dbos_durability_legacy_workflow_matches_wrapper_step_sequence(dbos: DBOS) -> None:
"""A legacy run delivers handler events exactly the way `DBOSAgent` did, so recovery replays.
Wrapper-era recordings contain only model and MCP steps: model events reached the handler
live inside the `__model.request_stream` step, and graph-level events through a direct
workflow-level call that consumed no step. Dispatching graph events through the capability's
`__event_stream_handler` step would insert step ids the recording doesn't have and fail
recovery with `DBOSUnexpectedStepError`.
"""
_legacy_handler_events.clear()
async def handled_tool() -> str:
return 'handled'
agent = Agent(
TestModel(),
name='legacy_workflow_handler',
tools=[handled_tool],
capabilities=[DBOSDurability(register_legacy_workflows=True)],
)
durability = DBOSDurability.from_agent(agent)
assert durability is not None
wfid = str(uuid.uuid4())
with SetWorkflowID(wfid):
result = await durability._legacy_run_workflow('Hello', event_stream_handler=_legacy_workflow_event_handler) # pyright: ignore[reportPrivateUsage]
assert result.output
events = [event for event, _ in _legacy_handler_events]
assert events
assert sum(isinstance(event, FunctionToolCallEvent) for event in events) == 1
# Model events arrive live inside the model-request step; graph events at workflow level.
assert all(in_step for event, in_step in _legacy_handler_events if isinstance(event, PartStartEvent))
assert not any(in_step for event, in_step in _legacy_handler_events if isinstance(event, FunctionToolCallEvent))
steps = await dbos.list_workflow_steps_async(wfid)
step_names = [step['function_name'] for step in steps]
assert 'legacy_workflow_handler__model.request_stream' in step_names
assert 'legacy_workflow_handler__event_stream_handler' not in step_names
async def test_unwrap_recorded_tool_call_result_handles_both_generations() -> None:
"""Recovery may replay step outputs recorded before control-flow wrapping (raw) or after (wrapped).
A unit test because a real recovery would need to kill and restart the process mid-workflow;
the step-recorded value is exactly what this helper receives on replay.
"""
async def raise_model_retry() -> str:
raise ModelRetry('again')
wrapped_result = await wrap_tool_call_result(raise_model_retry())
with pytest.raises(ModelRetry, match='again'):
unwrap_recorded_tool_call_result(wrapped_result)
assert unwrap_recorded_tool_call_result('raw recorded output') == 'raw recorded output'
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_mcp_model_retry_crosses_step_without_engine_retries(
dbos: DBOS, monkeypatch: pytest.MonkeyPatch
) -> None:
"""`ModelRetry` from an MCP tool crosses the step as a value; DBOS must not re-execute the call.
Previously the raw exception failed the step, so `retries_allowed` re-ran the side-effecting
MCP call and surfaced `DBOSMaxStepRetriesExceeded` instead of the agent's retry-prompt loop.
"""
calls = 0
mcp_toolset = MCPToolset(StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='retry_mcp')
async def raise_model_retry(
tool_name: str, tool_args: dict[str, Any], ctx: RunContext[None], tool: ToolsetTool[None]
) -> Any:
nonlocal calls
calls += 1
raise ModelRetry('try again')
monkeypatch.setattr(mcp_toolset, 'call_tool', raise_model_retry)
durable = dbosify_mcp_toolset(
mcp_toolset,
step_name_prefix='retry_mcp_agent',
step_config=StepConfig(retries_allowed=True, max_attempts=3),
)
run_context = RunContext(deps=None, model=TestModel(), usage=RunUsage())
tool = ToolsetTool(
toolset=durable,
tool_def=ToolDefinition(name='boom'),
max_retries=1,
args_validator=TOOL_SCHEMA_VALIDATOR,
)
@DBOS.workflow()
async def run_workflow() -> int:
with pytest.raises(ModelRetry, match='try again'):
await durable.call_tool('boom', {}, run_context, tool)
return calls
assert await run_workflow() == 1
async def test_dbos_dynamic_tool_model_retry_crosses_step_without_engine_retries(dbos: DBOS) -> None:
"""`ModelRetry` from a `DynamicToolset` tool crosses the step as a value, like MCP and function tools."""
calls = 0
async def retry_once() -> str:
nonlocal calls
calls += 1
if calls == 1:
raise ModelRetry('try again')
return 'done'
agent = Agent(
TestModel(),
name='retry_dynamic_agent',
toolsets=[DynamicToolset(lambda ctx: FunctionToolset([retry_once]), id='retry_dynamic')],
capabilities=[DBOSDurability(mcp_step_config=StepConfig(retries_allowed=True, max_attempts=3))],
)
@DBOS.workflow()
async def run_workflow() -> str:
return (await agent.run('run')).output
await run_workflow()
# The step recorded the `ModelRetry` as a value, so DBOS never re-ran it: exactly one
# retry, driven by the model, not `max_attempts=3`.
assert calls == 2
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_mcp_step_rejects_enqueue_in_workflow(dbos: DBOS, monkeypatch: pytest.MonkeyPatch) -> None:
"""The MCP step path guards enqueue too: a `process_tool_call=` hook receives the run context."""
mcp_toolset = MCPToolset(StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='enqueue_mcp')
async def enqueue_call_tool(
tool_name: str, tool_args: dict[str, Any], ctx: RunContext[None], tool: ToolsetTool[None]
) -> Any:
ctx.enqueue('later')
return 'done'
monkeypatch.setattr(mcp_toolset, 'call_tool', enqueue_call_tool)
durable = dbosify_mcp_toolset(mcp_toolset, step_name_prefix='enqueue_mcp_agent', step_config={})
run_context = RunContext(deps=None, model=TestModel(), usage=RunUsage())
tool = ToolsetTool(
toolset=durable,
tool_def=ToolDefinition(name='hook'),
max_retries=1,
args_validator=TOOL_SCHEMA_VALIDATOR,
)
@DBOS.workflow()
async def run_workflow() -> None:
await durable.call_tool('hook', {}, run_context, tool)
with pytest.raises(UserError, match='enqueued messages would be dropped'):
await run_workflow()
# Outside a workflow the step degrades to a plain call and enqueueing keeps working.
outside_context = RunContext(deps=None, model=TestModel(), usage=RunUsage(), pending_messages=[])
assert await durable.call_tool('hook', {}, outside_context, tool) == 'done'
assert len(outside_context.pending_messages or []) == 1
async def test_dbos_dynamic_tool_rejects_enqueue_in_workflow(dbos: DBOS) -> None:
"""`ctx.enqueue()` inside a step-wrapped dynamic tool raises instead of silently dropping.
Recovery replays the recorded step output without re-executing the tool, so in-step
enqueued messages would be lost. Outside a workflow the step degrades to a plain call
and enqueueing keeps working.
"""
async def enqueue(ctx: RunContext[object]) -> str:
ctx.enqueue('later')
return 'done'
agent = Agent(
TestModel(),
deps_type=object,
name='dbos_dynamic_enqueue',
toolsets=[DynamicToolset(lambda ctx: FunctionToolset([enqueue]), id='enqueue_dynamic')],
capabilities=[DBOSDurability()],
)
@DBOS.workflow()
async def run_workflow() -> None:
await agent.run('run')
with pytest.raises(UserError, match='enqueued messages would be dropped'):
await run_workflow()
await agent.run('run')
@dataclass(kw_only=True)
class DBOSCheckpointEvent(CapabilityEvent, namespace='dbos_test', name='checkpoint'):
label: str
@dataclass(kw_only=True)
class DBOSStepProgressEvent(CustomEvent, name='dbos_step_progress'):
label: str
async def test_dbos_workflow_level_emit_reaches_durable_handler(dbos: DBOS) -> None:
"""A capability event emitted at workflow level reaches the durability handler's own step."""
seen: list[str] = []
async def handler(ctx: RunContext[object], stream: AsyncIterable[AgentStreamEvent]) -> None:
async for event in stream:
if isinstance(event, DBOSCheckpointEvent):
seen.append(event.label)
class EmittingCapability(AbstractCapability[object]):
id = 'dbos_emitter'
async def before_run(self, ctx: RunContext[object]) -> None:
await ctx.emit(DBOSCheckpointEvent(label='start'))
agent = Agent(
TestModel(),
deps_type=object,
name='dbos_workflow_emit',
capabilities=[EmittingCapability(), DBOSDurability[object](event_stream_handler=handler)],
)
@DBOS.workflow()
async def run_workflow() -> None:
await agent.run('run')
await run_workflow()
assert seen == ['start']
async def test_dbos_step_emit_is_allowed(dbos: DBOS) -> None:
"""`ctx.emit()` from inside a step is allowed, unlike `ctx.enqueue()`.
Enqueueing is rejected because a replayed step drops the message and changes what the model
sees; an emitted event only notifies observers, so it's a side effect of running the step, like
a log line. That costs re-delivery on recovery, which replays the step's recorded output without
re-running the body -- pinned for the equivalent Prefect cache hit in
`test_prefect_task_wrapped_tool_emit_is_not_replayed`, since a real DBOS recovery would need the
process killed mid-workflow.
"""
observed: list[str] = []
@DBOS.step()
async def record(label: str) -> str:
ctx = get_current_run_context()
assert ctx is not None
await ctx.emit(DBOSStepProgressEvent(label=label))
return 'done'
async def emitter() -> str:
return await record('one')
async def observe(ctx: RunContext[object], stream: AsyncIterable[AgentStreamEvent]) -> None:
async for event in stream:
if isinstance(event, DBOSStepProgressEvent):
observed.append(event.label)
agent = Agent(
TestModel(),
deps_type=object,
name='dbos_step_emit',
tools=[emitter],
capabilities=[ProcessEventStream(observe), DBOSDurability[object]()],
)
@DBOS.workflow()
async def run_workflow() -> None:
await agent.run('run')
await run_workflow()
assert observed == ['one']
async def test_dbos_non_streaming_model_request_rejects_enqueue(dbos: DBOS) -> None:
enqueued = False
def request_with_enqueue(_: list[ModelMessage], __: AgentInfo) -> ModelResponse:
nonlocal enqueued
ctx = get_current_run_context()
assert ctx is not None
if not enqueued:
# Only the first request enqueues, so the outside-workflow run below terminates.
enqueued = True
ctx.enqueue('later')
return ModelResponse(parts=[TextPart('done')])
agent = Agent(
FunctionModel(request_with_enqueue),
name='dbos_model_enqueue',
capabilities=[DBOSDurability()],
)
@DBOS.workflow()
async def run_agent() -> None:
await agent.run('run')
with pytest.raises(UserError, match='enqueued messages would be dropped'):
await run_agent()
# Outside a workflow the step degrades to an inline call and enqueueing keeps working.
enqueued = False
assert (await agent.run('run')).output == 'done'
async def test_dbos_step_wrapped_tool_rejects_cancel_in_workflow(dbos: DBOS) -> None:
"""`ctx.cancel()` inside a step-wrapped (dynamic) tool raises instead of replay-diverging.
Recovery replays the recorded step output without re-executing the tool, so an in-step
cancellation would silently not happen again. Outside a workflow the step degrades to a
plain call and cancellation keeps working.
"""
async def cancel(ctx: RunContext[object]) -> str:
ctx.cancel()
# `cancel()` returns; the cancellation lands at the next await point, so this
# tool completes normally first and its (discarded) result is recorded.
return 'completed before the cancellation took effect'
agent = Agent(
TestModel(),
deps_type=object,
name='dbos_cancel_dynamic',
toolsets=[DynamicToolset(lambda ctx: FunctionToolset([cancel]), id='cancel_dynamic')],
capabilities=[DBOSDurability()],
)
@DBOS.workflow()
async def run_workflow() -> None:
await agent.run('run')
with pytest.raises(UserError, match='cancellation would silently not happen again'):
await run_workflow()
with pytest.raises(RunCancelled):
await agent.run('run')
async def test_dbos_plain_tool_cancel_works_in_workflow(dbos: DBOS) -> None:
"""A plain function tool is NOT step-wrapped under DBOS: it runs at workflow level, where
`cancel()` is live and replay-consistent (workflow-level code re-executes on recovery),
so cancellation works and surfaces as an ordinary `RunCancelled` application outcome."""
async def cancel(ctx: RunContext[object]) -> str:
ctx.cancel()
# `cancel()` returns; the cancellation lands at the next await point, so this
# tool completes normally first and its (discarded) result is recorded.
return 'completed before the cancellation took effect'
agent = Agent(
TestModel(),
deps_type=object,
name='dbos_cancel_plain',
tools=[cancel],
capabilities=[DBOSDurability()],
)
@DBOS.workflow()
async def run_workflow() -> None:
await agent.run('run')
with pytest.raises(RunCancelled):
await run_workflow()
async def test_dbos_dynamic_get_tools_rejects_enqueue_in_workflow(dbos: DBOS) -> None:
"""The dynamic-toolset discovery step guards enqueue: the user's factory receives the context."""
enqueue_errors: list[str] = []
def factory(ctx: RunContext[object]) -> FunctionToolset[object]:
try:
ctx.enqueue('later')
except UserError as e:
enqueue_errors.append(str(e))
return FunctionToolset([])
agent = Agent(
TestModel(),
deps_type=object,
name='dbos_dynamic_get_tools_enqueue',
toolsets=[DynamicToolset(factory, id='enqueue_factory')],
capabilities=[DBOSDurability()],
)
@DBOS.workflow()
async def run_workflow() -> None:
await agent.run('run')
await run_workflow()
assert enqueue_errors == snapshot(
[
"`ctx.enqueue()` is not supported inside a durable step: the durable runtime replays the step's recorded result without re-running your code, so the enqueued messages would be dropped. Enqueue messages from workflow-level code instead."
]
)
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_mcp_get_tools_and_instructions_reject_enqueue_in_workflow(
dbos: DBOS, monkeypatch: pytest.MonkeyPatch
) -> None:
"""The MCP discovery steps guard enqueue too, like the MCP call step."""
enqueue_errors: list[str] = []
mcp_toolset = MCPToolset(
StdioTransport(command='python', args=['-m', 'tests.mcp_server']),
id='enqueue_mcp_discovery',
include_instructions=True,
)
async def enqueue_get_tools(ctx: RunContext[None]) -> dict[str, ToolsetTool[None]]:
with pytest.raises(UserError) as exc_info:
ctx.enqueue('later')
enqueue_errors.append(str(exc_info.value))
return {}
async def enqueue_get_instructions(ctx: RunContext[None]) -> str:
with pytest.raises(UserError) as exc_info:
ctx.enqueue('later')
enqueue_errors.append(str(exc_info.value))
return ''
monkeypatch.setattr(mcp_toolset, 'get_tools', enqueue_get_tools)
monkeypatch.setattr(mcp_toolset, 'get_instructions', enqueue_get_instructions)
durable = dbosify_mcp_toolset(mcp_toolset, step_name_prefix='enqueue_mcp_discovery_agent', step_config={})
# A live queue, so the raise below can only come from the step's guard.
run_context = RunContext(deps=None, model=TestModel(), usage=RunUsage(), pending_messages=[])
@DBOS.workflow()
async def run_workflow() -> None:
await durable.get_tools(run_context)
await durable.get_instructions(run_context)
await run_workflow()
assert enqueue_errors == snapshot(
[
"`ctx.enqueue()` is not supported inside a durable step: the durable runtime replays the step's recorded result without re-running your code, so the enqueued messages would be dropped. Enqueue messages from workflow-level code instead.",
"`ctx.enqueue()` is not supported inside a durable step: the durable runtime replays the step's recorded result without re-running your code, so the enqueued messages would be dropped. Enqueue messages from workflow-level code instead.",
]
)
assert run_context.pending_messages == []
async def test_dbos_model_request_step_rejects_enqueue(dbos: DBOS) -> None:
"""The non-streaming model-request step guards enqueue like its streaming sibling.
`Model.request` takes no run context, so code inside the step (a custom model, a `models=`
wrapper, a `resolve_model_id` capability rebuilding it) reaches the run through
`get_current_run_context()`. Recovery replays the recorded step output without re-running
it, so an enqueue there would be dropped.
"""
enqueue_errors: list[str] = []
enqueued: list[str | None] = []
class AmbientEnqueueModel(TestModel):
async def request(
self,
messages: list[ModelMessage],
model_settings: ModelSettings | None,
model_request_parameters: ModelRequestParameters,
) -> ModelResponse:
ambient = get_current_run_context()
assert ambient is not None
# Only on the first request of a run: a successful enqueue triggers another request,
# and enqueueing from each of those would never terminate.
if not (enqueue_errors or enqueued):
try:
enqueued.append(ambient.enqueue('later'))
except UserError as e:
enqueue_errors.append(str(e))
return await super().request(messages, model_settings, model_request_parameters)
agent = Agent(AmbientEnqueueModel(), name='dbos_model_request_enqueue', capabilities=[DBOSDurability()])
@DBOS.workflow()
async def run_workflow() -> None:
await agent.run('go')
await run_workflow()
assert enqueued == []
assert enqueue_errors == snapshot(
[
"`ctx.enqueue()` is not supported inside a durable step: the durable runtime replays the step's recorded result without re-running your code, so the enqueued messages would be dropped. Enqueue messages from workflow-level code instead."
]
)
# Outside a workflow the step degrades to a plain call and enqueueing keeps working.
enqueue_errors.clear()
result = await agent.run('go')
assert enqueue_errors == []
assert len(enqueued) == 1
assert result.output == snapshot('success (no tool calls)')
async def test_dbos_cancel_suspended_response_step_rejects_enqueue(dbos: DBOS) -> None:
"""The suspended-response cancellation step guards enqueue too.
The teardown is a provider call inside its own step, so the same replay argument applies.
"""
enqueue_errors: list[str] = []
class AmbientEnqueueContinuationModel(ScriptedContinuationModel):
async def cancel_suspended_response(self, response: ModelResponse) -> None:
ambient = get_current_run_context()
assert ambient is not None
try:
ambient.enqueue('later')
except UserError as e:
enqueue_errors.append(str(e))
await super().cancel_suspended_response(response)
model = AmbientEnqueueContinuationModel(
responses=[
scripted_response(
texts=['still going '],
state='suspended',
provider_response_id='cont1',
input_tokens=10,
output_tokens=5,
),
scripted_response(
texts=['keeps going '],
state='suspended',
provider_response_id='cont2',
input_tokens=100,
output_tokens=50,
),
]
)
agent = Agent(model, name='dbos_cancel_enqueue', capabilities=[DBOSDurability()])
@DBOS.workflow()
async def run_workflow() -> None:
await agent.run('go', usage_limits=UsageLimits(total_tokens_limit=20))
with pytest.raises(UsageLimitExceeded, match='total_tokens_limit'):
await run_workflow()
assert [cancelled.provider_response_id for cancelled in model.cancelled] == ['cont2']
assert enqueue_errors == snapshot(
[
"`ctx.enqueue()` is not supported inside a durable step: the durable runtime replays the step's recorded result without re-running your code, so the enqueued messages would be dropped. Enqueue messages from workflow-level code instead."
]
)
async def test_dbos_durability_parallel_mode_applies_inside_run(dbos: DBOS) -> None:
"""The configured parallel-execution mode is active for the duration of the run."""
from pydantic_ai import tool_manager as _tm
captured: list[str] = []
def _capture_mode_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
captured.append(_tm._parallel_execution_mode_ctx_var.get()) # pyright: ignore[reportPrivateUsage]
return _durability_model_fn(messages, info)
capture_model = FunctionModel(_capture_mode_fn)
agent = Agent(
capture_model,
name='durability_parallel',
capabilities=[DBOSDurability(parallel_execution_mode='sequential')],
)
await agent.run('measure mode')
assert captured == ['sequential']
async def test_dbos_durability_outside_workflow() -> None:
"""DBOSDurability is transparent outside a DBOS workflow."""
agent = Agent(_durability_fn_model, name='durability_outside', capabilities=[DBOSDurability()])
result = await agent.run('Hello outside')
assert result.output == 'Echo: Hello outside'
async def test_dbos_durability_step_verification(dbos: DBOS) -> None:
"""Verify that model requests become DBOS steps."""
agent = Agent(_durability_fn_model, name='durability_steps', capabilities=[DBOSDurability()])
wfid = str(uuid.uuid4())
@DBOS.workflow()
async def run_durable_agent() -> str:
result = await agent.run('verify steps')
return result.output
with SetWorkflowID(wfid):
await run_durable_agent()
steps = await dbos.list_workflow_steps_async(wfid)
step_names = [step['function_name'] for step in steps]
assert 'durability_steps__model.request' in step_names
def test_dbos_durability_requires_agent_name() -> None:
"""DBOSDurability raises UserError when the agent has no name."""
with pytest.raises(UserError, match='unique `name`'):
Agent(_durability_fn_model, capabilities=[DBOSDurability()])
def test_dbos_durability_explicit_name_overrides_agent_name_and_supports_unnamed_agent() -> None:
named_agent = Agent(_durability_fn_model, name='agent-name', capabilities=[DBOSDurability(name='custom')])
bound = DBOSDurability.from_agent(named_agent)
assert bound is not None
assert bound.name == 'custom'
unnamed_agent = Agent(_durability_fn_model, capabilities=[DBOSDurability(name='unnamed-custom')])
unnamed_bound = DBOSDurability.from_agent(unnamed_agent)
assert unnamed_bound is not None
assert unnamed_bound.name == 'unnamed-custom'
def test_dbos_durability_requires_model() -> None:
"""DBOSDurability raises UserError when the agent has no model at all."""
with pytest.raises(UserError, match='needs to have a `model`'):
Agent(name='needs_model', capabilities=[DBOSDurability()])
def _dbos_alt_model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
return ModelResponse(parts=[TextPart(content='alt-response')])
# A module-level function (not a lambda): passing the instance to `agent.run(model=...)`
# makes it part of the DBOS workflow's pickled arguments.
_dbos_alt_model = FunctionModel(_dbos_alt_model_fn, model_name='alt')
_dbos_model_lifecycle_events: list[str] = []
class _DBOSLifecycleModel(LifecycleTrackingModel):
def __init__(self) -> None:
super().__init__(_dbos_model_lifecycle_events, event_prefix='model-')
@asynccontextmanager
async def request_stream(
self,
messages: list[ModelMessage],
model_settings: ModelSettings | None,
model_request_parameters: ModelRequestParameters,
run_context: RunContext[Any] | None = None,
) -> AsyncGenerator[StreamedResponse]:
_dbos_model_lifecycle_events.append('stream-enter')
try:
async with super().request_stream(
messages, model_settings, model_request_parameters, run_context
) as streamed:
yield streamed
finally:
_dbos_model_lifecycle_events.append('stream-exit')
def _resolve_dbos_lifecycle_model(_ctx: ModelResolutionContext[Any], _model_id: str) -> _DBOSLifecycleModel:
return _DBOSLifecycleModel()
async def test_dbos_durability_context_manages_resolved_models(dbos: DBOS) -> None:
"""Resolver-created models stay entered through ordinary and streamed DBOS steps."""
agent = Agent(
'lifecycle',
name='durability_resolved_model_lifecycle',
capabilities=[ResolveModelId(_resolve_dbos_lifecycle_model), DBOSDurability()],
)
@DBOS.workflow()
async def run_agent() -> str:
return (await agent.run('hello')).output
@DBOS.workflow()
async def stream_agent() -> str:
async with agent.run_stream('hello') as result:
return await result.get_output()
_dbos_model_lifecycle_events.clear()
with SetWorkflowID(str(uuid.uuid4())):
assert await run_agent() == 'ok'
assert _dbos_model_lifecycle_events == ['model-enter', 'request', 'model-exit:none']
_dbos_model_lifecycle_events.clear()
with SetWorkflowID(str(uuid.uuid4())):
assert await stream_agent() == 'ok'
assert _dbos_model_lifecycle_events == [
'model-enter',
'stream-enter',
'stream-exit',
'model-exit:none',
]
async def test_dbos_durability_runtime_registered_model(dbos: DBOS) -> None:
"""A model registered in `models=` can be selected at run time, by key or instance.
The `model_id` crosses the step boundary and the step rebuilds the model from the
registry, so the response is produced by the selected model inside the DBOS step.
"""
agent = Agent(
_durability_fn_model,
name='durability_runtime_registered',
capabilities=[DBOSDurability(models={'alt': _dbos_alt_model})],
)
@DBOS.workflow()
async def run_agent() -> tuple[str, str]:
by_key = await agent.run('hello', model='alt')
by_instance = await agent.run('hello', model=_dbos_alt_model)
return by_key.output, by_instance.output
assert await run_agent() == ('alt-response', 'alt-response')
async def test_dbos_durability_per_step_model_selector_uses_selected_model(dbos: DBOS) -> None:
"""A per-step selector's model identity, rather than the run-level model ID, crosses the step boundary."""
durability = DBOSDurability(models={'primary': _durability_fn_model, 'alt': _dbos_alt_model})
agent = Agent(
'primary',
name='durability_per_step_model_selector',
capabilities=[SelectModel(lambda ctx: _dbos_alt_model), durability],
)
@DBOS.workflow()
async def run_agent() -> str:
return (await agent.run('hello')).output
assert await run_agent() == 'alt-response'
async def test_dbos_durability_per_step_model_selector_preserves_alias(dbos: DBOS) -> None:
"""DBOS persists the selected registration alias instead of the resolved model's ID."""
durability = DBOSDurability(models={'primary': _durability_fn_model, 'alt': _dbos_alt_model})
agent = Agent(
'primary',
name='durability_per_step_model_selector_alias',
capabilities=[SelectModel(lambda ctx: 'alt'), durability],
)
@DBOS.workflow()
async def run_agent() -> str:
return (await agent.run('hello')).output
assert await run_agent() == 'alt-response'
async def test_dbos_durability_override_registered_model(dbos: DBOS) -> None:
"""A model set via `override(model=...)` round-trips the step boundary like a per-run `model=`."""
agent = Agent(
_durability_fn_model,
name='durability_override_registered',
capabilities=[DBOSDurability(models={'alt': _dbos_alt_model})],
)
@DBOS.workflow()
async def run_agent() -> str:
with agent.override(model='alt'):
result = await agent.run('hello')
return result.output
assert await run_agent() == 'alt-response'
async def test_dbos_durability_unregistered_model_instance_errors(dbos: DBOS) -> None:
"""An unregistered `Model` instance is rejected in the workflow, before any step runs.
A `Model` can't be serialized into a step, and rebuilding this one from its `model_id` would
build the same model name on the default provider — dropping the tenant's `base_url` and API
key, so the request would silently go to `api.openai.com` with the worker's credentials.
Registering the instance in `models=`, or passing a string a `ResolveModelId` capability builds
inside the step, are the two supported paths.
"""
agent = Agent(_durability_fn_model, name='durability_unregistered_instance', capabilities=[DBOSDurability()])
tenant_model = OpenAIChatModel(
'gpt-5.6-sol', provider=OpenAIProvider(api_key='tenant-key', base_url='https://tenant.example.com/v1')
)
@DBOS.workflow()
async def run_agent() -> None:
await agent.run('hello', model=tenant_model)
with pytest.raises(UserError) as exc_info:
await run_agent()
assert str(exc_info.value) == snapshot(
"The model instance 'openai:gpt-5.6-sol' was not registered with `DBOSDurability`, so it cannot be used inside a workflow. A `Model` instance cannot be serialized across the step boundary, and rebuilding it from its `model_id` would build a different model — the same model name on the provider the worker environment implies — so the request would go to another endpoint with other credentials. Register the instance in `models=` on `DBOSDurability` and reference it by key (or pass the registered instance), or pass a model-name string and build the instance from it with a `ResolveModelId` capability."
)
async def test_dbos_durability_unrebuildable_model_string_errors(dbos: DBOS) -> None:
"""A model-name string no resolver claims inside the step errors helpfully.
The alias resolves outside the step (so the run starts), but the resolver declines inside it —
as a worker whose configuration no longer knows the alias would — and `infer_model` can't build
it. Instead of a bare 'Unknown model' the step points at the escape hatches.
"""
def resolver(ctx: ModelResolutionContext[Any], model_id: str) -> FunctionModel | None:
# The resolved model never runs: the step's re-resolution declines and `infer_model` fails.
if model_id == 'stale-alias' and DBOS.step_id is None:
return FunctionModel(_dbos_alt_model_fn, model_name='stale-alias')
return None
agent = Agent(
_durability_fn_model,
name='durability_unrebuildable_string',
capabilities=[ResolveModelId(resolver), DBOSDurability()],
)
@DBOS.workflow()
async def run_agent() -> None:
await agent.run('hello', model='stale-alias')
with pytest.raises(UserError) as exc_info:
await run_agent()
assert str(exc_info.value) == snapshot(
"The model 'stale-alias' could not be rebuilt on the DBOS worker: it is not a model name `infer_model` can build, and no `resolve_model_id` capability claimed it. Register the instance in `models=` on `DBOSDurability` and reference it by key (or pass the registered instance), or pass a model-name string and build the instance from it with a `ResolveModelId` capability."
)
async def test_dbos_durability_string_default_model(dbos: DBOS) -> None:
"""A string default that `infer_model` can build is registered under its own id as well as `'default'`.
Passing the same string at run time then resolves to the bound default instance via the
registry, instead of building a fresh instance that would round-trip as a runtime override.
"""
agent = Agent('test', name='durability_string_default', capabilities=[DBOSDurability()])
@DBOS.workflow()
async def run_agent() -> tuple[str, str]:
default = await agent.run('hello')
by_string = await agent.run('hello', model='test')
return default.output, by_string.output
assert await run_agent() == ('success (no tool calls)', 'success (no tool calls)')
async def test_dbos_durability_wrapped_default_model(dbos: DBOS) -> None:
"""A `WrapperModel` around the agent's own model still takes the default's fast path.
Models are identity-matched after stripping wrapper layers (e.g. the `InstrumentedModel`
an `Instrumentation` capability wraps around the model before the request runs), so a
wrapped default isn't mistaken for a runtime override that would round-trip through
`infer_model` and fail to rebuild.
"""
agent = Agent(_durability_fn_model, name='durability_wrapped_default', capabilities=[DBOSDurability()])
@DBOS.workflow()
async def run_agent() -> str:
result = await agent.run('hello', model=WrapperModel(_durability_fn_model))
return result.output
assert await run_agent() == 'Echo: hello'
def _dbos_tenant_resolver(ctx: ModelResolutionContext[str], model_id: str) -> FunctionModel | None:
"""Resolve the 'tenant-model' alias to a model built from the run's deps.
Matches the alias exactly: the run's original model-id string (not the resolved
model's `'function:tenant-model'`) is what crosses the durable boundary, so the
worker-side re-resolution sees the same string the caller wrote.
"""
if model_id != 'tenant-model':
return None
tenant = ctx.deps
def fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
return ModelResponse(parts=[TextPart(content=f'tenant:{tenant}')])
return FunctionModel(fn, model_name='tenant-model')
async def test_dbos_durability_resolve_model_id_capability_is_deps_aware(dbos: DBOS) -> None:
"""A deps-aware `ResolveModelId` resolver rebuilds the model with the run's deps inside the step."""
agent = Agent(
_durability_fn_model,
name='durability_tenant',
deps_type=str,
capabilities=[ResolveModelId(_dbos_tenant_resolver), DBOSDurability()],
)
@DBOS.workflow()
async def run_agent() -> list[str]:
outputs = [(await agent.run('hi', model='tenant-model', deps=tenant)).output for tenant in ('acme', 'globex')]
# A string the resolver doesn't recognize defers to the default `infer_model` flow.
outputs.append((await agent.run('hi', model='test', deps='acme')).output)
return outputs
assert await run_agent() == ['tenant:acme', 'tenant:globex', 'success (no tool calls)']
def _dbos_broken_resolver(ctx: ModelResolutionContext[Any], model_id: str) -> FunctionModel | None:
if model_id == 'broken-model':
raise ValueError('resolver exploded')
# Only 'broken-model' flows through this test.
return None # pragma: no cover
async def test_dbos_durability_user_resolver_error_propagates(dbos: DBOS) -> None:
"""An exception raised by a user's resolver propagates unchanged.
Only the `infer_model` backstop gets translated into the could-not-be-rebuilt error
(see the unrebuildable test above); a resolver's own failure must surface as-is —
both from run-setup resolution and from re-resolution inside the step — so its
message isn't misattributed to the durable round-trip.
"""
agent = Agent(
_durability_fn_model,
name='durability_broken_resolver',
capabilities=[ResolveModelId(_dbos_broken_resolver), DBOSDurability()],
)
@DBOS.workflow()
async def run_agent() -> None:
await agent.run('hello', model='broken-model')
with pytest.raises(ValueError, match='resolver exploded'):
await run_agent()
async def test_dbos_durability_alias_default_model(dbos: DBOS) -> None:
"""An agent whose *default* model is an alias only a `ResolveModelId` capability can resolve.
`infer_model` can't build `'tenant-model'`, so binding registers no concrete default;
every request carries the raw alias string across the step boundary and the step
re-resolves it with the run's deps.
"""
agent = Agent(
'tenant-model',
name='durability_alias_default',
deps_type=str,
capabilities=[ResolveModelId(_dbos_tenant_resolver), DBOSDurability()],
)
@DBOS.workflow()
async def run_agent() -> str:
result = await agent.run('hi', deps='acme')
return result.output
assert await run_agent() == 'tenant:acme'
async def test_dbos_durability_per_step_model_selector_alias_resolved_by_capability(dbos: DBOS) -> None:
"""A selector alias that only a `ResolveModelId` capability can resolve crosses the step boundary.
Unlike a `models=`-registered alias, identity lookup can't recover this one: the selector's
string itself must be recorded as the request's model ID, or the step would receive the
resolved model's own ID (`function:tenant-model`), which neither the resolver nor
`infer_model` can rebuild.
"""
agent = Agent(
_durability_fn_model,
name='durability_selector_capability_alias',
deps_type=str,
capabilities=[
SelectModel(lambda ctx: 'tenant-model'),
ResolveModelId(_dbos_tenant_resolver),
DBOSDurability(),
],
)
@DBOS.workflow()
async def run_agent() -> str:
return (await agent.run('hi', deps='acme')).output
assert await run_agent() == 'tenant:acme'
@dataclass
class _WrapRequestModelCapability(AbstractCapability[str]):
"""Swaps a transparent `WrapperModel` around the request's model in `before_model_request`."""
async def before_model_request(
self, ctx: RunContext[str], request_context: ModelRequestContext
) -> ModelRequestContext:
request_context.model = WrapperModel(request_context.model)
return request_context
async def test_dbos_durability_wrapper_swap_keeps_alias_provenance(dbos: DBOS) -> None:
"""A transparent wrapper swapped in by `before_model_request` doesn't invalidate provenance.
Unwrapping both sides shows the request still targets the run's model, so the original
alias string crosses the step boundary and the worker re-resolves it deps-aware.
"""
agent = Agent(
'tenant-model',
name='durability_wrapper_swap_provenance',
deps_type=str,
capabilities=[_WrapRequestModelCapability(), ResolveModelId(_dbos_tenant_resolver), DBOSDurability()],
)
@DBOS.workflow()
async def run_agent() -> str:
return (await agent.run('hi', deps='acme')).output
assert await run_agent() == 'tenant:acme'
async def test_dbos_durability_allows_instrumented_default_model(dbos: DBOS) -> None:
"""An outer `Instrumentation` capability wraps the model, but the default model is still accepted.
`_find_model_id` peels any wrapper layers off the request's model before comparing
instances by identity, so an instrumented run still takes the default's `model_id=None`
fast path.
"""
agent = Agent(
_durability_fn_model,
name='durability_instrumented_default',
capabilities=[Instrumentation(settings=InstrumentationSettings()), DBOSDurability()],
)
@DBOS.workflow()
async def run_agent() -> str:
result = await agent.run('hello')
return result.output
assert await run_agent() == 'Echo: hello'
def test_dbos_durability_get_ordering() -> None:
"""DBOSDurability declares innermost ordering."""
from pydantic_ai.capabilities.abstract import CapabilityOrdering
durability = DBOSDurability()
ordering = durability.get_ordering()
assert ordering == CapabilityOrdering(position='innermost')
def test_dbos_durability_get_serialization_name() -> None:
"""DBOSDurability is not spec-serializable."""
assert DBOSDurability.get_serialization_name() is None
async def _durability_stream_fn(messages: list[ModelMessage], info: AgentInfo) -> AsyncIterator[str]:
# The first message carries the prompt and its first part is the `UserPromptPart`, so none of these branch.
for msg in reversed(messages): # pragma: no branch
for part in msg.parts: # pragma: no branch
if isinstance(part, UserPromptPart): # pragma: no branch
yield f'Echo: {part.content}'
return
yield 'no prompt' # pragma: no cover
async def test_dbos_durability_streaming_in_workflow(dbos: DBOS) -> None:
"""`ProcessEventStream` receives captured model events in workflow code."""
events_in_step: list[tuple[AgentStreamEvent, bool]] = []
async def handler(ctx: RunContext[object], stream: AsyncIterable[AgentStreamEvent]) -> None:
async for event in stream:
events_in_step.append((event, DBOS.step_id is not None))
stream_model = FunctionModel(_durability_model_fn, stream_function=_durability_stream_fn)
agent = Agent(
stream_model,
name='durability_streaming',
capabilities=[ProcessEventStream(handler), DBOSDurability()],
)
wfid = str(uuid.uuid4())
@DBOS.workflow()
async def run_durable_streaming_agent() -> str:
result = await agent.run('Hello streaming')
return result.output
with SetWorkflowID(wfid):
output = await run_durable_streaming_agent()
assert output == 'Echo: Hello streaming'
model_events_in_step = [
in_step for event, in_step in events_in_step if isinstance(event, (PartStartEvent, PartDeltaEvent))
]
assert model_events_in_step
assert not any(model_events_in_step)
steps = await dbos.list_workflow_steps_async(wfid)
step_names = [step['function_name'] for step in steps]
assert 'durability_streaming__model.request_stream' in step_names
async def _chunks_stream_fn(messages: list[ModelMessage], info: AgentInfo) -> AsyncIterator[str]:
yield 'Stream'
yield 'ed '
yield 'response'
async def test_dbos_durability_process_event_stream_fires_workflow_side(dbos: DBOS) -> None:
"""ProcessEventStream sees the real captured events replayed in the workflow."""
events_received: list[AgentStreamEvent] = []
async def collect(ctx: RunContext[object], stream: AsyncIterable[AgentStreamEvent]) -> None:
async for event in stream:
assert DBOS.step_id is None
events_received.append(event)
stream_model = FunctionModel(_durability_model_fn, stream_function=_chunks_stream_fn)
agent = Agent(
stream_model,
name='durability_process_stream',
capabilities=[ProcessEventStream(collect), DBOSDurability()],
)
@DBOS.workflow()
async def run_durable_agent() -> str:
result = await agent.run('Hello')
return result.output
output = await run_durable_agent()
assert output == 'Streamed response'
delta_events = [
e.delta.content_delta
for e in events_received
if isinstance(e, PartDeltaEvent) and isinstance(e.delta, TextPartDelta)
]
assert delta_events == ['ed ', 'response']
async def test_dbos_durability_buffers_caller_streams(dbos: DBOS) -> None:
agent = Agent(
TestModel(custom_output_text='hello world'),
name='durability_buffered_streams',
capabilities=[DBOSDurability()],
)
@DBOS.workflow()
async def run_durable_streams() -> tuple[list[str], str, list[str]]:
async with agent.run_stream('Hello') as stream:
chunks = [chunk async for chunk in stream.stream_text(debounce_by=None)]
output = await stream.get_output()
async with agent.run_stream_events('Hello') as event_stream:
events = [event async for event in event_stream]
deltas = [
event.delta.content_delta
for event in events
if isinstance(event, PartDeltaEvent) and isinstance(event.delta, TextPartDelta)
]
return chunks, output, deltas
assert await run_durable_streams() == (['hello ', 'hello world'], 'hello world', ['hello ', 'world'])
def test_dbos_durability_event_stream_handler_step_config(monkeypatch: pytest.MonkeyPatch) -> None:
registered_steps: list[tuple[str, dict[str, object]]] = []
def step(*, name: str, **config: object):
registered_steps.append((name, config))
def decorator(function: Any) -> Any:
return function
return decorator
monkeypatch.setattr(DBOS, 'step', step)
config = StepConfig(retries_allowed=True, max_attempts=3)
async def handler(ctx: RunContext[object], stream: AsyncIterable[AgentStreamEvent]) -> None:
async for _ in stream:
pass
Agent(
TestModel(),
name='event_handler_config',
capabilities=[DBOSDurability(event_stream_handler=handler, event_stream_handler_step_config=config)],
)
assert ('event_handler_config__event_stream_handler', config) in registered_steps
async def test_dbos_durability_event_stream_handler(dbos: DBOS) -> None:
events_in_boundary: list[tuple[AgentStreamEvent, bool]] = []
async def handler(ctx: RunContext[object], stream: AsyncIterable[AgentStreamEvent]) -> None:
async for event in stream:
events_in_boundary.append((event, DBOS.step_id is not None))
async def handled_tool() -> str:
return 'handled'
durability = DBOSDurability(event_stream_handler=handler)
agent = Agent(TestModel(), name='durability_handler', tools=[handled_tool], capabilities=[durability])
wfid = str(uuid.uuid4())
@DBOS.workflow()
async def run_durable_agent() -> str:
return (await agent.run('Hello')).output
with SetWorkflowID(wfid):
await run_durable_agent()
events = [event for event, _ in events_in_boundary]
assert events
# Model events run inside the model-request step; tool events run inside a dedicated
# event-handler step. Either way the handler always runs inside a checkpointed step.
assert all(in_boundary for _, in_boundary in events_in_boundary)
assert sum(isinstance(event, FunctionToolCallEvent) for event in events) == 1
assert sum(isinstance(event, FunctionToolResultEvent) for event in events) == 1
assert any(isinstance(event, PartStartEvent) for event in events)
assert any(isinstance(event, FinalResultEvent) for event in events)
steps = await dbos.list_workflow_steps_async(wfid)
step_names = [step['function_name'] for step in steps]
assert 'durability_handler__event_stream_handler' in step_names
async def test_dbos_durability_event_stream_handler_rejects_enqueue(dbos: DBOS) -> None:
"""An `event_stream_handler` that enqueues inside a durable step raises, like a tool would.
The handler runs inside a durable step for both model events (the model-request step) and
graph events (the `__event_stream_handler` step); either step's recorded result is replayed
without re-running it, so an enqueue would be dropped. The handler catches the error on every
event so the run still completes, exercising both delivery paths.
"""
enqueue_errors: list[str] = []
async def handler(ctx: RunContext[object], stream: AsyncIterable[AgentStreamEvent]) -> None:
async for _ in stream:
with pytest.raises(UserError, match='enqueued messages would be dropped') as exc_info:
ctx.enqueue('later')
# The ambient current context is guarded too, so reading it instead of the argument
# doesn't bypass the guard.
ambient = get_current_run_context()
assert ambient is not None
with pytest.raises(UserError, match='enqueued messages would be dropped'):
ambient.enqueue('later')
enqueue_errors.append(str(exc_info.value))
async def handled_tool() -> str:
return 'handled'
durability = DBOSDurability(event_stream_handler=handler)
agent = Agent(TestModel(), name='durability_handler_enqueue', tools=[handled_tool], capabilities=[durability])
@DBOS.workflow()
async def run_durable_agent() -> str:
return (await agent.run('Hello')).output
with SetWorkflowID(str(uuid.uuid4())):
await run_durable_agent()
# Guarded on both the model-event (model-request step) and graph-event (dispatch step) paths.
assert len(enqueue_errors) > 1
async def test_dbos_durability_event_stream_handler_outside_workflow(dbos: DBOS) -> None:
events: list[AgentStreamEvent] = []
async def handler(ctx: RunContext[object], stream: AsyncIterable[AgentStreamEvent]) -> None:
async for event in stream:
events.append(event)
durability = DBOSDurability(event_stream_handler=handler)
agent = Agent(TestModel(custom_output_text='done'), name='outside_handler', capabilities=[durability])
with agent.override():
await agent.run('Hello')
assert any(isinstance(event, PartStartEvent) for event in events)
def test_dbos_durability_without_handler_does_not_wrap_event_stream() -> None:
assert DBOSDurability().has_wrap_run_event_stream is False
async def test_dbos_durability_runtime_handler_receives_buffered_events(dbos: DBOS) -> None:
"""A per-run `event_stream_handler` passed to `agent.run()` inside a DBOS workflow
receives the events captured inside the step (rather than being silently dropped).
The buffered replay preserves real granular deltas — the per-run handler sees the
same multi-chunk stream the construction-time handler would see.
"""
events_received: list[AgentStreamEvent] = []
async def runtime_collect(ctx: RunContext[object], stream: AsyncIterable[AgentStreamEvent]) -> None:
async for event in stream:
events_received.append(event)
stream_model = FunctionModel(_durability_model_fn, stream_function=_chunks_stream_fn)
agent = Agent(
stream_model,
name='durability_runtime_handler',
capabilities=[DBOSDurability()],
)
@DBOS.workflow()
async def run_durable_agent() -> str:
result = await agent.run('Hello', event_stream_handler=runtime_collect)
return result.output
output = await run_durable_agent()
assert output == 'Streamed response'
# The runtime handler got real granular deltas (one PartDeltaEvent per chunk),
# not a single synthetic delta with the full text.
delta_events = [
e.delta.content_delta
for e in events_received
if isinstance(e, PartDeltaEvent) and isinstance(e.delta, TextPartDelta)
]
assert delta_events == ['ed ', 'response']
async def test_dbos_durability_mcp_toolset_wrapping(dbos: DBOS) -> None:
"""DBOSDurability discovers MCPToolset and creates DBOS wrappers."""
from pydantic_ai.durable_exec.dbos._mcp_toolset import DBOSMCPToolset
mcp_toolset = MCPToolset(
StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='my_mcp', init_timeout=20
)
agent = Agent(
_durability_fn_model,
name='durability_mcp',
toolsets=[mcp_toolset],
capabilities=[DBOSDurability()],
)
bound = DBOSDurability.from_agent(agent)
assert bound is not None
# The capability should have stored a DBOS wrapper keyed by the toolset id
assert 'my_mcp' in bound._toolsets_by_id # pyright: ignore[reportPrivateUsage]
assert isinstance(bound._toolsets_by_id['my_mcp'], DBOSMCPToolset) # pyright: ignore[reportPrivateUsage]
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_durability_mcp_operations_run_in_steps(dbos: DBOS) -> None:
seen_instructions: list[str] = []
async def local_tool() -> str:
return 'local result'
def call_then_answer(messages: list[ModelMessage], _: AgentInfo) -> ModelResponse:
seen_instructions.extend(
message.instructions for message in messages if isinstance(message, ModelRequest) and message.instructions
)
tool_returns = [part for message in messages for part in message.parts if isinstance(part, ToolReturnPart)]
if len(tool_returns) == 2:
return ModelResponse(parts=[TextPart('done')])
if len(tool_returns) == 1:
return ModelResponse(parts=[ToolCallPart('celsius_to_fahrenheit', {'celsius': 0}, tool_call_id='call-2')])
return ModelResponse(parts=[ToolCallPart('local_tool', {}, tool_call_id='call-1')])
agent = Agent(
FunctionModel(call_then_answer),
name='durability_mcp_operations',
toolsets=[
FunctionToolset(tools=[local_tool], id='functions'),
MCPToolset(
StdioTransport(command='python', args=['-m', 'tests.mcp_server']),
include_instructions=True,
id='mcp',
),
],
capabilities=[DBOSDurability()],
)
@DBOS.workflow()
async def run_agent() -> str:
return (await agent.run('Convert zero Celsius to Fahrenheit.')).output
wfid = str(uuid.uuid4())
with SetWorkflowID(wfid):
output = await run_agent()
assert output == 'done'
step_names = [step['function_name'] for step in await dbos.list_workflow_steps_async(wfid)]
assert step_names == [
'durability_mcp_operations__mcp_server__mcp.get_tools',
'durability_mcp_operations__mcp_server__mcp.get_instructions',
'durability_mcp_operations__model.request',
'durability_mcp_operations__mcp_server__mcp.get_instructions',
'durability_mcp_operations__model.request',
'durability_mcp_operations__mcp_server__mcp.call_tool',
'durability_mcp_operations__mcp_server__mcp.get_instructions',
'durability_mcp_operations__model.request',
]
# The instructions must actually reach the model, not just produce a step: they're captured
# during `__aenter__`, and DBOS's `enter-never` lifecycle means the step itself has to connect
# the server. Asserting only the step name let a silent `None` through.
assert 'Be a helpful assistant.' in seen_instructions
async def test_dbos_durability_rejects_idless_mcp_toolset(dbos: DBOS) -> None:
"""An `MCPToolset` without an `id` fails loudly at construction.
The DBOS step wrapper is swapped in by toolset ID at run time, so without one the
toolset's I/O would silently run un-checkpointed inside the DBOS workflow and
re-execute on recovery. Temporal raises the equivalent error for id-less leaves.
"""
mcp_toolset = MCPToolset(StdioTransport(command='python', args=['-m', 'tests.mcp_server']), init_timeout=20)
with pytest.raises(UserError, match=r"Toolsets that are 'leaves'.*unique `id`.*DBOS"):
Agent(
_durability_fn_model,
name='durability_idless_mcp',
toolsets=[mcp_toolset],
capabilities=[DBOSDurability()],
)
async def test_dbos_durability_wraps_capability_contributed_mcp_toolset(dbos: DBOS) -> None:
"""MCP toolsets contributed by other capabilities are wrapped as DBOS steps too.
Durability capabilities are in the `innermost` ordering tier, so `Agent.__init__` binds
them only after every other capability's contributed toolsets have been extracted into
`agent.toolsets`. Without that two-phase binding, this toolset would be invisible to
`for_agent` and run un-checkpointed inside the DBOS workflow.
"""
from pydantic_ai.durable_exec.dbos._mcp_toolset import DBOSMCPToolset
mcp_toolset = MCPToolset(
StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='cap_mcp', init_timeout=20
)
agent = Agent(
_durability_fn_model,
name='durability_cap_mcp',
capabilities=[Toolset(mcp_toolset), DBOSDurability()],
)
bound = DBOSDurability.from_agent(agent)
assert bound is not None
assert 'cap_mcp' in bound._toolsets_by_id # pyright: ignore[reportPrivateUsage]
assert isinstance(bound._toolsets_by_id['cap_mcp'], DBOSMCPToolset) # pyright: ignore[reportPrivateUsage]
async def test_dbos_durability_get_wrapper_toolset_with_mcp(dbos: DBOS) -> None:
"""DBOSDurability.get_wrapper_toolset replaces MCP toolsets by id."""
from pydantic_ai.durable_exec.dbos._mcp_toolset import DBOSMCPToolset
mcp_toolset = MCPToolset(
StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='swap_mcp', init_timeout=20
)
agent = Agent(
_durability_fn_model,
name='durability_swap',
toolsets=[mcp_toolset],
capabilities=[DBOSDurability()],
)
bound = DBOSDurability.from_agent(agent)
assert bound is not None
assert 'swap_mcp' in bound._toolsets_by_id # pyright: ignore[reportPrivateUsage]
# get_wrapper_toolset should replace the original MCP toolset with the DBOS wrapper
replaced = bound.get_wrapper_toolset(mcp_toolset)
assert replaced is not None
assert isinstance(replaced, DBOSMCPToolset)
async def test_dbos_durability_allows_runtime_function_toolset(dbos: DBOS) -> None:
"""A `FunctionToolset` added per-run is allowed and executes inline, like on `DBOSAgent`."""
def call_then_answer(messages: list[ModelMessage], _: AgentInfo) -> ModelResponse:
if any(isinstance(part, ToolReturnPart) for message in messages for part in message.parts):
return ModelResponse(parts=[TextPart('done')])
return ModelResponse(parts=[ToolCallPart('runtime_tool', {}, tool_call_id='call-1')])
agent = Agent(FunctionModel(call_then_answer), name='durability_runtime_fn', capabilities=[DBOSDurability()])
result = await agent.run(
'Call the runtime tool.', toolsets=[FunctionToolset(tools=[runtime_tool], id='runtime_fn')]
)
assert result.output == 'done'
async def test_dbos_durability_dynamic_capability_tool_runs_in_step(dbos: DBOS) -> None:
"""A dynamic capability's tool discovery and calls run in DBOS steps."""
calls: list[str] = []
def dynamic_tool() -> str:
calls.append('called')
return 'dynamic result'
def factory(ctx: RunContext[Any]) -> Capability[Any]:
return Capability(tools=[dynamic_tool])
agent = Agent(
TestModel(),
name='dbos_dynamic_capability',
capabilities=[DynamicCapability(factory, id='dyn'), DBOSDurability()],
)
@DBOS.workflow()
async def run_agent() -> str:
return (await agent.run('Call the tool')).output
wfid = str(uuid.uuid4())
with SetWorkflowID(wfid):
output = await run_agent()
assert output == '{"dynamic_tool":"dynamic result"}'
assert calls == ['called']
step_names = [step['function_name'] for step in await dbos.list_workflow_steps_async(wfid)]
assert 'dbos_dynamic_capability__dynamic_toolset__dyn.get_tools' in step_names
assert 'dbos_dynamic_capability__dynamic_toolset__dyn.call_tool' in step_names
async def test_dbos_durability_routes_concurrent_calls_to_same_dynamic_operation(dbos: DBOS) -> None:
"""Concurrent invocations share one registration without sharing call-local routing state."""
def factory(ctx: RunContext[str]) -> FunctionToolset[str]:
def read_route() -> str:
return ctx.deps
return FunctionToolset([read_route])
agent = Agent(
TestModel(call_tools=['read_route']),
deps_type=str,
name='dbos_concurrent_dynamic',
toolsets=[DynamicToolset(factory, id='routes')],
capabilities=[DBOSDurability()],
)
@DBOS.workflow()
async def run_agent() -> list[str]:
results = await asyncio.gather(agent.run('route', deps='left'), agent.run('route', deps='right'))
return [result.output for result in results]
assert await run_agent() == ['{"read_route":"left"}', '{"read_route":"right"}']
async def test_dbos_durability_nested_agent_in_tool_routes_to_inner_registrations(dbos: DBOS) -> None:
"""A workflow-level function tool can run another durable agent without ambient registry routing."""
def inner_factory(ctx: RunContext[str]) -> FunctionToolset[str]:
def inner_route() -> str:
return ctx.deps
return FunctionToolset([inner_route])
inner_agent = Agent(
TestModel(call_tools=['inner_route']),
deps_type=str,
name='dbos_nested_inner',
toolsets=[DynamicToolset(inner_factory, id='inner_routes')],
capabilities=[DBOSDurability()],
)
async def run_inner() -> str:
return (await inner_agent.run('route', deps='inner')).output
outer_agent = Agent(
TestModel(call_tools=['run_inner']),
name='dbos_nested_outer',
tools=[run_inner],
capabilities=[DBOSDurability()],
)
wfid = str(uuid.uuid4())
@DBOS.workflow()
async def run_outer() -> str:
return (await outer_agent.run('run nested')).output
with SetWorkflowID(wfid):
output = await run_outer()
assert 'inner_route' in output
step_names = [step['function_name'] for step in await dbos.list_workflow_steps_async(wfid)]
assert 'dbos_nested_inner__dynamic_toolset__inner_routes.get_tools' in step_names
assert 'dbos_nested_inner__dynamic_toolset__inner_routes.call_tool' in step_names
@pytest.mark.parametrize('metadata', [{'dbos': False}, {'': False}])
async def test_dbos_durability_ignores_per_tool_metadata(dbos: DBOS, metadata: dict[str, Any]) -> None:
"""DBOS takes no per-tool config: tool metadata never opts a tool out of its step.
DBOS registers a step once per name and its tool-call step names carry no tool name, so
per-tool config can't be honored. Metadata that other engines read (`{'dbos': False}` would
be an opt-out under Prefect/Temporal's `engine_spec.tool_config_key`) must leave the step in place --
dropping it would both un-checkpoint the call and shift the recorded step sequence, breaking
recovery of workflows recorded before this capability existed.
The empty-string key is the same contract from the other side: an engine with no
`engine_spec.tool_config_key` must not consult metadata at all, rather than reading key `''`.
"""
calls: list[str] = []
def opted_out_tool() -> str:
calls.append('called')
return 'dynamic result'
def factory(ctx: RunContext[Any]) -> Capability[Any]:
return Capability(tools=[Tool(opted_out_tool, metadata=metadata)])
agent = Agent(
TestModel(),
name='dbos_metadata_ignored',
capabilities=[DynamicCapability(factory, id='dyn'), DBOSDurability()],
)
@DBOS.workflow()
async def run_agent() -> str:
return (await agent.run('Call the tool')).output
wfid = str(uuid.uuid4())
with SetWorkflowID(wfid):
output = await run_agent()
assert output == '{"opted_out_tool":"dynamic result"}'
assert calls == ['called']
step_names = [step['function_name'] for step in await dbos.list_workflow_steps_async(wfid)]
assert 'dbos_metadata_ignored__dynamic_toolset__dyn.call_tool' in step_names
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_durability_mcp_tool_metadata_false_is_rejected(dbos: DBOS, monkeypatch: pytest.MonkeyPatch) -> None:
"""MCP tools perform I/O and cannot opt out of their durable DBOS step."""
mcp_toolset = MCPToolset(StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='dbos_mcp_opt_out')
tool = ToolsetTool(
toolset=mcp_toolset,
tool_def=ToolDefinition(name='inline', metadata={'dbos': False}),
max_retries=1,
args_validator=TOOL_SCHEMA_VALIDATOR,
)
async def get_tools(ctx: RunContext[None]) -> dict[str, ToolsetTool[None]]:
return {'inline': tool}
monkeypatch.setattr(mcp_toolset, 'get_tools', get_tools)
agent = Agent(
TestModel(call_tools='all'),
name='dbos_mcp_opt_out',
toolsets=[mcp_toolset],
capabilities=[DBOSDurability()],
)
@DBOS.workflow()
async def run_agent() -> str:
return (await agent.run('Hello')).output
with pytest.raises(UserError, match='MCP tools perform I/O'):
await run_agent()
def test_dbos_durability_dynamic_capability_requires_id(dbos: DBOS) -> None:
def factory(ctx: RunContext[Any]) -> Capability[Any]:
# Construction raises before the factory can run.
return Capability() # pragma: no cover
with pytest.raises(UserError, match=r"DynamicCapability\(\.\.\., id='user-tools'\)"):
Agent(
TestModel(),
name='dbos_dynamic_capability_no_id',
capabilities=[DynamicCapability(factory), DBOSDurability()],
)
def test_dbos_durability_bare_capability_func_requires_explicit_wrapper(dbos: DBOS) -> None:
"""A bare `CapabilityFunc` in `capabilities=` is wrapped in an id-less `DynamicCapability`,
so under durable execution it raises with a hint to wrap it explicitly."""
def factory(ctx: RunContext[Any]) -> Capability[Any]:
# Construction raises before the factory can run.
return Capability() # pragma: no cover
with pytest.raises(UserError, match=r'wrap it explicitly'):
Agent(
TestModel(),
name='dbos_bare_capability_func',
capabilities=[factory, DBOSDurability()],
)
@pytest.mark.subprocess(reason='connects to `tests.mcp_server` over stdio')
async def test_dbos_durability_dynamic_capability_mcp_runs_in_steps(dbos: DBOS) -> None:
def call_then_answer(messages: list[ModelMessage], _: AgentInfo) -> ModelResponse:
if any(isinstance(part, ToolReturnPart) for message in messages for part in message.parts):
return ModelResponse(parts=[TextPart('done')])
return ModelResponse(parts=[ToolCallPart('celsius_to_fahrenheit', {'celsius': 0}, tool_call_id='call-1')])
def factory(ctx: RunContext[Any]) -> Capability[Any]:
return Capability(
toolsets=[
MCPToolset(
StdioTransport(command='python', args=['-m', 'tests.mcp_server']),
id='dynamic_mcp',
)
]
)
agent = Agent(
FunctionModel(call_then_answer),
name='dbos_dynamic_mcp',
capabilities=[DynamicCapability(factory, id='dyn_mcp'), DBOSDurability()],
)
@DBOS.workflow()
async def run_agent() -> str:
return (await agent.run('Convert zero Celsius to Fahrenheit.')).output
wfid = str(uuid.uuid4())
with SetWorkflowID(wfid):
output = await run_agent()
assert output == 'done'
step_names = [step['function_name'] for step in await dbos.list_workflow_steps_async(wfid)]
assert 'dbos_dynamic_mcp__dynamic_toolset__dyn_mcp.get_tools' in step_names
assert 'dbos_dynamic_mcp__dynamic_toolset__dyn_mcp.call_tool' in step_names
async def test_dbos_durability_rejects_runtime_mcp_toolset(dbos: DBOS) -> None:
"""An `MCPToolset` added per-run is rejected: its I/O steps must be registered at construction."""
agent = Agent(_durability_fn_model, name='durability_runtime_mcp', capabilities=[DBOSDurability()])
@DBOS.workflow()
async def run_agent() -> None:
await agent.run(
'Hello',
toolsets=[MCPToolset(StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='runtime_mcp')],
)
with pytest.raises(UserError, match=r"MCPToolset 'runtime_mcp' cannot be added at runtime with DBOS"):
await run_agent()
def test_dbos_durability_rejects_runtime_dynamic_toolset_sync(dbos: DBOS) -> None:
"""A `DynamicToolset` added per-run is rejected, on `run_sync` as well as `run`."""
agent = Agent(_durability_fn_model, name='durability_runtime_dynamic', capabilities=[DBOSDurability()])
@DBOS.workflow()
def run_agent() -> None:
agent.run_sync('Hello', toolsets=[DynamicToolset(lambda _: FunctionToolset(), id='runtime_dynamic')])
with pytest.raises(
UserError,
match=r"DynamicToolset 'runtime_dynamic' cannot be added at runtime with DBOS",
):
run_agent()
async def test_dbos_durability_rejects_runtime_mcp_toolset_in_iter(dbos: DBOS) -> None:
"""`agent.iter(toolsets=...)` inside a user workflow is guarded like `run(toolsets=...)`.
The rejection lives in run setup (`get_wrapper_toolset`), which every entry point routes
through so `iter` inside a workflow cannot execute the MCP toolset's I/O un-checkpointed.
"""
agent = Agent(_durability_fn_model, name='durability_runtime_mcp_iter', capabilities=[DBOSDurability()])
@DBOS.workflow()
async def run_agent() -> None:
async with agent.iter(
'Hello',
toolsets=[MCPToolset(StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='iter_mcp')],
):
# Run setup raises before any node runs.
pass # pragma: no cover
with pytest.raises(UserError, match=r"MCPToolset 'iter_mcp' cannot be added at runtime with DBOS"):
await run_agent()
def _per_run_dynamic_factory(ctx: RunContext[Any]) -> FunctionToolset[Any]:
return FunctionToolset()
async def test_dbos_durability_rejects_per_run_capabilities(dbos: DBOS) -> None:
"""Capabilities added per-run inside a workflow are rejected; `Instrumentation` is exempt.
Construction-time capability toolsets are wrapped by `for_agent` (see the
capability-contributed test above); a per-run capability's toolset arrives after that
wrapping has happened, so it would run un-checkpointed inside the workflow. A
`DynamicToolset` with a module-level factory keeps the workflow input serializable.
"""
agent = Agent(_durability_fn_model, name='durability_per_run_cap_toolset', capabilities=[DBOSDurability()])
@DBOS.workflow()
async def run_with_toolset_capability() -> None:
await agent.run('Hello', capabilities=[Toolset(DynamicToolset(_per_run_dynamic_factory, id='per_run_dynamic'))])
with workflow_raises(
UserError,
snapshot(
'Capabilities added per-run inside a DBOS workflow are not supported: Toolset. DBOS registers '
'durable steps when a capability is bound to the agent, before the workflow starts. A capability '
'added per-run therefore has no registered durable steps for the toolsets it contributes or its own '
'`@durable_operation` methods. Attach all capabilities at agent construction time so '
'`DBOSDurability.for_agent()` can register their durable steps.'
),
):
await run_with_toolset_capability()
@DBOS.workflow()
async def run_with_instrumentation() -> str:
return (await agent.run('Hello', capabilities=[Instrumentation(InstrumentationSettings())])).output
assert await run_with_instrumentation() == snapshot('Echo: Hello')
async def test_dbos_durability_allows_per_run_capabilities_outside_workflow(dbos: DBOS) -> None:
"""Outside a workflow the capability is transparent, so per-run capabilities are fine."""
agent = Agent(_durability_fn_model, name='durability_per_run_cap_outside', capabilities=[DBOSDurability()])
result = await agent.run('Hello', capabilities=[Toolset(FunctionToolset(id='per_run_fn'))])
assert result.output == snapshot('Echo: Hello')
@pytest.mark.parametrize('kind', ['mcp', 'dynamic'])
async def test_dbos_durability_rejects_overridden_executing_toolsets(dbos: DBOS, kind: str) -> None:
toolsets = {
'mcp': MCPToolset(StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='override_mcp'),
'dynamic': DynamicToolset(_per_run_dynamic_factory, id='override_dynamic'),
}
agent = Agent(_durability_fn_model, name=f'durability_override_{kind}', capabilities=[DBOSDurability()])
@DBOS.workflow()
async def run_agent() -> None:
with agent.override(toolsets=[toolsets[kind]]):
await agent.run('Hello')
with pytest.raises(UserError, match=r'cannot be added at runtime .*`override'):
await run_agent()
async def test_dbos_durability_allows_overridden_function_toolset(dbos: DBOS) -> None:
calls: list[str] = []
def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
if any(isinstance(part, ToolReturnPart) for message in messages for part in message.parts):
return ModelResponse(parts=[TextPart('done')])
return ModelResponse(parts=[ToolCallPart('override_tool', {}, tool_call_id='call-1')])
def override_tool() -> str:
calls.append('override_tool')
return 'ok'
agent = Agent(FunctionModel(model_fn), name='durability_override_function', capabilities=[DBOSDurability()])
@DBOS.workflow()
async def run_agent() -> str:
with agent.override(toolsets=[FunctionToolset([override_tool], id='override_fn')]):
return (await agent.run('Hello')).output
assert await run_agent() == 'done'
assert calls == ['override_tool']
async def test_dbos_durability_rejects_runtime_toolset_reusing_registered_id(dbos: DBOS) -> None:
agent = Agent(
_durability_fn_model,
name='durability_runtime_id_collision',
toolsets=[DynamicToolset(_per_run_dynamic_factory, id='shared')],
capabilities=[DBOSDurability()],
)
colliding = FunctionToolset[Any](id='shared')
message = "A toolset added at run time has the same `id` 'shared' as one the agent was constructed with"
@DBOS.workflow()
async def run_with_override() -> None:
with agent.override(toolsets=[colliding]):
await agent.run('Hello')
with pytest.raises(UserError, match=message):
await run_with_override()
# Outside a workflow the capability is transparent: there is no durable unit to dispatch to, so
# the toolset that actually arrived is used as-is rather than the run being rejected.
assert await agent.run('Hello', toolsets=[colliding]) is not None
async def test_dbos_durability_rejects_duplicate_toolset_id(dbos: DBOS) -> None:
"""Two distinct MCP toolsets under one `id` are rejected at binding time.
The registry maps `id` → step wrapper, so a duplicate would silently replace the first
entry and route both toolsets' calls through the last one's steps.
"""
with pytest.raises(UserError, match="Two toolsets have the same `id` 'dup'"):
Agent(
_durability_fn_model,
name='durability_dup_mcp',
toolsets=[
MCPToolset(StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='dup'),
MCPToolset(StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='dup'),
],
capabilities=[DBOSDurability()],
)
async def test_dbos_durability_same_toolset_instance_reused(dbos: DBOS) -> None:
"""The same MCP toolset instance appearing twice maps to one wrapper, not an `id` conflict."""
mcp_toolset = MCPToolset(StdioTransport(command='python', args=['-m', 'tests.mcp_server']), id='shared_mcp')
agent = Agent(
_durability_fn_model,
name='durability_shared_mcp',
toolsets=[mcp_toolset, mcp_toolset],
capabilities=[DBOSDurability()],
)
bound = DBOSDurability.from_agent(agent)
assert bound is not None
assert list(bound._toolsets_by_id) == ['shared_mcp'] # pyright: ignore[reportPrivateUsage]
# --- Continuation chains (suspended → complete) run one step per segment ---
#
# When a model suspends a turn (Anthropic `pause_turn`, OpenAI background mode), the
# continuation loop in the innermost `model_request`/`model_request_stream` helpers runs
# workflow-side under `DBOSDurability`, dispatching each segment through its own model
# request step. These tests use a scripted model (no cassettes: `FunctionModel` can't emit
# suspended streaming segments, and VCR matchers wouldn't pin the chain shape).
async def test_dbos_durability_continuation_chain_in_workflow(dbos: DBOS) -> None:
"""A suspended → complete chain resolves across per-segment DBOS steps, as one merged response.
Usage is counted once (a continuation isn't a separate request step), and the workflow
record shows one model-request step per segment.
"""
model = ScriptedContinuationModel(
responses=[
scripted_response(
texts=['The answer '],
state='suspended',
provider_response_id='cont1',
input_tokens=5,
output_tokens=2,
),
scripted_response(texts=['is 42.'], provider_response_id='cont2', input_tokens=3, output_tokens=4),
]
)
agent = Agent(model, name='durability_continuation', capabilities=[DBOSDurability()])
results: list[AgentRunResult[str]] = []
@DBOS.workflow()
async def run_durable_agent() -> str:
result = await agent.run('go')
results.append(result)
return result.output
wfid = str(uuid.uuid4())
with SetWorkflowID(wfid):
output = await run_durable_agent()
assert output == 'The answer is 42.'
result = results[0]
response = result.all_messages()[-1]
assert isinstance(response, ModelResponse)
assert response.state == 'complete'
assert [part.content for part in response.parts if isinstance(part, TextPart)] == ['The answer ', 'is 42.']
assert result.usage.requests == 1
assert result.usage.input_tokens == 8
assert result.usage.output_tokens == 6
# Each segment ran in its own durable boundary: two model calls, two steps.
assert model.request_calls == 2
steps = await dbos.list_workflow_steps_async(wfid)
step_names = [step['function_name'] for step in steps]
assert step_names.count('durability_continuation__model.request') == 2
async def test_dbos_durability_continuation_usage_limit_cancels_suspended(dbos: DBOS) -> None:
"""A usage limit tripped between segments cancels the live suspended job in its own step.
The continuation loop runs workflow-side and checks the limit as each segment merges;
the provider teardown of the abandoned server-side job is I/O, so it crosses the
boundary through the dedicated cancellation step, and the error surfaces to workflow
code with its real type.
"""
model = ScriptedContinuationModel(
responses=[
scripted_response(
texts=['still going '],
state='suspended',
provider_response_id='cont1',
input_tokens=10,
output_tokens=5,
),
scripted_response(
texts=['keeps going '],
state='suspended',
provider_response_id='cont2',
input_tokens=100,
output_tokens=50,
),
]
)
agent = Agent(model, name='durability_continuation_usage_limit', capabilities=[DBOSDurability()])
@DBOS.workflow()
async def run_agent() -> None:
await agent.run('go', usage_limits=UsageLimits(total_tokens_limit=20))
wfid = str(uuid.uuid4())
with SetWorkflowID(wfid):
with pytest.raises(UsageLimitExceeded, match='total_tokens_limit'):
await run_agent()
assert model.request_calls == 2
# The over-budget merge was still suspended, so the live job was cancelled before raising.
assert [cancelled.provider_response_id for cancelled in model.cancelled] == ['cont2']
steps = await dbos.list_workflow_steps_async(wfid)
step_names = [step['function_name'] for step in steps]
assert step_names.count('durability_continuation_usage_limit__model.request') == 2
assert step_names.count('durability_continuation_usage_limit__model.cancel_suspended_response') == 1
async def test_dbos_cancellation_rejects_enqueue(dbos: DBOS) -> None:
enqueue_rejected = False
class EnqueueOnCancelModel(ScriptedContinuationModel):
async def cancel_suspended_response(self, response: ModelResponse) -> None:
nonlocal enqueue_rejected
ctx = get_current_run_context()
assert ctx is not None
with pytest.raises(UserError, match='enqueued messages would be dropped'):
ctx.enqueue('later')
enqueue_rejected = True
model = EnqueueOnCancelModel(
responses=[
scripted_response(
texts=['still going'],
state='suspended',
provider_response_id='cont1',
input_tokens=10,
output_tokens=5,
),
scripted_response(
texts=['over budget'],
state='suspended',
provider_response_id='cont2',
input_tokens=100,
output_tokens=50,
),
]
)
agent = Agent(model, name='dbos_cancel_enqueue', capabilities=[DBOSDurability()])
@DBOS.workflow()
async def run_agent() -> None:
await agent.run('continue', usage_limits=UsageLimits(total_tokens_limit=50))
with pytest.raises(UsageLimitExceeded, match='total_tokens_limit'):
await run_agent()
assert enqueue_rejected
async def test_dbos_durability_streaming_continuation_chain_in_workflow(dbos: DBOS) -> None:
"""A streamed suspended → complete chain is stitched across per-segment DBOS steps.
`ProcessEventStream` receives each captured segment in workflow code, and the
final response merges both segments' text with usage summed once.
"""
model = ScriptedContinuationModel(
segments=[
StreamSegment(
texts=['The answer '],
state='suspended',
provider_response_id='cont1',
input_tokens=5,
output_tokens=2,
),
StreamSegment(
texts=['is 42.'], state='complete', provider_response_id='cont2', input_tokens=3, output_tokens=4
),
]
)
events_received: list[AgentStreamEvent] = []
async def handler(ctx: RunContext[object], stream: AsyncIterable[AgentStreamEvent]) -> None:
async for event in stream:
events_received.append(event)
agent = Agent(
model,
name='durability_continuation_stream',
capabilities=[ProcessEventStream(handler), DBOSDurability()],
)
results: list[AgentRunResult[str]] = []
@DBOS.workflow()
async def run_durable_agent() -> str:
result = await agent.run('go')
results.append(result)
return result.output
wfid = str(uuid.uuid4())
with SetWorkflowID(wfid):
output = await run_durable_agent()
assert output == 'The answer is 42.'
result = results[0]
assert result.usage.requests == 1
assert result.usage.input_tokens == 8
assert result.usage.output_tokens == 6
indices = [
(type(event).__name__, event.index)
for event in events_received
if isinstance(event, (PartStartEvent, PartDeltaEvent))
]
assert indices == snapshot(
[('PartStartEvent', 0), ('PartDeltaEvent', 0), ('PartStartEvent', 1), ('PartDeltaEvent', 1)]
)
assert model.request_stream_calls == 2
steps = await dbos.list_workflow_steps_async(wfid)
step_names = [step['function_name'] for step in steps]
assert step_names.count('durability_continuation_stream__model.request_stream') == 2
async def test_dbos_durability_continuation_resume_from_history(dbos: DBOS) -> None:
"""A suspended history tail crosses the first DBOS step boundary.
The workflow-side continuation loop passes the tail as the last request message.
"""
model = ScriptedContinuationModel(
responses=[scripted_response(texts=['is 42.'], provider_response_id='cont2', input_tokens=3, output_tokens=4)]
)
agent = Agent(model, name='durability_continuation_resume', capabilities=[DBOSDurability()])
history: list[ModelMessage] = [
ModelRequest(parts=[UserPromptPart(content='go')]),
scripted_response(
texts=['The answer '], state='suspended', provider_response_id='cont1', input_tokens=5, output_tokens=2
),
]
results: list[AgentRunResult[str]] = []
@DBOS.workflow()
async def run_durable_agent() -> str:
result = await agent.run(message_history=history)
results.append(result)
return result.output
output = await run_durable_agent()
assert output == 'The answer is 42.'
result = results[0]
response = result.all_messages()[-1]
assert isinstance(response, ModelResponse)
assert response.state == 'complete'
assert [part.content for part in response.parts if isinstance(part, TextPart)] == ['The answer ', 'is 42.']
assert result.usage.requests == 1
assert result.usage.input_tokens == 8
assert result.usage.output_tokens == 6
# The continuation request ran inside the boundary — the seed wasn't re-generated.
assert model.request_calls == 1
async def test_dbos_agent_run_sync_from_sync_tool_is_rejected():
"""`DBOSAgent.run_sync()` dispatches through its own workflow, not `AbstractAgent.run_sync()`, so it carries its own guard."""
def call_tool(_: list[ModelMessage], __: AgentInfo) -> ModelResponse:
return ModelResponse(parts=[ToolCallPart('delegate', '{}')])
outer_agent = Agent(FunctionModel(call_tool))
@outer_agent.tool_plain
def delegate() -> str:
return simple_dbos_agent.run_sync('hello').output
with pytest.raises(UserError, match=r'cannot be used inside a synchronous tool'):
await outer_agent.run('delegate')
@pytest.mark.parametrize('deprecated_agent', [False, True])
async def test_dbos_mcp_server_keeps_one_session_per_workflow(dbos: DBOS, deprecated_agent: bool) -> None:
"""A statically attached `MCPToolset` connects once per workflow, not once per step (#8458).
The first step that needs the server connects it — inside a step, where DBOS retries a failed
connection — and the workflow holds the session for the steps that follow, so `cache_tools`
serves the later discovery steps and the `tools/list` the MCP SDK issues before a call finds a
warm session. Both APIs get it: the deprecated `DBOSAgent` builds its own MCP steps, and until
it is removed a `DBOSAgent` user should not be paying for a session per step.
Not a VCR test: an in-process server is what makes the round trips countable server-side, which
is what attributes the MCP SDK's own pre-call listing correctly.
"""
from .counting_mcp import counting_mcp_server, two_echo_calls_model
server, counts = counting_mcp_server()
toolset = MCPToolset(server, id=f'session_mcp_{deprecated_agent}')
agent = Agent(
two_echo_calls_model(),
name=f'dbos_mcp_session_{deprecated_agent}',
toolsets=[toolset],
capabilities=[] if deprecated_agent else [DBOSDurability()],
)
durable_agent = DBOSAgent(agent) if deprecated_agent else agent # pyright: ignore[reportDeprecated]
@DBOS.workflow(name=f'mcp_session_workflow_{deprecated_agent}')
async def run_workflow() -> str:
return (await durable_agent.run('go')).output
assert await run_workflow() == 'done'
assert counts == snapshot({'initialize': 1, 'tools/list': 1, 'tools/call': 2})
# The workflow closed the session it held; nothing keeps the server connected between runs.
assert not toolset.is_running
async def test_dbos_decide_span_nests_under_chat(
allow_model_requests: None, dbos: DBOS, capfire: CaptureLogfire
) -> None:
"""A decision model's `decide` span lands under the step, under `chat`, with the request's content policy.
The step runs in the workflow's process and context, so the policy the `chat` span set reaches it directly.
"""
agent = Agent(
ShipItDecisionModel(),
output_type=ShipIt,
name='dbos_decide',
capabilities=[DBOSDurability(), Instrumentation()],
)
@DBOS.workflow()
async def run_decision_agent() -> ShipIt:
return (await agent.run('The migration is reviewed and the tests pass.')).output
assert await run_decision_agent() == ShipIt(ship=True)
lineage, attributes = decide_span_lineage(capfire.exporter.exported_spans_as_dict())
assert lineage[:3] == snapshot(['dbos_decide__model.request', 'chat ship-it', 'invoke_agent dbos_decide'])
assert attributes['pydantic_ai.decision.state'] == 'The migration is reviewed and the tests pass.'
stable_ids_agent = Agent(TestModel(), name='stable_ids_agent', capabilities=[DBOSDurability()])
stable_ids_seen: list[tuple[str | None, str | None]] = []
class _ProcessCrash(BaseException):
"""Stands in for the process dying: DBOS leaves the workflow `PENDING` for recovery."""
@stable_ids_agent.tool
def record_ids(ctx: RunContext[object]) -> str:
stable_ids_seen.append((ctx.run_id, ctx.conversation_id))
return 'ok'
@DBOS.workflow()
async def stable_ids_workflow() -> list[tuple[str, str]]:
first = await stable_ids_agent.run('First.')
# Concurrent runs resolve their IDs before either starts a step.
results = [first, *await asyncio.gather(stable_ids_agent.run('Second.'), stable_ids_agent.run('Third.'))]
if len(stable_ids_seen) == 3:
raise _ProcessCrash
return [(result.run_id, result.conversation_id) for result in results]
async def test_dbos_default_ids_survive_recovery(dbos: DBOS) -> None:
"""Without `run_id=` or `conversation_id=`, a recovered workflow keeps the IDs its steps already saw."""
stable_ids_seen.clear()
workflow_id = f'stable-ids-{uuid.uuid4()}'
with SetWorkflowID(workflow_id), pytest.raises(_ProcessCrash):
await stable_ids_workflow()
seen = sorted(stable_ids_seen)
assert [run_id for run_id, _ in seen] == [f'{workflow_id}:1', f'{workflow_id}:3', f'{workflow_id}:3:1']
assert len({*seen[0], *seen[1], *seen[2]}) == 6
# Resuming the pending workflow re-executes it under the same ID, the way recovery does: model requests
# replay from their recorded steps, while the function tool, which runs in the workflow, runs again and
# sees the same IDs.
handle = await DBOS.resume_workflow_async(workflow_id)
assert sorted(tuple(ids) for ids in await handle.get_result()) == seen
assert sorted(stable_ids_seen[3:]) == seen
Reference in New Issue View Git Blame Copy Permalink