Propagate model reasoning through AgentRun protocols

Constraint: Reasoning exposure must follow MODEL_PARAMETER_RULES.thinking from the agent environment.

Rejected: Querying model metadata per invoke | The requirement pins behavior to agent env and would widen runtime behavior.

Confidence: high

Scope-risk: moderate

Directive: Keep protocol gating centralized so future protocol handlers do not leak reasoning when thinking is disabled.

Tested: uv run --extra server pytest tests/unittests/server -q; uv run --extra server ruff check agentrun/utils/reasoning.py agentrun/server/model.py agentrun/server/invoker.py agentrun/server/openai_protocol.py agentrun/server/agui_protocol.py scripts/smoke_reasoning_protocol.py tests/unittests/server/test_reasoning.py tests/unittests/server/test_openai_protocol.py tests/unittests/server/test_agui_protocol.py; real-model smoke with thinking=true and thinking=false for OpenAI and AG-UI

Change-Id: I3b1ae025db4c0d26631cf4b4bb8e322dec77ae18
Not-tested: Remote CI and hosted preprod endpoint
Signed-off-by: congxiao.wxx <congxiao.wxx@alibaba-inc.com>

Author: congxiao-wxx

agentrun/server/agui_protocol.py

@@ -8,6 +8,7 @@
 """
 
 from dataclasses import dataclass, field
+import json
 from typing import (
     Any,
     AsyncIterator,
@@ -30,6 +31,10 @@
 import pydash
 
 from ..utils.helper import merge, MergeOptions
+from ..utils.reasoning import (
+    get_reasoning_content,
+    is_thinking_enabled_from_env,
+)
 from .model import (
     AgentEvent,
     AgentRequest,
@@ -60,6 +65,14 @@ class TextState:
     message_id: str = field(default_factory=lambda: str(uuid.uuid4()))
 
 
+@dataclass
+class ReasoningState:
+    started: bool = False
+    message_started: bool = False
+    phase_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    message_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+
+
 @dataclass
 class ToolCallState:
     name: str = ""
@@ -72,6 +85,7 @@ class ToolCallState:
 @dataclass
 class StreamStateMachine:
     text: TextState = field(default_factory=TextState)
+    reasoning: ReasoningState = field(default_factory=ReasoningState)
     tool_call_states: Dict[str, ToolCallState] = field(default_factory=dict)
     tool_result_chunks: Dict[str, List[str]] = field(default_factory=dict)
     run_errored: bool = False
@@ -121,6 +135,43 @@ def cache_tool_result_chunk(self, tool_id: str, delta: str) -> None:
     def pop_tool_result_chunks(self, tool_id: str) -> str:
         return "".join(self.tool_result_chunks.pop(tool_id, []))
 
+    def ensure_reasoning_started(self) -> Iterator[str]:
+        if not self.reasoning.started:
+            yield _encode_reasoning_event(
+                "REASONING_START",
+                messageId=self.reasoning.phase_id,
+            )
+            self.reasoning.started = True
+        if not self.reasoning.message_started:
+            yield _encode_reasoning_event(
+                "REASONING_MESSAGE_START",
+                messageId=self.reasoning.message_id,
+                role="reasoning",
+            )
+            self.reasoning.message_started = True
+
+    def end_reasoning_if_open(self) -> Iterator[str]:
+        if self.reasoning.message_started:
+            yield _encode_reasoning_event(
+                "REASONING_MESSAGE_END",
+                messageId=self.reasoning.message_id,
+            )
+            self.reasoning.message_started = False
+        if self.reasoning.started:
+            yield _encode_reasoning_event(
+                "REASONING_END",
+                messageId=self.reasoning.phase_id,
+            )
+            self.reasoning = ReasoningState()
+
+
+def _encode_reasoning_event(event_type: str, **payload: Any) -> str:
+    return (
+        "data: "
+        + json.dumps({"type": event_type, **payload}, ensure_ascii=False)
+        + "\n\n"
+    )
+
 
 class AGUIProtocolHandler(BaseProtocolHandler):
     """AG-UI 协议处理器
