neatlogs 1.1.6 → 1.1.7

package/dist/vertex-ai.mjs

@@ -1,18 +1,15 @@
 // src/vertex-ai.ts
-import { trace as trace2, context as otelContext4, SpanStatusCode as SpanStatusCode3 } from "@opentelemetry/api";
+import { SpanStatusCode as SpanStatusCode3 } from "@opentelemetry/api";
 
 // src/core/auto-root.ts
 import {
-  trace as otelTrace3,
-  context as otelContext3
+  trace as otelTrace3
 } from "@opentelemetry/api";
 
 // src/core/context.ts
 import {
-  trace as otelTrace2,
-  context as otelContext2,
   createContextKey,
-  ROOT_CONTEXT,
+  ROOT_CONTEXT as ROOT_CONTEXT3,
   SpanStatusCode as SpanStatusCode2
 } from "@opentelemetry/api";
 
@@ -67,9 +64,6 @@ function getLogger() {
 // src/core/mask.ts
 var logger = getLogger();
 
-// src/core/end-user.ts
-import { trace as otelTrace, context as otelContext } from "@opentelemetry/api";
-
 // src/core/identity.ts
 import { AsyncLocalStorage as AsyncLocalStorage2 } from "async_hooks";
 var _GLOBAL_KEY = /* @__PURE__ */ Symbol.for("neatlogs.identity.storage");
@@ -85,6 +79,51 @@ function currentEndUserMetadata() {
   return _storage.getStore()?.endUserMetadata;
 }
 
+// src/core/provider.ts
+import { AsyncLocalStorage as AsyncLocalStorage3 } from "async_hooks";
+import {
+  ROOT_CONTEXT,
+  INVALID_SPAN_CONTEXT,
+  trace as otelTrace
+} from "@opentelemetry/api";
+var PRIVATE_SPAN_STORAGE_KEY = /* @__PURE__ */ Symbol.for(
+  "neatlogs.private_span_async_local_storage"
+);
+var PRIVATE_PROVIDER_STATE_KEY = /* @__PURE__ */ Symbol.for(
+  "neatlogs.private_provider_state"
+);
+var neatlogsGlobal = globalThis;
+var privateContextStorage = neatlogsGlobal[PRIVATE_SPAN_STORAGE_KEY] ?? (neatlogsGlobal[PRIVATE_SPAN_STORAGE_KEY] = new AsyncLocalStorage3());
+var providerState = neatlogsGlobal[PRIVATE_PROVIDER_STATE_KEY] ?? (neatlogsGlobal[PRIVATE_PROVIDER_STATE_KEY] = {
+  provider: null
+});
+var preInitTracer = {
+  startSpan() {
+    return otelTrace.wrapSpanContext(INVALID_SPAN_CONTEXT);
+  },
+  startActiveSpan(_name, arg2, arg3, arg4) {
+    const fn = typeof arg2 === "function" ? arg2 : typeof arg3 === "function" ? arg3 : arg4;
+    return fn(otelTrace.wrapSpanContext(INVALID_SPAN_CONTEXT));
+  }
+};
+function getNeatlogsTracer(name) {
+  return providerState.provider?.getTracer(name) ?? preInitTracer;
+}
+function getNeatlogsActiveContext() {
+  return privateContextStorage.getStore() ?? ROOT_CONTEXT;
+}
+function getActiveNeatlogsSpan() {
+  return otelTrace.getSpan(getNeatlogsActiveContext());
+}
+function getNeatlogsParentContext(baseContext) {
+  return baseContext ?? getNeatlogsActiveContext();
+}
+function withNeatlogsSpan(span, fn, baseContext) {
+  const base = baseContext ?? getNeatlogsActiveContext();
+  const ctx = otelTrace.setSpan(base, span);
+  return privateContextStorage.run(ctx, fn);
+}
+
 // src/core/end-user.ts
 var logger2 = getLogger();
 var END_USER_ID_KEY = "neatlogs.end_user.id";
@@ -133,7 +172,7 @@ function applySessionAttributes(span, sessionId, isRoot = true) {
 }
 
 // src/decorators/base.ts
-import { trace, SpanStatusCode } from "@opentelemetry/api";
+import { ROOT_CONTEXT as ROOT_CONTEXT2, SpanStatusCode } from "@opentelemetry/api";
 var logger4 = getLogger();
 
 // src/core/context.ts
@@ -198,11 +237,11 @@ var AutoRootTracer = class {
     const kind = String(
       options?.attributes?.["neatlogs.span.kind"] ?? ""
     ).toLowerCase();
-    const ctx = context ?? otelContext3.active();
-    const parent = otelTrace3.getSpan(ctx);
+    const ctx = context ?? getNeatlogsParentContext();
+    const parent = getActiveNeatlogsSpan();
     const needsRoot = autoRootEnabled() && !ROOT_KINDS.has(kind) && !(parent !== void 0 && parent.isRecording());
     if (!needsRoot) {
-      return this._tracer.startSpan(name, options, context);
+      return this._tracer.startSpan(name, options, ctx);
     }
     const root = this._tracer.startSpan(
       resolveRootWorkflowName(),
@@ -221,7 +260,7 @@ var AutoRootTracer = class {
   ));
 };
 function getProviderTracer(name) {
-  return new AutoRootTracer(otelTrace3.getTracer(name));
+  return new AutoRootTracer(getNeatlogsTracer(name));
 }
 
 // src/vertex-ai.ts
@@ -233,8 +272,8 @@ function wrapVertexAI(client) {
 }
 function traceTool(name, fn) {
   return async function tracedTool(args) {
-    const tracer = getProviderTracer(TRACER_NAME);
-    return tracer.startActiveSpan(
+    const tracer = getNeatlogsTracer(TRACER_NAME);
+    const span = tracer.startSpan(
       `tool.${name}`,
       {
         attributes: {
@@ -243,21 +282,21 @@ function traceTool(name, fn) {
           "input.value": safeStringify(args)
         }
       },
-      otelContext4.active(),
-      async (span) => {
-        try {
-          const result = await fn(args);
-          span.setAttribute("output.value", safeStringify(result));
-          span.setStatus({ code: SpanStatusCode3.OK });
-          return result;
-        } catch (err) {
-          recordError(span, err);
-          throw err;
-        } finally {
-          span.end();
-        }
-      }
+      getNeatlogsParentContext()
     );
+    return withNeatlogsSpan(span, async () => {
+      try {
+        const result = await fn(args);
+        span.setAttribute("output.value", safeStringify(result));
+        span.setStatus({ code: SpanStatusCode3.OK });
+        return result;
+      } catch (err) {
+        recordError(span, err);
+        throw err;
+      } finally {
+        span.end();
+      }
+    });
   };
 }
 function wrapNamespace(target, path) {
@@ -319,15 +358,14 @@ function tracedChatSend(original, chat, params, rest, isStream) {
       "neatlogs.llm.model_name": model,
       "neatlogs.llm.is_streaming": isStream
     }
-  }, otelContext4.active());
+  }, getNeatlogsParentContext());
   const message = params?.message ?? params;
   span.setAttribute("neatlogs.llm.input_messages.0.role", "user");
   span.setAttribute(
     "neatlogs.llm.input_messages.0.content",
     typeof message === "string" ? message : safeStringify(message)
   );
