From 9f6a9979faa31c3900fb2fb974d98b99b9f57a9e Mon Sep 17 00:00:00 2001
From: saltyaom <saltyaom@gmail.com>
Date: Tue, 7 Mar 2023 17:12:21 +0700
Subject: [PATCH] :wrench: broom: fix .ws merged type

---
 CHANGELOG.md |  4 ++++
 package.json |  2 +-
 src/index.ts | 25 ++++++++++++++++++++++---
 src/types.ts | 22 ++++++++++++++++++++++
 4 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b6506e3d..064b82ab 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,7 @@
+# 0.3.0-rc.2 - 7 Mar 2023
+Fix:
+- Missing merged return type for `.ws`
+
 # 0.3.0-rc.1 - 7 Mar 2023
 Fix:
 - Missing nanoid
diff --git a/package.json b/package.json
index e92d2292..a944c174 100644
--- a/package.json
+++ b/package.json
@@ -1,7 +1,7 @@
 {
   "name": "elysia",
   "description": "Fast, and friendly Bun web framework",
-  "version": "0.3.0-rc.1",
+  "version": "0.3.0-rc.2",
   "author": {
     "name": "saltyAom",
     "url": "https://github.com/SaltyAom",
diff --git a/src/index.ts b/src/index.ts
index e93c878c..396ddbee 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -54,7 +54,8 @@ import type {
 	MaybePromise,
 	IsNever,
 	MergeUnionObjects,
-	TypedRouteToEden
+	TypedRouteToEden,
+	TypedWSRouteToEden
 } from './types'
 import { type TSchema } from '@sinclair/typebox'
 import { ElysiaWSContext, ElysiaWSOptions, WSTypedSchema } from './ws'
@@ -1256,7 +1257,7 @@ export default class Elysia<Instance extends ElysiaInstance = ElysiaInstance> {
 	 *     })
 	 * ```
 	 */
-	ws<Path extends string = '', Schema = {}>(
+	ws<Path extends string = '', Schema extends TypedSchema = {}>(
 		/**
 		 * Path to register websocket to
 		 */
@@ -1270,7 +1271,25 @@ export default class Elysia<Instance extends ElysiaInstance = ElysiaInstance> {
 					Instance['meta'][typeof DEFS]
 			  >
 			: never
-	): this extends Elysia<infer Instance> ? Elysia<Instance> : ElysiaInstance {
+	): Elysia<{
+		request: Instance['request']
+		store: Instance['store']
+		schema: Instance['schema']
+		meta: Instance['meta'] &
+			Record<
+				typeof SCHEMA,
+				Record<
+					Path,
+					{
+						[method in 'subscribe']: TypedWSRouteToEden<
+							Schema,
+							Instance['meta'][typeof DEFS],
+							Path
+						>
+					}
+				>
+			>
+	}> {
 		if (!this.wsRouter)
 			throw new Error(
 				"Can't find WebSocket. Please register WebSocket plugin first by importing 'elysia/ws'"
diff --git a/src/types.ts b/src/types.ts
index 764513a0..42bc164c 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -319,6 +319,28 @@ export type TypedRouteToEden<
 	  }
 	: never
 
+export type TypedWSRouteToEden<
+	Schema extends TypedSchema = TypedSchema,
+	Definitions extends TypedSchema<string> = ElysiaInstance['meta'][typeof DEFS],
+	Path extends string = string,
+	Catch = unknown
+> = TypedSchemaToEden<
+	Schema,
+	Definitions
+> extends infer Typed extends AnyTypedSchema
+	? {
+			body: Typed['body']
+			headers: Typed['headers']
+			query: Typed['query']
+			params: undefined extends Typed['params']
+				? Record<ExtractPath<Path>, string>
+				: Typed['params']
+			response: undefined extends Typed['response']
+				? Catch
+				: Typed['response']['200']
+	  }
+	: never
+
 export type TypedSchemaToEden<
 	Schema extends TypedSchema,
 	Definitions extends ElysiaInstance['meta'][typeof DEFS]