diff --git a/src/agentexec/__init__.py b/src/agentexec/__init__.py index 01faa3a..613bf8b 100644 --- a/src/agentexec/__init__.py +++ b/src/agentexec/__init__.py @@ -30,7 +30,7 @@ async def search(agent_id: UUID, context: Input) -> Output: from agentexec.config import CONF from agentexec.core.db import Base from agentexec.core.queue import Priority, enqueue -from agentexec.core.results import gather, get_result +from agentexec.core.results import TaskFailedError, gather, get_result from agentexec.core.task import Task from agentexec import activity from agentexec.worker import Pool @@ -51,6 +51,7 @@ async def search(agent_id: UUID, context: Input) -> Output: "Tracker", "Pool", "Task", + "TaskFailedError", "Priority", "activity", "enqueue", diff --git a/src/agentexec/core/results.py b/src/agentexec/core/results.py index d23eaca..1502a68 100644 --- a/src/agentexec/core/results.py +++ b/src/agentexec/core/results.py @@ -16,6 +16,29 @@ DEFAULT_TIMEOUT: int = 300 # TODO improve this polling approach +class TaskFailure(BaseModel): + """Terminal failure record stored at the result key when retries are exhausted.""" + + task_name: str + agent_id: UUID + error: str + attempts: int + + +class TaskFailedError(Exception): + """A task permanently failed after exhausting its retries.""" + + def __init__(self, failure: TaskFailure) -> None: + self.task_name = failure.task_name + self.agent_id = failure.agent_id + self.error = failure.error + self.attempts = failure.attempts + super().__init__( + f"Task {failure.task_name} ({failure.agent_id}) failed " + f"after {failure.attempts} attempts: {failure.error}" + ) + + async def _get_result(agent_id: str | UUID) -> BaseModel | None: key = backend.format_key(*KEY_RESULT, str(agent_id)) data = await backend.state.get(key) @@ -23,11 +46,18 @@ async def _get_result(agent_id: str | UUID) -> BaseModel | None: async def get_result(task: Task, timeout: int = DEFAULT_TIMEOUT) -> BaseModel: - """Poll for a task result.""" + """Poll for a task result. + + Raises: + TaskFailedError: If the task permanently failed after exhausting retries. + TimeoutError: If no result is available within the timeout. + """ start = time.time() while time.time() - start < timeout: result = await _get_result(task.agent_id) + if isinstance(result, TaskFailure): + raise TaskFailedError(result) if result is not None: return result await asyncio.sleep(0.5) @@ -36,6 +66,15 @@ async def get_result(task: Task, timeout: int = DEFAULT_TIMEOUT) -> BaseModel: async def gather(*tasks: Task, timeout: int = DEFAULT_TIMEOUT) -> tuple[BaseModel, ...]: - """Wait for multiple tasks and return their results.""" + """Wait for multiple tasks and return their results. + + Raises on the first permanently failed task. The remaining tasks keep + running in their workers and their results stay retrievable — a + subsequent gather over the same tasks resolves finished ones instantly. + + Raises: + TaskFailedError: If any task permanently failed after exhausting retries. + TimeoutError: If any result is not available within the timeout. + """ results = await asyncio.gather(*[get_result(task, timeout) for task in tasks]) return tuple(results) diff --git a/src/agentexec/worker/pool.py b/src/agentexec/worker/pool.py index ddf5dc9..52029dd 100644 --- a/src/agentexec/worker/pool.py +++ b/src/agentexec/worker/pool.py @@ -11,7 +11,7 @@ from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from agentexec.config import CONF -from agentexec.state import backend +from agentexec.state import KEY_RESULT, backend import queue as stdlib_queue from agentexec import activity @@ -19,6 +19,7 @@ from agentexec.activity.handlers import IPCHandler from agentexec.core.db import configure_engine, dispose_engine from agentexec.core.queue import enqueue +from agentexec.core.results import TaskFailure from agentexec.core.task import Task, TaskDefinition, TaskHandler from agentexec import schedule @@ -224,6 +225,16 @@ async def _handle(self) -> None: priority=Priority.HIGH, ) else: + failure = TaskFailure( + task_name=task.task_name, + agent_id=task.agent_id, + error=error, + attempts=task.retry_count + 1, + ) + key = backend.format_key(*KEY_RESULT, str(task.agent_id)) + await backend.state.set( + key, backend.serialize(failure), ttl_seconds=CONF.result_ttl + ) logger.info( f"Task {task.task_name} failed " f"after {task.retry_count + 1} attempts, giving up: {error}" diff --git a/tests/test_results.py b/tests/test_results.py index 43c4197..1e3ef40 100644 --- a/tests/test_results.py +++ b/tests/test_results.py @@ -6,7 +6,7 @@ from pydantic import BaseModel import agentexec as ax -from agentexec.core.results import gather, get_result, _get_result +from agentexec.core.results import TaskFailedError, TaskFailure, gather, get_result, _get_result class SampleContext(BaseModel): @@ -150,6 +150,60 @@ async def mock_result(agent_id): assert results == expected +async def test_get_result_raises_on_task_failure(mock_get_result) -> None: + """A stored TaskFailure record raises TaskFailedError immediately.""" + task = ax.Task( + task_name="doomed_task", + context={"message": "test"}, + agent_id=uuid.uuid4(), + ) + mock_get_result.return_value = TaskFailure( + task_name="doomed_task", + agent_id=task.agent_id, + error="boom", + attempts=4, + ) + + with pytest.raises(TaskFailedError, match="failed after 4 attempts: boom") as exc_info: + await get_result(task, timeout=30) + + assert exc_info.value.task_name == "doomed_task" + assert exc_info.value.agent_id == task.agent_id + assert exc_info.value.error == "boom" + assert exc_info.value.attempts == 4 + # Raised on the first poll, not after the timeout. + mock_get_result.assert_called_once_with(task.agent_id) + + +async def test_gather_propagates_task_failure(mock_get_result) -> None: + """gather raises when any task has permanently failed.""" + ok_task = ax.Task( + task_name="ok_task", + context={"message": "ok"}, + agent_id=uuid.uuid4(), + ) + doomed_task = ax.Task( + task_name="doomed_task", + context={"message": "doom"}, + agent_id=uuid.uuid4(), + ) + + async def mock_result(agent_id): + if agent_id == ok_task.agent_id: + return SampleResult(status="ok", value=1) + return TaskFailure( + task_name="doomed_task", + agent_id=doomed_task.agent_id, + error="fatal", + attempts=4, + ) + + mock_get_result.side_effect = mock_result + + with pytest.raises(TaskFailedError, match="doomed_task"): + await gather(ok_task, doomed_task, timeout=30) + + async def test_get_result_with_complex_object(mock_get_result) -> None: task = ax.Task( task_name="test_task", diff --git a/tests/test_worker_logging.py b/tests/test_worker_logging.py index 5cd3ad0..56dd1ae 100644 --- a/tests/test_worker_logging.py +++ b/tests/test_worker_logging.py @@ -272,6 +272,7 @@ async def handler(agent_id: uuid.UUID, context: SampleContext): _, buf = _capture_handler() monkeypatch.setattr(ax.CONF, "max_task_retries", 3) + monkeypatch.setattr("agentexec.state.backend.state.set", AsyncMock()) import queue q = queue.Queue() diff --git a/tests/test_worker_pool.py b/tests/test_worker_pool.py index a65969b..24a8b09 100644 --- a/tests/test_worker_pool.py +++ b/tests/test_worker_pool.py @@ -423,8 +423,9 @@ async def mock_push(value, *, priority=None, partition_key=None): assert pushed[0]["priority"] is Priority.HIGH async def test_gives_up_after_max_retries(self, pool, monkeypatch, caplog): - """Failed task at max retries is not requeued.""" + """Failed task at max retries is not requeued and stores a failure record.""" from agentexec.worker.pool import TaskFailed + from agentexec.core.results import TaskFailure @pool.task("doomed_task") async def handler(agent_id: uuid.UUID, context: SampleContext): @@ -438,11 +439,17 @@ async def handler(agent_id: uuid.UUID, context: SampleContext): ) pushed = [] + stored = {} async def mock_push(value, *, priority=None, partition_key=None): pushed.append(value) + async def mock_set(key, value, ttl_seconds=None): + stored[key] = value + return True + monkeypatch.setattr("agentexec.state.backend.queue.push", mock_push) + monkeypatch.setattr("agentexec.state.backend.state.set", mock_set) monkeypatch.setattr(ax.CONF, "max_task_retries", 3) eh = self._make_handler(pool) @@ -455,6 +462,16 @@ async def mock_push(value, *, priority=None, partition_key=None): assert "4 attempts" in caplog.text assert "fatal" in caplog.text + from agentexec.state import KEY_RESULT + + result_key = backend.format_key(*KEY_RESULT, str(task.agent_id)) + failure = backend.deserialize(stored[result_key]) + assert isinstance(failure, TaskFailure) + assert failure.task_name == "doomed_task" + assert failure.agent_id == task.agent_id + assert failure.error == "fatal" + assert failure.attempts == 4 + async def test_retry_preserves_partition_key(self, pool, monkeypatch): """Requeued task uses the correct partition key from its definition.""" from agentexec.worker.pool import TaskFailed diff --git a/uv.lock b/uv.lock index adef82e..53c743b 100644 --- a/uv.lock +++ b/uv.lock @@ -4,7 +4,7 @@ requires-python = ">=3.12" [[package]] name = "agentexec" -version = "0.2.0rc2" +version = "0.2.0" source = { editable = "." } dependencies = [ { name = "croniter" },