diff --git a/bun.lock b/bun.lock index f1f6d509..00497b27 100644 --- a/bun.lock +++ b/bun.lock @@ -90,7 +90,6 @@ "react-textarea-autosize": "^8.5.9", "remark-gfm": "^4.0.1", "shiki": "^3.20.0", - "socket.io-client": "^4.8.1", "tailwind-merge": "^2.5.5", }, "devDependencies": { diff --git a/server/src/index.ts b/server/src/index.ts index f347b800..79d3bae4 100644 --- a/server/src/index.ts +++ b/server/src/index.ts @@ -16,12 +16,14 @@ import { HappyBot } from './telegram/bot' import { startWebServer } from './web/server' import { getOrCreateJwtSecret } from './web/jwtSecret' import { createSocketServer } from './socket/server' +import { SSEManager } from './sse/sseManager' import type { Server as BunServer } from 'bun' import type { WebSocketData } from '@socket.io/bun-engine' let syncEngine: SyncEngine | null = null let happyBot: HappyBot | null = null let webServer: BunServer | null = null +let sseManager: SSEManager | null = null async function main() { console.log('Happy Bot starting...') @@ -34,16 +36,17 @@ async function main() { const store = new Store(config.dbPath) const jwtSecret = await getOrCreateJwtSecret() + sseManager = new SSEManager(30_000) + const socketServer = createSocketServer({ store, - jwtSecret, onWebappEvent: (event: SyncEvent) => syncEngine?.handleRealtimeEvent(event), onSessionAlive: (payload) => syncEngine?.handleSessionAlive(payload), onSessionEnd: (payload) => syncEngine?.handleSessionEnd(payload), onMachineAlive: (payload) => syncEngine?.handleMachineAlive(payload) }) - syncEngine = new SyncEngine(store, socketServer.io, socketServer.rpcRegistry) + syncEngine = new SyncEngine(store, socketServer.io, socketServer.rpcRegistry, sseManager) // Initialize Telegram bot happyBot = new HappyBot({ syncEngine }) @@ -51,6 +54,7 @@ async function main() { // Start HTTP server for Telegram Mini App webServer = await startWebServer({ getSyncEngine: () => syncEngine, + getSseManager: () => sseManager, jwtSecret, socketEngine: socketServer.engine }) @@ -65,6 +69,7 @@ async function main() { console.log('\nShutting down...') await happyBot?.stop() syncEngine?.stop() + sseManager?.stop() webServer?.stop() process.exit(0) } diff --git a/server/src/socket/handlers/webapp.ts b/server/src/socket/handlers/webapp.ts deleted file mode 100644 index 2355ab82..00000000 --- a/server/src/socket/handlers/webapp.ts +++ /dev/null @@ -1,68 +0,0 @@ -import type { Socket } from 'socket.io' -import { z } from 'zod' - -const subscribeSchema = z.object({ - all: z.boolean().optional(), - sessionId: z.string().optional(), - machineId: z.string().optional() -}) - -export type WebappHandlersDeps = { - onSubscribe?: (socket: Socket, data: { sessionId?: string; machineId?: string }) => void -} - -export function registerWebappHandlers(socket: Socket, deps: WebappHandlersDeps): void { - type SubscriptionState = { - all: boolean - sessionId: string | null - machineId: string | null - } - - let state: SubscriptionState = { - all: false, - sessionId: null, - machineId: null - } - - socket.on('subscribe', (data: unknown) => { - const parsed = subscribeSchema.safeParse(data ?? {}) - if (!parsed.success) { - return - } - - const { all, sessionId, machineId } = parsed.data - - const next: SubscriptionState = { - all: Boolean(all), - sessionId: sessionId ?? null, - machineId: machineId ?? null - } - - if (state.all && !next.all) { - socket.leave('webapp:all') - } - if (!state.all && next.all) { - socket.join('webapp:all') - } - - if (state.sessionId && state.sessionId !== next.sessionId) { - socket.leave(`session:${state.sessionId}`) - } - - if (next.sessionId && next.sessionId !== state.sessionId) { - socket.join(`session:${next.sessionId}`) - } - - if (state.machineId && state.machineId !== next.machineId) { - socket.leave(`machine:${state.machineId}`) - } - - if (next.machineId && next.machineId !== state.machineId) { - socket.join(`machine:${next.machineId}`) - } - - state = next - - deps.onSubscribe?.(socket, { sessionId: next.sessionId ?? undefined, machineId: next.machineId ?? undefined }) - }) -} diff --git a/server/src/socket/server.ts b/server/src/socket/server.ts index 39146442..58dc07c7 100644 --- a/server/src/socket/server.ts +++ b/server/src/socket/server.ts @@ -2,20 +2,12 @@ import { Server as Engine } from '@socket.io/bun-engine' import { Server } from 'socket.io' import type { Store } from '../store' import { configuration } from '../configuration' -import { jwtVerify } from 'jose' -import { z } from 'zod' import { registerCliHandlers } from './handlers/cli' -import { registerWebappHandlers } from './handlers/webapp' import { RpcRegistry } from './rpcRegistry' import type { SyncEvent } from '../sync/syncEngine' -const webappJwtPayloadSchema = z.object({ - uid: z.number() -}) - export type SocketServerDeps = { store: Store - jwtSecret: Uint8Array onWebappEvent?: (event: SyncEvent) => void onSessionAlive?: (payload: { sid: string; time: number; thinking?: boolean; mode?: 'local' | 'remote' }) => void onSessionEnd?: (payload: { sid: string; time: number }) => void @@ -74,28 +66,5 @@ export function createSocketServer(deps: SocketServerDeps): { onWebappEvent: deps.onWebappEvent })) - const webappNs = io.of('/webapp') - webappNs.use(async (socket, next) => { - const auth = socket.handshake.auth as Record | undefined - const token = typeof auth?.token === 'string' ? auth.token : null - if (!token) { - return next(new Error('Missing token')) - } - - try { - const verified = await jwtVerify(token, deps.jwtSecret, { algorithms: ['HS256'] }) - const parsed = webappJwtPayloadSchema.safeParse(verified.payload) - if (!parsed.success) { - return next(new Error('Invalid token payload')) - } - - socket.data.telegramUserId = parsed.data.uid - next() - } catch { - return next(new Error('Invalid token')) - } - }) - webappNs.on('connection', (socket) => registerWebappHandlers(socket, {})) - return { io, engine, rpcRegistry } } diff --git a/server/src/sse/sseManager.ts b/server/src/sse/sseManager.ts new file mode 100644 index 00000000..3fe03d53 --- /dev/null +++ b/server/src/sse/sseManager.ts @@ -0,0 +1,121 @@ +import type { SyncEvent } from '../sync/syncEngine' + +export type SSESubscription = { + id: string + all: boolean + sessionId: string | null + machineId: string | null +} + +type SSEConnection = SSESubscription & { + send: (event: SyncEvent) => void | Promise + sendHeartbeat: () => void | Promise +} + +export class SSEManager { + private readonly connections: Map = new Map() + private heartbeatTimer: NodeJS.Timeout | null = null + private readonly heartbeatMs: number + + constructor(heartbeatMs = 30_000) { + this.heartbeatMs = heartbeatMs + } + + subscribe(options: { + id: string + all?: boolean + sessionId?: string | null + machineId?: string | null + send: (event: SyncEvent) => void | Promise + sendHeartbeat: () => void | Promise + }): SSESubscription { + const subscription: SSEConnection = { + id: options.id, + all: Boolean(options.all), + sessionId: options.sessionId ?? null, + machineId: options.machineId ?? null, + send: options.send, + sendHeartbeat: options.sendHeartbeat + } + + this.connections.set(subscription.id, subscription) + this.ensureHeartbeat() + return { + id: subscription.id, + all: subscription.all, + sessionId: subscription.sessionId, + machineId: subscription.machineId + } + } + + unsubscribe(id: string): void { + this.connections.delete(id) + if (this.connections.size === 0) { + this.stopHeartbeat() + } + } + + broadcast(event: SyncEvent): void { + for (const connection of this.connections.values()) { + if (!this.shouldSend(connection, event)) { + continue + } + + void Promise.resolve(connection.send(event)).catch(() => { + this.unsubscribe(connection.id) + }) + } + } + + stop(): void { + this.stopHeartbeat() + this.connections.clear() + } + + private ensureHeartbeat(): void { + if (this.heartbeatTimer || this.heartbeatMs <= 0) { + return + } + + this.heartbeatTimer = setInterval(() => { + for (const connection of this.connections.values()) { + void Promise.resolve(connection.sendHeartbeat()).catch(() => { + this.unsubscribe(connection.id) + }) + } + }, this.heartbeatMs) + } + + private stopHeartbeat(): void { + if (!this.heartbeatTimer) { + return + } + + clearInterval(this.heartbeatTimer) + this.heartbeatTimer = null + } + + private shouldSend(connection: SSEConnection, event: SyncEvent): boolean { + if (event.type === 'message-received') { + return Boolean(event.sessionId && connection.sessionId === event.sessionId) + } + + if (event.type === 'connection-changed') { + return true + } + + if (connection.all) { + return true + } + + if (event.sessionId && connection.sessionId === event.sessionId) { + return true + } + + if (event.machineId && connection.machineId === event.machineId) { + return true + } + + return false + } +} diff --git a/server/src/sync/syncEngine.ts b/server/src/sync/syncEngine.ts index e982c9ab..4689ec5c 100644 --- a/server/src/sync/syncEngine.ts +++ b/server/src/sync/syncEngine.ts @@ -11,6 +11,7 @@ import { z } from 'zod' import type { Server } from 'socket.io' import type { Store } from '../store' import type { RpcRegistry } from '../socket/rpcRegistry' +import type { SSEManager } from '../sse/sseManager' import { extractTodoWriteTodosFromMessageContent, TodosSchema, type TodoItem } from './todos' export type ConnectionStatus = 'disconnected' | 'connected' @@ -152,7 +153,8 @@ export class SyncEngine { constructor( private readonly store: Store, private readonly io: Server, - private readonly rpcRegistry: RpcRegistry + private readonly rpcRegistry: RpcRegistry, + private readonly sseManager: SSEManager ) { this.reloadAll() this.inactivityTimer = setInterval(() => this.expireInactive(), 5_000) @@ -196,22 +198,7 @@ export class SyncEngine { machineId: event.machineId } - const rooms = new Set() - if (webappEvent.sessionId) { - rooms.add(`session:${webappEvent.sessionId}`) - } - if (webappEvent.machineId) { - rooms.add(`machine:${webappEvent.machineId}`) - } - - if (webappEvent.type !== 'message-received') { - rooms.add('webapp:all') - } - - const webappNamespace = this.io.of('/webapp') - for (const room of rooms) { - webappNamespace.to(room).emit('update', webappEvent) - } + this.sseManager.broadcast(webappEvent) } getConnectionStatus(): ConnectionStatus { diff --git a/server/src/web/middleware/auth.ts b/server/src/web/middleware/auth.ts index c8990a1f..39f34a4a 100644 --- a/server/src/web/middleware/auth.ts +++ b/server/src/web/middleware/auth.ts @@ -22,10 +22,11 @@ export function createAuthMiddleware(jwtSecret: Uint8Array): MiddlewareHandler SSEManager | null): Hono { + const app = new Hono() + + app.get('/events', (c) => { + const manager = getSseManager() + if (!manager) { + return c.json({ error: 'Not connected' }, 503) + } + + const query = c.req.query() + const all = parseBoolean(query.all) + const sessionId = parseOptionalId(query.sessionId) + const machineId = parseOptionalId(query.machineId) + const subscriptionId = randomUUID() + + return streamSSE(c, async (stream) => { + manager.subscribe({ + id: subscriptionId, + all, + sessionId, + machineId, + send: (event) => stream.writeSSE({ data: JSON.stringify(event) }), + sendHeartbeat: async () => { + await stream.write(': heartbeat\n\n') + } + }) + + await new Promise((resolve) => { + const done = () => resolve() + c.req.raw.signal.addEventListener('abort', done, { once: true }) + stream.onAbort(done) + }) + + manager.unsubscribe(subscriptionId) + }) + }) + + return app +} diff --git a/server/src/web/server.ts b/server/src/web/server.ts index e8cff2b5..a842116f 100644 --- a/server/src/web/server.ts +++ b/server/src/web/server.ts @@ -8,11 +8,13 @@ import { configuration } from '../configuration' import type { SyncEngine } from '../sync/syncEngine' import { createAuthMiddleware, type WebAppEnv } from './middleware/auth' import { createAuthRoutes } from './routes/auth' +import { createEventsRoutes } from './routes/events' import { createSessionsRoutes } from './routes/sessions' import { createMessagesRoutes } from './routes/messages' import { createPermissionsRoutes } from './routes/permissions' import { createMachinesRoutes } from './routes/machines' import { createCliRoutes } from './routes/cli' +import type { SSEManager } from '../sse/sseManager' import type { Server as BunServer } from 'bun' import type { Server as SocketEngine } from '@socket.io/bun-engine' import type { WebSocketData } from '@socket.io/bun-engine' @@ -37,6 +39,7 @@ function findWebappDistDir(): { distDir: string; indexHtmlPath: string } { function createWebApp(options: { getSyncEngine: () => SyncEngine | null + getSseManager: () => SSEManager | null jwtSecret: Uint8Array }): Hono { const app = new Hono() @@ -58,6 +61,7 @@ function createWebApp(options: { app.route('/api', createAuthRoutes(options.jwtSecret)) app.use('/api/*', createAuthMiddleware(options.jwtSecret)) + app.route('/api', createEventsRoutes(options.getSseManager)) app.route('/api', createSessionsRoutes(options.getSyncEngine)) app.route('/api', createMessagesRoutes(options.getSyncEngine)) app.route('/api', createPermissionsRoutes(options.getSyncEngine)) @@ -100,11 +104,13 @@ function createWebApp(options: { export async function startWebServer(options: { getSyncEngine: () => SyncEngine | null + getSseManager: () => SSEManager | null jwtSecret: Uint8Array socketEngine: SocketEngine }): Promise> { const app = createWebApp({ getSyncEngine: options.getSyncEngine, + getSseManager: options.getSseManager, jwtSecret: options.jwtSecret }) diff --git a/web/index.html b/web/index.html index ad01fa5b..dc2fcb52 100644 --- a/web/index.html +++ b/web/index.html @@ -26,13 +26,6 @@ Hapi -
diff --git a/web/package.json b/web/package.json index 589a1c2e..f4d6adb4 100644 --- a/web/package.json +++ b/web/package.json @@ -27,7 +27,6 @@ "react-textarea-autosize": "^8.5.9", "remark-gfm": "^4.0.1", "shiki": "^3.20.0", - "socket.io-client": "^4.8.1", "tailwind-merge": "^2.5.5" }, "devDependencies": { diff --git a/web/src/App.tsx b/web/src/App.tsx index eb4a9221..4a29a17e 100644 --- a/web/src/App.tsx +++ b/web/src/App.tsx @@ -4,7 +4,7 @@ import { getTelegramWebApp, isTelegramApp } from '@/hooks/useTelegram' import { initializeTheme } from '@/hooks/useTheme' import { useAuth } from '@/hooks/useAuth' import { useAuthSource } from '@/hooks/useAuthSource' -import { useSocket } from '@/hooks/useSocket' +import { useSSE } from '@/hooks/useSSE' import { queryKeys } from '@/lib/query-keys' import { useMessages } from '@/hooks/queries/useMessages' import { useMachines } from '@/hooks/queries/useMachines' @@ -198,7 +198,7 @@ export function App() { void refetchMessages() }, [selectedSessionId, refetchMessages, refetchSession]) - const handleSocketConnect = useCallback(() => { + const handleSseConnect = useCallback(() => { void queryClient.invalidateQueries({ queryKey: queryKeys.sessions }) if (selectedSessionId) { void queryClient.invalidateQueries({ queryKey: queryKeys.session(selectedSessionId) }) @@ -206,9 +206,9 @@ export function App() { } }, [queryClient, selectedSessionId]) - const handleSocketEvent = useCallback(() => {}, []) + const handleSseEvent = useCallback(() => {}, []) - const socketSubscription = useMemo(() => { + const eventSubscription = useMemo(() => { if (screen.type === 'session') { return { sessionId: screen.sessionId } } @@ -218,12 +218,12 @@ export function App() { return { all: true } }, [screen]) - useSocket({ + useSSE({ enabled: Boolean(api && token), token: token ?? '', - subscription: socketSubscription, - onConnect: handleSocketConnect, - onEvent: handleSocketEvent, + subscription: eventSubscription, + onConnect: handleSseConnect, + onEvent: handleSseEvent, }) // Loading auth source diff --git a/web/src/hooks/useSocket.ts b/web/src/hooks/useSSE.ts similarity index 58% rename from web/src/hooks/useSocket.ts rename to web/src/hooks/useSSE.ts index 9f5773f0..4b6f205c 100644 --- a/web/src/hooks/useSocket.ts +++ b/web/src/hooks/useSSE.ts @@ -1,6 +1,5 @@ -import { useEffect, useRef } from 'react' +import { useEffect, useMemo, useRef } from 'react' import { useQueryClient, type InfiniteData } from '@tanstack/react-query' -import { io } from 'socket.io-client' import type { MessagesResponse, SyncEvent } from '@/types/api' import { queryKeys } from '@/lib/query-keys' import { upsertMessagesInCache } from '@/lib/messages' @@ -9,16 +8,32 @@ function isObject(value: unknown): value is Record { return Boolean(value) && typeof value === 'object' } -type SocketSubscription = { +type SSESubscription = { all?: boolean sessionId?: string machineId?: string } -export function useSocket(options: { +function buildEventsUrl(token: string, subscription: SSESubscription): string { + const params = new URLSearchParams() + params.set('token', token) + if (subscription.all) { + params.set('all', 'true') + } + if (subscription.sessionId) { + params.set('sessionId', subscription.sessionId) + } + if (subscription.machineId) { + params.set('machineId', subscription.machineId) + } + + return `/api/events?${params.toString()}` +} + +export function useSSE(options: { enabled: boolean token: string - subscription?: SocketSubscription + subscription?: SSESubscription onEvent: (event: SyncEvent) => void onConnect?: () => void onDisconnect?: (reason: string) => void @@ -29,8 +44,7 @@ export function useSocket(options: { const onConnectRef = useRef(options.onConnect) const onDisconnectRef = useRef(options.onDisconnect) const onErrorRef = useRef(options.onError) - const subscriptionRef = useRef(options.subscription ?? {}) - const socketRef = useRef | null>(null) + const eventSourceRef = useRef(null) useEffect(() => { onEventRef.current = options.onEvent @@ -48,47 +62,32 @@ export function useSocket(options: { onDisconnectRef.current = options.onDisconnect }, [options.onDisconnect]) - useEffect(() => { - subscriptionRef.current = options.subscription ?? {} - }, [options.subscription]) + const subscription = options.subscription ?? {} + const subscriptionKey = useMemo(() => { + return `${subscription.all ? '1' : '0'}|${subscription.sessionId ?? ''}|${subscription.machineId ?? ''}` + }, [subscription.all, subscription.sessionId, subscription.machineId]) useEffect(() => { if (!options.enabled) { - socketRef.current?.disconnect() - socketRef.current = null + eventSourceRef.current?.close() + eventSourceRef.current = null return } - const socket = io('/webapp', { - auth: { token: options.token }, - reconnection: true, - reconnectionDelay: 1000, - reconnectionDelayMax: 5000, - reconnectionAttempts: Infinity - }) - socketRef.current = socket - - const sendSubscribe = () => { - socket.emit('subscribe', subscriptionRef.current) - } - const handleConnect = () => { - sendSubscribe() - onConnectRef.current?.() - } - const handleDisconnect = (reason: string) => { - onDisconnectRef.current?.(reason) - } + const url = buildEventsUrl(options.token, subscription) + const eventSource = new EventSource(url) + eventSourceRef.current = eventSource const handleSyncEvent = (event: SyncEvent) => { if (event.type === 'message-received') { queryClient.setQueryData>( queryKeys.messages(event.sessionId), - (data) => upsertMessagesInCache(data, [event.message]), + (data) => upsertMessagesInCache(data, [event.message]) ) // Mark stale so the initial query still fetches history when it mounts. void queryClient.invalidateQueries({ queryKey: queryKeys.messages(event.sessionId), - refetchType: 'none', + refetchType: 'none' }) } @@ -111,44 +110,43 @@ export function useSocket(options: { onEventRef.current(event) } - socket.on('update', (event: unknown) => { - if (!isObject(event)) return - if (typeof event.type !== 'string') return - handleSyncEvent(event as SyncEvent) - }) + const handleMessage = (message: MessageEvent) => { + if (typeof message.data !== 'string') { + return + } - socket.on('connect_error', (error) => { + let parsed: unknown + try { + parsed = JSON.parse(message.data) + } catch { + return + } + + if (!isObject(parsed)) { + return + } + if (typeof parsed.type !== 'string') { + return + } + + handleSyncEvent(parsed as SyncEvent) + } + + eventSource.onmessage = handleMessage + eventSource.onopen = () => { + onConnectRef.current?.() + } + eventSource.onerror = (error) => { onErrorRef.current?.(error) - }) - - socket.on('error', (error) => { - onErrorRef.current?.(error) - }) - - socket.on('connect', handleConnect) - socket.on('disconnect', handleDisconnect) - sendSubscribe() + const reason = eventSource.readyState === EventSource.CLOSED ? 'closed' : 'error' + onDisconnectRef.current?.(reason) + } return () => { - socket.off('connect', handleConnect) - socket.off('disconnect', handleDisconnect) - socket.disconnect() - if (socketRef.current === socket) { - socketRef.current = null + eventSource.close() + if (eventSourceRef.current === eventSource) { + eventSourceRef.current = null } } - }, [options.enabled, options.token]) - - useEffect(() => { - if (!options.enabled) { - return - } - - const socket = socketRef.current - if (!socket) { - return - } - - socket.emit('subscribe', subscriptionRef.current) - }, [options.enabled, options.subscription]) + }, [options.enabled, options.token, subscriptionKey, queryClient]) }