Feat/websocket controller

core/controller-decorator/index.ts

@@ -1,4 +1,6 @@
 import './src/impl/http/HTTPControllerMetaBuilder';
+import './src/impl/websocket/WebSocketControllerMetaBuilder';
+import './src/impl/websocket-fetch/WebSocketFetchControllerMetaBuilder';
 import './src/impl/mcp/MCPControllerMetaBuilder';
 
 export * from '@eggjs/tegg-types/controller-decorator';
@@ -10,6 +12,12 @@ export * from './src/decorator/http/HTTPController';
 export * from './src/decorator/http/HTTPMethod';
 export * from './src/decorator/http/HTTPParam';
 export * from './src/decorator/http/Host';
+export * from './src/decorator/websocket/WebSocketController';
+export * from './src/decorator/websocket/WebSocketMethod';
+export * from './src/decorator/websocket/WebSocketParam';
+export * from './src/decorator/websocket-fetch/WebSocketFetchController';
+export * from './src/decorator/websocket-fetch/WebSocketFetchMethod';
+export * from './src/decorator/websocket-fetch/WebSocketFetchParam';
 export * from './src/decorator/mcp/MCPController';
 export * from './src/decorator/mcp/MCPPrompt';
 export * from './src/decorator/mcp/MCPResource';
@@ -19,6 +27,7 @@ export * from './src/builder/ControllerMetaBuilderFactory';
 export * from './src/util/ControllerMetadataUtil';
 export * from './src/util/MCPInfoUtil';
 export * from './src/util/HTTPPriorityUtil';
+export { default as WebSocketInfoUtil } from './src/util/WebSocketInfoUtil';
 
 export { default as ControllerInfoUtil } from './src/util/ControllerInfoUtil';
 export { default as MethodInfoUtil } from './src/util/MethodInfoUtil';

core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchController.ts

@@ -0,0 +1,29 @@
+import { PrototypeUtil, SingletonProto } from '@eggjs/core-decorator';
+import { StackUtil } from '@eggjs/tegg-common-util';
+import type { EggProtoImplClass, WebSocketFetchControllerParams } from '@eggjs/tegg-types';
+import { AccessLevel, ControllerType } from '@eggjs/tegg-types';
+import ControllerInfoUtil from '../../util/ControllerInfoUtil';
+import WebSocketInfoUtil from '../../util/WebSocketInfoUtil';
+
+export function WebSocketFetchController(param?: WebSocketFetchControllerParams) {
+  return function(constructor: EggProtoImplClass) {
+    ControllerInfoUtil.setControllerType(constructor, ControllerType.WEBSOCKET_FETCH);
+    if (param?.controllerName) {
+      ControllerInfoUtil.setControllerName(constructor, param.controllerName);
+    }
+    if (param?.path) {
+      WebSocketInfoUtil.setWebSocketPath(param.path, constructor);
+    }
+    if (param?.timeout !== undefined) {
+      ControllerInfoUtil.setControllerTimeout(param.timeout, constructor);
+    }
+
+    const func = SingletonProto({
+      accessLevel: AccessLevel.PUBLIC,
+      name: param?.protoName,
+    });
+    func(constructor);
+
+    PrototypeUtil.setFilePath(constructor, StackUtil.getCalleeFromStack(false, 5));
+  };
+}

core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchMethod.ts

