fix(typescript): Bugfixes to Typescript Websocket code generation (#6350

) * fix(typescript): Several bugfixes to generated `Socket.py` websocket file. * Fixes * US
fern-api · Mar 10, 2025 · 4325061 · 4325061
1 parent 75f3f7f
commit 4325061
Show file tree

Hide file tree

Showing 5 changed files with 77 additions and 61 deletions.
diff --git a/generators/typescript/sdk/CHANGELOG.md b/generators/typescript/sdk/CHANGELOG.md
@@ -5,6 +5,10 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.49.1] - 2025-03-10
+
+- Fix: This PR includes several fixes to the generated `Socket.py` file when websocket client code generation is enabled.
+
 ## [0.49.0] - 2025-03-06
 
 - Feat: This PR enables the Typescript generator to produce Websocket SDK endpoints. This can be enabled by adding the option `shouldGenerateWebsocketClients: true` to the Typescript generator config.

diff --git a/generators/typescript/sdk/VERSION b/generators/typescript/sdk/VERSION
@@ -1 +1 @@
-0.49.0
+0.49.1
diff --git a/generators/typescript/sdk/client-class-generator/src/GeneratedWebsocketSocketClassImpl.ts b/generators/typescript/sdk/client-class-generator/src/GeneratedWebsocketSocketClassImpl.ts
@@ -5,6 +5,7 @@ import {
     InterfaceDeclarationStructure,
     MethodDeclarationStructure,
     ModuleDeclarationStructure,
+    PropertyDeclarationStructure,
     Scope,
     StructureKind,
     TypeAliasDeclarationStructure,
@@ -126,13 +127,11 @@ export class GeneratedWebsocketSocketClassImpl implements GeneratedWebsocketSock
             closeMethod,
             tillSocketOpen,
             assertSocketIsOpen,
-            sendJson,
-            handleOpen,
-            handleMessage,
-            handleClose,
-            handleError
+            sendJson
         );
 
+        serviceClass.properties?.push(handleOpen, handleMessage, handleClose, handleError);
+
         context.sourceFile.addModule(serviceModule);
         context.sourceFile.addClass(serviceClass);
     }
