Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

tests/gateway/test_hygiene_deferred_work_drain.py

@@ -0,0 +1,150 @@
+"""Shutdown accounting for executor work detached by hygiene timeouts."""
+
+import asyncio
+import threading
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from gateway.run import run_codex_hygiene_compaction
+from tests.gateway.restart_test_helpers import make_restart_runner
+
+
+@pytest.mark.asyncio
+async def test_deferred_worker_remains_active_until_executor_future_finishes():
+    runner, _adapter = make_restart_runner()
+    runner._cleanup_agent_resources_off_loop = AsyncMock()
+    worker = asyncio.get_running_loop().create_future()
+    agent = MagicMock()
+
+    runner._defer_agent_cleanup_until_future_done(
+        worker,
+        agent,
+        context="test hygiene timeout",
+    )
+
+    assert runner._active_deferred_agent_worker_count() == 1
+    assert runner._active_work_count() == 1
+
+    worker.set_result(([], None))
+    await asyncio.gather(*runner._deferred_agent_cleanup_tasks)
+
+    assert runner._active_deferred_agent_worker_count() == 0
+    runner._cleanup_agent_resources_off_loop.assert_awaited_once_with(
+        agent,
+        context="test hygiene timeout",
+    )
+
+
+@pytest.mark.asyncio
+async def test_timed_out_codex_hygiene_worker_remains_visible_to_shutdown():
+    runner, _adapter = make_restart_runner()
+    started = threading.Event()
+    release = threading.Event()
+
+    class BlockingCodexAgent:
+        _codex_session = object()
+        context_compressor = SimpleNamespace(compression_count=0)
+
+        def _compress_context(self, *_args, **_kwargs):
+            started.set()
+            release.wait(timeout=5.0)
+
+    agent = BlockingCodexAgent()
+    runner._agent_cache = {"tg:123": (agent, 0.0)}
+    runner._agent_cache_lock = None
+
+    outcome = await run_codex_hygiene_compaction(
+        runner,
+        "tg:123",
+        "sess-1",
+        auto_mode="hermes",
+        history=[{"role": "user", "content": "hello"}],
+        approx_tokens=100,
+        timeout_seconds=0.01,
+        failure_cooldown_seconds=-1.0,
+    )
+
+    assert started.is_set()
+    assert outcome == "failed:timeout"
+    assert runner._active_deferred_agent_worker_count() == 1
+    with patch("gateway.run.request_hard_interrupt") as interrupt:
+        runner._interrupt_running_agents("gateway shutdown")
+    interrupt.assert_called_once_with(agent, "gateway shutdown")
+
+    release.set()
+    for _ in range(100):
+        if runner._active_deferred_agent_worker_count() == 0:
+            break
+        await asyncio.sleep(0.01)
+    assert runner._active_deferred_agent_worker_count() == 0
+
+
+@pytest.mark.asyncio
+async def test_shutdown_drain_waits_for_deferred_hygiene_worker():
+    runner, _adapter = make_restart_runner()
+    worker = asyncio.get_running_loop().create_future()
+    runner._deferred_agent_workers = {worker: MagicMock()}
+
+    async def finish_worker():
+        await asyncio.sleep(0.12)
+        worker.set_result(([], None))
+
+    finisher = asyncio.create_task(finish_worker())
+    _snapshot, timed_out = await runner._drain_active_agents(2.0)
+    await finisher
+
+    assert timed_out is False
+    assert _snapshot == {}
+
+
+@pytest.mark.asyncio
+async def test_deferred_hygiene_worker_times_out_and_receives_interrupt():
+    runner, _adapter = make_restart_runner()
+    worker = asyncio.get_running_loop().create_future()
+    agent = MagicMock()
+    runner._deferred_agent_workers = {worker: agent}
+
+    _snapshot, timed_out = await runner._drain_active_agents(0.01)
+
+    assert timed_out is True
+    assert _snapshot == {}
+    with patch("gateway.run.request_hard_interrupt") as interrupt:
+        runner._interrupt_running_agents("gateway shutdown")
+    interrupt.assert_called_once_with(agent, "gateway shutdown")
+
+    worker.cancel()
+
+
+@pytest.mark.asyncio
+async def test_stop_interrupts_deferred_worker_before_teardown():
+    runner, adapter = make_restart_runner()
+    runner._restart_drain_timeout = 0.01
+    worker = asyncio.get_running_loop().create_future()
+
+    class DeferredAgent:
+        def __init__(self):
+            self.interrupts = []
+
+        def hard_interrupt(self, reason):
+            self.interrupts.append(reason)
+            if not worker.done():
+                worker.set_result(([], None))
+
+    agent = DeferredAgent()
+    runner._deferred_agent_workers = {worker: agent}
+    adapter.disconnect = AsyncMock()
+
+    with (
+        patch("gateway.status.remove_pid_file"),
+        patch("gateway.status.write_runtime_status"),
+        patch("cron.scheduler.mark_job_run"),
+        patch("tools.process_registry.process_registry.kill_all", return_value=0),
+        patch("tools.terminal_tool.cleanup_all_environments"),
+        patch("tools.browser_tool.cleanup_all_browsers"),
+    ):
+        await runner.stop()
+
+    assert agent.interrupts == ["Gateway shutting down"]
+    assert worker.done()