bitfab 0.68.0 → 0.68.2

package/dist/{src-Dze8o7iU.cjs → src-CdQRfsrA.cjs}

@@ -1,6 +1,6 @@
 const require_bamlPeer = require("./bamlPeer-BUoZgNxZ.cjs");
-const require_http = require("./http-psFF1ptn.cjs");
-const require_autoTrace = require("./autoTrace-BikML0lX.cjs");
+const require_http = require("./http-Dqc-f3sp.cjs");
+const require_autoTrace = require("./autoTrace-D3-OZcmR.cjs");
 let superjson = require("superjson");
 superjson = require_bamlPeer.__toESM(superjson, 1);
 //#region src/categories.ts
@@ -1411,11 +1411,11 @@ var GradersClient = class {
 const JEV_PROVIDER = "typesafe";
 const JEV_WATCH_STORAGE_SYMBOL = Symbol.for("bitfab.jevWatchStorage");
 const repliesRecordedByAutoTrace = /* @__PURE__ */ new WeakSet();
-function isRecord$3(value) {
+function isRecord$4(value) {
 	return typeof value === "object" && value !== null;
 }
 function isJevReply(value) {
-	return isRecord$3(value) && typeof value.model === "string" && isRecord$3(value.usage);
+	return isRecord$4(value) && typeof value.model === "string" && isRecord$4(value.usage);
 }
 function tokenCount(usage, key) {
 	const count = usage[key];
@@ -1474,7 +1474,7 @@ function recordWatchedJevReply(span, watch, reply, durationMs) {
 	markEnclosingWatches(watch);
 }
 function isRecordedByAutoTrace(reply) {
-	return isRecord$3(reply) && repliesRecordedByAutoTrace.has(reply);
+	return isRecord$4(reply) && repliesRecordedByAutoTrace.has(reply);
 }
 function recordAutoTracedJev(invokeFn, span) {
 	const startedAt = Date.now();
@@ -1492,6 +1492,22 @@ function recordAutoTracedJev(invokeFn, span) {
 //#endregion
 //#region src/langgraph.ts
 const LANGSMITH_HIDDEN_TAG = "langsmith:hidden";
+function toolCallClaim(traceId, toolCallId) {
+	return `${traceId}\u0000${toolCallId}`;
+}
+var OutermostRunOutput = class {
+	constructor() {
+		this.name = "BitfabOutermostRunOutput";
+		this.awaitHandlers = true;
+		this.output = void 0;
+	}
+	handleChainStart(_chain, _inputs, runId) {
+		this.runId ?? (this.runId = runId);
+	}
+	handleChainEnd(outputs, runId) {
+		if (runId === this.runId) this.output = outputs;
+	}
+};
 const CHAIN_RUN_TYPES = /* @__PURE__ */ new Set([
 	"chain",
 	"parser",
@@ -1661,11 +1677,15 @@ function extractLangGraphMetadata(metadata) {
 var BitfabLangGraphCallbackHandler = class {
 	constructor(config) {
 		this.name = "BitfabLangGraphCallbackHandler";
+		this.awaitHandlers = true;
 		this.ignoreRetry = true;
 		this.ignoreRetriever = false;
 		this.ignoreCustomEvent = true;
 		this.runToSpan = /* @__PURE__ */ new Map();
 		this.invocations = /* @__PURE__ */ new Map();
+		this.claimedToolCalls = /* @__PURE__ */ new Set();
+		this.stepRunIdByNamespace = /* @__PURE__ */ new Map();
+		this.namespaceByStepRunId = /* @__PURE__ */ new Map();
 		this.ownsHttpClient = config._httpClient === void 0;
 		this.httpClient = config._httpClient ?? new require_http.HttpClient({
 			apiKey: config.apiKey,
@@ -1680,6 +1700,37 @@ var BitfabLangGraphCallbackHandler = class {
 	async close(timeoutMs) {
 		return this.ownsHttpClient ? this.httpClient.close(timeoutMs) : true;
 	}
+	claimToolCall(traceId, toolCallId) {
+		const claim = toolCallClaim(traceId, toolCallId);
+		this.claimedToolCalls.add(claim);
+		return () => {
+			this.claimedToolCalls.delete(claim);
+		};
+	}
+	visibleStepSpanId(checkpointNamespace, traceId) {
+		const runId = this.stepRunIdByNamespace.get(checkpointNamespace);
+		let span = runId === void 0 ? void 0 : this.runToSpan.get(runId);
+		while (span?.hidden === true) span = span.parentId ? this.runToSpan.get(span.parentId) : void 0;
+		return span !== void 0 && span.traceId === traceId ? span.spanId : null;
+	}
+	isClaimedToolCall(parentRunId, toolCallId) {
+		if (toolCallId === void 0) return false;
+		const traceId = (parentRunId ? this.runToSpan.get(parentRunId)?.traceId : void 0) ?? this.getActiveSpanContext?.()?.traceId;
+		return traceId !== void 0 && this.claimedToolCalls.has(toolCallClaim(traceId, toolCallId));
+	}
+	rememberStep(runId, metadata) {
+		const namespace = metadata?.langgraph_checkpoint_ns;
+		if (typeof namespace === "string") {
+			this.stepRunIdByNamespace.set(namespace, runId);
+			this.namespaceByStepRunId.set(runId, namespace);
+		}
+	}
+	forgetStep(runId) {
+		const namespace = this.namespaceByStepRunId.get(runId);
+		if (namespace === void 0) return;
+		this.namespaceByStepRunId.delete(runId);
+		if (this.stepRunIdByNamespace.get(namespace) === runId) this.stepRunIdByNamespace.delete(namespace);
+	}
 	startSpan(runId, parentRunId, name, spanType, inputData, metadata, tags) {
 		const parentSpan = parentRunId ? this.runToSpan.get(parentRunId) : void 0;
 		const willHide = tags?.includes(LANGSMITH_HIDDEN_TAG) === true;
@@ -1828,15 +1879,18 @@ var BitfabLangGraphCallbackHandler = class {
 			const idArr = serialized.id;
 			const name = runName ?? serialized.name ?? idArr?.[idArr.length - 1] ?? "chain";
 			this.startSpan(runId, parentRunId, String(name), "agent", inputs, metadata, tags);
+			this.rememberStep(runId, metadata);
 		} catch {}
 	}
 	async handleChainEnd(outputs, runId) {
 		try {
+			this.forgetStep(runId);
 			this.completeSpan(runId, outputs);
 		} catch {}
 	}
 	async handleChainError(error, runId) {
 		try {
+			this.forgetStep(runId);
 			if (error?.constructor?.name === "GraphBubbleUp") {
 				this.completeSpan(runId, void 0, void 0);
 				return;
@@ -1888,8 +1942,8 @@ var BitfabLangGraphCallbackHandler = class {
 		} catch {}
 	}
 	async handleLLMNewToken() {}
-	async handleToolStart(tool, input, runId, parentRunId, tags, metadata, runName) {
-		if (!this.captureTools) return;
+	async handleToolStart(tool, input, runId, parentRunId, tags, metadata, runName, toolCallId) {
+		if (!this.captureTools || this.isClaimedToolCall(parentRunId, toolCallId)) return;
 		try {
 			const name = runName ?? (tool ?? {}).name ?? "tool";
 			this.startSpan(runId, parentRunId, String(name), "function", input, metadata, tags);
@@ -1925,21 +1979,232 @@ var BitfabLangGraphCallbackHandler = class {
 	}
 };
 //#endregion
+//#region src/integrationFixture.ts
+const INTEGRATION_FIXTURE_CONTEXT_KEY = "__bitfabIntegrationFixture";
+const INTEGRATION_FIXTURE_VERSION = 1;
+function isRecord$3(value) {
+	return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+function createIntegrationFixtureContext(name, value) {
+	const serialized = serializeValue(value);
+	return { [INTEGRATION_FIXTURE_CONTEXT_KEY]: {
+		version: INTEGRATION_FIXTURE_VERSION,
+		name,
+		value: serialized.json,
+		...serialized.meta !== void 0 && { valueMeta: serialized.meta }
+	} };
+}
+function extractIntegrationFixtures(spanData) {
+	const contexts = spanData.contexts;
+	if (!Array.isArray(contexts)) return;
+	const fixtures = /* @__PURE__ */ new Map();
+	for (const context of contexts) {
+		if (!isRecord$3(context)) continue;
+		const fixture = context[INTEGRATION_FIXTURE_CONTEXT_KEY];
+		if (!isRecord$3(fixture) || fixture.version !== INTEGRATION_FIXTURE_VERSION || typeof fixture.name !== "string") continue;
+		fixtures.set(fixture.name, deserializeValue({
+			json: fixture.value,
+			...fixture.valueMeta !== void 0 && { meta: fixture.valueMeta }
+		}));
+	}
+	return fixtures.size > 0 ? fixtures : void 0;
+}
+//#endregion
+//#region src/langgraphRuns.ts
+function threadFixtureName(traceFunctionKey) {
+	return `langgraph-thread:${traceFunctionKey}`;
+}
+const CHECKPOINTER_CONFIG_KEY = "__pregel_checkpointer";
+const THREAD_CONFIG_KEYS = [
+	"thread_id",
+	"checkpoint_ns",
+	"checkpoint_id"
+];
+function isRecord$2(value) {
+	return typeof value === "object" && value !== null;
+}
+function toBase64(bytes) {
+	const buffer = globalThis.Buffer;
+	if (buffer !== void 0) return buffer.from(bytes).toString("base64");
+	let binary = "";
+	for (const byte of bytes) binary += String.fromCharCode(byte);
+	return btoa(binary);
+}
+function fromBase64(text) {
+	const buffer = globalThis.Buffer;
+	if (buffer !== void 0) return new Uint8Array(buffer.from(text, "base64"));
+	return Uint8Array.from(atob(text), (char) => char.charCodeAt(0));
+}
+async function newMemorySaver() {
+	try {
+		const { MemorySaver } = await require_bamlPeer.importOptionalPeer(["@langchain", "langgraph"]);
+		return new MemorySaver();
+	} catch {
+		throw new require_http.BitfabError("Restoring a LangGraph thread during replay requires @langchain/langgraph.", "https://docs.bitfab.ai/frameworks/langgraph");
+	}
+}
+function savingCheckpointer(graph) {
+	const own = isRecord$2(graph) ? graph.checkpointer : void 0;
+	const inner = isRecord$2(graph) && isRecord$2(graph.graph) ? graph.graph.checkpointer : void 0;
+	const checkpointer = own ?? inner;
+	return isRecord$2(checkpointer) && typeof checkpointer.getTuple === "function" ? checkpointer : null;
+}
+function threadConfig(config) {
+	const configurable = isRecord$2(config?.configurable) ? config.configurable : {};
+	return Object.fromEntries(THREAD_CONFIG_KEYS.filter((key) => key in configurable).map((key) => [key, configurable[key]]));
+}
+function withConfigurable(config, values) {
+	return {
+		...config ?? {},
+		configurable: {
+			...isRecord$2(config?.configurable) ? config.configurable : {},
+			...values
+		}
+	};
+}
+async function encodeTyped(saver, value) {
+	const [type, data] = await saver.serde.dumpsTyped(value);
+	return [type, toBase64(data)];
+}
+async function decodeTyped(saver, [type, data]) {
+	return saver.serde.loadsTyped(type, fromBase64(data));
+}
+async function encodeCheckpoint(tuple) {
+	if (tuple === void 0) return null;
+	const saver = await newMemorySaver();
+	return {
+		config: threadConfig(tuple.config),
+		checkpoint: await encodeTyped(saver, tuple.checkpoint),
+		metadata: await encodeTyped(saver, tuple.metadata ?? {}),
+		pendingWrites: await Promise.all((tuple.pendingWrites ?? []).map(async ([taskId, channel, value]) => [
+			taskId,
+			channel,
+			await encodeTyped(saver, value)
+		]))
+	};
+}
+async function restoredCheckpointer(recorded) {
+	const saver = await newMemorySaver();
+	if (recorded === null) return {
+		saver,
+		thread: {}
+	};
+	const checkpoint = await decodeTyped(saver, recorded.checkpoint);
+	const saved = await saver.put({ configurable: {
+		thread_id: recorded.config.thread_id,
+		checkpoint_ns: recorded.config.checkpoint_ns ?? ""
+	} }, checkpoint, await decodeTyped(saver, recorded.metadata), checkpoint.channel_versions ?? {});
+	const writesByTask = /* @__PURE__ */ new Map();
+	for (const [taskId, channel, value] of recorded.pendingWrites) writesByTask.set(taskId, [...writesByTask.get(taskId) ?? [], [channel, await decodeTyped(saver, value)]]);
+	for (const [taskId, writes] of writesByTask) await saver.putWrites(saved, writes, taskId);
+	return {
+		saver,
+		thread: recorded.config
+	};
+}
+var GraphRun = class {
+	constructor(host, graph, config) {
+		this.host = host;
+		this.graph = graph;
+		this.config = config;
+		this.finalState = new OutermostRunOutput();
+		this.traced = graph.withConfig({ callbacks: [host.callbackHandler, this.finalState] });
+		const active = host.activeSpan();
+		this.opensRoot = active?.rootTraceFunctionKey !== host.traceFunctionKey;
+		this.replaysFromRoot = !this.opensRoot && active?.isRoot === true && require_http.getReplayContext() !== null;
+	}
+	invoke(input) {
+		if (this.replaysFromRoot) return this.replayConfig().then(async ({ config, confirm }) => {
+			const result = await this.traced.invoke(await this.host.reviveRootInput(input), config);
+			await confirm();
+			return result;
+		});
+		if (!this.opensRoot) return this.host.reviveRootInput(input).then((rootInput) => this.traced.invoke(rootInput, this.config));
+		return this.host.rootSpan(async (rootInput) => {
+			await this.recordCheckpoint();
+			return this.traced.invoke(await this.host.reviveRootInput(rootInput), this.config);
+		})(this.host.encodeRootInput(input));
+	}
+	stream(input) {
+		return this.chunks(input, async (rootInput, config) => this.traced.stream(rootInput, config));
+	}
+	streamEvents(input, streamOptions) {
+		return this.chunks(input, async (rootInput, config) => this.traced.streamEvents(rootInput, config, streamOptions));
+	}
+	chunks(input, open) {
+		const { host, finalState } = this;
+		if (this.replaysFromRoot) {
+			const replayConfig = () => this.replayConfig();
+			return (async function* replayChunks() {
+				const { config, confirm } = await replayConfig();
+				yield* await open(await host.reviveRootInput(input), config);
+				await confirm();
+				const replay = require_http.getReplayContext();
+				if (replay !== null) replay.streamResult = finalState.output;
+			})();
+		}
+		if (!this.opensRoot) {
+			const config = this.config;
+			return (async function* nestedChunks() {
+				yield* await open(await host.reviveRootInput(input), config);
+			})();
+		}
+		const config = this.config;
+		const recordCheckpoint = () => this.recordCheckpoint();
+		return host.rootSpan(async function* rootChunks(rootInput) {
+			await recordCheckpoint();
+			yield* await open(await host.reviveRootInput(rootInput), config);
+		}, () => finalState.output)(host.encodeRootInput(input));
+	}
+	async recordCheckpoint() {
+		const checkpointer = savingCheckpointer(this.graph);
+		const thread = threadConfig(this.config);
+		if (checkpointer === null || thread.thread_id === void 0) return;
+		const encoded = await encodeCheckpoint(await checkpointer.getTuple({ configurable: thread }));
+		if (encoded !== null) this.host.addRootContext(createIntegrationFixtureContext(threadFixtureName(this.host.traceFunctionKey), encoded));
+	}
+	async replayConfig() {
+		if (savingCheckpointer(this.graph) === null) return {
+			config: this.config,
+			confirm: async () => {}
+		};
+		const { saver, thread } = await restoredCheckpointer(require_http.getReplayContext()?.integrationFixtures?.get(threadFixtureName(this.host.traceFunctionKey)) ?? null);
+		const configurable = isRecord$2(this.config?.configurable) ? this.config.configurable : {};
+		const replayThread = {
+			thread_id: thread.thread_id ?? configurable.thread_id ?? `bitfab-replay-${randomUuid()}`,
+			...thread.checkpoint_ns !== void 0 && { checkpoint_ns: thread.checkpoint_ns }
+		};
+		const restoredId = (await saver.getTuple({ configurable: replayThread }))?.checkpoint.id;
+		const confirm = async () => {
+			const latest = await saver.getTuple({ configurable: replayThread });
+			if (latest === void 0 || latest.checkpoint.id === restoredId) require_autoTrace.warnOnce("langgraph-replay-checkpointer-not-lent", "LangGraph did not use the in-memory checkpointer Bitfab lent to a replayed run, so the run may have read and written the graph's own checkpointer. Check that your @langchain/langgraph version still lends a checkpointer through the run config.");
+		};
+		return {
+			config: withConfigurable(this.config, {
+				...replayThread,
+				[CHECKPOINTER_CONFIG_KEY]: saver
+			}),
+			confirm
+		};
+	}
+};
+//#endregion
 //#region src/langgraphIntegration.ts
 const TOOL_RESULT_TAG = "__bitfabLangGraphToolResult";
-function isRecord$2(value) {
+const AGENT_MIDDLEWARE_BRAND = Symbol.for("AgentMiddleware");
+function isRecord$1(value) {
 	return typeof value === "object" && value !== null;
 }
 function isToolMessage(value) {
-	return isRecord$2(value) && value.lc_direct_tool_output === true && value.type === "tool" && (typeof value.content === "string" || Array.isArray(value.content));
+	return isRecord$1(value) && value.lc_direct_tool_output === true && value.type === "tool" && (typeof value.content === "string" || Array.isArray(value.content));
 }
 function isCommand(value) {
-	return isRecord$2(value) && value.lg_name === "Command";
+	return isRecord$1(value) && value.lg_name === "Command";
 }
 function encodeNested(value) {
 	if (isToolMessage(value) || isCommand(value)) return encodeNativeToolResult(value);
 	if (Array.isArray(value)) return value.map(encodeNested);
-	if (isRecord$2(value)) return Object.fromEntries(Object.entries(value).map(([key, entry]) => [key, encodeNested(entry)]));
+	if (isRecord$1(value)) return Object.fromEntries(Object.entries(value).map(([key, entry]) => [key, encodeNested(entry)]));
 	return value;
 }
 function encodeNativeToolResult(value) {
@@ -1956,17 +2221,24 @@ function encodeNativeToolResult(value) {
 	};
 	return {
 		[TOOL_RESULT_TAG]: "command",
-		graph: isRecord$2(value) && typeof value.graph === "string" ? value.graph : void 0,
-		update: isRecord$2(value) ? encodeNested(value.update) : void 0,
-		resume: isRecord$2(value) ? encodeNested(value.resume) : void 0,
-		goto: isRecord$2(value) ? encodeNested(value.goto) : void 0
+		graph: isRecord$1(value) && typeof value.graph === "string" ? value.graph : void 0,
+		update: isRecord$1(value) ? encodeNested(value.update) : void 0,
+		resume: isRecord$1(value) ? encodeNested(value.resume) : void 0,
+		goto: isRecord$1(value) ? encodeNested(value.goto) : void 0
 	};
 }
 function finalizeToolResult(value) {
 	return isToolMessage(value) || isCommand(value) ? encodeNativeToolResult(value) : value;
 }
 function isEncodedToolResult(value) {
-	return isRecord$2(value) && typeof value[TOOL_RESULT_TAG] === "string";
+	return isRecord$1(value) && typeof value[TOOL_RESULT_TAG] === "string";
+}
+async function loadPeer(specifierParts, packageName, docsPath) {
+	try {
+		return await require_bamlPeer.importOptionalPeer(specifierParts);
+	} catch {
+		throw new require_http.BitfabError(`This LangGraph integration method requires ${packageName}. Install it to use it.`, `https://docs.bitfab.ai${docsPath}`);
+	}
 }
 async function loadLangChainCore() {
 	try {
@@ -1989,7 +2261,7 @@ async function loadLangGraph() {
 async function reviveNested(value, toolCallId) {
 	if (isEncodedToolResult(value)) return reviveToolResult(value, toolCallId);
 	if (Array.isArray(value)) return Promise.all(value.map((entry) => reviveNested(entry, toolCallId)));
-	if (isRecord$2(value)) {
+	if (isRecord$1(value)) {
 		const entries = await Promise.all(Object.entries(value).map(async ([key, entry]) => [key, await reviveNested(entry, toolCallId)]));
 		return Object.fromEntries(entries);
 	}
@@ -2005,9 +2277,9 @@ async function reviveToolResult(value, toolCallId) {
 			...typeof value.id === "string" && { id: value.id },
 			...(value.status === "success" || value.status === "error") && { status: value.status },
 			...value.artifact !== void 0 && { artifact: await reviveNested(value.artifact, toolCallId) },
-			...isRecord$2(value.metadata) && { metadata: await reviveNested(value.metadata, toolCallId) },
-			...isRecord$2(value.additionalKwargs) && { additional_kwargs: await reviveNested(value.additionalKwargs, toolCallId) },
-			...isRecord$2(value.responseMetadata) && { response_metadata: await reviveNested(value.responseMetadata, toolCallId) }
+			...isRecord$1(value.metadata) && { metadata: await reviveNested(value.metadata, toolCallId) },
+			...isRecord$1(value.additionalKwargs) && { additional_kwargs: await reviveNested(value.additionalKwargs, toolCallId) },
+			...isRecord$1(value.responseMetadata) && { response_metadata: await reviveNested(value.responseMetadata, toolCallId) }
 		});
 	}
 	const { Command } = await loadLangGraph();
@@ -2024,37 +2296,148 @@ var BitfabLangGraphIntegration = class {
 		this.traceFunctionKey = config.traceFunctionKey;
 		this.callbackHandler = config.callbackHandler;
 		this.mockToolsOnReplay = config.options?.mockToolsOnReplay ?? true;
+		this.getActiveSpan = config.getActiveSpan ?? (() => null);
+		this.addRootContext = config.addContext ?? (() => {});
+	}
+	get middleware() {
+		this.agentMiddleware ?? (this.agentMiddleware = {
+			[AGENT_MIDDLEWARE_BRAND]: true,
+			name: "BitfabAgentMiddleware",
+			wrapToolCall: (request, handler) => this.interceptToolCall(request, handler)
+		});
+		return this.agentMiddleware;
+	}
+	instrument(graph) {
+		const run = (config) => new GraphRun(this.graphRunHost, graph, config);
+		return new Proxy(graph, { get: (target, property) => {
+			if (property === "invoke") return (input, config) => run(config).invoke(input);
+			if (property === "stream") return async (input, config) => run(config).stream(input);
+			if (property === "streamEvents") return (input, config, streamOptions) => run(config).streamEvents(input, streamOptions);
+			if (property === "batch") return (inputs, config) => Promise.all(inputs.map((input, index) => run(Array.isArray(config) ? config[index] : config).invoke(input)));
+			if (property === "withConfig") return (config) => this.instrument(target.withConfig(config));
+			const value = Reflect.get(target, property, target);
+			return typeof value === "function" ? value.bind(target) : value;
+		} });
+	}
+	get graphRunHost() {
+		this.runHost ?? (this.runHost = {
+			traceFunctionKey: this.traceFunctionKey,
+			callbackHandler: this.callbackHandler,
+			activeSpan: () => this.getActiveSpan(),
+			rootSpan: (fn, finalize) => this.client.withSpan(this.traceFunctionKey, {
+				name: this.traceFunctionKey,
+				type: "agent",
+				surface: "inherit",
+				instrumentation: "langgraph",
+				...finalize !== void 0 && { finalize }
+			}, fn),
+			addRootContext: (context) => this.addRootContext(context),
+			encodeRootInput: (input) => isCommand(input) ? encodeNativeToolResult(input) : input,
+			reviveRootInput: async (input) => isEncodedToolResult(input) ? await reviveToolResult(input, "") : input
+		});
+		return this.runHost;
+	}
+	async createDeepAgent(params = {}) {
+		const { createDeepAgent } = await loadPeer(["deepagents"], "deepagents", "/frameworks/deep-agents");
+		return this.instrument(createDeepAgent({
+			...params,
+			middleware: this.withMiddleware(params.middleware),
+			...params.subagents !== void 0 && { subagents: params.subagents.map((spec) => this.withSubagentMiddleware(spec)) }
+		}));
+	}
+	async createAgent(params) {
+		const { createAgent } = await loadPeer(["langchain"], "langchain", "/frameworks/langgraph");
+		return this.instrument(createAgent({
+			...params,
+			middleware: this.withMiddleware(params.middleware)
+		}));
+	}
+	withMiddleware(middleware = []) {
+		return middleware.includes(this.middleware) ? [...middleware] : [...middleware, this.middleware];
+	}
+	withSubagentMiddleware(spec) {
+		if (!isRecord$1(spec) || "runnable" in spec || "graphId" in spec) return spec;
+		return {
+			...spec,
+			middleware: this.withMiddleware(Array.isArray(spec.middleware) ? spec.middleware : [])
+		};
+	}
+	async interceptToolCall(request, handler) {
+		const toolName = request.toolCall.name;
+		const toolCallId = request.toolCall.id;
+		const shouldMock = this.shouldMock(toolName);
+		this.assertReplayToolResultExists(toolName, shouldMock);
+		const traceId = this.recordingTraceId(toolName);
+		if (traceId === null || !toolCallId) return await handler(request);
+		const stepSpanId = this.stepSpanIdFor(request.runtime, traceId);
+		return await this.runClaimedToolCall(traceId, toolCallId, toolName, shouldMock, request.toolCall.args, async () => await handler(request), stepSpanId);
+	}
+	stepSpanIdFor(config, traceId) {
+		const namespace = isRecord$1(config) && isRecord$1(config.configurable) ? config.configurable.checkpoint_ns : void 0;
+		return typeof namespace === "string" ? this.callbackHandler.visibleStepSpanId(namespace, traceId) : null;
+	}
+	recordingTraceId(toolName) {
+		const activeSpan = this.getActiveSpan();
+		if (activeSpan === null) {
+			require_autoTrace.warnOnce(`langgraph-tool-call-outside-root:${this.traceFunctionKey}`, `tool call "${toolName}" ran outside a Bitfab trace for "${this.traceFunctionKey}", so it was not recorded as a replayable span. Run the agent through instrument(agent), createAgent(), or createDeepAgent().`);
+			return null;
+		}
+		return activeSpan.traceId;
+	}
+	async runClaimedToolCall(traceId, toolCallId, toolName, shouldMock, args, run, stepSpanId) {
+		const release = this.callbackHandler.claimToolCall(traceId, toolCallId);
+		try {
+			const result = await this.client.withSpan(this.traceFunctionKey, {
+				name: toolName,
+				type: "function",
+				captureWhen: "nested",
+				mockOnReplay: shouldMock,
+				finalize: finalizeToolResult,
+				surface: "inherit",
+				instrumentation: "langgraph",
+				...stepSpanId !== null && { frameworkParentSpanId: stepSpanId }
+			}, async (_args) => await run())(args);
+			return isEncodedToolResult(result) ? await reviveToolResult(result, toolCallId) : result;
+		} finally {
+			release();
+		}
+	}
+	shouldMock(toolName) {
+		return typeof this.mockToolsOnReplay === "boolean" ? this.mockToolsOnReplay : this.mockToolsOnReplay.includes(toolName);
 	}
 	wrapTools(tools) {
 		return tools.map((tool) => this.wrapTool(tool));
 	}
 	createInvoker(graph) {
-		const configuredGraph = graph.withConfig({ callbacks: [this.callbackHandler] });
-		return (input, config) => {
-			return this.wrapInvoke((rootInput) => configuredGraph.invoke(rootInput, config))(input);
-		};
+		require_autoTrace.warnOnce("deprecated-create-invoker", "createInvoker(graph) is deprecated; use instrument(graph).invoke instead.");
+		const traced = this.instrument(graph);
+		return (input, config) => traced.invoke(input, config);
 	}
 	wrapTool(tool) {
 		const toolName = tool.name;
 		if (typeof toolName !== "string" || toolName.length === 0) throw new require_http.BitfabError("LangGraph replayable tools must have a name.", "https://docs.bitfab.ai/frameworks/langgraph");
-		const mockToolsOnReplay = this.mockToolsOnReplay;
-		const shouldMock = typeof mockToolsOnReplay === "boolean" ? mockToolsOnReplay : mockToolsOnReplay.includes(toolName);
+		const shouldMock = this.shouldMock(toolName);
 		const originalInvoke = tool.invoke.bind(tool);
 		return new Proxy(tool, { get: (target, property) => {
 			if (property === "invoke") return async (input, ...rest) => {
-				const toolCallId = isRecord$2(input) && typeof input.id === "string" ? input.id : "";
-				const args = isRecord$2(input) && "args" in input ? input.args : input;
+				const toolCallId = isRecord$1(input) && typeof input.id === "string" ? input.id : "";
+				const args = isRecord$1(input) && "args" in input ? input.args : input;
 				this.assertReplayToolResultExists(toolName, shouldMock);
-				const result = await this.client.withSpan(this.traceFunctionKey, {
-					name: toolName,
-					type: "function",
-					captureWhen: "nested",
-					mockOnReplay: shouldMock,
-					finalize: finalizeToolResult,
-					surface: "inherit",
-					instrumentation: "langgraph"
-				}, async (_args) => await originalInvoke(input, ...rest))(args);
-				return isEncodedToolResult(result) ? await reviveToolResult(result, toolCallId) : result;
+				const traceId = this.getActiveSpan()?.traceId ?? null;
+				const run = async () => await originalInvoke(input, ...rest);
+				if (traceId === null || toolCallId === "") {
+					const result = await this.client.withSpan(this.traceFunctionKey, {
+						name: toolName,
+						type: "function",
+						captureWhen: "nested",
+						mockOnReplay: shouldMock,
+						finalize: finalizeToolResult,
+						surface: "inherit",
+						instrumentation: "langgraph"
+					}, async (_args) => await run())(args);
+					return isEncodedToolResult(result) ? await reviveToolResult(result, toolCallId) : result;
+				}
+				return await this.runClaimedToolCall(traceId, toolCallId, toolName, shouldMock, args, run, this.stepSpanIdFor(rest[0], traceId));
 			};
 			const value = Reflect.get(target, property, target);
 			return typeof value === "function" ? value.bind(target) : value;
@@ -2505,37 +2888,6 @@ function looksBinary(s) {
 	return s.slice(0, 8e3).includes(NUL);
 }
 //#endregion
-//#region src/integrationFixture.ts
-const INTEGRATION_FIXTURE_CONTEXT_KEY = "__bitfabIntegrationFixture";
-const INTEGRATION_FIXTURE_VERSION = 1;
-function isRecord$1(value) {
-	return typeof value === "object" && value !== null && !Array.isArray(value);
-}
-function createIntegrationFixtureContext(name, value) {
-	const serialized = serializeValue(value);
-	return { [INTEGRATION_FIXTURE_CONTEXT_KEY]: {
-		version: INTEGRATION_FIXTURE_VERSION,
-		name,
-		value: serialized.json,
-		...serialized.meta !== void 0 && { valueMeta: serialized.meta }
-	} };
-}
-function extractIntegrationFixtures(spanData) {
-	const contexts = spanData.contexts;
-	if (!Array.isArray(contexts)) return;
-	const fixtures = /* @__PURE__ */ new Map();
-	for (const context of contexts) {
-		if (!isRecord$1(context)) continue;
-		const fixture = context[INTEGRATION_FIXTURE_CONTEXT_KEY];
-		if (!isRecord$1(fixture) || fixture.version !== INTEGRATION_FIXTURE_VERSION || typeof fixture.name !== "string") continue;
-		fixtures.set(fixture.name, deserializeValue({
-			json: fixture.value,
-			...fixture.valueMeta !== void 0 && { meta: fixture.valueMeta }
-		}));
-	}
-	return fixtures.size > 0 ? fixtures : void 0;
-}
-//#endregion
 //#region src/replayHeartbeat.ts
 const REPLAY_HEARTBEAT_INTERVAL_MS = 6e4;
 function startReplayHeartbeat(send) {
@@ -4686,25 +5038,45 @@ function runWithSpanStack(stack, fn) {
 		throw error;
 	}
 }
+function isAsyncIterable(value) {
+	return value !== null && typeof value === "object" && typeof value[Symbol.asyncIterator] === "function";
+}
+function drainingStreams(fn) {
+	return async (...args) => {
+		const result = await fn(...args);
+		if (!isAsyncIterable(result)) return result;
+		const chunks = [];
+		for await (const chunk of result) chunks.push(chunk);
+		const replay = require_http.getReplayContext();
+		return replay !== null && "streamResult" in replay ? replay.streamResult : chunks;
+	};
+}
 function isAsyncGenerator(value) {
 	if (value === null || typeof value !== "object") return false;
 	const candidate = value;
 	return typeof candidate.next === "function" && typeof candidate.return === "function" && typeof candidate.throw === "function" && typeof candidate[Symbol.asyncIterator] === "function";
 }
-function wrapAsyncGenerator(source, spanStack, sendSpan) {
+function wrapAsyncGenerator(source, spanStack, sendSpan, recordYielded) {
 	const yielded = [];
 	let returnValue;
 	let finalized = false;
 	const finalize = (errorMsg) => {
 		if (finalized) return;
 		finalized = true;
-		sendSpan({
-			result: {
-				yielded,
-				return: returnValue
-			},
-			...errorMsg && { error: errorMsg }
-		});
+		if (recordYielded === void 0 || errorMsg) {
+			sendSpan({
+				result: {
+					yielded,
+					return: returnValue
+				},
+				...errorMsg && { error: errorMsg }
+			});
+			return;
+		}
+		Promise.resolve().then(() => recordYielded(yielded)).then((output) => sendSpan({ result: output }), (error) => sendSpan({
+			result: void 0,
+			error: `finalize failed: ${error instanceof Error ? error.message : String(error)}`
+		}));
 	};
 	const step = (method, arg) => runWithSpanStack(spanStack, () => {
 		return source[method].call(source, arg);
@@ -4798,7 +5170,7 @@ function extractContextFromCollector(collector) {
 }
 const BOUND_FUNCTION_NAME_PREFIX = /^bound /;
 function discoveredSpanType(definition, nodeConfiguration) {
-	const recordsModelCall = definition.baml === true || definition.jev === true;
+	const recordsModelCall = definition.baml === true || definition.jev === true || definition.openrouter === true;
 	const ownType = recordsModelCall ? "llm" : "function";
 	if (nodeConfiguration === void 0) return ownType;
 	if (recordsModelCall && nodeConfiguration.defaultType === true) return ownType;
@@ -5199,7 +5571,7 @@ var Bitfab = class {
 						...nodeConfiguration?.finalize !== void 0 && { finalize: nodeConfiguration.finalize }
 					};
 					const isAsync = definition.async === true;
-					return openChild(entry, childOptions, inputs, isAsync, definition.baml === true ? () => invokeAutoTracedBaml(invokeFn, isAsync) : definition.jev === true ? () => recordAutoTracedJev(invokeFn, getCurrentSpan()) : invokeFn);
+					return openChild(entry, childOptions, inputs, isAsync, definition.baml === true ? () => invokeAutoTracedBaml(invokeFn, isAsync) : definition.jev === true ? () => recordAutoTracedJev(invokeFn, getCurrentSpan()) : definition.openrouter === true ? () => require_autoTrace.recordAutoTracedOpenRouter(invokeFn, inputs[0], getCurrentSpan()) : invokeFn);
 				},
 				invokeNestedRoot(nestedName, nestedFunctionName, inputs, invokeFn, entry, link, spanDetails) {
 					if (excluded.has(nestedName) || excluded.has(nestedFunctionName) || !recordsSpan(entry, nestedName, true)) return invokeFn();
@@ -5378,16 +5750,23 @@ var Bitfab = class {
 				const stack = getSpanStack();
 				return stack[stack.length - 1] ?? null;
 			},
-			captureTools: false,
 			_httpClient: this.httpClient
 		});
 		return new BitfabLangGraphIntegration({
 			client: this,
 			traceFunctionKey,
 			callbackHandler,
-			options
+			options,
+			getActiveSpan: () => {
+				const stack = getSpanStack();
+				return stack[stack.length - 1] ?? null;
+			},
+			addContext: (context) => getCurrentSpan().addContext(context)
 		});
 	}
+	getDeepAgentsIntegration(traceFunctionKey, options) {
+		return this.getLangGraphIntegration(traceFunctionKey, options);
+	}
 	getRivetIntegration(traceFunctionKey, options) {
 		return new BitfabRivetIntegration(traceFunctionKey, {
 			wrapEffect: (effect, mockOnReplay, fn, replayOutput) => {
@@ -5503,6 +5882,26 @@ var Bitfab = class {
 		});
 		return wrappedFn;
 	}
+	wrapOpenRouter(method, options) {
+		const methodName = method.name.replace(BOUND_FUNCTION_NAME_PREFIX, "");
+		if (!methodName) throw new require_http.BitfabError("wrapOpenRouter requires a named function (e.g., openai.chat.completions.create.bind(openai.chat.completions)).");
+		const wrappedFn = async (...args) => {
+			const span = getCurrentSpan();
+			const startedAt = Date.now();
+			const { watch, result } = require_autoTrace.watchOpenRouterUsage(() => method(...args));
+			const response = await result;
+			require_autoTrace.recordWatchedOpenRouterCompletion(span, watch, args[0], response, Date.now() - startedAt);
+			try {
+				options?.onResponse?.(response);
+			} catch {}
+			return response;
+		};
+		Object.defineProperty(wrappedFn, "name", {
+			configurable: true,
+			value: methodName
+		});
+		return wrappedFn;
+	}
 	withSpan(traceFunctionKey, optionsOrFn, maybeFn) {
 		this.simulationPlan.refresh();
 		const options = typeof optionsOrFn === "function" ? {} : optionsOrFn;
@@ -5554,7 +5953,9 @@ var Bitfab = class {
 				const seedCtxForTraceId = parentContext ? null : getSeedContext();
 				const traceId = parentContext?.traceId ?? replayCtxForTraceId?.traceId ?? seedCtxForTraceId?.traceId ?? randomUuid();
 				const spanId = randomUuid();
-				const parentSpanId = (parentContext !== void 0 && options.instrumentation === "trace" ? openAiAgentsParentSpanId(parentContext.traceId, parentContext.openAiAgentsSpanIds) : void 0) ?? parentContext?.spanId ?? null;
+				const openAiAgentsParent = parentContext !== void 0 && options.instrumentation === "trace" ? openAiAgentsParentSpanId(parentContext.traceId, parentContext.openAiAgentsSpanIds) : void 0;
+				const frameworkParent = parentContext !== void 0 ? options.frameworkParentSpanId : void 0;
+				const parentSpanId = openAiAgentsParent ?? frameworkParent ?? parentContext?.spanId ?? null;
 				const isRootSpan = parentSpanId === null;
 				if (options.enclosingTrace !== void 0) {
 					options.enclosingTrace.link.traceId = traceId;
@@ -5709,7 +6110,7 @@ var Bitfab = class {
 						});
 						throw error;
 					});
-					if (isAsyncGenerator(result)) return wrapAsyncGenerator(result, newStack, sendSpan);
+					if (isAsyncGenerator(result)) return wrapAsyncGenerator(result, newStack, sendSpan, options.finalize);
 					recordSpan(result);
 					return result;
 				};
@@ -6139,7 +6540,7 @@ var Bitfab = class {
 			await runWithSeedContext({ traceId }, async () => target(...args));
 		} finally {
 			try {
-				const { flushTraces } = await Promise.resolve().then(() => require("./http-psFF1ptn.cjs")).then((n) => n.http_exports);
+				const { flushTraces } = await Promise.resolve().then(() => require("./http-Dqc-f3sp.cjs")).then((n) => n.http_exports);
 				await flushTraces(3e4);
 			} finally {
 				unrecorded = activeTraceStates.delete(traceId);
@@ -6177,7 +6578,7 @@ var Bitfab = class {
 				type: "agent",
 				surface: "neutral"
 			};
-			replayFn = this.withSpan(traceFunctionKey, replayRootOptions, fn);
+			replayFn = this.withSpan(traceFunctionKey, replayRootOptions, drainingStreams(fn));
 		} else if (wrappedKey !== traceFunctionKey) throw new require_http.BitfabError(`Function is wrapped with trace function key '${wrappedKey}' but replay was called with '${traceFunctionKey}'. Pass matching keys, or pass the unwrapped function to replay it under the explicit key.`);
 		const { replay: doReplay } = await Promise.resolve().then(() => replay_exports);
 		return doReplay(this.httpClient, this.serviceUrl, traceFunctionKey, replayFn, options, this.mockOverrides);
@@ -6219,6 +6620,9 @@ var BitfabFunction = class {
 	getLangGraphIntegration(options) {
 		return this.client.getLangGraphIntegration(this.traceFunctionKey, options);
 	}
+	getDeepAgentsIntegration(options) {
+		return this.client.getDeepAgentsIntegration(this.traceFunctionKey, options);
+	}
 	getRivetIntegration(options) {
 		return this.client.getRivetIntegration(this.traceFunctionKey, options);
 	}
@@ -6228,6 +6632,9 @@ var BitfabFunction = class {
 	wrapJev(method, options) {
 		return this.client.wrapJev(method, options);
 	}
+	wrapOpenRouter(method, options) {
+		return this.client.wrapOpenRouter(method, options);
+	}
 };
 //#endregion
 //#region src/finalizers.ts
@@ -6330,7 +6737,7 @@ async function seedFromRegistry(registry, pipeline, cases, options = {}) {
 		metadata: seedCase.metadata,
 		sessionId: seedCase.sessionId
 	}));
-	const { flushTraces } = await Promise.resolve().then(() => require("./http-psFF1ptn.cjs")).then((n) => n.http_exports);
+	const { flushTraces } = await Promise.resolve().then(() => require("./http-Dqc-f3sp.cjs")).then((n) => n.http_exports);
 	await flushTraces(3e4);
 	return {
 		pipeline,
@@ -6574,4 +6981,4 @@ Object.defineProperty(exports, "serializeReplayResult", {
 	}
 });
 
-//# sourceMappingURL=src-Dze8o7iU.cjs.map
+//# sourceMappingURL=src-CdQRfsrA.cjs.map