2424from ag_ui .core import (
2525 BaseEvent ,
2626 EventType ,
27+ ReasoningMessageContentEvent ,
28+ ReasoningMessageEndEvent ,
29+ ReasoningMessageStartEvent ,
2730 RunAgentInput ,
2831 RunErrorEvent ,
2932 RunFinishedEvent ,
3033 RunStartedEvent ,
3134 StateSnapshotEvent ,
35+ TextMessageContentEvent ,
36+ TextMessageEndEvent ,
37+ TextMessageStartEvent ,
3238)
3339from ag_ui .core .types import Message as AGUIMessage
3440from ag_ui .encoder import EventEncoder
@@ -451,6 +457,190 @@ async def _gen():
451457 app .post (route , name = f"agui_agent_config_{ prefix .strip ('/' )} " )(_handler )
452458
453459
460+ # ---------------------------------------------------------------------------
461+ # Reasoning-aware AGUI handler
462+ # ---------------------------------------------------------------------------
463+ #
464+ # Agno's stock AGUI handler emits STEP_STARTED/STEP_FINISHED events for
465+ # reasoning, not the REASONING_MESSAGE_* events that CopilotKit expects.
466+ # And reasoning=True triggers a multi-call CoT loop that breaks with
467+ # aimock/fixture environments.
468+ #
469+ # This custom handler:
470+ # 1. Runs the agent with reasoning=False (single LLM call)
471+ # 2. Collects the streamed text content
472+ # 3. Parses <reasoning>...</reasoning> tags from the text
473+ # 4. Emits REASONING_MESSAGE_* events for the reasoning block
474+ # 5. Emits TEXT_MESSAGE_* events for the answer
475+ #
476+ # If no <reasoning> tags are found, the entire response is emitted as a
477+ # text message (graceful fallback for aimock fixtures that return plain
478+ # text containing reasoning keywords).
479+
480+ import re
481+
482+ _REASONING_PATTERN = re .compile (
483+ r"<reasoning>(.*?)</reasoning>" ,
484+ re .DOTALL | re .IGNORECASE ,
485+ )
486+
487+
488+ async def _run_reasoning_agent (
489+ agent : Union [Agent , RemoteAgent ], run_input : RunAgentInput
490+ ) -> AsyncIterator [BaseEvent ]:
491+ """Stream one reasoning agent run, synthesizing REASONING_MESSAGE events."""
492+ run_id = run_input .run_id or str (uuid .uuid4 ())
493+ thread_id = run_input .thread_id
494+
495+ try :
496+ user_input = extract_agui_user_input (run_input .messages or [])
497+
498+ yield RunStartedEvent (
499+ type = EventType .RUN_STARTED , thread_id = thread_id , run_id = run_id
500+ )
501+
502+ user_id : Optional [str ] = None
503+ if run_input .forwarded_props and isinstance (run_input .forwarded_props , dict ):
504+ user_id = run_input .forwarded_props .get ("user_id" )
505+
506+ session_state = validate_agui_state (run_input .state , thread_id ) or {}
507+
508+ response_stream = agent .arun ( # type: ignore[attr-defined]
509+ input = user_input ,
510+ session_id = thread_id ,
511+ stream = True ,
512+ stream_events = True ,
513+ user_id = user_id ,
514+ session_state = session_state ,
515+ run_id = run_id ,
516+ )
517+
518+ # Collect the full text from the agent stream — we need to see the
519+ # complete response to split reasoning from answer. We still forward
520+ # tool-call events in real-time (important for the reasoning-chain
521+ # demo that interleaves reasoning with tool rendering).
522+ full_text = ""
523+ tool_events : list [BaseEvent ] = []
524+
525+ async for event in async_stream_agno_response_as_agui_events (
526+ response_stream = response_stream , # type: ignore[arg-type]
527+ thread_id = thread_id ,
528+ run_id = run_id ,
529+ ):
530+ if event .type in (EventType .RUN_STARTED , EventType .RUN_FINISHED ):
531+ continue
532+ # Accumulate text content
533+ if event .type == EventType .TEXT_MESSAGE_CONTENT :
534+ full_text += event .delta # type: ignore[attr-defined]
535+ # Forward tool-call events immediately
536+ elif event .type in (
537+ EventType .TOOL_CALL_START ,
538+ EventType .TOOL_CALL_ARGS ,
539+ EventType .TOOL_CALL_END ,
540+ ):
541+ tool_events .append (event )
542+ # Skip text start/end — we'll re-emit with reasoning split
543+
544+ # Parse <reasoning>...</reasoning> tags
545+ match = _REASONING_PATTERN .search (full_text )
546+
547+ if match :
548+ reasoning_text = match .group (1 ).strip ()
549+ answer_text = (
550+ full_text [: match .start ()] + full_text [match .end () :]
551+ ).strip ()
552+ else :
553+ # Fallback: check for "Reasoning:" prefix pattern (aimock fixtures)
554+ lower = full_text .lower ()
555+ if lower .startswith ("reasoning:" ) or lower .startswith ("reasoning step" ):
556+ # Treat the whole text as containing reasoning — emit as
557+ # reasoning message so the ReasoningBlock renders, then emit
558+ # a short text answer.
559+ reasoning_text = full_text .strip ()
560+ answer_text = ""
561+ else :
562+ reasoning_text = ""
563+ answer_text = full_text .strip ()
564+
565+ # Emit reasoning message if we have reasoning content
566+ if reasoning_text :
567+ reasoning_msg_id = str (uuid .uuid4 ())
568+ yield ReasoningMessageStartEvent (
569+ type = EventType .REASONING_MESSAGE_START ,
570+ message_id = reasoning_msg_id ,
571+ role = "reasoning" ,
572+ )
573+ yield ReasoningMessageContentEvent (
574+ type = EventType .REASONING_MESSAGE_CONTENT ,
575+ message_id = reasoning_msg_id ,
576+ delta = reasoning_text ,
577+ )
578+ yield ReasoningMessageEndEvent (
579+ type = EventType .REASONING_MESSAGE_END ,
580+ message_id = reasoning_msg_id ,
581+ )
582+
583+ # Emit text message for the answer (or entire text if no reasoning)
584+ text_msg_id = str (uuid .uuid4 ())
585+ # Only emit a text message if there's actual content or tool events
586+ if answer_text or tool_events :
587+ yield TextMessageStartEvent (
588+ type = EventType .TEXT_MESSAGE_START ,
589+ message_id = text_msg_id ,
590+ role = "assistant" ,
591+ )
592+ if answer_text :
593+ yield TextMessageContentEvent (
594+ type = EventType .TEXT_MESSAGE_CONTENT ,
595+ message_id = text_msg_id ,
596+ delta = answer_text ,
597+ )
598+ yield TextMessageEndEvent (
599+ type = EventType .TEXT_MESSAGE_END ,
600+ message_id = text_msg_id ,
601+ )
602+
603+ # Emit any tool-call events that were collected
604+ for te in tool_events :
605+ yield te
606+
607+ yield RunFinishedEvent (
608+ type = EventType .RUN_FINISHED , thread_id = thread_id , run_id = run_id
609+ )
610+
611+ except asyncio .CancelledError : # noqa: TRY302
612+ raise
613+ except Exception as exc : # noqa: BLE001
614+ yield RunErrorEvent (type = EventType .RUN_ERROR , message = str (exc ))
615+
616+
617+ def _attach_reasoning_route (
618+ app : FastAPI , agent : Agent , prefix : str
619+ ) -> None :
620+ """Mount a reasoning-aware AGUI POST endpoint at `<prefix>/agui`."""
621+ encoder = EventEncoder ()
622+ route = f"{ prefix .rstrip ('/' )} /agui"
623+
624+ async def _handler (run_input : RunAgentInput ) -> StreamingResponse :
625+ async def _gen ():
626+ async for event in _run_reasoning_agent (agent , run_input ):
627+ yield encoder .encode (event )
628+
629+ return StreamingResponse (
630+ _gen (),
631+ media_type = "text/event-stream" ,
632+ headers = {
633+ "Cache-Control" : "no-cache" ,
634+ "Connection" : "keep-alive" ,
635+ "Access-Control-Allow-Origin" : "*" ,
636+ "Access-Control-Allow-Methods" : "POST, GET, OPTIONS" ,
637+ "Access-Control-Allow-Headers" : "*" ,
638+ },
639+ )
640+
641+ app .post (route , name = f"agui_reasoning_{ prefix .strip ('/' )} " )(_handler )
642+
643+
454644# ---------------------------------------------------------------------------
455645# AgentOS bootstrap
456646# ---------------------------------------------------------------------------
@@ -474,7 +664,8 @@ async def _gen():
474664 interfaces = [
475665 # main_agent is mounted separately below via _attach_hitl_aware_route
476666 # so it can forward tool results for HITL round-trips.
477- AGUI (agent = reasoning_agent , prefix = "/reasoning" ), # -> /reasoning/agui
667+ # reasoning_agent is mounted separately below via _attach_reasoning_route
668+ # so it can emit proper REASONING_MESSAGE_* AG-UI events.
478669 # No-tools agent for the MCP Apps cell. The CopilotKit runtime's
479670 # `mcpApps.servers` middleware injects MCP server tools at request
480671 # time, so the LLM only sees the MCP-provided toolset.
@@ -521,6 +712,11 @@ async def _gen():
521712_attach_state_aware_route (app , shared_state_rw_agent , "/shared-state-rw" )
522713_attach_state_aware_route (app , subagents_supervisor , "/subagents" )
523714
715+ # Reasoning-aware route — emits proper REASONING_MESSAGE_* AG-UI events
716+ # that CopilotKit renders via the reasoningMessage slot. Replaces the stock
717+ # AGUI(agent=reasoning_agent, prefix="/reasoning") interface.
718+ _attach_reasoning_route (app , reasoning_agent , "/reasoning" )
719+
524720# Agent Config Object cell — builds a per-request Agno Agent whose system
525721# prompt is composed from the CopilotKit provider's forwarded properties
526722# (tone / expertise / responseLength).
0 commit comments