@jeffusion/bungee-core 5.16.0 → 5.18.0

package/dist/plugins/token-metering/index.js

@@ -1986,6 +1986,209 @@ function classifyResponse(body, knownLlmRequest, event) {
   return;
 }
 
+// plugins/token-metering/server/websocket.ts
+import { randomUUID, createHash } from "crypto";
+var MAX_CONNECTIONS = 128;
+var MAX_RESPONSES = 1024;
+var MAX_CONNECTION_RESPONSES = 64;
+var MAX_TOMBSTONES = 256;
+var record2 = (value) => typeof value === "object" && value !== null && !Array.isArray(value);
+var identifier = (value) => typeof value === "string" && value.length > 0 && value.length <= 256 ? value : undefined;
+var count = (value) => typeof value === "number" && Number.isSafeInteger(value) && value >= 0;
+function metadata(event) {
+  return JSON.stringify([event.keyId ?? null, event.routeId, event.upstreamId, event.upstreamUrl, event.servingRevision]);
+}
+function responsesUrl(value) {
+  try {
+    return /\/(?:v1\/responses|backend-api\/codex\/responses)\/?$/.test(new URL(value).pathname);
+  } catch {
+    return false;
+  }
+}
+function stream(body) {
+  return identifier(body.stream_id) ?? (record2(body.response) ? identifier(body.response.stream_id) : undefined);
+}
+
+class ResponsesWebSocketMetering {
+  callbacks;
+  connections = new Map;
+  responseCount = 0;
+  closed = false;
+  constructor(callbacks) {
+    this.callbacks = callbacks;
+  }
+  async observe(event) {
+    try {
+      await this.observeActive(event);
+    } finally {
+      if (!event.isActive()) {
+        const connection = this.connections.get(event.connectionId);
+        if (connection) {
+          connection.disabled = true;
+          this.discard(connection);
+        }
+      }
+    }
+  }
+  async observeActive(event) {
+    if (this.closed || !event.isActive())
+      return;
+    if (event.phase === "open") {
+      if (!responsesUrl(event.upstreamUrl) || !this.callbacks.hasDemand() || this.connections.has(event.connectionId) || this.connections.size >= MAX_CONNECTIONS)
+        return;
+      this.connections.set(event.connectionId, { metadata: metadata(event), pending: [], active: new Map, finished: new Set, disabled: false });
+      return;
+    }
+    const connection = this.connections.get(event.connectionId);
+    if (!connection)
+      return;
+    if (connection.metadata !== metadata(event)) {
+      this.discard(connection);
+      this.connections.delete(event.connectionId);
+      return;
+    }
+    if (event.phase === "close" || event.phase === "incomplete") {
+      await this.interrupt(connection, event, event.phase === "close" ? "cancelled" : "failed");
+      if (event.phase === "close")
+        this.connections.delete(event.connectionId);
+      return;
+    }
+    if (connection.disabled || event.phase !== "message" || event.message.kind !== "text")
+      return;
+    const body = event.message.json();
+    if (!record2(body))
+      return;
+    if (event.direction === "client") {
+      if (body.type !== "response.create" || !(typeof body.input === "string" || Array.isArray(body.input) || identifier(body.previous_response_id) || identifier(body.model)))
+        return;
+      if (!this.hasCapacity(connection)) {
+        await this.interrupt(connection, event, "failed");
+        return;
+      }
+      const response = await this.start(event, identifier(body.model), stream(body));
+      if (response && event.isActive())
+        connection.pending.push(response);
+      return;
+    }
+    const type = identifier(body.type);
+    if (!type?.startsWith("response."))
+      return;
+    const envelope = record2(body.response) ? body.response : undefined;
+    const terminal = ["response.completed", "response.failed", "response.incomplete"].includes(type);
+    if (terminal && (!envelope || envelope.object !== "response" || !Array.isArray(envelope.output)))
+      return;
+    const responseId = identifier(envelope?.id) ?? identifier(body.response_id);
+    if (!responseId || connection.finished.has(responseId))
+      return;
+    let response = connection.active.get(responseId);
+    if (!response) {
+      if (!envelope || envelope.object !== "response" || !Array.isArray(envelope.output) || !["response.created", "response.in_progress", "response.completed", "response.failed", "response.incomplete"].includes(type))
+        return;
+      const key = stream(body);
+      let pendingIndex = connection.pending.findIndex((item) => item.stream === key);
+      if (pendingIndex < 0 && key !== undefined)
+        pendingIndex = connection.pending.findIndex((item) => item.stream === undefined);
+      if (pendingIndex >= 0 && !terminal)
+        response = connection.pending.splice(pendingIndex, 1)[0];
+      else {
+        if (!this.hasCapacity(connection)) {
+          await this.interrupt(connection, event, "failed");
+          return;
+        }
+        response = await this.start(event, identifier(envelope.model), key);
+      }
+      if (!response || !event.isActive())
+        return;
+      connection.active.set(responseId, response);
+    }
+    const usage = record2(envelope?.usage) ? envelope.usage : undefined;
+    if (!response.bound && !await this.bind(response, event, responseId))
+      return;
+    if (!await this.emit(response, event, { phase: "response", status: 200, protocol: "sse", body }))
+      return;
+    if (!terminal)
+      return;
+    connection.active.delete(responseId);
+    connection.finished.add(responseId);
+    const usageComplete = count(usage?.input_tokens) && count(usage?.output_tokens);
+    await this.finish(response, event, type === "response.completed" && usageComplete ? "completed" : "failed", !usageComplete);
+    if (connection.finished.size > MAX_TOMBSTONES)
+      connection.finished.delete(connection.finished.values().next().value);
+  }
+  hasCapacity(connection) {
+    return this.responseCount < MAX_RESPONSES && connection.pending.length + connection.active.size < MAX_CONNECTION_RESPONSES;
+  }
+  async start(event, model, streamId) {
+    const response = { requestId: randomUUID(), attemptId: randomUUID(), stream: streamId, model };
+    this.responseCount++;
+    this.callbacks.prepareRequest(response.requestId);
+    try {
+      if (await this.emit(response, event, { phase: "selected" }) && await this.emit(response, event, { phase: "request", url: event.upstreamUrl, body: { model, stream: true } }))
+        return response;
+    } catch (error) {
+      this.responseCount--;
+      this.callbacks.discardRequest(response.requestId, response.attemptId);
+      throw error;
+    }
+    this.responseCount--;
+    return;
+  }
+  async bind(response, event, responseId) {
+    this.callbacks.discardRequest(response.requestId, response.attemptId);
+    const uuid = (kind) => {
+      const hex = createHash("sha256").update(JSON.stringify(["bungee.ws.response.v1", event.connectionId, responseId, kind])).digest("hex");
+      return `${hex.slice(0, 8)}-${hex.slice(8, 12)}-5${hex.slice(13, 16)}-a${hex.slice(17, 20)}-${hex.slice(20, 32)}`;
+    };
+    response.requestId = uuid("request");
+    response.attemptId = uuid("attempt");
+    response.bound = true;
+    this.callbacks.prepareRequest(response.requestId);
+    return await this.emit(response, event, { phase: "selected" }) && await this.emit(response, event, { phase: "request", url: event.upstreamUrl, body: { model: response.model, stream: true } });
+  }
+  async emit(response, event, phase) {
+    if (this.closed || !event.isActive()) {
+      this.callbacks.discardRequest(response.requestId, response.attemptId);
+      return false;
+    }
+    await this.callbacks.emit({ requestId: response.requestId, attemptId: response.attemptId, keyId: event.keyId, routeId: event.routeId, upstreamId: event.upstreamId, isActive: event.isActive, ...phase });
+    if (this.closed || !event.isActive()) {
+      this.callbacks.discardRequest(response.requestId, response.attemptId);
+      return false;
+    }
+    return true;
+  }
+  async finish(response, event, outcome, incomplete) {
+    this.responseCount--;
+    if (incomplete && !await this.emit(response, event, { phase: "incomplete", reason: "raw-response-incomplete" }))
+      return;
+    if (!await this.emit(response, event, { phase: "end", sent: true, outcome }))
+      return;
+    await this.emit(response, event, { phase: "request-end" });
+  }
+  async interrupt(connection, event, outcome) {
+    connection.disabled = true;
+    const unfinished = [...connection.pending, ...connection.active.values()];
+    connection.pending.length = 0;
+    connection.active.clear();
+    for (const response of unfinished)
+      await this.finish(response, event, outcome, true);
+  }
+  discard(connection) {
+    for (const response of [...connection.pending, ...connection.active.values()]) {
+      this.responseCount--;
+      this.callbacks.discardRequest(response.requestId, response.attemptId);
+    }
+    connection.pending.length = 0;
+    connection.active.clear();
+  }
+  dispose() {
+    this.closed = true;
+    for (const connection of this.connections.values())
+      this.discard(connection);
+    this.connections.clear();
+  }
+}
+
 // plugins/token-metering/server/index.ts
 var MAX_ACTIVE_ATTEMPTS = 1024;
 var LOG_INTERVAL_MS = 60000;
