DEV Community

Cover image for Turn WebSockets into Async/Await Requests (AWS WebSocket API Gateway + Lambda)
Rishi
Rishi

Posted on • Originally published at tricksumo.com

Turn WebSockets into Async/Await Requests (AWS WebSocket API Gateway + Lambda)

Some time ago, I was building a chat application using AWS Websocket API gateway. Things were going smoothly. I created a WebSocket API Gateway, added $connect, $disconnect, and sendMessage/addGroup routes. From the frontend (React) side, everything was fire-and-forget. You send a message, and the onMessageHandler takes care of it 💪🏼

But then a new requirement of uploading files using S3 signed URLs came up. That's where I needed the Async/Await promise pattern. Now, one option was to create an HTTP API gateway and use it. But that meant a new connection, a new authorizer, and more setup. At that moment, I wished there was a way to use this existing WebSocket connection to get the signed URL ⭐

wish-movie-refernce

And that’s how this library "ws-await" was born!

It lets you:

  • establish a WebSocket connection
  • send normal fire-and-forget messages
  • send messages and wait for the response using async/await
  • handle reconnection with exponential backoff
  • auto-send heartbeat messages to keep the connection alive

How does it work?

For the async/await pattern, messages from the frontend are sent with a unique requestId. The client keeps a map of pending promises. On the lambda side, the backend reads the requestId and sends it back in the response.

Received messages are analyzed for requestId, and if requestId matches the id of any pending promise in the map, that promise is resolved. If a promise sits idle in the map for a period of more than ~30 seconds, it is rejected.

Steps to use:

Step 1: Install the library in your React project

npm install @tricksumo/ws-await zustand

Step 2: Import createSocket() and establish the connection. Then call ws.send("action") for fire and forget messages and await ws.request("action") for async/await pattern.

import { createSocket } from '@tricksumo/ws-await'
import { useEffect } from 'react'
const ws = createSocket({
  url: 'wss://id.execute-api.us-east-1.amazonaws.com/prod',
})
function App() {
  useEffect(() => {
    ws.connect()
    return () => {
      ws.disconnect()
    }
  }, [])
  const handleGetSignedURL = async () => {
  try {
    const response = await ws.request('getSignedURL', { fileType: 'image/png' })
    console.log('Signed URL:', response)
  } catch (err) {
    console.error('Request failed:', err)
  }
}
  return (
      <div>
        <button onClick={handleGetSignedURL}>
          Click to get signed URL
        </button>
      </div>
  )
}
Enter fullscreen mode Exit fullscreen mode

export default App

Step 3:  Your Lambda must echo the requestId back in its response.

export const handler = async (event) => {

  const { requestId, fileName, fileType } = JSON.parse(event.body || '{}')

  const signedUrl = await getPresignedUrl(fileName, fileType)

  return {
    statusCode: 200,
    body: JSON.stringify({
      signedUrl,
      requestId, // ← REQUIRED: echo it back or the Promise never resolves
    }),
  }

}
Enter fullscreen mode Exit fullscreen mode

Reading the connection state in the frontend.

import { useSocket } from '@tricksumo/ws-await'

function StatusBar() {
  const { isConnected, isConnecting, error } = useSocket()

  if (isConnecting) return <p>Connecting...</p>
  if (!isConnected)  return <p>Disconnected  {error?.message}</p>
  return <p>Connected</p>
}
Enter fullscreen mode Exit fullscreen mode

Conclusion

Initially, this logic was part of my chat application named Chatlings. But I thought it might help others, so I extracted it to create my first ever library 🙌🏼

Top comments (0)