neatlogs 1.1.19 → 1.1.21

package/dist/vertex-ai.mjs

@@ -166,11 +166,11 @@ function getActiveNeatlogsSpan() {
 function getNeatlogsParentContext(baseContext) {
   return baseContext ?? getNeatlogsActiveContext();
 }
-function withNeatlogsSpan(span, fn, baseContext) {
+function withNeatlogsSpan(span, fn, baseContext, rootSpan) {
   const base = baseContext ?? getNeatlogsActiveContext();
   let ctx = otelTrace.setSpan(base, span);
   if (base.getValue(NEATLOGS_ROOT_SPAN_KEY) === void 0) {
-    ctx = ctx.setValue(NEATLOGS_ROOT_SPAN_KEY, span);
+    ctx = ctx.setValue(NEATLOGS_ROOT_SPAN_KEY, rootSpan ?? span);
   }
   return privateContextStorage.run(ctx, fn);
 }
@@ -907,7 +907,8 @@ var USER_PROMPT_TEMPLATE_KEY = createContextKey2("neatlogs.user_prompt_template"
 var USER_PROMPT_VARIABLES_KEY = createContextKey2("neatlogs.user_prompt_variables");
 
 // src/core/auto-root.ts
-var ROOT_KINDS = /* @__PURE__ */ new Set(["workflow", "chain", "agent", "mcp_tool"]);
+var boundAutoRoots = /* @__PURE__ */ new WeakMap();
+var ROOT_KINDS = /* @__PURE__ */ new Set(["workflow", "chain", "agent", "mcp_tool", "llm"]);
 function autoRootEnabled() {
   const val = (process.env.NEATLOGS_AUTO_ROOT ?? "").trim().toLowerCase();
   return !["false", "0", "no", "off"].includes(val);
@@ -922,11 +923,12 @@ function resolveRootWorkflowName() {
   }
   return "workflow";
 }
-function stampAutoRootIdentity(root) {
+function stampRootIdentity(root) {
   applySessionAttributes(root, void 0, true);
   applyEndUserAttributes(root, void 0, void 0, true);
 }
-function wrapSpanWithRoot(child, root) {
+function bindSpanToAutoRoot(child, root) {
+  if (!root) return child;
   let ended = false;
   const proxy = new Proxy(child, {
     get(target, prop, _receiver) {
@@ -948,6 +950,7 @@ function wrapSpanWithRoot(child, root) {
       return typeof value === "function" ? value.bind(target) : value;
     }
   });
+  boundAutoRoots.set(proxy, root);
   aliasMediaCaptureSpan(proxy, child);
   return proxy;
 }
@@ -960,27 +963,41 @@ var AutoRootTracer = class {
     const kind = String(
       options?.attributes?.["neatlogs.span.kind"] ?? ""
     ).toLowerCase();
+    const hasExplicitContext = context !== void 0;
     const ctx = context ?? getNeatlogsParentContext();
     const parent = getActiveNeatlogsSpan();
-    const needsRoot = autoRootEnabled() && !ROOT_KINDS.has(kind) && !(parent !== void 0 && parent.isRecording());
+    const isParentless = !hasExplicitContext && !(parent !== void 0 && parent.isRecording());
+    const needsRoot = autoRootEnabled() && !ROOT_KINDS.has(kind) && isParentless;
     if (!needsRoot) {
-      return this._tracer.startSpan(name, options, ctx);
+      const span = this._tracer.startSpan(name, options, ctx);
+      if (isParentless && ROOT_KINDS.has(kind)) stampRootIdentity(span);
+      return span;
     }
     const root = this._tracer.startSpan(
       resolveRootWorkflowName(),
       { attributes: { "neatlogs.span.kind": "workflow", "neatlogs.auto_root": true } },
       ctx
     );
-    stampAutoRootIdentity(root);
+    stampRootIdentity(root);
     const childCtx = otelTrace3.setSpan(ctx, root);
     const child = this._tracer.startSpan(name, options, childCtx);
-    return wrapSpanWithRoot(child, root);
+    return bindSpanToAutoRoot(child, root);
   }
-  // startActiveSpan is used by traceTool (callback-scoped, always under an
-  // active span in practice) — pass straight through, no auto-root.
-  startActiveSpan = ((...args) => this._tracer.startActiveSpan(
-    ...args
-  ));
+  startActiveSpan = ((name, ...args) => {
+    const callback = args.at(-1);
+    if (typeof callback !== "function") {
+      throw new TypeError("startActiveSpan requires a callback");
+    }
+    const options = args.length >= 2 ? args[0] : void 0;
+    const parentContext = args.length >= 3 ? args[1] : getNeatlogsParentContext();
+    const span = args.length >= 3 ? this.startSpan(name, options, parentContext) : this.startSpan(name, options);
+    return withNeatlogsSpan(
+      span,
+      () => callback(span),
+      parentContext,
+      boundAutoRoots.get(span)
+    );
+  });
 };
 function getProviderTracer(name) {
   return new AutoRootTracer(getNeatlogsTracer(name));