From 18f2e0efbd410682ac486009bb95e05a873e2384 Mon Sep 17 00:00:00 2001 From: Jeroen Schweitzer Date: Tue, 14 Jul 2026 10:19:16 +0200 Subject: [PATCH] feat(agents): enforce one librarian timeout budget - add LIBRARIAN_TIMEOUT config (default 180s) and enforce it with asyncio.wait_for inside delegate_to_librarian, covering the live paths (steward direct delegation and SSE streaming) that had no cap - timeouts fail honestly: success=False with a curated butler sentence, detail in logs - set an explicit timeout on TatlockOllamaProvider's AsyncOpenAI client from OLLAMA_TIMEOUT instead of the SDK default (~600s per LLM call) - remove the contradictory unused 60s default from AgentRequest.timeout_seconds; coordination falls back to the configured budget Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 1 + src/agents/coordination.py | 5 +-- src/agents/delegation.py | 36 ++++++++++++++++++++-- src/agents/protocol.py | 15 +++++---- src/core/config.py | 14 ++++++--- src/ollama/provider.py | 11 +++++-- tests/agents/test_delegation.py | 24 +++++++++++++++ tests/agents/test_protocol.py | 3 +- tests/ollama/test_provider.py | 54 +++++++++++++++++++++++++++++++++ 9 files changed, 145 insertions(+), 18 deletions(-) create mode 100644 tests/ollama/test_provider.py diff --git a/CHANGELOG.md b/CHANGELOG.md index f9c834c..fc95e34 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- **One librarian timeout budget** - new `LIBRARIAN_TIMEOUT` (default 180s) enforced with `asyncio.wait_for` inside `delegate_to_librarian`, capping the previously uncapped live paths (steward direct delegation and streaming). The Ollama provider's AsyncOpenAI client now carries an explicit `OLLAMA_TIMEOUT` instead of the SDK's ~600s default, and the contradictory unused 60s default in `AgentRequest.timeout_seconds` was removed (None defers to the configured budget) - **Search degradation signaling** - The librarian client parses `source_counts` (plus the additive `source_status`/`degraded` fields when a newer library-desk sends them; absence is tolerated), and `hybrid_search` appends a one-line coverage note when a search is degraded or an enabled source leg contributed nothing, so outages are visible to the model and the user ### Fixed diff --git a/src/agents/coordination.py b/src/agents/coordination.py index 0215692..1039712 100644 --- a/src/agents/coordination.py +++ b/src/agents/coordination.py @@ -24,6 +24,7 @@ from src.agents.protocol import ( DelegationIntent, DelegationReason, ) +from src.core.config import config from src.core.household_registry import get_household_registry from src.core.logging_config import get_logger @@ -138,8 +139,8 @@ class CoordinationEngine: delegation_reason=intent.reason, ) - # Execute with timeout - timeout = request.timeout_seconds or 60 + # Execute with timeout (explicit request value or configured budget) + timeout = request.timeout_seconds or config.LIBRARIAN_TIMEOUT result = await asyncio.wait_for( executor( diff --git a/src/agents/delegation.py b/src/agents/delegation.py index 462aab9..8bd82b0 100644 --- a/src/agents/delegation.py +++ b/src/agents/delegation.py @@ -8,10 +8,12 @@ returns a structured result for synthesis. This implements the agent-as-tool pattern recommended by PydanticAI: agents call other agents via tool wrappers, keeping each agent focused. """ +import asyncio from collections.abc import AsyncGenerator from dataclasses import dataclass, field from enum import Enum +from src.core.config import config from src.core.logging_config import get_logger from src.core.tracing import SpanType, trace_span @@ -252,8 +254,14 @@ async def delegate_to_librarian( }, ) as span: try: - # Use run() not run_stream() - avoids Ollama bug - output = await run_librarian(task=task, context=context) + # Use run() not run_stream() - avoids Ollama bug. + # One timeout budget for the whole delegation - covers both + # live paths (steward direct delegation and streaming), which + # previously had no cap at all (SDK default ~600s per LLM call). + output = await asyncio.wait_for( + run_librarian(task=task, context=context), + timeout=config.LIBRARIAN_TIMEOUT, + ) logger.info( "delegation_to_librarian_completed", @@ -275,6 +283,30 @@ async def delegate_to_librarian( output=output, ) + except TimeoutError: + logger.error( + "delegation_to_librarian_timeout", + task=task[:50], + timeout_seconds=config.LIBRARIAN_TIMEOUT, + ) + + if span: + span.metadata["success"] = False + span.details["error"] = ( + f"timed out after {config.LIBRARIAN_TIMEOUT}s" + ) + + return DelegationResult( + expert_name="librarian", + task=task, + success=False, + output=( + "I'm afraid the research took longer than expected " + "and had to be abandoned, sir." + ), + error="The Librarian did not respond within the time budget.", + ) + except Exception as e: logger.error( "delegation_to_librarian_error", diff --git a/src/agents/protocol.py b/src/agents/protocol.py index 70dc00b..6161a5a 100644 --- a/src/agents/protocol.py +++ b/src/agents/protocol.py @@ -6,7 +6,7 @@ Defines standardized request/response formats for communication between: - Tatlock (coordination) → Expert agents (Librarian, Developer, etc.) """ from enum import Enum -from typing import Any, Optional +from typing import Any from pydantic import BaseModel, Field @@ -53,13 +53,16 @@ class AgentRequest(BaseModel): default="default", description="User identifier for multi-tenant operations" ) - max_tokens: Optional[int] = Field( + max_tokens: int | None = Field( default=None, description="Optional token limit for response" ) - timeout_seconds: Optional[int] = Field( - default=60, - description="Maximum time for task completion" + timeout_seconds: int | None = Field( + default=None, + description=( + "Maximum time for task completion; None uses the configured " + "expert budget (LIBRARIAN_TIMEOUT)" + ) ) @@ -103,7 +106,7 @@ class AgentResponse(BaseModel): default_factory=list, description="Sources or references used" ) - error_message: Optional[str] = Field( + error_message: str | None = Field( default=None, description="Error details if success=False" ) diff --git a/src/core/config.py b/src/core/config.py index 0970c12..2d3448d 100644 --- a/src/core/config.py +++ b/src/core/config.py @@ -42,7 +42,7 @@ class Environment(str, Enum): class Config(BaseSettings): """ Global application configuration. - + Loads from environment variables and .env file. Domain-specific configs should be in their respective modules. """ @@ -52,18 +52,18 @@ class Config(BaseSettings): case_sensitive=True, extra="ignore", ) - + # Application APP_NAME: str = "OpenAI-Compatible API" APP_VERSION: str = Field(default_factory=_get_version_from_pyproject) ENVIRONMENT: Environment = Environment.DEVELOPMENT DEBUG: bool = Field(default=False, description="Debug mode") - + # API Configuration API_HOST: str = Field(default="0.0.0.0", description="API host") API_PORT: int = Field(default=8000, description="API port") API_PREFIX: str = Field(default="/v1", description="API route prefix") - + # Anthropic Configuration (Claude - cloud fallback) ANTHROPIC_API_KEY: str | None = Field( default=None, @@ -125,6 +125,10 @@ class Config(BaseSettings): ) # Library-Desk Configuration (The Librarian backend) + LIBRARIAN_TIMEOUT: int = Field( + default=180, + description="Total time budget for a librarian delegation in seconds" + ) LIBRARY_DESK_HOST: HttpUrl = Field( default="http://localhost:8089", description="Library-Desk API URL" @@ -259,7 +263,7 @@ class Config(BaseSettings): def get_config() -> Config: """ Get cached configuration instance. - + Uses lru_cache to ensure config is loaded once and reused. """ return Config() diff --git a/src/ollama/provider.py b/src/ollama/provider.py index 7b04e48..d93485b 100644 --- a/src/ollama/provider.py +++ b/src/ollama/provider.py @@ -40,14 +40,21 @@ class TatlockOllamaProvider(OllamaProvider): # Override the client with our sanitized version self._openai_client = _SanitizedAsyncOpenAI(base_url=base_url) - logger.debug("tatlock_ollama_provider_created", base_url=base_url) + logger.debug( + "tatlock_ollama_provider_created", + base_url=base_url, + timeout=config.OLLAMA_TIMEOUT, + ) class _SanitizedAsyncOpenAI(AsyncOpenAI): """AsyncOpenAI client that sanitizes messages before sending.""" def __init__(self, **kwargs: Any): - # Ollama doesn't need an API key + # Ollama doesn't need an API key. Cap each LLM call at the + # configured Ollama timeout instead of the SDK default (~600s), + # so one stuck request cannot eat the whole delegation budget. + kwargs.setdefault("timeout", float(config.OLLAMA_TIMEOUT)) super().__init__(api_key="ollama", **kwargs) @property diff --git a/tests/agents/test_delegation.py b/tests/agents/test_delegation.py index 489c4da..a8a1703 100644 --- a/tests/agents/test_delegation.py +++ b/tests/agents/test_delegation.py @@ -193,6 +193,30 @@ class TestDelegateToLibrarian: assert "Connection refused" not in result.error assert "internal" not in result.error + @pytest.mark.asyncio + async def test_delegate_to_librarian_timeout(self, monkeypatch): + """Delegation is capped by LIBRARIAN_TIMEOUT and fails honestly.""" + import asyncio + + from src.core.config import config + + async def slow_run(task, context=""): + await asyncio.sleep(5) + return "too late" + + monkeypatch.setattr(config, "LIBRARIAN_TIMEOUT", 0.05) + + with patch( + "src.agents.librarian.agent.run_librarian", + new=slow_run, + ): + result = await delegate_to_librarian(task="Search for information") + + assert result.success is False + assert "longer than expected" in result.output + assert result.error is not None + assert "time budget" in result.error + @pytest.mark.asyncio async def test_delegate_to_librarian_preserves_task(self): """Test delegation result preserves original task.""" diff --git a/tests/agents/test_protocol.py b/tests/agents/test_protocol.py index e3df1a0..f8a762e 100644 --- a/tests/agents/test_protocol.py +++ b/tests/agents/test_protocol.py @@ -27,7 +27,8 @@ class TestAgentRequest: assert request.task == "Find information about Docker" assert request.context == "" - assert request.timeout_seconds == 60 + # No hardcoded default - None defers to the configured budget + assert request.timeout_seconds is None def test_request_with_context(self): """Test request with additional context.""" diff --git a/tests/ollama/test_provider.py b/tests/ollama/test_provider.py new file mode 100644 index 0000000..641eb47 --- /dev/null +++ b/tests/ollama/test_provider.py @@ -0,0 +1,54 @@ +""" +Tests for TatlockOllamaProvider configuration. + +The AsyncOpenAI client must carry an explicit timeout from +config.OLLAMA_TIMEOUT instead of the SDK default (~600s), so a stuck +LLM call cannot consume the whole delegation budget. +""" + +import pytest + +from src.core.config import config +from src.ollama.provider import TatlockOllamaProvider, _sanitize_messages + + +@pytest.mark.unit +class TestProviderTimeout: + """Timeout configuration on the underlying AsyncOpenAI client.""" + + def test_openai_client_timeout_from_config(self): + provider = TatlockOllamaProvider(base_url="http://localhost:11434/v1") + + assert provider._openai_client.timeout == float(config.OLLAMA_TIMEOUT) + + def test_timeout_is_not_sdk_default(self): + provider = TatlockOllamaProvider(base_url="http://localhost:11434/v1") + + # The OpenAI SDK defaults to 600s; the configured cap must win + assert provider._openai_client.timeout < 600 + + +@pytest.mark.unit +class TestMessageSanitization: + """Null content sanitization for Ollama compatibility.""" + + def test_null_content_with_tool_calls_becomes_empty_string(self): + messages = [ + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function"}], + } + ] + + sanitized = _sanitize_messages(messages) + + assert sanitized[0]["content"] == "" + + def test_regular_messages_unchanged(self): + messages = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Good day, sir."}, + ] + + assert _sanitize_messages(messages) == messages