@sentrial/sdk 0.4.0 → 0.4.3

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,182 +305,832 @@ 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
+    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 = "";
+      const toolCalls = [];
+      const msg = response.choices?.[0]?.message;
+      if (msg?.content) {
+        outputContent = msg.content;
+      }
+      if (msg?.tool_calls) {
+        for (const tc of msg.tool_calls) {
+          toolCalls.push({
+            name: tc.function?.name ?? "unknown",
+            arguments: tc.function?.arguments ?? "{}"
           });
-        } 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 cost = calculateOpenAICost({ model, inputTokens: promptTokens, outputTokens: completionTokens });
+      trackLLMCall({
+        provider: "openai",
+        model,
+        messages,
+        output: outputContent,
+        toolCalls,
+        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;
-  }
-  sleep(ms) {
-    return new Promise((resolve) => setTimeout(resolve, ms));
+  };
+  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;
   }
-  /**
-   * Create a new session
+  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 = "";
+      const toolCalls = [];
+      if (response.content) {
+        for (const block of response.content) {
+          if (block.type === "text") {
+            outputContent += block.text;
+          } else if (block.type === "tool_use") {
+            toolCalls.push({
+              name: block.name ?? "unknown",
+              arguments: JSON.stringify(block.input ?? {})
+            });
+          }
+        }
+      }
+      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,
+        toolCalls,
+        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;
+  const toolCallMap = /* @__PURE__ */ new Map();
+  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,
+      toolCalls: Array.from(toolCallMap.values()),
+      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;
+                if (delta?.content) fullContent += delta.content;
+                if (delta?.tool_calls) {
+                  for (const tc of delta.tool_calls) {
+                    const idx = tc.index ?? 0;
+                    const existing = toolCallMap.get(idx) ?? { name: "", arguments: "" };
+                    if (tc.function?.name) existing.name = tc.function.name;
+                    if (tc.function?.arguments) existing.arguments += tc.function.arguments;
+                    toolCallMap.set(idx, existing);
+                  }
+                }
+                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;
+  const toolCallsById = /* @__PURE__ */ new Map();
+  let currentBlockIdx = -1;
+  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,
+      toolCalls: Array.from(toolCallsById.values()),
+      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_start") {
+                  currentBlockIdx = event.index ?? currentBlockIdx + 1;
+                  if (event.content_block?.type === "tool_use") {
+                    toolCallsById.set(currentBlockIdx, {
+                      name: event.content_block.name ?? "unknown",
+                      arguments: ""
+                    });
+                  }
+                }
+                if (event.type === "content_block_delta") {
+                  if (event.delta?.text) {
+                    fullContent += event.delta.text;
+                  }
+                  if (event.delta?.type === "input_json_delta" && event.delta?.partial_json) {
+                    const idx = event.index ?? currentBlockIdx;
+                    const existing = toolCallsById.get(idx);
+                    if (existing) {
+                      existing.arguments += event.delta.partial_json;
+                    }
+                  }
+                }
+                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;
+  }
+  const toolOutput = {
+    content: params.output,
+    tokens: {
+      prompt: params.promptTokens,
+      completion: params.completionTokens,
+      total: params.totalTokens
+    },
+    cost_usd: params.cost
+  };
+  if (params.toolCalls && params.toolCalls.length > 0) {
+    toolOutput.tool_calls = params.toolCalls;
+  }
+  if (sessionId) {
+    client.trackToolCall({
+      sessionId,
+      toolName: `llm:${params.provider}:${params.model}`,
+      toolInput: {
+        messages: params.messages,
+        model: params.model,
+        provider: params.provider
+      },
+      toolOutput,
+      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,
+        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
    *
    * @param params - Session creation parameters
    * @returns Session ID, or null if the request failed and failSilently is true
@@ -444,6 +1163,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 +1182,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 +1195,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 +1211,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 +1241,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 +1290,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 +1323,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 +1354,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 +1411,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 +1439,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 +1456,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 +1492,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 +1569,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,19 +1598,20 @@ 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";
-  const provider = model.provider || guessProvider(modelId);
+  const rawProvider = model.provider || "";
+  const provider = rawProvider.split(".")[0] || guessProvider(modelId);
   return { modelId, provider };
 }
 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 +1637,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) {
@@ -858,6 +1651,15 @@ function extractInput(params) {
   }
   return "";
 }
+function normalizeUsage(usage) {
+  if (!usage) return void 0;
+  const u = usage;
+  const promptTokens = (u.inputTokens ?? u.promptTokens ?? 0) || 0;
+  const completionTokens = (u.outputTokens ?? u.completionTokens ?? 0) || 0;
+  const totalTokens = (u.totalTokens ?? promptTokens + completionTokens) || 0;
+  if (promptTokens === 0 && completionTokens === 0 && totalTokens === 0) return void 0;
+  return { promptTokens, completionTokens, totalTokens };
+}
 function wrapTools(tools, sessionId, client) {
   if (!tools) return void 0;
   const wrappedTools = {};
@@ -910,295 +1712,126 @@ function wrapToolsAsync(tools, sessionPromise, client) {
         ...tool,
         execute: async (...args) => {
           const startTime = Date.now();
-          const sid = await sessionPromise;
           try {
             const result = await originalExecute(...args);
             const durationMs = Date.now() - startTime;
-            if (sid) {
-              client.trackToolCall({
-                sessionId: sid,
-                toolName,
-                toolInput: args[0],
-                toolOutput: result,
-                reasoning: `Tool executed in ${durationMs}ms`
-              }).catch(() => {
-              });
-            }
-            return result;
-          } catch (error) {
-            const durationMs = Date.now() - startTime;
-            if (sid) {
-              client.trackToolCall({
-                sessionId: sid,
-                toolName,
-                toolInput: args[0],
-                toolOutput: {},
-                toolError: { message: error instanceof Error ? error.message : "Unknown error" },
-                reasoning: `Tool failed after ${durationMs}ms`
-              }).catch(() => {
-              });
-            }
-            throw error;
-          }
-        }
-      };
-    } else {
-      wrappedTools[toolName] = tool;
-    }
-  }
-  return wrappedTools;
-}
-function wrapGenerateText(originalFn, client, config) {
-  return async (params) => {
-    const startTime = Date.now();
-    const { modelId, provider } = extractModelInfo(params.model);
-    const input = extractInput(params);
-    const sessionId = await client.createSession({
-      name: `generateText: ${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: "generateText",
-        ...params.maxSteps ? { maxSteps: params.maxSteps } : {}
-      }
-    });
-    if (!sessionId) {
-      return originalFn(params);
-    }
-    await client.setInput(sessionId, input);
-    const wrappedParams = {
-      ...params,
-      tools: wrapTools(params.tools, sessionId, client)
-    };
-    try {
-      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 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({
-            sessionId,
-            eventType: "llm_call",
-            eventData: {
-              model: resolvedModelId,
-              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,
-              finish_reason: step.finishReason,
-              tool_calls: step.toolCalls?.map((tc) => tc.toolName)
-            }
-          });
-        }
-      } else {
-        await client.trackEvent({
-          sessionId,
-          eventType: "llm_call",
-          eventData: {
-            model: resolvedModelId,
-            provider,
-            prompt_tokens: promptTokens,
-            completion_tokens: completionTokens,
-            total_tokens: totalTokens,
-            finish_reason: result.finishReason,
-            tool_calls: result.toolCalls?.map((tc) => tc.toolName)
-          }
-        });
-      }
-      await client.completeSession({
-        sessionId,
-        success: true,
-        output: result.text,
-        durationMs,
-        estimatedCost: cost,
-        promptTokens,
-        completionTokens,
-        totalTokens
-      });
-      return result;
-    } catch (error) {
-      const durationMs = Date.now() - startTime;
-      await client.trackError({
-        sessionId,
-        errorType: error instanceof Error ? error.name : "Error",
-        errorMessage: error instanceof Error ? error.message : "Unknown error"
-      });
-      await client.completeSession({
-        sessionId,
-        success: false,
-        failureReason: error instanceof Error ? error.message : "Unknown error",
-        durationMs
-      });
-      throw error;
-    }
-  };
-}
-function wrapStreamText(originalFn, client, config) {
-  return (params) => {
-    const startTime = Date.now();
-    const { modelId, provider } = extractModelInfo(params.model);
-    const input = extractInput(params);
-    let sessionId = null;
-    const sessionPromise = (async () => {
-      try {
-        const id = await client.createSession({
-          name: `streamText: ${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: "streamText"
-          }
-        });
-        sessionId = id;
-        if (id) {
-          client.setInput(id, input).catch(() => {
-          });
-        }
-        return id;
-      } catch {
-        return null;
-      }
-    })();
-    const wrappedParams = {
-      ...params,
-      tools: params.tools ? wrapToolsAsync(params.tools, sessionPromise, client) : void 0
-    };
-    const result = originalFn(wrappedParams);
-    let tracked = false;
-    async function trackCompletion(fullText, error) {
-      if (tracked) return;
-      tracked = true;
-      const durationMs = Date.now() - startTime;
-      const sid = sessionId || await sessionPromise;
-      if (!sid) return;
-      if (error) {
-        await client.trackError({
-          sessionId: sid,
-          errorType: error.name || "Error",
-          errorMessage: error.message || "Unknown error"
-        });
-        await client.completeSession({
-          sessionId: sid,
-          success: false,
-          failureReason: error.message || "Unknown error",
-          durationMs
-        });
-        return;
-      }
-      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(() => {
-        });
-        return text;
-      }).catch((err) => {
-        trackCompletion("", err instanceof Error ? err : new Error(String(err))).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;
+            sessionPromise.then((sid) => {
+              if (sid) {
+                client.trackToolCall({
+                  sessionId: sid,
+                  toolName,
+                  toolInput: args[0],
+                  toolOutput: result,
+                  reasoning: `Tool executed in ${durationMs}ms`
+                }).catch(() => {
+                });
+              }
+            });
+            return result;
+          } catch (error) {
+            const durationMs = Date.now() - startTime;
+            sessionPromise.then((sid) => {
+              if (sid) {
+                client.trackToolCall({
+                  sessionId: sid,
+                  toolName,
+                  toolInput: args[0],
+                  toolOutput: {},
+                  toolError: { message: error instanceof Error ? error.message : "Unknown error" },
+                  reasoning: `Tool failed after ${durationMs}ms`
+                }).catch(() => {
+                });
+              }
+            });
+            throw error;
           }
-          await trackCompletion(fullText);
-        } catch (error) {
-          await trackCompletion(
-            "",
-            error instanceof Error ? error : new Error(String(error))
-          );
-          throw error;
         }
