diff --git a/src/praisonai-agents/praisonaiagents/gateway/protocols.py b/src/praisonai-agents/praisonaiagents/gateway/protocols.py index bf84041cb..de277f41b 100644 --- a/src/praisonai-agents/praisonaiagents/gateway/protocols.py +++ b/src/praisonai-agents/praisonaiagents/gateway/protocols.py @@ -187,6 +187,9 @@ class EventType(str, Enum): # Streaming events (relayed from agent's StreamEventEmitter) TOKEN_STREAM = "token_stream" TOOL_CALL_STREAM = "tool_call_stream" + REASONING_STREAM = "reasoning_stream" + TOOL_PROGRESS_STREAM = "tool_progress_stream" + STREAM_ERROR = "stream_error" STREAM_END = "stream_end" # System events diff --git a/src/praisonai-bot/praisonai_bot/gateway/server.py b/src/praisonai-bot/praisonai_bot/gateway/server.py index b1a1346f3..2a729891c 100644 --- a/src/praisonai-bot/praisonai_bot/gateway/server.py +++ b/src/praisonai-bot/praisonai_bot/gateway/server.py @@ -2761,6 +2761,9 @@ async def _handle_client_message(self, client_id: str, data: Dict[str, Any]) -> features["events"].extend([ EventType.TOKEN_STREAM.value, EventType.TOOL_CALL_STREAM.value, + EventType.REASONING_STREAM.value, + EventType.TOOL_PROGRESS_STREAM.value, + EventType.STREAM_ERROR.value, EventType.STREAM_END.value, ]) @@ -2988,11 +2991,29 @@ async def _handle_client_message(self, client_id: str, data: Dict[str, Any]) -> session_id=session_id, ) session.add_message(message) - + + # A missing agent never enqueues a turn, so it must resolve + # terminally here — otherwise an "accepted" ack would leave + # the client waiting forever for a "final" that never comes. + if self._agents.get(session.agent_id) is None: + await self._send_to_client(client_id, { + "type": "response", + "status": "final", + "content": "Agent not available", + "outcome": {"status": "error"}, + "session_id": session_id, + }) + return True + response = await self._process_agent_message(session, message) - + + # Provisional acknowledgement: the turn was accepted/enqueued + # but is not yet resolved. A distinct status lets clients tell + # "accepted" from the "final" answer (sent later by + # _run_session_queue) instead of string-sniffing the content. await self._send_to_client(client_id, { "type": "response", + "status": "accepted", "content": response, "session_id": session_id, }) @@ -3045,8 +3066,9 @@ async def _process_agent_message( """Process a message through the agent. If the agent has a stream_emitter, registers a callback that relays - token deltas to the connected WebSocket client in real-time via - TOKEN_STREAM / TOOL_CALL_STREAM / STREAM_END events. + live progress to the connected WebSocket client in real-time via + TOKEN_STREAM / TOOL_CALL_STREAM / REASONING_STREAM / + TOOL_PROGRESS_STREAM / STREAM_ERROR / STREAM_END events. """ agent = self._agents.get(session.agent_id) if not agent: @@ -3271,11 +3293,17 @@ async def _run_session_queue(self, session: GatewaySession, agent: Any, client_i if not content: break # Queue is empty, exit loop - # Wire streaming relay if agent has a stream_emitter + # Wire streaming relay if agent has a stream_emitter. + # ``relay_futures`` collects the cross-thread sends the relay + # schedules so we can drain them before the final frame, + # guaranteeing "final" never overtakes trailing stream events. relay_callback = None + relay_futures: List[Any] = [] emitter = getattr(agent, 'stream_emitter', None) if emitter is not None and client_id: - relay_callback = self._make_stream_relay(client_id, session) + relay_callback = self._make_stream_relay( + client_id, session, relay_futures + ) emitter.add_callback(relay_callback) # Issue #3467: run the turn under a cancel scope so it can be @@ -3285,6 +3313,7 @@ async def _run_session_queue(self, session: GatewaySession, agent: Any, client_i from praisonaiagents.agent.interrupt import InterruptController controller = InterruptController() timeout = getattr(self.config, "per_turn_timeout", 0.0) or 0.0 + outcome_status = "ok" try: gate = getattr(self, "_admission_gate", None) if gate is not None and getattr(gate, "enabled", False): @@ -3299,6 +3328,7 @@ async def _run_session_queue(self, session: GatewaySession, agent: Any, client_i ) except AdmissionRejected as rej: response = rej.message + outcome_status = "rejected" else: response = await self._drive_turn( session, agent, content, controller, timeout @@ -3306,6 +3336,7 @@ async def _run_session_queue(self, session: GatewaySession, agent: Any, client_i except Exception as e: logger.error(f"Agent error in queue processor: {e}") response = f"Error: {str(e)}" + outcome_status = "error" finally: # Always clean up the relay callback if relay_callback and emitter is not None: @@ -3314,6 +3345,15 @@ async def _run_session_queue(self, session: GatewaySession, agent: Any, client_i except (ValueError, AttributeError): pass + # Drain any stream sends the relay scheduled cross-thread so the + # final frame is enqueued strictly after every trailing stream + # event (token/reasoning/tool-progress/STREAM_END). + if relay_futures: + await asyncio.gather( + *[asyncio.wrap_future(f) for f in relay_futures], + return_exceptions=True, + ) + response_message = GatewayMessage( content=response, sender_id=session.agent_id, @@ -3321,18 +3361,31 @@ async def _run_session_queue(self, session: GatewaySession, agent: Any, client_i ) session.add_message(response_message) + # Final frame: distinct "final" status resolves the pending turn + # that the "accepted" ack opened, and carries a structured + # terminal outcome alongside the text (not just a bare string). await self._send_to_client(client_id, { "type": "response", + "status": "final", "content": response, + "outcome": {"status": outcome_status}, "session_id": session.session_id, }) finally: session.mark_executing(False) def _make_stream_relay( - self, client_id: str, session: "GatewaySession" + self, + client_id: str, + session: "GatewaySession", + pending: Optional[List[Any]] = None, ) -> Callable: - """Create a StreamCallback that relays events to a WS client.""" + """Create a StreamCallback that relays events to a WS client. + + When ``pending`` is provided, every cross-thread send future is + appended to it so the caller can await them before emitting a + terminal frame (ordering guarantee: final never precedes stream). + """ gateway = self # Capture the running loop while we are still on it. loop = asyncio.get_running_loop() @@ -3340,30 +3393,52 @@ def _make_stream_relay( def _relay(event) -> None: try: from praisonaiagents.streaming.events import StreamEventType - + event_type = getattr(event, 'type', None) if event_type is None: return - - # Map StreamEventType -> gateway EventType + + sid = session.session_id + # Map a *closed* set of StreamEventTypes -> gateway EventType so a + # WS UI can render live progress (thinking, tool progress) and + # streamed failures without sniffing message text. Unmapped + # events are dropped. if event_type == StreamEventType.DELTA_TEXT: - gw_type = EventType.TOKEN_STREAM + # A reasoning/thinking delta is surfaced under its own event + # so clients can show "thinking…" separately from the answer. + if getattr(event, 'is_reasoning', False): + gw_type = EventType.REASONING_STREAM + else: + gw_type = EventType.TOKEN_STREAM data = { "content": getattr(event, 'content', ''), - "session_id": session.session_id, + "session_id": sid, } elif event_type == StreamEventType.DELTA_TOOL_CALL: gw_type = EventType.TOOL_CALL_STREAM data = { "tool_call": getattr(event, 'tool_call', {}), - "session_id": session.session_id, + "session_id": sid, + } + elif event_type == StreamEventType.TOOL_PROGRESS: + gw_type = EventType.TOOL_PROGRESS_STREAM + data = { + "content": getattr(event, 'content', ''), + "metadata": getattr(event, 'metadata', None), + "session_id": sid, + } + elif event_type == StreamEventType.ERROR: + gw_type = EventType.STREAM_ERROR + data = { + "error": getattr(event, 'error', None), + "session_id": sid, } elif event_type == StreamEventType.STREAM_END: gw_type = EventType.STREAM_END - data = {"session_id": session.session_id} + data = {"session_id": sid} else: - return # Skip non-essential events - + return # Skip non-forwarded events + gw_event = GatewayEvent( type=gw_type, data=data, @@ -3372,10 +3447,12 @@ def _relay(event) -> None: ) # No get_event_loop() in the threaded callback. - asyncio.run_coroutine_threadsafe( + fut = asyncio.run_coroutine_threadsafe( gateway._send_to_client(client_id, gw_event.to_dict()), loop, ) + if pending is not None: + pending.append(fut) except Exception: logger.warning("Stream relay error (non-fatal)", exc_info=True) @@ -3519,6 +3596,9 @@ async def _send_to_client(self, client_id: str, data: Dict[str, Any]) -> None: "error", "token_stream", "tool_call_stream", + "reasoning_stream", + "tool_progress_stream", + "stream_error", ]: session_id = self._client_sessions.get(client_id) if session_id: