supabase-edge-function-continuous-stream 1.0.0

package/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Huy Pham
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

package/README.md

@@ -0,0 +1,54 @@
+# supabase-edge-function-continuous-stream
+
+Continuous WebSocket streaming for Supabase edge functions. Keeps a persistent connection alive across cold starts, supports warmup/context loading, and streams AI responses with automatic retry and worker TTL rotation.
+
+## Install
+
+```bash
+npm install supabase-edge-function-continuous-stream
+```
+
+## Entries
+
+| Import | React required? | Use for |
+|---|---|---|
+| `supabase-edge-function-continuous-stream` | No | `connectEdgeSocket`, `createStandardAiMessageHandler`, types |
+| `supabase-edge-function-continuous-stream/react` | Yes | `createUseEdgeStream` hook factory |
+
+## React usage
+
+Prefer the main entry (works with Next.js webpack and published npm installs):
+
+```ts
+import { createUseEdgeStream } from "supabase-edge-function-continuous-stream";
+```
+
+For explicit React-only imports:
+
+```ts
+import { createUseEdgeStream } from "supabase-edge-function-continuous-stream/react";
+```
+
+```ts
+import { DEFAULT_EDGE_WORKER_LIMITS } from "supabase-edge-function-continuous-stream";
+
+export const useEdgeStream = createUseEdgeStream({
+  getAccessToken: async () => {
+    const { data } = await supabase.auth.getSession();
+    if (!data.session?.access_token) throw new Error("Not authenticated");
+    return data.session.access_token;
+  },
+  getSupabaseUrl: () => process.env.NEXT_PUBLIC_SUPABASE_URL!,
+  workerLimits: DEFAULT_EDGE_WORKER_LIMITS,
+});
+```
+
+## Wire protocol
+
+1. Client opens WebSocket with JWT query param
+2. `client_warmup` → server responds `status: ready`
+3. `client_message` → streamed events → `complete`
+
+## License
+
+MIT

package/dist/chunk-GP2SEFBH.js

