feat(agents): extend cooperative early stop to codex and antigravity

docs/agents/ANTIGRAVITY.md

@@ -168,18 +168,15 @@ as every other agent.
 
 ## Known limitations
 
-1. **No cooperative early stop.** Antigravity does not support the `should_stop`
-   seam, so `run_limits.stop_early` is unsupported for this agent (it errors at
-   resolution, as it does for any non-Claude agent).
-2. **No endpoint routing.** Only `GEMINI_API_KEY` + `ANTIGRAVITY_MODEL` are read —
+1. **No endpoint routing.** Only `GEMINI_API_KEY` + `ANTIGRAVITY_MODEL` are read —
    there is no base-URL, project, region, or gateway override.
-3. **Default-model drift.** The runtime fallback (`gemini-3.5-flash`) may differ from
+2. **Default-model drift.** The runtime fallback (`gemini-3.5-flash`) may differ from
    what a given release's docs or example tasks pin; always set `agent.model`
    explicitly for reproducible runs.
-4. **`kill_sync()` is best-effort.** The SDK's cancel/disconnect are async-only, so
+3. **`kill_sync()` is best-effort.** The SDK's cancel/disconnect are async-only, so
    the watchdog's synchronous kill only flips agent state to `ERROR`; real teardown
    happens on the subsequent async `stop()`.
-5. **Process-global spawn lock.** The SDK spawns `localharness` via a subprocess with
+4. **Process-global spawn lock.** The SDK spawns `localharness` via a subprocess with
    no env-injection seam, so the agent transiently mutates `PATH` across the spawn
    under a process-wide lock. This serializes harness startup across concurrent
    tasks (it does not serialize the turns themselves).

docs/agents/CLAUDE_CODE.md

