@cofhe/hardhat-3-plugin 0.7.0 → 0.8.0

package/src/gas.ts

@@ -0,0 +1,340 @@
+import chalk from 'chalk';
+import fs from 'node:fs';
+import path from 'node:path';
+import { parseAbiItem, toEventSelector, toFunctionSelector, toFunctionSignature, type PublicClient } from 'viem';
+import { MOCKS_TASK_MANAGER_ADDRESS } from '@cofhe/sdk';
+import type { ArtifactManager } from 'hardhat/types/artifacts';
+
+/// Topic of MockCoFHE's `MockGasConsumed(uint256)` event, emitted once per block of
+/// mock-only work (op replication, decrypt-task storage, log building) with the gas it
+/// consumed. Summing these per transaction gives the mock overhead of that transaction.
+export const MOCK_GAS_CONSUMED_EVENT = parseAbiItem('event MockGasConsumed(uint256 gas)');
+export const MOCK_GAS_CONSUMED_TOPIC = toEventSelector(MOCK_GAS_CONSUMED_EVENT);
+
+/// Minimal receipt shape - compatible with viem and ethers v6 receipts.
+export type AdjustableGasReceipt = {
+  gasUsed: bigint | number | string;
+  logs?: ReadonlyArray<{ address: string; topics: ReadonlyArray<string>; data: string }>;
+};
+
+export type AdjustedGasBreakdown = {
+  /** Raw gas used by the transaction (includes mock overhead). */
+  gasUsed: bigint;
+  /** Gas consumed by mock-only work (FHE op replication, decrypt-task storage, mock logging). */
+  mockGas: bigint;
+  /** gasUsed - mockGas: an estimate of what the transaction would cost on a real CoFHE network. */
+  adjustedGasUsed: bigint;
+  /** Number of MockGasConsumed events - roughly the number of mock-only blocks executed. */
+  mockGasEvents: number;
+};
+
+/// Computes the gas breakdown of a transaction receipt by summing the MockGasConsumed
+/// events emitted by the mock task manager. Pure function of the receipt - no RPC calls.
+/// On a real network (no mock events in the logs) `adjustedGasUsed` equals `gasUsed`.
+/// Caveat: `mockGas` is measured pre-refund while `gasUsed` is post-refund, so transactions
+/// that earn large gas refunds (storage clearing) can slightly overstate the mock share;
+/// `adjustedGasUsed` is clamped at 0 for the pathological case.
+export const mock_getAdjustedGasBreakdown = (receipt: AdjustableGasReceipt): AdjustedGasBreakdown => {
+  let mockGas = 0n;
+  let mockGasEvents = 0;
+
+  for (const log of receipt.logs ?? []) {
+    if (
+      log.address.toLowerCase() === MOCKS_TASK_MANAGER_ADDRESS.toLowerCase() &&
+      log.topics[0]?.toLowerCase() === MOCK_GAS_CONSUMED_TOPIC
+    ) {
+      mockGas += BigInt(log.data);
+      mockGasEvents += 1;
+    }
+  }
+
+  const gasUsed = BigInt(receipt.gasUsed);
+  return { gasUsed, mockGas, adjustedGasUsed: mockGas > gasUsed ? 0n : gasUsed - mockGas, mockGasEvents };
+};
+
+/// Returns the receipt's gas usage excluding mock overhead - an estimate of what the
+/// transaction would cost on a real CoFHE network. See mock_getAdjustedGasBreakdown.
+export const mock_getAdjustedGasUsed = (receipt: AdjustableGasReceipt): bigint =>
+  mock_getAdjustedGasBreakdown(receipt).adjustedGasUsed;
+
+// =====================
+//  END-OF-RUN SUMMARY
+// =====================
+//
+// Hardhat 3 runs node:test files in worker subprocesses, and every network.connect()
+// creates its own in-process chain - so a single query at the end of the `test` task
+// (which runs in the parent process) can't see any of the transactions. Instead:
+//   1. Each connection is registered here (from the plugin's newConnection hook) when
+//      `cofhe.gasSummary` is enabled.
+//   2. On worker `beforeExit`, the gas rows of every registered connection are collected
+//      (eth_getLogs + receipts) and appended to a JSON file in the hardhat cache dir.
+//   3. The parent's `test` task override merges all row files and prints one table.
+
+const SUMMARY_DIR_NAME = 'cofhe-gas-summary';
+
+/// Serializable per-method aggregate (bigints as strings) passed from workers to the parent.
+export type MockGasSummaryRow = {
+  contract: string;
+  method: string;
+  calls: number;
+  totalGasUsed: string;
+  totalMockGas: string;
+};
+
+type RegisteredConnection = {
+  publicClient: PublicClient;
+  artifacts: ArtifactManager;
+  /** Block at which the mocks were deployed on this connection's chain - the summary only
+   *  scans from here, so a long-lived or forked chain isn't queried from genesis. */
+  fromBlock: bigint;
+};
+
+const registeredConnections: RegisteredConnection[] = [];
+let summaryDir: string | undefined;
+let exitHookInstalled = false;
+let exitHookRan = false;
+
+const shortAddress = (address: string) => `${address.slice(0, 6)}..${address.slice(-4)}`;
+
+/// Builds a selector -> "signature" map from every compiled artifact so summary rows can
+/// show method names instead of raw selectors.
+const buildSelectorNameMap = async (artifacts: ArtifactManager): Promise<Map<string, string>> => {
+  const selectorNames = new Map<string, string>();
+  for (const fqn of await artifacts.getAllFullyQualifiedNames()) {
+    try {
+      const artifact = await artifacts.readArtifact(fqn);
+      for (const item of artifact.abi) {
+        if (item.type !== 'function') continue;
+        const selector = toFunctionSelector(item);
+        if (!selectorNames.has(selector)) selectorNames.set(selector, toFunctionSignature(item));
+      }
+    } catch {
+      // Artifacts without a usable ABI are skipped
+    }
+  }
+  return selectorNames;
+};
+
+/// Resolves deployed addresses to contract names by matching on-chain code against
+/// artifact deployedBytecode. Best effort: contracts with immutables won't match and
+/// fall back to a shortened address.
+const buildAddressNameMap = async (
+  publicClient: PublicClient,
+  artifacts: ArtifactManager,
+  addresses: `0x${string}`[]
+): Promise<Map<string, string>> => {
+  const codeToName = new Map<string, string>();
+  for (const fqn of await artifacts.getAllFullyQualifiedNames()) {
+    try {
+      const artifact = await artifacts.readArtifact(fqn);
+      if (artifact.deployedBytecode && artifact.deployedBytecode !== '0x') {
+        codeToName.set(artifact.deployedBytecode.toLowerCase(), artifact.contractName);
+      }
+    } catch {
+      // skip
+    }
+  }
+
+  const addressNames = new Map<string, string>();
+  for (const address of addresses) {
+    try {
+      const code = (await publicClient.getCode({ address })) ?? '0x';
+      addressNames.set(address, codeToName.get(code.toLowerCase()) ?? shortAddress(address));
+    } catch {
+      addressNames.set(address, shortAddress(address));
+    }
+  }
+  return addressNames;
+};
+
+/// Collects per-method gas aggregates from one connection's chain.
+export const collectMockGasRows = async (
+  publicClient: PublicClient,
+  artifacts: ArtifactManager,
+  fromBlock: bigint = 0n
+): Promise<MockGasSummaryRow[]> => {
+  let logs;
+  try {
+    logs = await publicClient.getLogs({
+      address: MOCKS_TASK_MANAGER_ADDRESS,
+      event: MOCK_GAS_CONSUMED_EVENT,
+      fromBlock,
+      toBlock: 'latest',
+    });
+  } catch {
+    return [];
+  }
+  if (logs.length === 0) return [];
+
+  // Sum mock gas per transaction
+  const mockGasPerTx = new Map<`0x${string}`, bigint>();
+  for (const log of logs) {
+    if (log.transactionHash === null) continue;
+    mockGasPerTx.set(log.transactionHash, (mockGasPerTx.get(log.transactionHash) ?? 0n) + (log.args.gas ?? 0n));
+  }
+
+  // Aggregate per called contract + method
+  type Aggregate = { to: `0x${string}`; selector: string; calls: number; totalGasUsed: bigint; totalMockGas: bigint };
+  const aggregates = new Map<string, Aggregate>();
+  // Deployments (FHE ops in constructors) are attributed to the created contract under a
+  // '(deployment)' method label.
+  for (const [txHash, mockGas] of mockGasPerTx) {
+    try {
+      const receipt = await publicClient.getTransactionReceipt({ hash: txHash });
+      const tx = await publicClient.getTransaction({ hash: txHash });
+      const target = tx.to ?? receipt.contractAddress;
+      if (!target) continue;
+
+      const selector = tx.to ? (tx.input ?? '0x').slice(0, 10) : '(deployment)';
+      const key = `${target}:${selector}`;
+      const aggregate = aggregates.get(key) ?? { to: target, selector, calls: 0, totalGasUsed: 0n, totalMockGas: 0n };
+      aggregate.calls += 1;
+      aggregate.totalGasUsed += receipt.gasUsed;
+      aggregate.totalMockGas += mockGas;
+      aggregates.set(key, aggregate);
+    } catch {
+      // Skip transactions that can't be resolved
+    }
+  }
+  if (aggregates.size === 0) return [];
+
+  const selectorNames = await buildSelectorNameMap(artifacts);
+  const addressNames = await buildAddressNameMap(publicClient, artifacts, [
+    ...new Set([...aggregates.values()].map((a) => a.to)),
+  ]);
+
+  return [...aggregates.values()].map((a) => ({
+    contract: addressNames.get(a.to) ?? shortAddress(a.to),
+    method: selectorNames.get(a.selector) ?? a.selector,
+    calls: a.calls,
+    totalGasUsed: a.totalGasUsed.toString(),
+    totalMockGas: a.totalMockGas.toString(),
+  }));
+};
+
+/// Registers a connection for end-of-run gas summary collection, installing (once) a
+/// process beforeExit hook that dumps the aggregated rows into the hardhat cache dir
+/// for the parent `test` task to merge and print.
+export const registerGasSummaryConnection = (
+  publicClient: PublicClient,
+  artifacts: ArtifactManager,
+  cacheDir: string,
+  fromBlock: bigint
+) => {
+  registeredConnections.push({ publicClient, artifacts, fromBlock });
+  summaryDir = path.join(cacheDir, SUMMARY_DIR_NAME);
+
+  if (exitHookInstalled) return;
+  exitHookInstalled = true;
+
+  process.on('beforeExit', () => {
+    if (exitHookRan) return;
+    exitHookRan = true;
+
+    // Async work here keeps the event loop alive until it settles; the guard above
+    // prevents re-entry when beforeExit fires again afterwards.
+    void (async () => {
+      const rows: MockGasSummaryRow[] = [];
+      for (const { publicClient: client, artifacts: arts, fromBlock } of registeredConnections) {
+        try {
+          rows.push(...(await collectMockGasRows(client, arts, fromBlock)));
+        } catch {
+          // A closed or unusable connection contributes nothing
+        }
+      }
+      if (rows.length === 0 || summaryDir === undefined) return;
+
+      try {
+        fs.mkdirSync(summaryDir, { recursive: true });
+        const file = path.join(summaryDir, `rows-${process.pid}-${Date.now()}.json`);
+        fs.writeFileSync(file, JSON.stringify(rows), 'utf8');
+      } catch {
+        // Summary is best-effort; never fail the test run over it
+      }
+    })();
+  });
+};
+
+/// Removes any leftover row files (from previous or non-test runs).
+export const cleanGasSummaryDir = (cacheDir: string) => {
+  try {
+    fs.rmSync(path.join(cacheDir, SUMMARY_DIR_NAME), { recursive: true, force: true });
+  } catch {
+    // best-effort
+  }
+};
+
+/// Reads all worker row files, merges them by contract + method, prints the table, and
+/// cleans up. Called from the parent `test` task after the run completes.
+export const printMockGasSummaryFromDir = (cacheDir: string) => {
+  const dir = path.join(cacheDir, SUMMARY_DIR_NAME);
+  let files: string[];
+  try {
+    files = fs.readdirSync(dir).filter((f) => f.endsWith('.json'));
+  } catch {
+    return;
+  }
+
+  const merged = new Map<
+    string,
+    { contract: string; method: string; calls: number; gasUsed: bigint; mockGas: bigint }
+  >();
+  for (const file of files) {
+    try {
+      const rows: MockGasSummaryRow[] = JSON.parse(fs.readFileSync(path.join(dir, file), 'utf8'));
+      for (const r of rows) {
+        const key = `${r.contract}:${r.method}`;
+        const row = merged.get(key) ?? { contract: r.contract, method: r.method, calls: 0, gasUsed: 0n, mockGas: 0n };
+        row.calls += r.calls;
+        row.gasUsed += BigInt(r.totalGasUsed);
+        row.mockGas += BigInt(r.totalMockGas);
+        merged.set(key, row);
+      }
+    } catch {
+      // skip unreadable files
+    }
+  }
+  cleanGasSummaryDir(cacheDir);
+  if (merged.size === 0) return;
+
+  const rows = [...merged.values()]
+    .map((r) => {
+      const avgGasUsed = r.gasUsed / BigInt(r.calls);
+      const avgMockGas = r.mockGas / BigInt(r.calls);
+      return {
+        contract: r.contract,
+        method: r.method,
+        calls: r.calls,
+        avgGasUsed,
+        avgAdjusted: avgMockGas > avgGasUsed ? 0n : avgGasUsed - avgMockGas,
+        overheadPct: avgGasUsed === 0n ? 0 : Math.min(100, Number((avgMockGas * 100n) / avgGasUsed)),
+      };
+    })
+    .sort((a, b) => a.contract.localeCompare(b.contract) || a.method.localeCompare(b.method));
+
+  const headers = ['Contract', 'Method', 'Calls', 'Avg gas (mocks)', 'Avg gas (adjusted)', 'Mock overhead'];
+  const cells = rows.map((r) => [
+    r.contract,
+    r.method,
+    String(r.calls),
+    r.avgGasUsed.toLocaleString(),
+    r.avgAdjusted.toLocaleString(),
+    `${r.overheadPct}%`,
+  ]);
+  const widths = headers.map((h, i) => Math.max(h.length, ...cells.map((row) => row[i]!.length)));
+  const line = (row: string[]) => `│ ${row.map((cell, i) => cell.padEnd(widths[i]!)).join(' │ ')} │`;
+  const divider = (l: string, m: string, r: string) => `${l}${widths.map((w) => '─'.repeat(w + 2)).join(m)}${r}`;
+
+  console.log('');
+  console.log(chalk.bold('[COFHE-MOCKS] Gas summary') + chalk.dim(' — adjusted ≈ cost excluding mock-only overhead'));
+  console.log(divider('┌', '┬', '┐'));
+  console.log(line(headers));
+  console.log(divider('├', '┼', '┤'));
+  for (const row of cells) console.log(line(row));
+  console.log(divider('└', '┴', '┘'));
+  console.log(
+    chalk.dim(
+      'Raw numbers include on-chain replication of off-chain CoFHE work; adjusted numbers estimate real-network cost. estimateGas remains unadjusted.'
+    )
+  );
+};