@@ -2173,6 +2376,16 @@ var TokenMeteringPlugin = definePlugin(class {
   pending = new Set;
   requiredPending = new Map;
   closed = false;
+  websocket = new ResponsesWebSocketMetering({
+    emit: (event) => this.handleAttemptObservation(event, true),
+    prepareRequest: (requestId) => this.prepareRequest(requestId),
+    discardRequest: (requestId, attemptId) => {
+      this.attempts.delete(attemptId);
+      this.requests.delete(requestId);
+      this.droppedAttempts.delete(attemptId);
+    },
+    hasDemand: () => !this.closed && this.subscriptions.size > 0
+  });
   parsedResponses = 0;
   get activeAttempts() {
     return this.attempts.size;
@@ -2247,6 +2460,7 @@ var TokenMeteringPlugin = definePlugin(class {
   register(hooks) {
     hooks.onRequestInit.tapPromise("token-metering-demand", async (event) => this.prepareRequest(event.requestId));
     hooks.onAttemptObservation.tapPromise("token-metering", async (event) => this.handleAttemptObservation(event));
+    hooks.onWebSocketObservation.tapPromise("token-metering-websocket", async (event) => this.websocket.observe(event));
   }
   prepareRequest(requestId) {
     if (!this.requests.has(requestId))
@@ -2297,7 +2511,7 @@ var TokenMeteringPlugin = definePlugin(class {
       });
     }
   }
-  async handleAttemptObservation(event) {
+  async handleAttemptObservation(event, websocket = false) {
     if (!event.isActive()) {
       this.attempts.delete(event.attemptId);
       return;
@@ -2337,7 +2551,7 @@ var TokenMeteringPlugin = definePlugin(class {
       const body = parseRequestBody(event.body);
       if (!body)
         return;
-      const classification = classifyRequest(url, body);
+      const classification = websocket ? { llm: true, provider: "openai" } : classifyRequest(url, body);
       attempt.llm ||= classification.llm;
       if (!classification.llm)
         return;
@@ -2353,6 +2567,13 @@ var TokenMeteringPlugin = definePlugin(class {
     if (event.phase === "response") {
       attempt.responseSeen = true;
       const body = event.body;
+      if (websocket && isRecord2(body.response)) {
+        const model = detectModel(body.response);
+        if (model) {
+          attempt.model = model;
+          attempt.session?.consumeRequest({ body: { model } });
+        }
+      }
       attempt.responseFailed ||= event.status >= 400 || body.error !== undefined;
       const provider = attempt.provider === "unknown" ? classifyResponse(body, attempt.llm, event.envelope?.event) : attempt.provider;
       if (!provider)
@@ -2428,6 +2649,7 @@ var TokenMeteringPlugin = definePlugin(class {
   }
   async onDestroy() {
     this.closed = true;
+    this.websocket.dispose();
     await Promise.allSettled([...this.pending]);
     this.subscriptions.clear();
     this.requests.clear();
@@ -2442,4 +2664,4 @@ export {
   server_default as default
 };
 
-//# debugId=10FC3B2EECF76FCF64756E2164756E21
+//# debugId=E40F8525D6A44D4364756E2164756E21