fix(llmrails): track deferred startup and streaming followups (9/9)

nemoguardrails/rails/llm/startup/colang_flows.py

@@ -67,6 +67,21 @@ def _iter_colang_resources(resource: Traversable) -> Iterator[Traversable]:
             yield from _iter_colang_resources(child)
 
 
+def _flow_config_key(flow_config: dict) -> str | None:
+    return flow_config.get("id") or flow_config.get("name")
+
+
+def _extend_missing_flow_configs(config: RailsConfig, flow_configs: list[dict]) -> None:
+    existing_flow_keys = {_flow_config_key(flow_config) for flow_config in config.flows}
+    for flow_config in flow_configs:
+        flow_key = _flow_config_key(flow_config)
+        if flow_key is not None and flow_key in existing_flow_keys:
+            continue
+
+        config.flows.append(flow_config)
+        existing_flow_keys.add(flow_key)
+
+
 def load_default_colang_1_flows(config: RailsConfig) -> None:
     """Load the built-in Colang 1.0 LLM flows into the rails config."""
     if config.colang_version != "1.0":
@@ -79,7 +94,7 @@ def load_default_colang_1_flows(config: RailsConfig) -> None:
     for flow_config in default_flows:
         flow_config["is_system_flow"] = True
 
-    config.flows.extend(default_flows)
+    _extend_missing_flow_configs(config, default_flows)
 
 
 def load_guardrails_library_flows_and_bot_messages(config: RailsConfig) -> None:
@@ -101,7 +116,7 @@ def load_guardrails_library_flows_and_bot_messages(config: RailsConfig) -> None:
         for flow_config in content["flows"]:
             flow_config["is_system_flow"] = True
 
-        config.flows.extend(content["flows"])
+        _extend_missing_flow_configs(config, content["flows"])
 
         for message_id, utterances in content.get("bot_messages", {}).items():
             if message_id not in config.bot_messages:

nemoguardrails/rails/llm/streaming/generation_stream.py

@@ -127,16 +127,36 @@ async def _generation_task():
         base_iterator = streaming_handler
 
     async def wrapped_iterator():
+        stream_completed = False
+        terminal_output_rail_error = False
         try:
             async for chunk in base_iterator:
                 if chunk is not None:
+                    if output_rail_streaming_enabled and _is_stream_error_chunk(chunk):
+                        terminal_output_rail_error = True
                     yield chunk
+            stream_completed = True
         finally:
-            await task
+            await _finish_generation_task(
+                task,
+                stream_completed=stream_completed and not terminal_output_rail_error,
+            )
 
     return wrapped_iterator()
 
 
+def _is_stream_error_chunk(chunk: Union[str, dict]) -> bool:
+    if not isinstance(chunk, str):
+        return False
+
+    try:
+        chunk_data = json.loads(chunk)
+    except json.JSONDecodeError:
+        return False
+
+    return isinstance(chunk_data, dict) and "error" in chunk_data
+
+
 def _track_generation_task(rails: GenerationStreamSurface, task: asyncio.Task) -> None:
     """Track background stream tasks so they are not garbage collected."""
     task_holder = cast(Any, rails)
@@ -148,3 +168,21 @@ def task_done_callback(task):
         task_holder._active_tasks.discard(task)
 
     task.add_done_callback(task_done_callback)
+
+
+async def _finish_generation_task(
+    task: asyncio.Task,
+    *,
+    stream_completed: bool,
+) -> None:
+    """Finalize a background generation task for a stream consumer."""
+    cancelled_by_stream_close = False
+    if not stream_completed and not task.done():
+        task.cancel()
+        cancelled_by_stream_close = True
+
+    try:
+        await task
+    except asyncio.CancelledError:
+        if not cancelled_by_stream_close:
+            raise

nemoguardrails/rails/llm/streaming/streaming_output_rails.py