@@ -145,13 +145,14 @@ simulator force `[]` for the same reason.)
   UiPath CLI plugin discovery; when unset the sandbox derives it from the resolved
   `uip` binary. See [User Guide → Environment Variables](../USER_GUIDE.md#environment-variables).
 
-## Early stop (Claude-only)
+## Early stop
 
-Claude Code is the only agent that supports the cooperative early-stop seam. With
+Claude Code supports the cooperative early-stop seam, as do the
+[Codex](CODEX.md) and [Antigravity](ANTIGRAVITY.md) agents. With
 `run_limits.stop_early: true`, a single-shot run ends cleanly at the next
 tool-call boundary once its **armed** criteria (`stop_when: pass|fail|decided`) are
 decided — so a raised `max_turns` isn't wasted on a smoke run. Early stop errors at
-resolution for any other agent. See the
+resolution for any agent that does not declare `supports_cooperative_stop`. See the
 [Task Definition Guide](../TASK_DEFINITION_GUIDE.md) for the full contract.
 
 ## Telemetry

docs/agents/CODEX.md

@@ -216,6 +216,7 @@ The Codex SDK is synchronous. The agent uses `_run_async()` helper to detect and
 | **Session Resume** | `--resume {session_id}` | Via thread ID |
 | **Permissions** | `permission_mode` + `allowed_tools` | `permission_mode` → sandbox/approval + `allowed_tools`/`disallowed_tools` → thread config |
 | **Tool Enforcement** | Not enforced by Coder Eval wrapper | `enabled_tools` honored; `disabled_tools` NOT enforced by the SDK |
+| **Early stop** | Supported (cooperative `should_stop`, polled between messages) | Supported — polled after each streamed notification; the in-flight turn is interrupted best-effort |
 
 ## Known Limitations
 

src/coder_eval/agents/antigravity_agent.py

@@ -27,7 +27,7 @@
 from contextlib import AsyncExitStack
 from datetime import datetime
 from pathlib import Path
-from typing import Any
+from typing import Any, ClassVar
 
 from coder_eval.agent import Agent, AgentState
 from coder_eval.agents._logging import PrefixedAdapter
@@ -185,6 +185,10 @@ def _to_token_usage(usage: Any, model: str | None) -> TokenUsage:
 class AntigravityAgent(Agent[AntigravityAgentConfig]):
     """Implementation of the Agent interface for Google Antigravity (Gemini)."""
 
+    # The step loop has a between-steps guard where the cooperative
+    # ``should_stop`` check runs, so this agent supports early-stop-on-criterion.
+    supports_cooperative_stop: ClassVar[bool] = True
+
     def __init__(
         self,
         config: AntigravityAgentConfig,
@@ -413,10 +417,10 @@ async def communicate(
     ) -> TurnRecord:
         """Send a message to the Antigravity agent and receive its response.
 
-        ``should_stop`` is accepted for ``Agent.communicate`` override
-        compatibility and ignored — Antigravity does not support cooperative
-        early-stop (``supports_cooperative_stop`` is False), so the orchestrator
-        never passes it.
+        ``should_stop`` is the cooperative early-stop callback, polled after each
+        processed step. When it returns True the step loop breaks, the
+        conversation is cancelled (best-effort) and the turn finalizes cleanly as
+        ``STOPPED_EARLY`` (``crashed=False``).
 
         Drives one logical turn: ``conversation.send(prompt)`` then iterate
         ``receive_steps()`` until the turn goes idle, mapping the Gemini step
@@ -468,18 +472,39 @@ def _on_turn_timeout() -> None:
                 conversation = self._sdk_agent.conversation
                 try:
                     await conversation.send(user_input)
+                    # The cooperative should_stop poll runs AFTER process_step (the
+                    # emission that lets the watcher latch on the deciding tool
+                    # call) and BEFORE the next step is pulled — the deciding step
+                    # is kept, the next is not. No-op when should_stop is None.
                     async for step in conversation.receive_steps():
                         state.process_step(step)
+                        if should_stop is not None and should_stop():
+                            state.stopped_early_hit = True
+                            self._log.debug("Cooperative stop requested; ending step loop at this boundary")
+                            break
+                    if state.stopped_early_hit:
+                        # Best-effort server-side cancel, mirrors kill(); a raising
+                        # cancel() lands in the guarded handler below.
+                        with contextlib.suppress(Exception):
+                            await conversation.cancel()
                 except asyncio.CancelledError:
                     if state.timeout_hit:
                         self._finalize_and_raise_timeout(state.finalize, timeout or 0)
                     raise
                 except Exception as e:
                     if state.timeout_hit:
                         self._finalize_and_raise_timeout(state.finalize, timeout or 0, cause=e)
-                    self._finalize_and_raise_crash(
-                        state.finalize, truncate_crash_message(f"Antigravity turn failed: {e!s}"), cause=e
-                    )
+                    if state.stopped_early_hit:
+                        # The turn already stopped cleanly (e.g. the generator's
+                        # aclose() raised on the break); escalating to a crash
+                        # would trigger the orchestrator's retry with the watcher's
+                        # decision still latched → immediate stop-at-turn-0 on the
+                        # retry (wasted spend). Fall through to the clean tail.
+                        self._log.warning("Ignoring post-stop exception; finalizing as STOPPED_EARLY: %s", e)
+                    else:
+                        self._finalize_and_raise_crash(
+                            state.finalize, truncate_crash_message(f"Antigravity turn failed: {e!s}"), cause=e
+                        )
 
             if state.timeout_hit:
                 # Watchdog fired but the pump finished before the cancel landed.
@@ -493,13 +518,20 @@ def _on_turn_timeout() -> None:
                 state.finalize(AgentEndStatus.CRASHED, crashed=True, crash_reason="turn cancelled")
             raise
         except Exception as e:
-            self._finalize_and_raise_crash(
-                state.finalize, truncate_crash_message(f"Antigravity turn failed: {e!s}"), cause=e
-            )
+            if state.stopped_early_hit and not state.timeout_hit:
+                # Same retry-poisoning guard as the inner handler: a cooperative
+                # stop already happened, so finalize cleanly instead of crashing.
+                self._log.warning("Ignoring post-stop exception; finalizing as STOPPED_EARLY: %s", e)
+            else:
+                self._finalize_and_raise_crash(
+                    state.finalize, truncate_crash_message(f"Antigravity turn failed: {e!s}"), cause=e
+                )
 
         self._state = AgentState.WORKING
         self._end_turn_ok()
-        state.finalize(AgentEndStatus.COMPLETED, crashed=False, crash_reason=None)
+        # Precedence matches Claude: timeout (raised above) > stopped_early > completed.
+        status = AgentEndStatus.STOPPED_EARLY if state.stopped_early_hit else AgentEndStatus.COMPLETED
+        state.finalize(status, crashed=False, crash_reason=None)
         return collector.build_turn_record()
 
     async def stop(self) -> None:
@@ -589,6 +621,7 @@ def __init__(
         self.turn_start_time = turn_start_time
 
         self.timeout_hit = False
+        self.stopped_early_hit = False
         self.finalized = False
 
         self.total_usage = TokenUsage()

src/coder_eval/agents/codex_agent.py

@@ -12,7 +12,7 @@
 from collections.abc import Callable
 from datetime import datetime
 from pathlib import Path
-from typing import Any
+from typing import Any, ClassVar
 from urllib.parse import urlparse
 
 from coder_eval.agent import Agent, AgentState
@@ -280,6 +280,7 @@ def __init__(
         self.iteration = iteration
         self.turn_start_time = turn_start_time
         self.timeout_hit = False
+        self.stopped_early_hit = False
         self.finalized = False
 
         # Live pump scratch (set during streaming).
@@ -634,6 +635,10 @@ def finalize(self, status: AgentEndStatus, *, crashed: bool = False, crash_reaso
 class CodexAgent(Agent[CodexAgentConfig]):
     """Implementation of the Agent interface for OpenAI Codex using the Codex SDK."""
 
+    # The notification pump has a between-items guard where the cooperative
+    # ``should_stop`` check runs, so this agent supports early-stop-on-criterion.
+    supports_cooperative_stop: ClassVar[bool] = True
+
     def __init__(
         self,
         config: CodexAgentConfig,
@@ -740,10 +745,10 @@ async def communicate(
             stream_callback: Optional callback for real-time event streaming
             timeout: Hard wall-clock deadline in seconds
             max_turns: Hard cap on inner-loop turns (unused for Codex single-turn)
-            should_stop: Accepted for ``Agent.communicate`` override compatibility
-                and ignored — Codex does not support cooperative early-stop
-                (``supports_cooperative_stop`` is False), so the orchestrator
-                never passes it.
+            should_stop: Cooperative early-stop callback, polled after each
+                dispatched notification. When it returns True the pump breaks,
+                the in-flight turn is interrupted (best-effort) and the turn
+                finalizes cleanly as ``STOPPED_EARLY`` (``crashed=False``).
 
         Returns:
             TurnRecord containing the complete interaction
@@ -822,7 +827,7 @@ def _on_turn_timeout() -> None:
                     # return; a crash skips this, so finalize reads the crash
                     # defaults (None/"") and falls back to the captured messages.
                     state.result_turn, state.sdk_token_usage, state.result_text = await self._run_turn_with_streaming(
-                        state
+                        state, should_stop
                     )
                 except asyncio.CancelledError:
                     if state.timeout_hit:
@@ -831,9 +836,16 @@ def _on_turn_timeout() -> None:
                 except Exception as e:
                     if state.timeout_hit:
                         self._finalize_and_raise_timeout(state.finalize, timeout or 0, cause=e)
-                    self._finalize_and_raise_crash(
-                        state.finalize, truncate_crash_message(f"Codex turn failed: {e!s}"), cause=e
-                    )
+                    if state.stopped_early_hit:
+                        # The turn already stopped cleanly; escalating to a crash
+                        # would trigger the orchestrator's retry with the watcher's
+                        # decision still latched → immediate stop-at-turn-0 on the
+                        # retry (wasted spend). Fall through to the clean tail.
+                        self._log.warning("Ignoring post-stop exception; finalizing as STOPPED_EARLY: %s", e)
+                    else:
+                        self._finalize_and_raise_crash(
+                            state.finalize, truncate_crash_message(f"Codex turn failed: {e!s}"), cause=e
+                        )
 
             if state.timeout_hit:
                 # Watchdog fired but the pump finished before the cancel landed.
@@ -860,13 +872,22 @@ def _on_turn_timeout() -> None:
             # and _format_turn_result. Without this, such errors escape as a bare
             # exception: the orchestrator never drains pending_turn and _iteration
             # stays incremented, violating the pending-turn contract.
-            self._finalize_and_raise_crash(state.finalize, truncate_crash_message(f"Codex turn failed: {e!s}"), cause=e)
+            if state.stopped_early_hit and not state.timeout_hit:
+                # Same retry-poisoning guard as the inner handler: a cooperative
+                # stop already happened, so finalize cleanly instead of crashing.
+                self._log.warning("Ignoring post-stop exception; finalizing as STOPPED_EARLY: %s", e)
+            else:
+                self._finalize_and_raise_crash(
+                    state.finalize, truncate_crash_message(f"Codex turn failed: {e!s}"), cause=e
+                )
 
         self._state = AgentState.WORKING
         self._end_turn_ok()
 
         # The TurnRecord is the EventCollector's reduction of the emitted events.
-        state.finalize(AgentEndStatus.COMPLETED, crashed=False, crash_reason=None)
+        # Precedence matches Claude: timeout (raised above) > stopped_early > completed.
+        status = AgentEndStatus.STOPPED_EARLY if state.stopped_early_hit else AgentEndStatus.COMPLETED
+        state.finalize(status, crashed=False, crash_reason=None)
         return collector.build_turn_record()
 
     async def stop(self) -> None:
@@ -1380,14 +1401,21 @@ def _format_turn_result(self, turn_result: Any) -> str:
             self._log.warning(f"Failed to format turn result: {e}")
             return str(turn_result)
 
-    async def _run_turn_with_streaming(self, state: _CodexTurnState) -> tuple[Any, Any, str]:
+    async def _run_turn_with_streaming(
+        self, state: _CodexTurnState, should_stop: Callable[[], bool] | None = None
+    ) -> tuple[Any, Any, str]:
         """Drive ``turn.stream()`` through the per-turn state, emitting the standard
         event protocol; returns ``(turn_result, latest_token_usage, agent_text)``.
 
         The enclosing ``communicate()`` owns the TurnStart/TurnEnd/AgentEnd
         boundaries; this drives the inner notification pump. ``state`` accumulates
         commands, the assistant transcript, sub-agent spawns and per-generation
         tokens — all mutated in place so a mid-turn crash keeps the partial.
+
+        The cooperative ``should_stop`` poll runs AFTER ``state.dispatch`` (the
+        emission that lets the watcher latch on the deciding tool call) and BEFORE
+        the next notification is pulled — the deciding item is kept, the next is
+        not. No-op when ``should_stop is None`` (behaviorally identical to before).
         """
         # Create the turn handle (starts the turn but doesn't block) + event stream.
         turn_handle = await self._run_async(self.thread.turn, state.user_input)
@@ -1405,6 +1433,11 @@ async def _run_turn_with_streaming(self, state: _CodexTurnState) -> tuple[Any, A
                     break
                 if state.dispatch(notification):  # True on a valid turn/completed
                     break
+                if should_stop is not None and should_stop():
+                    state.stopped_early_hit = True
+                    self._log.debug("Cooperative stop requested; ending notification pump at this boundary")
+                    self._interrupt_active_turn()  # best-effort; stops server-side spend
+                    break
         finally:
             self._active_turn_handle = None
             # Close any orphan tool (item/started without item/completed), flush any
@@ -1415,7 +1448,7 @@ async def _run_turn_with_streaming(self, state: _CodexTurnState) -> tuple[Any, A
             with contextlib.suppress(Exception):
                 await self._run_async(stream.close)
 
-        if state.turn_result is None:
+        if state.turn_result is None and not state.stopped_early_hit:
             raise RuntimeError("Turn did not complete (no turn/completed notification received)")
 
         # Belt-and-suspenders: if streaming surfaced no assistant transcript,
@@ -1427,7 +1460,9 @@ async def _run_turn_with_streaming(self, state: _CodexTurnState) -> tuple[Any, A
         # and nest them under the spawning Agent call. The parent stream never
         # carries the child's commands (Limited persistence drops them), but its
         # rollout always persists the raw function_call/local_shell_call items.
-        if state.spawned_children:
+        # Skipped on a cooperative stop: children may have no rollout yet and the
+        # run is already decided — recovery adds nothing the armed gate uses.
+        if state.spawned_children and not state.stopped_early_hit:
             await self._recover_subagent_tool_calls(
                 state.spawned_children,
                 state.collab_results,

src/coder_eval/orchestration/early_stop.py

@@ -159,7 +159,7 @@ def validate_early_stop(task: TaskDefinition) -> None:
     if not supports:
         raise EarlyStopConfigError(
             "run_limits.stop_early requires an agent that supports cooperative stopping "
-            + f"(currently Claude Code); agent type {agent_type!r} does not."
+            + f"(claude-code, codex, antigravity); agent type {agent_type!r} does not."
         )
 
     # (2) Arming requires at least one stop criterion.

tests/test_codex_agent_live.py

@@ -105,6 +105,45 @@ async def test_codex_live_edits_file_and_records_telemetry(tmp_path):
     assert record.commands, "expected command telemetry for the file creation"
 
 
+@_live
+async def test_codex_live_cooperative_stop_ends_turn_promptly(tmp_path):
+    """A ``should_stop`` that flips True after the first ToolStart ends the turn
+    cleanly as STOPPED_EARLY instead of running the multi-step task out — the
+    only automated check of real ``handle.interrupt()`` behavior."""
+    from coder_eval.streaming.events import AgentEndEvent, AgentEndStatus, ToolStartEvent
+
+    class _Sink:
+        def __init__(self):
+            self.tool_started = False
+            self.ends = []
+
+        def on_event(self, event):
+            if isinstance(event, ToolStartEvent):
+                self.tool_started = True
+            if isinstance(event, AgentEndEvent):
+                self.ends.append(event)
+
+    sink = _Sink()
+    agent = _make_agent()
+    await agent.start(str(tmp_path))
+    try:
+        record = await agent.communicate(
+            "Run `echo one`, then `echo two`, then `echo three`, each as a separate shell command, "
+            "then create three files a.txt, b.txt and c.txt.",
+            timeout=180,
+            stream_callback=sink,
+            should_stop=lambda: sink.tool_started,
+        )
+    finally:
+        await agent.stop()
+
+    # Clean cooperative stop: no crash, no pending partial, STOPPED_EARLY status.
+    assert record.crashed is False
+    assert agent.pending_turn is None
+    assert sink.ends, "expected an AgentEndEvent"
+    assert sink.ends[-1].status == AgentEndStatus.STOPPED_EARLY
+
+
 @_live
 async def test_codex_live_token_usage_populated(tmp_path):
     """Token usage is captured from the SDK and attached to the TurnRecord."""