@@ -0,0 +1,62 @@
+import assert from 'node:assert';
+import { ControllerType, WebSocketFetchMethodType } from '@eggjs/tegg-types';
+import type { EggProtoImplClass, WebSocketFetchLifecycleMethodParams, WebSocketFetchMethodParams } from '@eggjs/tegg-types';
+import MethodInfoUtil from '../../util/MethodInfoUtil';
+import WebSocketInfoUtil from '../../util/WebSocketInfoUtil';
+
+function setWebSocketFetchMethodType(
+  target: any,
+  propertyKey: PropertyKey,
+  methodType: WebSocketFetchMethodType,
+  timeout?: number,
+) {
+  assert(typeof propertyKey === 'string',
+    `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`);
+  const controllerClazz = target.constructor as EggProtoImplClass;
+  const methodName = propertyKey as string;
+  MethodInfoUtil.setMethodControllerType(controllerClazz, methodName, ControllerType.WEBSOCKET_FETCH);
+  WebSocketInfoUtil.setWebSocketFetchMethodType(methodType, controllerClazz, methodName);
+  if (timeout !== undefined) {
+    MethodInfoUtil.setMethodTimeout(timeout, controllerClazz, methodName);
+  }
+  return { controllerClazz, methodName };
+}
+
+export function WebSocketFetchMethod(param?: WebSocketFetchMethodParams) {
+  return function(target: any, propertyKey: PropertyKey) {
+    const { controllerClazz, methodName } = setWebSocketFetchMethodType(
+      target,
+      propertyKey,
+      WebSocketFetchMethodType.DATA,
+      param?.timeout,
+    );
+    WebSocketInfoUtil.setWebSocketMethodPath('', controllerClazz, methodName);
+    if (param?.priority !== undefined) {
+      WebSocketInfoUtil.setWebSocketMethodPriority(param.priority, controllerClazz, methodName);
+    }
+  };
+}
+
+export function WebSocketFetchOnConnection(param?: WebSocketFetchLifecycleMethodParams) {
+  return function(target: any, propertyKey: PropertyKey) {
+    setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.CONNECTION, param?.timeout);
+  };
+}
+
+export function WebSocketFetchOnOpen(param?: WebSocketFetchLifecycleMethodParams) {
+  return function(target: any, propertyKey: PropertyKey) {
+    setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.OPEN, param?.timeout);
+  };
+}
+
+export function WebSocketFetchOnError(param?: WebSocketFetchLifecycleMethodParams) {
+  return function(target: any, propertyKey: PropertyKey) {
+    setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.ERROR, param?.timeout);
+  };
+}
+
+export function WebSocketFetchOnClose(param?: WebSocketFetchLifecycleMethodParams) {
+  return function(target: any, propertyKey: PropertyKey) {
+    setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.CLOSE, param?.timeout);
+  };
+}

core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchParam.ts

@@ -0,0 +1,47 @@
+import assert from 'node:assert';
+import { WebSocketParamType } from '@eggjs/tegg-types';
+import type { EggProtoImplClass } from '@eggjs/tegg-types';
+import WebSocketInfoUtil from '../../util/WebSocketInfoUtil';
+
+function setWebSocketFetchParamType(
+  target: any,
+  propertyKey: PropertyKey,
+  parameterIndex: number,
+  paramType: WebSocketParamType,
+) {
+  assert(typeof propertyKey === 'string',
+    `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`);
+  const methodName = propertyKey as string;
+  const controllerClazz = target.constructor as EggProtoImplClass;
+  WebSocketInfoUtil.setWebSocketMethodParamType(paramType, parameterIndex, controllerClazz, methodName);
+}
+
+export function WebSocketData() {
+  return function(target: any, propertyKey: PropertyKey, parameterIndex: number) {
+    setWebSocketFetchParamType(target, propertyKey, parameterIndex, WebSocketParamType.DATA);
+  };
+}
+
+export function WebSocketClose() {
+  return function(target: any, propertyKey: PropertyKey, parameterIndex: number) {
+    setWebSocketFetchParamType(target, propertyKey, parameterIndex, WebSocketParamType.CLOSE);
+  };
+}
+
+export function WebSocketError() {
+  return function(target: any, propertyKey: PropertyKey, parameterIndex: number) {
+    setWebSocketFetchParamType(target, propertyKey, parameterIndex, WebSocketParamType.ERROR);
+  };
+}
+
+export function WebSocketCloseCode() {
+  return function(target: any, propertyKey: PropertyKey, parameterIndex: number) {
+    setWebSocketFetchParamType(target, propertyKey, parameterIndex, WebSocketParamType.CLOSE_CODE);
+  };
+}
+
+export function WebSocketCloseReason() {
+  return function(target: any, propertyKey: PropertyKey, parameterIndex: number) {
+    setWebSocketFetchParamType(target, propertyKey, parameterIndex, WebSocketParamType.CLOSE_REASON);
+  };
+}

core/controller-decorator/src/decorator/websocket/WebSocketController.ts

