"""In-memory model for the chat panel.
Converts raw WebSocket events into logical chat messages that renderers understand.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Literal
@dataclass
class ChatItem:
"""Logical item shown in the chat history."""
kind: Literal[
"user_message",
"assistant_message",
"thinking_block",
"tool_started",
"tool_call",
"error",
"status",
]
content: str = ""
meta: dict = field(default_factory=dict)
id: str = ""
class ChatModel:
"""Accumulate stream deltas and tool events into ChatItems."""
def __init__(self) -> None:
self.items: list[ChatItem] = []
self._current_assistant: ChatItem | None = None
self._current_thinking: ChatItem | None = None
def add_user_message(self, text: str) -> None:
self.items.append(ChatItem(kind="user_message", content=text))
self._current_assistant = None
def handle_ws_event(self, msg: dict) -> ChatItem | None:
msg_type = msg.get("type")
if msg_type == "stream_start":
self._current_assistant = ChatItem(kind="assistant_message", content="")
self.items.append(self._current_assistant)
return None
if msg_type == "thinking_delta":
if self._current_thinking is None:
self._current_thinking = ChatItem(kind="thinking_block", content="")
self.items.append(self._current_thinking)
self._current_thinking.content += msg.get("delta", "")
return None
if msg_type == "thinking_end":
self._current_thinking = None
return None
if msg_type == "stream_delta":
if self._current_assistant is None:
self._current_assistant = ChatItem(kind="assistant_message", content="")
self.items.append(self._current_assistant)
self._current_assistant.content += msg.get("delta", "")
return None
if msg_type == "tool_started":
item = ChatItem(kind="tool_started", meta=msg)
self.items.append(item)
return item
if msg_type == "tool_call":
item = ChatItem(kind="tool_call", meta=msg)
self.items.append(item)
return item
if msg_type == "error":
item = ChatItem(kind="error", content=msg.get("message", ""))
self.items.append(item)
return item
if msg_type in ("stream_end", "context_compressed", "heartbeat", "session_sync"):
return None
# Unknown event — store as status for debugging.
item = ChatItem(kind="status", content=f"{msg_type}: {msg}")
self.items.append(item)
return item