@cortexkit/pi-claustrum-auth 0.4.1 → 0.4.2

package/dist/index.js

@@ -622,7 +622,42 @@ function hasBillingAccessor() {
   }
 }
 
-class BillingAttempt {
+class ResponseCompletion {
+  completed = false;
+  failed = false;
+  event(value) {
+    if (!object(value))
+      return;
+    if (value.type === "error" || value.type === "response.failed" || value.error)
+      this.failed = true;
+    if (value.type === "response.completed" || value.type === "response.incomplete")
+      this.completed = true;
+    if (Array.isArray(value.choices)) {
+      for (const choice of value.choices) {
+        if (!object(choice))
+          continue;
+        if (choice.finish_reason != null || object(choice.message))
+          this.completed = true;
+      }
+    }
+  }
+  bodyEnd(ended, status) {
+    return !this.failed && this.completed && ended && status >= 200 && status < 300;
+  }
+  cancel(status) {
+    return this.bodyEnd(this.completed, status);
+  }
+  json(value) {
+    this.event(value);
+    if (object(value) && !value.error)
+      this.completed = true;
+  }
+  done() {
+    this.completed = true;
+  }
+}
+
+class BillingAttempt extends ResponseCompletion {
   entry;
   provider;
   sessionId;
@@ -633,8 +668,6 @@ class BillingAttempt {
   pi;
   attemptId = randomUUID();
   reported = false;
-  completed = false;
-  failed = false;
   outputStarted = false;
   responseId;
   rawUsage;
@@ -642,6 +675,7 @@ class BillingAttempt {
   status;
   ended = false;
   constructor(entry, provider, sessionId, requestKind, model, receipt, attempt, pi = false) {
+    super();
     this.entry = entry;
     this.provider = provider;
     this.sessionId = sessionId;
@@ -652,24 +686,17 @@ class BillingAttempt {
     this.pi = pi;
   }
   event(value) {
+    super.event(value);
     if (!object(value))
       return;
-    if (value.type === "error" || value.type === "response.failed" || value.error)
-      this.failed = true;
-    if (value.type === "response.completed" || value.type === "response.incomplete")
-      this.completed = true;
     if (typeof value.type === "string" && value.type.endsWith(".delta"))
       this.outputStarted = true;
     if (Array.isArray(value.choices)) {
       for (const choice of value.choices) {
         if (!object(choice))
           continue;
-        if (choice.finish_reason != null)
-          this.completed = true;
         if (object(choice.delta) && Object.keys(choice.delta).some((key) => key !== "role"))
           this.outputStarted = true;
-        if (object(choice.message))
-          this.completed = true;
       }
     }
     const response = object(value.response) ? value.response : value;
@@ -687,17 +714,7 @@ class BillingAttempt {
     this.ended = ended;
     if (!this.pi || status < 200 || status >= 300)
       this.end();
-  }
-  cancel(status) {
-    this.bodyEnd(this.completed, status);
-  }
-  json(value) {
-    this.event(value);
-    if (object(value) && !value.error)
-      this.completed = true;
-  }
-  done() {
-    this.completed = true;
+    return super.bodyEnd(ended, status);
   }
   piEnd(responseId, stopReason) {
     this.responseId = responseId;
@@ -742,11 +759,21 @@ class BillingAttempt {
     });
   }
 }
-function observeBillingResponse(record, response) {
-  if (!record)
+function observeBillingResponse(record, response, onCompleted) {
+  if (!record && !onCompleted)
     return response;
+  const observer = record ?? new ResponseCompletion;
+  let notified = false;
+  const finish = (completed) => {
+    if (!completed || notified || !onCompleted)
+      return;
+    notified = true;
+    try {
+      onCompleted();
+    } catch {}
+  };
   if (!response.body) {
-    record.bodyEnd(true, response.status);
+    finish(observer.bodyEnd(true, response.status));
     return response;
   }
   const sse = /^text\/event-stream(?:\s*;|$)/i.test(response.headers.get("content-type") ?? "");
@@ -755,10 +782,10 @@ function observeBillingResponse(record, response) {
   let oversized = false;
   const body = sse ? response.body.pipeThrough(watchServerSentEvents((event) => {
     if (event.data === "[DONE]")
-      record.done();
+      observer.done();
     else {
       try {
-        record.event(JSON.parse(event.data));
+        observer.event(JSON.parse(event.data));
       } catch {}
     }
   }, () => {})) : response.body;
@@ -770,10 +797,10 @@ function observeBillingResponse(record, response) {
         if (next.done) {
           if (!sse && !oversized) {
             try {
-              record.json(JSON.parse(json + decoder.decode()));
+              observer.json(JSON.parse(json + decoder.decode()));
             } catch {}
           }
-          record.bodyEnd(true, response.status);
+          finish(observer.bodyEnd(true, response.status));
           controller.close();
         } else {
           if (!sse && !oversized) {
@@ -786,12 +813,12 @@ function observeBillingResponse(record, response) {
           controller.enqueue(next.value);
         }
       } catch (error) {
-        record.bodyEnd(false, response.status);
+        observer.bodyEnd(false, response.status);
         controller.error(error);
       }
     },
     cancel(reason) {
-      record.cancel(response.status);
+      finish(observer.cancel(response.status));
       return reader.cancel(reason);
     }
   }, { highWaterMark: 0 }), {
@@ -6760,6 +6787,23 @@ async function reportRoutedBalance(key, sessionId, service, status, body) {
     providerCode
   });
 }
+function routedCompletion(key, sessionId, response) {
+  if (!key || !sessionId || !response.ok)
+    return;
+  const headers = {};
+  response.headers.forEach((value, name) => {
+    const lower = name.toLowerCase();
+    if (lower.startsWith("x-ratelimit-") || lower.startsWith("ratelimit") || lower.startsWith("anthropic-ratelimit-") || lower === "retry-after")
+      headers[lower] = value;
+  });
+  return () => reportRouting({
+    sessionId,
+    accountId: key.accountId,
+    credentialId: key.credentialId,
+    outcome: "served",
+    headers
+  });
+}
 // ../../node_modules/.bun/@cortexkit+common-auth@0.47.0+d48131f39628582b/node_modules/@cortexkit/common-auth/dist/routing/sticky.js
 var QUOTA_STALENESS_MS = 15 * 60000;
 var MIN_RESET_HOURS = 1 / 60;
@@ -6972,7 +7016,7 @@ function createFetch(options) {
                 return reply;
               }
               billing = options.billing?.(receipt);
-              return observeBillingResponse(billing, reply);
+              return observeBillingResponse(billing, reply, routedCompletion(routed, options.routing?.sessionId, reply));
             } catch (error) {
               const failedBilling = billing ?? options.billing?.(receipt);
               failedBilling?.end(false, false);

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@cortexkit/pi-claustrum-auth",
-  "version": "0.4.1",
+  "version": "0.4.2",
   "description": "Claustrum vault credentials for Pi",
   "type": "module",
   "license": "MIT",