@@ -376,6 +427,10 @@ async def _format_stream(
         if state.run_errored:
             return
 
+        # 结束未结束的 reasoning 消息
+        for sse_data in state.end_reasoning_if_open():
+            yield sse_data
+
         # 结束所有未结束的工具调用
         for sse_data in state.end_all_tools(self._encoder):
             yield sse_data
@@ -399,8 +454,6 @@ def _process_event_with_boundaries(
         state: StreamStateMachine,
     ) -> Iterator[str]:
         """处理事件并注入边界事件"""
-        import json
-
         from ag_ui.core import CustomEvent as AguiCustomEvent
         from ag_ui.core import (
             RunErrorEvent,
@@ -413,6 +466,8 @@ def _process_event_with_boundaries(
             ToolCallStartEvent,
         )
 
+        thinking_enabled = is_thinking_enabled_from_env()
+
         # RAW 事件直接透传
         if event.event == EventType.RAW:
             raw_data = event.data.get("raw", "")
@@ -422,9 +477,46 @@ def _process_event_with_boundaries(
                 yield raw_data
             return
 
+        if event.event == EventType.REASONING:
+            if thinking_enabled:
+                reasoning_content = (
+                    event.data.get("delta")
+                    or get_reasoning_content(event.data)
+                    or ""
+                )
+                if reasoning_content:
+                    for sse_data in state.end_text_if_open(self._encoder):
+                        yield sse_data
+                    for sse_data in state.end_all_tools(self._encoder):
+                        yield sse_data
+                    for sse_data in state.ensure_reasoning_started():
+                        yield sse_data
+                    yield _encode_reasoning_event(
+                        "REASONING_MESSAGE_CONTENT",
+                        messageId=state.reasoning.message_id,
+                        delta=reasoning_content,
+                    )
+            return
+
         # TEXT 事件：在首个 TEXT 前注入 TEXT_MESSAGE_START
         # AG-UI 协议要求：发送 TEXT_MESSAGE_START 前必须先结束所有未结束的 TOOL_CALL
         if event.event == EventType.TEXT:
+            addition = self._strip_reasoning_from_addition(
+                event.addition, thinking_enabled
+            )
+            addition_reasoning = get_reasoning_content(event.addition or {})
+            if thinking_enabled and addition_reasoning:
+                for sse_data in state.ensure_reasoning_started():
+                    yield sse_data
+                yield _encode_reasoning_event(
+                    "REASONING_MESSAGE_CONTENT",
+                    messageId=state.reasoning.message_id,
+                    delta=addition_reasoning,
+                )
+
+            for sse_data in state.end_reasoning_if_open():
+                yield sse_data
+
             for sse_data in state.end_all_tools(self._encoder):
                 yield sse_data
 
@@ -435,13 +527,13 @@ def _process_event_with_boundaries(
                 message_id=state.text.message_id,
                 delta=event.data.get("delta", ""),
             )
-            if event.addition:
+            if addition:
                 event_dict = agui_event.model_dump(
                     by_alias=True, exclude_none=True
                 )
                 event_dict = self._apply_addition(
                     event_dict,
-                    event.addition,
+                    addition,
                     event.addition_merge_options,
                 )
                 json_str = json.dumps(event_dict, ensure_ascii=False)
@@ -455,6 +547,9 @@ def _process_event_with_boundaries(
             tool_id = event.data.get("id", "")
             tool_name = event.data.get("name", "")
 
+            for sse_data in state.end_reasoning_if_open():
+                yield sse_data
+
             for sse_data in state.end_text_if_open(self._encoder):
                 yield sse_data
 
@@ -491,6 +586,9 @@ def _process_event_with_boundaries(
             tool_name = event.data.get("name", "")
             tool_args = event.data.get("args", "")
 
+            for sse_data in state.end_reasoning_if_open():
+                yield sse_data
+
             for sse_data in state.end_text_if_open(self._encoder):
                 yield sse_data
 
@@ -541,6 +639,9 @@ def _process_event_with_boundaries(
             timeout = event.data.get("timeout")
             schema = event.data.get("schema")
 
+            for sse_data in state.end_reasoning_if_open():
+                yield sse_data
+
             for sse_data in state.end_text_if_open(self._encoder):
                 yield sse_data
 
@@ -601,6 +702,9 @@ def _process_event_with_boundaries(
             tool_id = event.data.get("id", "")
             tool_name = event.data.get("name", "")
 
+            for sse_data in state.end_reasoning_if_open():
+                yield sse_data
+
             for sse_data in state.end_text_if_open(self._encoder):
                 yield sse_data
 
@@ -767,6 +871,29 @@ def _apply_addition(
 
         return merge(event_data, addition, **(merge_options or {}))
 
+    def _strip_reasoning_from_addition(
+        self,
+        addition: Optional[Dict[str, Any]],
+        thinking_enabled: bool,
+    ) -> Optional[Dict[str, Any]]:
+        if not addition:
+            return addition
+
+        stripped = dict(addition)
+        stripped.pop("reasoning_content", None)
+        additional_kwargs = stripped.get("additional_kwargs")
+        if isinstance(additional_kwargs, dict):
+            additional_kwargs = dict(additional_kwargs)
+            additional_kwargs.pop("reasoning_content", None)
+            if additional_kwargs:
+                stripped["additional_kwargs"] = additional_kwargs
+            else:
+                stripped.pop("additional_kwargs", None)
+
+        if not thinking_enabled:
+            return stripped
+        return stripped or None
+
     async def _error_stream(self, message: str) -> AsyncIterator[str]:
         """生成错误事件流
 

agentrun/server/invoker.py

@@ -29,6 +29,7 @@
     InvokeAgentHandler,
     SyncInvokeAgentHandler,
 )
+from agentrun.utils.reasoning import get_reasoning_content
 
 
 class AgentInvoker:
@@ -124,6 +125,9 @@ async def invoke_stream(
                         # 处理用户返回的事件
                         for processed_event in self._process_user_event(item):
                             yield processed_event
+                    else:
+                        for processed_event in self._wrap_model_chunk(item):
+                            yield processed_event
             else:
                 # 非流式结果
                 results = self._wrap_non_stream(raw_result)
@@ -238,6 +242,11 @@ def _wrap_non_stream(self, result: Any) -> List[AgentEvent]:
                             data={"delta": item},
                         )
                     )
+                else:
+                    results.extend(self._wrap_model_chunk(item))
+
+        else:
+            results.extend(self._wrap_model_chunk(result))
 
         return results
 
@@ -267,6 +276,9 @@ async def _wrap_stream(
             elif isinstance(item, AgentEvent):
                 for processed_event in self._process_user_event(item):
                     yield processed_event
+            else:
+                for processed_event in self._wrap_model_chunk(item):
+                    yield processed_event
 
     async def _iterate_async(
         self, content: Union[Iterator[Any], AsyncIterator[Any]]
@@ -307,3 +319,31 @@ def _is_iterator(self, obj: Any) -> bool:
         if isinstance(obj, (str, bytes, dict, list, AgentEvent)):
             return False
         return hasattr(obj, "__iter__") or hasattr(obj, "__aiter__")
+
+    def _wrap_model_chunk(self, item: Any) -> List[AgentEvent]:
+        """Convert common model chunks into AgentEvent objects."""
+        events: List[AgentEvent] = []
+        reasoning_content = get_reasoning_content(item)
+        if reasoning_content:
+            events.append(
+                AgentEvent(
+                    event=EventType.REASONING,
+                    data={"delta": reasoning_content},
+                )
+            )
+
+        content = self._read_attr_or_key(item, "content")
+        if isinstance(content, str) and content:
+            events.append(
+                AgentEvent(
+                    event=EventType.TEXT,
+                    data={"delta": content},
+                )
+            )
+
+        return events
+
+    def _read_attr_or_key(self, obj: Any, key: str) -> Any:
+        if isinstance(obj, dict):
+            return obj.get(key)
+        return getattr(obj, key, None)

agentrun/server/model.py

@@ -14,7 +14,6 @@
     Iterator,
     List,
     Optional,
-    TYPE_CHECKING,
     Union,
 )
 
@@ -91,6 +90,7 @@ class Message(BaseModel):
     id: Optional[str] = None
     role: MessageRole
     content: Optional[Union[str, List[Dict[str, Any]]]] = None
+    reasoning_content: Optional[str] = None
     name: Optional[str] = None
     tool_calls: Optional[List[ToolCall]] = None
     tool_call_id: Optional[str] = None
@@ -125,6 +125,7 @@ class EventType(str, Enum):
     # 核心事件（用户主要使用）
     # =========================================================================
     TEXT = "TEXT"  # 文本内容块
+    REASONING = "REASONING"  # 模型思考内容块
     TOOL_CALL = "TOOL_CALL"  # 完整工具调用（含 id, name, args）
     TOOL_CALL_CHUNK = "TOOL_CALL_CHUNK"  # 工具调用参数片段（流式场景）
     TOOL_RESULT = "TOOL_RESULT"  # 工具执行结果（最终结果，标识流式输出结束）