@@ -49,7 +49,8 @@ async def run_output_rails_in_streaming(
     """
     buffer_strategy = get_buffer_strategy(output_rails_streaming_config)
     output_rails_flows_id = rails.config.rails.output.flows
-    stream_first = stream_first or output_rails_streaming_config.stream_first
+    if stream_first is None:
+        stream_first = output_rails_streaming_config.stream_first
     get_action_details = partial(get_action_details_from_flow_id, flows=rails.config.flows)
 
     parallel_mode = getattr(rails.config.rails.output, "parallel", False)
@@ -146,7 +147,10 @@ async def _run_parallel_output_rails(
         result, status = result_tuple
 
         if status != "success":
-            log.error(f"Parallel rails execution failed with status: {status}")
+            error_message = f"Parallel rails execution failed with status: {status}"
+            log.error(error_message)
+            yield _internal_rail_error("output rails", error_message)
+            return
         else:
             # if there are any stop events, content was blocked or internal error occurred
             result_events = getattr(result, "events", None)
@@ -155,7 +159,10 @@ async def _run_parallel_output_rails(
                 return
 
     except Exception as e:
-        log.error(f"Error in parallel rail execution: {e}")
+        error_message = f"Error in parallel rail execution: {e}"
+        log.error(error_message)
+        yield _internal_rail_error("output rails", error_message)
+        return
 
     # update explain info for parallel mode
     rails._explain_info = rails._ensure_explain_info()
@@ -186,9 +193,16 @@ async def _run_sequential_output_rails(
         action_result = await rails.runtime.action_dispatcher.execute_action(action_name, params)
         rails._explain_info = rails._ensure_explain_info()
 
-        # Use the mapping to decide if the result indicates blocked content.
         action_func = rails.runtime.action_dispatcher.get_action(action_name)
-        if is_output_blocked(action_result, action_func):
+        result, status = _sequential_action_result(action_result)
+        if status != "success":
+            error_message = f"Action {action_name} failed with status: {status}"
+            log.error(error_message)
+            yield _internal_rail_error(flow_id, error_message)
+            return
+
+        # Use the mapping to decide if the result indicates blocked content.
+        if is_output_blocked(result, action_func):
             yield _blocked_output_error(flow_id)
             return
 
@@ -322,6 +336,17 @@ def _internal_rail_error(flow_id: str, error_message: str) -> str:
     return json.dumps(error_data)
 
 
+def _sequential_action_result(action_result: Any) -> tuple[Any, str]:
+    if (
+        isinstance(action_result, tuple)
+        and len(action_result) == 2
+        and isinstance(action_result[1], str)
+        and action_result[1] in {"success", "failed"}
+    ):
+        return action_result
+    return action_result, "success"
+
+
 def _parallel_stop_error(stop_event: dict) -> str:
     blocked_flow = stop_event.get("flow_id", "output rails")
     error_type = stop_event.get("error_type")

tests/rails/llm/test_config_preparation.py

@@ -26,7 +26,7 @@ def test_config_preparation_startup_module_exports_public_helpers():
     ]
 
 
-def test_prepare_llmrails_config_loads_colang_flows_each_time_in_place():
+def test_prepare_llmrails_config_deduplicates_loaded_colang_flows():
     config = RailsConfig(models=[])
 
     first = prepare_llmrails_config(
@@ -41,7 +41,30 @@ def test_prepare_llmrails_config_loads_colang_flows_each_time_in_place():
     assert first is config
     assert second is config
     assert first_flow_count > 0
-    assert len(config.flows) > first_flow_count
+    assert len(config.flows) == first_flow_count
+
+
+def test_prepare_llmrails_config_can_reload_flows_after_caller_clears_them():
+    config = RailsConfig(models=[])
+
+    prepare_llmrails_config(config=config)
+    flow_count = len(config.flows)
+    config.flows = []
+
+    prepare_llmrails_config(config=config)
+
+    assert len(config.flows) == flow_count
+
+
+def test_prepare_llmrails_config_copy_does_not_duplicate_loaded_flows():
+    config = RailsConfig(models=[])
+
+    prepare_llmrails_config(config=config)
+    flow_count = len(config.flows)
+    prepared = prepare_llmrails_config(config=config, in_place=False)
+
+    assert prepared is not config
+    assert len(prepared.flows) == flow_count
 
 
 def test_prepare_llmrails_config_marks_rail_flows_each_time():

tests/rails/llm/test_generation_stream.py

@@ -326,39 +326,27 @@ async def test_generation_token_stream_pushes_error_payload_when_generation_fail
 
 
 @pytest.mark.asyncio
-async def test_generation_token_stream_waits_for_generation_when_consumer_closes_early():
+async def test_generation_token_stream_cancels_generation_when_consumer_closes_early():
     rails = SlowFakeRails()
     stream = generation_token_stream(rails, messages=[{"role": "user", "content": "Hi"}])
 
     assert await stream.__anext__() == "hello"
-    close_task = asyncio.create_task(cast(Any, stream).aclose())
+    await asyncio.wait_for(cast(Any, stream).aclose(), timeout=1)
+    await asyncio.wait_for(rails.cancelled.wait(), timeout=1)
     await asyncio.sleep(0)
 
-    assert not close_task.done()
-    assert not rails.cancelled.is_set()
-
-    rails.release.set()
-    await asyncio.wait_for(close_task, timeout=1)
-
-    assert not rails.cancelled.is_set()
     assert cast(set, getattr(rails, "_active_tasks")) == set()
 
 
 @pytest.mark.asyncio
-async def test_generation_token_stream_waits_for_generation_after_output_rail_error():
+async def test_generation_token_stream_cancels_generation_after_output_rail_error():
     rails = BlockingOutputRailsFakeRails()
     stream = generation_token_stream(rails, messages=[{"role": "user", "content": "Hi"}])
 
-    collect_task = asyncio.create_task(_collect(stream))
+    chunks = await asyncio.wait_for(_collect(stream), timeout=1)
+    await asyncio.wait_for(rails.cancelled.wait(), timeout=1)
     await asyncio.sleep(0)
 
-    assert not collect_task.done()
-    assert not rails.cancelled.is_set()
-
-    rails.release.set()
-    chunks = await asyncio.wait_for(collect_task, timeout=1)
-
     assert len(chunks) == 1
     assert json.loads(chunks[0])["error"]["code"] == "content_blocked"
-    assert not rails.cancelled.is_set()
     assert cast(set, getattr(rails, "_active_tasks")) == set()

tests/rails/llm/test_streaming_output_rails.py

@@ -298,7 +298,7 @@ async def test_output_rails_stream_yields_block_error_for_sequential_rails(monke
 
 
 @pytest.mark.asyncio
-async def test_output_rails_stream_uses_config_stream_first_when_explicit_false(monkeypatch):
+async def test_output_rails_stream_honors_explicit_stream_first_false(monkeypatch):
     _patch_buffer_strategy(
         monkeypatch,
         [ChunkBatch(processing_context=["blocked"], user_output_chunks=["blocked"])],
@@ -316,9 +316,8 @@ async def test_output_rails_stream_uses_config_stream_first_when_explicit_false(
         )
     )
 
-    assert len(chunks) == 2
-    assert chunks[0] == "blocked"
-    assert json.loads(chunks[1])["error"]["code"] == "content_blocked"
+    assert len(chunks) == 1
+    assert json.loads(chunks[0])["error"]["code"] == "content_blocked"
 
 
 @pytest.mark.asyncio
@@ -359,7 +358,7 @@ async def test_output_rails_stream_closes_external_generator_when_consumer_close
 
 
 @pytest.mark.asyncio
-async def test_output_rails_stream_continues_for_sequential_action_failure(monkeypatch):
+async def test_output_rails_stream_yields_internal_error_for_sequential_action_failure(monkeypatch):
     _patch_buffer_strategy(
         monkeypatch,
         [ChunkBatch(processing_context=["unsafe"], user_output_chunks=["unsafe"])],
@@ -376,7 +375,15 @@ async def test_output_rails_stream_continues_for_sequential_action_failure(monke
         )
     )
 
-    assert chunks == ["unsafe"]
+    assert len(chunks) == 1
+    assert json.loads(chunks[0]) == {
+        "error": {
+            "message": "Internal error in self check output rail: Action self_check_output failed with status: failed",
+            "type": "internal_error",
+            "param": "self check output",
+            "code": "rail_execution_failure",
+        }
+    }
     assert rails._explain_info == {"ensured": 1}
 
 
@@ -437,7 +444,7 @@ async def test_output_rails_stream_yields_parallel_stop_event_error(monkeypatch)
 
 
 @pytest.mark.asyncio
-async def test_output_rails_stream_continues_for_parallel_action_failure(monkeypatch):
+async def test_output_rails_stream_yields_internal_error_for_parallel_action_failure(monkeypatch):
     _patch_buffer_strategy(
         monkeypatch,
         [ChunkBatch(processing_context=["unsafe"], user_output_chunks=["unsafe"])],
@@ -454,12 +461,19 @@ async def test_output_rails_stream_continues_for_parallel_action_failure(monkeyp
         )
     )
 
-    assert chunks == ["unsafe"]
-    assert rails._explain_info == {"ensured": 1}
+    assert len(chunks) == 1
+    assert json.loads(chunks[0]) == {
+        "error": {
+            "message": "Internal error in output rails rail: Parallel rails execution failed with status: failed",
+            "type": "internal_error",
+            "param": "output rails",
+            "code": "rail_execution_failure",
+        }
+    }
 
 
 @pytest.mark.asyncio
-async def test_output_rails_stream_continues_for_parallel_exception(monkeypatch):
+async def test_output_rails_stream_yields_internal_error_for_parallel_exception(monkeypatch):
     _patch_buffer_strategy(
         monkeypatch,
         [ChunkBatch(processing_context=["unsafe"], user_output_chunks=["unsafe"])],
@@ -476,5 +490,12 @@ async def test_output_rails_stream_continues_for_parallel_exception(monkeypatch)
         )
     )
 
-    assert chunks == ["unsafe"]
-    assert rails._explain_info == {"ensured": 1}
+    assert len(chunks) == 1
+    assert json.loads(chunks[0]) == {
+        "error": {
+            "message": "Internal error in output rails rail: Error in parallel rail execution: parallel boom",
+            "type": "internal_error",
+            "param": "output rails",
+            "code": "rail_execution_failure",
+        }
+    }