@@ -0,0 +1,29 @@
+import { PrototypeUtil, SingletonProto } from '@eggjs/core-decorator';
+import { StackUtil } from '@eggjs/tegg-common-util';
+import type { EggProtoImplClass, WebSocketControllerParams } from '@eggjs/tegg-types';
+import { AccessLevel, ControllerType } from '@eggjs/tegg-types';
+import ControllerInfoUtil from '../../util/ControllerInfoUtil';
+import WebSocketInfoUtil from '../../util/WebSocketInfoUtil';
+
+export function WebSocketController(param?: WebSocketControllerParams) {
+  return function(constructor: EggProtoImplClass) {
+    ControllerInfoUtil.setControllerType(constructor, ControllerType.WEBSOCKET);
+    if (param?.controllerName) {
+      ControllerInfoUtil.setControllerName(constructor, param.controllerName);
+    }
+    if (param?.path) {
+      WebSocketInfoUtil.setWebSocketPath(param.path, constructor);
+    }
+    if (param?.timeout !== undefined) {
+      ControllerInfoUtil.setControllerTimeout(param.timeout, constructor);
+    }
+
+    const func = SingletonProto({
+      accessLevel: AccessLevel.PUBLIC,
+      name: param?.protoName,
+    });
+    func(constructor);
+
+    PrototypeUtil.setFilePath(constructor, StackUtil.getCalleeFromStack(false, 5));
+  };
+}

core/controller-decorator/src/decorator/websocket/WebSocketMethod.ts

@@ -0,0 +1,22 @@
+import assert from 'node:assert';
+import { ControllerType } from '@eggjs/tegg-types';
+import type { EggProtoImplClass, WebSocketMethodParams } from '@eggjs/tegg-types';
+import MethodInfoUtil from '../../util/MethodInfoUtil';
+import WebSocketInfoUtil from '../../util/WebSocketInfoUtil';
+
+export function WebSocketMethod(param: WebSocketMethodParams) {
+  return function(target: any, propertyKey: PropertyKey) {
+    assert(typeof propertyKey === 'string',
+      `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`);
+    const controllerClazz = target.constructor as EggProtoImplClass;
+    const methodName = propertyKey as string;
+    MethodInfoUtil.setMethodControllerType(controllerClazz, methodName, ControllerType.WEBSOCKET);
+    WebSocketInfoUtil.setWebSocketMethodPath(param.path, controllerClazz, methodName);
+    if (param.priority !== undefined) {
+      WebSocketInfoUtil.setWebSocketMethodPriority(param.priority, controllerClazz, methodName);
+    }
+    if (param.timeout !== undefined) {
+      MethodInfoUtil.setMethodTimeout(param.timeout, controllerClazz, methodName);
+    }
+  };
+}

core/controller-decorator/src/decorator/websocket/WebSocketParam.ts

@@ -0,0 +1,24 @@
+import assert from 'node:assert';
+import { WebSocketParamType } from '@eggjs/tegg-types';
+import type { EggProtoImplClass } from '@eggjs/tegg-types';
+import WebSocketInfoUtil from '../../util/WebSocketInfoUtil';
+
+export function WebSocketSocket() {
+  return function(target: any, propertyKey: PropertyKey, parameterIndex: number) {
+    assert(typeof propertyKey === 'string',
+      `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`);
+    const methodName = propertyKey as string;
+    const controllerClazz = target.constructor as EggProtoImplClass;
+    WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.SOCKET, parameterIndex, controllerClazz, methodName);
+  };
+}
+
+export function WebSocketStream() {
+  return function(target: any, propertyKey: PropertyKey, parameterIndex: number) {
+    assert(typeof propertyKey === 'string',
+      `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`);
+    const methodName = propertyKey as string;
+    const controllerClazz = target.constructor as EggProtoImplClass;
+    WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.STREAM, parameterIndex, controllerClazz, methodName);
+  };
+}

core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMetaBuilder.ts

