@mono-labs/dev 0.1.251

package/src/websocket/action-router.ts

@@ -0,0 +1,66 @@
+import type {
+	ActionHandler,
+	ActionHandlerContext,
+	ActionHandlerResult,
+	ConnectionId,
+	LocalRequestContext,
+	PostToConnectionFn,
+} from './types'
+import type { WebSocketUserContext } from './types'
+
+/**
+ * Routes incoming WebSocket messages by `action` field — mirrors
+ * API Gateway's `$request.body.action` route selection.
+ */
+export class ActionRouter {
+	private routes = new Map<string, ActionHandler>()
+	private defaultHandler: ActionHandler | null = null
+
+	/** Register a handler for a specific action (equivalent to webSocketApi.addRoute()) */
+	addRoute(action: string, handler: ActionHandler): void {
+		this.routes.set(action, handler)
+	}
+
+	/** Set the $default handler for unknown actions */
+	setDefaultHandler(handler: ActionHandler): void {
+		this.defaultHandler = handler
+	}
+
+	/** Parse incoming message JSON, extract `action`, and dispatch to the matching handler */
+	async route(
+		connectionId: ConnectionId,
+		rawBody: string,
+		postToConnection: PostToConnectionFn,
+		requestContext: LocalRequestContext,
+		userContext: WebSocketUserContext
+	): Promise<ActionHandlerResult> {
+		let parsed: { action?: string; [k: string]: unknown }
+		try {
+			parsed = JSON.parse(rawBody)
+		} catch {
+			return { statusCode: 400, body: JSON.stringify({ error: 'Invalid JSON' }) }
+		}
+
+		const action = parsed.action
+		if (!action || typeof action !== 'string') {
+			return { statusCode: 400, body: JSON.stringify({ error: 'Missing "action" field' }) }
+		}
+
+		const handler = this.routes.get(action) ?? this.defaultHandler
+		if (!handler) {
+			return {
+				statusCode: 400,
+				body: JSON.stringify({ error: `Unknown action: ${action}` }),
+			}
+		}
+
+		const ctx: ActionHandlerContext = {
+			connectionId,
+			requestContext: { ...requestContext, routeKey: action },
+			postToConnection,
+			userContext,
+		}
+
+		return handler(rawBody, ctx)
+	}
+}

package/src/websocket/connection-registry.ts

@@ -0,0 +1,67 @@
+import type { WebSocket } from 'ws'
+import type { ConnectionId } from './types'
+import type { WebSocketUserContext } from './types'
+
+/** Generates a unique connection ID (same pattern as socket-real.ts) */
+const makeConnectionId = (): ConnectionId =>
+	`${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 10)}`
+
+/**
+ * In-memory connection registry — replaces API Gateway's built-in connection management.
+ * Maps connectionId → WebSocket instance and connectionId → user context.
+ */
+export class ConnectionRegistry {
+	private connections = new Map<ConnectionId, WebSocket>()
+	private userContexts = new Map<ConnectionId, WebSocketUserContext>()
+
+	/** Register a WebSocket and return its assigned connectionId */
+	register(ws: WebSocket): ConnectionId {
+		const connectionId = makeConnectionId()
+		this.connections.set(connectionId, ws)
+		return connectionId
+	}
+
+	/** Remove a connection by ID */
+	unregister(connectionId: ConnectionId): void {
+		this.connections.delete(connectionId)
+		this.userContexts.delete(connectionId)
+	}
+
+	/** Look up a WebSocket by connectionId */
+	get(connectionId: ConnectionId): WebSocket | undefined {
+		return this.connections.get(connectionId)
+	}
+
+	/** Return all active connectionIds */
+	getAll(): ConnectionId[] {
+		return Array.from(this.connections.keys())
+	}
+
+	/** Store user context for a connection */
+	setUserContext(connectionId: ConnectionId, ctx: WebSocketUserContext): void {
+		this.userContexts.set(connectionId, ctx)
+	}
+
+	/** Retrieve user context for a connection */
+	getUserContext(connectionId: ConnectionId): WebSocketUserContext | undefined {
+		return this.userContexts.get(connectionId)
+	}
+
+	/** Find all connectionIds belonging to a specific user */
+	getConnectionsByUserId(userId: string): ConnectionId[] {
+		const result: ConnectionId[] = []
+		for (const [connId, ctx] of this.userContexts) {
+			if (ctx.userId === userId) result.push(connId)
+		}
+		return result
+	}
+
+	/** Find all connectionIds belonging to a specific organization */
+	getConnectionsByOrgId(orgId: string): ConnectionId[] {
+		const result: ConnectionId[] = []
+		for (const [connId, ctx] of this.userContexts) {
+			if (ctx.organizationId === orgId) result.push(connId)
+		}
+		return result
+	}
+}

