fix(trace): record streamed agent turns (start(stream=True))

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

@@ -29,6 +29,9 @@
 # into tool worker threads with the context, so each turn has its own counters and
 # overlapping turns on a shared agent never exhaust each other's limits.
 _turn_loop_guard = contextvars.ContextVar("praisonai_turn_loop_guard", default=None)
+# Agent whose streamed turn already records agent_start/agent_end, set while its
+# non-streaming chat() fallback runs so that turn is not recorded twice.
+_stream_traced_agent = contextvars.ContextVar("praisonai_stream_traced_agent", default=None)
 
 
 def _agent_end_status(db_status: str) -> str:
@@ -3554,7 +3557,9 @@ def chat(self, prompt: str, temperature: Optional[float] = None, tools: Optional
         # Emit context trace event (zero overhead when not set)
         from ..trace.context_events import get_context_emitter
         _trace_emitter = get_context_emitter()
-        _trace_emitter.agent_start(self.name, {"role": self.role, "goal": self.goal})
+        _trace_turn = _stream_traced_agent.get() is not self
+        if _trace_turn:
+            _trace_emitter.agent_start(self.name, {"role": self.role, "goal": self.goal})
         durable_context = None
         durable_token = None
         previous_run_state = self._get_db_run_state()
@@ -3623,7 +3628,8 @@ def chat(self, prompt: str, temperature: Optional[float] = None, tools: Optional
                 end_durable_run(durable_token)
                 durable_context.close()
             try:
-                _trace_emitter.agent_end(self.name, {"status": _agent_end_status(db_status)})
+                if _trace_turn:
+                    _trace_emitter.agent_end(self.name, {"status": _agent_end_status(db_status)})
             finally:
                 if getattr(self, "_active_turn_token", None) is _cancel:
                     self._active_turn_token = None
@@ -5431,9 +5437,26 @@ def _start_stream(self, prompt: str, **kwargs) -> Generator[str, None, None]:
         # via _enter_turn) so overlapping streams on a shared agent neither
         # exhaust nor inherit each other's per-turn limits.
         guard_token = _turn_loop_guard.set([self, None])
+        # Record the streamed turn like chat()/achat() do, so a live replay can
+        # show the agent running while its words stream (no-op when tracing is off).
+        from ..trace.context_events import get_context_emitter
+        _trace_emitter = get_context_emitter()
+        status = "completed"
+        if _trace_emitter._enabled:
+            _trace_emitter.agent_start(self.name, {"role": self.role, "goal": self.goal})
         try:
-            yield from self._start_stream_turn(prompt, **kwargs)
+            # _start_stream_turn returns "failed" when a guardrail or hook blocks
+            # the turn, matching chat(), which records a blocked turn as failed.
+            status = (yield from self._start_stream_turn(prompt, **kwargs)) or status
+        except (GeneratorExit, InterruptedError):
+            status = "cancelled"
+            raise
+        except BaseException:
+            status = "failed"
+            raise
         finally:
+            if _trace_emitter._enabled:
+                _trace_emitter.agent_end(self.name, {"status": status})
             try:
                 _turn_loop_guard.reset(guard_token)
             except (ValueError, LookupError):
@@ -5458,7 +5481,7 @@ def _start_stream_turn(self, prompt: str, **kwargs) -> Generator[str, None, None
             if not _in_ok:
                 logging.warning(f"Agent {getattr(self, 'name', '')}: input blocked by guardrail: {_in_err}")
                 yield f"[Input blocked by guardrail: {_in_err}]"
-                return
+                return "failed"
             prompt = _in_prompt
 
         # Same pre-turn checks as chat()/achat(): a streamed turn must not
@@ -5479,7 +5502,7 @@ def _start_stream_turn(self, prompt: str, **kwargs) -> Generator[str, None, None
             ))
             if self._hook_runner.is_blocked(hook_results):
                 logging.warning(f"Agent {self.name} execution blocked by BEFORE_AGENT hook")
-                return
+                return "failed"
             for res in hook_results:
                 if res.output and res.output.modified_input and "prompt" in res.output.modified_input:
                     prompt = res.output.modified_input["prompt"]
@@ -6299,6 +6322,9 @@ def _stream_fallback_chat(self, prompt: str, kwargs: Dict[str, Any], streaming_e
         else:
             fallback_kwargs.pop('stream', None)
 
+        # The streamed turn already records agent_start/agent_end; keep chat()
+        # from recording the same turn a second time.
+        _traced = _stream_traced_agent.set(self)
         try:
             return self.chat(prompt, **fallback_kwargs)
         except Exception as fallback_error:
@@ -6311,6 +6337,8 @@ def _stream_fallback_chat(self, prompt: str, kwargs: Dict[str, Any], streaming_e
             except Exception:
                 pass
             raise fallback_error from streaming_error
+        finally:
+            _stream_traced_agent.reset(_traced)
 
     def _create_llm_summarize_function(self):
         """

src/praisonai-agents/tests/unit/trace/test_unnamed_agent_trace.py

@@ -172,3 +172,89 @@ async def fake_dispatch(*args, **kwargs):
         assert d.called and s.get_events() == []
     finally:
         reset_context_emitter(token)
+
+
+# --- streamed turns (start(stream=True)) ---------------------------------------
+
+def _fake_stream(self, prompt, **kwargs):
+    yield "Solar "
+    yield "sails."
+
+
+def test_streamed_turn_records_agent_start_and_end(sink):
+    from praisonaiagents import Agent
+
+    agent = Agent(instructions="x")
+    with patch.object(type(agent), "_start_stream_turn", _fake_stream):
+        text = "".join(agent._start_stream("hi"))
+    assert text == "Solar sails."
+    kinds = [(e.event_type.value, e.node_id, e.data.get("status")) for e in sink.get_events()]
+    assert kinds == [("agent_start", "agent:Agent", None), ("agent_end", "agent:Agent", "completed")]
+
+
+def test_streamed_turn_stopped_early_is_cancelled(sink):
+    from praisonaiagents import Agent
+
+    agent = Agent(name="Writer", instructions="x")
+    with patch.object(type(agent), "_start_stream_turn", _fake_stream):
+        gen = agent._start_stream("hi")
+        next(gen)
+        gen.close()  # the page stopped reading
+    assert _of(sink, "agent_end")[0].data == {"status": "cancelled"}
+
+
+def test_streamed_turn_that_raises_is_failed(sink):
+    from praisonaiagents import Agent
+
+    def boom(self, prompt, **kwargs):
+        yield "x"
+        raise RuntimeError("provider down")
+
+    agent = Agent(name="Writer", instructions="x")
+    with patch.object(type(agent), "_start_stream_turn", boom):
+        with pytest.raises(RuntimeError):
+            list(agent._start_stream("hi"))
+    assert _of(sink, "agent_end")[0].data == {"status": "failed"}
+
+
+def test_streamed_turn_records_nothing_when_tracing_is_off(monkeypatch):
+    monkeypatch.setenv("OPENAI_API_KEY", "sk-test-not-used")
+    from praisonaiagents import Agent
+
+    s = ContextListSink()
+    token = set_context_emitter(ContextTraceEmitter(sink=s, session_id="t", enabled=False))
+    try:
+        agent = Agent(name="Writer", instructions="x")
+        with patch.object(type(agent), "_start_stream_turn", _fake_stream):
+            assert "".join(agent._start_stream("hi")) == "Solar sails."
+        assert s.get_events() == []
+    finally:
+        reset_context_emitter(token)
+
+
+def test_streamed_turn_blocked_by_guardrail_is_failed(sink):
+    from praisonaiagents import Agent
+
+    agent = Agent(name="Writer", instructions="x")
+    with patch.object(agent, "_validate_input_with_guardrail", return_value=(False, None, "no")):
+        assert "blocked" in "".join(agent._start_stream("hi"))
+    assert _of(sink, "agent_end")[0].data == {"status": "failed"}
+
+
+def test_streamed_turn_fallback_records_the_turn_once(sink):
+    from praisonaiagents import Agent
+
+    def falls_back(self, prompt, **kwargs):
+        yield self._stream_fallback_chat(prompt, kwargs, RuntimeError("stream broke"))
+
+    agent = Agent(name="Writer", instructions="x")
+    with patch.object(type(agent), "_start_stream_impl", falls_back), \
+            patch.object(agent, "_chat_completion", return_value=_fake_completion()):
+        assert "".join(agent._start_stream("hi")) == "Solar sails ride sunlight."
+    kinds = [e.event_type.value for e in sink.get_events() if e.event_type.value.startswith("agent_")]
+    assert kinds == ["agent_start", "agent_end"]
+    assert _of(sink, "agent_end")[0].data == {"status": "completed"}
+    # chat() outside a stream still records its own turn.
+    with patch.object(agent, "_chat_completion", return_value=_fake_completion()):
+        agent.chat("again")
+    assert len(_of(sink, "agent_start")) == 2