Fixes for problems Claude discovered in middle-layer (#1483)

Author: vadimpiven

.changeset/clean-carpets-hunt.md

@@ -0,0 +1,7 @@
+---
+"@milaboratories/pl-middle-layer": patch
+"@milaboratories/pl-client": patch
+"@milaboratories/pl-tree": patch
+---
+
+Fixes for problems Claude discovered

lib/node/pl-client/src/core/errors.ts

@@ -53,7 +53,7 @@ export function isAbortedError(err: unknown, nested: boolean = false): boolean {
   if (err instanceof DOMException && err.code === DOMException.ABORT_ERR) return true; // WebSocket error
   if ((err as any).name == "RpcError" && (err as any).code == "ABORTED") return true;
   if ((err as any).name == "RESTError" && (err as any).status.code == Code.ABORTED) return true;
-  if ((err as any).cause !== undefined && !nested) isAbortedError((err as any).cause, true);
+  if ((err as any).cause !== undefined && !nested) return isAbortedError((err as any).cause, true);
   return false;
 }
 

lib/node/pl-client/src/core/websocket_stream.ts

@@ -61,6 +61,9 @@ export class WebSocketBiDiStream<
   // Error tracking
   private lastError?: Error;
 
+  // Abort listener reference for cleanup
+  private readonly abortHandler?: () => void;
+
   // === Public API ===
 
   public readonly requests = {
@@ -122,7 +125,8 @@ export class WebSocketBiDiStream<
       return;
     }
 
-    this.options.abortSignal?.addEventListener("abort", () => this.close());
+    this.abortHandler = () => this.close();
+    this.options.abortSignal?.addEventListener("abort", this.abortHandler, { once: true });
     this.connect();
   }
 
@@ -192,6 +196,11 @@ export class WebSocketBiDiStream<
   private onClose(): void {
     this.progressConnectionState(ConnectionState.CLOSED);
 
+    // Clean up abort listener to prevent memory leaks
+    if (this.abortHandler) {
+      this.options.abortSignal?.removeEventListener("abort", this.abortHandler);
+    }
+
     // If abort signal was triggered, use that as the error source
     if (this.options.abortSignal?.aborted && !this.lastError) {
       const reason = this.options.abortSignal.reason;

lib/node/pl-middle-layer/src/middle_layer/middle_layer.ts

@@ -308,7 +308,7 @@ export class MiddleLayer {
       }),
       runtimeCapabilities,
       quickJs,
-      projectHelper: new ProjectHelper(quickJs),
+      projectHelper: new ProjectHelper(quickJs, logger),
       dispose: async () => {
         await retryHttpDispatcher.destroy();
         await driverKit.dispose();

lib/node/pl-middle-layer/src/middle_layer/project.ts

@@ -24,7 +24,14 @@ import { getBlockParameters, blockOutputs } from "./block";
 import type { FrontendData } from "../model/frontend";
 import type { ProjectStructure } from "../model/project_model";
 import { projectFieldName } from "../model/project_model";
-import { cachedDeserialize, notEmpty, type MiLogger } from "@milaboratories/ts-helpers";
+import {
+  cachedDeserialize,
+  notEmpty,
+  type MiLogger,
+  createInfiniteRetryState,
+  nextInfiniteRetryState,
+  type InfiniteRetryState,
+} from "@milaboratories/ts-helpers";
 import type { BlockPackInfo } from "../model/block_pack";
 import type {
   ProjectOverview,
@@ -116,6 +123,7 @@ export class Project {
   }
 
   private async refreshLoop(): Promise<void> {
+    let retryState: InfiniteRetryState | undefined;
     while (!this.destroyed) {
       try {
         await withProject(
@@ -128,7 +136,9 @@ export class Project {
           { name: "doRefresh", lockId: this.projectLockId },
         );
         await this.activeConfigs.getValue();
-        await setTimeout(this.env.ops.projectRefreshInterval, this.abortController.signal);
+        await setTimeout(this.env.ops.projectRefreshInterval, undefined, {
+          signal: this.abortController.signal,
+        });
 
         // Block computables housekeeping
         const overviewLight = await this.overviewLight.getValue();
@@ -141,25 +151,41 @@ export class Project {
             this.blockComputables.set(blockId, null);
           }
         }
+        retryState = undefined;
       } catch (e: unknown) {
         // If we're destroyed, exit gracefully regardless of error type
-        if (this.destroyed) {
-          // Log just in case, to help with debugging if something unexpected happens during shutdown
-          this.env.logger.warn(new Error("Error during refresh loop shutdown", { cause: e }));
-          break;
-        }
+        if (this.destroyed) break;
 
         if (isNotFoundError(e)) {
-          console.warn(
+          this.env.logger.warn(
             "project refresh routine terminated, because project was externally deleted",
           );
           break;
         } else if (isTimeoutOrCancelError(e)) {
           // Timeout during normal operation, continue the loop
         } else {
-          // TODO: This stops the refresh loop permanently, leaving the project broken.
-          // Need to decide how to handle this case.
-          throw new Error("Unexpected exception", { cause: e });
+          retryState = retryState
+            ? nextInfiniteRetryState(retryState)
+            : createInfiniteRetryState({
+                type: "exponentialWithMaxDelayBackoff",
+                initialDelay: 1000,
+                maxDelay: 60_000,
+                backoffMultiplier: 2,
+                jitter: 0,
+              });
+          this.env.logger.error(
+            new Error(`[refreshLoop] unexpected exception, retrying in ${retryState.nextDelay}ms`, {
+              cause: e,
+            }),
+          );
+          try {
+            await setTimeout(retryState.nextDelay, undefined, {
+              signal: this.abortController.signal,
+            });
+          } catch {
+            // Aborted during retry delay, will exit via while condition or destroyed check
+            break;
+          }
         }
       }
     }

lib/node/pl-middle-layer/src/model/project_helper.ts

@@ -16,6 +16,7 @@ import { LRUCache } from "lru-cache";
 import type { QuickJSWASMModule } from "quickjs-emscripten";
 import { executeSingleLambda } from "../js_render";
 import type { ResourceId } from "@milaboratories/pl-client";
+import { ConsoleLoggerAdapter, type MiLogger } from "@milaboratories/ts-helpers";
 
 type EnrichmentTargetsRequest = {
   blockConfig: () => BlockConfig;
@@ -55,7 +56,10 @@ export class ProjectHelper {
     },
   });
 
-  constructor(private readonly quickJs: QuickJSWASMModule) {}
+  constructor(
+    private readonly quickJs: QuickJSWASMModule,
+    public readonly logger: MiLogger = new ConsoleLoggerAdapter(),
+  ) {}
 
   // =============================================================================
   // Args Derivation from Storage (V3+)
@@ -182,7 +186,11 @@ export class ProjectHelper {
       ) as string;
       return result;
     } catch (e) {
-      console.error("[ProjectHelper.getInitialStorageInVM] Initial storage creation failed:", e);
+      this.logger.error(
+        new Error("[ProjectHelper.getInitialStorageInVM] Initial storage creation failed", {
+          cause: e,
+        }),
+      );
       throw new Error(`Block initial storage creation failed: ${e}`);
     }
   }
@@ -219,7 +227,9 @@ export class ProjectHelper {
       ) as string;
       return result;
     } catch (e) {
-      console.error("[ProjectHelper.applyStorageUpdateInVM] Storage update failed:", e);
+      this.logger.error(
+        new Error("[ProjectHelper.applyStorageUpdateInVM] Storage update failed", { cause: e }),
+      );
       throw new Error(`Block storage update failed: ${e}`);
     }
   }
