@sentrial/sdk 0.4.0 → 0.4.2

package/dist/index.js

@@ -68,7 +68,22 @@ var ValidationError = class extends SentrialError {
 };
 
 // src/redact.ts
-import { createHash } from "crypto";
+var _createHash;
+try {
+  const mod = eval("require")("crypto");
+  if (mod?.createHash) {
+    _createHash = mod.createHash;
+  }
+} catch {
+}
+function getCreateHash() {
+  if (!_createHash) {
+    throw new Error(
+      'Sentrial PII hash mode requires Node.js crypto module. Use mode "label" or "remove" in browser/edge environments.'
+    );
+  }
+  return _createHash;
+}
 var DEFAULT_FIELDS = [
   "userInput",
   "assistantOutput",
@@ -95,7 +110,7 @@ var BUILTIN_PATTERNS = {
   ipAddresses: { pattern: IP_ADDRESS_PATTERN, label: "IP_ADDRESS" }
 };
 function hashValue(value) {
-  return createHash("sha256").update(value).digest("hex").slice(0, 6);
+  return getCreateHash()("sha256").update(value).digest("hex").slice(0, 6);
 }
 function replaceMatch(match, label, mode) {
   switch (mode) {
@@ -157,6 +172,59 @@ function redactPayload(payload, config) {
   return result;
 }
 
+// src/async-context.ts
+var _AsyncLocalStorage = null;
+try {
+  const mod = eval("require")("node:async_hooks");
+  if (mod?.AsyncLocalStorage) {
+    _AsyncLocalStorage = mod.AsyncLocalStorage;
+  }
+} catch {
+}
+var NodeContextVar = class {
+  _storage;
+  _defaultValue;
+  constructor(defaultValue) {
+    this._storage = new _AsyncLocalStorage();
+    this._defaultValue = defaultValue;
+  }
+  get() {
+    const store = this._storage.getStore();
+    return store !== void 0 ? store : this._defaultValue;
+  }
+  set(value) {
+    const previous = this.get();
+    this._storage.enterWith(value);
+    return { _previous: previous };
+  }
+  reset(token) {
+    this._storage.enterWith(token._previous);
+  }
+};
+var SimpleContextVar = class {
+  _value;
+  constructor(defaultValue) {
+    this._value = defaultValue;
+  }
+  get() {
+    return this._value;
+  }
+  set(value) {
+    const previous = this._value;
+    this._value = value;
+    return { _previous: previous };
+  }
+  reset(token) {
+    this._value = token._previous;
+  }
+};
+function createContextVar(defaultValue) {
+  if (_AsyncLocalStorage) {
+    return new NodeContextVar(defaultValue);
+  }
+  return new SimpleContextVar(defaultValue);
+}
+
 // src/cost.ts
 var OPENAI_PRICING = {
   "gpt-5.2": { input: 5, output: 15 },
@@ -208,7 +276,8 @@ var GOOGLE_PRICING = {
   "gemini-1.0-pro": { input: 0.5, output: 1.5 }
 };
 function findModelKey(model, pricing) {
-  for (const key of Object.keys(pricing)) {
+  const keys = Object.keys(pricing).sort((a, b) => b.length - a.length);
+  for (const key of keys) {
     if (model.startsWith(key)) {
       return key;
     }
@@ -236,179 +305,783 @@ function calculateGoogleCost(params) {
   return calculateCost(inputTokens, outputTokens, GOOGLE_PRICING[modelKey]);
 }
 
-// src/types.ts
-var EventType = /* @__PURE__ */ ((EventType2) => {
-  EventType2["TOOL_CALL"] = "tool_call";
-  EventType2["LLM_DECISION"] = "llm_decision";
-  EventType2["STATE_CHANGE"] = "state_change";
-  EventType2["ERROR"] = "error";
-  return EventType2;
-})(EventType || {});
-
-// src/client.ts
-var DEFAULT_API_URL = "https://api.sentrial.com";
-var MAX_RETRIES = 3;
-var INITIAL_BACKOFF_MS = 500;
-var MAX_BACKOFF_MS = 8e3;
-var BACKOFF_MULTIPLIER = 2;
-var RETRYABLE_STATUS_CODES = /* @__PURE__ */ new Set([408, 429, 500, 502, 503, 504]);
-var REQUEST_TIMEOUT_MS = 1e4;
-var SentrialClient = class {
-  apiUrl;
-  apiKey;
-  failSilently;
-  piiConfig;
-  piiConfigNeedsHydration = false;
-  piiHydrationPromise;
-  currentState = {};
-  constructor(config = {}) {
-    this.apiUrl = (config.apiUrl ?? (typeof process !== "undefined" ? process.env?.SENTRIAL_API_URL : void 0) ?? DEFAULT_API_URL).replace(/\/$/, "");
-    this.apiKey = config.apiKey ?? (typeof process !== "undefined" ? process.env?.SENTRIAL_API_KEY : void 0);
-    this.failSilently = config.failSilently ?? true;
-    if (config.pii === true) {
-      this.piiConfig = { enabled: true };
-      this.piiConfigNeedsHydration = true;
-    } else if (config.pii && typeof config.pii === "object") {
-      this.piiConfig = config.pii;
-      this.piiConfigNeedsHydration = false;
-    }
+// src/wrappers.ts
+var _currentSessionId = createContextVar(null);
+var _currentClient = createContextVar(null);
+var _defaultClient = null;
+function setSessionContext(sessionId, client) {
+  _currentSessionId.set(sessionId);
+  if (client) {
+    _currentClient.set(client);
   }
-  /**
-   * Fetch the organization's PII config from the server.
-   *
-   * Called lazily on the first request when `pii: true` was passed to the constructor.
-   * Uses a single shared promise so concurrent requests don't trigger duplicate fetches.
-   */
-  async hydratePiiConfig() {
-    if (!this.piiConfigNeedsHydration) return;
-    if (this.piiHydrationPromise) {
-      await this.piiHydrationPromise;
-      return;
-    }
-    this.piiHydrationPromise = (async () => {
-      try {
-        const headers = {};
-        if (this.apiKey) {
-          headers["Authorization"] = `Bearer ${this.apiKey}`;
-        }
-        const controller = new AbortController();
-        const timeoutId = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS);
-        let response;
-        try {
-          response = await fetch(`${this.apiUrl}/api/sdk/pii-config`, {
-            method: "GET",
-            headers,
-            signal: controller.signal
-          });
-        } finally {
-          clearTimeout(timeoutId);
-        }
-        if (response.ok) {
-          const data = await response.json();
-          if (data.config) {
-            this.piiConfig = {
-              enabled: data.config.enabled,
-              mode: data.config.mode,
-              fields: data.config.fields,
-              builtinPatterns: data.config.builtinPatterns,
-              customPatterns: (data.config.customPatterns || []).map(
-                (cp) => ({
-                  pattern: new RegExp(cp.pattern, "g"),
-                  label: cp.label
-                })
-              ),
-              enhancedDetection: data.config.enhancedDetection
-            };
-          }
-        }
-      } catch {
-      }
-      this.piiConfigNeedsHydration = false;
-    })();
-    await this.piiHydrationPromise;
+}
+function clearSessionContext() {
+  _currentSessionId.set(null);
+  _currentClient.set(null);
+}
+function getSessionContext() {
+  return _currentSessionId.get();
+}
+function setDefaultClient(client) {
+  _defaultClient = client;
+}
+function _setSessionContextWithTokens(sessionId, client) {
+  const _sessionToken = _currentSessionId.set(sessionId);
+  const _clientToken = client ? _currentClient.set(client) : _currentClient.set(_currentClient.get());
+  return { _sessionToken, _clientToken };
+}
+function _restoreSessionContext(tokens) {
+  _currentSessionId.reset(tokens._sessionToken);
+  _currentClient.reset(tokens._clientToken);
+}
+function getTrackingClient() {
+  return _currentClient.get() ?? _defaultClient;
+}
+function wrapOpenAI(client, options = {}) {
+  const { trackWithoutSession = false } = options;
+  const chat = client.chat;
+  if (!chat?.completions?.create) {
+    console.warn("Sentrial: OpenAI client does not have chat.completions.create");
+    return client;
   }
-  /**
-   * Make an HTTP request with retry logic and exponential backoff.
-   *
-   * Retries on transient failures (network errors, timeouts, 429/5xx).
-   * Up to MAX_RETRIES attempts with exponential backoff.
-   */
-  async safeRequest(method, url, body) {
-    if (this.piiConfigNeedsHydration) {
-      await this.hydratePiiConfig();
+  const originalCreate = chat.completions.create.bind(chat.completions);
+  chat.completions.create = async function(...args) {
+    const startTime = Date.now();
+    const params = args[0] ?? {};
+    const messages = params.messages ?? [];
+    const model = params.model ?? "unknown";
+    const isStreaming = params.stream === true;
+    if (isStreaming && !params.stream_options?.include_usage) {
+      args[0] = { ...params, stream_options: { ...params.stream_options, include_usage: true } };
     }
-    let lastError;
-    let backoff = INITIAL_BACKOFF_MS;
-    for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
-      try {
-        const headers = {
-          "Content-Type": "application/json"
-        };
-        if (this.apiKey) {
-          headers["Authorization"] = `Bearer ${this.apiKey}`;
-        }
-        const finalBody = this.piiConfig && body && typeof body === "object" ? redactPayload(body, this.piiConfig) : body;
-        const controller = new AbortController();
-        const timeoutId = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS);
-        let response;
-        try {
-          response = await fetch(url, {
-            method,
-            headers,
-            body: finalBody ? JSON.stringify(finalBody) : void 0,
-            signal: controller.signal
-          });
-        } finally {
-          clearTimeout(timeoutId);
-        }
-        if (RETRYABLE_STATUS_CODES.has(response.status) && attempt < MAX_RETRIES) {
-          await this.sleep(backoff);
-          backoff = Math.min(backoff * BACKOFF_MULTIPLIER, MAX_BACKOFF_MS);
-          continue;
-        }
-        if (!response.ok) {
-          const errorBody = await response.text();
-          let errorData = {};
-          try {
-            errorData = JSON.parse(errorBody);
-          } catch {
-          }
-          const error = new ApiError(
-            errorData.error?.message || `HTTP ${response.status}: ${response.statusText}`,
-            response.status,
-            errorData.error?.code
-          );
-          if (this.failSilently) {
-            console.warn(`Sentrial: Request failed (${method} ${url}):`, error.message);
-            return null;
-          }
-          throw error;
-        }
-        return await response.json();
-      } catch (error) {
-        if (error instanceof ApiError) {
-          throw error;
-        }
-        lastError = error instanceof Error ? error : new Error(String(error));
-        if (attempt < MAX_RETRIES) {
-          await this.sleep(backoff);
-          backoff = Math.min(backoff * BACKOFF_MULTIPLIER, MAX_BACKOFF_MS);
-          continue;
-        }
+    try {
+      const response = await originalCreate(...args);
+      if (isStreaming) {
+        return wrapOpenAIStream(response, { startTime, messages, model, trackWithoutSession });
       }
+      const durationMs = Date.now() - startTime;
+      const promptTokens = response.usage?.prompt_tokens ?? 0;
+      const completionTokens = response.usage?.completion_tokens ?? 0;
+      const totalTokens = response.usage?.total_tokens ?? 0;
+      let outputContent = "";
+      if (response.choices?.[0]?.message?.content) {
+        outputContent = response.choices[0].message.content;
+      }
+      const cost = calculateOpenAICost({ model, inputTokens: promptTokens, outputTokens: completionTokens });
+      trackLLMCall({
+        provider: "openai",
+        model,
+        messages,
+        output: outputContent,
+        promptTokens,
+        completionTokens,
+        totalTokens,
+        cost,
+        durationMs,
+        trackWithoutSession
+      });
+      return response;
+    } catch (error) {
+      const durationMs = Date.now() - startTime;
+      trackLLMError({
+        provider: "openai",
+        model,
+        messages,
+        error,
+        durationMs,
+        trackWithoutSession
+      });
+      throw error;
     }
-    const networkError = new NetworkError(
-      lastError?.message ?? "Unknown network error",
-      lastError
-    );
-    if (this.failSilently) {
-      console.warn(`Sentrial: Request failed after ${MAX_RETRIES + 1} attempts (${method} ${url}):`, networkError.message);
-      return null;
-    }
-    throw networkError;
+  };
+  return client;
+}
+function wrapAnthropic(client, options = {}) {
+  const { trackWithoutSession = false } = options;
+  const messages = client.messages;
+  if (!messages?.create) {
+    console.warn("Sentrial: Anthropic client does not have messages.create");
+    return client;
   }
-  sleep(ms) {
-    return new Promise((resolve) => setTimeout(resolve, ms));
+  const originalCreate = messages.create.bind(messages);
+  messages.create = async function(...args) {
+    const startTime = Date.now();
+    const params = args[0] ?? {};
+    const inputMessages = params.messages ?? [];
+    const model = params.model ?? "unknown";
+    const system = params.system ?? "";
+    const isStreaming = params.stream === true;
+    try {
+      const response = await originalCreate(...args);
+      if (isStreaming) {
+        return wrapAnthropicStream(response, {
+          startTime,
+          messages: inputMessages,
+          model,
+          system,
+          trackWithoutSession
+        });
+      }
+      const durationMs = Date.now() - startTime;
+      const promptTokens = response.usage?.input_tokens ?? 0;
+      const completionTokens = response.usage?.output_tokens ?? 0;
+      const totalTokens = promptTokens + completionTokens;
+      let outputContent = "";
+      if (response.content) {
+        for (const block of response.content) {
+          if (block.type === "text") {
+            outputContent += block.text;
+          }
+        }
+      }
+      const cost = calculateAnthropicCost({ model, inputTokens: promptTokens, outputTokens: completionTokens });
+      const fullMessages = system ? [{ role: "system", content: system }, ...inputMessages] : inputMessages;
+      trackLLMCall({
+        provider: "anthropic",
+        model,
+        messages: fullMessages,
+        output: outputContent,
+        promptTokens,
+        completionTokens,
+        totalTokens,
+        cost,
+        durationMs,
+        trackWithoutSession
+      });
+      return response;
+    } catch (error) {
+      const durationMs = Date.now() - startTime;
+      trackLLMError({
+        provider: "anthropic",
+        model,
+        messages: inputMessages,
+        error,
+        durationMs,
+        trackWithoutSession
+      });
+      throw error;
+    }
+  };
+  return client;
+}
+function wrapGoogle(model, options = {}) {
+  const { trackWithoutSession = false } = options;
+  const originalGenerate = model.generateContent;
+  if (!originalGenerate) {
+    console.warn("Sentrial: Google model does not have generateContent");
+    return model;
+  }
+  model.generateContent = async function(...args) {
+    const startTime = Date.now();
+    const contents = args[0];
+    const modelName = model.model ?? "gemini-unknown";
+    const messages = googleContentsToMessages(contents);
+    try {
+      const response = await originalGenerate.apply(model, args);
+      const durationMs = Date.now() - startTime;
+      let promptTokens = 0;
+      let completionTokens = 0;
+      const usageMeta = response.response?.usageMetadata ?? response.usageMetadata;
+      if (usageMeta) {
+        promptTokens = usageMeta.promptTokenCount ?? 0;
+        completionTokens = usageMeta.candidatesTokenCount ?? 0;
+      }
+      const totalTokens = promptTokens + completionTokens;
+      let outputContent = "";
+      try {
+        outputContent = response.response?.text?.() ?? response.text?.() ?? "";
+      } catch {
+      }
+      const cost = calculateGoogleCost({ model: modelName, inputTokens: promptTokens, outputTokens: completionTokens });
+      trackLLMCall({
+        provider: "google",
+        model: modelName,
+        messages,
+        output: outputContent,
+        promptTokens,
+        completionTokens,
+        totalTokens,
+        cost,
+        durationMs,
+        trackWithoutSession
+      });
+      return response;
+    } catch (error) {
+      const durationMs = Date.now() - startTime;
+      trackLLMError({
+        provider: "google",
+        model: modelName,
+        messages,
+        error,
+        durationMs,
+        trackWithoutSession
+      });
+      throw error;
+    }
+  };
+  return model;
+}
+function googleContentsToMessages(contents) {
+  if (typeof contents === "string") {
+    return [{ role: "user", content: contents }];
+  }
+  if (Array.isArray(contents)) {
+    return contents.map((item) => {
+      if (typeof item === "string") {
+        return { role: "user", content: item };
+      }
+      if (item && typeof item === "object") {
+        return { role: item.role ?? "user", content: String(item.content ?? item) };
+      }
+      return { role: "user", content: String(item) };
+    });
+  }
+  return [{ role: "user", content: String(contents) }];
+}
+function wrapLLM(client, provider) {
+  if (provider === "openai" || client.chat?.completions?.create) {
+    return wrapOpenAI(client);
+  }
+  if (provider === "anthropic" || client.messages?.create) {
+    return wrapAnthropic(client);
+  }
+  if (provider === "google" || client.generateContent) {
+    return wrapGoogle(client);
+  }
+  console.warn("Sentrial: Unknown LLM client type. No auto-tracking applied.");
+  return client;
+}
+function wrapOpenAIStream(stream, ctx) {
+  let fullContent = "";
+  let usage = null;
+  let tracked = false;
+  const originalIterator = stream[Symbol.asyncIterator]?.bind(stream);
+  if (!originalIterator) return stream;
+  const trackResult = () => {
+    if (tracked) return;
+    tracked = true;
+    const durationMs = Date.now() - ctx.startTime;
+    const promptTokens = usage?.prompt_tokens ?? 0;
+    const completionTokens = usage?.completion_tokens ?? 0;
+    const totalTokens = usage?.total_tokens ?? promptTokens + completionTokens;
+    const cost = calculateOpenAICost({ model: ctx.model, inputTokens: promptTokens, outputTokens: completionTokens });
+    trackLLMCall({
+      provider: "openai",
+      model: ctx.model,
+      messages: ctx.messages,
+      output: fullContent,
+      promptTokens,
+      completionTokens,
+      totalTokens,
+      cost,
+      durationMs,
+      trackWithoutSession: ctx.trackWithoutSession
+    });
+  };
+  return new Proxy(stream, {
+    get(target, prop, receiver) {
+      if (prop === Symbol.asyncIterator) {
+        return function() {
+          const iter = originalIterator();
+          return {
+            async next() {
+              const result = await iter.next();
+              if (!result.done) {
+                const chunk = result.value;
+                const delta = chunk.choices?.[0]?.delta?.content;
+                if (delta) fullContent += delta;
+                if (chunk.usage) usage = chunk.usage;
+              } else {
+                trackResult();
+              }
+              return result;
+            },
+            async return(value) {
+              trackResult();
+              return iter.return?.(value) ?? { done: true, value: void 0 };
+            },
+            async throw(error) {
+              return iter.throw?.(error) ?? { done: true, value: void 0 };
+            }
+          };
+        };
+      }
+      return Reflect.get(target, prop, receiver);
+    }
+  });
+}
+function wrapAnthropicStream(stream, ctx) {
+  let fullContent = "";
+  let inputTokens = 0;
+  let outputTokens = 0;
+  let tracked = false;
+  const originalIterator = stream[Symbol.asyncIterator]?.bind(stream);
+  if (!originalIterator) return stream;
+  const trackResult = () => {
+    if (tracked) return;
+    tracked = true;
+    const durationMs = Date.now() - ctx.startTime;
+    const totalTokens = inputTokens + outputTokens;
+    const cost = calculateAnthropicCost({ model: ctx.model, inputTokens, outputTokens });
+    const fullMessages = ctx.system ? [{ role: "system", content: ctx.system }, ...ctx.messages] : ctx.messages;
+    trackLLMCall({
+      provider: "anthropic",
+      model: ctx.model,
+      messages: fullMessages,
+      output: fullContent,
+      promptTokens: inputTokens,
+      completionTokens: outputTokens,
+      totalTokens,
+      cost,
+      durationMs,
+      trackWithoutSession: ctx.trackWithoutSession
+    });
+  };
+  return new Proxy(stream, {
+    get(target, prop, receiver) {
+      if (prop === Symbol.asyncIterator) {
+        return function() {
+          const iter = originalIterator();
+          return {
+            async next() {
+              const result = await iter.next();
+              if (!result.done) {
+                const event = result.value;
+                if (event.type === "content_block_delta" && event.delta?.text) {
+                  fullContent += event.delta.text;
+                }
+                if (event.type === "message_start" && event.message?.usage) {
+                  inputTokens = event.message.usage.input_tokens ?? 0;
+                }
+                if (event.type === "message_delta" && event.usage) {
+                  outputTokens = event.usage.output_tokens ?? 0;
+                }
+              } else {
+                trackResult();
+              }
+              return result;
+            },
+            async return(value) {
+              trackResult();
+              return iter.return?.(value) ?? { done: true, value: void 0 };
+            },
+            async throw(error) {
+              return iter.throw?.(error) ?? { done: true, value: void 0 };
+            }
+          };
+        };
+      }
+      return Reflect.get(target, prop, receiver);
+    }
+  });
+}
+function trackLLMCall(params) {
+  const client = getTrackingClient();
+  if (!client) return;
+  const sessionId = _currentSessionId.get();
+  if (!sessionId && !params.trackWithoutSession) {
+    return;
+  }
+  if (sessionId) {
+    client.trackToolCall({
+      sessionId,
+      toolName: `llm:${params.provider}:${params.model}`,
+      toolInput: {
+        messages: params.messages,
+        model: params.model,
+        provider: params.provider
+      },
+      toolOutput: {
+        content: params.output,
+        tokens: {
+          prompt: params.promptTokens,
+          completion: params.completionTokens,
+          total: params.totalTokens
+        },
+        cost_usd: params.cost
+      },
+      reasoning: `LLM call to ${params.provider} ${params.model}`,
+      estimatedCost: params.cost,
+      tokenCount: params.totalTokens,
+      metadata: {
+        provider: params.provider,
+        model: params.model,
+        duration_ms: params.durationMs,
+        prompt_tokens: params.promptTokens,
+        completion_tokens: params.completionTokens
+      }
+    }).catch((err) => {
+      console.warn("Sentrial: Failed to track LLM call:", err.message);
+    });
+  } else if (params.trackWithoutSession) {
+    client.createSession({
+      name: `LLM: ${params.provider}/${params.model}`,
+      agentName: `${params.provider}-wrapper`,
+      userId: "anonymous"
+    }).then((sid) => {
+      if (!sid) return;
+      return client.trackToolCall({
+        sessionId: sid,
+        toolName: `llm:${params.provider}:${params.model}`,
+        toolInput: {
+          messages: params.messages,
+          model: params.model,
+          provider: params.provider
+        },
+        toolOutput: {
+          content: params.output,
+          tokens: {
+            prompt: params.promptTokens,
+            completion: params.completionTokens,
+            total: params.totalTokens
+          },
+          cost_usd: params.cost
+        },
+        estimatedCost: params.cost,
+        tokenCount: params.totalTokens,
+        metadata: {
+          provider: params.provider,
+          model: params.model,
+          duration_ms: params.durationMs
+        }
+      }).then(() => {
+        return client.completeSession({
+          sessionId: sid,
+          success: true,
+          estimatedCost: params.cost,
+          promptTokens: params.promptTokens,
+          completionTokens: params.completionTokens,
+          totalTokens: params.totalTokens,
+          durationMs: params.durationMs
+        });
+      });
+    }).catch((err) => {
+      console.warn("Sentrial: Failed to track standalone LLM call:", err.message);
+    });
+  }
+}
+function trackLLMError(params) {
+  const client = getTrackingClient();
+  if (!client) return;
+  const sessionId = _currentSessionId.get();
+  if (!sessionId && !params.trackWithoutSession) {
+    return;
+  }
+  if (!sessionId) return;
+  client.trackError({
+    sessionId,
+    errorMessage: params.error.message,
+    errorType: params.error.name,
+    toolName: `llm:${params.provider}:${params.model}`,
+    metadata: {
+      provider: params.provider,
+      model: params.model,
+      duration_ms: params.durationMs
+    }
+  }).catch((err) => {
+    console.warn("Sentrial: Failed to track LLM error:", err.message);
+  });
+}
+
+// src/batcher.ts
+var EventBatcher = class {
+  queue = [];
+  flushIntervalMs;
+  flushThreshold;
+  maxQueueSize;
+  timer = null;
+  sendFn;
+  flushing = false;
+  shutdownCalled = false;
+  exitHandler;
+  constructor(sendFn, config = {}) {
+    this.sendFn = sendFn;
+    this.flushIntervalMs = config.flushIntervalMs ?? 1e3;
+    this.flushThreshold = config.flushThreshold ?? 10;
+    this.maxQueueSize = config.maxQueueSize ?? 1e3;
+    this.timer = setInterval(() => {
+      void this.flush();
+    }, this.flushIntervalMs);
+    if (this.timer && typeof this.timer === "object" && "unref" in this.timer) {
+      this.timer.unref();
+    }
+    this.exitHandler = () => {
+      void this.shutdown();
+    };
+    if (typeof process !== "undefined" && process.on) {
+      process.on("beforeExit", this.exitHandler);
+    }
+  }
+  /**
+   * Enqueue an event for batched delivery.
+   *
+   * If the queue hits `flushThreshold`, an automatic flush is triggered.
+   * If the queue is full (`maxQueueSize`), the oldest event is dropped.
+   */
+  enqueue(method, url, body) {
+    if (this.shutdownCalled) return;
+    if (this.queue.length >= this.maxQueueSize) {
+      this.queue.shift();
+      if (typeof console !== "undefined") {
+        console.warn(
+          `Sentrial: Event queue full (${this.maxQueueSize}), dropping oldest event`
+        );
+      }
+    }
+    this.queue.push({ method, url, body });
+    if (this.queue.length >= this.flushThreshold) {
+      void this.flush();
+    }
+  }
+  /**
+   * Flush all queued events to the API.
+   *
+   * Drains the queue and fires all requests in parallel. Safe to call
+   * concurrently — only one flush runs at a time.
+   */
+  async flush() {
+    if (this.flushing || this.queue.length === 0) return;
+    this.flushing = true;
+    const batch = this.queue.splice(0, this.queue.length);
+    try {
+      await Promise.all(
+        batch.map(
+          (event) => this.sendFn(event.method, event.url, event.body).catch((err) => {
+            if (typeof console !== "undefined") {
+              console.warn("Sentrial: Batched event failed:", err);
+            }
+          })
+        )
+      );
+    } finally {
+      this.flushing = false;
+    }
+  }
+  /**
+   * Stop the batcher: clear the timer, flush remaining events, remove exit handler.
+   */
+  async shutdown() {
+    if (this.shutdownCalled) return;
+    this.shutdownCalled = true;
+    if (this.timer !== null) {
+      clearInterval(this.timer);
+      this.timer = null;
+    }
+    if (typeof process !== "undefined" && process.removeListener) {
+      process.removeListener("beforeExit", this.exitHandler);
+    }
+    this.flushing = false;
+    await this.flush();
+  }
+  /** Number of events currently queued. */
+  get size() {
+    return this.queue.length;
+  }
+};
+
+// src/types.ts
+var EventType = /* @__PURE__ */ ((EventType2) => {
+  EventType2["TOOL_CALL"] = "tool_call";
+  EventType2["LLM_DECISION"] = "llm_decision";
+  EventType2["STATE_CHANGE"] = "state_change";
+  EventType2["ERROR"] = "error";
+  return EventType2;
+})(EventType || {});
+
+// src/client.ts
+var DEFAULT_API_URL = "https://api.sentrial.com";
+var MAX_RETRIES = 3;
+var INITIAL_BACKOFF_MS = 500;
+var MAX_BACKOFF_MS = 8e3;
+var BACKOFF_MULTIPLIER = 2;
+var RETRYABLE_STATUS_CODES = /* @__PURE__ */ new Set([408, 429, 500, 502, 503, 504]);
+var REQUEST_TIMEOUT_MS = 1e4;
+var SentrialClient = class {
+  apiUrl;
+  apiKey;
+  failSilently;
+  piiConfig;
+  piiConfigNeedsHydration = false;
+  piiHydrationPromise;
+  _stateVar = createContextVar({});
+  batcher;
+  /** Per-session cost/token accumulator — populated by trackToolCall/trackDecision */
+  sessionAccumulators = /* @__PURE__ */ new Map();
+  get currentState() {
+    return this._stateVar.get();
+  }
+  set currentState(value) {
+    this._stateVar.set(value);
+  }
+  constructor(config = {}) {
+    this.apiUrl = (config.apiUrl ?? (typeof process !== "undefined" ? process.env?.SENTRIAL_API_URL : void 0) ?? DEFAULT_API_URL).replace(/\/$/, "");
+    this.apiKey = config.apiKey ?? (typeof process !== "undefined" ? process.env?.SENTRIAL_API_KEY : void 0);
+    this.failSilently = config.failSilently ?? true;
+    if (config.pii === true) {
+      this.piiConfig = { enabled: true };
+      this.piiConfigNeedsHydration = true;
+    } else if (config.pii && typeof config.pii === "object") {
+      this.piiConfig = config.pii;
+      this.piiConfigNeedsHydration = false;
+    }
+    if (config.batching?.enabled) {
+      this.batcher = new EventBatcher(
+        (method, url, body) => this.safeRequest(method, url, body),
+        config.batching
+      );
+    }
+  }
+  /**
+   * Fetch the organization's PII config from the server.
+   *
+   * Called lazily on the first request when `pii: true` was passed to the constructor.
+   * Uses a single shared promise so concurrent requests don't trigger duplicate fetches.
+   */
+  async hydratePiiConfig() {
+    if (!this.piiConfigNeedsHydration) return;
+    if (this.piiHydrationPromise) {
+      await this.piiHydrationPromise;
+      return;
+    }
+    this.piiHydrationPromise = (async () => {
+      try {
+        const headers = {};
+        if (this.apiKey) {
+          headers["Authorization"] = `Bearer ${this.apiKey}`;
+        }
+        const controller = new AbortController();
+        const timeoutId = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS);
+        let response;
+        try {
+          response = await fetch(`${this.apiUrl}/api/sdk/pii-config`, {
+            method: "GET",
+            headers,
+            signal: controller.signal
+          });
+        } finally {
+          clearTimeout(timeoutId);
+        }
+        if (response.ok) {
+          const data = await response.json();
+          if (data.config) {
+            this.piiConfig = {
+              enabled: data.config.enabled,
+              mode: data.config.mode,
+              fields: data.config.fields,
+              builtinPatterns: data.config.builtinPatterns,
+              customPatterns: (data.config.customPatterns || []).map(
+                (cp) => ({
+                  pattern: new RegExp(cp.pattern, "g"),
+                  label: cp.label
+                })
+              ),
+              enhancedDetection: data.config.enhancedDetection
+            };
+          }
+        }
+      } catch {
+      }
+      this.piiConfigNeedsHydration = false;
+    })();
+    await this.piiHydrationPromise;
+  }
+  /**
+   * Make an HTTP request with retry logic and exponential backoff.
+   *
+   * Retries on transient failures (network errors, timeouts, 429/5xx).
+   * Up to MAX_RETRIES attempts with exponential backoff.
+   */
+  async safeRequest(method, url, body) {
+    if (this.piiConfigNeedsHydration) {
+      await this.hydratePiiConfig();
+    }
+    let lastError;
+    let backoff = INITIAL_BACKOFF_MS;
+    for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
+      try {
+        const headers = {
+          "Content-Type": "application/json"
+        };
+        if (this.apiKey) {
+          headers["Authorization"] = `Bearer ${this.apiKey}`;
+        }
+        const finalBody = this.piiConfig && body && typeof body === "object" ? redactPayload(body, this.piiConfig) : body;
+        const controller = new AbortController();
+        const timeoutId = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS);
+        let response;
+        try {
+          response = await fetch(url, {
+            method,
+            headers,
+            body: finalBody ? JSON.stringify(finalBody) : void 0,
+            signal: controller.signal
+          });
+        } finally {
+          clearTimeout(timeoutId);
+        }
+        if (RETRYABLE_STATUS_CODES.has(response.status) && attempt < MAX_RETRIES) {
+          await this.sleep(backoff);
+          backoff = Math.min(backoff * BACKOFF_MULTIPLIER, MAX_BACKOFF_MS);
+          continue;
+        }
+        if (!response.ok) {
+          const errorBody = await response.text();
+          let errorData = {};
+          try {
+            errorData = JSON.parse(errorBody);
+          } catch {
+          }
+          const error = new ApiError(
+            errorData.error?.message || `HTTP ${response.status}: ${response.statusText}`,
+            response.status,
+            errorData.error?.code
+          );
+          if (this.failSilently) {
+            console.warn(`Sentrial: Request failed (${method} ${url}):`, error.message);
+            return null;
+          }
+          throw error;
+        }
+        return await response.json();
+      } catch (error) {
+        if (error instanceof ApiError) {
+          throw error;
+        }
+        lastError = error instanceof Error ? error : new Error(String(error));
+        if (attempt < MAX_RETRIES) {
+          await this.sleep(backoff);
+          backoff = Math.min(backoff * BACKOFF_MULTIPLIER, MAX_BACKOFF_MS);
+          continue;
+        }
+      }
+    }
+    const networkError = new NetworkError(
+      lastError?.message ?? "Unknown network error",
+      lastError
+    );
+    if (this.failSilently) {
+      console.warn(`Sentrial: Request failed after ${MAX_RETRIES + 1} attempts (${method} ${url}):`, networkError.message);
+      return null;
+    }
+    throw networkError;
+  }
+  sleep(ms) {
+    return new Promise((resolve) => setTimeout(resolve, ms));
+  }
+  accumulate(sessionId, cost, tokenCount, toolOutput) {
+    let acc = this.sessionAccumulators.get(sessionId);
+    if (!acc) {
+      acc = { cost: 0, promptTokens: 0, completionTokens: 0, totalTokens: 0 };
+      this.sessionAccumulators.set(sessionId, acc);
+    }
+    if (cost != null) acc.cost += cost;
+    if (tokenCount != null) acc.totalTokens += tokenCount;
+    const rawTokens = toolOutput?.tokens;
+    if (rawTokens && typeof rawTokens === "object" && !Array.isArray(rawTokens)) {
+      const tokens = rawTokens;
+      if (typeof tokens.prompt === "number") acc.promptTokens += tokens.prompt;
+      if (typeof tokens.completion === "number") acc.completionTokens += tokens.completion;
+    }
   }
   /**
    * Create a new session
@@ -444,6 +1117,7 @@ var SentrialClient = class {
    * @returns Event data
    */
   async trackToolCall(params) {
+    this.accumulate(params.sessionId, params.estimatedCost, params.tokenCount, params.toolOutput);
     const stateBefore = { ...this.currentState };
     this.currentState[`${params.toolName}_result`] = params.toolOutput;
     const payload = {
@@ -462,6 +1136,10 @@ var SentrialClient = class {
     if (params.traceId !== void 0) payload.traceId = params.traceId;
     if (params.spanId !== void 0) payload.spanId = params.spanId;
     if (params.metadata !== void 0) payload.metadata = params.metadata;
+    if (this.batcher) {
+      this.batcher.enqueue("POST", `${this.apiUrl}/api/sdk/events`, payload);
+      return null;
+    }
     return this.safeRequest("POST", `${this.apiUrl}/api/sdk/events`, payload);
   }
   /**
@@ -471,6 +1149,7 @@ var SentrialClient = class {
    * @returns Event data
    */
   async trackDecision(params) {
+    this.accumulate(params.sessionId, params.estimatedCost, params.tokenCount);
     const stateBefore = { ...this.currentState };
     const payload = {
       sessionId: params.sessionId,
@@ -486,6 +1165,10 @@ var SentrialClient = class {
     if (params.traceId !== void 0) payload.traceId = params.traceId;
     if (params.spanId !== void 0) payload.spanId = params.spanId;
     if (params.metadata !== void 0) payload.metadata = params.metadata;
+    if (this.batcher) {
+      this.batcher.enqueue("POST", `${this.apiUrl}/api/sdk/events`, payload);
+      return null;
+    }
     return this.safeRequest("POST", `${this.apiUrl}/api/sdk/events`, payload);
   }
   /**
@@ -512,6 +1195,10 @@ var SentrialClient = class {
     if (params.traceId !== void 0) payload.traceId = params.traceId;
     if (params.spanId !== void 0) payload.spanId = params.spanId;
     if (params.metadata !== void 0) payload.metadata = params.metadata;
+    if (this.batcher) {
+      this.batcher.enqueue("POST", `${this.apiUrl}/api/sdk/events`, payload);
+      return null;
+    }
     return this.safeRequest("POST", `${this.apiUrl}/api/sdk/events`, payload);
   }
   /**
@@ -557,6 +1244,10 @@ var SentrialClient = class {
     if (params.metadata) {
       payload.metadata = params.metadata;
     }
+    if (this.batcher) {
+      this.batcher.enqueue("POST", `${this.apiUrl}/api/sdk/events`, payload);
+      return null;
+    }
     return this.safeRequest("POST", `${this.apiUrl}/api/sdk/events`, payload);
   }
   /**
@@ -586,6 +1277,17 @@ var SentrialClient = class {
    * ```
    */
   async completeSession(params) {
+    if (this.batcher) {
+      await this.batcher.flush();
+    }
+    const acc = this.sessionAccumulators.get(params.sessionId);
+    if (acc) {
+      if (params.estimatedCost === void 0 && acc.cost > 0) params = { ...params, estimatedCost: acc.cost };
+      if (params.promptTokens === void 0 && acc.promptTokens > 0) params = { ...params, promptTokens: acc.promptTokens };
+      if (params.completionTokens === void 0 && acc.completionTokens > 0) params = { ...params, completionTokens: acc.completionTokens };
+      if (params.totalTokens === void 0 && acc.totalTokens > 0) params = { ...params, totalTokens: acc.totalTokens };
+      this.sessionAccumulators.delete(params.sessionId);
+    }
     const payload = {
       status: params.success !== false ? "completed" : "failed",
       success: params.success ?? true
@@ -606,6 +1308,27 @@ var SentrialClient = class {
       payload
     );
   }
+  /**
+   * Flush any queued events immediately.
+   *
+   * No-op if batching is not enabled.
+   */
+  async flush() {
+    if (this.batcher) {
+      await this.batcher.flush();
+    }
+  }
+  /**
+   * Shut down the event batcher, flushing remaining events.
+   *
+   * Call this before your process exits for a clean shutdown.
+   * No-op if batching is not enabled.
+   */
+  async shutdown() {
+    if (this.batcher) {
+      await this.batcher.shutdown();
+    }
+  }
   /**
    * Begin tracking an interaction (simplified API)
    *
@@ -642,13 +1365,18 @@ var SentrialClient = class {
     if (params.input) {
       this.currentState.input = params.input;
     }
+    let sessionTokens;
+    if (sessionId) {
+      sessionTokens = _setSessionContextWithTokens(sessionId, this);
+    }
     return new Interaction({
       client: this,
       sessionId,
       eventId,
       userId: params.userId,
       event: params.event,
-      userInput: params.input
+      userInput: params.input,
+      sessionTokens
     });
   }
   // Cost calculation static methods for convenience
@@ -665,12 +1393,15 @@ var Interaction = class {
   userId;
   /** Event name for this interaction */
   event;
+  startTime = Date.now();
   finished = false;
   success = true;
   failureReason;
   output;
   userInput;
   degraded;
+  /** Context tokens for restoring previous session context on finish() */
+  sessionTokens;
   constructor(config) {
     this.client = config.client;
     this.sessionId = config.sessionId;
@@ -679,6 +1410,7 @@ var Interaction = class {
     this.event = config.event;
     this.userInput = config.userInput;
     this.degraded = config.sessionId === null;
+    this.sessionTokens = config.sessionTokens;
   }
   /**
    * Set the output for this interaction
@@ -714,18 +1446,24 @@ var Interaction = class {
     }
     this.finished = true;
     const finalOutput = params.output ?? this.output;
-    return this.client.completeSession({
+    const result = await this.client.completeSession({
       sessionId: this.sessionId,
       success: params.success ?? this.success,
       failureReason: params.failureReason ?? this.failureReason,
       estimatedCost: params.estimatedCost,
       customMetrics: params.customMetrics,
+      durationMs: params.durationMs ?? Date.now() - this.startTime,
       promptTokens: params.promptTokens,
       completionTokens: params.completionTokens,
       totalTokens: params.totalTokens,
       userInput: this.userInput,
       assistantOutput: finalOutput
     });
+    if (this.sessionTokens) {
+      _restoreSessionContext(this.sessionTokens);
+      this.sessionTokens = void 0;
+    }
+    return result;
   }
   /**
    * Track a tool call within this interaction
@@ -785,16 +1523,24 @@ function configure(config) {
 function begin(params) {
   return getClient().begin(params);
 }
+async function flush() {
+  if (defaultClient) await defaultClient.flush();
+}
+async function shutdown() {
+  if (defaultClient) await defaultClient.shutdown();
+}
 var sentrial = {
   configure,
-  begin
+  begin,
+  flush,
+  shutdown
 };
 
 // src/vercel.ts
-var _defaultClient = null;
+var _defaultClient2 = null;
 var _globalConfig = {};
 function configureVercel(config) {
-  _defaultClient = new SentrialClient({
+  _defaultClient2 = new SentrialClient({
     apiKey: config.apiKey,
     apiUrl: config.apiUrl,
     failSilently: config.failSilently ?? true
@@ -806,10 +1552,10 @@ function configureVercel(config) {
   };
 }
 function getClient2() {
-  if (!_defaultClient) {
-    _defaultClient = new SentrialClient();
+  if (!_defaultClient2) {
+    _defaultClient2 = new SentrialClient();
   }
-  return _defaultClient;
+  return _defaultClient2;
 }
 function extractModelInfo(model) {
   const modelId = model.modelId || model.id || "unknown";
@@ -818,7 +1564,7 @@ function extractModelInfo(model) {
 }
 function guessProvider(modelId) {
   const id = modelId.toLowerCase();
-  if (id.includes("gpt") || id.includes("o1") || id.includes("o3") || id.includes("o4") || id.startsWith("chatgpt")) return "openai";
+  if (id.includes("gpt") || id.startsWith("o1") || id.startsWith("o3") || id.startsWith("o4") || id.startsWith("chatgpt")) return "openai";
   if (id.includes("claude")) return "anthropic";
   if (id.includes("gemini")) return "google";
   if (id.includes("mistral") || id.includes("mixtral") || id.includes("codestral") || id.includes("pixtral")) return "mistral";
@@ -844,7 +1590,7 @@ function calculateCostForCall(provider, modelId, promptTokens, completionTokens)
     case "mistral":
       return promptTokens / 1e6 * 2 + completionTokens / 1e6 * 6;
     default:
-      return promptTokens * 3e-6 + completionTokens * 6e-6;
+      return 0;
   }
 }
 function extractInput(params) {
@@ -977,15 +1723,14 @@ function wrapGenerateText(originalFn, client, config) {
       const result = await originalFn(wrappedParams);
       const durationMs = Date.now() - startTime;
       const resolvedModelId = result.response?.modelId || modelId;
-      const promptTokens = result.usage?.promptTokens || 0;
-      const completionTokens = result.usage?.completionTokens || 0;
-      const totalTokens = result.usage?.totalTokens || promptTokens + completionTokens;
+      const promptTokens = result.usage?.promptTokens ?? 0;
+      const completionTokens = result.usage?.completionTokens ?? 0;
+      const totalTokens = result.usage?.totalTokens ?? promptTokens + completionTokens;
       const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
       const steps = result.steps;
       if (steps && steps.length >= 1) {
-        for (let i = 0; i < steps.length; i++) {
-          const step = steps[i];
-          await client.trackEvent({
+        const stepPromises = steps.map(
+          (step, i) => client.trackEvent({
             sessionId,
             eventType: "llm_call",
             eventData: {
@@ -993,14 +1738,16 @@ function wrapGenerateText(originalFn, client, config) {
               provider,
               step: i + 1,
               total_steps: steps.length,
-              prompt_tokens: step.usage?.promptTokens || 0,
-              completion_tokens: step.usage?.completionTokens || 0,
-              total_tokens: step.usage?.totalTokens || 0,
+              prompt_tokens: step.usage?.promptTokens ?? 0,
+              completion_tokens: step.usage?.completionTokens ?? 0,
+              total_tokens: step.usage?.totalTokens ?? 0,
               finish_reason: step.finishReason,
               tool_calls: step.toolCalls?.map((tc) => tc.toolName)
             }
-          });
-        }
+          }).catch(() => {
+          })
+        );
+        await Promise.all(stepPromises);
       } else {
         await client.trackEvent({
           sessionId,
@@ -1078,8 +1825,10 @@ function wrapStreamText(originalFn, client, config) {
       tools: params.tools ? wrapToolsAsync(params.tools, sessionPromise, client) : void 0
     };
     const result = originalFn(wrappedParams);
+    const originalTextStream = result.textStream;
+    let fullText = "";
     let tracked = false;
-    async function trackCompletion(fullText, error) {
+    async function trackCompletion(text, error) {
       if (tracked) return;
       tracked = true;
       const durationMs = Date.now() - startTime;
@@ -1090,80 +1839,118 @@ function wrapStreamText(originalFn, client, config) {
           sessionId: sid,
           errorType: error.name || "Error",
           errorMessage: error.message || "Unknown error"
+        }).catch(() => {
         });
         await client.completeSession({
           sessionId: sid,
           success: false,
           failureReason: error.message || "Unknown error",
           durationMs
+        }).catch(() => {
         });
         return;
       }
+      let resolvedModelId = modelId;
+      try {
+        const resp = result.response ? await result.response : void 0;
+        if (resp?.modelId) resolvedModelId = resp.modelId;
+      } catch {
+      }
       let usage;
       try {
         usage = result.usage ? await result.usage : void 0;
       } catch {
       }
-      const promptTokens = usage?.promptTokens || 0;
-      const completionTokens = usage?.completionTokens || 0;
-      const totalTokens = usage?.totalTokens || promptTokens + completionTokens;
-      const cost = calculateCostForCall(provider, modelId, promptTokens, completionTokens);
-      await client.trackEvent({
-        sessionId: sid,
-        eventType: "llm_call",
-        eventData: {
-          model: modelId,
-          provider,
-          prompt_tokens: promptTokens,
-          completion_tokens: completionTokens,
-          total_tokens: totalTokens
-        }
-      });
-      await client.completeSession({
-        sessionId: sid,
-        success: true,
-        output: fullText,
-        durationMs,
-        estimatedCost: cost,
-        promptTokens,
-        completionTokens,
-        totalTokens
-      });
-    }
-    const textProp = result.text;
-    if (typeof textProp === "string") {
-      trackCompletion(textProp).catch(() => {
-      });
-    } else if (textProp != null && typeof textProp.then === "function") {
-      const originalTextPromise = textProp;
-      result.text = originalTextPromise.then((text) => {
-        trackCompletion(text).catch(() => {
+      let steps;
+      try {
+        steps = result.steps ? await result.steps : void 0;
+      } catch {
+      }
+      if (steps && steps.length >= 1) {
+        let totalPrompt = 0, totalCompletion = 0;
+        const stepPromises = steps.map((step, i) => {
+          const sp = step.usage?.promptTokens ?? 0;
+          const sc = step.usage?.completionTokens ?? 0;
+          totalPrompt += sp;
+          totalCompletion += sc;
+          return client.trackEvent({
+            sessionId: sid,
+            eventType: "llm_call",
+            eventData: {
+              model: resolvedModelId,
+              provider,
+              step: i + 1,
+              total_steps: steps.length,
+              prompt_tokens: sp,
+              completion_tokens: sc,
+              total_tokens: step.usage?.totalTokens ?? 0,
+              finish_reason: step.finishReason,
+              tool_calls: step.toolCalls?.map((tc) => tc.toolName)
+            }
+          }).catch(() => {
+          });
         });
-        return text;
-      }).catch((err) => {
-        trackCompletion("", err instanceof Error ? err : new Error(String(err))).catch(() => {
+        await Promise.all(stepPromises);
+        const promptTokens = usage?.promptTokens ?? totalPrompt;
+        const completionTokens = usage?.completionTokens ?? totalCompletion;
+        const totalTokens = usage?.totalTokens ?? promptTokens + completionTokens;
+        const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
+        await client.completeSession({
+          sessionId: sid,
+          success: true,
+          output: text,
+          durationMs,
+          estimatedCost: cost,
+          promptTokens,
+          completionTokens,
+          totalTokens
+        }).catch(() => {
         });
-        throw err;
-      });
-    } else {
-      const originalTextStream = result.textStream;
-      let fullText = "";
-      result.textStream = (async function* () {
-        try {
-          for await (const chunk of originalTextStream) {
-            fullText += chunk;
-            yield chunk;
+      } else {
+        const promptTokens = usage?.promptTokens ?? 0;
+        const completionTokens = usage?.completionTokens ?? 0;
+        const totalTokens = usage?.totalTokens ?? promptTokens + completionTokens;
+        const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
+        await client.trackEvent({
+          sessionId: sid,
+          eventType: "llm_call",
+          eventData: {
+            model: resolvedModelId,
+            provider,
+            prompt_tokens: promptTokens,
+            completion_tokens: completionTokens,
+            total_tokens: totalTokens
           }
-          await trackCompletion(fullText);
-        } catch (error) {
-          await trackCompletion(
-            "",
-            error instanceof Error ? error : new Error(String(error))
-          );
-          throw error;
-        }
-      })();
+        }).catch(() => {
+        });
+        await client.completeSession({
+          sessionId: sid,
+          success: true,
+          output: text,
+          durationMs,
+          estimatedCost: cost,
+          promptTokens,
+          completionTokens,
+          totalTokens
+        }).catch(() => {
+        });
+      }
     }
+    result.textStream = (async function* () {
+      try {
+        for await (const chunk of originalTextStream) {
+          fullText += chunk;
+          yield chunk;
+        }
+        await trackCompletion(fullText);
+      } catch (error) {
+        await trackCompletion(
+          fullText,
+          error instanceof Error ? error : new Error(String(error))
+        );
+        throw error;
+      }
+    })();
     return result;
   };
 }
@@ -1191,10 +1978,22 @@ function wrapGenerateObject(originalFn, client, config) {
       const result = await originalFn(params);
       const durationMs = Date.now() - startTime;
       const resolvedModelId = result.response?.modelId || modelId;
-      const promptTokens = result.usage?.promptTokens || 0;
-      const completionTokens = result.usage?.completionTokens || 0;
-      const totalTokens = result.usage?.totalTokens || promptTokens + completionTokens;
+      const promptTokens = result.usage?.promptTokens ?? 0;
+      const completionTokens = result.usage?.completionTokens ?? 0;
+      const totalTokens = result.usage?.totalTokens ?? promptTokens + completionTokens;
       const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
+      await client.trackEvent({
+        sessionId,
+        eventType: "llm_call",
+        eventData: {
+          model: resolvedModelId,
+          provider,
+          prompt_tokens: promptTokens,
+          completion_tokens: completionTokens,
+          total_tokens: totalTokens
+        }
+      }).catch(() => {
+      });
       await client.completeSession({
         sessionId,
         success: true,
@@ -1223,393 +2022,279 @@ function wrapGenerateObject(originalFn, client, config) {
     }
   };
 }
-function wrapStreamObject(originalFn, client, config) {
-  return (params) => {
-    const startTime = Date.now();
-    const { modelId, provider } = extractModelInfo(params.model);
-    const input = extractInput(params);
-    const sessionPromise = (async () => {
-      try {
-        const id = await client.createSession({
-          name: `streamObject: ${input.slice(0, 50)}${input.length > 50 ? "..." : ""}`,
-          agentName: config.defaultAgent ?? "vercel-ai-sdk",
-          userId: config.userId ?? "anonymous",
-          convoId: config.convoId,
-          metadata: {
-            model: modelId,
-            provider,
-            function: "streamObject"
-          }
-        });
-        if (id) {
-          client.setInput(id, input).catch(() => {
-          });
-        }
-        return id;
-      } catch {
-        return null;
-      }
-    })();
-    const result = originalFn(params);
-    if (result.object) {
-      const originalObjectPromise = result.object;
-      result.object = originalObjectPromise.then(async (obj) => {
-        const durationMs = Date.now() - startTime;
-        const sid = await sessionPromise;
-        if (sid) {
-          let usage;
-          try {
-            usage = result.usage ? await result.usage : void 0;
-          } catch {
-          }
-          const promptTokens = usage?.promptTokens || 0;
-          const completionTokens = usage?.completionTokens || 0;
-          const totalTokens = usage?.totalTokens || promptTokens + completionTokens;
-          const cost = calculateCostForCall(provider, modelId, promptTokens, completionTokens);
-          await client.completeSession({
-            sessionId: sid,
-            success: true,
-            output: JSON.stringify(obj),
-            durationMs,
-            estimatedCost: cost,
-            promptTokens,
-            completionTokens,
-            totalTokens
-          });
-        }
-        return obj;
-      }).catch(async (error) => {
-        const durationMs = Date.now() - startTime;
-        const sid = await sessionPromise;
-        if (sid) {
-          await client.trackError({
-            sessionId: sid,
-            errorType: error instanceof Error ? error.name : "Error",
-            errorMessage: error instanceof Error ? error.message : "Unknown error"
-          });
-          await client.completeSession({
-            sessionId: sid,
-            success: false,
-            failureReason: error instanceof Error ? error.message : "Unknown error",
-            durationMs
-          });
-        }
-        throw error;
-      });
-    }
-    return result;
-  };
-}
-function wrapAISDK(ai, options) {
-  const client = options?.client ?? getClient2();
-  const config = {
-    defaultAgent: options?.defaultAgent ?? _globalConfig.defaultAgent,
-    userId: options?.userId ?? _globalConfig.userId,
-    convoId: options?.convoId ?? _globalConfig.convoId
-  };
-  return {
-    generateText: ai.generateText ? wrapGenerateText(ai.generateText, client, config) : wrapGenerateText(
-      () => Promise.reject(new Error("generateText not available")),
-      client,
-      config
-    ),
-    streamText: ai.streamText ? wrapStreamText(ai.streamText, client, config) : wrapStreamText(() => ({ textStream: (async function* () {
-    })() }), client, config),
-    generateObject: ai.generateObject ? wrapGenerateObject(ai.generateObject, client, config) : wrapGenerateObject(
-      () => Promise.reject(new Error("generateObject not available")),
-      client,
-      config
-    ),
-    streamObject: ai.streamObject ? wrapStreamObject(ai.streamObject, client, config) : wrapStreamObject(() => ({}), client, config)
-  };
-}
-
-// src/wrappers.ts
-var _currentSessionId = null;
-var _currentClient = null;
-var _defaultClient2 = null;
-function setSessionContext(sessionId, client) {
-  _currentSessionId = sessionId;
-  if (client) {
-    _currentClient = client;
-  }
-}
-function clearSessionContext() {
-  _currentSessionId = null;
-  _currentClient = null;
-}
-function getSessionContext() {
-  return _currentSessionId;
-}
-function setDefaultClient(client) {
-  _defaultClient2 = client;
-}
-function getTrackingClient() {
-  return _currentClient ?? _defaultClient2;
-}
-function wrapOpenAI(client, options = {}) {
-  const { trackWithoutSession = false } = options;
-  const chat = client.chat;
-  if (!chat?.completions?.create) {
-    console.warn("Sentrial: OpenAI client does not have chat.completions.create");
-    return client;
-  }
-  const originalCreate = chat.completions.create.bind(chat.completions);
-  chat.completions.create = async function(...args) {
-    const startTime = Date.now();
-    const params = args[0] ?? {};
-    const messages = params.messages ?? [];
-    const model = params.model ?? "unknown";
-    try {
-      const response = await originalCreate(...args);
-      const durationMs = Date.now() - startTime;
-      const promptTokens = response.usage?.prompt_tokens ?? 0;
-      const completionTokens = response.usage?.completion_tokens ?? 0;
-      const totalTokens = response.usage?.total_tokens ?? 0;
-      let outputContent = "";
-      if (response.choices?.[0]?.message?.content) {
-        outputContent = response.choices[0].message.content;
-      }
-      const cost = calculateOpenAICost({ model, inputTokens: promptTokens, outputTokens: completionTokens });
-      trackLLMCall({
-        provider: "openai",
-        model,
-        messages,
-        output: outputContent,
-        promptTokens,
-        completionTokens,
-        totalTokens,
-        cost,
-        durationMs,
-        trackWithoutSession
-      });
-      return response;
-    } catch (error) {
-      const durationMs = Date.now() - startTime;
-      trackLLMError({
-        provider: "openai",
-        model,
-        messages,
-        error,
-        durationMs,
-        trackWithoutSession
-      });
-      throw error;
-    }
-  };
-  return client;
-}
-function wrapAnthropic(client, options = {}) {
-  const { trackWithoutSession = false } = options;
-  const messages = client.messages;
-  if (!messages?.create) {
-    console.warn("Sentrial: Anthropic client does not have messages.create");
-    return client;
-  }
-  const originalCreate = messages.create.bind(messages);
-  messages.create = async function(...args) {
-    const startTime = Date.now();
-    const params = args[0] ?? {};
-    const inputMessages = params.messages ?? [];
-    const model = params.model ?? "unknown";
-    const system = params.system ?? "";
-    try {
-      const response = await originalCreate(...args);
-      const durationMs = Date.now() - startTime;
-      const promptTokens = response.usage?.input_tokens ?? 0;
-      const completionTokens = response.usage?.output_tokens ?? 0;
-      const totalTokens = promptTokens + completionTokens;
-      let outputContent = "";
-      if (response.content) {
-        for (const block of response.content) {
-          if (block.type === "text") {
-            outputContent += block.text;
-          }
-        }
-      }
-      const cost = calculateAnthropicCost({ model, inputTokens: promptTokens, outputTokens: completionTokens });
-      const fullMessages = system ? [{ role: "system", content: system }, ...inputMessages] : inputMessages;
-      trackLLMCall({
-        provider: "anthropic",
-        model,
-        messages: fullMessages,
-        output: outputContent,
-        promptTokens,
-        completionTokens,
-        totalTokens,
-        cost,
-        durationMs,
-        trackWithoutSession
-      });
-      return response;
-    } catch (error) {
-      const durationMs = Date.now() - startTime;
-      trackLLMError({
-        provider: "anthropic",
-        model,
-        messages: inputMessages,
-        error,
-        durationMs,
-        trackWithoutSession
-      });
-      throw error;
-    }
-  };
-  return client;
-}
-function wrapGoogle(model, options = {}) {
-  const { trackWithoutSession = false } = options;
-  const originalGenerate = model.generateContent;
-  if (!originalGenerate) {
-    console.warn("Sentrial: Google model does not have generateContent");
-    return model;
-  }
-  model.generateContent = async function(...args) {
+function wrapStreamObject(originalFn, client, config) {
+  return (params) => {
     const startTime = Date.now();
-    const contents = args[0];
-    const modelName = model.model ?? "gemini-unknown";
-    const messages = googleContentsToMessages(contents);
-    try {
-      const response = await originalGenerate.apply(model, args);
+    const { modelId, provider } = extractModelInfo(params.model);
+    const input = extractInput(params);
+    const sessionPromise = (async () => {
+      try {
+        const id = await client.createSession({
+          name: `streamObject: ${input.slice(0, 50)}${input.length > 50 ? "..." : ""}`,
+          agentName: config.defaultAgent ?? "vercel-ai-sdk",
+          userId: config.userId ?? "anonymous",
+          convoId: config.convoId,
+          metadata: {
+            model: modelId,
+            provider,
+            function: "streamObject"
+          }
+        });
+        if (id) {
+          client.setInput(id, input).catch(() => {
+          });
+        }
+        return id;
+      } catch {
+        return null;
+      }
+    })();
+    const result = originalFn(params);
+    async function completeStreamObject(obj, error) {
       const durationMs = Date.now() - startTime;
-      let promptTokens = 0;
-      let completionTokens = 0;
-      if (response.usageMetadata) {
-        promptTokens = response.usageMetadata.promptTokenCount ?? 0;
-        completionTokens = response.usageMetadata.candidatesTokenCount ?? 0;
+      const sid = await sessionPromise;
+      if (!sid) return;
+      if (error) {
+        await client.trackError({
+          sessionId: sid,
+          errorType: error.name || "Error",
+          errorMessage: error.message || "Unknown error"
+        }).catch(() => {
+        });
+        await client.completeSession({
+          sessionId: sid,
+          success: false,
+          failureReason: error.message || "Unknown error",
+          durationMs
+        }).catch(() => {
+        });
+        return;
       }
-      const totalTokens = promptTokens + completionTokens;
-      let outputContent = "";
+      let usage;
       try {
-        outputContent = response.response?.text() ?? "";
+        usage = result.usage ? await result.usage : void 0;
       } catch {
       }
-      const cost = calculateGoogleCost({ model: modelName, inputTokens: promptTokens, outputTokens: completionTokens });
-      trackLLMCall({
-        provider: "google",
-        model: modelName,
-        messages,
-        output: outputContent,
+      const promptTokens = usage?.promptTokens ?? 0;
+      const completionTokens = usage?.completionTokens ?? 0;
+      const totalTokens = usage?.totalTokens ?? promptTokens + completionTokens;
+      const cost = calculateCostForCall(provider, modelId, promptTokens, completionTokens);
+      await client.trackEvent({
+        sessionId: sid,
+        eventType: "llm_call",
+        eventData: {
+          model: modelId,
+          provider,
+          prompt_tokens: promptTokens,
+          completion_tokens: completionTokens,
+          total_tokens: totalTokens
+        }
+      }).catch(() => {
+      });
+      await client.completeSession({
+        sessionId: sid,
+        success: true,
+        output: JSON.stringify(obj),
+        durationMs,
+        estimatedCost: cost,
         promptTokens,
         completionTokens,
-        totalTokens,
-        cost,
-        durationMs,
-        trackWithoutSession
+        totalTokens
+      }).catch(() => {
       });
-      return response;
-    } catch (error) {
-      const durationMs = Date.now() - startTime;
-      trackLLMError({
-        provider: "google",
-        model: modelName,
-        messages,
-        error,
-        durationMs,
-        trackWithoutSession
+    }
+    if (result.object) {
+      const originalObjectPromise = result.object;
+      result.object = originalObjectPromise.then(async (obj) => {
+        await completeStreamObject(obj);
+        return obj;
+      }).catch(async (error) => {
+        await completeStreamObject(void 0, error instanceof Error ? error : new Error(String(error)));
+        throw error;
+      });
+    } else if (result.usage) {
+      result.usage.then(async () => {
+        await completeStreamObject(void 0);
+      }).catch(async (error) => {
+        await completeStreamObject(void 0, error instanceof Error ? error : new Error(String(error)));
       });
-      throw error;
     }
+    return result;
   };
-  return model;
-}
-function googleContentsToMessages(contents) {
-  if (typeof contents === "string") {
-    return [{ role: "user", content: contents }];
-  }
-  if (Array.isArray(contents)) {
-    return contents.map((item) => {
-      if (typeof item === "string") {
-        return { role: "user", content: item };
-      }
-      if (item && typeof item === "object") {
-        return { role: item.role ?? "user", content: String(item.content ?? item) };
-      }
-      return { role: "user", content: String(item) };
-    });
-  }
-  return [{ role: "user", content: String(contents) }];
 }
-function wrapLLM(client, provider) {
-  if (provider === "openai" || client.chat?.completions?.create) {
-    return wrapOpenAI(client);
-  }
-  if (provider === "anthropic" || client.messages?.create) {
-    return wrapAnthropic(client);
-  }
-  if (provider === "google" || client.generateContent) {
-    return wrapGoogle(client);
-  }
-  console.warn("Sentrial: Unknown LLM client type. No auto-tracking applied.");
-  return client;
+function wrapAISDK(ai, options) {
+  const client = options?.client ?? getClient2();
+  const config = {
+    defaultAgent: options?.defaultAgent ?? _globalConfig.defaultAgent,
+    userId: options?.userId ?? _globalConfig.userId,
+    convoId: options?.convoId ?? _globalConfig.convoId
+  };
+  return {
+    generateText: ai.generateText ? wrapGenerateText(ai.generateText, client, config) : wrapGenerateText(
+      () => Promise.reject(new Error("generateText not available")),
+      client,
+      config
+    ),
+    streamText: ai.streamText ? wrapStreamText(ai.streamText, client, config) : wrapStreamText(() => ({ textStream: (async function* () {
+    })() }), client, config),
+    generateObject: ai.generateObject ? wrapGenerateObject(ai.generateObject, client, config) : wrapGenerateObject(
+      () => Promise.reject(new Error("generateObject not available")),
+      client,
+      config
+    ),
+    streamObject: ai.streamObject ? wrapStreamObject(ai.streamObject, client, config) : wrapStreamObject(() => ({}), client, config)
+  };
 }
-function trackLLMCall(params) {
-  const client = getTrackingClient();
-  if (!client) return;
-  const sessionId = _currentSessionId;
-  if (!sessionId && !params.trackWithoutSession) {
-    return;
-  }
-  if (sessionId) {
-    client.trackToolCall({
-      sessionId,
-      toolName: `llm:${params.provider}:${params.model}`,
-      toolInput: {
-        messages: params.messages,
-        model: params.model,
-        provider: params.provider
-      },
-      toolOutput: {
-        content: params.output,
-        tokens: {
-          prompt: params.promptTokens,
-          completion: params.completionTokens,
-          total: params.totalTokens
-        },
-        cost_usd: params.cost
-      },
-      reasoning: `LLM call to ${params.provider} ${params.model}`,
-      estimatedCost: params.cost,
-      tokenCount: params.totalTokens,
-      metadata: {
-        provider: params.provider,
-        model: params.model,
-        duration_ms: params.durationMs,
-        prompt_tokens: params.promptTokens,
-        completion_tokens: params.completionTokens
-      }
-    }).catch((err) => {
-      console.warn("Sentrial: Failed to track LLM call:", err.message);
+
+// src/claude-code.ts
+function wrapClaudeAgent(queryFn, wrapOptions) {
+  const {
+    client,
+    defaultAgent = "claude-agent",
+    userId = "anonymous",
+    convoId,
+    extraMetadata
+  } = wrapOptions;
+  return function wrappedQuery(params) {
+    const { prompt, options = {} } = params;
+    const startTime = Date.now();
+    let sessionId = null;
+    let resolveSessionReady;
+    const sessionReady = new Promise((resolve) => {
+      resolveSessionReady = resolve;
     });
-  }
-}
-function trackLLMError(params) {
-  const client = getTrackingClient();
-  if (!client) return;
-  const sessionId = _currentSessionId;
-  if (!sessionId && !params.trackWithoutSession) {
-    return;
-  }
-  if (sessionId) {
-    client.trackError({
-      sessionId,
-      errorMessage: params.error.message,
-      errorType: params.error.name,
-      toolName: `llm:${params.provider}:${params.model}`,
-      metadata: {
-        provider: params.provider,
-        model: params.model,
-        duration_ms: params.durationMs
+    const sessionName = typeof prompt === "string" ? `${defaultAgent}: ${prompt.slice(0, 100)}` : `${defaultAgent} session`;
+    const pendingToolCalls = [];
+    const sentrialToolHook = {
+      hooks: [
+        async (input, toolUseID, _opts) => {
+          await sessionReady;
+          if (!sessionId) return;
+          const toolOutput = input?.tool_response && typeof input.tool_response === "object" ? input.tool_response : { response: input?.tool_response ?? null };
+          const p = client.trackToolCall({
+            sessionId,
+            toolName: input?.tool_name ?? "unknown",
+            toolInput: input?.tool_input ?? {},
+            toolOutput,
+            metadata: { tool_use_id: toolUseID }
+          }).catch(() => {
+          });
+          pendingToolCalls.push(p);
+        }
+      ]
+    };
+    const sentrialToolFailureHook = {
+      hooks: [
+        async (input, toolUseID, _opts) => {
+          await sessionReady;
+          if (!sessionId) return;
+          const p = client.trackToolCall({
+            sessionId,
+            toolName: input?.tool_name ?? "unknown",
+            toolInput: input?.tool_input ?? {},
+            toolOutput: {},
+            toolError: { message: input?.error ?? "unknown error" },
+            metadata: { tool_use_id: toolUseID }
+          }).catch(() => {
+          });
+          pendingToolCalls.push(p);
+        }
+      ]
+    };
+    const mergedHooks = {
+      ...options.hooks ?? {}
+    };
+    const existingPostToolUse = mergedHooks.PostToolUse ?? [];
+    mergedHooks.PostToolUse = [...existingPostToolUse, sentrialToolHook];
+    const existingPostToolUseFailure = mergedHooks.PostToolUseFailure ?? [];
+    mergedHooks.PostToolUseFailure = [...existingPostToolUseFailure, sentrialToolFailureHook];
+    const mergedOptions = {
+      ...options,
+      hooks: mergedHooks
+    };
+    const generator = queryFn({ prompt, options: mergedOptions });
+    return (async function* () {
+      try {
+        for await (const message of generator) {
+          if (message.type === "system" && message.subtype === "init") {
+            const metadata = {
+              model: message.model,
+              tools: message.tools,
+              cwd: message.cwd,
+              mcp_servers: message.mcp_servers,
+              sdk_session_id: message.session_id,
+              ...extraMetadata ?? {}
+            };
+            try {
+              sessionId = await client.createSession({
+                name: sessionName,
+                agentName: defaultAgent,
+                userId,
+                convoId,
+                metadata
+              });
+            } catch {
+              sessionId = null;
+            }
+            resolveSessionReady();
+          }
+          if (message.type === "result" && sessionId) {
+            const isError = !!message.is_error;
+            const inputTokens = message.usage?.input_tokens ?? 0;
+            const outputTokens = message.usage?.output_tokens ?? 0;
+            let failureReason;
+            if (isError) {
+              if (message.errors && message.errors.length > 0) {
+                failureReason = message.errors.join("; ");
+              } else {
+                failureReason = message.subtype;
+              }
+            }
+            await Promise.allSettled(pendingToolCalls);
+            try {
+              await client.completeSession({
+                sessionId,
+                success: !isError,
+                failureReason,
+                estimatedCost: message.total_cost_usd,
+                promptTokens: inputTokens,
+                completionTokens: outputTokens,
+                totalTokens: inputTokens + outputTokens,
+                durationMs: message.duration_ms ?? Date.now() - startTime,
+                userInput: typeof prompt === "string" ? prompt : void 0,
+                output: message.result,
+                customMetrics: {
+                  num_turns: message.num_turns ?? 0,
+                  duration_api_ms: message.duration_api_ms ?? 0
+                }
+              });
+            } catch {
+            }
+          }
+          yield message;
+        }
+      } catch (error) {
+        if (sessionId) {
+          await Promise.allSettled(pendingToolCalls);
+          try {
+            await client.completeSession({
+              sessionId,
+              success: false,
+              failureReason: error instanceof Error ? error.message : String(error),
+              durationMs: Date.now() - startTime
+            });
+          } catch {
+          }
+        }
+        throw error;
       }
-    }).catch((err) => {
-      console.warn("Sentrial: Failed to track LLM error:", err.message);
-    });
-  }
+    })();
+  };
 }
 
 // src/decorators.ts
 var _defaultClient3 = null;
-var _currentInteraction = null;
+var _currentInteraction = createContextVar(null);
 function getClient3() {
   if (!_defaultClient3) {
     try {
@@ -1629,7 +2314,7 @@ function getCurrentSessionId() {
   return getSessionContext();
 }
 function getCurrentInteraction() {
-  return _currentInteraction;
+  return _currentInteraction.get();
 }
 function withTool(name, fn) {
   const isAsync = fn.constructor.name === "AsyncFunction";
@@ -1730,10 +2415,11 @@ function withSession(agentName, fn, options = {}) {
       input: userInput
     });
     const sessionId = interaction.getSessionId();
+    let sessionTokens;
     if (sessionId) {
-      setSessionContext(sessionId, client);
+      sessionTokens = _setSessionContextWithTokens(sessionId, client);
     }
-    _currentInteraction = interaction;
+    const interactionToken = _currentInteraction.set(interaction);
     try {
       const result = await fn(...args);
       let output;
@@ -1758,8 +2444,10 @@ function withSession(agentName, fn, options = {}) {
       });
       throw error;
     } finally {
-      clearSessionContext();
-      _currentInteraction = null;
+      if (sessionTokens) {
+        _restoreSessionContext(sessionTokens);
+      }
+      _currentInteraction.reset(interactionToken);
     }
   };
 }
@@ -1845,10 +2533,11 @@ function TrackSession(agentName, options) {
         input: userInput
       });
       const sessionId = interaction.getSessionId();
+      let sessionTokens;
       if (sessionId) {
-        setSessionContext(sessionId, client);
+        sessionTokens = _setSessionContextWithTokens(sessionId, client);
       }
-      _currentInteraction = interaction;
+      const interactionToken = _currentInteraction.set(interaction);
       try {
         const result = await originalMethod.apply(this, args);
         let output;
@@ -1873,8 +2562,10 @@ function TrackSession(agentName, options) {
         });
         throw error;
       } finally {
-        clearSessionContext();
-        _currentInteraction = null;
+        if (sessionTokens) {
+          _restoreSessionContext(sessionTokens);
+        }
+        _currentInteraction.reset(interactionToken);
       }
     };
     return descriptor;
@@ -1887,6 +2578,8 @@ var SessionContext = class {
   client;
   interaction = null;
   output;
+  sessionTokens;
+  interactionToken;
   constructor(options) {
     this.userId = options.userId;
     this.agent = options.agent;
@@ -1905,9 +2598,9 @@ var SessionContext = class {
     });
     const sessionId = this.interaction.getSessionId();
     if (sessionId) {
-      setSessionContext(sessionId, this.client);
+      this.sessionTokens = _setSessionContextWithTokens(sessionId, this.client);
     }
-    _currentInteraction = this.interaction;
+    this.interactionToken = _currentInteraction.set(this.interaction);
     return this;
   }
   /**
@@ -1927,8 +2620,12 @@ var SessionContext = class {
         failureReason: options?.error
       });
     }
-    clearSessionContext();
-    _currentInteraction = null;
+    if (this.sessionTokens) {
+      _restoreSessionContext(this.sessionTokens);
+    }
+    if (this.interactionToken) {
+      _currentInteraction.reset(this.interactionToken);
+    }
   }
   /**
    * Get the session ID
@@ -1982,30 +2679,31 @@ function serializeOutput(value) {
 }
 
 // src/context.ts
-var _experimentContext = null;
+var _experimentContext = createContextVar(null);
 function getSystemPrompt(defaultPrompt) {
-  if (_experimentContext?.systemPrompt) {
-    return _experimentContext.systemPrompt;
+  const ctx = _experimentContext.get();
+  if (ctx?.systemPrompt) {
+    return ctx.systemPrompt;
   }
   return defaultPrompt ?? "";
 }
 function getExperimentContext() {
-  return _experimentContext;
+  return _experimentContext.get();
 }
 function isExperimentMode() {
-  return _experimentContext !== null;
+  return _experimentContext.get() !== null;
 }
 function getVariantName() {
-  return _experimentContext?.variantName ?? null;
+  return _experimentContext.get()?.variantName ?? null;
 }
 function getExperimentId() {
-  return _experimentContext?.experimentId ?? null;
+  return _experimentContext.get()?.experimentId ?? null;
 }
 function setExperimentContext(context) {
-  _experimentContext = context;
+  _experimentContext.set(context);
 }
 function clearExperimentContext() {
-  _experimentContext = null;
+  _experimentContext.set(null);
 }
 
 // src/experiment.ts
@@ -2338,6 +3036,7 @@ var Experiment = class {
 };
 export {
   ApiError,
+  EventBatcher,
   EventType,
   Experiment,
   ExperimentRunTracker,
@@ -2357,6 +3056,7 @@ export {
   clearSessionContext,
   configure,
   configureVercel,
+  createContextVar,
   getCurrentInteraction,
   getCurrentSessionId,
   getExperimentContext,
@@ -2379,6 +3079,7 @@ export {
   withTool,
   wrapAISDK,
   wrapAnthropic,
+  wrapClaudeAgent,
   wrapGoogle,
   wrapLLM,
   wrapOpenAI