@@ -158,7 +157,7 @@ export class GeneratedWebsocketSocketClassImpl implements GeneratedWebsocketSock
             isExported: true,
             type: getTextOfTsNode(
                 ts.factory.createIntersectionTypeNode([
-                    this.getPublishMessageNode(context),
+                    this.getSubscribeMessageNode(context),
                     ts.factory.createTypeLiteralNode([
                         ts.factory.createPropertySignature(
                             undefined,
@@ -419,77 +418,76 @@ export class GeneratedWebsocketSocketClassImpl implements GeneratedWebsocketSock
         };
     }
 
-    private generateHandleMessage(context: SdkContext): MethodDeclarationStructure {
+    private generateHandleMessage(context: SdkContext): PropertyDeclarationStructure {
         const subscribeMessage = this.getSubscribeMessage();
+
+        const bodyLines: string[] = ["const data = JSON.parse(event.data);", ""];
+
+        if (this.includeSerdeLayer) {
+            const parsedResponseStatement = `const parsedResponse = ${getTextOfTsNode(
+                this.getParsedExpression(subscribeMessage.body, context)
+            )};`;
+            bodyLines.push(
+                parsedResponseStatement,
+                "if (parsedResponse.ok) {",
+                `    this.eventHandlers.${GeneratedWebsocketSocketClassImpl.MESSAGE_PARAMETER_NAME}?.({`,
+                "        ...parsedResponse.value,",
+                "        receivedAt: new Date()",
+                "    });",
+                "} else {",
+                "    this.eventHandlers.error?.(new Error(`Received unknown message type`));",
+                "}"
+            );
+        } else {
+            bodyLines.push("this.eventHandlers.message?.({", "    ...data,", "    receivedAt: new Date()", "});");
+        }
+
         return {
-            kind: StructureKind.Method,
+            kind: StructureKind.Property,
             name: "handleMessage",
             scope: Scope.Private,
-            parameters: [
-                {
-                    name: "event",
-                    type: "{ data: string }"
-                }
-            ],
-            statements: [
-                "const data = JSON.parse(event.data);",
-                "",
-                ...(this.includeSerdeLayer
-                    ? [
-                          `const parsedResponse = ${getTextOfTsNode(this.getParsedExpression(subscribeMessage.body, context))};`,
-                          "if (parsedResponse.ok) {",
-                          `    this.eventHandlers.${GeneratedWebsocketSocketClassImpl.MESSAGE_PARAMETER_NAME}?.({`,
-                          "        ...parsedResponse.value,",
-                          "        receivedAt: new Date()",
-                          "    });",
-                          "} else {",
-                          "    this.eventHandlers.error?.(new Error(`Received unknown message type`));",
-                          "}"
-                      ]
-                    : ["this.eventHandlers.message?.({", "    ...data,", "    receivedAt: new Date()", "});"])
-            ]
+            type: "(event: { data: string }) => void",
+            initializer: `event => {
+    ${bodyLines.map((line) => "    " + line).join("\n")}
+    }`
         };
     }
 
-    private generateHandleOpen(): MethodDeclarationStructure {
+    private generateHandleOpen(): PropertyDeclarationStructure {
         return {
-            kind: StructureKind.Method,
+            kind: StructureKind.Property,
             name: "handleOpen",
             scope: Scope.Private,
-            statements: [`this.${GeneratedWebsocketSocketClassImpl.EVENT_HANDLERS_PROPERTY_NAME}.open?.();`]
+            type: "() => void",
+            initializer: `() => {
+        this.${GeneratedWebsocketSocketClassImpl.EVENT_HANDLERS_PROPERTY_NAME}.open?.();
+    }`
         };
     }
 
-    private generateHandleClose(context: SdkContext): MethodDeclarationStructure {
+    private generateHandleClose(context: SdkContext): PropertyDeclarationStructure {
         return {
-            kind: StructureKind.Method,
+            kind: StructureKind.Property,
             name: "handleClose",
             scope: Scope.Private,
-            parameters: [
-                {
-                    name: "event",
-                    type: getTextOfTsNode(context.coreUtilities.websocket.CloseEvent._getReferenceToType())
-                }
-            ],
-            statements: [`this.${GeneratedWebsocketSocketClassImpl.EVENT_HANDLERS_PROPERTY_NAME}.close?.(event);`]
+            // Adjust the parameter type to match your CloseEvent shape
+            type: `(event: ${getTextOfTsNode(context.coreUtilities.websocket.CloseEvent._getReferenceToType())}) => void`,
+            initializer: `event => {
+        this.${GeneratedWebsocketSocketClassImpl.EVENT_HANDLERS_PROPERTY_NAME}.close?.(event);
+    }`
         };
     }
 
-    private generateHandleError(context: SdkContext): MethodDeclarationStructure {
+    private generateHandleError(context: SdkContext): PropertyDeclarationStructure {
         return {
-            kind: StructureKind.Method,
+            kind: StructureKind.Property,
             name: "handleError",
             scope: Scope.Private,
-            parameters: [
-                {
-                    name: "event",
-                    type: getTextOfTsNode(context.coreUtilities.websocket.ErrorEvent._getReferenceToType())
-                }
-            ],
-            statements: [
-                `const ${GeneratedWebsocketSocketClassImpl.MESSAGE_PARAMETER_NAME} = event.message ?? "core.ReconnectingWebSocket error";`,
-                `this.${GeneratedWebsocketSocketClassImpl.EVENT_HANDLERS_PROPERTY_NAME}.error?.(new Error(${GeneratedWebsocketSocketClassImpl.MESSAGE_PARAMETER_NAME}));`
-            ]
+            type: `(event: ${getTextOfTsNode(context.coreUtilities.websocket.ErrorEvent._getReferenceToType())}) => void`,
+            initializer: `event => {
+        const ${GeneratedWebsocketSocketClassImpl.MESSAGE_PARAMETER_NAME} = event.message ?? "core.ReconnectingWebSocket error";
+        this.${GeneratedWebsocketSocketClassImpl.EVENT_HANDLERS_PROPERTY_NAME}.error?.(new Error(${GeneratedWebsocketSocketClassImpl.MESSAGE_PARAMETER_NAME}));
+    }`
         };
     }
 
@@ -540,11 +538,11 @@ export class GeneratedWebsocketSocketClassImpl implements GeneratedWebsocketSock
     }
 
     private getPublishMessage(): WebSocketMessage {
-        return this.channel.messages.filter((message) => message.origin === "server")[0] as WebSocketMessage;
+        return this.channel.messages.filter((message) => message.origin === "client")[0] as WebSocketMessage;
     }
 
     private getSubscribeMessage(): WebSocketMessage {
-        return this.channel.messages.filter((message) => message.origin === "client")[0] as WebSocketMessage;
+        return this.channel.messages.filter((message) => message.origin === "server")[0] as WebSocketMessage;
     }
 
     private getPublishMessageNode(context: SdkContext): ts.TypeNode {
@@ -561,4 +559,19 @@ export class GeneratedWebsocketSocketClassImpl implements GeneratedWebsocketSock
                 return ts.factory.createTypeReferenceNode(getTextOfTsNode(generatedType.typeNode), undefined);
             })[0] as ts.TypeNode;
     }
+
+    private getSubscribeMessageNode(context: SdkContext): ts.TypeNode {
+        // TODO (Eden): At the moment, we're only extracting two messages in the IR: publish & subscribe.
+        // We'll need to update this if we want to support a different message array structure.
+        return this.channel.messages
+            .filter((message) => message.origin === "server")
+            .map((message) => {
+                if (message.body.type === "inlinedBody") {
+                    // TODO (Eden): Handle inlined body messages
+                    throw new Error("Inlined body messages are not supported yet");
+                }
+                const generatedType = context.type.getReferenceToType(message.body.bodyType);
+                return ts.factory.createTypeReferenceNode(getTextOfTsNode(generatedType.typeNode), undefined);
+            })[0] as ts.TypeNode;
+    }
 }
diff --git a/generators/typescript/utils/core-utilities/fetcher/src/websocket/ws.ts b/generators/typescript/utils/core-utilities/fetcher/src/websocket/ws.ts
@@ -381,8 +381,8 @@ export class ReconnectingWebSocket {
                 }
                 this._debug("connect", { url, protocols: this._protocols });
                 this._ws = this._protocols
-                    ? new WebSocket(url, this._protocols, this._headers)
-                    : new WebSocket(url, undefined, this._headers);
+                    ? new WebSocket(url, this._protocols, { headers: this._headers })
+                    : new WebSocket(url, undefined, { headers: this._headers });
                 this._ws!.binaryType = this._binaryType;
                 this._connectLock = false;
                 this._addListeners();

diff --git a/seed/ts-express/unions/.mock/ir.json b/seed/ts-express/unions/.mock/ir.json