@@ -249,7 +259,11 @@ export class ProjectHelper {
       ) as StringifiedJson<StorageDebugView>;
       return result;
     } catch (e) {
-      console.error("[ProjectHelper.getStorageDebugViewInVM] Get storage debug view failed:", e);
+      this.logger.error(
+        new Error("[ProjectHelper.getStorageDebugViewInVM] Get storage debug view failed", {
+          cause: e,
+        }),
+      );
       return undefined;
     }
   }
@@ -290,7 +304,9 @@ export class ProjectHelper {
       ) as MigrationResult;
       return result;
     } catch (e) {
-      console.error("[ProjectHelper.migrateStorageInVM] Migration failed:", e);
+      this.logger.error(
+        new Error("[ProjectHelper.migrateStorageInVM] Migration failed", { cause: e }),
+      );
       return { error: `VM execution failed: ${e}` };
     }
   }

lib/node/pl-middle-layer/src/mutator/project.ts

@@ -63,6 +63,8 @@ import {
   notEmpty,
   canonicalJsonBytes,
   cachedDecode,
+  type MiLogger,
+  ConsoleLoggerAdapter,
 } from "@milaboratories/ts-helpers";
 import type { ProjectHelper } from "../model/project_helper";
 import {
@@ -119,6 +121,7 @@ class BlockInfo {
     public readonly fields: BlockFieldStates,
     public readonly config: BlockConfig,
     public readonly source: BlockPackSpec,
+    private readonly logger: MiLogger = new ConsoleLoggerAdapter(),
   ) {}
 
   public check() {
@@ -200,7 +203,7 @@ class BlockInfo {
     try {
       return this.blockStorageC();
     } catch (e) {
-      console.error("Error getting blockStorage:", e);
+      this.logger.error(new Error(`Error getting blockStorage for ${this.id}`, { cause: e }));
       return undefined;
     }
   }
@@ -595,6 +598,9 @@ export class ProjectMutator {
   ): { args?: unknown; uiState?: unknown } {
     const info = this.getBlockInfo(blockId);
     const currentState = info.blockStorage as { args?: unknown; uiState?: unknown } | undefined;
+    if (currentState === undefined) {
+      throw new Error(`Cannot merge block state for ${blockId}: blockStorage is unavailable`);
+    }
     return { ...currentState, ...partialUpdate };
   }
 
@@ -653,9 +659,17 @@ export class ProjectMutator {
       );
       if (prerunArgs !== undefined) {
         this.setBlockFieldObj(blockId, "currentPrerunArgs", this.createJsonFieldValue(prerunArgs));
+      } else {
+        this.deleteBlockFields(blockId, "currentPrerunArgs");
       }
     } else {
+      if (info.fields.currentPrerunArgs !== undefined) {
+        this.projectHelper.logger.warn(
+          `[staging] ${blockId}: currentPrerunArgs cleared (args derivation failed)`,
+        );
+      }
       this.deleteBlockFields(blockId, "currentArgs");
+      this.deleteBlockFields(blockId, "currentPrerunArgs");
     }
   }
 
