fix(streaming): don't reuse stale across output-rail chunks (NVIDIA-NeMo#1935) (NVIDIA-NeMo#1943)

Signed-off-by: Aleksandr Popov <alexander230r@gmail.com>

Author: Alexander230

nemoguardrails/rails/llm/utils.py

@@ -94,6 +94,8 @@ def get_action_details_from_flow_id(
             and "execute" in element["_source_mapping"]["line_text"]
             and "action_name" in element
         ):
-            return element["action_name"], element["action_params"]
+            # Return a copy: action_params belongs to the shared flow config and
+            # callers may resolve $bot_message / $user_message placeholders into it.
+            return element["action_name"], dict(element["action_params"] or {})
 
     raise ValueError(f"No run_action element found for flow_id: {flow_id}")

tests/test_streaming_output_rails.py

@@ -542,3 +542,119 @@ async def self_check_output(**kwargs):
         tokens.append(token)
 
     assert "".join(tokens) == "This is a complete response in a single chunk."
+
+
+@pytest.mark.asyncio
+async def test_streaming_output_rails_no_stale_substituted_param():
+    """Output rails that take the bot message as a substituted kwarg
+    (text=$bot_message, as privateai / prompt_security / regex rails do) see
+    each streamed chunk's own text, not a stale value from an earlier chunk.
+    """
+    config = RailsConfig.from_content(
+        config={
+            "models": [],
+            "rails": {
+                "output": {
+                    "flows": ["capture output"],
+                    "streaming": {"enabled": True, "chunk_size": 4, "context_size": 2},
+                }
+            },
+            "streaming": False,
+        },
+        colang_content="""
+        define user express greeting
+          "hi"
+
+        define flow
+          user express greeting
+          bot tell joke
+
+        define subflow capture output
+          execute capture_output(text=$bot_message)
+        """,
+    )
+
+    seen = []
+
+    @action(name="capture_output", output_mapping=lambda result: not result)
+    async def capture_output(**params):
+        # the substituted `text` kwarg must match this chunk's bot_message
+        seen.append((params.get("text"), params["context"]["bot_message"]))
+        return True
+
+    chat = TestChat(
+        config,
+        llm_completions=[
+            '  express greeting\nbot express greeting\n  "hi"',
+            '  "one two three four five six"',
+        ],
+        streaming=True,
+    )
+    chat.app.register_action(capture_output, name="capture_output")
+
+    async for _ in chat.app.stream_async(messages=[{"role": "user", "content": "hi"}]):
+        pass
+
+    assert len(seen) >= 2  # response spans multiple chunks
+    assert all(text == bot_message for text, bot_message in seen)
+
+    await asyncio.gather(*asyncio.all_tasks() - {asyncio.current_task()})
+
+
+@pytest.mark.asyncio
+async def test_streaming_output_rails_substitutes_user_message_param():
+    """Output rails that take the user message as a substituted kwarg
+    (text=$user_message) receive the resolved user message on every streamed
+    chunk, not the literal "$user_message" placeholder.
+    """
+    config = RailsConfig.from_content(
+        config={
+            "models": [],
+            "rails": {
+                "output": {
+                    "flows": ["capture output"],
+                    "streaming": {"enabled": True, "chunk_size": 4, "context_size": 2},
+                }
+            },
+            "streaming": False,
+        },
+        colang_content="""
+        define user express greeting
+          "hi"
+
+        define flow
+          user express greeting
+          bot tell joke
+
+        define subflow capture output
+          execute capture_output(text=$user_message)
+        """,
+    )
+
+    seen = []
+
+    @action(name="capture_output", output_mapping=lambda result: not result)
+    async def capture_output(**params):
+        # the substituted `text` kwarg must be the resolved user message
+        seen.append((params.get("text"), params["context"]["user_message"]))
+        return True
+
+    chat = TestChat(
+        config,
+        llm_completions=[
+            '  express greeting\nbot express greeting\n  "hi"',
+            '  "one two three four five six"',
+        ],
+        streaming=True,
+    )
+    chat.app.register_action(capture_output, name="capture_output")
+
+    async for _ in chat.app.stream_async(messages=[{"role": "user", "content": "hi"}]):
+        pass
+
+    assert seen  # the output rail ran on at least one chunk
+    # $user_message was resolved, not passed through as the literal placeholder
+    assert all(text != "$user_message" for text, _ in seen)
+    assert all(text == user_message for text, user_message in seen)
+
+    await asyncio.gather(*asyncio.all_tasks() - {asyncio.current_task()})