elysia 2.0.0-beta.2 → 2.0.0-beta.4

package/dist/ws/context.js

@@ -8,6 +8,19 @@ function pickValidator(validators, defaultValidator, data) {
 	if (data instanceof require_error.ElysiaStatus) return validators[data.status];
 	return defaultValidator;
 }
+function trackWSSettling(data, result, onSettled) {
+	if (!(result instanceof Promise)) return;
+	data.settling = (data.settling ?? 0) + 1;
+	const release = () => {
+		data.settling--;
+		onSettled?.();
+	};
+	result.then(release, release);
+}
+const runWSLifecycle = (data, callback, ...args) => {
+	const lifecycle = data["~lifecycleRun"];
+	return lifecycle ? lifecycle(callback, ...args) : callback(...args);
+};
 function memoize(view, key, value) {
 	const self = view.raw.data.elysia ?? view;
 	Object.defineProperty(self, key, {
@@ -24,6 +37,15 @@ var ElysiaWS = class {
 		this.raw = raw;
 		if (context) for (const key in context) {
 			if (key === "ws" || key === "body") continue;
+			if (key === "__proto__") {
+				Object.defineProperty(this, key, {
+					value: context[key],
+					writable: true,
+					enumerable: true,
+					configurable: true
+				});
+				continue;
+			}
 			this[key] = context[key];
 		}
 	}
@@ -62,58 +84,71 @@ var ElysiaWS = class {
 	get data() {
 		return this.raw.data;
 	}
-	#prepare(data) {
-		if (data === void 0) return null;
-		if (data instanceof require_error.ElysiaStatus) return JSON.stringify({
-			status: data.status,
-			error: data.response
-		});
-		if (typeof data === "object") return JSON.stringify(data);
-		return data;
-	}
-	#validatedOrError(data) {
+	/**
+	* Validate, redact and serialize an outbound value, or return the
+	* `ValidationError` the payload was refused with.
+	*/
+	#frame(data) {
 		const connectionData = this.raw.data;
 		const v = pickValidator(connectionData?.validator, connectionData?.defaultValidator, data);
-		if (!v) return;
-		const value = data instanceof require_error.ElysiaStatus ? data.response : data;
-		if (v.Check(value)) return;
-		return new require_error.ValidationError("message", value, v.Errors(value)).message;
+		const status = data instanceof require_error.ElysiaStatus ? data : void 0;
+		let value = status ? status.response : data;
+		if (v) {
+			const encode = v.EncodeFrom;
+			let redacted;
+			if (typeof encode === "function" && !v.hasCodec && !v.mayReturnPromise) try {
+				redacted = encode.call(v, value, "message");
+			} catch {}
+			if (redacted !== void 0) value = redacted;
+			else if (!v.Check(value)) return new require_error.ValidationError("message", value, v.Errors(value));
+		}
+		if (status) return JSON.stringify({
+			status: status.status,
+			error: value
+		});
+		if (value === void 0) return null;
+		if (typeof value === "object") return JSON.stringify(value);
+		return value;
 	}
 	#send(data, compress) {
 		if (data === void 0) return 0;
 		if (data instanceof ArrayBuffer || ArrayBuffer.isView(data)) return this.raw.send(data, compress);
-		const err = this.#validatedOrError(data);
-		if (err !== void 0) return this.raw.send(err);
-		return this.raw.send(this.#prepare(data), compress);
+		const frame = this.#frame(data);
+		if (frame instanceof require_error.ValidationError) return this.raw.send(frame.message);
+		return this.raw.send(frame, compress);
 	}
 	#ping(data) {
 		if (data === void 0) return this.raw.ping();
 		if (data instanceof ArrayBuffer || ArrayBuffer.isView(data)) return this.raw.ping(data);
-		const err = this.#validatedOrError(data);
-		if (err !== void 0) return this.raw.send(err);
-		return this.raw.ping(this.#prepare(data));
+		const frame = this.#frame(data);
+		if (frame instanceof require_error.ValidationError) return this.raw.send(frame.message);
+		return this.raw.ping(frame);
 	}
 	#pong(data) {
 		if (data === void 0) return this.raw.pong();
 		if (data instanceof ArrayBuffer || ArrayBuffer.isView(data)) return this.raw.pong(data);
-		const err = this.#validatedOrError(data);
-		if (err !== void 0) return this.raw.send(err);
-		return this.raw.pong(this.#prepare(data));
+		const frame = this.#frame(data);
+		if (frame instanceof require_error.ValidationError) return this.raw.send(frame.message);
+		return this.raw.pong(frame);
 	}
 	#publish(topic, data, compress) {
 		if (data === void 0) return 0;
 		if (data instanceof ArrayBuffer || ArrayBuffer.isView(data)) return this.raw.publish(topic, data, compress);
-		const err = this.#validatedOrError(data);
-		if (err !== void 0) return this.raw.send(err);
-		return this.raw.publish(topic, this.#prepare(data), compress);
+		const frame = this.#frame(data);
+		if (frame instanceof require_error.ValidationError) return this.raw.send(frame.message);
+		return this.raw.publish(topic, frame, compress);
 	}
 	#close(code, reason) {
 		const data = this.raw.data;
 		if (!data.closeHandlerInvoked && data.close) {
 			data.closeHandlerInvoked = true;
 			try {
-				const result = data.close(this, code ?? 1e3, reason ?? "");
-				if (result instanceof Promise) result.then(() => this.raw.close(code, reason)).catch(() => this.raw.close(code, reason));
+				const result = runWSLifecycle(data, data.close, this, code ?? 1e3, reason ?? "");
+				if (result instanceof Promise) {
+					const closing = result.catch(() => {});
+					trackWSSettling(data, closing);
+					closing.then(() => this.raw.close(code, reason)).catch(() => this.raw.close(code, reason));
+				}
 				return;
 			} catch {}
 		}
@@ -149,4 +184,5 @@ function isGeneratorObject(value) {
 
 //#endregion
 exports.ElysiaWS = ElysiaWS;
-exports.isGeneratorObject = isGeneratorObject;
+exports.isGeneratorObject = isGeneratorObject;
+exports.trackWSSettling = trackWSSettling;

package/dist/ws/context.mjs

@@ -7,6 +7,19 @@ function pickValidator(validators, defaultValidator, data) {
 	if (data instanceof ElysiaStatus) return validators[data.status];
 	return defaultValidator;
 }
+function trackWSSettling(data, result, onSettled) {
+	if (!(result instanceof Promise)) return;
+	data.settling = (data.settling ?? 0) + 1;
+	const release = () => {
+		data.settling--;
+		onSettled?.();
+	};
+	result.then(release, release);
+}
+const runWSLifecycle = (data, callback, ...args) => {
+	const lifecycle = data["~lifecycleRun"];
+	return lifecycle ? lifecycle(callback, ...args) : callback(...args);
+};
 function memoize(view, key, value) {
 	const self = view.raw.data.elysia ?? view;
 	Object.defineProperty(self, key, {
@@ -23,6 +36,15 @@ var ElysiaWS = class {
 		this.raw = raw;
 		if (context) for (const key in context) {
 			if (key === "ws" || key === "body") continue;
+			if (key === "__proto__") {
+				Object.defineProperty(this, key, {
+					value: context[key],
+					writable: true,
+					enumerable: true,
+					configurable: true
+				});
+				continue;
+			}
 			this[key] = context[key];
 		}
 	}
@@ -61,58 +83,71 @@ var ElysiaWS = class {
 	get data() {
 		return this.raw.data;
 	}
-	#prepare(data) {
-		if (data === void 0) return null;
-		if (data instanceof ElysiaStatus) return JSON.stringify({
-			status: data.status,
-			error: data.response
-		});
-		if (typeof data === "object") return JSON.stringify(data);
-		return data;
-	}
-	#validatedOrError(data) {
+	/**
+	* Validate, redact and serialize an outbound value, or return the
+	* `ValidationError` the payload was refused with.
+	*/
+	#frame(data) {
 		const connectionData = this.raw.data;
 		const v = pickValidator(connectionData?.validator, connectionData?.defaultValidator, data);
-		if (!v) return;
-		const value = data instanceof ElysiaStatus ? data.response : data;
-		if (v.Check(value)) return;
-		return new ValidationError("message", value, v.Errors(value)).message;
+		const status = data instanceof ElysiaStatus ? data : void 0;
+		let value = status ? status.response : data;
+		if (v) {
+			const encode = v.EncodeFrom;
+			let redacted;
+			if (typeof encode === "function" && !v.hasCodec && !v.mayReturnPromise) try {
+				redacted = encode.call(v, value, "message");
+			} catch {}
+			if (redacted !== void 0) value = redacted;
+			else if (!v.Check(value)) return new ValidationError("message", value, v.Errors(value));
+		}
+		if (status) return JSON.stringify({
+			status: status.status,
+			error: value
+		});
+		if (value === void 0) return null;
+		if (typeof value === "object") return JSON.stringify(value);
+		return value;
 	}
 	#send(data, compress) {
 		if (data === void 0) return 0;
 		if (data instanceof ArrayBuffer || ArrayBuffer.isView(data)) return this.raw.send(data, compress);
-		const err = this.#validatedOrError(data);
-		if (err !== void 0) return this.raw.send(err);
-		return this.raw.send(this.#prepare(data), compress);
+		const frame = this.#frame(data);
+		if (frame instanceof ValidationError) return this.raw.send(frame.message);
+		return this.raw.send(frame, compress);
 	}
 	#ping(data) {
 		if (data === void 0) return this.raw.ping();
 		if (data instanceof ArrayBuffer || ArrayBuffer.isView(data)) return this.raw.ping(data);
-		const err = this.#validatedOrError(data);
-		if (err !== void 0) return this.raw.send(err);
-		return this.raw.ping(this.#prepare(data));
+		const frame = this.#frame(data);
+		if (frame instanceof ValidationError) return this.raw.send(frame.message);
+		return this.raw.ping(frame);
 	}
 	#pong(data) {
 		if (data === void 0) return this.raw.pong();
 		if (data instanceof ArrayBuffer || ArrayBuffer.isView(data)) return this.raw.pong(data);
-		const err = this.#validatedOrError(data);
-		if (err !== void 0) return this.raw.send(err);
-		return this.raw.pong(this.#prepare(data));
+		const frame = this.#frame(data);
+		if (frame instanceof ValidationError) return this.raw.send(frame.message);
+		return this.raw.pong(frame);
 	}
 	#publish(topic, data, compress) {
 		if (data === void 0) return 0;
 		if (data instanceof ArrayBuffer || ArrayBuffer.isView(data)) return this.raw.publish(topic, data, compress);
-		const err = this.#validatedOrError(data);
-		if (err !== void 0) return this.raw.send(err);
-		return this.raw.publish(topic, this.#prepare(data), compress);
+		const frame = this.#frame(data);
+		if (frame instanceof ValidationError) return this.raw.send(frame.message);
+		return this.raw.publish(topic, frame, compress);
 	}
 	#close(code, reason) {
 		const data = this.raw.data;
 		if (!data.closeHandlerInvoked && data.close) {
 			data.closeHandlerInvoked = true;
 			try {
-				const result = data.close(this, code ?? 1e3, reason ?? "");
-				if (result instanceof Promise) result.then(() => this.raw.close(code, reason)).catch(() => this.raw.close(code, reason));
+				const result = runWSLifecycle(data, data.close, this, code ?? 1e3, reason ?? "");
+				if (result instanceof Promise) {
+					const closing = result.catch(() => {});
+					trackWSSettling(data, closing);
+					closing.then(() => this.raw.close(code, reason)).catch(() => this.raw.close(code, reason));
+				}
 				return;
 			} catch {}
 		}
@@ -147,4 +182,4 @@ function isGeneratorObject(value) {
 }
 
 //#endregion
-export { ElysiaWS, isGeneratorObject };
+export { ElysiaWS, isGeneratorObject, trackWSSettling };

package/dist/ws/route.js

@@ -2,11 +2,12 @@ Object.defineProperty(exports, Symbol.toStringTag, { value: 'Module' });
 const require_universal_constants = require('../universal/constants.js');
 const require_utils = require('../utils.js');
 const require_type_bridge = require('../type/bridge.js');
-const require_universal_is_production = require('../universal/is-production.js');
 const require_error = require('../error.js');
 const require_adapter_web_standard_handler = require('../adapter/web-standard/handler.js');
 const require_parse_query = require('../parse-query.js');
 const require_generation = require('../generation.js');
+const require_cookie_config = require('../cookie/config.js');
+const require_cookie_utils = require('../cookie/utils.js');
 const require_handler_error = require('../handler/error.js');
 const require_compile_handler_utils = require('../compile/handler/utils.js');
 const require_validator_index = require('../validator/index.js');
@@ -181,9 +182,45 @@ function wsErrorFrame(error) {
 	}
 	return wsErrorFrameFallback(error);
 }
+/**
+* Frame served by an error that claims an RFC 9457 problem document, mirroring
+* the HTTP tiers in `handler/error.ts`: `value` replaces the whole frame,
+* `detail` fills the `detail` member, the message is the last resort.
+*
+* A frame carries no headers, so a `headers` annotation is ignored outright
+* rather than half-applied
+*/
+function wsProblemFrame(error) {
+	const status = require_handler_error.resolveStatus(error.status);
+	const served = typeof status === "number" && status >= 100 ? status : 500;
+	const stringify = (value) => typeof value === "object" && value !== null ? JSON.stringify(value) : String(value);
+	const problemFrame = (detail) => JSON.stringify(require_error.problemBody({
+		type: error.type ?? "about:blank",
+		detail,
+		status: served
+	}));
+	const serveMessage = () => problemFrame(require_handler_error.statusFallbackBody(error, served));
+	const serveDetail = () => {
+		const detail = require_handler_error.readAnnotation(error, "detail", true);
+		if (detail === void 0) return serveMessage();
+		if (detail instanceof Promise) return detail.then((resolved) => resolved === void 0 ? serveMessage() : problemFrame(resolved)).catch(() => wsLegacyFrame(error));
+		return problemFrame(detail);
+	};
+	const value = require_handler_error.readAnnotation(error, "value", true);
+	if (value === void 0) return serveDetail();
+	if (value instanceof Promise) return value.then((resolved) => resolved === void 0 ? serveDetail() : stringify(resolved)).catch(() => wsLegacyFrame(error));
+	return stringify(value);
+}
 function wsErrorFrameFallback(error) {
+	if (require_handler_error.claimsProblemType(error)) try {
+		return wsProblemFrame(error);
+	} catch {}
+	return wsLegacyFrame(error);
+}
+/** Frame an error that never self-described has always served */
+function wsLegacyFrame(error) {
 	if (error?.status) {
-		const body = error.response !== void 0 ? error.response : require_universal_is_production.isProduction() ? error.status >= 500 ? "Internal Server Error" : "" : error.message ?? "";
+		const body = require_handler_error.statusFallbackBody(error, error.status);
 		return typeof body === "object" ? JSON.stringify(body) : String(body);
 	}
 	return require_error.internalServerErrorBodyString(error);
@@ -218,13 +255,20 @@ const wsOptions = [
 ];
 function buildWSRoute(route, app) {
 	const hook = require_compile_handler_index.resolveWSLocalHook(require_compile_handler_index.localMacroRoot(route[7] ?? route[3] ?? app, app), route[4], app) ?? require_utils.nullObject();
+	const instance = route[3] ?? app;
+	const appHookChain = route[5];
+	const inheritedChain = route[6];
+	const composed = require_compile_handler_index.composeRouteHook(instance, route[4], appHookChain, inheritedChain, app, route[7]) ?? hook;
 	let validators;
-	const frozenEager = require_type_bridge.isBridgeLive() ? void 0 : require_compile_handler_frozen_validator.buildFrozenRouteValidator(hook, app, "WS", route[1]);
+	const frozenEager = require_type_bridge.isBridgeLive() ? void 0 : require_compile_handler_frozen_validator.buildFrozenRouteValidator(composed, app, "WS", route[1]);
 	if (frozenEager) validators = frozenEager;
 	else try {
-		validators = new require_validator_route.RouteValidator(hook, {
+		validators = new require_validator_route.RouteValidator(composed, {
 			models: require_generation.frozenRootOf(app)["~ext"]?.models,
 			app,
+			normalize: require_generation.frozenRootOf(app)["~config"]?.normalize,
+			sanitize: require_generation.frozenRootOf(app)["~config"]?.sanitize,
+			schemas: composed.schemas,
 			aot: {
 				method: "WS",
 				path: route[1]
@@ -233,18 +277,17 @@ function buildWSRoute(route, app) {
 		});
 	} catch (error) {
 		if (!require_compile_handler_frozen_validator.isBridgeNotInitialized(error)) throw error;
-		const frozen = require_compile_handler_frozen_validator.buildFrozenRouteValidator(hook, app, "WS", route[1]);
+		const frozen = require_compile_handler_frozen_validator.buildFrozenRouteValidator(composed, app, "WS", route[1]);
 		if (!frozen) throw error;
 		validators = frozen;
 	}
 	const responseValidator = validators.response;
 	const defaultResponseValidator = responseValidator ? responseValidator[200] ?? responseValidator[Object.keys(responseValidator)[0]] : void 0;
+	const cookieConfig = validators.cookie ? require_cookie_config.compileCookieConfig(composed.cookie, require_generation.frozenRootOf(app)["~config"]?.cookie) : void 0;
+	const cookieIsOptional = !!composed.cookie?.["~optional"];
 	const queryChannels = require_compile_handler_utils.getQueryParseChannels(validators.query?.schema);
 	const queryArray = queryChannels?.array;
 	const queryObject = queryChannels?.object;
-	const instance = route[3] ?? app;
-	const appHookChain = route[5];
-	const inheritedChain = route[6];
 	const flatAppHook = require_compile_handler_index.composeRouteHook(instance, void 0, appHookChain, inheritedChain, app, route[7]) ?? {};
 	const parseHooks = hook.parse == null ? [] : Array.isArray(hook.parse) ? hook.parse : [hook.parse];
 	const transforms = concatHooks(flatAppHook.transform, hook.transform);
@@ -272,6 +315,7 @@ function buildWSRoute(route, app) {
 				break;
 			}
 			if (r !== void 0) {
+				r = require_handler_error.adoptErrorType(r, error);
 				try {
 					await handleWSResponse(ws, r, mapResponses);
 				} catch {}
@@ -419,6 +463,15 @@ function buildWSRoute(route, app) {
 				if (r instanceof Promise) r = await r;
 				context.headers = r;
 			}
+			if (cookieConfig) {
+				const raw = await require_cookie_utils.parseCookieRaw(request.headers.get("cookie"), cookieConfig);
+				let r = raw;
+				if (!cookieIsOptional || Object.keys(raw).length) {
+					r = validateUpgradeChannel(validators.cookie, raw, "cookie");
+					if (r instanceof Promise) r = await r;
+				}
+				context.cookie = require_cookie_utils.buildCookieJar(context.set, r, cookieConfig);
+			}
 			for (let i = 0; i < transforms.length; i++) {
 				const r = transforms[i](context);
 				if (r instanceof Promise) await r;
@@ -430,7 +483,7 @@ function buildWSRoute(route, app) {
 				const deriveMode = upgradeDeriveModes?.[i];
 				if (deriveMode !== void 0 && !(r instanceof require_error.ElysiaStatus)) {
 					if (r && typeof r === "object") if (deriveMode) context = require_compile_handler_utils.replaceDeriveContext(context, r);
-					else Object.assign(context, r);
+					else require_utils.assignOwn(context, r);
 				} else if (r !== void 0) {
 					if (r instanceof Response) return r;
 					return require_adapter_web_standard_handler.mapResponse(r, context.set, context.request);
@@ -485,23 +538,40 @@ function resolveWSOptions(entries) {
 		const incoming = entry.value[key];
 		if (key in result && result[key] !== incoming && !warned?.has(key)) {
 			(warned ??= /* @__PURE__ */ new Set()).add(key);
-			console.warn(`[Elysia] Conflicting WebSocket option '${key}' across .use(websocket()) registrations; the shallowest (nearest-root) registration wins.`);
+			console.warn(`[Elysia] Conflicting WebSocket option '${key}' across .use(websocket()). Using nearest-root registration.`);
 		}
 		result[key] = incoming;
 	}
 	return result;
 }
+const wsLimits = [
+	"maxPayloadLength",
+	"backpressureLimit",
+	"idleTimeout"
+];
+const routeOwnedKeys = /* @__PURE__ */ new WeakMap();
 function accumulateWSOptions(target, routeOptions, _path) {
 	if (!require_utils.isNotEmpty(routeOptions)) return;
-	if (require_universal_constants.isBun) {
-		for (const key in routeOptions) if (key in target && target[key] !== routeOptions[key]) {
-			console.warn(`[Elysia] Conflicting per-route WebSocket option '${key}'\nBun uses one global WebSocket config per server, per-route values are not enforced (the last-registered route wins).`);
-			console.warn((/* @__PURE__ */ new Error()).stack);
+	let owned = routeOwnedKeys.get(target);
+	if (!owned) routeOwnedKeys.set(target, owned = /* @__PURE__ */ new Set());
+	for (const key in routeOptions) {
+		const incoming = routeOptions[key];
+		const current = target[key];
+		if (key in target && current !== incoming) {
+			if (require_universal_constants.isBun) {
+				console.warn(`[Elysia] Conflicting per-route WebSocket option '${key}'\nBun uses one global WebSocket config per server, per-route values are not enforced (for limits the strictest route wins, otherwise the last-registered route).`);
+				console.warn((/* @__PURE__ */ new Error()).stack);
+			}
+			if (owned.has(key) && typeof current === "number" && typeof incoming === "number" && current < incoming && wsLimits.includes(key)) continue;
 		}
+		owned.add(key);
+		target[key] = incoming;
 	}
-	Object.assign(target, routeOptions);
 }
+const MAX_INFLIGHT_MESSAGES = 256;
 function buildGlobalWSHandler() {
+	let lifecycle;
+	const runLifecycle = (callback, ...args) => lifecycle ? lifecycle.run(callback, ...args) : callback(...args);
 	function getElysia(ws) {
 		let elysia = ws.data.elysia;
 		if (!elysia) {
@@ -516,39 +586,117 @@ function buildGlobalWSHandler() {
 			ws.send(typeof error?.message === "string" ? error.message : "Internal Server Error");
 		} catch {}
 	};
-	return {
+	const releaseLifecycle = (ws) => {
+		if (!lifecycle?.closing && require_utils.isSocketQuiet(ws)) {
+			delete ws.data["~lifecycleRun"];
+			lifecycle?.sockets.delete(ws);
+		}
+	};
+	const release = (ws) => {
+		const data = ws.data;
+		const inflight = data.inflight;
+		if (inflight) {
+			data.inflight = inflight - 1;
+			if (inflight === 1 && data.closeHandlerInvoked) releaseLifecycle(ws);
+		}
+	};
+	function dispatch(ws, message) {
+		let result;
+		try {
+			result = ws.data.message(getElysia(ws), message);
+		} catch (error) {
+			release(ws);
+			sendErrMsg(ws, error);
+			return;
+		}
+		if (result instanceof Promise) result.then(() => release(ws), (error) => {
+			release(ws);
+			sendErrMsg(ws, error);
+		});
+		else release(ws);
+	}
+	const handler = {
 		message(ws, message) {
-			let result;
-			try {
-				result = ws.data.message?.(getElysia(ws), message);
-			} catch (error) {
-				sendErrMsg(ws, error);
+			if (lifecycle?.closing) return;
+			const data = ws.data;
+			if (!data.message) return;
+			const inflight = (data.inflight ?? 0) + 1;
+			if (inflight > MAX_INFLIGHT_MESSAGES) {
+				try {
+					ws.close(1013, "Too many in-flight messages");
+				} catch {}
 				return;
 			}
-			if (result instanceof Promise) result.catch((error) => sendErrMsg(ws, error));
+			data.inflight = inflight;
+			const opening = data.opening;
+			if (opening) return void opening.then(() => {
+				if (lifecycle?.closing || ws.readyState > 1) return release(ws);
+				dispatch(ws, message);
+			});
+			dispatch(ws, message);
 		},
 		open(ws) {
-			ws.data.open?.(getElysia(ws));
+			if (lifecycle) Object.defineProperty(ws.data, "~lifecycleRun", {
+				configurable: true,
+				value: lifecycle.run
+			});
+			lifecycle?.sockets.add(ws);
+			if (lifecycle?.closing) {
+				ws.data.closeHandlerInvoked = true;
+				try {
+					ws.terminate();
+				} catch {}
+				return;
+			}
+			const open = ws.data.open;
+			const result = open && runLifecycle(open, getElysia(ws));
+			if (result instanceof Promise) {
+				const clear = () => {
+					ws.data.opening = void 0;
+					releaseLifecycle(ws);
+				};
+				ws.data.opening = result.then(clear, clear);
+			}
 		},
 		drain(ws) {
+			if (lifecycle?.closing) return;
 			const elyWs = getElysia(ws);
 			drainWaiters(elyWs);
-			ws.data.drain?.(elyWs);
+			const drain = ws.data.drain;
+			require_ws_context.trackWSSettling(ws.data, drain && runLifecycle(drain, elyWs), () => releaseLifecycle(ws));
 		},
 		close(ws, code, reason) {
 			const elyWs = getElysia(ws);
 			drainWaiters(elyWs);
-			if (ws.data.closeHandlerInvoked) return;
+			if (ws.data.closeHandlerInvoked) return releaseLifecycle(ws);
 			ws.data.closeHandlerInvoked = true;
-			ws.data.close?.(elyWs, code, reason);
+			let result;
+			try {
+				const close = ws.data.close;
+				result = close && runLifecycle(close, elyWs, code, reason);
+			} finally {
+				if (!(result instanceof Promise)) releaseLifecycle(ws);
+			}
+			if (result instanceof Promise) {
+				const closing = result.catch(() => {});
+				require_ws_context.trackWSSettling(ws.data, closing, () => releaseLifecycle(ws));
+			}
 		},
 		ping(ws, data) {
-			ws.data.ping?.(getElysia(ws), data);
+			if (lifecycle?.closing) return;
+			const ping = ws.data.ping;
+			require_ws_context.trackWSSettling(ws.data, ping && runLifecycle(ping, getElysia(ws), data), () => releaseLifecycle(ws));
 		},
 		pong(ws, data) {
-			ws.data.pong?.(getElysia(ws), data);
+			if (lifecycle?.closing) return;
+			const pong = ws.data.pong;
+			require_ws_context.trackWSSettling(ws.data, pong && runLifecycle(pong, getElysia(ws), data), () => releaseLifecycle(ws));
 		}
 	};
+	Object.defineProperty(handler, "~lifecycle", { set(value) {
+		lifecycle = value;
+	} });
+	return handler;
 }
 
 //#endregion