@@ -0,0 +1,110 @@
+import assert from 'node:assert';
+import { PrototypeUtil } from '@eggjs/core-decorator';
+import { ObjectUtils } from '@eggjs/tegg-common-util';
+import { ClassUtil } from '@eggjs/tegg-metadata';
+import type { EggProtoImplClass } from '@eggjs/tegg-types';
+import { ControllerType, WebSocketFetchMethodType } from '@eggjs/tegg-types';
+import { ControllerMetaBuilderFactory } from '../../builder/ControllerMetaBuilderFactory';
+import { WebSocketFetchControllerMeta, WebSocketFetchMethodMeta } from '../../model';
+import ControllerInfoUtil from '../../util/ControllerInfoUtil';
+import { ControllerMetadataUtil } from '../../util/ControllerMetadataUtil';
+import { ControllerValidator } from '../../util/validator/ControllerValidator';
+import WebSocketInfoUtil from '../../util/WebSocketInfoUtil';
+import { WebSocketFetchControllerMethodMetaBuilder } from './WebSocketFetchControllerMethodMetaBuilder';
+
+export class WebSocketFetchControllerMetaBuilder {
+  private readonly clazz: EggProtoImplClass;
+
+  constructor(clazz: EggProtoImplClass) {
+    this.clazz = clazz;
+  }
+
+  private buildMethods() {
+    const methodNames = ObjectUtils.getProperties(this.clazz.prototype);
+    const methods: WebSocketFetchMethodMeta[] = [];
+    let connectionMethod: WebSocketFetchMethodMeta | undefined;
+    let openMethod: WebSocketFetchMethodMeta | undefined;
+    let errorMethod: WebSocketFetchMethodMeta | undefined;
+    let closeMethod: WebSocketFetchMethodMeta | undefined;
+
+    for (const methodName of methodNames) {
+      const builder = new WebSocketFetchControllerMethodMetaBuilder(this.clazz, methodName);
+      const methodMeta = builder.build();
+      if (!methodMeta) {
+        continue;
+      }
+      switch (methodMeta.methodType) {
+        case WebSocketFetchMethodType.DATA: {
+          methods.push(methodMeta);
+          break;
+        }
+        case WebSocketFetchMethodType.CONNECTION: {
+          this.assertSingleLifecycleMethod(connectionMethod, methodMeta.methodType);
+          connectionMethod = methodMeta;
+          break;
+        }
+        case WebSocketFetchMethodType.OPEN: {
+          this.assertSingleLifecycleMethod(openMethod, methodMeta.methodType);
+          openMethod = methodMeta;
+          break;
+        }
+        case WebSocketFetchMethodType.ERROR: {
+          this.assertSingleLifecycleMethod(errorMethod, methodMeta.methodType);
+          errorMethod = methodMeta;
+          break;
+        }
+        case WebSocketFetchMethodType.CLOSE: {
+          this.assertSingleLifecycleMethod(closeMethod, methodMeta.methodType);
+          closeMethod = methodMeta;
+          break;
+        }
+        default:
+          assert.fail('never arrive');
+      }
+    }
+    return { methods, connectionMethod, openMethod, errorMethod, closeMethod };
+  }
+
+  private assertSingleLifecycleMethod(methodMeta: WebSocketFetchMethodMeta | undefined, methodType: WebSocketFetchMethodType) {
+    if (!methodMeta) {
+      return;
+    }
+    const classDesc = ClassUtil.classDescription(this.clazz);
+    throw new Error(`build websocket fetch controller ${classDesc} failed: duplicate ${methodType} lifecycle method`);
+  }
+
+  build(): WebSocketFetchControllerMeta {
+    ControllerValidator.validate(this.clazz);
+    const controllerType = ControllerInfoUtil.getControllerType(this.clazz);
+    assert(controllerType === ControllerType.WEBSOCKET_FETCH, 'invalid controller type');
+    const webSocketPath = WebSocketInfoUtil.getWebSocketPath(this.clazz);
+    assert(webSocketPath, `build websocket fetch controller ${ClassUtil.classDescription(this.clazz)} failed: path is required`);
+    const middlewares = ControllerInfoUtil.getControllerMiddlewares(this.clazz);
+    const { methods, connectionMethod, openMethod, errorMethod, closeMethod } = this.buildMethods();
+    assert(methods.length === 1, 'websocket fetch controller must have exactly one @WebSocketFetchMethod');
+    const clazzName = this.clazz.name;
+    const controllerName = ControllerInfoUtil.getControllerName(this.clazz) || clazzName;
+    const property = PrototypeUtil.getProperty(this.clazz);
+    const protoName = property!.name as string;
+    const hosts = ControllerInfoUtil.getControllerHosts(this.clazz);
+    const timeout = ControllerInfoUtil.getControllerTimeout(this.clazz);
+    const metadata = new WebSocketFetchControllerMeta(
+      clazzName, protoName, controllerName, webSocketPath, middlewares, methods, hosts,
+      connectionMethod, openMethod, errorMethod, closeMethod, timeout);
+    ControllerMetadataUtil.setControllerMetadata(this.clazz, metadata);
+    for (const method of metadata.methods) {
+      const realPath = metadata.getMethodRealPath(method);
+      if (!realPath.startsWith('/')) {
+        const desc = ClassUtil.classDescription(this.clazz);
+        throw new Error(`class ${desc} method ${method.name} path ${realPath} not start with /`);
+      }
+    }
+    return metadata;
+  }
+
+  static create(clazz: EggProtoImplClass) {
+    return new WebSocketFetchControllerMetaBuilder(clazz);
+  }
+}
+
+ControllerMetaBuilderFactory.registerControllerMetaBuilder(ControllerType.WEBSOCKET_FETCH, WebSocketFetchControllerMetaBuilder.create);