Make a Type-safe and Runtime-safe WebSocket Communication with Zod
As we send and handle more and more WebSocket events, it would be great to bring some type safety to the mix. But type safety along won't cut it. We need the runtime safety as well.
Luckily, zod is to the rescue. In this one, we will use zod to create a chat message schema and use it to create createChangeMessage and parseChatMessage utilities to help us develop and mock our WebSocket API.
Resources
[00:00] Type safety is essential when it comes to any API communications. But in case of our WebSocket event handling, the messages we receive and send have the Any type. This is not very useful and can lead to a lot of runtime issues. But even if we cast this message to a narrower type, like the chat message payload we have, we still won't be completely safe. Although our message variable has the correct type now, in reality anything can be sent over the WebSocket protocol in addition to this object, like blobs or array buffers.
[00:29] And if that happens, our handling will fail. That is because type safety alone is not enough to ensure a proper WebSocket communication. We also need the runtime safety. Let's achieve that by using a library called Zod in TypeScript. Zod is a schema declaration and validation library and I already have it installed in my app.
[00:47] If you don't, open the terminal and run npm install zod to have it added as a dependency. Once you have it installed, head to the app directory and the utils directory and create a new module called message utils.ts. In that module, import z from Zod. Create a new variable called chatMessageSchema and assign it to z.object. This means that our chat messages will be objects.
[01:15] Z.object utility expects us to describe the exact shape of the object we shall be sending or receiving. I'm going to copy an actual chat message object right here and proceed with annotating it using Zod. So the type property is supposed to be the literal message string. And I will use z.literal to describe that. Next, the data property is an object.
[01:35] I will wrap it in z.object. The ID of the message can be any string, so I will use z.string to annotate it. I will turn the author property into another object, Also annotating its ID property as z.string and using it for the name property and the avatar property. The text of the message is also string, I will use the same utility here, and finally for the sent at, this is a timestamp, so it's z.number. Now our schema for the chat message is ready, and it's time to create some utilities to parse the message and also to create a new message.
[02:11] Declare a new function called parseChatMessage that accepts an argument data of type WebSocket data. You can grab this type from MSW. Inside this function, let's first check if the type of the data we're receiving is not a string, and if that's the case, let's ignore that message altogether. Next, let's create a variable called message and assign it the result of parsing this data string using JSON.parse. Then, call our chat message schema and its saveParse method, providing this message as an argument.
[02:48] And store this parsed result in the variable called result. If the schema parsing has errors, let's ignore this message. But if it doesn't, let's return the result.data as is. This is going to return the incoming message we provided. And finally, let's export this function to use it both in our mocks and the actual application.
[03:07] Now go to the index route of our app and replace this JSON.parse with our new parse chat message utility. We can drop this typecasting altogether. Now the message we receive is correctly annotated based on the types of our schema, but it can also be undefined in case of unsupported data types. So let's check for that. If the message is not defined, let's ignore it.
[03:27] We can finally remove this to-do And take a look that the message type we receive is actually type safe according to our schema. Now let's make it easier to create chat messages. Head back to the message utils module and declare a new function called createChatMessage. This function will accept an input argument, which is a partial type, from z.infer, and then the type of our chat message schema. We are using the z.infer utility to convert our actual chat message schema object from here to a TypeScript type to then use as a type of the input argument.
[04:02] Now inside this function, let's call chatMessageSchema and use parse, providing input as an argument. We are using the parse method and not the saveParse method because we actually want Zot to throw here if we're providing an invalid input. It's great that our input can satisfy the schema, but we still need to send a string over the WebSocket protocol. So in that case, let's return json.stringify and provide the input as an argument. And of course, export this utility function to use later.
[04:30] Now in our root loader, let's use the new create message utility. Let's head to the place where you're sending the message from the client, right here, and replace it with create chat message, providing the same message like this, and of course dropping the satisfies anymore. Our chat object can also have a media property right now for media uploads, but we don't have it in the schema right now. Let's head back to the message utils and add a new media property, which for now will be z.any. In the same way, let's head to handlers.ts and use our new create chat message utility here when we're recreating conversation between two mocked users.
[05:08] So create chat message providing this exact message as an argument. And the same for Emily. With these changes, we are using Zot to achieve the ultimate safety when handling our WebSocket messages. That includes our mocks right here, but also the actual chat implementation as well. And the cherry on top is that Zort provides us both type and runtime safety.