fix: wait for Relay ATIF finalization

adapters/claude/src/nemo_fabric_adapters/claude/adapter.py

@@ -32,6 +32,7 @@
 from claude_agent_sdk import HookMatcher
 from claude_agent_sdk._errors import MessageParseError
 from nemo_fabric_adapters.common import lifecycle
+from nemo_fabric_adapters.common import relay_artifacts
 from nemo_fabric_adapters.common import relay_gateway
 from nemo_fabric_adapters.common import relay_hooks
 from nemo_fabric_adapters.common import utils as common_utils
@@ -726,6 +727,8 @@ def child_environment(
 def _relay_output(
     output: dict[str, Any],
     relay: ClaudeRelaySettings,
+    *,
+    artifacts: list[dict[str, str]] | None = None,
 ) -> dict[str, Any]:
     output["relay_runtime"] = {
         "enabled": True,
@@ -735,8 +738,10 @@ def _relay_output(
         "gateway_url": relay.gateway.url,
         "gateway_log_path": str(relay.gateway.log_path),
     }
-    output["relay_artifacts"] = common_utils.collect_relay_artifacts(
-        relay.plugin_config
+    output["relay_artifacts"] = (
+        common_utils.collect_relay_artifacts(relay.plugin_config)
+        if artifacts is None
+        else artifacts
     )
     return output
 
@@ -882,7 +887,7 @@ async def invoke(self, invocation: dict[str, Any]) -> dict[str, Any]:
         if self._unusable:
             return _failure(
                 "claude_runtime_unavailable",
-                "Claude runtime cannot accept another invocation after an SDK failure",
+                "Claude runtime cannot accept another invocation after a runtime failure",
             )
 
         try:
@@ -891,12 +896,42 @@ async def invoke(self, invocation: dict[str, Any]) -> dict[str, Any]:
         except ClaudeAdapterError as error:
             output = adapter_failure(error)
         else:
+            relay = self._relay
+            atif_before = (
+                relay_artifacts.snapshot_atif_files(relay.plugin_config)
+                if relay is not None
+                and relay_artifacts.expects_local_atif(relay.plugin_config)
+                else None
+            )
             output = await self._run_query(
                 payload,
                 client,
                 prompt,
                 invocation_timeout,
             )
+            if (
+                output.get("completed")
+                and relay is not None
+                and atif_before is not None
+            ):
+                finalized = await relay_artifacts.wait_for_finalized_atif(
+                    relay.plugin_config, atif_before
+                )
+                if finalized is None:
+                    self._unusable = True
+                    return _relay_output(
+                        adapter_failure(
+                            AdapterRelayError(
+                                "claude_relay_atif_timeout",
+                                "NeMo Relay did not finalize an ATIF artifact before the deadline",
+                                metadata={
+                                    "timeout_seconds": relay_artifacts.ATIF_FINALIZATION_TIMEOUT_SECONDS,
+                                },
+                            )
+                        ),
+                        relay,
+                        artifacts=[],
+                    )
 
         if self._relay is not None:
             output = _relay_output(output, self._relay)

adapters/codex/src/nemo_fabric_adapters/codex/adapter.py

@@ -32,6 +32,7 @@
 
 import nemo_fabric_adapters.common.relay_gateway as relay_gateway
 import nemo_fabric_adapters.common.relay_hooks as relay_hooks
+import nemo_fabric_adapters.common.relay_artifacts as relay_artifacts
 import nemo_fabric_adapters.common.utils as common_utils
 from nemo_fabric_adapters.common import lifecycle
 
@@ -922,7 +923,12 @@ async def _invoke_thread(
         return sdk_failure(error), False
 
 
-def _relay_output(output: dict[str, Any], relay: CodexRelaySettings) -> dict[str, Any]:
+def _relay_output(
+    output: dict[str, Any],
+    relay: CodexRelaySettings,
+    *,
+    artifacts: list[dict[str, str]] | None = None,
+) -> dict[str, Any]:
     output["relay_runtime"] = {
         "enabled": True,
         "emitter": "codex-sdk/nemo-relay",
@@ -931,8 +937,10 @@ def _relay_output(output: dict[str, Any], relay: CodexRelaySettings) -> dict[str
         "gateway_url": relay.gateway.url,
         "gateway_log_path": str(relay.gateway.log_path),
     }
-    output["relay_artifacts"] = common_utils.collect_relay_artifacts(
-        relay.plugin_config
+    output["relay_artifacts"] = (
+        common_utils.collect_relay_artifacts(relay.plugin_config)
+        if artifacts is None
+        else artifacts
     )
     return output
 
@@ -1064,18 +1072,50 @@ async def invoke(self, invocation: dict[str, Any]) -> dict[str, Any]:
             "request": invocation.get("request"),
         }
         if self._unusable:
-            output = _failure(
+            return _failure(
                 "codex_runtime_unavailable",
-                "Codex runtime cannot accept another invocation after an SDK failure",
+                "Codex runtime cannot accept another invocation after a runtime failure",
             )
-            return _relay_output(output, self._relay) if self._relay else output
 
         try:
             request_prompt(payload)
             timeout_seconds(payload)
             _reasoning_effort(payload)
             _output_schema(payload)
+            relay = self._relay
+            atif_before = (
+                relay_artifacts.snapshot_atif_files(relay.plugin_config)
+                if relay is not None
+                and relay_artifacts.expects_local_atif(relay.plugin_config)
+                else None
+            )
             output, usable = await _invoke_thread(payload, self._thread)
+            if (
+                output.get("completed")
+                and relay is not None
+                and atif_before is not None
+            ):
+                finalized = await relay_artifacts.wait_for_finalized_atif(
+                    relay.plugin_config, atif_before
+                )
+                if finalized is None:
+                    self._unusable = True
+                    return _relay_output(
+                        adapter_failure(
+                            AdapterRelayError(
+                                "codex_relay_atif_timeout",
+                                "NeMo Relay did not finalize an ATIF artifact before the deadline",
+                                metadata={
+                                    "timeout_seconds": relay_artifacts.ATIF_FINALIZATION_TIMEOUT_SECONDS,
+                                },
+                            )
+                        ),
+                        relay,
+                        artifacts=[],
+                    )
+        except AdapterRelayError as error:
+            output = adapter_failure(error)
+            usable = False
         except CodexAdapterError as error:
             output = adapter_failure(error)
             usable = True

adapters/common/src/nemo_fabric_adapters/common/relay_artifacts.py

@@ -0,0 +1,109 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Readiness checks for artifacts written asynchronously by NeMo Relay."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from pathlib import Path
+from typing import Any
+
+from nemo_fabric_adapters.common import utils as common_utils
+
+ATIF_FINALIZATION_TIMEOUT_SECONDS = 5.0
+ATIF_POLL_INTERVAL_SECONDS = 0.05
+AtifFileFingerprint = tuple[int, int, int, int]
+AtifSnapshot = dict[Path, AtifFileFingerprint]
+
+
+def expects_local_atif(plugin_config: dict[str, Any]) -> bool:
+    """Return whether Relay is configured to write ATIF to the local runtime.
+
+    Relay treats a non-empty storage list as remote-only, so there is no local
+    artifact for an adapter to await in that configuration.
+    """
+
+    for component in plugin_config.get("components", []):
+        if (
+            not isinstance(component, dict)
+            or component.get("kind") != "observability"
+            or component.get("enabled", True) is False
+        ):
+            continue
+        config = component.get("config")
+        if not isinstance(config, dict):
+            continue
+        atif = config.get("atif")
+        if isinstance(atif, dict) and atif.get("enabled") and not atif.get("storage"):
+            return True
+    return False
+
+
+def _atif_fingerprint(path: Path) -> AtifFileFingerprint | None:
+    try:
+        status = path.stat()
+    except (OSError, RuntimeError, ValueError):
+        return None
+    return (status.st_dev, status.st_ino, status.st_size, status.st_mtime_ns)
+
+
+def snapshot_atif_files(plugin_config: dict[str, Any]) -> AtifSnapshot:
+    """Capture ATIF file metadata before an adapter invocation.
+
+    Relay requires ``{session_id}`` in the ATIF filename template and creates a
+    new session scope for each turn, so a new path is the normal case. Metadata
+    fingerprints also detect a writer that rewrites an existing path without
+    reading or hashing artifacts from prior turns.
+    """
+
+    snapshot: AtifSnapshot = {}
+    for artifact in common_utils.collect_relay_artifacts(plugin_config):
+        if artifact.get("kind") != "atif":
+            continue
+        path = Path(artifact["path"])
+        fingerprint = _atif_fingerprint(path)
+        if fingerprint is not None:
+            snapshot[path] = fingerprint
+    return snapshot
+
+
+def _finalized_atif_path(
+    plugin_config: dict[str, Any], before: AtifSnapshot
+) -> Path | None:
+    """Find a new or changed ATIF path containing a complete JSON object."""
+
+    current = snapshot_atif_files(plugin_config)
+    for path in sorted(current):
+        if before.get(path) == current[path]:
+            continue
+        try:
+            # Relay writes directly to the final path, so existence alone does
+            # not prove that the JSON payload has been written completely.
+            document = json.loads(path.read_bytes())
+        except (OSError, UnicodeDecodeError, json.JSONDecodeError):
+            continue
+        if isinstance(document, dict):
+            return path
+    return None
+
+
+async def wait_for_finalized_atif(
+    plugin_config: dict[str, Any],
+    before: AtifSnapshot,
+    *,
+    timeout_seconds: float = ATIF_FINALIZATION_TIMEOUT_SECONDS,
+    poll_interval_seconds: float = ATIF_POLL_INTERVAL_SECONDS,
+) -> Path | None:
+    """Wait for one new or changed, complete ATIF file until a deadline."""
+
+    loop = asyncio.get_running_loop()
+    deadline = loop.time() + timeout_seconds
+    while True:
+        if path := _finalized_atif_path(plugin_config, before):
+            return path
+        remaining = deadline - loop.time()
+        if remaining <= 0:
+            return None
+        await asyncio.sleep(min(poll_interval_seconds, remaining))