package/src/websocket/event-synthesizer.ts

@@ -0,0 +1,131 @@
+import type { APIGatewayProxyEvent } from 'aws-lambda'
+import type http from 'node:http'
+import type { ConnectionId, LocalRequestContext } from './types'
+
+const EMPTY_MULTI_VALUE: Record<string, string[]> = {}
+const EMPTY_HEADERS: Record<string, string> = {}
+
+/** Creates a minimal APIGatewayProxyEvent shell with the given requestContext fields */
+function baseEvent(
+	requestContext: Partial<APIGatewayProxyEvent['requestContext']>,
+	overrides?: Partial<APIGatewayProxyEvent>
+): APIGatewayProxyEvent {
+	return {
+		body: null,
+		headers: EMPTY_HEADERS,
+		multiValueHeaders: EMPTY_MULTI_VALUE,
+		httpMethod: 'GET',
+		isBase64Encoded: false,
+		path: '',
+		pathParameters: null,
+		queryStringParameters: null,
+		multiValueQueryStringParameters: null,
+		stageVariables: null,
+		resource: '',
+		requestContext: {
+			accountId: 'local',
+			apiId: 'local',
+			authorizer: null,
+			protocol: 'websocket',
+			httpMethod: 'GET',
+			identity: {
+				accessKey: null,
+				accountId: null,
+				apiKey: null,
+				apiKeyId: null,
+				caller: null,
+				clientCert: null,
+				cognitoAuthenticationProvider: null,
+				cognitoAuthenticationType: null,
+				cognitoIdentityId: null,
+				cognitoIdentityPoolId: null,
+				principalOrgId: null,
+				sourceIp: '127.0.0.1',
+				user: null,
+				userAgent: null,
+				userArn: null,
+			},
+			path: '',
+			stage: 'local',
+			requestId: `local-${Date.now()}`,
+			requestTimeEpoch: Date.now(),
+			resourceId: '',
+			resourcePath: '',
+			...requestContext,
+		},
+		...overrides,
+	}
+}
+
+/** Synthesize a $connect event */
+export function synthesizeConnectEvent(
+	connectionId: ConnectionId,
+	req: http.IncomingMessage,
+	config?: { domainName?: string; stage?: string }
+): APIGatewayProxyEvent {
+	const headers: Record<string, string> = {}
+	for (const [key, value] of Object.entries(req.headers)) {
+		if (typeof value === 'string') headers[key] = value
+		else if (Array.isArray(value)) headers[key] = value.join(', ')
+	}
+
+	return baseEvent(
+		{
+			connectionId,
+			domainName: config?.domainName ?? 'localhost',
+			stage: config?.stage ?? 'local',
+			routeKey: '$connect',
+			eventType: 'CONNECT',
+		} as any,
+		{ headers }
+	)
+}
+
+/** Synthesize a $disconnect event */
+export function synthesizeDisconnectEvent(
+	connectionId: ConnectionId,
+	config?: { domainName?: string; stage?: string }
+): APIGatewayProxyEvent {
+	return baseEvent({
+		connectionId,
+		domainName: config?.domainName ?? 'localhost',
+		stage: config?.stage ?? 'local',
+		routeKey: '$disconnect',
+		eventType: 'DISCONNECT',
+	} as any)
+}
+
+/** Synthesize a message event with body and routeKey */
+export function synthesizeMessageEvent(
+	connectionId: ConnectionId,
+	body: string,
+	routeKey: string,
+	config?: { domainName?: string; stage?: string }
+): APIGatewayProxyEvent {
+	return baseEvent(
+		{
+			connectionId,
+			domainName: config?.domainName ?? 'localhost',
+			stage: config?.stage ?? 'local',
+			routeKey,
+			eventType: 'MESSAGE',
+		} as any,
+		{ body }
+	)
+}
+
+/** Build a LocalRequestContext from parameters */
+export function buildRequestContext(
+	connectionId: ConnectionId,
+	routeKey: string,
+	eventType: LocalRequestContext['eventType'],
+	config?: { domainName?: string; stage?: string }
+): LocalRequestContext {
+	return {
+		connectionId,
+		domainName: config?.domainName ?? 'localhost',
+		stage: config?.stage ?? 'local',
+		routeKey,
+		eventType,
+	}
+}

