bitfab 0.68.2 → 0.68.3

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

@@ -1,5 +1,5 @@
 const require_bamlPeer = require("./bamlPeer-BUoZgNxZ.cjs");
-const require_http = require("./http-Dqc-f3sp.cjs");
+const require_http = require("./http-bBRhl6b9.cjs");
 const require_autoTrace = require("./autoTrace-D3-OZcmR.cjs");
 let superjson = require("superjson");
 superjson = require_bamlPeer.__toESM(superjson, 1);
@@ -1994,6 +1994,12 @@ function createIntegrationFixtureContext(name, value) {
 		...serialized.meta !== void 0 && { valueMeta: serialized.meta }
 	} };
 }
+function hasIntegrationFixture(contexts, name) {
+	return contexts.some((context) => {
+		const fixture = isRecord$3(context) ? context[INTEGRATION_FIXTURE_CONTEXT_KEY] : void 0;
+		return isRecord$3(fixture) && fixture.version === INTEGRATION_FIXTURE_VERSION && fixture.name === name;
+	});
+}
 function extractIntegrationFixtures(spanData) {
 	const contexts = spanData.contexts;
 	if (!Array.isArray(contexts)) return;
@@ -2001,7 +2007,7 @@ function extractIntegrationFixtures(spanData) {
 	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;
+		if (!isRecord$3(fixture) || fixture.version !== INTEGRATION_FIXTURE_VERSION || typeof fixture.name !== "string" || fixtures.has(fixture.name)) continue;
 		fixtures.set(fixture.name, deserializeValue({
 			json: fixture.value,
 			...fixture.valueMeta !== void 0 && { meta: fixture.valueMeta }
@@ -2035,14 +2041,42 @@ function fromBase64(text) {
 	if (buffer !== void 0) return new Uint8Array(buffer.from(text, "base64"));
 	return Uint8Array.from(atob(text), (char) => char.charCodeAt(0));
 }
-async function newMemorySaver() {
+async function loadLangGraphPeer() {
 	try {
-		const { MemorySaver } = await require_bamlPeer.importOptionalPeer(["@langchain", "langgraph"]);
-		return new MemorySaver();
+		return await require_bamlPeer.importOptionalPeer(["@langchain", "langgraph"]);
 	} catch {
 		throw new require_http.BitfabError("Restoring a LangGraph thread during replay requires @langchain/langgraph.", "https://docs.bitfab.ai/frameworks/langgraph");
 	}
 }
+async function newMemorySaver() {
+	const { MemorySaver } = await loadLangGraphPeer();
+	return new MemorySaver();
+}
+async function unstoredDeltaChannels(graph, tuple) {
+	const channels = isRecord$2(graph) ? graph.channels : void 0;
+	if (!isRecord$2(channels)) return [];
+	const { DeltaChannel } = await loadLangGraphPeer();
+	if (DeltaChannel === void 0) return [];
+	const versions = isRecord$2(tuple.checkpoint.channel_versions) ? tuple.checkpoint.channel_versions : {};
+	const values = isRecord$2(tuple.checkpoint.channel_values) ? tuple.checkpoint.channel_values : {};
+	return Object.entries(channels).filter(([name, spec]) => spec instanceof DeltaChannel && name in versions && !(name in values)).map(([name]) => name);
+}
+async function selfContained(graph, checkpointer, tuple) {
+	if (tuple === void 0) return;
+	const names = await unstoredDeltaChannels(graph, tuple);
+	if (names.length === 0) return tuple;
+	const { values } = await graph.getState(withConfigurable(tuple.config, { [CHECKPOINTER_CONFIG_KEY]: checkpointer }));
+	return {
+		...tuple,
+		checkpoint: {
+			...tuple.checkpoint,
+			channel_values: {
+				...isRecord$2(tuple.checkpoint.channel_values) ? tuple.checkpoint.channel_values : {},
+				...Object.fromEntries(names.filter((name) => name in values).map((name) => [name, values[name]]))
+			}
+		}
+	};
+}
 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;
@@ -2102,6 +2136,14 @@ async function restoredCheckpointer(recorded) {
 		thread: recorded.config
 	};
 }
+const REPLAY_THREADS = /* @__PURE__ */ new WeakMap();
+function replayThreadsOf(replay) {
+	const existing = REPLAY_THREADS.get(replay);
+	if (existing !== void 0) return existing;
+	const threads = /* @__PURE__ */ new Map();
+	REPLAY_THREADS.set(replay, threads);
+	return threads;
+}
 var GraphRun = class {
 	constructor(host, graph, config) {
 		this.host = host;
@@ -2111,19 +2153,21 @@ var GraphRun = class {
 		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;
+		this.replay = require_http.getReplayContext();
+		this.reportsStreamResult = this.replay !== null && !this.opensRoot && active?.isRoot === true;
+	}
+	get fixtureName() {
+		return threadFixtureName(this.host.traceFunctionKey);
 	}
 	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();
+		const run = async (rootInput) => {
+			const target = await this.runTarget();
+			await this.recordCheckpoint(target);
+			const result = await this.traced.invoke(await this.host.reviveRootInput(rootInput), target.config);
+			await target.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));
+		};
+		return this.opensRoot ? this.host.rootSpan(run)(this.host.encodeRootInput(input)) : run(input);
 	}
 	stream(input) {
 		return this.chunks(input, async (rootInput, config) => this.traced.stream(rootInput, config));
@@ -2132,61 +2176,67 @@ var GraphRun = class {
 		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 { host, finalState, reportsStreamResult } = this;
+		const runTarget = () => this.runTarget();
+		const recordCheckpoint = (target) => this.recordCheckpoint(target);
+		async function* run(rootInput) {
+			const target = await runTarget();
+			await recordCheckpoint(target);
+			yield* await open(await host.reviveRootInput(rootInput), target.config);
+			await target.confirm();
+			const replay = require_http.getReplayContext();
+			if (reportsStreamResult && replay !== null) replay.streamResult = finalState.output;
 		}
-		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 {
+		return this.opensRoot ? host.rootSpan(run, () => finalState.output)(host.encodeRootInput(input)) : run(input);
+	}
+	async recordCheckpoint({ config, checkpointer }) {
+		const thread = threadConfig(config);
+		const contexts = this.host.traceRootContexts();
+		if (checkpointer === null || thread.thread_id === void 0 || contexts === null || hasIntegrationFixture(contexts, this.fixtureName)) return;
+		const encoded = await encodeCheckpoint(await selfContained(this.graph, checkpointer, await checkpointer.getTuple({ configurable: thread })));
+		if (!hasIntegrationFixture(contexts, this.fixtureName)) contexts.push(createIntegrationFixtureContext(this.fixtureName, encoded));
+	}
+	async runTarget() {
+		const own = savingCheckpointer(this.graph);
+		if (this.replay === null || own === null) return {
 			config: this.config,
+			checkpointer: own,
 			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 { saver, thread } = await this.replayThread(this.replay);
+		const restoredId = (await saver.getTuple({ configurable: thread }))?.checkpoint.id;
 		const confirm = async () => {
-			const latest = await saver.getTuple({ configurable: replayThread });
+			const latest = await saver.getTuple({ configurable: thread });
 			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,
+				...thread,
 				[CHECKPOINTER_CONFIG_KEY]: saver
 			}),
+			checkpointer: saver,
 			confirm
 		};
 	}
+	replayThread(replay) {
+		const threads = replayThreadsOf(replay);
+		const existing = threads.get(this.fixtureName);
+		if (existing !== void 0) return existing;
+		const restored = this.restoreThread(replay);
+		threads.set(this.fixtureName, restored);
+		return restored;
+	}
+	async restoreThread(replay) {
+		const { saver, thread } = await restoredCheckpointer(replay.integrationFixtures?.get(this.fixtureName) ?? null);
+		const configurable = isRecord$2(this.config?.configurable) ? this.config.configurable : {};
+		return {
+			saver,
+			thread: {
+				thread_id: thread.thread_id ?? configurable.thread_id ?? `bitfab-replay-${randomUuid()}`,
+				...thread.checkpoint_ns !== void 0 && { checkpoint_ns: thread.checkpoint_ns }
+			}
+		};
+	}
 };
 //#endregion
 //#region src/langgraphIntegration.ts
@@ -2297,7 +2347,7 @@ var BitfabLangGraphIntegration = class {
 		this.callbackHandler = config.callbackHandler;
 		this.mockToolsOnReplay = config.options?.mockToolsOnReplay ?? true;
 		this.getActiveSpan = config.getActiveSpan ?? (() => null);
-		this.addRootContext = config.addContext ?? (() => {});
+		this.getTraceRootContexts = config.getTraceRootContexts ?? (() => null);
 	}
 	get middleware() {
 		this.agentMiddleware ?? (this.agentMiddleware = {
@@ -2331,7 +2381,7 @@ var BitfabLangGraphIntegration = class {
 				instrumentation: "langgraph",
 				...finalize !== void 0 && { finalize }
 			}, fn),
-			addRootContext: (context) => this.addRootContext(context),
+			traceRootContexts: () => this.getTraceRootContexts(),
 			encodeRootInput: (input) => isCommand(input) ? encodeNativeToolResult(input) : input,
 			reviveRootInput: async (input) => isEncodedToolResult(input) ? await reviveToolResult(input, "") : input
 		});
@@ -5240,6 +5290,14 @@ const noOpTrace = {
 	addContext() {},
 	drop() {}
 };
+function currentTraceRootSpan() {
+	const stack = getSpanStack();
+	const traceId = stack[stack.length - 1]?.traceId;
+	for (let index = stack.length - 1; index >= 0; index--) {
+		const context = stack[index];
+		if (context?.isRoot && context.traceId === traceId) return context;
+	}
+}
 function getCurrentSpan() {
 	const stack = getSpanStack();
 	const current = stack[stack.length - 1];
@@ -5761,7 +5819,7 @@ var Bitfab = class {
 				const stack = getSpanStack();
 				return stack[stack.length - 1] ?? null;
 			},
-			addContext: (context) => getCurrentSpan().addContext(context)
+			getTraceRootContexts: () => currentTraceRootSpan()?.details.contexts ?? null
 		});
 	}
 	getDeepAgentsIntegration(traceFunctionKey, options) {
@@ -5782,7 +5840,7 @@ var Bitfab = class {
 				return this.withSpan(traceFunctionKey, spanOptions, fn);
 			},
 			recordRootFixture: (context, recordedRoots) => {
-				const root = getSpanStack()[0];
+				const root = currentTraceRootSpan();
 				if (!root) return;
 				const rootId = `${root.traceId}:${root.spanId}`;
 				if (recordedRoots.has(rootId)) return;
@@ -6540,7 +6598,7 @@ var Bitfab = class {
 			await runWithSeedContext({ traceId }, async () => target(...args));
 		} finally {
 			try {
-				const { flushTraces } = await Promise.resolve().then(() => require("./http-Dqc-f3sp.cjs")).then((n) => n.http_exports);
+				const { flushTraces } = await Promise.resolve().then(() => require("./http-bBRhl6b9.cjs")).then((n) => n.http_exports);
 				await flushTraces(3e4);
 			} finally {
 				unrecorded = activeTraceStates.delete(traceId);
@@ -6737,7 +6795,7 @@ async function seedFromRegistry(registry, pipeline, cases, options = {}) {
 		metadata: seedCase.metadata,
 		sessionId: seedCase.sessionId
 	}));
-	const { flushTraces } = await Promise.resolve().then(() => require("./http-Dqc-f3sp.cjs")).then((n) => n.http_exports);
+	const { flushTraces } = await Promise.resolve().then(() => require("./http-bBRhl6b9.cjs")).then((n) => n.http_exports);
 	await flushTraces(3e4);
 	return {
 		pipeline,
@@ -6981,4 +7039,4 @@ Object.defineProperty(exports, "serializeReplayResult", {
 	}
 });
 
-//# sourceMappingURL=src-CdQRfsrA.cjs.map
+//# sourceMappingURL=src-BRR42NiJ.cjs.map