WebSocket と永続オブジェクトを使用した大規模なリアルタイム AI 推論

2024 年 10 月に、AI Gateway を使用して AI アプリケーションから数十億のログを保存することと、これを行うために Cloudflare の開発者プラットフォームを使用する方法について説明しました。

In October 2024, we covered how to store billions of logs from your AI application using AI Gateway, and how we used Cloudflare’s Developer Platform to do this.

2024 年 10 月に、AI Gateway を使用して AI アプリケーションから数十億のログを保存する方法と、これを行うために Cloudflare の開発者プラットフォームを使用する方法について説明しました。

With AI Gateway already processing over 3 billion logs and experiencing rapid growth, the number of connections to the platform continues to increase steadily. To help developers manage this scale more effectively, we wanted to offer an alternative to implementing HTTP/2 keep-alive to maintain persistent HTTP(S) connections, thereby avoiding the overhead of repeated handshakes and TLS negotiations with each new HTTP connection to AI Gateway. We understand that implementing HTTP/2 can present challenges, particularly when many libraries and tools may not support it by default and most modern programming languages have well-established WebSocket libraries available.

AI ゲートウェイはすでに 30 億を超えるログを処理し、急速に成長しており、プラットフォームへの接続数は着実に増加し続けています。開発者がこの規模をより効果的に管理できるように、HTTP/2 キープアライブを実装して永続的な HTTP(S) 接続を維持する代替手段を提供したいと考えました。これにより、AI Gateway への新しい HTTP 接続ごとに繰り返されるハンドシェイクと TLS ネゴシエーションのオーバーヘッドが回避されます。 。 HTTP/2 の実装には、特に多くのライブラリやツールがデフォルトでサポートしていない可能性があり、最新のプログラミング言語のほとんどが確立された WebSocket ライブラリを利用できる場合に、課題が生じる可能性があることを私たちは理解しています。

With this in mind, we used Cloudflare’s Developer Platform and Durable Objects (yes, again!) to build a WebSockets API that establishes a single, persistent connection, enabling continuous communication.

これを念頭に置いて、私たちはCloudflareの開発者プラットフォームとDurable Objects（はい、また！）を使用して、単一の永続的な接続を確立し、継続的な通信を可能にするWebSocket APIを構築しました。

Through this API, all AI providers supported by AI Gateway can be accessed via WebSocket, allowing you to maintain a single TCP connection between your client or server application and the AI Gateway. The best part? Even if your chosen provider doesn’t support WebSockets, we handle it for you, managing the requests to your preferred AI provider.

この API を通じて、AI Gateway でサポートされるすべての AI プロバイダーに WebSocket 経由でアクセスできるため、クライアントまたはサーバー アプリケーションと AI Gateway の間で単一の TCP 接続を維持できます。一番いいところは？選択したプロバイダーが WebSocket をサポートしていない場合でも、当社が WebSocket を処理し、好みの AI プロバイダーへのリクエストを管理します。

By connecting via WebSocket to AI Gateway, we make the requests to the inference service for you using the provider’s supported protocols (HTTPS, WebSocket, etc.), and you can keep the connection open to execute as many inference requests as you would like.

WebSocket 経由で AI Gateway に接続すると、プロバイダーがサポートするプロトコル (HTTPS、WebSocket など) を使用して推論サービスにリクエストが行われます。接続を開いたままにして、必要なだけ推論リクエストを実行できます。

To make your connection to AI Gateway more secure, we are also introducing authentication for AI Gateway. The new WebSockets API will require authentication. All you need to do is create a Cloudflare API token with the permission “AI Gateway: Run” and send that in the cf-aig-authorization header.

AI Gateway への接続をより安全にするために、AI Gateway の認証も導入します。新しい WebSocket API には認証が必要です。必要なのは、「AI Gateway: Run」権限を持つ Cloudflare API トークンを作成し、それを cf-aig-authorization ヘッダーで送信することだけです。

In the flow diagram above:

上のフロー図では次のようになります。

1. When Authenticated Gateway is enabled and a valid token is included, requests will pass successfully.

1. 認証されたゲートウェイが有効で、有効なトークンが含まれている場合、リクエストは正常に通過します。

2. If Authenticated Gateway is enabled, but a request does not contain the required cf-aig-authorization header with a valid token, the request will fail. This ensures only verified requests pass through the gateway.

2. 認証されたゲートウェイが有効になっていても、有効なトークンを含む必要な cf-aig-authorization ヘッダーがリクエストに含まれていない場合、リクエストは失敗します。これにより、検証されたリクエストのみがゲートウェイを通過することが保証されます。

3. When Authenticated Gateway is disabled, the cf-aig-authorization header is bypassed entirely, and any token — whether valid or invalid — is ignored.

3. Authenticated Gateway が無効になっている場合、cf-aig-authorization ヘッダーは完全にバイパスされ、有効か無効かにかかわらず、トークンは無視されます。

How we built it

どのように構築したか