package/src/hooks/hre.ts

@@ -7,6 +7,7 @@ import { privateKeyToAccount } from 'viem/accounts';
 import type { ArtifactManager } from 'hardhat/types/artifacts';
 
 import {
+  MockCoFHEAddressBookArtifact,
   MockTaskManagerArtifact,
   MockACLArtifact,
   MockZkVerifierArtifact,
@@ -19,6 +20,7 @@ import { hardhat as hardhatChain } from '@cofhe/sdk/chains';
 import { deployMocks, type DeployMocksArgs, type DeployedMockContracts } from '../deploy.js';
 import { mock_setLoggingEnabled, mock_withLogs } from '../logging.js';
 import { mock_getPlaintext, mock_expectPlaintext, getMockContractsNpmPaths } from '../utils.js';
+import { mock_getAdjustedGasBreakdown, mock_getAdjustedGasUsed, registerGasSummaryConnection } from '../gas.js';
 import type { CofheConnection } from '../type-extensions.js';
 
 // ─── Per-connection cofhe object factory ─────────────────────────────────────
@@ -55,6 +57,9 @@ function createCofheConnection(
       return createCofheClient(config);
     },
 
+    getAdjustedGasUsed: mock_getAdjustedGasUsed,
+    getAdjustedGasBreakdown: mock_getAdjustedGasBreakdown,
+
     async createClientWithBatteries(signerWalletClient?: WalletClient) {
       let signerClient: WalletClient;
       let address: `0x${string}`;
@@ -105,6 +110,10 @@ function createCofheConnection(
         await mock_expectPlaintext(publicClient, ctHash, expectedValue);
       },
 
+      MockCoFHEAddressBook: {
+        address: deployedMockContracts.MockCoFHEAddressBook,
+        abi: MockCoFHEAddressBookArtifact.abi,
+      },
       MockTaskManager: { address: deployedMockContracts.MockTaskManager, abi: MockTaskManagerArtifact.abi },
       MockACL: { address: deployedMockContracts.MockACL, abi: MockACLArtifact.abi },
       MockZkVerifier: { address: deployedMockContracts.MockZkVerifier, abi: MockZkVerifierArtifact.abi },
@@ -176,6 +185,20 @@ const hreHooks: Partial<HardhatRuntimeEnvironmentHooks> = {
           }
         );
 
+        // With gasSummary enabled, track this connection so its adjusted-gas rows are
+        // dumped at process exit for the `test` task to merge and print. Connections
+        // (and their chains) live in node:test worker processes, so collection has to
+        // happen here rather than in the test task itself.
+        // Only connections where mocks were actually deployed are tracked (deployMocks
+        // skips non-hardhat networks and returns an empty result) - registering external
+        // RPCs would trigger pointless log scans at exit and double-count shared chains.
+        // The current block is recorded so the scan never reaches back before the deploy
+        // (relevant for forks and long-lived chains).
+        if (hre.config.cofhe.gasSummary && deployedMockContracts.MockTaskManager !== undefined) {
+          const fromBlock = await publicClient.getBlockNumber();
+          registerGasSummaryConnection(publicClient, hre.artifacts, hre.config.paths.cache, fromBlock);
+        }
+
         (conn as any).cofhe = createCofheConnection(
           publicClient,
           walletClient,

package/src/index.ts

@@ -1,5 +1,5 @@
 import type { HardhatPlugin } from 'hardhat/types/plugins';
-import { task } from 'hardhat/config';
+import { overrideTask, task } from 'hardhat/config';
 import { ArgumentType } from 'hardhat/types/arguments';
 
 import { TASK_COFHE_USE_FAUCET, TASK_COFHE_SET_LOG_OPS } from './consts.js';
@@ -29,6 +29,11 @@ const plugin: HardhatPlugin = {
       .addFlag({ name: 'enable', description: 'Whether to enable logging' })
       .setAction(async () => import('./tasks/set-log-ops.js'))
       .build(),
+
+    // Prints the adjusted-gas summary after `hardhat test` when cofhe.gasSummary is enabled.
+    overrideTask('test')
+      .setAction(async () => import('./tasks/test-gas-summary.js'))
+      .build(),
   ],
 };
 
@@ -37,5 +42,6 @@ export default plugin;
 export * from './consts.js';
 export * from './deploy.js';
 export * from './fund.js';
+export * from './gas.js';
 export * from './logging.js';
 export * from './utils.js';

package/src/logging.ts

@@ -2,11 +2,11 @@ import type { PublicClient, WalletClient } from 'viem';
 import chalk from 'chalk';
 
 import { MockTaskManagerArtifact } from '@cofhe/mock-contracts';
-import { TASK_MANAGER_ADDRESS } from '@cofhe/sdk';
+import { MOCKS_TASK_MANAGER_ADDRESS } from '@cofhe/sdk';
 
 async function getLoggingEnabled(publicClient: PublicClient): Promise<boolean> {
   return publicClient.readContract({
-    address: TASK_MANAGER_ADDRESS,
+    address: MOCKS_TASK_MANAGER_ADDRESS,
     abi: MockTaskManagerArtifact.abi,
     functionName: 'logOps',
     args: [],
@@ -20,7 +20,7 @@ async function setLoggingEnabled(
 ): Promise<void> {
   const [account] = await walletClient.getAddresses();
   const hash = await walletClient.writeContract({
-    address: TASK_MANAGER_ADDRESS,
+    address: MOCKS_TASK_MANAGER_ADDRESS,
     abi: MockTaskManagerArtifact.abi,
     functionName: 'setLogOps',
     args: [enabled],

package/src/tasks/test-gas-summary.ts

@@ -0,0 +1,22 @@
+import type { TaskOverrideActionFunction } from 'hardhat/types/tasks';
+
+import { cleanGasSummaryDir, printMockGasSummaryFromDir } from '../gas.js';
+
+/// Wraps the `test` task to print the adjusted-gas summary when `cofhe.gasSummary` is
+/// enabled. Tests run in node:test worker processes, each with its own chains; workers
+/// dump their per-method gas rows into the hardhat cache dir at exit (see gas.ts), and
+/// this override merges and prints them once the run completes.
+const action: TaskOverrideActionFunction = async (args, hre, runSuper) => {
+  if (!hre.config.cofhe.gasSummary) {
+    return runSuper(args);
+  }
+
+  // Drop leftovers from previous or aborted runs so the table only reflects this run.
+  cleanGasSummaryDir(hre.config.paths.cache);
+
+  const result = await runSuper(args);
+  printMockGasSummaryFromDir(hre.config.paths.cache);
+  return result;
+};
+
+export default action;

package/src/type-extensions.ts

@@ -1,7 +1,9 @@
 import type { WalletClient } from 'viem';
 import type { CofheClient, CofheConfig, CofheInputConfig } from '@cofhe/sdk';
 import type { DeployMocksArgs, LogMocksDeploy } from './deploy.js';
+import type { AdjustableGasReceipt, AdjustedGasBreakdown } from './gas.js';
 import type {
+  MockCoFHEAddressBookArtifact,
   MockTaskManagerArtifact,
   MockACLArtifact,
   MockZkVerifierArtifact,
@@ -29,6 +31,22 @@ export interface CofheConnection {
    */
   createClientWithBatteries(walletClient?: WalletClient): Promise<CofheClient>;
 
+  /**
+   * Returns a transaction's gas usage excluding mock-only overhead (FHE op replication,
+   * decrypt-task storage, mock logging) - an estimate of what the transaction would cost
+   * on a real CoFHE network. Pure function of the receipt (sums the mock task manager's
+   * MockGasConsumed events from `receipt.logs`); on a real network the receipt carries no
+   * such events and the raw `gasUsed` is returned unchanged.
+   */
+  getAdjustedGasUsed(receipt: AdjustableGasReceipt): bigint;
+
+  /**
+   * Full gas breakdown of a transaction receipt: raw `gasUsed`, the mock-only `mockGas`,
+   * the `adjustedGasUsed` (raw minus mock), and the number of mock gas events.
+   * See getAdjustedGasUsed.
+   */
+  getAdjustedGasBreakdown(receipt: AdjustableGasReceipt): AdjustedGasBreakdown;
+
   mocks: {
     /** Deploy (or re-deploy) the mock contracts. */
     deployMocks(options?: DeployMocksArgs): Promise<void>;
@@ -48,6 +66,9 @@ export interface CofheConnection {
     /** Assert that a ciphertext hash represents the expected plaintext value. */
     expectPlaintext(ctHash: bigint | string, expectedValue: bigint): Promise<void>;
 
+    /** MockCoFHEAddressBook contract descriptor — spread into readContract / writeContract */
+    MockCoFHEAddressBook: { address: `0x${string}`; abi: typeof MockCoFHEAddressBookArtifact.abi };
+
     /** MockTaskManager contract descriptor — spread into readContract / writeContract */
     MockTaskManager: { address: `0x${string}`; abi: typeof MockTaskManagerArtifact.abi };
 
@@ -71,6 +92,14 @@ declare module 'hardhat/types/config' {
       logMocks?: boolean;
       /** Whether to show gas usage warnings for mock operations (default: true) */
       gasWarning?: boolean;
+      /**
+       * Print a per-method gas summary after `hardhat test`, showing raw gas next to
+       * adjusted gas (mock-only overhead excluded - an estimate of real-network cost).
+       * Collected per test worker from chain history: transactions rolled back by
+       * snapshots won't appear, and the table is best-effort after forced exits -
+       * use getAdjustedGasUsed(receipt) for exact numbers. (default: false)
+       */
+      gasSummary?: boolean;
       /**
        * Controls deploy-mocks console output.
        * - `''`   — silent, no output
@@ -85,6 +114,8 @@ declare module 'hardhat/types/config' {
     cofhe: {
       logMocks: boolean;
       gasWarning: boolean;
+      /** Print a per-method adjusted-gas summary after `hardhat test` (default: false) */
+      gasSummary: boolean;
       mocksDeployVerbosity: LogMocksDeploy;
     };
   }

package/src/utils.ts

@@ -1,5 +1,5 @@
 import type { PublicClient } from 'viem';
-import { TASK_MANAGER_ADDRESS, MOCKS_ZK_VERIFIER_ADDRESS } from '@cofhe/sdk';
+import { MOCKS_TASK_MANAGER_ADDRESS, MOCKS_ZK_VERIFIER_ADDRESS } from '@cofhe/sdk';
 import path from 'node:path';
 import fs from 'node:fs';
 import { createRequire } from 'node:module';
@@ -57,7 +57,7 @@ export async function mock_getPlaintext(
   }
 
   return publicClient.readContract({
-    address: TASK_MANAGER_ADDRESS,
+    address: MOCKS_TASK_MANAGER_ADDRESS,
     abi: [
       {
         name: 'mockStorage',
@@ -86,7 +86,7 @@ export async function mock_expectPlaintext(
   }
 
   const exists = (await publicClient.readContract({
-    address: TASK_MANAGER_ADDRESS,
+    address: MOCKS_TASK_MANAGER_ADDRESS,
     abi: [
       {
         name: 'inMockStorage',