Report tool failures with ToolFailed

CHANGELOG.md

@@ -3,9 +3,10 @@
 
 ### Changed
 
-- Bump `pydantic-ai` to 2.18.0.
-- `enable_thinking` maps onto Pydantic AI's unified `thinking` setting for the `anthropic`, `gemini`, `groq` and `bedrock` providers. The Anthropic thinking budget is now Pydantic AI's default of 10000 tokens, was 4096; `enable_thinking: false` disables Gemini thinking rather than only hiding thoughts; Groq maps reasoning effort rather than `groq_reasoning_format`. The `openai` and `ollama` providers still map to `openai_reasoning_effort`.
+- Require `pydantic-ai-slim>=2.18,<3`.
+- `enable_thinking` maps onto Pydantic AI's unified `thinking` setting for the `anthropic`, `gemini`, `groq` and `bedrock` providers. The Anthropic thinking budget is now Pydantic AI's default of 10000 tokens, was 4096, and `max_tokens` must exceed it on budget-based Claude models; `enable_thinking: false` disables Gemini thinking rather than only hiding thoughts; Groq maps reasoning effort rather than `groq_reasoning_format`; Bedrock Qwen with `enable_thinking: false` no longer sends `reasoning_config`, while Bedrock-served Claude keeps an explicit `thinking: disabled`. The `openai` and `ollama` providers still map to `openai_reasoning_effort`.
 - `provider: bedrock` with a proprietary OpenAI model such as `openai.o3-mini-v1:0` raises `UserError`: Bedrock Converse serves only the `gpt-oss` family. Use `provider: bedrock-mantle`.
+- Tool failures raise `pydantic_ai.ToolFailed` instead of returning failure text: search and code-execution limits, sandbox execution errors, and `get_document`/`summarize_document` misses.
 
 ### Fixed
 

haiku_rag_slim/haiku/rag/capabilities/_base.py

@@ -5,7 +5,7 @@ from pathlib import Path
 from typing import Any, cast
 
 from pydantic import BaseModel