-      })();
+      };
+    } else {
+      wrappedTools[toolName] = tool;
     }
-    return result;
-  };
+  }
+  return wrappedTools;
 }
-function wrapGenerateObject(originalFn, client, config) {
+function wrapGenerateText(originalFn, client, config) {
   return async (params) => {
     const startTime = Date.now();
     const { modelId, provider } = extractModelInfo(params.model);
     const input = extractInput(params);
     const sessionId = await client.createSession({
-      name: `generateObject: ${input.slice(0, 50)}${input.length > 50 ? "..." : ""}`,
+      name: `generateText: ${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: "generateObject"
+        function: "generateText",
+        ...params.maxSteps ? { maxSteps: params.maxSteps } : {}
       }
     });
     if (!sessionId) {
       return originalFn(params);
     }
     await client.setInput(sessionId, input);
+    const wrappedParams = {
+      ...params,
+      tools: wrapTools(params.tools, sessionId, client)
+    };
     try {
-      const result = await originalFn(params);
+      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 usage = normalizeUsage(result.usage);
+      const promptTokens = usage?.promptTokens ?? 0;
+      const completionTokens = usage?.completionTokens ?? 0;
+      const totalTokens = usage?.totalTokens ?? 0;
       const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
+      const steps = result.steps;
+      if (steps && steps.length >= 1) {
+        const stepPromises = steps.map((step) => {
+          const su = normalizeUsage(step.usage);
+          return client.trackToolCall({
+            sessionId,
+            toolName: `llm:${provider}:${resolvedModelId}`,
+            toolInput: { finishReason: step.finishReason },
+            toolOutput: {
+              text: step.text?.slice(0, 500),
+              tokens: {
+                prompt: su?.promptTokens ?? 0,
+                completion: su?.completionTokens ?? 0
+              }
+            },
+            estimatedCost: calculateCostForCall(provider, resolvedModelId, su?.promptTokens ?? 0, su?.completionTokens ?? 0),
+            tokenCount: su?.totalTokens,
+            metadata: {
+              tool_calls: step.toolCalls?.map((tc) => tc.toolName)
+            }
+          }).catch(() => {
+          });
+        });
+        await Promise.all(stepPromises);
+      } else {
+        await client.trackToolCall({
+          sessionId,
+          toolName: `llm:${provider}:${resolvedModelId}`,
+          toolInput: { finishReason: result.finishReason },
+          toolOutput: {
+            text: result.text?.slice(0, 500),
+            tokens: { prompt: promptTokens, completion: completionTokens }
+          },
+          estimatedCost: cost,
+          tokenCount: totalTokens,
+          metadata: {
+            tool_calls: result.toolCalls?.map((tc) => tc.toolName)
+          }
+        }).catch(() => {
+        });
+      }
       await client.completeSession({
         sessionId,
         success: true,
-        output: JSON.stringify(result.object),
+        output: result.text,
         durationMs,
         estimatedCost: cost,
         promptTokens,
@@ -1223,393 +1856,524 @@ function wrapGenerateObject(originalFn, client, config) {
     }
   };
 }
-function wrapStreamObject(originalFn, client, config) {
+function wrapStreamText(originalFn, client, config) {
   return (params) => {
     const startTime = Date.now();
     const { modelId, provider } = extractModelInfo(params.model);
     const input = extractInput(params);
-    const sessionPromise = (async () => {
+    let sessionId = null;
+    let resolveSessionReady;
+    const sessionReady = new Promise((resolve) => {
+      resolveSessionReady = resolve;
+    });
+    (async () => {
       try {
         const id = await client.createSession({
-          name: `streamObject: ${input.slice(0, 50)}${input.length > 50 ? "..." : ""}`,
+          name: `streamText: ${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"
+            function: "streamText"
           }
         });
+        sessionId = id;
         if (id) {
           client.setInput(id, input).catch(() => {
           });
         }
-        return id;
       } catch {
-        return null;
+        sessionId = null;
       }
+      resolveSessionReady();
     })();
