@firtoz/websocket-do 6.0.2 → 7.0.1

package/README.md

@@ -6,6 +6,8 @@
 
 Type-safe WebSocket session management for Cloudflare Durable Objects with Hono integration.
 
+> **⚠️ Early WIP Notice:** This package is in very early development and is **not production-ready**. It is TypeScript-only and may have breaking changes. While I (the maintainer) have limited time, I'm open to PRs for features, bug fixes, or additional support (like JS builds). Please feel free to try it out and contribute! See [CONTRIBUTING.md](../../CONTRIBUTING.md) for details.
+
 ## Features
 
 - 🔒 **Type-safe** - Full TypeScript support with generic types for messages and session data
@@ -67,44 +69,81 @@ interface SessionData {
 ### 2. Implement Your Session
 
 ```typescript
-import { BaseSession } from '@firtoz/websocket-do';
+import { BaseSession, type BaseSessionHandlers } from '@firtoz/websocket-do';
 import type { Context } from 'hono';
 
-class ChatSession extends BaseSession<
-  Env,
+// Define handlers for your session
+const chatSessionHandlers: BaseSessionHandlers<
   SessionData,
   ServerMessage,
-  ClientMessage
-> {
-  protected createData(ctx: Context<{ Bindings: Env }>): SessionData {
-    return {
-      userId: crypto.randomUUID(),
-      joinedAt: Date.now(),
-    };
-  }
-
-  async handleMessage(message: ClientMessage): Promise<void> {
+  ClientMessage,
+  Env
+> = {
+  createData: (ctx: Context<{ Bindings: Env }>) => ({
+    userId: crypto.randomUUID(),
+    joinedAt: Date.now(),
+  }),
+
+  handleMessage: async (message: ClientMessage) => {
+    // 'this' context will be the session instance
     switch (message.type) {
       case 'chat':
-        // Broadcast to all sessions
-        this.broadcast({
-          type: 'chat',
-          message: message.message,
-          from: this.data.userId,
-        });
+        // Access session via closure or bind
+        // Note: handlers receive session context when called
         break;
       case 'ping':
-        this.send({ type: 'welcome', userId: this.data.userId });
+        // Send messages
         break;
     }
-  }
+  },
 
-  async handleBufferMessage(message: ArrayBuffer): Promise<void> {
+  handleBufferMessage: async (message: ArrayBuffer) => {
     // Handle binary messages if needed
-  }
+  },
 
-  async handleClose(): Promise<void> {
-    console.log(`Session closed for user ${this.data.userId}`);
+  handleClose: async () => {
+    console.log('Session closed');
+  },
+};
+
+// Create session class (can be extended if needed)
+class ChatSession extends BaseSession<
+  SessionData,
+  ServerMessage,
+  ClientMessage,
+  Env
+> {
+  constructor(
+    websocket: WebSocket,
+    sessions: Map<WebSocket, ChatSession>
+  ) {
+    super(websocket, sessions, {
+      createData: (ctx) => ({
+        userId: crypto.randomUUID(),
+        joinedAt: Date.now(),
+      }),
+      handleMessage: async (message) => {
+        switch (message.type) {
+          case 'chat':
+            // Broadcast to all sessions
+            this.broadcast({
+              type: 'chat',
+              message: message.message,
+              from: this.data.userId,
+            });
+            break;
+          case 'ping':
+            this.send({ type: 'welcome', userId: this.data.userId });
+            break;
+        }
+      },
+      handleBufferMessage: async (message) => {
+        // Handle binary messages if needed
+      },
+      handleClose: async () => {
+        console.log(`Session closed for user ${this.data.userId}`);
+      },
+    });
   }
 }
 ```
@@ -115,7 +154,7 @@ class ChatSession extends BaseSession<
 import { BaseWebSocketDO } from '@firtoz/websocket-do';
 import { Hono } from 'hono';
 
-export class ChatRoomDO extends BaseWebSocketDO<Env, ChatSession> {
+export class ChatRoomDO extends BaseWebSocketDO<ChatSession, Env> {
   app = this.getBaseApp()
     .get('/info', (ctx) => {
       return ctx.json({
@@ -123,8 +162,12 @@ export class ChatRoomDO extends BaseWebSocketDO<Env, ChatSession> {
       });
     });
 
-  protected createSession(websocket: WebSocket): ChatSession {
-    return new ChatSession(websocket, this.sessions);
+  constructor(ctx: DurableObjectState, env: Env) {
+    super(ctx, env, {
+      createSession: (ctx, websocket) => {
+        return new ChatSession(websocket, this.sessions);
+      },
+    });
   }
 }
 ```
@@ -342,39 +385,42 @@ interface SessionData {
 
 // Implement validated session
 class ChatSession extends ZodSession<
-  Env,
   SessionData,
   ServerMessage,
-  ClientMessage
+  ClientMessage,
+  Env
 > {
-  protected clientSchema = ClientMessageSchema;
-  protected serverSchema = ServerMessageSchema;
-
-  protected createData(ctx: Context<{ Bindings: Env }>): SessionData {
-    return { name: 'Anonymous' };
-  }
-
-  async handleMessage(message: ClientMessage): Promise<void> {
-    // Message is already validated!
-    switch (message.type) {
-      case 'setName':
-        this.data.name = message.name;
-        this.update();
-        this.send({ type: 'nameChanged', newName: message.name });
-        break;
+  constructor(
+    websocket: WebSocket,
+    sessions: Map<WebSocket, ChatSession>,
+    options: ZodSessionOptions<ClientMessage, ServerMessage>
+  ) {
+    super(websocket, sessions, options, {
+      createData: (ctx) => ({ name: 'Anonymous' }),
       
-      case 'message':
-        this.broadcast({
-          type: 'message',
-          text: message.text,
-          from: this.data.name,
-        });
-        break;
-    }
-  }
-
-  async handleClose(): Promise<void> {
-    console.log(`${this.data.name} disconnected`);
+      handleValidatedMessage: async (message) => {
+        // Message is already validated!
+        switch (message.type) {
+          case 'setName':
+            this.data.name = message.name;
+            this.update();
+            this.send({ type: 'nameChanged', newName: message.name });
+            break;
+          
+          case 'message':
+            this.broadcast({
+              type: 'message',
+              text: message.text,
+              from: this.data.name,
+            });
+            break;
+        }
+      },
+
+      handleClose: async () => {
+        console.log(`${this.data.name} disconnected`);
+      },
+    });
   }
 }
 ```
@@ -383,34 +429,61 @@ class ChatSession extends ZodSession<
 
 ```typescript
 class ChatSession extends ZodSession<...> {
-  // Enable buffer mode for msgpack
-  protected enableBufferMessages = true;
-  
-  protected clientSchema = ClientMessageSchema;
-  protected serverSchema = ServerMessageSchema;
-  
-  // Messages automatically decoded from msgpack
-  async handleMessage(message: ClientMessage): Promise<void> {
-    // Handle validated message
+  constructor(
+    websocket: WebSocket,
+    sessions: Map<WebSocket, ChatSession>
+  ) {
+    super(websocket, sessions, {
+      clientSchema: ClientMessageSchema,
+      serverSchema: ServerMessageSchema,
+      enableBufferMessages: true, // Enable buffer mode for msgpack
+    }, {
+      createData: (ctx) => ({ name: 'Anonymous' }),
+      handleValidatedMessage: async (message) => {
+        // Messages automatically decoded from msgpack
+        // Handle validated message
+      },
+      handleClose: async () => {
+        console.log('Session closed');
+      },
+    });
   }
 }
 ```
 
 ## API Reference
 
-### `BaseWebSocketDO<TEnv, TSession>`
+### `BaseWebSocketDO<TSession, TEnv>`
 
-Abstract class for creating WebSocket-enabled Durable Objects.
+Base class for creating WebSocket-enabled Durable Objects. Uses composition instead of inheritance.
 
 #### Type Parameters
 
-- `TEnv` - Your Cloudflare Worker environment bindings
 - `TSession` - Your session class extending `BaseSession`
+- `TEnv` - Your Cloudflare Worker environment bindings
 
-#### Methods
+#### Constructor
+
+```typescript
+constructor(
+  ctx: DurableObjectState,
+  env: TEnv,
+  options: BaseWebSocketDOOptions<TSession, TEnv>
+)
+```
+
+#### Options Type
+
+```typescript
+type BaseWebSocketDOOptions<TSession, TEnv> = {
+  createSession: (
+    ctx: Context<{ Bindings: TEnv }> | undefined,
+    websocket: WebSocket
+  ) => TSession | Promise<TSession>;
+};
+```
 
-- `abstract createSession(websocket: WebSocket): TSession | Promise<TSession>`
-  - Factory method to create session instances
+#### Methods
 
 - `getBaseApp(): Hono`
   - Returns a base Hono app with `/websocket` endpoint configured
@@ -423,35 +496,44 @@ Abstract class for creating WebSocket-enabled Durable Objects.
 - `sessions: Map<WebSocket, TSession>` - Map of all active sessions
 - `app: Hono` - Your Hono application (must be implemented)
 
-### `BaseSession<TEnv, TData, TServerMessage, TClientMessage>`
+### `BaseSession<TData, TServerMessage, TClientMessage, TEnv>`
 
-Abstract class for managing individual WebSocket sessions.
+Concrete class for managing individual WebSocket sessions. Uses composition pattern with handlers.
 
 #### Type Parameters
 
-- `TEnv` - Your Cloudflare Worker environment bindings
 - `TData` - Type of data stored in the session
 - `TServerMessage` - Union type of messages sent to clients
 - `TClientMessage` - Union type of messages received from clients
+- `TEnv` - Your Cloudflare Worker environment bindings (default: `Cloudflare.Env`)
 
-#### Methods
+#### Constructor
 
-- `abstract createData(ctx: Context): TData`
-  - Creates initial session data
+```typescript
+constructor(
+  websocket: WebSocket,
+  sessions: Map<WebSocket, BaseSession<TData, TServerMessage, TClientMessage, TEnv>>,
+  handlers: BaseSessionHandlers<TData, TServerMessage, TClientMessage, TEnv>
+)
+```
 
-- `abstract handleMessage(message: TClientMessage): Promise<void>`
-  - Handles text messages from client
+#### Handlers Type
 
-- `abstract handleBufferMessage(message: ArrayBuffer): Promise<void>`
-  - Handles binary messages from client
+```typescript
+type BaseSessionHandlers<TData, TServerMessage, TClientMessage, TEnv> = {
+  createData: (ctx: Context<{ Bindings: TEnv }>) => TData;
+  handleMessage: (message: TClientMessage) => Promise<void>;
+  handleBufferMessage: (message: ArrayBuffer) => Promise<void>;
+  handleClose: () => Promise<void>;
+};
+```
 
-- `abstract handleClose(): Promise<void>`
-  - Cleanup when session closes
+#### Methods
 
-- `protected send(message: TServerMessage): void`
+- `send(message: TServerMessage): void`
   - Send message to this session's client
 
-- `protected broadcast(message: TServerMessage, excludeSelf?: boolean): void`
+- `broadcast(message: TServerMessage, excludeSelf?: boolean): void`
   - Send message to all connected sessions
 
 - `startFresh(ctx: Context): void`
@@ -490,7 +572,7 @@ Low-level wrapper for typed WebSocket operations.
 You can extend the base app with custom routes:
 
 ```typescript
-export class ChatRoomDO extends BaseWebSocketDO<Env, ChatSession> {
+export class ChatRoomDO extends BaseWebSocketDO<ChatSession, Env> {
   app = this.getBaseApp()
     .get('/stats', (ctx) => {
       const users = Array.from(this.sessions.values()).map(s => ({
@@ -509,6 +591,14 @@ export class ChatRoomDO extends BaseWebSocketDO<Env, ChatSession> {
       
       return ctx.json({ success: true });
     });
+
+  constructor(ctx: DurableObjectState, env: Env) {
+    super(ctx, env, {
+      createSession: (ctx, websocket) => {
+        return new ChatSession(websocket, this.sessions);
+      },
+    });
+  }
 }
 ```
 
@@ -517,25 +607,38 @@ export class ChatRoomDO extends BaseWebSocketDO<Env, ChatSession> {
 Session data is automatically serialized and persists across hibernation:
 
 ```typescript
-class GameSession extends BaseSession<Env, GameData, ServerMsg, ClientMsg> {
-  protected createData(ctx: Context): GameData {
-    return {
-      playerName: ctx.req.query('name') || 'Anonymous',
-      score: 0,
-      inventory: [],
-    };
-  }
-
-  async handleMessage(message: ClientMsg): Promise<void> {
-    if (message.type === 'collectItem') {
-      this.data.inventory.push(message.item);
-      this.data.score += 10;
-      
-      // Persist changes
-      this.update();
-      
-      this.send({ type: 'scoreUpdate', score: this.data.score });
-    }
+class GameSession extends BaseSession<GameData, ServerMsg, ClientMsg, Env> {
+  constructor(
+    websocket: WebSocket,
+    sessions: Map<WebSocket, GameSession>
+  ) {
+    super(websocket, sessions, {
+      createData: (ctx) => ({
+        playerName: ctx.req.query('name') || 'Anonymous',
+        score: 0,
+        inventory: [],
+      }),
+
+      handleMessage: async (message) => {
+        if (message.type === 'collectItem') {
+          this.data.inventory.push(message.item);
+          this.data.score += 10;
+          
+          // Persist changes
+          this.update();
+          
+          this.send({ type: 'scoreUpdate', score: this.data.score });
+        }
+      },
+
+      handleBufferMessage: async (message) => {
+        // Handle buffer messages if needed
+      },
+
+      handleClose: async () => {
+        console.log('Game session closed');
+      },
+    });
   }
 }
 ```
@@ -545,21 +648,40 @@ class GameSession extends BaseSession<Env, GameData, ServerMsg, ClientMsg> {
 Errors in message handlers are caught and logged, but don't crash the connection:
 
 ```typescript
-async handleMessage(message: ClientMessage): Promise<void> {
-  try {
-    // Your logic here
-    if (message.type === 'dangerous') {
-      throw new Error('Invalid operation');
-    }
-  } catch (error) {
-    // Send error to client
-    this.send({ 
-      type: 'error', 
-      message: error instanceof Error ? error.message : 'Unknown error' 
+class MySession extends BaseSession<...> {
+  constructor(
+    websocket: WebSocket,
+    sessions: Map<WebSocket, MySession>
+  ) {
+    super(websocket, sessions, {
+      createData: (ctx) => ({ /* ... */ }),
+
+      handleMessage: async (message) => {
+        try {
+          // Your logic here
+          if (message.type === 'dangerous') {
+            throw new Error('Invalid operation');
+          }
+        } catch (error) {
+          // Send error to client
+          this.send({ 
+            type: 'error', 
+            message: error instanceof Error ? error.message : 'Unknown error' 
+          });
+          
+          // Optionally close the connection
+          this.websocket.close(1008, 'Policy violation');
+        }
+      },
+
+      handleBufferMessage: async (message) => {
+        // Handle buffer messages
+      },
+
+      handleClose: async () => {
+        console.log('Session closed');
+      },
     });
-    
-    // Optionally close the connection
-    this.websocket.close(1008, 'Policy violation');
   }
 }
 ```
@@ -569,18 +691,25 @@ async handleMessage(message: ClientMessage): Promise<void> {
 This package exports the following:
 
 ### Classes
-- `BaseWebSocketDO` - Abstract base class for WebSocket Durable Objects
-- `BaseSession` - Abstract base class for WebSocket sessions
+- `BaseWebSocketDO` - Base class for WebSocket Durable Objects (composition-based)
+- `BaseSession` - Concrete session class with handler injection
 - `ZodWebSocketClient` - Type-safe WebSocket client with Zod validation
-- `ZodSession` - Session base class with Zod validation built-in
+- `ZodSession` - Concrete session class with Zod validation and handler injection
+- `ZodWebSocketDO` - Base class for WebSocket DOs with Zod validation
 - `WebsocketWrapper` - Low-level WebSocket wrapper with typed attachments
 
+### Types
+- `BaseSessionHandlers` - Handler interface for `BaseSession`
+- `BaseWebSocketDOOptions` - Options interface for `BaseWebSocketDO`
+- `ZodSessionHandlers` - Handler interface for `ZodSession`
+- `ZodSessionOptions` - Options interface for `ZodSession`
+- `ZodSessionOptionsOrFactory` - Options or factory function for `ZodSession`
+- `ZodWebSocketDOOptions` - Options interface for `ZodWebSocketDO`
+- `ZodWebSocketClientOptions` - Options interface for `ZodWebSocketClient`
+
 ### Utilities
 - `zodMsgpack` - Msgpack encode/decode with Zod validation
 
-### Types
-All classes export their type parameters and interfaces for custom implementations.
-
 ## Complete Example
 
 Here's a full example combining all features:
@@ -605,7 +734,7 @@ export type ClientMessage = z.infer<typeof ClientMessageSchema>;
 export type ServerMessage = z.infer<typeof ServerMessageSchema>;
 
 // do.ts - Server-side (Durable Object)
-import { BaseWebSocketDO, ZodSession } from '@firtoz/websocket-do';
+import { BaseWebSocketDO, ZodSession, type ZodSessionOptions } from '@firtoz/websocket-do';
 import { ClientMessageSchema, ServerMessageSchema } from './schemas';
 
 interface SessionData {
@@ -613,45 +742,47 @@ interface SessionData {
   joinedAt: number;
 }
 
-class ChatSession extends ZodSession<Env, SessionData, ServerMessage, ClientMessage> {
-  protected clientSchema = ClientMessageSchema;
-  protected serverSchema = ServerMessageSchema;
-  protected enableBufferMessages = true; // Use msgpack for efficiency
-
-  protected createData(): SessionData {
-    return {
-      name: 'Anonymous',
-      joinedAt: Date.now(),
-    };
-  }
-
-  async handleMessage(message: ClientMessage): Promise<void> {
-    switch (message.type) {
-      case 'setName':
-        const oldName = this.data.name;
-        this.data.name = message.name;
-        this.update();
-        
-        this.send({ type: 'nameChanged', newName: message.name });
-        this.broadcast({ type: 'userJoined', name: message.name }, true);
-        break;
-      
-      case 'message':
-        this.broadcast({
-          type: 'message',
-          text: message.text,
-          from: this.data.name,
-        });
-        break;
-    }
-  }
-
-  async handleClose(): Promise<void> {
-    console.log(`${this.data.name} disconnected`);
+class ChatSession extends ZodSession<SessionData, ServerMessage, ClientMessage, Env> {
+  constructor(
+    websocket: WebSocket,
+    sessions: Map<WebSocket, ChatSession>,
+    options: ZodSessionOptions<ClientMessage, ServerMessage>
+  ) {
+    super(websocket, sessions, options, {
+      createData: () => ({
+        name: 'Anonymous',
+        joinedAt: Date.now(),
+      }),
+
+      handleValidatedMessage: async (message) => {
+        switch (message.type) {
+          case 'setName':
+            const oldName = this.data.name;
+            this.data.name = message.name;
+            this.update();
+            
+            this.send({ type: 'nameChanged', newName: message.name });
+            this.broadcast({ type: 'userJoined', name: message.name }, true);
+            break;
+          
+          case 'message':
+            this.broadcast({
+              type: 'message',
+              text: message.text,
+              from: this.data.name,
+            });
+            break;
+        }
+      },
+
+      handleClose: async () => {
+        console.log(`${this.data.name} disconnected`);
+      },
+    });
   }
 }
 
-export class ChatRoomDO extends BaseWebSocketDO<Env, ChatSession> {
+export class ChatRoomDO extends BaseWebSocketDO<ChatSession, Env> {
   app = this.getBaseApp()
     .get('/info', (ctx) => {
       const users = Array.from(this.sessions.values()).map(s => ({
@@ -661,8 +792,16 @@ export class ChatRoomDO extends BaseWebSocketDO<Env, ChatSession> {
       return ctx.json({ users, count: users.length });
     });
 
-  protected createSession(websocket: WebSocket): ChatSession {
-    return new ChatSession(websocket, this.sessions);
+  constructor(ctx: DurableObjectState, env: Env) {
+    super(ctx, env, {
+      createSession: (ctx, websocket) => {
+        return new ChatSession(websocket, this.sessions, {
+          clientSchema: ClientMessageSchema,
+          serverSchema: ServerMessageSchema,
+          enableBufferMessages: true, // Use msgpack for efficiency
+        });
+      },
+    });
   }
 }
 

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@firtoz/websocket-do",
-	"version": "6.0.2",
+	"version": "7.0.1",
 	"description": "Type-safe WebSocket session management for Cloudflare Durable Objects with Hono integration",
 	"main": "./src/index.ts",
 	"module": "./src/index.ts",
@@ -45,14 +45,14 @@
 		"url": "https://github.com/firtoz/fullstack-toolkit/issues"
 	},
 	"peerDependencies": {
-		"@cloudflare/workers-types": "^4.20251008.0",
-		"@firtoz/hono-fetcher": "^2.3.1",
-		"hono": "^4.10.1"
+		"@cloudflare/workers-types": "^4.20251228.0",
+		"@firtoz/hono-fetcher": "^2.3.2",
+		"hono": "^4.11.4"
 	},
 	"optionalDependencies": {
-		"msgpackr": "^1.11.5",
-		"react": "^19.2.0",
-		"zod": "^4.1.12"
+		"msgpackr": "^1.11.8",
+		"react": "^19.2.3",
+		"zod": "^4.3.5"
 	},
 	"engines": {
 		"node": ">=18.0.0"
@@ -61,8 +61,8 @@
 		"access": "public"
 	},
 	"devDependencies": {
-		"@types/react": "^19.2.2",
-		"bun-types": "^1.3.0",
+		"@types/react": "^19.2.8",
+		"bun-types": "^1.3.6",
 		"typescript": "^5.9.3"
 	}
 }

package/src/BaseSession.ts

@@ -1,20 +1,73 @@
 import type { Context } from "hono";
 import { WebsocketWrapper } from "./WebsocketWrapper";
 
-export type SessionClientMessage<TSession extends BaseSession> =
-	TSession extends BaseSession<never, never, infer TClientMessage, never>
+// biome-ignore lint/suspicious/noExplicitAny: We are using any on purpose to allow any type of session.
+export type SessionData<TSession extends BaseSession<any, any, any, any>> =
+	TSession extends BaseSession<
+		infer TData,
+		infer _TServerMessage,
+		infer _TClientMessage,
+		infer _TEnv
+	>
+		? TData
+		: never;
+
+export type SessionClientMessage<
+	// biome-ignore lint/suspicious/noExplicitAny: We are using any on purpose to allow any type of session.
+	TSession extends BaseSession<any, any, any, any>,
+> =
+	TSession extends BaseSession<
+		infer _TData,
+		infer _TServerMessage,
+		infer TClientMessage,
+		infer _TEnv
+	>
 		? TClientMessage
 		: never;
 
-export abstract class BaseSession<
-	// biome-ignore lint/suspicious/noExplicitAny: Generic type parameter with flexible default
-	TEnv extends object = any,
-	// biome-ignore lint/suspicious/noExplicitAny: Generic type parameter with flexible default
-	TData = any,
-	// biome-ignore lint/suspicious/noExplicitAny: Generic type parameter with flexible default
-	TServerMessage = any,
-	// biome-ignore lint/suspicious/noExplicitAny: Generic type parameter with flexible default
-	TClientMessage = any,
+export type SessionServerMessage<
+	// biome-ignore lint/suspicious/noExplicitAny: We are using any on purpose to allow any type of session.
+	TSession extends BaseSession<any, any, any, any>,
+> =
+	TSession extends BaseSession<
+		infer _TData,
+		infer TServerMessage,
+		infer _TClientMessage,
+		infer _TEnv
+	>
+		? TServerMessage
+		: never;
+
+export type SessionEnv<
+	// biome-ignore lint/suspicious/noExplicitAny: We are using any on purpose to allow any type of session.
+	TSession extends BaseSession<any, any, any, any>,
+> =
+	TSession extends BaseSession<
+		infer _TData,
+		infer _TServerMessage,
+		infer _TClientMessage,
+		infer TEnv extends Cloudflare.Env
+	>
+		? TEnv
+		: never;
+
+export type BaseSessionHandlers<
+	TData,
+	_TServerMessage,
+	TClientMessage,
+	TEnv extends object = Cloudflare.Env,
+> = {
+	createData: (ctx: Context<{ Bindings: TEnv }>) => TData;
+	handleMessage: (message: TClientMessage) => Promise<void>;
+	handleBufferMessage: (message: ArrayBuffer) => Promise<void>;
+	handleClose: () => Promise<void>;
+};
+
+export class BaseSession<
+	TData,
+	TServerMessage,
+	TClientMessage,
+	TEnv extends object = Cloudflare.Env,
 > {
 	private _data!: TData;
 
@@ -27,19 +80,27 @@ export abstract class BaseSession<
 	}
 
 	private readonly wrapper: WebsocketWrapper<TData, TServerMessage>;
+	protected readonly handlers: BaseSessionHandlers<
+		TData,
+		TServerMessage,
+		TClientMessage,
+		TEnv
+	>;
 
 	constructor(
 		public websocket: WebSocket,
 		protected sessions: Map<
 			WebSocket,
-			BaseSession<TEnv, TData, TServerMessage, TClientMessage>
+			BaseSession<TData, TServerMessage, TClientMessage, TEnv>
 		>,
+		handlers: BaseSessionHandlers<TData, TServerMessage, TClientMessage, TEnv>,
 	) {
 		this.wrapper = new WebsocketWrapper<TData, TServerMessage>(websocket);
+		this.handlers = handlers;
 	}
 
 	public startFresh(ctx: Context<{ Bindings: TEnv }>) {
-		this.data = this.createData(ctx);
+		this.data = this.handlers.createData(ctx);
 		this.wrapper.serializeAttachment(this.data);
 	}
 
@@ -56,20 +117,26 @@ export abstract class BaseSession<
 		this.wrapper.serializeAttachment(this.data);
 	}
 
-	protected abstract createData(ctx: Context<{ Bindings: TEnv }>): TData;
-
-	protected broadcast(message: TServerMessage, excludeSelf = false) {
+	public broadcast(message: TServerMessage, excludeSelf = false) {
 		for (const session of this.sessions.values()) {
 			if (excludeSelf && session === this) continue;
 			session.send(message);
 		}
 	}
 
-	protected send(message: TServerMessage) {
+	public send(message: TServerMessage) {
 		this.wrapper.send(message);
 	}
 
-	abstract handleMessage(message: TClientMessage): Promise<void>;
-	abstract handleBufferMessage(message: ArrayBuffer): Promise<void>;
-	abstract handleClose(): Promise<void>;
+	async handleMessage(message: TClientMessage): Promise<void> {
+		return this.handlers.handleMessage(message);
+	}
+
+	async handleBufferMessage(message: ArrayBuffer): Promise<void> {
+		return this.handlers.handleBufferMessage(message);
+	}
+
+	async handleClose(): Promise<void> {
+		return this.handlers.handleClose();
+	}
 }

package/src/BaseWebSocketDO.ts

@@ -1,18 +1,48 @@
 import { DurableObject } from "cloudflare:workers";
 import type { DOWithHonoApp } from "@firtoz/hono-fetcher/honoDoFetcher";
 import { type Context, Hono } from "hono";
-import type { BaseSession, SessionClientMessage } from "./BaseSession";
+import type {
+	BaseSession,
+	SessionClientMessage,
+	SessionEnv,
+} from "./BaseSession";
+
+export type BaseWebSocketDOOptions<
+	// biome-ignore lint/suspicious/noExplicitAny: We are using any on purpose to allow any type of session.
+	TSession extends BaseSession<any, any, any, any>,
+	TEnv extends SessionEnv<TSession>,
+> = {
+	createSession: (
+		ctx: Context<{ Bindings: TEnv }> | undefined,
+		websocket: WebSocket,
+	) => TSession | Promise<TSession>;
+};
 
 export abstract class BaseWebSocketDO<
-		TEnv extends object,
-		TSession extends BaseSession<TEnv>,
+		// biome-ignore lint/suspicious/noExplicitAny: We are using any on purpose to allow any type of session.
+		TSession extends BaseSession<any, any, any, any> = BaseSession<
+			// biome-ignore lint/suspicious/noExplicitAny: We are using any on purpose to allow any type of session.
+			any,
+			// biome-ignore lint/suspicious/noExplicitAny: We are using any on purpose to allow any type of session.
+			any,
+			// biome-ignore lint/suspicious/noExplicitAny: We are using any on purpose to allow any type of session.
+			any,
+			// biome-ignore lint/suspicious/noExplicitAny: We are using any on purpose to allow any type of session.
+			any
+		>,
+		TEnv extends SessionEnv<TSession> = SessionEnv<TSession>,
 	>
 	extends DurableObject<TEnv>
 	implements DOWithHonoApp
 {
 	protected readonly sessions = new Map<WebSocket, TSession>();
+	abstract readonly app: Hono<{ Bindings: TEnv }>;
 
-	constructor(ctx: DurableObjectState, env: TEnv) {
+	constructor(
+		ctx: DurableObjectState,
+		env: TEnv,
+		private readonly options: BaseWebSocketDOOptions<TSession, TEnv>,
+	) {
 		super(ctx, env);
 
 		this.ctx.blockConcurrencyWhile(async () => {
@@ -20,7 +50,9 @@ export abstract class BaseWebSocketDO<
 			await Promise.all(
 				websockets.map(async (websocket) => {
 					try {
-						const session = await this.createSession(websocket);
+						// For resumed sessions, we don't have a Hono context
+						// Pass undefined and let implementers handle it
+						const session = await options.createSession(undefined, websocket);
 						session.resume();
 						this.sessions.set(websocket, session);
 					} catch (error) {
@@ -65,19 +97,13 @@ export abstract class BaseWebSocketDO<
 		);
 	}
 
-	abstract app: Hono<{ Bindings: TEnv }>;
-
-	protected abstract createSession(
-		websocket: WebSocket,
-	): TSession | Promise<TSession>;
-
 	async handleSession(
 		ctx: Context<{ Bindings: TEnv }>,
 		ws: WebSocket,
 	): Promise<void> {
 		this.ctx.acceptWebSocket(ws);
 		try {
-			const session = await this.createSession(ws);
+			const session = await this.options.createSession(ctx, ws);
 			session.startFresh(ctx);
 			this.sessions.set(ws, session);
 		} catch (error) {

package/src/ZodSession.ts

@@ -1,44 +1,77 @@
+import type { Context } from "hono";
 import type { ZodType } from "zod";
 import { BaseSession } from "./BaseSession";
 import { zodMsgpack } from "./zodMsgpack";
 
-export interface ZodSessionOptions<TClientMessage, TServerMessage> {
+export type ZodSessionOptions<TClientMessage, TServerMessage> = {
 	clientSchema: ZodType<TClientMessage>;
 	serverSchema: ZodType<TServerMessage>;
 	enableBufferMessages?: boolean;
-}
-
-export abstract class ZodSession<
-	// biome-ignore lint/suspicious/noExplicitAny: We need to allow any for the environment
-	TEnv extends object = any,
-	// biome-ignore lint/suspicious/noExplicitAny: We need to allow any for the data
-	TData = any,
-	// biome-ignore lint/suspicious/noExplicitAny: We need to allow any for the server message
-	TServerMessage = any,
-	// biome-ignore lint/suspicious/noExplicitAny: We need to allow any for the client message
-	TClientMessage = any,
-> extends BaseSession<TEnv, TData, TServerMessage, TClientMessage> {
-	protected readonly clientCodec: ReturnType<typeof zodMsgpack<TClientMessage>>;
-	protected readonly serverCodec: ReturnType<typeof zodMsgpack<TServerMessage>>;
+	sendProtocolError?: (
+		websocket: WebSocket,
+		errorMessage: string,
+	) => Promise<void>;
+};
+
+export type ZodSessionHandlers<
+	TData,
+	_TServerMessage,
+	TClientMessage,
+	TEnv extends object,
+> = {
+	createData: (ctx: Context<{ Bindings: TEnv }>) => TData;
+	handleValidatedMessage: (message: TClientMessage) => Promise<void>;
+	handleValidationError?: (
+		error: unknown,
+		originalMessage: unknown,
+	) => Promise<void>;
+	handleClose: () => Promise<void>;
+};
+
+export class ZodSession<
+	TData,
+	TServerMessage,
+	TClientMessage,
+	TEnv extends object = Cloudflare.Env,
+> extends BaseSession<TData, TServerMessage, TClientMessage, TEnv> {
+	private readonly clientCodec: ReturnType<typeof zodMsgpack<TClientMessage>>;
+	private readonly serverCodec: ReturnType<typeof zodMsgpack<TServerMessage>>;
 	protected readonly enableBufferMessages: boolean;
 
 	constructor(
 		websocket: WebSocket,
 		sessions: Map<
 			WebSocket,
-			ZodSession<TEnv, TData, TServerMessage, TClientMessage>
+			ZodSession<TData, TServerMessage, TClientMessage, TEnv>
+		>,
+		private readonly options: ZodSessionOptions<TClientMessage, TServerMessage>,
+		private readonly zodHandlers: ZodSessionHandlers<
+			TData,
+			TServerMessage,
+			TClientMessage,
+			TEnv
 		>,
-		protected options: ZodSessionOptions<TClientMessage, TServerMessage>,
 	) {
-		super(websocket, sessions);
+		super(websocket, sessions, {
+			createData: zodHandlers.createData,
+			handleMessage: async (message) => {
+				return this._internalHandleMessage(message);
+			},
+			handleBufferMessage: async (message) => {
+				return this._internalHandleBufferMessage(message);
+			},
+			handleClose: async () => {
+				return zodHandlers.handleClose();
+			},
+		});
 
 		this.clientCodec = zodMsgpack(options.clientSchema);
 		this.serverCodec = zodMsgpack(options.serverSchema);
 		this.enableBufferMessages = options.enableBufferMessages ?? false;
 	}
 
-	// Override the base handleMessage to add validation
-	async handleMessage(message: TClientMessage): Promise<void> {
+	// Internal method used by the base class handlers
+	private async _internalHandleMessage(message: TClientMessage): Promise<void> {
 		// If buffer messages are enabled, reject JSON messages
 		if (this.enableBufferMessages) {
 			console.error(
@@ -53,15 +86,17 @@ export abstract class ZodSession<
 		try {
 			// Validate the message using the client schema
 			const validatedMessage = this.options.clientSchema.parse(message);
-			await this.handleValidatedMessage(validatedMessage);
+			await this.zodHandlers.handleValidatedMessage(validatedMessage);
 		} catch (error) {
 			console.error("Invalid client message received:", error);
-			await this.handleValidationError(error, message);
+			await this._internalHandleValidationError(error, message);
 		}
 	}
 
-	// Override buffer message handling to support msgpack decoding
-	async handleBufferMessage(buffer: ArrayBuffer): Promise<void> {
+	// Internal method used by the base class handlers
+	private async _internalHandleBufferMessage(
+		buffer: ArrayBuffer,
+	): Promise<void> {
 		// If buffer messages are disabled, reject buffer messages
 		if (!this.enableBufferMessages) {
 			console.error(
@@ -75,15 +110,33 @@ export abstract class ZodSession<
 		try {
 			const bytes = new Uint8Array(buffer);
 			const decodedMessage = this.clientCodec.decode(bytes);
-			await this.handleValidatedMessage(decodedMessage);
+			await this.zodHandlers.handleValidatedMessage(decodedMessage);
 		} catch (error) {
 			console.error("Failed to decode buffer message:", error);
-			await this.handleValidationError(error, buffer);
+			await this._internalHandleValidationError(error, buffer);
+		}
+	}
+
+	// Internal validation error handler
+	private async _internalHandleValidationError(
+		error: unknown,
+		originalMessage: unknown,
+	): Promise<void> {
+		if (this.zodHandlers.handleValidationError) {
+			await this.zodHandlers.handleValidationError(error, originalMessage);
+		} else {
+			// Default implementation logs and continues
+			console.error(
+				"Validation error:",
+				error,
+				"Original message:",
+				originalMessage,
+			);
 		}
 	}
 
 	// Type-safe send method that automatically uses the correct format
-	protected send(message: TServerMessage): void {
+	public send(message: TServerMessage): void {
 		if (this.enableBufferMessages) {
 			this.sendBuffer(message);
 		} else {
@@ -118,12 +171,17 @@ export abstract class ZodSession<
 		}
 	}
 
-	// Send a protocol error message (always as JSON for compatibility)
+	// Send a protocol error message (always as JSON for compatibility by default)
 	private async sendProtocolError(errorMessage: string): Promise<void> {
 		try {
-			// Send a simple error object - no schema validation needed
-			if (this.websocket.readyState !== WebSocket.OPEN) return;
-			this.websocket.send(JSON.stringify({ error: errorMessage }));
+			// Use custom handler if provided, otherwise use default
+			if (this.options.sendProtocolError) {
+				await this.options.sendProtocolError(this.websocket, errorMessage);
+			} else {
+				// Default implementation: send a simple error object - no schema validation needed
+				if (this.websocket.readyState !== WebSocket.OPEN) return;
+				this.websocket.send(JSON.stringify({ error: errorMessage }));
+			}
 		} catch (error) {
 			console.error("Failed to send protocol error:", error);
 		}
@@ -131,7 +189,7 @@ export abstract class ZodSession<
 
 	// Type-safe broadcast that validates server messages
 	// Automatically uses the correct format based on session configuration
-	protected broadcast(message: TServerMessage, excludeSelf = false): void {
+	public broadcast(message: TServerMessage, excludeSelf = false): void {
 		for (const session of this.sessions.values()) {
 			if (excludeSelf && session === this) continue;
 			if (session instanceof ZodSession) {
@@ -139,23 +197,4 @@ export abstract class ZodSession<
 			}
 		}
 	}
-
-	// Abstract methods for implementers
-	protected abstract handleValidatedMessage(
-		message: TClientMessage,
-	): Promise<void>;
-
-	// Optional error handling - default implementation logs and continues
-	protected async handleValidationError(
-		error: unknown,
-		originalMessage: unknown,
-	): Promise<void> {
-		console.error(
-			"Validation error:",
-			error,
-			"Original message:",
-			originalMessage,
-		);
-		// Implementers can override this to send error responses to clients
-	}
 }

package/src/ZodWebSocketDO.ts

@@ -1,27 +1,83 @@
+import type { Context } from "hono";
+import type {
+	SessionClientMessage,
+	SessionEnv,
+	SessionServerMessage,
+} from "./BaseSession";
 import { BaseWebSocketDO } from "./BaseWebSocketDO";
 import type { ZodSession, ZodSessionOptions } from "./ZodSession";
 
-export abstract class ZodWebSocketDO<
-	TEnv extends object,
-	// biome-ignore lint/suspicious/noExplicitAny: We need to allow any for the session
-	TSession extends ZodSession<TEnv, any, any, any>,
-	// biome-ignore lint/suspicious/noExplicitAny: We need to allow any for the client message
-	TClientMessage = any,
-	// biome-ignore lint/suspicious/noExplicitAny: We need to allow any for the server message
-	TServerMessage = any,
-> extends BaseWebSocketDO<TEnv, TSession> {
-	protected abstract getZodOptions(): ZodSessionOptions<
+export type ZodSessionOptionsOrFactory<
+	TClientMessage,
+	TServerMessage,
+	TEnv extends Cloudflare.Env = Cloudflare.Env,
+> =
+	| ZodSessionOptions<TClientMessage, TServerMessage>
+	| ((
+			ctx: Context<{ Bindings: TEnv }> | undefined,
+			websocket: WebSocket,
+	  ) => ZodSessionOptions<TClientMessage, TServerMessage>);
+
+export type ZodWebSocketDOOptions<
+	// biome-ignore lint/suspicious/noExplicitAny: We are using any on purpose to allow any type of session.
+	TSession extends ZodSession<any, any, any, any>,
+	TClientMessage,
+	TServerMessage,
+	TEnv extends SessionEnv<TSession>,
+> = {
+	zodSessionOptions: ZodSessionOptionsOrFactory<
 		TClientMessage,
-		TServerMessage
+		TServerMessage,
+		TEnv
 	>;
+	createZodSession: (
+		ctx: Context<{ Bindings: TEnv }> | undefined,
+		websocket: WebSocket,
+		options: ZodSessionOptions<TClientMessage, TServerMessage>,
+	) => TSession | Promise<TSession>;
+};
 
-	protected abstract createZodSession(
+export abstract class ZodWebSocketDO<
+	// biome-ignore lint/suspicious/noExplicitAny: We are using any on purpose to allow any type of session.
+	TSession extends ZodSession<any, any, any, any>,
+	TClientMessage extends
+		SessionClientMessage<TSession> = SessionClientMessage<TSession>,
+	TServerMessage extends
+		SessionServerMessage<TSession> = SessionServerMessage<TSession>,
+	TEnv extends SessionEnv<TSession> = SessionEnv<TSession>,
+> extends BaseWebSocketDO<TSession, TEnv> {
+	protected readonly zodSessionOptions: ZodSessionOptionsOrFactory<
+		TClientMessage,
+		TServerMessage,
+		TEnv
+	>;
+	protected readonly createZodSessionFn: (
+		ctx: Context<{ Bindings: TEnv }> | undefined,
 		websocket: WebSocket,
 		options: ZodSessionOptions<TClientMessage, TServerMessage>,
-	): TSession | Promise<TSession>;
+	) => TSession | Promise<TSession>;
+
+	constructor(
+		ctx: DurableObjectState,
+		env: TEnv,
+		options: ZodWebSocketDOOptions<
+			TSession,
+			TClientMessage,
+			TServerMessage,
+			TEnv
+		>,
+	) {
+		super(ctx, env, {
+			createSession: (ctx, websocket) => {
+				const zodOptions =
+					typeof options.zodSessionOptions === "function"
+						? options.zodSessionOptions(ctx, websocket)
+						: options.zodSessionOptions;
 
-	protected createSession(websocket: WebSocket): TSession | Promise<TSession> {
-		const options = this.getZodOptions();
-		return this.createZodSession(websocket, options);
+				return options.createZodSession(ctx, websocket, zodOptions);
+			},
+		});
+		this.zodSessionOptions = options.zodSessionOptions;
+		this.createZodSessionFn = options.createZodSession;
 	}
 }

package/src/index.ts

@@ -1,10 +1,25 @@
-export { BaseSession, type SessionClientMessage } from "./BaseSession";
-export { BaseWebSocketDO } from "./BaseWebSocketDO";
+export {
+	BaseSession,
+	type BaseSessionHandlers,
+	type SessionClientMessage,
+} from "./BaseSession";
+export {
+	BaseWebSocketDO,
+	type BaseWebSocketDOOptions,
+} from "./BaseWebSocketDO";
 export { WebsocketWrapper } from "./WebsocketWrapper";
-export { ZodSession, type ZodSessionOptions } from "./ZodSession";
+export {
+	ZodSession,
+	type ZodSessionHandlers,
+	type ZodSessionOptions,
+} from "./ZodSession";
 export {
 	ZodWebSocketClient,
 	type ZodWebSocketClientOptions,
 } from "./ZodWebSocketClient";
-export { ZodWebSocketDO } from "./ZodWebSocketDO";
+export {
+	type ZodSessionOptionsOrFactory,
+	ZodWebSocketDO,
+	type ZodWebSocketDOOptions,
+} from "./ZodWebSocketDO";
 export { zodMsgpack } from "./zodMsgpack";