feat: add shareable State management lib

packages/snap-networks-utils/CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Add shared snap state helpers `IStateManager`, `State`, and `InMemoryState` (Tron-style write mutex plus blob/path locking). ([#288](https://github.com/MetaMask/internal-snaps/pull/288))
 - Add shared proof-of-ownership message parsing utilities, batch request/response structs, and batch request/response types. ([#268](https://github.com/MetaMask/internal-snaps/pull/268))
 - Add a `UuidStruct` Superstruct for validating UUID v4 strings. ([#243](https://github.com/MetaMask/internal-snaps/pull/243))
 - Add helpers `serialize`, `deserialize`, and `Serializable` for round-tripping `BigNumber`, `bigint`, `Uint8Array`, and `undefined` through snap state ([#197](https://github.com/MetaMask/internal-snaps/pull/197))

packages/snap-networks-utils/package.json

@@ -76,6 +76,7 @@
     "@metamask/snaps-sdk": "^12.0.1",
     "@metamask/superstruct": "^3.4.1",
     "@metamask/utils": "^11.11.0",
+    "async-mutex": "^0.5.0",
     "bignumber.js": "^9.3.1",
     "lodash": "^4.17.21"
   },

packages/snap-networks-utils/src/index.ts

@@ -62,6 +62,10 @@ export {
 } from './utils/originPermissions/createOriginPermissions';
 export type { CreateOriginPermissionsParams } from './utils/originPermissions/createOriginPermissions';
 export { validateOrigin } from './utils/originPermissions/validateOrigin';
+export { State } from './utils/state/State';
+export type { StateConfig } from './utils/state/State';
+export { InMemoryState } from './utils/state/InMemoryState';
+export type { IStateManager } from './utils/state/IStateManager';
 export {
   createSnapErrorHandling,
   createTrackError,

packages/snap-networks-utils/src/utils/state/IStateManager.ts

@@ -0,0 +1,76 @@
+import type { Serializable } from '../serialization/types';
+
+export type IStateManager<TStateValue extends Record<string, Serializable>> = {
+  /**
+   * Gets the whole state object.
+   *
+   * ⚠️ WARNING: Use with caution because it transfers the whole state, which might contain a lot of data.
+   * If you need to retrieve only a specific part of the state, use IStateManager.getKey instead.
+   */
+  get(): Promise<TStateValue>;
+  /**
+   * Gets the value of the passed key in the state object.
+   * The key is the JSON path to the value to get.
+   *
+   * @returns The value of the key, or undefined if the key does not exist.
+   */
+  getKey<TResponse extends Serializable>(
+    key: string,
+  ): Promise<TResponse | undefined>;
+  /**
+   * Sets the value of the passed key in the state object.
+   * The key is a JSON path to the value to set.
+   *
+   * @param key - The key to set, which is a JSON path to the location.
+   * @param value - The value to set.
+   */
+  setKey(key: string, value: Serializable): Promise<void>;
+  /**
+   * Atomically reads the current value at `key`, applies `updater`, and writes the result back.
+   *
+   * Implementations must ensure that no concurrent state write can interleave between the
+   * read and the write. This makes the method safe for updates where the next value depends
+   * on the current value, such as merging objects.
+   *
+   * Prefer this over a manual `getKey` + `setKey` sequence whenever the new value depends on
+   * the current one.
+   *
+   * @param key - The JSON-path key to update.
+   * @param updater - Receives the current value (or `undefined` when the key is absent) and
+   * returns the new value to store.
+   */
+  setKeyWith<TValue extends Serializable>(
+    key: string,
+    updater: (currentValue: TValue | undefined) => TValue,
+  ): Promise<void>;
+  /**
+   * Updates the whole state object.
+   *
+   * Typically used for bulk `set`s or `delete`s, because:
+   * - Atomicity: Using a single `state.update` ensures that all changes are applied atomically.
+   * - Performance: One round trip instead of many `setKey` / `deleteKey` calls.
+   * - State Consistency: Read once, modify in memory, write the complete updated state back.
+   *
+   * ⚠️ WARNING: Use with caution because:
+   * - it will override the whole state.
+   * - it transfers the whole state back and forth to the data store.
+   *
+   * For single updates, use `setKey` or `deleteKey` instead.
+   *
+   * @param updaterFunction - The function that updates the state.
+   * @returns The updated state.
+   */
+  update(
+    updaterFunction: (state: TStateValue) => TStateValue,
+  ): Promise<TStateValue>;
+  /**
+   * Deletes the value of the passed key in the state object.
+   * The key is a JSON path to the value to delete.
+   */
+  deleteKey(key: string): Promise<void>;
+  /**
+   * Deletes multiple keys in the state object in a single operation.
+   * The keys are JSON paths to the values to delete.
+   */
+  deleteKeys(keys: string[]): Promise<void>;
+};

packages/snap-networks-utils/src/utils/state/InMemoryState.test.ts

@@ -0,0 +1,40 @@
+import { InMemoryState } from './InMemoryState';
+
+type MockStateValue = {
+  users: { name: string; age?: number }[];
+};
+
+describe('InMemoryState', () => {
+  it('gets, sets, and deletes keys', async () => {
+    const state = new InMemoryState<MockStateValue>({
+      users: [{ name: 'John', age: 30 }],
+    });
+
+    expect(await state.get()).toStrictEqual({
+      users: [{ name: 'John', age: 30 }],
+    });
+    expect(await state.getKey('users.0.name')).toBe('John');
+
+    await state.setKey('users.0.name', 'Jane');
+    expect(await state.getKey('users.0.name')).toBe('Jane');
+
+    await state.setKeyWith<number>('users.0.age', (age) => (age ?? 0) + 1);
+    expect(await state.getKey('users.0.age')).toBe(31);
+
+    await state.deleteKey('users.0.age');
+    expect(await state.get()).toStrictEqual({ users: [{ name: 'Jane' }] });
+
+    await state.deleteKeys(['users']);
+    expect(await state.get()).toStrictEqual({});
+  });
+
+  it('replaces the whole state via update', async () => {
+    const state = new InMemoryState<MockStateValue>({
+      users: [{ name: 'John' }],
+    });
+
+    expect(
+      await state.update(() => ({ users: [{ name: 'Bob', age: 50 }] })),
+    ).toStrictEqual({ users: [{ name: 'Bob', age: 50 }] });
+  });
+});

packages/snap-networks-utils/src/utils/state/InMemoryState.ts

@@ -0,0 +1,57 @@
+import { get, set, unset } from 'lodash';
+
+import type { Serializable } from '../serialization/types';
+import type { IStateManager } from './IStateManager';
+
+/**
+ * A simple implementation of the `IStateManager` interface that relies on an in-memory
+ * state. Intended for tests.
+ */
+export class InMemoryState<
+  TStateValue extends Record<string, Serializable>,
+> implements IStateManager<TStateValue> {
+  #state: TStateValue;
+
+  constructor(initialState: TStateValue) {
+    this.#state = initialState;
+  }
+
+  async get(): Promise<TStateValue> {
+    return this.#state;
+  }
+
+  async getKey<TResponse extends Serializable>(
+    key: string,
+  ): Promise<TResponse | undefined> {
+    return get(this.#state, key) as TResponse | undefined;
+  }
+
+  async setKey(key: string, value: Serializable): Promise<void> {
+    set(this.#state, key, value);
+  }
+
+  async setKeyWith<TValue extends Serializable>(
+    key: string,
+    updater: (currentValue: TValue | undefined) => TValue,
+  ): Promise<void> {
+    const oldValue = get(this.#state, key) as TValue | undefined;
+    set(this.#state, key, updater(oldValue));
+  }
+
+  async update(
+    callback: (state: TStateValue) => TStateValue,
+  ): Promise<TStateValue> {
+    this.#state = callback(this.#state);
+    return this.#state;
+  }
+
+  async deleteKey(key: string): Promise<void> {
+    unset(this.#state, key);
+  }
+
+  async deleteKeys(keys: string[]): Promise<void> {
+    keys.forEach((key) => {
+      unset(this.#state, key);
+    });
+  }
+}