This is an automated email from the ASF dual-hosted git repository. vikramkoka pushed a commit to branch common_ai_managed_toolset in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 7bbbd2fdacce2cc81679c6d57a770a1091fc6663 Author: Vikram Koka <[email protected]> AuthorDate: Wed Aug 12 21:45:31 2026 -0700 Add BaseManagedAgentToolset for vendor-managed AI agents Cloud vendors now run agents on the user's behalf -- Snowflake Cortex Agents, Amazon Bedrock AgentCore, Azure AI Foundry hosted agents, Vertex AI Agent Engine. Their reasoning loops execute on the vendor's infrastructure, so they are not something AgentOperator can run, only something an Airflow task can consult. Each provider that has added support so far arrived at a different identity scheme, request shape, and invoke verb, so a Dag author who learns one learns nothing transferable about the next. Rather than unify them behind an operator facade -- their invocation semantics genuinely differ, from sync request/response to async jobs writing to object storage to stateful sessions -- converge them at the tool boundary, where a single call with validated arguments and a serialisable result is the same shape everywhere. Provider packages subclass this contract so credentials keep flowing through their own hooks and no new connection types are needed. Errors are deliberately sorted into three buckets rather than two: what the calling model can fix by rephrasing, what is terminal, and what is transient and belongs to Airflow's task retry. Collapsing the third into either of the others is how these implementations go wrong. Durable replay is opt-in per implementation, since a managed agent may act on systems Airflow cannot observe and replaying a cached answer could skip a side effect. Also promotes the private tool-result serialiser out of HookToolset into the shared tool_definition utils so both toolsets use one implementation. --- providers/common/ai/docs/toolsets.rst | 86 ++++++++++- providers/common/ai/provider.yaml | 1 + .../src/airflow/providers/common/ai/exceptions.py | 11 ++ .../providers/common/ai/toolsets/__init__.py | 11 +- .../airflow/providers/common/ai/toolsets/hook.py | 23 +-- .../providers/common/ai/toolsets/managed_agent.py | 168 +++++++++++++++++++++ .../providers/common/ai/utils/tool_definition.py | 17 +++ .../ai/tests/unit/common/ai/toolsets/test_hook.py | 16 +- .../unit/common/ai/toolsets/test_managed_agent.py | 164 ++++++++++++++++++++ 9 files changed, 471 insertions(+), 26 deletions(-) diff --git a/providers/common/ai/docs/toolsets.rst b/providers/common/ai/docs/toolsets.rst index 34a8b43f417..eb393713bd8 100644 --- a/providers/common/ai/docs/toolsets.rst +++ b/providers/common/ai/docs/toolsets.rst @@ -34,7 +34,13 @@ Three toolsets are included: `MCP servers <https://modelcontextprotocol.io/>`__ configured via Airflow connections. -All three implement pydantic-ai's +A fourth, +:class:`~airflow.providers.common.ai.toolsets.managed_agent.BaseManagedAgentToolset`, +is a base class rather than a usable toolset: provider packages subclass it to +expose a **vendor-managed agent** as a tool. See +:ref:`managed-agent-toolsets` below. + +All of them implement pydantic-ai's `AbstractToolset <https://ai.pydantic.dev/toolsets/>`__ interface and can be passed to any pydantic-ai ``Agent``, including via :class:`~airflow.providers.common.ai.operators.agent.AgentOperator`. @@ -839,3 +845,81 @@ Before deploying an agent task to production: 8. **Prompt injection**: Be cautious when the prompt includes untrusted data (user input, external API responses, upstream XCom). Consider sanitizing inputs before passing them to the agent. + +.. _managed-agent-toolsets: + +Managed Agent Toolsets +---------------------- + +Cloud vendors now run agents on your behalf — Snowflake Cortex Agents, Amazon +Bedrock AgentCore runtimes, Azure AI Foundry hosted agents, Vertex AI Agent +Engine. Their reasoning loops execute on the vendor's infrastructure, so they +are not something ``AgentOperator`` runs; they are something an Airflow task +*consults*. + +:class:`~airflow.providers.common.ai.toolsets.managed_agent.BaseManagedAgentToolset` +is the contract for exposing one of those as a tool. Each provider package +ships its own subclass, so credentials keep flowing through that provider's +existing hook and no new connection types are needed. + +A subclass implements two members: + +``agent_ref`` + Normalised identity of the remote agent — ``platform`` and ``name`` — logged + on every call so a run can be audited for which agents it consulted. + +``invoke(prompt)`` + Send the prompt, return the agent's answer. Return the *answer*, not the + transport envelope. + +Tool naming, argument validation, result serialisation, and logging are handled +by the base class, so every provider's implementation presents the same surface +to the calling model. + +.. note:: + + ``description`` is a required constructor argument. A remote agent's + competence cannot be introspected the way ``HookToolset`` reads a hook's + docstrings, and the description is the only basis the calling model has for + choosing between specialists. + +Error handling +"""""""""""""" + +Failures sort into three buckets, and conflating them is the most common way an +implementation goes wrong: + +.. list-table:: + :header-rows: 1 + :widths: 22 30 48 + + * - Raise + - When + - Who recovers + * - ``ModelRetry`` + - The agent rejected the request in a way rephrasing could fix. + - The calling model, bounded by its ``usage_limits``. + * - ``ManagedAgentInvocationError`` + - Terminal: bad credentials, missing agent, revoked quota. + - Nobody — the task fails fast instead of burning retries. + * - *let it propagate* + - Transient: 429, 5xx, connection reset, read timeout. + - Airflow's task-level retry. A rephrase does nothing for a 503. + +Durable execution +""""""""""""""""" + +``replayable`` is ``False`` by default. A managed agent may act on systems +Airflow cannot observe, so replaying a cached answer on retry could skip a side +effect. Implementations whose agent is read-only should set it to ``True`` to +avoid paying for the same invocation twice. + +Deferral +"""""""" + +A toolset call runs in the worker and cannot defer to the triggerer. Managed +agents with no request/response mode — an async job that writes output to +object storage, or a stateful session — are reachable through a toolset only by +blocking for the duration. That is acceptable for a tool call inside an agent's +reasoning, but it is not a substitute for that provider's own deferrable +operator when the Dag simply needs to submit work and wait. diff --git a/providers/common/ai/provider.yaml b/providers/common/ai/provider.yaml index b6133e2e2e4..5335de6179b 100644 --- a/providers/common/ai/provider.yaml +++ b/providers/common/ai/provider.yaml @@ -490,6 +490,7 @@ toolsets: - airflow.providers.common.ai.toolsets.mcp - airflow.providers.common.ai.toolsets.skills - airflow.providers.common.ai.toolsets.langchain_bridge + - airflow.providers.common.ai.toolsets.managed_agent retry-policies: - integration-name: Common AI diff --git a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py index d0412fb2a48..d02748d4a3a 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py +++ b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py @@ -37,3 +37,14 @@ class LLMFileAnalysisLimitExceededError(LLMFileAnalysisError): class LLMFileAnalysisMultimodalRequiredError(LLMFileAnalysisUnsupportedFormatError): """Raised when image/PDF inputs are used without ``multi_modal=True``.""" + + +class ManagedAgentInvocationError(RuntimeError): + """ + Raised when a managed agent cannot be reached and retrying will not help. + + Reserved for terminal conditions -- bad credentials, a missing agent, a + revoked quota. Transient failures should propagate unchanged so Airflow's + task-level retry handles them, and requests the model could fix by + rephrasing should raise ``pydantic_ai.exceptions.ModelRetry`` instead. + """ diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py index 6c30fa4a733..362692de249 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py @@ -19,8 +19,15 @@ from __future__ import annotations from airflow.providers.common.ai.toolsets.hook import HookToolset - -__all__ = ["HookToolset", "MCPToolset", "SQLToolset", "airflow_toolset_to_langchain_tools"] +from airflow.providers.common.ai.toolsets.managed_agent import BaseManagedAgentToolset + +__all__ = [ + "BaseManagedAgentToolset", + "HookToolset", + "MCPToolset", + "SQLToolset", + "airflow_toolset_to_langchain_tools", +] def __getattr__(name: str): diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py index 63037e1a4f8..13412b82c35 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py @@ -19,7 +19,6 @@ from __future__ import annotations import inspect -import json import re import types from typing import TYPE_CHECKING, Any, Union, get_args, get_origin, get_type_hints @@ -27,7 +26,11 @@ from typing import TYPE_CHECKING, Any, Union, get_args, get_origin, get_type_hin from pydantic_ai.tools import ToolDefinition from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool -from airflow.providers.common.ai.utils.tool_definition import build_args_validator, return_schema_kwargs +from airflow.providers.common.ai.utils.tool_definition import ( + build_args_validator, + return_schema_kwargs, + serialize_for_llm, +) if TYPE_CHECKING: from collections.abc import Callable @@ -109,7 +112,7 @@ class HookToolset(AbstractToolset[Any]): # sequential=True because hook methods perform synchronous I/O # (network calls, DB queries) and should not run concurrently. # return_schema is "string": call_tool serializes every result with - # _serialize_for_llm, so the tool always returns a (JSON-encoded) + # serialize_for_llm, so the tool always returns a (JSON-encoded) # string regardless of the method's own return annotation. This lets # code mode render `-> str` instead of `-> Any`. tool_def = ToolDefinition( @@ -137,7 +140,7 @@ class HookToolset(AbstractToolset[Any]): method_name = name.removeprefix(self._tool_name_prefix) if self._tool_name_prefix else name method: Callable[..., Any] = getattr(self._hook, method_name) result = method(**tool_args) - return _serialize_for_llm(result) + return serialize_for_llm(result) # --------------------------------------------------------------------------- @@ -260,15 +263,3 @@ def _parse_param_docs(docstring: str) -> dict[str, str]: params[m.group(1)] = " ".join(m.group(2).split()) return params - - -def _serialize_for_llm(value: Any) -> str: - """Convert a Python return value to a string suitable for an LLM.""" - if value is None: - return "null" - if isinstance(value, str): - return value - try: - return json.dumps(value, default=str) - except (TypeError, ValueError): - return str(value) diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/managed_agent.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/managed_agent.py new file mode 100644 index 00000000000..03ff1c54f3c --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/managed_agent.py @@ -0,0 +1,168 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import logging +from abc import abstractmethod +from typing import TYPE_CHECKING, Any + +from pydantic_ai.tools import ToolDefinition +from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool + +from airflow.providers.common.ai.utils.tool_definition import ( + build_args_validator, + return_schema_kwargs, + serialize_for_llm, +) + +if TYPE_CHECKING: + from pydantic_ai._run_context import RunContext + +log = logging.getLogger(__name__) + +_PROMPT_SCHEMA: dict[str, Any] = { + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "The question or instruction to send to this agent.", + } + }, + "required": ["prompt"], +} + + +class BaseManagedAgentToolset(AbstractToolset[Any]): + """ + Base class exposing a vendor-managed agent as a single pydantic-ai tool. + + A managed agent runs its own reasoning loop on the vendor's infrastructure + (Snowflake Cortex Agents, Amazon Bedrock AgentCore, Azure AI Foundry hosted + agents, Vertex AI Agent Engine). Airflow submits one request and reads one + answer, so the Airflow-side agent features -- toolsets, human-in-the-loop + review, durable step replay -- apply to the *calling* agent and never reach + inside the managed agent. + + Subclasses implement :meth:`agent_ref` and :meth:`invoke`. Tool naming, + argument validation, result serialisation and logging are handled here so + every provider's implementation presents the same surface to the model. + + :param tool_name: Name the calling model sees. A verb phrase naming the + specialist reads best, e.g. ``ask_bookings_analyst``. + :param description: What this agent knows and when to consult it. Required: + a remote agent's competence cannot be introspected, and this is the only + basis the calling model has for choosing between specialists. + :param timeout: Seconds to wait for a single invocation. ``None`` defers to + the platform default, which subclasses supply -- a number chosen here + would silently disagree with the vendor operator's documented timeout + for the same service. + """ + + #: Whether a completed invocation may be replayed from the durable cache + #: instead of re-invoked. Off by default because a managed agent may act on + #: systems Airflow cannot observe, so replaying a cached answer could skip a + #: side effect. Read-only agents should opt in. + replayable: bool = False + + def __init__( + self, + *, + tool_name: str, + description: str, + timeout: float | None = None, + ) -> None: + if not tool_name: + raise ValueError("tool_name must be a non-empty string.") + if not description or not description.strip(): + raise ValueError( + "description is required: the calling model uses it to decide which " + "specialist to consult, and it cannot be derived from the agent's identifier." + ) + self._tool_name = tool_name + self._description = description + self._timeout = timeout + + @property + @abstractmethod + def agent_ref(self) -> dict[str, str]: + """ + Normalised identity of the remote agent. + + Must contain ``platform`` and ``name``, e.g. + ``{"platform": "snowflake.cortex", "name": "ANALYTICS.REVENUE.BOOKINGS_ANALYST"}``. + Logged on every invocation so a run can be audited for which agents were + consulted. + """ + + @abstractmethod + async def invoke(self, prompt: str) -> Any: + """ + Send ``prompt`` to the remote agent and return the agent's answer. + + Return the answer, not the transport envelope -- whatever the calling + model should actually read. Unwrapping is the implementation's job. + + Failures sort into three buckets, and conflating them is the most common + way an implementation goes wrong: + + * ``pydantic_ai.exceptions.ModelRetry`` -- the remote agent rejected the + request in a way rephrasing could fix. The calling model sees the + message and tries again, bounded by its ``usage_limits``. + * :class:`~airflow.providers.common.ai.exceptions.ManagedAgentInvocationError` + -- terminal. Bad credentials, missing agent, revoked quota. Neither a + rephrase nor a task retry helps, so fail fast. + * Anything transient (429, 5xx, connection reset, read timeout) -- let it + propagate unchanged. Airflow's task-level retry is the right layer; a + rephrase does nothing for a 503. + + :param prompt: The question or instruction to send to the remote agent. + """ + + @property + def id(self) -> str: + return f"managed-agent-{self._tool_name}" + + async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]: + tool_def = ToolDefinition( + name=self._tool_name, + description=self._description, + parameters_json_schema=_PROMPT_SCHEMA, + # Each invocation is an independent request to a remote service, so + # consulting several specialists concurrently is both safe and the point. + sequential=False, + **return_schema_kwargs({"type": "string"}), + ) + return { + self._tool_name: ToolsetTool( + toolset=self, + tool_def=tool_def, + max_retries=1, + args_validator=build_args_validator(_PROMPT_SCHEMA), + ) + } + + async def call_tool( + self, + name: str, + tool_args: dict[str, Any], + ctx: RunContext[Any], + tool: ToolsetTool[Any], + ) -> Any: + ref = self.agent_ref + log.info("Consulting managed agent %s on %s", ref.get("name"), ref.get("platform")) + result = await self.invoke(tool_args["prompt"]) + return serialize_for_llm(result) diff --git a/providers/common/ai/src/airflow/providers/common/ai/utils/tool_definition.py b/providers/common/ai/src/airflow/providers/common/ai/utils/tool_definition.py index 9f4fb04b87b..28639701d94 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/utils/tool_definition.py +++ b/providers/common/ai/src/airflow/providers/common/ai/utils/tool_definition.py @@ -19,6 +19,7 @@ from __future__ import annotations import dataclasses +import json from typing import Any, Literal from pydantic_ai.tools import ToolDefinition @@ -45,6 +46,22 @@ def return_schema_kwargs(schema: dict[str, Any]) -> dict[str, Any]: return {} +def serialize_for_llm(value: Any) -> str: + """ + Convert a Python return value to a string suitable for an LLM. + + :param value: The tool's return value. + """ + if value is None: + return "null" + if isinstance(value, str): + return value + try: + return json.dumps(value, default=str) + except (TypeError, ValueError): + return str(value) + + def _fragment_to_core_schema(fragment: dict[str, Any]) -> core_schema.CoreSchema: any_of = fragment.get("anyOf") if isinstance(any_of, list): diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py b/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py index ae2d4c6f867..b8ce959db02 100644 --- a/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py +++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py @@ -27,9 +27,11 @@ from airflow.providers.common.ai.toolsets.hook import ( _build_json_schema_from_signature, _extract_description, _parse_param_docs, - _serialize_for_llm, ) -from airflow.providers.common.ai.utils.tool_definition import _SUPPORTS_RETURN_SCHEMA +from airflow.providers.common.ai.utils.tool_definition import ( + _SUPPORTS_RETURN_SCHEMA, + serialize_for_llm, +) class _FakeHook: @@ -320,20 +322,20 @@ class TestParseParamDocs: class TestSerializeForLlm: def test_string_passthrough(self): - assert _serialize_for_llm("hello") == "hello" + assert serialize_for_llm("hello") == "hello" def test_none_returns_null(self): - assert _serialize_for_llm(None) == "null" + assert serialize_for_llm(None) == "null" def test_dict_to_json(self): - result = _serialize_for_llm({"key": "value"}) + result = serialize_for_llm({"key": "value"}) assert result == '{"key": "value"}' def test_list_to_json(self): - result = _serialize_for_llm([1, 2, 3]) + result = serialize_for_llm([1, 2, 3]) assert result == "[1, 2, 3]" def test_non_serializable_falls_back_to_str(self): obj = object() - result = _serialize_for_llm(obj) + result = serialize_for_llm(obj) assert "object" in result diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_managed_agent.py b/providers/common/ai/tests/unit/common/ai/toolsets/test_managed_agent.py new file mode 100644 index 00000000000..cfc7904badd --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_managed_agent.py @@ -0,0 +1,164 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from typing import Any + +import pytest +from pydantic_ai.exceptions import ModelRetry +from pydantic_core import ValidationError + +from airflow.providers.common.ai.exceptions import ManagedAgentInvocationError +from airflow.providers.common.ai.toolsets.managed_agent import BaseManagedAgentToolset + + +class FakeManagedAgentToolset(BaseManagedAgentToolset): + """Minimal implementation standing in for a provider's concrete toolset.""" + + def __init__(self, *, result: Any = "the answer", raises: Exception | None = None, **kwargs): + kwargs.setdefault("tool_name", "ask_specialist") + kwargs.setdefault("description", "Answers questions about the thing.") + super().__init__(**kwargs) + self._result = result + self._raises = raises + self.prompts: list[str] = [] + + @property + def agent_ref(self) -> dict[str, str]: + return {"platform": "fake.cloud", "name": "specialist-1"} + + async def invoke(self, prompt: str) -> Any: + self.prompts.append(prompt) + if self._raises is not None: + raise self._raises + return self._result + + +class TestBaseManagedAgentToolsetConstruction: + def test_is_abstract(self): + with pytest.raises(TypeError, match="abstract"): + BaseManagedAgentToolset(tool_name="x", description="y") # type: ignore[abstract] + + @pytest.mark.parametrize( + "description", + ["", " ", "\n"], + ids=["empty", "whitespace", "newline"], + ) + def test_blank_description_rejected(self, description): + with pytest.raises(ValueError, match="description is required"): + FakeManagedAgentToolset(description=description) + + def test_empty_tool_name_rejected(self): + with pytest.raises(ValueError, match="tool_name must be a non-empty string"): + FakeManagedAgentToolset(tool_name="") + + def test_timeout_defaults_to_none_so_subclasses_supply_it(self): + assert FakeManagedAgentToolset()._timeout is None + + def test_subclass_may_default_its_own_timeout(self): + class WithPlatformDefault(FakeManagedAgentToolset): + def __init__(self, **kwargs): + kwargs.setdefault("timeout", 600.0) + super().__init__(**kwargs) + + assert WithPlatformDefault()._timeout == 600.0 + + def test_id_is_derived_from_tool_name(self): + assert FakeManagedAgentToolset(tool_name="ask_bookings").id == "managed-agent-ask_bookings" + + def test_not_replayable_by_default(self): + assert FakeManagedAgentToolset().replayable is False + + +class TestGetTools: + @pytest.mark.asyncio + async def test_exposes_exactly_one_tool_under_its_name(self): + toolset = FakeManagedAgentToolset(tool_name="ask_bookings") + tools = await toolset.get_tools(ctx=None) + assert list(tools) == ["ask_bookings"] + + @pytest.mark.asyncio + async def test_tool_definition_carries_name_description_and_prompt_schema(self): + toolset = FakeManagedAgentToolset(tool_name="ask_bookings", description="Knows bookings.") + tool_def = (await toolset.get_tools(ctx=None))["ask_bookings"].tool_def + + assert tool_def.name == "ask_bookings" + assert tool_def.description == "Knows bookings." + assert tool_def.parameters_json_schema["required"] == ["prompt"] + assert tool_def.parameters_json_schema["properties"]["prompt"]["type"] == "string" + + @pytest.mark.asyncio + async def test_not_sequential_so_specialists_can_be_consulted_concurrently(self): + toolset = FakeManagedAgentToolset() + tool_def = (await toolset.get_tools(ctx=None))["ask_specialist"].tool_def + assert tool_def.sequential is False + + @pytest.mark.asyncio + async def test_args_validator_rejects_a_missing_prompt(self): + toolset = FakeManagedAgentToolset() + validator = (await toolset.get_tools(ctx=None))["ask_specialist"].args_validator + with pytest.raises(ValidationError): + validator.validate_json("{}") + + +class TestCallTool: + async def _call(self, toolset, prompt="what is the number?"): + tools = await toolset.get_tools(ctx=None) + tool = tools[toolset._tool_name] + return await toolset.call_tool(toolset._tool_name, {"prompt": prompt}, None, tool) + + @pytest.mark.asyncio + async def test_passes_the_prompt_through_to_invoke(self): + toolset = FakeManagedAgentToolset() + await self._call(toolset, prompt="how many widgets?") + assert toolset.prompts == ["how many widgets?"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("result", "expected"), + [ + ("plain text", "plain text"), + (None, "null"), + ({"total": 42}, '{"total": 42}'), + ([1, 2], "[1, 2]"), + ], + ids=["str", "none", "dict", "list"], + ) + async def test_result_is_serialised_for_the_model(self, result, expected): + assert await self._call(FakeManagedAgentToolset(result=result)) == expected + + @pytest.mark.asyncio + async def test_logs_the_agent_it_consulted(self, caplog): + await self._call(FakeManagedAgentToolset()) + assert "specialist-1" in caplog.text + assert "fake.cloud" in caplog.text + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "error", + [ + ModelRetry("rephrase that"), + ManagedAgentInvocationError("bad credentials"), + RuntimeError("503 from upstream"), + ], + ids=["model_retry", "terminal", "transient"], + ) + async def test_invoke_errors_propagate_unchanged(self, error): + # The base class must not reclassify what invoke() raised: the three + # buckets are handled by different layers (model, task failure, task retry). + with pytest.raises(type(error), match=str(error)): + await self._call(FakeManagedAgentToolset(raises=error))
