[F120] release a stopped agent's claimed task immediately on budget-kill/shutdown

roboco/runtime/orchestrator.py

@@ -1057,7 +1057,13 @@ class AgentOrchestrator:
         # bg write, so log-and-continue rather than propagate.
         for agent_id in list(self._instances.keys()):
             try:
-                await self.stop_agent(agent_id)
+                # release_claim=True: on shutdown the orchestrator is going
+                # down and no agent will resume its task, so hand claimed
+                # tasks back to the pool now — they re-dispatch immediately on
+                # the next start instead of waiting for the reaper's TTL. A
+                # provider-parked agent is skipped inside stop_agent so its
+                # claim survives for the probe-resume loop across the restart.
+                await self.stop_agent(agent_id, release_claim=True)
             except Exception:
                 logger.exception(
                     "stop_agent raised during shutdown; continuing to drain",
@@ -4269,12 +4275,26 @@ class AgentOrchestrator:
         agent_id: str,
         graceful: bool = True,
         exit_reason: str = "stopped",
+        release_claim: bool = False,
     ) -> None:
         """Stop an agent container.
 
         Finalization (the HTTP call to the agent SDK's /usage/status endpoint)
         is performed BEFORE acquiring self._lock so that the network I/O does
         not block other operations that need the lock.
+
+        When ``release_claim`` is True the caller declares the stopped agent
+        will not continue its task (budget kill, orchestrator shutdown) and the
+        agent's claimed/in_progress task is handed back to the pool immediately
+        instead of waiting up to ``stale_claim_reap_seconds`` for the
+        stale-claim reaper to notice the dead heartbeat — closing the
+        SIGTERM-mid-verb gap where a task sat CLAIMED/IN_PROGRESS with no
+        running agent. Default False: the provider-park / waiting path
+        (``mark_waiting_long``) and interactive stops manage their own claim
+        lifecycle, so they opt out and the claim survives for the probe-resume
+        loop. A provider-parked agent (``rate_limit_lifted`` WaitingRecord) is
+        always skipped even when a caller opts in — its claim must survive so
+        probe-success revives the same agent on the same task.
         """
         # Finalize the spawn-session row before the container is removed so we
         # can still query the SDK's /usage/status endpoint.  This must happen
@@ -4286,6 +4306,11 @@ class AgentOrchestrator:
         if instance.container_id:
             await self._finalize_spawn_session(agent_id, exit_reason=exit_reason)
 
+        # Capture the task the agent was working on before the instance state
+        # is mutated, so a release_claim stop can hand it back to the pool once
+        # the container is gone. Only relevant when the caller opted in.
+        stopped_task_id = instance.current_task_id if release_claim else None
+
         async with self._lock:
             if agent_id not in self._instances:
                 return
@@ -4327,6 +4352,68 @@ class AgentOrchestrator:
 
             logger.info("Agent stopped", agent_id=agent_id)
 
+        # Hand the stopped agent's claimed task back to the pool now, instead
+        # of leaving it CLAIMED/IN_PROGRESS with no running agent for the
+        # reaper's full heartbeat TTL. Done outside self._lock (DB I/O) and
+        # best-effort: a failure logs a warning and the stale-claim reaper
+        # remains the backstop. Skipped for a provider-parked agent — the
+        # probe-resume loop owns its recovery and the claim must survive.
+        if stopped_task_id and not self._is_rate_limit_parked(agent_id):
+            await self._release_stopped_agent_claim(agent_id, stopped_task_id)
+
+    def _is_rate_limit_parked(self, agent_id: str) -> bool:
+        """True if the agent is provider-parked on a rate limit.
+
+        Mirrors the reaper's ``_assignee_is_provider_parked`` guard but keyed
+        by slug directly (no task row needed): a ``rate_limit_lifted``
+        WaitingRecord means the probe-resume loop owns this agent's recovery
+        and its claim must survive a stop. Defensive on a missing registry.
+        """
+        records = getattr(self, "_waiting_records", None)
+        if not records:
+            return False
+        record = records.get(agent_id)
+        return record is not None and record.waiting_for == "rate_limit_lifted"
+
+    async def _release_stopped_agent_claim(
+        self, agent_id: str, task_id_str: str
+    ) -> None:
+        """Force a stopped agent's claimed/in_progress task back to pending.
+
+        Reuses the hardened, idempotent, status-checked
+        ``TaskService.unclaim_for_reaper`` (the same path the stale-claim
+        reaper uses) so a task that already moved on (e.g. submitted to QA
+        before the stop) is a clean no-op. Opens its own short-lived session
+        outside ``self._lock``. Best-effort: a DB failure logs and the reaper
+        backstops on the next tick.
+        """
+        from roboco.db.base import get_session_factory
+        from roboco.services.task import TaskService
+        from roboco.utils.converters import require_uuid
+
+        try:
+            task_id = require_uuid(task_id_str)
+        except Exception:
+            return
+        try:
+            factory = get_session_factory()
+            async with factory() as db:
+                svc = TaskService(db)
+                await svc.unclaim_for_reaper(task_id)
+                await db.commit()
+            logger.info(
+                "stopped agent claim released to pool",
+                agent_id=agent_id,
+                task_id=task_id_str,
+            )
+        except Exception as exc:
+            logger.warning(
+                "stop_agent claim release failed; reaper will backstop",
+                agent_id=agent_id,
+                task_id=task_id_str,
+                error=str(exc),
+            )
+
     # =========================================================================
     # WAITING STATE MANAGEMENT
     # =========================================================================
