From 02888ed962ece3968b5702f51c4cdf01e087b98a Mon Sep 17 00:00:00 2001 From: Bar Kaduri Date: Sun, 1 Nov 2026 11:17:31 +0200 Subject: [PATCH] Keep NAT's Guardian calls off the event loop Round trips run in worker threads, and lifecycle hooks go out in order from one sender thread, so a slow Guardian does not stall the workflow. Signed-off-by: Bar Kaduri --- adapters/docs/nat/README.md | 2 +- adapters/docs/nat/mapping.md | 2 +- adapters/src/acs/integrations/nat/adapter.py | 170 +++++++++-- .../tests/integrations/nat/test_event_loop.py | 281 ++++++++++++++++++ 4 files changed, 431 insertions(+), 24 deletions(-) create mode 100644 adapters/tests/integrations/nat/test_event_loop.py diff --git a/adapters/docs/nat/README.md b/adapters/docs/nat/README.md index 0edffcdb..4c56c2df 100644 --- a/adapters/docs/nat/README.md +++ b/adapters/docs/nat/README.md @@ -122,7 +122,7 @@ NAT's middleware fires per-attachment-point: workflow, function_groups, and indi 2. **Run `acs-nat-wire --check` in CI** on every workflow YAML you ship. It exits non-zero if any attachment point is missing `acs_guardian`. Treat that as a build failure. 3. **Re-run `acs-nat-wire --write` after any YAML edit** that adds or restructures `function_groups`, `workflow`, or individual `functions`. The marker comments make it idempotent. -What this gives you: **the same hard-stop guarantee Cursor's `failClosed: true` gives.** The middleware fires synchronously *before* the function executes; a Guardian deny raises `ACSGuardianDenied` (or sets `InvocationAction.SKIP`) and NAT does not call the function. +What this gives you: **the same hard-stop guarantee Cursor's `failClosed: true` gives.** The middleware runs *before* the function executes and waits for the Guardian's decision, off the event loop; a Guardian deny raises `ACSGuardianDenied` (or sets `InvocationAction.SKIP`) and NAT does not call the function. **Backstop for the dynamic-registration caveat:** the adapter also subscribes to NAT's `IntermediateStepManager`, which fires for every call NAT routes regardless of middleware wiring. The Guardian therefore *observes* every call (audit trail + trace), even ones the middleware path missed. Subscribers can't *block* — they're notification-only — so this is detection, not enforcement. Audit Guardian logs for `toolCallRequest` envelopes that arrived via lifecycle but never via middleware; that delta is your coverage gap. diff --git a/adapters/docs/nat/mapping.md b/adapters/docs/nat/mapping.md index 850fc00c..326c08a5 100644 --- a/adapters/docs/nat/mapping.md +++ b/adapters/docs/nat/mapping.md @@ -76,7 +76,7 @@ NAT ships `nvidia-nat-security` with `defense_middleware` for content-level chec ## Concurrency -NAT can invoke several wrapped functions at once in one process. The middleware correlates each call through its own invocation context and keeps no per-call file state; the session's handshake decision is shared by every call. The Guardian's rolling `chain_hash` orders steps by arrival, so the order of concurrent steps is the order they reach the Guardian. +NAT can invoke several wrapped functions at once in one process. The builder opens the session off the event loop before NAT gets the middleware, unless `ACS_DISABLED=1`, and every call reuses that decision, a refusal or no answer included. Each step's Guardian round trip runs in a worker thread (`asyncio.to_thread`), so a slow Guardian leaves the event loop running while the call waits for its decision. Lifecycle hooks go out from one worker thread in the order NAT reports them, each step waits until the hooks reported before it have gone out, and when NAT shuts the middleware down, every queued hook goes out first, in order, including one reported during the shutdown. A hook reported after that is audited as `lifecycle_hook_not_sent`. The middleware correlates each call through its own invocation context and keeps no per-call file state. The Guardian's rolling `chain_hash` orders steps by arrival, so the order of concurrent steps is the order they reach the Guardian. ## Conformance posture diff --git a/adapters/src/acs/integrations/nat/adapter.py b/adapters/src/acs/integrations/nat/adapter.py index b034c8d7..88cf6bb5 100644 --- a/adapters/src/acs/integrations/nat/adapter.py +++ b/adapters/src/acs/integrations/nat/adapter.py @@ -45,12 +45,18 @@ """ from __future__ import annotations +import asyncio +import atexit +import collections import hashlib import json import os import sys import threading import uuid +import weakref +from collections.abc import Callable +from concurrent.futures import Future from dataclasses import asdict, is_dataclass, replace from typing import Any, Optional @@ -362,6 +368,86 @@ def _apply_overrides_to_context(context: Any, overrides: dict[str, Any]) -> None # ----- Middleware class ----- +class _LifecycleSender: + """Sends lifecycle hooks on one thread, in the order they are accepted. + + close() waits until every accepted hook has gone out, still accepting + hooks that arrive meanwhile, and then accepts no more. + """ + + def __init__(self, send: Callable[[str, dict], None]) -> None: + self._send = send + self._pending: collections.deque = collections.deque() + self._changed = threading.Condition() + self._busy = False + self._closed = False + self._thread: threading.Thread | None = None + _OPEN_SENDERS.add(self) + + def submit(self, method: str, payload: dict) -> bool: + """Queue a hook behind the ones already accepted; False once closed.""" + with self._changed: + if self._closed: + return False + self._pending.append((method, payload)) + self._wake() + return True + + def sent(self) -> Future: + """A future that completes once every hook accepted so far has gone out.""" + done: Future = Future() + with self._changed: + if self._pending or self._busy: + self._pending.append(done) + self._wake() + return done + done.set_result(None) + return done + + def close(self) -> None: + with self._changed: + while self._pending or self._busy: + self._changed.wait() + self._closed = True + self._changed.notify_all() + thread = self._thread + if thread is not None and thread is not threading.current_thread(): + thread.join() + + def _wake(self) -> None: + if self._thread is None: + self._thread = threading.Thread(target=self._run, name="acs-lifecycle", daemon=True) + self._thread.start() + self._changed.notify_all() + + def _run(self) -> None: + while True: + with self._changed: + while not self._pending and not self._closed: + self._changed.wait() + if not self._pending: + return + item = self._pending.popleft() + self._busy = True + try: + if not isinstance(item, Future): + self._send(*item) + elif item.set_running_or_notify_cancel(): + item.set_result(None) + except Exception: # noqa: BLE001 + # _send audits its own failures; nothing may stop the hooks queued behind this one. + pass + finally: + with self._changed: + self._busy = False + self._changed.notify_all() + + +# The lifecycle thread is a daemon, so interpreter exit sends what is still queued first. +_OPEN_SENDERS: weakref.WeakSet = weakref.WeakSet() +atexit.register(lambda: [sender.close() for sender in list(_OPEN_SENDERS)]) + + class ACSMiddleware(FunctionMiddleware): # type: ignore[misc, valid-type] """NAT middleware that defers each call's allow/deny/modify decision to an ACS Guardian.""" @@ -382,8 +468,11 @@ def __init__(self, config): ) target = getattr(config, "target_function_or_group", None) or "nat" self._agent_id = os.environ.get("ACS_AGENT_ID") or f"nat:{hashlib.sha256(target.encode()).hexdigest()[:8]}" - # The session's §4.1 decision, made by the first handshake and kept for the middleware's lifetime. + # The session's §4.1 decision, made by one handshake (in the builder, or on the first call) and reused by every call. self._session: Handshake | None = None + self._session_lock = threading.Lock() + # One thread sends lifecycle hooks off the event loop, in the order NAT reports them. + self._lifecycle = _LifecycleSender(self._send_lifecycle_hook) self._server_hello: dict | None = None self._lifecycle_subscribed = False self._lifecycle_subscription = None @@ -397,8 +486,15 @@ def __init__(self, config): self._frozen_ctx_lock = threading.Lock() def _ensure_handshake(self) -> Handshake: + """The session's decision, opened once under a lock so concurrent first calls share one handshake.""" if self._session is not None: return self._session + with self._session_lock: + if self._session is None: + self._session = self._open_session() + return self._session + + def _open_session(self) -> Handshake: methods = ["steps/toolCallRequest", "steps/toolCallResult"] if _HAS_LIFECYCLE: methods += [ @@ -417,9 +513,22 @@ def _ensure_handshake(self) -> Handshake: # §6.4: deny wins over proceed, and a negotiated default timeout replaces the local one. self._posture = self._posture.negotiated(session.server_hello) self._server_hello = dict(session.server_hello) - self._session = session return session + async def _session_decision(self) -> Handshake: + """The session's decision, opened off the event loop when the builder has not opened it yet.""" + if self._session is not None: + return self._session + return await asyncio.to_thread(self._ensure_handshake) + + async def _lifecycle_sent(self) -> None: + """Wait until the lifecycle hooks NAT reported earlier have gone out, so the Guardian sees steps in event order.""" + await asyncio.wrap_future(self._lifecycle.sent()) + + def close(self) -> None: + """Send the lifecycle hooks still queued, including any that arrive meanwhile, then accept no more.""" + self._lifecycle.close() + def _ensure_lifecycle_subscribed(self) -> None: """Subscribe to NAT's IntermediateStepManager so workflow-boundary events fire ACS sessionStart / userMessage / agentResponse / @@ -462,29 +571,46 @@ def _on_intermediate_step(self, step) -> None: if event_type == IntermediateStepType.WORKFLOW_START: # Workflow input becomes both sessionStart (boundary marker) # and userMessage (the input itself). - self._emit_lifecycle_hook( + self._queue_lifecycle_hook( "steps/sessionStart", payload={"platform_context": {"workflow_name": payload.name or ""}}) input_text = _stringify_step_data(payload.data) if input_text: - self._emit_lifecycle_hook( + self._queue_lifecycle_hook( "steps/userMessage", payload={"content": [{"type": "text", "value": input_text}]}) elif event_type == IntermediateStepType.WORKFLOW_END: # Workflow output becomes agentResponse; sessionEnd closes the boundary. output_text = _stringify_step_data(payload.data) if output_text: - self._emit_lifecycle_hook( + self._queue_lifecycle_hook( "steps/agentResponse", payload={"content": [{"type": "text", "value": output_text}]}) - self._emit_lifecycle_hook( + self._queue_lifecycle_hook( "steps/sessionEnd", payload={"reason": "completed"}) + def _queue_lifecycle_hook(self, method: str, payload: dict) -> None: + """Hand a lifecycle hook to the lifecycle thread, so NAT's step callback returns without waiting on the Guardian.""" + if not self._lifecycle.submit(method, payload): + audit_event("lifecycle_hook_not_sent", cause="middleware_closed", method=method, + session_id=self._session_id) + + def _send_lifecycle_hook(self, method: str, payload: dict) -> None: + try: + self._emit_lifecycle_hook(method, payload) + except Exception as e: # noqa: BLE001 + audit_event("lifecycle_hook_failed", cause="adapter_exception", error=str(e), + method=method, session_id=self._session_id) + def _emit_lifecycle_hook(self, method: str, payload: dict) -> None: """Build, sign, and POST a lifecycle hook; errors are audited but never interrupt the workflow (enforcement lives in pre_invoke/post_invoke).""" + if self._acs.disabled: + audit_event("acs_disabled_bypass", session_id=self._session_id, method=method, + detail="ACS_DISABLED=1 — lifecycle hook not sent") + return gated = self._acs.gate(self._ensure_handshake()) if gated is not None: # §4.1: a refused or unanswered session sends the Guardian nothing. @@ -544,7 +670,7 @@ async def pre_invoke(self, context): method="steps/toolCallRequest", detail="ACS_DISABLED=1 — pre_invoke bypassed ungoverned") return None - gated = self._acs.gate(self._ensure_handshake()) + gated = self._acs.gate(await self._session_decision()) self._ensure_lifecycle_subscribed() if gated is not None: # §4.1: a refused or unanswered session sends the Guardian nothing. @@ -566,7 +692,8 @@ async def pre_invoke(self, context): result = self._posture.resolve(Undecided("adapter_exception", {"error": str(e)}, f"adapter error: {e}")) else: - result = self._round_trip(request) + await self._lifecycle_sent() + result = await asyncio.to_thread(self._round_trip, request) if isinstance(result, Resolution): return self._settle_pre(context, result) response = result.response @@ -616,22 +743,13 @@ async def pre_invoke(self, context): f"interpreted; failing closed (see issue #32)") async def post_invoke(self, context): - """Record the result. Apply Guardian's verdict to the output. - - - allow: pass through. - - modify (with modified_content): replace context.output. - - deny: clear context.output to None and tag with reasoning. The - tool already ran (post_invoke fires after execution), so the - side effect cannot be undone — but downstream consumers see no - output. This matches Specification §6.4's output-redaction gate. - - unknown: respect default_deny — drop output if true. - """ + """Apply the verdict to the output: the tool already ran, so a deny or an unusable decision withholds the output (issue #32).""" if self._acs.disabled: audit_event("acs_disabled_bypass", session_id=self._session_id, method="steps/toolCallResult", detail="ACS_DISABLED=1 — post_invoke bypassed ungoverned") return None - gated = self._acs.gate(self._ensure_handshake()) + gated = self._acs.gate(await self._session_decision()) if gated is not None: return self._settle_post(context, gated) if not method_evaluated(self._server_hello, "steps/toolCallResult"): @@ -655,7 +773,8 @@ async def post_invoke(self, context): result = self._posture.resolve(Undecided("adapter_exception", {"error": str(e)}, f"adapter error: {e}")) else: - result = self._round_trip(request) + await self._lifecycle_sent() + result = await asyncio.to_thread(self._round_trip, request) if isinstance(result, Resolution): return self._settle_post(context, result) response = result.response @@ -784,5 +903,12 @@ def _build_request( if _NAT_AVAILABLE and _HAS_REGISTRATION: @register_middleware(config_type=ACSMiddlewareConfig) # type: ignore[misc] async def build_acs_middleware(config: "ACSMiddlewareConfig", builder): # type: ignore[name-defined] - """NAT factory entry point. Yields the middleware instance for NAT to wire up.""" - yield ACSMiddleware(config) + """NAT factory entry point: open the session off the event loop before NAT gets the middleware, + and send the queued lifecycle hooks when NAT shuts it down.""" + middleware = ACSMiddleware(config) + if not middleware._acs.disabled: + await asyncio.to_thread(middleware._ensure_handshake) + try: + yield middleware + finally: + await asyncio.to_thread(middleware.close) diff --git a/adapters/tests/integrations/nat/test_event_loop.py b/adapters/tests/integrations/nat/test_event_loop.py new file mode 100644 index 00000000..be5c5326 --- /dev/null +++ b/adapters/tests/integrations/nat/test_event_loop.py @@ -0,0 +1,281 @@ +"""NAT's Guardian calls run off the event loop: the builder opens the session before it yields the middleware, +every call reuses that decision, and lifecycle hooks go out in event order without blocking NAT.""" +from __future__ import annotations + +import asyncio +import os +import tempfile +import threading +import time +import types +import unittest +import uuid +from pathlib import Path +from unittest import mock + +try: + from nat.data_models.intermediate_step import (IntermediateStepPayload, IntermediateStepType, + StreamEventData) + from nat.middleware.middleware import FunctionMiddlewareContext + _NAT_OK = True +except ImportError: + _NAT_OK = False + +from acs.client import _handshake +from acs.integrations.nat import adapter as nat_adapter +from acs.integrations.nat.adapter import ACSGuardianDenied, ACSMiddleware +from acs.testing.harness import ProgrammableGuardian + +if _NAT_OK: + from acs.integrations.nat.adapter import ACSMiddlewareConfig, build_acs_middleware + +SECRET = "nat-event-loop-secret" +SLOW_S = 0.5 + + +def _slow(handler): + def slow(req: dict) -> dict: + time.sleep(SLOW_S) + return handler(req) + return slow + + +async def _ticking(coro): + """Run coro while counting event-loop turns of 10 ms; a blocked loop stops counting.""" + ticks = 0 + + async def tick() -> None: + nonlocal ticks + while True: + ticks += 1 + await asyncio.sleep(0.01) + task = asyncio.create_task(tick()) + try: + return await coro, ticks + finally: + task.cancel() + + +@unittest.skipUnless(_NAT_OK, "nvidia-nat-core not installed") +class GuardianCallsOffTheEventLoop(unittest.TestCase): + def setUp(self) -> None: + records = tempfile.TemporaryDirectory(prefix="acs-handshake-") + self.addCleanup(records.cleanup) + env = {k: v for k, v in os.environ.items() if not k.startswith("ACS_")} + for patcher in (mock.patch.object(_handshake, "_HANDSHAKE_CACHE_DIR", Path(records.name)), + mock.patch.object(_handshake, "_HANDSHAKE_TIMEOUT_S", 5.0), + mock.patch.dict(os.environ, {**env, "ACS_HMAC_SECRET": SECRET}, clear=True)): + patcher.start() + self.addCleanup(patcher.stop) + self.guardian = ProgrammableGuardian(hmac_secret=SECRET) + self.guardian.start() + self.addCleanup(self.guardian.stop) + self.executed = 0 + + def _config(self, url: str | None = None) -> ACSMiddlewareConfig: + return ACSMiddlewareConfig(guardian_url=url or self.guardian.url(), default_deny=False, + session_id=str(uuid.uuid4())) + + async def _invoke(self, mw: ACSMiddleware) -> bool: + """One protected call through the middleware; True when NAT blocked it.""" + async def target(command: str) -> str: + self.executed += 1 + return f"ran: {command}" + context = FunctionMiddlewareContext(name="Bash", config=None, description=None, input_schema=None, + single_output_schema=type(None), stream_output_schema=type(None)) + try: + await mw.function_middleware_invoke(command="echo hi", call_next=target, context=context) + except ACSGuardianDenied: + return True + return False + + def test_the_builder_opens_the_session_off_the_loop_before_yielding(self) -> None: + self.guardian.handlers["handshake/hello"] = _slow(self.guardian._default_handshake) + + async def build() -> list[str]: + async with build_acs_middleware(self._config(), None): + return self.guardian.methods_seen() + seen, ticks = asyncio.run(_ticking(build())) + self.assertEqual(seen, ["handshake/hello"], "the session is open before NAT gets the middleware") + self.assertGreater(ticks, 10, "the event loop kept running during the handshake") + + def test_concurrent_calls_share_one_handshake(self) -> None: + self.guardian.handlers["handshake/hello"] = _slow(self.guardian._default_handshake) + for name in ("builder", "constructed directly"): + with self.subTest(name): + self.guardian.reset() + self.executed = 0 + + async def run() -> list[bool]: + if name == "builder": + async with build_acs_middleware(self._config(), None) as mw: + return await asyncio.gather(*(self._invoke(mw) for _ in range(5))) + mw = ACSMiddleware(self._config()) + try: + return await asyncio.gather(*(self._invoke(mw) for _ in range(5))) + finally: + mw.close() + blocked, ticks = asyncio.run(_ticking(run())) + self.assertEqual(blocked, [False] * 5) + self.assertEqual(self.executed, 5) + self.assertEqual(self.guardian.methods_seen().count("handshake/hello"), 1) + self.assertGreater(ticks, 10, "the event loop kept running during the handshake") + + def test_a_slow_guardian_leaves_the_loop_running(self) -> None: + for method in ("steps/toolCallRequest", "steps/toolCallResult"): + with self.subTest(method): + self.guardian.handlers = {"handshake/hello": self.guardian._default_handshake, + "__default__": self.guardian._default_allow, + method: _slow(self.guardian._default_allow)} + self.executed = 0 + + async def run() -> tuple[bool, int]: + async with build_acs_middleware(self._config(), None) as mw: + return await _ticking(self._invoke(mw)) + blocked, ticks = asyncio.run(run()) + self.assertFalse(blocked) + self.assertEqual(self.executed, 1) + self.assertGreater(ticks, 20, "the event loop kept running during the round trip") + + def test_a_refused_or_unanswered_session_never_runs_the_function(self) -> None: + down = ProgrammableGuardian(hmac_secret=SECRET) + down._server.server_close() + self.guardian.handlers["handshake/hello"] = lambda req: {"code": -32000, "message": "agent not permitted"} + for name, url, startup in (("refused", self.guardian.url(), ""), + ("no answer, refuse", down.url(), "refuse")): + with self.subTest(name): + self.guardian.reset() + os.environ["ACS_STARTUP_POSTURE"] = startup + + async def run() -> list[bool]: + async with build_acs_middleware(self._config(url), None) as mw: + return await asyncio.gather(*(self._invoke(mw) for _ in range(3))) + self.assertEqual(asyncio.run(run()), [True] * 3) + self.assertEqual(self.executed, 0) + if name == "refused": + self.assertEqual(self.guardian.methods_seen(), ["handshake/hello"], + "every call reuses the refusal and sends nothing") + + def test_a_disabled_middleware_skips_the_handshake(self) -> None: + os.environ["ACS_DISABLED"] = "1" + + async def run() -> tuple[list[str], bool]: + async with build_acs_middleware(self._config(), None) as mw: + seen = self.guardian.methods_seen() + mw._on_intermediate_step(types.SimpleNamespace(payload=IntermediateStepPayload( + event_type=IntermediateStepType.WORKFLOW_START, name="wf", data=StreamEventData(input="q")))) + return seen, await self._invoke(mw) + seen, blocked = asyncio.run(run()) + self.assertEqual((seen, blocked, self.executed), ([], False, 1)) + self.assertEqual(self.guardian.methods_seen(), []) + + def test_lifecycle_hooks_go_out_in_event_order_without_blocking_nat(self) -> None: + self.guardian.handlers["steps/sessionStart"] = _slow(self.guardian._default_allow) + + async def run() -> tuple[float, bool]: + async with build_acs_middleware(self._config(), None) as mw: + # NAT's step manager calls this callback with each step, on the thread that pushed it. + start = IntermediateStepPayload(event_type=IntermediateStepType.WORKFLOW_START, name="wf", + data=StreamEventData(input="q")) + began = time.monotonic() + mw._on_intermediate_step(types.SimpleNamespace(payload=start)) + pushed = time.monotonic() - began + blocked = await self._invoke(mw) + mw._on_intermediate_step(types.SimpleNamespace(payload=IntermediateStepPayload( + event_type=IntermediateStepType.WORKFLOW_END, name="wf", + data=StreamEventData(output="a"), UUID=start.UUID))) + return pushed, blocked + pushed, blocked = asyncio.run(run()) + self.assertFalse(blocked) + self.assertLess(pushed, SLOW_S / 2, "NAT's step callback returned without waiting on the Guardian") + self.assertEqual(self.guardian.methods_seen(), + ["handshake/hello", "steps/sessionStart", "steps/userMessage", "steps/toolCallRequest", + "steps/toolCallResult", "steps/agentResponse", "steps/sessionEnd"], + "every hook reaches the Guardian in event order, sessionEnd included") + + + def test_a_hook_reported_during_a_call_goes_out_before_its_result(self) -> None: + self.guardian.handlers["steps/sessionStart"] = _slow(self.guardian._default_allow) + + async def run() -> bool: + async with build_acs_middleware(self._config(), None) as mw: + async def target(command: str) -> str: + mw._on_intermediate_step(types.SimpleNamespace(payload=IntermediateStepPayload( + event_type=IntermediateStepType.WORKFLOW_START, name="wf", data=StreamEventData(input="q")))) + return f"ran: {command}" + context = FunctionMiddlewareContext(name="Bash", config=None, description=None, input_schema=None, + single_output_schema=type(None), stream_output_schema=type(None)) + await mw.function_middleware_invoke(command="echo hi", call_next=target, context=context) + return True + self.assertTrue(asyncio.run(run())) + self.assertEqual(self.guardian.methods_seen(), + ["handshake/hello", "steps/toolCallRequest", "steps/sessionStart", "steps/userMessage", + "steps/toolCallResult"]) + + def test_a_hook_that_arrives_while_closing_waits_its_turn(self) -> None: + """close() drains in order: a late hook neither blocks NAT nor overtakes the hooks queued before it.""" + mw = ACSMiddleware(self._config()) + started, release = threading.Event(), threading.Event() + sent: list[tuple[str, str]] = [] + + def emit(method: str, payload: dict) -> None: + if method == "steps/sessionStart": + started.set() + release.wait(5) + sent.append((method, threading.current_thread().name)) + mw._emit_lifecycle_hook = emit + start = IntermediateStepPayload(event_type=IntermediateStepType.WORKFLOW_START, name="wf", + data=StreamEventData(input="q")) + end = IntermediateStepPayload(event_type=IntermediateStepType.WORKFLOW_END, name="wf", + data=StreamEventData(output="a"), UUID=start.UUID) + + async def run() -> tuple[float, bool]: + mw._on_intermediate_step(types.SimpleNamespace(payload=start)) + self.assertTrue(await asyncio.to_thread(started.wait, 5)) + closing = asyncio.create_task(asyncio.to_thread(mw.close)) + await asyncio.sleep(0.05) + began = time.monotonic() + mw._on_intermediate_step(types.SimpleNamespace(payload=end)) + callback = time.monotonic() - began + try: + await asyncio.wait_for(mw._lifecycle_sent(), 0.1) + barrier_waited = False + except asyncio.TimeoutError: + barrier_waited = True + release.set() + await closing + return callback, barrier_waited + with mock.patch("acs.integrations.nat.adapter.audit_event") as audit: + callback, barrier_waited = asyncio.run(run()) + mw._on_intermediate_step(types.SimpleNamespace(payload=end)) + self.assertLess(callback, 0.05, "NAT's step callback returned without sending anything itself") + self.assertTrue(barrier_waited, "a step waits for the hooks still queued, even while closing") + self.assertEqual([method for method, _ in sent], + ["steps/sessionStart", "steps/userMessage", "steps/agentResponse", "steps/sessionEnd"]) + self.assertTrue(all(thread == "acs-lifecycle" for _, thread in sent), sent) + self.assertEqual([c.args[0] for c in audit.call_args_list], ["lifecycle_hook_not_sent"] * 2, + "a hook after close is audited, not sent") + + def test_a_cancelled_wait_does_not_stop_the_hooks_behind_it(self) -> None: + mw = ACSMiddleware(self._config()) + # A sender whose thread died would hold up interpreter exit; let this test fail instead. + self.addCleanup(nat_adapter._OPEN_SENDERS.discard, mw._lifecycle) + release, sent = threading.Event(), [] + + def emit(method: str, payload: dict) -> None: + if method == "first": + release.wait(5) + sent.append(method) + mw._emit_lifecycle_hook = emit + mw._queue_lifecycle_hook("first", {}) + mw._lifecycle.sent().cancel() + mw._queue_lifecycle_hook("second", {}) + release.set() + closer = threading.Thread(target=mw.close, daemon=True) + closer.start() + closer.join(5) + self.assertFalse(closer.is_alive(), "close() returns once the queue drains") + self.assertEqual(sent, ["first", "second"]) + +if __name__ == "__main__": + unittest.main()