Add PS mutation (#131)

* add mutatable private state in live backend

* add tests

* fix doc

* add live mutation tests

* serialize rmw sequence

* use buffer in tests

* fix tsdoc links

* resolve updatePrivateState to the written state

* fix docs

* document per-instance scope of queue

* extract PrivateStateMutator

* add PrivateStateMutator tests

Author: andrew-fleming

packages/simulator/src/backend/Backend.ts

@@ -71,14 +71,18 @@ export interface Backend<P, L> {
   getContractState(): Promise<StateValue>;
 
   /**
-   * Replaces the private state. Dry mutates the in-memory context (used by
-   * per-module helpers like secret/nonce injection); live throws, because
-   * mid-test private-state mutation is the documented dry↔live asymmetry.
-   * Guard such specs with `isLiveBackend()`.
+   * Replaces the whole private state `P`. Dry mutates the in-memory context;
+   * live writes to the harness's private-state provider so the next impure
+   * `callTx` proves against it but only if the injected `LiveContext`
+   * implements `setPrivateState`; otherwise it throws
+   * `PRIVATE_STATE_MUTATION_UNSUPPORTED` (from the live backend).
+   *
+   * Async across backends for uniform `await`: dry resolves synchronously,
+   * live awaits the provider write.
    *
    * @param privateState - The new private state `P`.
    */
-  setPrivateState(privateState: P): void;
+  setPrivateState(privateState: P): Promise<void>;
 
   /**
    * Sets the caller identity for subsequent circuit calls.

packages/simulator/src/backend/DryBackend.ts

@@ -93,7 +93,7 @@ export class DryBackend<P, L> implements Backend<P, L> {
   }
 
   /** Mutates the in-memory private state (dry supports mid-test mutation). */
-  setPrivateState(privateState: P): void {
+  async setPrivateState(privateState: P): Promise<void> {
     this.sim.circuitContextManager.updatePrivateState(privateState);
   }
 

packages/simulator/src/core/PrivateStateMutator.ts

@@ -0,0 +1,81 @@
+/**
+ * Serializes read-modify-write access to a contract's private state.
+ *
+ * Built from a backend's `getPrivateState`/`setPrivateState` pair, it queues
+ * every mutation so a read-modify-write (`update`) can't interleave with another
+ * mutation and lose an update (a lost-update race). It is backend-agnostic: the
+ * dry, live, and future `sim` backends all reuse it by supplying their own
+ * read/write closures.
+ *
+ * The guarantee is scoped to this instance. On live, two `PrivateStateMutator`s
+ * (or two processes) targeting the same `privateStateProvider` + `privateStateId`
+ * still race — both read the provider, last write wins — since the queue is
+ * local. Cross-instance atomicity would need provider-side
+ * compare-and-set/versioning and is out of scope; tests drive one simulator per
+ * private state, so the per-instance queue is sufficient in practice.
+ *
+ * @template P - Private state type.
+ */
+export class PrivateStateMutator<P> {
+  /** Tail of the mutation queue; each op runs after the previous drains. */
+  #chain: Promise<unknown> = Promise.resolve();
+  readonly #read: () => Promise<P>;
+  readonly #write: (next: P) => Promise<void>;
+
+  /**
+   * @param read - Reads the current private state (backend `getPrivateState`).
+   * @param write - Replaces the whole private state (backend `setPrivateState`).
+   */
+  constructor(read: () => Promise<P>, write: (next: P) => Promise<void>) {
+    this.#read = read;
+    this.#write = write;
+  }
+
+  /**
+   * Runs `op` once the queue drains, serialized against every other mutation.
+   * The returned promise settles with `op`'s result (or rejection) for the
+   * caller; the queue tail deliberately swallows rejections so a failed op does
+   * not poison subsequent mutations.
+   *
+   * @param op - The operation to run once the queue drains.
+   * @returns `op`'s result.
+   */
+  enqueue<T>(op: () => Promise<T>): Promise<T> {
+    const run = this.#chain.then(op);
+    this.#chain = run.then(
+      () => undefined,
+      () => undefined,
+    );
+    return run;
+  }
+
+  /**
+   * Replaces the whole private state, serialized against other mutations.
+   *
+   * @param next - The new private state.
+   */
+  set(next: P): Promise<void> {
+    return this.enqueue(() => this.#write(next));
+  }
+
+  /**
+   * Read-modify-write, serialized end-to-end so the read and write are atomic
+   * against other mutations on this instance. Resolves to the state that was
+   * written, so callers need no follow-up read.
+   *
+   * @param updater - A partial patch to shallow-merge, or a function that
+   *   receives the current state and returns the next.
+   * @returns The private state that was written.
+   */
+  update(updater: Partial<P> | ((prev: P) => P)): Promise<P> {
+    return this.enqueue(async () => {
+      const prev = await this.#read();
+      const next =
+        typeof updater === 'function'
+          ? (updater as (p: P) => P)(prev)
+          : { ...prev, ...updater };
+      await this.#write(next);
+      return next;
+    });
+  }
+}

packages/simulator/src/factory/createSimulator.ts

@@ -1,5 +1,6 @@
 import type { Backend, BackendKind, CircuitKind } from '../backend/Backend.js';
 import { DryBackend, type SyncSimulator } from '../backend/DryBackend.js';
+import { PrivateStateMutator } from '../core/PrivateStateMutator.js';
 import type { LiveContext } from '../live/LiveContext.js';
 import { getRegisteredLiveBackend } from '../live/registry.js';
 import { Signers } from '../signers/Signers.js';
@@ -147,6 +148,12 @@ export function createSimulator<
     readonly _backend: Backend<P, L>;
     readonly _signers: Signers;
 
+    /**
+     * Serializes private-state read-modify-write for this instance. Internal;
+     * see {@link PrivateStateMutator} for the scope of the guarantee.
+     */
+    readonly _mutator: PrivateStateMutator<P>;
+
     /** Async circuit proxies; every call returns a promise. */
     readonly circuits: {
       pure: AsyncCircuits<ExtractPureCircuits<TContract>, P>;
@@ -163,6 +170,12 @@ export function createSimulator<
       this._backend = deps.backend;
       this.backendKind = deps.backend.kind;
       this._signers = deps.signers;
+      // Constructed here (not as a field initializer) so the closures capture
+      // the assigned `_backend` rather than an undefined one.
+      this._mutator = new PrivateStateMutator<P>(
+        () => this._backend.getPrivateState(),
+        (next) => this._backend.setPrivateState(next),
+      );
       this.circuits = {
         pure: buildProxy(
           this._backend,
@@ -241,14 +254,37 @@ export function createSimulator<
     }
 
     /**
-     * Replaces the private state (for per-module secret/nonce injection helpers).
-     * Dry mutates the in-memory context; live throws (mutation asymmetry) —
-     * guard such specs with `isLiveBackend()`.
+     * Replaces the whole private state. Dry mutates the in-memory context; live
+     * writes to the harness's private-state provider so the next impure call
+     * proves against it (throws if the `LiveContext` opted out of mutation).
+     * Serialized against other mutations via {@link _mutator}.
      *
      * @param privateState - The new private state.
      */
-    setPrivateState(privateState: P): void {
-      this._backend.setPrivateState(privateState);
+    setPrivateState(privateState: P): Promise<void> {
+      return this._mutator.set(privateState);
+    }
+
+    /**
+     * Ergonomic granular private-state mutation. Replaces the per-module
+     * `injectSecretKey`/`injectSecretNonce` helpers: a plain object shallow-
+     * merges onto the current state, a function receives the current state and
+     * returns the next.
+     *
+     * The read-modify-write is serialized (see {@link _mutator}) and resolves to
+     * the state that was written, so callers can `return sim.updatePrivateState(...)`
+     * without a follow-up `getPrivateState()`. Works on both dry (in-memory) and
+     * live (provider read then write); on live the current state must already
+     * exist (it is seeded at deploy).
+     *
+     * @example sim.updatePrivateState({ secretKey });
+     * @example sim.updatePrivateState((prev) => ({ ...prev, counter: prev.counter + 1n }));
+     *
+     * @param updater - A partial patch to merge, or an updater function.
+     * @returns The private state that was written.
+     */
+    updatePrivateState(updater: Partial<P> | ((prev: P) => P)): Promise<P> {
+      return this._mutator.update(updater);
     }
 
     /** The raw contract state value. */

packages/simulator/src/index.ts

@@ -7,6 +7,7 @@ export { DryBackend } from './backend/DryBackend.js';
 export { AbstractSimulator } from './core/AbstractSimulator.js';
 export { CircuitContextManager } from './core/CircuitContextManager.js';
 export { ContractSimulator } from './core/ContractSimulator.js';
+export { PrivateStateMutator } from './core/PrivateStateMutator.js';
 // --- Core simulator (one factory, two backends) ----------------------------
 // A dry import pulls zero midnight-js: the live adapter `LiveBackend`
 // is type-only here and reached at runtime only via the dynamic import inside
@@ -25,7 +26,10 @@ export {
   DEFAULT_INDEXER_LAG,
 } from './live/createLiveContext.js';
 export type { LiveBackend } from './live/LiveBackend.js';
-export { WITNESS_OVERRIDE_UNSUPPORTED } from './live/LiveBackend.js';
+export {
+  PRIVATE_STATE_MUTATION_UNSUPPORTED,
+  WITNESS_OVERRIDE_UNSUPPORTED,
+} from './live/LiveBackend.js';
 export type {
   DeployedTxHandle,
   FinalizedCallResult,

packages/simulator/src/live/LiveBackend.ts

@@ -8,9 +8,9 @@ import type { LiveContext } from './LiveContext.js';
 export const WITNESS_OVERRIDE_UNSUPPORTED =
   'witness override unsupported on live backend';
 
-/** The error thrown when private state is mutated on the live backend. */
+/** The error thrown when the live `LiveContext` cannot mutate private state. */
 export const PRIVATE_STATE_MUTATION_UNSUPPORTED =
-  'private-state mutation unsupported on live backend';
+  'private-state mutation unsupported: this LiveContext does not implement setPrivateState';
 
 /**
  * Dependencies the live adapter is constructed with by `createBackendSimulator`.
@@ -119,12 +119,16 @@ export class LiveBackend<P, L> implements Backend<P, L> {
   }
 
   /**
-   * Mid-test private-state mutation does not faithfully reproduce on live.
-   * Throws so such specs are explicitly guarded with `isLiveBackend()`
-   * rather than silently passing against unchanged state.
+   * Writes private state through the injected {@link LiveContext} so the next
+   * impure `callTx` proves against it. If the context opted out of mutation
+   * (no `setPrivateState`), throws {@link PRIVATE_STATE_MUTATION_UNSUPPORTED}
+   * so the spec fails loudly rather than proving against stale state.
    */
-  setPrivateState(_privateState: P): void {
-    throw new Error(PRIVATE_STATE_MUTATION_UNSUPPORTED);
+  async setPrivateState(privateState: P): Promise<void> {
+    if (!this.ctx.setPrivateState) {
+      throw new Error(PRIVATE_STATE_MUTATION_UNSUPPORTED);
+    }
+    await this.ctx.setPrivateState(privateState);
   }
 
   /**

packages/simulator/src/live/LiveContext.ts

@@ -66,4 +66,25 @@ export interface LiveContext<P> {
    * parity).
    */
   queryPrivateState(): Promise<P>;
+
+  /**
+   * Optional: writes the whole private state to the harness's private-state
+   * provider, so the NEXT impure `callTx` proves against it. Each `callTx`
+   * reads private state fresh from the provider (midnight-js `getStates` →
+   * `privateStateProvider.get`), so no handle invalidation is needed.
+   *
+   * Omit to opt out of mutation — {@link LiveBackend} then throws
+   * `PRIVATE_STATE_MUTATION_UNSUPPORTED`, so a spec that mutates fails loudly
+   * rather than silently proving against stale state.
+   *
+   * Faithful for any client-controlled field of `P` (secret keys, cached
+   * plaintexts, seeds, nonces) — injecting a hostile/stale value and asserting
+   * the resulting rejection or handled behavior mirrors a real client. The one
+   * thing it cannot do is fabricate on-chain state (`L`): a private state that
+   * presupposes an on-chain event that never happened will not make a happy
+   * path succeed on a live node.
+   *
+   * @param state - The new private state `P`.
+   */
+  setPrivateState?(state: P): Promise<void>;
 }

packages/simulator/src/live/createLiveContext.ts

@@ -54,7 +54,16 @@ export interface CreateLiveContextOptions<P> {
   privateStateId: string;
   /** Provider for reading on-chain public state. */
   publicDataProvider: PublicDataProvider;
-  /** Provider for reading private state (read parity). */
+  /**
+   * Provider for reading and (optionally) writing private state.
+   *
+   * Invariant: this MUST be the same provider instance wired into the
+   * `providersFor(alias)` bundle handed to `findDeployedContract`. The
+   * `setPrivateState` write below targets this provider, and the next
+   * `callTx` reads private state from the bundle's provider; if they are
+   * different instances the write is invisible to proving. `setContractAddress`
+   * is the harness's responsibility (already required for `get`).
+   */
   privateStateProvider: PrivateStateProvider<string, P>;
   /** Optional override of the indexer-lag policy. */
   indexerLag?: Partial<IndexerLagPolicy>;
@@ -179,5 +188,14 @@ export function createLiveContext<P>(
       }
       return state;
     },
+
+    /**
+     * Writes the whole private state to the provider under `privateStateId`.
+     * The next impure `callTx` reads it fresh (no handle-cache invalidation
+     * needed). See the invariant on {@link CreateLiveContextOptions.privateStateProvider}.
+     */
+    async setPrivateState(state: P): Promise<void> {
+      await options.privateStateProvider.set(options.privateStateId, state);
+    },
   };
 }