@@ -756,10 +770,7 @@ export class ProjectMutator {
           prerunArgsData,
         );
       } else {
-        // prerunArgs is undefined - check if we previously had one
-        if (info.fields.currentPrerunArgs !== undefined) {
-          prerunArgsChanged = true;
-        }
+        this.deleteBlockFields(req.blockId, "currentPrerunArgs");
       }
 
       blockChanged = true;
@@ -903,6 +914,7 @@ export class ProjectMutator {
       {},
       extractConfig(spec.blockPack.config),
       spec.blockPack.source,
+      this.projectHelper.logger,
     );
     this.blockInfos.set(blockId, info);
 
@@ -1006,7 +1018,13 @@ export class ProjectMutator {
   }
 
   private initializeBlockDuplicate(blockId: string, originalBlockInfo: BlockInfo) {
-    const info = new BlockInfo(blockId, {}, originalBlockInfo.config, originalBlockInfo.source);
+    const info = new BlockInfo(
+      blockId,
+      {},
+      originalBlockInfo.config,
+      originalBlockInfo.source,
+      this.projectHelper.logger,
+    );
 
     this.blockInfos.set(blockId, info);
 
@@ -1210,7 +1228,12 @@ export class ProjectMutator {
             "currentPrerunArgs",
             this.createJsonFieldValue(prerunArgs),
           );
+        } else {
+          this.deleteBlockFields(blockId, "currentPrerunArgs");
         }
+      } else {
+        this.deleteBlockFields(blockId, "currentArgs");
+        this.deleteBlockFields(blockId, "currentPrerunArgs");
       }
     };
 
@@ -1250,7 +1273,9 @@ export class ProjectMutator {
           );
         }
 
-        console.log(`[migrateBlockPack] Block ${blockId}: ${migrationResult.info}`);
+        this.projectHelper.logger.info(
+          `[migrateBlockPack] Block ${blockId}: ${migrationResult.info}`,
+        );
         applyStorageAndDeriveArgs(migrationResult.newStorageJson);
       } else {
         // Legacy blocks (modelAPIVersion 1): persist block pack, set prerunArgs = currentArgs
@@ -1413,9 +1438,14 @@ export 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++;
+        try {
+          this.renderStagingFor(blockId);
+          rendered++;
+        } catch (e) {
+          this.projectHelper.logger.error(
+            new Error(`[refreshStagings] renderStagingFor failed for ${blockId}`, { cause: e }),
+          );
+        }
       }
     });
     if (rendered > 0) this.resetStagingRefreshTimestamp();
@@ -1654,7 +1684,10 @@ export class ProjectMutator {
 
     const blockInfos = new Map<string, BlockInfo>();
     blockInfoStates.forEach(({ id, fields, blockConfig, blockPack }) =>
-      blockInfos.set(id, new BlockInfo(id, fields, notEmpty(blockConfig), notEmpty(blockPack))),
+      blockInfos.set(
+        id,
+        new BlockInfo(id, fields, notEmpty(blockConfig), notEmpty(blockPack), projectHelper.logger),
+      ),
     );
 
     // check consistency of project state

lib/node/pl-middle-layer/src/network_check/network_check.ts

@@ -68,6 +68,9 @@ interface NetworkReports {
 }
 
 export interface CheckNetworkOpts {
+  /** Signal to abort all network checks. */
+  signal?: AbortSignal;
+
   /** Platforma Backend pings options. */
   pingCheckDurationMs: number;
   pingTimeoutMs: number;