@@ -5769,7 +5856,11 @@ Start by:
                     halt_threshold=data.get("halt_threshold"),
                 )
                 try:
-                    await self.stop_agent(agent_id, graceful=True)
+                    # release_claim=True: a budget-exceeded agent is terminated
+                    # for cost overruns and will not continue its task, so hand
+                    # the claim back to the pool now instead of waiting for the
+                    # reaper's TTL.
+                    await self.stop_agent(agent_id, graceful=True, release_claim=True)
                 except Exception as e:
                     logger.warning(
                         "Failed to stop budget-exceeded agent",

tests/unit/runtime/test_orchestrator_write_hooks.py

@@ -20,7 +20,8 @@ Coverage:
 from __future__ import annotations
 
 import json
-from contextlib import asynccontextmanager
+from contextlib import ExitStack, asynccontextmanager
+from datetime import UTC, datetime
 from pathlib import Path
 from typing import Any, cast
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -31,8 +32,10 @@ from roboco.models.runtime import (
     AgentInstance,
     OrchestratorAgentConfig,
     OrchestratorAgentState,
+    WaitingRecord,
 )
 from roboco.runtime.orchestrator import AgentOrchestrator
+from roboco.utils.converters import require_uuid
 
 # ---------------------------------------------------------------------------
 # Module-level constants (ruff PLR2004: no magic values in comparisons)
@@ -547,6 +550,150 @@ async def test_stop_agent_finalizes_before_lock() -> None:
     assert finalized == [_AGENT_ID]
 
 
+# ---------------------------------------------------------------------------
+# stop_agent — release_claim (F120): hand a stopped agent's claimed task back
+# to the pool immediately instead of waiting for the stale-claim reaper's TTL.
+# ---------------------------------------------------------------------------
+
+
+def _stop_agent_patches(orch: AgentOrchestrator) -> Any:
+    """Stub Docker + finalize so stop_agent runs without real Docker/DB."""
+    mock_proc = MagicMock()
+    mock_proc.wait = AsyncMock()
+    return (
+        patch.object(orch, "_finalize_spawn_session", AsyncMock()),
+        patch("asyncio.create_subprocess_exec", AsyncMock(return_value=mock_proc)),
+        patch.object(orch, "_remove_container", AsyncMock()),
+    )
+
+
+async def test_stop_agent_releases_claim_when_release_claim_true() -> None:
+    """F120: stop_agent(release_claim=True) hands the agent's claimed task back
+    to the pool immediately. A SIGTERM/budget-kill mid-verb otherwise leaves
+    the task CLAIMED/IN_PROGRESS with no running agent for up to
+    stale_claim_reap_seconds (the reaper's heartbeat TTL)."""
+    orch = _make_orchestrator()
+    instance = _make_instance(_AGENT_ID)
+    instance.current_task_id = str(uuid4())
+    orch._instances[_AGENT_ID] = instance
+
+    released: list[str] = []
+
+    async def _fake_release(agent_id: str, task_id: str) -> None:
+        released.append((agent_id, task_id))  # type: ignore[arg-type]
+
+    with ExitStack() as stack:
+        for cm in _stop_agent_patches(orch):
+            stack.enter_context(cm)
+        stack.enter_context(
+            patch.object(
+                orch,
+                "_release_stopped_agent_claim",
+                side_effect=_fake_release,
+                create=True,
+            )
+        )
+        await orch.stop_agent(_AGENT_ID, graceful=True, release_claim=True)
+
+    assert released == [(_AGENT_ID, instance.current_task_id)]
+
+
+async def test_stop_agent_does_not_release_claim_by_default() -> None:
+    """Default stop_agent (release_claim=False) must NOT release the claim —
+    preserves the existing behavior for the provider-park / waiting path
+    (mark_waiting_long) and the interactive stops, which manage their own
+    claim lifecycle via the reaper's provider-park guard. No regression."""
+    orch = _make_orchestrator()
+    instance = _make_instance(_AGENT_ID)
+    instance.current_task_id = str(uuid4())
+    orch._instances[_AGENT_ID] = instance
+
+    released: list[str] = []
+
+    async def _fake_release(agent_id: str, task_id: str) -> None:
+        released.append((agent_id, task_id))  # type: ignore[arg-type]
+
+    with ExitStack() as stack:
+        for cm in _stop_agent_patches(orch):
+            stack.enter_context(cm)
+        stack.enter_context(
+            patch.object(
+                orch,
+                "_release_stopped_agent_claim",
+                side_effect=_fake_release,
+                create=True,
+            )
+        )
+        # Default — release_claim omitted.
+        await orch.stop_agent(_AGENT_ID, graceful=True)
+
+    assert released == [], "default stop_agent must not release the claim"
+
+
+async def test_stop_agent_skips_release_for_provider_parked_agent() -> None:
+    """F120: a provider-parked agent (rate_limit_lifted WaitingRecord) must NOT
+    have its claim released even when release_claim=True. The probe-resume loop
+    owns its recovery and the claim must survive so probe-success revives the
+    SAME agent on the SAME task — reaping would let another agent claim it."""
+    orch = _make_orchestrator()
+    instance = _make_instance(_AGENT_ID)
+    instance.current_task_id = str(uuid4())
+    orch._instances[_AGENT_ID] = instance
+    # Parked on a rate limit — the claim must survive.
+    orch._waiting_records[_AGENT_ID] = WaitingRecord(
+        agent_id=_AGENT_ID,
+        task_id=instance.current_task_id,
+        waiting_for="rate_limit_lifted",
+        waiting_since=datetime.now(UTC),
+    )
+
+    released: list[str] = []
+
+    async def _fake_release(agent_id: str, task_id: str) -> None:
+        released.append((agent_id, task_id))  # type: ignore[arg-type]
+
+    with ExitStack() as stack:
+        for cm in _stop_agent_patches(orch):
+            stack.enter_context(cm)
+        stack.enter_context(
+            patch.object(
+                orch,
+                "_release_stopped_agent_claim",
+                side_effect=_fake_release,
+                create=True,
+            )
+        )
+        await orch.stop_agent(_AGENT_ID, graceful=True, release_claim=True)
+
+    assert released == [], "provider-parked agent's claim must not be released"
+
+
+async def test_release_stopped_agent_claim_calls_unclaim_for_reaper() -> None:
+    """The release helper opens a fresh session and routes through the hardened
+    TaskService.unclaim_for_reaper (status-checked + idempotent), committing."""
+    orch = _make_orchestrator()
+    task_id = str(uuid4())
+
+    svc = MagicMock()
+    svc.unclaim_for_reaper = AsyncMock()
+
+    @asynccontextmanager
+    async def _factory_ctx() -> Any:
+        db = MagicMock()
+        db.commit = AsyncMock()
+        yield db
+
+    fake_factory = MagicMock(return_value=_factory_ctx())
+
+    with (
+        patch("roboco.db.base.get_session_factory", return_value=fake_factory),
+        patch("roboco.services.task.TaskService", return_value=svc),
+    ):
+        await orch._release_stopped_agent_claim(_AGENT_ID, task_id)
+
+    svc.unclaim_for_reaper.assert_awaited_once_with(require_uuid(task_id))
+
+
 # ---------------------------------------------------------------------------
 # _handle_stopped_container — self-exits finalize (stop_agent was not called)
 # ---------------------------------------------------------------------------