diff --git a/adapters/telegram/__tests__/telegram.test.ts b/adapters/telegram/__tests__/telegram.test.ts index cd4f22d7..7c01bf4b 100644 --- a/adapters/telegram/__tests__/telegram.test.ts +++ b/adapters/telegram/__tests__/telegram.test.ts @@ -2,6 +2,7 @@ import { describe, it, expect } from 'bun:test' import { splitMessage, formatPermissionRequest, truncateInput, escapeMarkdownV2 } from '../../common/format.js' import { parsePermitCallbackData } from '../../common/permission.js' import { + buildTelegramThinkingUpdate, formatTelegramOutboundText, formatTelegramStreamingText, 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', () => { it('formats Bash command request', () => { const result = formatPermissionRequest('Bash', { command: 'npm test' }, 'abcde') diff --git a/adapters/telegram/format.ts b/adapters/telegram/format.ts index 6b7fea27..4c9f5e50 100644 --- a/adapters/telegram/format.ts +++ b/adapters/telegram/format.ts @@ -11,6 +11,26 @@ export function formatTelegramStreamingText(text: string): string { 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 = { sealedChunks: string[] activeChunk: string diff --git a/adapters/telegram/index.ts b/adapters/telegram/index.ts index 9099c4cc..f38fbd46 100644 --- a/adapters/telegram/index.ts +++ b/adapters/telegram/index.ts @@ -19,6 +19,7 @@ import { splitMessage, } from '../common/format.js' import { + buildTelegramThinkingUpdate, formatTelegramOutboundText, formatTelegramStreamingText, planTelegramStreamingUpdate, @@ -68,6 +69,7 @@ attachmentStore.gc().catch((err) => { const placeholders = new Map() // Track accumulated text per chat for streaming const accumulatedText = new Map() +const accumulatedThinkingText = new Map() // Message buffers per chat const buffers = new Map() // Track chats waiting for project selection @@ -118,6 +120,7 @@ function getRuntimeState(chatId: string): ChatRuntimeState { function clearTransientChatState(chatId: string): void { placeholders.delete(chatId) accumulatedText.delete(chatId) + accumulatedThinkingText.delete(chatId) buffers.get(chatId)?.reset() const runtime = getRuntimeState(chatId) runtime.state = 'idle' @@ -408,11 +411,13 @@ async function handleServerMessage(chatId: string, msg: ServerMessage): Promise< const sent = await bot.api.sendMessage(numericChatId, 'πŸ’­ 思考中...') placeholders.set(chatId, { chatId, messageId: sent.message_id }) accumulatedText.set(chatId, '') + accumulatedThinkingText.set(chatId, '') } break case 'content_start': if (msg.blockType === 'text') { + accumulatedThinkingText.delete(chatId) if (!placeholders.has(chatId)) { const sent = await bot.api.sendMessage(numericChatId, '▍') placeholders.set(chatId, { chatId, messageId: sent.message_id }) @@ -443,6 +448,7 @@ async function handleServerMessage(chatId: string, msg: ServerMessage): Promise< case 'content_delta': if (msg.text) { + accumulatedThinkingText.delete(chatId) buf.append(msg.text) const newUploads = getTgWatcher(chatId).feed(msg.text) for (const pending of newUploads) { @@ -453,11 +459,16 @@ async function handleServerMessage(chatId: string, msg: ServerMessage): Promise< case 'thinking': if (placeholders.has(chatId)) { + const update = buildTelegramThinkingUpdate( + accumulatedThinkingText.get(chatId) ?? '', + msg.text, + ) + accumulatedThinkingText.set(chatId, update.fullText) try { await bot.api.editMessageText( numericChatId, placeholders.get(chatId)!.messageId, - `πŸ’­ ${msg.text.slice(0, 200)}...`, + update.messageText, ) } catch { /* ignore */ } } @@ -506,6 +517,7 @@ async function handleServerMessage(chatId: string, msg: ServerMessage): Promise< } placeholders.delete(chatId) accumulatedText.delete(chatId) + accumulatedThinkingText.delete(chatId) buffers.get(chatId)?.reset() } break @@ -513,6 +525,7 @@ async function handleServerMessage(chatId: string, msg: ServerMessage): Promise< case 'error': runtime.state = 'idle' runtime.verb = undefined + accumulatedThinkingText.delete(chatId) // 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. if (msg.message && /Invalid.*signature.*thinking/i.test(msg.message)) { diff --git a/src/server/__tests__/conversations.test.ts b/src/server/__tests__/conversations.test.ts index 16cec1e0..780a818f 100644 --- a/src/server/__tests__/conversations.test.ts +++ b/src/server/__tests__/conversations.test.ts @@ -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((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 () => { const providerService = new ProviderService() const providerA = await providerService.addProvider({ diff --git a/src/server/ws/handler.ts b/src/server/ws/handler.ts index 857a653d..feadc37f 100644 --- a/src/server/ws/handler.ts +++ b/src/server/ws/handler.ts @@ -107,8 +107,16 @@ export type WebSocketData = { serverHost: string } -// Active WebSocket sessions -const activeSessions = new Map>() +// Active WebSocket clients, grouped by session. Desktop, H5, and IM adapters can +// legitimately watch the same running session at the same time. +const activeSessions = new Map>>() +const clientOutputCallbacks = new Map< + ServerWebSocket, + { + sessionId: string + callback: (cliMsg: any) => void + } +>() export const handleWebSocket = { open(ws: ServerWebSocket) { @@ -135,11 +143,11 @@ export const handleWebSocket = { sessionCleanupTimers.delete(sessionId) } - activeSessions.set(sessionId, ws) + addActiveClient(sessionId, ws) if (prewarmedSessions.has(sessionId)) { bindPrewarmMetadataCapture(sessionId) } else { - rebindSessionOutput(sessionId, ws) + bindClientSessionOutput(sessionId, ws) } const msg: ServerMessage = { type: 'connected', sessionId } @@ -218,19 +226,23 @@ export const handleWebSocket = { } 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}`) return } + removeClientOutputCallback(ws) + + if (hasActiveClients(sessionId)) { + return + } + computerUseApprovalService.cancelSession(sessionId) - activeSessions.delete(sessionId) - conversationService.clearOutputCallbacks(sessionId) // Schedule delayed cleanup: if the client doesn't reconnect within 30 seconds, // stop the CLI subprocess to avoid leaking resources. const cleanupTimer = setTimeout(() => { sessionCleanupTimers.delete(sessionId) - if (!activeSessions.has(sessionId)) { + if (!hasActiveClients(sessionId)) { console.log(`[WS] Session ${sessionId} not reconnected after 30s, stopping CLI subprocess`) conversationService.stopSession(sessionId) cleanupSessionRuntimeState(sessionId) @@ -340,7 +352,7 @@ async function handleUserMessage( const shouldForwardCurrentTurnLocalCommand = createCurrentTurnLocalCommandForwarder(desktopSlashCommand) - rebindSessionOutput(sessionId, ws, { + bindAllClientSessionOutputs(sessionId, { shouldForward: (cliMsg) => { if (userMessageSent || (cliMsg.type === 'result' && cliMsg.is_error)) { return true @@ -1679,7 +1691,56 @@ function isCompactSummaryMessageContent(content: unknown): content is string { ) } -function rebindSessionOutput( +function addActiveClient( + sessionId: string, + ws: ServerWebSocket, +): void { + let clients = activeSessions.get(sessionId) + if (!clients) { + clients = new Set() + activeSessions.set(sessionId, clients) + } + clients.add(ws) +} + +function removeActiveClient( + sessionId: string, + ws: ServerWebSocket, +): 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): 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, ws: ServerWebSocket, options?: { @@ -1688,8 +1749,9 @@ function rebindSessionOutput( ) { if (!conversationService.hasSession(sessionId)) return - conversationService.clearOutputCallbacks(sessionId) - conversationService.onOutput(sessionId, (cliMsg) => { + removeClientOutputCallback(ws) + + const callback = (cliMsg: any) => { if (options?.shouldForward && !options.shouldForward(cliMsg)) { return } @@ -1702,7 +1764,10 @@ function rebindSessionOutput( if (cliMsg.type === 'result') { triggerTitleGeneration(ws, sessionId) } - }) + } + + clientOutputCallbacks.set(ws, { sessionId, callback }) + conversationService.onOutput(sessionId, callback) } type RuntimeSettings = { @@ -1926,9 +1991,12 @@ async function waitForRuntimeTransitionBeforeUserTurn( * Send a message to a specific session's WebSocket (for use by services) */ export function sendToSession(sessionId: string, message: ServerMessage): boolean { - const ws = activeSessions.get(sessionId) - if (!ws) return false - ws.send(JSON.stringify(message)) + const clients = activeSessions.get(sessionId) + if (!clients || clients.size === 0) return false + const payload = JSON.stringify(message) + for (const ws of clients) { + ws.send(payload) + } return true } @@ -1991,11 +2059,14 @@ export function closeSessionConnection(sessionId: string, reason = 'session clos conversationService.clearOutputCallbacks(sessionId) cleanupSessionRuntimeState(sessionId) - const ws = activeSessions.get(sessionId) - if (!ws) return false + const clients = activeSessions.get(sessionId) + if (!clients || clients.size === 0) return false activeSessions.delete(sessionId) - ws.close(1000, reason) + for (const ws of clients) { + clientOutputCallbacks.delete(ws) + ws.close(1000, reason) + } return true } @@ -2007,6 +2078,7 @@ export function __resetWebSocketHandlerStateForTests(): void { for (const timer of sessionCleanupTimers.values()) clearTimeout(timer) for (const timer of prewarmIdleTimers.values()) clearTimeout(timer) activeSessions.clear() + clientOutputCallbacks.clear() sessionCleanupTimers.clear() prewarmIdleTimers.clear() }