-  const ctx = trace2.setSpan(otelContext4.active(), span);
-  const result = otelContext4.with(ctx, () => original(params, ...rest));
+  const result = withNeatlogsSpan(span, () => original(params, ...rest));
   return Promise.resolve(result).then(
     (response) => {
       if (isStream) return wrapStream(response, span);
@@ -350,9 +388,8 @@ function tracedEmbedContent(original) {
         "neatlogs.embedding.model_name": opts?.model ?? "",
         "neatlogs.embedding.text": safeStringify(opts?.contents ?? "")
       }
-    }, otelContext4.active());
-    const ctx = trace2.setSpan(otelContext4.active(), span);
-    const result = otelContext4.with(ctx, () => original(opts, ...rest));
+    }, getNeatlogsParentContext());
+    const result = withNeatlogsSpan(span, () => original(opts, ...rest));
     return Promise.resolve(result).then(
       (response) => {
         const embeddings = response?.embeddings;
@@ -382,9 +419,8 @@ function tracedCountTokens(original) {
         "neatlogs.llm.task": "count_tokens",
         "neatlogs.llm.model_name": opts?.model ?? ""
       }
-    }, otelContext4.active());
-    const ctx = trace2.setSpan(otelContext4.active(), span);
-    const result = otelContext4.with(ctx, () => original(opts, ...rest));
+    }, getNeatlogsParentContext());
+    const result = withNeatlogsSpan(span, () => original(opts, ...rest));
     return Promise.resolve(result).then(
       (response) => {
         if (response?.totalTokens != null) span.setAttribute("neatlogs.llm.token_count.prompt", response.totalTokens);
@@ -411,10 +447,9 @@ function tracedGenerateContent(original, isStream) {
         "neatlogs.llm.model_name": model,
         "neatlogs.llm.is_streaming": isStream
       }
-    }, otelContext4.active());
+    }, getNeatlogsParentContext());
     setInputAttributes(span, opts);
-    const ctx = trace2.setSpan(otelContext4.active(), span);
-    const result = otelContext4.with(ctx, () => original(opts, ...rest));
+    const result = withNeatlogsSpan(span, () => original(opts, ...rest));
     return Promise.resolve(result).then(
       (response) => {
         if (isStream) {