@milaboratories/pl-middle-layer 1.45.5 → 1.46.1

package/dist/mutator/project.js

@@ -7,7 +7,7 @@ import { InitialBlockSettings } from '@milaboratories/pl-model-middle-layer';
 import Denque from 'denque';
 import { getPreparedExportTemplateEnvelope, exportContext } from './context_export.js';
 import { loadTemplate } from './template/template_loading.js';
-import { notEmpty, canonicalJsonBytes, canonicalJsonGzBytes, cachedDeserialize } from '@milaboratories/ts-helpers';
+import { notEmpty, canonicalJsonBytes, cachedDeserialize, cachedDecode } from '@milaboratories/ts-helpers';
 import { extractConfig, UiError } from '@platforma-sdk/model';
 import { getDebugFlags } from '../debug/index.js';
 
@@ -55,19 +55,57 @@ class BlockInfo {
             throw new Error('inconsistent stage cache fields');
         if (this.fields.blockPack === undefined)
             throw new Error('no block pack field');
-        if (this.fields.currentArgs === undefined)
-            throw new Error('no current args field');
+        if (this.fields.blockStorage === undefined)
+            throw new Error('no block storage field');
     }
-    currentArgsC = cached(() => this.fields.currentArgs.modCount, () => cachedDeserialize(this.fields.currentArgs.value));
+    currentArgsC = cached(() => this.fields.currentArgs?.modCount, () => {
+        const bin = this.fields.currentArgs?.value;
+        if (bin === undefined)
+            return undefined;
+        return cachedDeserialize(bin);
+    });
+    blockStorageC = cached(() => this.fields.blockStorage.modCount, () => {
+        const bin = this.fields.blockStorage?.value;
+        if (bin === undefined)
+            return undefined;
+        return cachedDeserialize(bin);
+    });
+    blockStorageJ = cached(() => this.fields.blockStorage.modCount, () => {
+        const bin = this.fields.blockStorage?.value;
+        if (bin === undefined)
+            return undefined;
+        return cachedDecode(bin);
+    });
     prodArgsC = cached(() => this.fields.prodArgs?.modCount, () => {
         const bin = this.fields.prodArgs?.value;
         if (bin === undefined)
             return undefined;
         return cachedDeserialize(bin);
     });
+    currentPrerunArgsC = cached(() => this.fields.currentPrerunArgs?.modCount, () => {
+        const bin = this.fields.currentPrerunArgs?.value;
+        if (bin === undefined)
+            return undefined;
+        return cachedDeserialize(bin);
+    });
     get currentArgs() {
         return this.currentArgsC();
     }
+    get blockStorage() {
+        try {
+            return this.blockStorageC();
+        }
+        catch (e) {
+            console.error('Error getting blockStorage:', e);
+            return undefined;
+        }
+    }
+    get blockStorageJson() {
+        return this.blockStorageJ();
+    }
+    get currentPrerunArgs() {
+        return this.currentPrerunArgsC();
+    }
     get stagingRendered() {
         return this.fields.stagingCtx !== undefined;
     }
@@ -79,12 +117,16 @@ class BlockInfo {
     }
     productionStaleC = cached(() => `${this.fields.currentArgs.modCount}_${this.fields.prodArgs?.modCount}`, () => this.fields.prodArgs === undefined
         || Buffer.compare(this.fields.currentArgs.value, this.fields.prodArgs.value) !== 0);
-    // get productionStale(): boolean {
-    //   return this.productionRendered && this.productionStaleC() && ;
-    // }
     get requireProductionRendering() {
         return !this.productionRendered || this.productionStaleC() || this.productionHasErrors;
     }
+    /** Returns true if staging should be re-rendered (stagingCtx is not set) */
+    get requireStagingRendering() {
+        // No staging needed if currentPrerunArgs is undefined (args derivation failed)
+        if (this.fields.currentPrerunArgs === undefined)
+            return false;
+        return !this.stagingRendered;
+    }
     get prodArgs() {
         return this.prodArgsC();
     }
@@ -159,7 +201,6 @@ class ProjectMutator {
             || this.renderingStateChanged);
     }
     get structure() {
-        // clone
         return JSON.parse(JSON.stringify(this.struct));
     }
     //
@@ -184,9 +225,13 @@ class ProjectMutator {
             args = bInfo.prodArgs;
         }
         else {
-            argsField = notEmpty(bInfo.fields.currentArgs);
+            argsField = bInfo.fields.currentArgs;
             args = bInfo.currentArgs;
         }
