fz-mirrors/AstrBot
1
0
Fork 0
mirror of https://github.com/AstrBotDevs/AstrBot synced 2026-07-19 10:22:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: handle ChatUI follow-up messages (#9309)

Browse Source
This commit is contained in:
Soulter
2026-07-18 16:27:39 +08:00
committed by GitHub
parent b853e9f41d
commit 991a8127fc
9 changed files with 629 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
astrbot/core/pipeline/process_stage/follow_up.py
View File

@@ -25,6 +25,7 @@ class FollowUpCapture:
ticket: FollowUpTicket
order_seq: int
monitor_task: asyncio.Task[None]
target_run_id: str | None = None
def _event_follow_up_text(event: AstrMessageEvent) -> str:
@@ -197,6 +198,9 @@ def try_capture_follow_up(event: AstrMessageEvent) -> FollowUpCapture | None:
ticket=ticket,
order_seq=order_seq,
monitor_task=monitor_task,
target_run_id=str(runner_event.message_obj.message_id)
if getattr(runner_event.message_obj, "message_id", None) is not None
else None,
)

4
astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py
View File

@@ -198,6 +198,10 @@ class InternalAgentSubStage(Stage):
follow_up_activated,
) = await prepare_follow_up_capture(follow_up_capture)
if follow_up_consumed_marked:
event.set_extra(
"_follow_up_captured",
{"target_run_id": follow_up_capture.target_run_id},
)
logger.info(
"Follow-up ticket already consumed, stopping processing. umo=%s, seq=%s",
event.unified_msg_origin,

26
astrbot/core/platform/sources/webchat/webchat_event.py
View File

@@ -164,9 +164,35 @@ class WebChatMessageEvent(AstrMessageEvent):
async def send(self, message: MessageChain | None) -> None:
message_id = self.message_obj.message_id
follow_up_capture = self.get_extra("_follow_up_captured")
if message is None and isinstance(follow_up_capture, dict):
request_id = str(message_id)
await webchat_queue_mgr.put_back_queue(
request_id,
{
"type": "follow_up_captured",
"data": follow_up_capture,
"streaming": False,
"message_id": message_id,
},
)
await WebChatMessageEvent._send(message_id, message, session_id=self.session_id)
await super().send(MessageChain([]))
async def send_typing(self) -> None:
"""Emit a run-start signal before an independent LLM request."""
message_id = self.message_obj.message_id
request_id = str(message_id)
await webchat_queue_mgr.put_back_queue(
request_id,
{
"type": "run_started",
"data": {"run_id": request_id},
"streaming": False,
"message_id": message_id,
},
)
async def send_streaming(self, generator, use_fallback: bool = False) -> None:
final_data = ""
reasoning_content = ""

119
astrbot/dashboard/services/live_chat_service.py
View File

@@ -56,6 +56,8 @@ class LiveChatSession:
self.temp_audio_path: str | None = None
self.chat_subscriptions: dict[str, str] = {}
self.chat_subscription_tasks: dict[str, asyncio.Task] = {}
self.chat_request_tasks: dict[str, asyncio.Task] = {}
self.interrupted_chat_requests: set[str] = set()
self.ws_send_lock = asyncio.Lock()
def start_speaking(self, stamp: str) -> None:
@@ -176,16 +178,73 @@ class LiveChatService:
live_session = self.create_session(username)
logger.info(f"[Live Chat] WebSocket 连接建立: {username}")
def finish_chat_request(
completed: asyncio.Task,
request_id: str,
) -> None:
"""Release a multiplexed request task and observe its exception.
Args:
completed: Finished WebSocket chat request task.
request_id: Request identifier used in the session task map.
"""
if live_session.chat_request_tasks.get(request_id) is completed:
live_session.chat_request_tasks.pop(request_id, None)
try:
completed.result()
except asyncio.CancelledError:
pass
except Exception as exc:
logger.error(
f"[Live Chat] WebSocket chat request failed: {exc}",
exc_info=True,
)
try:
while True:
message = await receive_json()
ct = force_ct or message.get("ct", "live")
if ct == "chat":
await self.handle_chat_message(
live_session,
message,
send_json,
)
if message.get("t") == "send":
request_id = str(message.get("message_id") or uuid.uuid4())
message["message_id"] = request_id
existing_task = live_session.chat_request_tasks.get(request_id)
if existing_task and not existing_task.done():
await self.send_chat_payload(
live_session,
{
"ct": "chat",
"t": "error",
"data": "Duplicate active message_id",
"code": "INVALID_MESSAGE_FORMAT",
"message_id": request_id,
},
send_json,
)
continue
task = asyncio.create_task(
self.handle_chat_message(
live_session,
message,
send_json,
),
name=f"chat_ws_request_{request_id}",
)
live_session.chat_request_tasks[request_id] = task
task.add_done_callback(
lambda completed, request_id=request_id: (
finish_chat_request(
completed,
request_id,
)
)
)
else:
await self.handle_chat_message(
live_session,
message,
send_json,
)
else:
await self.handle_live_message(
live_session,
@@ -356,7 +415,10 @@ class LiveChatService:
return request_id
async def cleanup_chat_subscriptions(self, session: LiveChatSession) -> None:
tasks = list(session.chat_subscription_tasks.values())
tasks = [
*session.chat_subscription_tasks.values(),
*session.chat_request_tasks.values(),
]
for task in tasks:
task.cancel()
if tasks:
@@ -366,6 +428,8 @@ class LiveChatService:
webchat_queue_mgr.remove_back_queue(request_id)
session.chat_subscriptions.clear()
session.chat_subscription_tasks.clear()
session.chat_request_tasks.clear()
session.interrupted_chat_requests.clear()
async def handle_chat_message(
self,
@@ -374,6 +438,12 @@ class LiveChatService:
send_json: SendJson,
) -> None:
msg_type = message.get("t")
message_id = str(message.get("message_id") or uuid.uuid4())
request_metadata = (
{"message_id": message_id}
if msg_type == "send" or message.get("message_id")
else {}
)
if msg_type == "bind":
chat_session_id = message.get("session_id")
@@ -406,7 +476,12 @@ class LiveChatService:
return
if msg_type == "interrupt":
session.should_interrupt = True
if message.get("message_id"):
session.interrupted_chat_requests.add(message_id)
else:
session.interrupted_chat_requests.update(
session.chat_request_tasks.keys()
)
await self.send_chat_payload(
session,
{
@@ -414,6 +489,7 @@ class LiveChatService:
"t": "error",
"data": "INTERRUPTED",
"code": "INTERRUPTED",
**request_metadata,
},
send_json,
)
@@ -432,22 +508,8 @@ class LiveChatService:
)
return
if session.is_processing:
await self.send_chat_payload(
session,
{
"ct": "chat",
"t": "error",
"data": "Session is busy",
"code": "PROCESSING_ERROR",
},
send_json,
)
return
payload = message.get("message")
session_id = message.get("session_id") or session.session_id
message_id = message.get("message_id") or str(uuid.uuid4())
selected_provider = message.get("selected_provider")
selected_model = message.get("selected_model")
selected_stt_provider = message.get("selected_stt_provider")
@@ -464,6 +526,7 @@ class LiveChatService:
"t": "error",
"data": "message must be list",
"code": "INVALID_MESSAGE_FORMAT",
**request_metadata,
},
send_json,
)
@@ -479,6 +542,7 @@ class LiveChatService:
"t": "error",
"data": "Message content is empty",
"code": "INVALID_MESSAGE_FORMAT",
**request_metadata,
},
send_json,
)
@@ -486,8 +550,6 @@ class LiveChatService:
await self.ensure_chat_subscription(session, session_id, send_json)
session.is_processing = True
session.should_interrupt = False
back_queue = webchat_queue_mgr.get_or_create_back_queue(message_id, session_id)
llm_checkpoint_id = str(uuid.uuid4())
@@ -532,6 +594,7 @@ class LiveChatService:
"created_at": to_utc_isoformat(saved_user_record.created_at),
"llm_checkpoint_id": llm_checkpoint_id,
},
**request_metadata,
},
send_json,
)
@@ -590,13 +653,14 @@ class LiveChatService:
"id": part["attachment_id"],
"type": part["type"],
},
**request_metadata,
},
send_json,
)
while True:
if session.should_interrupt:
session.should_interrupt = False
if message_id in session.interrupted_chat_requests:
session.interrupted_chat_requests.discard(message_id)
await flush_pending_bot_message()
break
@@ -624,6 +688,7 @@ class LiveChatService:
"ct": "chat",
"type": "agent_stats",
"data": parsed_agent_stats,
**request_metadata,
},
send_json,
)
@@ -703,6 +768,7 @@ class LiveChatService:
),
"llm_checkpoint_id": llm_checkpoint_id,
},
**request_metadata,
},
send_json,
)
@@ -719,6 +785,7 @@ class LiveChatService:
"t": "error",
"data": f"处理失败: {str(exc)}",
"code": "PROCESSING_ERROR",
**request_metadata,
},
send_json,
)
@@ -731,7 +798,7 @@ class LiveChatService:
f"[Live Chat] Failed to persist pending chat message: {exc}",
exc_info=True,
)
session.is_processing = False
session.interrupted_chat_requests.discard(message_id)
webchat_queue_mgr.remove_back_queue(message_id)
async def build_chat_message_parts(self, message: list[dict]) -> list[dict]:

