fix: isolate AG-UI stream events per run to stop cross-user token leak (fixes #6150)

Co-Authored-By: Claude Opus 5.5 <[email protected]>

src/praisonai-agents/praisonaiagents/agent/chat_mixin.py

@@ -7099,8 +7099,12 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
                     model_name=str(getattr(self, 'llm', '') or ''),
                 ),
             )
+            # Copy ContextVars into the worker so the re-scheduled model call
+            # keeps the caller's context (session identity, AG-UI run owner).
+            import contextvars
+            _ctx = contextvars.copy_context()
             response = await loop.run_in_executor(
-                None, lambda: manager.execute_model_call(request, _model_fn)
+                None, lambda: _ctx.run(manager.execute_model_call, request, _model_fn)
             )
             return self._unwrap_model_response(response)
 

src/praisonai-agents/praisonaiagents/ui/agui/streaming.py

@@ -5,6 +5,7 @@ Provides utilities for streaming PraisonAI agent responses as AG-UI events.
 """
 
 import asyncio
+import contextvars
 import json
 import uuid
 from dataclasses import dataclass, field
@@ -30,6 +31,12 @@ from praisonaiagents.ui.agui.types import (
 from praisonaiagents.ui.protocols import A2UI_MIME_TYPE
 from praisonaiagents.streaming.events import StreamEventType as _SE
 
+# Identifies which AG-UI run owns the current task/thread. The agent's
+# stream_emitter is shared by every request served by one Agent, and its
+# callbacks fire synchronously in the emitting context, so each run's callback
+# only accepts events emitted while its own owner token is active.
+_RUN_OWNER: contextvars.ContextVar = contextvars.ContextVar("agui_run_owner", default=None)
+
 
 @dataclass
 class EventBuffer:
@@ -466,20 +473,29 @@ async def async_stream_agent_response(
     text_started = False
     text_ended = False
     original_emitter = getattr(agent, "stream_emitter", None)
+    owner = object()
 
     def on_stream_event(event: Any) -> None:
+        # Drop events emitted by another concurrent run on the same agent.
+        if _RUN_OWNER.get() is not owner:
+            return
         loop.call_soon_threadsafe(queue.put_nowait, ("stream", event))
 
     if original_emitter is not None:
         original_emitter.add_callback(on_stream_event)
 
     async def run_chat() -> Any:
+        # Tasks run in a copy of the creator's context, so this set is local
+        # to the chat task and never leaks into the caller's context.
+        _RUN_OWNER.set(owner)
         try:
             if hasattr(agent, "achat"):
                 return await agent.achat(user_input, stream=True)
             if hasattr(agent, "chat"):
+                # run_in_executor does not copy ContextVars; carry the owner.
+                ctx = contextvars.copy_context()
                 return await loop.run_in_executor(
-                    None, lambda: agent.chat(user_input, stream=True)
+                    None, lambda: ctx.run(agent.chat, user_input, stream=True)
                 )
             raise ValueError("Agent must have 'chat' or 'achat' method")
         finally:

src/praisonai-agents/tests/agui/test_concurrent_isolation.py

@@ -0,0 +1,55 @@
+"""Concurrent AG-UI runs on one agent must not receive each other's tokens (#6150)."""
+
+import asyncio
+
+from praisonaiagents.streaming.events import StreamEvent, StreamEventEmitter, StreamEventType
+from praisonaiagents.ui.agui.streaming import async_stream_agent_response
+
+
+class _AsyncAgent:
+    def __init__(self):
+        self.stream_emitter = StreamEventEmitter()
+
+    async def achat(self, text, stream=True):
+        for tok in (f"[{text}-1]", f"[{text}-2]"):
+            self.stream_emitter.emit(StreamEvent(type=StreamEventType.DELTA_TEXT, content=tok))
+            await asyncio.sleep(0.02)
+        return text
+
+
+class _SyncAgent:
+    def __init__(self):
+        self.stream_emitter = StreamEventEmitter()
+
+    def chat(self, text, stream=True):
+        import time
+        for tok in (f"[{text}-1]", f"[{text}-2]"):
+            self.stream_emitter.emit(StreamEvent(type=StreamEventType.DELTA_TEXT, content=tok))
+            time.sleep(0.02)
+        return text
+
+
+async def _collect(agent, who):
+    out = []
+    async for ev in async_stream_agent_response(agent, who, "th-" + who, "run-" + who):
+        if getattr(ev, "delta", None):
+            out.append(ev.delta)
+    return "".join(out)
+
+
+def _run_pair(agent):
+    async def main():
+        return await asyncio.gather(_collect(agent, "ALICE"), _collect(agent, "BOB"))
+    return asyncio.run(main())
+
+
+def test_concurrent_async_runs_do_not_leak_tokens():
+    alice, bob = _run_pair(_AsyncAgent())
+    assert alice == "[ALICE-1][ALICE-2]"
+    assert bob == "[BOB-1][BOB-2]"
+
+
+def test_concurrent_sync_runs_do_not_leak_tokens():
+    alice, bob = _run_pair(_SyncAgent())
+    assert alice == "[ALICE-1][ALICE-2]"
+    assert bob == "[BOB-1][BOB-2]"