+        // Can't compute enrichment targets without args
+        if (argsField === undefined) {
+            return { args, enrichmentTargets: undefined };
+        }
         const blockPackField = notEmpty(bInfo.fields.blockPack);
         if (isResourceId(argsField.ref) && isResourceId(blockPackField.ref))
             return {
@@ -219,6 +264,8 @@ class ProjectMutator {
         return info;
     }
     createJsonFieldValueByContent(content) {
+        if (content === undefined)
+            throw new Error('content is undefined');
         const value = Buffer.from(content);
         const ref = this.tx.createValue(Pl.JsonObject, value);
         return { ref, value, status: 'Ready' };
@@ -319,6 +366,62 @@ class ProjectMutator {
             return this.deleteBlockFields(blockId, 'prodOutput', 'prodCtx', 'prodUiCtx', 'prodArgs', 'prodOutputPrevious', 'prodCtxPrevious', 'prodUiCtxPrevious');
         }
     }
+    /**
+     * Gets current block state and merges with partial updates.
+     * Used by legacy v1/v2 methods like setBlockArgs and setUiState.
+     *
+     * @param blockId The block to get state for
+     * @param partialUpdate Partial state to merge (e.g. { args } or { uiState })
+     * @returns Merged state in unified format { args, uiState }
+     */
+    mergeBlockState(blockId, partialUpdate) {
+        const info = this.getBlockInfo(blockId);
+        const currentState = info.blockStorage;
+        return { ...currentState, ...partialUpdate };
+    }
+    /**
+     * Sets raw block storage content directly (for testing purposes).
+     * This bypasses all normalization and VM transformations.
+     *
+     * @param blockId The block to set storage for
+     * @param rawStorageJson Raw storage as JSON string
+     */
+    setBlockStorageRaw(blockId, rawStorageJson) {
+        this.setBlockFieldObj(blockId, 'blockStorage', this.createJsonFieldValueByContent(rawStorageJson));
+        this.blocksWithChangedInputs.add(blockId);
+        this.updateLastModified();
+    }
+    /**
+     * Resets a v2+ block to its initial storage state.
+     * Gets initial storage from VM and derives args from it.
+     *
+     * For v1 blocks, use setStates() instead.
+     *
+     * @param blockId The block to reset
+     */
+    resetToInitialStorage(blockId) {
+        const info = this.getBlockInfo(blockId);
+        const blockConfig = info.config;
+        if (blockConfig.modelAPIVersion !== 2) {
+            throw new Error('resetToInitialStorage is only supported for model API version 2');
+        }
+        // Get initial storage from VM
+        const initialStorageJson = this.projectHelper.getInitialStorageInVM(blockConfig);
+        this.setBlockStorageRaw(blockId, initialStorageJson);
+        // Derive args from storage - set or clear currentArgs based on derivation result
+        const deriveArgsResult = this.projectHelper.deriveArgsFromStorage(blockConfig, initialStorageJson);
+        if (!deriveArgsResult.error) {
+            this.setBlockFieldObj(blockId, 'currentArgs', this.createJsonFieldValue(deriveArgsResult.value));
+            // Derive prerunArgs from storage
+            const prerunArgs = this.projectHelper.derivePrerunArgsFromStorage(blockConfig, initialStorageJson);
+            if (prerunArgs !== undefined) {
+                this.setBlockFieldObj(blockId, 'currentPrerunArgs', this.createJsonFieldValue(prerunArgs));
+            }
+        }
+        else {
+            this.deleteBlockFields(blockId, 'currentArgs');
+        }
+    }
     /** Optimally sets inputs for multiple blocks in one go */
     setStates(requests) {
         const changedArgs = [];
@@ -326,32 +429,76 @@ class ProjectMutator {
         for (const req of requests) {
             const info = this.getBlockInfo(req.blockId);
             let blockChanged = false;
-            for (const stateKey of ['args', 'uiState']) {
-                if (!(stateKey in req))
-                    continue;
-                const statePart = req[stateKey];
-                if (statePart === undefined || statePart === null)
-                    throw new Error(`Can't set ${stateKey} to null or undefined, please omit the key if you don't want to change it`);
-                const fieldName = stateKey === 'args' ? 'currentArgs' : 'uiState';
-                let data;
-                let gzipped = false;
-                if (stateKey === 'args') {
-                    // don't gzip args, workflow code can't uncompress gzip yet
-                    data = canonicalJsonBytes(statePart);
+            const blockConfig = info.config;
+            // modelAPIVersion === 2 means BlockModelV3 with .args() lambda for deriving args
+            if (req.modelAPIVersion !== blockConfig.modelAPIVersion) {
+                throw new Error(`Model API version mismatch for block ${req.blockId}: ${req.modelAPIVersion} !== ${blockConfig.modelAPIVersion}`);
+            }
+            // Derive args from storage using the block's config.args() callback
+            let args;
+            let prerunArgs;
+            if (req.modelAPIVersion === 2) {
+                const currentStorageJson = info.blockStorageJson;
+                if (currentStorageJson === undefined) {
+                    throw new Error(`Block ${req.blockId} has no blockStorage - this should not happen`);
+                }
+                // Apply the state update to storage
+                const updatedStorageJson = this.projectHelper.applyStorageUpdateInVM(blockConfig, currentStorageJson, req.payload);
+                this.setBlockFieldObj(req.blockId, 'blockStorage', this.createJsonFieldValueByContent(updatedStorageJson));
+                // Derive args directly from storage (VM extracts data internally)
+                const derivedArgsResult = this.projectHelper.deriveArgsFromStorage(blockConfig, updatedStorageJson);
+                if (derivedArgsResult.error) {
+                    args = undefined;
+                    prerunArgs = undefined;
                 }
                 else {
-                    const { data: binary, isGzipped } = canonicalJsonGzBytes(statePart);
-                    data = binary;
-                    gzipped = isGzipped;
+                    args = derivedArgsResult.value;
+                    // Derive prerunArgs from storage, or fall back to args
+                    prerunArgs = this.projectHelper.derivePrerunArgsFromStorage(blockConfig, updatedStorageJson);
                 }
-                if (Buffer.compare(info.fields[fieldName].value, data) === 0)
-                    continue;
-                // console.log('setting', fieldName, gzipped, data.length);
-                const statePartRef = this.tx.createValue(gzipped ? Pl.JsonGzObject : Pl.JsonObject, data);
-                this.setBlockField(req.blockId, fieldName, statePartRef, 'Ready', data);
-                blockChanged = true;
-                if (stateKey === 'args')
-                    changedArgs.push(req.blockId);
+            }
+            else {
+                this.setBlockFieldObj(req.blockId, 'blockStorage', this.createJsonFieldValue(req.state));
+                if (req.state !== null && typeof req.state === 'object' && 'args' in req.state) {
+                    args = req.state.args;
+                }
+                else {
+                    args = req.state;
+                }
+                // For the legacy blocks, prerunArgs = args (same as production args)
+                prerunArgs = args;
+            }
+            // Set or clear currentArgs based on derivation result
+            if (args !== undefined) {
+                const currentArgsData = canonicalJsonBytes(args);
+                const argsPartRef = this.tx.createValue(Pl.JsonObject, currentArgsData);
+                this.setBlockField(req.blockId, 'currentArgs', argsPartRef, 'Ready', currentArgsData);
+            }
+            else {
+                this.deleteBlockFields(req.blockId, 'currentArgs');
+            }
+            // Set currentPrerunArgs field and check if it actually changed
+            let prerunArgsChanged = false;
+            if (prerunArgs !== undefined) {
+                const prerunArgsData = canonicalJsonBytes(prerunArgs);
+                const oldPrerunArgsData = info.fields.currentPrerunArgs?.value;
+                // Check if prerunArgs actually changed
+                if (oldPrerunArgsData === undefined || Buffer.compare(oldPrerunArgsData, prerunArgsData) !== 0) {
+                    prerunArgsChanged = true;
+                }
+                const prerunArgsRef = this.tx.createValue(Pl.JsonObject, prerunArgsData);
+                this.setBlockField(req.blockId, 'currentPrerunArgs', prerunArgsRef, 'Ready', prerunArgsData);
+            }
+            else {
+                // prerunArgs is undefined - check if we previously had one
+                if (info.fields.currentPrerunArgs !== undefined) {
+                    prerunArgsChanged = true;
+                }
+            }
+            blockChanged = true;
+            // Only add to changedArgs if prerunArgs changed - this controls staging reset
+            if (prerunArgsChanged) {
+                changedArgs.push(req.blockId);
             }
             if (blockChanged) {
                 // will be assigned our author marker
@@ -393,15 +540,25 @@ class ProjectMutator {
     exportCtx(ctx) {
         return exportContext(this.tx, Pl.unwrapHolder(this.tx, this.ctxExportTplHolder), ctx);
     }
+    /**
+     * Renders staging for a block using currentPrerunArgs.
+     * If currentPrerunArgs is not set (prerunArgs returned undefined), skips staging for this block.
+     */
     renderStagingFor(blockId) {
         this.resetStaging(blockId);
         const info = this.getBlockInfo(blockId);
+        // If currentPrerunArgs is not set (prerunArgs returned undefined), skip staging for this block
+        const prerunArgsRef = info.fields.currentPrerunArgs?.ref;
+        if (prerunArgsRef === undefined) {
+            return;
+        }
         const ctx = this.createStagingCtx(this.getStagingGraph().nodes.get(blockId).upstream);
         if (this.getBlock(blockId).renderingMode !== 'Heavy')
             throw new Error('not supported yet');
         const tpl = info.getTemplate(this.tx);
+        // Use currentPrerunArgs for staging rendering
         const results = createRenderHeavyBlock(this.tx, tpl, {
-            args: info.fields.currentArgs.ref,
+            args: prerunArgsRef,
             blockId: this.tx.createValue(Pl.JsonString, JSON.stringify(blockId)),
             isProduction: this.tx.createValue(Pl.JsonBool, JSON.stringify(false)),
             context: ctx,
@@ -418,6 +575,10 @@ class ProjectMutator {
     renderProductionFor(blockId) {
         this.resetProduction(blockId);
         const info = this.getBlockInfo(blockId);
+        // Can't render production if currentArgs is not set
+        if (info.fields.currentArgs === undefined) {
+            throw new Error(`Can't render production for block ${blockId}: currentArgs not set`);
+        }
         const ctx = this.createProdCtx(this.getPendingProductionGraph().nodes.get(blockId).upstream);
         if (this.getBlock(blockId).renderingMode === 'Light')
             throw new Error('Can\'t render production for light block.');
@@ -448,10 +609,47 @@ class ProjectMutator {
         this.setBlockField(blockId, 'blockPack', Pl.wrapInHolder(this.tx, bp), 'NotReady');
         // settings
         this.setBlockFieldObj(blockId, 'blockSettings', this.createJsonFieldValue(InitialBlockSettings));
-        // args
-        this.setBlockFieldObj(blockId, 'currentArgs', this.createJsonFieldValueByContent(spec.args));
-        // uiState
-        this.setBlockFieldObj(blockId, 'uiState', this.createJsonFieldValueByContent(spec.uiState ?? '{}'));
+        const blockConfig = info.config;
+        let args;
+        let prerunArgs;
+        let storageToWrite;
+        if (spec.storageMode === 'fromModel') {
+            // Model API v2+: get initial storage and derive args from it
+            storageToWrite = this.projectHelper.getInitialStorageInVM(blockConfig);
+            // Derive args directly from storage (VM extracts data internally)
+            const deriveArgsResult = this.projectHelper.deriveArgsFromStorage(blockConfig, storageToWrite);
+            if (deriveArgsResult.error) {
+                args = undefined;
+                prerunArgs = undefined;
+            }
+            else {
+                args = deriveArgsResult.value;
+                prerunArgs = this.projectHelper.derivePrerunArgsFromStorage(blockConfig, storageToWrite);
+            }
+        }
+        else if (spec.storageMode === 'legacy') {
+            // Model API v1: use legacyState from spec
+            const parsedState = JSON.parse(spec.legacyState);
+            args = parsedState.args;
+            if (args === undefined) {
+                throw new Error('args is undefined in legacyState');
+            }
+            prerunArgs = args;
+            storageToWrite = spec.legacyState;
+        }
+        else {
+            throw new Error(`Unknown storageMode: ${spec.storageMode}`);
+        }
+        // currentArgs
+        if (args !== undefined) {
+            this.setBlockFieldObj(blockId, 'currentArgs', this.createJsonFieldValue(args));
+        }
+        // currentPrerunArgs
+        if (prerunArgs !== undefined) {
+            this.setBlockFieldObj(blockId, 'currentPrerunArgs', this.createJsonFieldValue(prerunArgs));
+        }
+        // blockStorage
+        this.setBlockFieldObj(blockId, 'blockStorage', this.createJsonFieldValueByContent(storageToWrite));
         // checking structure
         info.check();
     }
@@ -627,14 +825,53 @@ class ProjectMutator {
     //
     // Block-pack migration
     //
-    migrateBlockPack(blockId, spec, newArgsAndUiState) {
+    migrateBlockPack(blockId, spec, newClearState) {
         const info = this.getBlockInfo(blockId);
+        const newConfig = extractConfig(spec.config);
         this.setBlockField(blockId, 'blockPack', Pl.wrapInHolder(this.tx, createBlockPack(this.tx, spec)), 'NotReady');
-        if (newArgsAndUiState !== undefined) {
-            // this will also reset all downstream stagings
-            this.setStates([{ blockId, args: newArgsAndUiState.args, uiState: newArgsAndUiState.uiState }]);
+        if (newClearState !== undefined) {
+            // State is being reset - no migration needed
+            const supportsStorageFromVM = newConfig.modelAPIVersion === 2;
+            if (supportsStorageFromVM) {
+                // V2+: Get initial storage directly from VM and derive args from it
+                const initialStorageJson = this.projectHelper.getInitialStorageInVM(newConfig);
+                this.setBlockStorageRaw(blockId, initialStorageJson);
+                // Derive args from storage - only set currentArgs if derivation succeeds
+                const deriveArgsResult = this.projectHelper.deriveArgsFromStorage(newConfig, initialStorageJson);
+                if (!deriveArgsResult.error) {
+                    this.setBlockFieldObj(blockId, 'currentArgs', this.createJsonFieldValue(deriveArgsResult.value));
+                    // Derive prerunArgs from storage
+                    const prerunArgs = this.projectHelper.derivePrerunArgsFromStorage(newConfig, initialStorageJson);
+                    if (prerunArgs !== undefined) {
+                        this.setBlockFieldObj(blockId, 'currentPrerunArgs', this.createJsonFieldValue(prerunArgs));
+                    }
+                }
+                this.blocksWithChangedInputs.add(blockId);
+                this.updateLastModified();
+            }
+            else {
+                // V1: Use setStates with legacy state format
+                this.setStates([{ modelAPIVersion: 1, blockId, state: newClearState.state }]);
+            }
         }
         else {
+            // State is being preserved - run migrations if needed via VM
+            // Only Model API v2 blocks support migrations
+            const supportsStateMigrations = newConfig.modelAPIVersion === 2;
+            if (supportsStateMigrations) {
+                const currentStorageJson = info.blockStorageJson;
+                const migrationResult = this.projectHelper.migrateStorageInVM(newConfig, currentStorageJson);
+                if (migrationResult.error !== undefined) {
+                    console.error(`[migrateBlockPack] Block ${blockId} migration error: ${migrationResult.error}`);
+                }
+                else {
+                    console.log(`[migrateBlockPack] Block ${blockId}: ${migrationResult.info}`);
+                    if (migrationResult.warn) {
+                        console.warn(`[migrateBlockPack] Block ${blockId} migration warning: ${migrationResult.warn}`);
+                    }
+                    this.setBlockStorageRaw(blockId, migrationResult.newStorageJson);
+                }
+            }
             // resetting staging outputs for all downstream blocks
             this.getStagingGraph().traverse('downstream', [blockId], ({ id }) => this.resetStaging(id));
         }
@@ -735,13 +972,16 @@ class ProjectMutator {
         const stagingGraph = this.getStagingGraph();
         stagingGraph.nodes.forEach((node) => {
             const info = this.getBlockInfo(node.id);
-            let lag = info.stagingRendered ? 0 : 1;
+            // Use requireStagingRendering to check both: staging exists AND prerunArgs hasn't changed
+            const requiresRendering = info.requireStagingRendering;
+            let lag = requiresRendering ? 1 : 0;
             node.upstream.forEach((upstream) => {
                 const upstreamLag = lags.get(upstream);
                 if (upstreamLag === 0)
                     return;
                 lag = Math.max(upstreamLag + 1, lag);
             });
+            if (!requiresRendering && info.stagingRendered) ;
             cb(node.id, lag);
             lags.set(node.id, lag);
         });
@@ -758,6 +998,7 @@ class ProjectMutator {
                 // meaning staging already rendered
                 return;
             if (lagThreshold === undefined || lag <= lagThreshold) {
+                // console.log(`[refreshStagings] RENDER staging for ${blockId} (lag=${lag})`);
                 this.renderStagingFor(blockId);
                 rendered++;
             }