@platforma-sdk/model 1.51.6 → 1.52.0

package/src/block_model.ts

@@ -0,0 +1,440 @@
+import type {
+  BlockRenderingMode,
+  BlockSection,
+  OutputWithStatus,
+  PlRef,
+  BlockCodeKnownFeatureFlags,
+  BlockConfigContainer,
+} from '@milaboratories/pl-model-common';
+import { getPlatformaInstance, isInUI, createAndRegisterRenderLambda, createRenderLambda } from './internal';
+import type { DataModel } from './block_migrations';
+import type { PlatformaV3 } from './platforma';
+import type { InferRenderFunctionReturn, RenderFunction } from './render';
+import { RenderCtx } from './render';
+import { PlatformaSDKVersion } from './version';
+// Import storage VM integration (side-effect: registers internal callbacks)
+import './block_storage_vm';
+import type {
+  ConfigRenderLambda,
+  DeriveHref,
+  ConfigRenderLambdaFlags,
+  InferOutputsFromLambdas,
+} from './bconfig';
+import {
+  downgradeCfgOrLambda,
+  isConfigLambda,
+} from './bconfig';
+import type { PlatformaExtended } from './platforma';
+
+type SectionsExpectedType = readonly BlockSection[];
+
+type NoOb = Record<string, never>;
+
+interface BlockModelV3Config<
+  Args,
+  OutputsCfg extends Record<string, ConfigRenderLambda>,
+  Data,
+> {
+  renderingMode: BlockRenderingMode;
+  dataModel: DataModel<Data>;
+  outputs: OutputsCfg;
+  sections: ConfigRenderLambda;
+  title: ConfigRenderLambda | undefined;
+  subtitle: ConfigRenderLambda | undefined;
+  tags: ConfigRenderLambda | undefined;
+  enrichmentTargets: ConfigRenderLambda | undefined;
+  featureFlags: BlockCodeKnownFeatureFlags;
+  args: ConfigRenderLambda<Args> | undefined;
+  prerunArgs: ConfigRenderLambda<unknown> | undefined;
+}
+
+/** Options for creating a BlockModelV3 */
+export interface BlockModelV3Options<Data extends Record<string, unknown>> {
+  /** The data model that defines initial data and migrations */
+  dataModel: DataModel<Data>;
+  /** Rendering mode for the block (default: 'Heavy') */
+  renderingMode?: BlockRenderingMode;
+}
+
+/** Main entry point that each block should use in it's "config" module. Don't forget
+ * to call {@link done()} at the end of configuration. Value returned by this builder must be
+ * exported as constant with name "platforma" from the "config" module.
+ * API version is 3 (for UI) and 2 (for model) */
+export class BlockModelV3<
+  Args,
+  OutputsCfg extends Record<string, ConfigRenderLambda>,
+  Data extends Record<string, unknown> = Record<string, unknown>,
+  Href extends `/${string}` = '/',
+> {
+  private constructor(
+    private readonly config: BlockModelV3Config<Args, OutputsCfg, Data>,
+  ) {}
+
+  public static readonly INITIAL_BLOCK_FEATURE_FLAGS: BlockCodeKnownFeatureFlags = {
+    supportsLazyState: true,
+    requiresUIAPIVersion: 3,
+    requiresModelAPIVersion: 2,
+  };
+
+  /**
+   * Creates a new BlockModelV3 builder with the specified data model and options.
+   *
+   * @example
+   * const dataModel = DataModel.create<BlockData>(() => ({ numbers: [], labels: [] }));
+   * BlockModelV3.create({ dataModel })
+   *   .args((data) => ({ numbers: data.numbers }))
+   *   .sections(() => [{ type: 'link', href: '/', label: 'Main' }])
+   *   .done();
+   *
+   * @param options Configuration options including required data model
+   */
+  public static create<Data extends Record<string, unknown>>(
+    options: BlockModelV3Options<Data>,
+  ): BlockModelV3<NoOb, {}, Data> {
+    const { dataModel, renderingMode = 'Heavy' } = options;
+
+    // Register data model callbacks for VM use (initialData and upgrade)
+    dataModel.registerCallbacks();
+
+    return new BlockModelV3<NoOb, {}, Data>({
+      renderingMode,
+      dataModel,
+      outputs: {},
+      // Register default sections callback (returns empty array)
+      sections: createAndRegisterRenderLambda({ handle: 'sections', lambda: () => [] }, true),
+      title: undefined,
+      subtitle: undefined,
+      tags: undefined,
+      enrichmentTargets: undefined,
+      featureFlags: { ...BlockModelV3.INITIAL_BLOCK_FEATURE_FLAGS },
+      args: undefined,
+      prerunArgs: undefined,
+    });
+  }
+
+  /**
+   * Add output cell wrapped with additional status information to the configuration
+   *
+   * @param key output cell name, that can be later used to retrieve the rendered value
+   * @param rf  callback calculating output value using context, that allows to access
+   *            workflows outputs and interact with platforma drivers
+   * @param flags additional flags that may alter lambda rendering procedure
+   * */
+  public output<const Key extends string, const RF extends RenderFunction<Args, Data>>(
+    key: Key,
+    rf: RF,
+    flags: ConfigRenderLambdaFlags & { withStatus: true }
+  ): BlockModelV3<
+    Args,
+      OutputsCfg & { [K in Key]: ConfigRenderLambda<InferRenderFunctionReturn<RF>> & { withStatus: true } },
+      Data,
+      Href
+  >;
+  /**
+   * Add output cell to the configuration
+   *
+   * @param key output cell name, that can be later used to retrieve the rendered value
+   * @param rf  callback calculating output value using context, that allows to access
+   *            workflows outputs and interact with platforma drivers
+   * @param flags additional flags that may alter lambda rendering procedure
+   * */
+  public output<const Key extends string, const RF extends RenderFunction<Args, Data>>(
+    key: Key,
+    rf: RF,
+    flags?: ConfigRenderLambdaFlags
+  ): BlockModelV3<
+    Args,
+    OutputsCfg & { [K in Key]: ConfigRenderLambda<InferRenderFunctionReturn<RF>> },
+    Data,
+    Href
+  >;
+  public output(
+    key: string,
+    cfgOrRf: RenderFunction<Args, Data>,
+    flags: ConfigRenderLambdaFlags = {},
+  ): BlockModelV3<Args, OutputsCfg, Data, Href> {
+    return new BlockModelV3({
+      ...this.config,
+      outputs: {
+        ...this.config.outputs,
+        [key]: createAndRegisterRenderLambda({ handle: `output#${key}`, lambda: () => cfgOrRf(new RenderCtx()), ...flags }),
+      },
+    });
+  }
+
+  /** Shortcut for {@link output} with retentive flag set to true. */
+  public retentiveOutput<const Key extends string, const RF extends RenderFunction<Args, Data>>(
+    key: Key,
+    rf: RF,
+  ): BlockModelV3<
+      Args,
+    OutputsCfg & { [K in Key]: ConfigRenderLambda<InferRenderFunctionReturn<RF>> },
+    Data,
+    Href
+    > {
+    return this.output(key, rf, { retentive: true });
+  }
+
+  /** Shortcut for {@link output} with withStatus flag set to true. */
+  public outputWithStatus<const Key extends string, const RF extends RenderFunction<Args, Data>>(
+    key: Key,
+    rf: RF,
+  ) {
+    return this.output(key, rf, { withStatus: true });
+  }
+
+  /**
+   * Sets a function to derive block args from data.
+   * This is called during setData to compute the args that will be used for block execution.
+   *
+   * @example
+   * .args<BlockArgs>((data) => ({ numbers: data.numbers }))
+   *
+   * @example
+   * .args<BlockArgs>((data) => {
+   *   if (data.numbers.length === 0) throw new Error('Numbers required'); // block not ready
+   *   return { numbers: data.numbers };
+   * })
+   */
+  public args<Args>(lambda: (data: Data) => Args): BlockModelV3<Args, OutputsCfg, Data, Href> {
+    return new BlockModelV3<Args, OutputsCfg, Data, Href>({
+      ...this.config,
+      args: createAndRegisterRenderLambda<Args>({ handle: 'args', lambda }),
+    });
+  }
+
+  /**
+   * Sets a function to derive pre-run args from data (optional).
+   * This is called during setData to compute the args that will be used for staging/pre-run phase.
+   *
+   * If not defined, defaults to using the args() function result.
+   * If defined, uses its return value for the staging / prerun phase.
+   *
+   * The staging / prerun phase runs only if currentPrerunArgs differs from the executed
+   * version of prerunArgs (same comparison logic as currentArgs vs prodArgs).
+   *
+   * @example
+   * .prerunArgs((data) => ({ numbers: data.numbers }))
+   *
+   * @example
+   * .prerunArgs((data) => {
+   *   // Return undefined to skip staging for this block
+   *   if (!data.isReady) return undefined;
+   *   return { numbers: data.numbers };
+   * })
+   */
+  public prerunArgs(fn: (data: Data) => unknown): BlockModelV3<Args, OutputsCfg, Data, Href> {
+    return new BlockModelV3<Args, OutputsCfg, Data, Href>({
+      ...this.config,
+      prerunArgs: createAndRegisterRenderLambda({ handle: 'prerunArgs', lambda: fn }),
+    });
+  }
+
+  /** Sets the lambda to generate list of sections in the left block overviews panel. */
+  public sections<
+    const Ret extends SectionsExpectedType,
+    const RF extends RenderFunction<Args, Data, Ret>,
+  >(rf: RF): BlockModelV3<Args, OutputsCfg, Data, DeriveHref<ReturnType<RF>>> {
+    return new BlockModelV3<Args, OutputsCfg, Data, DeriveHref<ReturnType<RF>>>({
+      ...this.config,
+      // Replace the default sections callback with the user-provided one
+      sections: createAndRegisterRenderLambda({ handle: 'sections', lambda: () => rf(new RenderCtx()) }, true),
+    });
+  }
+
+  /** Sets a rendering function to derive block title, shown for the block in the left blocks-overview panel. */
+  public title(
+    rf: RenderFunction<Args, Data, string>,
+  ): BlockModelV3<Args, OutputsCfg, Data, Href> {
+    return new BlockModelV3<Args, OutputsCfg, Data, Href>({
+      ...this.config,
+      title: createAndRegisterRenderLambda({ handle: 'title', lambda: () => rf(new RenderCtx()) }),
+    });
+  }
+
+  public subtitle(
+    rf: RenderFunction<Args, Data, string>,
+  ): BlockModelV3<Args, OutputsCfg, Data, Href> {
+    return new BlockModelV3<Args, OutputsCfg, Data, Href>({
+      ...this.config,
+      subtitle: createAndRegisterRenderLambda({ handle: 'subtitle', lambda: () => rf(new RenderCtx()) }),
+    });
+  }
+
+  public tags(
+    rf: RenderFunction<Args, Data, string[]>,
+  ): BlockModelV3<Args, OutputsCfg, Data, Href> {
+    return new BlockModelV3<Args, OutputsCfg, Data, Href>({
+      ...this.config,
+      tags: createAndRegisterRenderLambda({ handle: 'tags', lambda: () => rf(new RenderCtx()) }),
+    });
+  }
+
+  /** Sets or overrides feature flags for the block. */
+  public withFeatureFlags(flags: Partial<BlockCodeKnownFeatureFlags>): BlockModelV3<Args, OutputsCfg, Data, Href> {
+    return new BlockModelV3<Args, OutputsCfg, Data, Href>({
+      ...this.config,
+      featureFlags: { ...this.config.featureFlags, ...flags },
+    });
+  }
+
+  /**
+   * Defines how to derive list of upstream references this block is meant to enrich with its exports from block args.
+   * Influences dependency graph construction.
+   */
+  public enriches(
+    lambda: (args: Args) => PlRef[],
+  ): BlockModelV3<Args, OutputsCfg, Data, Href> {
+    return new BlockModelV3<Args, OutputsCfg, Data, Href>({
+      ...this.config,
+      enrichmentTargets: createAndRegisterRenderLambda({ handle: 'enrichmentTargets', lambda: lambda }),
+    });
+  }
+
+  /** Renders all provided block settings into a pre-configured platforma API
+   * instance, that can be used in frontend to interact with block data, and
+   * other features provided by the platforma to the block. */
+  public done(): PlatformaExtended<PlatformaV3<
+    Args,
+    InferOutputsFromLambdas<OutputsCfg>,
+    Data,
+    Href
+  >> {
+    return this.withFeatureFlags({
+      ...this.config.featureFlags,
+    })._done();
+  }
+
+  public _done(): PlatformaExtended<PlatformaV3<
+    Args,
+    InferOutputsFromLambdas<OutputsCfg>,
+    Data,
+    Href
+  >> {
+    if (this.config.args === undefined) throw new Error('Args rendering function not set.');
+
+    const apiVersion = 3;
+
+    const migrationCount = this.config.dataModel.migrationCount;
+
+    const blockConfig: BlockConfigContainer = {
+      v4: {
+        configVersion: 4,
+        modelAPIVersion: 2,
+        sdkVersion: PlatformaSDKVersion,
+        renderingMode: this.config.renderingMode,
+        args: this.config.args,
+        prerunArgs: this.config.prerunArgs,
+        // Reference to __pl_data_initial callback registered by DataModel
+        initialData: createRenderLambda({ handle: '__pl_data_initial' }),
+        sections: this.config.sections,
+        title: this.config.title,
+        outputs: this.config.outputs,
+        enrichmentTargets: this.config.enrichmentTargets,
+        featureFlags: this.config.featureFlags,
+        // Generate migration descriptors (indices for metadata)
+        migrations: migrationCount > 0
+          ? Array.from({ length: migrationCount }, (_, i) => ({ index: i }))
+          : undefined,
+      },
+
+      // fields below are added to allow previous desktop versions read generated configs
+      sdkVersion: PlatformaSDKVersion,
+      renderingMode: this.config.renderingMode,
+      sections: this.config.sections,
+      outputs: Object.fromEntries(
+        Object.entries(this.config.outputs).map(([key, value]) => [key, downgradeCfgOrLambda(value)]),
+      ),
+    };
+
+    globalThis.platformaApiVersion = apiVersion;
+
+    if (!isInUI())
+    // we are in the configuration rendering routine, not in actual UI
+      return { config: blockConfig } as any;
+    // normal operation inside the UI
+    else return {
+      ...getPlatformaInstance({ sdkVersion: PlatformaSDKVersion, apiVersion }),
+      blockModelInfo: {
+        outputs: Object.fromEntries(
+          Object.entries(this.config.outputs)
+            .map(([key, value]) => [key, {
+              withStatus: Boolean(isConfigLambda(value) && value.withStatus),
+            }]),
+        ),
+      },
+    } as any;
+  }
+}
+
+// Type tests for BlockModelV3
+
+export type Expect<T extends true> = T;
+
+export type Equal<X, Y> =
+  (<T>() => T extends X ? 1 : 2) extends (<T>() => T extends Y ? 1 : 2) ? true : false;
+
+export type Merge<A, B> = {
+  [K in keyof A | keyof B]: K extends keyof B ? B[K] : K extends keyof A ? A[K] : never;
+};
+
+// Helper types for testing
+type _TestArgs = { inputFile: string; threshold: number };
+type _TestData = { selectedTab: string };
+type _TestOutputs = { result: ConfigRenderLambda<string>; count: ConfigRenderLambda<number> };
+
+// Test: Merge type works correctly
+type _MergeTest1 = Expect<Equal<Merge<{ a: 1 }, { b: 2 }>, { a: 1; b: 2 }>>;
+type _MergeTest2 = Expect<Equal<Merge<{ a: 1 }, { a: 2 }>, { a: 2 }>>;
+type _MergeTest3 = Expect<Equal<Merge<{ a: 1; b: 1 }, { b: 2; c: 3 }>, { a: 1; b: 2; c: 3 }>>;
+
+// Test: create() returns a BlockModelV3 instance
+// Note: Due to function overloads, ReturnType uses the last overload signature.
+// We verify the structure is correct using a simpler assignability test.
+type _CreateResult = ReturnType<typeof BlockModelV3.create>;
+type _CreateIsBlockModelV3 = _CreateResult extends BlockModelV3<infer _A, infer _O, infer _S> ? true : false;
+type _CreateTest = Expect<_CreateIsBlockModelV3>;
+
+// Test: BlockModelV3Config interface structure
+type _ConfigTest = Expect<Equal<
+  BlockModelV3Config<_TestArgs, _TestOutputs, _TestData>,
+  {
+    renderingMode: BlockRenderingMode;
+    args: ConfigRenderLambda<_TestArgs> | undefined;
+    prerunArgs: ConfigRenderLambda<unknown> | undefined;
+    dataModel: DataModel<_TestData>;
+    outputs: _TestOutputs;
+    sections: ConfigRenderLambda;
+    title: ConfigRenderLambda | undefined;
+    subtitle: ConfigRenderLambda | undefined;
+    tags: ConfigRenderLambda | undefined;
+    enrichmentTargets: ConfigRenderLambda | undefined;
+    featureFlags: BlockCodeKnownFeatureFlags;
+  }
+>>;
+
+// Test: Default Href is '/'
+type _HrefDefaultTest = BlockModelV3<_TestArgs, {}, _TestData> extends BlockModelV3<_TestArgs, {}, _TestData, '/'> ? true : false;
+type _VerifyHrefDefault = Expect<_HrefDefaultTest>;
+
+// Test: Custom Href can be specified
+type _CustomHref = '/settings' | '/main';
+type _HrefCustomBuilder = BlockModelV3<_TestArgs, {}, _TestData, _CustomHref>;
+type _HrefCustomTest = _HrefCustomBuilder extends BlockModelV3<_TestArgs, {}, _TestData, _CustomHref> ? true : false;
+type _VerifyHrefCustom = Expect<_HrefCustomTest>;
+
+// Test: Output type accumulation with & intersection
+type _OutputsAccumulation = { a: ConfigRenderLambda<string> } & { b: ConfigRenderLambda<number> };
+type _VerifyOutputsHaveKeys = Expect<Equal<keyof _OutputsAccumulation, 'a' | 'b'>>;
+
+// Test: Builder with all type parameters specified compiles
+type _FullBuilder = BlockModelV3<_TestArgs, _TestOutputs, _TestData, '/main'>;
+type _FullBuilderTest = _FullBuilder extends BlockModelV3<_TestArgs, _TestOutputs, _TestData, '/main'> ? true : false;
+type _VerifyFullBuilder = Expect<_FullBuilderTest>;
+
+// Test: InferOutputsFromLambdas maps outputs correctly
+type _InferOutputsTest = InferOutputsFromLambdas<{ myOutput: ConfigRenderLambda<number> }>;
+type _VerifyInferOutputs = Expect<Equal<
+  _InferOutputsTest,
+  { myOutput: OutputWithStatus<number> & { __unwrap: true } }
+>>;

