AI Gateway Websockets (Beta): Example results in errors

[nikothomas]

Existing documentation URL(s)

• https://developers.cloudflare.com/ai-gateway/websockets-api/non-realtime-api/

What changes are you suggesting?

When running the existing example. I ran into two problems:

1. The WebSocket sends messages before connecting: The example calls ws.send() immediately after creating the WebSocket, but this throws an error because the connection hasn't been established yet.

2. Authentication keeps failing: Even when I fixed the connection issue, I kept getting authentication errors. I discovered that the cf-aig-authorization header needs to be included in both the WebSocket connection AND in the request payload headers, but the example only shows it in the connection headers.

Here are the errors I got:

Connection error:

Error: WebSocket is not open: readyState 0 (CONNECTING)
    at WebSocket.send (/workspaces/sh-backend/sh-cloudflare-ai-gateway/node_modules/ws/lib/websocket.js:450:13)
    at file:///workspaces/sh-backend/sh-cloudflare-ai-gateway/example.js:12:4

Authentication error:

{
  "type": "universal.created",
  "metadata": {
    "contentType": "application/json"
  },
  "response": {
    "result": null,
    "success": false,
    "errors": [
      {
        "code": 10000,
        "message": "Authentication error"
      }
    ],
    "messages": []
  }
}

Additional information

Here's a working example:

import WebSocket from "ws";

const ws = new WebSocket(
  "wss://gateway.ai.cloudflare.com/v1/my-account-id/my-gateway/",
  {
    headers: {
      "cf-aig-authorization": "Bearer AI_GATEWAY_TOKEN",
    },
  },
);

// Wait for connection before sending
ws.on("open", function open() {
  console.log("WebSocket connected");
  
  ws.send(
    JSON.stringify({
      type: "universal.create",
      request: {
        eventId: "my-request",
        provider: "workers-ai",
        endpoint: "@cf/meta/llama-3.1-8b-instruct",
        headers: {
          Authorization: "Bearer WORKERS_AI_TOKEN",
          "Content-Type": "application/json",
          "cf-aig-authorization": "Bearer AI_GATEWAY_TOKEN", // This was missing
        },
        query: {
          prompt: "tell me a joke",
        },
      },
    }),
  );
});

// Add error handling
ws.on("error", function error(err) {
  console.error("WebSocket error:", err);
});

ws.on("message", function incoming(message) {
  console.log(message.toString());
});

The key fixes:

1. Wrap ws.send() in the open event handler
2. Add cf-aig-authorization to the request payload headers
3. Add error handling for better debugging