neatlogs 1.1.6 → 1.1.7

package/dist/bedrock.mjs

@@ -1,18 +1,15 @@
 // src/bedrock.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/bedrock.ts
@@ -229,8 +268,8 @@ var TRACER_NAME = "neatlogs.bedrock";
 var PROVIDER = "bedrock";
 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: {
@@ -239,21 +278,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 wrapBedrock(client) {
@@ -298,7 +337,7 @@ function startSpan(name, modelId, isStream) {
       "neatlogs.llm.model_name": String(modelId ?? ""),
       "neatlogs.llm.is_streaming": isStream
     }
-  }, otelContext4.active());
+  }, getNeatlogsParentContext());
 }
 function tracedConverse(originalSend, command, input, rest, isStream) {
   const span = startSpan(
@@ -307,8 +346,7 @@ function tracedConverse(originalSend, command, input, rest, isStream) {
     isStream
   );
   setConverseInput(span, input);
-  const ctx = trace2.setSpan(otelContext4.active(), span);
-  const promise = otelContext4.with(ctx, () => originalSend(command, ...rest));
+  const promise = withNeatlogsSpan(span, () => originalSend(command, ...rest));
   return promise.then(
     (response) => {
       if (isStream) {
@@ -488,7 +526,7 @@ function tracedInvokeModel(originalSend, command, input, rest, isStream) {
         "neatlogs.llm.provider": PROVIDER,
         "neatlogs.embedding.model_name": String(input?.modelId ?? "")
       }
-    }, otelContext4.active());
+    }, getNeatlogsParentContext());
     const text = bodyIn?.inputText ?? bodyIn?.texts ?? bodyIn?.input_text;
     if (text) {
       span.setAttribute("neatlogs.embedding.text", typeof text === "string" ? text : safeStringify(text));
@@ -501,8 +539,7 @@ function tracedInvokeModel(originalSend, command, input, rest, isStream) {
     );
     setInvokeInput(span, vendor, bodyIn);
   }
-  const ctx = trace2.setSpan(otelContext4.active(), span);
-  const promise = otelContext4.with(ctx, () => originalSend(command, ...rest));
+  const promise = withNeatlogsSpan(span, () => originalSend(command, ...rest));
   return promise.then(
     (response) => {
       if (isStream) {