60559d4e0e
_attach_child now mirrors a pending parent stop with the same split interrupt() uses for its own fan-out (hard -> hard_interrupt, soft -> interrupt), so a redirect is not turned into a cancel on a child that was attached late. _restore_parent_cancellation collapses to re-attaching the rejected unit's children: the replay is the attach step's job now. Test fixture: _Batch gained origin_session_history_delivery on main after the salvaged PR was written. Co-authored-by: illidan <noequal666@gmail.com>
306 lines
13 KiB
Python
306 lines
13 KiB
Python
"""A rejected background batch retains synchronous cancellation ownership."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import queue
|
|
import threading
|
|
import time
|
|
from concurrent.futures import Future
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from agent.interrupt_control import InterruptControlMixin
|
|
from agent.turn_context import _bind_interrupt_scope
|
|
from tools import async_delegation
|
|
from tools.delegate_tool_dispatch import _Batch, _dispatch_background
|
|
from tools.interrupt import is_interrupted, set_interrupt
|
|
from tools.process_registry import process_registry
|
|
|
|
|
|
class _Parent(InterruptControlMixin):
|
|
def __init__(self):
|
|
self.session_id = "capacity-interrupt-parent"
|
|
self._active_children = []
|
|
self._active_children_lock = threading.Lock()
|
|
self._execution_thread_id = None
|
|
self._interrupt_requested = False
|
|
self._hard_interrupt_requested = threading.Event()
|
|
self.quiet_mode = True
|
|
|
|
|
|
class _ControlledChild(_Parent):
|
|
"""Replace model work while retaining real interrupt and worker-start semantics."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.session_id = "capacity-interrupt-child"
|
|
self._delegate_role = "leaf"
|
|
self._delegate_depth = 1
|
|
self._delegate_saved_tool_names = []
|
|
self._credential_pool = None
|
|
self._subagent_id = None
|
|
self.tool_progress_callback = None
|
|
self.model = "test-model"
|
|
self.started = threading.Event()
|
|
self.stop_received = threading.Event()
|
|
self.unwinding = threading.Event()
|
|
self.allow_finish = threading.Event()
|
|
self.finished = threading.Event()
|
|
self.closed = threading.Event()
|
|
self.close_count = 0
|
|
self.closed_while_running = False
|
|
self.observed_interrupt = None
|
|
|
|
def interrupt(self, message=None, **kwargs):
|
|
accepted = super().interrupt(message, **kwargs)
|
|
self.stop_received.set()
|
|
return accepted
|
|
|
|
def hard_interrupt(self, message=None, **kwargs):
|
|
super().hard_interrupt(message, **kwargs)
|
|
self.stop_received.set()
|
|
|
|
def run_conversation(self, **_kwargs):
|
|
# A stop can arrive before this thread exists. Use the real turn-start
|
|
# binding so the pending agent interrupt must reach the tool thread too.
|
|
_bind_interrupt_scope(self, lambda: SimpleNamespace(_set_interrupt=set_interrupt))
|
|
self.started.set()
|
|
try:
|
|
assert self.stop_received.wait(30), "child never received cancellation"
|
|
assert self._interrupt_requested
|
|
assert is_interrupted(), "stop did not reach the child execution thread"
|
|
self.observed_interrupt = (
|
|
self._interrupt_message, self._hard_interrupt_requested.is_set(),
|
|
)
|
|
self.unwinding.set()
|
|
assert self.allow_finish.wait(30), "test did not release child cleanup"
|
|
return {
|
|
"final_response": "", "completed": False, "interrupted": True,
|
|
"api_calls": 0, "messages": [],
|
|
}
|
|
finally:
|
|
self.clear_interrupt()
|
|
self.finished.set()
|
|
|
|
def get_activity_summary(self):
|
|
return {"api_call_count": 0}
|
|
|
|
def close(self):
|
|
self.closed_while_running |= not self.finished.is_set()
|
|
self.close_count += 1
|
|
self.closed.set()
|
|
|
|
|
|
def _batch(parent, *children):
|
|
tasks = [{"goal": f"wait until cancelled {i}"} for i in range(len(children))]
|
|
parent._active_children.extend(children)
|
|
return _Batch(
|
|
task_list=tasks, children=[(i, tasks[i], child) for i, child in enumerate(children)], parent_agent=parent,
|
|
creds={"model": children[0].model}, context=None, top_role="leaf", max_children=len(children),
|
|
live_deleg_id=None, live_writers=[], live_paths=[], origin_wake_sid="",
|
|
origin_ui_session_id="", origin_owner_transport=None,
|
|
origin_owner_session_record=None, origin_session_history_delivery=False, overall_start=time.monotonic(),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def registry_state(tmp_path, monkeypatch):
|
|
monkeypatch.setenv("HERMES_HOME", str(tmp_path))
|
|
monkeypatch.delenv("HERMES_IGNORE_USER_CONFIG", raising=False)
|
|
(tmp_path / "config.yaml").write_text(
|
|
"delegation:\n max_concurrent_children: 1\n worktree_isolation: false\n",
|
|
encoding="utf-8",
|
|
)
|
|
async_delegation._reset_for_tests()
|
|
completion_queue = queue.Queue()
|
|
monkeypatch.setattr(process_registry, "completion_queue", completion_queue)
|
|
yield completion_queue
|
|
# Test bodies release their gates and join their workers before registry teardown.
|
|
if async_delegation._executor is not None:
|
|
async_delegation._executor.shutdown(wait=True)
|
|
async_delegation._reset_for_tests()
|
|
|
|
|
|
@pytest.mark.parametrize("rejection", ["capacity", "schedule_failure", "partial_schedule_failure"])
|
|
@pytest.mark.parametrize("stop_timing", ["running", "during_admission"])
|
|
@pytest.mark.parametrize("stop_kind", ["soft", "hard"])
|
|
def test_rejected_background_child_stops_with_parent(
|
|
registry_state, monkeypatch, tmp_path, rejection, stop_timing, stop_kind,
|
|
):
|
|
parent, child = _Parent(), _ControlledChild()
|
|
background_child = pending_child = None
|
|
if rejection == "partial_schedule_failure":
|
|
# Three independent units: one accepted, one rejected, one not yet submitted.
|
|
# The model-facing batch width is legal under the configured limit.
|
|
(tmp_path / "config.yaml").write_text(
|
|
"delegation:\n max_concurrent_children: 3\n worktree_isolation: false\n"
|
|
" independent_completions: true\n",
|
|
encoding="utf-8",
|
|
)
|
|
background_child, pending_child = _ControlledChild(), _ControlledChild()
|
|
background_child.session_id += "-background"
|
|
pending_child.session_id += "-pending"
|
|
batch = _batch(parent, background_child, child, pending_child)
|
|
else:
|
|
batch = _batch(parent, child)
|
|
occupied = threading.Event()
|
|
release_occupier = threading.Event()
|
|
admission_started = threading.Event()
|
|
continue_admission = threading.Event()
|
|
outcome = Future()
|
|
|
|
def occupy_slot():
|
|
occupied.set()
|
|
assert release_occupier.wait(30)
|
|
return {"status": "completed", "summary": "slot released"}
|
|
|
|
if rejection == "capacity":
|
|
accepted = async_delegation.dispatch_async_delegation(
|
|
goal="occupy the only slot", context=None, toolsets=None, role="leaf",
|
|
model=child.model, session_key="other-session", runner=occupy_slot,
|
|
max_async_children=1,
|
|
)
|
|
assert accepted["status"] == "dispatched"
|
|
assert occupied.wait(5)
|
|
elif rejection == "schedule_failure":
|
|
class RejectingExecutor:
|
|
def submit(self, *_args, **_kwargs):
|
|
raise RuntimeError("executor shut down")
|
|
|
|
monkeypatch.setattr(async_delegation, "_get_executor", lambda _n: RejectingExecutor())
|
|
else:
|
|
executor = async_delegation._get_executor(3)
|
|
|
|
class PartiallyRejectingExecutor:
|
|
submitted = 0
|
|
|
|
def submit(self, *args, **kwargs):
|
|
self.submitted += 1
|
|
if self.submitted == 2:
|
|
raise RuntimeError("unit submission failed")
|
|
return executor.submit(*args, **kwargs)
|
|
|
|
partial_executor = PartiallyRejectingExecutor()
|
|
monkeypatch.setattr(async_delegation, "_get_executor", lambda _n: partial_executor)
|
|
|
|
dispatch = async_delegation.dispatch_async_delegation_batch
|
|
admissions = 0
|
|
accepted_ids = []
|
|
|
|
def pause_admission(**kwargs):
|
|
nonlocal admissions
|
|
admissions += 1
|
|
rejected_admission = 2 if background_child is not None else 1
|
|
if admissions == rejected_admission:
|
|
admission_started.set()
|
|
assert continue_admission.wait(5)
|
|
result = dispatch(**kwargs)
|
|
if result.get("status") == "dispatched":
|
|
accepted_ids.append(result["delegation_id"])
|
|
return result
|
|
|
|
monkeypatch.setattr(async_delegation, "dispatch_async_delegation_batch", pause_admission)
|
|
|
|
def run_dispatch():
|
|
try:
|
|
outcome.set_result(json.loads(_dispatch_background(batch)))
|
|
except BaseException as exc:
|
|
outcome.set_exception(exc)
|
|
|
|
worker = threading.Thread(target=run_dispatch, daemon=True)
|
|
worker.start()
|
|
try:
|
|
assert admission_started.wait(5)
|
|
if background_child is not None:
|
|
assert background_child.started.wait(5)
|
|
request_stop = parent.hard_interrupt if stop_kind == "hard" else parent.interrupt
|
|
stop_message = "user correction or stop request"
|
|
if stop_timing == "during_admission":
|
|
request_stop(stop_message)
|
|
continue_admission.set()
|
|
assert child.started.wait(5)
|
|
if stop_timing == "running":
|
|
request_stop(stop_message)
|
|
|
|
assert child.stop_received.wait(5), "fallback lost parent cancellation ownership"
|
|
assert child.unwinding.wait(5)
|
|
assert child.observed_interrupt == (stop_message, stop_kind == "hard")
|
|
if background_child is not None:
|
|
assert not background_child.stop_received.is_set()
|
|
assert pending_child.stop_received.is_set(), "unsubmitted unit lost parent cancellation ownership"
|
|
assert not pending_child.started.is_set()
|
|
assert not outcome.done(), "dispatch returned while its child still owned resources"
|
|
assert child.close_count == 0
|
|
child.allow_finish.set()
|
|
result = outcome.result(timeout=5)
|
|
if background_child is None:
|
|
assert "SYNCHRONOUSLY" in result["note"]
|
|
assert result["results"][0]["status"] == "interrupted"
|
|
else:
|
|
assert result["status"] == "dispatched"
|
|
assert result["inline_results"][0]["status"] == "interrupted"
|
|
assert not background_child.stop_received.is_set()
|
|
assert pending_child.unwinding.wait(5)
|
|
assert pending_child.observed_interrupt == (stop_message, stop_kind == "hard")
|
|
assert child.finished.is_set()
|
|
assert child.close_count == 1
|
|
assert not child.closed_while_running
|
|
assert parent._active_children == []
|
|
finally:
|
|
continue_admission.set()
|
|
if not child.finished.is_set():
|
|
child.hard_interrupt("test teardown")
|
|
child.allow_finish.set()
|
|
worker.join(timeout=5)
|
|
release_occupier.set()
|
|
if background_child is not None:
|
|
async_delegation.interrupt_for_session(parent_session_id=parent.session_id)
|
|
for extra in (background_child, pending_child):
|
|
extra.allow_finish.set()
|
|
assert extra.closed.wait(5)
|
|
assert extra.finished.is_set()
|
|
assert extra.close_count == 1
|
|
assert not extra.closed_while_running
|
|
completed_ids = {registry_state.get(timeout=5)["delegation_id"] for _ in accepted_ids}
|
|
assert completed_ids == set(accepted_ids)
|
|
if rejection == "capacity":
|
|
completion = registry_state.get(timeout=5)
|
|
assert completion["delegation_id"] == accepted["delegation_id"]
|
|
assert not worker.is_alive()
|
|
|
|
|
|
def test_accepted_background_child_keeps_registry_cancellation_ownership(registry_state):
|
|
parent, child = _Parent(), _ControlledChild()
|
|
try:
|
|
result = json.loads(_dispatch_background(_batch(parent, child)))
|
|
assert result["status"] == "dispatched"
|
|
assert child.started.wait(5)
|
|
parent.interrupt()
|
|
# Parent interrupt fan-out is synchronous; observing it return establishes
|
|
# that a detached child did not receive it without a timing-based wait.
|
|
assert parent._interrupt_requested
|
|
assert not child.stop_received.is_set()
|
|
assert not child.finished.is_set()
|
|
parent.hard_interrupt("stop the current parent turn")
|
|
assert parent._hard_interrupt_requested.is_set()
|
|
assert not child.stop_received.is_set()
|
|
assert async_delegation.interrupt_for_session(parent_session_id=parent.session_id) == 1
|
|
assert child.unwinding.wait(5)
|
|
assert child.observed_interrupt[1] is True
|
|
assert child.close_count == 0
|
|
child.allow_finish.set()
|
|
completion = registry_state.get(timeout=5)
|
|
assert completion["delegation_id"] == result["delegation_id"]
|
|
assert completion["results"][0]["status"] == "interrupted"
|
|
assert child.finished.is_set()
|
|
assert child.close_count == 1
|
|
assert not child.closed_while_running
|
|
assert parent._active_children == []
|
|
finally:
|
|
if not child.finished.is_set():
|
|
child.hard_interrupt("test teardown")
|
|
child.allow_finish.set()
|
|
assert child.closed.wait(5)
|