package/src/websocket/index.ts

@@ -0,0 +1,146 @@
+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 type { ConnectionId, SocketAdapterConfig } from './types'
+
+export type { ConnectionId, PostToConnectionFn, SocketAdapterConfig } from './types'
+export { ConnectionRegistry } from './connection-registry'
+export { LocalGatewayClient } from './local-gateway-client'
+export { ActionRouter } from './action-router'
+
+/**
+ * Attaches a full socket adapter to a WebSocketServer instance.
+ * Maps 1:1 to each API Gateway WebSocket event so local dev
+ * behaves identically to the deployed system.
+ */
+export function attachSocketAdapter(wss: WebSocketServer, config?: SocketAdapterConfig) {
+	const debug = config?.debug ?? false
+	const domainName = config?.domainName ?? 'localhost'
+	const stage = config?.stage ?? 'local'
+
+	const connectionRegistry = new ConnectionRegistry()
+	const gatewayClient = new LocalGatewayClient(connectionRegistry)
+	const postToConnection = gatewayClient.asFunction()
+	const actionRouter = new ActionRouter()
+
+	// Register consumer-provided routes
+	if (config?.routes) {
+		for (const [action, handler] of Object.entries(config.routes)) {
+			actionRouter.addRoute(action, handler)
+		}
+	}
+
+	// $default handler for unknown actions
+	if (config?.defaultHandler) {
+		actionRouter.setDefaultHandler(config.defaultHandler)
+	} else {
+		actionRouter.setDefaultHandler(async (body, ctx) => {
+			let action = 'unknown'
+			try {
+				action = JSON.parse(body).action ?? 'unknown'
+			} catch {}
+			return {
+				statusCode: 400,
+				body: JSON.stringify({ error: `Unknown action: ${action}` }),
+			}
+		})
+	}
+
+	// Use consumer-provided handlers or sensible defaults
+	const connectHandler = config?.connectHandler ?? (async () => ({
+		response: { statusCode: 200 },
+		userContext: undefined,
+	}))
+
+	const disconnectHandler = config?.disconnectHandler ?? (async () => {})
+
+	// Reverse lookup: WebSocket → connectionId
+	const wsToConnectionId = new WeakMap<WebSocket, ConnectionId>()
+
+	wss.on('connection', async (ws: WebSocket, req) => {
+		// 1. Extract token from query string
+		const url = new URL(req.url ?? '', `http://${req.headers.host ?? 'localhost'}`)
+		const token = url.searchParams.get('token') ?? undefined
+
+		// 2. Register connection and assign connectionId
+		const connectionId = connectionRegistry.register(ws)
+		wsToConnectionId.set(ws, connectionId)
+
+		if (debug) console.log(`[socket-adapter] $connect connectionId=${connectionId}`)
+		if (debug) console.log(`[socket-adapter] token ${token ? 'present' : 'MISSING'}`)
+
+		// 3. Authenticate via connectHandler
+		const { response, userContext } = await connectHandler(connectionId, { token })
+		if (debug) console.log(`[socket-adapter] connectHandler result:`, response)
+		if (debug && userContext) {
+			console.log(`[socket-adapter] authenticated userId=${userContext.userId} orgId=${userContext.organizationId}`)
+		}
+
+		// 4. Reject if connectHandler returns non-200
+		if (response.statusCode !== 200) {
+			if (debug) console.log(`[socket-adapter] rejected connectionId=${connectionId} status=${response.statusCode}`)
+			ws.close(1008, 'Authentication failed')
+			connectionRegistry.unregister(connectionId)
+			wsToConnectionId.delete(ws)
+			return
+		}
+
+		// 5. Store user context in the registry (if provided)
+		if (userContext) {
+			connectionRegistry.setUserContext(connectionId, userContext)
+		}
+
+		// 6. Send welcome message to client
+		const welcomeMessage: Record<string, unknown> = { type: 'connected', connectionId }
+		if (userContext) welcomeMessage.userId = userContext.userId
+		ws.send(JSON.stringify(welcomeMessage))
+		if (debug) console.log(`[socket-adapter] welcome sent to ${connectionId}${userContext ? ` userId=${userContext.userId}` : ''}`)
+
+		// 7. Route incoming messages through ActionRouter
+		ws.on('message', async (raw) => {
+			const rawBody = raw.toString()
+
+			if (debug) console.log(`[socket-adapter] message from ${connectionId}:`, rawBody)
+
+			const requestContext = buildRequestContext(connectionId, '$default', 'MESSAGE', {
+				domainName,
+				stage,
+			})
+
+			const resolvedUserContext = userContext ?? connectionRegistry.getUserContext(connectionId) ?? {
+				userId: 'anonymous',
+				organizationId: 'anonymous',
+			}
+
+			const result = await actionRouter.route(connectionId, rawBody, postToConnection, requestContext, resolvedUserContext)
+
+			// Send the handler result back to the sender
+			if (result.body) {
+				ws.send(result.body)
+			}
+		})
+
+		// 8. Handle disconnect
+		ws.on('close', async (code, reason) => {
+			if (debug) {
+				console.log(
+					`[socket-adapter] $disconnect connectionId=${connectionId} code=${code} reason=${reason.toString()}`
+				)
+			}
+
+			await disconnectHandler(connectionId)
+			connectionRegistry.unregister(connectionId)
+			wsToConnectionId.delete(ws)
+		})
+	})
+
+	return {
+		postToConnection,
+		connectionRegistry,
+		actionRouter,
+		getConnectionId: (ws: WebSocket) => wsToConnectionId.get(ws),
+	}
+}

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

