fix: keep IM session streams subscribed across desktop views

Telegram and desktop can attach to the same session while a turn is streaming. The WebSocket handler now tracks output callbacks per client and broadcasts session messages instead of replacing the previous subscriber. Telegram thinking deltas are accumulated before editing the placeholder so the preview does not collapse to the latest tiny chunk.

Constraint: Desktop and IM adapters may observe the same active session concurrently
Rejected: Keep a single session output callback | a later desktop view can steal Telegram's live content and completion events
Confidence: high
Scope-risk: moderate
Directive: Do not collapse session output callbacks back to one callback without a multi-client streaming regression
Tested: bun run check:server (828 pass); bun run check:adapters (358 pass); git diff --check
Not-tested: Live Telegram Bot API smoke against the packaged app
This commit is contained in:
程序员阿江(Relakkes) 2026-05-23 16:19:39 +08:00
parent 28ddcd0afd
commit 673523f3fd
5 changed files with 221 additions and 20 deletions

View File

@ -2,6 +2,7 @@ import { describe, it, expect } from 'bun:test'
import { splitMessage, formatPermissionRequest, truncateInput, escapeMarkdownV2 } from '../../common/format.js' import { splitMessage, formatPermissionRequest, truncateInput, escapeMarkdownV2 } from '../../common/format.js'
import { parsePermitCallbackData } from '../../common/permission.js' import { parsePermitCallbackData } from '../../common/permission.js'
import { import {
buildTelegramThinkingUpdate,
formatTelegramOutboundText, formatTelegramOutboundText,
formatTelegramStreamingText, formatTelegramStreamingText,
planTelegramStreamingUpdate, planTelegramStreamingUpdate,
@ -96,6 +97,23 @@ describe('Telegram message formatting', () => {
}) })
}) })
describe('Telegram thinking updates', () => {
it('accumulates thinking deltas before formatting the preview', () => {
const first = buildTelegramThinkingUpdate('', 'The user')
const second = buildTelegramThinkingUpdate(first.fullText, ' wants a long answer')
expect(second.fullText).toBe('The user wants a long answer')
expect(second.messageText).toBe('💭 The user wants a long answer...')
})
it('caps long thinking previews while keeping the full accumulated text', () => {
const result = buildTelegramThinkingUpdate('abcdef', 'ghij', 6)
expect(result.fullText).toBe('abcdefghij')
expect(result.messageText).toBe('💭 abcdef...')
})
})
describe('permission request formatting', () => { describe('permission request formatting', () => {
it('formats Bash command request', () => { it('formats Bash command request', () => {
const result = formatPermissionRequest('Bash', { command: 'npm test' }, 'abcde') const result = formatPermissionRequest('Bash', { command: 'npm test' }, 'abcde')

View File

@ -11,6 +11,26 @@ export function formatTelegramStreamingText(text: string): string {
return `${formatTelegramOutboundText(text)}` return `${formatTelegramOutboundText(text)}`
} }
const DEFAULT_THINKING_PREVIEW_LIMIT = 1000
export type TelegramThinkingUpdate = {
fullText: string
messageText: string
}
export function buildTelegramThinkingUpdate(
currentText: string,
deltaText: string,
previewLimit = DEFAULT_THINKING_PREVIEW_LIMIT,
): TelegramThinkingUpdate {
const fullText = currentText + deltaText
const preview = fullText.slice(0, Math.max(0, previewLimit)).trimStart()
return {
fullText,
messageText: preview ? `💭 ${preview}...` : '💭 思考中...',
}
}
export type TelegramStreamingUpdate = { export type TelegramStreamingUpdate = {
sealedChunks: string[] sealedChunks: string[]
activeChunk: string activeChunk: string

View File

@ -19,6 +19,7 @@ import {
splitMessage, splitMessage,
} from '../common/format.js' } from '../common/format.js'
import { import {
buildTelegramThinkingUpdate,
formatTelegramOutboundText, formatTelegramOutboundText,
formatTelegramStreamingText, formatTelegramStreamingText,
planTelegramStreamingUpdate, planTelegramStreamingUpdate,
@ -68,6 +69,7 @@ attachmentStore.gc().catch((err) => {
const placeholders = new Map<string, { chatId: string; messageId: number }>() const placeholders = new Map<string, { chatId: string; messageId: number }>()
// Track accumulated text per chat for streaming // Track accumulated text per chat for streaming
const accumulatedText = new Map<string, string>() const accumulatedText = new Map<string, string>()
const accumulatedThinkingText = new Map<string, string>()
// Message buffers per chat // Message buffers per chat
const buffers = new Map<string, MessageBuffer>() const buffers = new Map<string, MessageBuffer>()
// Track chats waiting for project selection // Track chats waiting for project selection
@ -118,6 +120,7 @@ function getRuntimeState(chatId: string): ChatRuntimeState {
function clearTransientChatState(chatId: string): void { function clearTransientChatState(chatId: string): void {
placeholders.delete(chatId) placeholders.delete(chatId)
accumulatedText.delete(chatId) accumulatedText.delete(chatId)
accumulatedThinkingText.delete(chatId)
buffers.get(chatId)?.reset() buffers.get(chatId)?.reset()
const runtime = getRuntimeState(chatId) const runtime = getRuntimeState(chatId)
runtime.state = 'idle' runtime.state = 'idle'
@ -408,11 +411,13 @@ async function handleServerMessage(chatId: string, msg: ServerMessage): Promise<
const sent = await bot.api.sendMessage(numericChatId, '💭 思考中...') const sent = await bot.api.sendMessage(numericChatId, '💭 思考中...')
placeholders.set(chatId, { chatId, messageId: sent.message_id }) placeholders.set(chatId, { chatId, messageId: sent.message_id })
accumulatedText.set(chatId, '') accumulatedText.set(chatId, '')
accumulatedThinkingText.set(chatId, '')
} }
break break
case 'content_start': case 'content_start':
if (msg.blockType === 'text') { if (msg.blockType === 'text') {
accumulatedThinkingText.delete(chatId)
if (!placeholders.has(chatId)) { if (!placeholders.has(chatId)) {
const sent = await bot.api.sendMessage(numericChatId, '▍') const sent = await bot.api.sendMessage(numericChatId, '▍')
placeholders.set(chatId, { chatId, messageId: sent.message_id }) placeholders.set(chatId, { chatId, messageId: sent.message_id })
@ -443,6 +448,7 @@ async function handleServerMessage(chatId: string, msg: ServerMessage): Promise<
case 'content_delta': case 'content_delta':
if (msg.text) { if (msg.text) {
accumulatedThinkingText.delete(chatId)
buf.append(msg.text) buf.append(msg.text)
const newUploads = getTgWatcher(chatId).feed(msg.text) const newUploads = getTgWatcher(chatId).feed(msg.text)
for (const pending of newUploads) { for (const pending of newUploads) {
@ -453,11 +459,16 @@ async function handleServerMessage(chatId: string, msg: ServerMessage): Promise<
case 'thinking': case 'thinking':
if (placeholders.has(chatId)) { if (placeholders.has(chatId)) {
const update = buildTelegramThinkingUpdate(
accumulatedThinkingText.get(chatId) ?? '',
msg.text,
)
accumulatedThinkingText.set(chatId, update.fullText)
try { try {
await bot.api.editMessageText( await bot.api.editMessageText(
numericChatId, numericChatId,
placeholders.get(chatId)!.messageId, placeholders.get(chatId)!.messageId,
`💭 ${msg.text.slice(0, 200)}...`, update.messageText,
) )
} catch { /* ignore */ } } catch { /* ignore */ }
} }
@ -506,6 +517,7 @@ async function handleServerMessage(chatId: string, msg: ServerMessage): Promise<
} }
placeholders.delete(chatId) placeholders.delete(chatId)
accumulatedText.delete(chatId) accumulatedText.delete(chatId)
accumulatedThinkingText.delete(chatId)
buffers.get(chatId)?.reset() buffers.get(chatId)?.reset()
} }
break break
@ -513,6 +525,7 @@ async function handleServerMessage(chatId: string, msg: ServerMessage): Promise<
case 'error': case 'error':
runtime.state = 'idle' runtime.state = 'idle'
runtime.verb = undefined runtime.verb = undefined
accumulatedThinkingText.delete(chatId)
// Auto-recover from stale thinking block signatures by creating a fresh session. // Auto-recover from stale thinking block signatures by creating a fresh session.
// This happens when the API key or provider changed since the session was created. // This happens when the API key or provider changed since the session was created.
if (msg.message && /Invalid.*signature.*thinking/i.test(msg.message)) { if (msg.message && /Invalid.*signature.*thinking/i.test(msg.message)) {

View File

@ -2528,6 +2528,84 @@ describe('WebSocket Chat Integration', () => {
}) })
}) })
it('should stream one active turn to multiple connected clients', async () => {
await withMockStreamDelay(150, async () => {
const sessionId = `chat-multi-client-${crypto.randomUUID()}`
const firstMessages: any[] = []
const secondMessages: any[] = []
await new Promise<void>((resolve, reject) => {
let secondConnected = false
let firstComplete = false
let secondComplete = false
let ws2: WebSocket | null = null
const timeout = setTimeout(() => {
ws1.close()
ws2?.close()
reject(new Error(`Timed out waiting for both clients to complete for session ${sessionId}`))
}, 10_000)
const cleanup = () => {
if (!firstComplete || !secondComplete) return
clearTimeout(timeout)
ws1.close()
ws2?.close()
resolve()
}
const handleFailure = (message: string) => {
clearTimeout(timeout)
ws1.close()
ws2?.close()
reject(new Error(message))
}
const ws1 = new WebSocket(`${wsUrl}/ws/${sessionId}`)
ws1.onmessage = (event) => {
const msg = JSON.parse(event.data as string)
firstMessages.push(msg)
if (msg.type === 'connected') {
ws1.send(JSON.stringify({ type: 'user_message', content: 'multi client stream' }))
return
}
if (msg.type === 'thinking' && !secondConnected) {
secondConnected = true
ws2 = new WebSocket(`${wsUrl}/ws/${sessionId}`)
ws2.onmessage = (secondEvent) => {
const secondMsg = JSON.parse(secondEvent.data as string)
secondMessages.push(secondMsg)
if (secondMsg.type === 'error') {
handleFailure(secondMsg.message)
return
}
if (secondMsg.type === 'message_complete') {
secondComplete = true
cleanup()
}
}
ws2.onerror = () => handleFailure(`Second WebSocket error for session ${sessionId}`)
}
if (msg.type === 'message_complete') {
firstComplete = true
cleanup()
}
}
ws1.onerror = () => handleFailure(`First WebSocket error for session ${sessionId}`)
})
expect(firstMessages.some((msg) => msg.type === 'content_delta')).toBe(true)
expect(firstMessages.some((msg) => msg.type === 'message_complete')).toBe(true)
expect(secondMessages.some((msg) => msg.type === 'connected')).toBe(true)
expect(secondMessages.some((msg) => msg.type === 'content_delta')).toBe(true)
expect(secondMessages.some((msg) => msg.type === 'message_complete')).toBe(true)
})
})
it('should keep using the selected runtime config across the whole session until changed', async () => { it('should keep using the selected runtime config across the whole session until changed', async () => {
const providerService = new ProviderService() const providerService = new ProviderService()
const providerA = await providerService.addProvider({ const providerA = await providerService.addProvider({

View File

@ -107,8 +107,16 @@ export type WebSocketData = {
serverHost: string serverHost: string
} }
// Active WebSocket sessions // Active WebSocket clients, grouped by session. Desktop, H5, and IM adapters can
const activeSessions = new Map<string, ServerWebSocket<WebSocketData>>() // legitimately watch the same running session at the same time.
const activeSessions = new Map<string, Set<ServerWebSocket<WebSocketData>>>()
const clientOutputCallbacks = new Map<
ServerWebSocket<WebSocketData>,
{
sessionId: string
callback: (cliMsg: any) => void
}
>()
export const handleWebSocket = { export const handleWebSocket = {
open(ws: ServerWebSocket<WebSocketData>) { open(ws: ServerWebSocket<WebSocketData>) {
@ -135,11 +143,11 @@ export const handleWebSocket = {
sessionCleanupTimers.delete(sessionId) sessionCleanupTimers.delete(sessionId)
} }
activeSessions.set(sessionId, ws) addActiveClient(sessionId, ws)
if (prewarmedSessions.has(sessionId)) { if (prewarmedSessions.has(sessionId)) {
bindPrewarmMetadataCapture(sessionId) bindPrewarmMetadataCapture(sessionId)
} else { } else {
rebindSessionOutput(sessionId, ws) bindClientSessionOutput(sessionId, ws)
} }
const msg: ServerMessage = { type: 'connected', sessionId } const msg: ServerMessage = { type: 'connected', sessionId }
@ -218,19 +226,23 @@ export const handleWebSocket = {
} }
console.log(`[WS] Client disconnected from session: ${sessionId} (${code}: ${reason})`) console.log(`[WS] Client disconnected from session: ${sessionId} (${code}: ${reason})`)
if (activeSessions.get(sessionId) !== ws) { if (!removeActiveClient(sessionId, ws)) {
console.log(`[WS] Ignoring stale client disconnect for session: ${sessionId}`) console.log(`[WS] Ignoring stale client disconnect for session: ${sessionId}`)
return return
} }
removeClientOutputCallback(ws)
if (hasActiveClients(sessionId)) {
return
}
computerUseApprovalService.cancelSession(sessionId) computerUseApprovalService.cancelSession(sessionId)
activeSessions.delete(sessionId)
conversationService.clearOutputCallbacks(sessionId)
// Schedule delayed cleanup: if the client doesn't reconnect within 30 seconds, // Schedule delayed cleanup: if the client doesn't reconnect within 30 seconds,
// stop the CLI subprocess to avoid leaking resources. // stop the CLI subprocess to avoid leaking resources.
const cleanupTimer = setTimeout(() => { const cleanupTimer = setTimeout(() => {
sessionCleanupTimers.delete(sessionId) sessionCleanupTimers.delete(sessionId)
if (!activeSessions.has(sessionId)) { if (!hasActiveClients(sessionId)) {
console.log(`[WS] Session ${sessionId} not reconnected after 30s, stopping CLI subprocess`) console.log(`[WS] Session ${sessionId} not reconnected after 30s, stopping CLI subprocess`)
conversationService.stopSession(sessionId) conversationService.stopSession(sessionId)
cleanupSessionRuntimeState(sessionId) cleanupSessionRuntimeState(sessionId)
@ -340,7 +352,7 @@ async function handleUserMessage(
const shouldForwardCurrentTurnLocalCommand = const shouldForwardCurrentTurnLocalCommand =
createCurrentTurnLocalCommandForwarder(desktopSlashCommand) createCurrentTurnLocalCommandForwarder(desktopSlashCommand)
rebindSessionOutput(sessionId, ws, { bindAllClientSessionOutputs(sessionId, {
shouldForward: (cliMsg) => { shouldForward: (cliMsg) => {
if (userMessageSent || (cliMsg.type === 'result' && cliMsg.is_error)) { if (userMessageSent || (cliMsg.type === 'result' && cliMsg.is_error)) {
return true return true
@ -1679,7 +1691,56 @@ function isCompactSummaryMessageContent(content: unknown): content is string {
) )
} }
function rebindSessionOutput( function addActiveClient(
sessionId: string,
ws: ServerWebSocket<WebSocketData>,
): void {
let clients = activeSessions.get(sessionId)
if (!clients) {
clients = new Set()
activeSessions.set(sessionId, clients)
}
clients.add(ws)
}
function removeActiveClient(
sessionId: string,
ws: ServerWebSocket<WebSocketData>,
): boolean {
const clients = activeSessions.get(sessionId)
if (!clients?.has(ws)) return false
clients.delete(ws)
if (clients.size === 0) {
activeSessions.delete(sessionId)
}
return true
}
function hasActiveClients(sessionId: string): boolean {
return (activeSessions.get(sessionId)?.size ?? 0) > 0
}
function removeClientOutputCallback(ws: ServerWebSocket<WebSocketData>): void {
const entry = clientOutputCallbacks.get(ws)
if (!entry) return
conversationService.removeOutputCallback(entry.sessionId, entry.callback)
clientOutputCallbacks.delete(ws)
}
function bindAllClientSessionOutputs(
sessionId: string,
options?: {
shouldForward?: (cliMsg: any) => boolean
},
): void {
const clients = activeSessions.get(sessionId)
if (!clients) return
for (const ws of clients) {
bindClientSessionOutput(sessionId, ws, options)
}
}
function bindClientSessionOutput(
sessionId: string, sessionId: string,
ws: ServerWebSocket<WebSocketData>, ws: ServerWebSocket<WebSocketData>,
options?: { options?: {
@ -1688,8 +1749,9 @@ function rebindSessionOutput(
) { ) {
if (!conversationService.hasSession(sessionId)) return if (!conversationService.hasSession(sessionId)) return
conversationService.clearOutputCallbacks(sessionId) removeClientOutputCallback(ws)
conversationService.onOutput(sessionId, (cliMsg) => {
const callback = (cliMsg: any) => {
if (options?.shouldForward && !options.shouldForward(cliMsg)) { if (options?.shouldForward && !options.shouldForward(cliMsg)) {
return return
} }
@ -1702,7 +1764,10 @@ function rebindSessionOutput(
if (cliMsg.type === 'result') { if (cliMsg.type === 'result') {
triggerTitleGeneration(ws, sessionId) triggerTitleGeneration(ws, sessionId)
} }
}) }
clientOutputCallbacks.set(ws, { sessionId, callback })
conversationService.onOutput(sessionId, callback)
} }
type RuntimeSettings = { type RuntimeSettings = {
@ -1926,9 +1991,12 @@ async function waitForRuntimeTransitionBeforeUserTurn(
* Send a message to a specific session's WebSocket (for use by services) * Send a message to a specific session's WebSocket (for use by services)
*/ */
export function sendToSession(sessionId: string, message: ServerMessage): boolean { export function sendToSession(sessionId: string, message: ServerMessage): boolean {
const ws = activeSessions.get(sessionId) const clients = activeSessions.get(sessionId)
if (!ws) return false if (!clients || clients.size === 0) return false
ws.send(JSON.stringify(message)) const payload = JSON.stringify(message)
for (const ws of clients) {
ws.send(payload)
}
return true return true
} }
@ -1991,11 +2059,14 @@ export function closeSessionConnection(sessionId: string, reason = 'session clos
conversationService.clearOutputCallbacks(sessionId) conversationService.clearOutputCallbacks(sessionId)
cleanupSessionRuntimeState(sessionId) cleanupSessionRuntimeState(sessionId)
const ws = activeSessions.get(sessionId) const clients = activeSessions.get(sessionId)
if (!ws) return false if (!clients || clients.size === 0) return false
activeSessions.delete(sessionId) activeSessions.delete(sessionId)
for (const ws of clients) {
clientOutputCallbacks.delete(ws)
ws.close(1000, reason) ws.close(1000, reason)
}
return true return true
} }
@ -2007,6 +2078,7 @@ export function __resetWebSocketHandlerStateForTests(): void {
for (const timer of sessionCleanupTimers.values()) clearTimeout(timer) for (const timer of sessionCleanupTimers.values()) clearTimeout(timer)
for (const timer of prewarmIdleTimers.values()) clearTimeout(timer) for (const timer of prewarmIdleTimers.values()) clearTimeout(timer)
activeSessions.clear() activeSessions.clear()
clientOutputCallbacks.clear()
sessionCleanupTimers.clear() sessionCleanupTimers.clear()
prewarmIdleTimers.clear() prewarmIdleTimers.clear()
} }