@ash-cloud/ash-ai 0.1.18 → 0.1.19

package/dist/index.cjs

@@ -4639,6 +4639,8 @@ echo "[warmup] Warmup complete!"
       lastMaintenanceAt = null;
       metricsCallback;
       startPromise = null;
+      /** Registered warmup specs by tag (e.g. agentId -> spec) */
+      warmupSpecs = /* @__PURE__ */ new Map();
       /** Consecutive warmup failure count (reset on success) */
       consecutiveFailures = 0;
       /** Timestamp of last warmup attempt — used for backoff */
@@ -4695,11 +4697,13 @@ echo "[warmup] Warmup complete!"
         if (this.maintenanceIntervalId.unref) {
           this.maintenanceIntervalId.unref();
         }
+        const initialSpecs = this.selectSpecsForReplenishment(this.config.minPoolSize);
         console.log(`[POOL] Spawning ${this.config.minPoolSize} initial sandbox(es)...`);
         const warmupPromises = [];
         for (let i = 0; i < this.config.minPoolSize; i++) {
+          const spec = initialSpecs[i];
           warmupPromises.push(
-            this.warmSandbox().then((sandbox) => {
+            this.warmSandbox(spec).then((sandbox) => {
               this.pool.set(sandbox.sandboxId, sandbox);
               console.log(`[POOL] Initial sandbox ready: ${sandbox.sandboxId}`);
             }).catch((error) => {
@@ -4743,9 +4747,11 @@ echo "[warmup] Warmup complete!"
       }
       /**
        * Acquire a pre-warmed sandbox for a session.
+       * If preferTag is provided, tries to find a sandbox warmed for that tag first.
+       * Falls back to a generic (untagged) sandbox if no tag match is found.
        * If no eligible sandbox is available, creates one on-demand.
        */
-      async acquire(sessionId) {
+      async acquire(sessionId, preferTag) {
         if (!this.running) {
           throw new Error("Sandbox pool is not running");
         }
@@ -4755,26 +4761,34 @@ echo "[warmup] Warmup complete!"
             return pooled;
           }
         }
-        const available = this.getAvailableSandbox();
+        const available = this.getAvailableSandbox(preferTag);
         if (available) {
           available.assignedTo = sessionId;
-          console.log(`[POOL] Acquired sandbox ${available.sandboxId} for session ${sessionId}`);
+          const tagInfo = available.warmupTag ? ` [tag=${available.warmupTag}, agentSetupComplete=${available.agentSetupComplete}]` : " [generic]";
+          console.log(`[POOL] Acquired sandbox ${available.sandboxId} for session ${sessionId}${tagInfo}`);
           this.emitMetric("sandbox_assigned", {
             sandboxId: available.sandboxId,
             sessionId,
-            poolAvailable: this.getAvailableCount()
+            poolAvailable: this.getAvailableCount(),
+            warmupTag: available.warmupTag,
+            agentSetupComplete: available.agentSetupComplete,
+            preferTag
           });
           this.triggerReplenishment();
           return available;
         }
-        console.log(`[POOL] No available sandbox, creating on-demand for session ${sessionId}...`);
-        const sandbox = await this.warmSandbox();
+        const spec = preferTag ? this.warmupSpecs.get(preferTag) : void 0;
+        console.log(`[POOL] No available sandbox, creating on-demand for session ${sessionId}${spec ? ` [spec=${preferTag}]` : ""}...`);
+        const sandbox = await this.warmSandbox(spec);
         sandbox.assignedTo = sessionId;
         this.pool.set(sandbox.sandboxId, sandbox);
         this.emitMetric("sandbox_assigned", {
           sandboxId: sandbox.sandboxId,
           sessionId,
-          onDemand: true
+          onDemand: true,
+          warmupTag: sandbox.warmupTag,
+          agentSetupComplete: sandbox.agentSetupComplete,
+          preferTag
         });
         return sandbox;
       }
@@ -4805,11 +4819,14 @@ echo "[warmup] Warmup complete!"
         let available = 0;
         let assigned = 0;
         let ineligible = 0;
+        const availableByTag = {};
         for (const pooled of this.pool.values()) {
           if (pooled.assignedTo) {
             assigned++;
           } else if (pooled.eligible) {
             available++;
+            const tagKey = pooled.warmupTag || "generic";
+            availableByTag[tagKey] = (availableByTag[tagKey] || 0) + 1;
           } else {
             ineligible++;
           }
@@ -4824,6 +4841,8 @@ echo "[warmup] Warmup complete!"
           lastMaintenanceAt: this.lastMaintenanceAt,
           consecutiveFailures: this.consecutiveFailures,
           warmupSuspended: this.consecutiveFailures >= _SandboxPool.MAX_CONSECUTIVE_FAILURES,
+          registeredSpecs: this.warmupSpecs.size,
+          availableByTag,
           config: {
             minPoolSize: this.config.minPoolSize,
             maxPoolSize: this.config.maxPoolSize,
@@ -4838,13 +4857,51 @@ echo "[warmup] Warmup complete!"
       onMetrics(callback) {
         this.metricsCallback = callback;
       }
+      /**
+       * Register a warmup spec so the pool can pre-warm agent-specific sandboxes.
+       * If a spec with the same tag already exists, it is replaced.
+       * Triggers replenishment to warm a sandbox for this spec.
+       */
+      registerWarmupSpec(spec) {
+        const isNew = !this.warmupSpecs.has(spec.tag);
+        this.warmupSpecs.set(spec.tag, spec);
+        console.log(`[POOL] ${isNew ? "Registered" : "Updated"} warmup spec: ${spec.tag} (priority=${spec.priority})`);
+        this.emitMetric("spec_registered", {
+          tag: spec.tag,
+          priority: spec.priority,
+          isNew,
+          totalSpecs: this.warmupSpecs.size
+        });
+        if (this.running) {
+          this.triggerReplenishment();
+        }
+      }
+      /**
+       * Remove a warmup spec. Existing tagged sandboxes remain but won't be replaced.
+       */
+      unregisterWarmupSpec(tag) {
+        this.warmupSpecs.delete(tag);
+        console.log(`[POOL] Unregistered warmup spec: ${tag}`);
+      }
+      /**
+       * Update the priority of a warmup spec (e.g. for MRU tracking).
+       * Higher priority = more likely to get a warm sandbox during replenishment.
+       */
+      updateSpecPriority(tag, priority) {
+        const spec = this.warmupSpecs.get(tag);
+        if (spec) {
+          spec.priority = priority;
+        }
+      }
       // ===========================================================================
       // PRIVATE METHODS
       // ===========================================================================
       /**
-       * Create and warm a new sandbox
+       * Create and warm a new sandbox.
+       * If a spec is provided, runs the spec's setup function after SDK installation.
+       * On spec setup failure, the sandbox remains generic (graceful degradation).
        */
-      async warmSandbox() {
+      async warmSandbox(spec) {
         const warmupId = `warming-${Date.now()}-${Math.random().toString(36).slice(2)}`;
         this.warmingInProgress.add(warmupId);
         this.emitMetric("warmup_started", { warmupId });
@@ -4938,6 +4995,37 @@ echo "[warmup] Warmup complete!"
               throw new Error(`Warmup failed: ${stderr}`);
             }
           }
+          let warmupTag;
+          let agentSetupComplete = false;
+          if (spec) {
+            console.log(`[POOL] Running spec setup for tag=${spec.tag} on sandbox ${sandbox.sandboxId}...`);
+            this.emitMetric("spec_setup_started", { tag: spec.tag, sandboxId: sandbox.sandboxId });
+            const specStartTime = Date.now();
+            try {
+              await spec.setup(sandbox);
+              warmupTag = spec.tag;
+              agentSetupComplete = true;
+              const specDuration = Date.now() - specStartTime;
+              console.log(`[POOL] Spec setup completed for tag=${spec.tag} on sandbox ${sandbox.sandboxId} (${specDuration}ms)`);
+              this.emitMetric("spec_setup_completed", {
+                tag: spec.tag,
+                sandboxId: sandbox.sandboxId,
+                durationMs: specDuration
+              });
+            } catch (specError) {
+              const specDuration = Date.now() - specStartTime;
+              const specErrorMessage = specError instanceof Error ? specError.message : "Unknown";
+              console.warn(
+                `[POOL] Spec setup failed for tag=${spec.tag} on sandbox ${sandbox.sandboxId} (${specDuration}ms): ${specErrorMessage}. Sandbox stays generic.`
+              );
+              this.emitMetric("spec_setup_failed", {
+                tag: spec.tag,
+                sandboxId: sandbox.sandboxId,
+                durationMs: specDuration,
+                error: specErrorMessage
+              });
+            }
+          }
           const warmupTime = Date.now() - startTime;
           const now = Date.now();
           const pooled = {
@@ -4947,14 +5035,19 @@ echo "[warmup] Warmup complete!"
             expiresAt: now + this.config.sandboxTimeout * 1e3,
             sdkInstalled: true,
             eligible: true,
-            lastHeartbeat: now
+            lastHeartbeat: now,
+            warmupTag,
+            agentSetupComplete
           };
-          console.log(`[POOL] Warmup completed for ${sandbox.sandboxId} (took ${warmupTime}ms)${useTarball ? " [tarball]" : ""}`);
+          const tagInfo = warmupTag ? ` [tag=${warmupTag}]` : "";
+          console.log(`[POOL] Warmup completed for ${sandbox.sandboxId} (took ${warmupTime}ms)${useTarball ? " [tarball]" : ""}${tagInfo}`);
           this.consecutiveFailures = 0;
           this.emitMetric("warmup_completed", {
             sandboxId: pooled.sandboxId,
             warmupTimeMs: warmupTime,
-            usedTarball: useTarball
+            usedTarball: useTarball,
+            warmupTag,
+            agentSetupComplete
           });
           this.emitMetric("sandbox_created", { sandboxId: pooled.sandboxId });
           return pooled;
@@ -5035,18 +5128,29 @@ echo "[warmup] Warmup complete!"
         }
       }
       /**
-       * Get an available eligible sandbox for assignment
+       * Get an available eligible sandbox for assignment.
+       * If preferTag is provided:
+       * - First tries to find a sandbox tagged for that specific agent
+       * - Falls back to a generic (untagged) sandbox
+       * - Tagged sandboxes for OTHER agents are NOT used as fallback (reserved)
        */
-      getAvailableSandbox() {
-        let best = null;
+      getAvailableSandbox(preferTag) {
+        let bestTagged = null;
+        let bestGeneric = null;
         for (const pooled of this.pool.values()) {
           if (!pooled.assignedTo && pooled.eligible) {
-            if (!best || pooled.expiresAt > best.expiresAt) {
-              best = pooled;
+            if (preferTag && pooled.warmupTag === preferTag) {
+              if (!bestTagged || pooled.expiresAt > bestTagged.expiresAt) {
+                bestTagged = pooled;
+              }
+            } else if (!pooled.warmupTag) {
+              if (!bestGeneric || pooled.expiresAt > bestGeneric.expiresAt) {
+                bestGeneric = pooled;
+              }
             }
           }
         }
-        return best;
+        return bestTagged || bestGeneric;
       }
       /**
        * Get count of available sandboxes
@@ -5113,13 +5217,16 @@ echo "[warmup] Warmup complete!"
           this.config.maxPoolSize - this.pool.size - warmingCount
         );
         if (needed <= 0) return;
+        const specAssignments = this.selectSpecsForReplenishment(needed);
         console.log(`[POOL] Spawning ${needed} sandbox(es) to maintain pool...`);
         const promises = [];
         for (let i = 0; i < needed; i++) {
+          const spec = specAssignments[i];
           promises.push(
-            this.warmSandbox().then((sandbox) => {
+            this.warmSandbox(spec).then((sandbox) => {
               this.pool.set(sandbox.sandboxId, sandbox);
-              console.log(`[POOL] Replenishment sandbox ready: ${sandbox.sandboxId}`);
+              const tagInfo = sandbox.warmupTag ? ` [tag=${sandbox.warmupTag}]` : "";
+              console.log(`[POOL] Replenishment sandbox ready: ${sandbox.sandboxId}${tagInfo}`);
             }).catch((error) => {
               console.error("[POOL] Failed to warm replenishment sandbox:", error);
             })
@@ -5130,6 +5237,39 @@ echo "[warmup] Warmup complete!"
         }
         await Promise.all(promises);
       }
+      /**
+       * Decide which specs to apply to new sandboxes during replenishment.
+       * Strategy: cover uncovered specs first (highest priority), then fill remaining as generic.
+       * Returns an array of length `needed`, where each element is a spec or undefined (generic).
+       */
+      selectSpecsForReplenishment(needed) {
+        if (this.warmupSpecs.size === 0) {
+          return new Array(needed).fill(void 0);
+        }
+        const uncoveredSpecs = [];
+        for (const spec of this.warmupSpecs.values()) {
+          let hasCoverage = false;
+          for (const pooled of this.pool.values()) {
+            if (pooled.warmupTag === spec.tag && pooled.eligible && !pooled.assignedTo) {
+              hasCoverage = true;
+              break;
+            }
+          }
+          if (!hasCoverage) {
+            uncoveredSpecs.push(spec);
+          }
+        }
+        uncoveredSpecs.sort((a, b) => b.priority - a.priority);
+        const assignments = [];
+        for (const spec of uncoveredSpecs) {
+          if (assignments.length >= needed) break;
+          assignments.push(spec);
+        }
+        while (assignments.length < needed) {
+          assignments.push(void 0);
+        }
+        return assignments;
+      }
       /**
        * Destroy a sandbox and clean up
        */
@@ -5361,7 +5501,7 @@ function removeExpiredSandbox(sessionId) {
   }
 }
 async function getOrCreateSandbox(options) {
-  const { sessionId, runtime = "node22", timeout = 300, vcpus = 4, existingSandboxId } = options;
+  const { sessionId, runtime = "node22", timeout = 300, vcpus = 4, existingSandboxId, preferTag } = options;
   ensureCleanupRunning();
   ensureHeartbeatRunning();
   const { Sandbox } = await import('@vercel/sandbox');
@@ -5378,6 +5518,8 @@ async function getOrCreateSandbox(options) {
         sdkInstalled: cached.sdkInstalled,
         startupScriptRan: cached.startupScriptRan,
         startupScriptHash: cached.startupScriptHash,
+        installScriptRan: cached.installScriptRan,
+        installScriptHash: cached.installScriptHash,
         isNew: false,
         configFileUrl: cached.configFileUrl,
         configInstalledAt: cached.configInstalledAt
@@ -5386,41 +5528,6 @@ async function getOrCreateSandbox(options) {
     console.log("[SANDBOX] Cached sandbox has expired (HTTP 410 or similar), creating new sandbox");
     removeExpiredSandbox(sessionId);
   }
-  const pool = await ensureSandboxPoolInitialized();
-  if (pool && pool.isRunning()) {
-    try {
-      console.log("[SANDBOX] Attempting to acquire from pre-warmed pool...");
-      const pooled = await pool.acquire(sessionId);
-      console.log("[SANDBOX] Acquired pre-warmed sandbox:", pooled.sandboxId);
-      const now2 = Date.now();
-      const entry2 = {
-        sandbox: pooled.sandbox,
-        sessionId,
-        createdAt: pooled.createdAt,
-        lastUsedAt: now2,
-        sdkInstalled: pooled.sdkInstalled,
-        startupScriptRan: false
-        // User script hasn't run yet
-      };
-      sandboxCache.set(sessionId, entry2);
-      return {
-        sandbox: pooled.sandbox,
-        sandboxId: pooled.sandboxId,
-        sdkInstalled: pooled.sdkInstalled,
-        startupScriptRan: false,
-        startupScriptHash: void 0,
-        isNew: false,
-        // Not new - came from pool
-        configFileUrl: void 0,
-        configInstalledAt: void 0
-      };
-    } catch (error) {
-      console.warn(
-        "[SANDBOX] Failed to acquire from pool, falling back to on-demand creation:",
-        error instanceof Error ? error.message : "Unknown error"
-      );
-    }
-  }
   if (existingSandboxId) {
     console.log("[SANDBOX] Attempting to reconnect to existing sandbox:", existingSandboxId);
     try {
@@ -5441,7 +5548,9 @@ async function getOrCreateSandbox(options) {
           // We assume SDK is installed since this is an existing sandbox
           // The caller can verify and re-mark if needed
           sdkInstalled: true,
-          startupScriptRan: true
+          startupScriptRan: true,
+          installScriptRan: true
+          // Assume ran for reconnected sandboxes
         };
         sandboxCache.set(sessionId, entry2);
         return {
@@ -5453,13 +5562,17 @@ async function getOrCreateSandbox(options) {
           // Assume ran for reconnected sandboxes
           startupScriptHash: void 0,
           // Unknown — caller should not re-run based on hash mismatch alone
+          installScriptRan: true,
+          // Assume ran for reconnected sandboxes
+          installScriptHash: void 0,
+          // Unknown — same logic as startup script
           isNew: false,
           configFileUrl: void 0,
           configInstalledAt: now2
           // Assume config was installed — prevents unnecessary re-install on reconnection
         };
       } else {
-        console.log("[SANDBOX] Reconnected sandbox failed health check, will create new");
+        console.log("[SANDBOX] Reconnected sandbox failed health check, will try pool or create new");
       }
     } catch (error) {
       console.log(
@@ -5469,6 +5582,47 @@ async function getOrCreateSandbox(options) {
       );
     }
   }
+  const pool = await ensureSandboxPoolInitialized();
+  if (pool && pool.isRunning()) {
+    try {
+      console.log(`[SANDBOX] Attempting to acquire from pre-warmed pool...${preferTag ? ` [preferTag=${preferTag}]` : ""}`);
+      const pooled = await pool.acquire(sessionId, preferTag);
+      const tagInfo = pooled.warmupTag ? ` [tag=${pooled.warmupTag}, agentSetup=${pooled.agentSetupComplete}]` : "";
+      console.log(`[SANDBOX] Acquired pre-warmed sandbox: ${pooled.sandboxId}${tagInfo}`);
+      const agentSetupDone = pooled.agentSetupComplete === true;
+      const now2 = Date.now();
+      const entry2 = {
+        sandbox: pooled.sandbox,
+        sessionId,
+        createdAt: pooled.createdAt,
+        lastUsedAt: now2,
+        sdkInstalled: pooled.sdkInstalled,
+        startupScriptRan: agentSetupDone,
+        installScriptRan: agentSetupDone
+      };
+      sandboxCache.set(sessionId, entry2);
+      return {
+        sandbox: pooled.sandbox,
+        sandboxId: pooled.sandboxId,
+        sdkInstalled: pooled.sdkInstalled,
+        startupScriptRan: agentSetupDone,
+        startupScriptHash: void 0,
+        installScriptRan: agentSetupDone,
+        installScriptHash: void 0,
+        isNew: false,
+        // Not new - came from pool
+        configFileUrl: void 0,
+        configInstalledAt: agentSetupDone ? now2 : void 0,
+        warmupTag: pooled.warmupTag,
+        agentSetupComplete: pooled.agentSetupComplete
+      };
+    } catch (error) {
+      console.warn(
+        "[SANDBOX] Failed to acquire from pool, falling back to on-demand creation:",
+        error instanceof Error ? error.message : "Unknown error"
+      );
+    }
+  }
   console.log("[SANDBOX] Creating new sandbox for session:", sessionId);
   const baseTarballUrl = process.env.SANDBOX_BASE_TARBALL_URL;
   const useTarball = !!baseTarballUrl;
@@ -5502,7 +5656,8 @@ async function getOrCreateSandbox(options) {
     lastUsedAt: now,
     // If we used tarball, SDK is pre-installed
     sdkInstalled: useTarball,
-    startupScriptRan: false
+    startupScriptRan: false,
+    installScriptRan: false
   };
   sandboxCache.set(sessionId, entry);
   return {
@@ -5512,6 +5667,8 @@ async function getOrCreateSandbox(options) {
     sdkInstalled: useTarball,
     startupScriptRan: false,
     startupScriptHash: void 0,
+    installScriptRan: false,
+    installScriptHash: void 0,
     isNew: true,
     configFileUrl: void 0,
     configInstalledAt: void 0
@@ -5551,6 +5708,14 @@ function markStartupScriptRan(sessionId, scriptHash) {
     cached.lastUsedAt = Date.now();
   }
 }
+function markInstallScriptRan(sessionId, scriptHash) {
+  const cached = sandboxCache.get(sessionId);
+  if (cached) {
+    cached.installScriptRan = true;
+    cached.installScriptHash = scriptHash;
+    cached.lastUsedAt = Date.now();
+  }
+}
 function needsStartupScriptRerun(sessionId, newScript) {
   const cached = sandboxCache.get(sessionId);
   if (!cached) return true;
@@ -5618,6 +5783,8 @@ function getCachedSandbox(sessionId) {
     sdkInstalled: cached.sdkInstalled,
     startupScriptRan: cached.startupScriptRan,
     startupScriptHash: cached.startupScriptHash,
+    installScriptRan: cached.installScriptRan,
+    installScriptHash: cached.installScriptHash,
     isNew: false,
     configFileUrl: cached.configFileUrl,
     configInstalledAt: cached.configInstalledAt
@@ -19986,6 +20153,7 @@ exports.loadWorkspaceState = loadWorkspaceState;
 exports.mapClaudeOptionsToGemini = mapClaudeOptionsToGemini;
 exports.mapToolToActionType = mapToolToActionType;
 exports.markConfigInstalled = markConfigInstalled;
+exports.markInstallScriptRan = markInstallScriptRan;
 exports.markSdkInstalled = markSdkInstalled;
 exports.markStartupScriptRan = markStartupScriptRan;
 exports.mcpAuthToHeaders = mcpAuthToHeaders;