We recently used Durable Objects (DOs) to scale our logging solution for AI Gateway, so using WebSockets within the same DOs was a natural fit.

私たちは最近、Durable Objects (DO) を使用して AI Gateway のロギング ソリューションを拡張したため、同じ DO 内で WebSocket を使用するのが自然でした。

When a new WebSocket connection is received by our Cloudflare Workers, we implement authentication in two ways to support the diverse capabilities of WebSocket clients. The primary method involves validating a Cloudflare API token through the cf-aig-authorization header, ensuring the token is valid for the connecting account and gateway.

Cloudflareワーカーが新しいWebSocket接続を受信すると、WebSocketクライアントの多様な機能をサポートするために2つの方法で認証を実装します。主な方法には、cf-aig-authorization ヘッダーを介して Cloudflare API トークンを検証し、接続するアカウントとゲートウェイに対してトークンが有効であることを確認することが含まれます。

However, due to limitations in browser WebSocket implementations, we also support authentication via the “sec-websocket-protocol” header. Browser WebSocket clients don't allow for custom headers in their standard API, complicating the addition of authentication tokens in requests. While we don’t recommend that you store API keys in a browser, we decided to add this method to add more flexibility to all WebSocket clients.

ただし、ブラウザの WebSocket 実装には制限があるため、「sec-websocket-protocol」ヘッダーを介した認証もサポートされています。ブラウザ WebSocket クライアントは標準 API でカスタム ヘッダーを許可しないため、リクエストへの認証トークンの追加が複雑になります。 API キーをブラウザーに保存することはお勧めしませんが、すべての WebSocket クライアントの柔軟性を高めるためにこの方法を追加することにしました。

After this initial verification step, we upgrade the connection to the Durable Object, meaning that it will now handle all the messages for the connection. Before the new connection is fully accepted, we generate a random UUID, so this connection is identifiable among all the messages received by the Durable Object. During an open connection, any AI Gateway settings passed via headers — such as cf-aig-skip-cache (which bypasses caching when set to true) — are stored and applied to all requests in the session. However, these headers can still be overridden on a per-request basis, just like with the Universal Endpoint today.

この最初の検証ステップの後、接続を永続オブジェクトにアップグレードします。つまり、接続のすべてのメッセージが処理されるようになります。新しい接続が完全に受け入れられる前に、ランダムな UUID が生成されるため、Durable Object が受信するすべてのメッセージの中でこの接続を識別できるようになります。接続が開いている間、ヘッダー経由で渡された AI Gateway 設定 (cf-aig-skip-cache (true に設定するとキャッシュをバイパスする) など) が保存され、セッション内のすべてのリクエストに適用されます。ただし、これらのヘッダーは、現在のユニバーサル エンドポイントと同様に、リクエストごとにオーバーライドできます。

How it works

仕組み

Once the connection is established, the Durable Object begins listening for incoming messages. From this point on, users can send messages in the AI Gateway universal format via WebSocket, simplifying the transition of your application from an existing HTTP setup to WebSockets-based communication.

接続が確立されると、Durable Object は受信メッセージの待機を開始します。この時点から、ユーザーは WebSocket 経由で AI Gateway ユニバーサル形式でメッセージを送信できるようになり、既存の HTTP セットアップから WebSocket ベースの通信へのアプリケーションの移行が簡素化されます。

When a new message reaches the Durable Object, it’s processed using the same code that powers the HTTP Universal Endpoint, enabling seamless code reuse across Workers and Durable Objects — one of the key benefits of building on Cloudflare.

新しいメッセージがDurable Objectに到達すると、HTTP Universal Endpointを動作させるのと同じコードを使用して処理され、ワー​​カーとDurable Object間でのシームレスなコードの再利用が可能になります。これは、Cloudflare上に構築することの主な利点の1つです。

For non-streaming requests, the response is wrapped in a JSON envelope, allowing us to include additional information beyond the AI inference itself, such as the AI Gateway log ID for that request.

非ストリーミング リクエストの場合、応答は JSON エンベロープでラップされ、そのリクエストの AI ゲートウェイ ログ ID など、AI 推論自体を超える追加情報を含めることができます。

Here’s an example response for the request above:

上記のリクエストに対する応答の例を次に示します。

For streaming requests, AI Gateway sends an initial message with request metadata telling the developer the stream is starting.

ストリーミング リクエストの場合、AI Gateway は、開発者にストリームが開始していることを伝えるリクエスト メタデータを含む初期メッセージを送信します。

After this initial message, all streaming chunks are relayed in real-time to the WebSocket connection as they arrive from the inference provider. Note that only the eventId field is included in the metadata for these streaming chunks (more info on what this new field is below).

この最初のメッセージの後、すべてのストリーミング チャンクは、推論プロバイダーから到着すると、リアルタイムで WebSocket 接続に中継されます。これらのストリーミング チャンクのメタデータには、eventId フィールドのみが含まれることに注意してください (この新しいフィールドの内容については、以下で詳しく説明します)。

This approach serves two purposes:

このアプローチには次の 2 つの目的があります。