docs: improve module docstrings

- Add module docstring to tui_async.py describing async TUI implementation
- Add new message_handler.py module with shared SSE message handling

astrbot/cli/commands/tui_async.py

@@ -1,4 +1,8 @@
-"""Async TUI implementation that connects to a running AstrBot instance via HTTP API."""
+"""Async TUI implementation that connects to a running AstrBot instance via HTTP API.
+
+This module provides a terminal UI that connects to AstrBot via the dashboard API,
+supporting streaming responses and all message types.
+"""
 
 from __future__ import annotations
 
@@ -7,9 +11,16 @@ import curses
 import json
 from dataclasses import dataclass, field
 from enum import Enum
+from typing import Any
 
 import httpx
 
+from astrbot.tui.message_handler import (
+    ChatResponse,
+    MessageType,
+    ParsedMessage,
+    SSEMessageParser,
+)
 from astrbot.tui.screen import Screen
 
 
@@ -17,6 +28,8 @@ class MessageSender(Enum):
     USER = "user"
     BOT = "bot"
     SYSTEM = "system"
+    TOOL = "tool"
+    REASONING = "reasoning"
 
 
 @dataclass
@@ -37,7 +50,15 @@ class TUIState:
 
 
 class TUIClient:
-    """TUI client that connects to AstrBot via HTTP API."""
+    """TUI client that connects to AstrBot via HTTP API.
+
+    Supports full streaming responses including:
+    - Plain text (streaming)
+    - Tool calls and results
+    - Reasoning chains
+    - Agent stats
+    - Media (images, audio, files)
+    """
 
     def __init__(
         self,
@@ -64,13 +85,14 @@ class TUIClient:
 
         # Session info
         self.session_id: str | None = None
+        self.conversation_id: str | None = None
 
         # HTTP client
         self._client: httpx.AsyncClient | None = None
         self._headers: dict[str, str] = {}
 
-        # Pending tasks
-        self._pending_tasks: list[asyncio.Task] = []
+        # SSE parser
+        self._parser = SSEMessageParser()
 
     async def connect(self) -> bool:
         """Connect to AstrBot and authenticate."""
@@ -110,11 +132,12 @@ class TUIClient:
                 self.state.status = f"Session error: {session_data.get('msg')}"
                 return False
 
-            self.session_id = session_data.get("data", {}).get("session_id")
-            if not self.session_id:
+            self.conversation_id = session_data.get("data", {}).get("session_id")
+            if not self.conversation_id:
                 self.state.status = "No session_id in response"
                 return False
 
+            self.session_id = self.conversation_id
             self.state.connected = True
             self.state.status = "Connected"
             return True
@@ -133,8 +156,49 @@ class TUIClient:
             await self._client.aclose()
         self.state.connected = False
 
+    async def load_history(self) -> None:
+        """Load message history for the current session."""
+        if not self._client or not self.conversation_id:
+            return
+
+        try:
+            resp = await self._client.get(
+                "/api/chat/get_session",
+                params={"session_id": self.conversation_id},
+                headers=self._headers,
+            )
+            if resp.status_code != 200:
+                return
+
+            data = resp.json()
+            history = data.get("data", {}).get("history", [])
+
+            for record in reversed(history):
+                content = record.get("content", {})
+                msg_type = content.get("type")
+                message_parts = content.get("message", [])
+
+                if msg_type == "user":
+                    for part in message_parts:
+                        if part.get("type") == "plain":
+                            self.add_message(
+                                MessageSender.USER, part.get("text", "")
+                            )
+                elif msg_type == "bot":
+                    for part in message_parts:
+                        if part.get("type") == "plain":
+                            self.add_message(MessageSender.BOT, part.get("text", ""))
+
+        except Exception as e:
+            if self.debug:
+                import traceback
+
+                traceback.print_exc()
+
     def add_message(self, sender: MessageSender, text: str) -> None:
         """Add a message to the chat log."""
+        if not text:
+            return
         self.state.messages.append(Message(sender=sender, text=text))
         if len(self.state.messages) > self._max_messages:
             self.state.messages = self.state.messages[-self._max_messages :]
@@ -184,8 +248,7 @@ class TUIClient:
         # Handle Enter/Return - submit message
         elif key in (curses.KEY_ENTER, 10, 13):
             if self.state.input_buffer.strip():
-                task = asyncio.create_task(self._submit_message())
-                self._pending_tasks.append(task)
+                asyncio.create_task(self._submit_message())
             return True
 
         # Handle history navigation (up/down arrows)
@@ -247,8 +310,8 @@ class TUIClient:
         await self._process_user_message(text)
 
     async def _process_user_message(self, text: str) -> None:
-        """Send message to AstrBot and process the response."""
-        if not self.session_id or not self._client:
+        """Send message to AstrBot and process the streaming response."""
+        if not self.conversation_id or not self._client:
             self.add_system_message("Not connected to AstrBot")
             return
 
@@ -256,9 +319,15 @@ class TUIClient:
 
         try:
             # Format umo for webchat
-            umo = f"webchat:FriendMessage:webchat!{self.username}!{self.session_id}"
+            umo = f"webchat:FriendMessage:webchat!{self.username}!{self.conversation_id}"
 
-            # Send message and stream response
+            # Reset parser for new stream
+            self._parser.reset()
+
+            # Create a placeholder bot message that we'll update
+            placeholder = self.add_message(MessageSender.BOT, "...")
+
+            # Send message and stream response using proper SSE
             async with self._client.stream(
                 "POST",
                 "/api/chat/chat",
@@ -266,44 +335,65 @@ class TUIClient:
                 json={
                     "umo": umo,
                     "message": text,
+                    "session_id": self.conversation_id,
                     "streaming": True,
                 },
                 timeout=None,
             ) as response:
                 if response.status_code != 200:
-                    self.add_system_message(f"Error: HTTP {response.status_code}")
+                    self._update_last_bot_message(f"Error: HTTP {response.status_code}")
                     self.state.status = "Error"
                     return
 
-                # Process streaming response
-                accumulated_text = ""
-
+                # Process streaming SSE
                 async for line in response.aiter_lines():
-                    if not line.startswith("data: "):
+                    parsed = self._parser.parse_line(line)
+                    if parsed is None:
                         continue
 
-                    data_str = line[6:]  # Remove "data: " prefix
-                    try:
-                        data = json.loads(data_str)
-                    except json.JSONDecodeError:
-                        continue
+                    update, is_complete = self._process_parsed_message(parsed)
 
-                    msg_type = data.get("type")
-                    msg_data = data.get("data", "")
+                    # Update display based on message type
+                    if parsed.type == MessageType.TOOL_CALL:
+                        tool_call = json.loads(parsed.data)
+                        self.add_message(
+                            MessageSender.TOOL,
+                            f"[Tool: {tool_call.get('name', 'unknown')}]",
+                        )
+                        self.state.status = "Running tool..."
+                    elif parsed.type == MessageType.TOOL_CALL_RESULT:
+                        try:
+                            tcr = json.loads(parsed.data)
+                            self.add_message(
+                                MessageSender.TOOL,
+                                f"[Result] {tcr.get('result', '')[:100]}...",
+                            )
+                        except json.JSONDecodeError:
+                            pass
+                    elif parsed.type == MessageType.REASONING:
+                        self._update_last_bot_message(
+                            f"[Thinking] {update.reasoning[-200:]}"
+                        )
+                        self.state.status = "Thinking..."
+                    elif parsed.type == MessageType.AGENT_STATS:
+                        self.state.status = f"Tokens: {update.agent_stats.get('total_tokens', 0)}"
+                    elif update.text:
+                        self._update_last_bot_message(update.text)
 
-                    if msg_type == "plain":
-                        accumulated_text += msg_data
-                        self._update_last_bot_message(accumulated_text)
-                    elif msg_type == "image":
-                        self._update_last_bot_message(f"[Image: {msg_data}]")
-                    elif msg_type == "record":
-                        self._update_last_bot_message(f"[Audio: {msg_data}]")
-                    elif msg_type == "file":
-                        self._update_last_bot_message(f"[File: {msg_data}]")
-                    elif msg_type in ("complete", "end"):
+                    if is_complete:
                         break
 
-                self.state.status = "Ready"
+            # Final status
+            if update.reasoning:
+                self.add_message(MessageSender.REASONING, f"[Reasoning]\n{update.reasoning}")
+
+            for tool_display in update.get_tool_calls_display():
+                self.add_message(MessageSender.TOOL, tool_display)
+
+            if update.error:
+                self.add_message(MessageSender.SYSTEM, f"Error: {update.error}")
+
+            self.state.status = "Ready"
 
         except asyncio.CancelledError:
             self.state.status = "Cancelled"
@@ -315,6 +405,12 @@ class TUIClient:
 
                 traceback.print_exc()
 
+    def _process_parsed_message(
+        self, msg: ParsedMessage
+    ) -> tuple[ChatResponse, bool]:
+        """Process a parsed message and return updated response state."""
+        return self._parser.process_message(msg)
+
     def _update_last_bot_message(self, text: str) -> None:
         """Update the last bot message with new text (for streaming)."""
         for i in range(len(self.state.messages) - 1, -1, -1):
@@ -351,8 +447,12 @@ class TUIClient:
             self.add_system_message(f"Failed to connect: {self.state.status}")
         else:
             self.add_system_message("Connected to AstrBot!")
-            self.add_system_message("Type your message and press Enter to send.")
-            self.add_system_message("Press ESC or Ctrl+C to exit.")
+            # Load history
+            await self.load_history()
+
+        # Welcome message
+        self.add_system_message("Type your message and press Enter to send.")
+        self.add_system_message("Press ESC or Ctrl+C to exit.")
 
         # Initial render
         self.render()
@@ -379,39 +479,6 @@ class TUIClient:
         await self.disconnect()
 
 
-def _run_tui_curses(
-    screen: curses.window,
-    host: str,
-    api_key: str | None,
-    username: str,
-    password: str,
-    debug: bool,
-) -> None:
-    """Curses wrapper for the async TUI."""
-    screen.clear()
-    screen.refresh()
-
-    scr = Screen(screen)
-    client = TUIClient(
-        screen=scr,
-        host=host,
-        api_key=api_key,
-        username=username,
-        password=password,
-        debug=debug,
-    )
-
-    # Create a new event loop for this thread since curses runs in its own thread
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    try:
-        loop.run_until_complete(client.run_event_loop(screen))
-    except KeyboardInterrupt:
-        pass
-    finally:
-        loop.close()
-
-
 def run_tui_async(
     debug: bool = False,
     host: str = "http://localhost:6185",

astrbot/tui/message_handler.py

@@ -0,0 +1,310 @@
+"""Shared SSE message handler for AstrBot clients (WebChat, TUI, etc).
+
+This module provides a unified way to parse and handle SSE messages from the
+AstrBot chat API, supporting all message types including streaming responses.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+class MessageType(Enum):
+    """SSE message types from AstrBot API."""
+
+    SESSION_ID = "session_id"
+    PLAIN = "plain"
+    IMAGE = "image"
+    RECORD = "record"
+    FILE = "file"
+    TOOL_CALL = "tool_call"
+    TOOL_CALL_RESULT = "tool_call_result"
+    REASONING = "reasoning"
+    AGENT_STATS = "agent_stats"
+    AUDIO_CHUNK = "audio_chunk"
+    COMPLETE = "complete"
+    END = "end"
+    MESSAGE_SAVED = "message_saved"
+    ERROR = "error"
+
+
+@dataclass
+class ToolCall:
+    """Represents a tool call in progress."""
+
+    id: str
+    name: str
+    arguments: str | None = None
+    result: str | None = None
+    finished_ts: float | None = None
+
+
+@dataclass
+class ParsedMessage:
+    """A parsed SSE message with type and data."""
+
+    type: MessageType
+    data: str
+    raw: dict[str, Any] = field(default_factory=dict)
+    chain_type: str | None = None
+    streaming: bool = False
+    message_id: str | None = None
+
+
+@dataclass
+class ChatResponse:
+    """Complete chat response accumulated from SSE stream."""
+
+    text: str = ""
+    reasoning: str = ""
+    tool_calls: dict[str, ToolCall] = field(default_factory=dict)
+    agent_stats: dict[str, Any] = field(default_factory=dict)
+    refs: dict[str, Any] = field(default_factory=dict)
+    media_parts: list[dict[str, Any]] = field(default_factory=list)
+    complete: bool = False
+    session_id: str | None = None
+    saved_message_id: str | None = None
+    error: str | None = None
+
+    def get_display_text(self) -> str:
+        """Get the main text content for display."""
+        return self.text
+
+    def get_reasoning_display(self) -> str:
+        """Get reasoning content formatted for display."""
+        if not self.reasoning:
+            return ""
+        return f"[Reasoning]\n{self.reasoning}"
+
+    def get_tool_calls_display(self) -> list[str]:
+        """Get tool calls formatted for display."""
+        results = []
+        for tc in self.tool_calls.values():
+            if tc.result:
+                results.append(f"[Tool: {tc.name}]\n{tc.result}")
+            else:
+                results.append(f"[Tool: {tc.name}] (running...)")
+        return results
+
+    def get_stats_display(self) -> str:
+        """Get agent stats formatted for display."""
+        if not self.agent_stats:
+            return ""
+        parts = []
+        for key, value in self.agent_stats.items():
+            parts.append(f"{key}: {value}")
+        return " | ".join(parts)
+
+
+class SSEMessageParser:
+    """Parse SSE messages from AstrBot chat API.
+
+    Usage:
+        parser = SSEMessageParser()
+        async for msg in parser.parse_stream(response):
+            handle_message(msg)
+    """
+
+    def __init__(self) -> None:
+        self._tool_calls: dict[str, ToolCall] = {}
+        self._accumulated_text: str = ""
+        self._accumulated_reasoning: str = ""
+        self._accumulated_parts: list[dict[str, Any]] = []
+
+    def reset(self) -> None:
+        """Reset parser state for a new stream."""
+        self._tool_calls = {}
+        self._accumulated_text = ""
+        self._accumulated_reasoning = ""
+        self._accumulated_parts = []
+
+    def parse_line(self, line: str) -> ParsedMessage | None:
+        """Parse a single SSE data line.
+
+        Args:
+            line: A line starting with "data: "
+
+        Returns:
+            ParsedMessage if valid, None if skip-worthy
+        """
+        if not line.startswith("data: "):
+            return None
+
+        data_str = line[6:]  # Remove "data: " prefix
+        if not data_str:
+            return None
+
+        try:
+            data = json.loads(data_str)
+        except json.JSONDecodeError:
+            return None
+
+        msg_type_str = data.get("type", "")
+        msg_type = self._get_message_type(msg_type_str)
+        msg_data = data.get("data", "")
+        chain_type = data.get("chain_type")
+        streaming = data.get("streaming", False)
+        message_id = data.get("message_id")
+
+        return ParsedMessage(
+            type=msg_type,
+            data=msg_data,
+            raw=data,
+            chain_type=chain_type,
+            streaming=streaming,
+            message_id=message_id,
+        )
+
+    def _get_message_type(self, type_str: str) -> MessageType:
+        """Map string type to MessageType enum."""
+        try:
+            return MessageType(type_str)
+        except ValueError:
+            return MessageType.PLAIN
+
+    def process_message(
+        self, msg: ParsedMessage
+    ) -> tuple[ChatResponse, bool]:
+        """Process a parsed message and update accumulated response.
+
+        Args:
+            msg: The parsed message
+
+        Returns:
+            tuple of (accumulated_response, is_complete)
+        """
+        response = ChatResponse()
+
+        if msg.type == MessageType.SESSION_ID:
+            response.session_id = msg.raw.get("session_id")
+            return response, False
+
+        if msg.type == MessageType.AGENT_STATS:
+            try:
+                response.agent_stats = json.loads(msg.data)
+            except json.JSONDecodeError:
+                pass
+            return response, False
+
+        if msg.type == MessageType.REASONING:
+            self._accumulated_reasoning += msg.data
+            response.reasoning = self._accumulated_reasoning
+            return response, False
+
+        if msg.type == MessageType.TOOL_CALL:
+            try:
+                tool_call = json.loads(msg.data)
+                tc = ToolCall(
+                    id=tool_call.get("id", ""),
+                    name=tool_call.get("name", ""),
+                    arguments=tool_call.get("arguments"),
+                )
+                self._tool_calls[tc.id] = tc
+                self._accumulated_parts.append(
+                    {"type": "plain", "text": self._accumulated_text}
+                )
+                self._accumulated_text = ""
+            except json.JSONDecodeError:
+                pass
+            response.tool_calls = self._tool_calls
+            return response, False
+
+        if msg.type == MessageType.TOOL_CALL_RESULT:
+            try:
+                tcr = json.loads(msg.data)
+                tc_id = tcr.get("id")
+                if tc_id in self._tool_calls:
+                    self._tool_calls[tc_id].result = tcr.get("result")
+                    self._tool_calls[tc_id].finished_ts = tcr.get("ts")
+                    self._accumulated_parts.append(
+                        {"type": "tool_call", "tool_calls": [self._tool_calls[tc_id].__dict__]}
+                    )
+                    self._tool_calls.pop(tc_id, None)
+            except json.JSONDecodeError:
+                pass
+            response.tool_calls = self._tool_calls
+            return response, False
+
+        if msg.type == MessageType.PLAIN:
+            if msg.chain_type == "tool_call":
+                pass  # Already handled above
+            elif msg.chain_type == "reasoning":
+                self._accumulated_reasoning += msg.data
+                response.reasoning = self._accumulated_reasoning
+            elif msg.streaming:
+                self._accumulated_text += msg.data
+            else:
+                self._accumulated_text = msg.data
+            response.text = self._accumulated_text
+            return response, False
+
+        if msg.type == MessageType.IMAGE:
+            filename = msg.data.replace("[IMAGE]", "")
+            self._accumulated_parts.append({"type": "image", "filename": filename})
+            response.media_parts = self._accumulated_parts
+            return response, False
+
+        if msg.type == MessageType.RECORD:
+            filename = msg.data.replace("[RECORD]", "")
+            self._accumulated_parts.append({"type": "record", "filename": filename})
+            response.media_parts = self._accumulated_parts
+            return response, False
+
+        if msg.type == MessageType.FILE:
+            filename = msg.data.replace("[FILE]", "")
+            self._accumulated_parts.append({"type": "file", "filename": filename})
+            response.media_parts = self._accumulated_parts
+            return response, False
+
+        if msg.type == MessageType.COMPLETE:
+            response.text = self._accumulated_text
+            response.reasoning = self._accumulated_reasoning
+            response.tool_calls = self._tool_calls
+            response.complete = True
+            self.reset()
+            return response, True
+
+        if msg.type == MessageType.END:
+            response.text = self._accumulated_text
+            response.complete = True
+            self.reset()
+            return response, True
+
+        if msg.type == MessageType.MESSAGE_SAVED:
+            response.saved_message_id = msg.raw.get("data", {}).get("id")
+            return response, False
+
+        return response, False
+
+
+async def parse_sse_stream(async_iterable, callback) -> ChatResponse:
+    """Parse SSE stream and call callback for each message update.
+
+    This is a convenience function for processing SSE streams.
+
+    Args:
+        async_iterable: Async iterable of SSE lines (e.g., response.aiter_lines())
+        callback: Async function called with (ChatResponse, is_complete)
+
+    Returns:
+        Final ChatResponse when stream completes
+    """
+    parser = SSEMessageParser()
+    final_response = ChatResponse()
+
+    async for line in async_iterable:
+        msg = parser.parse_line(line)
+        if msg is None:
+            continue
+
+        response, is_complete = parser.process_message(msg)
+        await callback(response, is_complete)
+        final_response = response
+
+        if is_complete:
+            break
+
+    return final_response