@mono-labs/dev 0.1.251 → 0.1.256

package/src/websocket/index.ts

@@ -3,13 +3,20 @@ import type { WebSocket, WebSocketServer } from 'ws'
 import { ActionRouter } from './action-router'
 import { ConnectionRegistry } from './connection-registry'
 import { buildRequestContext } from './event-synthesizer'
-import { LocalGatewayClient } from './local-gateway-client'
+import { SocketGatewayClient } from './socket-gateway-client'
+import { InMemoryChannelStore, RedisChannelStore } from './channel-store'
+import { SocketEmitter } from './socket-emitter'
+import { initCacheRelay } from '../cache-relay'
 import type { ConnectionId, SocketAdapterConfig } from './types'
 
-export type { ConnectionId, PostToConnectionFn, SocketAdapterConfig } from './types'
+export type { ConnectionId, PostToConnectionFn, SocketAdapterConfig, RedisConfig } from './types'
 export { ConnectionRegistry } from './connection-registry'
-export { LocalGatewayClient } from './local-gateway-client'
+export { SocketGatewayClient } from './socket-gateway-client'
 export { ActionRouter } from './action-router'
+export { InMemoryChannelStore, RedisChannelStore } from './channel-store'
+export type { ChannelStore } from './channel-store'
+export { SocketEmitter } from './socket-emitter'
+export type { EmitTarget } from './socket-emitter'
 
 /**
  * Attaches a full socket adapter to a WebSocketServer instance.
@@ -22,10 +29,26 @@ export function attachSocketAdapter(wss: WebSocketServer, config?: SocketAdapter
 	const stage = config?.stage ?? 'local'
 
 	const connectionRegistry = new ConnectionRegistry()
-	const gatewayClient = new LocalGatewayClient(connectionRegistry)
+	const gatewayClient = new SocketGatewayClient(connectionRegistry)
 	const postToConnection = gatewayClient.asFunction()
 	const actionRouter = new ActionRouter()
 
+	// Create channel store
+	let channelStore = config?.channelStore
+	if (!channelStore) {
+		if (config?.useRedis) {
+			const host = config.redis?.host ?? 'localhost'
+			const port = config.redis?.port ?? 6379
+			initCacheRelay(`${host}:${port}`)
+			channelStore = new RedisChannelStore({ keyPrefix: config.redis?.keyPrefix })
+		} else {
+			channelStore = new InMemoryChannelStore()
+		}
+	}
+
+	// Create socket emitter
+	const socketEmitter = new SocketEmitter({ postToConnection, connectionRegistry, channelStore })
+
 	// Register consumer-provided routes
 	if (config?.routes) {
 		for (const [action, handler] of Object.entries(config.routes)) {
@@ -131,6 +154,7 @@ export function attachSocketAdapter(wss: WebSocketServer, config?: SocketAdapter
 				)
 			}
 
+			await channelStore.removeAll(connectionId)
 			await disconnectHandler(connectionId)
 			connectionRegistry.unregister(connectionId)
 			wsToConnectionId.delete(ws)
@@ -141,6 +165,8 @@ export function attachSocketAdapter(wss: WebSocketServer, config?: SocketAdapter
 		postToConnection,
 		connectionRegistry,
 		actionRouter,
+		channelStore,
+		socketEmitter,
 		getConnectionId: (ws: WebSocket) => wsToConnectionId.get(ws),
 	}
 }

package/src/websocket/socket-emitter.ts

@@ -0,0 +1,64 @@
+import type { ConnectionRegistry } from './connection-registry'
+import type { ChannelStore } from './channel-store'
+import type { ConnectionId, PostToConnectionFn } from './types'
+
+export type EmitTarget =
+	| { userId: string }
+	| { orgId: string }
+	| { connectionId: string }
+	| { channel: string }
+	| 'broadcast'
+
+export class SocketEmitter {
+	private postToConnection: PostToConnectionFn
+	private connectionRegistry: ConnectionRegistry
+	private channelStore: ChannelStore
+
+	constructor(deps: {
+		postToConnection: PostToConnectionFn
+		connectionRegistry: ConnectionRegistry
+		channelStore: ChannelStore
+	}) {
+		this.postToConnection = deps.postToConnection
+		this.connectionRegistry = deps.connectionRegistry
+		this.channelStore = deps.channelStore
+	}
+
+	async emit(target: EmitTarget, data: unknown): Promise<void> {
+		const connectionIds = await this.resolveConnectionIds(target)
+		if (connectionIds.length === 0) return
+
+		const payload = typeof data === 'string' ? data : JSON.stringify(data)
+
+		await Promise.allSettled(
+			connectionIds.map(async (connId) => {
+				try {
+					await this.postToConnection(connId, payload)
+				} catch (err: unknown) {
+					const e = err as { statusCode?: number; name?: string }
+					if (e?.statusCode === 410 || e?.name === 'GoneException') return
+					console.error(`[SocketEmitter] failed to send to ${connId}:`, err)
+				}
+			})
+		)
+	}
+
+	private async resolveConnectionIds(target: EmitTarget): Promise<ConnectionId[]> {
+		if (target === 'broadcast') {
+			return this.connectionRegistry.getAll()
+		}
+		if ('connectionId' in target) {
+			return [target.connectionId]
+		}
+		if ('userId' in target) {
+			return this.connectionRegistry.getConnectionsByUserId(target.userId)
+		}
+		if ('orgId' in target) {
+			return this.connectionRegistry.getConnectionsByOrgId(target.orgId)
+		}
+		if ('channel' in target) {
+			return this.channelStore.getSubscribers(target.channel)
+		}
+		return []
+	}
+}

package/src/websocket/socket-gateway-client.ts

@@ -0,0 +1,83 @@
+import type { ConnectionId, PostToConnectionFn } from './types'
+import type { ConnectionRegistry } from './connection-registry'
+
+/**
+ * Unified gateway client that auto-detects local vs production mode.
+ *
+ * - Pass a `ConnectionRegistry` → local mode (sends via in-memory WebSocket)
+ * - Pass a URL string → API Gateway mode (sends via AWS SDK PostToConnectionCommand)
+ */
+export class SocketGatewayClient {
+	private mode: 'local' | 'apigateway'
+	private registry?: ConnectionRegistry
+	private endpoint?: string
+
+	constructor(backend: ConnectionRegistry | string) {
+		if (typeof backend === 'string') {
+			this.mode = 'apigateway'
+			this.endpoint = backend
+		} else {
+			this.mode = 'local'
+			this.registry = backend
+		}
+	}
+
+	async postToConnection(connectionId: ConnectionId, data: unknown): Promise<void> {
+		if (this.mode === 'local') {
+			return this.postLocal(connectionId, data)
+		}
+		return this.postApiGateway(connectionId, data)
+	}
+
+	private async postLocal(connectionId: ConnectionId, data: unknown): Promise<void> {
+		const { WebSocket } = require('ws') as typeof import('ws')
+		const ws = this.registry!.get(connectionId)
+
+		if (!ws || ws.readyState !== WebSocket.OPEN) {
+			const error = new Error(`GoneException: Connection ${connectionId} is no longer available`)
+			;(error as any).statusCode = 410
+			;(error as any).name = 'GoneException'
+			throw error
+		}
+
+		const payload = typeof data === 'string' ? data : JSON.stringify(data)
+		ws.send(payload)
+	}
+
+	private async postApiGateway(connectionId: ConnectionId, data: unknown): Promise<void> {
+		let ApiGatewayManagementApiClient: any, PostToConnectionCommand: any
+		try {
+			const mod = require('@aws-sdk/client-apigatewaymanagementapi')
+			ApiGatewayManagementApiClient = mod.ApiGatewayManagementApiClient
+			PostToConnectionCommand = mod.PostToConnectionCommand
+		} catch {
+			throw new Error(
+				'API Gateway mode requires "@aws-sdk/client-apigatewaymanagementapi". '
+				+ 'Install it with: npm install @aws-sdk/client-apigatewaymanagementapi'
+			)
+		}
+
+		const client = new ApiGatewayManagementApiClient({ endpoint: this.endpoint })
+		const payload = typeof data === 'string' ? data : JSON.stringify(data)
+
+		try {
+			await client.send(new PostToConnectionCommand({
+				ConnectionId: connectionId,
+				Data: new TextEncoder().encode(payload),
+			}))
+		} catch (err: any) {
+			if (err.name === 'GoneException' || err.$metadata?.httpStatusCode === 410) {
+				const error = new Error(`GoneException: Connection ${connectionId} is no longer available`)
+				;(error as any).statusCode = 410
+				;(error as any).name = 'GoneException'
+				throw error
+			}
+			throw err
+		}
+	}
+
+	/** Returns a bound PostToConnectionFn for dependency injection */
+	asFunction(): PostToConnectionFn {
+		return this.postToConnection.bind(this)
+	}
+}

