mirror of
https://github.com/NanmiCoder/cc-haha
synced 2026-07-18 13:23:33 +08:00
Merge Telegram stream subscription fix into local main
Bring the IM streaming repair from the ea41 worktree into local main so Telegram keeps receiving live content when the desktop opens the same active session. The merged fix also accumulates Telegram thinking deltas instead of replacing the preview with the latest short chunk. Constraint: Local main already contains independent worktree label cleanup changes, so this merge preserves both histories Rejected: Cherry-pick the fix | the user asked to merge the worktree back into local main Confidence: high Scope-risk: moderate Directive: Keep multi-client session streaming covered before changing WebSocket output callback ownership 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:
commit
77552375e9
@ -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')
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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<string, { chatId: string; messageId: number }>()
|
||||
// Track accumulated text per chat for streaming
|
||||
const accumulatedText = new Map<string, string>()
|
||||
const accumulatedThinkingText = new Map<string, string>()
|
||||
// Message buffers per chat
|
||||
const buffers = new Map<string, MessageBuffer>()
|
||||
// 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)) {
|
||||
|
||||
@ -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 () => {
|
||||
const providerService = new ProviderService()
|
||||
const providerA = await providerService.addProvider({
|
||||
|
||||
@ -107,8 +107,16 @@ export type WebSocketData = {
|
||||
serverHost: string
|
||||
}
|
||||
|
||||
// Active WebSocket sessions
|
||||
const activeSessions = new Map<string, ServerWebSocket<WebSocketData>>()
|
||||
// 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<string, Set<ServerWebSocket<WebSocketData>>>()
|
||||
const clientOutputCallbacks = new Map<
|
||||
ServerWebSocket<WebSocketData>,
|
||||
{
|
||||
sessionId: string
|
||||
callback: (cliMsg: any) => void
|
||||
}
|
||||
>()
|
||||
|
||||
export const handleWebSocket = {
|
||||
open(ws: ServerWebSocket<WebSocketData>) {
|
||||
@ -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<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,
|
||||
ws: ServerWebSocket<WebSocketData>,
|
||||
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()
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user