16
dashboard/src/components/chat/StandaloneChat.vue
View File

@@ -68,7 +68,10 @@
type="button"
@click="openImage(partUrl(part))"
>
<img :src="partUrl(part)" :alt="part.filename || 'image'" />
<img
:src="partUrl(part)"
:alt="part.filename || 'image'"
/>
</button>
<audio
@@ -144,7 +147,9 @@
</template>
</div>
<pre v-else class="unknown-part">{{ formatJson(part) }}</pre>
<pre v-else class="unknown-part">{{
formatJson(part)
}}</pre>
</template>
</template>
</template>
@@ -332,7 +337,11 @@ async function sendCurrentMessage() {
const parts = buildOutgoingParts(text);
const messageId = crypto.randomUUID?.() || `${Date.now()}-${Math.random()}`;
const selection = inputRef.value?.getCurrentSelection();
const { botRecord } = createLocalExchange({ sessionId, messageId, parts });
const { userRecord, botRecord } = createLocalExchange({
sessionId,
messageId,
parts,
});
draft.value = "";
clearStaged({ revokeUrls: false });
@@ -347,6 +356,7 @@ async function sendCurrentMessage() {
enableStreaming: enableStreaming.value,
selectedProvider: selection?.providerId || "",
selectedModel: selection?.modelName || "",
userRecord,
botRecord,
});
}