package/src/block_storage.test.ts

@@ -0,0 +1,258 @@
+import { describe, expect, it } from 'vitest';
+import {
+  BLOCK_STORAGE_KEY,
+  BLOCK_STORAGE_SCHEMA_VERSION,
+  createBlockStorage,
+  defaultBlockStorageHandlers,
+  getFromStorage,
+  getPluginData,
+  getPluginNames,
+  getStorageData,
+  getStorageDataVersion,
+  isBlockStorage,
+  mergeBlockStorageHandlers,
+  normalizeBlockStorage,
+  removePluginData,
+  setPluginData,
+  updateStorageData,
+  updateStorageDataVersion,
+  updateStorage,
+} from './block_storage';
+
+describe('BlockStorage', () => {
+  describe('BLOCK_STORAGE_KEY and BLOCK_STORAGE_SCHEMA_VERSION', () => {
+    it('should have correct key constant', () => {
+      expect(typeof BLOCK_STORAGE_KEY).toBe('string');
+      expect(BLOCK_STORAGE_KEY).toBe('__pl_a7f3e2b9__');
+    });
+
+    it('should have correct schema version', () => {
+      expect(BLOCK_STORAGE_SCHEMA_VERSION).toBe('v1');
+    });
+  });
+
+  describe('isBlockStorage', () => {
+    it('should return true for valid BlockStorage with discriminator', () => {
+      const storage = createBlockStorage({});
+      expect(isBlockStorage(storage)).toBe(true);
+    });
+
+    it('should return true for BlockStorage with plugin data', () => {
+      const storage = setPluginData(createBlockStorage({ foo: 'bar' }), 'test', { data: 123 });
+      expect(isBlockStorage(storage)).toBe(true);
+    });
+
+    it('should return false for null', () => {
+      expect(isBlockStorage(null)).toBe(false);
+    });
+
+    it('should return false for undefined', () => {
+      expect(isBlockStorage(undefined)).toBe(false);
+    });
+
+    it('should return false for primitive values', () => {
+      expect(isBlockStorage(42)).toBe(false);
+      expect(isBlockStorage('string')).toBe(false);
+      expect(isBlockStorage(true)).toBe(false);
+    });
+
+    it('should return false for objects without discriminator', () => {
+      expect(isBlockStorage({ __dataVersion: 1, __data: {} })).toBe(false);
+    });
+
+    it('should return false for objects with wrong discriminator value', () => {
+      expect(isBlockStorage({ [BLOCK_STORAGE_KEY]: 'wrong', __dataVersion: 1, __data: {} })).toBe(false);
+    });
+  });
+
+  describe('createBlockStorage', () => {
+    it('should create storage with discriminator key and default values', () => {
+      const storage = createBlockStorage();
+      expect(storage[BLOCK_STORAGE_KEY]).toBe(BLOCK_STORAGE_SCHEMA_VERSION);
+      expect(storage.__dataVersion).toBe(1);
+      expect(storage.__data).toEqual({});
+    });
+
+    it('should create storage with custom initial data', () => {
+      const data = { numbers: [1, 2, 3] };
+      const storage = createBlockStorage(data);
+      expect(storage[BLOCK_STORAGE_KEY]).toBe(BLOCK_STORAGE_SCHEMA_VERSION);
+      expect(storage.__dataVersion).toBe(1);
+      expect(storage.__data).toEqual(data);
+    });
+
+    it('should create storage with custom version', () => {
+      const storage = createBlockStorage({ foo: 'bar' }, 5);
+      expect(storage[BLOCK_STORAGE_KEY]).toBe(BLOCK_STORAGE_SCHEMA_VERSION);
+      expect(storage.__dataVersion).toBe(5);
+      expect(storage.__data).toEqual({ foo: 'bar' });
+    });
+  });
+
+  describe('normalizeBlockStorage', () => {
+    it('should return BlockStorage as-is', () => {
+      const storage = createBlockStorage({ data: 'test' }, 2);
+      const normalized = normalizeBlockStorage(storage);
+      expect(normalized).toEqual(storage);
+    });
+
+    it('should wrap legacy data in BlockStorage structure', () => {
+      const legacyData = { numbers: [1, 2, 3], name: 'test' };
+      const normalized = normalizeBlockStorage(legacyData);
+      expect(normalized[BLOCK_STORAGE_KEY]).toBe(BLOCK_STORAGE_SCHEMA_VERSION);
+      expect(normalized.__dataVersion).toBe(1);
+      expect(normalized.__data).toEqual(legacyData);
+    });
+
+    it('should wrap primitive legacy data', () => {
+      const normalized = normalizeBlockStorage('simple string');
+      expect(normalized[BLOCK_STORAGE_KEY]).toBe(BLOCK_STORAGE_SCHEMA_VERSION);
+      expect(normalized.__dataVersion).toBe(1);
+      expect(normalized.__data).toBe('simple string');
+    });
+
+    it('should wrap null legacy data', () => {
+      const normalized = normalizeBlockStorage(null);
+      expect(normalized[BLOCK_STORAGE_KEY]).toBe(BLOCK_STORAGE_SCHEMA_VERSION);
+      expect(normalized.__dataVersion).toBe(1);
+      expect(normalized.__data).toBeNull();
+    });
+  });
+
+  describe('Data access functions', () => {
+    const storage = createBlockStorage({ count: 42 }, 3);
+
+    it('getStorageData should return the data', () => {
+      expect(getStorageData(storage)).toEqual({ count: 42 });
+    });
+
+    it('getStorageDataVersion should return the version', () => {
+      expect(getStorageDataVersion(storage)).toBe(3);
+    });
+
+    it('updateStorageData should return new storage with updated data', () => {
+      const newStorage = updateStorageData(storage, { operation: 'update-data', value: { count: 100 } });
+      expect(newStorage.__data).toEqual({ count: 100 });
+      expect(newStorage.__dataVersion).toBe(3);
+      expect(newStorage[BLOCK_STORAGE_KEY]).toBe(BLOCK_STORAGE_SCHEMA_VERSION);
+      // Original should be unchanged
+      expect(storage.__data).toEqual({ count: 42 });
+    });
+
+    it('updateStorageDataVersion should return new storage with updated version', () => {
+      const newStorage = updateStorageDataVersion(storage, 5);
+      expect(newStorage.__dataVersion).toBe(5);
+      expect(newStorage.__data).toEqual({ count: 42 });
+      expect(newStorage[BLOCK_STORAGE_KEY]).toBe(BLOCK_STORAGE_SCHEMA_VERSION);
+      // Original should be unchanged
+      expect(storage.__dataVersion).toBe(3);
+    });
+  });
+
+  describe('Plugin data functions', () => {
+    const baseStorage = createBlockStorage({});
+
+    it('setPluginData should add plugin data', () => {
+      const storage = setPluginData(baseStorage, 'table', { columns: ['a', 'b'] });
+      expect(storage['@plugin/table']).toEqual({ columns: ['a', 'b'] });
+    });
+
+    it('getPluginData should retrieve plugin data', () => {
+      const storage = setPluginData(baseStorage, 'chart', { type: 'bar' });
+      expect(getPluginData(storage, 'chart')).toEqual({ type: 'bar' });
+    });
+
+    it('getPluginData should return undefined for missing plugin', () => {
+      expect(getPluginData(baseStorage, 'nonexistent')).toBeUndefined();
+    });
+
+    it('removePluginData should remove plugin data', () => {
+      let storage = setPluginData(baseStorage, 'toRemove', { data: 'test' });
+      storage = setPluginData(storage, 'toKeep', { other: 'data' });
+      const result = removePluginData(storage, 'toRemove');
+      expect(result['@plugin/toRemove']).toBeUndefined();
+      expect(result['@plugin/toKeep']).toEqual({ other: 'data' });
+    });
+
+    it('getPluginNames should return all plugin names', () => {
+      let storage = createBlockStorage({});
+      storage = setPluginData(storage, 'alpha', {});
+      storage = setPluginData(storage, 'beta', {});
+      storage = setPluginData(storage, 'gamma', {});
+      const names = getPluginNames(storage);
+      expect(names.sort()).toEqual(['alpha', 'beta', 'gamma']);
+    });
+
+    it('getPluginNames should return empty array when no plugins', () => {
+      expect(getPluginNames(baseStorage)).toEqual([]);
+    });
+  });
+
+  describe('Generic storage access', () => {
+    const storage = createBlockStorage('hello', 2);
+
+    it('getFromStorage should get __data', () => {
+      expect(getFromStorage(storage, '__data')).toBe('hello');
+    });
+
+    it('getFromStorage should get __dataVersion', () => {
+      expect(getFromStorage(storage, '__dataVersion')).toBe(2);
+    });
+
+    it('updateStorage should update any key', () => {
+      const updated = updateStorage(storage, '__data', 'world');
+      expect(updated.__data).toBe('world');
+      expect(storage.__data).toBe('hello'); // immutable
+    });
+  });
+
+  describe('BlockStorageHandlers', () => {
+    describe('defaultBlockStorageHandlers', () => {
+      it('transformStateForStorage should replace data', () => {
+        const storage = createBlockStorage('old');
+        const result = defaultBlockStorageHandlers.transformStateForStorage(storage, 'new');
+        expect(result.__data).toBe('new');
+        expect(result.__dataVersion).toBe(1);
+      });
+
+      it('deriveStateForArgs should return data directly', () => {
+        const storage = createBlockStorage({ data: 'test' });
+        expect(defaultBlockStorageHandlers.deriveStateForArgs(storage)).toEqual({ data: 'test' });
+      });
+
+      it('migrateStorage should update version only', () => {
+        const storage = createBlockStorage({ data: 'test' });
+        const result = defaultBlockStorageHandlers.migrateStorage(storage, 1, 3);
+        expect(result.__dataVersion).toBe(3);
+        expect(result.__data).toEqual({ data: 'test' });
+      });
+    });
+
+    describe('mergeBlockStorageHandlers', () => {
+      it('should return defaults when no custom handlers provided', () => {
+        const handlers = mergeBlockStorageHandlers();
+        expect(handlers.transformStateForStorage).toBe(
+          defaultBlockStorageHandlers.transformStateForStorage,
+        );
+        expect(handlers.deriveStateForArgs).toBe(defaultBlockStorageHandlers.deriveStateForArgs);
+        expect(handlers.migrateStorage).toBe(defaultBlockStorageHandlers.migrateStorage);
+      });
+
+      it('should override with custom handlers', () => {
+        const customTransform = <T>(storage: ReturnType<typeof createBlockStorage<T>>, data: T) => ({
+          ...storage,
+          __data: data,
+          __dataVersion: storage.__dataVersion + 1,
+        });
+
+        const handlers = mergeBlockStorageHandlers({
+          transformStateForStorage: customTransform,
+        });
+
+        expect(handlers.transformStateForStorage).toBe(customTransform);
+        expect(handlers.deriveStateForArgs).toBe(defaultBlockStorageHandlers.deriveStateForArgs);
+      });
+    });
+  });
+});
+