@aztec/foundation 0.0.0-test.1 → 0.0.1-commit.001888fc

package/src/log/libp2p_logger.ts

@@ -2,15 +2,17 @@ import type { ComponentLogger, Logger } from '@libp2p/interface';
 
 import { getLogLevelFromFilters } from './log-filters.js';
 import type { LogLevel } from './log-levels.js';
-import { logFilters, logger } from './pino-logger.js';
+import { type LoggerBindings, logFilters, logger } from './pino-logger.js';
 
 /**
  * Creates a libp2p compatible logger that wraps our pino logger.
  * This adapter implements the ComponentLogger interface required by libp2p.
+ * @param namespace - Base namespace for the logger
+ * @param bindings - Optional bindings to pass to the logger (actor, instanceId)
  */
-export function createLibp2pComponentLogger(namespace: string): ComponentLogger {
+export function createLibp2pComponentLogger(namespace: string, bindings?: LoggerBindings): ComponentLogger {
   return {
-    forComponent: (component: string) => createLibp2pLogger(`${namespace}:${component}`),
+    forComponent: (component: string) => createLibp2pLogger(`${namespace}:${component}`, bindings),
   };
 }
 
@@ -24,16 +26,21 @@ function replaceFormatting(message: string) {
   return message.replace(/(%p|%a)/g, '%s');
 }
 
-function createLibp2pLogger(component: string): Logger {
+function createLibp2pLogger(component: string, bindings?: LoggerBindings): Logger {
   // Create a direct pino logger instance for libp2p that supports string interpolation
-  const log = logger.child({ module: component }, { level: getLogLevelFromFilters(logFilters, component) });
+  const actor = bindings?.actor;
+  const instanceId = bindings?.instanceId;
+  const log = logger.child(
+    { module: component, ...(actor && { actor }), ...(instanceId && { instanceId }) },
+    { level: getLogLevelFromFilters(logFilters, component) },
+  );
 
   const logIfEnabled = (level: LogLevel, message: string, ...args: unknown[]) => {
     if (!log.isLevelEnabled(level)) {
       return;
     }
 
-    log[level](replaceFormatting(message), ...args);
+    log[level](replaceFormatting(message), ...formatArgs(message, args));
   };
 
   // Default log level is trace as this is super super noisy
@@ -45,7 +52,7 @@ function createLibp2pLogger(component: string): Logger {
     enabled: log.isLevelEnabled('debug'),
     error(message: string, ...args: unknown[]) {
       // We write error outputs as debug as they are often expected, e.g. connection errors can happen in happy paths
-      logIfEnabled('debug', message, ...args);
+      logIfEnabled('debug', `error: ${message}`, ...args);
     },
 
     debug(message: string, ...args: unknown[]) {
@@ -65,3 +72,22 @@ function createLibp2pLogger(component: string): Logger {
     },
   });
 }
+
+function formatArgs(message: string, args: unknown[]) {
+  if (!args) {
+    return args;
+  }
+  return args.map(arg => {
+    if (
+      typeof arg === 'object' &&
+      arg &&
+      'err' in arg &&
+      arg.err instanceof Error &&
+      'type' in arg.err &&
+      arg.err.type === 'AbortError'
+    ) {
+      delete arg.err; // Remove the AbortError from the logs
+    }
+    return arg;
+  });
+}

package/src/log/log-filters.ts

@@ -19,22 +19,40 @@ export function getLogLevelFromFilters(filters: LogFilters, module: string): Log
   return undefined;
 }
 
-export function assertLogLevel(level: string): asserts level is LogLevel {
-  if (!LogLevels.includes(level as LogLevel)) {
-    throw new Error(`Invalid log level: ${level}`);
+/**
+ * Parses the LOG_LEVEL env string into a default level and per-module filter overrides.
+ *
+ * Format: `<default_level>;<level>:<module1>,<module2>;<level>:<module3>;...`
+ * - First segment (before the first `;`) is the default log level for all modules.
+ * - Remaining segments are `level:module` pairs: apply the given level to the listed modules (comma-separated).
+ * - Later filters override earlier ones for overlapping module matches.
+ * - The `aztec:` prefix is stripped from module names; spaces are trimmed.
+ *
+ * @example
+ * ```ts
+ * parseLogLevel('debug;warn:module1,module2;error:module3', 'info')
+ * // => ['debug', [['module3', 'error'], ['module2', 'warn'], ['module1', 'warn']]]
+ * ```
+ */
+export function parseLogLevelEnvVar(
+  logLevelEnvVar: string | undefined,
+  defaultLevel: LogLevel,
+): [LogLevel, LogFilters] {
+  if (!logLevelEnvVar) {
+    return [defaultLevel, []];
   }
+  const [level] = logLevelEnvVar.split(';', 1);
+  assertValidLogLevel(level);
+  return [level, parseFilters(logLevelEnvVar.slice(level.length + 1))];
 }
 
-export function parseEnv(env: string | undefined, defaultLevel: LogLevel): [LogLevel, LogFilters] {
-  if (!env) {
-    return [defaultLevel, []];
+function assertValidLogLevel(level: string): asserts level is LogLevel {
+  if (!LogLevels.includes(level as LogLevel)) {
+    throw new Error(`Invalid log level: ${level}`);
   }
-  const [level] = env.split(';', 1);
-  assertLogLevel(level);
-  return [level, parseFilters(env.slice(level.length + 1))];
 }
 
-export function parseFilters(definition: string | undefined): LogFilters {
+function parseFilters(definition: string | undefined): LogFilters {
   if (!definition) {
     return [];
   }
@@ -48,7 +66,7 @@ export function parseFilters(definition: string | undefined): LogFilters {
       throw new Error(`Invalid log filter statement: ${statement}`);
     }
     const sanitizedLevel = level.trim().toLowerCase();
-    assertLogLevel(sanitizedLevel);
+    assertValidLogLevel(sanitizedLevel);
     for (const module of modules.split(',')) {
       filters.push([
         module

package/src/log/log_fn.ts

@@ -1,5 +1,5 @@
 /** Structured log data to include with the message. */
-export type LogData = Record<string, string | number | bigint | boolean | { toString(): string } | undefined>;
+export type LogData = Record<string, string | number | bigint | boolean | { toString(): string } | undefined | null>;
 
 /** A callable logger instance. */
 export type LogFn = (msg: string, data?: unknown) => void;

package/src/log/noir_debug_log_util.ts

@@ -0,0 +1,21 @@
+interface Printable {
+  toString(): string;
+}
+
+/**
+ * Format a debug string filling in `'{0}'` entries with their
+ * corresponding values from the args array, amd `'{}'` with the whole array.
+ *
+ * @param formatStr - str of form `'this is a string with some entries like {0} and {1}'`
+ * @param args - array of fields to fill in the string format entries with
+ * @returns formatted string
+ */
+export function applyStringFormatting(formatStr: string, args: Printable[]): string {
+  return formatStr
+    .replace(/{(\d+)}/g, (match, index) => {
+      return typeof args[index] === 'undefined' ? match : args[index].toString();
+    })
+    .replace(/{}/g, (_match, _index) => {
+      return args.toString();
+    });
+}

package/src/log/pino-logger-server.ts

@@ -0,0 +1,25 @@
+import { AsyncLocalStorage } from 'node:async_hooks';
+
+import { type LoggerBindings, addLogBindingsHandler, removeLogBindingsHandler } from './pino-logger.js';
+
+/** AsyncLocalStorage for logger bindings context propagation (Node.js only). */
+const bindingsStorage = new AsyncLocalStorage<LoggerBindings>();
+
+/** Returns the current bindings from AsyncLocalStorage, if any. */
+export function getBindings(): LoggerBindings | undefined {
+  return bindingsStorage.getStore();
+}
+
+/**
+ * Runs a callback within a bindings context. All loggers created within the callback
+ * will automatically inherit the bindings (actor, instanceId) via the log bindings handler.
+ */
+export async function withLoggerBindings<T>(bindings: LoggerBindings, callback: () => Promise<T>): Promise<T> {
+  const handler = () => bindingsStorage.getStore();
+  addLogBindingsHandler(handler);
+  try {
+    return await bindingsStorage.run(bindings, callback);
+  } finally {
+    removeLogBindingsHandler(handler);
+  }
+}

package/src/log/pino-logger.ts

@@ -1,19 +1,63 @@
-import { createColors, isColorSupported } from 'colorette';
+import { type Color, createColors, isColorSupported } from 'colorette';
 import isNode from 'detect-node';
 import { pino, symbols } from 'pino';
 import type { Writable } from 'stream';
 import { inspect } from 'util';
 
 import { compactArray } from '../collection/array.js';
-import { type EnvVar, parseBooleanEnv } from '../config/index.js';
+import type { EnvVar } from '../config/index.js';
+import { parseBooleanEnv } from '../config/parse-env.js';
+import { convertBigintsToStrings } from './bigint-utils.js';
 import { GoogleCloudLoggerConfig } from './gcloud-logger-config.js';
-import { getLogLevelFromFilters, parseEnv } from './log-filters.js';
+import { getLogLevelFromFilters, parseLogLevelEnvVar } from './log-filters.js';
 import type { LogLevel } from './log-levels.js';
 import type { LogData, LogFn } from './log_fn.js';
 
-export function createLogger(module: string): Logger {
-  module = logNameHandlers.reduce((moduleName, handler) => handler(moduleName), module.replace(/^aztec:/, ''));
-  const pinoLogger = logger.child({ module }, { level: getLogLevelFromFilters(logFilters, module) });
+/** Optional bindings to pass to createLogger for additional context. */
+export type LoggerBindings = {
+  /** Actor label shown in logs (e.g., 'MAIN', 'prover-node'). */
+  actor?: string;
+  /** Instance identifier for distinguishing multiple instances of the same component. */
+  instanceId?: string;
+};
+
+// Allow global hooks for providing default bindings.
+// Used by withLoggerBindings in pino-logger-server to propagate bindings via AsyncLocalStorage.
+type LogBindingsHandler = () => LoggerBindings | undefined;
+const logBindingsHandlers: LogBindingsHandler[] = [];
+
+export function addLogBindingsHandler(handler: LogBindingsHandler): void {
+  logBindingsHandlers.push(handler);
+}
+
+export function removeLogBindingsHandler(handler: LogBindingsHandler) {
+  const index = logBindingsHandlers.indexOf(handler);
+  if (index !== -1) {
+    logBindingsHandlers.splice(index, 1);
+  }
+}
+
+function getBindingsFromHandlers(): LoggerBindings | undefined {
+  for (const handler of logBindingsHandlers) {
+    const bindings = handler();
+    if (bindings) {
+      return bindings;
+    }
+  }
+  return undefined;
+}
+
+export function createLogger(module: string, bindings?: LoggerBindings): Logger {
+  module = module.replace(/^aztec:/, '');
+
+  const resolvedBindings = { ...getBindingsFromHandlers(), ...bindings };
+  const actor = resolvedBindings?.actor;
+  const instanceId = resolvedBindings?.instanceId;
+
+  const pinoLogger = logger.child(
+    { module, ...(actor && { actor }), ...(instanceId && { instanceId }) },
+    { level: getLogLevelFromFilters(logFilters, module) },
+  );
 
   // We check manually for isLevelEnabled to avoid calling processLogData unnecessarily.
   // Note that isLevelEnabled is missing from the browser version of pino.
@@ -37,14 +81,30 @@ export function createLogger(module: string): Logger {
     debug: (msg: string, data?: unknown) => logFn('debug', msg, data),
     /** Log as trace. Use for when we want to denial-of-service any recipient of the logs. */
     trace: (msg: string, data?: unknown) => logFn('trace', msg, data),
+    /** Level of the logger */
     level: pinoLogger.level as LogLevel,
     /** Whether the given level is enabled for this logger. */
     isLevelEnabled: (level: LogLevel) => isLevelEnabled(pinoLogger, level),
     /** Module name for the logger. */
     module,
+    /** Creates another logger by extending this logger module name and preserving bindings. */
+    createChild: (childModule: string) => createLogger(`${module}:${childModule}`, { actor, instanceId }),
+    /** Returns the bindings (actor, instanceId) for this logger. */
+    getBindings: () => ({ actor, instanceId }),
   };
 }
 
+/**
+ * Returns a logger for the given module. If loggerOrBindings is already a Logger, returns it directly.
+ * Otherwise, creates a new logger with the given module name and bindings.
+ */
+export function resolveLogger(module: string, loggerOrBindings?: Logger | LoggerBindings): Logger {
+  if (loggerOrBindings && 'info' in loggerOrBindings) {
+    return loggerOrBindings as Logger;
+  }
+  return createLogger(module, loggerOrBindings);
+}
+
 // Allow global hooks for processing log data.
 // Used for injecting OTEL trace_id in telemetry client.
 type LogDataHandler = (data: LogData) => LogData;
@@ -58,31 +118,6 @@ function processLogData(data: LogData): LogData {
   return logDataHandlers.reduce((accum, handler) => handler(accum), data);
 }
 
-// Allow global hooks for tweaking module names.
-// Used in tests to add a uid to modules, so we can differentiate multiple nodes in the same process.
-type LogNameHandler = (module: string) => string;
-const logNameHandlers: LogNameHandler[] = [];
-
-export function addLogNameHandler(handler: LogNameHandler): void {
-  logNameHandlers.push(handler);
-}
-
-export function removeLogNameHandler(handler: LogNameHandler) {
-  const index = logNameHandlers.indexOf(handler);
-  if (index !== -1) {
-    logNameHandlers.splice(index, 1);
-  }
-}
-
-/** Creates all loggers within the given callback with the suffix appended to the module name. */
-export async function withLogNameSuffix<T>(suffix: string, callback: () => Promise<T>): Promise<T> {
-  const logNameHandler = (module: string) => `${module}:${suffix}`;
-  addLogNameHandler(logNameHandler);
-  const result = await callback();
-  removeLogNameHandler(logNameHandler);
-  return result;
-}
-
 // Patch isLevelEnabled missing from pino/browser.
 function isLevelEnabled(logger: pino.Logger<'verbose', boolean>, level: LogLevel): boolean {
   return typeof logger.isLevelEnabled === 'function'
@@ -92,18 +127,48 @@ function isLevelEnabled(logger: pino.Logger<'verbose', boolean>, level: LogLevel
 
 // Load log levels from environment variables.
 const defaultLogLevel = process.env.NODE_ENV === 'test' ? 'silent' : 'info';
-export const [logLevel, logFilters] = parseEnv(process.env.LOG_LEVEL, defaultLogLevel);
+export const [logLevel, logFilters] = parseLogLevelEnvVar(process.env.LOG_LEVEL, defaultLogLevel);
 
 // Define custom logging levels for pino.
 const customLevels = { verbose: 25 };
 
 // Global pino options, tweaked for google cloud if running there.
 const useGcloudLogging = parseBooleanEnv(process.env['USE_GCLOUD_LOGGING' satisfies EnvVar]);
+
+const redactedPaths = [
+  'validatorPrivateKeys',
+  // for both the validator and the prover
+  'publisherPrivateKeys',
+  'peerIdPrivateKey',
+  // bot keys
+  'l1PrivateKey',
+  'senderPrivateKey',
+  // blob sink
+  'l1ConsensusHostApiKeys',
+  // sensitive options used in the CLI
+  'privateKey',
+  'mnemonic',
+  'l1Mnemonic',
+  'l1PrivateKey',
+];
+
 const pinoOpts: pino.LoggerOptions<keyof typeof customLevels> = {
   customLevels,
   messageKey: 'msg',
   useOnlyCustomLevels: false,
   level: logLevel,
+  redact: {
+    paths: [
+      ...redactedPaths,
+      ...redactedPaths.map(p => `config.${p}`),
+      ...redactedPaths.map(p => `cfg.${p}`),
+      ...redactedPaths.map(p => `options.${p}`),
+      ...redactedPaths.map(p => `opts.${p}`),
+    ],
+  },
+  formatters: {
+    log: obj => convertBigintsToStrings(obj) as Record<string, unknown>,
+  },
   ...(useGcloudLogging ? GoogleCloudLoggerConfig : {}),
 };
 
@@ -115,22 +180,90 @@ export const levels = {
 // Transport options for pretty logging to stderr via pino-pretty.
 const colorEnv = process.env['FORCE_COLOR' satisfies EnvVar];
 const useColor = colorEnv === undefined ? isColorSupported : parseBooleanEnv(colorEnv);
-const { bold, reset } = createColors({ useColor });
-export const pinoPrettyOpts = {
+const { bold, reset, cyan, magenta, yellow, blue, green, magentaBright, yellowBright, blueBright, greenBright } =
+  createColors({ useColor });
+
+// Per-actor coloring: each unique actor gets a different color for easier visual distinction.
+// Disabled when LOG_NO_COLOR_PER_ACTOR is set to a truthy value.
+const useColorPerActor = useColor && !parseBooleanEnv(process.env['LOG_NO_COLOR_PER_ACTOR' satisfies EnvVar]);
+const actorColors: Color[] = [yellow, magenta, blue, green, magentaBright, yellowBright, blueBright, greenBright];
+const actorColorMap = new Map<string, Color>();
+let nextColorIndex = 0;
+
+/** Returns the color function assigned to a given actor, assigning a new one if needed. */
+export function getActorColor(actor: string): Color {
+  let color = actorColorMap.get(actor);
+  if (!color) {
+    color = actorColors[nextColorIndex % actorColors.length];
+    actorColorMap.set(actor, color);
+    nextColorIndex++;
+  }
+  return color;
+}
+
+/** Resets the actor-to-color mapping. Useful for testing. */
+export function resetActorColors(): void {
+  actorColorMap.clear();
+  nextColorIndex = 0;
+}
+
+// String template for messageFormat (used in worker threads and when per-actor coloring is disabled).
+const messageFormatString = `${bold('{module}')}{if actor} ${cyan('{actor}')}{end}{if instanceId} ${reset(cyan('{instanceId}'))}{end} ${reset('{msg}')}`;
+
+// Function for messageFormat when per-actor coloring is enabled (can only be used in-process, not worker threads).
+type LogObject = { actor?: string; module?: string; instanceId?: string; msg?: string };
+
+/** Formats a log message with per-actor coloring. Actor, module, and instanceId share the same color. */
+export function formatLogMessage(log: LogObject, messageKey: string): string {
+  const actor = log.actor;
+  const module = log.module ?? '';
+  const instanceId = log.instanceId;
+  const msg = log[messageKey as keyof LogObject] ?? '';
+
+  // Use actor color for actor, module, and instanceId when actor is present
+  const color = actor ? getActorColor(actor) : cyan;
+
+  let result = bold(color(module));
+  if (actor) {
+    result += ' ' + color(actor);
+  }
+  if (instanceId) {
+    result += ' ' + reset(color(instanceId));
+  }
+  result += ' ' + reset(String(msg));
+  return result;
+}
+
+// Base options for pino-pretty (shared between transport and direct use).
+const pinoPrettyBaseOpts = {
   destination: 2,
   sync: true,
   colorize: useColor,
-  ignore: 'module,pid,hostname,trace_id,span_id,trace_flags,severity',
-  messageFormat: `${bold('{module}')} ${reset('{msg}')}`,
+  ignore: 'module,actor,instanceId,pid,hostname,trace_id,span_id,trace_flags,severity',
   customLevels: 'fatal:60,error:50,warn:40,info:30,verbose:25,debug:20,trace:10',
   customColors: 'fatal:bgRed,error:red,warn:yellow,info:green,verbose:magenta,debug:blue,trace:gray',
   minimumLevel: 'trace' as const,
   singleLine: !parseBooleanEnv(process.env['LOG_MULTILINE' satisfies EnvVar]),
 };
 
+/**
+ * Pino-pretty options for direct use (e.g., jest/setup.mjs).
+ * Includes function-based messageFormat for per-actor coloring when enabled.
+ */
+export const pinoPrettyOpts = {
+  ...pinoPrettyBaseOpts,
+  messageFormat: useColorPerActor ? formatLogMessage : messageFormatString,
+};
+
+// Transport options use string template only (functions can't be serialized to worker threads).
+const prettyTransportOpts = {
+  ...pinoPrettyBaseOpts,
+  messageFormat: messageFormatString,
+};
+
 const prettyTransport: pino.TransportTargetOptions = {
   target: 'pino-pretty',
-  options: pinoPrettyOpts,
+  options: prettyTransportOpts,
   level: 'trace',
 };
 
@@ -221,7 +354,7 @@ export function registerLoggingStream(stream: Writable): void {
 }
 
 /** Log function that accepts an exception object */
-type ErrorLogFn = (msg: string, err?: Error | unknown, data?: LogData) => void;
+type ErrorLogFn = (msg: string, err?: unknown, data?: LogData) => void;
 
 /**
  * Logger that supports multiple severity levels.
@@ -230,6 +363,9 @@ export type Logger = { [K in LogLevel]: LogFn } & { /** Error log function */ er
   level: LogLevel;
   isLevelEnabled: (level: LogLevel) => boolean;
   module: string;
+  createChild: (childModule: string) => Logger;
+  /** Returns the bindings (actor, instanceId) for this logger. */
+  getBindings: () => LoggerBindings;
 };
 
 /**
@@ -238,6 +374,6 @@ export type Logger = { [K in LogLevel]: LogFn } & { /** Error log function */ er
  * @param err - Error to log
  * @returns A string with both the log message and the error message.
  */
-function formatErr(msg: string, err?: Error | unknown): string {
+function formatErr(msg: string, err?: unknown): string {
   return err ? `${msg}: ${inspect(err)}` : msg;
 }

package/src/message/index.ts

@@ -31,7 +31,11 @@ interface TypedMessageLike {
 }
 
 export class TypedMessage<T, B> {
-  public constructor(public readonly msgType: T, public readonly header: MessageHeader, public readonly value: B) {}
+  public constructor(
+    public readonly msgType: T,
+    public readonly header: MessageHeader,
+    public readonly value: B,
+  ) {}
 
   static fromMessagePack<T, B>(data: TypedMessageLike): TypedMessage<T, B> {
     return new TypedMessage<T, B>(data['msgType'] as T, MessageHeader.fromMessagePack(data['header']), data['value']);

package/src/mutex/mutex_database.ts

@@ -3,10 +3,9 @@
  * Provides functionality for acquiring, extending, and releasing locks on resources to ensure exclusive access and prevent conflicts in concurrent applications.
  */
 export interface MutexDatabase {
-  // eslint-disable-next-line jsdoc/require-jsdoc
   acquireLock(name: string, timeout: number): Promise<boolean>;
-  // eslint-disable-next-line jsdoc/require-jsdoc
+
   extendLock(name: string, timeout: number): Promise<void>;
-  // eslint-disable-next-line jsdoc/require-jsdoc
+
   releaseLock(name: string): Promise<void>;
 }

package/src/number/index.ts

@@ -0,0 +1,14 @@
+/** Bounds the given value between the min and max provided (both inclusive) */
+export function boundInclusive(value: number, min: number, max: number): number {
+  if (min > max) {
+    throw new Error('Minimum bound cannot be greater than maximum bound');
+  }
+
+  if (value < min) {
+    return min;
+  } else if (value > max) {
+    return max;
+  } else {
+    return value;
+  }
+}

package/src/profiler/index.ts

@@ -0,0 +1 @@
+export * from './profiler.js';

package/src/profiler/profiler.ts

@@ -0,0 +1,125 @@
+import { AsyncLocalStorage } from 'node:async_hooks';
+import * as fs from 'node:fs';
+import * as path from 'node:path';
+import { performance } from 'node:perf_hooks';
+
+interface Span {
+  label: string;
+  start: number;
+  dur: number;
+  count: number;
+  children: Span[];
+  parent: Span | undefined;
+}
+
+interface ProfileData {
+  spans: SerializedSpan[];
+  timestamp: string;
+  totalTime: number;
+}
+
+interface SerializedSpan {
+  label: string;
+  dur: number;
+  count: number;
+  children: SerializedSpan[];
+}
+
+const als = new AsyncLocalStorage<Span>();
+const roots: Span[] = [];
+
+function reset(): void {
+  roots.length = 0;
+}
+
+// Strip out circular references (parent) and unused fields (start) for JSON serialization
+function serializeSpans(spans: Span[]): SerializedSpan[] {
+  return spans.map(span => ({
+    label: span.label,
+    dur: span.dur,
+    count: span.count,
+    children: serializeSpans(span.children),
+  }));
+}
+
+let i = 0;
+function save(): void {
+  if (roots.length === 0) {
+    return;
+  }
+
+  // Find max single execution time across all spans (dur/count since dur is accumulated)
+  const findMaxSingleDuration = (spans: Span[]): number => {
+    let max = 0;
+    for (const span of spans) {
+      const singleDur = span.dur / span.count;
+      max = Math.max(max, singleDur);
+      if (span.children.length > 0) {
+        max = Math.max(max, findMaxSingleDuration(span.children));
+      }
+    }
+    return max;
+  };
+
+  const profileData: ProfileData = {
+    spans: serializeSpans(roots),
+    timestamp: new Date().toISOString(),
+    totalTime: findMaxSingleDuration(roots),
+  };
+
+  const profilePath = path.join(process.cwd(), `profile-${i++}.json`);
+  process.stdout.write(`Writing profile data to ${profilePath}\n`);
+  fs.writeFileSync(profilePath, JSON.stringify(profileData, null, 2));
+}
+
+// Hook into Jest to save after each test
+if (typeof afterEach === 'function') {
+  afterEach(() => {
+    save();
+    reset();
+  });
+}
+
+// Also save on process exit for non-Jest environments
+process.on('exit', () => {
+  save();
+});
+
+// Wrapper for async functions to maintain context properly
+async function runAsync<ReturnType>(label: string, fn: () => Promise<ReturnType>): Promise<ReturnType> {
+  const parent = als.getStore();
+
+  // Check if we already have a span with this label in the current context
+  let existingSpan: Span | undefined;
+  if (parent) {
+    existingSpan = parent.children.find(c => c.label === label);
+  } else {
+    existingSpan = roots.find(r => r.label === label);
+  }
+
+  let span: Span;
+  if (existingSpan) {
+    // Reuse existing span and increment count
+    span = existingSpan;
+    span.count++;
+  } else {
+    // Create new span
+    span = { label, start: performance.now(), dur: 0, count: 1, children: [], parent };
+    if (parent) {
+      parent.children.push(span);
+    } else {
+      roots.push(span);
+    }
+  }
+
+  const startTime = performance.now();
+  const result: ReturnType = await als.run(span, fn);
+  const elapsed = performance.now() - startTime;
+
+  // Add to total duration (for averaging)
+  span.dur += elapsed;
+
+  return result;
+}
+
+export const profiler = { reset, runAsync };