diff --git a/.agents/skills/contribute-adapter/SKILL.md b/.agents/skills/contribute-adapter/SKILL.md index cd50fe23c..ad7b7786b 100644 --- a/.agents/skills/contribute-adapter/SKILL.md +++ b/.agents/skills/contribute-adapter/SKILL.md @@ -83,8 +83,11 @@ Decide the following before implementation: `failed: true`, and structured `error` (`code`, `message`, `retryable`, and optional `metadata`); Fabric normalizes it into a failed `RunResult`. -> **TODO:** Revisit this output contract when Fabric adds streaming support; -> update this guidance and affected adapter evidence then. +- NVIDIA NeMo Relay streaming through `Runtime.invoke_stream()` does not change + the adapter lifecycle stdout contract. Fabric injects an SDK-owned NDJSON ATOF + endpoint before the adapter starts, NeMo Relay sends raw records out of band, + and the adapter still returns exactly one terminal lifecycle response. Do not + emit stream records on adapter stdout. - Scope workspace, generated config, state, sessions, and artifacts to the resolved runtime context. Stateful adapters must isolate Fabric runtime IDs. diff --git a/README.md b/README.md index 0af92e46a..ac7de763c 100644 --- a/README.md +++ b/README.md @@ -177,7 +177,7 @@ flowchart TB Adapter["Selected NeMo Fabric adapter"] Harness["Agent harness runtime\nHermes Agent | Codex | Claude Code | LangChain Deep Agents | custom"] Artifacts["Normalized results and artifacts\nresponse | logs | patches | telemetry refs"] - Relay["NeMo Relay\nATOF | ATIF | OTel | OpenInference when enabled"] + Relay["NVIDIA NeMo Relay\nATOF | ATIF | OTel | OpenInference when enabled"] Consumer --> Core Config --> Core @@ -198,7 +198,8 @@ Use the following resources to learn about NeMo Fabric: - [Example Notebooks](examples/notebooks/README.md) provide a guided tour of the Python SDK. - [Python SDK guide](docs/sdk/python.mdx): typed configuration, planning, - diagnostics, requests, multi-turn runtimes, parallelism, results, and errors. + diagnostics, requests, multi-turn runtimes, NeMo Relay streaming, + parallelism, results, and errors. - [Experimentation CLI](docs/experimentation/cli.mdx): presets, maintained examples, editable application scaffolds, and explicit non-goals. - [Getting Started overview](docs/about-nemo-fabric/overview.mdx): interface diff --git a/adapters/common/src/nemo_fabric_adapters/common/utils.py b/adapters/common/src/nemo_fabric_adapters/common/utils.py index e8dad5d3b..32944ed20 100644 --- a/adapters/common/src/nemo_fabric_adapters/common/utils.py +++ b/adapters/common/src/nemo_fabric_adapters/common/utils.py @@ -184,6 +184,10 @@ def merge_unique(*values: Any) -> list[str]: return merged +def without_none(mapping: dict[str, Any]) -> dict[str, Any]: + return {key: value for key, value in mapping.items() if value is not None} + + def dump_yaml(value: dict[str, Any]) -> str: try: import yaml diff --git a/adapters/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py b/adapters/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py index 2afc77939..37c827bb7 100644 --- a/adapters/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py +++ b/adapters/deepagents/src/nemo_fabric_adapters/deepagents/adapter.py @@ -579,6 +579,7 @@ async def invoke(self, invocation: dict[str, Any]) -> dict[str, Any]: user_message = request.get("input") or "" if not isinstance(user_message, str): user_message = json.dumps(user_message, sort_keys=True) + request_id = request.get("request_id") result_state: Any = None events: list[dict[str, Any]] = [] @@ -590,7 +591,9 @@ async def invoke(self, invocation: dict[str, Any]) -> dict[str, Any]: callback_handler = self._callback_handler_type() async with self._relay_plugin.plugin(self._relay_plugin_config): with self._relay_scope.scope( - "deepagents-request", self._relay_scope_type.Agent + "deepagents-request", + self._relay_scope_type.Agent, + metadata={"nemo_fabric_request_id": request_id}, ): ( result_state, diff --git a/adapters/hermes/src/nemo_fabric_adapters/hermes/adapter.py b/adapters/hermes/src/nemo_fabric_adapters/hermes/adapter.py index b89876eff..ced7c25e7 100755 --- a/adapters/hermes/src/nemo_fabric_adapters/hermes/adapter.py +++ b/adapters/hermes/src/nemo_fabric_adapters/hermes/adapter.py @@ -31,6 +31,26 @@ LOGGER = logging.getLogger(__name__) +def _fabric_stream_sink_enabled(config: dict[str, Any] | None) -> bool: + if config is None: + return False + for component in config.get("components") or []: + if not isinstance(component, dict) or component.get("kind") != "observability": + continue + component_config = component.get("config") + if not isinstance(component_config, dict): + continue + atof = component_config.get("atof") + if not isinstance(atof, dict): + continue + if any( + isinstance(sink, dict) and sink.get("name") == "nemo-fabric-stream" + for sink in atof.get("sinks") or [] + ): + return True + return False + + def validate_hermes_telemetry_provider(payload: dict[str, Any]) -> None: providers = common_utils.telemetry_providers(payload) if any(provider != "relay" for provider in providers): @@ -342,20 +362,41 @@ async def invoke(self, invocation: dict[str, Any]) -> dict[str, Any]: user_message = request.get("input") or "" if not isinstance(user_message, str): user_message = json.dumps(user_message, sort_keys=True) - try: - self._relay_session_pending = self._relay_plugin_config is not None - self._relay_finalize_hook_invoked = False - result, adapter_stdout = _invoke_hermes_turn( + + def invoke_turn() -> tuple[dict[str, Any], str]: + return _invoke_hermes_turn( agent=self._agent, settings=self._settings, user_message=user_message, conversation_history=self._conversation_history, ) - finally: - # Hermes' Relay plugin materializes ATIF when its session-finalize - # hook runs. Finalize the telemetry session for each Fabric - # invocation while retaining the native AIAgent and SessionDB. - self._finalize_relay_session() + + self._relay_session_pending = self._relay_plugin_config is not None + self._relay_finalize_hook_invoked = False + if _fabric_stream_sink_enabled(self._relay_plugin_config): + from nemo_relay import ScopeType, scope + + with scope.scope( + "nemo-fabric-invocation", + ScopeType.Agent, + metadata={ + "nemo_fabric_request_id": request.get("request_id"), + }, + ): + try: + result, adapter_stdout = invoke_turn() + finally: + # The Hermes plugin pushes its session below this correlation + # scope, so finalize that session before popping the parent. + self._finalize_relay_session() + else: + try: + result, adapter_stdout = invoke_turn() + finally: + # Hermes' Relay plugin materializes ATIF when its session-finalize + # hook runs. Finalize the telemetry session for each Fabric + # invocation while retaining the native AIAgent and SessionDB. + self._finalize_relay_session() messages = result.get("messages") or [] if isinstance(messages, list): self._conversation_history = messages diff --git a/docs/about-nemo-fabric/overview.mdx b/docs/about-nemo-fabric/overview.mdx index 7a8e1f9ef..3aeba8962 100644 --- a/docs/about-nemo-fabric/overview.mdx +++ b/docs/about-nemo-fabric/overview.mdx @@ -119,7 +119,7 @@ Harness installation and credential requirements differ by adapter. The contains the complete Hermes Agent environment recipe. Refer to the [Python SDK guide](../sdk/python.mdx) for planning, -diagnostics, typed requests, and multi-turn runtime examples. +diagnostics, typed requests, multi-turn runtimes, and NVIDIA NeMo Relay streaming. ## Choose Your Interface @@ -127,6 +127,7 @@ diagnostics, typed requests, and multi-turn runtime examples. | --- | --- | --- | | Python SDK | Your application owns job config, runtime lifecycle, or multi-turn state | [Client API](../reference/api/python-library-reference/nemo_fabric.client.md) | | Runtime API | You need multiple ordered turns over one live harness runtime | [Runtime](../reference/api/python-library-reference/nemo_fabric.runtime.md) | +| Streaming API | You need live ATOF records generated by NeMo Relay during a runtime turn | [Streaming](../reference/api/python-library-reference/nemo_fabric.streaming.md) | | `nemo-fabric` CLI | You are experimenting with harnesses, running maintained examples, or troubleshooting configs | [Experimentation CLI](../experimentation/cli.mdx) | | JSON Schema | You are building editors, validation, code generation, or another language binding | Committed schemas in the [repository](https://github.com/NVIDIA/NeMo-Fabric/tree/main/schemas) | @@ -141,8 +142,9 @@ obtain complete typed configs from built-in presets or maintained examples. observability settings without mutating the base config. 3. **Plan and diagnose** to resolve the adapter and check capabilities and requirements before spending work on a runtime. -4. **Run or start a runtime** through the shared start, invoke, and stop - lifecycle contract. +4. **Run or start a runtime** through the shared lifecycle contract. To consume + live ATOF records, enable NeMo Relay, start the runtime with `streaming=True`, + and call `Runtime.invoke_stream()`. 5. **Consume evidence** from `RunResult`: output, structured failure details, artifacts, events, and telemetry references. @@ -161,6 +163,12 @@ obtain complete typed configs from built-in presets or maintained examples. > Invoke multiple ordered turns and stop runtime handles safely. + + Consume live, raw NeMo Relay ATOF records and retrieve the terminal run result. + # API Overview @@ -12,6 +12,7 @@ SPDX-License-Identifier: Apache-2.0 */} - [`nemo_fabric.client`](./nemo_fabric.client.md#module-nemo_fabricclient): Native Python client for resolving and running NeMo Fabric agents. - [`nemo_fabric.runtime`](./nemo_fabric.runtime.md#module-nemo_fabricruntime): Runtime lifecycle support for the Fabric Python SDK. +- [`nemo_fabric.streaming`](./nemo_fabric.streaming.md#module-nemo_fabricstreaming): NVIDIA NeMo Relay streaming support for the NVIDIA NeMo Fabric Python SDK. - [`nemo_fabric.models`](./nemo_fabric.models.md#module-nemo_fabricmodels): Pydantic SDK models for NeMo Fabric configuration and requests. - [`nemo_fabric.types`](./nemo_fabric.types.md#module-nemo_fabrictypes): Public data contracts for the NeMo Fabric Python SDK. - [`nemo_fabric.errors`](./nemo_fabric.errors.md#module-nemo_fabricerrors): Public exception hierarchy for the NeMo Fabric Python SDK. @@ -21,6 +22,7 @@ SPDX-License-Identifier: Apache-2.0 */} - [`client.Fabric`](./nemo_fabric.client.md#class-fabric): Primary Python entrypoint for NeMo Fabric. - [`runtime.Runtime`](./nemo_fabric.runtime.md#class-runtime): One logical, stateful harness execution. - [`runtime.RuntimeStatus`](./nemo_fabric.runtime.md#class-runtimestatus): Lifecycle state of a runtime. +- [`streaming.InvokeStream`](./nemo_fabric.streaming.md#class-invokestream): Async iterator of raw ATOF records for one runtime invocation. - [`models.EnvironmentConfig`](./nemo_fabric.models.md#class-environmentconfig): Execution environment configuration supplied by the consumer. - [`models.FabricBaseModel`](./nemo_fabric.models.md#class-fabricbasemodel): Base class for SDK-facing Pydantic models. - [`models.FabricConfig`](./nemo_fabric.models.md#class-fabricconfig): SDK-facing typed Fabric agent configuration. diff --git a/docs/reference/api/python-library-reference/nemo_fabric.client.md b/docs/reference/api/python-library-reference/nemo_fabric.client.md index c719a3ead..d3a612c17 100644 --- a/docs/reference/api/python-library-reference/nemo_fabric.client.md +++ b/docs/reference/api/python-library-reference/nemo_fabric.client.md @@ -3,8 +3,8 @@ title: "Client" slug: "/reference/api/python-library-reference/client" description: "Resolve, plan, diagnose, and run agents with NeMo Fabric." --- -{/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 */} + # module `nemo_fabric.client` Native Python client for resolving and running NeMo Fabric agents. @@ -145,13 +145,14 @@ Execute one complete start, invoke, and stop lifecycle. start_runtime( config: 'FabricConfig', base_dir: 'str | PathLike[str] | None' = None, - overrides: 'Mapping[str, Any] | None' = None + overrides: 'Mapping[str, Any] | None' = None, + streaming: 'bool' = False ) → Runtime ``` Start a stateful runtime for one or more ordered invocations. -Each call starts a new logical runtime. Runtime-scoped overrides are recursively merged below invocation-scoped overrides. +Each call starts a new logical runtime. Runtime-scoped overrides are recursively merged below invocation-scoped overrides. Set ``streaming=True`` with NVIDIA NeMo Relay enabled to provision the SDK-owned ATOF endpoint used by ``Runtime.invoke_stream()``. @@ -160,6 +161,7 @@ Each call starts a new logical runtime. Runtime-scoped overrides are recursively - `config`: Complete typed ``FabricConfig``. - `base_dir`: Base directory for resolving relative paths. - `overrides`: JSON-compatible overrides applied to every invocation in the runtime unless superseded by invocation overrides. + - `streaming`: Whether to provision NeMo Relay ATOF streaming for ``Runtime.invoke_stream()``. @@ -170,7 +172,7 @@ Each call starts a new logical runtime. Runtime-scoped overrides are recursively **Raises:** - - `FabricConfigError`: If inputs or overrides are invalid. + - `FabricConfigError`: If inputs or overrides are invalid, or streaming is requested without NeMo Relay enabled. - `FabricNativeUnavailableError`: If the native extension is not installed. - `FabricRuntimeError`: If runtime startup fails. diff --git a/docs/reference/api/python-library-reference/nemo_fabric.errors.md b/docs/reference/api/python-library-reference/nemo_fabric.errors.md index 44bad6a22..5c5dad186 100644 --- a/docs/reference/api/python-library-reference/nemo_fabric.errors.md +++ b/docs/reference/api/python-library-reference/nemo_fabric.errors.md @@ -3,8 +3,8 @@ title: "Errors" slug: "/reference/api/python-library-reference/errors" description: "Structured exception hierarchy for config, capability, state, and runtime failures." --- -{/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 */} + # module `nemo_fabric.errors` Public exception hierarchy for the NeMo Fabric Python SDK. diff --git a/docs/reference/api/python-library-reference/nemo_fabric.models.md b/docs/reference/api/python-library-reference/nemo_fabric.models.md index 01dbf039a..459027ebe 100644 --- a/docs/reference/api/python-library-reference/nemo_fabric.models.md +++ b/docs/reference/api/python-library-reference/nemo_fabric.models.md @@ -3,8 +3,8 @@ title: "Models" slug: "/reference/api/python-library-reference/models" description: "Pydantic authoring models for NeMo Fabric config and request inputs." --- -{/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 */} + # module `nemo_fabric.models` Pydantic SDK models for NeMo Fabric configuration and requests. @@ -609,7 +609,7 @@ Return a detached JSON-compatible mapping for Rust/core calls. ## class `RelayConfigPolicy` -NeMo Relay config validation policy. +NVIDIA NeMo Relay config validation policy. --- @@ -669,6 +669,7 @@ Return a detached JSON-compatible mapping for Rust/core calls. ## class `RelayAtofFileSinkConfig` + NeMo Relay ATOF file sink configuration. @@ -729,6 +730,7 @@ Return a detached JSON-compatible mapping for Rust/core calls. ## class `RelayAtofStreamSinkConfig` + NeMo Relay ATOF stream sink configuration. diff --git a/docs/reference/api/python-library-reference/nemo_fabric.runtime.md b/docs/reference/api/python-library-reference/nemo_fabric.runtime.md index e0e1474f2..0953f16ff 100644 --- a/docs/reference/api/python-library-reference/nemo_fabric.runtime.md +++ b/docs/reference/api/python-library-reference/nemo_fabric.runtime.md @@ -3,8 +3,8 @@ title: "Runtime" slug: "/reference/api/python-library-reference/runtime" description: "Drive stateful multi-turn execution through the Runtime API." --- -{/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 */} + # module `nemo_fabric.runtime` Runtime lifecycle support for the Fabric Python SDK. @@ -64,6 +64,12 @@ Return the unique identifier for this started runtime lifecycle. Return the current ``ACTIVE``, ``STOPPED``, or ``FAILED`` state. +--- + +### property supports_streaming + +Return whether NVIDIA NeMo Relay ATOF streaming is enabled. + --- @@ -103,6 +109,30 @@ Run one turn on this runtime. --- +### method `invoke_stream` + +```python +invoke_stream( + input: 'Any' = None, + request: 'RunRequest | None' = None +) → InvokeStream +``` + +Start one turn and stream raw NeMo Relay ATOF records as they arrive. + +``input`` and ``request`` are mutually exclusive. The returned :class:`InvokeStream` yields raw ATOF dictionaries. Await ``stream.result()`` for the terminal normalized :class:`RunResult`. + + + +**Raises:** + + - `FabricCapabilityError`: If the runtime was not started with NeMo Relay enabled and ``streaming=True``. + - `FabricConfigError`: If request fields conflict or are not JSON-compatible. + - `FabricStateError`: If another turn or stream is active. + +--- + + ### method `stop` ```python diff --git a/docs/reference/api/python-library-reference/nemo_fabric.streaming.md b/docs/reference/api/python-library-reference/nemo_fabric.streaming.md new file mode 100644 index 000000000..1cee9c623 --- /dev/null +++ b/docs/reference/api/python-library-reference/nemo_fabric.streaming.md @@ -0,0 +1,52 @@ +--- +title: "Streaming" +slug: "/reference/api/python-library-reference/streaming" +description: "Consume raw NVIDIA NeMo Relay ATOF records and terminal invocation results." +--- + + +# module `nemo_fabric.streaming` +NeMo Relay streaming support for the NVIDIA NeMo Fabric Python SDK. + + + +--- + + +## class `InvokeStream` +Async iterator of raw ATOF records for one runtime invocation. + +Consume the final normalized result separately with :meth:`result`. If iteration stops early, call :meth:`aclose` before starting another turn. + + + + +--- + + +### method `aclose` + +```python +async def aclose() → None +``` + +Stop iteration and drain this turn without cancelling the invocation. + +--- + + +### method `result` + +```python +async def result() → RunResult +``` + +Return the terminal normalized result without adding it to the stream. + + + + +--- + +_This file was automatically generated via [lazydocs](https://github.com/ml-tooling/lazydocs)._ diff --git a/docs/reference/api/python-library-reference/nemo_fabric.types.md b/docs/reference/api/python-library-reference/nemo_fabric.types.md index 5de0e451d..6b05213cf 100644 --- a/docs/reference/api/python-library-reference/nemo_fabric.types.md +++ b/docs/reference/api/python-library-reference/nemo_fabric.types.md @@ -3,8 +3,8 @@ title: "Types" slug: "/reference/api/python-library-reference/types" description: "Typed config, request, plan, result, artifact, telemetry, and runtime contracts." --- -{/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: Apache-2.0 */} + # module `nemo_fabric.types` Public data contracts for the NeMo Fabric Python SDK. @@ -576,7 +576,7 @@ Reference to external or persisted telemetry for a run. **Attributes:** - - `provider`: Telemetry provider, such as Relay. + - `provider`: Telemetry provider, such as NVIDIA NeMo Relay. - `kind`: Reference kind, such as ``trace``. - `uri`: Optional location of persisted telemetry. - `trace_id`: Optional provider trace identifier. diff --git a/docs/sdk/python.mdx b/docs/sdk/python.mdx index a99c8b351..1dd5d4bbe 100644 --- a/docs/sdk/python.mdx +++ b/docs/sdk/python.mdx @@ -138,10 +138,10 @@ use the descriptor contract values. | Adapter ID | Models | Tools / Blocked Tools | MCP | Skills | Subagents | Telemetry | Persistent Local Host | Remote Service | | --- | --- | --- | --- | --- | --- | --- | --- | --- | -| `nvidia.fabric.claude` | Anthropic provider and Claude models | `allowed_tools` adapter setting / normalized `tools.blocked` | Normalized | Normalized | Not exposed | Relay: `atif`, `otel`, and `openinference` through hooks and gateway | Yes: connected `ClaudeSDKClient`, session, and optional Relay gateway | Not implemented | -| `nvidia.fabric.codex` | Built-in OpenAI and custom NVIDIA Responses providers with Codex-compatible models | Not normalized | Normalized: stdio, HTTP, and streamable HTTP | Normalized: `SKILL.md` directories | Not exposed | Relay: `atif`, `otel`, and `openinference` through hooks and gateway; native `otel` | Yes: `AsyncCodex` app-server client, thread, and optional Relay gateway | Not implemented | -| `nvidia.fabric.langchain.deepagents` | NVIDIA, OpenAI, OpenAI-compatible, and other LangChain providers | Built-ins and MCP / normalized middleware block list | Normalized | Normalized | Constrained: declarative local subagents inherit parent capabilities | Relay SDK: `atif`, `otel`, and `openinference`; native `otel` and `openinference` | Yes: compiled graph and async LangGraph checkpointer | Not implemented | -| `nvidia.fabric.hermes` | Normalized provider, model, and base URL | Toolsets / normalized disabled toolsets | Normalized | Normalized | Not exposed | Relay plugin: `atif`, `otel`, and `openinference` | Yes: `AIAgent`, `SessionDB`, and Relay plugin context | Not implemented | +| `nvidia.fabric.claude` | Anthropic provider and Claude models | `allowed_tools` adapter setting / normalized `tools.blocked` | Normalized | Normalized | Not exposed | NVIDIA NeMo Relay: `atif`, `otel`, and `openinference` through hooks and gateway | Yes: connected `ClaudeSDKClient`, session, and optional NeMo Relay gateway | Not implemented | +| `nvidia.fabric.codex` | Built-in OpenAI and custom NVIDIA Responses providers with Codex-compatible models | Not normalized | Normalized: stdio, HTTP, and streamable HTTP | Normalized: `SKILL.md` directories | Not exposed | NeMo Relay: `atif`, `otel`, and `openinference` through hooks and gateway; native `otel` | Yes: `AsyncCodex` app-server client, thread, and optional NeMo Relay gateway | Not implemented | +| `nvidia.fabric.langchain.deepagents` | NVIDIA, OpenAI, OpenAI-compatible, and other LangChain providers | Built-ins and MCP / normalized middleware block list | Normalized | Normalized | Constrained: declarative local subagents inherit parent capabilities | NeMo Relay SDK: `atif`, `otel`, and `openinference`; native `otel` and `openinference` | Yes: compiled graph and async LangGraph checkpointer | Not implemented | +| `nvidia.fabric.hermes` | Normalized provider, model, and base URL | Toolsets / normalized disabled toolsets | Normalized | Normalized | Not exposed | NeMo Relay plugin: `atif`, `otel`, and `openinference` | Yes: `AIAgent`, `SessionDB`, and NeMo Relay plugin context | Not implemented | "Normalized" means the adapter accepts the corresponding `FabricConfig` field and maps it to the harness. "Not normalized" does not mean that the @@ -189,8 +189,8 @@ capability_config.enable_relay( Config helpers edit the typed config before planning or starting a runtime. They do not modify already-started runtimes. Use `remove_mcp_server(name)` and `remove_skill_path(path)` to remove capabilities from a copied config. -Telemetry is enabled by adding entries to `telemetry.providers`; Relay-specific -settings live in the top-level `relay` block. +Telemetry is enabled by adding entries to `telemetry.providers`; settings +specific to NeMo Relay live in the top-level `relay` block. For evaluation or deployment variations, use ordinary Python functions and copies of the typed config. Supply the complete final config to NeMo Fabric. @@ -212,8 +212,8 @@ def review_agent_config(base, *, github_mcp: bool, relay: bool): variant = review_agent_config(config, github_mcp=True, relay=True) ``` -Relay observability is represented directly in the SDK config's top-level -`relay` block. ATOF uses Relay 0.6 file and stream sinks: +NeMo Relay observability is represented directly in the SDK config's top-level +`relay` block. ATOF uses NeMo Relay 0.6 file and stream sinks: ```python from nemo_fabric import ( @@ -245,7 +245,7 @@ relay_config.enable_relay( ) ``` -Additional Relay plugin components can be supplied generically with +Additional NeMo Relay plugin components can be supplied generically with `RelayComponentConfig` when their component package is available in the runtime environment. @@ -270,8 +270,9 @@ multiple independent runtimes. | `Fabric.plan(config, base_dir=...)` | No | You need to inspect the selected adapter, capability mapping, and runtime capabilities before running. | Does not start a runtime. | | `Fabric.doctor(config, base_dir=...)` | Yes | You need preflight diagnostics for adapter availability, config support, and environment assumptions. | Checks may touch runtime dependencies. | | `Fabric.run(config, base_dir=..., input=...)` | Yes | You need one complete start, invoke, result, stop lifecycle. | `base_dir` is optional. Pass a `RunRequest` instead when the invocation needs IDs, context, or overrides. | -| `Fabric.start_runtime(config, ...)` | Yes | You need state across multiple ordered invocations. | Returns a `Runtime`. Use it as an async context manager. | +| `Fabric.start_runtime(config, ..., streaming=False)` | Yes | You need state across multiple ordered invocations. | Returns a `Runtime`. Use `streaming=True` with NeMo Relay enabled to provision streaming. | | `Runtime.invoke(...)` | Yes | You need one turn on an existing runtime. | A runtime permits one active invocation at a time. | +| `Runtime.invoke_stream(...)` | No | You need live ATOF records generated by NeMo Relay for one turn. | Returns an async `InvokeStream`; await `stream.result()` for the terminal `RunResult`. | | `Runtime.stop()` | Yes | You need to stop or detach from the runtime. | Called automatically when using `async with`. | ## Single-Invocation Runs @@ -334,6 +335,117 @@ on one live thread, Deep Agents invokes one compiled graph and checkpointer, Hermes Agent reuses one agent and session database, and Claude keeps one connected SDK client. Harness-native identifiers remain adapter-internal. +## NeMo Relay Streaming + +Enable NeMo Relay before starting a runtime to consume raw ATOF records while a +turn runs. `Runtime.supports_streaming` reports whether this NeMo Relay path is +available. It is separate from `RuntimeCapabilities.streaming`, which describes +adapter-native progressive output. This separation is intentional: +`Runtime.invoke_stream()` exposes only ATOF records generated by NeMo Relay. A future +normalized Fabric streaming contract will address adapter-native progressive +output, such as Codex app-server message and item deltas. + +The following example streams one invocation and collects its terminal result: + +```python +import asyncio + +from nemo_fabric import ( + EnvironmentConfig, + Fabric, + FabricConfig, + HarnessConfig, + MetadataConfig, + ModelConfig, +) + + +async def main() -> None: + config = FabricConfig( + metadata=MetadataConfig(name="streaming-agent"), + harness=HarnessConfig(adapter_id="nvidia.fabric.codex"), + models={ + "default": ModelConfig(provider="openai", model="openai/gpt-5.4") + }, + environment=EnvironmentConfig(provider="local", workspace="."), + ).enable_relay() + + async with await Fabric().start_runtime(config, streaming=True) as runtime: + stream = runtime.invoke_stream(input="Review the latest patch") + async for record in stream: + print(record) + result = await stream.result() + print(result.status, result.output) + + +asyncio.run(main()) +``` + +`invoke_stream(...)` is synchronous and starts the invocation in the background. +The returned `InvokeStream` is an async iterator of raw NeMo Relay ATOF +dictionaries. The terminal `RunResult` stays out of band and is available only +through `await stream.result()`. Treat that terminal result as authoritative: +replace any provisional rendering instead of appending it. Calling +`stream.result()` does not finalize unread stream records. Calling +`Runtime.stop()` after the result completes finalizes the stream and discards +those unread records. + +Streaming has the following v0.1 constraints: + +- Streaming requires two explicit settings: enable NeMo Relay in `FabricConfig` + and pass `streaming=True` to `start_runtime(...)`. The `streaming=True` flag + does not enable NeMo Relay by itself; without NeMo Relay telemetry, + `start_runtime(...)` raises `FabricConfigError`. With both settings, the SDK + binds its HTTP listener and injects the reserved ATOF stream sink. +- The default `streaming=False` leaves the existing NeMo Relay configuration + unchanged; it neither enables nor disables ATOF. It also avoids the additional + listener HTTP and JSON parsing cost for runtimes that use only ATIF, + OpenTelemetry, OpenInference, or application-configured ATOF sinks. +- One invocation can be active on a runtime. Fully consume the stream or call + `await stream.aclose()` before starting another turn. +- Breaking an `async for` loop does not finalize the stream. Call + `await stream.aclose()` explicitly. It waits for the invocation to finish and + discards unread records; it does not cancel the harness turn. +- The end of async iteration only means that no more ATOF records are available. + It does not indicate invocation success. Always await `stream.result()`. + Invocation exceptions raise from that call, while harness-reported failures + are represented by the returned result's status and error fields. +- The SDK limits its queue to 1,024 records and 16 MiB of encoded record data. + It rejects individual records larger than 1 MiB and applies TCP backpressure. + NeMo Relay can drop records if a consumer stalls longer than its delivery + timeout. +- ATOF granularity depends on the NeMo Relay integration. Gateway harnesses expose + per-delta event structure, but current ATOF records keep token text in the + terminal scope. In-process harnesses expose scope-level progress. +- Reconstruct nested and parallel work with `uuid` and `parent_uuid`. Stream + order alone does not define the Deep Agents scope tree. +- The listener correlates records to one NeMo Relay scope tree. Deep Agents and + Hermes roots carry the Fabric request ID. For Claude and Codex, the SDK + matches NeMo Relay turn scopes by their role and 1-based turn index. Records + outside the matched root and its descendants are discarded, so delayed + records from another turn do not enter the active stream. If the gateway + turn sequence does not align with the Fabric invocation sequence, the SDK + yields no uncorrelated records and emits a `RuntimeWarning` after natural + stream exhaustion. A short drain window only collects late records from the + matched tree; it does not define turn ownership. +- The SDK listener binds to the address in `NEMO_FABRIC_STREAMING_HOST`, which + defaults to `127.0.0.1`. Override it when a Claude or Codex gateway must reach + the SDK through another network interface. Because the listener accepts ATOF + HTTP posts, restrict access to the configured interface. If async iteration + reaches the post-turn drain timeout without a NeMo Relay connection, or if + NeMo Relay sends data but no record matches the active turn, the SDK emits + one `RuntimeWarning` for that failure mode. It also warns if a NeMo Relay + upload terminates before completing its chunked request body because the + yielded stream can be incomplete. A caller that only awaits + `stream.result()` still receives the terminal result and does not run this + warning check. + +Claude and Codex streaming use the NeMo Relay `nemo-relay` gateway CLI and require a +stream-sink-capable release, version 0.6.0 or later. Follow the +[NeMo Relay CLI installation instructions](/getting-started/install#nemo-relay-cli) +to provision it. Hermes Agent and Deep Agents use their in-process NeMo Relay +integrations. + ## Application-Owned Parallelism Applications create independent runtimes when they want parallel work. NeMo Fabric @@ -386,7 +498,7 @@ Important fields: | `output` | Harness output normalized to the configured output schema. | | `error` | Structured failure metadata when available. | | `artifacts` | Output files, logs, patches, native artifacts, and other materialized references. | -| `telemetry` | References to Relay or other telemetry streams produced by the run. | +| `telemetry` | References to NeMo Relay or other telemetry streams produced by the run. | | `events` | Ordered normalized lifecycle and invocation events. | | `metadata` | Result-specific structured metadata. | | `runtime_id`, `invocation_id`, `request_id` | IDs for correlation across runtimes, logs, telemetry, and artifacts. | @@ -454,7 +566,7 @@ Runtime compatibility checks should validate: - selected adapter version; - selected harness version or version range; - required environment variables or secret references; -- optional capability support such as Relay, MCP, or tool exposure. +- optional capability support such as NeMo Relay, MCP, or tool exposure. ## Custom Fields And Adapter Settings diff --git a/justfile b/justfile index 5cecbf962..580a88d55 100644 --- a/justfile +++ b/justfile @@ -388,7 +388,15 @@ docs: PATH="{{ REPO_ROOT }}/.venv/bin:$PATH" bash scripts/generate_api_docs.sh uv run --no-sync python scripts/docs/generate_rust_library_reference.py npx --prefix docs --no-install fern check --warnings - npx --prefix docs --no-install fern docs broken-links --strict + fern_validation_root="$(mktemp -d)" + trap 'rm -rf "$fern_validation_root"' EXIT + uv run --no-sync python scripts/docs/sync_fern_docs_branch.py sync-dev \ + --source-root "$REPO_ROOT" \ + --target-root "$fern_validation_root" + ( + cd "$fern_validation_root/fern" + "$REPO_ROOT/docs/node_modules/.bin/fern" docs broken-links --strict + ) # Launch Jupyter Lab for the onboarding notebooks under examples/notebooks/. # Jupyter is fetched on demand so it stays out of the project lockfile. diff --git a/python/src/nemo_fabric/__init__.py b/python/src/nemo_fabric/__init__.py index e9836d2b7..54c1f40aa 100644 --- a/python/src/nemo_fabric/__init__.py +++ b/python/src/nemo_fabric/__init__.py @@ -37,6 +37,7 @@ from nemo_fabric.models import ToolsConfig from nemo_fabric.runtime import Runtime from nemo_fabric.runtime import RuntimeStatus +from nemo_fabric.streaming import InvokeStream from nemo_fabric.types import AdapterInfo from nemo_fabric.types import ArtifactManifest from nemo_fabric.types import ArtifactRef @@ -67,6 +68,7 @@ "FabricError", "FabricEvent", "HarnessConfig", + "InvokeStream", "McpConfig", "McpServerConfig", "MetadataConfig", diff --git a/python/src/nemo_fabric/client.py b/python/src/nemo_fabric/client.py index f6d23c10a..8666672ae 100644 --- a/python/src/nemo_fabric/client.py +++ b/python/src/nemo_fabric/client.py @@ -25,6 +25,11 @@ _run_native_lifecycle, _run_request_payload, ) +from nemo_fabric.streaming import ( + _AtofStreamListener, + _relay_enabled, + _with_stream_sink, +) from nemo_fabric.types import ( DoctorReport, RunPlan, @@ -184,37 +189,69 @@ async def start_runtime( *, base_dir: str | os.PathLike[str] | None = None, overrides: Mapping[str, Any] | None = None, + streaming: bool = False, ) -> Runtime: """Start a stateful runtime for one or more ordered invocations. Each call starts a new logical runtime. Runtime-scoped overrides are - recursively merged below invocation-scoped overrides. + recursively merged below invocation-scoped overrides. Set + ``streaming=True`` with NVIDIA NeMo Relay enabled to provision the SDK-owned + ATOF endpoint used by ``Runtime.invoke_stream()``. Args: config: Complete typed ``FabricConfig``. base_dir: Base directory for resolving relative paths. overrides: JSON-compatible overrides applied to every invocation in the runtime unless superseded by invocation overrides. + streaming: Whether to provision NeMo Relay ATOF streaming for + ``Runtime.invoke_stream()``. Returns: An active ``Runtime``. Use it as an asynchronous context manager to guarantee runtime shutdown. Raises: - FabricConfigError: If inputs or overrides are invalid. + FabricConfigError: If inputs or overrides are invalid, or streaming + is requested without NeMo Relay enabled. FabricNativeUnavailableError: If the native extension is not installed. FabricRuntimeError: If runtime startup fails. """ runtime_overrides = _json_mapping(overrides, "runtime overrides") - plan = await _call_blocking(lambda: self.plan(config, base_dir=base_dir)) - native = self._require_native_module("start_runtime") + stream_listener: _AtofStreamListener | None = None + runtime_config = config + if streaming and not _relay_enabled(config): + raise FabricConfigError("streaming requires Relay telemetry to be enabled") + if streaming: + try: + stream_listener = await _AtofStreamListener().start() + runtime_config = _with_stream_sink(config, stream_listener.url) + except Exception as error: + if stream_listener is not None: + await stream_listener.close() + raise FabricRuntimeError( + str(error), + stage="start", + code="stream_listener_start_failed", + ) from error + + try: + plan = await _call_blocking( + lambda: self.plan(runtime_config, base_dir=base_dir) + ) + native = self._require_native_module("start_runtime") + except BaseException: + if stream_listener is not None: + await stream_listener.close() + raise started_runtime: dict[str, Any] | None = None def start() -> dict[str, Any]: nonlocal started_runtime - started_runtime = json.loads(native.start_runtime(json.dumps(plan.to_mapping()))) + started_runtime = json.loads( + native.start_runtime(json.dumps(plan.to_mapping())) + ) return started_runtime try: @@ -232,16 +269,23 @@ def start() -> dict[str, Any]: ) except Exception: pass + if stream_listener is not None: + await stream_listener.close() raise except FabricError: + if stream_listener is not None: + await stream_listener.close() raise except Exception as error: + if stream_listener is not None: + await stream_listener.close() raise FabricRuntimeError(str(error), stage="start") from error return Runtime( client=self, plan=plan, runtime=runtime, overrides=runtime_overrides, + stream_listener=stream_listener, ) def _native_module(self) -> Any | None: diff --git a/python/src/nemo_fabric/models.py b/python/src/nemo_fabric/models.py index b93d67143..cb85d109b 100644 --- a/python/src/nemo_fabric/models.py +++ b/python/src/nemo_fabric/models.py @@ -24,7 +24,9 @@ from pydantic import BaseModel from pydantic import ConfigDict from pydantic import Field +from pydantic import SerializerFunctionWrapHandler from pydantic import field_validator +from pydantic import model_serializer from pydantic import model_validator @@ -233,7 +235,7 @@ def remove_server(self, name: str) -> Self: class RelayConfigPolicy(FabricBaseModel): - """NeMo Relay config validation policy.""" + """NVIDIA NeMo Relay config validation policy.""" unknown_component: Literal["ignore", "warn", "error"] = "warn" unknown_field: Literal["ignore", "warn", "error"] = "warn" @@ -255,12 +257,23 @@ class RelayAtofStreamSinkConfig(FabricBaseModel): type: Literal["stream"] = "stream" url: str transport: Literal["http_post", "websocket", "ndjson"] = "http_post" - headers: dict[str, str] = Field(default_factory=dict, exclude_if=lambda value: not value) - header_env: dict[str, str] = Field(default_factory=dict, exclude_if=lambda value: not value) + headers: dict[str, str] = Field(default_factory=dict) + header_env: dict[str, str] = Field(default_factory=dict) timeout_millis: int = 3000 field_name_policy: Literal["preserve", "replace_dots"] = "preserve" name: str | None = None + @model_serializer(mode="wrap") + def _omit_empty_header_maps(self, handler: SerializerFunctionWrapHandler) -> dict[str, Any]: + """Omit optional header maps when they are empty.""" + + data = handler(self) + if not self.headers: + data.pop("headers", None) + if not self.header_env: + data.pop("header_env", None) + return data + class RelayAtofConfig(FabricBaseModel): """NeMo Relay ATOF export configuration.""" diff --git a/python/src/nemo_fabric/runtime.py b/python/src/nemo_fabric/runtime.py index 3fd736b4e..d5ed465a5 100644 --- a/python/src/nemo_fabric/runtime.py +++ b/python/src/nemo_fabric/runtime.py @@ -14,8 +14,15 @@ from pydantic import ValidationError -from nemo_fabric.errors import FabricConfigError, FabricError, FabricRuntimeError, FabricStateError +from nemo_fabric.errors import ( + FabricCapabilityError, + FabricConfigError, + FabricError, + FabricRuntimeError, + FabricStateError, +) from nemo_fabric.models import RunRequest +from nemo_fabric.streaming import InvokeStream, _AtofStreamListener from nemo_fabric.types import RunPlan, RunResult, RuntimeHandle @@ -51,6 +58,7 @@ def __init__( plan: RunPlan | Mapping[str, Any], runtime: RuntimeHandle | Mapping[str, Any], overrides: Mapping[str, Any] | None = None, + stream_listener: _AtofStreamListener | None = None, ) -> None: """lazydocs: ignore""" @@ -64,6 +72,8 @@ def __init__( self._invocations: list[dict[str, Any]] = [] self._status = RuntimeStatus.ACTIVE self._current_task: asyncio.Task[Any] | None = None + self._current_stream: InvokeStream | None = None + self._stream_listener = stream_listener self._closing = False @property @@ -96,6 +106,12 @@ def runtime_id(self) -> str: return self._runtime.runtime_id + @property + def supports_streaming(self) -> bool: + """Return whether NVIDIA NeMo Relay ATOF streaming is enabled.""" + + return self._stream_listener is not None + async def invoke( self, *, @@ -125,18 +141,22 @@ async def invoke( normalized result. """ - if self._status is not RuntimeStatus.ACTIVE: - raise FabricStateError(f"cannot invoke a {self._status.value} runtime") - if self._closing: - raise FabricStateError("cannot invoke while runtime shutdown is in progress") - if self._current_task is not None: - raise FabricStateError("runtime is already running an invocation") + if self._current_stream is not None and not self._current_stream._finalized: + raise FabricStateError( + "a streaming invocation is active; fully consume it or call " + "`await stream.aclose()` before starting another turn" + ) + self._ensure_invocable() + payload = _run_request_payload(input=input, request=request) + return await self._invoke_payload(payload) + + async def _invoke_payload( + self, + payload: dict[str, Any], + ) -> RunResult: + self._ensure_invocable() self._current_task = asyncio.current_task() try: - payload = _run_request_payload( - input=input, - request=request, - ) merged = _merge_overrides(self._overrides, payload.get("overrides")) if merged: payload["overrides"] = merged @@ -203,6 +223,60 @@ def stop_after_cancel() -> Any: finally: self._current_task = None + def invoke_stream( + self, + *, + input: Any = None, + request: RunRequest | None = None, + ) -> InvokeStream: + """Start one turn and stream raw NeMo Relay ATOF records as they arrive. + + ``input`` and ``request`` are mutually exclusive. The returned + :class:`InvokeStream` yields raw ATOF dictionaries. Await + ``stream.result()`` for the terminal normalized :class:`RunResult`. + + Raises: + FabricCapabilityError: If the runtime was not started with NeMo Relay + enabled and ``streaming=True``. + FabricConfigError: If request fields conflict or are not + JSON-compatible. + FabricStateError: If another turn or stream is active. + """ + + if self._stream_listener is None: + raise FabricCapabilityError( + "streaming requires Relay telemetry and " + "start_runtime(..., streaming=True)", + stage="invoke", + code="streaming_unavailable", + details={"capability": "streaming"}, + ) + if self._current_stream is not None and not self._current_stream._finalized: + raise FabricStateError( + "a streaming invocation is active; fully consume it or call " + "`await stream.aclose()` before starting another turn" + ) + self._ensure_invocable() + payload = _run_request_payload(input=input, request=request) + stream = InvokeStream( + self._invoke_payload(payload), + self._stream_listener, + request_id=payload["request_id"], + turn_index=len(self._invocations) + 1, + ) + self._current_stream = stream + return stream + + def _ensure_invocable(self) -> None: + if self._status is not RuntimeStatus.ACTIVE: + raise FabricStateError(f"cannot invoke a {self._status.value} runtime") + if self._closing: + raise FabricStateError( + "cannot invoke while runtime shutdown is in progress" + ) + if self._current_task is not None: + raise FabricStateError("runtime is already running an invocation") + async def stop(self) -> None: """Destroy an idle runtime exactly once. @@ -218,6 +292,13 @@ async def stop(self) -> None: if self._status is RuntimeStatus.STOPPED: return + if self._current_stream is not None and not self._current_stream._finalized: + if not self._current_stream._task.done(): + raise FabricStateError( + "cannot stop while a streaming invocation is active; await " + "`stream.result()` and then call `await stream.aclose()`" + ) + await self._current_stream.aclose() if self._current_task is not None: raise FabricStateError("cannot stop while a turn is in flight") if self._closing: @@ -252,6 +333,8 @@ def stop() -> Any: self._status = RuntimeStatus.STOPPED finally: self._closing = False + if self._stream_listener is not None: + await self._stream_listener.close() def _absorb(self, result: RunResult) -> None: self._invocations.append( @@ -276,11 +359,16 @@ async def __aexit__( traceback: object, ) -> None: try: + if self._current_stream is not None and not self._current_stream._finalized: + await self._current_stream.aclose() await self.stop() except Exception as cleanup_error: if exc is None: raise exc.add_note(f"runtime cleanup failed: {cleanup_error}") + finally: + if self._stream_listener is not None: + await self._stream_listener.close() def _json_mapping(value: Mapping[str, Any] | None, name: str) -> dict[str, Any]: diff --git a/python/src/nemo_fabric/streaming.py b/python/src/nemo_fabric/streaming.py new file mode 100644 index 000000000..d5f0f2a83 --- /dev/null +++ b/python/src/nemo_fabric/streaming.py @@ -0,0 +1,662 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""NVIDIA NeMo Relay streaming support for the NVIDIA NeMo Fabric Python SDK.""" + +from __future__ import annotations + +import asyncio +import json +import os +import warnings +from collections.abc import Coroutine +from contextlib import suppress +from typing import Any + +from nemo_fabric.models import ( + FabricConfig, + RelayAtofConfig, + RelayAtofFileSinkConfig, + RelayAtofStreamSinkConfig, + RelayConfig, + RelayObservabilityConfig, +) +from nemo_fabric.types import RunResult + +_DRAIN_SECONDS = 0.25 +_MAX_RECORD_BYTES = 1024 * 1024 +_QUEUE_MAX_BYTES = 16 * 1024 * 1024 +_QUEUE_MAXSIZE = 1024 +_READ_SIZE = 64 * 1024 +_STREAMING_HOST_ENV = "NEMO_FABRIC_STREAMING_HOST" +_STREAM_SINK_NAME = "nemo-fabric-stream" + + +class _RecordTooLarge(ValueError): + pass + + +class _AtofRecordQueue: + def __init__(self, *, maxsize: int, max_bytes: int) -> None: + self._queue: asyncio.Queue[tuple[dict[str, Any], int]] = asyncio.Queue( + maxsize=maxsize + ) + self._max_bytes = max_bytes + self._queued_bytes = 0 + self._space_available = asyncio.Event() + self._space_available.set() + + def empty(self) -> bool: + return self._queue.empty() + + def full(self) -> bool: + return self._queue.full() or self._queued_bytes >= self._max_bytes + + async def put( + self, + record: dict[str, Any], + *, + byte_size: int | None = None, + ) -> None: + size = byte_size if byte_size is not None else _record_size(record) + if size > self._max_bytes: + raise _RecordTooLarge + while self._queued_bytes + size > self._max_bytes: + self._space_available.clear() + await self._space_available.wait() + self._queued_bytes += size + try: + await self._queue.put((record, size)) + except BaseException: + self._queued_bytes -= size + self._space_available.set() + raise + + def put_nowait( + self, + record: dict[str, Any], + *, + byte_size: int | None = None, + ) -> None: + size = byte_size if byte_size is not None else _record_size(record) + if size > self._max_bytes: + raise _RecordTooLarge + if self._queued_bytes + size > self._max_bytes: + raise asyncio.QueueFull + self._queue.put_nowait((record, size)) + self._queued_bytes += size + + async def get(self) -> dict[str, Any]: + record, size = await self._queue.get() + self._release(size) + return record + + def get_nowait(self) -> dict[str, Any]: + record, size = self._queue.get_nowait() + self._release(size) + return record + + def _release(self, size: int) -> None: + self._queued_bytes -= size + self._space_available.set() + + +class InvokeStream: + """Async iterator of raw ATOF records for one runtime invocation. + + Consume the final normalized result separately with :meth:`result`. If + iteration stops early, call :meth:`aclose` before starting another turn. + """ + + def __init__( + self, + invoke: Coroutine[Any, Any, RunResult], + listener: _AtofStreamListener, + *, + request_id: str | None = None, + turn_index: int | None = None, + ) -> None: + """lazydocs: ignore""" + + self._listener = listener + self._closed = False + self._finalized = False + self._pending_record: dict[str, Any] | None = None + listener.begin_stream(request_id=request_id, turn_index=turn_index) + try: + self._task = asyncio.create_task(invoke) + except BaseException: + listener.end_stream() + invoke.close() + raise + + def __aiter__(self) -> InvokeStream: + """Return this stream as its asynchronous iterator.""" + + return self + + async def __anext__(self) -> dict[str, Any]: + """Return the next raw ATOF record.""" + + queue = self._listener.records + while True: + if self._closed: + await self._finalize() + raise StopAsyncIteration + if self._pending_record is not None: + record = self._pending_record + self._pending_record = None + return record + if not queue.empty(): + return queue.get_nowait() + if self._task.done(): + try: + return await asyncio.wait_for(queue.get(), _DRAIN_SECONDS) + except TimeoutError: + await self._finalize(warn_if_unavailable=True) + raise StopAsyncIteration from None + + getter = asyncio.create_task(queue.get()) + try: + await asyncio.wait( + {getter, self._task}, + return_when=asyncio.FIRST_COMPLETED, + ) + except asyncio.CancelledError: + if not getter.done(): + getter.cancel() + try: + self._pending_record = await getter + except asyncio.CancelledError: + pass + raise + if getter.done() and not getter.cancelled(): + return getter.result() + getter.cancel() + with suppress(asyncio.CancelledError): + await getter + + async def result(self) -> RunResult: + """Return the terminal normalized result without adding it to the stream.""" + + return await asyncio.shield(self._task) + + async def aclose(self) -> None: + """Stop iteration and drain this turn without cancelling the invocation.""" + + self._closed = True + await self._finalize() + + async def _finalize(self, *, warn_if_unavailable: bool = False) -> None: + if self._finalized: + return + queue = self._listener.records + while not self._task.done(): + getter = asyncio.create_task(queue.get()) + try: + await asyncio.wait( + {getter, self._task}, + return_when=asyncio.FIRST_COMPLETED, + ) + finally: + if not getter.done(): + getter.cancel() + with suppress(asyncio.CancelledError): + await getter + + invocation_completed = False + try: + await asyncio.shield(self._task) + invocation_completed = True + except asyncio.CancelledError: + if not self._task.cancelled(): + raise + except Exception: + pass + + loop = asyncio.get_running_loop() + deadline = loop.time() + _DRAIN_SECONDS + while True: + while not queue.empty(): + queue.get_nowait() + remaining = deadline - loop.time() + if remaining <= 0: + break + try: + await asyncio.wait_for(queue.get(), remaining) + except TimeoutError: + break + self._pending_record = None + self._listener.end_stream() + self._finalized = True + if invocation_completed and warn_if_unavailable: + self._listener.warn_if_unavailable() + + +class _AtofStreamListener: + """Receive chunked NDJSON ATOF records on an SDK-owned endpoint.""" + + def __init__( + self, + *, + host: str | None = None, + port: int = 0, + maxsize: int = _QUEUE_MAXSIZE, + max_bytes: int = _QUEUE_MAX_BYTES, + max_record_bytes: int = _MAX_RECORD_BYTES, + ) -> None: + self._host = host or os.environ.get(_STREAMING_HOST_ENV) or "127.0.0.1" + self._port = port + self._queue = _AtofRecordQueue(maxsize=maxsize, max_bytes=max_bytes) + self._max_record_bytes = min(max_record_bytes, max_bytes) + self._server: asyncio.Server | None = None + self._bound_port: int | None = None + self._accepting = False + self._request_id: str | None = None + self._turn_index: int | None = None + self._turn_root_uuid: str | None = None + self._turn_scope_uuids: set[str] = set() + self._saw_atof_data = False + self._matched_turn_root = False + self._active_atof_connections = 0 + self._saw_atof_connection = False + self._lost_atof_connection = False + self._warned_unconnected = False + self._warned_uncorrelated = False + self._warned_interrupted = False + self._tasks: set[asyncio.Task[None]] = set() + self._writers: set[asyncio.StreamWriter] = set() + + @property + def url(self) -> str: + """Return the listener endpoint after startup.""" + + if self._bound_port is None: + raise RuntimeError("ATOF stream listener is not started") + return f"http://{self._host}:{self._bound_port}/atof" + + @property + def records(self) -> _AtofRecordQueue: + """Return the bounded record queue used by the active stream.""" + + return self._queue + + async def start(self) -> _AtofStreamListener: + """Bind the HTTP server.""" + + self._server = await asyncio.start_server( + self._connected, + self._host, + self._port, + ) + socket = self._server.sockets[0] + self._bound_port = int(socket.getsockname()[1]) + return self + + def begin_stream( + self, + *, + request_id: str | None = None, + turn_index: int | None = None, + ) -> None: + """Route subsequent records to the active invocation queue.""" + + if self._accepting: + raise RuntimeError("ATOF stream listener already has an active consumer") + while not self._queue.empty(): + self._queue.get_nowait() + self._request_id = request_id + self._turn_index = turn_index + self._turn_root_uuid = None + self._turn_scope_uuids.clear() + self._saw_atof_data = False + self._matched_turn_root = request_id is None and turn_index is None + self._saw_atof_connection = self._active_atof_connections > 0 + self._lost_atof_connection = False + self._accepting = True + + def end_stream(self) -> None: + """Discard records until another streaming invocation begins.""" + + self._accepting = False + self._request_id = None + self._turn_index = None + self._turn_root_uuid = None + self._turn_scope_uuids.clear() + + def warn_if_unavailable(self) -> None: + """Warn once when Relay is unreachable or turn correlation fails.""" + + if not self._saw_atof_connection or ( + self._lost_atof_connection + and not self._saw_atof_data + and self._active_atof_connections == 0 + ): + if self._warned_unconnected: + return + self._warned_unconnected = True + warnings.warn( + "No Relay ATOF connection reached the SDK listener. " + "Relay-backed streaming yielded no records. Claude and Codex " + f"gateways must be able to reach {self._host}.", + RuntimeWarning, + stacklevel=3, + ) + return + if ( + self._saw_atof_data + and not self._matched_turn_root + and not self._warned_uncorrelated + ): + self._warned_uncorrelated = True + warnings.warn( + "Relay ATOF data reached the SDK listener, but no record matched " + "the active Fabric turn. Relay-backed streaming yielded no " + "records. Verify the Relay turn correlation metadata and record " + "size limits.", + RuntimeWarning, + stacklevel=3, + ) + return + if ( + self._lost_atof_connection + and self._matched_turn_root + and self._active_atof_connections == 0 + and not self._warned_interrupted + ): + self._warned_interrupted = True + warnings.warn( + "The Relay ATOF connection closed during the active Fabric " + "turn. Relay-backed streaming may be incomplete.", + RuntimeWarning, + stacklevel=3, + ) + + def _connected( + self, + reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, + ) -> None: + task = asyncio.create_task(self._handle_client(reader, writer)) + self._tasks.add(task) + task.add_done_callback(self._task_done) + + def _task_done(self, task: asyncio.Task[None]) -> None: + self._tasks.discard(task) + if not task.cancelled(): + task.exception() + + async def _handle_client( + self, + reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, + ) -> None: + self._writers.add(writer) + is_atof_connection = False + is_chunked = False + chunked_body_completed = False + try: + request = await reader.readuntil(b"\r\n\r\n") + request_line, *header_lines = request[:-4].split(b"\r\n") + method, target, _ = request_line.decode("ascii").split(" ", 2) + headers = _http_headers(header_lines) + if method != "POST" or target.split("?", 1)[0] != "/atof": + await _write_response(writer, 404, "Not Found") + return + is_atof_connection = True + self._active_atof_connections += 1 + if self._accepting: + self._saw_atof_connection = True + if headers.get("expect", "").lower() == "100-continue": + writer.write(b"HTTP/1.1 100 Continue\r\n\r\n") + await writer.drain() + + buffer = bytearray() + is_chunked = "chunked" in headers.get("transfer-encoding", "").lower() + if is_chunked: + await self._read_chunked(reader, buffer) + chunked_body_completed = True + elif "content-length" in headers: + await self._read_sized( + reader, + buffer, + int(headers["content-length"]), + ) + else: + await _write_response(writer, 411, "Length Required") + return + await self._emit(buffer) + await _write_response(writer, 200, "OK") + except _RecordTooLarge: + with suppress(ConnectionError): + await _write_response(writer, 413, "Content Too Large") + except (ValueError, UnicodeDecodeError, asyncio.IncompleteReadError): + with suppress(ConnectionError): + await _write_response(writer, 400, "Bad Request") + except ConnectionError: + pass + except asyncio.CancelledError: + raise + finally: + if is_atof_connection: + self._active_atof_connections -= 1 + if is_chunked and not chunked_body_completed and self._accepting: + self._lost_atof_connection = True + self._writers.discard(writer) + writer.close() + with suppress(ConnectionError): + await writer.wait_closed() + + async def _read_chunked( + self, + reader: asyncio.StreamReader, + buffer: bytearray, + ) -> None: + while True: + size_line = await reader.readline() + size = int(size_line.split(b";", 1)[0].strip(), 16) + if size == 0: + while True: + trailer = await reader.readline() + if trailer in (b"\r\n", b"\n"): + return + if trailer == b"": + raise ValueError("incomplete HTTP chunk trailers") + remaining = size + while remaining: + chunk = await reader.readexactly(min(_READ_SIZE, remaining)) + remaining -= len(chunk) + await self._feed(buffer, chunk) + if await reader.readexactly(2) != b"\r\n": + raise ValueError("invalid HTTP chunk terminator") + + async def _read_sized( + self, + reader: asyncio.StreamReader, + buffer: bytearray, + size: int, + ) -> None: + if size < 0: + raise ValueError("negative HTTP content length") + remaining = size + while remaining: + chunk = await reader.readexactly(min(_READ_SIZE, remaining)) + remaining -= len(chunk) + await self._feed(buffer, chunk) + + async def _feed(self, buffer: bytearray, chunk: bytes) -> None: + if chunk and self._accepting: + self._saw_atof_data = True + buffer.extend(chunk) + while True: + newline = buffer.find(b"\n") + if newline < 0: + if len(buffer) > self._max_record_bytes: + raise _RecordTooLarge + return + if newline > self._max_record_bytes: + raise _RecordTooLarge + line = bytes(buffer[:newline]) + del buffer[: newline + 1] + await self._emit(line) + + async def _emit(self, line: bytes | bytearray) -> None: + stripped = bytes(line).strip() + if not stripped or not self._accepting: + return + if len(stripped) > self._max_record_bytes: + raise _RecordTooLarge + try: + record = json.loads(stripped) + except json.JSONDecodeError: + return + if isinstance(record, dict): + if self._belongs_to_active_turn(record): + await self._queue.put(record, byte_size=len(stripped)) + + def _belongs_to_active_turn(self, record: dict[str, Any]) -> bool: + if self._request_id is None and self._turn_index is None: + return True + + uuid = record.get("uuid") + if not isinstance(uuid, str): + return False + if self._turn_root_uuid is None: + if not self._matches_turn_root(record): + return False + self._turn_root_uuid = uuid + self._turn_scope_uuids.add(uuid) + self._matched_turn_root = True + return True + + if uuid in self._turn_scope_uuids: + return True + parent_uuid = record.get("parent_uuid") + if ( + not isinstance(parent_uuid, str) + or parent_uuid not in self._turn_scope_uuids + ): + return False + if record.get("kind") == "scope" and record.get("scope_category") == "start": + self._turn_scope_uuids.add(uuid) + return True + + def _matches_turn_root(self, record: dict[str, Any]) -> bool: + if record.get("kind") != "scope" or record.get("scope_category") != "start": + return False + metadata = record.get("metadata") + if not isinstance(metadata, dict): + return False + if ( + self._request_id is not None + and metadata.get("nemo_fabric_request_id") == self._request_id + ): + return True + return ( + self._turn_index is not None + and metadata.get("nemo_relay_scope_role") == "turn" + and metadata.get("turn_index") == self._turn_index + ) + + async def close(self) -> None: + """Stop accepting records and close active HTTP connections.""" + + self._accepting = False + if self._server is not None: + self._server.close() + await self._server.wait_closed() + self._server = None + for writer in tuple(self._writers): + writer.close() + for task in tuple(self._tasks): + task.cancel() + if self._tasks: + await asyncio.gather(*self._tasks, return_exceptions=True) + self._writers.clear() + self._bound_port = None + + +def _relay_enabled(config: FabricConfig) -> bool: + telemetry = config.telemetry + return telemetry is not None and "relay" in telemetry.providers + + +def _record_size(record: dict[str, Any]) -> int: + return len(json.dumps(record, separators=(",", ":"), ensure_ascii=False).encode()) + + +def _sink_name( + sink: RelayAtofFileSinkConfig | RelayAtofStreamSinkConfig | dict[str, Any], +) -> str | None: + name = sink.get("name") if isinstance(sink, dict) else getattr(sink, "name", None) + return name if isinstance(name, str) else None + + +def _with_stream_sink(config: FabricConfig, url: str) -> FabricConfig: + copied = config.model_copy(deep=True) + if copied.relay is None: + relay = RelayConfig() + elif isinstance(copied.relay, RelayConfig): + relay = copied.relay + else: + relay = RelayConfig.model_validate(copied.relay) + + if relay.observability is None: + observability = RelayObservabilityConfig() + elif isinstance(relay.observability, RelayObservabilityConfig): + observability = relay.observability + else: + observability = RelayObservabilityConfig.model_validate(relay.observability) + + if observability.atof is None: + atof = RelayAtofConfig() + elif isinstance(observability.atof, RelayAtofConfig): + atof = observability.atof + else: + atof = RelayAtofConfig.model_validate(observability.atof) + + if atof.enabled: + sinks = [ + sink for sink in atof.sinks or () if _sink_name(sink) != _STREAM_SINK_NAME + ] + else: + atof = RelayAtofConfig(enabled=True) + sinks = [] + sinks.append( + RelayAtofStreamSinkConfig( + name=_STREAM_SINK_NAME, + url=url, + transport="ndjson", + ) + ) + atof.sinks = sinks + observability.atof = atof + relay.observability = observability + copied.relay = relay + return copied + + +def _http_headers(lines: list[bytes]) -> dict[str, str]: + headers: dict[str, str] = {} + for line in lines: + name, separator, value = line.partition(b":") + if not separator: + raise ValueError("invalid HTTP header") + headers[name.decode("ascii").strip().lower()] = value.decode("ascii").strip() + return headers + + +async def _write_response( + writer: asyncio.StreamWriter, + status: int, + reason: str, +) -> None: + body = reason.encode("ascii") + writer.write( + f"HTTP/1.1 {status} {reason}\r\n".encode("ascii") + + f"Content-Length: {len(body)}\r\n".encode("ascii") + + b"Connection: close\r\n" + + b"Content-Type: text/plain\r\n\r\n" + + body + ) + await writer.drain() diff --git a/python/src/nemo_fabric/types.py b/python/src/nemo_fabric/types.py index f9a4cb1e0..6e329a2e4 100644 --- a/python/src/nemo_fabric/types.py +++ b/python/src/nemo_fabric/types.py @@ -1083,7 +1083,7 @@ class TelemetryRef(FabricMapping): """Reference to external or persisted telemetry for a run. Attributes: - provider: Telemetry provider, such as Relay. + provider: Telemetry provider, such as NVIDIA NeMo Relay. kind: Reference kind, such as ``trace``. uri: Optional location of persisted telemetry. trace_id: Optional provider trace identifier. diff --git a/scripts/docs/sync_fern_docs_branch.py b/scripts/docs/sync_fern_docs_branch.py index 409feb4b2..29362f827 100644 --- a/scripts/docs/sync_fern_docs_branch.py +++ b/scripts/docs/sync_fern_docs_branch.py @@ -42,6 +42,14 @@ "# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n" "# SPDX-License-Identifier: Apache-2.0\n\n" ) +MARKDOWN_SPDX_HEADER = ( + "" +) +MDX_SPDX_HEADER = ( + "{/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n" + "SPDX-License-Identifier: Apache-2.0 */}" +) def read_yaml(path: Path) -> Any: @@ -72,6 +80,17 @@ def docs_ignore(_directory: str, names: list[str]) -> set[str]: return {name for name in names if name in COPY_EXCLUDES} +def prepare_generated_python_references_for_fern(pages_directory: Path) -> None: + references = pages_directory / "reference" / "api" / "python-library-reference" + if not references.is_dir(): + return + for path in references.glob("*.md"): + text = path.read_text(encoding="utf-8") + updated = text.replace(MARKDOWN_SPDX_HEADER, MDX_SPDX_HEADER, 1) + if updated != text: + path.write_text(updated, encoding="utf-8") + + def prefixed_doc_path(value: str, pages_directory: str) -> str: if value.startswith(("http://", "https://", "/", "../")): return value @@ -159,6 +178,7 @@ def sync_dev(source_root: Path, target_root: Path) -> None: if pages_dev.exists(): shutil.rmtree(pages_dev) shutil.copytree(source_docs, pages_dev, ignore=docs_ignore) + prepare_generated_python_references_for_fern(pages_dev) navigation = read_yaml(source_docs / "index.yml") write_yaml( @@ -214,6 +234,7 @@ def release_version(target_root: Path, tag: str, source_root: Path) -> None: version_yml.unlink() shutil.copytree(source_docs, pages_version, ignore=docs_ignore) + prepare_generated_python_references_for_fern(pages_version) update_github_links(pages_version, tag) version_navigation = rewrite_doc_references( read_yaml(source_docs / "index.yml"), f"pages-{display_tag}" diff --git a/scripts/generate_api_docs.sh b/scripts/generate_api_docs.sh index b8628f68f..51add4669 100755 --- a/scripts/generate_api_docs.sh +++ b/scripts/generate_api_docs.sh @@ -23,13 +23,14 @@ PYTHONPATH="python/src" lazydocs \ --overview-file "index.md" \ "nemo_fabric.client" \ "nemo_fabric.runtime" \ + "nemo_fabric.streaming" \ "nemo_fabric.models" \ "nemo_fabric.types" \ "nemo_fabric.errors" -# Make the lazydocs output MDX-safe for Fern (Fern parses .md as MDX): +# Normalize the lazydocs output for Fern: # - drop source badges (relative links don't resolve on the site) -# - strip HTML comments (), which are invalid in MDX +# - strip lazydocs HTML comments before adding the generated SPDX header # - remove trailing whitespace emitted by lazydocs perl -ni -e 'print unless m{img\.shields\.io/badge/-source}' "$out"/*.md perl -0pi -e 's///gs' "$out"/*.md @@ -39,8 +40,12 @@ perl -0pi -e 's/\A\s+//' "$out"/*.md # lazydocs nests properties at h4 directly under h2 class sections. Flatten # those headings to h3 so generated pages satisfy markdown heading order. perl -pi -e 's/^#### (property<\/kbd>)/### $1/' "$out"/*.md -# lazydocs emits the ToolsConfig class heading without a separating blank line. -perl -0pi -e 's/(^## class<\/kbd> `ToolsConfig`\n)(?!\n)/$1\n/m' "$out/nemo_fabric.models.md" +# lazydocs emits some class headings without a separating blank line. +perl -0pi -e 's/(^## class<\/kbd> `(ToolsConfig|RelayAtofFileSinkConfig|RelayAtofStreamSinkConfig)`\n)(?!\n)/$1\n/gm' \ + "$out/nemo_fabric.models.md" +# lazydocs omits the async marker from generated method signatures. +perl -0pi -e 's/(### method<\/kbd> `(aclose|result)`\n\n```python\n)\2\(/${1}async def ${2}(/g' \ + "$out/nemo_fabric.streaming.md" add_frontmatter() { local file="$1" @@ -52,8 +57,8 @@ add_frontmatter() { { printf -- '---\ntitle: "%s"\nslug: "%s"\ndescription: "%s"\n---\n' \ "$title" "$slug" "$description" - printf '%s\n' '{/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.' - printf '%s\n\n' 'SPDX-License-Identifier: Apache-2.0 */}' + printf '%s\n' '' command cat "$file" } > "$temporary" mv "$temporary" "$file" @@ -74,6 +79,11 @@ add_frontmatter \ "Runtime" \ "Drive stateful multi-turn execution through the Runtime API." \ "/reference/api/python-library-reference/runtime" +add_frontmatter \ + "$out/nemo_fabric.streaming.md" \ + "Streaming" \ + "Consume raw NVIDIA NeMo Relay ATOF records and terminal invocation results." \ + "/reference/api/python-library-reference/streaming" add_frontmatter \ "$out/nemo_fabric.models.md" \ "Models" \ @@ -90,6 +100,12 @@ add_frontmatter \ "Structured exception hierarchy for config, capability, state, and runtime failures." \ "/reference/api/python-library-reference/errors" +# Use the full product name on first mention in each generated page and the +# shortened product name on subsequent mentions. +for file in "$out"/*.md; do + perl -0pi -e 'my $seen = 0; s/NVIDIA NeMo Relay/++$seen == 1 ? $& : "NeMo Relay"/ge' "$file" +done + # Drop the mkdocs-specific .pages file lazydocs emits; Fern does not use it. rm -f "$out"/.pages diff --git a/skills/integrations/consumer/nemo-fabric-integrate/SKILL.md b/skills/integrations/consumer/nemo-fabric-integrate/SKILL.md index ea7d17977..5b9195a4d 100644 --- a/skills/integrations/consumer/nemo-fabric-integrate/SKILL.md +++ b/skills/integrations/consumer/nemo-fabric-integrate/SKILL.md @@ -121,6 +121,34 @@ Pick the smallest lifecycle the consumer needs: (`stop()` can raise `FabricRuntimeError`; see Consume Results And Handle Errors). A runtime accepts one active invocation at a time; overlapping calls raise `FabricStateError`. +- **NVIDIA NeMo Relay stream** — live, raw ATOF records plus a terminal normalized + result. Enable NeMo Relay, pass `streaming=True` to `start_runtime(...)`, call + `runtime.invoke_stream(...)`, iterate the returned `InvokeStream`, and then + await `stream.result()`. Iteration ending does not indicate invocation + success; invocation exceptions raise from `result()`, while harness-reported + failures remain normalized `RunResult` values. If iteration stops early, + call `await stream.aclose()` before starting another turn. `aclose()` waits + for the turn to finish; it does not cancel the harness invocation. The SDK + intentionally exposes ATOF records generated by NeMo Relay only; adapter-native progressive + output is deferred to a future normalized Fabric contract. The listener + limits each record to 1 MiB and its queue to 1,024 records or 16 MiB of + encoded data. It correlates records through the Fabric request ID for + in-process harnesses or NeMo Relay's turn-scope role and 1-based turn index for + gateway harnesses, then yields only the matched scope tree. Delayed + prior-turn records therefore do not enter the next stream. If gateway and + Fabric turn sequences do not align, the SDK discards the uncorrelated records + and emits a `RuntimeWarning` after natural stream exhaustion. The listener + binds to `NEMO_FABRIC_STREAMING_HOST`, which defaults to `127.0.0.1`. + Override it when the gateway must reach the SDK through another network + interface, and restrict access to that interface. If async iteration reaches + its post-turn drain timeout without a NeMo Relay connection, or receives data + without a matching turn root, the SDK emits one `RuntimeWarning` for that + failure mode; callers that only await `stream.result()` do not run that + warning check. The SDK also warns when a NeMo Relay upload terminates before + completing its chunked request body because yielded records can be incomplete. + The `streaming=True` flag does not enable NeMo Relay by itself. Without + `streaming=True`, startup leaves the NeMo Relay configuration unchanged and + does not inject the SDK-owned ATOF stream sink. The selected adapter owns the execution topology. The bundled Claude, Codex, Deep Agents, and Hermes Agent adapters retain their native client, graph/checkpointer, @@ -151,12 +179,29 @@ async def main() -> None: first = await runtime.invoke(input="Inspect the repository") second = await runtime.invoke(input="Now review the latest patch") + # NeMo Relay streaming + streaming_config = config.model_copy(deep=True).enable_relay() + async with await fabric.start_runtime( + streaming_config, + base_dir=base, + streaming=True, + ) as runtime: + stream = runtime.invoke_stream(input="Review the latest patch") + async for record in stream: + print(record) + streamed_result = await stream.result() + asyncio.run(main()) ``` -NeMo Fabric owns no queue, worker pool, retry policy, or concurrency limit. For -parallel work, start independent runtimes and let the consumer decide how many. +NeMo Fabric owns no application scheduling queue, worker pool, retry policy, or +global concurrency policy. Each runtime still permits only one active +invocation; start independent runtimes for parallel work. The NeMo Relay +streaming path uses an internal bounded transport queue and TCP backpressure +only to carry one invocation's ATOF records. Treat `stream.result()` as +authoritative, and reconstruct nested work from ATOF `uuid` and `parent_uuid` +fields rather than stream order. ## Validate Before Running @@ -240,7 +285,7 @@ result-field and error inventory, and - [ ] The consumer config object is translated directly into an in-memory `FabricConfig`. - [ ] Only public `nemo_fabric` symbols are imported; no `_native` or adapter internals. - [ ] The consumer config is built in memory and passed directly to NeMo Fabric. -- [ ] The right lifecycle is chosen: `run(...)` for a single invocation, `start_runtime(...)` with `async with` for multi-turn. +- [ ] The right lifecycle is chosen: `run(...)` for a single invocation, `start_runtime(...)` with `async with` for multi-turn, or `invoke_stream(...)` for raw NeMo Relay ATOF. - [ ] `plan(...)` and `doctor(...)` validate adapter selection, capabilities, and environment before execution. - [ ] Installation, adapter dependencies, and credentials are owned by the environment, not consumer code. - [ ] `RunResult` status, error, and events are inspected before output; artifacts and telemetry are captured. @@ -259,6 +304,7 @@ Link to these canonical sources instead of duplicating them: stubs are authoritative for exact signatures, fields, and defaults): [client](https://github.com/NVIDIA/NeMo-Fabric/blob/main/docs/reference/api/python-library-reference/nemo_fabric.client.md), [runtime](https://github.com/NVIDIA/NeMo-Fabric/blob/main/docs/reference/api/python-library-reference/nemo_fabric.runtime.md), + [streaming](https://github.com/NVIDIA/NeMo-Fabric/blob/main/docs/reference/api/python-library-reference/nemo_fabric.streaming.md), [models](https://github.com/NVIDIA/NeMo-Fabric/blob/main/docs/reference/api/python-library-reference/nemo_fabric.models.md), [types](https://github.com/NVIDIA/NeMo-Fabric/blob/main/docs/reference/api/python-library-reference/nemo_fabric.types.md), [errors](https://github.com/NVIDIA/NeMo-Fabric/blob/main/docs/reference/api/python-library-reference/nemo_fabric.errors.md) diff --git a/skills/integrations/consumer/nemo-fabric-integrate/references/config-mapping.md b/skills/integrations/consumer/nemo-fabric-integrate/references/config-mapping.md index 9a93b6784..6d163ea64 100644 --- a/skills/integrations/consumer/nemo-fabric-integrate/references/config-mapping.md +++ b/skills/integrations/consumer/nemo-fabric-integrate/references/config-mapping.md @@ -24,7 +24,7 @@ Import these from the top-level `nemo_fabric` package: | `McpConfig` / `McpServerConfig` | MCP servers and exposure. | | `SkillConfig` | Skill directories. | | `TelemetryConfig` | Telemetry providers. | -| `RelayConfig` and `Relay*Config` | Relay observability under the top-level `relay` block. | +| `RelayConfig` and `Relay*Config` | NVIDIA NeMo Relay observability under the top-level `relay` block. | The [models reference](https://github.com/NVIDIA/NeMo-Fabric/blob/main/docs/reference/api/python-library-reference/nemo_fabric.models.md) indexes the public config models. The generated pages omit constructor fields and @@ -38,7 +38,7 @@ methods that edit the typed config in place and return it: - `add_skill_path(path)` / `remove_skill_path(path)` - `add_mcp_server(name, *, transport, url, exposure, ...)` / `remove_mcp_server(name)` -- `enable_relay(...)` for Relay observability in the `relay` block +- `enable_relay(...)` for NeMo Relay observability in the `relay` block ```python config = FabricConfig( @@ -65,7 +65,7 @@ def with_relay(base: FabricConfig) -> FabricConfig: Use this function-and-copy pattern for every variant; keep all variation in ordinary Python. -For ATOF, author the Relay 0.6 sink model directly. Put +For ATOF, author the NeMo Relay 0.6 sink model directly. Put `RelayAtofFileSinkConfig` and `RelayAtofStreamSinkConfig` instances in `RelayAtofConfig.sinks`, and set `RelayAtofConfig.enabled=True`. diff --git a/skills/integrations/consumer/nemo-fabric-integrate/references/results-and-errors.md b/skills/integrations/consumer/nemo-fabric-integrate/references/results-and-errors.md index d8a00abcf..9743f53b7 100644 --- a/skills/integrations/consumer/nemo-fabric-integrate/references/results-and-errors.md +++ b/skills/integrations/consumer/nemo-fabric-integrate/references/results-and-errors.md @@ -17,7 +17,7 @@ Every invocation that reaches the adapter boundary returns a normalized | `error` | Structured `ErrorInfo`, or `None` — may be `None` even when `status` is not `succeeded`, so do not use it as the success signal. | | `output` | Harness output normalized to the configured output schema. | | `artifacts` | Output files, logs, patches, and other materialized references. | -| `telemetry` | References to Relay or other telemetry streams from the run. | +| `telemetry` | References to NVIDIA NeMo Relay or other telemetry streams from the run. | | `events` | Ordered normalized lifecycle and invocation events. | | `metadata` | Result-specific structured metadata. | | `runtime_id`, `invocation_id`, `request_id` | Correlation IDs across runtimes, logs, telemetry, and artifacts. | diff --git a/skills/integrations/consumer/nemo-fabric-integrate/references/sdk-api-inventory.md b/skills/integrations/consumer/nemo-fabric-integrate/references/sdk-api-inventory.md index 284e3799c..f75d94f91 100644 --- a/skills/integrations/consumer/nemo-fabric-integrate/references/sdk-api-inventory.md +++ b/skills/integrations/consumer/nemo-fabric-integrate/references/sdk-api-inventory.md @@ -10,6 +10,7 @@ lifecycle context manager — and can plan, diagnose, or start multiple independent runtimes. The generated [client reference](https://github.com/NVIDIA/NeMo-Fabric/blob/main/docs/reference/api/python-library-reference/nemo_fabric.client.md) and [runtime reference](https://github.com/NVIDIA/NeMo-Fabric/blob/main/docs/reference/api/python-library-reference/nemo_fabric.runtime.md) +and [streaming reference](https://github.com/NVIDIA/NeMo-Fabric/blob/main/docs/reference/api/python-library-reference/nemo_fabric.streaming.md) document the public methods, but they omit `async` and keyword-only markers — this inventory records those, and the installed `nemo_fabric` package ships type information (`py.typed`) for exact signatures. @@ -23,7 +24,7 @@ The following table lists the `Fabric` methods and when to use each. | `plan(config, *, base_dir=...)` | No | You need the selected adapter, capability routing, and runtime capabilities before running. | `RunPlan` | | `doctor(config, *, base_dir=...)` | Yes | You need preflight diagnostics for adapter availability, config support, and environment assumptions. | `DoctorReport` | | `run(config, *, base_dir=..., input=... \| request=...)` | Yes | You need one complete start, invoke, result, and stop cycle. | `RunResult` | -| `start_runtime(config, *, base_dir=..., overrides=...)` | Yes | You need state across multiple ordered invocations. | `Runtime` | +| `start_runtime(config, *, base_dir=..., overrides=..., streaming=False)` | Yes | You need state across multiple ordered invocations. Pass `streaming=True` with NVIDIA NeMo Relay enabled to provision `invoke_stream(...)`. | `Runtime` | `input` and `request` on `run(...)` are mutually exclusive. Use `input=...` for the common case; use `request=RunRequest(...)` when the invocation needs a @@ -36,8 +37,10 @@ The following table lists the `Runtime` members for driving a stateful runtime. | Member | Async | Notes | | --- | --- | --- | | `invoke(*, input=... \| request=...)` | Yes | One turn on an active runtime. One active invocation at a time; overlap raises `FabricStateError`. | +| `invoke_stream(*, input=... \| request=...)` | No | Start one NeMo Relay turn and return an async `InvokeStream` of raw ATOF records. Await `stream.result()` for the terminal `RunResult`. | | `stop()` | Yes | Stop the runtime. Called automatically by `async with`. | | `status` | No | `RuntimeStatus`: `ACTIVE`, `STOPPED`, or `FAILED`. | +| `supports_streaming` | No | `True` when NeMo Relay ATOF streaming was enabled at runtime startup. | | `runtime_id` | No | Opaque identifier for this runtime lifecycle. | | `messages` / `invocations` | No | Copied harness history and per-turn IDs. | @@ -57,6 +60,7 @@ invocations: ```text FabricConfig -> plan() -> RunPlan -> start_runtime() -> Runtime -> invoke() -> RunResult + \-> invoke_stream() -> InvokeStream ``` - `Fabric` is a lightweight facade; it holds no started state and needs no diff --git a/tests/adapters/test_adapaters_common_utils.py b/tests/adapters/test_adapaters_common_utils.py index 2a8c498aa..b9ed3ba2a 100644 --- a/tests/adapters/test_adapaters_common_utils.py +++ b/tests/adapters/test_adapaters_common_utils.py @@ -271,6 +271,16 @@ def test_normalize_list(value: object, expected: list[str]): assert common_utils.normalize_list(value) == expected +def test_without_none_preserves_falsey_values(): + assert common_utils.without_none( + {"zero": 0, "false": False, "empty": "", "missing": None} + ) == { + "zero": 0, + "false": False, + "empty": "", + } + + def test_load_relay_plugin_config_wraps_and_normalizes_bare_observability_config( tmp_path: Path, ): diff --git a/tests/adapters/test_deepagents.py b/tests/adapters/test_deepagents.py index d718050b4..15041f190 100644 --- a/tests/adapters/test_deepagents.py +++ b/tests/adapters/test_deepagents.py @@ -165,7 +165,7 @@ def make(tmp_path: Path, *, runtime_id: str = "run-1") -> dict[str, Any]: "invocation_id": "inv-1", "environment": {"workspace": str(tmp_path)}, }, - "request": {"input": "hello"}, + "request": {"input": "hello", "request_id": "request-1"}, "capability_plan": {}, } @@ -201,10 +201,15 @@ class ScopeType: Agent = "agent" @contextlib.contextmanager - def scope_ctx(name: str, scope_type: object, **_: object) -> Iterator[None]: + def scope_ctx( + name: str, + scope_type: object, + **kwargs: object, + ) -> Iterator[None]: # Record every scope entered so tests can assert the top-level # ``deepagents-request`` Agent scope wraps the invocation. calls.setdefault("scopes", []).append((name, scope_type)) + calls.setdefault("scope_metadata", []).append(kwargs.get("metadata")) yield class NemoRelayDeepAgentsCallbackHandler: @@ -369,6 +374,7 @@ async def test_relay_telemetry_wraps_agent_and_reports_artifacts( # the top-level invocation is wrapped in the deepagents-request Agent scope # ("agent" is the fake ScopeType.Agent sentinel from the fake_relay fixture) assert fake_relay["scopes"] == [("deepagents-request", "agent")] + assert fake_relay["scope_metadata"] == [{"nemo_fabric_request_id": "request-1"}] # the Deep Agents callback handler is added to the LangGraph run config so # LangGraph scopes and human-in-the-loop interrupt/resume marks are captured assert fake_relay["callback_handler"] in (fake_sdks["config"] or {}).get( @@ -422,6 +428,7 @@ async def test_native_telemetry_exports_without_artifacts( assert "relay-mw" in fake_sdks["create_kwargs"]["middleware"] # the scope + callback handler apply to any observability-enabled run, native included assert fake_relay["scopes"] == [("deepagents-request", "agent")] + assert fake_relay["scope_metadata"] == [{"nemo_fabric_request_id": "request-1"}] assert fake_relay["callback_handler"] in (fake_sdks["config"] or {}).get( "callbacks", [] ) diff --git a/tests/adapters/test_hermes_config_builder.py b/tests/adapters/test_hermes_config_builder.py new file mode 100644 index 000000000..ece011e46 --- /dev/null +++ b/tests/adapters/test_hermes_config_builder.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Dependency-free tests for Hermes configuration construction.""" + +import sys + +import pytest + +if sys.version_info >= (3, 14): + pytest.skip( + "Hermes adapter requires Python 3.13 or earlier", + allow_module_level=True, + ) + +from nemo_fabric_adapters.hermes import adapter + + +def test_build_hermes_config_omits_unset_values_without_hermes_agent(): + payload = { + "config": { + "harness": {"settings": {}}, + "models": { + "default": { + "provider": "nvidia", + "model": "nvidia/test-model", + } + }, + } + } + + config = adapter.build_hermes_config(payload) + + assert config["model"] == { + "provider": "nvidia", + "default": "nvidia/test-model", + "base_url": "https://integrate.api.nvidia.com/v1", + } + assert config["agent"] == {} diff --git a/tests/adapters/test_hermes_streaming.py b/tests/adapters/test_hermes_streaming.py new file mode 100644 index 000000000..f9e5081fb --- /dev/null +++ b/tests/adapters/test_hermes_streaming.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Dependency-free tests for Hermes Relay streaming integration.""" + +import sys +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace + +import pytest + +if sys.version_info >= (3, 14): + pytest.skip( + "Hermes adapter requires Python 3.13 or earlier", + allow_module_level=True, + ) + +from nemo_fabric_adapters.hermes import adapter + + +@pytest.mark.parametrize( + ("relay_plugin_config", "expected_metadata"), + [ + ( + { + "components": [ + { + "kind": "observability", + "config": { + "atof": { + "enabled": True, + "sinks": [ + { + "type": "stream", + "name": "nemo-fabric-stream", + "url": "http://127.0.0.1:1234/atof", + } + ], + } + }, + } + ] + }, + [{"nemo_fabric_request_id": "request-1"}], + ), + ({"components": []}, []), + ], + ids=["streaming", "non-streaming"], +) +async def test_relay_invocation_scope_carries_fabric_request_id( + monkeypatch, + tmp_path: Path, + relay_plugin_config: dict[str, object], + expected_metadata: list[object], +): + events: list[str] = [] + runtime = adapter.HermesRuntime() + runtime._started = True + runtime._start_payload = {} + runtime._runtime_id = "runtime-1" + runtime._agent = SimpleNamespace( + session_id="runtime-1", + model="test-model", + platform="fabric", + ) + runtime._invoke_hook = lambda *_args, **_kwargs: events.append("finalize") + runtime._relay_plugin_config = relay_plugin_config + runtime._hermes_home = tmp_path + runtime._hermes_config_path = tmp_path / "config.yaml" + runtime._enabled_toolsets = [] + + def invoke_turn(**_kwargs: object): + events.append("turn") + return ( + { + "response": "done", + "completed": True, + "failed": False, + "messages": [], + }, + "", + ) + + monkeypatch.setattr(adapter, "_invoke_hermes_turn", invoke_turn) + monkeypatch.setattr( + adapter.common_utils, + "collect_relay_artifacts", + lambda _config: [], + ) + + from nemo_relay import scope, subscribers + + captured_metadata: list[object] = [] + + @contextmanager + def capture_scope(*_args: object, **kwargs: object): + captured_metadata.append(kwargs["metadata"]) + events.append("scope-enter") + try: + yield + finally: + events.append("scope-exit") + + monkeypatch.setattr(scope, "scope", capture_scope) + monkeypatch.setattr(subscribers, "flush", lambda: None) + + await runtime.invoke( + { + "runtime_context": {"runtime_id": "runtime-1"}, + "request": {"input": "hello", "request_id": "request-1"}, + } + ) + + assert captured_metadata == expected_metadata + if expected_metadata: + assert events == ["scope-enter", "turn", "finalize", "scope-exit"] + else: + assert events == ["turn", "finalize"] diff --git a/tests/docs/test_python_api_docs.py b/tests/docs/test_python_api_docs.py index f8873df3c..1ce52b8b3 100644 --- a/tests/docs/test_python_api_docs.py +++ b/tests/docs/test_python_api_docs.py @@ -21,6 +21,7 @@ MODULE_SLUGS = { "nemo_fabric.client": "/reference/api/python-library-reference/client", "nemo_fabric.runtime": "/reference/api/python-library-reference/runtime", + "nemo_fabric.streaming": "/reference/api/python-library-reference/streaming", "nemo_fabric.models": "/reference/api/python-library-reference/models", "nemo_fabric.types": "/reference/api/python-library-reference/types", "nemo_fabric.errors": "/reference/api/python-library-reference/errors", @@ -88,6 +89,36 @@ def test_generated_reference_uses_valid_heading_order() -> None: assert "#### property" not in text, page +def test_generated_reference_uses_markdown_spdx_comments() -> None: + for page in REFERENCE_DIR.glob("*.md"): + text = page.read_text(encoding="utf-8") + assert "" in text, page + assert "{/* SPDX-FileCopyrightText:" not in text, page + + +def test_generated_reference_uses_full_relay_name_once_per_page() -> None: + for page in REFERENCE_DIR.glob("*.md"): + text = page.read_text(encoding="utf-8") + if "NeMo Relay" in text: + assert text.count("NVIDIA NeMo Relay") == 1, page + + +def test_streaming_reference_hides_constructor_and_preserves_async_methods(): + reference = (REFERENCE_DIR / "nemo_fabric.streaming.md").read_text(encoding="utf-8") + + assert "### method `__init__`" not in reference + assert "async def aclose()" in reference + assert "async def result()" in reference + + +def test_relay_sink_reference_uses_valid_class_heading_spacing(): + reference = (REFERENCE_DIR / "nemo_fabric.models.md").read_text(encoding="utf-8") + + for class_name in ("RelayAtofFileSinkConfig", "RelayAtofStreamSinkConfig"): + assert f"## class `{class_name}`\n\n" in reference + + def test_landing_page_routes_new_users_through_the_product() -> None: landing = LANDING_PAGE.read_text(encoding="utf-8") navigation = NAVIGATION.read_text(encoding="utf-8") @@ -108,6 +139,7 @@ def test_landing_page_routes_new_users_through_the_product() -> None: for destination in ( "/reference/api/python-library-reference/client", "/reference/api/python-library-reference/runtime", + "/reference/api/python-library-reference/streaming", "/reference/api/python-library-reference/types", "/reference/api/python-library-reference/errors", ): diff --git a/tests/e2e/test_claude.py b/tests/e2e/test_claude.py index 986705b47..669848f27 100644 --- a/tests/e2e/test_claude.py +++ b/tests/e2e/test_claude.py @@ -8,6 +8,7 @@ import json import os import sys +import warnings from pathlib import Path import pytest @@ -299,13 +300,27 @@ async def test_live_claude_relay_session(tmp_path): nemo_relay_command=relay_command, ) - async with await Fabric().start_runtime(config, base_dir=tmp_path) as runtime: - first = await runtime.invoke(input="Remember token FABRIC-CLAUDE-RELAY-7") - second = await runtime.invoke( - input="Reply only with the token I asked you to remember" - ) + with warnings.catch_warnings(): + warnings.simplefilter("error", RuntimeWarning) + async with await Fabric().start_runtime( + config, + base_dir=tmp_path, + streaming=True, + ) as runtime: + first_stream = runtime.invoke_stream( + input="Remember token FABRIC-CLAUDE-RELAY-7" + ) + first_records = [record async for record in first_stream] + first = await first_stream.result() + second_stream = runtime.invoke_stream( + input="Reply only with the token I asked you to remember" + ) + second_records = [record async for record in second_stream] + second = await second_stream.result() results = (first.to_mapping(), second.to_mapping()) + assert first_records + assert second_records assert first.status == second.status == "succeeded", results assert first.output["session_id"] == second.output["session_id"], results assert first.metadata["host_pid"] == second.metadata["host_pid"], results diff --git a/tests/e2e/test_codex.py b/tests/e2e/test_codex.py index ac059d4d7..c59c462cb 100644 --- a/tests/e2e/test_codex.py +++ b/tests/e2e/test_codex.py @@ -12,6 +12,7 @@ import os import shutil import uuid +import warnings import pytest from _utils.utils import assert_semantic_relay_artifacts @@ -112,13 +113,25 @@ async def _run_relay(relay_command: str) -> None: assert_semantic_relay_artifacts(result["output"], "FABRIC_CODEX_RELAY_OK") nonce = f"fabric-relay-{uuid.uuid4().hex[:8]}" - async with await client.start_runtime(config, base_dir=BASE_DIR) as runtime: - first = await runtime.invoke(input=f"Remember this value: {nonce}") - second = await runtime.invoke( - input="Reply with only the value I asked you to remember." - ) + with warnings.catch_warnings(): + warnings.simplefilter("error", RuntimeWarning) + async with await client.start_runtime( + config, + base_dir=BASE_DIR, + streaming=True, + ) as runtime: + first_stream = runtime.invoke_stream(input=f"Remember this value: {nonce}") + first_records = [record async for record in first_stream] + first = await first_stream.result() + second_stream = runtime.invoke_stream( + input="Reply with only the value I asked you to remember." + ) + second_records = [record async for record in second_stream] + second = await second_stream.result() results = (first.to_mapping(), second.to_mapping()) + assert first_records + assert second_records assert first["status"] == second["status"] == "succeeded", results assert first["output"]["thread_id"] == second["output"]["thread_id"], results assert nonce in second["output"]["response"], second.to_mapping() diff --git a/tests/e2e/test_deepagents.py b/tests/e2e/test_deepagents.py index 5a90d3d07..f7b957bce 100644 --- a/tests/e2e/test_deepagents.py +++ b/tests/e2e/test_deepagents.py @@ -12,6 +12,7 @@ import importlib.util import os import uuid +import warnings import pytest @@ -72,11 +73,24 @@ async def test_deepagents_persistent_host_with_relay_and_mock_model( artifacts=tmp_path / "artifacts", ) - async with await Fabric().start_runtime(config, base_dir=tmp_path) as runtime: - first = await runtime.invoke(input="first") - second = await runtime.invoke(input="second") + with warnings.catch_warnings(): + warnings.simplefilter("error", RuntimeWarning) + async with await Fabric().start_runtime( + config, + base_dir=tmp_path, + streaming=True, + ) as runtime: + first_stream = runtime.invoke_stream(input="first") + first_records = [record async for record in first_stream] + first = await first_stream.result() + + second_stream = runtime.invoke_stream(input="second") + second_records = [record async for record in second_stream] + second = await second_stream.result() results = (first.to_mapping(), second.to_mapping()) + assert first_records + assert second_records assert first["status"] == second["status"] == "succeeded", results assert first["metadata"]["host_pid"] == second["metadata"]["host_pid"], results assert first["output"]["thread_id"] == second["output"]["thread_id"], results diff --git a/tests/e2e/test_hermes_e2e.py b/tests/e2e/test_hermes_e2e.py index eb63e298e..232705854 100644 --- a/tests/e2e/test_hermes_e2e.py +++ b/tests/e2e/test_hermes_e2e.py @@ -6,6 +6,7 @@ import json import os import sys +import warnings from pathlib import Path from types import ModuleType @@ -30,13 +31,23 @@ async def test_hermes_persistent_host_reuses_native_session( config = hermes_config() config.harness.settings["base_url"] = f"{api_server}/v1" - async with await Fabric().start_runtime( - config, base_dir=code_review_agent_dir - ) as runtime: - first = await runtime.invoke(input="first") - second = await runtime.invoke(input="second") + with warnings.catch_warnings(): + warnings.simplefilter("error", RuntimeWarning) + async with await Fabric().start_runtime( + config, + base_dir=code_review_agent_dir, + streaming=True, + ) as runtime: + first_stream = runtime.invoke_stream(input="first") + first_records = [record async for record in first_stream] + first = await first_stream.result() + second_stream = runtime.invoke_stream(input="second") + second_records = [record async for record in second_stream] + second = await second_stream.result() results = (first.to_mapping(), second.to_mapping()) + assert first_records + assert second_records assert first["status"] == second["status"] == "succeeded", results assert first["metadata"]["adapter_runner"] == "persistent_local_host", results assert first["metadata"]["host_pid"] == second["metadata"]["host_pid"], results diff --git a/tests/python/test_sdk_contract.py b/tests/python/test_sdk_contract.py index e37ee1487..8b8b53000 100644 --- a/tests/python/test_sdk_contract.py +++ b/tests/python/test_sdk_contract.py @@ -299,16 +299,34 @@ def test_fabric_config_authors_first_class_relay_observability(): } -def test_relay_atof_stream_sink_omits_empty_header_maps(): - sink = RelayAtofStreamSinkConfig(url="https://example.test/events") - - assert sink.to_mapping() == { - "type": "stream", - "url": "https://example.test/events", - "transport": "http_post", - "timeout_millis": 3000, - "field_name_policy": "preserve", - } +@pytest.mark.parametrize( + ("headers", "header_env"), + [ + ({}, {}), + ({"authorization": "Bearer test"}, {}), + ({}, {"authorization": "RELAY_AUTHORIZATION"}), + ], +) +def test_relay_atof_stream_sink_header_maps_round_trip( + headers: dict[str, str], + header_env: dict[str, str], +): + config = RelayAtofConfig( + enabled=True, + sinks=[ + RelayAtofStreamSinkConfig( + url="https://example.test/events", + headers=headers, + header_env=header_env, + ) + ], + ) + + mapping = config.to_mapping() + sink = mapping["sinks"][0] + assert sink.get("headers", {}) == headers + assert sink.get("header_env", {}) == header_env + assert RelayAtofConfig.from_mapping(mapping).to_mapping() == mapping def test_fabric_config_enable_relay_preserves_omitted_fields(): diff --git a/tests/python/test_streaming.py b/tests/python/test_streaming.py new file mode 100644 index 000000000..3ebc68633 --- /dev/null +++ b/tests/python/test_streaming.py @@ -0,0 +1,1020 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Behavior tests for Relay-backed SDK streaming.""" + +from __future__ import annotations + +import asyncio +import json +import os +import threading +import warnings +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +from nemo_fabric import ( + Fabric, + FabricCapabilityError, + FabricConfig, + FabricConfigError, + FabricStateError, + HarnessConfig, + InvokeStream, + MetadataConfig, + RelayAtifConfig, + RelayAtofConfig, + RelayAtofFileSinkConfig, + RelayAtofStreamSinkConfig, + RelayObservabilityConfig, + RunRequest, + RunResult, +) +from nemo_fabric import client as client_mod +from nemo_fabric.streaming import _AtofStreamListener, _with_stream_sink + + +def _config(*, relay: bool = False) -> FabricConfig: + config = FabricConfig( + metadata=MetadataConfig(name="demo"), + harness=HarnessConfig(adapter_id="test.fabric.shim"), + ) + if relay: + config.enable_relay( + observability=RelayObservabilityConfig( + atof=RelayAtofConfig( + enabled=True, + sinks=[ + RelayAtofStreamSinkConfig( + name="user-stream", + url="https://example.com/events", + ) + ], + ) + ) + ) + return config + + +def _plan(config: dict[str, Any]) -> dict[str, Any]: + return { + "agent_name": "demo", + "base_dir": ".", + "config": config, + "adapter_descriptor": { + "descriptor": { + "adapter_id": "test.fabric.shim", + "harness": "hermes", + "adapter_kind": "python", + } + }, + "capabilities": { + "service": False, + "streaming": False, + "updates": False, + "cancellation": False, + }, + } + + +def _runtime() -> dict[str, Any]: + return { + "runtime_id": "runtime-1", + "runtime_binding": "fabric-runtime-binding-test", + "agent_name": "demo", + "harness": "hermes", + "adapter_kind": "python", + "adapter_id": "test.fabric.shim", + "environment": { + "environment_id": "environment-1", + "provider": "local", + "control_location": "external_control", + "ownership": "caller_owned", + }, + } + + +def _result(request: dict[str, Any], runtime: dict[str, Any]) -> dict[str, Any]: + return { + "agent_name": "demo", + "harness": "hermes", + "adapter_kind": "python", + "adapter_id": "test.fabric.shim", + "runtime_id": runtime["runtime_id"], + "invocation_id": f"invocation-{request['request_id']}", + "request_id": request["request_id"], + "status": "succeeded", + "output": {"response": "done"}, + "artifacts": {"artifacts": []}, + "events": [], + } + + +@pytest.fixture(name="mock_native") +def mock_native_fixture() -> MagicMock: + mock_native = MagicMock() + mock_native.plan_config.side_effect = lambda config_json, base_dir: json.dumps( + _plan(json.loads(config_json)) + ) + mock_native.start_runtime.return_value = json.dumps(_runtime()) + mock_native.invoke_runtime.side_effect = ( + lambda plan_json, runtime_json, request_json: json.dumps( + _result(json.loads(request_json), json.loads(runtime_json)) + ) + ) + mock_native.stop_runtime.return_value = "[]" + return mock_native + + +@pytest.fixture(name="native_client") +def native_client_fixture( + monkeypatch: pytest.MonkeyPatch, + mock_native: MagicMock, +) -> Fabric: + monkeypatch.setattr(client_mod, "_native", mock_native) + return Fabric() + + +async def _post_chunked(url: str, records: list[dict[str, Any]]) -> None: + host_port = url.removeprefix("http://").split("/", 1)[0] + host, port = host_port.split(":", 1) + reader, writer = await asyncio.open_connection(host, int(port)) + writer.write( + b"POST /atof HTTP/1.1\r\n" + + f"Host: {host_port}\r\n".encode() + + b"Transfer-Encoding: chunked\r\n" + + b"Content-Type: application/x-ndjson\r\n\r\n" + ) + for record in records: + payload = json.dumps(record).encode() + b"\n" + writer.write(f"{len(payload):x}\r\n".encode() + payload + b"\r\n") + await writer.drain() + writer.write(b"0\r\n\r\n") + await writer.drain() + assert await reader.readline() == b"HTTP/1.1 200 OK\r\n" + writer.close() + await writer.wait_closed() + + +async def _open_chunked_upload(url: str) -> asyncio.StreamWriter: + host_port = url.removeprefix("http://").split("/", 1)[0] + host, port = host_port.split(":", 1) + _reader, writer = await asyncio.open_connection(host, int(port)) + writer.write( + b"POST /atof HTTP/1.1\r\n" + + f"Host: {host_port}\r\n".encode() + + b"Transfer-Encoding: chunked\r\n" + + b"Content-Type: application/x-ndjson\r\n\r\n" + ) + await writer.drain() + return writer + + +async def _post_content_length( + url: str, + records: list[dict[str, Any]], + *, + expect_continue: bool = False, +) -> None: + host_port = url.removeprefix("http://").split("/", 1)[0] + host, port = host_port.split(":", 1) + reader, writer = await asyncio.open_connection(host, int(port)) + payload = b"".join(json.dumps(record).encode() + b"\n" for record in records) + expect = b"Expect: 100-continue\r\n" if expect_continue else b"" + writer.write( + b"POST /atof HTTP/1.1\r\n" + + f"Host: {host_port}\r\n".encode() + + f"Content-Length: {len(payload)}\r\n".encode() + + expect + + b"Content-Type: application/x-ndjson\r\n\r\n" + ) + await writer.drain() + if expect_continue: + assert await reader.readline() == b"HTTP/1.1 100 Continue\r\n" + assert await reader.readline() == b"\r\n" + writer.write(payload) + await writer.drain() + assert await reader.readline() == b"HTTP/1.1 200 OK\r\n" + writer.close() + await writer.wait_closed() + + +async def _request_status(url: str, request: bytes) -> bytes: + host_port = url.removeprefix("http://").split("/", 1)[0] + host, port = host_port.split(":", 1) + reader, writer = await asyncio.open_connection(host, int(port)) + writer.write(request) + await writer.drain() + status = await reader.readline() + writer.close() + await writer.wait_closed() + return status + + +async def _wait_for(event: threading.Event, timeout: float = 2.0) -> bool: + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while not event.is_set() and loop.time() < deadline: + await asyncio.sleep(0.001) + return event.is_set() + + +async def test_start_runtime_injects_stream_sink_without_mutating_config( + native_client: Fabric, + mock_native: MagicMock, +): + config = _config(relay=True) + + runtime = await native_client.start_runtime(config, streaming=True) + + planned = json.loads(mock_native.plan_config.call_args.args[0]) + sinks = planned["relay"]["observability"]["atof"]["sinks"] + assert sinks[0] == { + "type": "stream", + "url": "https://example.com/events", + "transport": "http_post", + "timeout_millis": 3000, + "field_name_policy": "preserve", + "name": "user-stream", + } + assert sinks[1]["type"] == "stream" + assert sinks[1]["name"] == "nemo-fabric-stream" + assert sinks[1]["transport"] == "ndjson" + assert sinks[1]["url"].startswith("http://127.0.0.1:") + assert runtime.supports_streaming is True + assert len(config.relay.observability.atof.sinks) == 1 + + await runtime.stop() + + +async def test_start_runtime_uses_streaming_host_environment_variable( + native_client: Fabric, + mock_native: MagicMock, +): + os.environ["NEMO_FABRIC_STREAMING_HOST"] = "0.0.0.0" + + runtime = await native_client.start_runtime(_config(relay=True), streaming=True) + + planned = json.loads(mock_native.plan_config.call_args.args[0]) + stream_sink = planned["relay"]["observability"]["atof"]["sinks"][-1] + assert stream_sink["url"].startswith("http://0.0.0.0:") + + await runtime.stop() + + +async def test_start_runtime_without_streaming_preserves_disabled_atof( + native_client: Fabric, + mock_native: MagicMock, +): + config = _config(relay=True) + atof = config.relay.observability.atof + atof.enabled = False + atof.sinks = [RelayAtofFileSinkConfig(output_directory="./disabled")] + + runtime = await native_client.start_runtime(config) + + planned = json.loads(mock_native.plan_config.call_args.args[0]) + planned_atof = planned["relay"]["observability"]["atof"] + assert planned_atof["enabled"] is False + assert len(planned_atof["sinks"]) == 1 + assert planned_atof["sinks"][0]["type"] == "file" + assert runtime.supports_streaming is False + assert config.relay.observability.atof.enabled is False + assert config.relay.observability.atof.sinks[0].output_directory == "./disabled" + + await runtime.stop() + + +async def test_start_runtime_without_streaming_does_not_add_atof( + native_client: Fabric, + mock_native: MagicMock, +): + config = _config() + config.enable_relay( + observability=RelayObservabilityConfig( + atif=RelayAtifConfig(enabled=True), + ) + ) + + runtime = await native_client.start_runtime(config) + + planned = json.loads(mock_native.plan_config.call_args.args[0]) + observability = planned["relay"]["observability"] + assert "atof" not in observability + assert observability["atif"]["enabled"] is True + assert runtime.supports_streaming is False + + await runtime.stop() + + +async def test_start_runtime_streaming_enables_only_reserved_atof_sink( + native_client: Fabric, + mock_native: MagicMock, +): + config = _config(relay=True) + atof = config.relay.observability.atof + atof.enabled = False + atof.sinks = [RelayAtofFileSinkConfig(output_directory="./disabled")] + + runtime = await native_client.start_runtime(config, streaming=True) + + planned = json.loads(mock_native.plan_config.call_args.args[0]) + planned_atof = planned["relay"]["observability"]["atof"] + assert planned_atof["enabled"] is True + assert len(planned_atof["sinks"]) == 1 + assert planned_atof["sinks"][0]["type"] == "stream" + assert planned_atof["sinks"][0]["name"] == "nemo-fabric-stream" + assert planned_atof["sinks"][0]["url"].startswith("http://127.0.0.1:") + assert runtime.supports_streaming is True + assert config.relay.observability.atof.enabled is False + assert config.relay.observability.atof.sinks[0].output_directory == "./disabled" + + await runtime.stop() + + +async def test_start_runtime_rejects_streaming_without_relay( + native_client: Fabric, +): + with pytest.raises( + FabricConfigError, + match="streaming requires Relay telemetry", + ): + await native_client.start_runtime(_config(), streaming=True) + + +def test_with_stream_sink_replaces_reserved_sink_and_preserves_user_sinks(): + config = _config(relay=True) + + first = _with_stream_sink(config, "http://127.0.0.1:4100/atof") + second = _with_stream_sink(first, "http://127.0.0.1:4200/atof") + + sinks = second.relay.observability.atof.sinks + assert [sink.name for sink in sinks] == [ + "user-stream", + "nemo-fabric-stream", + ] + assert sinks[-1].url == "http://127.0.0.1:4200/atof" + assert len(config.relay.observability.atof.sinks) == 1 + + +async def test_invoke_stream_yields_raw_records_and_returns_result_out_of_band( + native_client: Fabric, + mock_native: MagicMock, +): + runtime = await native_client.start_runtime(_config(relay=True), streaming=True) + endpoint = json.loads(mock_native.plan_config.call_args.args[0])["relay"][ + "observability" + ]["atof"]["sinks"][-1]["url"] + request = RunRequest(input="hello", request_id="request-stream") + records = [ + { + "kind": "scope", + "scope_category": "start", + "uuid": "scope-1", + "name": "request", + "metadata": {"nemo_fabric_request_id": request.request_id}, + }, + {"kind": "mark", "uuid": "mark-1", "parent_uuid": "scope-1"}, + ] + + stream = runtime.invoke_stream(request=request) + assert isinstance(stream, InvokeStream) + await _post_content_length(endpoint, records) + streamed = [record async for record in stream] + result = await stream.result() + + assert streamed == records + assert isinstance(result, RunResult) + assert result.output["response"] == "done" + assert all(not isinstance(record, RunResult) for record in streamed) + await runtime.stop() + + +async def test_invoke_stream_correlates_relay_gateway_turn_indexes( + native_client: Fabric, + mock_native: MagicMock, +): + runtime = await native_client.start_runtime(_config(relay=True), streaming=True) + endpoint = json.loads(mock_native.plan_config.call_args.args[0])["relay"][ + "observability" + ]["atof"]["sinks"][-1]["url"] + + for turn_index in (1, 2): + records = [ + { + "kind": "scope", + "scope_category": "start", + "uuid": f"turn-{turn_index}", + "metadata": { + "nemo_relay_scope_role": "turn", + "turn_index": turn_index, + }, + }, + { + "kind": "mark", + "uuid": f"mark-{turn_index}", + "parent_uuid": f"turn-{turn_index}", + }, + ] + + stream = runtime.invoke_stream(input=f"turn {turn_index}") + await _post_content_length(endpoint, records) + + assert [record async for record in stream] == records + assert (await stream.result()).status == "succeeded" + + await runtime.stop() + + +async def test_stream_must_be_finalized_before_another_turn( + native_client: Fabric, + mock_native: MagicMock, +): + runtime = await native_client.start_runtime(_config(relay=True), streaming=True) + endpoint = json.loads(mock_native.plan_config.call_args.args[0])["relay"][ + "observability" + ]["atof"]["sinks"][-1]["url"] + request = RunRequest(input="first", request_id="request-first") + first = { + "kind": "scope", + "scope_category": "start", + "uuid": "first", + "metadata": {"nemo_fabric_request_id": request.request_id}, + } + stream = runtime.invoke_stream(request=request) + await _post_content_length(endpoint, [first]) + + async for record in stream: + assert record == first + break + + with pytest.raises(FabricStateError, match="streaming invocation is active"): + runtime.invoke_stream(input="second") + + await stream.aclose() + second = runtime.invoke_stream(input="second") + with pytest.warns(RuntimeWarning, match="No Relay ATOF connection"): + assert [record async for record in second] == [] + assert (await second.result()).status == "succeeded" + await runtime.stop() + + +async def test_invoke_stream_validates_request_before_returning_stream( + native_client: Fabric, +): + runtime = await native_client.start_runtime(_config(relay=True), streaming=True) + request = RunRequest(input="request") + + with pytest.raises(FabricConfigError, match="mutually exclusive"): + runtime.invoke_stream(input="input", request=request) + + stream = runtime.invoke_stream(input="valid") + with pytest.warns(RuntimeWarning, match="No Relay ATOF connection"): + assert [record async for record in stream] == [] + assert (await stream.result()).status == "succeeded" + await runtime.stop() + + +async def test_invoke_stream_warns_only_once_when_relay_never_connects( + native_client: Fabric, +): + runtime = await native_client.start_runtime(_config(relay=True), streaming=True) + + first = runtime.invoke_stream(input="first") + with pytest.warns(RuntimeWarning, match="able to reach 127.0.0.1"): + assert [record async for record in first] == [] + assert (await first.result()).status == "succeeded" + + second = runtime.invoke_stream(input="second") + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + assert [record async for record in second] == [] + assert caught == [] + assert (await second.result()).status == "succeeded" + await runtime.stop() + + +async def test_invoke_stream_does_not_warn_after_relay_connects( + native_client: Fabric, + mock_native: MagicMock, +): + runtime = await native_client.start_runtime(_config(relay=True), streaming=True) + endpoint = json.loads(mock_native.plan_config.call_args.args[0])["relay"][ + "observability" + ]["atof"]["sinks"][-1]["url"] + + stream = runtime.invoke_stream(input="empty stream") + await _post_content_length(endpoint, []) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + assert [record async for record in stream] == [] + + assert caught == [] + assert (await stream.result()).status == "succeeded" + await runtime.stop() + + +async def test_invoke_stream_warns_after_long_lived_relay_upload_disconnects( + native_client: Fabric, + mock_native: MagicMock, +): + runtime = await native_client.start_runtime(_config(relay=True), streaming=True) + endpoint = json.loads(mock_native.plan_config.call_args.args[0])["relay"][ + "observability" + ]["atof"]["sinks"][-1]["url"] + writer = await _open_chunked_upload(endpoint) + + first = runtime.invoke_stream(input="connected stream") + root = { + "kind": "scope", + "scope_category": "start", + "uuid": "turn-1", + "metadata": { + "nemo_relay_scope_role": "turn", + "turn_index": 1, + }, + } + payload = json.dumps(root).encode() + b"\n" + writer.write(f"{len(payload):x}\r\n".encode() + payload + b"\r\n") + await writer.drain() + + assert [record async for record in first] == [root] + assert (await first.result()).status == "succeeded" + + writer.close() + await writer.wait_closed() + listener = runtime._stream_listener + while listener._active_atof_connections: + await asyncio.sleep(0) + + second = runtime.invoke_stream(input="disconnected stream") + with pytest.warns(RuntimeWarning, match="No Relay ATOF connection"): + assert [record async for record in second] == [] + + assert (await second.result()).status == "succeeded" + await runtime.stop() + + +async def test_invoke_stream_warns_when_long_lived_upload_drops_during_turn( + native_client: Fabric, + mock_native: MagicMock, +): + runtime = await native_client.start_runtime(_config(relay=True), streaming=True) + endpoint = json.loads(mock_native.plan_config.call_args.args[0])["relay"][ + "observability" + ]["atof"]["sinks"][-1]["url"] + writer = await _open_chunked_upload(endpoint) + listener = runtime._stream_listener + while not listener._active_atof_connections: + await asyncio.sleep(0) + + stream = runtime.invoke_stream(input="dropped stream") + writer.close() + await writer.wait_closed() + while listener._active_atof_connections: + await asyncio.sleep(0) + + with pytest.warns(RuntimeWarning, match="No Relay ATOF connection"): + assert [record async for record in stream] == [] + + assert (await stream.result()).status == "succeeded" + await runtime.stop() + + +async def test_invoke_stream_warns_when_long_lived_upload_truncates_turn( + native_client: Fabric, + mock_native: MagicMock, +): + runtime = await native_client.start_runtime(_config(relay=True), streaming=True) + endpoint = json.loads(mock_native.plan_config.call_args.args[0])["relay"][ + "observability" + ]["atof"]["sinks"][-1]["url"] + writer = await _open_chunked_upload(endpoint) + listener = runtime._stream_listener + while not listener._active_atof_connections: + await asyncio.sleep(0) + + stream = runtime.invoke_stream(input="truncated stream") + root = { + "kind": "scope", + "scope_category": "start", + "uuid": "turn-1", + "metadata": { + "nemo_relay_scope_role": "turn", + "turn_index": 1, + }, + } + payload = json.dumps(root).encode() + b"\n" + writer.write(f"{len(payload):x}\r\n".encode() + payload + b"\r\n") + await writer.drain() + while listener.records.empty(): + await asyncio.sleep(0) + + writer.write(b"0\r\n") + await writer.drain() + writer.close() + await writer.wait_closed() + while listener._active_atof_connections: + await asyncio.sleep(0) + + with pytest.warns(RuntimeWarning, match="streaming may be incomplete"): + assert [record async for record in stream] == [root] + + assert (await stream.result()).status == "succeeded" + await runtime.stop() + + +async def test_invoke_stream_does_not_warn_when_chunked_upload_ends_cleanly( + native_client: Fabric, + mock_native: MagicMock, +): + runtime = await native_client.start_runtime(_config(relay=True), streaming=True) + endpoint = json.loads(mock_native.plan_config.call_args.args[0])["relay"][ + "observability" + ]["atof"]["sinks"][-1]["url"] + root = { + "kind": "scope", + "scope_category": "start", + "uuid": "turn-1", + "metadata": { + "nemo_relay_scope_role": "turn", + "turn_index": 1, + }, + } + + stream = runtime.invoke_stream(input="ended stream") + await _post_chunked(endpoint, [root]) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + assert [record async for record in stream] == [root] + + assert caught == [] + assert (await stream.result()).status == "succeeded" + await runtime.stop() + + +async def test_invoke_stream_warns_when_records_do_not_match_active_turn( + native_client: Fabric, + mock_native: MagicMock, +): + runtime = await native_client.start_runtime(_config(relay=True), streaming=True) + endpoint = json.loads(mock_native.plan_config.call_args.args[0])["relay"][ + "observability" + ]["atof"]["sinks"][-1]["url"] + stream = runtime.invoke_stream(input="unmatched stream") + await _post_content_length( + endpoint, + [ + { + "kind": "scope", + "scope_category": "start", + "uuid": "unexpected-turn", + "metadata": { + "nemo_relay_scope_role": "turn", + "turn_index": 99, + }, + } + ], + ) + + with pytest.warns(RuntimeWarning, match="no record matched the active Fabric turn"): + assert [record async for record in stream] == [] + + assert (await stream.result()).status == "succeeded" + + second = runtime.invoke_stream(input="another unmatched stream") + await _post_content_length( + endpoint, + [ + { + "kind": "scope", + "scope_category": "start", + "uuid": "another-unexpected-turn", + "metadata": { + "nemo_relay_scope_role": "turn", + "turn_index": 99, + }, + } + ], + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + assert [record async for record in second] == [] + + assert caught == [] + assert (await second.result()).status == "succeeded" + await runtime.stop() + + +async def test_stop_finalizes_completed_stream_after_result( + native_client: Fabric, +): + runtime = await native_client.start_runtime(_config(relay=True), streaming=True) + stream = runtime.invoke_stream(input="hello") + + assert (await stream.result()).status == "succeeded" + assert stream._finalized is False + + await runtime.stop() + + assert stream._finalized is True + + +@pytest.mark.parametrize("relay", [False, True]) +async def test_invoke_stream_requires_streaming_enabled_at_startup( + native_client: Fabric, + relay: bool, +): + runtime = await native_client.start_runtime(_config(relay=relay)) + + assert runtime.supports_streaming is False + with pytest.raises( + FabricCapabilityError, + match=r"requires Relay telemetry.*streaming=True", + ) as caught: + runtime.invoke_stream(input="hello") + + assert caught.value.code == "streaming_unavailable" + assert caught.value.details == {"capability": "streaming"} + await runtime.stop() + + +async def test_context_manager_finalizes_unconsumed_stream( + native_client: Fabric, +): + async with await native_client.start_runtime( + _config(relay=True), streaming=True + ) as runtime: + stream = runtime.invoke_stream(input="hello") + + assert (await stream.result()).status == "succeeded" + + +async def test_cancelled_aclose_keeps_turn_active_and_result_awaitable( + native_client: Fabric, + mock_native: MagicMock, +): + started = threading.Event() + release = threading.Event() + + def invoke(plan_json: str, runtime_json: str, request_json: str) -> str: + started.set() + assert release.wait(timeout=2) + return json.dumps(_result(json.loads(request_json), json.loads(runtime_json))) + + mock_native.invoke_runtime.side_effect = invoke + runtime = await native_client.start_runtime(_config(relay=True), streaming=True) + stream = runtime.invoke_stream(input="hello") + assert await _wait_for(started) + + closing = asyncio.create_task(stream.aclose()) + await asyncio.sleep(0) + closing.cancel() + with pytest.raises(asyncio.CancelledError): + await closing + with pytest.raises(FabricStateError, match="streaming invocation is active"): + runtime.invoke_stream(input="too soon") + with pytest.raises( + FabricStateError, + match="streaming invocation is active", + ): + await runtime.stop() + + release.set() + await stream.aclose() + assert (await stream.result()).status == "succeeded" + await runtime.stop() + + +async def test_cancelled_anext_does_not_consume_next_record(): + listener = await _AtofStreamListener(maxsize=2).start() + invocation_finished = asyncio.Event() + + async def invoke() -> RunResult: + await invocation_finished.wait() + return RunResult.from_mapping(_result({"request_id": "request-1"}, _runtime())) + + stream = InvokeStream(invoke(), listener) + pending = asyncio.create_task(stream.__anext__()) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + + records = [{"uuid": "first"}, {"uuid": "second"}] + await _post_chunked(listener.url, records) + + assert await stream.__anext__() == records[0] + assert await stream.__anext__() == records[1] + + invocation_finished.set() + await stream.aclose() + await listener.close() + + +async def test_cancelled_anext_retains_record_consumed_during_cancellation(): + listener = await _AtofStreamListener(maxsize=1).start() + invocation_finished = asyncio.Event() + + async def invoke() -> RunResult: + await invocation_finished.wait() + return RunResult.from_mapping(_result({"request_id": "request-1"}, _runtime())) + + stream = InvokeStream(invoke(), listener) + record = {"uuid": "first"} + + async def cancel_after_getter_completes( + tasks: set[asyncio.Task[Any]], + *, + return_when: str, + ) -> None: + assert return_when == asyncio.FIRST_COMPLETED + getter = next(task for task in tasks if task is not stream._task) + listener.records.put_nowait(record) + assert await getter == record + raise asyncio.CancelledError + + with ( + patch.object(asyncio, "wait", new=cancel_after_getter_completes), + pytest.raises(asyncio.CancelledError), + ): + await stream.__anext__() + + assert await stream.__anext__() == record + + invocation_finished.set() + await stream.aclose() + await listener.close() + + +async def test_aclose_drains_backpressure_while_invocation_finishes(): + listener = await _AtofStreamListener(maxsize=1).start() + producer_finished = asyncio.Event() + + async def invoke() -> RunResult: + await producer_finished.wait() + return RunResult.from_mapping(_result({"request_id": "request-1"}, _runtime())) + + stream = InvokeStream(invoke(), listener) + records = [{"uuid": f"record-{index}"} for index in range(3)] + + async def produce() -> None: + await _post_chunked(listener.url, records) + producer_finished.set() + + producer = asyncio.create_task(produce()) + while not listener.records.full(): + await asyncio.sleep(0) + assert not producer.done() + + await asyncio.wait_for(stream.aclose(), timeout=1) + + assert producer.done() + assert (await stream.result()).status == "succeeded" + await listener.close() + + +@pytest.mark.parametrize( + "current_metadata", + [ + {"nemo_fabric_request_id": "request-2"}, + {"nemo_relay_scope_role": "turn", "turn_index": 2}, + ], +) +async def test_listener_correlates_records_to_active_turn( + current_metadata: dict[str, Any], +): + listener = await _AtofStreamListener(maxsize=4).start() + listener.begin_stream(request_id="request-2", turn_index=2) + current = [ + { + "kind": "scope", + "scope_category": "start", + "uuid": "current", + "metadata": current_metadata, + }, + { + "kind": "scope", + "scope_category": "start", + "uuid": "child", + "parent_uuid": "current", + }, + {"kind": "mark", "uuid": "mark", "parent_uuid": "child"}, + ] + + await _post_chunked( + listener.url, + [ + { + "kind": "scope", + "scope_category": "start", + "uuid": "previous", + "metadata": {"nemo_fabric_request_id": "request-1"}, + }, + {"kind": "mark", "uuid": "late", "parent_uuid": "previous"}, + *current, + {"kind": "mark", "uuid": "unrelated", "parent_uuid": "previous"}, + ], + ) + + assert [await listener.records.get() for _ in current] == current + assert listener.records.empty() + listener.end_stream() + await listener.close() + + +async def test_listener_applies_byte_budget_backpressure(): + record = {"uuid": "record", "payload": "x" * 16} + record_size = len(json.dumps(record).encode()) + listener = await _AtofStreamListener( + maxsize=10, + max_bytes=record_size, + max_record_bytes=record_size, + ).start() + listener.begin_stream() + + producer = asyncio.create_task(_post_chunked(listener.url, [record, record])) + while listener.records.empty(): + await asyncio.sleep(0) + assert not producer.done() + + assert await listener.records.get() == record + assert await listener.records.get() == record + await producer + listener.end_stream() + await listener.close() + + +async def test_listener_rejects_oversized_record(): + listener = await _AtofStreamListener(max_record_bytes=32).start() + listener.begin_stream(request_id="request-1") + payload = json.dumps({"payload": "x" * 64}).encode() + b"\n" + request = ( + b"POST /atof HTTP/1.1\r\n" + + f"Content-Length: {len(payload)}\r\n".encode() + + b"Content-Type: application/x-ndjson\r\n\r\n" + + payload + ) + + assert await _request_status(listener.url, request) == ( + b"HTTP/1.1 413 Content Too Large\r\n" + ) + listener.end_stream() + with pytest.warns(RuntimeWarning, match="no record matched the active Fabric turn"): + listener.warn_if_unavailable() + await listener.close() + + +async def test_listener_accepts_atof_record_larger_than_default_read_limits(): + listener = await _AtofStreamListener(maxsize=2).start() + listener.begin_stream() + record = {"uuid": "large", "payload": "x" * (600 * 1024)} + + await _post_chunked(listener.url, [record]) + + assert await listener.records.get() == record + listener.end_stream() + await listener.close() + + +async def test_listener_accepts_content_length_and_100_continue(): + listener = await _AtofStreamListener(maxsize=2).start() + listener.begin_stream() + records = [{"uuid": "first"}, {"uuid": "second"}] + + await _post_content_length(listener.url, records, expect_continue=True) + + assert [await listener.records.get(), await listener.records.get()] == records + listener.end_stream() + await listener.close() + + +@pytest.mark.parametrize( + ("raw_request", "expected"), + [ + ( + b"GET /atof HTTP/1.1\r\nContent-Length: 0\r\n\r\n", + b"HTTP/1.1 404 Not Found\r\n", + ), + ( + b"POST /atof HTTP/1.1\r\n\r\n", + b"HTTP/1.1 411 Length Required\r\n", + ), + ( + b"POST /atof HTTP/1.1\r\nContent-Length: invalid\r\n\r\n", + b"HTTP/1.1 400 Bad Request\r\n", + ), + ], +) +async def test_listener_rejects_invalid_http_requests( + raw_request: bytes, + expected: bytes, +): + listener = await _AtofStreamListener().start() + + assert await _request_status(listener.url, raw_request) == expected + + await listener.close() diff --git a/tests/scripts/test_sync_fern_docs_branch.py b/tests/scripts/test_sync_fern_docs_branch.py index aec4e150a..42ada5c01 100644 --- a/tests/scripts/test_sync_fern_docs_branch.py +++ b/tests/scripts/test_sync_fern_docs_branch.py @@ -29,6 +29,12 @@ def _source_tree(root: Path) -> None: ) (docs / "_source").mkdir() (docs / "_source" / "ignored.md").write_text("ignored\n", encoding="utf-8") + python_reference = docs / "reference" / "api" / "python-library-reference" + python_reference.mkdir(parents=True) + (python_reference / "client.md").write_text( + f"{sync_fern_docs_branch.MARKDOWN_SPDX_HEADER}\n\n# Client\n", + encoding="utf-8", + ) _write_yaml( docs / "index.yml", { @@ -90,6 +96,14 @@ def test_sync_dev_rewrites_navigation_and_preserves_versions(tmp_path: Path): assert not (target_fern / "pages-dev" / "index.yml").exists() assert not (target_fern / "pages-dev" / "_source").exists() assert (target_fern / "fern.config.json").is_file() + assert ( + target_fern + / "pages-dev" + / "reference" + / "api" + / "python-library-reference" + / "client.md" + ).read_text(encoding="utf-8").startswith(sync_fern_docs_branch.MDX_SPDX_HEADER) navigation = sync_fern_docs_branch.read_yaml(target_fern / "versions" / "dev.yml") assert navigation["navigation"][0]["path"] == "../pages-dev/guide.mdx" @@ -179,3 +193,11 @@ def test_release_version_promotes_stable_snapshot(tmp_path: Path): assert "blob/0.2.0/README.md" in ( target_fern / "pages-v0.2.0" / "guide.mdx" ).read_text(encoding="utf-8") + assert ( + target_fern + / "pages-v0.2.0" + / "reference" + / "api" + / "python-library-reference" + / "client.md" + ).read_text(encoding="utf-8").startswith(sync_fern_docs_branch.MDX_SPDX_HEADER)