From 164b14595f57da4677fc35b9c02d32dd18ffb771 Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Wed, 2 Sep 2026 16:26:36 -0500 Subject: [PATCH 1/6] route requests by session-header presence Signed-off-by: Manjesh Mogallapalli --- .../src/nemo_agents_plugin/fabric/runtime.py | 34 ++++++- .../src/nemo_agents_plugin/fabric/server.py | 56 +++++++---- .../fabric/session_manager.py | 71 ++++++++++---- .../tests/unit/test_fabric_runtime.py | 68 ++++++++++++-- .../tests/unit/test_fabric_server.py | 92 +++++++------------ .../tests/unit/test_fabric_session_manager.py | 64 +++++++++++-- 6 files changed, 273 insertions(+), 112 deletions(-) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py index b95f6b76d5..d26d9a5f8b 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py @@ -15,12 +15,13 @@ import asyncio from collections.abc import AsyncIterator, Mapping, Sequence +from contextlib import asynccontextmanager from dataclasses import dataclass, field from pathlib import Path from typing import Any # CI type-checks this plugin via ty extra-paths without installing nemo-agents deps. -from nemo_fabric import ( # ty: ignore[unresolved-import] +from nemo_fabric import ( Fabric, FabricConfig, FabricError, @@ -202,6 +203,37 @@ async def run_fabric_agent_once( return _normalize_fabric_run_result(result) +@asynccontextmanager +async def stream_fabric_agent_once( + request: FabricOneShotRequest, + *, + fabric: Any | None = None, +) -> AsyncIterator[FabricRuntimeStream]: + """Start an ephemeral Fabric runtime and keep it alive for one stream.""" + fabric_client = fabric or Fabric() + + try: + async with await fabric_client.start_runtime( + request.fabric_config, + base_dir=request.base_dir, + overrides=request.overrides, + streaming=True, + ) as runtime: + yield stream_fabric_runtime( + runtime, + FabricInvocationRequest( + input=request.input, + request_id=request.request_id, + caller_context=request.caller_context, + timeout_seconds=request.timeout_seconds, + ), + ) + except FabricError as error: + raise FabricRuntimeExecutionError( + f"Fabric runtime streaming failed: {error}", + ) from error + + async def _invoke_fabric_agent_once( request: FabricOneShotRequest, *, diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py index 9603b6916c..46e7f10101 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py @@ -14,6 +14,7 @@ from contextlib import AbstractAsyncContextManager, asynccontextmanager from dataclasses import dataclass from datetime import UTC, datetime +from functools import partial from pathlib import Path from typing import Annotated, Any @@ -73,7 +74,7 @@ def __post_init__(self) -> None: def _to_fabric_invocation_request( request: ChatCompletionRequest, *, - session_id: str, + session_id: str | None, ) -> FabricInvocationRequest: """Translate the chat transcript into a Platform-owned Fabric request.""" messages = request.messages @@ -83,9 +84,10 @@ def _to_fabric_invocation_request( if len(messages) == 1 else "\n\n".join(f"{message.role}: {message.content}" for message in messages) ) + caller_context = {"session_id": session_id} if session_id is not None else {} return FabricInvocationRequest( input=input_text, - caller_context={"session_id": session_id}, + caller_context=caller_context, ) @@ -305,29 +307,40 @@ async def chat_completions( response: Response, session_id: Annotated[str | None, Header(alias=SESSION_ID_HEADER)] = None, ) -> ChatCompletionResponse | StreamingResponse: - try: - session = await app.state.session_manager.resolve_session(session_id) - except FabricSessionNotFoundError as error: - raise HTTPException(status_code=404, detail=str(error)) from error - except FabricSessionStartError as error: - raise HTTPException(status_code=503, detail=str(error)) from error + response_headers: dict[str, str] | None = None + if session_id is None: + invocation_request = _to_fabric_invocation_request(request, session_id=None) + invoke = app.state.session_manager.invoke_once + stream = app.state.session_manager.stream_once + else: + try: + session = await app.state.session_manager.resolve_session(session_id) + except FabricSessionNotFoundError as error: + raise HTTPException(status_code=404, detail=str(error)) from error + except FabricSessionStartError as error: + raise HTTPException(status_code=503, detail=str(error)) from error + invocation_request = _to_fabric_invocation_request(request, session_id=session.session_id) + response_headers = _session_headers(session.session_id) + invoke = partial(app.state.session_manager.invoke_session, session) + stream = partial(app.state.session_manager.stream_session, session) - invocation_request = _to_fabric_invocation_request(request, session_id=session.session_id) if request.stream: - stream_context = app.state.session_manager.stream_session(session, invocation_request) + stream_context = stream(invocation_request) try: fabric_stream = await stream_context.__aenter__() + except FabricSessionStartError as error: + raise HTTPException(status_code=503, detail=str(error), headers=response_headers) from error except FabricSessionNotFoundError as error: raise HTTPException( status_code=404, detail=str(error), - headers=_session_headers(session.session_id), + headers=response_headers, ) from error except FabricRuntimeExecutionError as error: raise HTTPException( status_code=502, detail=str(error), - headers=_session_headers(session.session_id), + headers=response_headers, ) from error return StreamingResponse( _iter_streaming_chat_completion( @@ -337,35 +350,37 @@ async def chat_completions( model=_request_model_name(request), ), media_type="text/event-stream", - headers=_session_headers(session.session_id), + headers=response_headers, ) try: - result = await app.state.session_manager.invoke_session(session, invocation_request) + result = await invoke(invocation_request) + except FabricSessionStartError as error: + raise HTTPException(status_code=503, detail=str(error), headers=response_headers) from error except FabricSessionNotFoundError as error: raise HTTPException( status_code=404, detail=str(error), - headers=_session_headers(session.session_id), + headers=response_headers, ) from error except FabricRuntimeTimeoutError as error: raise HTTPException( status_code=504, detail=str(error), - headers=_session_headers(session.session_id), + headers=response_headers, ) from error except FabricRuntimeExecutionError as error: raise HTTPException( status_code=502, detail=str(error), - headers=_session_headers(session.session_id), + headers=response_headers, ) from error if result.status != "succeeded": raise HTTPException( status_code=502, detail=_failed_result_detail(result), - headers=_session_headers(session.session_id), + headers=response_headers, ) try: @@ -374,10 +389,11 @@ async def chat_completions( raise HTTPException( status_code=502, detail=str(error), - headers=_session_headers(session.session_id), + headers=response_headers, ) from error - response.headers[SESSION_ID_HEADER] = session.session_id + if response_headers is not None: + response.headers.update(response_headers) return completion @app.delete("/v1/sessions/{session_id}", status_code=204) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py index 46d9c9806f..bfe9f16759 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py @@ -16,9 +16,12 @@ from nemo_agents_plugin.fabric.environment import ensure_local_workspace_dir from nemo_agents_plugin.fabric.runtime import ( FabricInvocationRequest, + FabricOneShotRequest, FabricRuntimeResult, FabricRuntimeStream, invoke_fabric_runtime, + run_fabric_agent_once, + stream_fabric_agent_once, stream_fabric_runtime, ) from nemo_agents_plugin.fabric.session_registry import ( @@ -82,16 +85,11 @@ def __init__( async def open_session(self, *, session_id: str | None = None) -> FabricRuntimeSession: """Materialize a Fabric config, start its runtime, and register the session.""" - try: - fabric_config = translate_agent_config(self._agent_config) - except FabricTranslationError as error: - raise FabricSessionStartError(f"Fabric config translation failed: {error}") from error - - await asyncio.to_thread(ensure_local_workspace_dir, self._agent_config, self._base_dir) + fabric_config = await self._materialize_fabric_config(streaming=True) fabric = self._fabric or Fabric() try: runtime = await fabric.start_runtime( - _prepare_serving_fabric_config(fabric_config), + fabric_config, base_dir=self._base_dir, streaming=True, ) @@ -108,10 +106,8 @@ async def open_session(self, *, session_id: str | None = None) -> FabricRuntimeS logger.exception("Failed to stop Fabric runtime after session registration failed.") raise - async def resolve_session(self, session_id: str | None) -> FabricRuntimeSession: + async def resolve_session(self, session_id: str) -> FabricRuntimeSession: """Resolve a session, lazily starting a runtime under a supplied Platform ID.""" - if session_id is None: - return await self.open_session() if session_id in self._closed_session_ids: raise FabricSessionNotFoundError(f"Fabric session '{session_id}' was not found.") @@ -135,6 +131,20 @@ async def resolve_session(self, session_id: str | None) -> FabricRuntimeSession: finally: self._release_session_creation_gate(session_id, creation_gate) + async def invoke_once(self, request: FabricInvocationRequest) -> FabricRuntimeResult: + """Run one request on an ephemeral runtime without registering a session.""" + async with self._invocation_slot(): + one_shot_request = await self._to_one_shot_request(request, streaming=False) + return await run_fabric_agent_once(one_shot_request, fabric=self._fabric) + + @asynccontextmanager + async def stream_once(self, request: FabricInvocationRequest) -> AsyncIterator[FabricRuntimeStream]: + """Stream one request from an ephemeral, unregistered runtime.""" + async with self._invocation_slot(): + one_shot_request = await self._to_one_shot_request(request, streaming=True) + async with stream_fabric_agent_once(one_shot_request, fabric=self._fabric) as stream: + yield stream + async def invoke_session( self, session: FabricRuntimeSession, @@ -145,9 +155,7 @@ async def invoke_session( if session.closing: raise FabricSessionNotFoundError(f"Fabric session '{session.session_id}' was not found.") try: - if self._invocation_semaphore is None: - return await invoke_fabric_runtime(session.runtime, request) - async with self._invocation_semaphore: + async with self._invocation_slot(): return await invoke_fabric_runtime(session.runtime, request) finally: await self._session_registry.refresh_activity(session) @@ -163,11 +171,8 @@ async def stream_session( if session.closing: raise FabricSessionNotFoundError(f"Fabric session '{session.session_id}' was not found.") try: - if self._invocation_semaphore is None: + async with self._invocation_slot(): yield stream_fabric_runtime(session.runtime, request) - else: - async with self._invocation_semaphore: - yield stream_fabric_runtime(session.runtime, request) finally: await self._session_registry.refresh_activity(session) @@ -243,6 +248,38 @@ async def _stop_session(self, session: FabricRuntimeSession) -> None: except FabricError as error: raise FabricSessionStopError(f"Fabric runtime shutdown failed: {error}") from error + async def _materialize_fabric_config(self, *, streaming: bool) -> FabricConfig: + try: + fabric_config = translate_agent_config(self._agent_config) + except FabricTranslationError as error: + raise FabricSessionStartError(f"Fabric config translation failed: {error}") from error + + await asyncio.to_thread(ensure_local_workspace_dir, self._agent_config, self._base_dir) + return _prepare_serving_fabric_config(fabric_config) if streaming else fabric_config + + async def _to_one_shot_request( + self, + request: FabricInvocationRequest, + *, + streaming: bool, + ) -> FabricOneShotRequest: + return FabricOneShotRequest( + fabric_config=await self._materialize_fabric_config(streaming=streaming), + base_dir=self._base_dir, + input=request.input, + request_id=request.request_id, + caller_context=request.caller_context, + timeout_seconds=request.timeout_seconds, + ) + + @asynccontextmanager + async def _invocation_slot(self) -> AsyncIterator[None]: + if self._invocation_semaphore is None: + yield + return + async with self._invocation_semaphore: + yield + def _prepare_serving_fabric_config(fabric_config: FabricConfig) -> FabricConfig: """Enable serving-owned Relay support without mutating the translated config.""" diff --git a/plugins/nemo-agents/tests/unit/test_fabric_runtime.py b/plugins/nemo-agents/tests/unit/test_fabric_runtime.py index a7d1a6d0a9..c632ceeb38 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_runtime.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_runtime.py @@ -18,9 +18,10 @@ FabricRuntimeTimeoutError, invoke_fabric_runtime, run_fabric_agent_once, + stream_fabric_agent_once, stream_fabric_runtime, ) -from nemo_fabric import FabricConfig # ty: ignore[unresolved-import] +from nemo_fabric import FabricConfig class _FabricMapping: @@ -148,14 +149,16 @@ async def start_runtime( *, base_dir: Path | str, overrides: dict[str, Any] | None = None, + streaming: bool = False, ) -> _FakeRuntime: - self.start_calls.append( - { - "base_dir": base_dir, - "fabric_config": fabric_config, - "overrides": overrides, - } - ) + call = { + "base_dir": base_dir, + "fabric_config": fabric_config, + "overrides": overrides, + } + if streaming: + call["streaming"] = True + self.start_calls.append(call) if self.start_error is not None: raise self.start_error return self.runtime @@ -278,6 +281,55 @@ async def test_normalizes_fabric_mapping_fields_to_plain_values(self) -> None: assert result.metadata == {"adapter_runner": "python"} +@pytest.mark.asyncio +class TestStreamFabricAgentOnce: + async def test_keeps_ephemeral_runtime_alive_for_stream_context(self) -> None: + fabric_config = cast(FabricConfig, object()) + fake_runtime = _FakeRuntime() + fake_fabric = _FakeFabric(runtime=fake_runtime) + request = FabricOneShotRequest( + fabric_config=fabric_config, + base_dir=Path("/tmp/agent"), + input="hello", + request_id="platform-request-1", + caller_context={"source": "server"}, + ) + + async with stream_fabric_agent_once(request, fabric=fake_fabric) as stream: + assert fake_runtime.entered is True + assert fake_runtime.exited is False + assert [record async for record in stream.records()] == [{"type": "span", "message": "thinking"}] + result = await stream.result() + + assert fake_runtime.exited is True + assert fake_fabric.start_calls == [ + { + "base_dir": Path("/tmp/agent"), + "fabric_config": fabric_config, + "overrides": None, + "streaming": True, + } + ] + fabric_request = fake_runtime.invoke_stream_requests[0] + assert fabric_request.input == "hello" + assert fabric_request.request_id == "platform-request-1" + assert fabric_request.context == {"source": "server"} + assert result.response == "hello" + + async def test_cleans_up_runtime_when_stream_start_fails(self) -> None: + fake_runtime = _FakeRuntime(invoke_error=fabric_runtime.FabricError("stream unavailable")) + request = FabricOneShotRequest( + fabric_config=cast(FabricConfig, object()), + base_dir=Path("/tmp/agent"), + ) + + with pytest.raises(FabricRuntimeExecutionError, match="stream unavailable"): + async with stream_fabric_agent_once(request, fabric=_FakeFabric(runtime=fake_runtime)): + pass + + assert fake_runtime.exited is True + + @pytest.mark.asyncio class TestInvokeFabricRuntime: async def test_invokes_active_runtime_without_changing_its_lifecycle(self) -> None: diff --git a/plugins/nemo-agents/tests/unit/test_fabric_server.py b/plugins/nemo-agents/tests/unit/test_fabric_server.py index 4f0870e95d..a5e0455568 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_server.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_server.py @@ -300,23 +300,20 @@ def test_rejects_invalid_serving_settings() -> None: FabricServingSettings(session_cleanup_interval_seconds=0) -def test_chat_completion_without_session_id_opens_and_returns_session( +def test_chat_completion_without_session_id_invokes_once_without_creating_session( tmp_path: Path, mock_validate_agent_config: list[tuple[AgentConfig, Path]], monkeypatch: pytest.MonkeyPatch, ) -> None: config_path = _write_agent_config(tmp_path) app = create_fabric_serving_app(config_path) - resolve_calls: list[str | None] = [] - invocation_calls: list[tuple[Any, Any]] = [] - runtime = object() + invocation_calls: list[Any] = [] - async def resolve_session(session_id: str | None) -> Any: - resolve_calls.append(session_id) - return SimpleNamespace(session_id="session-1", runtime=runtime) + async def resolve_session(session_id: str) -> Any: + pytest.fail(f"headerless request resolved session {session_id!r}") - async def invoke_session(session: Any, request: Any) -> FabricRuntimeResult: - invocation_calls.append((session, request)) + async def invoke_once(request: Any) -> FabricRuntimeResult: + invocation_calls.append(request) return FabricRuntimeResult( status="succeeded", output={"response": "hello", "usage": {"total_tokens": 3}}, @@ -326,14 +323,14 @@ async def invoke_session(session: Any, request: Any) -> FabricRuntimeResult: with TestClient(app) as client: monkeypatch.setattr(app.state.session_manager, "resolve_session", resolve_session) - monkeypatch.setattr(app.state.session_manager, "invoke_session", invoke_session) + monkeypatch.setattr(app.state.session_manager, "invoke_once", invoke_once) response = client.post( "/v1/chat/completions", json={"messages": [{"role": "user", "content": "hello"}]}, ) assert response.status_code == 200 - assert response.headers[SESSION_ID_HEADER] == "session-1" + assert SESSION_ID_HEADER not in response.headers assert response.json() == { "id": "invocation-1", "object": "chat.completion", @@ -347,11 +344,9 @@ async def invoke_session(session: Any, request: Any) -> FabricRuntimeResult: ], "usage": {"total_tokens": 3}, } - assert resolve_calls == [None] - resolved_session, invocation_request = invocation_calls[0] - assert resolved_session.runtime is runtime + invocation_request = invocation_calls[0] assert invocation_request.input == "hello" - assert invocation_request.caller_context == {"session_id": "session-1"} + assert invocation_request.caller_context == {} def test_chat_completion_with_session_id_reuses_session( @@ -361,9 +356,9 @@ def test_chat_completion_with_session_id_reuses_session( ) -> None: config_path = _write_agent_config(tmp_path) app = create_fabric_serving_app(config_path) - resolve_calls: list[str | None] = [] + resolve_calls: list[str] = [] - async def resolve_session(session_id: str | None) -> Any: + async def resolve_session(session_id: str) -> Any: resolve_calls.append(session_id) return SimpleNamespace(session_id="session-1", runtime=object()) @@ -400,22 +395,18 @@ def test_chat_completion_maps_runtime_errors( ) -> None: app = create_fabric_serving_app(_write_agent_config(tmp_path)) - async def resolve_session(session_id: str | None) -> Any: - return SimpleNamespace(session_id="session-1", runtime=object()) - - async def invoke_session(session: Any, request: Any) -> FabricRuntimeResult: + async def invoke_once(request: Any) -> FabricRuntimeResult: raise error with TestClient(app) as client: - monkeypatch.setattr(app.state.session_manager, "resolve_session", resolve_session) - monkeypatch.setattr(app.state.session_manager, "invoke_session", invoke_session) + monkeypatch.setattr(app.state.session_manager, "invoke_once", invoke_once) response = client.post( "/v1/chat/completions", json={"messages": [{"role": "user", "content": "hello"}]}, ) assert response.status_code == status_code - assert response.headers[SESSION_ID_HEADER] == "session-1" + assert SESSION_ID_HEADER not in response.headers assert response.json() == {"detail": str(error)} @@ -435,7 +426,7 @@ def test_chat_completion_maps_session_resolution_errors( ) -> None: app = create_fabric_serving_app(_write_agent_config(tmp_path)) - async def resolve_session(session_id: str | None) -> Any: + async def resolve_session(session_id: str) -> Any: raise error with TestClient(app) as client: @@ -451,41 +442,35 @@ async def resolve_session(session_id: str | None) -> Any: assert response.json() == {"detail": str(error)} -def test_streaming_chat_completion_maps_stream_start_errors( +def test_streaming_chat_completion_without_session_maps_stream_start_errors( tmp_path: Path, mock_validate_agent_config: list[tuple[AgentConfig, Path]], monkeypatch: pytest.MonkeyPatch, ) -> None: app = create_fabric_serving_app(_write_agent_config(tmp_path)) - async def resolve_session(session_id: str | None) -> Any: - return SimpleNamespace(session_id="session-1", runtime=object()) - - def stream_session(session: Any, request: Any) -> _FakeStreamContext: + def stream_once(request: Any) -> _FakeStreamContext: return _FakeStreamContext(enter_error=FabricRuntimeExecutionError("stream failed")) with TestClient(app) as client: - monkeypatch.setattr(app.state.session_manager, "resolve_session", resolve_session) - monkeypatch.setattr(app.state.session_manager, "stream_session", stream_session) + monkeypatch.setattr(app.state.session_manager, "stream_once", stream_once) response = client.post( "/v1/chat/completions", json={"messages": [{"role": "user", "content": "hello"}], "stream": True}, ) assert response.status_code == 502 - assert response.headers[SESSION_ID_HEADER] == "session-1" + assert SESSION_ID_HEADER not in response.headers assert response.json() == {"detail": "stream failed"} -def test_streaming_chat_completion_returns_openai_sse_response( +def test_streaming_chat_completion_without_session_uses_one_shot_stream( tmp_path: Path, mock_validate_agent_config: list[tuple[AgentConfig, Path]], monkeypatch: pytest.MonkeyPatch, ) -> None: app = create_fabric_serving_app(_write_agent_config(tmp_path)) - resolve_calls: list[str | None] = [] - stream_calls: list[tuple[Any, Any]] = [] - runtime = object() + stream_calls: list[Any] = [] fabric_stream = _FakeFabricStream( [ {"kind": "scope", "scope_category": "start", "name": "request"}, @@ -496,17 +481,12 @@ def test_streaming_chat_completion_returns_openai_sse_response( ) stream_context = _FakeStreamContext(fabric_stream) - async def resolve_session(session_id: str | None) -> Any: - resolve_calls.append(session_id) - return SimpleNamespace(session_id="session-1", runtime=runtime) - - def stream_session(session: Any, request: Any) -> _FakeStreamContext: - stream_calls.append((session, request)) + def stream_once(request: Any) -> _FakeStreamContext: + stream_calls.append(request) return stream_context with TestClient(app) as client: - monkeypatch.setattr(app.state.session_manager, "resolve_session", resolve_session) - monkeypatch.setattr(app.state.session_manager, "stream_session", stream_session) + monkeypatch.setattr(app.state.session_manager, "stream_once", stream_once) with client.stream( "POST", "/v1/chat/completions", @@ -516,12 +496,10 @@ def stream_session(session: Any, request: Any) -> _FakeStreamContext: assert response.status_code == 200 assert "text/event-stream" in response.headers["content-type"] - assert response.headers[SESSION_ID_HEADER] == "session-1" - assert resolve_calls == [None] - resolved_session, invocation_request = stream_calls[0] - assert resolved_session.runtime is runtime + assert SESSION_ID_HEADER not in response.headers + invocation_request = stream_calls[0] assert invocation_request.input == "hello" - assert invocation_request.caller_context == {"session_id": "session-1"} + assert invocation_request.caller_context == {} assert stream_context.exit_calls == 1 assert fabric_stream.aclose_calls == 0 @@ -538,9 +516,9 @@ def test_streaming_chat_completion_reuses_supplied_session_id( monkeypatch: pytest.MonkeyPatch, ) -> None: app = create_fabric_serving_app(_write_agent_config(tmp_path)) - resolve_calls: list[str | None] = [] + resolve_calls: list[str] = [] - async def resolve_session(session_id: str | None) -> Any: + async def resolve_session(session_id: str) -> Any: resolve_calls.append(session_id) return SimpleNamespace(session_id="session-1", runtime=object()) @@ -639,25 +617,21 @@ def test_chat_completion_maps_failed_run_result( ) -> None: app = create_fabric_serving_app(_write_agent_config(tmp_path)) - async def resolve_session(session_id: str | None) -> Any: - return SimpleNamespace(session_id="session-1", runtime=object()) - - async def invoke_session(session: Any, request: Any) -> FabricRuntimeResult: + async def invoke_once(request: Any) -> FabricRuntimeResult: return FabricRuntimeResult( status="failed", error={"stage": "invoke", "message": "adapter failed"}, ) with TestClient(app) as client: - monkeypatch.setattr(app.state.session_manager, "resolve_session", resolve_session) - monkeypatch.setattr(app.state.session_manager, "invoke_session", invoke_session) + monkeypatch.setattr(app.state.session_manager, "invoke_once", invoke_once) response = client.post( "/v1/chat/completions", json={"messages": [{"role": "user", "content": "hello"}]}, ) assert response.status_code == 502 - assert response.headers[SESSION_ID_HEADER] == "session-1" + assert SESSION_ID_HEADER not in response.headers assert response.json() == {"detail": "adapter failed"} diff --git a/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py b/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py index 941aa98f37..551b9a37c1 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio +from contextlib import asynccontextmanager from pathlib import Path from typing import Any, cast @@ -150,14 +151,21 @@ async def fail_registration(runtime: Any, *, session_id: str | None = None) -> N @pytest.mark.asyncio -async def test_resolve_session_opens_session_when_id_is_absent( +async def test_invoke_once_uses_unregistered_runtime( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setattr(session_manager, "translate_agent_config", lambda config: _FakeFabricConfig()) - runtime = _FakeRuntime() - fabric = _FakeFabric(runtime) + fabric_config = _FakeFabricConfig() + monkeypatch.setattr(session_manager, "translate_agent_config", lambda config: fabric_config) + captured: list[tuple[Any, Any]] = [] + + async def run_fabric_agent_once(request: Any, *, fabric: Any) -> FabricRuntimeResult: + captured.append((request, fabric)) + return FabricRuntimeResult(status="succeeded", response=request.input) + + monkeypatch.setattr(session_manager, "run_fabric_agent_once", run_fabric_agent_once) registry = FabricSessionRegistry() + fabric = object() manager = FabricSessionManager( _agent_config(), base_dir=tmp_path, @@ -165,10 +173,52 @@ async def test_resolve_session_opens_session_when_id_is_absent( fabric=cast(Any, fabric), ) - session = await manager.resolve_session(None) + result = await manager.invoke_once(FabricInvocationRequest(input="hello")) - assert session.runtime is runtime - assert len(fabric.start_calls) == 1 + assert result.response == "hello" + request, captured_fabric = captured[0] + assert request.fabric_config is fabric_config + assert request.input == "hello" + assert request.caller_context == {} + assert captured_fabric is fabric + assert await registry.count() == 0 + + +@pytest.mark.asyncio +async def test_stream_once_uses_relay_enabled_unregistered_runtime( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + fabric_config = _FakeFabricConfig() + monkeypatch.setattr(session_manager, "translate_agent_config", lambda config: fabric_config) + captured: list[tuple[Any, Any]] = [] + expected_stream = object() + + @asynccontextmanager + async def stream_fabric_agent_once(request: Any, *, fabric: Any) -> Any: + captured.append((request, fabric)) + yield expected_stream + + monkeypatch.setattr(session_manager, "stream_fabric_agent_once", stream_fabric_agent_once) + registry = FabricSessionRegistry() + fabric = object() + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + fabric=cast(Any, fabric), + ) + + async with manager.stream_once(FabricInvocationRequest(input="hello")) as stream: + assert stream is expected_stream + + request, captured_fabric = captured[0] + assert request.fabric_config is not fabric_config + assert request.fabric_config.copied is True + assert request.fabric_config.relay_enabled is True + assert request.input == "hello" + assert captured_fabric is fabric + assert await registry.count() == 0 @pytest.mark.asyncio From 10213326dbae5e0383a236470051f71665a38f96 Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Wed, 2 Sep 2026 16:38:40 -0500 Subject: [PATCH 2/6] implementing the respectinve invocactions and preserving compat Signed-off-by: Manjesh Mogallapalli --- .../src/nemo_agents_plugin/fabric/runtime.py | 65 +++++---- .../src/nemo_agents_plugin/fabric/server.py | 37 +++-- .../fabric/session_manager.py | 10 +- .../tests/unit/test_fabric_runtime.py | 64 ++++++++- .../tests/unit/test_fabric_server.py | 89 +++++++++++- .../tests/unit/test_fabric_session_manager.py | 132 ++++++++++++++++++ 6 files changed, 351 insertions(+), 46 deletions(-) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py index d26d9a5f8b..a2b2c6d83d 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py @@ -132,6 +132,10 @@ class FabricRuntimeExecutionError(RuntimeError): """Raised when Fabric cannot return a normalized runtime result.""" +class FabricRuntimeStartError(FabricRuntimeExecutionError): + """Raised when Fabric cannot start a runtime.""" + + class FabricRuntimeTimeoutError(FabricRuntimeExecutionError): """Raised when a Fabric runtime invocation times out.""" @@ -186,19 +190,25 @@ async def run_fabric_agent_once( """Start an ephemeral Fabric runtime, invoke it once, and stop it.""" fabric_client = fabric or Fabric() + runtime = await _start_one_shot_runtime(request, fabric=fabric_client) + try: - result = await asyncio.wait_for( - _invoke_fabric_agent_once(request, fabric=fabric_client), - timeout=request.timeout_seconds, - ) - except TimeoutError as error: - raise FabricRuntimeTimeoutError( - _timeout_error_message(request.timeout_seconds), - ) from error + async with runtime: + try: + result = await asyncio.wait_for( + runtime.invoke(request=_with_platform_invocation_context(request)), + timeout=request.timeout_seconds, + ) + except TimeoutError as error: + raise FabricRuntimeTimeoutError( + _timeout_error_message(request.timeout_seconds), + ) from error + except FabricError as error: + raise FabricRuntimeExecutionError( + f"Fabric runtime invocation failed: {error}", + ) from error except FabricError as error: - raise FabricRuntimeExecutionError( - f"Fabric runtime invocation failed: {error}", - ) from error + raise FabricRuntimeExecutionError(f"Fabric runtime cleanup failed: {error}") from error return _normalize_fabric_run_result(result) @@ -212,13 +222,10 @@ async def stream_fabric_agent_once( """Start an ephemeral Fabric runtime and keep it alive for one stream.""" fabric_client = fabric or Fabric() + runtime = await _start_one_shot_runtime(request, fabric=fabric_client, streaming=True) + try: - async with await fabric_client.start_runtime( - request.fabric_config, - base_dir=request.base_dir, - overrides=request.overrides, - streaming=True, - ) as runtime: + async with runtime: yield stream_fabric_runtime( runtime, FabricInvocationRequest( @@ -229,22 +236,24 @@ async def stream_fabric_agent_once( ), ) except FabricError as error: - raise FabricRuntimeExecutionError( - f"Fabric runtime streaming failed: {error}", - ) from error + raise FabricRuntimeExecutionError(f"Fabric runtime cleanup failed: {error}") from error -async def _invoke_fabric_agent_once( +async def _start_one_shot_runtime( request: FabricOneShotRequest, *, fabric: Any, -) -> RunResult: - async with await fabric.start_runtime( - request.fabric_config, - base_dir=request.base_dir, - overrides=request.overrides, - ) as runtime: - return await runtime.invoke(request=_with_platform_invocation_context(request)) + streaming: bool = False, +) -> Runtime: + try: + return await fabric.start_runtime( + request.fabric_config, + base_dir=request.base_dir, + overrides=request.overrides, + streaming=streaming, + ) + except FabricError as error: + raise FabricRuntimeStartError(f"Fabric runtime startup failed: {error}") from error def _with_platform_invocation_context(request: FabricInvocationRequest | FabricOneShotRequest) -> RunRequest: diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py index 46e7f10101..37033a360a 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py @@ -26,6 +26,7 @@ FabricInvocationRequest, FabricRuntimeExecutionError, FabricRuntimeResult, + FabricRuntimeStartError, FabricRuntimeStream, FabricRuntimeTimeoutError, ) @@ -50,6 +51,7 @@ openai_chat_completion_error_sse, ) from nemo_agents_plugin.session_protocol import SESSION_ID_HEADER +from starlette.types import Receive, Scope, Send logger = logging.getLogger(__name__) @@ -140,7 +142,7 @@ def _iter_streaming_chat_completion( *, completion_id: str, model: str, -) -> AsyncIterator[str]: +) -> _StreamingChatCompletionIterator: return _StreamingChatCompletionIterator( stream_context, fabric_stream, @@ -212,6 +214,20 @@ async def _iter_events(self) -> AsyncGenerator[str, None]: yield openai_chat_completion_error_sse(error) +class _FabricStreamingResponse(StreamingResponse): + """Close the Fabric stream when response delivery ends or disconnects.""" + + def __init__(self, iterator: _StreamingChatCompletionIterator, **kwargs: Any) -> None: + super().__init__(iterator, **kwargs) + self._iterator = iterator + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + try: + await super().__call__(scope, receive, send) + finally: + await self._iterator.aclose() + + async def _close_interrupted_stream(fabric_stream: FabricRuntimeStream) -> None: try: await fabric_stream.aclose() @@ -328,7 +344,7 @@ async def chat_completions( stream_context = stream(invocation_request) try: fabric_stream = await stream_context.__aenter__() - except FabricSessionStartError as error: + except FabricRuntimeStartError as error: raise HTTPException(status_code=503, detail=str(error), headers=response_headers) from error except FabricSessionNotFoundError as error: raise HTTPException( @@ -342,20 +358,21 @@ async def chat_completions( detail=str(error), headers=response_headers, ) from error - return StreamingResponse( - _iter_streaming_chat_completion( - stream_context, - fabric_stream, - completion_id=f"chatcmpl-{uuid.uuid4().hex}", - model=_request_model_name(request), - ), + iterator = _iter_streaming_chat_completion( + stream_context, + fabric_stream, + completion_id=f"chatcmpl-{uuid.uuid4().hex}", + model=_request_model_name(request), + ) + return _FabricStreamingResponse( + iterator, media_type="text/event-stream", headers=response_headers, ) try: result = await invoke(invocation_request) - except FabricSessionStartError as error: + except FabricRuntimeStartError as error: raise HTTPException(status_code=503, detail=str(error), headers=response_headers) from error except FabricSessionNotFoundError as error: raise HTTPException( diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py index bfe9f16759..35a95e128f 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py @@ -18,6 +18,7 @@ FabricInvocationRequest, FabricOneShotRequest, FabricRuntimeResult, + FabricRuntimeStartError, FabricRuntimeStream, invoke_fabric_runtime, run_fabric_agent_once, @@ -42,7 +43,7 @@ DEFAULT_SESSION_CLEANUP_INTERVAL_SECONDS = 5 * 60 -class FabricSessionStartError(RuntimeError): +class FabricSessionStartError(FabricRuntimeStartError): """Raised when a Fabric runtime cannot be started for a Platform session.""" @@ -85,7 +86,10 @@ def __init__( async def open_session(self, *, session_id: str | None = None) -> FabricRuntimeSession: """Materialize a Fabric config, start its runtime, and register the session.""" - fabric_config = await self._materialize_fabric_config(streaming=True) + try: + fabric_config = await self._materialize_fabric_config(streaming=True) + except FabricRuntimeStartError as error: + raise FabricSessionStartError(str(error)) from error fabric = self._fabric or Fabric() try: runtime = await fabric.start_runtime( @@ -252,7 +256,7 @@ async def _materialize_fabric_config(self, *, streaming: bool) -> FabricConfig: try: fabric_config = translate_agent_config(self._agent_config) except FabricTranslationError as error: - raise FabricSessionStartError(f"Fabric config translation failed: {error}") from error + raise FabricRuntimeStartError(f"Fabric config translation failed: {error}") from error await asyncio.to_thread(ensure_local_workspace_dir, self._agent_config, self._base_dir) return _prepare_serving_fabric_config(fabric_config) if streaming else fabric_config diff --git a/plugins/nemo-agents/tests/unit/test_fabric_runtime.py b/plugins/nemo-agents/tests/unit/test_fabric_runtime.py index c632ceeb38..9fd59e3175 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_runtime.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_runtime.py @@ -15,6 +15,7 @@ FabricInvocationRequest, FabricOneShotRequest, FabricRuntimeExecutionError, + FabricRuntimeStartError, FabricRuntimeTimeoutError, invoke_fabric_runtime, run_fabric_agent_once, @@ -59,12 +60,15 @@ def __init__( result: Any | None = None, invoke_error: Exception | None = None, invoke_delay: float = 0.0, + exit_error: Exception | None = None, ) -> None: self.result = result if result is not None else _FakeRunResult() self.invoke_error = invoke_error self.invoke_delay = invoke_delay + self.exit_error = exit_error self.entered = False self.exited = False + self.exit_calls = 0 self.runtime_id = "runtime-1" self.invoke_requests: list[Any] = [] self.invoke_stream_requests: list[Any] = [] @@ -81,6 +85,9 @@ async def __aexit__( traceback: object, ) -> None: self.exited = True + self.exit_calls += 1 + if self.exit_error is not None: + raise self.exit_error async def invoke(self, *, request: Any) -> Any: self.invoke_requests.append(request) @@ -190,6 +197,7 @@ async def test_starts_invokes_and_cleans_up_ephemeral_runtime(self) -> None: ] assert fake_runtime.entered is True assert fake_runtime.exited is True + assert fake_runtime.exit_calls == 1 fabric_request = fake_runtime.invoke_requests[0] assert fabric_request.input == {"prompt": "hi"} assert fabric_request.request_id == "platform-request-1" @@ -213,6 +221,7 @@ async def test_wraps_timeout(self) -> None: await run_fabric_agent_once(request, fabric=fake_fabric) assert fake_runtime.exited is True + assert fake_runtime.exit_calls == 1 async def test_wraps_runtime_timeout_without_configured_deadline(self) -> None: timeout_error = TimeoutError("adapter timed out") @@ -228,16 +237,44 @@ async def test_wraps_runtime_timeout_without_configured_deadline(self) -> None: assert exc_info.value.__cause__ is timeout_error - async def test_wraps_fabric_lifecycle_errors(self) -> None: + async def test_wraps_runtime_start_errors(self) -> None: fake_fabric = _FakeFabric(start_error=fabric_runtime.FabricError("native unavailable")) request = FabricOneShotRequest( fabric_config=cast(FabricConfig, object()), base_dir=Path("/tmp/agent"), ) - with pytest.raises(FabricRuntimeExecutionError, match="Fabric runtime invocation failed: native unavailable"): + with pytest.raises(FabricRuntimeStartError, match="Fabric runtime startup failed: native unavailable"): await run_fabric_agent_once(request, fabric=fake_fabric) + async def test_cleans_up_once_after_cancellation(self) -> None: + fake_runtime = _FakeRuntime(invoke_delay=60) + request = FabricOneShotRequest( + fabric_config=cast(FabricConfig, object()), + base_dir=Path("/tmp/agent"), + ) + invocation = asyncio.create_task(run_fabric_agent_once(request, fabric=_FakeFabric(runtime=fake_runtime))) + while not fake_runtime.invoke_requests: + await asyncio.sleep(0) + + invocation.cancel() + with pytest.raises(asyncio.CancelledError): + await invocation + + assert fake_runtime.exit_calls == 1 + + async def test_maps_cleanup_errors_after_successful_invocation(self) -> None: + fake_runtime = _FakeRuntime(exit_error=fabric_runtime.FabricError("stop failed")) + request = FabricOneShotRequest( + fabric_config=cast(FabricConfig, object()), + base_dir=Path("/tmp/agent"), + ) + + with pytest.raises(FabricRuntimeExecutionError, match="Fabric runtime cleanup failed: stop failed"): + await run_fabric_agent_once(request, fabric=_FakeFabric(runtime=fake_runtime)) + + assert fake_runtime.exit_calls == 1 + async def test_failed_run_result_is_returned_as_normalized_result(self) -> None: failed_result = _FakeRunResult( status="failed", @@ -302,6 +339,7 @@ async def test_keeps_ephemeral_runtime_alive_for_stream_context(self) -> None: result = await stream.result() assert fake_runtime.exited is True + assert fake_runtime.exit_calls == 1 assert fake_fabric.start_calls == [ { "base_dir": Path("/tmp/agent"), @@ -328,6 +366,28 @@ async def test_cleans_up_runtime_when_stream_start_fails(self) -> None: pass assert fake_runtime.exited is True + assert fake_runtime.exit_calls == 1 + + async def test_cleans_up_once_after_cancellation(self) -> None: + fake_runtime = _FakeRuntime() + request = FabricOneShotRequest( + fabric_config=cast(FabricConfig, object()), + base_dir=Path("/tmp/agent"), + ) + stream_started = asyncio.Event() + + async def consume() -> None: + async with stream_fabric_agent_once(request, fabric=_FakeFabric(runtime=fake_runtime)): + stream_started.set() + await asyncio.Event().wait() + + invocation = asyncio.create_task(consume()) + await stream_started.wait() + invocation.cancel() + with pytest.raises(asyncio.CancelledError): + await invocation + + assert fake_runtime.exit_calls == 1 @pytest.mark.asyncio diff --git a/plugins/nemo-agents/tests/unit/test_fabric_server.py b/plugins/nemo-agents/tests/unit/test_fabric_server.py index a5e0455568..c8de1cb42f 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_server.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_server.py @@ -21,6 +21,7 @@ from nemo_agents_plugin.fabric.runtime import ( FabricRuntimeExecutionError, FabricRuntimeResult, + FabricRuntimeStartError, FabricRuntimeStream, FabricRuntimeTimeoutError, ) @@ -32,6 +33,7 @@ FabricSessionStopError, ) from nemo_agents_plugin.fabric.session_registry import FabricSessionNotFoundError, FabricSessionRegistry +from starlette.requests import ClientDisconnect class _FakeStreamContext: @@ -382,6 +384,7 @@ async def invoke_session(session: Any, request: Any) -> FabricRuntimeResult: @pytest.mark.parametrize( ("error", "status_code"), [ + (FabricRuntimeStartError("startup failed"), 503), (FabricRuntimeTimeoutError("timed out"), 504), (FabricRuntimeExecutionError("invoke failed"), 502), ], @@ -410,6 +413,33 @@ async def invoke_once(request: Any) -> FabricRuntimeResult: assert response.json() == {"detail": str(error)} +def test_session_chat_completion_runtime_error_returns_session_header( + tmp_path: Path, + mock_validate_agent_config: list[tuple[AgentConfig, Path]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + app = create_fabric_serving_app(_write_agent_config(tmp_path)) + + async def resolve_session(session_id: str) -> Any: + return SimpleNamespace(session_id=session_id, runtime=object()) + + async def invoke_session(session: Any, request: Any) -> FabricRuntimeResult: + raise FabricRuntimeTimeoutError("timed out") + + with TestClient(app) as client: + monkeypatch.setattr(app.state.session_manager, "resolve_session", resolve_session) + monkeypatch.setattr(app.state.session_manager, "invoke_session", invoke_session) + response = client.post( + "/v1/chat/completions", + headers={SESSION_ID_HEADER: "session-1"}, + json={"messages": [{"role": "user", "content": "hello"}]}, + ) + + assert response.status_code == 504 + assert response.headers[SESSION_ID_HEADER] == "session-1" + assert response.json() == {"detail": "timed out"} + + @pytest.mark.parametrize( ("error", "status_code"), [ @@ -442,15 +472,24 @@ async def resolve_session(session_id: str) -> Any: assert response.json() == {"detail": str(error)} +@pytest.mark.parametrize( + ("error", "status_code"), + [ + (FabricRuntimeStartError("startup failed"), 503), + (FabricRuntimeExecutionError("stream failed"), 502), + ], +) def test_streaming_chat_completion_without_session_maps_stream_start_errors( tmp_path: Path, mock_validate_agent_config: list[tuple[AgentConfig, Path]], monkeypatch: pytest.MonkeyPatch, + error: Exception, + status_code: int, ) -> None: app = create_fabric_serving_app(_write_agent_config(tmp_path)) def stream_once(request: Any) -> _FakeStreamContext: - return _FakeStreamContext(enter_error=FabricRuntimeExecutionError("stream failed")) + return _FakeStreamContext(enter_error=error) with TestClient(app) as client: monkeypatch.setattr(app.state.session_manager, "stream_once", stream_once) @@ -459,9 +498,9 @@ def stream_once(request: Any) -> _FakeStreamContext: json={"messages": [{"role": "user", "content": "hello"}], "stream": True}, ) - assert response.status_code == 502 + assert response.status_code == status_code assert SESSION_ID_HEADER not in response.headers - assert response.json() == {"detail": "stream failed"} + assert response.json() == {"detail": str(error)} def test_streaming_chat_completion_without_session_uses_one_shot_stream( @@ -610,6 +649,33 @@ async def test_streaming_chat_completion_closes_stream_before_first_event() -> N assert stream_context.exit_calls == 1 +@pytest.mark.asyncio +async def test_streaming_response_closes_stream_after_client_disconnect() -> None: + fabric_stream = _FakeFabricStream([{"data": {"choices": [{"delta": {"content": "partial"}}]}}]) + stream_context = _FakeStreamContext(fabric_stream) + iterator = server._iter_streaming_chat_completion( + stream_context, + cast(FabricRuntimeStream, fabric_stream), + completion_id="chatcmpl-test", + model="test-model", + ) + response = server._FabricStreamingResponse(iterator, media_type="text/event-stream") + + async def receive() -> Any: + return {"type": "http.disconnect"} + + async def send(message: dict[str, Any]) -> None: + if message["type"] == "http.response.body": + raise OSError("client disconnected") + + scope = cast(Any, {"type": "http", "asgi": {"spec_version": "2.4"}}) + with pytest.raises(ClientDisconnect): + await response(scope, receive, cast(Any, send)) + + assert fabric_stream.aclose_calls == 1 + assert stream_context.exit_calls == 1 + + def test_chat_completion_maps_failed_run_result( tmp_path: Path, mock_validate_agent_config: list[tuple[AgentConfig, Path]], @@ -669,6 +735,23 @@ def test_chat_completion_request_serializes_full_transcript() -> None: assert invocation_request.caller_context == {"session_id": "session-1"} +def test_stateless_chat_completion_request_serializes_full_transcript_without_session_context() -> None: + request = ChatCompletionRequest.model_validate( + { + "messages": [ + {"role": "system", "content": "Be concise."}, + {"role": "assistant", "content": "How can I help?"}, + {"role": "user", "content": "Say hello."}, + ] + } + ) + + invocation_request = server._to_fabric_invocation_request(request, session_id=None) + + assert invocation_request.input == ("system: Be concise.\n\nassistant: How can I help?\n\nuser: Say hello.") + assert invocation_request.caller_context == {} + + def test_close_session_stops_registered_runtime( tmp_path: Path, mock_validate_agent_config: list[tuple[AgentConfig, Path]], diff --git a/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py b/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py index 551b9a37c1..e06cc16dd1 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_session_manager.py @@ -875,6 +875,138 @@ async def invoke_fabric_runtime(runtime: Any, request: FabricInvocationRequest) assert result.response == "next" +@pytest.mark.asyncio +async def test_invoke_once_limits_runtime_construction( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(session_manager, "translate_agent_config", lambda config: _FakeFabricConfig()) + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=FabricSessionRegistry(), + max_concurrent_invocations=1, + ) + first_started = asyncio.Event() + release_first = asyncio.Event() + invocation_order: list[str] = [] + + async def run_fabric_agent_once(request: Any, *, fabric: Any) -> FabricRuntimeResult: + invocation_order.append(request.input) + if request.input == "first": + first_started.set() + await release_first.wait() + return FabricRuntimeResult(status="succeeded", response=request.input) + + monkeypatch.setattr(session_manager, "run_fabric_agent_once", run_fabric_agent_once) + first = asyncio.create_task(manager.invoke_once(FabricInvocationRequest(input="first"))) + await first_started.wait() + second = asyncio.create_task(manager.invoke_once(FabricInvocationRequest(input="second"))) + await asyncio.sleep(0) + + assert invocation_order == ["first"] + + release_first.set() + results = await asyncio.gather(first, second) + + assert invocation_order == ["first", "second"] + assert [result.response for result in results] == ["first", "second"] + + +@pytest.mark.asyncio +async def test_one_shot_and_session_invocations_share_concurrency_limit( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(session_manager, "translate_agent_config", lambda config: _FakeFabricConfig()) + registry = FabricSessionRegistry() + session = await registry.register(cast(Any, _FakeRuntime()), session_id="session-1") + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=registry, + max_concurrent_invocations=1, + ) + session_started = asyncio.Event() + release_session = asyncio.Event() + invocation_order: list[str] = [] + + async def invoke_fabric_runtime(runtime: Any, request: FabricInvocationRequest) -> FabricRuntimeResult: + invocation_order.append(request.input) + session_started.set() + await release_session.wait() + return FabricRuntimeResult(status="succeeded", response=request.input) + + async def run_fabric_agent_once(request: Any, *, fabric: Any) -> FabricRuntimeResult: + invocation_order.append(request.input) + return FabricRuntimeResult(status="succeeded", response=request.input) + + monkeypatch.setattr(session_manager, "invoke_fabric_runtime", invoke_fabric_runtime) + monkeypatch.setattr(session_manager, "run_fabric_agent_once", run_fabric_agent_once) + stateful = asyncio.create_task(manager.invoke_session(session, FabricInvocationRequest(input="stateful"))) + await session_started.wait() + one_shot = asyncio.create_task(manager.invoke_once(FabricInvocationRequest(input="one-shot"))) + await asyncio.sleep(0) + + assert invocation_order == ["stateful"] + + release_session.set() + await asyncio.gather(stateful, one_shot) + + assert invocation_order == ["stateful", "one-shot"] + + +@pytest.mark.asyncio +async def test_stream_once_holds_concurrency_limit_until_runtime_cleanup( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(session_manager, "translate_agent_config", lambda config: _FakeFabricConfig()) + manager = FabricSessionManager( + _agent_config(), + base_dir=tmp_path, + session_registry=FabricSessionRegistry(), + max_concurrent_invocations=1, + ) + stream_started = asyncio.Event() + release_stream = asyncio.Event() + stream_cleaned_up = asyncio.Event() + invocation_order: list[str] = [] + + @asynccontextmanager + async def stream_fabric_agent_once(request: Any, *, fabric: Any) -> Any: + invocation_order.append(request.input) + try: + yield object() + finally: + stream_cleaned_up.set() + + async def run_fabric_agent_once(request: Any, *, fabric: Any) -> FabricRuntimeResult: + assert stream_cleaned_up.is_set() + invocation_order.append(request.input) + return FabricRuntimeResult(status="succeeded", response=request.input) + + monkeypatch.setattr(session_manager, "stream_fabric_agent_once", stream_fabric_agent_once) + monkeypatch.setattr(session_manager, "run_fabric_agent_once", run_fabric_agent_once) + + async def consume_stream() -> None: + async with manager.stream_once(FabricInvocationRequest(input="stream")): + stream_started.set() + await release_stream.wait() + + streaming = asyncio.create_task(consume_stream()) + await stream_started.wait() + one_shot = asyncio.create_task(manager.invoke_once(FabricInvocationRequest(input="one-shot"))) + await asyncio.sleep(0) + + assert invocation_order == ["stream"] + + release_stream.set() + await asyncio.gather(streaming, one_shot) + + assert invocation_order == ["stream", "one-shot"] + + @pytest.mark.asyncio async def test_zero_concurrency_limit_allows_parallel_sessions( tmp_path: Path, From dfd079b42683ee7f75ddda1dd0cb44a599df1062 Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Wed, 2 Sep 2026 16:44:18 -0500 Subject: [PATCH 3/6] remaining tests and docs Signed-off-by: Manjesh Mogallapalli --- plugins/nemo-agents/README.md | 5 ++ .../tests/unit/test_fabric_runtime.py | 83 ++++++++++++++++++- .../tests/unit/test_fabric_server.py | 12 ++- 3 files changed, 95 insertions(+), 5 deletions(-) diff --git a/plugins/nemo-agents/README.md b/plugins/nemo-agents/README.md index 19e35ef098..ab82e1e0d7 100644 --- a/plugins/nemo-agents/README.md +++ b/plugins/nemo-agents/README.md @@ -781,6 +781,11 @@ http://127.0.0.1:8080/apis/agents/v2/workspaces/default/agents/react-agent/-/v1/ You can call it directly with any OpenAI-compatible client using the same path. +Requests without ``X-Nemo-Session-Id`` use a one-shot Fabric runtime that is +stopped when the response or response stream completes. To retain runtime +context across turns, send a stable session ID in that header; the registered +runtime then follows the Platform session lifecycle. + The agent is still running — continue to the [Evaluation](#evaluation) section below, or see [Cleanup](#cleanup-optional) to tear everything down. diff --git a/plugins/nemo-agents/tests/unit/test_fabric_runtime.py b/plugins/nemo-agents/tests/unit/test_fabric_runtime.py index 9fd59e3175..6a745d499e 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_runtime.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_runtime.py @@ -61,6 +61,7 @@ def __init__( invoke_error: Exception | None = None, invoke_delay: float = 0.0, exit_error: Exception | None = None, + stream: _FakeInvokeStream | None = None, ) -> None: self.result = result if result is not None else _FakeRunResult() self.invoke_error = invoke_error @@ -72,7 +73,7 @@ def __init__( self.runtime_id = "runtime-1" self.invoke_requests: list[Any] = [] self.invoke_stream_requests: list[Any] = [] - self.stream: _FakeInvokeStream | None = None + self.stream = stream async def __aenter__(self) -> "_FakeRuntime": self.entered = True @@ -101,7 +102,7 @@ def invoke_stream(self, *, request: Any) -> "_FakeInvokeStream": self.invoke_stream_requests.append(request) if self.invoke_error is not None: raise self.invoke_error - self.stream = _FakeInvokeStream(result=self.result) + self.stream = self.stream or _FakeInvokeStream(result=self.result) return self.stream @@ -171,6 +172,21 @@ async def start_runtime( return self.runtime +class _SequenceFabric: + def __init__(self, runtimes: list[_FakeRuntime]) -> None: + self._runtimes = iter(runtimes) + + async def start_runtime( + self, + fabric_config: Any, + *, + base_dir: Path | str, + overrides: dict[str, Any] | None = None, + streaming: bool = False, + ) -> _FakeRuntime: + return next(self._runtimes) + + @pytest.mark.asyncio class TestRunFabricAgentOnce: async def test_starts_invokes_and_cleans_up_ephemeral_runtime(self) -> None: @@ -247,6 +263,35 @@ async def test_wraps_runtime_start_errors(self) -> None: with pytest.raises(FabricRuntimeStartError, match="Fabric runtime startup failed: native unavailable"): await run_fabric_agent_once(request, fabric=fake_fabric) + async def test_uses_a_separate_runtime_for_each_invocation(self) -> None: + first_runtime = _FakeRuntime() + second_runtime = _FakeRuntime() + fabric = _SequenceFabric([first_runtime, second_runtime]) + request = FabricOneShotRequest( + fabric_config=cast(FabricConfig, object()), + base_dir=Path("/tmp/agent"), + ) + + await run_fabric_agent_once(request, fabric=fabric) + await run_fabric_agent_once(request, fabric=fabric) + + assert len(first_runtime.invoke_requests) == 1 + assert len(second_runtime.invoke_requests) == 1 + assert first_runtime.exit_calls == 1 + assert second_runtime.exit_calls == 1 + + async def test_cleans_up_once_after_invocation_error(self) -> None: + fake_runtime = _FakeRuntime(invoke_error=fabric_runtime.FabricError("invoke failed")) + request = FabricOneShotRequest( + fabric_config=cast(FabricConfig, object()), + base_dir=Path("/tmp/agent"), + ) + + with pytest.raises(FabricRuntimeExecutionError, match="Fabric runtime invocation failed: invoke failed"): + await run_fabric_agent_once(request, fabric=_FakeFabric(runtime=fake_runtime)) + + assert fake_runtime.exit_calls == 1 + async def test_cleans_up_once_after_cancellation(self) -> None: fake_runtime = _FakeRuntime(invoke_delay=60) request = FabricOneShotRequest( @@ -291,6 +336,7 @@ async def test_failed_run_result_is_returned_as_normalized_result(self) -> None: assert result.status == "failed" assert result.error == {"stage": "invoke", "message": "adapter failed"} + assert fake_fabric.runtime.exit_calls == 1 async def test_normalizes_fabric_mapping_fields_to_plain_values(self) -> None: fake_result = _FakeRunResult( @@ -368,6 +414,39 @@ async def test_cleans_up_runtime_when_stream_start_fails(self) -> None: assert fake_runtime.exited is True assert fake_runtime.exit_calls == 1 + async def test_cleans_up_runtime_after_stream_result_error(self) -> None: + stream = _FakeInvokeStream(result_error=fabric_runtime.FabricError("result failed")) + fake_runtime = _FakeRuntime(stream=stream) + request = FabricOneShotRequest( + fabric_config=cast(FabricConfig, object()), + base_dir=Path("/tmp/agent"), + ) + + with pytest.raises(FabricRuntimeExecutionError, match="Fabric runtime streaming failed: result failed"): + async with stream_fabric_agent_once(request, fabric=_FakeFabric(runtime=fake_runtime)) as runtime_stream: + await runtime_stream.result() + + assert fake_runtime.exit_calls == 1 + + async def test_cleans_up_runtime_after_stream_result_timeout(self) -> None: + class _SlowInvokeStream(_FakeInvokeStream): + async def result(self) -> Any: + await asyncio.sleep(1.0) + return await super().result() + + fake_runtime = _FakeRuntime(stream=_SlowInvokeStream()) + request = FabricOneShotRequest( + fabric_config=cast(FabricConfig, object()), + base_dir=Path("/tmp/agent"), + timeout_seconds=0.01, + ) + + with pytest.raises(FabricRuntimeTimeoutError, match="timed out after 0.01s"): + async with stream_fabric_agent_once(request, fabric=_FakeFabric(runtime=fake_runtime)) as runtime_stream: + await runtime_stream.result() + + assert fake_runtime.exit_calls == 1 + async def test_cleans_up_once_after_cancellation(self) -> None: fake_runtime = _FakeRuntime() request = FabricOneShotRequest( diff --git a/plugins/nemo-agents/tests/unit/test_fabric_server.py b/plugins/nemo-agents/tests/unit/test_fabric_server.py index c8de1cb42f..6d9591fc4c 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_server.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_server.py @@ -330,9 +330,15 @@ async def invoke_once(request: Any) -> FabricRuntimeResult: "/v1/chat/completions", json={"messages": [{"role": "user", "content": "hello"}]}, ) + repeated_response = client.post( + "/v1/chat/completions", + json={"messages": [{"role": "user", "content": "again"}]}, + ) assert response.status_code == 200 + assert repeated_response.status_code == 200 assert SESSION_ID_HEADER not in response.headers + assert SESSION_ID_HEADER not in repeated_response.headers assert response.json() == { "id": "invocation-1", "object": "chat.completion", @@ -346,9 +352,9 @@ async def invoke_once(request: Any) -> FabricRuntimeResult: ], "usage": {"total_tokens": 3}, } - invocation_request = invocation_calls[0] - assert invocation_request.input == "hello" - assert invocation_request.caller_context == {} + assert [request.input for request in invocation_calls] == ["hello", "again"] + assert all(request.caller_context == {} for request in invocation_calls) + assert asyncio.run(app.state.session_registry.count()) == 0 def test_chat_completion_with_session_id_reuses_session( From 4a2459495b0c81b0e73583e017b85e46906eba4b Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Wed, 2 Sep 2026 17:17:42 -0500 Subject: [PATCH 4/6] self review Signed-off-by: Manjesh Mogallapalli --- .../src/nemo_agents_plugin/fabric/runtime.py | 33 +++++- .../src/nemo_agents_plugin/fabric/server.py | 23 +++- .../tests/unit/test_fabric_runtime.py | 63 +++++++++++ .../tests/unit/test_fabric_server.py | 100 +++++++++++++++++- 4 files changed, 211 insertions(+), 8 deletions(-) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py index a2b2c6d83d..2a7422718e 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py @@ -189,15 +189,26 @@ async def run_fabric_agent_once( ) -> FabricRuntimeResult: """Start an ephemeral Fabric runtime, invoke it once, and stop it.""" fabric_client = fabric or Fabric() + deadline = None if request.timeout_seconds is None else asyncio.get_running_loop().time() + request.timeout_seconds - runtime = await _start_one_shot_runtime(request, fabric=fabric_client) + try: + runtime = await asyncio.wait_for( + _start_one_shot_runtime(request, fabric=fabric_client), + timeout=_remaining_timeout(deadline), + ) + except TimeoutError as error: + raise FabricRuntimeTimeoutError( + _timeout_error_message(request.timeout_seconds), + ) from error + runtime_entered = False try: async with runtime: + runtime_entered = True try: result = await asyncio.wait_for( runtime.invoke(request=_with_platform_invocation_context(request)), - timeout=request.timeout_seconds, + timeout=_remaining_timeout(deadline), ) except TimeoutError as error: raise FabricRuntimeTimeoutError( @@ -208,11 +219,17 @@ async def run_fabric_agent_once( f"Fabric runtime invocation failed: {error}", ) from error except FabricError as error: - raise FabricRuntimeExecutionError(f"Fabric runtime cleanup failed: {error}") from error + raise _runtime_context_error(error, entered=runtime_entered) from error return _normalize_fabric_run_result(result) +def _remaining_timeout(deadline: float | None) -> float | None: + if deadline is None: + return None + return max(0.0, deadline - asyncio.get_running_loop().time()) + + @asynccontextmanager async def stream_fabric_agent_once( request: FabricOneShotRequest, @@ -224,8 +241,10 @@ async def stream_fabric_agent_once( runtime = await _start_one_shot_runtime(request, fabric=fabric_client, streaming=True) + runtime_entered = False try: async with runtime: + runtime_entered = True yield stream_fabric_runtime( runtime, FabricInvocationRequest( @@ -236,7 +255,13 @@ async def stream_fabric_agent_once( ), ) except FabricError as error: - raise FabricRuntimeExecutionError(f"Fabric runtime cleanup failed: {error}") from error + raise _runtime_context_error(error, entered=runtime_entered) from error + + +def _runtime_context_error(error: FabricError, *, entered: bool) -> FabricRuntimeExecutionError: + if entered: + return FabricRuntimeExecutionError(f"Fabric runtime cleanup failed: {error}") + return FabricRuntimeStartError(f"Fabric runtime startup failed: {error}") async def _start_one_shot_runtime( diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py index 37033a360a..532793c8fb 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py @@ -55,6 +55,8 @@ logger = logging.getLogger(__name__) +_FABRIC_STREAM_CLEANUP_TIMEOUT_SECONDS = 5.0 + @dataclass(frozen=True, slots=True) class FabricServingSettings: @@ -168,13 +170,14 @@ def __init__( self._model = model self._events: AsyncGenerator[str, None] | None = None self._close_fabric_stream_on_exit = True + self._close_task: asyncio.Task[None] | None = None self._closed = False def __aiter__(self) -> AsyncIterator[str]: return self async def __anext__(self) -> str: - if self._closed: + if self._closed or self._close_task is not None: raise StopAsyncIteration if self._events is None: self._events = self._iter_events() @@ -190,12 +193,17 @@ async def __anext__(self) -> str: async def aclose(self) -> None: if self._closed: return - self._closed = True + if self._close_task is None: + self._close_task = asyncio.create_task(self._cleanup()) + await asyncio.shield(self._close_task) + + async def _cleanup(self) -> None: if self._events is not None: await self._events.aclose() if self._close_fabric_stream_on_exit: await _close_interrupted_stream(self._fabric_stream) await self._stream_context.__aexit__(None, None, None) + self._closed = True async def _iter_events(self) -> AsyncGenerator[str, None]: try: @@ -225,7 +233,16 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: try: await super().__call__(scope, receive, send) finally: - await self._iterator.aclose() + try: + await asyncio.wait_for( + asyncio.shield(self._iterator.aclose()), + timeout=_FABRIC_STREAM_CLEANUP_TIMEOUT_SECONDS, + ) + except TimeoutError: + logger.warning( + "Timed out waiting for Fabric stream cleanup after %gs; cleanup will continue in the background.", + _FABRIC_STREAM_CLEANUP_TIMEOUT_SECONDS, + ) async def _close_interrupted_stream(fabric_stream: FabricRuntimeStream) -> None: diff --git a/plugins/nemo-agents/tests/unit/test_fabric_runtime.py b/plugins/nemo-agents/tests/unit/test_fabric_runtime.py index 6a745d499e..48349134dd 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_runtime.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_runtime.py @@ -60,12 +60,14 @@ def __init__( result: Any | None = None, invoke_error: Exception | None = None, invoke_delay: float = 0.0, + enter_error: Exception | None = None, exit_error: Exception | None = None, stream: _FakeInvokeStream | None = None, ) -> None: self.result = result if result is not None else _FakeRunResult() self.invoke_error = invoke_error self.invoke_delay = invoke_delay + self.enter_error = enter_error self.exit_error = exit_error self.entered = False self.exited = False @@ -76,6 +78,8 @@ def __init__( self.stream = stream async def __aenter__(self) -> "_FakeRuntime": + if self.enter_error is not None: + raise self.enter_error self.entered = True return self @@ -146,9 +150,11 @@ def __init__( *, runtime: _FakeRuntime | None = None, start_error: Exception | None = None, + start_delay: float = 0.0, ) -> None: self.runtime = runtime if runtime is not None else _FakeRuntime() self.start_error = start_error + self.start_delay = start_delay self.start_calls: list[dict[str, Any]] = [] async def start_runtime( @@ -167,6 +173,8 @@ async def start_runtime( if streaming: call["streaming"] = True self.start_calls.append(call) + if self.start_delay: + await asyncio.sleep(self.start_delay) if self.start_error is not None: raise self.start_error return self.runtime @@ -239,6 +247,21 @@ async def test_wraps_timeout(self) -> None: assert fake_runtime.exited is True assert fake_runtime.exit_calls == 1 + async def test_timeout_includes_runtime_startup(self) -> None: + fake_runtime = _FakeRuntime() + fake_fabric = _FakeFabric(runtime=fake_runtime, start_delay=1.0) + request = FabricOneShotRequest( + fabric_config=cast(FabricConfig, object()), + base_dir=Path("/tmp/agent"), + timeout_seconds=0.01, + ) + + with pytest.raises(FabricRuntimeTimeoutError, match="timed out after 0.01s"): + await run_fabric_agent_once(request, fabric=fake_fabric) + + assert fake_runtime.entered is False + assert fake_runtime.exit_calls == 0 + async def test_wraps_runtime_timeout_without_configured_deadline(self) -> None: timeout_error = TimeoutError("adapter timed out") fake_runtime = _FakeRuntime(invoke_error=timeout_error) @@ -263,6 +286,19 @@ async def test_wraps_runtime_start_errors(self) -> None: with pytest.raises(FabricRuntimeStartError, match="Fabric runtime startup failed: native unavailable"): await run_fabric_agent_once(request, fabric=fake_fabric) + async def test_wraps_runtime_context_entry_errors_as_start_errors(self) -> None: + fake_runtime = _FakeRuntime(enter_error=fabric_runtime.FabricError("harness unavailable")) + request = FabricOneShotRequest( + fabric_config=cast(FabricConfig, object()), + base_dir=Path("/tmp/agent"), + ) + + with pytest.raises(FabricRuntimeStartError, match="Fabric runtime startup failed: harness unavailable"): + await run_fabric_agent_once(request, fabric=_FakeFabric(runtime=fake_runtime)) + + assert fake_runtime.entered is False + assert fake_runtime.exit_calls == 0 + async def test_uses_a_separate_runtime_for_each_invocation(self) -> None: first_runtime = _FakeRuntime() second_runtime = _FakeRuntime() @@ -414,6 +450,33 @@ async def test_cleans_up_runtime_when_stream_start_fails(self) -> None: assert fake_runtime.exited is True assert fake_runtime.exit_calls == 1 + async def test_wraps_runtime_context_entry_errors_as_start_errors(self) -> None: + fake_runtime = _FakeRuntime(enter_error=fabric_runtime.FabricError("harness unavailable")) + request = FabricOneShotRequest( + fabric_config=cast(FabricConfig, object()), + base_dir=Path("/tmp/agent"), + ) + + with pytest.raises(FabricRuntimeStartError, match="Fabric runtime startup failed: harness unavailable"): + async with stream_fabric_agent_once(request, fabric=_FakeFabric(runtime=fake_runtime)): + pass + + assert fake_runtime.entered is False + assert fake_runtime.exit_calls == 0 + + async def test_maps_runtime_context_exit_errors_as_cleanup_errors(self) -> None: + fake_runtime = _FakeRuntime(exit_error=fabric_runtime.FabricError("stop failed")) + request = FabricOneShotRequest( + fabric_config=cast(FabricConfig, object()), + base_dir=Path("/tmp/agent"), + ) + + with pytest.raises(FabricRuntimeExecutionError, match="Fabric runtime cleanup failed: stop failed"): + async with stream_fabric_agent_once(request, fabric=_FakeFabric(runtime=fake_runtime)): + pass + + assert fake_runtime.exit_calls == 1 + async def test_cleans_up_runtime_after_stream_result_error(self) -> None: stream = _FakeInvokeStream(result_error=fabric_runtime.FabricError("result failed")) fake_runtime = _FakeRuntime(stream=stream) diff --git a/plugins/nemo-agents/tests/unit/test_fabric_server.py b/plugins/nemo-agents/tests/unit/test_fabric_server.py index 6d9591fc4c..c281b1be8c 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_server.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_server.py @@ -37,10 +37,20 @@ class _FakeStreamContext: - def __init__(self, stream: Any = None, enter_error: BaseException | None = None) -> None: + def __init__( + self, + stream: Any = None, + enter_error: BaseException | None = None, + *, + exit_checkpoint: bool = False, + exit_waiter: asyncio.Event | None = None, + ) -> None: self.stream = stream self.enter_error = enter_error + self.exit_checkpoint = exit_checkpoint + self.exit_waiter = exit_waiter self.exit_calls = 0 + self.exit_completions = 0 async def __aenter__(self) -> Any: if self.enter_error is not None: @@ -49,6 +59,11 @@ async def __aenter__(self) -> Any: async def __aexit__(self, exc_type: object, exc: object, traceback: object) -> None: self.exit_calls += 1 + if self.exit_checkpoint: + await asyncio.sleep(0) + if self.exit_waiter is not None: + await self.exit_waiter.wait() + self.exit_completions += 1 class _FakeFabricStream: @@ -57,20 +72,30 @@ def __init__( records: list[dict[str, Any]] | None = None, *, result: FabricRuntimeResult | None = None, + block_after_records: bool = False, + close_checkpoint: bool = False, ) -> None: self._records = records or [] self._result = result or FabricRuntimeResult(status="succeeded", response="done") + self._block_after_records = block_after_records + self._close_checkpoint = close_checkpoint self.aclose_calls = 0 + self.aclose_completions = 0 async def records(self) -> Any: for record in self._records: yield record + if self._block_after_records: + await asyncio.Event().wait() async def result(self) -> FabricRuntimeResult: return self._result async def aclose(self) -> None: self.aclose_calls += 1 + if self._close_checkpoint: + await asyncio.sleep(0) + self.aclose_completions += 1 @pytest.fixture() @@ -682,6 +707,79 @@ async def send(message: dict[str, Any]) -> None: assert stream_context.exit_calls == 1 +@pytest.mark.asyncio +async def test_streaming_response_completes_cleanup_after_asgi_cancellation() -> None: + fabric_stream = _FakeFabricStream( + [{"data": {"choices": [{"delta": {"content": "partial"}}]}}], + block_after_records=True, + close_checkpoint=True, + ) + stream_context = _FakeStreamContext(fabric_stream, exit_checkpoint=True) + iterator = server._iter_streaming_chat_completion( + stream_context, + cast(FabricRuntimeStream, fabric_stream), + completion_id="chatcmpl-test", + model="test-model", + ) + response = server._FabricStreamingResponse(iterator, media_type="text/event-stream") + body_sent = asyncio.Event() + + async def receive() -> Any: + await body_sent.wait() + return {"type": "http.disconnect"} + + async def send(message: dict[str, Any]) -> None: + if message["type"] == "http.response.body" and message.get("body"): + body_sent.set() + + scope = cast(Any, {"type": "http", "asgi": {"spec_version": "2.3"}}) + await asyncio.wait_for(response(scope, receive, cast(Any, send)), timeout=1) + + assert fabric_stream.aclose_calls == 1 + assert fabric_stream.aclose_completions == 1 + assert stream_context.exit_calls == 1 + assert stream_context.exit_completions == 1 + + +@pytest.mark.asyncio +async def test_streaming_response_bounds_cleanup_wait_without_cancelling_cleanup( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + cleanup_release = asyncio.Event() + fabric_stream = _FakeFabricStream() + stream_context = _FakeStreamContext(fabric_stream, exit_waiter=cleanup_release) + iterator = server._iter_streaming_chat_completion( + stream_context, + cast(FabricRuntimeStream, fabric_stream), + completion_id="chatcmpl-test", + model="test-model", + ) + response = server._FabricStreamingResponse(iterator, media_type="text/event-stream") + monkeypatch.setattr(server, "_FABRIC_STREAM_CLEANUP_TIMEOUT_SECONDS", 0.01) + + async def receive() -> Any: + return {"type": "http.disconnect"} + + async def send(message: dict[str, Any]) -> None: + if message["type"] == "http.response.body": + raise OSError("client disconnected") + + scope = cast(Any, {"type": "http", "asgi": {"spec_version": "2.4"}}) + with caplog.at_level("WARNING"): + with pytest.raises(ClientDisconnect): + await response(scope, receive, cast(Any, send)) + + assert stream_context.exit_calls == 1 + assert stream_context.exit_completions == 0 + assert "cleanup will continue in the background" in caplog.text + + cleanup_release.set() + await iterator.aclose() + + assert stream_context.exit_completions == 1 + + def test_chat_completion_maps_failed_run_result( tmp_path: Path, mock_validate_agent_config: list[tuple[AgentConfig, Path]], From 12d244d17961299f0191ef24de29160686e9bf82 Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Wed, 2 Sep 2026 18:00:16 -0500 Subject: [PATCH 5/6] self review pt.2 Signed-off-by: Manjesh Mogallapalli --- .../src/nemo_agents_plugin/fabric/runtime.py | 20 +--- .../src/nemo_agents_plugin/fabric/server.py | 53 +++++++--- .../tests/unit/test_fabric_runtime.py | 43 ++++++-- .../tests/unit/test_fabric_server.py | 98 ++++++++++++++++++- 4 files changed, 176 insertions(+), 38 deletions(-) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py index 2a7422718e..256bf5f4dd 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/runtime.py @@ -189,17 +189,7 @@ async def run_fabric_agent_once( ) -> FabricRuntimeResult: """Start an ephemeral Fabric runtime, invoke it once, and stop it.""" fabric_client = fabric or Fabric() - deadline = None if request.timeout_seconds is None else asyncio.get_running_loop().time() + request.timeout_seconds - - try: - runtime = await asyncio.wait_for( - _start_one_shot_runtime(request, fabric=fabric_client), - timeout=_remaining_timeout(deadline), - ) - except TimeoutError as error: - raise FabricRuntimeTimeoutError( - _timeout_error_message(request.timeout_seconds), - ) from error + runtime = await _start_one_shot_runtime(request, fabric=fabric_client) runtime_entered = False try: @@ -208,7 +198,7 @@ async def run_fabric_agent_once( try: result = await asyncio.wait_for( runtime.invoke(request=_with_platform_invocation_context(request)), - timeout=_remaining_timeout(deadline), + timeout=request.timeout_seconds, ) except TimeoutError as error: raise FabricRuntimeTimeoutError( @@ -224,12 +214,6 @@ async def run_fabric_agent_once( return _normalize_fabric_run_result(result) -def _remaining_timeout(deadline: float | None) -> float | None: - if deadline is None: - return None - return max(0.0, deadline - asyncio.get_running_loop().time()) - - @asynccontextmanager async def stream_fabric_agent_once( request: FabricOneShotRequest, diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py index 532793c8fb..c69f75c1a9 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py @@ -171,6 +171,8 @@ def __init__( self._events: AsyncGenerator[str, None] | None = None self._close_fabric_stream_on_exit = True self._close_task: asyncio.Task[None] | None = None + self._close_deadline: float | None = None + self._close_observer_added = False self._closed = False def __aiter__(self) -> AsyncIterator[str]: @@ -186,6 +188,9 @@ async def __anext__(self) -> str: except StopAsyncIteration: await self.aclose() raise + except asyncio.CancelledError: + self._start_cleanup() + raise except BaseException: await self.aclose() raise @@ -193,9 +198,44 @@ async def __anext__(self) -> str: async def aclose(self) -> None: if self._closed: return + close_task = self._start_cleanup() + assert self._close_deadline is not None + remaining = max(0.0, self._close_deadline - asyncio.get_running_loop().time()) + try: + await asyncio.wait_for(asyncio.shield(close_task), timeout=remaining) + except TimeoutError: + self._observe_background_cleanup(close_task) + + def _start_cleanup(self) -> asyncio.Task[None]: if self._close_task is None: self._close_task = asyncio.create_task(self._cleanup()) - await asyncio.shield(self._close_task) + self._close_deadline = asyncio.get_running_loop().time() + _FABRIC_STREAM_CLEANUP_TIMEOUT_SECONDS + return self._close_task + + def _observe_background_cleanup(self, close_task: asyncio.Task[None]) -> None: + if self._close_observer_added: + return + self._close_observer_added = True + close_task.add_done_callback(self._log_background_cleanup_result) + logger.warning( + "Timed out waiting for Fabric stream cleanup after %gs; cleanup will continue in the background.", + _FABRIC_STREAM_CLEANUP_TIMEOUT_SECONDS, + ) + + def _log_background_cleanup_result(self, close_task: asyncio.Task[None]) -> None: + if close_task.cancelled(): + logger.warning( + "Background Fabric stream cleanup was cancelled for completion %s.", + self._completion_id, + ) + return + error = close_task.exception() + if error is not None: + logger.error( + "Background Fabric stream cleanup failed for completion %s.", + self._completion_id, + exc_info=(type(error), error, error.__traceback__), + ) async def _cleanup(self) -> None: if self._events is not None: @@ -233,16 +273,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: try: await super().__call__(scope, receive, send) finally: - try: - await asyncio.wait_for( - asyncio.shield(self._iterator.aclose()), - timeout=_FABRIC_STREAM_CLEANUP_TIMEOUT_SECONDS, - ) - except TimeoutError: - logger.warning( - "Timed out waiting for Fabric stream cleanup after %gs; cleanup will continue in the background.", - _FABRIC_STREAM_CLEANUP_TIMEOUT_SECONDS, - ) + await asyncio.shield(self._iterator.aclose()) async def _close_interrupted_stream(fabric_stream: FabricRuntimeStream) -> None: diff --git a/plugins/nemo-agents/tests/unit/test_fabric_runtime.py b/plugins/nemo-agents/tests/unit/test_fabric_runtime.py index 48349134dd..3c6d9163cb 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_runtime.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_runtime.py @@ -61,6 +61,7 @@ def __init__( invoke_error: Exception | None = None, invoke_delay: float = 0.0, enter_error: Exception | None = None, + enter_delay: float = 0.0, exit_error: Exception | None = None, stream: _FakeInvokeStream | None = None, ) -> None: @@ -68,6 +69,7 @@ def __init__( self.invoke_error = invoke_error self.invoke_delay = invoke_delay self.enter_error = enter_error + self.enter_delay = enter_delay self.exit_error = exit_error self.entered = False self.exited = False @@ -78,6 +80,8 @@ def __init__( self.stream = stream async def __aenter__(self) -> "_FakeRuntime": + if self.enter_delay: + await asyncio.sleep(self.enter_delay) if self.enter_error is not None: raise self.enter_error self.entered = True @@ -247,20 +251,21 @@ async def test_wraps_timeout(self) -> None: assert fake_runtime.exited is True assert fake_runtime.exit_calls == 1 - async def test_timeout_includes_runtime_startup(self) -> None: - fake_runtime = _FakeRuntime() - fake_fabric = _FakeFabric(runtime=fake_runtime, start_delay=1.0) + async def test_runtime_setup_does_not_consume_invocation_timeout(self) -> None: + fake_runtime = _FakeRuntime(enter_delay=0.06, invoke_delay=0.06) request = FabricOneShotRequest( fabric_config=cast(FabricConfig, object()), base_dir=Path("/tmp/agent"), - timeout_seconds=0.01, + timeout_seconds=0.1, ) - with pytest.raises(FabricRuntimeTimeoutError, match="timed out after 0.01s"): - await run_fabric_agent_once(request, fabric=fake_fabric) + result = await run_fabric_agent_once( + request, + fabric=_FakeFabric(runtime=fake_runtime, start_delay=0.06), + ) - assert fake_runtime.entered is False - assert fake_runtime.exit_calls == 0 + assert result.status == "succeeded" + assert fake_runtime.exit_calls == 1 async def test_wraps_runtime_timeout_without_configured_deadline(self) -> None: timeout_error = TimeoutError("adapter timed out") @@ -436,6 +441,28 @@ async def test_keeps_ephemeral_runtime_alive_for_stream_context(self) -> None: assert fabric_request.context == {"source": "server"} assert result.response == "hello" + async def test_runtime_setup_does_not_consume_result_timeout(self) -> None: + class _SlowInvokeStream(_FakeInvokeStream): + async def result(self) -> Any: + await asyncio.sleep(0.06) + return await super().result() + + fake_runtime = _FakeRuntime(stream=_SlowInvokeStream(), enter_delay=0.06) + request = FabricOneShotRequest( + fabric_config=cast(FabricConfig, object()), + base_dir=Path("/tmp/agent"), + timeout_seconds=0.1, + ) + + async with stream_fabric_agent_once( + request, + fabric=_FakeFabric(runtime=fake_runtime, start_delay=0.06), + ) as stream: + result = await stream.result() + + assert result.status == "succeeded" + assert fake_runtime.exit_calls == 1 + async def test_cleans_up_runtime_when_stream_start_fails(self) -> None: fake_runtime = _FakeRuntime(invoke_error=fabric_runtime.FabricError("stream unavailable")) request = FabricOneShotRequest( diff --git a/plugins/nemo-agents/tests/unit/test_fabric_server.py b/plugins/nemo-agents/tests/unit/test_fabric_server.py index c281b1be8c..c7e9b8f6c4 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_server.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_server.py @@ -44,11 +44,13 @@ def __init__( *, exit_checkpoint: bool = False, exit_waiter: asyncio.Event | None = None, + exit_error: BaseException | None = None, ) -> None: self.stream = stream self.enter_error = enter_error self.exit_checkpoint = exit_checkpoint self.exit_waiter = exit_waiter + self.exit_error = exit_error self.exit_calls = 0 self.exit_completions = 0 @@ -63,6 +65,8 @@ async def __aexit__(self, exc_type: object, exc: object, traceback: object) -> N await asyncio.sleep(0) if self.exit_waiter is not None: await self.exit_waiter.wait() + if self.exit_error is not None: + raise self.exit_error self.exit_completions += 1 @@ -775,11 +779,103 @@ async def send(message: dict[str, Any]) -> None: assert "cleanup will continue in the background" in caplog.text cleanup_release.set() - await iterator.aclose() + assert iterator._close_task is not None + await iterator._close_task assert stream_context.exit_completions == 1 +@pytest.mark.asyncio +async def test_streaming_response_bounds_cleanup_after_direct_task_cancellation( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + cleanup_release = asyncio.Event() + fabric_stream = _FakeFabricStream( + [{"data": {"choices": [{"delta": {"content": "partial"}}]}}], + block_after_records=True, + ) + stream_context = _FakeStreamContext(fabric_stream, exit_waiter=cleanup_release) + iterator = server._iter_streaming_chat_completion( + stream_context, + cast(FabricRuntimeStream, fabric_stream), + completion_id="chatcmpl-test", + model="test-model", + ) + response = server._FabricStreamingResponse(iterator, media_type="text/event-stream") + body_sent = asyncio.Event() + monkeypatch.setattr(server, "_FABRIC_STREAM_CLEANUP_TIMEOUT_SECONDS", 0.01) + + async def receive() -> Any: + await asyncio.Event().wait() + + async def send(message: dict[str, Any]) -> None: + if message["type"] == "http.response.body" and message.get("body"): + body_sent.set() + + scope = cast(Any, {"type": "http", "asgi": {"spec_version": "2.4"}}) + response_task = asyncio.create_task(response(scope, receive, cast(Any, send))) + await body_sent.wait() + response_task.cancel() + + with caplog.at_level("WARNING"): + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(response_task, timeout=1) + + assert stream_context.exit_calls == 1 + assert stream_context.exit_completions == 0 + assert "cleanup will continue in the background" in caplog.text + + cleanup_release.set() + assert iterator._close_task is not None + await iterator._close_task + + assert stream_context.exit_completions == 1 + + +@pytest.mark.asyncio +async def test_streaming_response_logs_cleanup_failure_after_timeout( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + cleanup_release = asyncio.Event() + cleanup_error = RuntimeError("cleanup failed") + fabric_stream = _FakeFabricStream() + stream_context = _FakeStreamContext( + fabric_stream, + exit_waiter=cleanup_release, + exit_error=cleanup_error, + ) + iterator = server._iter_streaming_chat_completion( + stream_context, + cast(FabricRuntimeStream, fabric_stream), + completion_id="chatcmpl-test", + model="test-model", + ) + response = server._FabricStreamingResponse(iterator, media_type="text/event-stream") + monkeypatch.setattr(server, "_FABRIC_STREAM_CLEANUP_TIMEOUT_SECONDS", 0.01) + + async def receive() -> Any: + return {"type": "http.disconnect"} + + async def send(message: dict[str, Any]) -> None: + if message["type"] == "http.response.body": + raise OSError("client disconnected") + + scope = cast(Any, {"type": "http", "asgi": {"spec_version": "2.4"}}) + with caplog.at_level("WARNING"): + with pytest.raises(ClientDisconnect): + await response(scope, receive, cast(Any, send)) + + cleanup_release.set() + assert iterator._close_task is not None + with pytest.raises(RuntimeError, match="cleanup failed"): + await iterator._close_task + await asyncio.sleep(0) + + assert "Background Fabric stream cleanup failed for completion chatcmpl-test." in caplog.text + + def test_chat_completion_maps_failed_run_result( tmp_path: Path, mock_validate_agent_config: list[tuple[AgentConfig, Path]], From 6ce736b9b3672419706569dc8f87dbae36a4a2de Mon Sep 17 00:00:00 2001 From: Manjesh Mogallapalli Date: Wed, 2 Sep 2026 18:12:18 -0500 Subject: [PATCH 6/6] nit Signed-off-by: Manjesh Mogallapalli --- .../src/nemo_agents_plugin/fabric/server.py | 8 +++++ .../tests/unit/test_fabric_server.py | 34 +++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py index c69f75c1a9..1b7dead6f2 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py @@ -205,6 +205,12 @@ async def aclose(self) -> None: await asyncio.wait_for(asyncio.shield(close_task), timeout=remaining) except TimeoutError: self._observe_background_cleanup(close_task) + except Exception: + self._closed = True + logger.exception( + "Fabric stream cleanup failed for completion %s.", + self._completion_id, + ) def _start_cleanup(self) -> asyncio.Task[None]: if self._close_task is None: @@ -224,6 +230,7 @@ def _observe_background_cleanup(self, close_task: asyncio.Task[None]) -> None: def _log_background_cleanup_result(self, close_task: asyncio.Task[None]) -> None: if close_task.cancelled(): + self._closed = True logger.warning( "Background Fabric stream cleanup was cancelled for completion %s.", self._completion_id, @@ -231,6 +238,7 @@ def _log_background_cleanup_result(self, close_task: asyncio.Task[None]) -> None return error = close_task.exception() if error is not None: + self._closed = True logger.error( "Background Fabric stream cleanup failed for completion %s.", self._completion_id, diff --git a/plugins/nemo-agents/tests/unit/test_fabric_server.py b/plugins/nemo-agents/tests/unit/test_fabric_server.py index c7e9b8f6c4..06dddddc5d 100644 --- a/plugins/nemo-agents/tests/unit/test_fabric_server.py +++ b/plugins/nemo-agents/tests/unit/test_fabric_server.py @@ -711,6 +711,38 @@ async def send(message: dict[str, Any]) -> None: assert stream_context.exit_calls == 1 +@pytest.mark.asyncio +async def test_streaming_response_cleanup_failure_does_not_mask_client_disconnect( + caplog: pytest.LogCaptureFixture, +) -> None: + cleanup_error = RuntimeError("cleanup failed") + fabric_stream = _FakeFabricStream() + stream_context = _FakeStreamContext(fabric_stream, exit_error=cleanup_error) + iterator = server._iter_streaming_chat_completion( + stream_context, + cast(FabricRuntimeStream, fabric_stream), + completion_id="chatcmpl-test", + model="test-model", + ) + response = server._FabricStreamingResponse(iterator, media_type="text/event-stream") + + async def receive() -> Any: + return {"type": "http.disconnect"} + + async def send(message: dict[str, Any]) -> None: + if message["type"] == "http.response.body": + raise OSError("client disconnected") + + scope = cast(Any, {"type": "http", "asgi": {"spec_version": "2.4"}}) + with caplog.at_level("ERROR"): + with pytest.raises(ClientDisconnect): + await response(scope, receive, cast(Any, send)) + + assert "Fabric stream cleanup failed for completion chatcmpl-test." in caplog.text + assert iterator._closed + await iterator.aclose() + + @pytest.mark.asyncio async def test_streaming_response_completes_cleanup_after_asgi_cancellation() -> None: fabric_stream = _FakeFabricStream( @@ -874,6 +906,8 @@ async def send(message: dict[str, Any]) -> None: await asyncio.sleep(0) assert "Background Fabric stream cleanup failed for completion chatcmpl-test." in caplog.text + assert iterator._closed + await iterator.aclose() def test_chat_completion_maps_failed_run_result(