package/src/websocket/types.ts

@@ -54,6 +54,13 @@ export type ConnectHandlerFn = (
 /** Handler called on $disconnect */
 export type DisconnectHandlerFn = (connectionId: ConnectionId) => Promise<void>
 
+/** Redis connection configuration */
+export interface RedisConfig {
+	host?: string
+	port?: number
+	keyPrefix?: string
+}
+
 /** Configuration for the socket adapter */
 export interface SocketAdapterConfig {
 	domainName?: string
@@ -63,4 +70,7 @@ export interface SocketAdapterConfig {
 	disconnectHandler?: DisconnectHandlerFn
 	routes?: Record<string, ActionHandler>
 	defaultHandler?: ActionHandler
+	channelStore?: import('./channel-store').ChannelStore
+	useRedis?: boolean
+	redis?: RedisConfig
 }

package/src/websocket/local-gateway-client.ts

@@ -1,31 +0,0 @@
-import { WebSocket } from 'ws'
-import type { ConnectionRegistry } from './connection-registry'
-import type { ConnectionId, PostToConnectionFn } from './types'
-
-/**
- * Local replacement for @aws-sdk/client-apigatewaymanagementapi.
- * Looks up the WebSocket in ConnectionRegistry and sends data directly.
- */
-export class LocalGatewayClient {
-	constructor(private registry: ConnectionRegistry) {}
-
-	/** Send data to a specific connection. Throws GoneException (410) if not found. */
-	async postToConnection(connectionId: ConnectionId, data: unknown): Promise<void> {
-		const ws = this.registry.get(connectionId)
-
-		if (!ws || ws.readyState !== WebSocket.OPEN) {
-			const error = new Error(`GoneException: Connection ${connectionId} is no longer available`)
-			;(error as any).statusCode = 410
-			;(error as any).name = 'GoneException'
-			throw error
-		}
-
-		const payload = typeof data === 'string' ? data : JSON.stringify(data)
-		ws.send(payload)
-	}
-
-	/** Returns a bound PostToConnectionFn for dependency injection */
-	asFunction(): PostToConnectionFn {
-		return this.postToConnection.bind(this)
-	}
-}