From 72f1fb68a1409d0c07f77245f154dbffc7a55707 Mon Sep 17 00:00:00 2001 From: Brian Krabach Date: Tue, 22 Sep 2026 19:10:07 -0700 Subject: [PATCH] Prevent dispatch workers from reopening after cleanup --- .../handlers/logging_handler.py | 28 +- .../tests/test_dispatcher_lifecycle.py | 345 ++++++++++++++++++ 2 files changed, 372 insertions(+), 1 deletion(-) create mode 100644 modules/hook-context-intelligence/tests/test_dispatcher_lifecycle.py diff --git a/modules/hook-context-intelligence/amplifier_module_hook_context_intelligence/handlers/logging_handler.py b/modules/hook-context-intelligence/amplifier_module_hook_context_intelligence/handlers/logging_handler.py index 7dd4f3ed..e0e93572 100644 --- a/modules/hook-context-intelligence/amplifier_module_hook_context_intelligence/handlers/logging_handler.py +++ b/modules/hook-context-intelligence/amplifier_module_hook_context_intelligence/handlers/logging_handler.py @@ -343,6 +343,7 @@ def __init__( maxsize=self._queue_capacity ) self._worker_task: asyncio.Task[None] | None = None + self._close_task: asyncio.Task[None] | None = None self._consecutive_failures = 0 # backoff driver only — never disables self._degraded_warned = False # Wall-clock start (time.monotonic()) of the CURRENT sustained-degraded @@ -473,7 +474,7 @@ def enqueue(self, event: str, data: dict[str, Any]) -> bool: """Enqueue an event for dispatch. HOT PATH — zero awaits, zero I/O. Returns ``True`` if the event was queued for delivery, ``False`` if it - was dropped because the queue is full. The caller uses this to tell + was dropped because the queue is full or closing. The caller uses this to tell "delivered to the server pipeline" apart from "dropped" — which, when the disk is ALSO full, is the difference between a stale local log and outright permanent data loss. @@ -488,6 +489,10 @@ def enqueue(self, event: str, data: dict[str, Any]) -> bool: the event. Any mutation would silently change the idempotency key and defeat server-side dedup. """ + # Cleanup is terminal. Events emitted or callbacks resumed after + # close() must not recreate a worker or HTTP client. + if self._close_task is not None: + return False self._ensure_worker() try: self._queue.put_nowait((event, data)) @@ -1255,6 +1260,19 @@ async def _post(self, event: str, data: dict[str, Any]) -> str: return _classify_http_outcome(response.status_code) async def close(self) -> None: + """Stop admission and await one bounded cleanup, even across callers. + + Cancelling a caller does not cancel cleanup. A later close() joins the + same operation; it cannot reopen the dispatcher or double-close a client. + """ + if self._close_task is None: + self._close_task = asyncio.create_task(self._drain_and_close()) + self._close_task.add_done_callback( + partial(_retrieve_task_exception, context=f"{self._name} dispatcher close") + ) + await asyncio.shield(self._close_task) + + async def _drain_and_close(self) -> None: """Drain, cancel worker, close client. Emits a loud WARNING when shutting down with undelivered events or in a @@ -1399,6 +1417,7 @@ def __init__(self, resolver: Any) -> None: self._parent_id: str = getattr(resolver, "parent_id", "") or "" self._resolve_instance_id: str = getattr(resolver, "resolve_instance_id", "") or "" self._dispatchers: list[_DestinationDispatcher] = [] + self._closed = False # Disk-pressure circuit breaker state (see _DISK_* constants). # _disk_backoff_seconds == 0.0 means healthy; > 0.0 means the breaker is # open and this is the current cooldown length. _disk_retry_at is the @@ -1440,6 +1459,12 @@ async def set_dispatchers(self, dispatchers: list[_DestinationDispatcher]) -> No and the carry-forward is a best-effort enhancement that must not raise for those -- it simply no-ops when the attributes aren't present. """ + # A ready/filter callback that resumes after final cleanup must not + # install new destinations. Close its candidates and fail explicitly + # so the caller cannot report them as active. + if self._closed: + await asyncio.gather(*(d.close() for d in dispatchers), return_exceptions=True) + raise RuntimeError("LoggingHandler is closed") old = self._dispatchers old_heartbeats_by_name: dict[str, bool] = {} for d in old: @@ -1853,6 +1878,7 @@ def _finalize_metadata(self, session_dir: Path, data: dict[str, Any]) -> None: # -- lifecycle management ------------------------------------------------ async def close(self) -> None: """Close all destination dispatchers concurrently.""" + self._closed = True await asyncio.gather(*(d.close() for d in self._dispatchers), return_exceptions=True) # -- metadata freshness update ------------------------------------------ diff --git a/modules/hook-context-intelligence/tests/test_dispatcher_lifecycle.py b/modules/hook-context-intelligence/tests/test_dispatcher_lifecycle.py new file mode 100644 index 00000000..cfd3d87e --- /dev/null +++ b/modules/hook-context-intelligence/tests/test_dispatcher_lifecycle.py @@ -0,0 +1,345 @@ +"""Offline lifecycle regressions for session:end emitted after hook cleanup. + +The HTTP sink only observes delivery and supplies scheduling barriers. These +tests prove component lifetime behavior, not acceptance by a real CI server. +""" + +from __future__ import annotations + +import asyncio +import errno +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import httpx +import pytest + +from amplifier_module_hook_context_intelligence.handlers.logging_handler import ( + LoggingHandler, + _DestinationDispatcher, +) + + +class _Sink: + def __init__(self) -> None: + self.events: list[str] = [] + self.clients: list[httpx.AsyncClient] = [] + self.request_started = asyncio.Event() + self.release_request = asyncio.Event() + self.release_request.set() + self.close_started = asyncio.Event() + self.release_close = asyncio.Event() + self.release_close.set() + self.closes = 0 + self.request_cancelled = False + + +class _Transport(httpx.AsyncBaseTransport): + def __init__(self, sink: _Sink) -> None: + self.sink = sink + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + self.sink.events.append(json.loads(request.content)["event"]) + self.sink.request_started.set() + try: + await self.sink.release_request.wait() + except asyncio.CancelledError: + self.sink.request_cancelled = True + raise + return httpx.Response(200) + + async def aclose(self) -> None: + self.sink.closes += 1 + self.sink.close_started.set() + await self.sink.release_close.wait() + + +@pytest.fixture +def sink(monkeypatch: pytest.MonkeyPatch) -> _Sink: + sink = _Sink() + client_type = httpx.AsyncClient + + def client(**kwargs: Any) -> httpx.AsyncClient: + result = client_type(transport=_Transport(sink), **kwargs) + sink.clients.append(result) + return result + + # Every client, including an incorrectly recreated one, stays offline. + monkeypatch.setattr(httpx, "AsyncClient", client) + return sink + + +def _dispatcher(tmp_path: Path) -> _DestinationDispatcher: + return _DestinationDispatcher( + name="lifecycle-test", + url="https://ci.invalid", + api_key="fixture-key", + workspace="fixture", + dispatch_timeout=1.0, + failure_threshold=3, + queue_capacity=8, + close_drain_timeout=0.1, + storage_path=tmp_path, + ) + + +@pytest.fixture +async def dispatcher(sink: _Sink, tmp_path: Path) -> Any: + dispatcher = _dispatcher(tmp_path) + yield dispatcher + sink.release_request.set() + sink.release_close.set() + await dispatcher.close() + + +@pytest.mark.parametrize("used_client", [False, True]) +async def test_session_end_after_cleanup_stays_local_without_reopening( + dispatcher: _DestinationDispatcher, sink: _Sink, tmp_path: Path, used_client: bool +) -> None: + """Core cleanup-before-end order retains JSONL, but cannot deliver remotely. + + The same boundary covers callbacks already scheduled when cleanup starts. + Remote terminal-event delivery requires the emitter to run before cleanup. + """ + session_dir = tmp_path / "session" / "context-intelligence" + handler = LoggingHandler( + SimpleNamespace(session_dir=lambda _: session_dir, working_dir="fixture") + ) + await handler.set_dispatchers([dispatcher]) + if used_client: + await handler("session:start", {"session_id": "session", "timestamp": "t0"}) + await asyncio.wait_for(sink.request_started.wait(), timeout=1) + previous_worker = dispatcher._worker_task + deliver = asyncio.Event() + + async def queued_hook() -> None: + await deliver.wait() + await handler("session:end", {"session_id": "session", "timestamp": "t1"}) + + callback = asyncio.create_task(queued_hook()) + try: + await handler.close() + deliver.set() + await callback + await asyncio.sleep(0) # let any erroneously spawned worker run + + records = [ + json.loads(line) for line in (session_dir / "events.jsonl").read_text().splitlines() + ] + assert records[-1]["event"] == "session:end" + assert sink.events == (["session:start"] if used_client else []) + assert len(sink.clients) == int(used_client) + assert all(client.is_closed for client in sink.clients) + assert previous_worker is None or previous_worker.done() + assert dispatcher._worker_task is None + finally: + deliver.set() + await callback + + +async def test_accepted_events_drain_in_order_and_close_once( + dispatcher: _DestinationDispatcher, sink: _Sink +) -> None: + assert dispatcher.enqueue("session:start", {"session_id": "s1"}) + assert dispatcher.enqueue("session:end", {"session_id": "s1"}) + worker = dispatcher._worker_task + + await dispatcher.close() + await dispatcher.close() + + assert sink.events == ["session:start", "session:end"] + assert sink.closes == 1 + assert all(client.is_closed for client in sink.clients) + assert worker is not None and worker.done() + assert dispatcher._worker_task is None + await asyncio.wait_for(dispatcher._queue.join(), timeout=1) + + +async def test_closed_dispatcher_rejects_admission_without_starting_worker( + dispatcher: _DestinationDispatcher, sink: _Sink +) -> None: + await dispatcher.close() + + assert dispatcher.enqueue("session:end", {"session_id": "s1"}) is False + assert dispatcher._queue.empty() + assert dispatcher._worker_task is None + assert sink.clients == [] + + +async def test_close_rejects_new_events_while_draining_accepted_event( + dispatcher: _DestinationDispatcher, sink: _Sink, monkeypatch: pytest.MonkeyPatch +) -> None: + sink.release_request.clear() + draining = asyncio.Event() + join = dispatcher._queue.join + + async def observed_join() -> None: + draining.set() + await join() + + monkeypatch.setattr(dispatcher._queue, "join", observed_join) + assert dispatcher.enqueue("session:start", {"session_id": "s1"}) + await asyncio.wait_for(sink.request_started.wait(), timeout=1) + closer = asyncio.create_task(dispatcher.close()) + try: + await asyncio.wait_for(draining.wait(), timeout=1) + assert dispatcher.enqueue("session:end", {"session_id": "s1"}) is False + finally: + sink.release_request.set() + await closer + + assert sink.events == ["session:start"] + assert sink.closes == 1 + assert dispatcher._worker_task is None + + +async def test_concurrent_close_waiters_share_client_cleanup( + dispatcher: _DestinationDispatcher, sink: _Sink +) -> None: + sink.release_close.clear() + assert dispatcher.enqueue("session:end", {"session_id": "s1"}) + first = asyncio.create_task(dispatcher.close()) + await asyncio.wait_for(sink.close_started.wait(), timeout=1) + second_started = asyncio.Event() + + async def second_close() -> None: + second_started.set() + await dispatcher.close() + + second = asyncio.create_task(second_close()) + try: + await second_started.wait() + assert not first.done() + assert not second.done(), "every close waiter must wait for the shared cleanup" + finally: + sink.release_close.set() + await asyncio.gather(first, second) + + assert sink.closes == 1 + assert dispatcher._worker_task is None + + +async def test_cancelled_close_waiter_does_not_cancel_cleanup( + dispatcher: _DestinationDispatcher, sink: _Sink +) -> None: + sink.release_request.clear() + assert dispatcher.enqueue("session:end", {"session_id": "s1"}) + await asyncio.wait_for(sink.request_started.wait(), timeout=1) + worker = dispatcher._worker_task + closer = asyncio.create_task(dispatcher.close()) + await asyncio.sleep(0) # enter close before cancelling its caller + closer.cancel() + with pytest.raises(asyncio.CancelledError): + await closer + + # Cleanup must finish without needing another close call to rescue it. + sink.release_request.set() + await asyncio.wait_for(sink.close_started.wait(), timeout=1) + await dispatcher.close() + + assert sink.events == ["session:end"] + assert sink.closes == 1 + assert worker is not None and worker.done() + assert dispatcher._worker_task is None + assert dispatcher.enqueue("after-cancel", {"session_id": "s1"}) is False + + +async def test_late_dispatcher_install_cannot_reopen_closed_handler( + dispatcher: _DestinationDispatcher, sink: _Sink, tmp_path: Path +) -> None: + """A delayed ready/filter callback must not install new post-cleanup workers.""" + handler = LoggingHandler(SimpleNamespace(session_dir=lambda _: tmp_path, working_dir="fixture")) + await handler.set_dispatchers([dispatcher]) + await handler.close() + replacement = _dispatcher(tmp_path) + try: + with pytest.raises(RuntimeError, match="closed"): + await handler.set_dispatchers([replacement]) + assert replacement.enqueue("session:end", {"session_id": "s1"}) is False + await handler("session:end", {"session_id": "s1"}) + assert json.loads((tmp_path / "events.jsonl").read_text())["event"] == "session:end" + assert sink.clients == [] + assert replacement._worker_task is None + finally: + await replacement.close() + await handler.close() + + +async def test_live_dispatcher_replacement_still_delivers( + dispatcher: _DestinationDispatcher, sink: _Sink, tmp_path: Path +) -> None: + handler = LoggingHandler(SimpleNamespace(session_dir=lambda _: tmp_path, working_dir="fixture")) + await handler.set_dispatchers([dispatcher]) + await handler("session:start", {"session_id": "s1"}) + previous_worker = dispatcher._worker_task + replacement = _dispatcher(tmp_path) + try: + await handler.set_dispatchers([replacement]) + assert previous_worker is not None and previous_worker.done() + await handler("session:end", {"session_id": "s1"}) + await handler.close() + + assert sink.events == ["session:start", "session:end"] + assert sink.closes == 2 + assert replacement._worker_task is None + assert all(client.is_closed for client in sink.clients) + finally: + await replacement.close() + await handler.close() + + +async def test_closed_dispatch_is_not_reported_as_delivery_when_disk_is_full( + dispatcher: _DestinationDispatcher, + sink: _Sink, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + handler = LoggingHandler(SimpleNamespace(session_dir=lambda _: tmp_path, working_dir="fixture")) + await handler.set_dispatchers([dispatcher]) + await handler.close() + + def disk_full(*args: Any) -> None: + raise OSError(errno.ENOSPC, "fixture disk full") + + monkeypatch.setattr(handler, "_write_session_to_disk", disk_full) + result = await handler("session:end", {"session_id": "s1"}) + + assert not handler._disk_episode_delivered + assert result.user_message_level == "warning" + assert result.user_message is not None + assert sink.clients == [] + assert dispatcher._worker_task is None + + +async def test_drain_timeout_cancels_inflight_and_cannot_restart( + dispatcher: _DestinationDispatcher, + sink: _Sink, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + issues: list[tuple[str, str]] = [] + monkeypatch.setattr(dispatcher, "_record_forwarding_issue", lambda *args: issues.append(args)) + sink.release_request.clear() + assert dispatcher.enqueue("session:end", {"session_id": "s1"}) + await asyncio.wait_for(sink.request_started.wait(), timeout=1) + worker = dispatcher._worker_task + + first = asyncio.create_task(dispatcher.close()) + await asyncio.sleep(0) + second = asyncio.create_task(dispatcher.close()) + first.cancel() + with pytest.raises(asyncio.CancelledError): + await first + await asyncio.wait_for(second, timeout=1) + await dispatcher.close() + + assert sink.request_cancelled + assert worker is not None and worker.done() + assert dispatcher._worker_task is None + assert sink.closes == 1 + assert dispatcher.enqueue("after-timeout", {"session_id": "s1"}) is False + assert sink.events == ["session:end"] + assert [kind for kind, _ in issues] == ["shutdown_undelivered"] + assert sum("shutdown:" in record.message for record in caplog.records) == 1