-from pydantic_ai import ModelRetry, RunContext
+from pydantic_ai import ModelRetry, RunContext, ToolFailed
 from pydantic_ai.capabilities import AbstractCapability
 from pydantic_ai.messages import (
     InstructionPart,
@@ -198,16 +198,21 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
             self.outer_state[self.state_namespace] = self.state.model_dump(mode="json")
 
     async def _with_state(self, operation: Any) -> Any:
-        """Execute an operation and copy its state back to the host dependencies."""
-        result = await operation
-        self._sync_state()
-        return result
+        """Execute an operation and copy its state back to the host dependencies.
+
+        A failing tool still syncs, so evidence it gathered before the failure
+        reaches the host.
+        """
+        try:
+            return await operation
+        finally:
+            self._sync_state()
 
     async def _search(self, query: str, limit: int | None) -> str | ToolReturn:
         assert self.state is not None
         self.search_count += 1
         if self.search_count > self.config.qa.max_searches:
-            return (
+            raise ToolFailed(
                 "Search limit reached. Answer the question using "
                 "the results you already have."
             )
@@ -227,7 +232,10 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
     async def _cite(self, chunk_ids: list[str]) -> str:
         assert self.state is not None
         if not chunk_ids:
-            return "Registered 0 citations (empty chunk_ids)."
+            raise ModelRetry(
+                "No citations registered: chunk_ids was empty. Pass the chunk_ids "
+                "you want to cite, copied verbatim from search results."
+            )
 
         all_results: list[SearchResult] = []
         state = cast(Any, self.state)

haiku_rag_slim/haiku/rag/capabilities/analysis.py

@@ -4,7 +4,7 @@ from pathlib import Path
 from typing import Any
 
 from pydantic import BaseModel, Field
-from pydantic_ai import RunContext
+from pydantic_ai import RunContext, ToolFailed
 from pydantic_ai.messages import ToolReturn
 from pydantic_ai.toolsets import FunctionToolset
 
@@ -74,7 +74,7 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
         assert self.state is not None
         self.execute_count += 1
         if self.execute_count > self.config.analysis.max_executions:
-            return (
+            raise ToolFailed(
                 "Code-execution limit reached. Give your final answer now from what "
                 "you already have; do not call analysis_execute_code again."
             )
@@ -96,9 +96,9 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
                 success=result.success,
             )
         )
-        if result.success:
-            return result.stdout or "No output."
-        return f"Error: {result.stderr}\n\nOutput: {result.stdout}"
+        if not result.success:
+            raise ToolFailed(f"{result.stderr}\n\nOutput: {result.stdout}")
+        return result.stdout or "No output."
 
     def get_toolset(self) -> FunctionToolset[Any]:
         async def analysis_search(

haiku_rag_slim/haiku/rag/tools/document.py

@@ -1,5 +1,5 @@
 from pydantic import BaseModel
-from pydantic_ai import Agent, FunctionToolset, RunContext
+from pydantic_ai import Agent, FunctionToolset, RunContext, ToolFailed
 
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config.models import AppConfig
@@ -123,14 +123,14 @@ def create_document_toolset(
             query: The document title or URI to look up.
 
         Returns:
-            Document content and metadata, or not found message.
+            Document content and metadata.
         """
         client = ctx.deps.client
 
         doc = await find_document(client, query)
 
         if doc is None:
-            return f"Document not found: {query}"
+            raise ToolFailed(f"Document not found: {query}")
 
         return (
             f"**{doc.title or 'Untitled'}**\n\n"
@@ -147,14 +147,14 @@ def create_document_toolset(
             query: The document title or URI to summarize.
 
         Returns:
-            Generated summary or not found message.
+            Generated summary.
         """
         client = ctx.deps.client
 
         doc = await find_document(client, query)
 
         if doc is None:
-            return f"Document not found: {query}"
+            raise ToolFailed(f"Document not found: {query}")
 
         summary_model = get_model(config.qa.model, config)
         summary_agent: Agent[None, str] = Agent(

haiku_rag_slim/haiku/rag/tools/search.py

@@ -3,7 +3,7 @@ from collections.abc import Callable
 from io import BytesIO
 
 from PIL import Image
-from pydantic_ai import FunctionToolset, RunContext
+from pydantic_ai import FunctionToolset, RunContext, ToolFailed
 from pydantic_ai.messages import BinaryContent, ToolReturn
 
 from haiku.rag.config.models import AppConfig
@@ -68,8 +68,9 @@ def create_search_toolset(
         tool_name: Name for the search tool. Defaults to "search".
         on_results: Optional callback invoked with search results after each search.
             Useful for accumulating results externally (e.g., for citation resolution).
-        max_searches: Maximum number of searches allowed. When exceeded, returns
-            a message directing the agent to answer with existing results.
+        max_searches: Maximum number of searches allowed. When exceeded, the
+            tool fails with a message directing the agent to answer with
+            existing results.
 
     Returns:
         FunctionToolset with a search tool.
@@ -99,7 +100,7 @@ def create_search_toolset(
         rid = ctx.run_id or ""
         search_counts[rid] = search_counts.get(rid, 0) + 1
         if max_searches is not None and search_counts[rid] > max_searches:
-            return (
+            raise ToolFailed(
                 "Search limit reached. "
                 "Answer the question using the results you already have."
             )

tests/capabilities/test_capabilities.py

@@ -4,7 +4,7 @@ from typing import Any, cast
 from unittest.mock import AsyncMock, patch
 
 import pytest
-from pydantic_ai import Agent, RunContext
+from pydantic_ai import Agent, ModelRetry, RunContext, ToolFailed
 from pydantic_ai.messages import (
     ModelRequest,
     ModelResponse,
@@ -291,13 +291,11 @@ async def test_search_and_empty_citation_limits(temp_db_path):
     capability = create_rag(db_path=temp_db_path, config=config)
     capability.state = RAGState()
 
-    result = await capability._search("anything", None)
+    with pytest.raises(ToolFailed, match="Search limit reached"):
+        await capability._search("anything", None)
 
-    assert (
-        result
-        == "Search limit reached. Answer the question using the results you already have."
-    )
-    assert await capability._cite([]) == "Registered 0 citations (empty chunk_ids)."
+    with pytest.raises(ModelRetry, match="chunk_ids was empty"):
+        await capability._cite([])
 
 
 @pytest.mark.asyncio
@@ -372,6 +370,75 @@ async def test_analysis_records_new_sandbox_search_results(temp_db_path):
     ]
 
 
+@pytest.mark.asyncio
+async def test_failed_tool_reaches_the_model_and_the_run_continues(temp_db_path):
+    """A `ToolFailed` tool leaves a failed result in history and answers anyway."""
+    config = AppConfig()
+    config.qa.max_searches = 0
+    calls = 0
+
+    def model_function(_messages, _info):
+        nonlocal calls
+        calls += 1
+        if calls == 1:
+            return ModelResponse(parts=[ToolCallPart("rag_search", {"query": "x"})])
+        return ModelResponse(parts=[TextPart("answered from what I had")])
+
+    agent = Agent(
+        FunctionModel(model_function),
+        deps_type=Deps,
+        capabilities=[
+            create_rag(db_path=temp_db_path, config=config, defer_loading=False)
+        ],
+    )
+
+    result = await agent.run("question", deps=Deps())
+
+    assert result.output == "answered from what I had"
+    failed = [
+        part
+        for message in result.all_messages()
+        for part in message.parts
+        if isinstance(part, ToolReturnPart) and part.outcome == "failed"
+    ]
+    assert [part.tool_name for part in failed] == ["rag_search"]
+    assert "Search limit reached" in str(failed[0].content)
+
+
+@pytest.mark.asyncio
+async def test_analysis_execution_limit_fails_the_tool(temp_db_path):
+    config = AppConfig()
+    config.analysis.max_executions = 0
+    capability = create_analysis(db_path=temp_db_path, config=config)
+    capability.state = AnalysisState()
+
+    with pytest.raises(ToolFailed, match="Code-execution limit reached"):
+        await capability._execute_code("print('done')")
+
+
+@pytest.mark.asyncio
+async def test_analysis_sandbox_failure_records_execution_and_fails_the_tool(
+    temp_db_path,
+):
+    capability = create_analysis(db_path=temp_db_path, config=AppConfig())
+    capability.state = AnalysisState()
+    capability.outer_state = {}
+    sandbox = AsyncMock()
+    sandbox.execute.return_value = SandboxResult(
+        stdout="partial", stderr="NameError: undefined", success=False
+    )
+    sandbox._search_results = []
+    capability.sandbox = cast(Sandbox, sandbox)
+
+    with pytest.raises(ToolFailed, match="NameError: undefined"):
+        await capability._with_state(capability._execute_code("boom"))
+
+    entry = capability.state.executions[-1]
+    assert entry.success is False
+    assert entry.stderr == "NameError: undefined"
+    assert capability.outer_state["analysis"]["executions"][-1]["code"] == "boom"
+
+
 @pytest.mark.asyncio
 async def test_native_agent_composition_initializes_host_state(temp_db_path):
     capability = create_rag(

tests/tools/test_document.py

@@ -2,6 +2,7 @@ from pathlib import Path
 from types import SimpleNamespace
 
 import pytest
+from pydantic_ai import ToolFailed
 
 from haiku.rag.tools.document import (
     DocumentInfo,
@@ -130,9 +131,9 @@ class TestDocumentToolExecution:
 
         get_tool = toolset.tools["get_document"]
         ctx = make_ctx(doc_client)
-        result = await get_tool.function(ctx, "nonexistent")
 
-        assert "Document not found" in result
+        with pytest.raises(ToolFailed, match="Document not found: nonexistent"):
+            await get_tool.function(ctx, "nonexistent")
 
     @pytest.mark.asyncio
     async def test_list_documents_with_base_filter(self, doc_client, doc_config):
@@ -183,9 +184,9 @@ class TestSummarizeDocumentTool:
 
         summarize_tool = toolset.tools["summarize_document"]
         ctx = make_ctx(doc_client)
-        result = await summarize_tool.function(ctx, "nonexistent document")
 
-        assert "Document not found" in result
+        with pytest.raises(ToolFailed, match="Document not found: nonexistent"):
+            await summarize_tool.function(ctx, "nonexistent document")
 
     @pytest.mark.vcr()
     @pytest.mark.asyncio

tests/tools/test_search.py

@@ -2,6 +2,7 @@ from pathlib import Path
 from types import SimpleNamespace
 
 import pytest
+from pydantic_ai import ToolFailed
 
 from haiku.rag.store.models import SearchResult
 from haiku.rag.tools.search import create_search_toolset
@@ -139,26 +140,22 @@ class TestSearchMaxSearches:
         search_tool = toolset.tools["search"]
         ctx = make_ctx(search_client)
 
-        result1 = await search_tool.function(ctx, "Python")
-        assert "Search limit reached" not in result1
-
-        result2 = await search_tool.function(ctx, "JavaScript")
-        assert "Search limit reached" not in result2
+        assert await search_tool.function(ctx, "Python")
+        assert await search_tool.function(ctx, "JavaScript")
 
     @pytest.mark.asyncio
-    async def test_searches_beyond_limit_return_cap_message(
+    async def test_searches_beyond_limit_fail_the_tool(
         self, search_client, search_config
     ):
-        """Searches beyond max_searches return limit message."""
+        """Searches beyond max_searches fail with the limit message."""
         toolset = create_search_toolset(search_config, max_searches=1)
         search_tool = toolset.tools["search"]
         ctx = make_ctx(search_client)
 
-        result1 = await search_tool.function(ctx, "Python")
-        assert "Search limit reached" not in result1
+        assert await search_tool.function(ctx, "Python")
 
-        result2 = await search_tool.function(ctx, "JavaScript")
-        assert "Search limit reached" in result2
+        with pytest.raises(ToolFailed, match="Search limit reached"):
+            await search_tool.function(ctx, "JavaScript")
 
     @pytest.mark.asyncio
     async def test_counter_resets_across_runs(self, search_client, search_config):
@@ -167,15 +164,13 @@ class TestSearchMaxSearches:
         search_tool = toolset.tools["search"]
 
         ctx_run1 = make_ctx(search_client, run_id="run-1")
-        result = await search_tool.function(ctx_run1, "Python")
-        assert "Search limit reached" not in result
+        assert await search_tool.function(ctx_run1, "Python")
 
-        result2 = await search_tool.function(ctx_run1, "JavaScript")
-        assert "Search limit reached" in result2
+        with pytest.raises(ToolFailed, match="Search limit reached"):
+            await search_tool.function(ctx_run1, "JavaScript")
 
         ctx_run2 = make_ctx(search_client, run_id="run-2")
-        result3 = await search_tool.function(ctx_run2, "Python")
-        assert "Search limit reached" not in result3
+        assert await search_tool.function(ctx_run2, "Python")
 
     @pytest.mark.asyncio
     async def test_no_limit_by_default(self, search_client, search_config):
@@ -185,8 +180,7 @@ class TestSearchMaxSearches:
         ctx = make_ctx(search_client)
 
         for _ in range(5):
-            result = await search_tool.function(ctx, "Python")
-            assert "Search limit reached" not in result
+            assert await search_tool.function(ctx, "Python")
 
 
 @pytest.fixture