@@ -0,0 +1,462 @@
+// src/constants.ts
+var MAX_RETRIES = 200;
+var INITIAL_RETRY_DELAY_MS = 1e3;
+var CONNECTION_TIMEOUT_MS = 1e4;
+
+// src/connection.ts
+async function connectEdgeSocket(deps) {
+  const {
+    functionPath,
+    getAccessToken,
+    getSupabaseUrl,
+    wsRef,
+    isExplicitDisconnectRef,
+    isConnectingRef,
+    retryCountRef,
+    socketOpenedAtRef,
+    isWarmupReadyRef,
+    warmupWaitersRef,
+    passiveOnServerActionRef,
+    activeRequestRef,
+    lastWarmupPayloadRef,
+    setIsConnected,
+    closeSocket,
+    sendWarmupPayload,
+    settleWarmupWaiters
+  } = deps;
+  if (wsRef.current?.readyState === WebSocket.OPEN) return Promise.resolve();
+  if (isConnectingRef.current) {
+    return new Promise((resolve, reject) => {
+      const checkInterval = setInterval(() => {
+        if (wsRef.current?.readyState === WebSocket.OPEN) {
+          clearInterval(checkInterval);
+          resolve();
+        } else if (wsRef.current?.readyState === WebSocket.CLOSED) {
+          clearInterval(checkInterval);
+          reject(new Error("Connection closed during check"));
+        }
+      }, 100);
+    });
+  }
+  isConnectingRef.current = true;
+  isExplicitDisconnectRef.current = false;
+  try {
+    const accessToken = await getAccessToken();
+    const supabaseUrl = getSupabaseUrl();
+    if (!supabaseUrl) throw new Error("Missing Supabase URL");
+    const wsUrl = supabaseUrl.replace(
+      /^https?:\/\//,
+      (m) => m === "https://" ? "wss://" : "ws://"
+    );
+    const functionUrl = `${wsUrl}/functions/v1/${functionPath}`;
+    return await new Promise((resolveConnection, rejectConnection) => {
+      let connectionTimeout = null;
+      if (wsRef.current) {
+        try {
+          wsRef.current.onopen = null;
+          wsRef.current.onmessage = null;
+          wsRef.current.onerror = null;
+          wsRef.current.onclose = null;
+          if (wsRef.current.readyState === WebSocket.OPEN || wsRef.current.readyState === WebSocket.CONNECTING) {
+            wsRef.current.close();
+          }
+        } catch {
+        }
+      }
+      const urlWithToken = new URL(functionUrl);
+      urlWithToken.searchParams.set("jwt", accessToken);
+      wsRef.current = new WebSocket(urlWithToken.toString());
+      if (retryCountRef.current === 0) {
+        connectionTimeout = setTimeout(() => {
+          if (wsRef.current && wsRef.current.readyState !== WebSocket.OPEN) {
+            isExplicitDisconnectRef.current = false;
+            closeSocket();
+          }
+        }, CONNECTION_TIMEOUT_MS);
+      }
+      wsRef.current.onopen = () => {
+        if (connectionTimeout) clearTimeout(connectionTimeout);
+        retryCountRef.current = 0;
+        isConnectingRef.current = false;
+        setIsConnected(true);
+        socketOpenedAtRef.current = Date.now();
+        if (lastWarmupPayloadRef.current) {
+          sendWarmupPayload();
+        }
+        resolveConnection();
+      };
+      wsRef.current.onerror = () => {
+      };
+      const retryConnect = () => {
+        connectEdgeSocket(deps).then(resolveConnection).catch(rejectConnection);
+      };
+      wsRef.current.onclose = () => {
+        if (connectionTimeout) clearTimeout(connectionTimeout);
+        isConnectingRef.current = false;
+        setIsConnected(false);
+        isWarmupReadyRef.current = false;
+        const willRetry = !isExplicitDisconnectRef.current && retryCountRef.current < MAX_RETRIES;
+        if (!willRetry) {
+          settleWarmupWaiters(new Error("WebSocket closed"));
+        }
+        if (isExplicitDisconnectRef.current) {
+          rejectConnection(
+            new Error("WebSocket closed by server or aborted")
+          );
+        } else if (retryCountRef.current < MAX_RETRIES) {
+          retryCountRef.current++;
+          const delay = INITIAL_RETRY_DELAY_MS * Math.pow(2, retryCountRef.current - 1);
+          setTimeout(retryConnect, delay);
+        } else {
+          rejectConnection(
+            new Error(
+              `WebSocket connection failed after ${retryCountRef.current} retries`
+            )
+          );
+        }
+      };
+      wsRef.current.onmessage = (event) => {
+        try {
+          const message = JSON.parse(event.data);
+          if (message.type === "status" && message.data === "ready") {
+            isWarmupReadyRef.current = true;
+            settleWarmupWaiters();
+          }
+          if (message.type === "error" && warmupWaitersRef.current.length > 0 && !isWarmupReadyRef.current) {
+            isWarmupReadyRef.current = false;
+            settleWarmupWaiters(
+              new Error(String(message.data ?? "Warmup failed"))
+            );
+          }
+          const passive = passiveOnServerActionRef.current;
+          if (passive && !activeRequestRef.current) {
+            const isWarmupControl = message.type === "status" && (message.data === "ready" || message.data === "context");
+            if (!isWarmupControl) {
+              passive(message.type, message.data);
+            }
+          }
+          if (!activeRequestRef.current) return;
+          const { ctx, handler } = activeRequestRef.current;
+          handler(message, ctx);
+        } catch (error) {
+          if (!activeRequestRef.current) return;
+          const ctx = activeRequestRef.current.ctx;
+          ctx.reject(
+            new Error(
+              `Failed to parse WebSocket message: ${error instanceof Error ? error.message : String(error)}`
+            )
+          );
+        }
+      };
+    });
+  } catch (err) {
+    isConnectingRef.current = false;
+    throw err;
+  }
+}
+
+// src/handler.ts
+function createStandardAiMessageHandler(defaults = {}, callbacks) {
+  const state = {};
+  const buildResponse = (overrides) => ({
+    ...defaults,
+    ...state,
+    ...overrides
+  });
+  return (message, ctx) => {
+    if (message.type === "response_text" && typeof message.data === "string") {
+      state.reply = message.data;
+    }
+    if (typeof message.data === "object" && message.data !== null && message.type !== "error") {
+      Object.assign(state, message.data);
+    }
+    if (message.type === "error") {
+      const raw = message.data ?? "Server error";
+      const userMsg = callbacks?.toUserMessage?.(raw);
+      if (userMsg) callbacks?.onServerAction?.("error", userMsg);
+      ctx.reject(new Error(raw));
+      return;
+    }
+    if (message.type === "complete") {
+      ctx.clearOverallTimeout();
+      callbacks?.onServerAction?.("complete", message.data);
+      if (!ctx.isResolved()) {
+        ctx.resolve(buildResponse(message.data));
+      }
+      return;
+    }
+    callbacks?.onServerAction?.(message.type, message.data);
+  };
+}
+
+// src/createUseEdgeStream.ts
+import { useCallback, useEffect, useRef, useState } from "react";
+function createUseEdgeStream(deps) {
+  const {
+    getAccessToken,
+    getSupabaseUrl,
+    workerLimits,
+    toUserMessage,
+    invalidateTags
+  } = deps;
+  const { edgeWorkerTtlMs, edgeRotateThresholdMs, overallTimeoutMs } = workerLimits;
+  return function useEdgeStream(config) {
+    const [isLoading, setIsLoading] = useState(false);
+    const [error, setError] = useState(null);
+    const [data, setData] = useState(void 0);
+    const wsRef = useRef(null);
+    const isExplicitDisconnectRef = useRef(false);
+    const [isConnected, setIsConnected] = useState(false);
+    const isConnectingRef = useRef(false);
+    const lastWarmupPayloadRef = useRef(null);
+    const isWarmupReadyRef = useRef(false);
+    const warmupWaitersRef = useRef([]);
+    const retryCountRef = useRef(0);
+    const passiveOnServerActionRef = useRef(null);
+    const socketOpenedAtRef = useRef(null);
+    const settleWarmupWaiters = useCallback((warmupError) => {
+      const waiters = warmupWaitersRef.current;
+      warmupWaitersRef.current = [];
+      if (warmupError) {
+        waiters.forEach((w) => w.reject(warmupError));
+        return;
+      }
+      waiters.forEach((w) => w.resolve());
+    }, []);
+    const waitForWarmupReady = useCallback(() => {
+      if (isWarmupReadyRef.current) return Promise.resolve();
+      return new Promise((resolve, reject) => {
+        warmupWaitersRef.current.push({ resolve, reject });
+      });
+    }, []);
+    const sendWarmupPayload = useCallback(() => {
+      if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
+        return;
+      }
+      if (!lastWarmupPayloadRef.current) return;
+      isWarmupReadyRef.current = false;
+      wsRef.current.send(
+        JSON.stringify({
+          type: "client_warmup",
+          data: lastWarmupPayloadRef.current
+        })
+      );
+    }, []);
+    const activeRequestRef = useRef(null);
+    useEffect(() => {
+      return () => {
+        isExplicitDisconnectRef.current = true;
+        if (wsRef.current) {
+          try {
+            wsRef.current.close();
+          } catch {
+          }
+        }
+      };
+    }, []);
+    const closeSocket = useCallback(() => {
+      if (!wsRef.current) return;
+      try {
+        wsRef.current.close();
+      } catch {
+      }
+      socketOpenedAtRef.current = null;
+    }, []);
+    const abort = useCallback(() => {
+      isExplicitDisconnectRef.current = true;
+      closeSocket();
+      setIsLoading(false);
+      if (activeRequestRef.current) {
+        activeRequestRef.current.reject(new Error("Request aborted"));
+        if (activeRequestRef.current.overallTimeout) {
+          clearTimeout(activeRequestRef.current.overallTimeout);
+        }
+        activeRequestRef.current = null;
+      }
+    }, [closeSocket]);
+    const connectWebSocket = useCallback(
+      () => connectEdgeSocket({
+        functionPath: config.functionPath,
+        getAccessToken,
+        getSupabaseUrl,
+        wsRef,
+        isExplicitDisconnectRef,
+        isConnectingRef,
+        retryCountRef,
+        socketOpenedAtRef,
+        isWarmupReadyRef,
+        warmupWaitersRef,
+        passiveOnServerActionRef,
+        activeRequestRef,
+        lastWarmupPayloadRef,
+        setIsConnected,
+        closeSocket,
+        sendWarmupPayload,
+        settleWarmupWaiters
+      }),
+      [
+        config.functionPath,
+        closeSocket,
+        sendWarmupPayload,
+        settleWarmupWaiters
+      ]
+    );
+    const rotateConnectionIfNeeded = useCallback(
+      async (getExpiresAt) => {
+        const fallbackExpires = socketOpenedAtRef.current ? socketOpenedAtRef.current + edgeWorkerTtlMs : null;
+        const expiresAt = getExpiresAt?.() ?? fallbackExpires;
+        if (!expiresAt) return;
+        if (expiresAt - Date.now() >= edgeRotateThresholdMs) return;
+        isExplicitDisconnectRef.current = false;
+        closeSocket();
+        isWarmupReadyRef.current = false;
+        await connectWebSocket();
+      },
+      [closeSocket, connectWebSocket, edgeRotateThresholdMs, edgeWorkerTtlMs]
+    );
+    const warmup = useCallback(
+      async (payload, options) => {
+        if (options?.onServerAction) {
+          passiveOnServerActionRef.current = options.onServerAction;
+        }
+        lastWarmupPayloadRef.current = payload;
+        isWarmupReadyRef.current = false;
+        const wasOpen = wsRef.current?.readyState === WebSocket.OPEN;
+        await connectWebSocket();
+        if (!isWarmupReadyRef.current) {
+          if (wasOpen) sendWarmupPayload();
+          await waitForWarmupReady();
+        }
+      },
+      [connectWebSocket, sendWarmupPayload, waitForWarmupReady]
+    );
+    const send = useCallback(
+      async (payload, options) => {
+        if (!isLoading) {
+          await rotateConnectionIfNeeded(options.getWorkerExpiresAt);
+        }
+        setIsLoading(true);
+        setError(null);
+        setData(void 0);
+        isExplicitDisconnectRef.current = false;
+        let resolved = false;
+        const ctx = {
+          resolve: (value) => {
+            resolved = true;
+            setData(value);
+            setIsLoading(false);
+            if (options.invalidateTags && invalidateTags) {
+              const shouldInvalidate = options.invalidateTags.condition ? options.invalidateTags.condition(value) : true;
+              if (shouldInvalidate) {
+                invalidateTags(options.invalidateTags.tags);
+              }
+            }
+            if (activeRequestRef.current?.overallTimeout) {
+              clearTimeout(activeRequestRef.current.overallTimeout);
+            }
+            if (activeRequestRef.current?.resolve) {
+              activeRequestRef.current.resolve(value);
+            }
+            activeRequestRef.current = null;
+          },
+          reject: (err) => {
+            resolved = true;
+            setError(err);
+            setIsLoading(false);
+            if (activeRequestRef.current?.overallTimeout) {
+              clearTimeout(activeRequestRef.current.overallTimeout);
+            }
+            if (activeRequestRef.current?.reject) {
+              activeRequestRef.current.reject(err);
+            }
+            activeRequestRef.current = null;
+          },
+          closeSocket,
+          clearOverallTimeout: () => {
+            if (activeRequestRef.current?.overallTimeout) {
+              clearTimeout(activeRequestRef.current.overallTimeout);
+            }
+          },
+          isResolved: () => resolved
+        };
+        return new Promise((resolvePromise, rejectPromise) => {
+          const overallTimeout = setTimeout(() => {
+            if (!resolved) {
+              const err = new Error("WebSocket request timeout");
+              setError(err);
+              setIsLoading(false);
+              rejectPromise(err);
+              activeRequestRef.current = null;
+            }
+          }, overallTimeoutMs);
+          const handler = options.onMessage || createStandardAiMessageHandler(options.defaults || {}, {
+            onServerAction: options.onServerAction,
+            toUserMessage
+          });
+          activeRequestRef.current = {
+            payload,
+            options,
+            ctx,
+            resolve: resolvePromise,
+            reject: rejectPromise,
+            overallTimeout,
+            handler
+          };
+          connectWebSocket().then(async () => {
+            if (!activeRequestRef.current || activeRequestRef.current.sent) {
+              return;
+            }
+            if (!lastWarmupPayloadRef.current) {
+              throw new Error("Warmup required");
+            }
+            if (!isWarmupReadyRef.current) {
+              sendWarmupPayload();
+              await waitForWarmupReady();
+            }
+            if (wsRef.current?.readyState === WebSocket.OPEN && activeRequestRef.current && !activeRequestRef.current.sent) {
+              activeRequestRef.current.sent = true;
+              wsRef.current.send(
+                JSON.stringify({ type: "client_message", data: payload })
+              );
+            }
+          }).catch((err) => {
+            clearTimeout(overallTimeout);
+            setError(err);
+            setIsLoading(false);
+            rejectPromise(err);
+            activeRequestRef.current = null;
+          });
+        });
+      },
+      [
+        connectWebSocket,
+        closeSocket,
+        sendWarmupPayload,
+        waitForWarmupReady,
+        rotateConnectionIfNeeded,
+        isLoading,
+        overallTimeoutMs,
+        toUserMessage,
+        invalidateTags
+      ]
+    );
+    return {
+      warmup,
+      send,
+      abort,
+      isLoading,
+      error,
+      data,
+      isConnected
+    };
+  };
+}
+
+export {
+  MAX_RETRIES,
+  INITIAL_RETRY_DELAY_MS,
+  CONNECTION_TIMEOUT_MS,
+  connectEdgeSocket,
+  createStandardAiMessageHandler,
+  createUseEdgeStream
+};