-    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;
+    const userOnStepFinish = params.onStepFinish;
+    const userOnFinish = params.onFinish;
+    const userOnError = params.onError;
+    const wrappedParams = {
+      ...params,
+      tools: params.tools ? wrapToolsAsync(params.tools, sessionReady.then(() => sessionId), client) : void 0,
+      onStepFinish: async (step) => {
+        await sessionReady;
+        if (sessionId) {
+          const su = normalizeUsage(step.usage);
+          const resolvedStepModel = step.response?.modelId ?? modelId;
+          client.trackToolCall({
+            sessionId,
+            toolName: `llm:${provider}:${resolvedStepModel}`,
+            toolInput: { finishReason: step.finishReason },
+            toolOutput: {
+              text: step.text?.slice(0, 500),
+              tokens: {
+                prompt: su?.promptTokens ?? 0,
+                completion: su?.completionTokens ?? 0
+              }
+            },
+            estimatedCost: calculateCostForCall(provider, resolvedStepModel, su?.promptTokens ?? 0, su?.completionTokens ?? 0),
+            tokenCount: su?.totalTokens,
+            metadata: {
+              tool_calls: step.toolCalls?.map((tc) => tc.toolName)
+            }
+          }).catch(() => {
+          });
+        }
+        if (userOnStepFinish) {
           try {
-            usage = result.usage ? await result.usage : void 0;
+            userOnStepFinish(step);
           } catch {
           }
-          const promptTokens = usage?.promptTokens || 0;
-          const completionTokens = usage?.completionTokens || 0;
-          const totalTokens = usage?.totalTokens || promptTokens + completionTokens;
-          const cost = calculateCostForCall(provider, modelId, promptTokens, completionTokens);
+        }
+      },
+      onFinish: async (event) => {
+        await sessionReady;
+        if (sessionId) {
+          const durationMs = Date.now() - startTime;
+          const usage = normalizeUsage(event.totalUsage ?? event.usage);
+          const resolvedModelId = event.response?.modelId ?? modelId;
+          const text = event.text ?? "";
+          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.completeSession({
-            sessionId: sid,
+            sessionId,
             success: true,
-            output: JSON.stringify(obj),
+            output: text,
             durationMs,
             estimatedCost: cost,
             promptTokens,
             completionTokens,
             totalTokens
+          }).catch(() => {
           });
         }
-        return obj;
-      }).catch(async (error) => {
-        const durationMs = Date.now() - startTime;
-        const sid = await sessionPromise;
-        if (sid) {
+        if (userOnFinish) {
+          try {
+            userOnFinish(event);
+          } catch {
+          }
+        }
+      },
+      onError: async (event) => {
+        await sessionReady;
+        if (sessionId) {
+          const durationMs = Date.now() - startTime;
+          const msg = event.error?.message ?? "Unknown error";
           await client.trackError({
-            sessionId: sid,
-            errorType: error instanceof Error ? error.name : "Error",
-            errorMessage: error instanceof Error ? error.message : "Unknown error"
+            sessionId,
+            errorType: event.error?.name ?? "Error",
+            errorMessage: msg
+          }).catch(() => {
           });
           await client.completeSession({
-            sessionId: sid,
+            sessionId,
             success: false,
-            failureReason: error instanceof Error ? error.message : "Unknown error",
+            failureReason: msg,
             durationMs
+          }).catch(() => {
           });
         }
-        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;
+        if (userOnError) {
+          try {
+            userOnError(event);
+          } catch {
+          }
+        }
       }
-      const cost = calculateOpenAICost({ model, inputTokens: promptTokens, outputTokens: completionTokens });
-      trackLLMCall({
-        provider: "openai",
-        model,
-        messages,
-        output: outputContent,
-        promptTokens,
-        completionTokens,
-        totalTokens,
-        cost,
-        durationMs,
-        trackWithoutSession
-      });
-      return response;
+    };
+    try {
+      return originalFn(wrappedParams);
     } catch (error) {
-      const durationMs = Date.now() - startTime;
-      trackLLMError({
-        provider: "openai",
-        model,
-        messages,
-        error,
-        durationMs,
-        trackWithoutSession
+      sessionReady.then(() => {
+        if (sessionId) {
+          const durationMs = Date.now() - startTime;
+          const msg = error instanceof Error ? error.message : String(error);
+          client.trackError({ sessionId, errorType: error instanceof Error ? error.name : "Error", errorMessage: msg }).catch(() => {
+          });
+          client.completeSession({ sessionId, success: false, failureReason: msg, durationMs }).catch(() => {
+          });
+        }
       });
       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) {
+function wrapGenerateObject(originalFn, client, config) {
+  return async (params) => {
     const startTime = Date.now();
-    const params = args[0] ?? {};
-    const inputMessages = params.messages ?? [];
-    const model = params.model ?? "unknown";
-    const system = params.system ?? "";
+    const { modelId, provider } = extractModelInfo(params.model);
+    const input = extractInput(params);
+    const sessionId = await client.createSession({
+      name: `generateObject: ${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: "generateObject"
+      }
+    });
+    if (!sessionId) {
+      return originalFn(params);
+    }
+    await client.setInput(sessionId, input);
     try {
-      const response = await originalCreate(...args);
+      const result = await originalFn(params);
       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,
+      const resolvedModelId = result.response?.modelId || modelId;
+      const usage = normalizeUsage(result.usage);
+      const promptTokens = usage?.promptTokens ?? 0;
+      const completionTokens = usage?.completionTokens ?? 0;
+      const totalTokens = usage?.totalTokens ?? 0;
+      const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
+      await client.trackToolCall({
+        sessionId,
+        toolName: `llm:${provider}:${resolvedModelId}`,
+        toolInput: { function: "generateObject" },
+        toolOutput: {
+          object: result.object,
+          tokens: { prompt: promptTokens, completion: completionTokens }
+        },
+        estimatedCost: cost,
+        tokenCount: totalTokens
+      }).catch(() => {
+      });
+      await client.completeSession({
+        sessionId,
+        success: true,
+        output: JSON.stringify(result.object),
+        durationMs,
+        estimatedCost: cost,
         promptTokens,
         completionTokens,
-        totalTokens,
-        cost,
-        durationMs,
-        trackWithoutSession
+        totalTokens
       });
-      return response;
+      return result;
     } catch (error) {
       const durationMs = Date.now() - startTime;
-      trackLLMError({
-        provider: "anthropic",
-        model,
-        messages: inputMessages,
-        error,
-        durationMs,
-        trackWithoutSession
+      await client.trackError({
+        sessionId,
+        errorType: error instanceof Error ? error.name : "Error",
+        errorMessage: error instanceof Error ? error.message : "Unknown error"
+      });
+      await client.completeSession({
+        sessionId,
+        success: false,
+        failureReason: error instanceof Error ? error.message : "Unknown error",
+        durationMs
       });
       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 durationMs = Date.now() - startTime;
-      let promptTokens = 0;
-      let completionTokens = 0;
-      if (response.usageMetadata) {
-        promptTokens = response.usageMetadata.promptTokenCount ?? 0;
-        completionTokens = response.usageMetadata.candidatesTokenCount ?? 0;
-      }
-      const totalTokens = promptTokens + completionTokens;
-      let outputContent = "";
+    const { modelId, provider } = extractModelInfo(params.model);
+    const input = extractInput(params);
+    let sessionId = null;
+    let resolveSessionReady;
+    const sessionReady = new Promise((resolve) => {
+      resolveSessionReady = resolve;
+    });
+    (async () => {
       try {
-        outputContent = response.response?.text() ?? "";
+        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"
+          }
+        });
+        sessionId = id;
+        if (id) {
+          client.setInput(id, input).catch(() => {
+          });
+        }
       } catch {
+        sessionId = null;
       }
-      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;
+      resolveSessionReady();
+    })();
+    const userOnFinish = params.onFinish;
+    const userOnError = params.onError;
+    const wrappedParams = {
+      ...params,
+      onFinish: async (event) => {
+        await sessionReady;
+        if (sessionId) {
+          const durationMs = Date.now() - startTime;
+          const usage = normalizeUsage(event.usage);
+          const resolvedModelId = event.response?.modelId ?? modelId;
+          const promptTokens = usage?.promptTokens ?? 0;
+          const completionTokens = usage?.completionTokens ?? 0;
+          const totalTokens = usage?.totalTokens ?? 0;
+          const cost = calculateCostForCall(provider, resolvedModelId, promptTokens, completionTokens);
+          if (event.error) {
+            const errMsg = event.error instanceof Error ? event.error.message : String(event.error);
+            await client.trackError({
+              sessionId,
+              errorType: "SchemaValidationError",
+              errorMessage: errMsg
+            }).catch(() => {
+            });
+          }
+          await client.trackToolCall({
+            sessionId,
+            toolName: `llm:${provider}:${resolvedModelId}`,
+            toolInput: { function: "streamObject" },
+            toolOutput: {
+              object: event.object,
+              tokens: { prompt: promptTokens, completion: completionTokens }
+            },
+            estimatedCost: cost,
+            tokenCount: totalTokens
+          }).catch(() => {
+          });
+          await client.completeSession({
+            sessionId,
+            success: !event.error,
+            output: event.object != null ? JSON.stringify(event.object) : void 0,
+            failureReason: event.error ? String(event.error) : void 0,
+            durationMs,
+            estimatedCost: cost,
+            promptTokens,
+            completionTokens,
+            totalTokens
+          }).catch(() => {
+          });
+        }
+        if (userOnFinish) {
+          try {
+            userOnFinish(event);
+          } catch {
+          }
+        }
+      },
+      onError: async (event) => {
+        await sessionReady;
+        if (sessionId) {
+          const durationMs = Date.now() - startTime;
+          const msg = event.error?.message ?? "Unknown error";
+          await client.trackError({
+            sessionId,
+            errorType: event.error?.name ?? "Error",
+            errorMessage: msg
+          }).catch(() => {
+          });
+          await client.completeSession({
+            sessionId,
+            success: false,
+            failureReason: msg,
+            durationMs
+          }).catch(() => {
+          });
+        }
+        if (userOnError) {
+          try {
+            userOnError(event);
+          } catch {
+          }
+        }
+      }
+    };
+    try {
+      return originalFn(wrappedParams);
     } catch (error) {
-      const durationMs = Date.now() - startTime;
-      trackLLMError({
-        provider: "google",
-        model: modelName,
-        messages,
-        error,
-        durationMs,
-        trackWithoutSession
+      sessionReady.then(() => {
+        if (sessionId) {
+          const durationMs = Date.now() - startTime;
+          const msg = error instanceof Error ? error.message : String(error);
+          client.trackError({ sessionId, errorType: error instanceof Error ? error.name : "Error", errorMessage: msg }).catch(() => {
+          });
+          client.completeSession({ sessionId, success: false, failureReason: msg, durationMs }).catch(() => {
+          });
+        }
       });
       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 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 +2393,7 @@ function getCurrentSessionId() {
   return getSessionContext();
 }
 function getCurrentInteraction() {
-  return _currentInteraction;
+  return _currentInteraction.get();
 }
 function withTool(name, fn) {
   const isAsync = fn.constructor.name === "AsyncFunction";
@@ -1730,10 +2494,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 +2523,10 @@ function withSession(agentName, fn, options = {}) {
       });
       throw error;
     } finally {
-      clearSessionContext();
-      _currentInteraction = null;
+      if (sessionTokens) {
+        _restoreSessionContext(sessionTokens);
+      }
+      _currentInteraction.reset(interactionToken);
     }
   };
 }
@@ -1845,10 +2612,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 +2641,10 @@ function TrackSession(agentName, options) {
         });
         throw error;
       } finally {
-        clearSessionContext();
-        _currentInteraction = null;
+        if (sessionTokens) {
+          _restoreSessionContext(sessionTokens);
+        }
+        _currentInteraction.reset(interactionToken);
       }
     };
     return descriptor;
@@ -1887,6 +2657,8 @@ var SessionContext = class {
   client;
   interaction = null;
   output;
+  sessionTokens;
+  interactionToken;
   constructor(options) {
     this.userId = options.userId;
     this.agent = options.agent;
@@ -1905,9 +2677,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 +2699,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 +2758,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 +3115,7 @@ var Experiment = class {
 };
 export {
   ApiError,
+  EventBatcher,
   EventType,
   Experiment,
   ExperimentRunTracker,
@@ -2357,6 +3135,7 @@ export {
   clearSessionContext,
   configure,
   configureVercel,
+  createContextVar,
   getCurrentInteraction,
   getCurrentSessionId,
   getExperimentContext,
@@ -2379,6 +3158,7 @@ export {
   withTool,
   wrapAISDK,
   wrapAnthropic,
+  wrapClaudeAgent,
   wrapGoogle,
   wrapLLM,
   wrapOpenAI