@@ -0,0 +1,31 @@
+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)
+	}
+}

package/src/websocket/types.ts

@@ -0,0 +1,66 @@
+import type { WebSocket } from 'ws'
+
+/** User context attached to a WebSocket connection */
+export interface WebSocketUserContext {
+	userId: string
+	organizationId: string
+	[key: string]: unknown
+}
+
+/** Unique identifier for a WebSocket connection (mirrors API Gateway connectionId) */
+export type ConnectionId = string
+
+/** Function signature for sending data to a connection (replaces PostToConnectionCommand) */
+export type PostToConnectionFn = (connectionId: ConnectionId, data: unknown) => Promise<void>
+
+/** Synthesized equivalent of event.requestContext for local dev */
+export interface LocalRequestContext {
+	connectionId: ConnectionId
+	domainName: string
+	stage: string
+	routeKey: string
+	eventType: 'CONNECT' | 'DISCONNECT' | 'MESSAGE'
+}
+
+/** Context passed to action handlers */
+export interface ActionHandlerContext {
+	connectionId: ConnectionId
+	requestContext: LocalRequestContext
+	postToConnection: PostToConnectionFn
+	userContext: WebSocketUserContext
+}
+
+/** Result returned from an action handler */
+export interface ActionHandlerResult {
+	statusCode: number
+	body?: string
+}
+
+/** Handler function for a routed action */
+export type ActionHandler = (
+	body: string,
+	ctx: ActionHandlerContext
+) => Promise<ActionHandlerResult>
+
+/** Handler called on $connect — returns response + optional user context */
+export type ConnectHandlerFn = (
+	connectionId: ConnectionId,
+	params: { token?: string },
+) => Promise<{
+	response: { statusCode: number; body?: string }
+	userContext?: WebSocketUserContext
+}>
+
+/** Handler called on $disconnect */
+export type DisconnectHandlerFn = (connectionId: ConnectionId) => Promise<void>
+
+/** Configuration for the socket adapter */
+export interface SocketAdapterConfig {
+	domainName?: string
+	stage?: string
+	debug?: boolean
+	connectHandler?: ConnectHandlerFn
+	disconnectHandler?: DisconnectHandlerFn
+	routes?: Record<string, ActionHandler>
+	defaultHandler?: ActionHandler
+}

package/tsconfig.json

@@ -0,0 +1,19 @@
+{
+  "compilerOptions": {
+    "rootDir": "./src",
+    "outDir": "./dist",
+    "target": "ES2022",
+    "module": "CommonJS",
+    "moduleResolution": "Node",
+    "esModuleInterop": true,
+    "allowSyntheticDefaultImports": true,
+    "declaration": true,
+    "declarationMap": true,
+    "composite": true,
+    "strict": true,
+    "skipLibCheck": true,
+    "resolveJsonModule": true
+  },
+  "include": ["src/**/*"],
+  "exclude": ["node_modules", "dist"]
+}