bitfab 0.54.2 → 0.54.4

package/dist/index.cjs

@@ -51,7 +51,7 @@ var __version__, __packageName__;
 var init_version_generated = __esm({
   "src/version.generated.ts"() {
     "use strict";
-    __version__ = "0.54.2";
+    __version__ = "0.54.4";
     __packageName__ = "bitfab";
   }
 });
@@ -66,6 +66,99 @@ var init_constants = __esm({
   }
 });
 
+// src/traceCompletion.ts
+var TraceCompletion;
+var init_traceCompletion = __esm({
+  "src/traceCompletion.ts"() {
+    "use strict";
+    TraceCompletion = class {
+      constructor() {
+        this.traces = /* @__PURE__ */ new Map();
+        this.closed = /* @__PURE__ */ new Map();
+      }
+      state(traceId) {
+        let state = this.traces.get(traceId);
+        if (state === void 0) {
+          state = { spanIds: /* @__PURE__ */ new Set(), active: /* @__PURE__ */ new Set() };
+          this.traces.set(traceId, state);
+        }
+        return state;
+      }
+      start(traceId, spanId) {
+        if (this.closed.has(traceId)) {
+          return;
+        }
+        const state = this.state(traceId);
+        state.spanIds.add(spanId);
+        state.active.add(spanId);
+      }
+      open(traceId) {
+        if (!this.closed.has(traceId)) {
+          this.state(traceId);
+        }
+      }
+      hold(traceId, key) {
+        this.state(traceId).active.add(key);
+      }
+      record(traceId, spanId) {
+        if (!this.closed.has(traceId)) {
+          this.state(traceId).spanIds.add(spanId);
+        }
+      }
+      end(traceId, spanId) {
+        const state = this.traces.get(traceId);
+        if (state !== void 0) {
+          state.active.delete(spanId);
+          this.drain(traceId, state);
+        }
+      }
+      abort(traceId, spanId) {
+        const state = this.traces.get(traceId);
+        if (state === void 0) {
+          return;
+        }
+        state.active.delete(spanId);
+        state.spanIds.delete(spanId);
+        if (state.active.size === 0 && state.spanIds.size === 0 && state.complete === void 0) {
+          this.traces.delete(traceId);
+          return;
+        }
+        this.drain(traceId, state);
+      }
+      close(traceId, complete, dropped = false) {
+        const count = this.closed.get(traceId);
+        if (count !== void 0) {
+          complete(count);
+          return;
+        }
+        const state = this.traces.get(traceId);
+        if (state === void 0) {
+          return;
+        }
+        state.complete = complete;
+        if (dropped) {
+          state.active.clear();
+        }
+        this.drain(traceId, state);
+      }
+      drain(traceId, state) {
+        if (state.active.size !== 0 || state.complete === void 0) {
+          return;
+        }
+        this.traces.delete(traceId);
+        this.closed.set(traceId, state.spanIds.size);
+        if (this.closed.size > 1024) {
+          const oldest = this.closed.keys().next().value;
+          if (oldest !== void 0) {
+            this.closed.delete(oldest);
+          }
+        }
+        state.complete(state.spanIds.size);
+      }
+    };
+  }
+});
+
 // src/readEnv.ts
 function readEnv(name) {
   if (typeof process !== "undefined" && process.env) {
@@ -2152,6 +2245,7 @@ var REPLAY_DB_BRANCH_REQUEST_TIMEOUT_MS, REPLAY_COMPLETE_REQUEST_TIMEOUT_MS, OTL
 var init_http = __esm({
   "src/http.ts"() {
     "use strict";
+    init_traceCompletion();
     init_compress();
     init_constants();
     init_errors();
@@ -2214,6 +2308,7 @@ var init_http = __esm({
         // must not wait on another client's slow finalize: a false `close()` failure
         // caused by unrelated work is worse than no signal at all.
         this.deferredWork = /* @__PURE__ */ new Set();
+        this.traceCompletion = new TraceCompletion();
         this.closed = false;
         this.apiKey = config.apiKey;
         this.serviceUrl = config.serviceUrl;
@@ -2647,6 +2742,10 @@ var init_http = __esm({
         );
       }
       sendExternalSpan(payload) {
+        const ref2 = carrierRef(payload);
+        if (ref2?.spanId !== void 0) {
+          this.traceCompletion.record(ref2.traceId, ref2.spanId);
+        }
         const gate = this.externalSpanTransform;
         if (gate === void 0) {
           this.submitExternalSpan(payload);
@@ -2678,13 +2777,28 @@ var init_http = __esm({
        */
       sendExternalTrace(rawPayload) {
         const payload = mergeCallerMetadataIntoTracePayload(rawPayload);
-        const gate = this.externalTraceTransform;
-        if (gate === void 0) {
-          this.submitExternalTrace(payload);
+        const traceId = sourceTraceIdOf(payload);
+        if (payload.completed === true && traceId !== void 0) {
+          this.traceCompletion.close(
+            traceId,
+            (expectedSpanCount) => this.sendCountedExternalTrace({ ...payload, expectedSpanCount }),
+            payload.dropped === true
+          );
           return;
         }
+        if (traceId !== void 0) {
+          this.traceCompletion.open(traceId);
+        }
+        this.sendCountedExternalTrace(payload);
+      }
+      sendCountedExternalTrace(payload) {
         try {
-          gate(payload, (ready) => this.submitExternalTrace(ready));
+          const gate = this.externalTraceTransform;
+          if (gate === void 0) {
+            this.submitExternalTrace(payload);
+          } else {
+            gate(payload, (ready) => this.submitExternalTrace(ready));
+          }
         } catch (error) {
           warnDroppedExternalTrace(error);
         }
@@ -4503,6 +4617,7 @@ var BitfabClaudeAgentHandler = class {
       spanInfo.dropped = [...inputDropped];
     }
     this.runToSpan.set(spanId, spanInfo);
+    this.httpClient.traceCompletion.start(spanInfo.traceId, spanInfo.spanId);
     return spanInfo;
   }
   completeSpan(spanId, output, error, extraContexts) {
@@ -4524,6 +4639,7 @@ var BitfabClaudeAgentHandler = class {
       spanInfo.contexts.push(extraContexts);
     }
     this.sendSpan(spanInfo);
+    this.httpClient.traceCompletion.end(spanInfo.traceId, spanInfo.spanId);
   }
   sendSpan(spanInfo) {
     const spanData = {
@@ -6181,6 +6297,9 @@ var BitfabLangGraphCallbackHandler = class {
       spanInfo.hidden = true;
     }
     this.runToSpan.set(runId, spanInfo);
+    if (spanInfo.hidden !== true) {
+      this.httpClient.traceCompletion.start(spanInfo.traceId, spanInfo.spanId);
+    }
     if (isRootInvocation) {
       this.sendTraceStart(spanInfo);
     }
@@ -6212,6 +6331,7 @@ var BitfabLangGraphCallbackHandler = class {
       this.sendTraceCompletion(spanInfo, invocation?.activeContext ?? null);
       this.invocations.delete(runId);
     }
+    this.httpClient.traceCompletion.end(spanInfo.traceId, spanInfo.spanId);
   }
   sendSpan(spanInfo) {
     const spanData = {
@@ -6933,6 +7053,8 @@ var BitfabOpenAITracingProcessor = class {
     this.activeTraces = {};
     this.activeSpanMappings = {};
     this.canonicalTraceIds = {};
+    this.startedSpans = /* @__PURE__ */ new Map();
+    this.endedTraces = /* @__PURE__ */ new Set();
     this.ownsHttpClient = config._httpClient === void 0;
     this.httpClient = config._httpClient ?? new HttpClient({
       apiKey: config.apiKey,
@@ -6971,6 +7093,10 @@ var BitfabOpenAITracingProcessor = class {
     }
     const canonicalTraceId = activeContext?.traceId ?? this.getCanonicalTraceId(trace.traceId);
     this.canonicalTraceIds[trace.traceId] = canonicalTraceId;
+    this.httpClient.traceCompletion.hold(
+      activeContext?.traceId ?? trace.traceId,
+      `openai:${trace.traceId}`
+    );
     this.sendTrace(trace, {
       id: canonicalTraceId,
       sourceTraceId: activeContext?.traceId
@@ -6988,16 +7114,33 @@ var BitfabOpenAITracingProcessor = class {
       id: mapping?.traceId ?? this.getCanonicalTraceId(trace.traceId),
       sourceTraceId: mapping?.traceId
     });
-    delete this.activeSpanMappings[trace.traceId];
-    delete this.canonicalTraceIds[trace.traceId];
-    delete this.activeTraces[trace.traceId];
+    this.endedTraces.add(trace.traceId);
+    this.httpClient.traceCompletion.end(
+      mapping?.traceId ?? trace.traceId,
+      `openai:${trace.traceId}`
+    );
+    this.cleanupTrace(trace.traceId);
+  }
+  cleanupTrace(traceId) {
+    if (!this.endedTraces.has(traceId) || [...this.startedSpans.values()].includes(traceId)) {
+      return;
+    }
+    this.endedTraces.delete(traceId);
+    delete this.activeSpanMappings[traceId];
+    delete this.canonicalTraceIds[traceId];
+    delete this.activeTraces[traceId];
   }
   /**
    * Called when a span is started. Span payloads are authoritative completed
    * snapshots, so start notifications do not cross the transport boundary.
    */
   // biome-ignore lint/suspicious/noExplicitAny: OpenAI Agents SDK uses any for span data
-  async onSpanStart(_span) {
+  async onSpanStart(span) {
+    if (span.spanId !== void 0 && span.traceId !== null && span.traceId !== void 0) {
+      this.startedSpans.set(span.spanId, span.traceId);
+      const traceId = this.activeSpanMappings[span.traceId]?.traceId ?? span.traceId;
+      this.httpClient.traceCompletion.start(traceId, span.spanId);
+    }
   }
   /**
    * Called when a span is ended.
@@ -7006,7 +7149,17 @@ var BitfabOpenAITracingProcessor = class {
    */
   // biome-ignore lint/suspicious/noExplicitAny: OpenAI Agents SDK uses any for span data
   async onSpanEnd(span) {
-    this.sendSpan(span);
+    const originalTraceId = span.traceId ?? "";
+    const traceId = this.activeSpanMappings[originalTraceId]?.traceId ?? originalTraceId;
+    try {
+      this.sendSpan(span);
+    } finally {
+      if (span.spanId !== void 0) {
+        this.httpClient.traceCompletion.end(traceId, span.spanId);
+        this.startedSpans.delete(span.spanId);
+      }
+      this.cleanupTrace(originalTraceId);
+    }
   }
   /**
    * Called when a trace is being flushed.
@@ -8712,6 +8865,7 @@ var Bitfab = class {
       let newStack;
       let executeWithContext;
       let registeredTraceId;
+      let startedSpan;
       try {
         const currentStack = getSpanStack();
         const parentContext = currentStack[currentStack.length - 1];
@@ -8763,6 +8917,8 @@ var Bitfab = class {
           });
           registeredTraceId = traceId;
         }
+        self.httpClient.traceCompletion.start(traceId, spanId);
+        startedSpan = { traceId, spanId };
         const functionName = fn.name !== "" ? fn.name : void 0;
         const rootTraceFunctionKey = activeTraceStates.get(traceId)?.traceFunctionKey ?? traceFunctionKey;
         const baseSpanParams = {
@@ -8837,6 +8993,8 @@ var Bitfab = class {
               retireTraceMetadata(traceId);
             }
           } catch {
+          } finally {
+            self.httpClient.traceCompletion.end(traceId, spanId);
           }
         };
         const recordSpan = (result) => {
@@ -8853,6 +9011,14 @@ var Bitfab = class {
             void sendSpan({ result });
           }
         };
+        const resolveMockOrAbort = async (resolve) => {
+          try {
+            return await resolve();
+          } catch (error) {
+            self.httpClient.traceCompletion.abort(traceId, spanId);
+            throw error;
+          }
+        };
         executeWithContext = () => {
           let result;
           try {
@@ -8909,7 +9075,7 @@ var Bitfab = class {
               );
             }
             return (async () => {
-              const output = await pending;
+              const output = await resolveMockOrAbort(() => pending);
               void sendSpan({
                 result: output,
                 mocked: true,
@@ -8979,7 +9145,7 @@ var Bitfab = class {
                 );
               }
               return runWithSpanStack(newStack, async () => {
-                const resolved2 = await resolution;
+                const resolved2 = await resolveMockOrAbort(() => resolution);
                 if (resolved2.matched) {
                   void sendSpan({
                     result: resolved2.output,
@@ -8995,7 +9161,9 @@ var Bitfab = class {
                   );
                 }
                 if (shouldMockWithBaseStrategy) {
-                  const output = await resolveRecordedOutput();
+                  const output = await resolveMockOrAbort(
+                    () => resolveRecordedOutput()
+                  );
                   void sendSpan({
                     result: output,
                     mocked: true,
@@ -9025,6 +9193,12 @@ var Bitfab = class {
           }
         }
       } catch (setupError) {
+        if (startedSpan) {
+          self.httpClient.traceCompletion.abort(
+            startedSpan.traceId,
+            startedSpan.spanId
+          );
+        }
         if (registeredTraceId) {
           activeTraceStates.delete(registeredTraceId);
           retireTraceMetadata(registeredTraceId);