caplyr 0.2.3 → 0.3.0

package/dist/index.js

@@ -36,20 +36,32 @@ var LogShipper = class {
   constructor(config) {
     this.buffer = [];
     this.timer = null;
+    // Store bound handlers so we can remove them on shutdown
+    this.processHandlers = [];
     this.endpoint = config.endpoint ?? "https://api.caplyr.com";
     this.apiKey = config.apiKey;
     this.batchSize = config.batchSize ?? 10;
     this.flushInterval = config.flushInterval ?? 3e4;
+    this.maxBufferSize = 1e3;
     this.onError = config.onError;
     this.timer = setInterval(() => this.flush(), this.flushInterval);
+    this.timer.unref?.();
     if (typeof process !== "undefined" && process.on) {
-      const flushAndExit = () => {
+      const onBeforeExit = () => {
+        this.flush();
+      };
+      const onSignal = () => {
         this.flush().finally(() => {
         });
       };
-      process.on("beforeExit", () => this.flush());
-      process.on("SIGTERM", flushAndExit);
-      process.on("SIGINT", flushAndExit);
+      process.on("beforeExit", onBeforeExit);
+      process.on("SIGTERM", onSignal);
+      process.on("SIGINT", onSignal);
+      this.processHandlers = [
+        { event: "beforeExit", handler: onBeforeExit },
+        { event: "SIGTERM", handler: onSignal },
+        { event: "SIGINT", handler: onSignal }
+      ];
     }
   }
   /**
@@ -57,6 +69,10 @@ var LogShipper = class {
    * Auto-flushes when batch size is reached.
    */
   push(log) {
+    if (this.buffer.length >= this.maxBufferSize) {
+      const excess = this.buffer.length - this.maxBufferSize + 1;
+      this.buffer.splice(0, excess);
+    }
     this.buffer.push(log);
     if (this.buffer.length >= this.batchSize) {
       this.flush();
@@ -88,7 +104,18 @@ var LogShipper = class {
     }
   }
   /**
-   * Stop the periodic flush timer.
+   * Remove process signal handlers registered in the constructor.
+   */
+  removeProcessHandlers() {
+    if (typeof process !== "undefined" && process.removeListener) {
+      for (const { event, handler } of this.processHandlers) {
+        process.removeListener(event, handler);
+      }
+    }
+    this.processHandlers = [];
+  }
+  /**
+   * Stop the periodic flush timer and remove signal handlers.
    * Call this when tearing down the SDK.
    */
   destroy() {
@@ -96,17 +123,53 @@ var LogShipper = class {
       clearInterval(this.timer);
       this.timer = null;
     }
+    this.removeProcessHandlers();
     this.flush();
   }
+  /**
+   * Await the final log flush, stop timers, and remove signal handlers.
+   * Preferred over destroy() for clean shutdown.
+   */
   async shutdown() {
     if (this.timer) {
       clearInterval(this.timer);
       this.timer = null;
     }
+    this.removeProcessHandlers();
     await this.flush();
   }
 };
 
+// src/mutex.ts
+var Mutex = class {
+  constructor() {
+    this.queue = [];
+    this.locked = false;
+  }
+  async acquire() {
+    if (!this.locked) {
+      this.locked = true;
+      return this.createRelease();
+    }
+    return new Promise((resolve) => {
+      this.queue.push(() => resolve(this.createRelease()));
+    });
+  }
+  createRelease() {
+    let released = false;
+    return () => {
+      if (released) return;
+      released = true;
+      const next = this.queue.shift();
+      if (next) {
+        next();
+      } else {
+        this.locked = false;
+      }
+    };
+  }
+};
+
 // src/heartbeat.ts
 var Heartbeat = class {
   constructor(config) {
@@ -124,6 +187,8 @@ var Heartbeat = class {
     };
     /** Current protection status */
     this.status = "ACTIVE";
+    /** Mutex for serializing budget check → API call → trackSpend */
+    this.budgetMutex = new Mutex();
     /** Local budget limits set via config (not from server) */
     this.localDailyLimit = null;
     this.localMonthlyLimit = null;
@@ -155,6 +220,7 @@ var Heartbeat = class {
   start() {
     this.beat();
     this.timer = setInterval(() => this.beat(), this.interval);
+    this.timer.unref?.();
   }
   /**
    * Send a single heartbeat and update local state.
@@ -174,21 +240,20 @@ var Heartbeat = class {
         throw new Error(`Heartbeat failed: ${res.status}`);
       }
       const data = await res.json();
-      const localDailyUsed = this.budgetStatus.daily_used;
-      const localMonthlyUsed = this.budgetStatus.monthly_used;
       const serverDailyUsed = Number(data.daily_used) || 0;
       const serverMonthlyUsed = Number(data.monthly_used) || 0;
       const serverDailyLimit = data.daily_limit != null ? Number(data.daily_limit) : null;
       const serverMonthlyLimit = data.monthly_limit != null ? Number(data.monthly_limit) : null;
-      this.budgetStatus = {
-        ...data,
-        // Use whichever spend is higher — server or local tracking
-        daily_used: Math.max(serverDailyUsed, localDailyUsed),
-        monthly_used: Math.max(serverMonthlyUsed, localMonthlyUsed),
-        // Use the stricter (lower) limit — local config takes priority if lower
-        daily_limit: this.pickStricterLimit(serverDailyLimit, this.localDailyLimit),
-        monthly_limit: this.pickStricterLimit(serverMonthlyLimit, this.localMonthlyLimit)
-      };
+      const snapshotDaily = this.budgetStatus.daily_used;
+      const snapshotMonthly = this.budgetStatus.monthly_used;
+      const mergedDaily = Math.max(serverDailyUsed, snapshotDaily);
+      const mergedMonthly = Math.max(serverMonthlyUsed, snapshotMonthly);
+      this.budgetStatus.daily_used = mergedDaily + (this.budgetStatus.daily_used - snapshotDaily);
+      this.budgetStatus.monthly_used = mergedMonthly + (this.budgetStatus.monthly_used - snapshotMonthly);
+      this.budgetStatus.daily_limit = this.pickStricterLimit(serverDailyLimit, this.localDailyLimit);
+      this.budgetStatus.monthly_limit = this.pickStricterLimit(serverMonthlyLimit, this.localMonthlyLimit);
+      this.budgetStatus.status = data.status;
+      this.budgetStatus.kill_switch_active = data.kill_switch_active;
       this.consecutiveFailures = 0;
       const newStatus = data.kill_switch_active ? "OFF" : data.status;
       if (newStatus !== this.status) {
@@ -382,18 +447,98 @@ function wrapAnthropic(client, config, shipper, heartbeat) {
               });
             }
           }
-          if (config.mode === "cost_protect") {
-            if (heartbeat.isMonthlyBudgetExceeded() || heartbeat.isDailyBudgetExceeded()) {
-              blocked = true;
-              enforcementReason = heartbeat.isDailyBudgetExceeded() ? "daily_budget_exceeded" : "monthly_budget_exceeded";
-              const blockError = {
-                code: "BUDGET_EXCEEDED",
-                message: `AI budget exceeded. ${enforcementReason.replace(/_/g, " ")}.`,
-                budget_used: heartbeat.budgetStatus.monthly_used,
-                budget_limit: heartbeat.budgetStatus.monthly_limit ?? 0,
-                retry_after: getNextResetTime(enforcementReason),
-                dashboard_url: dashboardUrl
-              };
+          const release = await heartbeat.budgetMutex.acquire();
+          try {
+            if (config.mode === "cost_protect") {
+              if (heartbeat.isMonthlyBudgetExceeded() || heartbeat.isDailyBudgetExceeded()) {
+                blocked = true;
+                enforcementReason = heartbeat.isDailyBudgetExceeded() ? "daily_budget_exceeded" : "monthly_budget_exceeded";
+                const blockError = {
+                  code: "BUDGET_EXCEEDED",
+                  message: `AI budget exceeded. ${enforcementReason.replace(/_/g, " ")}.`,
+                  budget_used: heartbeat.budgetStatus.monthly_used,
+                  budget_limit: heartbeat.budgetStatus.monthly_limit ?? 0,
+                  retry_after: getNextResetTime(enforcementReason),
+                  dashboard_url: dashboardUrl
+                };
+                shipper.push({
+                  id: generateId(),
+                  timestamp: startTime,
+                  provider: "anthropic",
+                  model,
+                  input_tokens: 0,
+                  output_tokens: 0,
+                  cost: 0,
+                  latency_ms: Date.now() - startTime,
+                  endpoint_tag: config.endpoint_tag,
+                  downgraded: false,
+                  blocked: true,
+                  enforcement_reason: enforcementReason
+                });
+                throw Object.assign(new Error(blockError.message), {
+                  caplyr: blockError
+                });
+              }
+              if (heartbeat.isDowngradeThresholdReached(downgradeThreshold)) {
+                const fallback = config.fallback ?? getDefaultFallback(model);
+                if (fallback && fallback !== model) {
+                  originalModel = model;
+                  model = fallback;
+                  downgraded = true;
+                  enforcementReason = "auto_downgrade_threshold";
+                  config.onEnforcement?.({
+                    type: "downgrade",
+                    timestamp: Date.now(),
+                    reason: `Budget at ${Math.round(downgradeThreshold * 100)}% \u2014 downgraded ${originalModel} \u2192 ${model}`,
+                    original_model: originalModel,
+                    fallback_model: model,
+                    budget_used: heartbeat.budgetStatus.monthly_used,
+                    budget_limit: heartbeat.budgetStatus.monthly_limit ?? 0,
+                    estimated_savings: 0
+                    // Calculated after response
+                  });
+                }
+              }
+            }
+            const requestParams = downgraded ? { ...params, model } : params;
+            try {
+              const response = await target.create.call(
+                target,
+                requestParams,
+                options
+              );
+              const latency = Date.now() - startTime;
+              const inputTokens = response?.usage?.input_tokens ?? 0;
+              const outputTokens = response?.usage?.output_tokens ?? 0;
+              const cost = calculateCost(model, inputTokens, outputTokens);
+              heartbeat.trackSpend(cost);
+              let estimatedSavings = 0;
+              if (downgraded && originalModel) {
+                const originalCost = calculateCost(
+                  originalModel,
+                  inputTokens,
+                  outputTokens
+                );
+                estimatedSavings = originalCost - cost;
+              }
+              shipper.push({
+                id: generateId(),
+                timestamp: startTime,
+                provider: "anthropic",
+                model,
+                input_tokens: inputTokens,
+                output_tokens: outputTokens,
+                cost,
+                latency_ms: latency,
+                endpoint_tag: config.endpoint_tag,
+                downgraded,
+                original_model: originalModel,
+                blocked: false,
+                enforcement_reason: enforcementReason
+              });
+              return response;
+            } catch (err) {
+              if (err?.caplyr) throw err;
               shipper.push({
                 id: generateId(),
                 timestamp: startTime,
@@ -404,90 +549,15 @@ function wrapAnthropic(client, config, shipper, heartbeat) {
                 cost: 0,
                 latency_ms: Date.now() - startTime,
                 endpoint_tag: config.endpoint_tag,
-                downgraded: false,
-                blocked: true,
-                enforcement_reason: enforcementReason
+                downgraded,
+                original_model: originalModel,
+                blocked: false,
+                enforcement_reason: "provider_error"
               });
-              throw Object.assign(new Error(blockError.message), {
-                caplyr: blockError
-              });
-            }
-            if (heartbeat.isDowngradeThresholdReached(downgradeThreshold)) {
-              const fallback = config.fallback ?? getDefaultFallback(model);
-              if (fallback && fallback !== model) {
-                originalModel = model;
-                model = fallback;
-                downgraded = true;
-                enforcementReason = "auto_downgrade_threshold";
-                config.onEnforcement?.({
-                  type: "downgrade",
-                  timestamp: Date.now(),
-                  reason: `Budget at ${Math.round(downgradeThreshold * 100)}% \u2014 downgraded ${originalModel} \u2192 ${model}`,
-                  original_model: originalModel,
-                  fallback_model: model,
-                  budget_used: heartbeat.budgetStatus.monthly_used,
-                  budget_limit: heartbeat.budgetStatus.monthly_limit ?? 0,
-                  estimated_savings: 0
-                  // Calculated after response
-                });
-              }
+              throw err;
             }
-          }
-          const requestParams = downgraded ? { ...params, model } : params;
-          try {
-            const response = await target.create.call(
-              target,
-              requestParams,
-              options
-            );
-            const latency = Date.now() - startTime;
-            const inputTokens = response?.usage?.input_tokens ?? 0;
-            const outputTokens = response?.usage?.output_tokens ?? 0;
-            const cost = calculateCost(model, inputTokens, outputTokens);
-            heartbeat.trackSpend(cost);
-            let estimatedSavings = 0;
-            if (downgraded && originalModel) {
-              const originalCost = calculateCost(
-                originalModel,
-                inputTokens,
-                outputTokens
-              );
-              estimatedSavings = originalCost - cost;
-            }
-            shipper.push({
-              id: generateId(),
-              timestamp: startTime,
-              provider: "anthropic",
-              model,
-              input_tokens: inputTokens,
-              output_tokens: outputTokens,
-              cost,
-              latency_ms: latency,
-              endpoint_tag: config.endpoint_tag,
-              downgraded,
-              original_model: originalModel,
-              blocked: false,
-              enforcement_reason: enforcementReason
-            });
-            return response;
-          } catch (err) {
-            if (err?.caplyr) throw err;
-            shipper.push({
-              id: generateId(),
-              timestamp: startTime,
-              provider: "anthropic",
-              model,
-              input_tokens: 0,
-              output_tokens: 0,
-              cost: 0,
-              latency_ms: Date.now() - startTime,
-              endpoint_tag: config.endpoint_tag,
-              downgraded,
-              original_model: originalModel,
-              blocked: false,
-              enforcement_reason: "provider_error"
-            });
-            throw err;
+          } finally {
+            release();
           }
         };
       }
@@ -564,18 +634,89 @@ function wrapOpenAI(client, config, shipper, heartbeat) {
               });
             }
           }
-          if (config.mode === "cost_protect") {
-            if (heartbeat.isMonthlyBudgetExceeded() || heartbeat.isDailyBudgetExceeded()) {
-              blocked = true;
-              enforcementReason = heartbeat.isDailyBudgetExceeded() ? "daily_budget_exceeded" : "monthly_budget_exceeded";
-              const blockError = {
-                code: "BUDGET_EXCEEDED",
-                message: `AI budget exceeded. ${enforcementReason.replace(/_/g, " ")}.`,
-                budget_used: heartbeat.budgetStatus.monthly_used,
-                budget_limit: heartbeat.budgetStatus.monthly_limit ?? 0,
-                retry_after: getNextResetTime2(enforcementReason),
-                dashboard_url: dashboardUrl
-              };
+          const release = await heartbeat.budgetMutex.acquire();
+          try {
+            if (config.mode === "cost_protect") {
+              if (heartbeat.isMonthlyBudgetExceeded() || heartbeat.isDailyBudgetExceeded()) {
+                blocked = true;
+                enforcementReason = heartbeat.isDailyBudgetExceeded() ? "daily_budget_exceeded" : "monthly_budget_exceeded";
+                const blockError = {
+                  code: "BUDGET_EXCEEDED",
+                  message: `AI budget exceeded. ${enforcementReason.replace(/_/g, " ")}.`,
+                  budget_used: heartbeat.budgetStatus.monthly_used,
+                  budget_limit: heartbeat.budgetStatus.monthly_limit ?? 0,
+                  retry_after: getNextResetTime2(enforcementReason),
+                  dashboard_url: dashboardUrl
+                };
+                shipper.push({
+                  id: generateId2(),
+                  timestamp: startTime,
+                  provider: "openai",
+                  model,
+                  input_tokens: 0,
+                  output_tokens: 0,
+                  cost: 0,
+                  latency_ms: Date.now() - startTime,
+                  endpoint_tag: config.endpoint_tag,
+                  downgraded: false,
+                  blocked: true,
+                  enforcement_reason: enforcementReason
+                });
+                throw Object.assign(new Error(blockError.message), {
+                  caplyr: blockError
+                });
+              }
+              if (heartbeat.isDowngradeThresholdReached(downgradeThreshold)) {
+                const fallback = config.fallback ?? getDefaultFallback(model);
+                if (fallback && fallback !== model) {
+                  originalModel = model;
+                  model = fallback;
+                  downgraded = true;
+                  enforcementReason = "auto_downgrade_threshold";
+                  config.onEnforcement?.({
+                    type: "downgrade",
+                    timestamp: Date.now(),
+                    reason: `Budget at ${Math.round(downgradeThreshold * 100)}% \u2014 downgraded ${originalModel} \u2192 ${model}`,
+                    original_model: originalModel,
+                    fallback_model: model,
+                    budget_used: heartbeat.budgetStatus.monthly_used,
+                    budget_limit: heartbeat.budgetStatus.monthly_limit ?? 0,
+                    estimated_savings: 0
+                  });
+                }
+              }
+            }
+            const requestParams = downgraded ? { ...params, model } : params;
+            try {
+              const response = await target.create.call(
+                target,
+                requestParams,
+                options
+              );
+              const latency = Date.now() - startTime;
+              const usage = response?.usage;
+              const inputTokens = usage?.prompt_tokens ?? 0;
+              const outputTokens = usage?.completion_tokens ?? 0;
+              const cost = calculateCost(model, inputTokens, outputTokens);
+              heartbeat.trackSpend(cost);
+              shipper.push({
+                id: generateId2(),
+                timestamp: startTime,
+                provider: "openai",
+                model,
+                input_tokens: inputTokens,
+                output_tokens: outputTokens,
+                cost,
+                latency_ms: latency,
+                endpoint_tag: config.endpoint_tag,
+                downgraded,
+                original_model: originalModel,
+                blocked: false,
+                enforcement_reason: enforcementReason
+              });
+              return response;
+            } catch (err) {
+              if (err?.caplyr) throw err;
               shipper.push({
                 id: generateId2(),
                 timestamp: startTime,
@@ -586,81 +727,15 @@ function wrapOpenAI(client, config, shipper, heartbeat) {
                 cost: 0,
                 latency_ms: Date.now() - startTime,
                 endpoint_tag: config.endpoint_tag,
-                downgraded: false,
-                blocked: true,
-                enforcement_reason: enforcementReason
-              });
-              throw Object.assign(new Error(blockError.message), {
-                caplyr: blockError
+                downgraded,
+                original_model: originalModel,
+                blocked: false,
+                enforcement_reason: "provider_error"
               });
+              throw err;
             }
-            if (heartbeat.isDowngradeThresholdReached(downgradeThreshold)) {
-              const fallback = config.fallback ?? getDefaultFallback(model);
-              if (fallback && fallback !== model) {
-                originalModel = model;
-                model = fallback;
-                downgraded = true;
-                enforcementReason = "auto_downgrade_threshold";
-                config.onEnforcement?.({
-                  type: "downgrade",
-                  timestamp: Date.now(),
-                  reason: `Budget at ${Math.round(downgradeThreshold * 100)}% \u2014 downgraded ${originalModel} \u2192 ${model}`,
-                  original_model: originalModel,
-                  fallback_model: model,
-                  budget_used: heartbeat.budgetStatus.monthly_used,
-                  budget_limit: heartbeat.budgetStatus.monthly_limit ?? 0,
-                  estimated_savings: 0
-                });
-              }
-            }
-          }
-          const requestParams = downgraded ? { ...params, model } : params;
-          try {
-            const response = await target.create.call(
-              target,
-              requestParams,
-              options
-            );
-            const latency = Date.now() - startTime;
-            const usage = response?.usage;
-            const inputTokens = usage?.prompt_tokens ?? 0;
-            const outputTokens = usage?.completion_tokens ?? 0;
-            const cost = calculateCost(model, inputTokens, outputTokens);
-            heartbeat.trackSpend(cost);
-            shipper.push({
-              id: generateId2(),
-              timestamp: startTime,
-              provider: "openai",
-              model,
-              input_tokens: inputTokens,
-              output_tokens: outputTokens,
-              cost,
-              latency_ms: latency,
-              endpoint_tag: config.endpoint_tag,
-              downgraded,
-              original_model: originalModel,
-              blocked: false,
-              enforcement_reason: enforcementReason
-            });
-            return response;
-          } catch (err) {
-            if (err?.caplyr) throw err;
-            shipper.push({
-              id: generateId2(),
-              timestamp: startTime,
-              provider: "openai",
-              model,
-              input_tokens: 0,
-              output_tokens: 0,
-              cost: 0,
-              latency_ms: Date.now() - startTime,
-              endpoint_tag: config.endpoint_tag,
-              downgraded,
-              original_model: originalModel,
-              blocked: false,
-              enforcement_reason: "provider_error"
-            });
-            throw err;
+          } finally {
+            release();
           }
         };
       }
@@ -723,6 +798,7 @@ function protect(client, config) {
     ...config
   };
   let shared = instances.get(resolvedConfig.apiKey);
+  const isExisting = !!shared;
   if (!shared) {
     const shipper2 = new LogShipper(resolvedConfig);
     const heartbeat2 = new Heartbeat(resolvedConfig);
@@ -732,7 +808,19 @@ function protect(client, config) {
   }
   const { shipper, heartbeat } = shared;
   if (resolvedConfig.budget) {
-    const budgetConfig = typeof resolvedConfig.budget === "number" ? { monthly: resolvedConfig.budget } : resolvedConfig.budget;
+    const raw = typeof resolvedConfig.budget === "number" ? { monthly: resolvedConfig.budget } : resolvedConfig.budget;
+    const budgetConfig = { ...raw };
+    if (isExisting) {
+      const current = heartbeat.budgetStatus;
+      if (budgetConfig.daily !== void 0) {
+        const existing = current.daily_limit;
+        budgetConfig.daily = existing !== null ? Math.min(existing, budgetConfig.daily) : budgetConfig.daily;
+      }
+      if (budgetConfig.monthly !== void 0) {
+        const existing = current.monthly_limit;
+        budgetConfig.monthly = existing !== null ? Math.min(existing, budgetConfig.monthly) : budgetConfig.monthly;
+      }
+    }
     heartbeat.applyLocalLimits(budgetConfig);
   }
   const provider = detectProvider(client);
@@ -774,14 +862,16 @@ async function shutdown(apiKey) {
     const shared = instances.get(apiKey);
     if (shared) {
       shared.heartbeat.destroy();
-      shared.shipper.destroy();
+      await shared.shipper.shutdown();
       instances.delete(apiKey);
     }
   } else {
+    const shutdowns = [];
     for (const [key, shared] of instances) {
       shared.heartbeat.destroy();
-      shared.shipper.destroy();
+      shutdowns.push(shared.shipper.shutdown());
     }
+    await Promise.all(shutdowns);
     instances.clear();
   }
 }

package/dist/index.mjs

@@ -3,20 +3,32 @@ var LogShipper = class {
   constructor(config) {
     this.buffer = [];
     this.timer = null;
+    // Store bound handlers so we can remove them on shutdown
+    this.processHandlers = [];
     this.endpoint = config.endpoint ?? "https://api.caplyr.com";
     this.apiKey = config.apiKey;
     this.batchSize = config.batchSize ?? 10;
     this.flushInterval = config.flushInterval ?? 3e4;
+    this.maxBufferSize = 1e3;
     this.onError = config.onError;
     this.timer = setInterval(() => this.flush(), this.flushInterval);
+    this.timer.unref?.();
     if (typeof process !== "undefined" && process.on) {
-      const flushAndExit = () => {
+      const onBeforeExit = () => {
+        this.flush();
+      };
+      const onSignal = () => {
         this.flush().finally(() => {
         });
       };
-      process.on("beforeExit", () => this.flush());
-      process.on("SIGTERM", flushAndExit);
-      process.on("SIGINT", flushAndExit);
+      process.on("beforeExit", onBeforeExit);
+      process.on("SIGTERM", onSignal);
+      process.on("SIGINT", onSignal);
+      this.processHandlers = [
+        { event: "beforeExit", handler: onBeforeExit },
+        { event: "SIGTERM", handler: onSignal },
+        { event: "SIGINT", handler: onSignal }
+      ];
     }
   }
   /**
@@ -24,6 +36,10 @@ var LogShipper = class {
    * Auto-flushes when batch size is reached.
    */
   push(log) {
+    if (this.buffer.length >= this.maxBufferSize) {
+      const excess = this.buffer.length - this.maxBufferSize + 1;
+      this.buffer.splice(0, excess);
+    }
     this.buffer.push(log);
     if (this.buffer.length >= this.batchSize) {
       this.flush();
@@ -55,7 +71,18 @@ var LogShipper = class {
     }
   }
   /**
-   * Stop the periodic flush timer.
+   * Remove process signal handlers registered in the constructor.
+   */
+  removeProcessHandlers() {
+    if (typeof process !== "undefined" && process.removeListener) {
+      for (const { event, handler } of this.processHandlers) {
+        process.removeListener(event, handler);
+      }
+    }
+    this.processHandlers = [];
+  }
+  /**
+   * Stop the periodic flush timer and remove signal handlers.
    * Call this when tearing down the SDK.
    */
   destroy() {
@@ -63,17 +90,53 @@ var LogShipper = class {
       clearInterval(this.timer);
       this.timer = null;
     }
+    this.removeProcessHandlers();
     this.flush();
   }
+  /**
+   * Await the final log flush, stop timers, and remove signal handlers.
+   * Preferred over destroy() for clean shutdown.
+   */
   async shutdown() {
     if (this.timer) {
       clearInterval(this.timer);
       this.timer = null;
     }
+    this.removeProcessHandlers();
     await this.flush();
   }
 };
 
+// src/mutex.ts
+var Mutex = class {
+  constructor() {
+    this.queue = [];
+    this.locked = false;
+  }
+  async acquire() {
+    if (!this.locked) {
+      this.locked = true;
+      return this.createRelease();
+    }
+    return new Promise((resolve) => {
+      this.queue.push(() => resolve(this.createRelease()));
+    });
+  }
+  createRelease() {
+    let released = false;
+    return () => {
+      if (released) return;
+      released = true;
+      const next = this.queue.shift();
+      if (next) {
+        next();
+      } else {
+        this.locked = false;
+      }
+    };
+  }
+};
+
 // src/heartbeat.ts
 var Heartbeat = class {
   constructor(config) {
@@ -91,6 +154,8 @@ var Heartbeat = class {
     };
     /** Current protection status */
     this.status = "ACTIVE";
+    /** Mutex for serializing budget check → API call → trackSpend */
+    this.budgetMutex = new Mutex();
     /** Local budget limits set via config (not from server) */
     this.localDailyLimit = null;
     this.localMonthlyLimit = null;
@@ -122,6 +187,7 @@ var Heartbeat = class {
   start() {
     this.beat();
     this.timer = setInterval(() => this.beat(), this.interval);
+    this.timer.unref?.();
   }
   /**
    * Send a single heartbeat and update local state.
@@ -141,21 +207,20 @@ var Heartbeat = class {
         throw new Error(`Heartbeat failed: ${res.status}`);
       }
       const data = await res.json();
-      const localDailyUsed = this.budgetStatus.daily_used;
-      const localMonthlyUsed = this.budgetStatus.monthly_used;
       const serverDailyUsed = Number(data.daily_used) || 0;
       const serverMonthlyUsed = Number(data.monthly_used) || 0;
       const serverDailyLimit = data.daily_limit != null ? Number(data.daily_limit) : null;
       const serverMonthlyLimit = data.monthly_limit != null ? Number(data.monthly_limit) : null;
-      this.budgetStatus = {
-        ...data,
-        // Use whichever spend is higher — server or local tracking
-        daily_used: Math.max(serverDailyUsed, localDailyUsed),
-        monthly_used: Math.max(serverMonthlyUsed, localMonthlyUsed),
-        // Use the stricter (lower) limit — local config takes priority if lower
-        daily_limit: this.pickStricterLimit(serverDailyLimit, this.localDailyLimit),
-        monthly_limit: this.pickStricterLimit(serverMonthlyLimit, this.localMonthlyLimit)
-      };
+      const snapshotDaily = this.budgetStatus.daily_used;
+      const snapshotMonthly = this.budgetStatus.monthly_used;
+      const mergedDaily = Math.max(serverDailyUsed, snapshotDaily);
+      const mergedMonthly = Math.max(serverMonthlyUsed, snapshotMonthly);
+      this.budgetStatus.daily_used = mergedDaily + (this.budgetStatus.daily_used - snapshotDaily);
+      this.budgetStatus.monthly_used = mergedMonthly + (this.budgetStatus.monthly_used - snapshotMonthly);
+      this.budgetStatus.daily_limit = this.pickStricterLimit(serverDailyLimit, this.localDailyLimit);
+      this.budgetStatus.monthly_limit = this.pickStricterLimit(serverMonthlyLimit, this.localMonthlyLimit);
+      this.budgetStatus.status = data.status;
+      this.budgetStatus.kill_switch_active = data.kill_switch_active;
       this.consecutiveFailures = 0;
       const newStatus = data.kill_switch_active ? "OFF" : data.status;
       if (newStatus !== this.status) {
@@ -349,18 +414,98 @@ function wrapAnthropic(client, config, shipper, heartbeat) {
               });
             }
           }
-          if (config.mode === "cost_protect") {
-            if (heartbeat.isMonthlyBudgetExceeded() || heartbeat.isDailyBudgetExceeded()) {
-              blocked = true;
-              enforcementReason = heartbeat.isDailyBudgetExceeded() ? "daily_budget_exceeded" : "monthly_budget_exceeded";
-              const blockError = {
-                code: "BUDGET_EXCEEDED",
-                message: `AI budget exceeded. ${enforcementReason.replace(/_/g, " ")}.`,
-                budget_used: heartbeat.budgetStatus.monthly_used,
-                budget_limit: heartbeat.budgetStatus.monthly_limit ?? 0,
-                retry_after: getNextResetTime(enforcementReason),
-                dashboard_url: dashboardUrl
-              };
+          const release = await heartbeat.budgetMutex.acquire();
+          try {
+            if (config.mode === "cost_protect") {
+              if (heartbeat.isMonthlyBudgetExceeded() || heartbeat.isDailyBudgetExceeded()) {
+                blocked = true;
+                enforcementReason = heartbeat.isDailyBudgetExceeded() ? "daily_budget_exceeded" : "monthly_budget_exceeded";
+                const blockError = {
+                  code: "BUDGET_EXCEEDED",
+                  message: `AI budget exceeded. ${enforcementReason.replace(/_/g, " ")}.`,
+                  budget_used: heartbeat.budgetStatus.monthly_used,
+                  budget_limit: heartbeat.budgetStatus.monthly_limit ?? 0,
+                  retry_after: getNextResetTime(enforcementReason),
+                  dashboard_url: dashboardUrl
+                };
+                shipper.push({
+                  id: generateId(),
+                  timestamp: startTime,
+                  provider: "anthropic",
+                  model,
+                  input_tokens: 0,
+                  output_tokens: 0,
+                  cost: 0,
+                  latency_ms: Date.now() - startTime,
+                  endpoint_tag: config.endpoint_tag,
+                  downgraded: false,
+                  blocked: true,
+                  enforcement_reason: enforcementReason
+                });
+                throw Object.assign(new Error(blockError.message), {
+                  caplyr: blockError
+                });
+              }
+              if (heartbeat.isDowngradeThresholdReached(downgradeThreshold)) {
+                const fallback = config.fallback ?? getDefaultFallback(model);
+                if (fallback && fallback !== model) {
+                  originalModel = model;
+                  model = fallback;
+                  downgraded = true;
+                  enforcementReason = "auto_downgrade_threshold";
+                  config.onEnforcement?.({
+                    type: "downgrade",
+                    timestamp: Date.now(),
+                    reason: `Budget at ${Math.round(downgradeThreshold * 100)}% \u2014 downgraded ${originalModel} \u2192 ${model}`,
+                    original_model: originalModel,
+                    fallback_model: model,
+                    budget_used: heartbeat.budgetStatus.monthly_used,
+                    budget_limit: heartbeat.budgetStatus.monthly_limit ?? 0,
+                    estimated_savings: 0
+                    // Calculated after response
+                  });
+                }
+              }
+            }
+            const requestParams = downgraded ? { ...params, model } : params;
+            try {
+              const response = await target.create.call(
+                target,
+                requestParams,
+                options
+              );
+              const latency = Date.now() - startTime;
+              const inputTokens = response?.usage?.input_tokens ?? 0;
+              const outputTokens = response?.usage?.output_tokens ?? 0;
+              const cost = calculateCost(model, inputTokens, outputTokens);
+              heartbeat.trackSpend(cost);
+              let estimatedSavings = 0;
+              if (downgraded && originalModel) {
+                const originalCost = calculateCost(
+                  originalModel,
+                  inputTokens,
+                  outputTokens
+                );
+                estimatedSavings = originalCost - cost;
+              }
+              shipper.push({
+                id: generateId(),
+                timestamp: startTime,
+                provider: "anthropic",
+                model,
+                input_tokens: inputTokens,
+                output_tokens: outputTokens,
+                cost,
+                latency_ms: latency,
+                endpoint_tag: config.endpoint_tag,
+                downgraded,
+                original_model: originalModel,
+                blocked: false,
+                enforcement_reason: enforcementReason
+              });
+              return response;
+            } catch (err) {
+              if (err?.caplyr) throw err;
               shipper.push({
                 id: generateId(),
                 timestamp: startTime,
@@ -371,90 +516,15 @@ function wrapAnthropic(client, config, shipper, heartbeat) {
                 cost: 0,
                 latency_ms: Date.now() - startTime,
                 endpoint_tag: config.endpoint_tag,
-                downgraded: false,
-                blocked: true,
-                enforcement_reason: enforcementReason
+                downgraded,
+                original_model: originalModel,
+                blocked: false,
+                enforcement_reason: "provider_error"
               });
-              throw Object.assign(new Error(blockError.message), {
-                caplyr: blockError
-              });
-            }
-            if (heartbeat.isDowngradeThresholdReached(downgradeThreshold)) {
-              const fallback = config.fallback ?? getDefaultFallback(model);
-              if (fallback && fallback !== model) {
-                originalModel = model;
-                model = fallback;
-                downgraded = true;
-                enforcementReason = "auto_downgrade_threshold";
-                config.onEnforcement?.({
-                  type: "downgrade",
-                  timestamp: Date.now(),
-                  reason: `Budget at ${Math.round(downgradeThreshold * 100)}% \u2014 downgraded ${originalModel} \u2192 ${model}`,
-                  original_model: originalModel,
-                  fallback_model: model,
-                  budget_used: heartbeat.budgetStatus.monthly_used,
-                  budget_limit: heartbeat.budgetStatus.monthly_limit ?? 0,
-                  estimated_savings: 0
-                  // Calculated after response
-                });
-              }
+              throw err;
             }
-          }
-          const requestParams = downgraded ? { ...params, model } : params;
-          try {
-            const response = await target.create.call(
-              target,
-              requestParams,
-              options
-            );
-            const latency = Date.now() - startTime;
-            const inputTokens = response?.usage?.input_tokens ?? 0;
-            const outputTokens = response?.usage?.output_tokens ?? 0;
-            const cost = calculateCost(model, inputTokens, outputTokens);
-            heartbeat.trackSpend(cost);
-            let estimatedSavings = 0;
-            if (downgraded && originalModel) {
-              const originalCost = calculateCost(
-                originalModel,
-                inputTokens,
-                outputTokens
-              );
-              estimatedSavings = originalCost - cost;
-            }
-            shipper.push({
-              id: generateId(),
-              timestamp: startTime,
-              provider: "anthropic",
-              model,
-              input_tokens: inputTokens,
-              output_tokens: outputTokens,
-              cost,
-              latency_ms: latency,
-              endpoint_tag: config.endpoint_tag,
-              downgraded,
-              original_model: originalModel,
-              blocked: false,
-              enforcement_reason: enforcementReason
-            });
-            return response;
-          } catch (err) {
-            if (err?.caplyr) throw err;
-            shipper.push({
-              id: generateId(),
-              timestamp: startTime,
-              provider: "anthropic",
-              model,
-              input_tokens: 0,
-              output_tokens: 0,
-              cost: 0,
-              latency_ms: Date.now() - startTime,
-              endpoint_tag: config.endpoint_tag,
-              downgraded,
-              original_model: originalModel,
-              blocked: false,
-              enforcement_reason: "provider_error"
-            });
-            throw err;
+          } finally {
+            release();
           }
         };
       }
@@ -531,18 +601,89 @@ function wrapOpenAI(client, config, shipper, heartbeat) {
               });
             }
           }
-          if (config.mode === "cost_protect") {
-            if (heartbeat.isMonthlyBudgetExceeded() || heartbeat.isDailyBudgetExceeded()) {
-              blocked = true;
-              enforcementReason = heartbeat.isDailyBudgetExceeded() ? "daily_budget_exceeded" : "monthly_budget_exceeded";
-              const blockError = {
-                code: "BUDGET_EXCEEDED",
-                message: `AI budget exceeded. ${enforcementReason.replace(/_/g, " ")}.`,
-                budget_used: heartbeat.budgetStatus.monthly_used,
-                budget_limit: heartbeat.budgetStatus.monthly_limit ?? 0,
-                retry_after: getNextResetTime2(enforcementReason),
-                dashboard_url: dashboardUrl
-              };
+          const release = await heartbeat.budgetMutex.acquire();
+          try {
+            if (config.mode === "cost_protect") {
+              if (heartbeat.isMonthlyBudgetExceeded() || heartbeat.isDailyBudgetExceeded()) {
+                blocked = true;
+                enforcementReason = heartbeat.isDailyBudgetExceeded() ? "daily_budget_exceeded" : "monthly_budget_exceeded";
+                const blockError = {
+                  code: "BUDGET_EXCEEDED",
+                  message: `AI budget exceeded. ${enforcementReason.replace(/_/g, " ")}.`,
+                  budget_used: heartbeat.budgetStatus.monthly_used,
+                  budget_limit: heartbeat.budgetStatus.monthly_limit ?? 0,
+                  retry_after: getNextResetTime2(enforcementReason),
+                  dashboard_url: dashboardUrl
+                };
+                shipper.push({
+                  id: generateId2(),
+                  timestamp: startTime,
+                  provider: "openai",
+                  model,
+                  input_tokens: 0,
+                  output_tokens: 0,
+                  cost: 0,
+                  latency_ms: Date.now() - startTime,
+                  endpoint_tag: config.endpoint_tag,
+                  downgraded: false,
+                  blocked: true,
+                  enforcement_reason: enforcementReason
+                });
+                throw Object.assign(new Error(blockError.message), {
+                  caplyr: blockError
+                });
+              }
+              if (heartbeat.isDowngradeThresholdReached(downgradeThreshold)) {
+                const fallback = config.fallback ?? getDefaultFallback(model);
+                if (fallback && fallback !== model) {
+                  originalModel = model;
+                  model = fallback;
+                  downgraded = true;
+                  enforcementReason = "auto_downgrade_threshold";
+                  config.onEnforcement?.({
+                    type: "downgrade",
+                    timestamp: Date.now(),
+                    reason: `Budget at ${Math.round(downgradeThreshold * 100)}% \u2014 downgraded ${originalModel} \u2192 ${model}`,
+                    original_model: originalModel,
+                    fallback_model: model,
+                    budget_used: heartbeat.budgetStatus.monthly_used,
+                    budget_limit: heartbeat.budgetStatus.monthly_limit ?? 0,
+                    estimated_savings: 0
+                  });
+                }
+              }
+            }
+            const requestParams = downgraded ? { ...params, model } : params;
+            try {
+              const response = await target.create.call(
+                target,
+                requestParams,
+                options
+              );
+              const latency = Date.now() - startTime;
+              const usage = response?.usage;
+              const inputTokens = usage?.prompt_tokens ?? 0;
+              const outputTokens = usage?.completion_tokens ?? 0;
+              const cost = calculateCost(model, inputTokens, outputTokens);
+              heartbeat.trackSpend(cost);
+              shipper.push({
+                id: generateId2(),
+                timestamp: startTime,
+                provider: "openai",
+                model,
+                input_tokens: inputTokens,
+                output_tokens: outputTokens,
+                cost,
+                latency_ms: latency,
+                endpoint_tag: config.endpoint_tag,
+                downgraded,
+                original_model: originalModel,
+                blocked: false,
+                enforcement_reason: enforcementReason
+              });
+              return response;
+            } catch (err) {
+              if (err?.caplyr) throw err;
               shipper.push({
                 id: generateId2(),
                 timestamp: startTime,
@@ -553,81 +694,15 @@ function wrapOpenAI(client, config, shipper, heartbeat) {
                 cost: 0,
                 latency_ms: Date.now() - startTime,
                 endpoint_tag: config.endpoint_tag,
-                downgraded: false,
-                blocked: true,
-                enforcement_reason: enforcementReason
-              });
-              throw Object.assign(new Error(blockError.message), {
-                caplyr: blockError
+                downgraded,
+                original_model: originalModel,
+                blocked: false,
+                enforcement_reason: "provider_error"
               });
+              throw err;
             }
-            if (heartbeat.isDowngradeThresholdReached(downgradeThreshold)) {
-              const fallback = config.fallback ?? getDefaultFallback(model);
-              if (fallback && fallback !== model) {
-                originalModel = model;
-                model = fallback;
-                downgraded = true;
-                enforcementReason = "auto_downgrade_threshold";
-                config.onEnforcement?.({
-                  type: "downgrade",
-                  timestamp: Date.now(),
-                  reason: `Budget at ${Math.round(downgradeThreshold * 100)}% \u2014 downgraded ${originalModel} \u2192 ${model}`,
-                  original_model: originalModel,
-                  fallback_model: model,
-                  budget_used: heartbeat.budgetStatus.monthly_used,
-                  budget_limit: heartbeat.budgetStatus.monthly_limit ?? 0,
-                  estimated_savings: 0
-                });
-              }
-            }
-          }
-          const requestParams = downgraded ? { ...params, model } : params;
-          try {
-            const response = await target.create.call(
-              target,
-              requestParams,
-              options
-            );
-            const latency = Date.now() - startTime;
-            const usage = response?.usage;
-            const inputTokens = usage?.prompt_tokens ?? 0;
-            const outputTokens = usage?.completion_tokens ?? 0;
-            const cost = calculateCost(model, inputTokens, outputTokens);
-            heartbeat.trackSpend(cost);
-            shipper.push({
-              id: generateId2(),
-              timestamp: startTime,
-              provider: "openai",
-              model,
-              input_tokens: inputTokens,
-              output_tokens: outputTokens,
-              cost,
-              latency_ms: latency,
-              endpoint_tag: config.endpoint_tag,
-              downgraded,
-              original_model: originalModel,
-              blocked: false,
-              enforcement_reason: enforcementReason
-            });
-            return response;
-          } catch (err) {
-            if (err?.caplyr) throw err;
-            shipper.push({
-              id: generateId2(),
-              timestamp: startTime,
-              provider: "openai",
-              model,
-              input_tokens: 0,
-              output_tokens: 0,
-              cost: 0,
-              latency_ms: Date.now() - startTime,
-              endpoint_tag: config.endpoint_tag,
-              downgraded,
-              original_model: originalModel,
-              blocked: false,
-              enforcement_reason: "provider_error"
-            });
-            throw err;
+          } finally {
+            release();
           }
         };
       }
@@ -690,6 +765,7 @@ function protect(client, config) {
     ...config
   };
   let shared = instances.get(resolvedConfig.apiKey);
+  const isExisting = !!shared;
   if (!shared) {
     const shipper2 = new LogShipper(resolvedConfig);
     const heartbeat2 = new Heartbeat(resolvedConfig);
@@ -699,7 +775,19 @@ function protect(client, config) {
   }
   const { shipper, heartbeat } = shared;
   if (resolvedConfig.budget) {
-    const budgetConfig = typeof resolvedConfig.budget === "number" ? { monthly: resolvedConfig.budget } : resolvedConfig.budget;
+    const raw = typeof resolvedConfig.budget === "number" ? { monthly: resolvedConfig.budget } : resolvedConfig.budget;
+    const budgetConfig = { ...raw };
+    if (isExisting) {
+      const current = heartbeat.budgetStatus;
+      if (budgetConfig.daily !== void 0) {
+        const existing = current.daily_limit;
+        budgetConfig.daily = existing !== null ? Math.min(existing, budgetConfig.daily) : budgetConfig.daily;
+      }
+      if (budgetConfig.monthly !== void 0) {
+        const existing = current.monthly_limit;
+        budgetConfig.monthly = existing !== null ? Math.min(existing, budgetConfig.monthly) : budgetConfig.monthly;
+      }
+    }
     heartbeat.applyLocalLimits(budgetConfig);
   }
   const provider = detectProvider(client);
@@ -741,14 +829,16 @@ async function shutdown(apiKey) {
     const shared = instances.get(apiKey);
     if (shared) {
       shared.heartbeat.destroy();
-      shared.shipper.destroy();
+      await shared.shipper.shutdown();
       instances.delete(apiKey);
     }
   } else {
+    const shutdowns = [];
     for (const [key, shared] of instances) {
       shared.heartbeat.destroy();
-      shared.shipper.destroy();
+      shutdowns.push(shared.shipper.shutdown());
     }
+    await Promise.all(shutdowns);
     instances.clear();
   }
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "caplyr",
-  "version": "0.2.3",
+  "version": "0.3.0",
   "description": "AI Cost Control Plane — budget guardrails, auto-downgrade, and kill switch for AI API calls",
   "main": "dist/index.js",
   "module": "dist/index.mjs",