diff --git a/app/agent/errors.py b/app/agent/errors.py new file mode 100644 index 0000000..84def7d --- /dev/null +++ b/app/agent/errors.py @@ -0,0 +1,60 @@ +"""Agent-agnostic error classification. + +Used by agent graphs (``error_node`` retry/fallback decisions) and by the +harness scheduler (retry policy). Kept dependency-free - no langgraph, no +agent state - so the scheduler can import it without pulling in the full +agent stack (``app.agent.memory.__init__`` imports langgraph). +""" + +from __future__ import annotations + +from enum import Enum +from typing import Sequence + + +class ErrorCategory(Enum): + RETRYABLE = "retryable" + NON_RETRYABLE = "non_retryable" + FATAL = "fatal" + + +_RETRYABLE_PATTERNS: Sequence[str] = [ + "timeout", + "connection", + "try again", + "rate limit", + "too many", + "temporarily", + "service unavailable", + "eof", + "reset", + "retryable", + "deadline exceeded", + "too many requests", + "internal server error", + "503", + "502", + "500", +] + +_FATAL_PATTERNS: Sequence[str] = [ + "authentication", + "unauthorized", + "invalid api key", + "permission denied", + "forbidden", + "account suspended", + "access denied", +] + + +def classify_error(error_message: str) -> ErrorCategory: + """Categorise an error string into RETRYABLE, NON_RETRYABLE, or FATAL.""" + error_lower = error_message.lower() + for p in _FATAL_PATTERNS: + if p in error_lower: + return ErrorCategory.FATAL + for p in _RETRYABLE_PATTERNS: + if p in error_lower: + return ErrorCategory.RETRYABLE + return ErrorCategory.NON_RETRYABLE diff --git a/app/agent/memory/handlers.py b/app/agent/memory/handlers.py index 733470b..5f04146 100644 --- a/app/agent/memory/handlers.py +++ b/app/agent/memory/handlers.py @@ -7,8 +7,7 @@ import asyncio import logging -from enum import Enum -from typing import Sequence +from app.agent.errors import ErrorCategory, classify_error # noqa: F401 (re-exported) from app.agent.memory.state import AgentState @@ -24,55 +23,6 @@ def build_fallback(state: AgentState) -> dict: "error": state.error or "unknown error", } - -class ErrorCategory(Enum): - RETRYABLE = "retryable" - NON_RETRYABLE = "non_retryable" - FATAL = "fatal" - - -_RETRYABLE_PATTERNS: Sequence[str] = [ - "timeout", - "connection", - "try again", - "rate limit", - "too many", - "temporarily", - "service unavailable", - "eof", - "reset", - "retryable", - "deadline exceeded", - "too many requests", - "internal server error", - "503", - "502", - "500", -] - -_FATAL_PATTERNS: Sequence[str] = [ - "authentication", - "unauthorized", - "invalid api key", - "permission denied", - "forbidden", - "account suspended", - "access denied", -] - - -def classify_error(error_message: str) -> ErrorCategory: - """Categorise an error string into RETRYABLE, NON_RETRYABLE, or FATAL.""" - error_lower = error_message.lower() - for p in _FATAL_PATTERNS: - if p in error_lower: - return ErrorCategory.FATAL - for p in _RETRYABLE_PATTERNS: - if p in error_lower: - return ErrorCategory.RETRYABLE - return ErrorCategory.NON_RETRYABLE - - async def backoff_delay(attempt: int, base_seconds: float = 1.0) -> None: """Exponential backoff: sleep base * 2^attempt seconds (capped at 10).""" delay = min(base_seconds * (2**attempt), 10.0) diff --git a/app/harness/scheduling/agent/scheduler.py b/app/harness/scheduling/agent/scheduler.py index 3cdec5a..8ebf4e3 100644 --- a/app/harness/scheduling/agent/scheduler.py +++ b/app/harness/scheduling/agent/scheduler.py @@ -15,10 +15,11 @@ import time import uuid from dataclasses import dataclass, field +from enum import Enum from typing import Any, Protocol, runtime_checkable +from app.agent.errors import ErrorCategory, classify_error from app.agent.lifecycle import AgentLifecycleManager -from app.agent.memory.handlers import ErrorCategory, classify_error logger = logging.getLogger(__name__) @@ -79,6 +80,19 @@ class AgentConfig: retry: AgentRetryConfig | None = None +class TicketState(str, Enum): + """Lifecycle states for an :class:`InvocationTicket`. + + ``str`` mixin keeps backwards-compatible string comparison + (``ticket.state == "queued"``) while making the states self-documenting. + """ + + QUEUED = "queued" + RUNNING = "running" + DONE = "done" + CANCELLED = "cancelled" + + @dataclass class InvocationTicket: """One queued or scheduled invocation.""" @@ -96,6 +110,11 @@ class InvocationTicket: error: str | None = None cancelled: bool = False retry_count: int = 0 + # Lifecycle state used by cancel(): "queued" -> "running" -> "done", + # or "queued" -> "cancelled". Only tickets still in "queued" state + # can be cancelled - once a worker has started execution the scheduler + # cannot interrupt the underlying ``lifecycle.invoke``. + state: TicketState = TicketState.QUEUED @dataclass @@ -136,6 +155,11 @@ def __init__(self, lifecycle: AgentLifecycleManager) -> None: self._lifecycle = lifecycle self._slots: dict[str, AgentSlot] = {} self._scheduled: dict[str, asyncio.Task] = {} + # Registry of every live ticket keyed by job_id, so ``cancel()`` + # can locate a queued/scheduled job without scanning asyncio.Queue + # (which is not iterable). Entries are removed when the worker + # finishes or skips the ticket. + self._tickets: dict[str, InvocationTicket] = {} self._running = False # ── lifecycle ──────────────────────────────────────────────────── @@ -181,6 +205,10 @@ async def shutdown(self) -> None: if workers: await asyncio.gather(*workers, return_exceptions=True) + # Drop any tickets that never reached a worker (e.g. delayed jobs + # cancelled while sleeping) so a reused scheduler starts clean. + self._tickets.clear() + logger.info("[AGENT_SCHED] shutdown complete") # ── configuration ──────────────────────────────────────────────── @@ -229,6 +257,7 @@ async def invoke( timeout=timeout, input=input, ) + self._tickets[ticket.job_id] = ticket # Scheduled / delayed path if delay_seconds is not None or at_timestamp is not None: @@ -262,6 +291,7 @@ async def invoke( session_id, slot.config.max_queue, ) + self._tickets.pop(ticket.job_id, None) return {"error": "queue full, try again later"} await event.wait() @@ -277,33 +307,48 @@ async def invoke( async def cancel(self, job_id: str) -> bool: """Cancel a queued or scheduled invocation by *job_id*. - Returns True if the job was found and cancelled. + Returns True only when the job was still pending - i.e. enqueued but + not yet picked up by a worker, or a delayed job still waiting to + fire. Running or already-finished jobs return False: the scheduler + cannot interrupt an in-flight ``lifecycle.invoke``. + + On success the waiting ``invoke()`` caller is unblocked (its ticket + event is set) and receives ``{"cancelled": True, "job_id": ...}``. + + Note: this method is ``async`` for API symmetry with ``invoke()`` but + contains no ``await`` - it runs atomically w.r.t. the event loop, so + the state check and the cancellation flip cannot be interleaved by a + worker flipping ``state`` to ``running``. """ - # Check scheduled tasks + ticket = self._tickets.get(job_id) + if ticket is None: + return False + + # Only queued (not-yet-running) jobs are cancellable. + if ticket.state != TicketState.QUEUED: + return False + + ticket.cancelled = True + ticket.state = TicketState.CANCELLED + if ticket.event is not None: + ticket.event.set() + + # If a delayed job is still sleeping in _schedule_one, cancel the + # timer task too. If it already fired into the queue, the worker + # will see ``ticket.cancelled`` and skip it. task = self._scheduled.get(job_id) if task is not None and not task.done(): task.cancel() - logger.debug("[AGENT_SCHED] cancelled scheduled job=%s", job_id) - return True - - # Check queues — O(n) per agent type. Acceptable for small queues. - for name, slot in self._slots.items(): - cancelled = await self._cancel_in_queue(slot, job_id) - if cancelled: - logger.debug( - "[AGENT_SCHED] cancelled queued job=%s agent=%s", job_id, name - ) - return True - return False + # Drop the ticket from the registry. For delayed jobs that never + # reached the queue this is the only cleanup path; for queued jobs + # the worker also pops (idempotent). + self._tickets.pop(job_id, None) - async def _cancel_in_queue(self, slot: AgentSlot, job_id: str) -> bool: - """Scan the queue and cancel a ticket by job_id.""" - # asyncio.Queue is not iterable — we can't scan it directly. - # Instead we mark the ticket as cancelled when it's dequeued. - # For scheduled tasks we already handle it above. - # For queued-but-not-yet-processed, we can only cancel via scheduled path. - return False + logger.debug( + "[AGENT_SCHED] cancelled job=%s agent=%s", job_id, ticket.agent_name + ) + return True # ── health ─────────────────────────────────────────────────────── @@ -384,10 +429,12 @@ async def _schedule_one(self, ticket: InvocationTicket, fire_at: float) -> None: return if not self._running or ticket.cancelled: + self._tickets.pop(ticket.job_id, None) return slot = self._slots.get(ticket.agent_name) if slot is None: + self._tickets.pop(ticket.job_id, None) return event = asyncio.Event() @@ -400,6 +447,7 @@ async def _schedule_one(self, ticket: InvocationTicket, fire_at: float) -> None: ticket.agent_name, ticket.session_id, ) + self._tickets.pop(ticket.job_id, None) return await event.wait() @@ -424,26 +472,54 @@ async def _drain_queue(self, agent_name: str) -> None: if not self._running: break + if ticket.event is None: + ticket.event = asyncio.Event() + if ticket.cancelled: - if ticket.event: - ticket.event.set() + ticket.state = TicketState.CANCELLED + ticket.event.set() + self._tickets.pop(ticket.job_id, None) continue async with slot.semaphore: - slot.active += 1 - ticket.event = asyncio.Event() if ticket.event is None else ticket.event - - result = await self._execute_with_retry(ticket, slot) - + # Re-check cancellation after acquiring the slot: under + # multiple workers a job can be cancelled while waiting for + # the semaphore (single-worker is unaffected since acquire + # does not yield when the slot is free). if ticket.cancelled: - result = {"cancelled": True, "job_id": ticket.job_id} - - if ticket.error: - result = {"error": ticket.error} - - ticket.result = result - ticket.event.set() - slot.active -= 1 + ticket.state = TicketState.CANCELLED + if ticket.event is not None: + ticket.event.set() + self._tickets.pop(ticket.job_id, None) + continue + slot.active += 1 + ticket.state = TicketState.RUNNING + try: + result = await self._execute_with_retry(ticket, slot) + + if ticket.cancelled: + result = {"cancelled": True, "job_id": ticket.job_id} + + if ticket.error: + result = {"error": ticket.error} + + ticket.result = result + if ticket.state != TicketState.CANCELLED: + ticket.state = TicketState.DONE + finally: + # Guarantee the caller is unblocked and the slot count + # is released even if the worker is cancelled + # mid-execution (CancelledError is BaseException and is + # NOT caught by _execute_with_retry's ``except Exception``). + # If result was never set the worker was killed - surface + # a sentinel so the caller can tell this apart from an + # empty-but-successful result. + if ticket.result is None: + ticket.result = {"error": "worker cancelled"} + if ticket.event is not None: + ticket.event.set() + slot.active -= 1 + self._tickets.pop(ticket.job_id, None) logger.debug("[AGENT_SCHED] worker stopped agent=%s", agent_name) diff --git a/app/test/test_agent_scheduler.py b/app/test/test_agent_scheduler.py new file mode 100644 index 0000000..d815d62 --- /dev/null +++ b/app/test/test_agent_scheduler.py @@ -0,0 +1,352 @@ +"""Unit tests for AgentScheduler - cancel semantics and cancellation safety. + +Covers two HIGH-severity behaviours that were previously broken: + +1. ``cancel(job_id)`` now actually cancels queued / scheduled jobs (the old + ``_cancel_in_queue`` always returned False and ``ticket.cancelled`` was + never set anywhere). +2. A worker task cancelled mid-execution must still set the ticket event + (so the ``invoke()`` caller does not hang forever) and decrement + ``slot.active`` (so the slot count does not leak). + +All tests use a fake lifecycle to avoid real agent / LLM / DB dependencies. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from app.harness.scheduling.agent.scheduler import ( + AgentConfig, + AgentRetryConfig, + AgentScheduler, + InvocationTicket, + TicketState, +) + + +# --------------------------------------------------------------------------- +# Fake lifecycle - controllable invoke timing +# --------------------------------------------------------------------------- + + +class _FakeLifecycle: + """Stand-in for ``AgentLifecycleManager`` with a gate on ``invoke``. + + ``invoke`` blocks until ``gate`` is set, so tests can occupy the single + concurrency slot and queue / cancel work behind it. + """ + + def __init__(self) -> None: + self.gate = asyncio.Event() + self.started = asyncio.Event() + self.invoke_count = 0 + + async def invoke(self, agent_name, session_id, *, timeout=None, **input): + self.started.set() + await self.gate.wait() # block until the test releases us + self.invoke_count += 1 + return {"result": f"{agent_name}:{session_id}"} + + async def health(self): + return {"registered_agents": ["chat"]} + + +async def _drain() -> None: + """Yield control to let queued tasks / workers make progress.""" + await asyncio.sleep(0) + + +# --------------------------------------------------------------------------- +# cancel() - unknown / running +# --------------------------------------------------------------------------- + + +class TestCancelUnknown: + @pytest.mark.asyncio + async def test_cancel_unknown_job_returns_false(self) -> None: + sched = AgentScheduler(_FakeLifecycle()) + await sched.start() + try: + assert await sched.cancel("does-not-exist") is False + finally: + await sched.shutdown() + + +class TestCancelRunning: + @pytest.mark.asyncio + async def test_cancel_running_job_returns_false(self) -> None: + lc = _FakeLifecycle() + sched = AgentScheduler(lc) + sched.set_config("chat", AgentConfig(max_concurrent=1)) + await sched.start() + try: + t1 = asyncio.create_task( + sched.invoke("chat", session_id="s1", query="q") + ) + await lc.started.wait() # worker is inside lifecycle.invoke + + running = [ + t for t in sched._tickets.values() if t.session_id == "s1" + ] + assert len(running) == 1 + assert running[0].state == TicketState.RUNNING + + assert await sched.cancel(running[0].job_id) is False + + lc.gate.set() + await asyncio.wait_for(t1, timeout=2) + finally: + await sched.shutdown() + + +# --------------------------------------------------------------------------- +# cancel() - queued job (occupies slot, cancel the waiting one) +# --------------------------------------------------------------------------- + + +class TestCancelQueued: + @pytest.mark.asyncio + async def test_cancel_queued_job_skips_execution(self) -> None: + lc = _FakeLifecycle() + sched = AgentScheduler(lc) + sched.set_config("chat", AgentConfig(max_concurrent=1)) + await sched.start() + try: + # t1 occupies the only slot (gate closed). + t1 = asyncio.create_task( + sched.invoke("chat", session_id="s1", query="q1") + ) + await lc.started.wait() + + # t2 queues behind it. + t2 = asyncio.create_task( + sched.invoke("chat", session_id="s2", query="q2") + ) + await _drain() + + queued = [ + t for t in sched._tickets.values() if t.session_id == "s2" + ] + assert len(queued) == 1 + assert queued[0].state == TicketState.QUEUED + + ok = await sched.cancel(queued[0].job_id) + assert ok is True + + # t2 must return promptly with a cancelled payload, no hang. + r2 = await asyncio.wait_for(t2, timeout=2) + assert r2 == {"cancelled": True, "job_id": queued[0].job_id} + + # Release t1 and confirm only s1 executed. + lc.gate.set() + r1 = await asyncio.wait_for(t1, timeout=2) + assert r1 == {"result": "chat:s1"} + assert lc.invoke_count == 1 + finally: + await sched.shutdown() + + +# --------------------------------------------------------------------------- +# cancel() - scheduled (delayed) job, before it fires +# --------------------------------------------------------------------------- + + +class TestCancelScheduled: + @pytest.mark.asyncio + async def test_cancel_scheduled_before_fire(self) -> None: + lc = _FakeLifecycle() + sched = AgentScheduler(lc) + sched.set_config("chat", AgentConfig(max_concurrent=1)) + await sched.start() + try: + res = await sched.invoke( + "chat", session_id="s1", delay_seconds=10, query="q" + ) + assert res["scheduled"] is True + job_id = res["job_id"] + + ok = await sched.cancel(job_id) + assert ok is True + await _drain() + assert lc.invoke_count == 0 # never fired + finally: + await sched.shutdown() + + @pytest.mark.asyncio + async def test_cancel_scheduled_after_fired_into_queue(self) -> None: + """A scheduled job that has already fired into a full queue can + still be cancelled before the worker picks it up.""" + lc = _FakeLifecycle() + sched = AgentScheduler(lc) + sched.set_config("chat", AgentConfig(max_concurrent=1)) + await sched.start() + try: + # Occupy the slot. + t1 = asyncio.create_task( + sched.invoke("chat", session_id="s1", query="q1") + ) + await lc.started.wait() + + res = await sched.invoke( + "chat", session_id="s2", delay_seconds=0.05, query="q2" + ) + job_id = res["job_id"] + await asyncio.sleep(0.15) # let the delay elapse -> enqueued + + ok = await sched.cancel(job_id) + assert ok is True + + lc.gate.set() + await asyncio.wait_for(t1, timeout=2) + assert lc.invoke_count == 1 # s2 cancelled, never ran + finally: + await sched.shutdown() + + +# --------------------------------------------------------------------------- +# Worker cancellation safety (HIGH #2) +# --------------------------------------------------------------------------- + + +class TestWorkerCancellation: + @pytest.mark.asyncio + async def test_worker_cancel_unblocks_caller_and_releases_slot(self) -> None: + lc = _FakeLifecycle() + sched = AgentScheduler(lc) + sched.set_config("chat", AgentConfig(max_concurrent=1)) + await sched.start() + try: + t1 = asyncio.create_task( + sched.invoke("chat", session_id="s1", query="q") + ) + await lc.started.wait() # worker is inside _execute_with_retry + + assert sched._slots["chat"].active == 1 + + worker = sched._slots["chat"].worker_task + worker.cancel() # simulate shutdown / external cancellation + + # Caller must NOT hang - event must be set in finally. + await asyncio.wait_for(t1, timeout=2) + + # Slot count must not leak. + assert sched._slots["chat"].active == 0 + finally: + await sched.shutdown() + + +# --------------------------------------------------------------------------- +# InvocationTicket state contract +# --------------------------------------------------------------------------- + + +class TestTicketState: + def test_initial_state_is_queued(self) -> None: + ticket = InvocationTicket( + job_id="abc", + agent_name="chat", + session_id="s1", + input={}, + ) + assert ticket.state == TicketState.QUEUED + assert ticket.cancelled is False + + +# --------------------------------------------------------------------------- +# _execute_with_retry - retry classification (lazy import path exercised) +# --------------------------------------------------------------------------- + + +class _FlakyLifecycle: + """Lifecycle that raises *fail_times* times then succeeds.""" + + def __init__(self, error_factory, fail_times: int = 0) -> None: + self.error_factory = error_factory + self.fail_times = fail_times + self.calls = 0 + + async def invoke(self, agent_name, session_id, *, timeout=None, **input): + self.calls += 1 + if self.calls <= self.fail_times: + raise self.error_factory() + return {"result": "ok"} + + async def health(self): + return {"registered_agents": ["chat"]} + + +def _retry_cfg() -> AgentRetryConfig: + return AgentRetryConfig(max_retries=2, backoff_base=0.01, backoff_max=0.05) + + +class TestExecuteWithRetry: + @pytest.mark.asyncio + async def test_retryable_error_retries_then_succeeds(self) -> None: + lc = _FlakyLifecycle(lambda: RuntimeError("connection timeout"), fail_times=1) + sched = AgentScheduler(lc) + sched.set_config( + "chat", AgentConfig(max_concurrent=1, retry=_retry_cfg()) + ) + await sched.start() + try: + result = await asyncio.wait_for( + sched.invoke("chat", session_id="s1", query="q"), timeout=5 + ) + assert result == {"result": "ok"} + assert lc.calls == 2 + finally: + await sched.shutdown() + + @pytest.mark.asyncio + async def test_non_retryable_error_does_not_retry(self) -> None: + lc = _FlakyLifecycle(lambda: RuntimeError("unknown glitch"), fail_times=99) + sched = AgentScheduler(lc) + sched.set_config( + "chat", AgentConfig(max_concurrent=1, retry=_retry_cfg()) + ) + await sched.start() + try: + result = await asyncio.wait_for( + sched.invoke("chat", session_id="s1", query="q"), timeout=5 + ) + assert "error" in result + assert lc.calls == 1 + finally: + await sched.shutdown() + + @pytest.mark.asyncio + async def test_fatal_error_does_not_retry(self) -> None: + lc = _FlakyLifecycle(lambda: RuntimeError("invalid api key"), fail_times=99) + sched = AgentScheduler(lc) + sched.set_config( + "chat", AgentConfig(max_concurrent=1, retry=_retry_cfg()) + ) + await sched.start() + try: + result = await asyncio.wait_for( + sched.invoke("chat", session_id="s1", query="q"), timeout=5 + ) + assert "error" in result + assert lc.calls == 1 + finally: + await sched.shutdown() + + @pytest.mark.asyncio + async def test_retry_exhausted_returns_error(self) -> None: + lc = _FlakyLifecycle(lambda: RuntimeError("connection timeout"), fail_times=99) + sched = AgentScheduler(lc) + sched.set_config( + "chat", AgentConfig(max_concurrent=1, retry=_retry_cfg()) + ) + await sched.start() + try: + result = await asyncio.wait_for( + sched.invoke("chat", session_id="s1", query="q"), timeout=5 + ) + assert "error" in result + assert lc.calls == 3 # 1 initial + 2 retries + finally: + await sched.shutdown()