From b595f2cd08afe128ae3321597c9c4c7b95429c3d Mon Sep 17 00:00:00 2001 From: Bar Kaduri Date: Sun, 1 Nov 2026 11:59:58 +0200 Subject: [PATCH] Report every NAT run's lifecycle from its start An acs_lifecycle exporter subscribes to each run's step stream before WORKFLOW_START, so every run reports its sessionStart and userMessage. Signed-off-by: Bar Kaduri --- adapters/docs/nat/README.md | 4 +- adapters/docs/nat/mapping.md | 6 +- adapters/src/acs/integrations/nat/adapter.py | 81 ++++++++-- adapters/src/acs/integrations/nat/wire.py | 15 +- .../integrations/nat/test_failure_modes.py | 14 +- .../integrations/nat/test_runner_lifecycle.py | 139 ++++++++++++++++++ 6 files changed, 227 insertions(+), 32 deletions(-) create mode 100644 adapters/tests/integrations/nat/test_runner_lifecycle.py diff --git a/adapters/docs/nat/README.md b/adapters/docs/nat/README.md index 4c56c2df..37f734c0 100644 --- a/adapters/docs/nat/README.md +++ b/adapters/docs/nat/README.md @@ -16,7 +16,7 @@ NAT exposes a `FunctionMiddleware` abstraction that wraps any function — tools The handshake/hello fires once per session, cached in process state so subsequent events skip the round-trip. -Independently of the middleware path, the adapter subscribes to NAT's `IntermediateStepManager` for `WORKFLOW_START` / `WORKFLOW_END` events and emits `steps/sessionStart` + `steps/userMessage` (on workflow start with input) and `steps/agentResponse` + `steps/sessionEnd` (on workflow end with output). The observer is auto-subscribed on the first `pre_invoke` call. **The lifecycle stream is observation-only; it cannot block calls.** Enforcement is the middleware path's job. +Independently of the middleware path, the adapter reports each run's `WORKFLOW_START` / `WORKFLOW_END` events as `steps/sessionStart` + `steps/userMessage` (on workflow start with input) and `steps/agentResponse` + `steps/sessionEnd` (on workflow end with output). When NAT builds the middleware, it registers an `acs_lifecycle` telemetry exporter, which NAT starts for every run before `WORKFLOW_START`. A middleware built without NAT's builder subscribes on each run's first call instead, so it reports that run from the call on. **The lifecycle stream is observation-only; it cannot block calls.** Enforcement is the middleware path's job. ### Decision honoring (§6.4) @@ -124,7 +124,7 @@ NAT's middleware fires per-attachment-point: workflow, function_groups, and indi 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. +**The lifecycle observer reports workflow boundaries only.** It sends `sessionStart`, `userMessage`, `agentResponse` and `sessionEnd` for each run and no tool calls, so a function the middleware is not wired to reaches the Guardian neither as a decision nor as an observation. Steps 1 and 2 close that gap: define every function in YAML, and run `acs-nat-wire --check` in CI. ## Smoke tests diff --git a/adapters/docs/nat/mapping.md b/adapters/docs/nat/mapping.md index 326c08a5..145a80b2 100644 --- a/adapters/docs/nat/mapping.md +++ b/adapters/docs/nat/mapping.md @@ -14,8 +14,8 @@ The ACS adapter is implemented as a `FunctionMiddleware` that, for each wrapped |---|---|---| | Any wrapped function `pre_invoke` (tools, LLMs, retrievers, memory — all of them) | `steps/toolCallRequest` | **Yes** — every wrapped call, with the function name as the tool name | | Any wrapped function `post_invoke` | `steps/toolCallResult` | **Yes** | -| Workflow start (via `IntermediateStepManager` subscription) | `steps/sessionStart` + `steps/userMessage` | **Yes** — observation-only, see below | -| Workflow end (same subscription) | `steps/agentResponse` + `steps/sessionEnd` | **Yes** — observation-only, see below | +| Workflow start (via the `acs_lifecycle` exporter NAT starts for every run) | `steps/sessionStart` + `steps/userMessage` | **Yes** — observation-only, see below | +| Workflow end (same exporter) | `steps/agentResponse` + `steps/sessionEnd` | **Yes** — observation-only, see below | | Retriever calls as `steps/knowledgeRetrieval` | — | **No.** Retrievers surface as `toolCallRequest` with the retriever's function name; the Guardian dispatches on name. A dedicated `knowledgeRetrieval` emission is possible future work. | | Memory read / write as `steps/memoryContextRetrieval` / `steps/memoryStore` | — | **No.** Same status: memory functions surface as generic tool calls. Documented gap — a Guardian keying policy on the memory-specific hooks will not see them from this adapter. | | Sub-workflow invocation as `steps/subagentStart` / `steps/subagentStop` | — | **No.** Nested workflow functions surface as generic tool calls; NAT exposes no distinct spawn boundary to the middleware. ACS-Core requires `subagentStart` only of a subagent-capable client, and the Guardian determines that capability by deployment-defined means. The requirement is vacuous for NAT only when sub-workflows are genuinely not distinguishable at this seam. Otherwise this is a mapping gap. | @@ -24,7 +24,7 @@ The Guardian can dispatch on the tool name to apply retriever-, memory-, or suba ### Lifecycle hooks are observation-only -The adapter's `_on_intermediate_step` subscriber emits `steps/sessionStart` + `steps/userMessage` on `WORKFLOW_START` and `steps/agentResponse` + `steps/sessionEnd` on `WORKFLOW_END` via `_emit_lifecycle_hook`. **This path is fire-and-forget.** NAT's `IntermediateStepManager` subscription model is a notification stream — a subscriber callback cannot veto the event after the fact. So although ACS-Core §hooks.md lists `agentResponse` as decision-eligible (ALLOW / DENY / MODIFY), the adapter cannot actually block or rewrite the workflow's output through this path. The envelopes are for trace + audit; enforcement on outputs must be placed at `pre_invoke` / `post_invoke` of the function that produced them, or as a separate review function in the YAML the workflow calls before returning. +The adapter's `_on_intermediate_step` callback emits `steps/sessionStart` + `steps/userMessage` on `WORKFLOW_START` and `steps/agentResponse` + `steps/sessionEnd` on `WORKFLOW_END` via `_emit_lifecycle_hook`. NAT creates a new step stream for each run, so the builder registers an `acs_lifecycle` telemetry exporter that NAT starts for every run before `WORKFLOW_START`; a middleware built without NAT's builder subscribes on each run's first call instead and reports that run from the call on. **This path is fire-and-forget.** NAT's `IntermediateStepManager` subscription model is a notification stream — a subscriber callback cannot veto the event after the fact. So although ACS-Core §hooks.md lists `agentResponse` as decision-eligible (ALLOW / DENY / MODIFY), the adapter cannot actually block or rewrite the workflow's output through this path. The envelopes are for trace + audit; enforcement on outputs must be placed at `pre_invoke` / `post_invoke` of the function that produced them, or as a separate review function in the YAML the workflow calls before returning. **Deliberately not mapped — the skill lifecycle set (`steps/skillRegister`, `steps/skillLoad`, `steps/skillUnload`).** NAT composes workflows from functions declared in the workflow YAML; the middleware surface exposes no runtime skill registration, activation, or unload events. Function *invocations* are governed at `pre_invoke`/`post_invoke` (`steps/toolCallRequest`/`toolCallResult`); lifecycle wiring becomes possible when NAT exposes lifecycle events for loadable components. diff --git a/adapters/src/acs/integrations/nat/adapter.py b/adapters/src/acs/integrations/nat/adapter.py index 88cf6bb5..b8d940eb 100644 --- a/adapters/src/acs/integrations/nat/adapter.py +++ b/adapters/src/acs/integrations/nat/adapter.py @@ -122,13 +122,24 @@ try: from nat.data_models.intermediate_step import IntermediateStepType # type: ignore[import-not-found] from nat.builder.context import Context as _NATContext # type: ignore[import-not-found] + from nat.builder.context import ContextState as _NATContextState # type: ignore[import-not-found] _HAS_LIFECYCLE = True except ImportError: IntermediateStepType = None # type: ignore[assignment] _NATContext = None # type: ignore[assignment] + _NATContextState = None # type: ignore[assignment] _HAS_LIFECYCLE = False -# InvocationAction.SKIP is on the dev branch; not in NAT 1.7.0 release. +# A telemetry exporter is the hook NAT starts for each run, subscribed before the run's WORKFLOW_START. +try: + from nat.cli.register_workflow import register_telemetry_exporter # type: ignore[import-not-found] + from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig # type: ignore[import-not-found] + from nat.observability.exporter.base_exporter import BaseExporter # type: ignore[import-not-found] + _HAS_EXPORTER = True +except ImportError: + _HAS_EXPORTER = False + +# InvocationAction.SKIP is absent from NAT 1.7.0. try: from nat.middleware.middleware import InvocationAction # type: ignore[attr-defined] _HAS_INVOCATION_ACTION = True @@ -474,14 +485,15 @@ def __init__(self, config): # 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 + # Set when the builder registered a lifecycle exporter, which NAT starts for every run. + self._lifecycle_exported = False + # The run streams this middleware subscribed to itself, one subscription each; a stream leaves with its run. + self._lifecycle_streams: weakref.WeakSet = weakref.WeakSet() # Locks the check-then-set in _ensure_lifecycle_subscribed # against double subscription from parallel pre_invoke calls. self._lifecycle_lock = threading.Lock() # WeakKeyDictionary fallback for frozen contexts — id(context) # would collide once GC recycles addresses. - import weakref self._frozen_ctx_rids: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() self._frozen_ctx_lock = threading.Lock() @@ -530,30 +542,36 @@ def close(self) -> None: self._lifecycle.close() def _ensure_lifecycle_subscribed(self) -> None: - """Subscribe to NAT's IntermediateStepManager so workflow-boundary - events fire ACS sessionStart / userMessage / agentResponse / - sessionEnd; idempotent and locked against double subscription. + """Without the builder's exporter, subscribe to the current run's step stream on the run's first call. + + NAT builds a new stream for each run and pushes WORKFLOW_START before + any call, so such a run reports from its first call on; the builder's + exporter subscribes before the start. Locked against double subscription. """ - if self._lifecycle_subscribed or not _HAS_LIFECYCLE: + if self._lifecycle_exported or not _HAS_LIFECYCLE: + return + try: + stream = _NATContextState.get().event_stream.get() + except Exception: # noqa: BLE001 + return + if stream in self._lifecycle_streams: return with self._lifecycle_lock: - # Re-check inside the lock. - if self._lifecycle_subscribed: + if stream in self._lifecycle_streams: return try: - ctx = _NATContext.get() - mgr = ctx.intermediate_step_manager + mgr = _NATContext.get().intermediate_step_manager except Exception: # noqa: BLE001 # No active Context; function-call hooks still fire via FunctionMiddleware. return try: - self._lifecycle_subscription = mgr.subscribe( + mgr.subscribe( on_next=self._on_intermediate_step, on_error=lambda e: audit_event( "lifecycle_subscription_error", session_id=self._session_id, error=str(e)), ) - self._lifecycle_subscribed = True + self._lifecycle_streams.add(stream) except Exception as e: # noqa: BLE001 audit_event("lifecycle_subscribe_failed", session_id=self._session_id, error=str(e)) @@ -908,7 +926,42 @@ async def build_acs_middleware(config: "ACSMiddlewareConfig", builder): # type: middleware = ACSMiddleware(config) if not middleware._acs.disabled: await asyncio.to_thread(middleware._ensure_handshake) + add_exporter = getattr(builder, "add_telemetry_exporter", None) + if add_exporter is not None and _HAS_EXPORTER and _HAS_LIFECYCLE: + key = str(uuid.uuid4()) + _LIFECYCLE_MIDDLEWARES[key] = middleware + await add_exporter(f"acs_lifecycle_{key}", ACSLifecycleExporterConfig(middleware_key=key)) + middleware._lifecycle_exported = True try: yield middleware finally: await asyncio.to_thread(middleware.close) + + +# ----- Per-run lifecycle exporter ----- + +# The middleware each registered lifecycle exporter feeds, by the key in its config. +_LIFECYCLE_MIDDLEWARES: weakref.WeakValueDictionary = weakref.WeakValueDictionary() + +if _NAT_AVAILABLE and _HAS_REGISTRATION and _HAS_EXPORTER and _HAS_LIFECYCLE: + class ACSLifecycleExporterConfig(TelemetryExporterBaseConfig, name="acs_lifecycle"): # type: ignore[misc, valid-type, call-arg] + """Added by `acs_guardian` for its own middleware; it is not configured in YAML.""" + middleware_key: str + + class _LifecycleExporter(BaseExporter): # type: ignore[misc, valid-type] + """Hands every step of a run to the middleware's lifecycle callback; NAT subscribes it before WORKFLOW_START.""" + + def __init__(self, middleware: ACSMiddleware) -> None: + super().__init__() + self._middleware = middleware + + def export(self, event) -> None: + self._middleware._on_intermediate_step(event) + + @register_telemetry_exporter(config_type=ACSLifecycleExporterConfig) # type: ignore[misc] + async def build_acs_lifecycle_exporter(config: "ACSLifecycleExporterConfig", builder): # type: ignore[name-defined] + """Build the exporter that feeds one `acs_guardian` middleware.""" + middleware = _LIFECYCLE_MIDDLEWARES.get(config.middleware_key) + if middleware is None: + raise ValueError("acs_lifecycle is added by the acs_guardian middleware; it is not configured on its own") + yield _LifecycleExporter(middleware) diff --git a/adapters/src/acs/integrations/nat/wire.py b/adapters/src/acs/integrations/nat/wire.py index 313017dd..dc588c35 100644 --- a/adapters/src/acs/integrations/nat/wire.py +++ b/adapters/src/acs/integrations/nat/wire.py @@ -11,12 +11,10 @@ `# acs-adapter-wired` comment so a re-run is a no-op and `--unwire` removes exactly what we added (and nothing else). -NAT's framework-wide lifecycle stream (`IntermediateStepManager`, -subscribed in acs_adapter.py:_ensure_lifecycle_subscribed) gives the -Guardian observability of every call even when middleware wiring is -incomplete — that's a backstop, not a substitute for wire+lint, because -lifecycle subscribers can only OBSERVE calls, not BLOCK them. Enforcement -still needs the middleware wired. +The adapter's lifecycle observer reports workflow start and end only, +never tool calls, so a function this tool does not wire reaches the +Guardian neither as a decision nor as an observation. Wire every +function, and lint in CI. Modes ===== @@ -42,9 +40,8 @@ - Sub-workflows loaded at runtime from other files. - Custom middleware classes that fork their own call path. -For those, rely on the lifecycle-subscription backstop the adapter -emits and audit the Guardian for "saw call via lifecycle but never via -middleware" findings. +Move those into the YAML, so this tool can wire them and --check can +see them. """ from __future__ import annotations diff --git a/adapters/tests/integrations/nat/test_failure_modes.py b/adapters/tests/integrations/nat/test_failure_modes.py index 8b620aed..64be7585 100644 --- a/adapters/tests/integrations/nat/test_failure_modes.py +++ b/adapters/tests/integrations/nat/test_failure_modes.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio +import contextvars import json import os import socket @@ -278,16 +279,21 @@ class FakeCtx: # Patch Context.get() so both threads see our FakeCtx import unittest.mock as mock + from nat.utils.reactive.subject import Subject with mock.patch.object(adapter_mod, "_NATContext") as patched_ctx: patched_ctx.get.return_value = FakeCtx() + # Both calls belong to one run, so both see that run's step stream. + token = ContextState.get().event_stream.set(Subject()) + self.addCleanup(ContextState.get().event_stream.reset, token) + runs = [contextvars.copy_context(), contextvars.copy_context()] barrier = threading.Barrier(2) - def runner(): + def runner(run): barrier.wait() # release both threads simultaneously - mw._ensure_lifecycle_subscribed() + run.run(mw._ensure_lifecycle_subscribed) - t1 = threading.Thread(target=runner) - t2 = threading.Thread(target=runner) + t1 = threading.Thread(target=runner, args=(runs[0],)) + t2 = threading.Thread(target=runner, args=(runs[1],)) t1.start(); t2.start() t1.join(); t2.join() diff --git a/adapters/tests/integrations/nat/test_runner_lifecycle.py b/adapters/tests/integrations/nat/test_runner_lifecycle.py new file mode 100644 index 00000000..f4e4dd49 --- /dev/null +++ b/adapters/tests/integrations/nat/test_runner_lifecycle.py @@ -0,0 +1,139 @@ +"""Every NAT run reports its whole lifecycle: the builder registers an exporter that NAT starts for each run before +WORKFLOW_START, so a run's first events and a later run's events all reach the Guardian.""" +from __future__ import annotations + +import asyncio +import contextvars +import os +import tempfile +import unittest +import uuid +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + +try: + from nat.builder.context import ContextState + from nat.builder.workflow_builder import WorkflowBuilder + from nat.data_models.intermediate_step import IntermediateStepPayload, IntermediateStepType, StreamEventData + from nat.middleware.middleware import FunctionMiddlewareContext + from nat.observability.exporter_manager import ExporterManager + from nat.runtime.runner import Runner + from nat.utils.reactive.subject import Subject + _NAT_OK = True +except ImportError: + _NAT_OK = False + +from acs.client import _handshake +from acs.testing.harness import ProgrammableGuardian + +if _NAT_OK: + from acs.integrations.nat.adapter import (ACSLifecycleExporterConfig, ACSMiddleware, ACSMiddlewareConfig, + build_acs_middleware) + +SECRET = "nat-runner-lifecycle-secret" +RUN = ["steps/sessionStart", "steps/userMessage", "steps/toolCallRequest", "steps/toolCallResult", + "steps/agentResponse", "steps/sessionEnd"] + + +class _Entry: + """A workflow entry function that makes one wrapped call through the middleware, as NAT's Runner invokes it.""" + has_single_output = True + config = SimpleNamespace(name="lifecycle-workflow", type="local") + instance_name = display_name = "lifecycle-workflow" + + def __init__(self, middleware: ACSMiddleware, executed: list[str]) -> None: + self.middleware, self.executed = middleware, executed + + async def ainvoke(self, message: str, to_type=None) -> str: + async def target(command: str) -> str: + self.executed.append(command) + return f"ran {command}" + context = FunctionMiddlewareContext(name="Read", config=None, description=None, input_schema=None, + single_output_schema=str, stream_output_schema=str) + return await self.middleware.function_middleware_invoke(command=message, call_next=target, context=context) + + +@unittest.skipUnless(_NAT_OK, "nvidia-nat-core not installed") +class EveryRunReportsItsLifecycle(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.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) + + def _runs(self, *messages: str, with_builder: bool = True) -> list[str]: + """Run each message through NAT's Runner with the exporters the middleware's builder registered.""" + executed: list[str] = [] + self.exporters: list[str] = [] + + async def run() -> None: + config = ACSMiddlewareConfig(guardian_url=self.guardian.url(), default_deny=False, + session_id=str(uuid.uuid4())) + async with WorkflowBuilder() as builder, \ + build_acs_middleware(config, builder if with_builder else None) as middleware: + # Workflow.run hands each Runner the exporters its builder collected. + exporters = {name: configured.instance for name, configured in builder._telemetry_exporters.items()} + self.exporters = [name.split("_")[0] + "_" + name.split("_")[1] for name in exporters] + manager = ExporterManager.from_exporters(exporters) + for message in messages: + async with Runner(message, _Entry(middleware, executed), ContextState.get(), manager) as runner: + await runner.result() + await middleware._lifecycle_sent() + asyncio.run(run()) + return executed + + def test_two_consecutive_runs_each_report_every_hook_in_order(self) -> None: + self.assertEqual(self._runs("first", "second"), ["first", "second"]) + self.assertEqual(self.exporters, ["acs_lifecycle"]) + self.assertEqual(self.guardian.methods_seen(), ["handshake/hello"] + RUN + RUN) + + def test_without_the_builder_each_run_reports_from_its_first_call(self) -> None: + self.assertEqual(self._runs("first", "second", with_builder=False), ["first", "second"]) + self.assertEqual(self.guardian.methods_seen(), ["handshake/hello"] + RUN[2:] + RUN[2:], + "the late subscription follows each run's new stream") + + def test_without_the_builder_overlapping_runs_are_each_subscribed_once(self) -> None: + for order in ("AA", "ABA", "ABAB"): + with self.subTest(order): + middleware = ACSMiddleware(ACSMiddlewareConfig(guardian_url=self.guardian.url(), + session_id=str(uuid.uuid4()))) + sent: list[str] = [] + middleware._emit_lifecycle_hook = lambda method, payload, run=None: sent.append(method) + streams = {name: Subject() for name in "AB"} + runs = {name: contextvars.copy_context() for name in "AB"} + for name, run in runs.items(): + run.run(ContextState.get().event_stream.set, streams[name]) + for name in order: + runs[name].run(middleware._ensure_lifecycle_subscribed) + end = IntermediateStepPayload(event_type=IntermediateStepType.WORKFLOW_END, name="wf", + data=StreamEventData(output="a")) + for name in sorted(set(order)): + streams[name].on_next(SimpleNamespace(payload=end)) + middleware.close() + self.assertEqual(sent, ["steps/agentResponse", "steps/sessionEnd"] * len(set(order)), + "each run's end is reported once") + + def test_a_disabled_middleware_sends_nothing_across_runs(self) -> None: + os.environ["ACS_DISABLED"] = "1" + self.assertEqual(self._runs("only"), ["only"]) + self.assertEqual(self.exporters, [], "a disabled middleware registers no exporter") + self.assertEqual(self.guardian.methods_seen(), []) + + + def test_a_lifecycle_exporter_configured_on_its_own_is_refused(self) -> None: + async def build() -> None: + async with WorkflowBuilder() as builder: + await builder.add_telemetry_exporter("acs", ACSLifecycleExporterConfig(middleware_key="unknown")) + with self.assertRaisesRegex(ValueError, "added by the acs_guardian middleware"): + asyncio.run(build()) + + +if __name__ == "__main__": + unittest.main()