259
dashboard/src/composables/useMessages.ts
View File

@@ -81,6 +81,7 @@ interface ActiveChatRun {
interface ActiveConnection {
sessionId: string;
messageId: string;
runId?: string;
transport: TransportMode;
abort?: AbortController;
ws?: WebSocket;
@@ -88,6 +89,10 @@ interface ActiveConnection {
userRecord?: ChatRecord;
completed?: boolean;
errorShown?: boolean;
botVisible?: boolean;
deferredBeforeBot?: ChatRecord;
followUpCaptured?: boolean;
followUpTargetRunId?: string;
}
interface SendMessageStreamOptions {
@@ -132,6 +137,7 @@ export function useMessages(options: UseMessagesOptions) {
const activeConnections = reactive<Record<string, ActiveConnection>>({});
const chatWebSockets: Record<string, WebSocket> = {};
const closingChatWebSockets = new WeakSet<WebSocket>();
const deferredBotAnchors = new WeakMap<ChatRecord, ChatRecord>();
const attachmentBlobCache = new Map<string, Promise<string>>();
const sessionProjects = reactive<Record<string, ChatSessionProject | null>>(
{},
@@ -152,7 +158,9 @@ export function useMessages(options: UseMessagesOptions) {
});
function isSessionRunning(sessionId: string) {
return Boolean(activeConnections[sessionId]);
return Object.values(activeConnections).some(
(connection) => connection.sessionId === sessionId,
);
}
function isUserMessage(msg: ChatRecord) {
@@ -170,14 +178,17 @@ export function useMessages(options: UseMessagesOptions) {
return [];
}
function isMessageStreaming(msg: ChatRecord, msgIndex: number) {
if (
!options.currentSessionId.value ||
!isSessionRunning(options.currentSessionId.value)
) {
return false;
}
return !isUserMessage(msg) && msgIndex === activeMessages.value.length - 1;
function isMessageStreaming(msg: ChatRecord, _msgIndex: number) {
const sessionId = options.currentSessionId.value;
if (!sessionId || isUserMessage(msg)) return false;
return Object.values(activeConnections).some(
(connection) =>
connection.sessionId === sessionId &&
connection.botVisible !== false &&
(connection.botRecord === msg ||
(connection.botRecord?.id != null &&
String(connection.botRecord.id) === String(msg.id))),
);
}
async function resolvePartMedia(part: MessagePart): Promise<void> {
@@ -268,7 +279,7 @@ export function useMessages(options: UseMessagesOptions) {
activeRuns: ActiveChatRun[],
) {
const run = activeRuns[0];
if (!run?.run_id || activeConnections[sessionId]) return;
if (!run?.run_id || isSessionRunning(sessionId)) return;
const checkpointId = run.llm_checkpoint_id || null;
const records = (messagesBySession[sessionId] || []).filter((record) => {
@@ -301,16 +312,16 @@ export function useMessages(options: UseMessagesOptions) {
loadedSessions[sessionId] = true;
messagesBySession[sessionId] = messagesBySession[sessionId] || [];
const userRecord: ChatRecord = {
const userRecord = reactive<ChatRecord>({
id: `local-user-${messageId}`,
created_at: new Date().toISOString(),
content: {
type: "user",
message: parts.map(stripUploadOnlyFields),
},
};
});
const botRecord: ChatRecord = {
const botRecord = reactive<ChatRecord>({
id: `local-bot-${messageId}`,
created_at: new Date().toISOString(),
content: {
@@ -319,14 +330,33 @@ export function useMessages(options: UseMessagesOptions) {
reasoning: "",
isLoading: true,
},
};
messagesBySession[sessionId].push(userRecord, botRecord);
});
const sessionMessages = messagesBySession[sessionId];
if (isSessionRunning(sessionId)) {
const activeBotConnections = Object.values(activeConnections).filter(
(connection) =>
connection.sessionId === sessionId &&
connection.botVisible !== false &&
connection.botRecord &&
sessionMessages.includes(connection.botRecord),
);
const activeBotRecord =
activeBotConnections[activeBotConnections.length - 1]?.botRecord;
if (activeBotRecord) {
const activeBotIndex = sessionMessages.indexOf(activeBotRecord);
sessionMessages.splice(activeBotIndex, 0, userRecord);
deferredBotAnchors.set(botRecord, activeBotRecord);
} else {
sessionMessages.push(userRecord);
}
} else {
sessionMessages.push(userRecord, botRecord);
}
return {
userRecord: sessionMessages[sessionMessages.length - 2],
botRecord: sessionMessages[sessionMessages.length - 1],
userRecord,
botRecord,
};
}
@@ -464,12 +494,15 @@ export function useMessages(options: UseMessagesOptions) {
};
const abort = new AbortController();
activeConnections[sessionId] = {
const connection: ActiveConnection = {
sessionId,
messageId: String(botRecord.id),
transport: "sse",
abort,
botRecord,
botVisible: true,
};
activeConnections[connection.messageId] = connection;
try {
const response = await fetchWithAuth(
@@ -495,7 +528,7 @@ export function useMessages(options: UseMessagesOptions) {
throw new Error(payload?.message || "Regenerate failed.");
}
await readSseStream(response.body, (payload) => {
processStreamPayload(botRecord, payload);
processStreamPayload(botRecord, payload, undefined, connection);
options.onStreamUpdate?.(sessionId);
});
} catch (error) {
@@ -507,7 +540,9 @@ export function useMessages(options: UseMessagesOptions) {
console.error("Regenerate failed:", error);
}
} finally {
delete activeConnections[sessionId];
if (activeConnections[connection.messageId]?.abort === abort) {
delete activeConnections[connection.messageId];
}
await options.onSessionsChanged?.();
}
}
@@ -522,8 +557,8 @@ export function useMessages(options: UseMessagesOptions) {
connection.abort?.abort();
});
Object.values(chatWebSockets).forEach(closeTrackedWebSocket);
Object.keys(activeConnections).forEach((sessionId) => {
delete activeConnections[sessionId];
Object.keys(activeConnections).forEach((messageId) => {
delete activeConnections[messageId];
});
Object.keys(chatWebSockets).forEach((sessionId) => {
delete chatWebSockets[sessionId];
@@ -580,12 +615,17 @@ export function useMessages(options: UseMessagesOptions) {
llmCheckpointId: string | null = null,
) {
const abort = new AbortController();
activeConnections[sessionId] = {
const connection: ActiveConnection = {
sessionId,
messageId,
transport: "sse",
abort,
botRecord,
userRecord,
botVisible: messagesBySession[sessionId]?.includes(botRecord) ?? false,
deferredBeforeBot: deferredBotAnchors.get(botRecord),
};
activeConnections[messageId] = connection;
fetchWithAuth(chatApi.sendStreamUrl(), {
method: "POST",
@@ -608,18 +648,21 @@ export function useMessages(options: UseMessagesOptions) {
throw new Error(`SSE connection failed: ${response.status}`);
}
await readSseStream(response.body, (payload) => {
processStreamPayload(botRecord, payload, userRecord);
processStreamPayload(botRecord, payload, userRecord, connection);
options.onStreamUpdate?.(sessionId);
});
})
.catch((error) => {
if (abort.signal.aborted) return;
ensureBotRecordVisible(connection);
appendPlain(botRecord, `\n\n${String(error?.message || error)}`);
console.error("SSE chat failed:", error);
})
.finally(async () => {
delete activeConnections[sessionId];
await options.onSessionsChanged?.();
if (activeConnections[messageId]?.abort === abort) {
delete activeConnections[messageId];
await options.onSessionsChanged?.();
}
});
}
@@ -629,13 +672,16 @@ export function useMessages(options: UseMessagesOptions) {
botRecord: ChatRecord,
) {
const abort = new AbortController();
activeConnections[sessionId] = {
const connection: ActiveConnection = {
sessionId,
messageId: runId,
runId,
transport: "sse",
abort,
botRecord,
botVisible: true,
};
activeConnections[runId] = connection;
void (async () => {
let receivedEnd = false;
@@ -666,7 +712,7 @@ export function useMessages(options: UseMessagesOptions) {
}
await readSseStream(response.body, (payload) => {
processStreamPayload(botRecord, payload);
processStreamPayload(botRecord, payload, undefined, connection);
options.onStreamUpdate?.(sessionId);
const payloadType = payload?.type || payload?.t;
if (payloadType === "end") receivedEnd = true;
@@ -696,10 +742,12 @@ export function useMessages(options: UseMessagesOptions) {
console.error("Resume chat stream failed:", lastError);
}
const ownsConnection = activeConnections[sessionId]?.abort === abort;
if (ownsConnection) delete activeConnections[sessionId];
const ownsConnection = activeConnections[runId]?.abort === abort;
if (ownsConnection) delete activeConnections[runId];
if (!abort.signal.aborted && ownsConnection) {
await loadSessionMessages(sessionId, true, false);
if (!isSessionRunning(sessionId)) {
await loadSessionMessages(sessionId, true, false);
}
await options.onSessionsChanged?.();
}
})();
@@ -717,7 +765,7 @@ export function useMessages(options: UseMessagesOptions) {
) {
const ws = getOrCreateChatWebSocket(sessionId);
activeConnections[sessionId] = {
const connection: ActiveConnection = {
sessionId,
messageId,
transport: "websocket",
@@ -726,7 +774,10 @@ export function useMessages(options: UseMessagesOptions) {
userRecord,
completed: false,
errorShown: false,
botVisible: messagesBySession[sessionId]?.includes(botRecord) ?? false,
deferredBeforeBot: deferredBotAnchors.get(botRecord),
};
activeConnections[messageId] = connection;
sendWebSocketPayload(sessionId, messageId, {
ct: "chat",
@@ -740,6 +791,15 @@ export function useMessages(options: UseMessagesOptions) {
});
}
function getWebSocketConnections(sessionId: string, ws?: WebSocket) {
return Object.values(activeConnections).filter(
(connection) =>
connection.sessionId === sessionId &&
connection.transport === "websocket" &&
(!ws || connection.ws === ws),
);
}
function getOrCreateChatWebSocket(sessionId: string) {
const existing = chatWebSockets[sessionId];
if (
@@ -758,9 +818,10 @@ export function useMessages(options: UseMessagesOptions) {
handleWebSocketMessage(sessionId, event);
};
ws.onerror = () => {
const connection = activeConnections[sessionId];
if (connection?.transport === "websocket" && connection.botRecord) {
for (const connection of getWebSocketConnections(sessionId, ws)) {
if (!connection.botRecord) continue;
connection.errorShown = true;
ensureBotRecordVisible(connection);
appendPlain(connection.botRecord, "\n\nWebSocket connection failed.");
}
};
@@ -769,20 +830,20 @@ export function useMessages(options: UseMessagesOptions) {
delete chatWebSockets[sessionId];
}
const connection = activeConnections[sessionId];
if (connection?.transport !== "websocket" || connection.ws !== ws) {
return;
const connections = getWebSocketConnections(sessionId, ws);
for (const connection of connections) {
if (
!connection.completed &&
!connection.errorShown &&
!closingChatWebSockets.has(ws) &&
connection.botRecord
) {
ensureBotRecordVisible(connection);
appendPlain(connection.botRecord, "\n\nWebSocket connection closed.");
}
delete activeConnections[connection.messageId];
}
if (
!connection.completed &&
!connection.errorShown &&
!closingChatWebSockets.has(ws) &&
connection.botRecord
) {
appendPlain(connection.botRecord, "\n\nWebSocket connection closed.");
}
delete activeConnections[sessionId];
await options.onSessionsChanged?.();
if (connections.length) await options.onSessionsChanged?.();
};
return ws;
}
@@ -794,7 +855,7 @@ export function useMessages(options: UseMessagesOptions) {
) {
const ws = getOrCreateChatWebSocket(sessionId);
const send = () => {
const connection = activeConnections[sessionId];
const connection = activeConnections[messageId];
if (
connection?.transport !== "websocket" ||
connection.messageId !== messageId ||
@@ -807,6 +868,7 @@ export function useMessages(options: UseMessagesOptions) {
} catch (error) {
connection.errorShown = true;
if (connection.botRecord) {
ensureBotRecordVisible(connection);
appendPlain(connection.botRecord, "\n\nWebSocket connection failed.");
}
console.error("Failed to send WebSocket payload:", error);
@@ -826,17 +888,25 @@ export function useMessages(options: UseMessagesOptions) {
}
function handleWebSocketMessage(sessionId: string, event: MessageEvent) {
const connection = activeConnections[sessionId];
if (connection?.transport !== "websocket" || !connection.botRecord) {
return;
}
try {
const payload = JSON.parse(event.data);
const payloadMessageId =
payload?.message_id == null ? "" : String(payload.message_id);
let connection = payloadMessageId
? activeConnections[payloadMessageId]
: undefined;
if (!connection) {
const candidates = getWebSocketConnections(sessionId);
connection = candidates[0];
}
if (connection?.transport !== "websocket" || !connection.botRecord) {
return;
}
processStreamPayload(
connection.botRecord,
payload,
connection.userRecord,
connection,
);
options.onStreamUpdate?.(sessionId);
if (payload.type === "end" || payload.t === "end") {
@@ -848,7 +918,7 @@ export function useMessages(options: UseMessagesOptions) {
}
async function finishWebSocketStream(sessionId: string, messageId: string) {
const connection = activeConnections[sessionId];
const connection = activeConnections[messageId];
if (
connection?.transport !== "websocket" ||
connection.messageId !== messageId
@@ -856,7 +926,7 @@ export function useMessages(options: UseMessagesOptions) {
return;
}
connection.completed = true;
delete activeConnections[sessionId];
delete activeConnections[messageId];
await options.onSessionsChanged?.();
}
@@ -870,10 +940,48 @@ export function useMessages(options: UseMessagesOptions) {
}
}
function ensureBotRecordVisible(connection: ActiveConnection) {
const { botRecord, userRecord } = connection;
if (!botRecord) return;
const records = messagesBySession[connection.sessionId] || [];
if (records.includes(botRecord)) {
connection.botVisible = true;
return;
}
if (!userRecord) return;
const userIndex = records.indexOf(userRecord);
let insertionAnchor = connection.deferredBeforeBot;
if (insertionAnchor) {
for (const candidate of Object.values(activeConnections)) {
if (candidate.messageId === connection.messageId) break;
if (
candidate.sessionId === connection.sessionId &&
candidate.deferredBeforeBot === connection.deferredBeforeBot &&
candidate.botVisible &&
candidate.botRecord &&
records.includes(candidate.botRecord)
) {
insertionAnchor = candidate.botRecord;
}
}
}
const anchorIndex = insertionAnchor ? records.indexOf(insertionAnchor) : -1;
if (anchorIndex >= 0) {
records.splice(anchorIndex + 1, 0, botRecord);
} else if (userIndex >= 0) {
records.splice(userIndex + 1, 0, botRecord);
} else {
records.push(botRecord);
}
connection.botVisible = true;
}
function processStreamPayload(
botRecord: ChatRecord,
payload: any,
userRecord?: ChatRecord,
connection?: ActiveConnection,
) {
const normalized =
payload?.ct === "chat"
@@ -883,7 +991,46 @@ export function useMessages(options: UseMessagesOptions) {
const chainType = normalized?.chain_type;
const data = normalized?.data ?? "";
if (msgType === "follow_up_captured") {
if (connection) {
connection.followUpCaptured = true;
connection.followUpTargetRunId = String(data?.target_run_id || "");
const target = Object.values(activeConnections).find(
(candidate) =>
candidate.runId === connection.followUpTargetRunId &&
candidate.botRecord,
);
const records = messagesBySession[connection.sessionId] || [];
if (target?.botRecord && connection.userRecord) {
const userIndex = records.indexOf(connection.userRecord);
const targetIndex = records.indexOf(target.botRecord);
if (targetIndex >= 0) {
if (userIndex > targetIndex) {
records.splice(userIndex, 1);
const updatedTargetIndex = records.indexOf(target.botRecord);
records.splice(updatedTargetIndex, 0, connection.userRecord);
} else if (userIndex < 0) {
records.splice(targetIndex, 0, connection.userRecord);
}
connection.deferredBeforeBot = target.botRecord;
} else if (userIndex < 0) {
records.push(connection.userRecord);
}
}
}
return;
}
if (msgType === "run_started") {
if (connection) {
connection.runId = String(data?.run_id || connection.messageId);
ensureBotRecordVisible(connection);
}
return;
}
if (msgType === "session_id" || msgType === "session_bound") return;
if (connection && msgType !== "user_message_saved" && msgType !== "end") {
ensureBotRecordVisible(connection);
}
if (msgType === "run_snapshot") {
const snapshot = data && typeof data === "object" ? data : {};
const snapshotRecord = normalizeHistoryRecord({

51
tests/test_chat_route.py
View File

@@ -385,3 +385,54 @@ async def test_legacy_chat_stream_keeps_existing_event_shape(chat_service_instan
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_chat_stream_forwards_follow_up_status_by_default(
chat_service_instance,
):
service = chat_service_instance
session_id = "follow-up-status-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"
assert _decode_sse_event(await anext(stream))["type"] == "user_message_saved"
status_payload = {
"type": "follow_up_captured",
"data": {"target_run_id": "original-run"},
"streaming": False,
"message_id": run.run_id,
}
await chat_service.webchat_queue_mgr.put_back_queue(
run.run_id,
status_payload,
)
assert _decode_sse_event(await asyncio.wait_for(anext(stream), 1)) == (
status_payload
)
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:
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)

155
tests/unit/test_live_chat_service.py
View File

@@ -1,8 +1,12 @@
import asyncio
from datetime import UTC, datetime
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
from starlette.websockets import WebSocketDisconnect
from astrbot.core.platform.sources.webchat.webchat_queue_mgr import webchat_queue_mgr
from astrbot.dashboard.services.live_chat_service import LiveChatService
@@ -140,3 +144,154 @@ async def test_run_websocket_session_handles_disconnect_without_error_log(
}
]
assert service.sessions == {}
@pytest.mark.asyncio
async def test_run_websocket_session_multiplexes_chat_requests_by_default(
monkeypatch,
):
service = _service()
started = asyncio.Event()
started_requests: list[str] = []
messages = iter(
[
{
"ct": "chat",
"t": "send",
"session_id": "chat-session",
"message_id": "request-1",
},
{
"ct": "chat",
"t": "send",
"session_id": "chat-session",
"message_id": "request-2",
},
]
)
monkeypatch.setattr(service, "authenticate_token", lambda _token: "alice")
async def handle_chat_message(_session, message, _send_json) -> None:
started_requests.append(message["message_id"])
if len(started_requests) == 2:
started.set()
await asyncio.Event().wait()
monkeypatch.setattr(service, "handle_chat_message", handle_chat_message)
async def receive_json() -> dict:
try:
return next(messages)
except StopIteration as exc:
await asyncio.wait_for(started.wait(), timeout=1)
raise WebSocketDisconnect(1000) from exc
async def send_json(_payload: dict) -> None:
pass
async def close(_code: int, _reason: str) -> None:
raise AssertionError("close should not be called")
await service.run_websocket_session(
token="valid",
force_ct=None,
receive_json=receive_json,
send_json=send_json,
close=close,
)
assert started_requests == ["request-1", "request-2"]
assert service.sessions == {}
@pytest.mark.asyncio
async def test_handle_chat_interrupt_without_message_id_targets_all_requests():
service = _service()
session = service.create_session("alice")
sent: list[dict] = []
tasks = {
"request-1": asyncio.create_task(asyncio.Event().wait()),
"request-2": asyncio.create_task(asyncio.Event().wait()),
}
session.chat_request_tasks.update(tasks)
async def send_json(payload: dict) -> None:
sent.append(payload)
try:
await service.handle_chat_message(
session,
{"t": "interrupt"},
send_json,
)
assert session.interrupted_chat_requests == set(tasks)
assert sent == [
{
"ct": "chat",
"t": "error",
"data": "INTERRUPTED",
"code": "INTERRUPTED",
}
]
finally:
await service.cleanup_session(session)
@pytest.mark.asyncio
async def test_handle_chat_message_scopes_events_to_request_by_default():
service = _service()
session = service.create_session("alice")
session_id = "multiplexed-chat-session"
message_id = "request-1"
sent: list[dict] = []
service.platform_history_mgr.insert = AsyncMock(
return_value=SimpleNamespace(id=1, created_at=datetime.now(UTC))
)
service.build_chat_message_parts = AsyncMock(
return_value=[{"type": "plain", "text": "hello"}]
)
service.ensure_chat_subscription = AsyncMock(return_value="subscription-1")
async def send_json(payload: dict) -> None:
sent.append(payload)
task = asyncio.create_task(
service.handle_chat_message(
session,
{
"t": "send",
"session_id": session_id,
"message_id": message_id,
"message": [{"type": "plain", "text": "hello"}],
},
send_json,
)
)
try:
input_queue = webchat_queue_mgr.get_or_create_queue(session_id)
await asyncio.wait_for(input_queue.get(), timeout=1)
await webchat_queue_mgr.put_back_queue(
message_id,
{
"type": "end",
"data": "",
"streaming": False,
"message_id": message_id,
},
)
await asyncio.wait_for(task, timeout=1)
assert sent[0]["type"] == "user_message_saved"
assert sent[0]["message_id"] == message_id
assert sent[-1]["type"] == "end"
assert sent[-1]["message_id"] == message_id
finally:
if not task.done():
task.cancel()
await asyncio.gather(task, return_exceptions=True)
await service.cleanup_session(session)
webchat_queue_mgr.remove_queues(session_id)

80
tests/unit/test_webchat_follow_up_status.py Normal file
View File

@@ -0,0 +1,80 @@
import pytest
from astrbot.core.platform.astrbot_message import AstrBotMessage, MessageMember
from astrbot.core.platform.message_type import MessageType
from astrbot.core.platform.platform_metadata import PlatformMetadata
from astrbot.core.platform.sources.webchat.webchat_event import WebChatMessageEvent
from astrbot.core.platform.sources.webchat.webchat_queue_mgr import webchat_queue_mgr
def _event(message_id: str) -> WebChatMessageEvent:
"""Create an isolated WebChat event for stream-status tests.
Args:
message_id: Request identifier used by the response queue.
Returns:
Configured WebChat event.
"""
message = AstrBotMessage()
message.type = MessageType.FRIEND_MESSAGE
message.self_id = "webchat"
message.session_id = "session-1"
message.message_id = message_id
message.sender = MessageMember("alice", "Alice")
message.message = []
message.message_str = "hello"
return WebChatMessageEvent(
"hello",
message,
PlatformMetadata(name="webchat", description="webchat", id="webchat"),
"webchat!alice!session-1",
)
@pytest.mark.asyncio
async def test_webchat_run_started_is_emitted_by_default():
event = _event("default-request")
queue = webchat_queue_mgr.get_or_create_back_queue("default-request")
try:
await event.send_typing()
await event.send(None)
assert await queue.get() == {
"type": "run_started",
"data": {"run_id": "default-request"},
"streaming": False,
"message_id": "default-request",
}
assert await queue.get() == {
"type": "end",
"data": "",
"streaming": False,
"message_id": "default-request",
}
assert queue.empty()
finally:
webchat_queue_mgr.remove_back_queue("default-request")
@pytest.mark.asyncio
async def test_webchat_follow_up_captured_is_emitted_by_default():
event = _event("follow-up-request")
queue = webchat_queue_mgr.get_or_create_back_queue("follow-up-request")
try:
event.set_extra(
"_follow_up_captured",
{"target_run_id": "original-run"},
)
await event.send(None)
assert await queue.get() == {
"type": "follow_up_captured",
"data": {"target_run_id": "original-run"},
"streaming": False,
"message_id": "follow-up-request",
}
assert (await queue.get())["type"] == "end"
finally:
webchat_queue_mgr.remove_back_queue("follow-up-request")