fix: typings for messgeSchema() funciton

Author: koistya

package.json

@@ -1,6 +1,6 @@
 {
   "name": "bun-ws-router",
-  "version": "0.0.1",
+  "version": "0.0.2",
   "description": "A simple and efficient WebSocket router for Bun with Zod-based message validation.",
   "keywords": [
     "bun",

schema.ts

@@ -1,7 +1,7 @@
 /* SPDX-FileCopyrightText: 2025-present Kriasoft */
 /* SPDX-License-Identifier: MIT */
 
-import type { ZodObject, ZodRawShape, ZodTypeAny } from "zod";
+import type { ZodLiteral, ZodObject, ZodRawShape, ZodTypeAny } from "zod";
 import { z } from "zod";
 
 /**
@@ -23,37 +23,164 @@ export const MessageSchema = z.object({
   meta: MessageMetadataSchema,
 });
 
+// -----------------------------------------------------------------------
+// Type Definitions
+// -----------------------------------------------------------------------
+
+/**
+ * Type for a message schema with no payload
+ */
+export type BaseMessageSchema<T extends string> = ZodObject<{
+  type: ZodLiteral<T>;
+  meta: typeof MessageMetadataSchema;
+}>;
+
+/**
+ * Type for a message schema with a payload
+ */
+export type PayloadMessageSchema<
+  T extends string,
+  P extends ZodTypeAny
+> = ZodObject<{
+  type: ZodLiteral<T>;
+  meta: typeof MessageMetadataSchema;
+  payload: P;
+}>;
+
+/**
+ * Type for a message schema with a custom meta object
+ */
+export type MessageSchemaWithCustomMeta<
+  T extends string,
+  M extends ZodRawShape
+> = ZodObject<{
+  type: ZodLiteral<T>;
+  meta: ZodObject<typeof MessageMetadataSchema.shape & M>;
+}>;
+
+/**
+ * Type for a message schema with a payload and custom meta object
+ */
+export type PayloadMessageSchemaWithCustomMeta<
+  T extends string,
+  P extends ZodTypeAny,
+  M extends ZodRawShape
+> = ZodObject<{
+  type: ZodLiteral<T>;
+  meta: ZodObject<typeof MessageMetadataSchema.shape & M>;
+  payload: P;
+}>;
+
+// -----------------------------------------------------------------------
+// Function Overloads
+// -----------------------------------------------------------------------
+
+/**
+ * Creates a message schema with a literal type but no payload or custom metadata
+ */
+export function messageSchema<T extends string>(
+  messageType: T
+): BaseMessageSchema<T>;
+
+/**
+ * Creates a message schema with a literal type and a payload schema (object form)
+ */
+export function messageSchema<
+  T extends string,
+  P extends Record<string, ZodTypeAny>
+>(messageType: T, payload: P): PayloadMessageSchema<T, ZodObject<P>>;
+
+/**
+ * Creates a message schema with a literal type and a payload schema (ZodType form)
+ */
+export function messageSchema<T extends string, P extends ZodTypeAny>(
+  messageType: T,
+  payload: P
+): PayloadMessageSchema<T, P>;
+
+/**
+ * Creates a message schema with a literal type and custom metadata
+ */
+export function messageSchema<T extends string, M extends ZodRawShape>(
+  messageType: T,
+  payload: undefined,
+  meta: ZodObject<M>
+): MessageSchemaWithCustomMeta<T, M>;
+
+/**
+ * Creates a message schema with a literal type, payload (object form), and custom metadata
+ */
+export function messageSchema<
+  T extends string,
+  P extends Record<string, ZodTypeAny>,
+  M extends ZodRawShape
+>(
+  messageType: T,
+  payload: P,
+  meta: ZodObject<M>
+): PayloadMessageSchemaWithCustomMeta<T, ZodObject<P>, M>;
+
+/**
+ * Creates a message schema with a literal type, payload (ZodType form), and custom metadata
+ */
+export function messageSchema<
+  T extends string,
+  P extends ZodTypeAny,
+  M extends ZodRawShape
+>(
+  messageType: T,
+  payload: P,
+  meta: ZodObject<M>
+): PayloadMessageSchemaWithCustomMeta<T, P, M>;
+
+// -----------------------------------------------------------------------
+// Implementation
+// -----------------------------------------------------------------------
+
 /**
  * A helper function to create specific WebSocket message schemas.
  * It extends the base `MessageSchema`, setting a literal type, adding a payload schema,
  * and optionally extending the metadata schema.
+ *
+ * Implementation for all overloads
  */
 export function messageSchema<
-  Payload extends Record<string, ZodTypeAny> | ZodTypeAny | undefined,
-  Metadata extends ZodRawShape | undefined = undefined
+  T extends string,
+  P extends Record<string, ZodTypeAny> | ZodTypeAny | undefined = undefined,
+  M extends ZodRawShape = {}
 >(
-  messageType: string,
-  schema?: Payload,
-  meta?: Metadata extends ZodRawShape ? ZodObject<Metadata> : undefined
-) {
+  messageType: T,
+  payload?: P,
+  meta?: M extends {} ? undefined : ZodObject<M>
+): P extends undefined
+  ? M extends {}
+    ? BaseMessageSchema<T>
+    : MessageSchemaWithCustomMeta<T, M>
+  : P extends Record<string, ZodTypeAny>
+  ? M extends {}
+    ? PayloadMessageSchema<T, ZodObject<P>>
+    : PayloadMessageSchemaWithCustomMeta<T, ZodObject<P>, M>
+  : M extends {}
+  ? PayloadMessageSchema<T, P & ZodTypeAny>
+  : PayloadMessageSchemaWithCustomMeta<T, P & ZodTypeAny, M> {
   // Create base schema with type and meta
   const baseSchema = MessageSchema.extend({
     type: z.literal(messageType),
     meta: meta
-      ? MessageMetadataSchema.extend(meta.shape)
+      ? MessageMetadataSchema.extend((meta as ZodObject<any>).shape)
       : MessageMetadataSchema,
   });
 
   // If no payload schema provided, return without payload
-  if (schema === undefined) {
-    return baseSchema;
+  if (payload === undefined) {
+    return baseSchema as any;
   }
 
   // Add payload to schema based on input type
   return baseSchema.extend({
     payload:
-      schema instanceof z.ZodType
-        ? schema
-        : z.object(schema as Record<string, ZodTypeAny>),
-  });
+      payload instanceof z.ZodType
+        ? payload
+        : z.object(payload as Record<string, ZodTypeAny>),
+  }) as any;
 }

types.ts

@@ -2,7 +2,8 @@
 /* SPDX-License-Identifier: MIT */
 
 import type { HeadersInit, Server, ServerWebSocket } from "bun";
-import { z, ZodLiteral, type ZodRawShape } from "zod";
+import type { ZodObject, ZodType, ZodTypeAny } from "zod";
+import { z, ZodLiteral } from "zod";
 import { MessageMetadataSchema } from "./schema";
 
 export type WebSocketRouterOptions = {
@@ -22,7 +23,7 @@ export type UpgradeOptions<T> = {
 export type SendFunction = <Schema extends MessageSchemaType>(
   schema: Schema,
   data: Schema["shape"] extends { payload: infer P }
-    ? P extends z.ZodTypeAny
+    ? P extends ZodTypeAny
       ? z.infer<P>
       : unknown
     : unknown,
@@ -34,7 +35,7 @@ export type MessageContext<Schema extends MessageSchemaType, Data> = {
   meta: z.infer<Schema["shape"]["meta"]>;
   send: SendFunction;
 } & (Schema["shape"] extends { payload: infer P }
-  ? P extends z.ZodTypeAny
+  ? P extends ZodTypeAny
     ? { payload: z.infer<P> }
     : {}
   : {});
@@ -43,12 +44,11 @@ export type MessageHandler<Schema extends MessageSchemaType, Data> = (
   context: MessageContext<Schema, Data>
 ) => void | Promise<void>;
 
-export type MessageSchemaType = z.ZodObject<
-  {
-    type: ZodLiteral<string>; // Must have a literal string type
-    meta: typeof MessageMetadataSchema; // Must have compatible meta
-  } & ZodRawShape // Allows for other fields like 'payload'
->;
+export type MessageSchemaType = ZodObject<{
+  type: ZodLiteral<string>; // Must have a literal string type
+  meta: ZodType<z.infer<typeof MessageMetadataSchema>>;
+  payload?: ZodTypeAny;
+}>;
 
 export type MessageHandlerEntry = {
   schema: MessageSchemaType;