@juspay/yama 1.5.0 → 1.6.0

package/dist/types/index.d.ts

@@ -569,6 +569,7 @@ export interface TokenBudgetManagerInterface {
     getAvailableBudget(): number;
     getTotalBudget(): number;
     getUsedTokens(): number;
+    preAllocateAllBatches(allocations: Map<number, number>): boolean;
 }
 export declare class GuardianError extends Error {
     code: string;
@@ -584,4 +585,40 @@ export declare class ProviderError extends GuardianError {
 export declare class ValidationError extends GuardianError {
     constructor(message: string, context?: any);
 }
+export declare enum CacheErrorCode {
+    CACHE_SYSTEM_FAILURE = "CACHE_SYSTEM_FAILURE",
+    CACHE_MEMORY_EXHAUSTED = "CACHE_MEMORY_EXHAUSTED",
+    CACHE_INITIALIZATION_FAILED = "CACHE_INITIALIZATION_FAILED",
+    CACHE_STORAGE_FULL = "CACHE_STORAGE_FULL",
+    CACHE_STORAGE_PERMISSION = "CACHE_STORAGE_PERMISSION",
+    CACHE_STORAGE_CORRUPTION = "CACHE_STORAGE_CORRUPTION",
+    CACHE_NETWORK_CONNECTION = "CACHE_NETWORK_CONNECTION",
+    CACHE_NETWORK_TIMEOUT = "CACHE_NETWORK_TIMEOUT",
+    CACHE_NETWORK_AUTH = "CACHE_NETWORK_AUTH",
+    CACHE_CONFIG_INVALID = "CACHE_CONFIG_INVALID",
+    CACHE_CONFIG_MISSING = "CACHE_CONFIG_MISSING",
+    CACHE_OPERATION_FAILED = "CACHE_OPERATION_FAILED",
+    CACHE_SERIALIZATION_ERROR = "CACHE_SERIALIZATION_ERROR",
+    CACHE_KEY_INVALID = "CACHE_KEY_INVALID"
+}
+export declare abstract class CacheError extends GuardianError {
+    operation?: string | undefined;
+    key?: string | undefined;
+    constructor(code: CacheErrorCode, message: string, operation?: string | undefined, key?: string | undefined, context?: any);
+}
+export declare class CacheSystemError extends CacheError {
+    constructor(message: string, operation?: string, key?: string, context?: any);
+}
+export declare class CacheStorageError extends CacheError {
+    constructor(code: CacheErrorCode | undefined, message: string, operation?: string, key?: string, context?: any);
+}
+export declare class CacheNetworkError extends CacheError {
+    constructor(code: CacheErrorCode | undefined, message: string, operation?: string, key?: string, context?: any);
+}
+export declare class CacheConfigurationError extends CacheError {
+    constructor(code: CacheErrorCode | undefined, message: string, operation?: string, key?: string, context?: any);
+}
+export declare class CacheOperationError extends CacheError {
+    constructor(code: CacheErrorCode | undefined, message: string, operation?: string, key?: string, context?: any);
+}
 //# sourceMappingURL=index.d.ts.map

package/dist/types/index.js

@@ -34,6 +34,71 @@ export class ValidationError extends GuardianError {
     }
 }
 // ============================================================================
+// Cache Error Types
+// ============================================================================
+export var CacheErrorCode;
+(function (CacheErrorCode) {
+    // System-level cache errors
+    CacheErrorCode["CACHE_SYSTEM_FAILURE"] = "CACHE_SYSTEM_FAILURE";
+    CacheErrorCode["CACHE_MEMORY_EXHAUSTED"] = "CACHE_MEMORY_EXHAUSTED";
+    CacheErrorCode["CACHE_INITIALIZATION_FAILED"] = "CACHE_INITIALIZATION_FAILED";
+    // Storage-related errors
+    CacheErrorCode["CACHE_STORAGE_FULL"] = "CACHE_STORAGE_FULL";
+    CacheErrorCode["CACHE_STORAGE_PERMISSION"] = "CACHE_STORAGE_PERMISSION";
+    CacheErrorCode["CACHE_STORAGE_CORRUPTION"] = "CACHE_STORAGE_CORRUPTION";
+    // Network-related errors (for future Redis support)
+    CacheErrorCode["CACHE_NETWORK_CONNECTION"] = "CACHE_NETWORK_CONNECTION";
+    CacheErrorCode["CACHE_NETWORK_TIMEOUT"] = "CACHE_NETWORK_TIMEOUT";
+    CacheErrorCode["CACHE_NETWORK_AUTH"] = "CACHE_NETWORK_AUTH";
+    // Configuration errors
+    CacheErrorCode["CACHE_CONFIG_INVALID"] = "CACHE_CONFIG_INVALID";
+    CacheErrorCode["CACHE_CONFIG_MISSING"] = "CACHE_CONFIG_MISSING";
+    // Operation errors
+    CacheErrorCode["CACHE_OPERATION_FAILED"] = "CACHE_OPERATION_FAILED";
+    CacheErrorCode["CACHE_SERIALIZATION_ERROR"] = "CACHE_SERIALIZATION_ERROR";
+    CacheErrorCode["CACHE_KEY_INVALID"] = "CACHE_KEY_INVALID";
+})(CacheErrorCode || (CacheErrorCode = {}));
+export class CacheError extends GuardianError {
+    operation;
+    key;
+    constructor(code, message, operation, key, context) {
+        super(code, message, context);
+        this.operation = operation;
+        this.key = key;
+        this.name = "CacheError";
+    }
+}
+export class CacheSystemError extends CacheError {
+    constructor(message, operation, key, context) {
+        super(CacheErrorCode.CACHE_SYSTEM_FAILURE, message, operation, key, context);
+        this.name = "CacheSystemError";
+    }
+}
+export class CacheStorageError extends CacheError {
+    constructor(code = CacheErrorCode.CACHE_STORAGE_FULL, message, operation, key, context) {
+        super(code, message, operation, key, context);
+        this.name = "CacheStorageError";
+    }
+}
+export class CacheNetworkError extends CacheError {
+    constructor(code = CacheErrorCode.CACHE_NETWORK_CONNECTION, message, operation, key, context) {
+        super(code, message, operation, key, context);
+        this.name = "CacheNetworkError";
+    }
+}
+export class CacheConfigurationError extends CacheError {
+    constructor(code = CacheErrorCode.CACHE_CONFIG_INVALID, message, operation, key, context) {
+        super(code, message, operation, key, context);
+        this.name = "CacheConfigurationError";
+    }
+}
+export class CacheOperationError extends CacheError {
+    constructor(code = CacheErrorCode.CACHE_OPERATION_FAILED, message, operation, key, context) {
+        super(code, message, operation, key, context);
+        this.name = "CacheOperationError";
+    }
+}
+// ============================================================================
 // Export all types - Main file, no re-exports needed
 // ============================================================================
 //# sourceMappingURL=index.js.map

package/dist/utils/Cache.d.ts

@@ -8,7 +8,7 @@ export declare class Cache implements ICache {
     private statsData;
     constructor(options?: CacheOptions);
     /**
-     * Get value from cache
+     * Get value from cache with resilient error handling
      */
     get<T>(key: string): T | undefined;
     /**
@@ -39,15 +39,21 @@ export declare class Cache implements ICache {
         misses: number;
         keys: number;
         size: number;
+        cacheErrors: number;
+        nonCacheErrors: number;
     };
     /**
      * Get detailed cache statistics from node-cache
      */
     getDetailedStats(): any;
     /**
-     * Get or set pattern - common caching pattern
+     * Get or set pattern with automatic fallback on cache failures
      */
     getOrSet<T>(key: string, fetchFn: () => Promise<T>, ttl?: number): Promise<T>;
+    /**
+     * Resilient get or set pattern that bypasses cache entirely on cache system failures
+     */
+    getOrSetResilient<T>(key: string, fetchFn: () => Promise<T>, ttl?: number): Promise<T>;
     /**
      * Cache with tags for group invalidation
      */

package/dist/utils/Cache.js

@@ -3,12 +3,146 @@
  * Provides intelligent caching for PR data, file contents, and AI responses
  */
 import NodeCache from "node-cache";
+import { CacheError, } from "../types/index.js";
 import { logger } from "./Logger.js";
+/**
+ * Enhanced cache error detection utility
+ * Provides multi-layer error classification to avoid false positives
+ */
+class CacheErrorDetector {
+    /**
+     * Detect if an error is cache-related using multiple strategies
+     */
+    static isCacheError(error, operation, key) {
+        // Strategy 1: Check error type/class (most reliable)
+        if (error instanceof CacheError) {
+            return true;
+        }
+        // Strategy 2: Check for specific cache error patterns in NodeCache
+        if (error instanceof Error) {
+            const errorMessage = error.message.toLowerCase();
+            const stackTrace = error.stack?.toLowerCase() || "";
+            // Check for NodeCache-specific error patterns
+            const nodeCachePatterns = [
+                /node_modules\/node-cache/,
+                /cache\.js:\d+/,
+                /nodecache/,
+            ];
+            const isNodeCacheError = nodeCachePatterns.some((pattern) => pattern.test(stackTrace));
+            if (isNodeCacheError) {
+                return true;
+            }
+            // Strategy 3: Check for specific cache-related error messages (more targeted)
+            const cacheSpecificPatterns = [
+                /cache.*(?:full|exhausted|limit)/,
+                /memory.*(?:cache|allocation).*(?:failed|error)/,
+                /storage.*(?:cache|quota).*(?:exceeded|full)/,
+                /cache.*(?:initialization|setup).*(?:failed|error)/,
+                /ttl.*(?:invalid|expired)/,
+                /cache.*(?:key|value).*(?:invalid|malformed)/,
+            ];
+            const hasCacheSpecificError = cacheSpecificPatterns.some((pattern) => pattern.test(errorMessage));
+            if (hasCacheSpecificError) {
+                return true;
+            }
+            // Strategy 4: Context-aware detection
+            if (operation && key) {
+                // If we're in a cache operation and get memory/storage errors, likely cache-related
+                const cacheOperations = [
+                    "get",
+                    "set",
+                    "del",
+                    "clear",
+                    "has",
+                    "getorset",
+                    "getorsetresilient",
+                ];
+                const isCacheOperation = cacheOperations.includes(operation.toLowerCase());
+                const contextualPatterns = [
+                    /^out of memory$/,
+                    /storage quota exceeded/,
+                    /disk full/,
+                ];
+                if (isCacheOperation &&
+                    contextualPatterns.some((pattern) => pattern.test(errorMessage))) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+    /**
+     * Classify cache error for better handling and logging
+     */
+    static classifyError(error, operation, key) {
+        if (!this.isCacheError(error, operation, key)) {
+            return {
+                isCache: false,
+                category: "unknown",
+                confidence: "high",
+                reason: "Not identified as cache-related error",
+            };
+        }
+        if (error instanceof CacheError) {
+            const category = error.code.includes("STORAGE")
+                ? "storage"
+                : error.code.includes("NETWORK")
+                    ? "network"
+                    : error.code.includes("SYSTEM")
+                        ? "system"
+                        : "operation";
+            return {
+                isCache: true,
+                category,
+                confidence: "high",
+                reason: `Explicit cache error: ${error.code}`,
+            };
+        }
+        if (error instanceof Error) {
+            const message = error.message.toLowerCase();
+            const stack = error.stack?.toLowerCase() || "";
+            // High confidence patterns
+            if (/node_modules\/node-cache/.test(stack)) {
+                return {
+                    isCache: true,
+                    category: "system",
+                    confidence: "high",
+                    reason: "NodeCache stack trace detected",
+                };
+            }
+            // Medium confidence patterns
+            if (/cache.*(?:full|exhausted)/.test(message)) {
+                return {
+                    isCache: true,
+                    category: "storage",
+                    confidence: "medium",
+                    reason: "Cache capacity error pattern",
+                };
+            }
+            if (/memory.*cache.*failed/.test(message)) {
+                return {
+                    isCache: true,
+                    category: "system",
+                    confidence: "medium",
+                    reason: "Memory allocation error in cache context",
+                };
+            }
+        }
+        return {
+            isCache: true,
+            category: "unknown",
+            confidence: "low",
+            reason: "Fallback detection",
+        };
+    }
+}
 export class Cache {
     cache;
     statsData = {
         hits: 0,
         misses: 0,
+        cacheErrors: 0,
+        nonCacheErrors: 0,
     };
     constructor(options = {}) {
         const { ttl = 3600, // 1 hour default
@@ -33,18 +167,25 @@ export class Cache {
         });
     }
     /**
-     * Get value from cache
+     * Get value from cache with resilient error handling
      */
     get(key) {
-        const value = this.cache.get(key);
-        if (value !== undefined) {
-            this.statsData.hits++;
-            logger.debug(`Cache HIT: ${key}`);
-            return value;
+        try {
+            const value = this.cache.get(key);
+            if (value !== undefined) {
+                this.statsData.hits++;
+                logger.debug(`Cache HIT: ${key}`);
+                return value;
+            }
+            else {
+                this.statsData.misses++;
+                logger.debug(`Cache MISS: ${key}`);
+                return undefined;
+            }
         }
-        else {
+        catch (error) {
             this.statsData.misses++;
-            logger.debug(`Cache MISS: ${key}`);
+            logger.warn(`Cache GET error for ${key}, treating as miss:`, error);
             return undefined;
         }
     }
@@ -105,6 +246,8 @@ export class Cache {
             misses: this.statsData.misses,
             keys: this.cache.keys().length,
             size: this.cache.getStats().keys,
+            cacheErrors: this.statsData.cacheErrors,
+            nonCacheErrors: this.statsData.nonCacheErrors,
         };
     }
     /**
@@ -114,9 +257,10 @@ export class Cache {
         return this.cache.getStats();
     }
     /**
-     * Get or set pattern - common caching pattern
+     * Get or set pattern with automatic fallback on cache failures
      */
     async getOrSet(key, fetchFn, ttl) {
+        // Try to get from cache with resilient error handling
         const cached = this.get(key);
         if (cached !== undefined) {
             return cached;
@@ -124,7 +268,13 @@ export class Cache {
         try {
             logger.debug(`Cache FETCH: ${key}`);
             const value = await fetchFn();
-            this.set(key, value, ttl);
+            // Try to cache the result, but don't fail if caching fails
+            try {
+                this.set(key, value, ttl);
+            }
+            catch (cacheError) {
+                logger.warn(`Cache SET failed for ${key}, continuing without cache:`, cacheError);
+            }
             return value;
         }
         catch (error) {
@@ -132,6 +282,36 @@ export class Cache {
             throw error;
         }
     }
+    /**
+     * Resilient get or set pattern that bypasses cache entirely on cache system failures
+     */
+    async getOrSetResilient(key, fetchFn, ttl) {
+        try {
+            // Try normal cache flow first
+            return await this.getOrSet(key, fetchFn, ttl);
+        }
+        catch (error) {
+            // Use enhanced error detection to determine if this is a cache-related error
+            const errorClassification = CacheErrorDetector.classifyError(error, "getOrSet", key);
+            if (errorClassification.isCache) {
+                // Track cache error statistics
+                this.statsData.cacheErrors++;
+                logger.warn(`Cache system error detected for ${key} (${errorClassification.confidence} confidence: ${errorClassification.reason}), bypassing cache entirely`, {
+                    error: error instanceof Error ? error.message : String(error),
+                    category: errorClassification.category,
+                    confidence: errorClassification.confidence,
+                    key,
+                    operation: "getOrSet",
+                });
+                // Bypass cache completely and just fetch the data
+                return await fetchFn();
+            }
+            // Track non-cache errors for debugging
+            this.statsData.nonCacheErrors++;
+            // Re-throw non-cache errors
+            throw error;
+        }
+    }
     /**
      * Cache with tags for group invalidation
      */

package/dist/utils/ConfigManager.js

@@ -80,7 +80,7 @@ export class ConfigManager {
                     },
                 ],
                 autoFormat: true,
-                systemPrompt: "You are an Expert Technical Writer specializing in pull request documentation. Your role is to:\n\n📝 CLARITY FIRST: Create clear, comprehensive PR descriptions that help reviewers understand the changes\n🎥 STORY TELLING: Explain the 'why' behind changes, not just the 'what'\n📈 STRUCTURED: Follow consistent formatting with required sections\n🔗 CONTEXTUAL: Link changes to business value and technical rationale\n\nCRITICAL INSTRUCTION: Return ONLY the enhanced PR description content as clean markdown. Do NOT include any meta-commentary, explanations about what you're doing, or introductory text like \"I will enhance...\" or \"Here is the enhanced description:\". \n\nOutput the enhanced description directly without any wrapper text or explanations.",
+                systemPrompt: "You are an Expert Technical Writer specializing in pull request documentation. Your role is to:\n\n📝 CLARITY FIRST: Create clear, comprehensive PR descriptions that help reviewers understand the changes\n🎥 STORY TELLING: Explain the 'why' behind changes, not just the 'what'\n📈 STRUCTURED: Follow consistent formatting with required sections\n🔗 CONTEXTUAL: Link changes to business value and technical rationale\n\nCRITICAL INSTRUCTION: Return ONLY the enhanced PR description content as clean markdown.\n- DO NOT add meta-commentary like \"No description provided\" or \"Here is the enhanced description\"\n- DO NOT add explanatory text about what you're doing\n- START directly with the actual PR content (sections, lists, explanations)\n- If there's no existing description, just write the new sections without mentioning it\n\nOutput the enhanced description directly without any wrapper text or explanations.",
                 outputTemplate: "# PR Title Enhancement (if needed)\n\n## Summary\n[Clear overview of what this PR accomplishes]\n\n## Changes Made\n[Specific technical changes - be precise]\n\n## Testing\n[How the changes were tested]\n\n## Impact\n[Business/technical impact and considerations]\n\n## Additional Notes\n[Any deployment notes, follow-ups, or special considerations]",
                 enhancementInstructions: 'Return ONLY the enhanced PR description as clean markdown. Do NOT include any explanatory text, meta-commentary, or phrases like "Here is the enhanced description:" or "I will enhance...".\n\nStart directly with the enhanced description content using this structure:',
             },
@@ -244,11 +244,10 @@ export class ConfigManager {
      * Apply provider-aware token limits using shared utility
      */
     applyProviderTokenLimits(config) {
-        const provider = config.providers.ai.provider || 'auto';
+        const provider = config.providers.ai.provider || "auto";
         const configuredTokens = config.providers.ai.maxTokens;
         // Use the shared utility to validate and adjust token limits
-        const validatedTokens = validateProviderTokenLimit(provider, configuredTokens, false // Use standard limits for configuration
-        );
+        const validatedTokens = validateProviderTokenLimit(provider, configuredTokens, false);
         config.providers.ai.maxTokens = validatedTokens;
         return config;
     }

package/dist/utils/ParallelProcessing.d.ts

@@ -46,6 +46,9 @@ export declare class TokenBudgetManager implements TokenBudgetManagerInterface {
     private usedTokens;
     private batchAllocations;
     private reservedTokens;
+    private preAllocationMode;
+    private preAllocatedBatches;
+    private batchStates;
     constructor(totalBudget: number);
     /**
      * Allocate tokens for a specific batch
@@ -92,6 +95,31 @@ export declare class TokenBudgetManager implements TokenBudgetManagerInterface {
      * Reset the budget manager (for testing or reuse)
      */
     reset(): void;
+    /**
+     * Pre-allocate tokens for all batches upfront
+     * This ensures all batches have guaranteed token allocation before processing starts
+     */
+    preAllocateAllBatches(allocations: Map<number, number>): boolean;
+    /**
+     * Mark a batch as failed and handle cleanup
+     */
+    markBatchFailed(batchIndex: number, error?: string): void;
+    /**
+     * Get the current state of a batch
+     */
+    getBatchState(batchIndex: number): "pending" | "processing" | "completed" | "failed" | undefined;
+    /**
+     * Check if pre-allocation mode is active
+     */
+    isPreAllocationMode(): boolean;
+    /**
+     * Get all batch states for debugging
+     */
+    getAllBatchStates(): Map<number, "pending" | "processing" | "completed" | "failed">;
+    /**
+     * Disable pre-allocation mode and clean up
+     */
+    disablePreAllocationMode(): void;
     /**
      * Update the total budget (useful for dynamic adjustment)
      */

package/dist/utils/ParallelProcessing.js

@@ -77,12 +77,18 @@ export class TokenBudgetManager {
     usedTokens = 0;
     batchAllocations = new Map();
     reservedTokens = 0; // Tokens allocated but not yet used
+    // NEW: Pre-allocation mode tracking
+    preAllocationMode = false;
+    preAllocatedBatches = new Set();
+    batchStates = new Map();
     constructor(totalBudget) {
         if (totalBudget <= 0) {
             throw new Error("Token budget must be greater than 0");
         }
-        this.totalBudget = totalBudget;
-        logger.debug(`TokenBudgetManager created with budget of ${totalBudget} tokens`);
+        // Floor the budget to ensure integer arithmetic and avoid floating-point precision issues.
+        // The fractional part is discarded, so the budget is always rounded down.
+        this.totalBudget = Math.floor(totalBudget);
+        logger.debug(`TokenBudgetManager created with budget of ${this.totalBudget} tokens (original: ${totalBudget})`);
     }
     /**
      * Allocate tokens for a specific batch
@@ -93,8 +99,31 @@ export class TokenBudgetManager {
             logger.warn(`Invalid token estimate for batch ${batchIndex}: ${estimatedTokens}`);
             return false;
         }
-        // Check if we already have an allocation for this batch
+        // NEW: Handle pre-allocation mode
+        if (this.preAllocationMode && this.preAllocatedBatches.has(batchIndex)) {
+            // Check if batch is already being processed
+            const currentState = this.batchStates.get(batchIndex);
+            if (currentState === "processing") {
+                logger.warn(`Batch ${batchIndex} is already being processed`);
+                return false;
+            }
+            if (currentState !== "pending") {
+                logger.warn(`Batch ${batchIndex} is not in pending state (current: ${currentState})`);
+                return false;
+            }
+            // In pre-allocation mode, just mark batch as processing and return success
+            this.batchStates.set(batchIndex, "processing");
+            logger.debug(`Batch ${batchIndex} using pre-allocated tokens (${this.batchAllocations.get(batchIndex)} tokens)`);
+            return true;
+        }
+        // Check if we already have an allocation for this batch (non-pre-allocation mode)
         if (this.batchAllocations.has(batchIndex)) {
+            // Check if batch is already being processed
+            const currentState = this.batchStates.get(batchIndex);
+            if (currentState === "processing") {
+                logger.warn(`Batch ${batchIndex} is already being processed`);
+                return false;
+            }
             logger.warn(`Batch ${batchIndex} already has token allocation`);
             return false;
         }
@@ -108,6 +137,7 @@ export class TokenBudgetManager {
         // Allocate the tokens
         this.reservedTokens += estimatedTokens;
         this.batchAllocations.set(batchIndex, estimatedTokens);
+        this.batchStates.set(batchIndex, "processing");
         logger.debug(`Allocated ${estimatedTokens} tokens for batch ${batchIndex} ` +
             `(${this.getAvailableBudget()} remaining)`);
         return true;
@@ -122,10 +152,19 @@ export class TokenBudgetManager {
             logger.warn(`No token allocation found for batch ${batchIndex}`);
             return;
         }
+        // Update batch state to completed only if not already failed
+        const currentState = this.batchStates.get(batchIndex);
+        if (currentState !== "failed") {
+            this.batchStates.set(batchIndex, "completed");
+        }
         // Move from reserved to used (assuming the tokens were actually used)
         this.reservedTokens -= allocated;
         this.usedTokens += allocated;
         this.batchAllocations.delete(batchIndex);
+        // Clean up pre-allocation tracking
+        if (this.preAllocationMode) {
+            this.preAllocatedBatches.delete(batchIndex);
+        }
         logger.debug(`Released ${allocated} tokens from batch ${batchIndex} ` +
             `(${this.getAvailableBudget()} now available)`);
     }
@@ -182,6 +221,72 @@ export class TokenBudgetManager {
         this.batchAllocations.clear();
         logger.debug("TokenBudgetManager reset");
     }
+    /**
+     * Pre-allocate tokens for all batches upfront
+     * This ensures all batches have guaranteed token allocation before processing starts
+     */
+    preAllocateAllBatches(allocations) {
+        const totalRequired = Array.from(allocations.values()).reduce((sum, tokens) => sum + tokens, 0);
+        if (totalRequired > this.totalBudget) {
+            logger.error(`Pre-allocation failed: total required (${totalRequired}) exceeds budget (${this.totalBudget})`);
+            return false;
+        }
+        // Clear any existing allocations and reset state
+        this.batchAllocations.clear();
+        this.reservedTokens = 0;
+        this.batchStates.clear();
+        this.preAllocatedBatches.clear();
+        // Enable pre-allocation mode
+        this.preAllocationMode = true;
+        // Reserve all tokens upfront
+        allocations.forEach((tokens, batchIndex) => {
+            this.batchAllocations.set(batchIndex, tokens);
+            this.reservedTokens += tokens;
+            this.preAllocatedBatches.add(batchIndex);
+            this.batchStates.set(batchIndex, "pending");
+        });
+        logger.info(`Pre-allocated ${totalRequired} tokens across ${allocations.size} batches ` +
+            `(${this.getAvailableBudget()} remaining)`);
+        return true;
+    }
+    /**
+     * Mark a batch as failed and handle cleanup
+     */
+    markBatchFailed(batchIndex, error) {
+        this.batchStates.set(batchIndex, "failed");
+        if (error) {
+            logger.debug(`Batch ${batchIndex} marked as failed: ${error}`);
+        }
+        else {
+            logger.debug(`Batch ${batchIndex} marked as failed`);
+        }
+    }
+    /**
+     * Get the current state of a batch
+     */
+    getBatchState(batchIndex) {
+        return this.batchStates.get(batchIndex);
+    }
+    /**
+     * Check if pre-allocation mode is active
+     */
+    isPreAllocationMode() {
+        return this.preAllocationMode;
+    }
+    /**
+     * Get all batch states for debugging
+     */
+    getAllBatchStates() {
+        return new Map(this.batchStates);
+    }
+    /**
+     * Disable pre-allocation mode and clean up
+     */
+    disablePreAllocationMode() {
+        this.preAllocationMode = false;
+        this.preAllocatedBatches.clear();
+        logger.debug("Pre-allocation mode disabled");
+    }
     /**
      * Update the total budget (useful for dynamic adjustment)
      */