furyhawk/pydantic-ai-harness
1
0
Fork 0
mirror of https://github.com/pydantic/pydantic-ai-harness.git synced 2026-07-21 02:45:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(subagents): per-delegate run controls (usage limits, timeout, call budget, soft-failure) (#283)

Browse Source 
* feat(subagents): per-delegate run controls via SubAgentLimits

Delegations through SubAgents ran with no per-child budgets, so an
orchestrator that needs one child capped at N requests, another bounded by
a wall-clock deadline, or a per-run delegation cap had to keep a parallel
hand-rolled delegation path (motivating consumer: Pydanty's run_delegate).

Add a `limits: Mapping[str, SubAgentLimits]` knob. `SubAgentLimits` carries
`usage_limits` (isolated child accounting so the budget is the child's own,
reached softly rather than via a propagating UsageLimitExceeded),
`timeout_seconds` (wall-clock cancel with a soft steering message),
`max_calls` (per-run delegation budget, counts keyed by run_id and cleared
in wrap_run), and `on_failure` (a single steering message that overrides the
default soft text and flips child failures from a parent ModelRetry to a
soft tool result). Absent names keep current behavior.

Closes #282

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* review: validate limits keys, exclude run-state from eq, cover parallel budget

Adversarial self-review follow-ups on the per-delegate controls:

- SubAgents.__post_init__ raises when `limits` names a sub-agent absent from
  `agents`. A typo'd key would otherwise silently run that delegate unbudgeted,
  the exact failure the budgets exist to prevent.
- `_call_counts` is marked compare=False so accumulating delegation counts
  during a run don't change SubAgents equality.
- Add a test for two delegations issued in one parent step: the synchronous
  count still caps them at max_calls. Document that max_calls is scoped per
  run_id (so nested trees budget independently) and that a timed-out child's
  event stream is cut without a terminal event.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: David SF <david.sanchez@pydantic.dev>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
David SF
2026-06-11 22:50:11 -05:00
committed by GitHub
co-authored by David SF Claude Opus 4.8
parent 6f2aa11438
commit 34dcbf8352
5 changed files with 455 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pydantic_ai_harness/experimental/subagents/README.md
+36 -1
View File
@@ -63,9 +63,43 @@ print(result.output)
- **Capabilities can be shared.** `shared_capabilities` are applied to every sub-agent run -- e.g. give all sub-agents a common guardrail, memory, or planning capability without rebuilding each `Agent`.
- **Sub-agent events can be streamed.** Pass an `event_stream_handler` and it's forwarded to each sub-agent run, so the sub-agent's model-streaming and tool events surface to the caller (the handler receives the sub-agent's own `RunContext`).
## Per-delegate run controls
`limits` maps a sub-agent name to a `SubAgentLimits`, giving one delegate its own budgets without touching the others. A name absent from `limits` runs with the `SubAgents` defaults.
```python
from pydantic_ai.usage import UsageLimits
from pydantic_ai_harness.experimental.subagents import SubAgentLimits, SubAgents
# reproducer and librarian are Agent instances, as in the example above.
orchestrator = Agent(
'anthropic:claude-opus-4-7',
capabilities=[
SubAgents(
agents={'reproducer': reproducer, 'librarian': librarian},
limits={
'reproducer': SubAgentLimits(usage_limits=UsageLimits(request_limit=35), timeout_seconds=600, max_calls=1),
'librarian': SubAgentLimits(usage_limits=UsageLimits(request_limit=18), timeout_seconds=300, max_calls=2),
},
)
],
)
```
| Field | Effect |
|---|---|
| `usage_limits` | A request/token budget for one delegation. The child runs with its own usage accounting, so the budget counts only that child's requests and tokens (not the parent's or siblings'), even when `forward_usage=True`. The tradeoff: that child's tokens no longer aggregate into the parent's `usage`. Reaching the budget is a soft outcome (see below), not a run-stopping `UsageLimitExceeded`. |
| `timeout_seconds` | A wall-clock budget for one delegation. When the child exceeds it, its run is cancelled and the parent gets a soft steering message instead of hanging on the child. The cancelled child's `event_stream_handler` (if any) stops receiving events without a terminal event. |
| `max_calls` | The maximum number of delegations to this sub-agent per parent run. Once reached, further delegations return a soft budget-exhausted message without running the child. Counts are scoped to one `Agent.run` (a `run_id`) and cleared when it ends, so each parent run and each level of a nested tree budgets independently. |
| `on_failure` | A steering message returned to the parent for any soft degradation of this delegate, in place of the built-in default. Setting it also makes child failures soft (see below). |
## Failure handling
If a sub-agent run fails with a *soft* model error (`ModelRetry`, `UnexpectedModelBehavior`, e.g. it exhausted its own retries), the failure is converted into a `ModelRetry` for the parent -- so the parent's model sees `Sub-agent '<name>' failed: …` and can react. Hard errors (e.g. `UsageLimitExceeded`) propagate to stop the whole run.
A *soft outcome* returns a steering message to the parent as a normal tool result, so its model reads the message and decides what to do next (rather than immediately re-delegating, which a `ModelRetry` invites). A timeout, a reached `usage_limits` budget, and an exhausted `max_calls` budget are always soft. When `on_failure` is set, the message it carries replaces the built-in default for these outcomes.
A sub-agent run that fails with a *soft model error* (`ModelRetry`, `UnexpectedModelBehavior`, e.g. it exhausted its own retries) is, by default, converted into a `ModelRetry` for the parent -- so the parent's model sees `Sub-agent '<name>' failed: ...` and can react. Set `on_failure` for that delegate to make its failures soft instead: the child error returns the `on_failure` message as a normal tool result.
Hard errors propagate to stop the whole run. A `UsageLimitExceeded` from a child that has *no* per-delegate `usage_limits` (so it shares the parent's accounting) means the whole tree is out of budget and propagates; a child reaching its *own* `usage_limits` is soft, as above.
## Discovery
@@ -82,6 +116,7 @@ SubAgents(
shared_capabilities=(),# capabilities applied to every sub-agent run
event_stream_handler=None, # forwarded to each sub-agent run to stream its events
tool_name='delegate_task',
limits={}, # Mapping[str, SubAgentLimits] -- per-delegate run controls
)
```
pydantic_ai_harness/experimental/subagents/__init__.py
+2 -2
View File
@@ -2,8 +2,8 @@
from pydantic_ai_harness.experimental._warn import warn_experimental
from pydantic_ai_harness.experimental.subagents._capability import SubAgents
from pydantic_ai_harness.experimental.subagents._toolset import SubAgentToolset
from pydantic_ai_harness.experimental.subagents._toolset import SubAgentLimits, SubAgentToolset
warn_experimental('subagents')
__all__ = ['SubAgentToolset', 'SubAgents']
__all__ = ['SubAgentLimits', 'SubAgentToolset', 'SubAgents']
pydantic_ai_harness/experimental/subagents/_capability.py
+39 -4
View File
@@ -6,12 +6,12 @@ from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
from pydantic_ai.agent import AbstractAgent, EventStreamHandler
from pydantic_ai.capabilities import AbstractCapability, AgentCapability
from pydantic_ai.tools import AgentDepsT
from pydantic_ai.agent import AbstractAgent, AgentRunResult, EventStreamHandler
from pydantic_ai.capabilities import AbstractCapability, AgentCapability, WrapRunHandler
from pydantic_ai.tools import AgentDepsT, RunContext
from pydantic_ai.toolsets import AgentToolset
from pydantic_ai_harness.experimental.subagents._toolset import SubAgentToolset
from pydantic_ai_harness.experimental.subagents._toolset import SubAgentLimits, SubAgentToolset
if TYPE_CHECKING:
from pydantic_ai._instructions import AgentInstructions
@@ -33,6 +33,11 @@ class SubAgents(AbstractCapability[AgentDepsT]):
applied to every sub-agent run (`shared_capabilities`), and sub-agent events
can be streamed to a handler (`event_stream_handler`).
Per-delegate run controls (a `usage_limits` budget, a wall-clock
`timeout_seconds`, a per-run `max_calls` budget, and an `on_failure` steering
message) are configured via `limits`, keyed by sub-agent name. See
`SubAgentLimits`.
```python
from pydantic_ai import Agent
from pydantic_ai_harness.experimental.subagents import SubAgents
@@ -78,6 +83,34 @@ class SubAgents(AbstractCapability[AgentDepsT]):
tool_name: str = 'delegate_task'
"""Name of the delegate tool exposed to the model."""
limits: Mapping[str, SubAgentLimits] = field(default_factory=dict[str, SubAgentLimits])
"""Per-delegate run controls, keyed by sub-agent name. A name absent here runs
with no per-delegate controls (the `SubAgents` defaults). See `SubAgentLimits`
for `usage_limits`, `timeout_seconds`, `max_calls`, and `on_failure`."""
_call_counts: dict[str, dict[str, int]] = field(
default_factory=dict[str, 'dict[str, int]'], init=False, repr=False, compare=False
)
"""Run-scoped delegation counts (run_id -> name -> count), shared with the
toolset and cleared per run in `wrap_run`. Backs `SubAgentLimits.max_calls`."""
def __post_init__(self) -> None:
# A limits entry for an unknown name would silently run unbudgeted, defeating
# the point of configuring a budget; fail fast on the likely typo instead.
unknown = set(self.limits) - set(self.agents)
if unknown:
raise ValueError(
f'`limits` names sub-agents not in `agents`: {", ".join(sorted(unknown))}. '
f'Available sub-agents: {", ".join(sorted(self.agents)) or "(none)"}.'
)
async def wrap_run(self, ctx: RunContext[AgentDepsT], *, handler: WrapRunHandler) -> AgentRunResult[Any]:
"""Run the parent agent, then drop this run's delegation counts so they don't accumulate."""
try:
return await handler()
finally:
self._call_counts.pop(ctx.run_id or '', None)
def get_instructions(self) -> AgentInstructions[AgentDepsT] | None:
"""Static, cache-stable listing of the available sub-agents."""
if not self.agents:
@@ -105,6 +138,8 @@ class SubAgents(AbstractCapability[AgentDepsT]):
shared_capabilities=self.shared_capabilities,
event_stream_handler=self.event_stream_handler,
tool_name=self.tool_name,
limits=self.limits,
call_counts=self._call_counts,
)
@classmethod
pydantic_ai_harness/experimental/subagents/_toolset.py
+116 -12
View File
@@ -2,14 +2,52 @@
from __future__ import annotations
import asyncio
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any
from pydantic_ai.agent import AbstractAgent, EventStreamHandler
from pydantic_ai.capabilities import AgentCapability
from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior
from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior, UsageLimitExceeded
from pydantic_ai.tools import AgentDepsT, RunContext
from pydantic_ai.toolsets import AbstractToolset, FunctionToolset
from pydantic_ai.usage import UsageLimits
@dataclass(frozen=True)
class SubAgentLimits:
"""Per-delegate run controls for one named sub-agent.
Every field is optional; an unset field leaves the corresponding behaviour
at the `SubAgents` default. Pass a mapping of sub-agent name to
`SubAgentLimits` as `SubAgents(limits=...)`.
"""
usage_limits: UsageLimits | None = None
"""Request/token budget for one delegation. When set, the child runs with
its own usage accounting so the budget counts only the child's own requests
and tokens (not the parent's or siblings'), even when `forward_usage=True`.
The tradeoff: that child's tokens no longer aggregate into the parent's
`usage`. Hitting this budget is a soft outcome (steering message), not a
run-stopping `UsageLimitExceeded`."""
timeout_seconds: float | None = None
"""Wall-clock budget for one delegation. When the child exceeds it, the run
is cancelled and the parent gets a soft steering message instead of hanging
on the child."""
max_calls: int | None = None
"""Maximum number of delegations to this sub-agent per parent run. Once
reached, further delegations return a soft budget-exhausted message without
running the child."""
on_failure: str | None = None
"""Steering message returned to the parent for any soft degradation of this
delegate (timeout, child failure, usage budget reached, call budget
exhausted), in place of the built-in default. Setting it also makes child
failures soft: a child error returns this message as a normal tool result
instead of raising a parent `ModelRetry`."""
class SubAgentToolset(FunctionToolset[AgentDepsT]):
@@ -20,7 +58,7 @@ class SubAgentToolset(FunctionToolset[AgentDepsT]):
`deps` are forwarded; its `usage` is shared when enabled; its tools are
inherited when enabled; any `shared_capabilities` are applied to every
sub-agent run; and sub-agent events are streamed to `event_stream_handler`
when one is set.
when one is set. Per-delegate run controls come from `limits`.
"""
def __init__(
@@ -32,6 +70,8 @@ class SubAgentToolset(FunctionToolset[AgentDepsT]):
shared_capabilities: Sequence[AgentCapability[AgentDepsT]],
event_stream_handler: EventStreamHandler[AgentDepsT] | None,
tool_name: str,
limits: Mapping[str, SubAgentLimits],
call_counts: dict[str, dict[str, int]],
) -> None:
super().__init__()
self._agents: dict[str, AbstractAgent[AgentDepsT, Any]] = dict(agents)
@@ -40,6 +80,10 @@ class SubAgentToolset(FunctionToolset[AgentDepsT]):
self._shared_capabilities = list(shared_capabilities)
self._event_stream_handler = event_stream_handler
self._tool_name = tool_name
self._limits = dict(limits)
# Run-scoped delegation counts, keyed by run_id then sub-agent name.
# Shared with the capability, which clears each run's entry in wrap_run.
self._call_counts = call_counts
self.add_function(self.delegate_task, name=tool_name)
def _inherited_toolsets(self, ctx: RunContext[AgentDepsT]) -> list[AbstractToolset[AgentDepsT]] | None:
@@ -49,6 +93,16 @@ class SubAgentToolset(FunctionToolset[AgentDepsT]):
return None
return [toolset.filtered(lambda _ctx, tool_def: tool_def.name != self._tool_name) for toolset in agent.toolsets]
def _budget_exhausted(self, ctx: RunContext[AgentDepsT], agent_name: str, max_calls: int) -> bool:
"""Increment this run's delegation count for `agent_name` and report whether it is over budget.
Runs synchronously before any await, so concurrent delegations in one run
count without a lock.
"""
counts = self._call_counts.setdefault(ctx.run_id or '', {})
counts[agent_name] = counts.get(agent_name, 0) + 1
return counts[agent_name] > max_calls
async def delegate_task(self, ctx: RunContext[AgentDepsT], agent_name: str, task: str) -> str:
"""Delegate a self-contained task to a named sub-agent and return its result.
@@ -65,20 +119,70 @@ class SubAgentToolset(FunctionToolset[AgentDepsT]):
if agent is None:
available = ', '.join(sorted(self._agents))
raise ModelRetry(f'Unknown sub-agent {agent_name!r}. Available sub-agents: {available}.')
limits = self._limits.get(agent_name)
if (
limits is not None
and limits.max_calls is not None
and self._budget_exhausted(ctx, agent_name, limits.max_calls)
):
return self._steer(
limits,
f'Delegate budget for {agent_name!r} is exhausted for this run '
f'({limits.max_calls} call(s)). Synthesize from existing evidence and '
f'choose the next action; do not delegate to {agent_name!r} again.',
)
toolsets = self._inherited_toolsets(ctx) if self._inherit_tools else None
capabilities = self._shared_capabilities or None
usage = ctx.usage if self._forward_usage else None
usage_limits: UsageLimits | None
if limits is not None and limits.usage_limits is not None:
# Isolated accounting so the per-child budget counts only this child.
own_budget = True
usage = None
usage_limits = limits.usage_limits
else:
own_budget = False
usage = ctx.usage if self._forward_usage else None
usage_limits = None
run = agent.run(
task,
deps=ctx.deps,
usage=usage,
usage_limits=usage_limits,
toolsets=toolsets,
capabilities=capabilities,
event_stream_handler=self._event_stream_handler,
)
timeout = limits.timeout_seconds if limits is not None else None
try:
result = await agent.run(
task,
deps=ctx.deps,
usage=usage,
toolsets=toolsets,
capabilities=capabilities,
event_stream_handler=self._event_stream_handler,
result = await (asyncio.wait_for(run, timeout) if timeout is not None else run)
except asyncio.TimeoutError:
return self._steer(
limits,
f'Sub-agent {agent_name!r} exceeded its {timeout}s time budget. '
f'Treat this as a recoverable observation and decide from existing evidence.',
)
except UsageLimitExceeded:
if own_budget:
return self._steer(
limits,
f'Sub-agent {agent_name!r} reached its usage budget. '
f'Treat this as a recoverable observation and decide from existing evidence.',
)
# A shared/parent usage limit means the whole tree is out of budget.
raise
except (ModelRetry, UnexpectedModelBehavior) as exc:
# Soft sub-agent failures come back to the parent as a retry it can
# react to; hard limits (e.g. UsageLimitExceeded) propagate.
if limits is not None and limits.on_failure is not None:
return limits.on_failure
# Soft sub-agent failures come back to the parent as a retry it can react to.
raise ModelRetry(f'Sub-agent {agent_name!r} failed: {exc}') from exc
return str(result.output)
@staticmethod
def _steer(limits: SubAgentLimits | None, default: str) -> str:
"""A soft steering message: the delegate's `on_failure` override, else `default`."""
if limits is not None and limits.on_failure is not None:
return limits.on_failure
return default
tests/experimental/subagents/test_subagents.py
+262 -1
View File
@@ -2,6 +2,7 @@
from __future__ import annotations
import asyncio
from collections.abc import AsyncIterable
from dataclasses import dataclass, field
from typing import Any
@@ -22,8 +23,9 @@ from pydantic_ai.messages import (
from pydantic_ai.models.function import AgentInfo, FunctionModel
from pydantic_ai.models.test import TestModel
from pydantic_ai.tools import AgentDepsT, RunContext
from pydantic_ai.usage import UsageLimits
from pydantic_ai_harness.experimental.subagents import SubAgents, SubAgentToolset
from pydantic_ai_harness.experimental.subagents import SubAgentLimits, SubAgents, SubAgentToolset
@dataclass
@@ -76,6 +78,35 @@ def _delegate_then_finish(agent_name: str, *, retries_before: int = 0) -> Functi
return FunctionModel(model_fn)
def _delegate_n_then_finish(agent_name: str, n: int) -> FunctionModel:
"""A parent model that delegates to `agent_name` `n` times, then replies with text."""
calls = {'n': 0}
def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
calls['n'] += 1
if calls['n'] <= n:
return ModelResponse(
parts=[
ToolCallPart(
'delegate_task', {'agent_name': agent_name, 'task': 't'}, tool_call_id=f'c{calls["n"]}'
)
]
)
return ModelResponse(parts=[TextPart('all done')])
return FunctionModel(model_fn)
def _delegate_returns(result: Any) -> list[str]:
"""The `delegate_task` tool-return contents from a run result, in order."""
return [
str(part.content)
for message in result.all_messages()
for part in message.parts
if isinstance(part, ToolReturnPart) and part.tool_name == 'delegate_task'
]
class TestConstruction:
def test_serialization_name_is_none(self) -> None:
assert SubAgents.get_serialization_name() is None
@@ -324,3 +355,233 @@ class TestDelegation:
assert result.output == 'all done'
assert captured['deps'] == 'PARENT' # deps still forwarded
assert captured['usage_is_parent'] is False # usage isolated
class TestRunControls:
async def test_usage_limits_isolate_child_accounting(self) -> None:
captured: dict[str, Any] = {}
parent_usage: dict[str, Any] = {}
worker = Agent(TestModel(custom_output_text='W'), name='worker')
@worker.instructions
def _capture(ctx: RunContext[None]) -> str: # pyright: ignore[reportUnusedFunction]
captured['usage_is_parent'] = ctx.usage is parent_usage.get('usage')
return ''
parent: Agent[None, str] = Agent(
_delegate_then_finish('worker'),
capabilities=[
SubAgents(
agents={'worker': worker},
limits={'worker': SubAgentLimits(usage_limits=UsageLimits(request_limit=5))},
)
],
)
@parent.instructions
def _remember_usage(ctx: RunContext[None]) -> str: # pyright: ignore[reportUnusedFunction]
parent_usage['usage'] = ctx.usage
return ''
result = await parent.run('go')
assert result.output == 'all done'
# A per-child usage_limits forces isolated accounting even though forward_usage defaults to True.
assert captured['usage_is_parent'] is False
async def test_usage_budget_reached_is_soft(self) -> None:
counter = {'n': 0}
def worker_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
counter['n'] += 1
if counter['n'] == 1:
return ModelResponse(parts=[ToolCallPart('noop', {}, tool_call_id='n1')])
return ModelResponse(parts=[TextPart('done')]) # pragma: no cover - blocked by the request budget
worker = Agent(FunctionModel(worker_fn), name='worker')
@worker.tool_plain
def noop() -> str: # pyright: ignore[reportUnusedFunction]
return 'x'
parent: Agent[None, str] = Agent(
_delegate_then_finish('worker'),
capabilities=[
SubAgents(
agents={'worker': worker},
limits={'worker': SubAgentLimits(usage_limits=UsageLimits(request_limit=1))},
)
],
)
result = await parent.run('go')
assert result.output == 'all done'
returns = _delegate_returns(result)
# The child's own budget being hit is recoverable, not a run-stopping UsageLimitExceeded.
assert len(returns) == 1
assert "Sub-agent 'worker' reached its usage budget" in returns[0]
async def test_shared_usage_limit_still_propagates(self) -> None:
# No per-child limit -> the child shares accounting and a parent-level usage
# limit remains a hard stop for the whole tree.
worker = Agent(TestModel(custom_output_text='W'), name='worker')
parent: Agent[None, str] = Agent(
_delegate_then_finish('worker'),
capabilities=[SubAgents(agents={'worker': worker})],
)
with pytest.raises(UsageLimitExceeded):
await parent.run('go', usage_limits=UsageLimits(request_limit=1))
async def test_timeout_returns_soft_message(self) -> None:
async def slow_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
await asyncio.sleep(1)
return ModelResponse(parts=[TextPart('late')]) # pragma: no cover - cancelled by the timeout
worker = Agent(FunctionModel(slow_fn), name='worker')
parent: Agent[None, str] = Agent(
_delegate_then_finish('worker'),
capabilities=[
SubAgents(
agents={'worker': worker},
limits={'worker': SubAgentLimits(timeout_seconds=0.01)},
)
],
)
result = await parent.run('go')
assert result.output == 'all done'
returns = _delegate_returns(result)
assert len(returns) == 1
assert "Sub-agent 'worker' exceeded its 0.01s time budget" in returns[0]
async def test_max_calls_exhausted_returns_soft_and_skips_child(self) -> None:
runs = {'n': 0}
def worker_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
runs['n'] += 1
return ModelResponse(parts=[TextPart('W')])
worker = Agent(FunctionModel(worker_fn), name='worker')
parent: Agent[None, str] = Agent(
_delegate_n_then_finish('worker', 2),
capabilities=[SubAgents(agents={'worker': worker}, limits={'worker': SubAgentLimits(max_calls=1)})],
)
result = await parent.run('go')
assert result.output == 'all done'
returns = _delegate_returns(result)
assert len(returns) == 2
assert returns[0] == 'W'
assert "Delegate budget for 'worker' is exhausted" in returns[1]
assert runs['n'] == 1 # the over-budget delegation never ran the child
async def test_call_counts_reset_between_runs(self) -> None:
worker = Agent(TestModel(custom_output_text='W'), name='worker')
capability = SubAgents(agents={'worker': worker}, limits={'worker': SubAgentLimits(max_calls=1)})
# Two parents share the one capability (and its run-scoped count store); each
# gets a fresh delegate-once model so the second run actually delegates again.
first = await Agent(_delegate_then_finish('worker'), capabilities=[capability]).run('go')
second = await Agent(_delegate_then_finish('worker'), capabilities=[capability]).run('go')
# A fresh run starts the budget over: each delegation succeeds, none is exhausted.
assert _delegate_returns(first) == ['W']
assert _delegate_returns(second) == ['W']
# wrap_run clears each run's counts, so the store does not accumulate.
assert capability._call_counts == {}
async def test_on_failure_makes_child_failure_soft(self) -> None:
def boom(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
raise UnexpectedModelBehavior('kaboom')
boomer = Agent(FunctionModel(boom), name='boomer')
parent: Agent[None, str] = Agent(
_delegate_then_finish('boomer'),
capabilities=[
SubAgents(
agents={'boomer': boomer},
limits={'boomer': SubAgentLimits(on_failure='steer: use existing evidence')},
)
],
)
result = await parent.run('go')
assert result.output == 'all done'
# on_failure returns a normal tool result, so there is no RetryPrompt.
retries = [
part
for message in result.all_messages()
for part in message.parts
if isinstance(part, RetryPromptPart) and part.tool_name == 'delegate_task'
]
assert retries == []
assert _delegate_returns(result) == ['steer: use existing evidence']
async def test_on_failure_overrides_default_steering(self) -> None:
worker = Agent(TestModel(custom_output_text='W'), name='worker')
parent: Agent[None, str] = Agent(
_delegate_n_then_finish('worker', 2),
capabilities=[
SubAgents(
agents={'worker': worker},
limits={'worker': SubAgentLimits(max_calls=1, on_failure='custom budget note')},
)
],
)
result = await parent.run('go')
assert result.output == 'all done'
returns = _delegate_returns(result)
assert returns[1] == 'custom budget note'
async def test_limits_without_budget_run_normally(self) -> None:
# A limits entry with only an unrelated field set must not alter the happy path.
worker = Agent(TestModel(custom_output_text='W'), name='worker')
parent: Agent[None, str] = Agent(
_delegate_then_finish('worker'),
capabilities=[SubAgents(agents={'worker': worker}, limits={'worker': SubAgentLimits(timeout_seconds=30)})],
)
result = await parent.run('go')
assert _delegate_returns(result) == ['W']
async def test_max_calls_counts_parallel_delegations_in_one_step(self) -> None:
runs = {'n': 0}
def worker_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
runs['n'] += 1
return ModelResponse(parts=[TextPart('W')])
worker = Agent(FunctionModel(worker_fn), name='worker')
# Two delegations issued in a single parent model step run concurrently; the
# synchronous increment in _budget_exhausted must still cap them at max_calls=1.
step = {'n': 0}
def parent_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
step['n'] += 1
if step['n'] == 1:
return ModelResponse(
parts=[
ToolCallPart('delegate_task', {'agent_name': 'worker', 'task': 't'}, tool_call_id='a'),
ToolCallPart('delegate_task', {'agent_name': 'worker', 'task': 't'}, tool_call_id='b'),
]
)
return ModelResponse(parts=[TextPart('all done')])
parent: Agent[None, str] = Agent(
FunctionModel(parent_fn),
capabilities=[SubAgents(agents={'worker': worker}, limits={'worker': SubAgentLimits(max_calls=1)})],
)
result = await parent.run('go')
assert result.output == 'all done'
returns = _delegate_returns(result)
# Exactly one delegation ran the child; the other was over budget.
assert runs['n'] == 1
assert len(returns) == 2
assert 'W' in returns
assert any("Delegate budget for 'worker' is exhausted" in r for r in returns)
class TestLimitsValidation:
def test_limits_key_not_in_agents_raises(self) -> None:
worker = Agent(TestModel(), name='worker')
with pytest.raises(ValueError, match='names sub-agents not in `agents`: ghost'):
SubAgents(agents={'worker': worker}, limits={'ghost': SubAgentLimits(max_calls=1)})
def test_limits_subset_of_agents_is_accepted(self) -> None:
worker = Agent(TestModel(), name='worker')
capability = SubAgents(agents={'worker': worker}, limits={'worker': SubAgentLimits(max_calls=1)})
assert 'worker' in capability.limits