opensource/AstrBot
3
0
Fork 0
mirror of https://github.com/AstrBotDevs/AstrBot.git synced 2026-08-30 17:33:24 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: resume ChatUI streams after refresh (#9259)

Browse Source 
* fix: resume ChatUI streams after refresh

* fix: hide chat run service errors

* fix: harden resumed chat streams
This commit is contained in:
Soulter
2026-07-13 23:05:03 +08:00
committed by GitHub
parent cc0ec5e5c5
commit 6070da9ae2
16 changed files with 1294 additions and 360 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrbot/core/astr_agent_run_util.py
+18
View File
@@ -246,6 +246,24 @@ async def run_agent(
if stream_to_general and resp.type == "streaming_delta":
continue
if (
resp.type == "err"
and agent_runner.streaming
and not stream_to_general
):
chain = (
resp.data.get("chain") if isinstance(resp.data, dict) else None
)
if not isinstance(chain, MessageChain):
logger.error(
"Agent runner returned an error response without a message chain."
)
chain = MessageChain().message(
"Error occurred during AI execution."
)
yield chain
continue
if stream_to_general or not agent_runner.streaming:
if can_buffer_llm_result and resp.type == "llm_result":
buffered_llm_chains.append(resp.data["chain"])
astrbot/core/platform/sources/webchat/webchat_event.py
+29 -28
View File
@@ -20,15 +20,6 @@ from .webchat_queue_mgr import webchat_queue_mgr
attachments_dir = os.path.join(get_astrbot_data_path(), "attachments")
def _extract_conversation_id(session_id: str) -> str:
"""Extract raw webchat conversation id from event/session id."""
if session_id.startswith("webchat!"):
parts = session_id.split("!", 2)
if len(parts) == 3:
return parts[2]
return session_id
class WebChatMessageEvent(AstrMessageEvent):
def __init__(self, message_str, message_obj, platform_meta, session_id) -> None:
super().__init__(message_str, message_obj, platform_meta, session_id)
@@ -43,13 +34,9 @@ class WebChatMessageEvent(AstrMessageEvent):
emit_complete: bool = False,
) -> str | None:
request_id = str(message_id)
conversation_id = _extract_conversation_id(session_id)
web_chat_back_queue = webchat_queue_mgr.get_or_create_back_queue(
request_id,
conversation_id,
)
if not message:
await web_chat_back_queue.put(
await webchat_queue_mgr.put_back_queue(
request_id,
{
"type": "end",
"data": "",
@@ -63,7 +50,8 @@ class WebChatMessageEvent(AstrMessageEvent):
for comp in message.chain:
if isinstance(comp, Plain):
data = comp.text
await web_chat_back_queue.put(
accepted = await webchat_queue_mgr.put_back_queue(
request_id,
{
"type": "plain",
"data": data,
@@ -72,8 +60,11 @@ class WebChatMessageEvent(AstrMessageEvent):
"message_id": message_id,
},
)
if not accepted:
return None
elif isinstance(comp, Json):
await web_chat_back_queue.put(
accepted = await webchat_queue_mgr.put_back_queue(
request_id,
{
"type": "plain",
"data": json.dumps(comp.data, ensure_ascii=False),
@@ -82,6 +73,8 @@ class WebChatMessageEvent(AstrMessageEvent):
"message_id": message_id,
},
)
if not accepted:
return None
elif isinstance(comp, Image):
# save image to local
image_base64 = await comp.convert_to_base64()
@@ -95,7 +88,8 @@ class WebChatMessageEvent(AstrMessageEvent):
path = os.path.join(attachments_dir, filename)
await asyncio.to_thread(Path(path).write_bytes, image_bytes)
data = f"[IMAGE]{filename}"
await web_chat_back_queue.put(
accepted = await webchat_queue_mgr.put_back_queue(
request_id,
{
"type": "image",
"data": data,
@@ -103,6 +97,8 @@ class WebChatMessageEvent(AstrMessageEvent):
"message_id": message_id,
},
)
if not accepted:
return None
elif isinstance(comp, Record):
# save record to local
filename = f"{str(uuid.uuid4())}.wav"
@@ -111,7 +107,8 @@ class WebChatMessageEvent(AstrMessageEvent):
record_bytes = base64.b64decode(record_base64)
await asyncio.to_thread(Path(path).write_bytes, record_bytes)
data = f"[RECORD]{filename}"
await web_chat_back_queue.put(
accepted = await webchat_queue_mgr.put_back_queue(
request_id,
{
"type": "record",
"data": data,
@@ -119,6 +116,8 @@ class WebChatMessageEvent(AstrMessageEvent):
"message_id": message_id,
},
)
if not accepted:
return None
elif isinstance(comp, File):
# save file to local
file_path = await comp.get_file()
@@ -135,7 +134,8 @@ class WebChatMessageEvent(AstrMessageEvent):
dest_path = os.path.join(attachments_dir, filename)
shutil.copy2(file_path, dest_path)
data = f"[FILE]{filename}|{original_name}"
await web_chat_back_queue.put(
accepted = await webchat_queue_mgr.put_back_queue(
request_id,
{
"type": "file",
"data": data,
@@ -143,11 +143,14 @@ class WebChatMessageEvent(AstrMessageEvent):
"message_id": message_id,
},
)
if not accepted:
return None
else:
logger.debug(f"webchat 忽略: {comp.type}")
if emit_complete:
await web_chat_back_queue.put(
await webchat_queue_mgr.put_back_queue(
request_id,
{
"type": "complete",
"data": data,
@@ -169,11 +172,6 @@ class WebChatMessageEvent(AstrMessageEvent):
reasoning_content = ""
message_id = self.message_obj.message_id
request_id = str(message_id)
conversation_id = _extract_conversation_id(self.session_id)
web_chat_back_queue = webchat_queue_mgr.get_or_create_back_queue(
request_id,
conversation_id,
)
async for chain in generator:
# 处理音频流（Live Mode）
if chain.type == "audio_chunk":
@@ -196,7 +194,9 @@ class WebChatMessageEvent(AstrMessageEvent):
if text:
payload["text"] = text
await web_chat_back_queue.put(payload)
accepted = await webchat_queue_mgr.put_back_queue(request_id, payload)
if not accepted:
return
continue
# if chain.type == "break" and final_data:
@@ -224,7 +224,8 @@ class WebChatMessageEvent(AstrMessageEvent):
else:
final_data += r
await web_chat_back_queue.put(
await webchat_queue_mgr.put_back_queue(
request_id,
{
"type": "complete", # complete means we return the final result
"data": final_data,
astrbot/core/platform/sources/webchat/webchat_queue_mgr.py
+43
View File
@@ -10,6 +10,7 @@ class WebChatQueueMgr:
"""Conversation ID to asyncio.Queue mapping"""
self.back_queues: dict[str, asyncio.Queue] = {}
"""Request ID to asyncio.Queue mapping for responses"""
self._back_queue_close_events: dict[str, asyncio.Event] = {}
self._conversation_back_requests: dict[str, set[str]] = {}
self._request_conversation: dict[str, str] = {}
self._queue_close_events: dict[str, asyncio.Event] = {}
@@ -36,6 +37,7 @@ class WebChatQueueMgr:
self.back_queues[request_id] = asyncio.Queue(
maxsize=self.back_queue_maxsize
)
self._back_queue_close_events[request_id] = asyncio.Event()
if conversation_id:
self._request_conversation[request_id] = conversation_id
if conversation_id not in self._conversation_back_requests:
@@ -43,8 +45,49 @@ class WebChatQueueMgr:
self._conversation_back_requests[conversation_id].add(request_id)
return self.back_queues[request_id]
async def put_back_queue(self, request_id: str, data: object) -> bool:
"""Write a response while the request queue remains active.
Args:
request_id: Response request identifier.
data: Payload to enqueue.
Returns:
Whether an active response queue accepted the payload.
"""
queue = self.back_queues.get(request_id)
close_event = self._back_queue_close_events.get(request_id)
if queue is None or close_event is None or close_event.is_set():
return False
try:
queue.put_nowait(data)
return True
except asyncio.QueueFull:
pass
put_task = asyncio.create_task(queue.put(data))
close_task = asyncio.create_task(close_event.wait())
try:
done, _ = await asyncio.wait(
{put_task, close_task},
return_when=asyncio.FIRST_COMPLETED,
)
if close_task in done:
return False
await put_task
return True
finally:
for task in (put_task, close_task):
if not task.done():
task.cancel()
await asyncio.gather(put_task, close_task, return_exceptions=True)
def remove_back_queue(self, request_id: str):
"""Remove back queue for the given request ID"""
close_event = self._back_queue_close_events.pop(request_id, None)
if close_event is not None:
close_event.set()
self.back_queues.pop(request_id, None)
conversation_id = self._request_conversation.pop(request_id, None)
if conversation_id:
astrbot/dashboard/api/chat.py
+22
View File
@@ -178,6 +178,28 @@ async def stop_chat_session(
return await _run(lambda: service.stop_session(auth.username, session_id))
@router.get("/chat/runs/{run_id}/stream")
async def resume_chat_run(
run_id: str,
auth: AuthContext = Depends(require_chat_scope),
service: ChatService = Depends(get_service),
):
try:
stream = await service.build_chat_run_stream(auth.username, run_id)
except ChatServiceError:
return JSONResponse(error("Chat run is unavailable"))
return StreamingResponse(
stream,
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Transfer-Encoding": "chunked",
"Connection": "keep-alive",
},
)
@router.patch("/chat/sessions/{session_id}/messages/{message_id}")
async def update_chat_message(
session_id: str,
astrbot/dashboard/services/chat_service.py
+435 -265
View File
@@ -6,8 +6,8 @@ import os
import re
import uuid
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from copy import deepcopy
from dataclasses import dataclass, field
from pathlib import Path, PurePosixPath
from typing import Any
@@ -32,6 +32,7 @@ from astrbot.core.utils.media_utils import (
)
SSE_HEARTBEAT = ": heartbeat\n\n"
CHAT_RUN_SUBSCRIBER_QUEUE_SIZE = 256
def sanitize_upload_filename(filename: str | None) -> str:
@@ -44,29 +45,6 @@ def sanitize_upload_filename(filename: str | None) -> str:
return name
@asynccontextmanager
async def track_conversation(convs: dict, conv_id: str):
convs[conv_id] = True
try:
yield
finally:
convs.pop(conv_id, None)
async def poll_webchat_stream_result(back_queue, username: str):
try:
result = await asyncio.wait_for(back_queue.get(), timeout=1)
except asyncio.TimeoutError:
return None, False
except asyncio.CancelledError:
logger.debug(f"[WebChat] 用户 {username} 断开聊天长连接。")
return None, True
except Exception as e:
logger.error(f"WebChat stream error: {e}")
return None, False
return result, False
def normalize_reasoning_message_parts(
message_parts: list[dict] | None,
reasoning: str = "",
@@ -492,6 +470,25 @@ class ChatServiceError(Exception):
pass
@dataclass(slots=True)
class ChatRunState:
"""State owned by a WebChat generation independently of its subscribers."""
run_id: str
username: str
session_id: str
llm_checkpoint_id: str
platform_history_id: str
back_queue: asyncio.Queue
subscribers: set[asyncio.Queue] = field(default_factory=set)
message_parts: list[dict] = field(default_factory=list)
agent_stats: dict = field(default_factory=dict)
refs: dict = field(default_factory=dict)
revision: int = 0
status: str = "running"
task: asyncio.Task[None] | None = None
class ChatService:
def __init__(
self,
@@ -509,6 +506,8 @@ class ChatService:
self.platform_history_mgr = core_lifecycle.platform_message_history_manager
self.umop_config_router = core_lifecycle.umop_config_router
self.running_convs: dict[str, bool] = {}
self.chat_runs: dict[str, ChatRunState] = {}
self.chat_runs_by_session: dict[str, set[str]] = {}
async def build_user_message_parts(self, message: str | list) -> list[dict]:
return await build_webchat_message_parts(
@@ -646,6 +645,14 @@ class ChatService:
for thread_id in thread_ids:
unified_msg_origin = build_thread_unified_msg_origin(creator, thread_id)
active_event_registry.request_agent_stop_all(unified_msg_origin)
tasks = []
for run_id in list(self.chat_runs_by_session.get(thread_id, set())):
run = self.chat_runs.get(run_id)
if run and run.task and not run.task.done():
run.task.cancel()
tasks.append(run.task)
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
await self.conv_mgr.delete_conversations_by_user_id(unified_msg_origin)
await self.platform_history_mgr.delete(
platform_id="webchat_thread",
@@ -724,6 +731,352 @@ class ChatService:
llm_checkpoint_id=llm_checkpoint_id,
)
def get_active_chat_runs(self, username: str, session_id: str) -> list[dict]:
"""Return resumable runs owned by a user in one chat session.
Args:
username: Authenticated run owner.
session_id: WebChat session or thread identifier.
Returns:
Active run snapshots in creation order.
"""
snapshots = []
for run in self.chat_runs.values():
if run.username != username or run.session_id != session_id:
continue
snapshots.append(
{
"run_id": run.run_id,
"session_id": run.session_id,
"llm_checkpoint_id": run.llm_checkpoint_id,
"status": run.status,
"revision": run.revision,
"content": build_bot_history_content(
deepcopy(run.message_parts),
agent_stats=deepcopy(run.agent_stats),
refs=deepcopy(run.refs),
),
}
)
return snapshots
@staticmethod
def _publish_chat_run(run: ChatRunState, payload: dict) -> None:
"""Publish one output event without coupling the run to subscribers.
Args:
run: Chat run producing the event.
payload: Existing WebChat event payload.
"""
run.revision += 1
item = (run.revision, payload)
for subscriber in list(run.subscribers):
try:
subscriber.put_nowait(item)
except asyncio.QueueFull:
# End slow streams so they can reconnect from a fresh snapshot.
run.subscribers.discard(subscriber)
while not subscriber.empty():
subscriber.get_nowait()
subscriber.put_nowait(None)
def _subscribe_chat_run(
self,
run: ChatRunState,
*,
include_snapshot: bool,
saved_user_record=None,
) -> AsyncIterator[str]:
"""Create an SSE subscriber for a running chat generation.
Args:
run: Chat run to observe.
include_snapshot: Whether to begin with accumulated run state.
saved_user_record: Newly persisted user record for the legacy stream.
Returns:
SSE iterator detached from the generation task lifecycle.
"""
subscriber: asyncio.Queue = asyncio.Queue(
maxsize=CHAT_RUN_SUBSCRIBER_QUEUE_SIZE
)
run.subscribers.add(subscriber)
snapshot = None
if include_snapshot:
snapshot = {
"run_id": run.run_id,
"session_id": run.session_id,
"llm_checkpoint_id": run.llm_checkpoint_id,
"status": run.status,
"revision": run.revision,
"content": build_bot_history_content(
deepcopy(run.message_parts),
agent_stats=deepcopy(run.agent_stats),
refs=deepcopy(run.refs),
),
}
snapshot_revision = run.revision
async def stream():
try:
if snapshot is not None:
payload = {"type": "run_snapshot", "data": snapshot}
yield f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
else:
session_info = {
"type": "session_id",
"data": None,
"session_id": run.session_id,
}
yield f"data: {json.dumps(session_info, ensure_ascii=False)}\n\n"
if saved_user_record:
user_saved_info = {
"type": "user_message_saved",
"data": {
"id": saved_user_record.id,
"created_at": to_utc_isoformat(
saved_user_record.created_at
),
"llm_checkpoint_id": run.llm_checkpoint_id,
},
}
yield f"data: {json.dumps(user_saved_info, ensure_ascii=False)}\n\n"
while True:
try:
item = await asyncio.wait_for(subscriber.get(), timeout=1)
except asyncio.TimeoutError:
yield SSE_HEARTBEAT
continue
if item is None:
break
revision, payload = item
if include_snapshot and revision <= snapshot_revision:
continue
yield f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
finally:
run.subscribers.discard(subscriber)
return stream()
async def build_chat_run_stream(
self,
username: str,
run_id: str,
) -> AsyncIterator[str]:
"""Attach a new SSE subscriber to an active chat run.
Args:
username: Authenticated run owner.
run_id: Active run identifier.
Returns:
SSE iterator beginning with a full accumulated snapshot.
Raises:
ChatServiceError: If the run is absent or owned by another user.
"""
run = self.chat_runs.get(run_id)
if run is None:
raise ChatServiceError(f"Chat run {run_id} not found")
if run.username != username:
raise ChatServiceError("Permission denied")
return self._subscribe_chat_run(run, include_snapshot=True)
async def _consume_chat_run(self, run: ChatRunState) -> None:
"""Drain runner output, persist it, and fan it out to subscribers.
Args:
run: Chat run owning the producer queue and durable state.
"""
pending_accumulator = BotMessageAccumulator()
display_accumulator = BotMessageAccumulator()
pending_agent_stats = {}
pending_refs = {}
async def flush_pending_bot_message():
nonlocal pending_accumulator, pending_agent_stats, pending_refs
if not (
pending_accumulator.has_content() or pending_refs or pending_agent_stats
):
return None
message_parts_to_save = pending_accumulator.build_message_parts(
include_pending_tool_calls=True
)
plain_text = collect_plain_text_from_message_parts(message_parts_to_save)
try:
extracted_refs = extract_web_search_refs(
plain_text,
message_parts_to_save,
)
except Exception as exc:
logger.exception(
f"Failed to extract web search refs: {exc}",
exc_info=True,
)
extracted_refs = pending_refs
run.refs = extracted_refs
saved_record = await self.save_bot_message(
run.session_id,
message_parts_to_save,
pending_agent_stats,
extracted_refs,
run.llm_checkpoint_id,
run.platform_history_id,
)
pending_accumulator = BotMessageAccumulator()
pending_agent_stats = {}
pending_refs = {}
return saved_record
self.running_convs[run.session_id] = True
try:
while True:
result = await run.back_queue.get()
if not result:
continue
if result.get("message_id") and str(result["message_id"]) != run.run_id:
logger.warning("webchat stream message_id mismatch")
continue
result_text = result.get("data", "")
msg_type = result.get("type")
streaming = result.get("streaming", False)
chain_type = result.get("chain_type")
if chain_type == "agent_stats":
try:
run.agent_stats = json.loads(result_text)
except (TypeError, json.JSONDecodeError):
run.agent_stats = {}
pending_agent_stats = run.agent_stats
self._publish_chat_run(
run,
{"type": "agent_stats", "data": run.agent_stats},
)
continue
attachment_saved_payload = None
if msg_type == "plain":
for accumulator in (pending_accumulator, display_accumulator):
accumulator.add_plain(
result_text,
chain_type=chain_type,
streaming=streaming,
)
elif msg_type in {"image", "record", "file", "video"}:
prefix = {
"image": "[IMAGE]",
"record": "[RECORD]",
"file": "[FILE]",
"video": "[VIDEO]",
}[msg_type]
filename = str(result_text).replace(prefix, "", 1)
display_name = None
if msg_type in {"file", "video"} and "|" in filename:
filename, display_name = filename.split("|", 1)
part = await self.create_attachment_from_file(
filename,
msg_type,
display_name=display_name,
)
for accumulator in (pending_accumulator, display_accumulator):
accumulator.add_attachment(part)
if part and part.get("attachment_id") and part.get("type"):
attachment_saved_payload = {
"type": "attachment_saved",
"data": {
"id": part["attachment_id"],
"type": part["type"],
},
}
snapshot_accumulator = deepcopy(display_accumulator)
run.message_parts = snapshot_accumulator.build_message_parts(
include_pending_tool_calls=True
)
self._publish_chat_run(run, result)
if attachment_saved_payload:
self._publish_chat_run(run, attachment_saved_payload)
should_save = False
if msg_type == "end":
should_save = bool(
pending_accumulator.has_content()
or pending_refs
or pending_agent_stats
)
elif (streaming and msg_type == "complete") or not streaming:
if chain_type not in ("tool_call", "tool_call_result"):
should_save = True
if should_save:
saved_record = await flush_pending_bot_message()
if saved_record:
self._publish_chat_run(
run,
{
"type": "message_saved",
"data": {
"id": saved_record.id,
"created_at": to_utc_isoformat(
saved_record.created_at
),
"llm_checkpoint_id": run.llm_checkpoint_id,
},
},
)
if msg_type == "end":
run.status = "completed"
break
except asyncio.CancelledError:
run.status = "stopped"
except Exception as exc:
run.status = "failed"
logger.exception(f"WebChat run unexpected error: {exc}", exc_info=True)
self._publish_chat_run(
run,
{"type": "error", "data": "WebChat run failed"},
)
finally:
try:
saved_record = await asyncio.shield(flush_pending_bot_message())
if saved_record:
self._publish_chat_run(
run,
{
"type": "message_saved",
"data": {
"id": saved_record.id,
"created_at": to_utc_isoformat(saved_record.created_at),
"llm_checkpoint_id": run.llm_checkpoint_id,
},
},
)
except Exception as exc:
logger.exception(
f"Failed to persist pending webchat message: {exc}",
exc_info=True,
)
webchat_queue_mgr.remove_back_queue(run.run_id)
if self.chat_runs.get(run.run_id) is run:
self.chat_runs.pop(run.run_id, None)
run_ids = self.chat_runs_by_session.get(run.session_id)
if run_ids is not None:
run_ids.discard(run.run_id)
if not run_ids:
self.chat_runs_by_session.pop(run.session_id, None)
self.running_convs.pop(run.session_id, None)
for subscriber in list(run.subscribers):
while not subscriber.empty():
subscriber.get_nowait()
subscriber.put_nowait(None)
run.subscribers.clear()
async def build_chat_stream(
self,
username: str,
@@ -755,248 +1108,8 @@ class ChatService:
message_id = str(uuid.uuid4())
llm_checkpoint_id = post_data.get("_llm_checkpoint_id") or str(uuid.uuid4())
skip_user_history = bool(post_data.get("_skip_user_history"))
back_queue = webchat_queue_mgr.get_or_create_back_queue(
message_id,
webchat_conv_id,
)
saved_user_record = None
async def stream():
client_disconnected = False
message_accumulator = BotMessageAccumulator()
agent_stats = {}
refs = {}
async def flush_pending_bot_message():
nonlocal message_accumulator, agent_stats, refs
if not (message_accumulator.has_content() or refs or agent_stats):
return None
message_parts_to_save = message_accumulator.build_message_parts(
include_pending_tool_calls=True
)
plain_text = collect_plain_text_from_message_parts(
message_parts_to_save
)
try:
extracted_refs = extract_web_search_refs(
plain_text,
message_parts_to_save,
)
except Exception as e:
logger.exception(
f"Failed to extract web search refs: {e}",
exc_info=True,
)
extracted_refs = refs
saved_record = await self.save_bot_message(
webchat_conv_id,
message_parts_to_save,
agent_stats,
extracted_refs,
llm_checkpoint_id,
platform_history_id,
)
message_accumulator = BotMessageAccumulator()
agent_stats = {}
refs = {}
return saved_record
def build_attachment_saved_event(part: dict | None) -> str | None:
if not part or not part.get("attachment_id") or not part.get("type"):
return None
payload = {
"type": "attachment_saved",
"data": {
"id": part["attachment_id"],
"type": part["type"],
},
}
return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
try:
session_info = {
"type": "session_id",
"data": None,
"session_id": webchat_conv_id,
}
yield f"data: {json.dumps(session_info, ensure_ascii=False)}\n\n"
if saved_user_record and not client_disconnected:
user_saved_info = {
"type": "user_message_saved",
"data": {
"id": saved_user_record.id,
"created_at": to_utc_isoformat(
saved_user_record.created_at
),
"llm_checkpoint_id": llm_checkpoint_id,
},
}
yield f"data: {json.dumps(user_saved_info, ensure_ascii=False)}\n\n"
async with track_conversation(self.running_convs, webchat_conv_id):
while True:
result, should_break = await poll_webchat_stream_result(
back_queue, username
)
if should_break:
client_disconnected = True
break
if not result:
if not client_disconnected:
yield SSE_HEARTBEAT
continue
if (
"message_id" in result
and result["message_id"] != message_id
):
logger.warning("webchat stream message_id mismatch")
continue
result_text = result["data"]
msg_type = result.get("type")
streaming = result.get("streaming", False)
chain_type = result.get("chain_type")
if chain_type == "agent_stats":
stats_info = {
"type": "agent_stats",
"data": json.loads(result_text),
}
yield f"data: {json.dumps(stats_info, ensure_ascii=False)}\n\n"
agent_stats = stats_info["data"]
continue
try:
if not client_disconnected:
yield f"data: {json.dumps(result, ensure_ascii=False)}\n\n"
except Exception as e:
if not client_disconnected:
logger.debug(
f"[WebChat] 用户 {username} 断开聊天长连接。 {e}"
)
client_disconnected = True
try:
if not client_disconnected:
await asyncio.sleep(0.05)
except asyncio.CancelledError:
logger.debug(f"[WebChat] 用户 {username} 断开聊天长连接。")
client_disconnected = True
if msg_type == "plain":
message_accumulator.add_plain(
result_text,
chain_type=chain_type,
streaming=streaming,
)
elif msg_type == "image":
filename = result_text.replace("[IMAGE]", "")
part = await self.create_attachment_from_file(
filename, "image"
)
message_accumulator.add_attachment(part)
if attachment_saved_event := build_attachment_saved_event(
part
):
yield attachment_saved_event
elif msg_type == "record":
filename = result_text.replace("[RECORD]", "")
part = await self.create_attachment_from_file(
filename, "record"
)
message_accumulator.add_attachment(part)
if attachment_saved_event := build_attachment_saved_event(
part
):
yield attachment_saved_event
elif msg_type == "file":
filename = result_text.replace("[FILE]", "", 1)
display_name = None
if "|" in filename:
filename, display_name = filename.split("|", 1)
part = await self.create_attachment_from_file(
filename,
"file",
display_name=display_name,
)
message_accumulator.add_attachment(part)
if attachment_saved_event := build_attachment_saved_event(
part
):
yield attachment_saved_event
elif msg_type == "video":
filename = result_text.replace("[VIDEO]", "")
part = await self.create_attachment_from_file(
filename, "video"
)
message_accumulator.add_attachment(part)
if attachment_saved_event := build_attachment_saved_event(
part
):
yield attachment_saved_event
should_save = False
if msg_type == "end":
should_save = message_accumulator.has_content() or bool(
refs or agent_stats
)
elif (streaming and msg_type == "complete") or not streaming:
if chain_type not in ("tool_call", "tool_call_result"):
should_save = True
if should_save:
saved_record = await flush_pending_bot_message()
if saved_record and not client_disconnected:
saved_info = {
"type": "message_saved",
"data": {
"id": saved_record.id,
"created_at": to_utc_isoformat(
saved_record.created_at
),
"llm_checkpoint_id": llm_checkpoint_id,
},
}
try:
yield f"data: {json.dumps(saved_info, ensure_ascii=False)}\n\n"
except Exception:
pass
if msg_type == "end":
break
except BaseException as e:
logger.exception(f"WebChat stream unexpected error: {e}", exc_info=True)
finally:
try:
await flush_pending_bot_message()
except Exception as e:
logger.exception(
f"Failed to persist pending webchat message: {e}",
exc_info=True,
)
webchat_queue_mgr.remove_back_queue(message_id)
chat_queue = webchat_queue_mgr.get_or_create_queue(webchat_conv_id)
await chat_queue.put(
(
username,
webchat_conv_id,
{
"message": message_parts,
"selected_provider": selected_provider,
"selected_model": selected_model,
"enable_streaming": enable_streaming,
"message_id": message_id,
"llm_checkpoint_id": llm_checkpoint_id,
"thread_selected_text": thread_selected_text,
},
),
)
message_parts_for_storage = strip_message_parts_path_fields(message_parts)
if not skip_user_history:
saved_user_record = await self.platform_history_mgr.insert(
@@ -1008,7 +1121,53 @@ class ChatService:
llm_checkpoint_id=llm_checkpoint_id,
)
return stream()
back_queue = webchat_queue_mgr.get_or_create_back_queue(
message_id,
webchat_conv_id,
)
run = ChatRunState(
run_id=message_id,
username=username,
session_id=webchat_conv_id,
llm_checkpoint_id=llm_checkpoint_id,
platform_history_id=platform_history_id,
back_queue=back_queue,
)
self.chat_runs[message_id] = run
self.chat_runs_by_session.setdefault(webchat_conv_id, set()).add(message_id)
stream = self._subscribe_chat_run(
run,
include_snapshot=False,
saved_user_record=saved_user_record,
)
run.task = asyncio.create_task(
self._consume_chat_run(run),
name=f"webchat_run_{message_id}",
)
try:
chat_queue = webchat_queue_mgr.get_or_create_queue(webchat_conv_id)
await chat_queue.put(
(
username,
webchat_conv_id,
{
"message": message_parts,
"selected_provider": selected_provider,
"selected_model": selected_model,
"enable_streaming": enable_streaming,
"message_id": message_id,
"llm_checkpoint_id": llm_checkpoint_id,
"thread_selected_text": thread_selected_text,
},
),
)
except BaseException:
run.task.cancel()
await asyncio.gather(run.task, return_exceptions=True)
raise
return stream
async def stop_session(self, username: str, session_id: str) -> dict:
session = await self.db.get_platform_session_by_id(session_id)
@@ -1039,6 +1198,15 @@ class ChatService:
f"{session.platform_id}:{message_type}:"
f"{session.platform_id}!{username}!{session_id}"
)
active_event_registry.request_agent_stop_all(unified_msg_origin)
tasks = []
for run_id in list(self.chat_runs_by_session.get(session_id, set())):
run = self.chat_runs.get(run_id)
if run and run.task and not run.task.done():
run.task.cancel()
tasks.append(run.task)
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
await self.conv_mgr.delete_conversations_by_user_id(unified_msg_origin)
history_list = await self.platform_history_mgr.get(
@@ -1237,6 +1405,7 @@ class ChatService:
"history": [serialize_history_entry(history) for history in history_ls],
"threads": [serialize_thread(thread) for thread in threads],
"is_running": self.running_convs.get(session_id, False),
"active_runs": self.get_active_chat_runs(username, session_id),
}
if project_info:
response_data["project"] = {
@@ -1352,6 +1521,7 @@ class ChatService:
"thread": serialize_thread(thread),
"history": [serialize_history_entry(history) for history in history_ls],
"is_running": self.running_convs.get(thread_id, False),
"active_runs": self.get_active_chat_runs(username, thread_id),
}
async def get_thread_from_dashboard_query(
dashboard/src/api/generated/openapi-v1/sdk.gen.ts
+11 -1
View File
File diff suppressed because one or more lines are too long
dashboard/src/api/generated/openapi-v1/types.gen.ts
+10
View File
@@ -1371,6 +1371,16 @@ export type StopChatSessionResponse = (SuccessEnvelope);
export type StopChatSessionError = unknown;
export type ResumeChatRunData = {
path: {
run_id: string;
};
};
export type ResumeChatRunResponse = (unknown);
export type ResumeChatRunError = unknown;
export type UpdateChatMessageData = {
body: ChatMessagePatchRequest;
path: {
dashboard/src/api/v1.ts
+3
View File
@@ -795,6 +795,9 @@ export const chatApi = {
sendStreamUrl() {
return '/api/v1/chat';
},
resumeRunStreamUrl(runId: string) {
return `/api/v1/chat/runs/${encodeURIComponent(runId)}/stream`;
},
liveWebSocketUrl(token: string, host = window.location.host) {
const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
return `${protocol}//${host}/api/v1/live-chat/ws?token=${encodeURIComponent(token)}`;
dashboard/src/composables/useMessages.ts
+203 -32
View File
@@ -69,6 +69,15 @@ export interface ChatSessionProject {
emoji?: string;
}
interface ActiveChatRun {
run_id: string;
session_id: string;
llm_checkpoint_id?: string | null;
status?: string;
revision?: number;
content?: ChatContent;
}
interface ActiveConnection {
sessionId: string;
messageId: string;
@@ -226,9 +235,13 @@ export function useMessages(options: UseMessagesOptions) {
await Promise.all(tasks);
}
async function loadSessionMessages(sessionId: string) {
async function loadSessionMessages(
sessionId: string,
resumeRuns = true,
showLoading = true,
) {
if (!sessionId) return;
loadingMessages.value = true;
if (showLoading) loadingMessages.value = true;
try {
const response = await chatApi.getSession(sessionId);
const payload = response.data?.data || {};
@@ -239,14 +252,47 @@ export function useMessages(options: UseMessagesOptions) {
messagesBySession[sessionId] = records;
sessionProjects[sessionId] = normalizeSessionProject(payload.project);
loadedSessions[sessionId] = true;
if (resumeRuns && Array.isArray(payload.active_runs)) {
await restoreNextActiveRun(sessionId, payload.active_runs);
}
} catch (error) {
console.error("Failed to load session messages:", error);
messagesBySession[sessionId] = messagesBySession[sessionId] || [];
} finally {
loadingMessages.value = false;
if (showLoading) loadingMessages.value = false;
}
}
async function restoreNextActiveRun(
sessionId: string,
activeRuns: ActiveChatRun[],
) {
const run = activeRuns[0];
if (!run?.run_id || activeConnections[sessionId]) return;
const checkpointId = run.llm_checkpoint_id || null;
const records = (messagesBySession[sessionId] || []).filter((record) => {
return !(
checkpointId &&
record.llm_checkpoint_id === checkpointId &&
messageContent(record).type === "bot"
);
});
const botRecord = normalizeHistoryRecord({
id: `active-run-${run.run_id}`,
content: run.content || { type: "bot", message: [] },
llm_checkpoint_id: checkpointId,
created_at: new Date().toISOString(),
});
botRecord.content.isLoading = botRecord.content.message.length === 0;
records.push(botRecord);
messagesBySession[sessionId] = records;
const restoredRecords = messagesBySession[sessionId];
const reactiveBotRecord = restoredRecords[restoredRecords.length - 1];
await resolveRecordMedia([reactiveBotRecord]);
startResumeStream(sessionId, run.run_id, reactiveBotRecord);
}
function createLocalExchange({
sessionId,
messageId,
@@ -335,7 +381,9 @@ export function useMessages(options: UseMessagesOptions) {
content: content as unknown as Record<string, unknown>,
});
const payload = response.data?.data || {};
const updated = payload.message ? normalizeHistoryRecord(payload.message) : null;
const updated = payload.message
? normalizeHistoryRecord(payload.message)
: null;
if (updated) {
Object.assign(record, updated);
await resolveRecordMedia([record]);
@@ -424,17 +472,20 @@ export function useMessages(options: UseMessagesOptions) {
};
try {
const response = await fetchWithAuth(chatApi.regenerateMessageUrl(sessionId, targetMessageId), {
method: "POST",
headers: {
"Content-Type": "application/json",
const response = await fetchWithAuth(
chatApi.regenerateMessageUrl(sessionId, targetMessageId),
{
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify({
selected_provider: selectedProvider,
selected_model: selectedModel,
}),
signal: abort.signal,
},
body: JSON.stringify({
selected_provider: selectedProvider,
selected_model: selectedModel,
}),
signal: abort.signal,
});
);
if (!response.ok || !response.body) {
throw new Error(`Regenerate failed: ${response.status}`);
}
@@ -449,7 +500,10 @@ export function useMessages(options: UseMessagesOptions) {
});
} catch (error) {
if (!abort.signal.aborted) {
appendPlain(botRecord, `\n\n${String((error as Error)?.message || error)}`);
appendPlain(
botRecord,
`\n\n${String((error as Error)?.message || error)}`,
);
console.error("Regenerate failed:", error);
}
} finally {
@@ -485,7 +539,10 @@ export function useMessages(options: UseMessagesOptions) {
const normalizedContent: ChatContent = {
type: content.type || (record.sender_id === "bot" ? "bot" : "user"),
message: normalizedMessage,
reasoning: extractReasoningText(normalizedMessage, content.reasoning || ""),
reasoning: extractReasoningText(
normalizedMessage,
content.reasoning || "",
),
agentStats: content.agentStats || content.agent_stats,
refs: content.refs,
};
@@ -566,6 +623,88 @@ export function useMessages(options: UseMessagesOptions) {
});
}
function startResumeStream(
sessionId: string,
runId: string,
botRecord: ChatRecord,
) {
const abort = new AbortController();
activeConnections[sessionId] = {
sessionId,
messageId: runId,
transport: "sse",
abort,
botRecord,
};
void (async () => {
let receivedEnd = false;
let lastError: unknown = null;
for (
let attempt = 0;
attempt < 5 && !abort.signal.aborted;
attempt += 1
) {
let retryable = true;
try {
const response = await fetchWithAuth(
chatApi.resumeRunStreamUrl(runId),
{
headers: { Accept: "text/event-stream" },
signal: abort.signal,
},
);
const contentType = response.headers.get("content-type") || "";
if (
!response.ok ||
!response.body ||
!contentType.includes("text/event-stream")
) {
retryable = response.status >= 500;
throw new Error(`Resume stream failed: ${response.status}`);
}
await readSseStream(response.body, (payload) => {
processStreamPayload(botRecord, payload);
options.onStreamUpdate?.(sessionId);
const payloadType = payload?.type || payload?.t;
if (payloadType === "end") receivedEnd = true;
});
if (receivedEnd) break;
lastError = new Error("Resume stream closed before completion.");
} catch (error) {
if (abort.signal.aborted) return;
lastError = error;
}
if (!retryable || attempt === 4 || abort.signal.aborted) break;
await new Promise<void>((resolve) => {
const timeout = window.setTimeout(resolve, 250 * 2 ** attempt);
abort.signal.addEventListener(
"abort",
() => {
window.clearTimeout(timeout);
resolve();
},
{ once: true },
);
});
}
if (!receivedEnd && lastError && !abort.signal.aborted) {
console.error("Resume chat stream failed:", lastError);
}
const ownsConnection = activeConnections[sessionId]?.abort === abort;
if (ownsConnection) delete activeConnections[sessionId];
if (!abort.signal.aborted && ownsConnection) {
await loadSessionMessages(sessionId, true, false);
await options.onSessionsChanged?.();
}
})();
}
function startWebSocketStream(
sessionId: string,
messageId: string,
@@ -694,7 +833,11 @@ export function useMessages(options: UseMessagesOptions) {
try {
const payload = JSON.parse(event.data);
processStreamPayload(connection.botRecord, payload, connection.userRecord);
processStreamPayload(
connection.botRecord,
payload,
connection.userRecord,
);
options.onStreamUpdate?.(sessionId);
if (payload.type === "end" || payload.t === "end") {
void finishWebSocketStream(sessionId, connection.messageId);
@@ -741,6 +884,21 @@ export function useMessages(options: UseMessagesOptions) {
const data = normalized?.data ?? "";
if (msgType === "session_id" || msgType === "session_bound") return;
if (msgType === "run_snapshot") {
const snapshot = data && typeof data === "object" ? data : {};
const snapshotRecord = normalizeHistoryRecord({
id: `active-run-${snapshot.run_id || "unknown"}`,
content: snapshot.content || { type: "bot", message: [] },
llm_checkpoint_id: snapshot.llm_checkpoint_id || null,
});
snapshotRecord.content.isLoading =
snapshot.status === "running" &&
snapshotRecord.content.message.length === 0;
botRecord.content = snapshotRecord.content;
botRecord.llm_checkpoint_id = snapshotRecord.llm_checkpoint_id;
void resolveRecordMedia([botRecord]);
return;
}
if (msgType === "user_message_saved") {
if (userRecord) {
userRecord.id = data?.id || userRecord.id;
@@ -811,9 +969,13 @@ export function useMessages(options: UseMessagesOptions) {
.replace("[VIDEO]", "");
const separatorIndex = rawFilename.indexOf("|");
const storedFilename =
separatorIndex >= 0 ? rawFilename.slice(0, separatorIndex) : rawFilename;
separatorIndex >= 0
? rawFilename.slice(0, separatorIndex)
: rawFilename;
const displayFilename =
separatorIndex >= 0 ? rawFilename.slice(separatorIndex + 1) : storedFilename;
separatorIndex >= 0
? rawFilename.slice(separatorIndex + 1)
: storedFilename;
const filename = displayFilename || storedFilename;
const mediaPart: MessagePart = { type: msgType, filename };
if (storedFilename && storedFilename !== filename) {
@@ -905,7 +1067,10 @@ export function normalizeMessageParts(
fallbackReasoning = "",
): MessagePart[] {
const normalizedParts = normalizePartsInternal(parts);
if (fallbackReasoning && !normalizedParts.some((part) => part.type === "think")) {
if (
fallbackReasoning &&
!normalizedParts.some((part) => part.type === "think")
) {
normalizedParts.unshift({ type: "think", think: fallbackReasoning });
}
return normalizedParts;
@@ -951,16 +1116,18 @@ export function reasoningActivityTitle(
counts: ReturnType<typeof reasoningActivityCounts>,
tm: (key: string, params?: Record<string, string | number>) => string,
) {
return [
counts.thinkCount > 0
? tm("reasoning.thinkSummary", { count: counts.thinkCount })
: "",
counts.toolCount > 0
? tm("reasoning.toolSummary", { count: counts.toolCount })
: "",
]
.filter(Boolean)
.join(tm("reasoning.summarySeparator")) || tm("reasoning.thinking");
return (
[
counts.thinkCount > 0
? tm("reasoning.thinkSummary", { count: counts.thinkCount })
: "",
counts.toolCount > 0
? tm("reasoning.toolSummary", { count: counts.toolCount })
: "",
]
.filter(Boolean)
.join(tm("reasoning.summarySeparator")) || tm("reasoning.thinking")
);
}
export function thinkingParts(content: ChatContent): MessagePart[] {
@@ -1132,7 +1299,8 @@ export function upsertToolCall(record: ChatRecord, toolCall: any) {
const targetId = toolCall.id;
if (targetId != null) {
for (const part of record.content.message) {
if (part.type !== "tool_call" || !Array.isArray(part.tool_calls)) continue;
if (part.type !== "tool_call" || !Array.isArray(part.tool_calls))
continue;
const matched = part.tool_calls.find((item) => item.id === targetId);
if (matched) {
Object.assign(matched, toolCall);
@@ -1140,7 +1308,10 @@ export function upsertToolCall(record: ChatRecord, toolCall: any) {
}
}
}
record.content.message.push({ type: "tool_call", tool_calls: [{ ...toolCall }] });
record.content.message.push({
type: "tool_call",
tool_calls: [{ ...toolCall }],
});
}
export function finishToolCall(record: ChatRecord, result: any) {
docs/public/openapi.json
+25
View File
@@ -1799,6 +1799,31 @@
}
}
},
"/api/v1/chat/runs/{run_id}/stream": {
"get": {
"tags": [
"Chat"
],
"summary": "Resume an active webchat run as an SSE stream",
"operationId": "resumeChatRun",
"x-astrbot-scope": "chat",
"parameters": [
{
"name": "run_id",
"in": "path",
"required": true,
"schema": {
"type": "string"
}
}
],
"responses": {
"200": {
"description": "Resumed server-sent chat stream or an error envelope"
}
}
}
},
"/api/v1/chat/sessions/{session_id}/messages/{message_id}": {
"patch": {
"tags": [
openspec/openapi-v1.yaml
+16
View File
@@ -1267,6 +1267,22 @@ paths:
"200":
$ref: "#/components/responses/Ok"
/api/v1/chat/runs/{run_id}/stream:
get:
tags: [Chat]
summary: Resume an active webchat run as an SSE stream
operationId: resumeChatRun
x-astrbot-scope: chat
parameters:
- name: run_id
in: path
required: true
schema:
type: string
responses:
"200":
description: Resumed server-sent chat stream or an error envelope
/api/v1/chat/sessions/{session_id}/messages/{message_id}:
patch:
tags: [Chat]
tests/test_astr_agent_run_util.py
+74
View File
@@ -0,0 +1,74 @@
from types import SimpleNamespace
import pytest
from astrbot.core.agent.response import AgentResponse
from astrbot.core.astr_agent_run_util import run_agent
from astrbot.core.message.message_event_result import MessageChain
class _FakeEvent:
"""Minimal event surface used by the agent stream bridge."""
def is_stopped(self) -> bool:
return False
def get_extra(self, key: str):
del key
return None
def get_platform_name(self) -> str:
return "test"
class _StreamingErrorRunner:
"""Agent runner that finishes with one provider error response."""
streaming = True
req = None
def __init__(self, error_text: str) -> None:
self.error_text = error_text
self.finished = False
self.run_context = SimpleNamespace(context=SimpleNamespace(event=_FakeEvent()))
async def step(self):
self.finished = True
yield AgentResponse(
type="err",
data={"chain": MessageChain().message(self.error_text)},
)
def done(self) -> bool:
return self.finished
class _MalformedStreamingErrorRunner(_StreamingErrorRunner):
"""Agent runner that returns an invalid provider error payload."""
async def step(self):
self.finished = True
yield AgentResponse(type="err", data={})
@pytest.mark.asyncio
async def test_run_agent_forwards_streaming_provider_error():
error_text = (
"LLM 响应错误: Not found the model k2.7-code-highspeed or Permission denied"
)
runner = _StreamingErrorRunner(error_text)
chains = [chain async for chain in run_agent(runner)]
assert len(chains) == 1
assert chains[0].get_plain_text() == error_text
@pytest.mark.asyncio
async def test_run_agent_replaces_malformed_streaming_provider_error():
runner = _MalformedStreamingErrorRunner("unused")
chains = [chain async for chain in run_agent(runner)]
assert len(chains) == 1
assert chains[0].get_plain_text() == "Error occurred during AI execution."
tests/test_chat_route.py
+361 -30
View File
@@ -1,56 +1,387 @@
import asyncio
import json
from datetime import UTC, datetime
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
from astrbot.dashboard.services.chat_service import poll_webchat_stream_result
from astrbot.dashboard.api.chat import resume_chat_run
from astrbot.dashboard.services import chat_service
from astrbot.dashboard.services.chat_service import ChatService, ChatServiceError
class _QueueThatRaises:
def __init__(self, exc: BaseException):
self._exc = exc
async def get(self):
raise self._exc
@pytest.fixture
def chat_service_instance(monkeypatch, tmp_path):
"""Create a ChatService with isolated persistence dependencies."""
monkeypatch.setattr(chat_service, "get_astrbot_data_path", lambda: str(tmp_path))
platform_history_mgr = Mock()
platform_history_mgr.insert = AsyncMock(
return_value=SimpleNamespace(
id=1,
created_at=datetime.now(UTC),
)
)
core_lifecycle = SimpleNamespace(
conversation_manager=Mock(),
platform_message_history_manager=platform_history_mgr,
umop_config_router=Mock(),
)
service = ChatService(Mock(), core_lifecycle)
service.build_user_message_parts = AsyncMock(
return_value=[{"type": "plain", "text": "hello"}]
)
service.save_bot_message = AsyncMock(
return_value=SimpleNamespace(
id=2,
created_at=datetime.now(UTC),
)
)
return service
class _QueueWithResult:
def __init__(self, result):
self._result = result
def _decode_sse_event(event: str) -> dict:
"""Decode one JSON SSE event emitted by ChatService.
async def get(self):
return self._result
Args:
event: Complete SSE event text.
Returns:
Decoded event payload.
"""
return json.loads(event.removeprefix("data: ").strip())
@pytest.mark.asyncio
async def test_poll_webchat_stream_result_breaks_on_cancelled_error():
result, should_break = await poll_webchat_stream_result(
_QueueThatRaises(asyncio.CancelledError()),
"alice",
async def test_resume_chat_run_does_not_expose_service_error():
service = SimpleNamespace(
build_chat_run_stream=AsyncMock(
side_effect=ChatServiceError("internal stack trace details")
)
)
auth = SimpleNamespace(username="alice")
assert result is None
assert should_break is True
response = await resume_chat_run("missing-run", auth, service)
assert response.status_code == 200
assert json.loads(response.body) == {
"status": "error",
"message": "Chat run is unavailable",
}
@pytest.mark.asyncio
async def test_poll_webchat_stream_result_continues_on_generic_exception():
result, should_break = await poll_webchat_stream_result(
_QueueThatRaises(RuntimeError("boom")),
async def test_chat_stream_disconnect_does_not_own_run_lifecycle(
chat_service_instance,
):
service = chat_service_instance
session_id = "disconnect-session"
stream = await service.build_chat_stream(
"alice",
{"message": "hello", "session_id": session_id},
)
run = next(iter(service.chat_runs.values()))
assert result is None
assert should_break is False
try:
assert _decode_sse_event(await anext(stream))["type"] == "session_id"
await stream.aclose()
assert not run.subscribers
assert run.task is not None and not run.task.done()
await chat_service.webchat_queue_mgr.put_back_queue(
run.run_id,
{
"type": "plain",
"data": "completed after refresh",
"streaming": True,
"message_id": run.run_id,
},
)
await chat_service.webchat_queue_mgr.put_back_queue(
run.run_id,
{
"type": "complete",
"data": "completed after refresh",
"streaming": True,
"message_id": run.run_id,
},
)
await chat_service.webchat_queue_mgr.put_back_queue(
run.run_id,
{
"type": "end",
"data": "",
"streaming": False,
"message_id": run.run_id,
},
)
await asyncio.wait_for(run.task, timeout=1)
saved_parts = service.save_bot_message.await_args.args[1]
assert saved_parts == [{"type": "plain", "text": "completed after refresh"}]
assert run.run_id not in service.chat_runs
finally:
if run.task and not run.task.done():
run.task.cancel()
await asyncio.gather(run.task, return_exceptions=True)
chat_service.webchat_queue_mgr.remove_queues(session_id)
@pytest.mark.asyncio
async def test_poll_webchat_stream_result_returns_queue_payload():
payload = {"type": "end", "data": ""}
result, should_break = await poll_webchat_stream_result(
_QueueWithResult(payload),
async def test_resumed_stream_starts_with_full_snapshot(chat_service_instance):
service = chat_service_instance
session_id = "resume-session"
legacy_stream = await service.build_chat_stream(
"alice",
{"message": "hello", "session_id": session_id},
)
run = next(iter(service.chat_runs.values()))
assert result == payload
assert should_break is False
try:
await anext(legacy_stream)
await legacy_stream.aclose()
await chat_service.webchat_queue_mgr.put_back_queue(
run.run_id,
{
"type": "plain",
"data": "before refresh",
"streaming": True,
"message_id": run.run_id,
},
)
for _ in range(10):
if run.message_parts:
break
await asyncio.sleep(0)
active_runs = service.get_active_chat_runs("alice", session_id)
assert [active_run["run_id"] for active_run in active_runs] == [run.run_id]
resumed_stream = await service.build_chat_run_stream("alice", run.run_id)
snapshot_event = _decode_sse_event(await anext(resumed_stream))
assert snapshot_event["type"] == "run_snapshot"
assert snapshot_event["data"]["content"]["message"] == [
{"type": "plain", "text": "before refresh"}
]
await chat_service.webchat_queue_mgr.put_back_queue(
run.run_id,
{
"type": "plain",
"data": " and after refresh",
"streaming": True,
"message_id": run.run_id,
},
)
next_event = _decode_sse_event(await asyncio.wait_for(anext(resumed_stream), 1))
assert next_event["data"] == " and after refresh"
await resumed_stream.aclose()
for payload in (
{
"type": "complete",
"data": "before refresh and after refresh",
"streaming": True,
"message_id": run.run_id,
},
{
"type": "end",
"data": "",
"streaming": False,
"message_id": run.run_id,
},
):
await chat_service.webchat_queue_mgr.put_back_queue(run.run_id, payload)
await asyncio.wait_for(run.task, timeout=1)
finally:
if run.task and not run.task.done():
run.task.cancel()
await asyncio.gather(run.task, return_exceptions=True)
chat_service.webchat_queue_mgr.remove_queues(session_id)
@pytest.mark.asyncio
async def test_active_chat_runs_keep_creation_order(chat_service_instance):
service = chat_service_instance
session_id = "ordered-runs-session"
streams = []
try:
streams.append(
await service.build_chat_stream(
"alice",
{"message": "first", "session_id": session_id},
)
)
first_run_id = next(iter(service.chat_runs))
streams.append(
await service.build_chat_stream(
"alice",
{"message": "follow-up", "session_id": session_id},
)
)
active_runs = service.get_active_chat_runs("alice", session_id)
assert active_runs[0]["run_id"] == first_run_id
assert len(active_runs) == 2
finally:
for stream in streams:
await stream.aclose()
tasks = [run.task for run in service.chat_runs.values() if run.task]
for task in tasks:
task.cancel()
await asyncio.gather(*tasks, return_exceptions=True)
chat_service.webchat_queue_mgr.remove_queues(session_id)
@pytest.mark.asyncio
async def test_slow_chat_run_subscriber_is_closed_at_buffer_limit(
chat_service_instance,
):
service = chat_service_instance
session_id = "slow-subscriber-session"
stream = await service.build_chat_stream(
"alice",
{"message": "hello", "session_id": session_id},
)
run = next(iter(service.chat_runs.values()))
subscriber = next(iter(run.subscribers))
try:
for index in range(chat_service.CHAT_RUN_SUBSCRIBER_QUEUE_SIZE + 1):
service._publish_chat_run(
run,
{"type": "plain", "data": str(index), "streaming": True},
)
assert subscriber.maxsize == chat_service.CHAT_RUN_SUBSCRIBER_QUEUE_SIZE
assert subscriber.qsize() == 1
assert not run.subscribers
assert _decode_sse_event(await anext(stream))["type"] == "session_id"
assert _decode_sse_event(await anext(stream))["type"] == "user_message_saved"
with pytest.raises(StopAsyncIteration):
await anext(stream)
finally:
await stream.aclose()
if run.task and not run.task.done():
run.task.cancel()
await asyncio.gather(run.task, return_exceptions=True)
chat_service.webchat_queue_mgr.remove_queues(session_id)
@pytest.mark.asyncio
async def test_resume_during_attachment_save_does_not_skip_attachment(
chat_service_instance,
):
service = chat_service_instance
session_id = "attachment-race-session"
legacy_stream = await service.build_chat_stream(
"alice",
{"message": "hello", "session_id": session_id},
)
run = next(iter(service.chat_runs.values()))
attachment_started = asyncio.Event()
release_attachment = asyncio.Event()
async def create_attachment(filename, attach_type, display_name=None):
"""Pause attachment persistence to exercise the resume race.
Args:
filename: Stored attachment filename.
attach_type: WebChat attachment type.
display_name: Optional client-facing filename.
Returns:
Persisted attachment metadata.
"""
del display_name
attachment_started.set()
await release_attachment.wait()
return {
"attachment_id": "attachment-1",
"filename": filename,
"type": attach_type,
}
service.create_attachment_from_file = create_attachment
try:
await anext(legacy_stream)
await legacy_stream.aclose()
await chat_service.webchat_queue_mgr.put_back_queue(
run.run_id,
{
"type": "image",
"data": "[IMAGE]result.png",
"streaming": True,
"message_id": run.run_id,
},
)
await asyncio.wait_for(attachment_started.wait(), timeout=1)
resumed_stream = await service.build_chat_run_stream("alice", run.run_id)
snapshot_event = _decode_sse_event(await anext(resumed_stream))
assert snapshot_event["data"]["content"]["message"] == []
release_attachment.set()
image_event = _decode_sse_event(
await asyncio.wait_for(anext(resumed_stream), timeout=1)
)
assert image_event["type"] == "image"
assert image_event["data"] == "[IMAGE]result.png"
await resumed_stream.aclose()
await chat_service.webchat_queue_mgr.put_back_queue(
run.run_id,
{
"type": "end",
"data": "",
"streaming": False,
"message_id": run.run_id,
},
)
await asyncio.wait_for(run.task, timeout=1)
finally:
release_attachment.set()
if run.task and not run.task.done():
run.task.cancel()
await asyncio.gather(run.task, return_exceptions=True)
chat_service.webchat_queue_mgr.remove_queues(session_id)
@pytest.mark.asyncio
async def test_legacy_chat_stream_keeps_existing_event_shape(chat_service_instance):
service = chat_service_instance
session_id = "legacy-session"
stream = await service.build_chat_stream(
"alice",
{"message": "hello", "session_id": session_id},
)
run = next(iter(service.chat_runs.values()))
try:
assert _decode_sse_event(await anext(stream)) == {
"type": "session_id",
"data": None,
"session_id": session_id,
}
assert _decode_sse_event(await anext(stream))["type"] == "user_message_saved"
plain_payload = {
"type": "plain",
"data": "unchanged",
"streaming": True,
"message_id": run.run_id,
}
await chat_service.webchat_queue_mgr.put_back_queue(
run.run_id,
plain_payload,
)
assert (
_decode_sse_event(await asyncio.wait_for(anext(stream), 1)) == plain_payload
)
finally:
await stream.aclose()
if run.task and not run.task.done():
run.task.cancel()
await asyncio.gather(run.task, return_exceptions=True)
chat_service.webchat_queue_mgr.remove_queues(session_id)
tests/test_platform_image_format_preservation.py
+7 -2
View File
@@ -62,11 +62,16 @@ async def test_webchat_image_attachment_uses_detected_extension(tmp_path, monkey
"data:image/png;base64," + base64.b64encode(image_buffer.getvalue()).decode()
)
queue = asyncio.Queue()
async def put_back_queue(_request_id, payload):
await queue.put(payload)
return True
monkeypatch.setattr(webchat_event, "attachments_dir", str(tmp_path))
monkeypatch.setattr(
webchat_event.webchat_queue_mgr,
"get_or_create_back_queue",
lambda *_args: queue,
"put_back_queue",
put_back_queue,
)
await webchat_event.WebChatMessageEvent._send(
tests/test_webchat_queue_lifecycle.py
+30
View File
@@ -0,0 +1,30 @@
import asyncio
import pytest
from astrbot.core.platform.sources.webchat.webchat_queue_mgr import WebChatQueueMgr
@pytest.mark.asyncio
async def test_removed_back_queue_unblocks_pending_writer():
queue_manager = WebChatQueueMgr(back_queue_maxsize=1)
request_id = "request-1"
queue = queue_manager.get_or_create_back_queue(request_id, "conversation-1")
await queue.put({"type": "plain", "data": "first"})
blocked_writer = asyncio.create_task(
queue_manager.put_back_queue(
request_id,
{"type": "plain", "data": "second"},
)
)
await asyncio.sleep(0)
assert not blocked_writer.done()
queue_manager.remove_back_queue(request_id)
assert await asyncio.wait_for(blocked_writer, timeout=1) is False
assert not await queue_manager.put_back_queue(
request_id,
{"type": "plain", "data": "late"},
)
tests/unit/test_webchat_message_parts.py
+7 -2
View File
@@ -16,6 +16,11 @@ from astrbot.core.platform.sources.webchat.message_parts_helper import (
async def test_webchat_file_send_keeps_original_filename(tmp_path, monkeypatch):
"""WebChat file payloads should carry both stored and display filenames."""
queue = asyncio.Queue()
async def put_back_queue(_request_id, payload):
await queue.put(payload)
return True
attachments_dir = tmp_path / "attachments"
attachments_dir.mkdir()
source_file = tmp_path / "source.txt"
@@ -23,8 +28,8 @@ async def test_webchat_file_send_keeps_original_filename(tmp_path, monkeypatch):
monkeypatch.setattr(webchat_event, "attachments_dir", str(attachments_dir))
monkeypatch.setattr(
webchat_event.webchat_queue_mgr,
"get_or_create_back_queue",
lambda *_args: queue,
"put_back_queue",
put_back_queue,
)
await webchat_event.WebChatMessageEvent._send(